prismatic

prismatic

Prismatic VLM component for VLA decoders.

PrismaticVLM

PrismaticVLM(input_keys, pretrained, frozen, model_name=value, repository_id=PRISMATIC_REPOSITORY_ID, attention_type=value, model_dtype=None, max_text_length=None, lora_config=None, gradient_checkpointing=False)

Bases: GenerativeVLM

Raw Prismatic VLM checkpoint loader for interleaved VLA decoders.

Load or initialize a raw Prismatic VLM.

Parameters:

Name	Type	Description	Default
input_keys	str | list[str]	RGB camera keys consumed by the VLM.	required
pretrained	bool	Whether to load the raw Prismatic checkpoint.	required
frozen	bool	Whether to freeze all model weights.	required
model_name	str	Prismatic checkpoint folder name, or a local checkpoint directory containing config.json and checkpoints.	value
repository_id	str	HuggingFace repository containing raw Prismatic checkpoint folders.	PRISMATIC_REPOSITORY_ID
attention_type	str	HuggingFace attention implementation for the language model.	value
model_dtype	str | None	Optional precision string for model parameter dtype.	None
max_text_length	int | None	Optional text sequence length. Defaults to the raw Prismatic llm_max_length field.	None
lora_config	LoRAAdaptation | None	Optional LoRA adapter configuration for the language model.	None
gradient_checkpointing	bool	Whether to enable activation checkpointing in the language model during training.	False

Source code in src/versatil/models/decoding/generative_language_models/vision_language/prismatic.py

def __init__(
    self,
    input_keys: str | list[str],
    pretrained: bool,
    frozen: bool,
    model_name: str = PrismaticModelType.PRISM_DINOSIGLIP_224PX_7B.value,
    repository_id: str = PRISMATIC_REPOSITORY_ID,
    attention_type: str = AttentionImplementation.SDPA.value,
    model_dtype: str | None = None,
    max_text_length: int | None = None,
    lora_config: LoRAAdaptation | None = None,
    gradient_checkpointing: bool = False,
) -> None:
    """Load or initialize a raw Prismatic VLM.

    Args:
        input_keys: RGB camera keys consumed by the VLM.
        pretrained: Whether to load the raw Prismatic checkpoint.
        frozen: Whether to freeze all model weights.
        model_name: Prismatic checkpoint folder name, or a local checkpoint
            directory containing ``config.json`` and ``checkpoints``.
        repository_id: HuggingFace repository containing raw Prismatic
            checkpoint folders.
        attention_type: HuggingFace attention implementation for the
            language model.
        model_dtype: Optional precision string for model parameter dtype.
        max_text_length: Optional text sequence length. Defaults to the
            raw Prismatic ``llm_max_length`` field.
        lora_config: Optional LoRA adapter configuration for the language
            model.
        gradient_checkpointing: Whether to enable activation checkpointing
            in the language model during training.
    """
    super().__init__(
        input_keys=input_keys,
        pretrained=pretrained,
        frozen=frozen,
        model_dtype=model_dtype,
        max_text_length=max_text_length,
    )
    self.model_name = model_name
    self.repository_id = repository_id
    self.lora_config = lora_config
    self.gradient_checkpointing = gradient_checkpointing
    model_config = self._load_model_config(
        model_name=model_name,
        repository_id=repository_id,
    )
    self.prismatic_config = model_config
    self.vision_backbone_id = str(model_config["vision_backbone_id"])
    self.llm_backbone_id = str(model_config["llm_backbone_id"])
    self.arch_specifier = str(model_config["arch_specifier"])
    self.vision_backbone_type = self._resolve_vision_backbone_type(
        vision_backbone_id=self.vision_backbone_id
    )
    self.vision_backbone_types = PRISMATIC_VISION_BACKBONES[
        self.vision_backbone_type
    ]
    self.image_size = PRISMATIC_VISION_IMAGE_SIZES[self.vision_backbone_type]
    # Raw DinoSigLIP checkpoints do not store the frozen vision towers, so
    # the timm towers must load their own pretrained weights.
    self.vision_encoders = self._build_vision_encoders(pretrained=pretrained)
    self.num_image_tokens_per_camera = self._resolve_num_image_tokens()
    self.vision_embedding_dimension = sum(
        int(encoder.feature_dim) for encoder in self.vision_encoders
    )
    self.language_model = self._build_language_model(
        llm_backbone_id=self.llm_backbone_id,
        attention_type=attention_type,
    )
    self.hidden_dimension = int(self.language_model.config.hidden_size)
    self.projector = self._build_projector(
        arch_specifier=self.arch_specifier,
        vision_dimension=self.vision_embedding_dimension,
        language_dimension=self.hidden_dimension,
    )
    self.max_text_length = (
        max_text_length
        if max_text_length is not None
        else int(model_config["llm_max_length"])
    )
    if pretrained:
        checkpoint_path = self._resolve_checkpoint_path(
            model_name=model_name,
            repository_id=repository_id,
        )
        self._load_prismatic_checkpoint(checkpoint_path=checkpoint_path)
    if lora_config is not None and lora_config.enabled:
        # PEFT mutates custom modules in place; assigning the returned wrapper
        # here would recursively register this module as its own child.
        apply_lora_config(model=self, lora_config=lora_config, frozen=frozen)
    if gradient_checkpointing:
        self.language_model.gradient_checkpointing_enable()
        self.language_model.config.use_cache = False
    if frozen:
        super()._freeze_weights()
    self._apply_model_dtype()

get_vocab_size

get_vocab_size()

Return the Prismatic language vocabulary size.

Source code in src/versatil/models/decoding/generative_language_models/vision_language/prismatic.py

def get_vocab_size(self) -> int:
    """Return the Prismatic language vocabulary size."""
    return int(self.language_model.config.vocab_size)

resize_token_embeddings

resize_token_embeddings(vocabulary_size)

Resize the Prismatic causal language-model token embeddings and output head.

Source code in src/versatil/models/decoding/generative_language_models/vision_language/prismatic.py

def resize_token_embeddings(self, vocabulary_size: int) -> None:
    """Resize the Prismatic causal language-model token embeddings and output head."""
    self.language_model.resize_token_embeddings(vocabulary_size)

forward_language_model

forward_language_model(input_ids=None, inputs_embeds=None, attention_mask=None, past_key_values=None, use_cache=False, cache_position=None, position_ids=None, output_hidden_states=True)

Run the Prismatic language model over token IDs or embeddings.

Parameters:

Name	Type	Description	Default
input_ids	Tensor | None	Optional token IDs with shape (B, S).	None
inputs_embeds	Tensor | None	Optional token embeddings with shape (B, S, D).	None
attention_mask	Tensor | None	Optional language-model attention mask.	None
past_key_values	Cache | tuple[tuple[Tensor, ...], ...] | None	Optional cached key/value tensors.	None
use_cache	bool	Whether to return/update cached key/value tensors.	False
cache_position	Tensor | None	Optional HuggingFace KV-cache slots for the tokens in this call. During cached decoding, if the prefix has length P, the next token uses cache slot P so its key/value is appended after the prefix.	None
position_ids	Tensor | None	Optional positions for the language model positional encoding, with shape (B, S). These should count real tokens, not padding: [PAD, PAD, t0, t1] should pass [0, 0, 0, 1] so t0 and t1 get positions 0 and 1.	None
output_hidden_states	bool	Whether to return hidden states.	True

Returns:

Type	Description
CausalLanguageModelOutput	Causal language-model output with logits shape (B, S, V).

Source code in src/versatil/models/decoding/generative_language_models/vision_language/prismatic.py

def forward_language_model(
    self,
    input_ids: torch.Tensor | None = None,
    inputs_embeds: torch.Tensor | None = None,
    attention_mask: torch.Tensor | None = None,
    past_key_values: Cache | tuple[tuple[torch.Tensor, ...], ...] | None = None,
    use_cache: bool = False,
    cache_position: torch.Tensor | None = None,
    position_ids: torch.Tensor | None = None,
    output_hidden_states: bool = True,
) -> CausalLanguageModelOutput:
    """Run the Prismatic language model over token IDs or embeddings.

    Args:
        input_ids: Optional token IDs with shape ``(B, S)``.
        inputs_embeds: Optional token embeddings with shape ``(B, S, D)``.
        attention_mask: Optional language-model attention mask.
        past_key_values: Optional cached key/value tensors.
        use_cache: Whether to return/update cached key/value tensors.
        cache_position: Optional HuggingFace KV-cache slots for the tokens
            in this call. During cached decoding, if the prefix has length
            ``P``, the next token uses cache slot ``P`` so its key/value is
            appended after the prefix.
        position_ids: Optional positions for the language model positional
            encoding, with shape ``(B, S)``. These should count real tokens,
            not padding: ``[PAD, PAD, t0, t1]`` should pass
            ``[0, 0, 0, 1]`` so ``t0`` and ``t1`` get positions ``0`` and
            ``1``.
        output_hidden_states: Whether to return hidden states.

    Returns:
        Causal language-model output with logits shape ``(B, S, V)``.
    """
    language_model_inputs = {
        "input_ids": input_ids,
        "inputs_embeds": inputs_embeds,
        "attention_mask": attention_mask,
        "past_key_values": past_key_values,
        "use_cache": use_cache,
        "output_hidden_states": output_hidden_states,
        "return_dict": True,
    }
    if cache_position is not None:
        language_model_inputs["cache_position"] = cache_position
    if position_ids is not None:
        language_model_inputs["position_ids"] = position_ids
    return self.language_model(**language_model_inputs)

get_text_config

get_text_config()

Return the Prismatic text model configuration.

Source code in src/versatil/models/decoding/generative_language_models/vision_language/prismatic.py

def get_text_config(self) -> PretrainedConfig:
    """Return the Prismatic text model configuration."""
    return self._get_language_model().config