Add files using upload-large-folder tool

capvector-oft/prismatic/models/__init__.py

@@ -0,0 +1,2 @@
+from .load import available_model_names, available_models, get_model_description, load, load_vla
+from .materialize import get_llm_backbone_and_tokenizer, get_vision_backbone_and_transform, get_vlm

capvector-oft/prismatic/models/backbones/llm/__init__.py

@@ -0,0 +1,4 @@
+from .base_llm import LLMBackbone
+from .llama2 import LLaMa2LLMBackbone
+from .mistral import MistralLLMBackbone
+from .phi import PhiLLMBackbone

capvector-oft/prismatic/models/backbones/llm/base_llm.py

@@ -0,0 +1,223 @@
+"""
+base_llm.py
+
+Abstract class definition of a large (autoregressive) language model backbone (LLM), with full annotations of class
+methods, utility functions, and initialization logic.
+
+We also define the generic HFLLMBackbone class here, providing a default interface for loading any HF
+AutoModelForCausalLM (e.g., LLamaForCausalLM). In general, we make the assumption that any given LLM backbone implements
+the AutoModelForCausalLM API (though we may add Seq2Seq models in the future).
+
+We make this assumption to keep the LLM handling in this codebase relatively lightweight, and to inherit all the nice HF
+utilities around different types of decoding/generation strategies.
+"""
+
+import warnings
+from abc import ABC, abstractmethod
+from functools import partial
+from typing import Callable, List, Optional, Sequence, Type
+
+import torch
+import torch.nn as nn
+from torch.distributed.fsdp.wrap import transformer_auto_wrap_policy
+from transformers import AutoConfig, AutoTokenizer, PreTrainedModel, PreTrainedTokenizerBase
+from transformers.modeling_outputs import CausalLMOutputWithPast
+
+from prismatic.models.backbones.llm.prompting import PromptBuilder
+from prismatic.overwatch import initialize_overwatch
+
+# Suppress HF Deprecation Warnings
+warnings.filterwarnings("ignore", category=FutureWarning)
+
+# Initialize Overwatch =>> Wraps `logging.Logger`
+overwatch = initialize_overwatch(__name__)
+
+
+# === Abstract Base Class for arbitrary HF LLM Backbones ===
+class LLMBackbone(nn.Module, ABC):
+    def __init__(self, llm_backbone_id: str) -> None:
+        super().__init__()
+        self.identifier = llm_backbone_id
+
+        # Instance attributes for an LLM Backbone
+        self.llm: PreTrainedModel = None
+        self.tokenizer: PreTrainedTokenizerBase = None
+
+    def get_tokenizer(self) -> PreTrainedTokenizerBase:
+        return self.tokenizer
+
+    @abstractmethod
+    def get_fsdp_wrapping_policy(self) -> Callable: ...
+
+    @abstractmethod
+    def enable_gradient_checkpointing(self) -> None: ...
+
+    @abstractmethod
+    def forward(
+        self,
+        input_ids: Optional[torch.LongTensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_values: Optional[List[torch.FloatTensor]] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        labels: Optional[torch.LongTensor] = None,
+        use_cache: Optional[bool] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+    ) -> CausalLMOutputWithPast:
+        """Run a forward pass through the LLM given targets (labels), returning the scalar Cross-Entropy Loss"""
+        raise NotImplementedError
+
+    @abstractmethod
+    def embed_input_ids(self, input_ids: torch.LongTensor) -> torch.Tensor: ...
+
+    @property
+    @abstractmethod
+    def prompt_builder_fn(self) -> Type[PromptBuilder]: ...
+
+    @property
+    @abstractmethod
+    def transformer_layer_cls(self) -> Type[nn.Module]: ...
+
+    @property
+    @abstractmethod
+    def half_precision_dtype(self) -> torch.dtype: ...
+
+    @property
+    @abstractmethod
+    def last_layer_finetune_modules(self) -> Sequence[nn.Module]: ...
+
+    @property
+    def embed_dim(self) -> int:
+        return self.llm.config.hidden_size
+
+    @property
+    def pad_token_id(self) -> int:
+        return self.tokenizer.pad_token_id
+
+
+# === Abstract Base Class for Arbitrary HF Causal LLMs ===
+class HFCausalLLMBackbone(LLMBackbone, ABC):
+    def __init__(
+        self,
+        llm_backbone_id: str,
+        llm_family: str,
+        llm_cls: Type[PreTrainedModel],
+        hf_hub_path: str,
+        llm_max_length: int = 2048,
+        hf_token: Optional[str] = None,
+        inference_mode: bool = False,
+        use_flash_attention_2: bool = False,
+    ) -> None:
+        super().__init__(llm_backbone_id)
+        self.llm_family = llm_family
+        self.llm_max_length = llm_max_length
+        self.inference_mode = inference_mode
+
+        # Initialize LLM (downloading from HF Hub if necessary) --> `llm_cls` is the actual {Model}ForCausalLM class!
+        #   => Note: We're eschewing use of the AutoModel API so that we can be more explicit about LLM-specific details
+        if not self.inference_mode:
+            overwatch.info(f"Loading [bold]{llm_family}[/] LLM from [underline]`{hf_hub_path}`[/]", ctx_level=1)
+            self.llm = llm_cls.from_pretrained(
+                hf_hub_path,
+                token=hf_token,
+                use_flash_attention_2=use_flash_attention_2 if not self.inference_mode else False,
+                # The following parameters are set to prevent `UserWarnings` from HF; we want greedy decoding!
+                do_sample=False,
+                temperature=1.0,
+                top_p=1.0,
+            )
+
+        # [Contract] `inference_mode` means we're loading from a pretrained checkpoint; no need to load base weights!
+        else:
+            overwatch.info(f"Building empty [bold]{llm_family}[/] LLM from [underline]`{hf_hub_path}`[/]", ctx_level=1)
+            llm_config = AutoConfig.from_pretrained(hf_hub_path, token=hf_token)
+            self.llm = llm_cls._from_config(llm_config)
+
+        # Lightweight Handling (with extended explanation) for setting some LLM Parameters
+        #   => Set `decoder.use_cache = False` --> incompatible with gradient checkpointing (+ training in general)
+        #
+        #      Reference: https://discuss.huggingface.co/t/what-is-the-purpose-of-use-cache-in-decoder/958
+        self.llm.config.use_cache = False if not self.inference_mode else True
+
+        #   => Turns out that when gradient checkpointing is on and the underlying LLM has no "trainable" parameters
+        #      (requires_grad is False), backprop will fail; setting `enable_input_requires_grad()` registers a new
+        #      forward hook that fixes this =>> also totally safe for the "full finetuning" setting!
+        if not self.inference_mode:
+            self.llm.enable_input_require_grads()
+
+        # Load (Fast) Tokenizer
+        overwatch.info(f"Loading [bold]{llm_family}[/] (Fast) Tokenizer via the AutoTokenizer API", ctx_level=1)
+        self.tokenizer = AutoTokenizer.from_pretrained(
+            hf_hub_path, model_max_length=self.llm_max_length, token=hf_token, padding_side="right"
+        )
+
+        # Validation =>> Our VLM logic currently operates under the assumption that the tokenization of a new input
+        #                starts with a <BOS> token unless `add_special_tokens = False`; for these models, we empirically
+        #                find that adding image patches *after* the BOS leads to much better performance.
+        #
+        # As a result we explicitly validate that a tokenizer conforms to the expected behavior; if you're reading this
+        # line, it's probably because you're adding a new LLM with a different tokenizer behavior. If so, feel free to
+        # override the `SPECIAL_CASES` set below, but make sure to make the appropriate changes in the `datasets.py`
+        # and VLM `forward()` logic!
+        SPECIAL_CASES = {
+            # Phi-2 Tokenizer doesn't add any BOS tokens by default, and sets BOS == EOS == "<|endoftext|>"
+            #   =>> We'll prepend BOS to first input (to play nicely with image token insertion logic; verified that
+            #       this works well with base LLM generation.
+            #   =>> Like Llama-2 Tokenizers -- we'll add a special PAD token for training purposes.
+            "phi-2-3b",
+        }
+        if self.identifier in SPECIAL_CASES:
+            return
+
+        # Note =>> this assert should hold for all Llama-derived tokenizers (`LlamaTokenizerFast` ==> includes Mistral!
+        assert (self.tokenizer("Test 123", add_special_tokens=True).input_ids[0] == self.tokenizer.bos_token_id) and (
+            self.tokenizer("Test 123", add_special_tokens=False).input_ids[0] != self.tokenizer.bos_token_id
+        ), (
+            f"Default Tokenizer of type `{type(self.tokenizer)}` does not automatically prefix inputs with BOS token!\n"
+            "Please read the comment in `base_llm.py` for more information!"
+        )
+
+    def get_fsdp_wrapping_policy(self) -> Callable:
+        """Return a `transformer_auto_wrap_policy` where we wrap each instance of `self.transformer_layer_cls`"""
+        transformer_block_policy = partial(
+            transformer_auto_wrap_policy, transformer_layer_cls={self.transformer_layer_cls}
+        )
+
+        return transformer_block_policy
+
+    def enable_gradient_checkpointing(self) -> None:
+        """Dispatch to underlying LLM instance's `gradient_checkpointing_enable`; defined for all `PretrainedModel`."""
+        self.llm.gradient_checkpointing_enable()
+
+    def embed_input_ids(self, input_ids: torch.LongTensor) -> torch.Tensor:
+        return self.llm.get_input_embeddings()(input_ids)
+
+    # [Contract] Should match the `forward` call of the underlying `llm` instance!
+    def forward(
+        self,
+        input_ids: Optional[torch.LongTensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_values: Optional[List[torch.FloatTensor]] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        labels: Optional[torch.LongTensor] = None,
+        use_cache: Optional[bool] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+    ) -> CausalLMOutputWithPast:
+        output: CausalLMOutputWithPast = self.llm(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            past_key_values=past_key_values,
+            inputs_embeds=inputs_embeds,
+            labels=labels,
+            use_cache=use_cache,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            return_dict=return_dict,
+        )
+        return output

capvector-oft/prismatic/models/backbones/llm/llama2.py

@@ -0,0 +1,102 @@
+"""
+llama2.py
+
+Class definition for all LLMs derived from LlamaForCausalLM.
+"""
+
+from typing import Optional, Sequence, Type
+
+import torch
+from torch import nn as nn
+from transformers import LlamaForCausalLM
+from transformers.models.llama.modeling_llama import LlamaDecoderLayer
+
+from prismatic.models.backbones.llm.base_llm import HFCausalLLMBackbone
+from prismatic.models.backbones.llm.prompting import (
+    LLaMa2ChatPromptBuilder,
+    PromptBuilder,
+    PurePromptBuilder,
+    VicunaV15ChatPromptBuilder,
+)
+
+# Registry =>> Support LLaMa-2 Models (from HF Transformers)
+# fmt: off
+LLAMA2_MODELS = {
+    # === Pure Meta LLaMa-2 (non-instruct/chat-tuned) Models ===
+    "llama2-7b-pure": {
+        "llm_family": "llama2", "llm_cls": LlamaForCausalLM, "hf_hub_path": "meta-llama/Llama-2-7b-hf"
+    },
+
+    "llama2-13b-pure": {
+        "llm_family": "llama2", "llm_cls": LlamaForCausalLM, "hf_hub_path": "meta-llama/Llama-2-13b-hf"
+    },
+
+    # === Meta LLaMa-2 Chat Models ===
+    "llama2-7b-chat": {
+        "llm_family": "llama2", "llm_cls": LlamaForCausalLM, "hf_hub_path": "meta-llama/Llama-2-7b-chat-hf"
+    },
+
+    "llama2-13b-chat": {
+        "llm_family": "llama2", "llm_cls": LlamaForCausalLM, "hf_hub_path": "meta-llama/Llama-2-13b-chat-hf"
+    },
+
+    # === Vicuna v1.5 Chat Models ===
+    "vicuna-v15-7b": {
+        "llm_family": "llama2", "llm_cls": LlamaForCausalLM, "hf_hub_path": "lmsys/vicuna-7b-v1.5"
+    },
+
+    "vicuna-v15-13b": {
+        "llm_family": "llama2", "llm_cls": LlamaForCausalLM, "hf_hub_path": "lmsys/vicuna-13b-v1.5"
+    },
+}
+# fmt: on
+
+
+class LLaMa2LLMBackbone(HFCausalLLMBackbone):
+    def __init__(
+        self,
+        llm_backbone_id: str,
+        llm_max_length: int = 2048,
+        hf_token: Optional[str] = None,
+        inference_mode: bool = False,
+        use_flash_attention_2: bool = True,
+    ) -> None:
+        super().__init__(
+            llm_backbone_id,
+            llm_max_length=llm_max_length,
+            hf_token=hf_token,
+            inference_mode=inference_mode,
+            use_flash_attention_2=use_flash_attention_2,
+            **LLAMA2_MODELS[llm_backbone_id],
+        )
+
+        # [Special Case] LLaMa-2 PAD Token Handling --> for clarity, we add an extra token (and resize)
+        self.tokenizer.add_special_tokens({"pad_token": "<PAD>"})
+        self.llm.config.pad_token_id = self.tokenizer.pad_token_id
+        self.llm.resize_token_embeddings(len(self.tokenizer), pad_to_multiple_of=64)
+
+    @property
+    def prompt_builder_fn(self) -> Type[PromptBuilder]:
+        if self.identifier.startswith("llama2-") and self.identifier.endswith("-pure"):
+            return PurePromptBuilder
+
+        elif self.identifier.startswith("llama2-") and self.identifier.endswith("-chat"):
+            return LLaMa2ChatPromptBuilder
+
+        elif self.identifier.startswith("vicuna"):
+            return VicunaV15ChatPromptBuilder
+
+        raise ValueError(f"No PromptBuilder defined for LLM Backbone `{self.identifier}`")
+
+    @property
+    def transformer_layer_cls(self) -> Type[nn.Module]:
+        return LlamaDecoderLayer
+
+    @property
+    def half_precision_dtype(self) -> torch.dtype:
+        """LLaMa-2 was trained in BF16; see https://huggingface.co/docs/transformers/main/model_doc/llama2."""
+        return torch.bfloat16
+
+    @property
+    def last_layer_finetune_modules(self) -> Sequence[nn.Module]:
+        return (self.llm.model.embed_tokens, self.llm.model.layers[-1], self.llm.lm_head)

capvector-oft/prismatic/models/backbones/llm/mistral.py

@@ -0,0 +1,72 @@
+"""
+mistral.py
+
+Class definition for all LLMs derived from MistralForCausalLM.
+"""
+
+from typing import Optional, Type
+
+import torch
+from torch import nn as nn
+from transformers import MistralForCausalLM
+from transformers.models.mistral.modeling_mistral import MistralDecoderLayer
+
+from prismatic.models.backbones.llm.base_llm import HFCausalLLMBackbone
+from prismatic.models.backbones.llm.prompting import MistralInstructPromptBuilder, PromptBuilder, PurePromptBuilder
+
+# Registry =>> Support Mistral Models (from HF Transformers)
+# fmt: off
+MISTRAL_MODELS = {
+    # === Base Mistral v0.1 ===
+    "mistral-v0.1-7b-pure": {
+        "llm_family": "mistral", "llm_cls": MistralForCausalLM, "hf_hub_path": "mistralai/Mistral-7B-v0.1"
+    },
+
+    # === Mistral Instruct v0.1 ===
+    "mistral-v0.1-7b-instruct": {
+        "llm_family": "mistral", "llm_cls": MistralForCausalLM, "hf_hub_path": "mistralai/Mistral-7B-Instruct-v0.1"
+    }
+}
+# fmt: on
+
+
+class MistralLLMBackbone(HFCausalLLMBackbone):
+    def __init__(
+        self,
+        llm_backbone_id: str,
+        llm_max_length: int = 2048,
+        hf_token: Optional[str] = None,
+        inference_mode: bool = False,
+        use_flash_attention_2: bool = True,
+    ) -> None:
+        super().__init__(
+            llm_backbone_id,
+            llm_max_length=llm_max_length,
+            hf_token=hf_token,
+            inference_mode=inference_mode,
+            use_flash_attention_2=use_flash_attention_2,
+            **MISTRAL_MODELS[llm_backbone_id],
+        )
+
+        # [Special Case] Mistral PAD Token Handling --> for clarity, we add an extra token (and resize)
+        self.tokenizer.add_special_tokens({"pad_token": "<PAD>"})
+        self.llm.config.pad_token_id = self.tokenizer.pad_token_id
+        self.llm.resize_token_embeddings(len(self.tokenizer), pad_to_multiple_of=64)
+
+    @property
+    def prompt_builder_fn(self) -> Type[PromptBuilder]:
+        if self.identifier.endswith("-pure"):
+            return PurePromptBuilder
+
+        elif self.identifier.endswith("-instruct"):
+            return MistralInstructPromptBuilder
+
+        raise ValueError(f"No PromptBuilder defined for LLM Backbone `{self.identifier}`")
+
+    @property
+    def transformer_layer_cls(self) -> Type[nn.Module]:
+        return MistralDecoderLayer
+
+    @property
+    def half_precision_dtype(self) -> torch.dtype:
+        return torch.bfloat16

capvector-oft/prismatic/models/backbones/llm/phi.py

@@ -0,0 +1,64 @@
+"""
+phi.py
+
+Class definition for all LLMs derived from PhiForCausalLM.
+"""
+
+from typing import Optional, Type
+
+import torch
+from torch import nn as nn
+from transformers import PhiForCausalLM
+from transformers.models.phi.modeling_phi import PhiDecoderLayer
+
+from prismatic.models.backbones.llm.base_llm import HFCausalLLMBackbone
+from prismatic.models.backbones.llm.prompting import PhiPromptBuilder, PromptBuilder
+
+# Registry ==> Support Phi Models (from HF Transformers)
+# fmt: off
+PHI_MODELS = {
+    # === Phi-2 ===
+    "phi-2-3b": {
+        "llm_family": "phi", "llm_cls": PhiForCausalLM, "hf_hub_path": "microsoft/phi-2"
+    }
+}
+# fmt: on
+
+
+class PhiLLMBackbone(HFCausalLLMBackbone):
+    def __init__(
+        self,
+        llm_backbone_id: str,
+        llm_max_length: int = 2048,
+        hf_token: Optional[str] = None,
+        inference_mode: bool = False,
+        use_flash_attention_2: bool = True,
+    ) -> None:
+        super().__init__(
+            llm_backbone_id,
+            llm_max_length=llm_max_length,
+            hf_token=hf_token,
+            inference_mode=inference_mode,
+            use_flash_attention_2=use_flash_attention_2,
+            **PHI_MODELS[llm_backbone_id],
+        )
+
+        # [Special Case] Phi PAD Token Handling --> for clarity, we add an extra token (and resize)
+        self.tokenizer.add_special_tokens({"pad_token": "<|pad|>"})
+        self.llm.config.pad_token_id = self.tokenizer.pad_token_id
+        self.llm.resize_token_embeddings(len(self.tokenizer), pad_to_multiple_of=64)
+
+    @property
+    def prompt_builder_fn(self) -> Type[PromptBuilder]:
+        if self.identifier.startswith("phi-2"):
+            return PhiPromptBuilder
+
+        raise ValueError(f"No PromptBuilder defined for LLM Backbone `{self.identifier}`")
+
+    @property
+    def transformer_layer_cls(self) -> Type[nn.Module]:
+        return PhiDecoderLayer
+
+    @property
+    def half_precision_dtype(self) -> torch.dtype:
+        return torch.bfloat16

capvector-oft/prismatic/models/backbones/llm/prompting/__init__.py

@@ -0,0 +1,5 @@
+from .base_prompter import PromptBuilder, PurePromptBuilder
+from .llama2_chat_prompter import LLaMa2ChatPromptBuilder
+from .mistral_instruct_prompter import MistralInstructPromptBuilder
+from .phi_prompter import PhiPromptBuilder
+from .vicuna_v15_prompter import VicunaV15ChatPromptBuilder

capvector-oft/prismatic/models/backbones/llm/prompting/base_prompter.py

@@ -0,0 +1,73 @@
+"""
+base_prompter.py
+
+Abstract class definition of a multi-turn prompt builder for ensuring consistent formatting for chat-based LLMs.
+"""
+
+from abc import ABC, abstractmethod
+from typing import Optional
+
+
+class PromptBuilder(ABC):
+    def __init__(self, model_family: str, system_prompt: Optional[str] = None) -> None:
+        self.model_family = model_family
+
+        # Only some models define a system prompt => let subclasses handle this logic!
+        self.system_prompt = system_prompt
+
+    @abstractmethod
+    def add_turn(self, role: str, message: str) -> str: ...
+
+    @abstractmethod
+    def get_potential_prompt(self, user_msg: str) -> None: ...
+
+    @abstractmethod
+    def get_prompt(self) -> str: ...
+
+
+class PurePromptBuilder(PromptBuilder):
+    def __init__(self, model_family: str, system_prompt: Optional[str] = None) -> None:
+        super().__init__(model_family, system_prompt)
+
+        # TODO (siddk) =>> Can't always assume LlamaTokenizer --> FIX ME!
+        self.bos, self.eos = "<s>", "</s>"
+
+        # Get role-specific "wrap" functions
+        self.wrap_human = lambda msg: f"In: {msg}\nOut: "
+        self.wrap_gpt = lambda msg: f"{msg if msg != '' else ' '}{self.eos}"
+
+        # === `self.prompt` gets built up over multiple turns ===
+        self.prompt, self.turn_count = "", 0
+
+    def add_turn(self, role: str, message: str) -> str:
+        assert (role == "human") if (self.turn_count % 2 == 0) else (role == "gpt")
+        message = message.replace("<image>", "").strip()
+
+        if (self.turn_count % 2) == 0:
+            human_message = self.wrap_human(message)
+            wrapped_message = human_message
+        else:
+            gpt_message = self.wrap_gpt(message)
+            wrapped_message = gpt_message
+
+        # Update Prompt
+        self.prompt += wrapped_message
+
+        # Bump Turn Counter
+        self.turn_count += 1
+
+        # Return "wrapped_message" (effective string added to context)
+        return wrapped_message
+
+    def get_potential_prompt(self, message: str) -> None:
+        # Assumes that it's always the user's (human's) turn!
+        prompt_copy = str(self.prompt)
+
+        human_message = self.wrap_human(message)
+        prompt_copy += human_message
+
+        return prompt_copy.removeprefix(self.bos).rstrip()
+
+    def get_prompt(self) -> str:
+        # Remove prefix <bos> (if exists) because it gets auto-inserted by tokenizer!
+        return self.prompt.removeprefix(self.bos).rstrip()

capvector-oft/prismatic/models/backbones/llm/prompting/llama2_chat_prompter.py

@@ -0,0 +1,91 @@
+"""
+llama2_prompter.py
+
+Defines a PromptBuilder for building LLaMa-2 Chat Prompts --> not sure if this is "optimal", but this is the pattern
+that's used by HF and other online tutorials.
+
+Reference: https://huggingface.co/blog/llama2#how-to-prompt-llama-2
+"""
+
+from typing import Optional
+
+from prismatic.models.backbones.llm.prompting.base_prompter import PromptBuilder
+
+# Default System Prompt for Prismatic Models
+SYS_PROMPTS = {
+    "prismatic": (
+        "You are a helpful language and vision assistant. "
+        "You are able to understand the visual content that the user provides, "
+        "and assist the user with a variety of tasks using natural language."
+    ),
+    "openvla": (
+        "You are a helpful language and vision assistant. "
+        "You are able to understand the visual content that the user provides, "
+        "and assist the user with a variety of tasks using natural language."
+    ),
+}
+
+
+def format_system_prompt(system_prompt: str) -> str:
+    return f"<<SYS>\n{system_prompt.strip()}\n<</SYS>>\n\n"
+
+
+class LLaMa2ChatPromptBuilder(PromptBuilder):
+    def __init__(self, model_family: str, system_prompt: Optional[str] = None) -> None:
+        super().__init__(model_family, system_prompt)
+        self.system_prompt = format_system_prompt(
+            SYS_PROMPTS[self.model_family] if system_prompt is None else system_prompt
+        )
+
+        # LLaMa-2 Specific
+        self.bos, self.eos = "<s>", "</s>"
+
+        # Get role-specific "wrap" functions
+        self.wrap_human = lambda msg: f"[INST] {msg} [/INST] "
+        self.wrap_gpt = lambda msg: f"{msg if msg != '' else ' '}{self.eos}"
+
+        # === `self.prompt` gets built up over multiple turns ===
+        self.prompt, self.turn_count = "", 0
+
+    def add_turn(self, role: str, message: str) -> str:
+        assert (role == "human") if (self.turn_count % 2 == 0) else (role == "gpt")
+        message = message.replace("<image>", "").strip()
+
+        # Special Handling for "system" prompt (turn_count == 0)
+        if self.turn_count == 0:
+            sys_message = self.wrap_human(self.system_prompt + message)
+            wrapped_message = sys_message
+        elif (self.turn_count % 2) == 0:
+            human_message = self.wrap_human(message)
+            wrapped_message = human_message
+        else:
+            gpt_message = self.wrap_gpt(message)
+            wrapped_message = gpt_message
+
+        # Update Prompt
+        self.prompt += wrapped_message
+
+        # Bump Turn Counter
+        self.turn_count += 1
+
+        # Return "wrapped_message" (effective string added to context)
+        return wrapped_message
+
+    def get_potential_prompt(self, message: str) -> None:
+        # Assumes that it's always the user's (human's) turn!
+        prompt_copy = str(self.prompt)
+
+        # Special Handling for "system" prompt (turn_count == 0)
+        if self.turn_count == 0:
+            sys_message = self.wrap_human(self.system_prompt + message)
+            prompt_copy += sys_message
+
+        else:
+            human_message = self.wrap_human(message)
+            prompt_copy += human_message
+
+        return prompt_copy.removeprefix(self.bos).rstrip()
+
+    def get_prompt(self) -> str:
+        # Remove prefix <bos> because it gets auto-inserted by tokenizer!
+        return self.prompt.removeprefix(self.bos).rstrip()

capvector-oft/prismatic/models/backbones/llm/prompting/mistral_instruct_prompter.py

@@ -0,0 +1,60 @@
+"""
+mistral_instruct_prompter.py
+
+Defines a PromptBuilder for building Mistral Instruct Chat Prompts --> recommended pattern used by HF / online tutorial.s
+
+Reference: https://huggingface.co/mistralai/Mistral-7B-Instruct-v0.1#instruction-format
+"""
+
+from typing import Optional
+
+from prismatic.models.backbones.llm.prompting.base_prompter import PromptBuilder
+
+
+class MistralInstructPromptBuilder(PromptBuilder):
+    def __init__(self, model_family: str, system_prompt: Optional[str] = None) -> None:
+        super().__init__(model_family, system_prompt)
+
+        # Note =>> Mistral Tokenizer is an instance of `LlamaTokenizer(Fast)`
+        #      =>> Mistral Instruct *does not* use a System Prompt
+        self.bos, self.eos = "<s>", "</s>"
+
+        # Get role-specific "wrap" functions
+        self.wrap_human = lambda msg: f"[INST] {msg} [/INST] "
+        self.wrap_gpt = lambda msg: f"{msg if msg != '' else ' '}{self.eos}"
+
+        # === `self.prompt` gets built up over multiple turns ===
+        self.prompt, self.turn_count = "", 0
+
+    def add_turn(self, role: str, message: str) -> str:
+        assert (role == "human") if (self.turn_count % 2 == 0) else (role == "gpt")
+        message = message.replace("<image>", "").strip()
+
+        if (self.turn_count % 2) == 0:
+            human_message = self.wrap_human(message)
+            wrapped_message = human_message
+        else:
+            gpt_message = self.wrap_gpt(message)
+            wrapped_message = gpt_message
+
+        # Update Prompt
+        self.prompt += wrapped_message
+
+        # Bump Turn Counter
+        self.turn_count += 1
+
+        # Return "wrapped_message" (effective string added to context)
+        return wrapped_message
+
+    def get_potential_prompt(self, message: str) -> None:
+        # Assumes that it's always the user's (human's) turn!
+        prompt_copy = str(self.prompt)
+
+        human_message = self.wrap_human(message)
+        prompt_copy += human_message
+
+        return prompt_copy.removeprefix(self.bos).rstrip()
+
+    def get_prompt(self) -> str:
+        # Remove prefix <bos> because it gets auto-inserted by tokenizer!
+        return self.prompt.removeprefix(self.bos).rstrip()

capvector-oft/prismatic/models/backbones/llm/prompting/phi_prompter.py

@@ -0,0 +1,65 @@
+"""
+phi_prompter.py
+
+Defines a PromptBuilder for building Phi-2 Input/Output Prompts --> recommended pattern used by HF / Microsoft.
+Also handles Phi special case BOS token additions.
+
+Reference: https://huggingface.co/microsoft/phi-2#qa-format
+"""
+
+from typing import Optional
+
+from prismatic.models.backbones.llm.prompting.base_prompter import PromptBuilder
+
+
+class PhiPromptBuilder(PromptBuilder):
+    def __init__(self, model_family: str, system_prompt: Optional[str] = None) -> None:
+        super().__init__(model_family, system_prompt)
+
+        # Note =>> Phi Tokenizer is an instance of `CodeGenTokenizer(Fast)`
+        #      =>> By default, does *not* append <BOS> / <EOS> tokens --> we handle that here (IMPORTANT)!
+        self.bos, self.eos = "<|endoftext|>", "<|endoftext|>"
+
+        # Get role-specific "wrap" functions
+        #   =>> Note that placement of <bos>/<eos> were based on experiments generating from Phi-2 in Input/Output mode
+        self.wrap_human = lambda msg: f"Input: {msg}\nOutput: "
+        self.wrap_gpt = lambda msg: f"{msg if msg != '' else ' '}\n{self.eos}"
+
+        # === `self.prompt` gets built up over multiple turns ===
+        self.prompt, self.turn_count = "", 0
+
+    def add_turn(self, role: str, message: str) -> str:
+        assert (role == "human") if (self.turn_count % 2 == 0) else (role == "gpt")
+        message = message.replace("<image>", "").strip()
+
+        # Special Handling for "first" input --> prepend a <BOS> token (expected by Prismatic)
+        if self.turn_count == 0:
+            bos_human_message = f"{self.bos}{self.wrap_human(message)}"
+            wrapped_message = bos_human_message
+        elif (self.turn_count % 2) == 0:
+            human_message = self.wrap_human(message)
+            wrapped_message = human_message
+        else:
+            gpt_message = self.wrap_gpt(message)
+            wrapped_message = gpt_message
+
+        # Update Prompt
+        self.prompt += wrapped_message
+
+        # Bump Turn Counter
+        self.turn_count += 1
+
+        # Return "wrapped_message" (effective string added to context)
+        return wrapped_message
+
+    def get_potential_prompt(self, message: str) -> None:
+        # Assumes that it's always the user's (human's) turn!
+        prompt_copy = str(self.prompt)
+
+        human_message = self.wrap_human(message)
+        prompt_copy += human_message
+
+        return prompt_copy.rstrip()
+
+    def get_prompt(self) -> str:
+        return self.prompt.rstrip()

capvector-oft/prismatic/models/backbones/llm/prompting/vicuna_v15_prompter.py

@@ -0,0 +1,82 @@
+"""
+vicuna_v15_prompter.py
+
+Defines a PromptBuilder for building Vicuna-v1.5 Chat Prompts.
+
+Reference: https://huggingface.co/lmsys/vicuna-13b-v1.5
+"""
+
+from typing import Optional
+
+from prismatic.models.backbones.llm.prompting.base_prompter import PromptBuilder
+
+# Default System Prompt for LLaVa Models
+SYS_PROMPTS = {
+    "prismatic": (
+        "A chat between a curious user and an artificial intelligence assistant. "
+        "The assistant gives helpful, detailed, and polite answers to the user's questions."
+    ),
+    "openvla": (
+        "A chat between a curious user and an artificial intelligence assistant. "
+        "The assistant gives helpful, detailed, and polite answers to the user's questions."
+    ),
+}
+
+
+class VicunaV15ChatPromptBuilder(PromptBuilder):
+    def __init__(self, model_family: str, system_prompt: Optional[str] = None) -> None:
+        super().__init__(model_family, system_prompt)
+        self.system_prompt = (SYS_PROMPTS[self.model_family] if system_prompt is None else system_prompt).strip() + " "
+
+        # LLaMa-2 Specific
+        self.bos, self.eos = "<s>", "</s>"
+
+        # Get role-specific "wrap" functions
+        self.wrap_human = lambda msg: f"USER: {msg} ASSISTANT: "
+        self.wrap_gpt = lambda msg: f"{msg if msg != '' else ' '}{self.eos}"
+
+        # === `self.prompt` gets built up over multiple turns ===
+        self.prompt, self.turn_count = "", 0
+
+    def add_turn(self, role: str, message: str) -> str:
+        assert (role == "human") if (self.turn_count % 2 == 0) else (role == "gpt")
+        message = message.replace("<image>", "").strip()
+
+        # Special Handling for "system" prompt (turn_count == 0)
+        if self.turn_count == 0:
+            sys_message = self.system_prompt + self.wrap_human(message)
+            wrapped_message = sys_message
+        elif (self.turn_count % 2) == 0:
+            human_message = self.wrap_human(message)
+            wrapped_message = human_message
+        else:
+            gpt_message = self.wrap_gpt(message)
+            wrapped_message = gpt_message
+
+        # Update Prompt
+        self.prompt += wrapped_message
+
+        # Bump Turn Counter
+        self.turn_count += 1
+
+        # Return "wrapped_message" (effective string added to context)
+        return wrapped_message
+
+    def get_potential_prompt(self, message: str) -> None:
+        # Assumes that it's always the user's (human's) turn!
+        prompt_copy = str(self.prompt)
+
+        # Special Handling for "system" prompt (turn_count == 0)
+        if self.turn_count == 0:
+            sys_message = self.system_prompt + self.wrap_human(message)
+            prompt_copy += sys_message
+
+        else:
+            human_message = self.wrap_human(message)
+            prompt_copy += human_message
+
+        return prompt_copy.removeprefix(self.bos).rstrip()
+
+    def get_prompt(self) -> str:
+        # Remove prefix <bos> (if exists) because it gets auto-inserted by tokenizer!
+        return self.prompt.removeprefix(self.bos).rstrip()

capvector-oft/prismatic/models/backbones/vision/__init__.py

@@ -0,0 +1,7 @@
+from .base_vision import ImageTransform, VisionBackbone
+from .clip_vit import CLIPViTBackbone
+from .dinoclip_vit import DinoCLIPViTBackbone
+from .dinosiglip_vit import DinoSigLIPViTBackbone
+from .dinov2_vit import DinoV2ViTBackbone
+from .in1k_vit import IN1KViTBackbone
+from .siglip_vit import SigLIPViTBackbone

capvector-oft/prismatic/models/backbones/vision/base_vision.py

@@ -0,0 +1,207 @@
+"""
+base_vision.py
+
+Abstract class definition of a Vision Backbone (Visual Featurizer), with full annotations of class methods, utility
+functions, and initialization logic.
+
+We also define the generic TimmViTBackbone class here, providing a default interface for loading any TIMM Vision
+Transformer model for feature extraction.
+"""
+
+from abc import ABC, abstractmethod
+from dataclasses import dataclass
+from functools import partial
+from typing import Any, Callable, Dict, Optional, Protocol, Tuple, Union
+
+import timm
+import torch
+import torch.nn as nn
+import torchvision.transforms.functional as TVF
+from PIL.Image import Image
+from timm.models.vision_transformer import Block, VisionTransformer
+from torch.distributed.fsdp.wrap import _module_wrap_policy, _or_policy, transformer_auto_wrap_policy
+from torchvision.transforms import Compose, Resize
+
+
+# === Utility Functions for Monkey-Patching ===
+def unpack_tuple(fn: Callable[[Any], Tuple[Any]]) -> Callable[[Any], Any]:
+    def wrapper(*args: Any, **kwargs: Any) -> Any:
+        result = fn(*args, **kwargs)
+        return result[0] if isinstance(result, tuple) else result
+
+    return wrapper
+
+
+# === Interface for an Image Transform ===
+class ImageTransform(Protocol):
+    def __call__(self, img: Image, **kwargs: str) -> Union[torch.Tensor, Dict[str, torch.Tensor]]: ...
+
+
+# === Custom Torchvision Image Transforms ===
+@dataclass
+class LetterboxPad:
+    padding_fill_value: Tuple[int, int, int]
+
+    def __call__(self, image: Image) -> Image:
+        """Given a PIL.Image, pad to square by adding a symmetric border around the height/width."""
+        (w, h), max_wh = image.size, max(image.size)
+        horizontal_pad, vertical_pad = int((max_wh - w) / 2), int((max_wh - h) / 2)
+        padding = (horizontal_pad, vertical_pad, horizontal_pad, vertical_pad)
+        return TVF.pad(image, padding, fill=self.padding_fill_value, padding_mode="constant")
+
+
+# === Abstract Base Class for arbitrary Vision Backbones ===
+class VisionBackbone(nn.Module, ABC):
+    def __init__(self, vision_backbone_id: str, image_resize_strategy: str, default_image_size: int = 224) -> None:
+        super().__init__()
+        self.identifier: str = vision_backbone_id
+        self.image_resize_strategy: str = image_resize_strategy
+        self.default_image_size: int = default_image_size
+
+        # Instance attributes for a Vision Backbone
+        self.featurizer: nn.Module = None
+        self.image_transform: ImageTransform = None
+
+    def get_image_transform(self) -> ImageTransform:
+        return self.image_transform
+
+    @abstractmethod
+    def get_fsdp_wrapping_policy(self) -> Callable: ...
+
+    @abstractmethod
+    def forward(self, pixel_values: torch.Tensor) -> torch.Tensor:
+        """Run a forward pass through the featurizer given a set of processed images, returning patch/grid features."""
+        raise NotImplementedError
+
+    @property
+    @abstractmethod
+    def default_image_resolution(self) -> Tuple[int, int, int]: ...
+
+    @property
+    @abstractmethod
+    def embed_dim(self) -> int: ...
+
+    @property
+    @abstractmethod
+    def num_patches(self) -> int: ...
+
+    @property
+    @abstractmethod
+    def half_precision_dtype(self) -> torch.dtype: ...
+
+
+# === Abstract Base Class for Arbitrary TIMM Vision Transformer Backbones ===
+class TimmViTBackbone(VisionBackbone, ABC):
+    def __init__(
+        self,
+        vision_backbone_id: str,
+        timm_path_or_url: str,
+        image_resize_strategy: str,
+        default_image_size: int = 224,
+        override_act_layer: Optional[str] = None,
+    ) -> None:
+        super().__init__(vision_backbone_id, image_resize_strategy, default_image_size=default_image_size)
+        self.timm_path_or_url = timm_path_or_url
+        self.override_act_layer = override_act_layer
+        self.dtype = torch.bfloat16
+
+        # Initialize Featurizer (ViT) by downloading from HF / TIMM Hub if necessary
+        if self.override_act_layer is None:
+            self.featurizer: VisionTransformer = timm.create_model(
+                self.timm_path_or_url, pretrained=True, num_classes=0, img_size=self.default_image_size
+            )
+        else:
+            self.featurizer: VisionTransformer = timm.create_model(
+                self.timm_path_or_url,
+                pretrained=True,
+                num_classes=0,
+                img_size=self.default_image_size,
+                act_layer=self.override_act_layer,
+            )
+        self.featurizer.eval()
+
+        # Monkey-Patch the `forward()` function of the featurizer to ensure FSDP-compatibility
+        #   => Note: By default set `get_intermediate_layers` to return the *SECOND-TO-LAST* layer patches!
+        #   => TODO (siddk) Remove after resolution of https://github.com/pytorch/pytorch/issues/109385
+        self.featurizer.forward = unpack_tuple(
+            partial(self.featurizer.get_intermediate_layers, n={len(self.featurizer.blocks) - 2})
+        )
+
+        # Validation =>> for now, this class *only* supports TIMM Vision Transformers (but can be extended!)
+        assert isinstance(self.featurizer, VisionTransformer), (
+            "Featurizer is not a TIMM VisionTransformer; if you would like to support a new visual representation, "
+            "file an issue or implement the requisite logic (see `prismatic/models/backbones/vision/base_vision.py`)!"
+        )
+
+        # Get Config =>> Note :: Override default image size to ensure correct image transform
+        self.data_cfg = timm.data.resolve_model_data_config(self.featurizer)
+        self.data_cfg["input_size"] = (3, self.default_image_size, self.default_image_size)
+
+        # Initialize Default Image Transform --> Modified by `self.image_resize_strategy`
+        default_image_transform = timm.data.create_transform(**self.data_cfg, is_training=False)
+
+        # Fix =>> SigLIP & IN1K default transforms resize to *larger* than `self.default_image_size` (crops image)!
+        if "siglip" in self.timm_path_or_url or "in1k" in self.timm_path_or_url:
+            assert isinstance(default_image_transform, Compose), "Unexpected `default_image_transform`!"
+            assert isinstance(default_image_transform.transforms[0], Resize)
+            default_image_transform = Compose(
+                [
+                    Resize(self.default_image_size, interpolation=default_image_transform.transforms[0].interpolation),
+                    *default_image_transform.transforms[1:],
+                ]
+            )
+
+        # Switch on `image_resize_strategy`
+        if self.image_resize_strategy == "resize-naive":
+            assert isinstance(default_image_transform, Compose), "Unexpected `default_image_transform`!"
+            assert isinstance(default_image_transform.transforms[0], Resize)
+
+            target_size = (self.default_image_size, self.default_image_size)
+            self.image_transform = Compose(
+                [
+                    Resize(target_size, interpolation=default_image_transform.transforms[0].interpolation),
+                    *default_image_transform.transforms[1:],
+                ]
+            )
+
+        elif self.image_resize_strategy == "resize-crop":
+            self.image_transform = default_image_transform
+
+        elif self.image_resize_strategy == "letterbox":
+            assert isinstance(default_image_transform, Compose), "Unexpected `default_image_transform`!"
+            assert "mean" in self.data_cfg, "TIMM `data_cfg` missing image normalization mean!"
+
+            # Compute Padding Fill Value (rescaled normalization mean if applicable)
+            fill = tuple([int(x * 255) for x in self.data_cfg["mean"]])
+
+            # Build New Transform
+            self.image_transform = Compose([LetterboxPad(fill), *default_image_transform.transforms])
+
+        else:
+            raise ValueError(f"Image Resize Strategy `{self.image_resize_strategy}` is not supported!")
+
+    def get_fsdp_wrapping_policy(self) -> Callable:
+        """Return a simple FSDP policy that wraps each ViT block and then the _entire_ featurizer."""
+        vit_wrap_policy = partial(_module_wrap_policy, module_classes={VisionTransformer})
+        transformer_block_policy = partial(transformer_auto_wrap_policy, transformer_layer_cls={Block})
+        return partial(_or_policy, policies=[vit_wrap_policy, transformer_block_policy])
+
+    def forward(self, pixel_values: Union[torch.Tensor, Dict[str, torch.Tensor]]) -> torch.Tensor:
+        """Runs transformed image/pixel tensor through vision backbone, returning _all_ patch features."""
+        return self.featurizer(pixel_values)
+
+    @property
+    def default_image_resolution(self) -> Tuple[int, int, int]:
+        return self.data_cfg["input_size"]
+
+    @property
+    def embed_dim(self) -> int:
+        return self.featurizer.embed_dim
+
+    @property
+    def num_patches(self) -> int:
+        return self.featurizer.patch_embed.num_patches
+
+    @property
+    def half_precision_dtype(self) -> torch.dtype:
+        return self.dtype

capvector-oft/prismatic/models/backbones/vision/clip_vit.py

@@ -0,0 +1,27 @@
+"""
+clip_vit.py
+"""
+
+from prismatic.models.backbones.vision.base_vision import TimmViTBackbone
+
+# Registry =>> Supported CLIP Vision Backbones (from TIMM)
+CLIP_VISION_BACKBONES = {
+    "clip-vit-b": "vit_base_patch16_clip_224.openai",
+    "clip-vit-l": "vit_large_patch14_clip_224.openai",
+    "clip-vit-l-336px": "vit_large_patch14_clip_336.openai",
+}
+
+
+# [IMPORTANT] By Default, TIMM initialized OpenAI CLIP models with the standard GELU activation from PyTorch.
+#             HOWEVER =>> Original OpenAI models were trained with the quick_gelu *approximation* -- while it's
+#                         a decent approximation, the resulting features are *worse*; this was a super tricky bug
+#                         to identify, but luckily there's an easy fix (`override_act_layer`)
+class CLIPViTBackbone(TimmViTBackbone):
+    def __init__(self, vision_backbone_id: str, image_resize_strategy: str, default_image_size: int = 224) -> None:
+        super().__init__(
+            vision_backbone_id,
+            CLIP_VISION_BACKBONES[vision_backbone_id],
+            image_resize_strategy,
+            default_image_size=default_image_size,
+            override_act_layer="quick_gelu" if CLIP_VISION_BACKBONES[vision_backbone_id].endswith(".openai") else None,
+        )

capvector-oft/prismatic/models/backbones/vision/dinoclip_vit.py

@@ -0,0 +1,147 @@
+"""
+dinoclip_vit.py
+
+Vision backbone that returns concatenated features from both DINOv2 and CLIP.
+"""
+
+from dataclasses import dataclass
+from functools import partial
+from typing import Callable, Dict, Tuple
+
+import timm
+import torch
+from PIL import Image
+from timm.models.vision_transformer import Block, VisionTransformer
+from torch.distributed.fsdp.wrap import _module_wrap_policy, _or_policy, transformer_auto_wrap_policy
+from torchvision.transforms import Compose, Resize
+
+from prismatic.models.backbones.vision.base_vision import ImageTransform, LetterboxPad, VisionBackbone, unpack_tuple
+
+# Registry =>> Supported DinoCLIP Pairs (as TIMM identifiers)
+DINOCLIP_VISION_BACKBONES = {
+    "dinoclip-vit-l-336px": {
+        "dino": "vit_large_patch14_reg4_dinov2.lvd142m",
+        "clip": "vit_large_patch14_clip_336.openai",
+    },
+}
+
+
+@dataclass
+class DinoCLIPImageTransform:
+    dino_image_transform: ImageTransform
+    clip_image_transform: ImageTransform
+    is_prismatic: bool = True
+
+    def __call__(self, img: Image, **kwargs: str) -> Dict[str, torch.Tensor]:
+        return {"dino": self.dino_image_transform(img, **kwargs), "clip": self.clip_image_transform(img, **kwargs)}
+
+
+class DinoCLIPViTBackbone(VisionBackbone):
+    def __init__(self, vision_backbone_id: str, image_resize_strategy: str, default_image_size: int = 224) -> None:
+        super().__init__(vision_backbone_id, image_resize_strategy, default_image_size=default_image_size)
+        self.dino_timm_path_or_url = DINOCLIP_VISION_BACKBONES[vision_backbone_id]["dino"]
+        self.clip_timm_path_or_url = DINOCLIP_VISION_BACKBONES[vision_backbone_id]["clip"]
+
+        # Initialize both Featurizers (ViTs) by downloading from HF / TIMM Hub if necessary
+        self.dino_featurizer: VisionTransformer = timm.create_model(
+            self.dino_timm_path_or_url, pretrained=True, num_classes=0, img_size=self.default_image_size
+        )
+        self.dino_featurizer.eval()
+
+        self.clip_featurizer: VisionTransformer = timm.create_model(
+            self.clip_timm_path_or_url, pretrained=True, num_classes=0, img_size=self.default_image_size
+        )
+        self.clip_featurizer.eval()
+
+        # Monkey-Patch the `forward()` function of the featurizers to ensure FSDP-compatibility
+        #   => Note: By default set `get_intermediate_layers` to return the *SECOND-TO-LAST* layer patches!
+        #   => TODO (siddk) Remove after resolution of https://github.com/pytorch/pytorch/issues/109385
+        self.dino_featurizer.forward = unpack_tuple(
+            partial(self.dino_featurizer.get_intermediate_layers, n={len(self.dino_featurizer.blocks) - 2})
+        )
+        self.clip_featurizer.forward = unpack_tuple(
+            partial(self.clip_featurizer.get_intermediate_layers, n={len(self.clip_featurizer.blocks) - 2})
+        )
+
+        # Get Configs for _both_ Featurizers =>> Note :: Override default image size for larger resolution models
+        self.dino_data_cfg = timm.data.resolve_model_data_config(self.dino_featurizer)
+        self.dino_data_cfg["input_size"] = (3, self.default_image_size, self.default_image_size)
+
+        self.clip_data_cfg = timm.data.resolve_model_data_config(self.clip_featurizer)
+        self.clip_data_cfg["input_size"] = (3, self.default_image_size, self.default_image_size)
+
+        # Initialize *both* Transforms
+        default_dino_transform = timm.data.create_transform(**self.dino_data_cfg, is_training=False)
+        default_clip_transform = timm.data.create_transform(**self.clip_data_cfg, is_training=False)
+        if self.image_resize_strategy == "resize-naive":
+            assert isinstance(default_dino_transform, Compose), "Unexpected `default_dino_image_transform`!"
+            assert isinstance(default_clip_transform, Compose), "Unexpected `default_clip_image_transform`!"
+            assert isinstance(default_dino_transform.transforms[0], Resize)
+            assert isinstance(default_clip_transform.transforms[0], Resize)
+
+            target_size = (self.default_image_size, self.default_image_size)
+            dino_transform = Compose(
+                [
+                    Resize(target_size, interpolation=default_dino_transform.transforms[0].interpolation),
+                    *default_dino_transform.transforms[1:],
+                ]
+            )
+            clip_transform = Compose(
+                [
+                    Resize(target_size, interpolation=default_clip_transform.transforms[0].interpolation),
+                    *default_clip_transform.transforms[1:],
+                ]
+            )
+
+            self.image_transform = DinoCLIPImageTransform(dino_transform, clip_transform)
+
+        elif self.image_resize_strategy == "resize-crop":
+            self.image_transform = DinoCLIPImageTransform(default_dino_transform, default_clip_transform)
+
+        elif self.image_resize_strategy == "letterbox":
+            assert isinstance(default_dino_transform, Compose), "Unexpected `default_dino_transform`!"
+            assert isinstance(default_clip_transform, Compose), "Unexpected `default_clip_transform`!"
+            assert "mean" in self.dino_data_cfg and "mean" in self.clip_data_cfg, "DinoCLIP `data_cfg` missing `mean`!"
+
+            # Compute Padding Fill Value(s) (rescaled normalization mean if applicable)
+            dino_fill = tuple([int(x * 255) for x in self.dino_data_cfg["mean"]])
+            clip_fill = tuple([int(x * 255) for x in self.clip_data_cfg["mean"]])
+
+            # Build New Transform
+            self.image_transform = DinoCLIPImageTransform(
+                Compose([LetterboxPad(dino_fill), *default_dino_transform.transforms]),
+                Compose([LetterboxPad(clip_fill), *default_clip_transform.transforms]),
+            )
+
+        else:
+            raise ValueError(f"Image Resize Strategy `{self.image_resize_strategy}` is not supported!")
+
+    def get_fsdp_wrapping_policy(self) -> Callable:
+        """Return a simple FSDP policy that wraps each ViT block and then both of the _entire_ featurizers."""
+        vit_wrap_policy = partial(_module_wrap_policy, module_classes={VisionTransformer})
+        transformer_block_policy = partial(transformer_auto_wrap_policy, transformer_layer_cls={Block})
+        return partial(_or_policy, policies=[vit_wrap_policy, transformer_block_policy])
+
+    def forward(self, pixel_values: Dict[str, torch.Tensor]) -> torch.Tensor:
+        """Runs the transformed image/pixel tensors through each vision backbone, returning concatenated patches."""
+        dino_patches = self.dino_featurizer(pixel_values["dino"])
+        clip_patches = self.clip_featurizer(pixel_values["clip"])
+
+        return torch.cat([dino_patches, clip_patches], dim=2)
+
+    @property
+    def default_image_resolution(self) -> Tuple[int, int, int]:
+        return self.dino_data_cfg["input_size"]
+
+    @property
+    def embed_dim(self) -> int:
+        return self.dino_featurizer.embed_dim + self.clip_featurizer.embed_dim
+
+    @property
+    def num_patches(self) -> int:
+        assert self.dino_featurizer.patch_embed.num_patches == self.clip_featurizer.patch_embed.num_patches
+        return self.dino_featurizer.patch_embed.num_patches
+
+    @property
+    def half_precision_dtype(self) -> torch.dtype:
+        return torch.bfloat16

capvector-oft/prismatic/models/backbones/vision/dinosiglip_vit.py

@@ -0,0 +1,164 @@
+"""
+dinosiglip_vit.py
+
+Vision backbone that returns concatenated features from both DINOv2 and SigLIP.
+"""
+
+from dataclasses import dataclass
+from functools import partial
+from typing import Callable, Dict, Tuple
+
+import timm
+import torch
+from PIL import Image
+from timm.models.vision_transformer import Block, VisionTransformer
+from torch.distributed.fsdp.wrap import _module_wrap_policy, _or_policy, transformer_auto_wrap_policy
+from torchvision.transforms import Compose, Resize
+
+from prismatic.models.backbones.vision.base_vision import ImageTransform, LetterboxPad, VisionBackbone, unpack_tuple
+
+# Registry =>> Supported DinoSigLIP Pairs (as TIMM identifiers)
+DINOSigLIP_VISION_BACKBONES = {
+    "dinosiglip-vit-so-224px": {
+        "dino": "vit_large_patch14_reg4_dinov2.lvd142m",
+        "siglip": "vit_so400m_patch14_siglip_224",
+    },
+    "dinosiglip-vit-so-384px": {
+        "dino": "vit_large_patch14_reg4_dinov2.lvd142m",
+        "siglip": "vit_so400m_patch14_siglip_384",
+    },
+}
+
+
+@dataclass
+class DinoSigLIPImageTransform:
+    dino_image_transform: ImageTransform
+    siglip_image_transform: ImageTransform
+    is_prismatic: bool = True
+
+    def __call__(self, img: Image, **kwargs: str) -> Dict[str, torch.Tensor]:
+        return {"dino": self.dino_image_transform(img, **kwargs), "siglip": self.siglip_image_transform(img, **kwargs)}
+
+
+class DinoSigLIPViTBackbone(VisionBackbone):
+    def __init__(self, vision_backbone_id: str, image_resize_strategy: str, default_image_size: int = 224) -> None:
+        super().__init__(vision_backbone_id, image_resize_strategy, default_image_size=default_image_size)
+        self.dino_timm_path_or_url = DINOSigLIP_VISION_BACKBONES[vision_backbone_id]["dino"]
+        self.siglip_timm_path_or_url = DINOSigLIP_VISION_BACKBONES[vision_backbone_id]["siglip"]
+
+        # Initialize both Featurizers (ViTs) by downloading from HF / TIMM Hub if necessary
+        self.dino_featurizer: VisionTransformer = timm.create_model(
+            self.dino_timm_path_or_url, pretrained=True, num_classes=0, img_size=self.default_image_size
+        )
+        self.dino_featurizer.eval()
+
+        self.siglip_featurizer: VisionTransformer = timm.create_model(
+            self.siglip_timm_path_or_url, pretrained=True, num_classes=0, img_size=self.default_image_size
+        )
+        self.siglip_featurizer.eval()
+
+        # Monkey-Patch the `forward()` function of the featurizers to ensure FSDP-compatibility
+        #   => Note: By default set `get_intermediate_layers` to return the *SECOND-TO-LAST* layer patches!
+        #   => TODO (siddk) Remove after resolution of https://github.com/pytorch/pytorch/issues/109385
+        self.dino_featurizer.forward = unpack_tuple(
+            partial(self.dino_featurizer.get_intermediate_layers, n={len(self.dino_featurizer.blocks) - 2})
+        )
+        self.siglip_featurizer.forward = unpack_tuple(
+            partial(self.siglip_featurizer.get_intermediate_layers, n={len(self.siglip_featurizer.blocks) - 2})
+        )
+
+        # Get Configs for _both_ Featurizers =>> Note :: Override default image size for larger resolution models
+        self.dino_data_cfg = timm.data.resolve_model_data_config(self.dino_featurizer)
+        self.dino_data_cfg["input_size"] = (3, self.default_image_size, self.default_image_size)
+
+        self.siglip_data_cfg = timm.data.resolve_model_data_config(self.siglip_featurizer)
+        self.siglip_data_cfg["input_size"] = (3, self.default_image_size, self.default_image_size)
+
+        # Initialize *both* Transforms
+        default_dino_transform = timm.data.create_transform(**self.dino_data_cfg, is_training=False)
+        default_siglip_transform = timm.data.create_transform(**self.siglip_data_cfg, is_training=False)
+
+        # Fix =>> SigLIP default transform resizes to *larger* than `self.default_image_size` (crops image)!!
+        assert isinstance(default_siglip_transform, Compose), "Unexpected `default_image_transform`!"
+        assert isinstance(default_siglip_transform.transforms[0], Resize)
+        default_siglip_transform = Compose(
+            [
+                Resize(self.default_image_size, interpolation=default_siglip_transform.transforms[0].interpolation),
+                *default_siglip_transform.transforms[1:],
+            ]
+        )
+
+        if self.image_resize_strategy == "resize-naive":
+            assert isinstance(default_dino_transform, Compose), "Unexpected `default_dino_image_transform`!"
+            assert isinstance(default_siglip_transform, Compose), "Unexpected `default_siglip_image_transform`!"
+            assert isinstance(default_dino_transform.transforms[0], Resize)
+            assert isinstance(default_siglip_transform.transforms[0], Resize)
+
+            target_size = (self.default_image_size, self.default_image_size)
+            dino_transform = Compose(
+                [
+                    Resize(target_size, interpolation=default_dino_transform.transforms[0].interpolation),
+                    *default_dino_transform.transforms[1:],
+                ]
+            )
+            siglip_transform = Compose(
+                [
+                    Resize(target_size, interpolation=default_siglip_transform.transforms[0].interpolation),
+                    *default_siglip_transform.transforms[1:],
+                ]
+            )
+
+            self.image_transform = DinoSigLIPImageTransform(dino_transform, siglip_transform)
+
+        elif self.image_resize_strategy == "resize-crop":
+            self.image_transform = DinoSigLIPImageTransform(default_dino_transform, default_siglip_transform)
+
+        elif self.image_resize_strategy == "letterbox":
+            assert isinstance(default_dino_transform, Compose), "Unexpected `default_dino_transform`!"
+            assert isinstance(default_siglip_transform, Compose), "Unexpected `default_siglip_transform`!"
+            assert (
+                "mean" in self.dino_data_cfg and "mean" in self.siglip_data_cfg
+            ), "DinoSigLIP `data_cfg` missing `mean`!"
+
+            # Compute Padding Fill Value(s) (rescaled normalization mean if applicable)
+            dino_fill = tuple([int(x * 255) for x in self.dino_data_cfg["mean"]])
+            siglip_fill = tuple([int(x * 255) for x in self.siglip_data_cfg["mean"]])
+
+            # Build New Transform
+            self.image_transform = DinoSigLIPImageTransform(
+                Compose([LetterboxPad(dino_fill), *default_dino_transform.transforms]),
+                Compose([LetterboxPad(siglip_fill), *default_siglip_transform.transforms]),
+            )
+
+        else:
+            raise ValueError(f"Image Resize Strategy `{self.image_resize_strategy}` is not supported!")
+
+    def get_fsdp_wrapping_policy(self) -> Callable:
+        """Return a simple FSDP policy that wraps each ViT block and then both of the _entire_ featurizers."""
+        vit_wrap_policy = partial(_module_wrap_policy, module_classes={VisionTransformer})
+        transformer_block_policy = partial(transformer_auto_wrap_policy, transformer_layer_cls={Block})
+        return partial(_or_policy, policies=[vit_wrap_policy, transformer_block_policy])
+
+    def forward(self, pixel_values: Dict[str, torch.Tensor]) -> torch.Tensor:
+        """Runs the transformed image/pixel tensors through each vision backbone, returning concatenated patches."""
+        dino_patches = self.dino_featurizer(pixel_values["dino"])
+        siglip_patches = self.siglip_featurizer(pixel_values["siglip"])
+
+        return torch.cat([dino_patches, siglip_patches], dim=2)
+
+    @property
+    def default_image_resolution(self) -> Tuple[int, int, int]:
+        return self.dino_data_cfg["input_size"]
+
+    @property
+    def embed_dim(self) -> int:
+        return self.dino_featurizer.embed_dim + self.siglip_featurizer.embed_dim
+
+    @property
+    def num_patches(self) -> int:
+        assert self.dino_featurizer.patch_embed.num_patches == self.siglip_featurizer.patch_embed.num_patches
+        return self.dino_featurizer.patch_embed.num_patches
+
+    @property
+    def half_precision_dtype(self) -> torch.dtype:
+        return torch.bfloat16

capvector-oft/prismatic/models/backbones/vision/dinov2_vit.py

@@ -0,0 +1,19 @@
+"""
+dinov2_vit.py
+"""
+
+from prismatic.models.backbones.vision.base_vision import TimmViTBackbone
+
+# Registry =>> Supported DINOv2 Vision Backbones (from TIMM) =>> Note:: Using DINOv2 w/ Registers!
+#   => Reference: https://arxiv.org/abs/2309.16588
+DINOv2_VISION_BACKBONES = {"dinov2-vit-l": "vit_large_patch14_reg4_dinov2.lvd142m"}
+
+
+class DinoV2ViTBackbone(TimmViTBackbone):
+    def __init__(self, vision_backbone_id: str, image_resize_strategy: str, default_image_size: int = 224) -> None:
+        super().__init__(
+            vision_backbone_id,
+            DINOv2_VISION_BACKBONES[vision_backbone_id],
+            image_resize_strategy,
+            default_image_size=default_image_size,
+        )

capvector-oft/prismatic/models/backbones/vision/in1k_vit.py

@@ -0,0 +1,22 @@
+"""
+in1k_vit.py
+
+Vision Transformers trained / finetuned on ImageNet (ImageNet-21K =>> ImageNet-1K)
+"""
+
+from prismatic.models.backbones.vision.base_vision import TimmViTBackbone
+
+# Registry =>> Supported Vision Backbones (from TIMM)
+IN1K_VISION_BACKBONES = {
+    "in1k-vit-l": "vit_large_patch16_224.augreg_in21k_ft_in1k",
+}
+
+
+class IN1KViTBackbone(TimmViTBackbone):
+    def __init__(self, vision_backbone_id: str, image_resize_strategy: str, default_image_size: int = 224) -> None:
+        super().__init__(
+            vision_backbone_id,
+            IN1K_VISION_BACKBONES[vision_backbone_id],
+            image_resize_strategy,
+            default_image_size=default_image_size,
+        )

capvector-oft/prismatic/models/backbones/vision/siglip_vit.py

@@ -0,0 +1,24 @@
+"""
+siglip_vit.py
+"""
+
+from prismatic.models.backbones.vision.base_vision import TimmViTBackbone
+
+# Registry =>> Supported SigLIP Vision Backbones (from TIMM) =>> Note:: Using SigLIP w/ Patch = 14 (but SO400M Arch)
+SIGLIP_VISION_BACKBONES = {
+    "siglip-vit-b16-224px": "vit_base_patch16_siglip_224",
+    "siglip-vit-b16-256px": "vit_base_patch16_siglip_256",
+    "siglip-vit-b16-384px": "vit_base_patch16_siglip_384",
+    "siglip-vit-so400m": "vit_so400m_patch14_siglip_224",
+    "siglip-vit-so400m-384px": "vit_so400m_patch14_siglip_384",
+}
+
+
+class SigLIPViTBackbone(TimmViTBackbone):
+    def __init__(self, vision_backbone_id: str, image_resize_strategy: str, default_image_size: int = 224) -> None:
+        super().__init__(
+            vision_backbone_id,
+            SIGLIP_VISION_BACKBONES[vision_backbone_id],
+            image_resize_strategy,
+            default_image_size=default_image_size,
+        )

capvector-oft/prismatic/models/load.py

@@ -0,0 +1,226 @@
+"""
+load.py
+
+Entry point for loading pretrained VLMs for inference; exposes functions for listing available models (with canonical
+IDs, mappings to paper experiments, and short descriptions), as well as for loading models (from disk or HF Hub).
+"""
+
+import json
+import os
+from pathlib import Path
+from typing import List, Optional, Union
+
+from huggingface_hub import HfFileSystem, hf_hub_download
+
+from prismatic.conf import ModelConfig
+from prismatic.models.materialize import get_llm_backbone_and_tokenizer, get_vision_backbone_and_transform
+from prismatic.models.registry import GLOBAL_REGISTRY, MODEL_REGISTRY
+from prismatic.models.vlas import OpenVLA
+from prismatic.models.vlms import PrismaticVLM
+from prismatic.overwatch import initialize_overwatch
+from prismatic.vla.action_tokenizer import ActionTokenizer
+
+# Initialize Overwatch =>> Wraps `logging.Logger`
+overwatch = initialize_overwatch(__name__)
+
+
+# === HF Hub Repository ===
+HF_HUB_REPO = "TRI-ML/prismatic-vlms"
+VLA_HF_HUB_REPO = "openvla/openvla-dev"
+
+
+# === Available Models ===
+def available_models() -> List[str]:
+    return list(MODEL_REGISTRY.keys())
+
+
+def available_model_names() -> List[str]:
+    return list(GLOBAL_REGISTRY.items())
+
+
+def get_model_description(model_id_or_name: str) -> str:
+    if model_id_or_name not in GLOBAL_REGISTRY:
+        raise ValueError(f"Couldn't find `{model_id_or_name = }; check `prismatic.available_model_names()`")
+
+    # Print Description & Return
+    print(json.dumps(description := GLOBAL_REGISTRY[model_id_or_name]["description"], indent=2))
+
+    return description
+
+
+# === Load Pretrained Model ===
+def load(
+    model_id_or_path: Union[str, Path],
+    hf_token: Optional[str] = None,
+    cache_dir: Optional[Union[str, Path]] = None,
+    load_for_training: bool = False,
+) -> PrismaticVLM:
+    """Loads a pretrained PrismaticVLM from either local disk or the HuggingFace Hub."""
+    if os.path.isdir(model_id_or_path):
+        overwatch.info(f"Loading from local path `{(run_dir := Path(model_id_or_path))}`")
+
+        # Get paths for `config.json` and pretrained checkpoint
+        config_json, checkpoint_pt = run_dir / "config.json", run_dir / "checkpoints" / "latest-checkpoint.pt"
+        assert config_json.exists(), f"Missing `config.json` for `{run_dir = }`"
+        assert checkpoint_pt.exists(), f"Missing checkpoint for `{run_dir = }`"
+    else:
+        if model_id_or_path not in GLOBAL_REGISTRY:
+            raise ValueError(f"Couldn't find `{model_id_or_path = }; check `prismatic.available_model_names()`")
+
+        overwatch.info(f"Downloading `{(model_id := GLOBAL_REGISTRY[model_id_or_path]['model_id'])} from HF Hub")
+        with overwatch.local_zero_first():
+            config_json = hf_hub_download(repo_id=HF_HUB_REPO, filename=f"{model_id}/config.json", cache_dir=cache_dir)
+            checkpoint_pt = hf_hub_download(
+                repo_id=HF_HUB_REPO, filename=f"{model_id}/checkpoints/latest-checkpoint.pt", cache_dir=cache_dir
+            )
+
+    # Load Model Config from `config.json`
+    with open(config_json, "r") as f:
+        model_cfg = json.load(f)["model"]
+
+    # = Load Individual Components necessary for Instantiating a VLM =
+    #   =>> Print Minimal Config
+    overwatch.info(
+        f"Found Config =>> Loading & Freezing [bold blue]{model_cfg['model_id']}[/] with:\n"
+        f"             Vision Backbone =>> [bold]{model_cfg['vision_backbone_id']}[/]\n"
+        f"             LLM Backbone    =>> [bold]{model_cfg['llm_backbone_id']}[/]\n"
+        f"             Arch Specifier  =>> [bold]{model_cfg['arch_specifier']}[/]\n"
+        f"             Checkpoint Path =>> [underline]`{checkpoint_pt}`[/]"
+    )
+
+    # Load Vision Backbone
+    overwatch.info(f"Loading Vision Backbone [bold]{model_cfg['vision_backbone_id']}[/]")
+    vision_backbone, image_transform = get_vision_backbone_and_transform(
+        model_cfg["vision_backbone_id"],
+        model_cfg["image_resize_strategy"],
+    )
+
+    # Load LLM Backbone --> note `inference_mode = True` by default when calling `load()`
+    overwatch.info(f"Loading Pretrained LLM [bold]{model_cfg['llm_backbone_id']}[/] via HF Transformers")
+    llm_backbone, tokenizer = get_llm_backbone_and_tokenizer(
+        model_cfg["llm_backbone_id"],
+        llm_max_length=model_cfg.get("llm_max_length", 2048),
+        hf_token=hf_token,
+        inference_mode=not load_for_training,
+    )
+
+    # Load VLM using `from_pretrained` (clobbers HF syntax... eventually should reconcile)
+    overwatch.info(f"Loading VLM [bold blue]{model_cfg['model_id']}[/] from Checkpoint")
+    vlm = PrismaticVLM.from_pretrained(
+        checkpoint_pt,
+        model_cfg["model_id"],
+        vision_backbone,
+        llm_backbone,
+        arch_specifier=model_cfg["arch_specifier"],
+        freeze_weights=not load_for_training,
+    )
+
+    return vlm
+
+
+# === Load Pretrained VLA Model ===
+def load_vla(
+    model_id_or_path: Union[str, Path],
+    hf_token: Optional[str] = None,
+    cache_dir: Optional[Union[str, Path]] = None,
+    load_for_training: bool = False,
+    step_to_load: Optional[int] = None,
+    model_type: str = "pretrained",
+) -> OpenVLA:
+    """Loads a pretrained OpenVLA from either local disk or the HuggingFace Hub."""
+
+    # TODO (siddk, moojink) :: Unify semantics with `load()` above; right now, `load_vla()` assumes path points to
+    #   checkpoint `.pt` file, rather than the top-level run directory!
+    if os.path.isfile(model_id_or_path):
+        overwatch.info(f"Loading from local checkpoint path `{(checkpoint_pt := Path(model_id_or_path))}`")
+
+        # [Validate] Checkpoint Path should look like `.../<RUN_ID>/checkpoints/<CHECKPOINT_PATH>.pt`
+        assert (checkpoint_pt.suffix == ".pt") and (checkpoint_pt.parent.name == "checkpoints"), "Invalid checkpoint!"
+        run_dir = checkpoint_pt.parents[1]
+
+        # Get paths for `config.json`, `dataset_statistics.json` and pretrained checkpoint
+        config_json, dataset_statistics_json = run_dir / "config.json", run_dir / "dataset_statistics.json"
+        assert config_json.exists(), f"Missing `config.json` for `{run_dir = }`"
+        assert dataset_statistics_json.exists(), f"Missing `dataset_statistics.json` for `{run_dir = }`"
+
+    # Otherwise =>> try looking for a match on `model_id_or_path` on the HF Hub (`VLA_HF_HUB_REPO`)
+    else:
+        # Search HF Hub Repo via fsspec API
+        overwatch.info(f"Checking HF for `{(hf_path := str(Path(VLA_HF_HUB_REPO) / model_type / model_id_or_path))}`")
+        if not (tmpfs := HfFileSystem()).exists(hf_path):
+            raise ValueError(f"Couldn't find valid HF Hub Path `{hf_path = }`")
+
+        # Identify Checkpoint to Load (via `step_to_load`)
+        step_to_load = f"{step_to_load:06d}" if step_to_load is not None else None
+        valid_ckpts = tmpfs.glob(f"{hf_path}/checkpoints/step-{step_to_load if step_to_load is not None else ''}*.pt")
+        if (len(valid_ckpts) == 0) or (step_to_load is not None and len(valid_ckpts) != 1):
+            raise ValueError(f"Couldn't find a valid checkpoint to load from HF Hub Path `{hf_path}/checkpoints/")
+
+        # Call to `glob` will sort steps in ascending order (if `step_to_load` is None); just grab last element
+        target_ckpt = Path(valid_ckpts[-1]).name
+
+        overwatch.info(f"Downloading Model `{model_id_or_path}` Config & Checkpoint `{target_ckpt}`")
+        with overwatch.local_zero_first():
+            relpath = Path(model_type) / model_id_or_path
+            config_json = hf_hub_download(
+                repo_id=VLA_HF_HUB_REPO, filename=f"{(relpath / 'config.json')!s}", cache_dir=cache_dir
+            )
+            dataset_statistics_json = hf_hub_download(
+                repo_id=VLA_HF_HUB_REPO, filename=f"{(relpath / 'dataset_statistics.json')!s}", cache_dir=cache_dir
+            )
+            checkpoint_pt = hf_hub_download(
+                repo_id=VLA_HF_HUB_REPO, filename=f"{(relpath / 'checkpoints' / target_ckpt)!s}", cache_dir=cache_dir
+            )
+
+    # Load VLA Config (and corresponding base VLM `ModelConfig`) from `config.json`
+    with open(config_json, "r") as f:
+        vla_cfg = json.load(f)["vla"]
+        model_cfg = ModelConfig.get_choice_class(vla_cfg["base_vlm"])()
+
+    # Load Dataset Statistics for Action Denormalization
+    with open(dataset_statistics_json, "r") as f:
+        norm_stats = json.load(f)
+
+    # = Load Individual Components necessary for Instantiating a VLA (via base VLM components) =
+    #   =>> Print Minimal Config
+    overwatch.info(
+        f"Found Config =>> Loading & Freezing [bold blue]{model_cfg.model_id}[/] with:\n"
+        f"             Vision Backbone =>> [bold]{model_cfg.vision_backbone_id}[/]\n"
+        f"             LLM Backbone    =>> [bold]{model_cfg.llm_backbone_id}[/]\n"
+        f"             Arch Specifier  =>> [bold]{model_cfg.arch_specifier}[/]\n"
+        f"             Checkpoint Path =>> [underline]`{checkpoint_pt}`[/]"
+    )
+
+    # Load Vision Backbone
+    overwatch.info(f"Loading Vision Backbone [bold]{model_cfg.vision_backbone_id}[/]")
+    vision_backbone, image_transform = get_vision_backbone_and_transform(
+        model_cfg.vision_backbone_id,
+        model_cfg.image_resize_strategy,
+    )
+
+    # Load LLM Backbone --> note `inference_mode = True` by default when calling `load()`
+    overwatch.info(f"Loading Pretrained LLM [bold]{model_cfg.llm_backbone_id}[/] via HF Transformers")
+    llm_backbone, tokenizer = get_llm_backbone_and_tokenizer(
+        model_cfg.llm_backbone_id,
+        llm_max_length=model_cfg.llm_max_length,
+        hf_token=hf_token,
+        inference_mode=not load_for_training,
+    )
+
+    # Create Action Tokenizer
+    action_tokenizer = ActionTokenizer(llm_backbone.get_tokenizer())
+
+    # Load VLM using `from_pretrained` (clobbers HF syntax... eventually should reconcile)
+    overwatch.info(f"Loading VLA [bold blue]{model_cfg.model_id}[/] from Checkpoint")
+    vla = OpenVLA.from_pretrained(
+        checkpoint_pt,
+        model_cfg.model_id,
+        vision_backbone,
+        llm_backbone,
+        arch_specifier=model_cfg.arch_specifier,
+        freeze_weights=not load_for_training,
+        norm_stats=norm_stats,
+        action_tokenizer=action_tokenizer,
+    )
+
+    return vla

capvector-oft/prismatic/models/materialize.py

@@ -0,0 +1,130 @@
+"""
+materialize.py
+
+Factory class for initializing Vision Backbones, LLM Backbones, and VLMs from a set registry; provides and exports
+individual functions for clear control flow.
+"""
+
+from typing import Optional, Tuple
+
+from transformers import PreTrainedTokenizerBase
+
+from prismatic.models.backbones.llm import LLaMa2LLMBackbone, LLMBackbone, MistralLLMBackbone, PhiLLMBackbone
+from prismatic.models.backbones.vision import (
+    CLIPViTBackbone,
+    DinoCLIPViTBackbone,
+    DinoSigLIPViTBackbone,
+    DinoV2ViTBackbone,
+    ImageTransform,
+    IN1KViTBackbone,
+    SigLIPViTBackbone,
+    VisionBackbone,
+)
+from prismatic.models.vlms import PrismaticVLM
+
+# === Registries =>> Maps ID --> {cls(), kwargs} :: Different Registries for Vision Backbones, LLM Backbones, VLMs ===
+# fmt: off
+
+# === Vision Backbone Registry ===
+VISION_BACKBONES = {
+    # === 224px Backbones ===
+    "clip-vit-l": {"cls": CLIPViTBackbone, "kwargs": {"default_image_size": 224}},
+    "siglip-vit-so400m": {"cls": SigLIPViTBackbone, "kwargs": {"default_image_size": 224}},
+    "dinov2-vit-l": {"cls": DinoV2ViTBackbone, "kwargs": {"default_image_size": 224}},
+    "in1k-vit-l": {"cls": IN1KViTBackbone, "kwargs": {"default_image_size": 224}},
+    "dinosiglip-vit-so-224px": {"cls": DinoSigLIPViTBackbone, "kwargs": {"default_image_size": 224}},
+
+    # === Assorted CLIP Backbones ===
+    "clip-vit-b": {"cls": CLIPViTBackbone, "kwargs": {"default_image_size": 224}},
+    "clip-vit-l-336px": {"cls": CLIPViTBackbone, "kwargs": {"default_image_size": 336}},
+
+    # === Assorted SigLIP Backbones ===
+    "siglip-vit-b16-224px": {"cls": SigLIPViTBackbone, "kwargs": {"default_image_size": 224}},
+    "siglip-vit-b16-256px": {"cls": SigLIPViTBackbone, "kwargs": {"default_image_size": 256}},
+    "siglip-vit-b16-384px": {"cls": SigLIPViTBackbone, "kwargs": {"default_image_size": 384}},
+    "siglip-vit-so400m-384px": {"cls": SigLIPViTBackbone, "kwargs": {"default_image_size": 384}},
+
+    # === Fused Backbones ===
+    "dinoclip-vit-l-336px": {"cls": DinoCLIPViTBackbone, "kwargs": {"default_image_size": 336}},
+    "dinosiglip-vit-so-384px": {"cls": DinoSigLIPViTBackbone, "kwargs": {"default_image_size": 384}},
+}
+
+
+# === Language Model Registry ===
+LLM_BACKBONES = {
+    # === LLaMa-2 Pure (Non-Chat) Backbones ===
+    "llama2-7b-pure": {"cls": LLaMa2LLMBackbone, "kwargs": {}},
+    "llama2-13b-pure": {"cls": LLaMa2LLMBackbone, "kwargs": {}},
+
+    # === LLaMa-2 Chat Backbones ===
+    "llama2-7b-chat": {"cls": LLaMa2LLMBackbone, "kwargs": {}},
+    "llama2-13b-chat": {"cls": LLaMa2LLMBackbone, "kwargs": {}},
+
+    # === Vicuna-v1.5 Backbones ===
+    "vicuna-v15-7b": {"cls": LLaMa2LLMBackbone, "kwargs": {}},
+    "vicuna-v15-13b": {"cls": LLaMa2LLMBackbone, "kwargs": {}},
+
+    # === Mistral v0.1 Backbones ===
+    "mistral-v0.1-7b-pure": {"cls": MistralLLMBackbone, "kwargs": {}},
+    "mistral-v0.1-7b-instruct": {"cls": MistralLLMBackbone, "kwargs": {}},
+
+    # === Phi-2 Backbone ===
+    "phi-2-3b": {"cls": PhiLLMBackbone, "kwargs": {}},
+}
+
+# fmt: on
+
+
+def get_vision_backbone_and_transform(
+    vision_backbone_id: str, image_resize_strategy: str
+) -> Tuple[VisionBackbone, ImageTransform]:
+    """Instantiate a Vision Backbone, returning both the nn.Module wrapper class and default Image Transform."""
+    if vision_backbone_id in VISION_BACKBONES:
+        vision_cfg = VISION_BACKBONES[vision_backbone_id]
+        vision_backbone: VisionBackbone = vision_cfg["cls"](
+            vision_backbone_id, image_resize_strategy, **vision_cfg["kwargs"]
+        )
+        image_transform = vision_backbone.get_image_transform()
+        return vision_backbone, image_transform
+
+    else:
+        raise ValueError(f"Vision Backbone `{vision_backbone_id}` is not supported!")
+
+
+def get_llm_backbone_and_tokenizer(
+    llm_backbone_id: str,
+    llm_max_length: int = 2048,
+    hf_token: Optional[str] = None,
+    inference_mode: bool = False,
+) -> Tuple[LLMBackbone, PreTrainedTokenizerBase]:
+    if llm_backbone_id in LLM_BACKBONES:
+        llm_cfg = LLM_BACKBONES[llm_backbone_id]
+        llm_backbone: LLMBackbone = llm_cfg["cls"](
+            llm_backbone_id,
+            llm_max_length=llm_max_length,
+            hf_token=hf_token,
+            inference_mode=inference_mode,
+            **llm_cfg["kwargs"],
+        )
+        tokenizer = llm_backbone.get_tokenizer()
+        return llm_backbone, tokenizer
+
+    else:
+        raise ValueError(f"LLM Backbone `{llm_backbone_id}` is not supported!")
+
+
+def get_vlm(
+    model_id: str,
+    arch_specifier: str,
+    vision_backbone: VisionBackbone,
+    llm_backbone: LLMBackbone,
+    enable_mixed_precision_training: bool = True,
+) -> PrismaticVLM:
+    """Lightweight wrapper around initializing a VLM, mostly for future-proofing (if one wants to add a new VLM)."""
+    return PrismaticVLM(
+        model_id,
+        vision_backbone,
+        llm_backbone,
+        enable_mixed_precision_training=enable_mixed_precision_training,
+        arch_specifier=arch_specifier,
+    )

capvector-oft/prismatic/models/projectors.py

@@ -0,0 +1,49 @@
+"""Implementation of additional projectors for additional inputs to the VLA models."""
+import torch
+import torch.nn as nn
+
+
+class ProprioProjector(nn.Module):
+    """
+    Projects proprio state inputs into the LLM's embedding space.
+    """
+    def __init__(self, llm_dim: int, proprio_dim: int) -> None:
+        super().__init__()
+        self.llm_dim = llm_dim
+        self.proprio_dim = proprio_dim
+
+        self.fc1 = nn.Linear(self.proprio_dim, self.llm_dim, bias=True)
+        self.fc2 = nn.Linear(self.llm_dim, self.llm_dim, bias=True)
+        self.act_fn1 = nn.GELU()
+
+    def forward(self, proprio: torch.Tensor = None) -> torch.Tensor:
+        # proprio: (bsz, proprio_dim)
+        projected_features = self.fc1(proprio)
+        projected_features = self.act_fn1(projected_features)
+        projected_features = self.fc2(projected_features)
+        return projected_features
+
+
+class NoisyActionProjector(nn.Module):
+    """
+    [Diffusion] Projects noisy action inputs into the LLM's embedding space.
+
+    Note that since each action is tokenized into 7 tokens in OpenVLA (rather
+    than having 1 token per action), each noisy action token will have dimension 1
+    instead of 7.
+    """
+    def __init__(self, llm_dim: int) -> None:
+        super().__init__()
+        self.llm_dim = llm_dim
+        self.action_token_dim = 1
+
+        self.fc1 = nn.Linear(self.action_token_dim, self.llm_dim, bias=True)
+        self.fc2 = nn.Linear(self.llm_dim, self.llm_dim, bias=True)
+        self.act_fn1 = nn.GELU()
+
+    def forward(self, noisy_actions: torch.Tensor = None) -> torch.Tensor:
+        # noisy_actions: (bsz, num_action_tokens=chunk_len*action_dim, 1)
+        projected_features = self.fc1(noisy_actions)
+        projected_features = self.act_fn1(projected_features)
+        projected_features = self.fc2(projected_features)
+        return projected_features

capvector-oft/prismatic/models/registry.py

@@ -0,0 +1,691 @@
+"""
+registry.py
+
+Exhaustive list of pretrained VLMs (with full descriptions / links to corresponding names and sections of paper).
+"""
+
+# === Pretrained Model Registry ===
+# fmt: off
+MODEL_REGISTRY = {
+    # === LLaVa v1.5 Reproductions ===
+    "reproduction-llava-v15+7b": {
+        "model_id": "reproduction-llava-v15+7b",
+        "names": ["LLaVa v1.5 7B (Reproduction)"],
+        "description": {
+            "name": "LLaVa v1.5 7B (Reproduction)",
+            "optimization_procedure": "multi-stage",
+            "visual_representation": "CLIP ViT-L/14 @ 336px",
+            "image_processing": "Letterbox",
+            "language_model": "Vicuña v1.5 7B",
+            "datasets": ["LLaVa v1.5 Instruct"],
+            "train_epochs": 1,
+        }
+    },
+    "reproduction-llava-v15+13b": {
+        "model_id": "reproduction-llava-v15+13b",
+        "names": ["LLaVa v1.5 13B (Reproduction)"],
+        "description": {
+            "name": "LLaVa v1.5 13B (Reproduction)",
+            "optimization_procedure": "multi-stage",
+            "visual_representation": "CLIP ViT-L/14 @ 336px",
+            "image_processing": "Letterbox",
+            "language_model": "Vicuña v1.5 13B",
+            "datasets": ["LLaVa v1.5 Instruct"],
+            "train_epochs": 1,
+        }
+    },
+
+    # === Section 4.1 :: Optimization Procedure ===
+    "one-stage+7b": {
+        "model_id": "one-stage+7b",
+        "names": [
+            "One-Stage 7B",
+            "Single-Stage 7B",
+            "Frozen ViT (Single-Stage)",
+            "CLIP ViT-L 336px (Letterbox)",
+            "CLIP ViT-L 336px",
+            "Vicuña v1.5 7B",
+            "1 Epoch",
+            "Base",
+        ],
+        "description": {
+            "name": "Single-Stage 7B",
+            "optimization_procedure": "single-stage",
+            "visual_representation": "CLIP ViT-L/14 @ 336px",
+            "image_processing": "Letterbox",
+            "language_model": "Vicuña v1.5 7B",
+            "datasets": ["LLaVa v1.5 Instruct"],
+            "train_epochs": 1,
+        }
+    },
+    "one-stage+13b": {
+        "model_id": "one-stage+13b",
+        "names": [
+            "One-Stage 13B",
+            "Single-Stage 13B",
+            "Vicuña v1.5 13B",
+        ],
+        "description": {
+            "name": "Single-Stage 13B",
+            "optimization_procedure": "single-stage",
+            "visual_representation": "CLIP ViT-L/14 @ 336px",
+            "image_processing": "Letterbox",
+            "language_model": "Vicuña v1.5 13B",
+            "datasets": ["LLaVa v1.5 Instruct"],
+            "train_epochs": 1,
+        }
+    },
+
+    "full-ft-multi-stage+7b": {
+        "model_id": "full-ft-multi-stage+7b",
+        "names": ["Finetune ViT (Multi-Stage)"],
+        "description": {
+            "name": "Finetune ViT (Multi-Stage)",
+            "optimization_procedure": "multi-stage-full-finetune",
+            "visual_representation": "CLIP ViT-L/14 @ 336px",
+            "image_processing": "Letterbox",
+            "language_model": "Vicuña v1.5 7B",
+            "datasets": ["LLaVa v1.5 Instruct"],
+            "train_epochs": 1,
+        }
+    },
+    "full-ft-one-stage+7b": {
+        "model_id": "full-ft-one-stage+7b",
+        "names": ["Finetune ViT (Single-Stage)"],
+        "description": {
+            "name": "Finetune ViT (Single-Stage)",
+            "optimization_procedure": "single-stage-full-finetune",
+            "visual_representation": "CLIP ViT-L/14 @ 336px",
+            "image_processing": "Letterbox",
+            "language_model": "Vicuña v1.5 7B",
+            "datasets": ["LLaVa v1.5 Instruct"],
+            "train_epochs": 1,
+        }
+    },
+
+    # === Section 4.2 :: Image Processing and Visual Representations ===
+    "in1k-224px+7b": {
+        "model_id": "in1k-224px+7b",
+        "names": ["IN1K ViT-L 224px"],
+        "description": {
+            "name": "IN1K ViT-L 224px",
+            "optimization_procedure": "single-stage",
+            "visual_representation": "ImageNet-21K+1K ViT-L/16 @ 224px",
+            "image_processing": "Letterbox",
+            "language_model": "Vicuña v1.5 7B",
+            "datasets": ["LLaVa v1.5 Instruct"],
+            "train_epochs": 1,
+        },
+    },
+    "dinov2-224px+7b": {
+        "model_id": "dinov2-224px+7b",
+        "names": ["DINOv2 ViT-L 224px"],
+        "description": {
+            "name": "DINOv2 ViT-L 224px",
+            "optimization_procedure": "single-stage",
+            "visual_representation": "DINOv2 ViT-L/14 @ 224px",
+            "image_processing": "Letterbox",
+            "language_model": "Vicuña v1.5 7B",
+            "datasets": ["LLaVa v1.5 Instruct"],
+            "train_epochs": 1,
+        },
+    },
+    "clip-224px+7b": {
+        "model_id": "clip-224px+7b",
+        "names": ["CLIP ViT-L 224px"],
+        "description": {
+            "name": "CLIP ViT-L 224px",
+            "optimization_procedure": "single-stage",
+            "visual_representation": "CLIP ViT-L/14 @ 224px",
+            "image_processing": "Letterbox",
+            "language_model": "Vicuña v1.5 7B",
+            "datasets": ["LLaVa v1.5 Instruct"],
+            "train_epochs": 1,
+        },
+    },
+    "siglip-224px+7b": {
+        "model_id": "siglip-224px+7b",
+        "names": ["SigLIP ViT-SO 224px"],
+        "description": {
+            "name": "SigLIP ViT-SO 224px",
+            "optimization_procedure": "single-stage",
+            "visual_representation": "SigLIP ViT-SO/14 @ 224px",
+            "image_processing": "Letterbox",
+            "language_model": "Vicuña v1.5 7B",
+            "datasets": ["LLaVa v1.5 Instruct"],
+            "train_epochs": 1,
+        },
+    },
+
+    "clip-336px-resize-crop+7b": {
+        "model_id": "clip-336px-resize-crop+7b",
+        "names": ["CLIP ViT-L 336px (Resize Crop)"],
+        "description": {
+            "name": "CLIP ViT-L 336px (Resize Crop)",
+            "optimization_procedure": "single-stage",
+            "visual_representation": "CLIP ViT-L/14 @ 336px",
+            "image_processing": "Resize Crop",
+            "language_model": "Vicuña v1.5 7B",
+            "datasets": ["LLaVa v1.5 Instruct"],
+            "train_epochs": 1,
+        }
+    },
+    "clip-336px-resize-naive+7b": {
+        "model_id": "clip-336px-resize-naive+7b",
+        "names": ["CLIP ViT-L 336px (Naive Resize)", "CLIP 336px (Naive Resize)"],
+        "description": {
+            "name": "CLIP ViT-L 336px (Naive Resize)",
+            "optimization_procedure": "single-stage",
+            "visual_representation": "CLIP ViT-L/14 @ 336px",
+            "image_processing": "Naive Resize",
+            "language_model": "Vicuña v1.5 7B",
+            "datasets": ["LLaVa v1.5 Instruct"],
+            "train_epochs": 1,
+        }
+    },
+    "siglip-384px-letterbox+7b": {
+        "model_id": "siglip-384px-letterbox+7b",
+        "names": ["SigLIP ViT-SO 384px (Letterbox)", "SigLIP ViT-SO 384px"],
+        "description": {
+            "name": "SigLIP ViT-SO 384px (Letterbox)",
+            "optimization_procedure": "single-stage",
+            "visual_representation": "SigLIP ViT-SO/14 @ 384px",
+            "image_processing": "Letterbox",
+            "language_model": "Vicuña v1.5 7B",
+            "datasets": ["LLaVa v1.5 Instruct"],
+            "train_epochs": 1,
+        }
+    },
+    "siglip-384px-resize-crop+7b": {
+        "model_id": "siglip-384px-resize-crop+7b",
+        "names": ["SigLIP ViT-SO 384px (Resize Crop)"],
+        "description": {
+            "name": "SigLIP ViT-SO 384px (Resize Crop)",
+            "optimization_procedure": "single-stage",
+            "visual_representation": "SigLIP ViT-SO/14 @ 384px",
+            "image_processing": "Resize Crop",
+            "language_model": "Vicuña v1.5 7B",
+            "datasets": ["LLaVa v1.5 Instruct"],
+            "train_epochs": 1,
+        }
+    },
+    "siglip-384px-resize-naive+7b": {
+        "model_id": "siglip-384px-resize-naive+7b",
+        "names": ["SigLIP ViT-SO 384px (Naive Resize)", "SigLIP 384px (Naive Resize)"],
+        "description": {
+            "name": "SigLIP ViT-SO 384px (Naive Resize)",
+            "optimization_procedure": "single-stage",
+            "visual_representation": "SigLIP ViT-SO/14 @ 384px",
+            "image_processing": "Naive Resize",
+            "language_model": "Vicuña v1.5 7B",
+            "datasets": ["LLaVa v1.5 Instruct"],
+            "train_epochs": 1,
+        }
+    },
+
+    "dinoclip-336px-letterbox+7b": {
+        "model_id": "dinoclip-336px-letterbox+7b",
+        "names": ["DINOv2 + CLIP 336px (Letterbox)"],
+        "description": {
+            "name": "DINOv2 + CLIP 336px (Letterbox)",
+            "optimization_procedure": "single-stage",
+            "visual_representation": "DINOv2 ViT-L/14 + CLIP ViT-L/14 @ 336px",
+            "image_processing": "Letterbox",
+            "language_model": "Vicuña v1.5 7B",
+            "datasets": ["LLaVa v1.5 Instruct"],
+            "train_epochs": 1,
+        }
+    },
+    "dinoclip-336px-resize-naive+7b": {
+        "model_id": "dinoclip-336px-resize-naive+7b",
+        "names": ["DINOv2 + CLIP 336px (Naive Resize)"],
+        "description": {
+            "name": "DINOv2 + CLIP 336px (Naive Resize)",
+            "optimization_procedure": "single-stage",
+            "visual_representation": "DINOv2 ViT-L/14 + CLIP ViT-L/14 @ 336px",
+            "image_processing": "Naive Resize",
+            "language_model": "Vicuña v1.5 7B",
+            "datasets": ["LLaVa v1.5 Instruct"],
+            "train_epochs": 1,
+        }
+    },
+    "dinosiglip-384px-letterbox+7b": {
+        "model_id": "dinosiglip-384px-letterbox+7b",
+        "names": ["DINOv2 + SigLIP 384px (Letterbox)"],
+        "description": {
+            "name": "DINOv2 + SigLIP 384px (Letterbox)",
+            "optimization_procedure": "single-stage",
+            "visual_representation": "DINOv2 ViT-L/14 + SigLIP ViT-L/14 @ 384px",
+            "image_processing": "Letterbox",
+            "language_model": "Vicuña v1.5 7B",
+            "datasets": ["LLaVa v1.5 Instruct"],
+            "train_epochs": 1,
+        }
+    },
+    "dinosiglip-384px-resize-naive+7b": {
+        "model_id": "dinosiglip-384px-resize-naive+7b",
+        "names": ["DINOv2 + SigLIP 384px (Naive Resize)"],
+        "description": {
+            "name": "DINOv2 + SigLIP 384px (Naive Resize)",
+            "optimization_procedure": "single-stage",
+            "visual_representation": "DINOv2 ViT-L/14 + SigLIP ViT-L/14 @ 384px",
+            "image_processing": "Naive Resize",
+            "language_model": "Vicuña v1.5 7B",
+            "datasets": ["LLaVa v1.5 Instruct"],
+            "train_epochs": 1,
+        }
+    },
+
+    # === Section 4.3 :: Language Models ===
+    "llama2+7b": {
+        "model_id": "llama2+7b",
+        "names": ["Llama-2 7B"],
+        "description": {
+            "name": "Llama-2 7B",
+            "optimization_procedure": "single-stage",
+            "visual_representation": "CLIP ViT-L/14 @ 336px",
+            "image_processing": "Letterbox",
+            "language_model": "Llama-2 7B",
+            "datasets": ["LLaVa v1.5 Instruct"],
+            "train_epochs": 1,
+        },
+    },
+    "llama2+13b": {
+        "model_id": "llama2+13b",
+        "names": ["Llama-2 13B"],
+        "description": {
+            "name": "Llama-2 13B",
+            "optimization_procedure": "single-stage",
+            "visual_representation": "CLIP ViT-L/14 @ 336px",
+            "image_processing": "Letterbox",
+            "language_model": "Llama-2 13B",
+            "datasets": ["LLaVa v1.5 Instruct"],
+            "train_epochs": 1,
+        },
+    },
+
+    "vicuna-no-cotraining+7b": {
+        "model_id": "vicuna-no-cotraining+7b",
+        "names": ["Vicuña v1.5 7B (No Co-training)"],
+        "description": {
+            "name": "Vicuña v1.5 7B (No Co-training)",
+            "optimization_procedure": "single-stage",
+            "visual_representation": "CLIP ViT-L/14 @ 336px",
+            "image_processing": "Letterbox",
+            "language_model": "Vicuña v1.5 7B",
+            "datasets": ["LLaVa v1.5 Multimodal-Only"],
+            "train_epochs": 1,
+        },
+    },
+    "llama2-no-cotraining+7b": {
+        "model_id": "llama2-no-cotraining+7b",
+        "names": ["Llama-2 7B (No Co-training)"],
+        "description": {
+            "name": "Llama-2 7B (No Co-training)",
+            "optimization_procedure": "single-stage",
+            "visual_representation": "CLIP ViT-L/14 @ 336px",
+            "image_processing": "Letterbox",
+            "language_model": "Llama-2 7B",
+            "datasets": ["LLaVa v1.5 Multimodal-Only"],
+            "train_epochs": 1,
+        },
+    },
+
+    # === Section 4.4 :: Scaling Properties ===
+    "train-1.25-epochs+7b": {
+        "model_id": "train-1.25-epochs+7b",
+        "names": ["1.25 Epochs"],
+        "description": {
+            "name": "1.25 Epochs",
+            "optimization_procedure": "single-stage",
+            "visual_representation": "CLIP ViT-L/14 @ 336px",
+            "image_processing": "Letterbox",
+            "language_model": "Vicuña v1.5 7B",
+            "datasets": ["LLaVa v1.5 Instruct"],
+            "train_epochs": 1.25,
+        }
+    },
+    "train-1.5-epochs+7b": {
+        "model_id": "train-1.5-epochs+7b",
+        "names": ["1.5 Epochs"],
+        "description": {
+            "name": "1.5 Epochs",
+            "optimization_procedure": "single-stage",
+            "visual_representation": "CLIP ViT-L/14 @ 336px",
+            "image_processing": "Letterbox",
+            "language_model": "Vicuña v1.5 7B",
+            "datasets": ["LLaVa v1.5 Instruct"],
+            "train_epochs": 1.5,
+        }
+    },
+    "train-2-epochs+7b": {
+        "model_id": "train-2-epochs+7b",
+        "names": ["2 Epochs"],
+        "description": {
+            "name": "2 Epochs",
+            "optimization_procedure": "single-stage",
+            "visual_representation": "CLIP ViT-L/14 @ 336px",
+            "image_processing": "Letterbox",
+            "language_model": "Vicuña v1.5 7B",
+            "datasets": ["LLaVa v1.5 Instruct"],
+            "train_epochs": 2,
+        }
+    },
+    "train-3-epochs+7b": {
+        "model_id": "train-3-epochs+7b",
+        "names": ["3 Epochs"],
+        "description": {
+            "name": "3 Epochs",
+            "optimization_procedure": "single-stage",
+            "visual_representation": "CLIP ViT-L/14 @ 336px",
+            "image_processing": "Letterbox",
+            "language_model": "Vicuña v1.5 7B",
+            "datasets": ["LLaVa v1.5 Instruct"],
+            "train_epochs": 3,
+        }
+    },
+
+    "llava-lvis4v+7b": {
+        "model_id": "llava-lvis4v+7b",
+        "names": ["Base + LVIS-4V"],
+        "description": {
+            "name": "Base + LVIS-4V",
+            "optimization_procedure": "single-stage",
+            "visual_representation": "CLIP ViT-L/14 @ 336px",
+            "image_processing": "Letterbox",
+            "language_model": "Vicuña v1.5 7B",
+            "datasets": ["LLaVa v1.5 Instruct", "LVIS-Instruct-4V"],
+            "train_epochs": 1,
+        }
+    },
+    "llava-lrv+7b": {
+        "model_id": "llava-lrv+7b",
+        "names": ["Base + LRV"],
+        "description": {
+            "name": "Base + LRV",
+            "optimization_procedure": "single-stage",
+            "visual_representation": "CLIP ViT-L/14 @ 336px",
+            "image_processing": "Letterbox",
+            "language_model": "Vicuña v1.5 7B",
+            "datasets": ["LLaVa v1.5 Instruct", "LRV-Instruct"],
+            "train_epochs": 1,
+        }
+    },
+    "llava-lvis4v-lrv+7b": {
+        "model_id": "llava-lvis4v-lrv+7b",
+        "names": ["Base + LVIS-4V + LRV"],
+        "description": {
+            "name": "Base + LVIS-4V + LRV",
+            "optimization_procedure": "single-stage",
+            "visual_representation": "CLIP ViT-L/14 @ 336px",
+            "image_processing": "Letterbox",
+            "language_model": "Vicuña v1.5 7B",
+            "datasets": ["LLaVa v1.5 Instruct", "LVIS-Instruct-4V", "LRV-Instruct"],
+            "train_epochs": 1,
+        }
+    },
+
+    # ===
+
+    # === CLIP Prism Models ===
+    "prism-clip-controlled+7b": {
+        "model_id": "prism-clip-controlled+7b",
+        "names": ["Prism-CLIP 7B (Controlled)"],
+        "description": {
+            "name": "CLIP Prism 7B (Controlled)",
+            "optimization_procedure": "single-stage",
+            "visual_representation": "CLIP ViT-L/14 @ 336px",
+            "image_processing": "Naive Resize",
+            "language_model": "Llama-2 7B",
+            "datasets": ["LLaVa v1.5 Instruct"],
+            "train_epochs": 1,
+        }
+    },
+    "prism-clip-controlled+13b": {
+        "model_id": "prism-clip-controlled+13b",
+        "names": ["Prism-CLIP 13B (Controlled)"],
+        "description": {
+            "name": "CLIP Prism 13B (Controlled)",
+            "optimization_procedure": "single-stage",
+            "visual_representation": "CLIP ViT-L/14 @ 336px",
+            "image_processing": "Naive Resize",
+            "language_model": "Llama-2 13B",
+            "datasets": ["LLaVa v1.5 Instruct"],
+            "train_epochs": 1,
+        }
+    },
+    "prism-clip+7b": {
+        "model_id": "prism-clip+7b",
+        "names": ["Prism-CLIP 7B"],
+        "description": {
+            "name": "CLIP Prism 7B",
+            "optimization_procedure": "single-stage",
+            "visual_representation": "CLIP ViT-L/14 @ 336px",
+            "image_processing": "Naive Resize",
+            "language_model": "Llama-2 7B",
+            "datasets": ["LLaVa v1.5 Instruct", "LVIS-Instruct-4V", "LRV-Instruct"],
+            "train_epochs": 2,
+        },
+    },
+    "prism-clip+13b": {
+        "model_id": "prism-clip+13b",
+        "names": ["Prism-CLIP 13B"],
+        "description": {
+            "name": "CLIP Prism 13B",
+            "optimization_procedure": "single-stage",
+            "visual_representation": "CLIP ViT-L/14 @ 336px",
+            "image_processing": "Naive Resize",
+            "language_model": "Llama-2 13B",
+            "datasets": ["LLaVa v1.5 Instruct", "LVIS-Instruct-4V", "LRV-Instruct"],
+            "train_epochs": 2,
+        },
+    },
+
+    # === SigLIP Prism Models ==
+    "prism-siglip-controlled+7b": {
+        "model_id": "prism-siglip-controlled+7b",
+        "names": ["Prism-SigLIP 7B (Controlled)"],
+        "description": {
+            "name": "SigLIP Prism 7B (Controlled)",
+            "optimization_procedure": "single-stage",
+            "visual_representation": "SigLIP ViT-SO/14 @ 384px",
+            "image_processing": "Naive Resize",
+            "language_model": "Llama-2 7B",
+            "datasets": ["LLaVa v1.5 Instruct"],
+            "train_epochs": 1,
+        }
+    },
+    "prism-siglip-controlled+13b": {
+        "model_id": "prism-siglip-controlled+7b",
+        "names": ["Prism-SigLIP 13B (Controlled)"],
+        "description": {
+            "name": "SigLIP Prism 13B (Controlled)",
+            "optimization_procedure": "single-stage",
+            "visual_representation": "SigLIP ViT-SO/14 @ 384px",
+            "image_processing": "Naive Resize",
+            "language_model": "Llama-2 13B",
+            "datasets": ["LLaVa v1.5 Instruct"],
+            "train_epochs": 1,
+        }
+    },
+    "prism-siglip+7b": {
+        "model_id": "prism-siglip+7b",
+        "names": ["Prism-SigLIP 7B"],
+        "description": {
+            "name": "SigLIP Prism 7B",
+            "optimization_procedure": "single-stage",
+            "visual_representation": "SigLIP ViT-SO/14 @ 384px",
+            "image_processing": "Naive Resize",
+            "language_model": "Llama-2 7B",
+            "datasets": ["LLaVa v1.5 Instruct", "LVIS-Instruct-4V", "LRV-Instruct"],
+            "train_epochs": 2,
+        }
+    },
+    "prism-siglip+13b": {
+        "model_id": "prism-siglip+13b",
+        "names": ["Prism-SigLIP 13B"],
+        "description": {
+            "name": "SigLIP Prism 13B",
+            "optimization_procedure": "single-stage",
+            "visual_representation": "SigLIP ViT-SO/14 @ 384px",
+            "image_processing": "Naive Resize",
+            "language_model": "Llama-2 13B",
+            "datasets": ["LLaVa v1.5 Instruct", "LVIS-Instruct-4V", "LRV-Instruct"],
+            "train_epochs": 2,
+        }
+    },
+
+    # === DINOSigLIP Prism Models ===
+    "prism-dinosiglip-controlled+7b": {
+        "model_id": "prism-dinosiglip-controlled+7b",
+        "names": ["Prism-DINOSigLIP 7B (Controlled)", "Prism 7B (Controlled)"],
+        "description": {
+            "name": "DINOSigLIP Prism 7B (Controlled)",
+            "optimization_procedure": "single-stage",
+            "visual_representation": "DINOv2 ViT-L/14 + SigLIP ViT-SO/14 @ 384px",
+            "image_processing": "Naive Resize",
+            "language_model": "Llama-2 7B",
+            "datasets": ["LLaVa v1.5 Instruct"],
+            "train_epochs": 1,
+        }
+    },
+    "prism-dinosiglip-controlled+13b": {
+        "model_id": "prism-dinosiglip-controlled+13b",
+        "names": ["Prism-DINOSigLIP 13B (Controlled)", "Prism 13B (Controlled)"],
+        "description": {
+            "name": "DINOSigLIP Prism 13B (Controlled)",
+            "optimization_procedure": "single-stage",
+            "visual_representation": "DINOv2 ViT-L/14 + SigLIP ViT-SO/14 @ 384px",
+            "image_processing": "Naive Resize",
+            "language_model": "Llama-2 13B",
+            "datasets": ["LLaVa v1.5 Instruct"],
+            "train_epochs": 1,
+        }
+    },
+    "prism-dinosiglip+7b": {
+        "model_id": "prism-dinosiglip+7b",
+        "names": ["Prism-DINOSigLIP 7B"],
+        "description": {
+            "name": "DINOSigLIP Prism 7B",
+            "optimization_procedure": "single-stage",
+            "visual_representation": "DINOv2 ViT-L/14 + SigLIP ViT-SO/14 @ 384px",
+            "image_processing": "Naive Resize",
+            "language_model": "Llama-2 7B",
+            "datasets": ["LLaVa v1.5 Instruct", "LVIS-Instruct-4V", "LRV-Instruct"],
+            "train_epochs": 2,
+        },
+    },
+    "prism-dinosiglip+13b": {
+        "model_id": "prism-dinosiglip+13b",
+        "names": ["Prism-DINOSigLIP 13B"],
+        "description": {
+            "name": "DINOSigLIP Prism 13B",
+            "optimization_procedure": "single-stage",
+            "visual_representation": "DINOv2 ViT-L/14 + SigLIP ViT-SO/14 @ 384px",
+            "image_processing": "Naive Resize",
+            "language_model": "Llama-2 13B",
+            "datasets": ["LLaVa v1.5 Instruct", "LVIS-Instruct-4V", "LRV-Instruct"],
+            "train_epochs": 2,
+        },
+    },
+
+    # === DINOSigLIP 224px Prism Models ===
+    "prism-dinosiglip-224px-controlled+7b": {
+        "model_id": "prism-dinosiglip-224px-controlled+7b",
+        "names": ["Prism-DINOSigLIP 224px 7B (Controlled)"],
+        "description": {
+            "name": "DINOSigLIP 224px 7B (Controlled)",
+            "optimization_procedure": "single-stage",
+            "visual_representation": "DINOv2 ViT-L/14 + SigLIP ViT-SO 14 @ 224px",
+            "image_processing": "Naive Resize",
+            "language_model": "Llama-2 7B",
+            "datasets": ["LLaVa v1.5 Instruct"],
+            "train_epochs": 1,
+        }
+    },
+    "prism-dinosiglip-224px+7b": {
+        "model_id": "prism-dinosiglip-224px+7b",
+        "names": ["Prism-DINOSigLIP 224px 7B"],
+        "description": {
+            "name": "DINOSigLIP 224px 7B",
+            "optimization_procedure": "single-stage",
+            "visual_representation": "DINOv2 ViT-L/14 + SigLIP ViT-SO 14 @ 224px",
+            "image_processing": "Naive Resize",
+            "language_model": "Llama-2 7B",
+            "datasets": ["LLaVa v1.5 Instruct", "LVIS-Instruct-4V", "LRV-Instruct"],
+            "train_epochs": 2,
+        }
+    },
+
+    # === Additional LLM Backbones ===
+    "llama2-chat+7b": {
+        "model_id": "llama2-chat+7b",
+        "names": ["Llama-2 Chat 7B"],
+        "description": {
+            "name": "Llama-2 Chat 7B",
+            "optimization_procedure": "single-stage",
+            "visual_representation": "CLIP ViT-L/14 @ 336px",
+            "image_processing": "Letterbox",
+            "language_model": "Llama-2 Chat 7B",
+            "datasets": ["LLaVa v1.5 Instruct"],
+            "train_epochs": 1,
+        }
+    },
+    "llama2-chat+13b": {
+        "model_id": "llama2-chat+13b",
+        "names": ["Llama-2 Chat 13B"],
+        "description": {
+            "name": "Llama-2 Chat 13B",
+            "optimization_procedure": "single-stage",
+            "visual_representation": "CLIP ViT-L/14 @ 336px",
+            "image_processing": "Letterbox",
+            "language_model": "Llama-2 Chat 13B",
+            "datasets": ["LLaVa v1.5 Instruct"],
+            "train_epochs": 1,
+        }
+    },
+    "mistral-v0.1+7b": {
+        "model_id": "mistral-v0.1+7b",
+        "names": ["Mistral v0.1 7B"],
+        "description": {
+            "name": "Mistral v0.1 7B",
+            "optimization_procedure": "single-stage",
+            "visual_representation": "CLIP ViT-L/14 @ 336px",
+            "image_processing": "Letterbox",
+            "language_model": "Mistral v0.1 7B",
+            "datasets": ["LLaVa v1.5 Instruct"],
+            "train_epochs": 1,
+        }
+    },
+    "mistral-instruct-v0.1+7b": {
+        "model_id": "mistral-instruct-v0.1+7b",
+        "names": ["Mistral Instruct v0.1 7B"],
+        "description": {
+            "name": "Mistral Instruct v0.1 7B",
+            "optimization_procedure": "single-stage",
+            "visual_representation": "CLIP ViT-L/14 @ 336px",
+            "image_processing": "Letterbox",
+            "language_model": "Mistral Instruct v0.1 7B",
+            "datasets": ["LLaVa v1.5 Instruct"],
+            "train_epochs": 1,
+        }
+    },
+    "phi-2+3b": {
+        "model_id": "phi-2+3b",
+        "names": ["Phi-2 3B"],
+        "description": {
+            "name": "Phi-2 3B",
+            "optimization_procedure": "single-stage",
+            "visual_representation": "CLIP ViT-L/14 @ 336px",
+            "image_processing": "Letterbox",
+            "language_model": "Phi-2 3B",
+            "datasets": ["LLaVa v1.5 Instruct"],
+            "train_epochs": 1,
+        }
+    },
+}
+
+# Build Global Registry (Model ID, Name) -> Metadata
+GLOBAL_REGISTRY = {name: v for k, v in MODEL_REGISTRY.items() for name in [k] + v["names"]}
+
+# fmt: on

capvector-oft/prismatic/models/vlas/__init__.py

@@ -0,0 +1 @@
+from .openvla import OpenVLA

capvector-oft/prismatic/models/vlas/openvla.py

@@ -0,0 +1,131 @@
+"""
+openvla.py
+
+PyTorch Module defining OpenVLA as a lightweight wrapper around a PrismaticVLM; defines custom logic around
+discretizing actions with the ActionTokenizer.
+"""
+
+from typing import Dict, List, Optional
+
+import numpy as np
+import torch
+from PIL import Image
+from transformers import LlamaTokenizerFast
+
+from prismatic.models.vlms.prismatic import PrismaticVLM
+from prismatic.overwatch import initialize_overwatch
+from prismatic.vla.action_tokenizer import ActionTokenizer
+
+# Initialize Overwatch =>> Wraps `logging.Logger`
+overwatch = initialize_overwatch(__name__)
+
+
+class OpenVLA(PrismaticVLM):
+    def __init__(
+        self,
+        *args,
+        norm_stats: Dict[str, Dict[str, Dict[str, Dict[str, List[float]]]]],
+        action_tokenizer: ActionTokenizer,
+        **kwargs,
+    ) -> None:
+        super().__init__(*args, **kwargs)
+        self.norm_stats = norm_stats
+        self.action_tokenizer = action_tokenizer
+
+    @torch.inference_mode()
+    def predict_action(
+        self, image: Image, instruction: str, unnorm_key: Optional[str] = None, **kwargs: str
+    ) -> np.ndarray:
+        """
+        Core function for VLA inference; maps input image and task instruction to continuous action (de-tokenizes).
+
+        @param image: PIL Image as [height, width, 3]
+        @param instruction: Task instruction string
+        @param unnorm_key: Optional dataset name for retrieving un-normalizing statistics; if None, checks that model
+                           was trained only on a single dataset, and retrieves those statistics.
+
+        @return Unnormalized (continuous) action vector --> end-effector deltas.
+        """
+        image_transform, tokenizer = self.vision_backbone.image_transform, self.llm_backbone.tokenizer
+
+        # Build VLA Prompt
+        prompt_builder = self.get_prompt_builder()
+        prompt_builder.add_turn(role="human", message=f"What action should the robot take to {instruction.lower()}?")
+        prompt_text = prompt_builder.get_prompt()
+
+        # Prepare Inputs
+        input_ids = tokenizer(prompt_text, truncation=True, return_tensors="pt").input_ids.to(self.device)
+        if isinstance(tokenizer, LlamaTokenizerFast):
+            # If the special empty token ('') does not already appear after the colon (':') token in the prompt
+            # (after "OUT:" or "ASSISTANT:"), insert it to match the inputs seen at training time
+            if not torch.all(input_ids[:, -1] == 29871):
+                input_ids = torch.cat(
+                    (input_ids, torch.unsqueeze(torch.Tensor([29871]).long(), dim=0).to(input_ids.device)), dim=1
+                )
+        else:
+            raise ValueError(f"Unsupported `tokenizer` type = {type(tokenizer)}")
+
+        # Preprocess Image
+        pixel_values = image_transform(image)
+        if isinstance(pixel_values, torch.Tensor):
+            pixel_values = pixel_values[None, ...].to(self.device)
+        elif isinstance(pixel_values, dict):
+            pixel_values = {k: v[None, ...].to(self.device) for k, v in pixel_values.items()}
+        else:
+            raise ValueError(f"Unsupported `pixel_values` type = {type(pixel_values)}")
+
+        # Invoke super().generate --> taps into `GenerationMixin` which (redirects) to `forward()`
+        autocast_dtype = self.llm_backbone.half_precision_dtype
+        with torch.autocast("cuda", dtype=autocast_dtype, enabled=self.enable_mixed_precision_training):
+            # fmt: off
+            generated_ids = super(PrismaticVLM, self).generate(
+                input_ids=input_ids,                            # Shape: [1, seq]
+                pixel_values=pixel_values,                      # Shape: [1, 3, res, res] or Dict[str, ...]
+                max_new_tokens=self.get_action_dim(unnorm_key),
+                **kwargs
+            )
+            # fmt: on
+
+        # Extract predicted action tokens and translate into (normalized) continuous actions
+        predicted_action_token_ids = generated_ids[0, -self.get_action_dim(unnorm_key) :]
+        normalized_actions = self.action_tokenizer.decode_token_ids_to_actions(predicted_action_token_ids.cpu().numpy())
+
+        # Un-normalize Actions
+        action_norm_stats = self.get_action_stats(unnorm_key)
+        mask = action_norm_stats.get("mask", np.ones_like(action_norm_stats["q01"], dtype=bool))
+        action_high, action_low = np.array(action_norm_stats["q99"]), np.array(action_norm_stats["q01"])
+        actions = np.where(
+            mask,
+            0.5 * (normalized_actions + 1) * (action_high - action_low) + action_low,
+            normalized_actions,
+        )
+
+        return actions
+
+    @staticmethod
+    def _check_unnorm_key(norm_stats: Dict, unnorm_key: str) -> str:
+        if unnorm_key is None:
+            assert len(norm_stats) == 1, (
+                f"Your model was trained on more than one dataset, please pass a `unnorm_key` from the following "
+                f"options to choose the statistics used for un-normalizing actions: {norm_stats.keys()}"
+            )
+            unnorm_key = next(iter(norm_stats.keys()))
+
+        # Error Handling
+        assert (
+            unnorm_key in norm_stats
+        ), f"The `unnorm_key` you chose is not in the set of available statistics; choose from: {norm_stats.keys()}"
+
+        return unnorm_key
+
+    def get_action_dim(self, unnorm_key: Optional[str] = None) -> int:
+        """Dimensionality of the policy's action space."""
+        unnorm_key = self._check_unnorm_key(self.norm_stats, unnorm_key)
+
+        return len(self.norm_stats[unnorm_key]["action"]["q01"])
+
+    def get_action_stats(self, unnorm_key: Optional[str] = None) -> Dict:
+        """Dimensionality of the policy's action space."""
+        unnorm_key = self._check_unnorm_key(self.norm_stats, unnorm_key)
+
+        return self.norm_stats[unnorm_key]["action"]

capvector-oft/prismatic/models/vlms/__init__.py

@@ -0,0 +1 @@
+from .prismatic import PrismaticVLM

capvector-oft/prismatic/models/vlms/base_vlm.py

@@ -0,0 +1,108 @@
+"""
+base_vlm.py
+
+Abstract class definition of a Vision-Language Model (VLM), with full annotations of class methods, utility functions,
+and initialization logic. This is mostly to future-proof the codebase; while all our experiments instantiate
+from PrismaticVLM, theoretically, this base class should be general enough to cover almost all models (e.g., IDEFICS,
+PALI, Fuyu) in the future.
+
+We use Abstract base classes *sparingly* -- mostly as a way to encapsulate any redundant logic or nested inheritance
+(e.g., dependence on nn.Module, HF PretrainedModel, etc.). For other abstract objects (e.g., Tokenizers/Transforms),
+prefer Protocol definitions instead.
+"""
+
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from pathlib import Path
+from typing import Callable, List, Optional
+
+import torch
+import torch.nn as nn
+from transformers import GenerationMixin, PretrainedConfig
+from transformers.modeling_outputs import CausalLMOutputWithPast
+
+from prismatic.models.backbones.llm import LLMBackbone
+from prismatic.models.backbones.llm.prompting import PromptBuilder
+from prismatic.models.backbones.vision import VisionBackbone
+
+
+# === Abstract Base Class for arbitrary Vision-Language Models ===
+class VLM(nn.Module, GenerationMixin, ABC):
+    def __init__(
+        self,
+        model_family: str,
+        model_id: str,
+        vision_backbone: VisionBackbone,
+        llm_backbone: LLMBackbone,
+        enable_mixed_precision_training: bool = True,
+    ) -> None:
+        super().__init__()
+        self.model_family, self.model_id = model_family, model_id
+        self.vision_backbone, self.llm_backbone = vision_backbone, llm_backbone
+        self.enable_mixed_precision_training = enable_mixed_precision_training
+
+        # Instance Attributes for a generic VLM
+        self.all_module_keys, self.trainable_module_keys = None, None
+
+        # === GenerationMixin Expected Attributes =>> *DO NOT MODIFY* ===
+        self.generation_config = self.llm_backbone.llm.generation_config
+        self.main_input_name = "input_ids"
+
+    @property
+    def device(self) -> torch.device:
+        """Borrowed from `transformers.modeling_utils.py` -- checks parameter device; assumes model on *ONE* device!"""
+        return next(self.parameters()).device
+
+    @classmethod
+    @abstractmethod
+    def from_pretrained(
+        cls,
+        pretrained_checkpoint: Path,
+        model_family: str,
+        model_id: str,
+        vision_backbone: VisionBackbone,
+        llm_backbone: LLMBackbone,
+        **kwargs: str,
+    ) -> VLM: ...
+
+    @abstractmethod
+    def get_prompt_builder(self, system_prompt: Optional[str] = None) -> PromptBuilder: ...
+
+    @abstractmethod
+    def freeze_backbones(self, stage: str) -> None: ...
+
+    @abstractmethod
+    def load_from_checkpoint(self, stage: str, run_dir: Path, pretrained_checkpoint: Optional[Path] = None) -> None: ...
+
+    @abstractmethod
+    def get_fsdp_wrapping_policy(self) -> Callable: ...
+
+    @abstractmethod
+    def forward(
+        self,
+        input_ids: Optional[torch.LongTensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        pixel_values: Optional[torch.FloatTensor] = None,
+        labels: Optional[torch.LongTensor] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        past_key_values: Optional[List[torch.FloatTensor]] = None,
+        use_cache: Optional[bool] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+        multimodal_indices: Optional[torch.LongTensor] = None,
+    ) -> CausalLMOutputWithPast: ...
+
+    # === GenerationMixin Expected Properties & Methods (DO NOT MODIFY) ===
+    @staticmethod
+    def can_generate() -> bool:
+        return True
+
+    @property
+    def config(self) -> PretrainedConfig:
+        return self.llm_backbone.llm.config
+
+    # => Beam Search Utility
+    def _reorder_cache(self, past_key_values, beam_idx):
+        return self.llm_backbone.llm._reorder_cache(past_key_values, beam_idx)

capvector-oft/prismatic/models/vlms/prismatic.py

@@ -0,0 +1,621 @@
+"""
+prismatic.py
+
+PyTorch Module defining a PrismaticVLM, our general interface for defining the various different VLMs in our work.
+
+Notes:
+    - For now, we don't subclass `transformers.PretrainedModel` (or CausalLM). Instead, we assume a very limited subset
+      of the {Model}ForCausalLM API that enables dispatch to the underlying LLM's `generate` utilities (feeding inputs
+      through our custom projection shim).
+"""
+
+from __future__ import annotations
+
+from functools import partial
+from pathlib import Path
+from typing import Callable, Dict, List, Optional, Type, Union
+
+import torch
+from PIL import Image
+from torch.distributed.fsdp.wrap import _module_wrap_policy, _or_policy
+from transformers.modeling_outputs import CausalLMOutputWithPast
+
+from prismatic.models.backbones.llm import LLMBackbone
+from prismatic.models.backbones.llm.prompting import PromptBuilder
+from prismatic.models.backbones.vision import VisionBackbone
+from prismatic.models.vlms.base_vlm import VLM
+from prismatic.overwatch import initialize_overwatch
+from prismatic.util.nn_utils import FusedMLPProjector, LinearProjector, MLPProjector
+
+# Initialize Overwatch =>> Wraps `logging.Logger`
+overwatch = initialize_overwatch(__name__)
+
+
+# HuggingFace Default / LLaMa-2 IGNORE_INDEX (for labels)
+IGNORE_INDEX = -100
+
+
+class PrismaticVLM(VLM):
+    def __init__(
+        self,
+        model_id: str,
+        vision_backbone: VisionBackbone,
+        llm_backbone: LLMBackbone,
+        enable_mixed_precision_training: bool = True,
+        arch_specifier: str = "gelu-mlp",
+        **kwargs,
+    ) -> None:
+        super().__init__(
+            "prismatic",
+            model_id,
+            vision_backbone,
+            llm_backbone,
+            enable_mixed_precision_training=enable_mixed_precision_training,
+        )
+
+        # Set Weight Initialization Seed for Projector Consistency
+        torch.manual_seed(vision_backbone.embed_dim)
+
+        # Initialize Projection (Adapter) based on `arch_specifier`
+        self.arch_specifier = arch_specifier
+        if arch_specifier == "linear":
+            self.projector = LinearProjector(vision_backbone.embed_dim, llm_backbone.embed_dim)
+        elif arch_specifier.endswith("fused-gelu-mlp"):
+            self.projector = FusedMLPProjector(vision_backbone.embed_dim, llm_backbone.embed_dim)
+        elif arch_specifier.endswith("gelu-mlp"):
+            self.projector = MLPProjector(vision_backbone.embed_dim, llm_backbone.embed_dim)
+        else:
+            raise ValueError(f"PrismaticVLM with `{arch_specifier = }` is not supported!")
+
+        # Trackers
+        self.vision_backbone_requires_grad = False
+
+        # Set Module Keys =>> used in Checkpoint Saving / Model Loading
+        self.all_module_keys = ["vision_backbone", "llm_backbone", "projector"]
+        self.trainable_module_keys = []
+
+        # === Generation Utilities ===
+        #   => For computing likelihoods --> get tokens corresponding to "True", "False" and "Yes", "No"
+        self.string2idx = {}
+        for trigger_string in ["True", "False", "Yes", "No"] + [chr(ord("A") + i) for i in range(26)]:
+            token_idx_list = self.llm_backbone.tokenizer.encode(trigger_string, add_special_tokens=False)
+            assert len(token_idx_list) == 1, f'String "{trigger_string}" is tokenized as more than one token!'
+            self.string2idx[trigger_string] = token_idx_list[0]
+
+    @classmethod
+    def from_pretrained(
+        cls,
+        pretrained_checkpoint: Path,
+        model_id: str,
+        vision_backbone: VisionBackbone,
+        llm_backbone: LLMBackbone,
+        enable_mixed_precision_training: bool = True,
+        arch_specifier: str = "gelu-mlp",
+        freeze_weights: bool = True,
+        **kwargs,
+    ) -> PrismaticVLM:
+        """Initialize a PrismaticVLM from a pretrained checkpoint, freezing all weights, tailored for inference."""
+        vlm = cls(
+            model_id,
+            vision_backbone,
+            llm_backbone,
+            enable_mixed_precision_training=enable_mixed_precision_training,
+            arch_specifier=arch_specifier,
+            **kwargs,
+        )
+
+        # Load from Checkpoint (Custom --> should load both *projector* and *llm* weights)
+        model_state_dict = torch.load(pretrained_checkpoint, map_location="cpu")["model"]
+        assert (
+            "projector" in model_state_dict and "llm_backbone" in model_state_dict
+        ), "PrismaticVLM `from_pretrained` expects checkpoint with keys for `projector` AND `llm_backbone`!"
+
+        vlm.projector.load_state_dict(model_state_dict["projector"])
+        vlm.llm_backbone.load_state_dict(model_state_dict["llm_backbone"])
+        if "vision_backbone" in model_state_dict.keys():
+            vlm.vision_backbone.load_state_dict(model_state_dict["vision_backbone"])
+
+        # Freeze Weights
+        if freeze_weights:
+            vlm.requires_grad_(False)
+            vlm.eval()
+
+        return vlm
+
+    def get_prompt_builder(self, system_prompt: Optional[str] = None) -> PromptBuilder:
+        prompt_initializer: Type[PromptBuilder] = self.llm_backbone.prompt_builder_fn
+        return prompt_initializer(self.model_family, system_prompt=system_prompt)
+
+    def freeze_backbones(self, stage: str) -> None:
+        """
+        This function sets `requires_grad_` on each of the component modules explicitly, depending on stage.
+
+        We support two separate stages --> "align" and "finetune".
+            => "align" --> vision_backbone*, llm_backbone* are frozen; only the `projector` is trained.
+            => "finetune" --> vision_backbone* is frozen; both `projector` and `llm_backbone` are trained.
+
+        :param stage: Pretraining stage in < "align" | "finetune" | "full-finetune" | "vla-train" | "vla-full-train" >
+        """
+        if stage == "align":
+            self.vision_backbone.requires_grad_(False)
+            self.llm_backbone.requires_grad_(False)
+            self.projector.requires_grad_(True)
+
+            # Add to `self.trainable_module_keys`
+            self.trainable_module_keys = ["projector"]
+
+            # Update Trackers
+            self.vision_backbone_requires_grad = False
+
+            # Explicitly Log Frozen / Trainable Components
+            overwatch.info(f"[Frozen]    🥶 =>> Vision Backbone `{self.vision_backbone.identifier}`", ctx_level=1)
+            overwatch.info(f"[Frozen]    🥶 =>> LLM Backbone `{self.llm_backbone.identifier}`", ctx_level=1)
+            overwatch.info(f"[TRAINABLE] 🔥 =>> Projector `{self.arch_specifier}`", ctx_level=1)
+
+        elif stage in {"finetune", "vla-train"}:
+            self.vision_backbone.requires_grad_(False)
+            self.llm_backbone.requires_grad_(True)
+            self.projector.requires_grad_(True)
+
+            # Add to `self.trainable_module_keys`
+            self.trainable_module_keys = ["projector", "llm_backbone"]
+
+            # Update Trackers
+            self.vision_backbone_requires_grad = False
+
+            # Explicitly Log Frozen / Unfrozen Components
+            overwatch.info(f"[Frozen]    🥶 =>> Vision Backbone `{self.vision_backbone.identifier}`", ctx_level=1)
+            overwatch.info(f"[TRAINABLE] 🔥 =>> LLM Backbone `{self.llm_backbone.identifier}`", ctx_level=1)
+            overwatch.info(f"[TRAINABLE] 🔥 =>> Projector `{self.arch_specifier}`", ctx_level=1)
+
+        elif stage in {"full-finetune", "vla-full-train"}:
+            self.vision_backbone.dtype = torch.float32
+            self.vision_backbone.requires_grad_(True)
+            self.llm_backbone.requires_grad_(True)
+            self.projector.requires_grad_(True)
+
+            # Add to `self.trainable_module_keys`
+            self.trainable_module_keys = ["vision_backbone", "projector", "llm_backbone"]
+
+            # Update Trackers
+            self.vision_backbone_requires_grad = True
+
+            # Explicitly Log Frozen / Unfrozen Components
+            overwatch.info(f"[TRAINABLE] 🔥 =>> Vision Backbone `{self.vision_backbone.identifier}`", ctx_level=1)
+            overwatch.info(f"[TRAINABLE] 🔥 =>> LLM Backbone `{self.llm_backbone.identifier}`", ctx_level=1)
+            overwatch.info(f"[TRAINABLE] 🔥 =>> Projector `{self.arch_specifier}`", ctx_level=1)
+
+        elif stage in {"last-layer-finetune", "vla-last-layer-train"}:
+            self.vision_backbone.requires_grad_(False)
+            self.projector.requires_grad_(False)
+            self.llm_backbone.requires_grad_(False)
+
+            # Unfreeze final LLM layer
+            for module in self.llm_backbone.last_layer_finetune_modules:
+                module.requires_grad_(True)
+
+            # Add to `self.trainable_module_keys`
+            self.trainable_module_keys = ["llm_backbone"]
+
+            # Update Trackers
+            self.vision_backbone_requires_grad = False
+
+            # Explicitly Log Frozen / Unfrozen Components
+            # fmt: off
+            overwatch.info(f"[Frozen]                    🥶   =>> Vision Backbone `{self.vision_backbone.identifier}`", ctx_level=1)  # noqa: E501
+            overwatch.info(f"[Frozen, except last layer] 🥶🔥 =>> LLM Backbone `{self.llm_backbone.identifier}`", ctx_level=1)  # noqa: E501
+            overwatch.info(f"[Frozen]                    🥶   =>> Projector `{self.arch_specifier}`", ctx_level=1)
+            # fmt: on
+
+        elif stage in {"vla-sandwich-train"}:
+            self.vision_backbone.dtype = torch.float32
+            self.vision_backbone.requires_grad_(True)
+            self.projector.requires_grad_(True)
+            self.llm_backbone.requires_grad_(False)
+
+            # Unfreeze final LLM layer
+            for module in self.llm_backbone.last_layer_finetune_modules:
+                module.requires_grad_(True)
+
+            # Add to `self.trainable_module_keys`
+            self.trainable_module_keys = ["vision_backbone", "projector", "llm_backbone"]
+
+            # Update Trackers
+            self.vision_backbone_requires_grad = True
+
+            # Explicitly Log Frozen / Unfrozen Components
+            # fmt: off
+            overwatch.info(f"[TRAINABLE]                 🔥   =>> Vision Backbone `{self.vision_backbone.identifier}`", ctx_level=1)  # noqa: E501
+            overwatch.info(f"[Frozen, except last layer] 🥶🔥 =>> LLM Backbone `{self.llm_backbone.identifier}`", ctx_level=1)  # noqa: E501
+            overwatch.info(f"[TRAINABLE]                 🔥   =>> Projector `{self.arch_specifier}`", ctx_level=1)
+            # fmt: on
+
+        else:
+            raise ValueError(f"Stage `{stage}` is not supported for LLaVa! Try < align | finetune >")
+
+        overwatch.debug("##################################################")
+        overwatch.debug("#####      Trainable Network Parameters:     #####")
+        overwatch.debug("##################################################")
+        for name, param in self.named_parameters():
+            if param.requires_grad:
+                overwatch.debug(name)
+
+    def load_from_checkpoint(self, stage: str, run_dir: Path, pretrained_checkpoint: Optional[Path] = None) -> None:
+        """Load weights from checkpoint (if required by the given stage)."""
+        assert stage in {"align", "finetune", "full-finetune"}, f"Stage {stage} is not supported!"
+
+        # If we're running a `no-align` architecture, we're good!
+        if self.arch_specifier.startswith("no-align"):
+            overwatch.info(
+                f"PrismaticVLM with `{self.arch_specifier = }` does not require pretrained weights!", ctx_level=1
+            )
+            return
+
+        # Otherwise, handle stage-specific logic!
+        if stage == "align":
+            overwatch.info("Stage `align` does not require pretrained weights =>> Starting Training", ctx_level=1)
+            return
+
+        # Otherwise, load from `pretrained_checkpoint` or match on `run_dir` (s/+stage-finetune/+stage-align/g)
+        overwatch.info("Stage `finetune` requires `align` pretrained weights", ctx_level=1)
+
+        # Config specifies path to a checkpoint to load
+        if pretrained_checkpoint is not None:
+            overwatch.info(f"Loading from Provided Checkpoint `{pretrained_checkpoint}`", ctx_level=1)
+            model_state_dict = torch.load(pretrained_checkpoint)["model"]
+            self.projector.load_state_dict(model_state_dict["projector"])
+
+            return
+
+        # [Contract] If no `pretrained_checkpoint`, assume `align` lives in the run directory; string substitution!
+        model, scale, _, seed = run_dir.name.split("+")
+        align_dirs = [
+            d
+            for d in run_dir.parent.iterdir()
+            if (d.name.startswith(f"{model}+{scale}") and d.name.endswith(f"+stage-align+{seed}"))
+        ]
+        assert len(align_dirs) == 1, "Multiple or No Valid Pretrained Directories Exist -- Double Check `runs`!"
+        if (pretrained_checkpoint := (align_dirs[0] / "checkpoints" / "latest-checkpoint.pt")).exists():
+            overwatch.info(f"Loading from Discovered Checkpoint `{pretrained_checkpoint}`", ctx_level=1)
+            model_state_dict = torch.load(pretrained_checkpoint)["model"]
+            self.projector.load_state_dict(model_state_dict["projector"])
+        else:
+            raise ValueError(f"Could not find valid `align` checkpoint at {pretrained_checkpoint}!")
+
+    def get_fsdp_wrapping_policy(self) -> Callable:
+        """Return an FSDP _or_policy over the policies returned by each individual backbone (and our VLM policy)."""
+        vision_fsdp_wrapping_policy = self.vision_backbone.get_fsdp_wrapping_policy()
+        llm_fsdp_wrapping_policy = self.llm_backbone.get_fsdp_wrapping_policy()
+
+        # Get Prismatic Wrapping Policy =>> just a module wrapping policy around `self.projector`
+        prismatic_fsdp_wrapping_policy = partial(
+            _module_wrap_policy,
+            module_classes={LinearProjector, MLPProjector, FusedMLPProjector},
+        )
+
+        # Return union (_or_) over constituent policies
+        #   => Note: there is *not* a fall-through policy; any module that isn't covered by the above constituents will
+        #            automatically be folded into the root VLM FSDP instance.
+        return partial(
+            _or_policy,
+            policies=[
+                vision_fsdp_wrapping_policy,
+                llm_fsdp_wrapping_policy,
+                prismatic_fsdp_wrapping_policy,
+            ],
+        )
+
+    # Note =>> We're not explicitly subclassing `PreTrainedModel` because we don't need the bloat; however, `forward()`
+    #          *must* match the signature of a `{Model}ForCausalLM` so that we can inherit from `GenerationMixin`
+
+    # ruff: noqa: C901
+    def forward(
+        self,
+        input_ids: Optional[torch.LongTensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        pixel_values: Optional[torch.FloatTensor] = None,
+        labels: Optional[torch.LongTensor] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        past_key_values: Optional[List[torch.FloatTensor]] = None,
+        use_cache: Optional[bool] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+        multimodal_indices: Optional[torch.LongTensor] = None,
+    ) -> CausalLMOutputWithPast:
+        """Run a forward pass through the VLM, returning a CausalLMOutputWithPast instance (contains loss)."""
+
+        # Handle Inference (leverage cache, short-circuit on just LLM forward)
+        if input_ids.shape[1] == 1 and past_key_values is not None:
+            # We're leveraging the cache, so just redirect to `self.llm_backbone` with `input_ids` and `past_key_values`
+            output = self.llm_backbone(
+                input_ids=input_ids,
+                attention_mask=None,
+                position_ids=None,
+                past_key_values=past_key_values,
+                inputs_embeds=None,
+                labels=None,
+                use_cache=use_cache,
+                output_attentions=output_attentions,
+                output_hidden_states=output_hidden_states,
+                return_dict=return_dict,
+            )
+            return output
+
+        elif input_ids.shape[1] == 1 or pixel_values is None:
+            raise RuntimeError("Invalid `forward()` call!")
+
+        # Handle Multimodal Indices is None --> pretend like the batch is fully multimodal (always image + text)!
+        if multimodal_indices is None:
+            multimodal_indices = torch.arange(len(input_ids), dtype=torch.long, device=input_ids.device)
+
+        # Handle Multimodal Indices is Empty (len == 0) --> simple unimodal forward
+        elif len(multimodal_indices) == 0:
+            return self.llm_backbone(
+                input_ids=input_ids,
+                attention_mask=attention_mask,
+                position_ids=None,
+                past_key_values=past_key_values,
+                inputs_embeds=None,
+                labels=labels,
+                use_cache=use_cache,
+                output_attentions=output_attentions,
+                output_hidden_states=output_hidden_states,
+                return_dict=return_dict,
+            )
+
+        # Run Visual Feature Extraction
+        with torch.set_grad_enabled(self.vision_backbone_requires_grad):
+            if isinstance(pixel_values, dict):
+                patch_features = self.vision_backbone({k: pixel_values[k][multimodal_indices] for k in pixel_values})
+            else:
+                patch_features = self.vision_backbone(pixel_values[multimodal_indices])
+
+        # Projection Logic :: [bsz, num_patches, llm_embed_dim] =>> num_patches = (2 *) (256 + 1) for ViT-L + CLS
+        projected_patch_embeddings = self.projector(patch_features)
+        projected_patch_attention_mask = None
+        if attention_mask is not None:
+            projected_patch_attention_mask = torch.full(
+                (projected_patch_embeddings.shape[0], projected_patch_embeddings.shape[1]),
+                True,
+                dtype=attention_mask.dtype,
+                device=attention_mask.device,
+            )
+
+        # Get Input Embeddings from LLM Backbone :: [bsz, input_seq_len, llm_embed_dim]
+        input_embeddings = self.llm_backbone.embed_input_ids(input_ids)
+
+        # Build Multimodal Embeddings (and build resulting attention mask)
+        multimodal_embeddings = torch.cat(
+            [
+                input_embeddings[multimodal_indices, :1, :],
+                projected_patch_embeddings,
+                input_embeddings[multimodal_indices, 1:, :],
+            ],
+            dim=1,
+        )
+        multimodal_attention_mask = None
+        if attention_mask is not None:
+            multimodal_attention_mask = torch.cat(
+                [
+                    attention_mask[multimodal_indices, :1],
+                    projected_patch_attention_mask,
+                    attention_mask[multimodal_indices, 1:],
+                ],
+                dim=1,
+            )
+
+        # [Contract] We assume the first token of `labels` (associated with <BOS>) is already marked as "IGNORE"
+        #   => We'll ignore the per-token outputs for each of the patch embeddings as well!
+        multimodal_labels = None
+        if labels is not None:
+            projected_patch_labels = torch.full(
+                (projected_patch_embeddings.shape[0], projected_patch_embeddings.shape[1]),
+                IGNORE_INDEX,
+                dtype=labels.dtype,
+                device=labels.device,
+            )
+            multimodal_labels = torch.cat(
+                [labels[multimodal_indices, :1], projected_patch_labels, labels[multimodal_indices, 1:]], dim=1
+            )
+
+        # === Add Unimodal Handling ===
+
+        # Create Fused Embeddings, Attention Mask, and Labels by Merging with "unimodal" Inputs (if applicable)
+        unimodal_indices = torch.tensor(
+            [idx for idx in range(len(input_ids)) if idx not in multimodal_indices],
+            dtype=torch.long,
+            device=multimodal_indices.device,
+        )
+
+        # No "unimodal" data --> Fused == Multimodal
+        if len(unimodal_indices) == 0:
+            fused_embeddings = multimodal_embeddings
+            fused_attention_mask = multimodal_attention_mask
+            fused_labels = multimodal_labels
+
+        else:
+            # Otherwise --> Merge w/ unimodal data
+
+            # This doesn't matter --> but in the "normal" case this is the embedding of the <PAD> token
+            #   => NOTE :: Verified that `zeros/randn/empty/<PAD> embedding` all return the same result!
+            unimodal_embeddings_pad = torch.zeros(
+                (len(unimodal_indices), projected_patch_embeddings.shape[1], input_embeddings.shape[2]),
+                dtype=input_embeddings.dtype,
+                device=input_embeddings.device,
+            )
+            unimodal_attention_pad = torch.full(
+                (len(unimodal_indices), projected_patch_embeddings.shape[1]),
+                False,
+                dtype=attention_mask.dtype,
+                device=attention_mask.device,
+            )
+            unimodal_labels_pad = torch.full(
+                (len(unimodal_indices), projected_patch_embeddings.shape[1]),
+                IGNORE_INDEX,
+                dtype=labels.dtype,
+                device=labels.device,
+            )
+
+            unimodal_embeddings = torch.cat([input_embeddings[unimodal_indices], unimodal_embeddings_pad], dim=1)
+            unimodal_attention_mask = torch.cat([attention_mask[unimodal_indices], unimodal_attention_pad], dim=1)
+            unimodal_labels = torch.cat([labels[unimodal_indices], unimodal_labels_pad], dim=1)
+
+            # Create "Fused" Tensors by Stacking Multimodal & Unimodal
+            fused_embeddings = torch.vstack([multimodal_embeddings, unimodal_embeddings])
+            fused_attention_mask = torch.vstack([multimodal_attention_mask, unimodal_attention_mask])
+            fused_labels = torch.vstack([multimodal_labels, unimodal_labels])
+
+        # Run LLM Forward --> returns CausalLMOutputWithPast!
+        return self.llm_backbone(
+            input_ids=None,
+            attention_mask=fused_attention_mask,
+            position_ids=None,
+            past_key_values=past_key_values,
+            inputs_embeds=fused_embeddings,
+            labels=fused_labels,
+            use_cache=use_cache,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            return_dict=return_dict,
+        )
+
+    # === GenerationMixin Methods ===
+    #   => Note: The following methods override the functionality of `transformers.GenerationMixin`; these expect the
+    #            contract in each of the function signatures, and also expect our `forward` function to roughly take
+    #            the same arguments as the underlying LLM (see `LlamaModelForCausalLM` as an example)
+
+    def prepare_inputs_for_generation(
+        self,
+        input_ids: Optional[torch.LongTensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        pixel_values: Optional[torch.FloatTensor] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        past_key_values: Optional[List[torch.FloatTensor]] = None,
+        use_cache: Optional[bool] = None,
+        **kwargs: torch.Tensor,
+    ) -> Dict[str, torch.Tensor]:
+        """Borrowed from `LlamaForCausalLM` --> in general, just handles caching logic during generation."""
+        if past_key_values:
+            input_ids = input_ids[:, -1:]
+
+        # if `inputs_embeds` are passed, we only want to use them in the 1st generation step
+        if inputs_embeds is not None and past_key_values is None:
+            model_inputs = {"inputs_embeds": inputs_embeds}
+        else:
+            model_inputs = {"input_ids": input_ids}
+
+        # Make sure `pixel_values` are preserved in `model_inputs`
+        model_inputs.update(
+            {
+                "attention_mask": attention_mask,
+                "pixel_values": pixel_values,
+                "past_key_values": past_key_values,
+                "use_cache": use_cache,
+            }
+        )
+
+        return model_inputs
+
+    @torch.inference_mode()
+    def generate_batch(
+        self,
+        pixel_values: Union[torch.Tensor, Dict[str, torch.Tensor]],
+        texts: List[str],
+        return_string_probabilities: Optional[List[str]] = None,
+        **kwargs: str,
+    ) -> Union[List[str], List[List[float]]]:
+        # For now, only support generation with a batch size of 1 for simplicity
+        tokenizer = self.llm_backbone.tokenizer
+
+        # Prepare Inputs
+        batch_input_ids = [
+            tokenizer(text, truncation=True, return_tensors="pt").input_ids.to(self.device) for text in texts
+        ]
+        if isinstance(pixel_values, torch.Tensor):
+            pixel_values = pixel_values[None, ...].to(self.device)
+        elif isinstance(pixel_values, dict):
+            pixel_values = {k: v[None, ...].to(self.device) for k, v in pixel_values.items()}
+        else:
+            raise ValueError(f"Unsupported `pixel_values` type = {type(pixel_values)}")
+
+        # Create Output Lists
+        gen_texts, gen_probabilities = [], []
+
+        # Invoke super().generate --> taps into `GenerationMixin` which (redirects) to `forward()`
+        autocast_dtype = self.llm_backbone.half_precision_dtype
+        with torch.autocast("cuda", dtype=autocast_dtype, enabled=self.enable_mixed_precision_training):
+            for idx, input_ids in enumerate(batch_input_ids):
+                if isinstance(pixel_values, torch.Tensor):
+                    pixel_values = pixel_values[idx]
+                elif isinstance(pixel_values, dict):
+                    pixel_values = {k: pixel_values[k][idx] for k in pixel_values}
+                else:
+                    raise ValueError(f"Unsupported `pixel_values` type = {type(pixel_values)}")
+
+                # Handle `return_string_probabilities`
+                if return_string_probabilities is None:
+                    full_out_ids = super().generate(input_ids=input_ids, pixel_values=pixel_values, **kwargs)
+                    gen_ids = full_out_ids[0, input_ids.shape[1] :]
+
+                    # Decode `gen_ids` and strip any <EOS> tokens
+                    gen_texts.append(tokenizer.decode(gen_ids, skip_special_tokens=True).strip())
+
+                else:
+                    full_out_dict = super().generate(
+                        input_ids=input_ids,
+                        pixel_values=pixel_values,
+                        output_scores=True,
+                        return_dict_in_generate=True,
+                        **kwargs,
+                    )
+
+                    # Generation pattern should usually be [TOKEN] <EOS> for True/False and Yes/No Generations
+                    gen_ids = full_out_dict.sequences[0, input_ids.shape[1] :]
+
+                    # [Debug] Verify that the first token generated is in `self.string2idx.values()`
+                    # assert gen_ids[0] in self.string2idx.values(), "Generated ID not in mapping!"
+
+                    # Decode `gen_ids` and strip any <EOS> tokens
+                    gen_texts.append(tokenizer.decode(gen_ids, skip_special_tokens=True).strip())
+
+                    # Get all token probabilities --> softmax over logits
+                    token_probs = torch.softmax(full_out_dict.scores[0][0], dim=0)
+
+                    # Get *normalized* probabilities for all values in `return_token_probabilities`
+                    slice_idxs = torch.tensor([self.string2idx[s] for s in return_string_probabilities])
+                    string_probs_unnormalized = token_probs[slice_idxs]
+                    string_probs = string_probs_unnormalized / string_probs_unnormalized.sum()
+                    gen_probabilities.append(string_probs.cpu().numpy().tolist())
+
+        return gen_texts if return_string_probabilities is None else gen_probabilities
+
+    @torch.inference_mode()
+    def generate(self, image: Image, prompt_text: str, **kwargs: str) -> str:
+        # For now, only support generation with a batch size of 1 for simplicity
+        image_transform, tokenizer = self.vision_backbone.image_transform, self.llm_backbone.tokenizer
+
+        # Prepare Inputs
+        input_ids = tokenizer(prompt_text, truncation=True, return_tensors="pt").input_ids.to(self.device)
+        pixel_values = image_transform(image)
+        if isinstance(pixel_values, torch.Tensor):
+            pixel_values = pixel_values[None, ...].to(self.device)
+        elif isinstance(pixel_values, dict):
+            pixel_values = {k: v[None, ...].to(self.device) for k, v in pixel_values.items()}
+        else:
+            raise ValueError(f"Unsupported `pixel_values` type = {type(pixel_values)}")
+
+        # Invoke super().generate --> taps into `GenerationMixin` which (redirects) to `forward()`
+        autocast_dtype = self.llm_backbone.half_precision_dtype
+        with torch.autocast("cuda", dtype=autocast_dtype, enabled=self.enable_mixed_precision_training):
+            # fmt: off
+            generated_ids = super().generate(
+                input_ids=input_ids,            # Shape: [1, seq]
+                pixel_values=pixel_values,      # Shape: [1, 3, res, res] or Dict[str, Shape[1, 3, res, res]]
+                **kwargs
+            )
+            # fmt: on
+
+        generated_text = tokenizer.decode(generated_ids[0, input_ids.shape[1] :], skip_special_tokens=True).strip()
+
+        return generated_text

capvector-oft/prismatic/overwatch/__init__.py

@@ -0,0 +1 @@
+from .overwatch import initialize_overwatch

capvector-oft/prismatic/overwatch/overwatch.py

@@ -0,0 +1,147 @@
+"""
+overwatch.py
+
+Utility class for creating a centralized/standardized logger (built on Rich) and accelerate handler.
+"""
+
+import logging
+import logging.config
+import os
+from contextlib import nullcontext
+from logging import LoggerAdapter
+from typing import Any, Callable, ClassVar, Dict, MutableMapping, Tuple, Union
+
+# Overwatch Default Format String
+RICH_FORMATTER, DATEFMT = "| >> %(message)s", "%m/%d [%H:%M:%S]"
+
+# Set Logging Configuration
+LOG_CONFIG = {
+    "version": 1,
+    "disable_existing_loggers": True,
+    "formatters": {"simple-console": {"format": RICH_FORMATTER, "datefmt": DATEFMT}},
+    "handlers": {
+        "console": {
+            "class": "rich.logging.RichHandler",
+            "formatter": "simple-console",
+            "markup": True,
+            "rich_tracebacks": True,
+            "show_level": True,
+            "show_path": True,
+            "show_time": True,
+        }
+    },
+    "root": {"level": "INFO", "handlers": ["console"]},
+}
+logging.config.dictConfig(LOG_CONFIG)
+
+
+# === Custom Contextual Logging Logic ===
+class ContextAdapter(LoggerAdapter):
+    CTX_PREFIXES: ClassVar[Dict[int, str]] = {**{0: "[*] "}, **{idx: "|=> ".rjust(4 + (idx * 4)) for idx in [1, 2, 3]}}
+
+    def process(self, msg: str, kwargs: MutableMapping[str, Any]) -> Tuple[str, MutableMapping[str, Any]]:
+        ctx_level = kwargs.pop("ctx_level", 0)
+        return f"{self.CTX_PREFIXES[ctx_level]}{msg}", kwargs
+
+
+class DistributedOverwatch:
+    def __init__(self, name: str) -> None:
+        """Initializer for an Overwatch object that wraps logging & `accelerate.PartialState`."""
+        from accelerate import PartialState
+
+        # Note that PartialState is always safe to initialize regardless of `accelerate launch` or `torchrun`
+        #   =>> However, might be worth actually figuring out if we need the `accelerate` dependency at all!
+        self.logger, self.distributed_state = ContextAdapter(logging.getLogger(name), extra={}), PartialState()
+
+        # Logger Delegation (for convenience; would be nice to just compose & dynamic dispatch eventually)
+        self.debug = self.logger.debug
+        self.info = self.logger.info
+        self.warning = self.logger.warning
+        self.error = self.logger.error
+        self.critical = self.logger.critical
+
+        # Logging Defaults =>> only Log `INFO` on Main Process, `ERROR` on others!
+        self.logger.setLevel(logging.INFO if self.distributed_state.is_main_process else logging.ERROR)
+
+    @property
+    def rank_zero_only(self) -> Callable[..., Any]:
+        return self.distributed_state.on_main_process
+
+    @property
+    def local_zero_only(self) -> Callable[..., Any]:
+        return self.distributed_state.on_local_main_process
+
+    @property
+    def rank_zero_first(self) -> Callable[..., Any]:
+        return self.distributed_state.main_process_first
+
+    @property
+    def local_zero_first(self) -> Callable[..., Any]:
+        return self.distributed_state.local_main_process_first
+
+    def is_rank_zero(self) -> bool:
+        return self.distributed_state.is_main_process
+
+    def rank(self) -> int:
+        return self.distributed_state.process_index
+
+    def local_rank(self) -> int:
+        return self.distributed_state.local_process_index
+
+    def world_size(self) -> int:
+        return self.distributed_state.num_processes
+
+
+class PureOverwatch:
+    def __init__(self, name: str) -> None:
+        """Initializer for an Overwatch object that just wraps logging."""
+        self.logger = ContextAdapter(logging.getLogger(name), extra={})
+
+        # Logger Delegation (for convenience; would be nice to just compose & dynamic dispatch eventually)
+        self.debug = self.logger.debug
+        self.info = self.logger.info
+        self.warning = self.logger.warning
+        self.error = self.logger.error
+        self.critical = self.logger.critical
+
+        # Logging Defaults =>> INFO
+        self.logger.setLevel(logging.INFO)
+
+    @staticmethod
+    def get_identity_ctx() -> Callable[..., Any]:
+        def identity(fn: Callable[..., Any]) -> Callable[..., Any]:
+            return fn
+
+        return identity
+
+    @property
+    def rank_zero_only(self) -> Callable[..., Any]:
+        return self.get_identity_ctx()
+
+    @property
+    def local_zero_only(self) -> Callable[..., Any]:
+        return self.get_identity_ctx()
+
+    @property
+    def rank_zero_first(self) -> Callable[..., Any]:
+        return nullcontext
+
+    @property
+    def local_zero_first(self) -> Callable[..., Any]:
+        return nullcontext
+
+    @staticmethod
+    def is_rank_zero() -> bool:
+        return True
+
+    @staticmethod
+    def rank() -> int:
+        return 0
+
+    @staticmethod
+    def world_size() -> int:
+        return 1
+
+
+def initialize_overwatch(name: str) -> Union[DistributedOverwatch, PureOverwatch]:
+    return DistributedOverwatch(name) if int(os.environ.get("WORLD_SIZE", -1)) != -1 else PureOverwatch(name)

capvector-oft/prismatic/preprocessing/__init__.py

@@ -0,0 +1,2 @@
+from .download import convert_to_jpg, download_extract
+from .materialize import get_dataset_and_collator

capvector-oft/prismatic/preprocessing/datasets/__init__.py

@@ -0,0 +1 @@
+from .datasets import AlignDataset, FinetuneDataset

capvector-oft/prismatic/preprocessing/datasets/datasets.py

@@ -0,0 +1,200 @@
+"""
+datasets.py
+
+PyTorch Dataset Definitions for Prismatic models; supports processing for both the `align` and `finetune` stages, with
+utilities for formatting conversations during the `finetune` stage subject to the given LLM backbone's expected
+formatting (e.g., SYS_PROMPT + USER: ... ASSISTANT: ... for Vicuña v1.5 Chat models).
+
+We currently only support Map-style Datasets; assumes that all files (annotations, images) are on local disk, and that
+random access image reading is relatively cheap/fast.
+"""
+
+import copy
+import json
+from pathlib import Path
+from typing import Dict, List, Tuple, Type
+
+import torch
+from PIL import Image
+from torch.utils.data import Dataset
+from transformers import CodeGenTokenizerFast, LlamaTokenizerFast, PreTrainedTokenizerBase
+
+from prismatic.models.backbones.llm.prompting import PromptBuilder
+from prismatic.models.backbones.vision import ImageTransform
+
+# HuggingFace Default / LLaMa-2 IGNORE_INDEX (for labels)
+IGNORE_INDEX = -100
+
+
+class AlignDataset(Dataset[Dict[str, torch.Tensor]]):
+    def __init__(
+        self,
+        chat_json: Path,
+        image_dir: Path,
+        image_transform: ImageTransform,
+        tokenizer: PreTrainedTokenizerBase,
+    ) -> None:
+        super().__init__()
+        self.chat_json, self.image_dir = chat_json, image_dir
+        self.image_transform, self.tokenizer = image_transform, tokenizer
+        self.dataset_type = "align"
+
+        # Create Prompt Template
+        self.prompt_template = "{caption}" + self.tokenizer.eos_token
+
+        # Load Chat JSON
+        with open(self.chat_json, "r") as f:
+            self.examples = json.load(f)
+
+    def __getitem__(self, idx: int) -> Dict[str, torch.Tensor]:
+        """
+        Following the *actual* code executed from the LLaVa codebase, during the "align" phase, we actually discard
+        the "prompt" from the human, and instead directly predict the caption from the image.
+
+        As a concrete example given the "raw data" for the first example:
+            example = self.examples[0]["conversations"]` = {
+                [
+                    {"from": "human", "value": "Render a clear and concise summary of the photo.\n<image>"},
+                    {"from": "gpt", "value": "select luxury furniture 3 - inch gel memory foam mattress topper"}
+                ]
+            }
+
+        Return =>> self.tokenizer("<image> select luxury furniture 3 - inch gel memory foam mattress topper\n")
+
+        :param idx: Index to retrieve from the dataset.
+
+        :return: Dictionary of {"pixel_values": torch.Tensor, "input_ids": torch.Tensor, "labels": torch.Tensor}
+        """
+        image_path, conversation = Path(self.examples[idx]["image"]), self.examples[idx]["conversations"]
+        assert (len(conversation) == 2) and ("<image>" not in conversation[-1]["value"]), "Unexpected text!"
+
+        # Format Caption --> {caption}{eos_token}
+        caption = self.prompt_template.format(caption=conversation[-1]["value"].strip())
+
+        # We treat image patches as "tokens = [p1 p2 p3, ...]"; we need to specify ordering of text/patch tokens.
+        #   => Critically, we find that inserting *after* the BOS token leads to the strongest performance!
+        #       - input_ids = "<s> p1 p2 p3 ... <caption_text> \n"
+        #       - labels = "IGNORE IGNORE ..." (copy `input_ids` replacing <s> and p{1...K} with IGNORE)
+        #
+        # IMPORTANT => IF WE'RE USING HF LLM.forward(... labels=labels), SHIFTING HAPPENS _INSIDE_ MODEL!
+        input_ids = self.tokenizer(caption, truncation=True, return_tensors="pt").input_ids[0]
+        labels = copy.deepcopy(input_ids)
+
+        # Set the <BOS> token's label to IGNORE_INDEX (since we're inserting the image patches right after)
+        labels[0] = IGNORE_INDEX
+
+        # Process Image --> get "pixel_values" (will either be a torch.Tensor OR a Dict[str,torch.Tensor])
+        pixel_values = self.image_transform(Image.open(self.image_dir / image_path).convert("RGB"))
+
+        return dict(pixel_values=pixel_values, input_ids=input_ids, labels=labels)
+
+    def get_modality_lengths(self, n_image_patches: int) -> List[Tuple[bool, int]]:
+        """Get a list of modalities (unimodal / text-only vs. multimodal) and length of conversations per example."""
+        modality_lengths = []
+        for example in self.examples:
+            is_multimodal = "image" in example
+            n_words = sum([len(turn["value"].replace("<image>", "").split()) for turn in example["conversations"]])
+            modality_lengths.append((is_multimodal, (n_image_patches + n_words) if is_multimodal else n_words))
+        return modality_lengths
+
+    def __len__(self) -> int:
+        return len(self.examples)
+
+
+class FinetuneDataset(Dataset[Dict[str, torch.Tensor]]):
+    def __init__(
+        self,
+        instruct_json: Path,
+        image_dir: Path,
+        image_transform: ImageTransform,
+        tokenizer: PreTrainedTokenizerBase,
+        prompt_builder_fn: Type[PromptBuilder],
+    ) -> None:
+        super().__init__()
+        self.instruct_json, self.image_dir = instruct_json, image_dir
+        self.image_transform, self.tokenizer = image_transform, tokenizer
+        self.prompt_builder_fn = prompt_builder_fn
+        self.dataset_type = "finetune"
+
+        # Load Instruct JSON
+        with open(self.instruct_json, "r") as f:
+            self.examples = json.load(f)
+
+    # === Unimodal + Multimodal Handling ===
+    def __getitem__(self, idx: int) -> Dict[str, torch.Tensor]:
+        """
+        Unlike the *align* stage handling, for the *finetune* stage, we actually need to handle multiple "turns" of
+        dialog grounded in a single image.
+
+        To do this, we leverage the `prompt_builder_fn` which instantiates a PromptBuilder object. By calling the
+        methods for adding turns and getting a prompt, we ensure proper formatting and consistency for each example.
+
+        :param idx: Index to retrieve from the dataset.
+
+        :return: Dictionary of {"pixel_values": torch.Tensor, "input_ids": torch.Tensor, "labels": torch.Tensor}
+        """
+        conversation = self.examples[idx]["conversations"]
+
+        # Create Prompt Builder --> add each message sequentially
+        prompt_builder, input_ids, labels = self.prompt_builder_fn(model_family="prismatic"), [], []
+        for turn_idx, turn in enumerate(conversation):
+            # Get "effective" string added to prompt --> handle whitespace for tokenizer type!
+            msg = prompt_builder.add_turn(turn["from"], turn["value"])
+
+            # Llama Tokenizer (Fast) adds extra character if a string ends in whitespace --> strip if non-empty!
+            if isinstance(self.tokenizer, LlamaTokenizerFast):
+                msg = msg.rstrip()
+
+            # Phi-2 Tokenizer == CodeGenTokenizer (Fast) -- no special handling!
+            elif isinstance(self.tokenizer, CodeGenTokenizerFast):
+                pass
+
+            else:
+                raise ValueError(f"Tokenizer of type `{type(self.tokenizer)}` is not explicitly handled!")
+
+            # Tokenize Input IDs
+            turn_input_ids = self.tokenizer(msg, add_special_tokens=turn_idx == 0).input_ids
+
+            # [CRITICAL] We do not want to take the loss for the "USER: <msg>" prompts =>> just the responses!
+            turn_labels = (
+                [IGNORE_INDEX for _ in range(len(turn_input_ids))] if (turn_idx % 2) == 0 else list(turn_input_ids)
+            )
+
+            # Add to Trackers
+            input_ids.extend(turn_input_ids)
+            labels.extend(turn_labels)
+
+        # Tensorize =>> Set the <BOS> token's label to IGNORE_INDEX (since we're inserting the image patches after)
+        #   - IMPORTANT => IF WE'RE USING HF LLM.forward(... labels=labels), SHIFTING HAPPENS _INSIDE_ MODEL!
+        input_ids, labels = torch.tensor(input_ids), torch.tensor(labels)
+
+        # Handle Truncation (if necessary)
+        input_ids, labels = input_ids[: self.tokenizer.model_max_length], labels[: self.tokenizer.model_max_length]
+
+        # === Handle "unimodal" (language-only) vs. "multimodal" ===
+        if "image" in self.examples[idx]:
+            image_path = Path(self.examples[idx]["image"])
+
+            # Set the <BOS> token's label to IGNORE_INDEX (since we're inserting the image patches right after)
+            labels[0] = IGNORE_INDEX
+
+            # Process Image --> get "pixel_values" (will either be a torch.Tensor OR a Dict[str,torch.Tensor])
+            pixel_values = self.image_transform(Image.open(self.image_dir / image_path).convert("RGB"))
+
+            return dict(pixel_values=pixel_values, input_ids=input_ids, labels=labels)
+
+        else:
+            # No image --> return `pixel_values` = None; Collator will do the smart batch handling for us!
+            return dict(pixel_values=None, input_ids=input_ids, labels=labels)
+
+    def get_modality_lengths(self) -> List[Tuple[bool, int]]:
+        """Get a list of modalities (unimodal / text-only vs. multimodal) and length of conversations per example."""
+        modality_lengths = []
+        for example in self.examples:
+            is_multimodal = "image" in example
+            n_words = sum([len(turn["value"].split()) for turn in example["conversations"]])
+            modality_lengths.append((is_multimodal, n_words))
+        return modality_lengths
+
+    def __len__(self) -> int:
+        return len(self.examples)

capvector-oft/prismatic/preprocessing/download.py

@@ -0,0 +1,207 @@
+"""
+download.py
+
+Utility functions for downloading and extracting various datasets to (local) disk.
+"""
+
+import os
+import shutil
+from pathlib import Path
+from typing import Dict, List, TypedDict
+from zipfile import ZipFile
+
+import requests
+from PIL import Image
+from rich.progress import BarColumn, DownloadColumn, MofNCompleteColumn, Progress, TextColumn, TransferSpeedColumn
+from tqdm import tqdm
+
+from prismatic.overwatch import initialize_overwatch
+
+# Initialize Overwatch =>> Wraps `logging.Logger`
+overwatch = initialize_overwatch(__name__)
+
+
+# === Dataset Registry w/ Links ===
+# fmt: off
+DatasetComponent = TypedDict(
+    "DatasetComponent",
+    {"name": str, "extract": bool, "extract_type": str, "url": str, "do_rename": bool},
+    total=False
+)
+
+DATASET_REGISTRY: Dict[str, List[DatasetComponent]] = {
+    # === LLaVa v1.5 Dataset(s) ===
+
+    # Note =>> This is the full suite of datasets included in the LLaVa 1.5 "finetuning" stage; all the LLaVa v1.5
+    #          models are finetuned on this split. We use this dataset for all experiments in our paper.
+    "llava-laion-cc-sbu-558k": [
+        {
+            "name": "chat.json",        # Contains the "chat" traces :: {"human" => <prompt>, "gpt" => <caption>}
+            "extract": False,
+            "url": "https://huggingface.co/datasets/liuhaotian/LLaVA-Pretrain/resolve/main/blip_laion_cc_sbu_558k.json",
+            "do_rename": True,
+        },
+        {
+            "name": "images",           # Contains the LLaVa Processed Images (jpgs, 224x224 resolution)
+            "extract": True,
+            "extract_type": "directory",
+            "url": "https://huggingface.co/datasets/liuhaotian/LLaVA-Pretrain/resolve/main/images.zip",
+            "do_rename": False,
+        }
+    ],
+
+    "llava-v1.5-instruct": [
+        {
+            "name": "llava_v1_5_mix665k.json",
+            "extract": False,
+            "url": (
+                "https://huggingface.co/datasets/liuhaotian/LLaVA-Instruct-150K/resolve/main/llava_v1_5_mix665k.json"
+            ),
+            "do_rename": True,
+        },
+        {
+            "name": "coco/train2017",       # Visual Instruct Tuning images are all sourced from COCO Train 2017
+            "extract": True,
+            "extract_type": "directory",
+            "url": "http://images.cocodataset.org/zips/train2017.zip",
+            "do_rename": True,
+        },
+        {
+            "name": "gqa/images",
+            "extract": True,
+            "extract_type": "directory",
+            "url": "https://downloads.cs.stanford.edu/nlp/data/gqa/images.zip",
+            "do_rename": True,
+        },
+        {
+            "name": "ocr_vqa/images",
+            "extract": True,
+            "extract_type": "directory",
+            "url": "https://huggingface.co/datasets/qnguyen3/ocr_vqa/resolve/main/ocr_vqa.zip",
+            "do_rename": True,
+        },
+        {
+            "name": "textvqa/train_images",
+            "extract": True,
+            "extract_type": "directory",
+            "url": "https://dl.fbaipublicfiles.com/textvqa/images/train_val_images.zip",
+            "do_rename": True,
+        },
+        {
+            "name": "vg/VG_100K",
+            "extract": True,
+            "extract_type": "directory",
+            "url": "https://cs.stanford.edu/people/rak248/VG_100K_2/images.zip",
+            "do_rename": True,
+        },
+        {
+            "name": "vg/VG_100K_2",
+            "extract": True,
+            "extract_type": "directory",
+            "url": "https://cs.stanford.edu/people/rak248/VG_100K_2/images2.zip",
+            "do_rename": True,
+        },
+    ]
+}
+# fmt: on
+
+
+def convert_to_jpg(image_dir: Path) -> None:
+    """Handling for OCR-VQA Images specifically; iterates through directory, converts all GIFs/PNGs."""
+    overwatch.info(f"Converting all Images in `{image_dir}` to JPG")
+
+    for image_fn in tqdm(list(image_dir.iterdir())):
+        if image_fn.suffix in {".jpg", ".jpeg"} or (jpg_fn := image_dir / f"{image_fn.stem}.jpg").exists():
+            continue
+
+        if image_fn.suffix == ".gif":
+            gif = Image.open(image_fn)
+            gif.seek(0)
+            gif.convert("RGB").save(jpg_fn)
+        elif image_fn.suffix == ".png":
+            Image.open(image_fn).convert("RGB").save(jpg_fn)
+        else:
+            raise ValueError(f"Unexpected image format `{image_fn.suffix}`")
+
+
+def download_with_progress(url: str, download_dir: Path, chunk_size_bytes: int = 1024) -> Path:
+    """Utility function for downloading files from the internet, with a handy Rich-based progress bar."""
+    overwatch.info(f"Downloading {(dest_path := download_dir / Path(url).name)} from `{url}`", ctx_level=1)
+    if dest_path.exists():
+        return dest_path
+
+    # Otherwise --> fire an HTTP Request, with `stream = True`
+    response = requests.get(url, stream=True)
+
+    # Download w/ Transfer-Aware Progress
+    #   => Reference: https://github.com/Textualize/rich/blob/master/examples/downloader.py
+    with Progress(
+        TextColumn("[bold]{task.description} - {task.fields[fname]}"),
+        BarColumn(bar_width=None),
+        "[progress.percentage]{task.percentage:>3.1f}%",
+        "•",
+        DownloadColumn(),
+        "•",
+        TransferSpeedColumn(),
+        transient=True,
+    ) as dl_progress:
+        dl_tid = dl_progress.add_task(
+            "Downloading", fname=dest_path.name, total=int(response.headers.get("content-length", "None"))
+        )
+        with open(dest_path, "wb") as f:
+            for data in response.iter_content(chunk_size=chunk_size_bytes):
+                dl_progress.advance(dl_tid, f.write(data))
+
+    return dest_path
+
+
+def extract_with_progress(archive_path: Path, download_dir: Path, extract_type: str, cleanup: bool = False) -> Path:
+    """Utility function for extracting compressed archives, with a handy Rich-based progress bar."""
+    assert archive_path.suffix == ".zip", "Only `.zip` compressed archives are supported for now!"
+    overwatch.info(f"Extracting {archive_path.name} to `{download_dir}`", ctx_level=1)
+
+    # Extract w/ Progress
+    with Progress(
+        TextColumn("[bold]{task.description} - {task.fields[aname]}"),
+        BarColumn(bar_width=None),
+        "[progress.percentage]{task.percentage:>3.1f}%",
+        "•",
+        MofNCompleteColumn(),
+        transient=True,
+    ) as ext_progress:
+        with ZipFile(archive_path) as zf:
+            ext_tid = ext_progress.add_task("Extracting", aname=archive_path.name, total=len(members := zf.infolist()))
+            extract_path = Path(zf.extract(members[0], download_dir))
+            if extract_type == "file":
+                assert len(members) == 1, f"Archive `{archive_path}` with extract type `{extract_type} has > 1 member!"
+            elif extract_type == "directory":
+                for member in members[1:]:
+                    zf.extract(member, download_dir)
+                    ext_progress.advance(ext_tid)
+            else:
+                raise ValueError(f"Extract type `{extract_type}` for archive `{archive_path}` is not defined!")
+
+    # Cleanup (if specified)
+    if cleanup:
+        archive_path.unlink()
+
+    return extract_path
+
+
+def download_extract(dataset_id: str, root_dir: Path) -> None:
+    """Download all files for a given dataset (querying registry above), extracting archives if necessary."""
+    os.makedirs(download_dir := root_dir / "download" / dataset_id, exist_ok=True)
+
+    # Download Files => Single-Threaded, with Progress Bar
+    dl_tasks = [d for d in DATASET_REGISTRY[dataset_id] if not (download_dir / d["name"]).exists()]
+    for dl_task in dl_tasks:
+        dl_path = download_with_progress(dl_task["url"], download_dir)
+
+        # Extract Files (if specified) --> Note (assumes ".zip" ONLY!)
+        if dl_task["extract"]:
+            dl_path = extract_with_progress(dl_path, download_dir, dl_task["extract_type"])
+            dl_path = dl_path.parent if dl_path.is_file() else dl_path
+
+        # Rename Path --> dl_task["name"]
+        if dl_task["do_rename"]:
+            shutil.move(dl_path, download_dir / dl_task["name"])

capvector-oft/prismatic/preprocessing/materialize.py

@@ -0,0 +1,69 @@
+"""
+materialize.py
+
+Factory class for initializing pretraining datasets on a per-VLM basis; provides and exports individual functions for
+clear control flow.
+"""
+
+from typing import Tuple, Type
+
+from torch.utils.data import Dataset
+from transformers import PreTrainedTokenizerBase
+
+from prismatic.conf import DatasetConfig
+from prismatic.models.backbones.llm.prompting import PromptBuilder
+from prismatic.models.backbones.vision import ImageTransform
+from prismatic.preprocessing.datasets import AlignDataset, FinetuneDataset
+from prismatic.util.data_utils import PaddedCollatorForLanguageModeling
+
+# Dataset Initializers =>> Maps Stage --> cls()
+DATASET_INITIALIZER = {"align": AlignDataset, "finetune": FinetuneDataset, "full-finetune": FinetuneDataset}
+
+
+def get_dataset_and_collator(
+    stage: str,
+    dataset_cfg: DatasetConfig,
+    image_transform: ImageTransform,
+    tokenizer: PreTrainedTokenizerBase,
+    prompt_builder_fn: Type[PromptBuilder],
+    default_image_resolution: Tuple[int, int, int],
+    padding_side: str = "right",
+) -> Tuple[Dataset, PaddedCollatorForLanguageModeling]:
+    dataset_cls = DATASET_INITIALIZER[stage]
+    dataset_root_dir = dataset_cfg.dataset_root_dir
+    collator = PaddedCollatorForLanguageModeling(
+        tokenizer.model_max_length, tokenizer.pad_token_id, default_image_resolution, padding_side=padding_side
+    )
+
+    # Switch on `stage`
+    if stage == "align":
+        annotation_json, image_dir = dataset_cfg.align_stage_components
+        dataset = dataset_cls(
+            dataset_root_dir / annotation_json, dataset_root_dir / image_dir, image_transform, tokenizer
+        )
+        return dataset, collator
+
+    elif stage == "finetune":
+        annotation_json, image_dir = dataset_cfg.finetune_stage_components
+        dataset = dataset_cls(
+            dataset_root_dir / annotation_json,
+            dataset_root_dir / image_dir,
+            image_transform,
+            tokenizer,
+            prompt_builder_fn=prompt_builder_fn,
+        )
+        return dataset, collator
+
+    elif stage == "full-finetune":
+        annotation_json, image_dir = dataset_cfg.finetune_stage_components
+        dataset = dataset_cls(
+            dataset_root_dir / annotation_json,
+            dataset_root_dir / image_dir,
+            image_transform,
+            tokenizer,
+            prompt_builder_fn=prompt_builder_fn,
+        )
+        return dataset, collator
+
+    else:
+        raise ValueError(f"Stage `{stage}` is not supported!")

capvector-oft/prismatic/training/__init__.py

@@ -0,0 +1,2 @@
+from .materialize import get_train_strategy
+from .metrics import Metrics, VLAMetrics

capvector-oft/prismatic/training/materialize.py

@@ -0,0 +1,66 @@
+"""
+materialize.py
+
+Factory class defining functions for instantiating various Training Strategies, supporting different VLMs, backbones,
+and strategy configurations.
+"""
+
+from typing import Callable, Optional
+
+import torch
+
+from prismatic.models.vlms import PrismaticVLM
+from prismatic.training.strategies import FSDPStrategy, TrainingStrategy
+
+# Registry =>> Maps ID --> {cls(), kwargs} :: supports FSDP for now, but DDP handler is also implemented!
+TRAIN_STRATEGIES = {
+    "fsdp-shard-grad-op": {"cls": FSDPStrategy, "kwargs": {"sharding_strategy": "shard-grad-op"}},
+    "fsdp-full-shard": {"cls": FSDPStrategy, "kwargs": {"sharding_strategy": "full-shard"}},
+}
+
+
+def get_train_strategy(
+    train_strategy: str,
+    vlm: PrismaticVLM,
+    device_id: int,
+    stage: str,
+    epochs: int,
+    max_steps: Optional[int],
+    global_batch_size: int,
+    per_device_batch_size: int,
+    learning_rate: float,
+    weight_decay: float,
+    max_grad_norm: float,
+    lr_scheduler_type: str,
+    warmup_ratio: float,
+    enable_gradient_checkpointing: bool = True,
+    enable_mixed_precision_training: bool = True,
+    reduce_in_full_precision: bool = False,
+    mixed_precision_dtype: torch.dtype = torch.bfloat16,
+    worker_init_fn: Optional[Callable[[int], None]] = None,
+) -> TrainingStrategy:
+    if train_strategy in TRAIN_STRATEGIES:
+        strategy_cfg = TRAIN_STRATEGIES[train_strategy]
+        strategy = strategy_cfg["cls"](
+            vlm=vlm,
+            device_id=device_id,
+            stage=stage,
+            epochs=epochs,
+            max_steps=max_steps,
+            global_batch_size=global_batch_size,
+            per_device_batch_size=per_device_batch_size,
+            learning_rate=learning_rate,
+            weight_decay=weight_decay,
+            max_grad_norm=max_grad_norm,
+            lr_scheduler_type=lr_scheduler_type,
+            warmup_ratio=warmup_ratio,
+            enable_gradient_checkpointing=enable_gradient_checkpointing,
+            enable_mixed_precision_training=enable_mixed_precision_training,
+            reduce_in_full_precision=reduce_in_full_precision,
+            mixed_precision_dtype=mixed_precision_dtype,
+            worker_init_fn=worker_init_fn,
+            **strategy_cfg["kwargs"],
+        )
+        return strategy
+    else:
+        raise ValueError(f"Train Strategy `{train_strategy}` is not supported!")

capvector-oft/prismatic/training/metrics.py

@@ -0,0 +1,348 @@
+"""
+metrics.py
+
+Utility classes defining a Metrics container and multiple Trackers to enable model/stage-specific logging to various
+endpoints (e.g., JSONL local logs, Weights & Biases).
+"""
+
+import time
+from collections import defaultdict, deque
+from pathlib import Path
+from typing import Any, Dict, Optional, Protocol, Tuple, Union
+
+import jsonlines
+import numpy as np
+import torch
+import wandb
+
+from prismatic.overwatch import initialize_overwatch
+
+# Initialize Overwatch =>> Wraps `logging.Logger`
+overwatch = initialize_overwatch(__name__)
+
+
+# === Define Tracker Interface ===
+class Tracker(Protocol):
+    def write_hyperparameters(self) -> None: ...
+
+    def write(self, global_step: int, metrics: Dict[str, Union[int, float]]) -> None: ...
+
+    def finalize(self) -> None: ...
+
+
+# === Individual Tracker Definitions ===
+class JSONLinesTracker:
+    def __init__(self, run_id: str, run_dir: Path, hparams: Dict[str, Any]) -> None:
+        self.run_id, self.run_dir, self.hparams = run_id, run_dir, hparams
+
+    @overwatch.rank_zero_only
+    def write_hyperparameters(self) -> None:
+        with jsonlines.open(self.run_dir / "run-metrics.jsonl", mode="w", sort_keys=True) as js_tracker:
+            js_tracker.write({"run_id": self.run_id, "hparams": self.hparams})
+
+    @overwatch.rank_zero_only
+    def write(self, _: int, metrics: Dict[str, Union[int, float]]) -> None:
+        with jsonlines.open(self.run_dir / f"{self.run_id}.jsonl", mode="a", sort_keys=True) as js_tracker:
+            js_tracker.write(metrics)
+
+    def finalize(self) -> None:
+        return
+
+
+class WeightsBiasesTracker:
+    def __init__(
+        self,
+        run_id: str,
+        run_dir: Path,
+        hparams: Dict[str, Any],
+        project: str = "prismatic",
+        entity: Optional[str] = None,
+        group: str = "align",
+    ) -> None:
+        self.run_id, self.run_dir, self.hparams = run_id, run_dir, hparams
+
+        # Get W&B-Specific Initialization Parameters
+        self.project, self.entity, self.group, self.wandb_dir = project, entity, group, self.run_dir
+
+        # Call W&B.init()
+        self.initialize()
+
+    @overwatch.rank_zero_only
+    def initialize(self) -> None:
+        wandb.init(
+            name=self.run_id,
+            dir=self.wandb_dir,
+            config=self.hparams,
+            project=self.project,
+            entity=self.entity,
+            group=self.group,
+        )
+
+    @overwatch.rank_zero_only
+    def write_hyperparameters(self) -> None:
+        wandb.config = self.hparams
+
+    @overwatch.rank_zero_only
+    def write(self, global_step: int, metrics: Dict[str, Union[int, float]]) -> None:
+        wandb.log(metrics, step=global_step)
+
+    @staticmethod
+    def finalize() -> None:
+        if overwatch.is_rank_zero():
+            wandb.finish()
+
+        # A job gets 210 seconds to get its affairs in order
+        time.sleep(210)
+
+
+# === Core Metrics Container :: Initializes Trackers => Compiles/Pushes Metrics ===
+
+
+class Metrics:
+    def __init__(
+        self,
+        active_trackers: Tuple[str, ...],
+        run_id: str,
+        run_dir: Path,
+        hparams: Dict[str, Any],
+        stage: str,
+        wandb_project: str = "prismatic",
+        wandb_entity: Optional[str] = None,
+        grad_accumulation_steps: int = 1,
+        window_size: int = 128,
+    ) -> None:
+        self.run_id, self.run_dir, self.hparams, self.stage = run_id, run_dir, hparams, stage
+
+        # Initialize Trackers
+        self.trackers = []
+        for tracker_type in active_trackers:
+            if tracker_type == "jsonl":
+                tracker = JSONLinesTracker(run_id, run_dir, hparams)
+            elif tracker_type == "wandb":
+                tracker = WeightsBiasesTracker(
+                    run_id, run_dir, hparams, project=wandb_project, entity=wandb_entity, group=self.stage
+                )
+            else:
+                raise ValueError(f"Tracker with type `{tracker_type} is not supported!")
+
+            # Add Hyperparameters --> add to `self.trackers`
+            tracker.write_hyperparameters()
+            self.trackers.append(tracker)
+
+        # Create Universal Metrics Buffers
+        self.global_step, self.start_time, self.step_start_time = 0, time.time(), time.time()
+        self.state = {
+            "loss_raw": deque(maxlen=grad_accumulation_steps),
+            "loss": deque(maxlen=window_size),
+            "step_time": deque(maxlen=window_size),
+            "lr": [],
+        }
+
+    def log(self, global_step: int, metrics: Dict[str, Union[int, float]]) -> None:
+        for tracker in self.trackers:
+            tracker.write(global_step, metrics)
+
+    def get_status(self, loss: Optional[torch.Tensor] = None) -> str:
+        lr = self.state["lr"][-1] if len(self.state["lr"]) > 0 else 0
+        if loss is None:
+            return f"=>> [Global Step] {self.global_step:06d} =>> LR :: {lr:.6f}"
+
+        # Otherwise, embed `loss` in status report!
+        return f"=>> [Global Step] {self.global_step:06d} =>> LR :: {lr:.6f} -- Loss :: {loss:.4f}"
+
+    def commit(
+        self, *, global_step: Optional[int] = None, lr: Optional[float] = None, update_step_time: bool = False, **kwargs
+    ) -> None:
+        """Update all metrics in `self.state` by iterating through special positional arguments & kwargs."""
+        if global_step is not None:
+            self.global_step = global_step
+
+        # For all other variables --> only track on rank zero!
+        if not overwatch.is_rank_zero():
+            return
+
+        # Special Positional Arguments
+        if lr is not None:
+            self.state["lr"].append(lr)
+
+        if update_step_time:
+            self.state["step_time"].append(time.time() - self.step_start_time)
+            self.step_start_time = time.time()
+
+        # Generic Keyword Arguments
+        for key, value in kwargs.items():
+            if key == "loss":
+                loss_val = value.detach()
+                self.state["loss_raw"].append(loss_val)
+                self.state["loss"].append(loss_val)
+            else:
+                self.state[key].append(value.detach())
+
+    @overwatch.rank_zero_only
+    def push(self) -> str:
+        # Note :: Raw Loss is an Average over Gradient Accumulation Steps --> No Smoothing!
+        loss_raw = torch.stack(list(self.state["loss_raw"])).mean().item()
+        loss = torch.stack(list(self.state["loss"])).mean().item()
+        step_time, lr = np.mean(list(self.state["step_time"])), self.state["lr"][-1]
+        status = self.get_status(loss)
+
+        # Fire to Trackers
+        prefix = self.stage.capitalize()
+        self.log(
+            self.global_step,
+            metrics={
+                f"{prefix}/Step": self.global_step,
+                f"{prefix}/Loss": loss,
+                f"{prefix}/Loss (Raw)": loss_raw,
+                f"{prefix}/Learning Rate": lr,
+                f"{prefix}/Step Time": step_time,
+            },
+        )
+        return status
+
+    def finalize(self) -> str:
+        for tracker in self.trackers:
+            tracker.finalize()
+
+
+class VLAMetrics:
+    def __init__(
+        self,
+        active_trackers: Tuple[str, ...],
+        run_id: str,
+        run_dir: Path,
+        hparams: Dict[str, Any],
+        wandb_project: str = "openvla",
+        wandb_entity: Optional[str] = "stanford-voltron",
+        grad_accumulation_steps: int = 1,
+        window_size: int = 1,
+        resume_step: Optional[int] = None,
+        resume_epoch: Optional[int] = None,
+    ) -> None:
+        self.run_id, self.run_dir, self.hparams = run_id, run_dir, hparams
+
+        # Initialize Trackers
+        self.trackers = []
+        for tracker_type in active_trackers:
+            if tracker_type == "jsonl":
+                tracker = JSONLinesTracker(run_id, run_dir, hparams)
+            elif tracker_type == "wandb":
+                tracker = WeightsBiasesTracker(
+                    run_id, run_dir, hparams, project=wandb_project, entity=wandb_entity, group="vla-train"
+                )
+            else:
+                raise ValueError(f"Tracker with type `{tracker_type} is not supported!")
+
+            # Add Hyperparameters --> add to `self.trackers`
+            tracker.write_hyperparameters()
+            self.trackers.append(tracker)
+
+        # Create Universal Metrics Buffers
+        self.global_step = 0 if resume_step is None else resume_step
+        self.epoch = 0 if resume_epoch is None else resume_epoch
+        self.start_time, self.step_start_time = time.time(), time.time()
+        self.state = {
+            "loss_raw": deque(maxlen=grad_accumulation_steps),
+            "loss": deque(maxlen=window_size),
+            "l1_loss": deque(maxlen=window_size),
+            "action_accuracy": deque(maxlen=window_size),
+            "step_time": deque(maxlen=window_size),
+            "lr": [],
+        }
+
+        # Created metrics buffers for individual tracked datasets
+        self.dataset_trackers = defaultdict(lambda: VLAMetrics([], "", "", {}))
+
+    def log(self, global_step: int, metrics: Dict[str, Union[int, float]]) -> None:
+        for tracker in self.trackers:
+            tracker.write(global_step, metrics)
+
+    def get_status(self, loss: Optional[torch.Tensor] = None) -> str:
+        lr = self.state["lr"][-1] if len(self.state["lr"]) > 0 else 0
+        if loss is None:
+            return f"=>> [Epoch {self.epoch:03d}] Global Step {self.global_step:06d} =>> LR :: {lr:.6f}"
+
+        # Otherwise, embed `loss` in status report!
+        return f"=>> [Epoch {self.epoch:03d}] Global Step {self.global_step:06d} =>> LR :: {lr:.6f} - Loss :: {loss:.4f}"
+
+    def commit(
+        self,
+        *,
+        global_step: Optional[int] = None,
+        epoch: Optional[int] = None,
+        lr: Optional[float] = None,
+        update_step_time: bool = False,
+        **kwargs,
+    ) -> None:
+        """Update all metrics in `self.state` by iterating through special positional arguments & kwargs."""
+        if global_step is not None:
+            self.global_step = global_step
+
+        if epoch is not None:
+            self.epoch = epoch
+
+        # For all other variables --> only track on rank zero!
+        if not overwatch.is_rank_zero():
+            return
+
+        # Special Positional Arguments
+        if lr is not None:
+            self.state["lr"].append(lr)
+
+        if update_step_time:
+            self.state["step_time"].append(time.time() - self.step_start_time)
+            self.step_start_time = time.time()
+
+        # Generic Keyword Arguments
+        for key, value in kwargs.items():
+            if key == "loss":
+                loss_val = value.detach()
+                self.state["loss_raw"].append(loss_val)
+                self.state["loss"].append(loss_val)
+            else:
+                self.state[key].append(value.detach())
+
+    def commit_for_dataset(self, dataset_name: str, **kwargs) -> None:
+        self.dataset_trackers[dataset_name].commit(**kwargs)
+
+    @overwatch.rank_zero_only
+    def push(self) -> str:
+        # Note :: Raw Loss is an Average over Gradient Accumulation Steps --> No Smoothing!
+        loss_raw = torch.stack(list(self.state["loss_raw"])).mean().item()
+        loss = torch.stack(list(self.state["loss"])).mean().item()
+        l1_loss = torch.stack(list(self.state["l1_loss"])).mean().item()
+        action_accuracy = torch.stack(list(self.state["action_accuracy"])).mean().item()
+        step_time, lr = np.mean(list(self.state["step_time"])), self.state["lr"][-1]
+        status = self.get_status(loss)
+
+        # Get metrics per dataset
+        dataset_metrics = {}
+        for ds, tracker in self.dataset_trackers.items():
+            dataset_metrics.update(
+                {
+                    f"{ds}/L1 Loss": torch.stack(list(tracker.state["l1_loss"])).mean().item(),
+                    f"{ds}/Action Token Accuracy": torch.stack(list(tracker.state["action_accuracy"])).mean().item(),
+                }
+            )
+
+        # Fire to Trackers
+        prefix = "VLA Train"
+        self.log(
+            self.global_step,
+            metrics={
+                f"{prefix}/Step": self.global_step,
+                f"{prefix}/Epoch": self.epoch,
+                f"{prefix}/Loss": loss,
+                f"{prefix}/L1 Loss": l1_loss,
+                f"{prefix}/Action Token Accuracy": action_accuracy,
+                f"{prefix}/Loss (Raw)": loss_raw,
+                f"{prefix}/Learning Rate": lr,
+                f"{prefix}/Step Time": step_time,
+                **dataset_metrics,
+            },
+        )
+        return status
+
+    def finalize(self) -> str:
+        for tracker in self.trackers:
+            tracker.finalize()

capvector-oft/prismatic/training/strategies/__init__.py

@@ -0,0 +1,3 @@
+from .base_strategy import TrainingStrategy
+from .ddp import DDPStrategy
+from .fsdp import FSDPStrategy

capvector-oft/prismatic/training/strategies/base_strategy.py

@@ -0,0 +1,417 @@
+"""
+base_strategy.py
+
+Abstract class definition of a (distributed) training strategy, with full annotations of class methods, utility
+functions, and initialization logic.
+
+Training Strategies (DDP, FSDP-Grad, FSDP-Full) tend to have a lot of repeated components; this class does a lot of
+heavy lifting.
+"""
+
+from abc import ABC, abstractmethod
+from pathlib import Path
+from typing import Callable, Optional
+
+import numpy as np
+import torch
+import torch.distributed as dist
+from torch.utils.data import DataLoader, Dataset, DistributedSampler, IterableDataset
+from tqdm import tqdm
+from transformers.modeling_outputs import CausalLMOutputWithPast
+
+from prismatic.models.vlms import PrismaticVLM
+from prismatic.overwatch import initialize_overwatch
+from prismatic.training.metrics import Metrics, VLAMetrics
+from prismatic.training.train_utils import (
+    compute_actions_l1_loss,
+    compute_token_accuracy,
+    get_current_action_mask,
+    get_next_actions_mask,
+)
+from prismatic.util import check_bloat16_supported
+from prismatic.util.batching_utils import SplitModalitySampler
+from prismatic.util.data_utils import PaddedCollatorForActionPrediction, PaddedCollatorForLanguageModeling
+from prismatic.vla.action_tokenizer import ActionTokenizer
+
+# HuggingFace Default / LLaMa-2 IGNORE_INDEX (for labels)
+from prismatic.vla.constants import ACTION_DIM, ACTION_TOKEN_BEGIN_IDX, NUM_ACTIONS_CHUNK, IGNORE_INDEX
+NEWLINE_INDEX = 13  # '\n'
+STOP_INDEX = 2  # '</s>'
+
+# Initialize Overwatch =>> Wraps `logging.Logger`
+overwatch = initialize_overwatch(__name__)
+
+
+# === Abstract Base Class for an arbitrary Training Strategy ===
+class TrainingStrategy(ABC):
+    def __init__(
+        self,
+        vlm: PrismaticVLM,
+        device_id: int,
+        stage: str,
+        epochs: int,
+        max_steps: Optional[int],
+        global_batch_size: int,
+        per_device_batch_size: int,
+        learning_rate: float,
+        weight_decay: float,
+        max_grad_norm: float,
+        lr_scheduler_type: str,
+        warmup_ratio: float,
+        enable_gradient_checkpointing: bool = True,
+        enable_mixed_precision_training: bool = True,
+        reduce_in_full_precision: bool = False,
+        mixed_precision_dtype: torch.dtype = torch.bfloat16,
+        worker_init_fn: Optional[Callable[[int], None]] = None,
+        **_: str,
+    ) -> None:
+        self.vlm, self.device_id, self.stage = vlm, device_id, stage
+
+        # Get relevant VLM instance parameters before they get (potentially) wrapped
+        self.all_module_keys, self.trainable_module_keys = self.vlm.all_module_keys, self.vlm.trainable_module_keys
+        self.llm_transformer_layer_cls = self.vlm.llm_backbone.transformer_layer_cls
+
+        # Optimization Parameters
+        self.epochs, self.max_steps = epochs, max_steps
+        self.global_batch_size, self.per_device_batch_size = global_batch_size, per_device_batch_size
+
+        self.learning_rate, self.weight_decay, self.max_grad_norm = learning_rate, weight_decay, max_grad_norm
+        self.lr_scheduler_type, self.warmup_ratio = lr_scheduler_type, warmup_ratio
+
+        # Generic Strategy Parameters
+        self.enable_gradient_checkpointing = enable_gradient_checkpointing
+        self.enable_mixed_precision_training = enable_mixed_precision_training
+        self.reduce_in_full_precision = reduce_in_full_precision
+        self.mixed_precision_dtype = mixed_precision_dtype
+
+        # DataLoader Parameters
+        self.worker_init_fn = worker_init_fn
+
+        # Optimizers & Scheduler (initialized in `run_setup`)
+        self.optimizer, self.lr_scheduler = None, None
+
+        # Lightweight Validation
+        assert (
+            self.global_batch_size % self.per_device_batch_size == 0
+        ), "Per-device batch size must evenly divide global batch size!"
+        self.grad_accumulation_steps = self.global_batch_size // self.per_device_batch_size // overwatch.world_size()
+        if self.enable_mixed_precision_training:
+            assert self.mixed_precision_dtype == torch.bfloat16, "Only BF16 mixed precision training is supported!"
+            assert check_bloat16_supported(), "BFloat16 is not supported on this hardware; unset `mixed_precision`"
+
+    @abstractmethod
+    def save_checkpoint(
+        self,
+        run_dir: Path,
+        global_step: int,
+        epoch: int,
+        train_loss: Optional[float] = None,
+        only_trainable: bool = True,
+    ) -> None: ...
+
+    @abstractmethod
+    def run_setup(self, run_dir: Path, n_train_examples: int) -> None: ...
+
+    @abstractmethod
+    def clip_grad_norm(self) -> None: ...
+
+    def run_training(
+        self,
+        dataset: Dataset,
+        collator: PaddedCollatorForLanguageModeling,
+        metrics: Metrics,
+        stage: str = "finetune",
+        batch_construction_strategy: str = "split-modality",
+        seed: int = 7,
+    ) -> None:
+        """Run the training loop for the given `dataset` and `collator`; log losses, results to `metrics`"""
+        if "finetune" in stage and batch_construction_strategy == "split-modality":
+            # Instantiate the split-modality sampler; if you want to extend with other batch construction schemes,
+            #   (e.g., grouping by length) =>> can easily add them here!
+            modality_lengths = dataset.get_modality_lengths()
+            sampler = SplitModalitySampler(
+                dataset,
+                modality_lengths,
+                global_batch_size=self.global_batch_size,
+                num_replicas=overwatch.world_size(),
+                rank=overwatch.rank(),
+                seed=seed,
+                drop_last=False,
+            )
+
+        else:
+            sampler = DistributedSampler(
+                dataset,
+                num_replicas=overwatch.world_size(),
+                rank=overwatch.rank(),
+                shuffle=True,
+                seed=seed,
+                drop_last=False,
+            )
+
+        # Create a DataLoader with the initialized sampler, per-device-bsz, and collator
+        dataloader = DataLoader(
+            dataset,
+            batch_size=self.per_device_batch_size,
+            sampler=sampler,
+            collate_fn=collator,
+            num_workers=2,
+            worker_init_fn=self.worker_init_fn,
+        )
+
+        # Max Steps vs. Epochs Computation
+        steps_per_epoch = len(dataloader) // self.grad_accumulation_steps
+        if self.max_steps is not None and steps_per_epoch < self.max_steps:
+            # Just set `epochs` to some large number --> we'll short-circuit based on steps anyway
+            self.epochs = 100
+
+        # === Train ===
+        status = metrics.get_status()
+        with tqdm(
+            total=(
+                (self.epochs * (len(dataloader) // self.grad_accumulation_steps))
+                if self.max_steps is None
+                else self.max_steps
+            ),
+            desc=status,
+            leave=False,
+            disable=not overwatch.is_rank_zero(),
+        ) as progress:
+            for epoch in range(self.epochs):
+                self.vlm.train()
+                sampler.set_epoch(epoch)
+
+                # Zero-Gradients (just in case)
+                self.optimizer.zero_grad()
+
+                # Note that we'll unpack batch (and let AMP/FSDP do its thing) in the VLM.forward() call
+                #   => Basically, if we're using mixed precision (or not), autocast()/FSDP will move to device!
+                for train_idx, batch in enumerate(dataloader):
+                    # [Contract] self.vlm.forward() must automatically compute `loss` and return!
+                    with torch.autocast(
+                        "cuda",
+                        dtype=self.mixed_precision_dtype,
+                        enabled=self.enable_mixed_precision_training,
+                    ):
+                        output: CausalLMOutputWithPast = self.vlm(
+                            input_ids=batch["input_ids"],
+                            attention_mask=batch["attention_mask"],
+                            pixel_values=batch["pixel_values"],
+                            labels=batch["labels"],
+                            multimodal_indices=batch["multimodal_indices"],
+                        )
+                        loss = output.loss
+
+                    # Commit Loss (Prior to Gradient Accumulation Normalization)
+                    metrics.commit(loss=loss)
+
+                    # Normalize Loss to account for Gradient Accumulation --> Backward!
+                    # [IMPORTANT] Technically speaking, doing gradient accumulation in this way is "incorrect"; this is
+                    #             because in general, each batch has a *different number of masked out tokens* (because
+                    #             we're instruct-tuning). Taking the mean over two unbalanced means != the right thing!
+                    #
+                    #             HOWEVER -- at least at the 7B scale, the "naive" approach is just as performant as
+                    #             the "correct" implementation, without adding extra complexity.
+                    #
+                    # That being said =>> at the 13B scale, *no matter what we tried, ANY gradient accumulation is just
+                    #   really bad for downstream performance. Initial investigation shows that BF16 accumulation
+                    #   just really tanks in precision... and don't have a good/clean way to fix this. Would love for
+                    #   someone to PR and fix this (and I'd greatly appreciate it!!!)
+                    normalized_loss = loss / self.grad_accumulation_steps
+                    normalized_loss.backward()
+
+                    # Step =>> Only if Done w/ Gradient Accumulation
+                    if (train_idx + 1) % self.grad_accumulation_steps == 0:
+                        metrics.commit(update_step_time=True)
+
+                        # Clip Gradients --> this is custom, per-strategy because of DDP vs. FSDP locality-assumptions
+                        self.clip_grad_norm()
+
+                        # Optimizer & LR Scheduler Step
+                        self.optimizer.step()
+                        self.lr_scheduler.step()
+                        self.optimizer.zero_grad()
+
+                        # Push Metrics
+                        metrics.commit(global_step=metrics.global_step + 1, lr=self.lr_scheduler.get_last_lr()[0])
+                        status = metrics.push()
+
+                        # Check for Termination & Save Final Checkpoint (in case `max_steps` is not None)
+                        if self.max_steps is not None and metrics.global_step >= self.max_steps:
+                            self.save_checkpoint(metrics.run_dir, metrics.global_step, epoch, loss.item())
+                            dist.barrier()
+
+                            return
+
+                        # Update Progress Bar
+                        progress.update()
+                        progress.set_description(status)
+
+            # Save checkpoint at end each epoch (if `self.max_steps` is None)
+            if self.max_steps is None:
+                self.save_checkpoint(metrics.run_dir, metrics.global_step, epoch, loss.item())
+                dist.barrier()
+
+    # === VLA Training ===
+
+    def run_vla_training(
+        self,
+        vla_dataset: IterableDataset,
+        collator: PaddedCollatorForActionPrediction,
+        action_tokenizer: ActionTokenizer,
+        metrics: VLAMetrics,
+        save_interval: int = 2500,
+        save_full_model: bool = True,
+    ) -> None:
+        """Run the VLA training loop for the given `dataset` and `collator`; log losses, action metrics to `metrics`."""
+        assert isinstance(vla_dataset, IterableDataset), "VLA training expects an IterableDataset!"
+        assert self.grad_accumulation_steps == 1, "VLA training does not support gradient accumulation!"
+
+        # Create a DataLoader =>> Set `num_workers` to 0; RLDS loader handles parallelism!
+        dataloader = DataLoader(
+            vla_dataset,
+            batch_size=self.per_device_batch_size,
+            sampler=None,
+            collate_fn=collator,
+            num_workers=0,
+            worker_init_fn=self.worker_init_fn,
+        )
+
+        # === Train ===
+        status = metrics.get_status()
+        with tqdm(
+            total=(self.epochs * len(dataloader)) if self.max_steps is None else self.max_steps,
+            desc=status,
+            leave=False,
+            disable=not overwatch.is_rank_zero(),
+        ) as progress:
+            self.vlm.train()
+
+            # Zero Gradients (just in case)
+            self.optimizer.zero_grad()
+
+            # [Contract] DataLoader wraps RLDS Loader (`.as_numpy_iterator() =>> implicit `.repeat()`)
+            #   => This means looping over the DataLoader is basically "infinite" (so no outer loop over epochs).
+            #      Slightly breaks default PyTorch semantics, which is why we adaptively compute `epoch` below.
+            for batch in dataloader:
+                # Note that we'll unpack batch (and let AMP/FSDP do its thing) in the VLM.forward() call
+                #   => Basically, if we're using mixed precision (or not), autocast()/FSDP will move to device!
+                with torch.autocast(
+                    "cuda", dtype=self.mixed_precision_dtype, enabled=self.enable_mixed_precision_training
+                ):
+                    # [Contract] self.vlm.forward() must automatically compute `loss` and return!
+                    output: CausalLMOutputWithPast = self.vlm(
+                        input_ids=batch["input_ids"],
+                        attention_mask=batch["attention_mask"],
+                        pixel_values=batch["pixel_values"],
+                        labels=batch["labels"],
+                    )
+                    loss = output.loss
+
+                # Commit Loss =>> Backward!
+                metrics.commit(loss=loss)
+                loss.backward()
+
+                # Get predicted and ground-truth token IDs
+                predicted_token_ids = output.logits[:, self.vlm.vision_backbone.num_patches : -1].argmax(dim=2)
+                ground_truth_token_ids = batch["labels"][:, 1:].to(predicted_token_ids.device)
+
+                #######################################################################
+                # === Compute Current Action Token Accuracy & L1 Loss ===
+                #######################################################################
+
+                # Get current action mask: Target the first ACTION_DIM non-ignore tokens
+                current_action_mask = get_current_action_mask(ground_truth_token_ids)
+
+                # Compute Accuracy
+                action_accuracy = compute_token_accuracy(predicted_token_ids, ground_truth_token_ids, mask=current_action_mask)
+
+                # Compute L1 Loss on Predicted (Continuous) Actions
+                action_l1_loss = compute_actions_l1_loss(action_tokenizer, predicted_token_ids, ground_truth_token_ids, mask=current_action_mask)
+
+                #######################################################################
+                # === Compute Next Actions Token Accuracy & L1 Loss ===
+                #######################################################################
+
+                # Get next actions mask: Target all tokens after the first ACTION_DIM non-ignore tokens (excluding the last token, which is the stop token)
+                next_actions_mask = get_next_actions_mask(ground_truth_token_ids)
+
+                # Compute Accuracy
+                next_actions_accuracy = compute_token_accuracy(predicted_token_ids, ground_truth_token_ids, mask=next_actions_mask)
+
+                # Compute L1 Loss on Predicted (Continuous) Actions
+                next_actions_l1_loss = compute_actions_l1_loss(action_tokenizer, predicted_token_ids, ground_truth_token_ids, mask=next_actions_mask)
+
+                #######################################################################
+                # === Log ===
+                #######################################################################
+
+                # Commit Metrics
+                metrics.commit(
+                    action_accuracy=action_accuracy,
+                    l1_loss=action_l1_loss,
+                    next_actions_accuracy=next_actions_accuracy,
+                    next_actions_l1_loss=next_actions_l1_loss,
+                    update_step_time=True,
+                )
+
+                # Compute metrics per dataset --> only on rank_zero since we don't log them on other workers anyways
+                if overwatch.is_rank_zero():
+                    datasets = set(batch["dataset_names"])
+                    if len(datasets) > 1:
+                        for ds in datasets:
+                            ds_mask = torch.tensor([elem == ds for elem in batch["dataset_names"]])
+                            action_accuracy_ds = correct_preds[ds_mask].sum().float() / mask[ds_mask].sum().float()
+                            pred_continuous_actions_ds = torch.tensor(
+                                action_tokenizer.decode_token_ids_to_actions(
+                                    predicted_token_ids[ds_mask][mask[ds_mask]].cpu().numpy()
+                                )
+                            )
+                            continuous_actions_gt_ds = torch.tensor(
+                                action_tokenizer.decode_token_ids_to_actions(
+                                    ground_truth_token_ids[ds_mask][mask[ds_mask]].cpu().numpy()
+                                )
+                            )
+                            action_l1_loss_ds = torch.nn.functional.l1_loss(
+                                pred_continuous_actions_ds, continuous_actions_gt_ds
+                            )
+                            metrics.commit_for_dataset(
+                                dataset_name=ds.decode(),
+                                action_accuracy=action_accuracy_ds,
+                                l1_loss=action_l1_loss_ds,
+                                next_actions_accuracy=next_actions_accuracy,
+                                next_actions_l1_loss=next_actions_l1_loss,
+                            )
+
+                # === Gradient Step ===
+
+                # Clip Gradients --> this is custom, per-strategy because of DDP vs. FSDP locality assumptions
+                self.clip_grad_norm()
+
+                # Optimizer & LR Scheduler Step
+                self.optimizer.step()
+                self.lr_scheduler.step()
+                self.optimizer.zero_grad()
+
+                # Compute epoch value using number of completed gradient steps
+                epoch = (metrics.global_step + 1) // (len(vla_dataset) // self.global_batch_size)
+
+                # Push Metrics
+                metrics.commit(global_step=metrics.global_step + 1, epoch=epoch, lr=self.lr_scheduler.get_last_lr()[0])
+                status = metrics.push()
+
+                # Check for Save Interval or Max Steps & Save Checkpoint
+                if (terminate := (self.max_steps is not None and metrics.global_step >= self.max_steps)) or (
+                    (metrics.global_step % save_interval) == 0
+                ):
+                    self.save_checkpoint(
+                        metrics.run_dir, metrics.global_step, epoch, loss.item(), only_trainable=not save_full_model
+                    )
+                    dist.barrier()
+
+                    if terminate:
+                        return
+
+                # Update Progress Bar
+                progress.update()
+                progress.set_description(status)

capvector-oft/prismatic/training/strategies/ddp.py

@@ -0,0 +1,128 @@
+"""
+ddp.py
+
+Core class definition for a strategy implementing Torch native Distributed Data Parallel Training; note that on most
+GPU hardware and LLM backbones >= 5-7B parameters, DDP training will OOM, which is why we opt for FSDP.
+"""
+
+import shutil
+from pathlib import Path
+from typing import Optional
+
+import torch
+from torch.nn.parallel import DistributedDataParallel as DDP
+from torch.optim import AdamW
+from transformers.optimization import get_constant_schedule, get_cosine_schedule_with_warmup
+
+from prismatic.overwatch import initialize_overwatch
+from prismatic.training.strategies.base_strategy import TrainingStrategy
+
+# Initialize Overwatch =>> Wraps `logging.Logger`
+overwatch = initialize_overwatch(__name__)
+
+
+class DDPStrategy(TrainingStrategy):
+    @overwatch.rank_zero_only
+    def save_checkpoint(
+        self,
+        run_dir: Path,
+        global_step: int,
+        epoch: int,
+        train_loss: Optional[float] = None,
+        only_trainable: bool = True,
+    ) -> None:
+        """Save a checkpoint to the `run_dir` only containing the state_dicts for trainable parameters by default."""
+        assert isinstance(self.vlm, DDP), "save_checkpoint assumes VLM is already wrapped in DDP!"
+
+        # Splinter State Dictionary by Top-Level Submodules (or subset, if `only_trainable`)
+        model_state_dicts = {
+            mkey: getattr(self.vlm.module, mkey).state_dict()
+            for mkey in (self.trainable_module_keys if only_trainable else self.all_module_keys)
+        }
+        optimizer_state_dict = self.optimizer.state_dict()
+
+        # Set Checkpoint Path =>> Embed *minimal* training statistics!
+        checkpoint_dir = run_dir / "checkpoints"
+        if train_loss is None:
+            checkpoint_path = checkpoint_dir / f"step-{global_step:06d}-epoch-{epoch:02d}-loss=inf.pt"
+        else:
+            checkpoint_path = checkpoint_dir / f"step-{global_step:06d}-epoch-{epoch:02d}-loss={train_loss:.4f}.pt"
+
+        # Save Checkpoint & Copy Latest to `latest-checkpoint.pt`
+        torch.save({"model": model_state_dicts, "optimizer": optimizer_state_dict}, checkpoint_path)
+        shutil.copy(checkpoint_path, checkpoint_dir / "latest-checkpoint.pt")
+
+    def run_setup(self, run_dir: Path, n_train_examples: int) -> None:
+        # Gradient Checkpointing Setup
+        if self.enable_gradient_checkpointing:
+            # For Gradient Checkpointing --> we make the assumption that the "bulk" of activation memory is taken up
+            #     by the LLM; because we also make the explicit assumption that each LLM is derived from a HF
+            #     pretrained model, the only thing we *need* to do (technically) is call `gradient_checkpoint_enable`
+            #     on `self.llm_backbone`.
+            #
+            # What does it actually do? --> runs the *generic* custom_forward + torch.utils.checkpoint.checkpoint logic
+            #   => github.com/huggingface/transformers/.../models/llama/modeling_llama.py#L692-L706
+            #
+            # Additional Reference (to better understand gradient checkpointing in PyTorch writ large)
+            #   => github.com/prigoyal/pytorch_memonger/blob/master/tutorial/Checkpointing_for_PyTorch_models.ipynb
+            overwatch.info("Enabling Gradient Checkpointing on LLM Backbone", ctx_level=1)
+            self.vlm.llm_backbone.gradient_checkpointing_enable()
+
+        # Move to Device =>> Note parameters are in full precision (*mixed precision* will only autocast as appropriate)
+        overwatch.info("Placing Entire VLM (Vision Backbone, LLM Backbone, Projector Weights) on GPU", ctx_level=1)
+        self.vlm.to(self.device_id)
+
+        # Wrap with Distributed Data Parallel
+        #   => Note: By default, wrapping naively with DDP(self.vlm) will initialize a *separate* buffer on GPU that
+        #            is the same size/dtype as the model parameters; this will *double* GPU memory!
+        # - stackoverflow.com/questions/68949954/model-takes-twice-the-memory-footprint-with-distributed-data-parallel
+        overwatch.info("Wrapping VLM with Distributed Data Parallel", ctx_level=1)
+        self.vlm = DDP(self.vlm, device_ids=[self.device_id], gradient_as_bucket_view=True)
+
+        # Create Optimizer and LR Scheduler =>> note that most of the LR Schedulers we use require `max_steps/epochs`
+        #   => Optimizer should only operate on parameters that are *unfrozen* / trainable!
+        trainable_params = [param for param in self.vlm.parameters() if param.requires_grad]
+        if self.max_steps is None:
+            num_training_steps = (n_train_examples * self.epochs) // self.global_batch_size
+        else:
+            num_training_steps = self.max_steps
+
+        if self.lr_scheduler_type == "linear-warmup+cosine-decay":
+            # Set warmup steps (floor) based on `warmup_ratio` (should be 0.03 - 0.05)
+            num_warmup_steps = int(num_training_steps * self.warmup_ratio)
+
+            assert self.weight_decay == 0, "DDP training does not currently support `weight_decay` > 0!"
+            self.optimizer = AdamW(trainable_params, lr=self.learning_rate, weight_decay=self.weight_decay)
+            self.lr_scheduler = get_cosine_schedule_with_warmup(self.optimizer, num_warmup_steps, num_training_steps)
+            for param_group in self.optimizer.param_groups:
+                param_group["lr"] = 0.0
+
+        elif self.lr_scheduler_type == "constant":
+            num_warmup_steps = 0
+
+            assert self.weight_decay == 0, "DDP training does not currently support `weight_decay` > 0!"
+            self.optimizer = AdamW(trainable_params, lr=self.learning_rate, weight_decay=self.weight_decay)
+            self.lr_scheduler = get_constant_schedule(self.optimizer)
+
+        else:
+            raise ValueError(f"Learning Rate Schedule with type `{self.lr_scheduler_type}` is not supported!")
+
+        # Finalize Setup =>> Log
+        overwatch.info(
+            "DDP Strategy =>> Finalized Training Setup:\n"
+            f"         |-> Global (Effective) Batch Size = {self.global_batch_size}\n"
+            f"         |-> Per-Device Batch Size = {self.per_device_batch_size}\n"
+            f"         |-> Distributed World Size = {overwatch.world_size()}\n"
+            f"         |-> Gradient Accumulation Steps = {self.grad_accumulation_steps}\n\n"
+            f"         |-> LLM Backbone Gradient Checkpointing = {self.enable_gradient_checkpointing}\n"
+            f"         |-> Use Native AMP = {self.enable_mixed_precision_training} ({self.mixed_precision_dtype})\n\n"
+            f"         |-> Default AdamW LR = {self.learning_rate}\n"
+            f"         |-> AdamW Weight Decay = {self.weight_decay}\n"
+            f"         |-> LR Scheduler Type = {self.lr_scheduler_type}\n"
+            f"         |-> LR Scheduler Warmup Steps (Ratio) = {num_warmup_steps} ({self.warmup_ratio})\n"
+            f"         |-> Dataset Size = {n_train_examples} Examples\n"
+            f"         |-> Max Steps = {num_training_steps}\n"
+        )
+
+    def clip_grad_norm(self) -> None:
+        torch.nn.utils.clip_grad_norm_(self.vlm.parameters(), max_norm=self.max_grad_norm)

capvector-oft/prismatic/training/strategies/fsdp.py

@@ -0,0 +1,270 @@
+"""
+fsdp.py
+
+Core class definition for a strategy implementing Torch native Fully Sharded Data Parallel Training (with support for
+fine-grained control over wrapping policies and mixed precision per component).
+"""
+
+import math
+from collections import OrderedDict
+from functools import partial
+from pathlib import Path
+from typing import Callable, Optional
+
+import torch
+import torch.distributed as dist
+import torch.nn as nn
+from torch.distributed.algorithms._checkpoint.checkpoint_wrapper import (
+    CheckpointImpl,
+    apply_activation_checkpointing,
+    checkpoint_wrapper,
+)
+from torch.distributed.fsdp import (
+    FullStateDictConfig,
+    MixedPrecision,
+    ShardingStrategy,
+    StateDictType,
+)
+from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
+from torch.optim import AdamW
+from transformers.optimization import get_constant_schedule, get_cosine_schedule_with_warmup
+
+from prismatic.models.vlms import PrismaticVLM
+from prismatic.overwatch import initialize_overwatch
+from prismatic.training.strategies.base_strategy import TrainingStrategy
+
+# Initialize Overwatch =>> Wraps `logging.Logger`
+overwatch = initialize_overwatch(__name__)
+
+
+class FSDPStrategy(TrainingStrategy):
+    def __init__(
+        self,
+        vlm: PrismaticVLM,
+        device_id: int,
+        stage: str,
+        epochs: int,
+        max_steps: Optional[int],
+        global_batch_size: int,
+        per_device_batch_size: int,
+        learning_rate: float,
+        weight_decay: float,
+        max_grad_norm: float,
+        lr_scheduler_type: str,
+        warmup_ratio: float,
+        enable_gradient_checkpointing: bool = True,
+        enable_mixed_precision_training: bool = True,
+        reduce_in_full_precision: bool = False,
+        mixed_precision_dtype: torch.dtype = torch.bfloat16,
+        worker_init_fn: Optional[Callable[[int], None]] = None,
+        sharding_strategy: str = "shard-grad-op",
+        state_dict_type: StateDictType = StateDictType.FULL_STATE_DICT,
+    ) -> None:
+        super().__init__(
+            vlm=vlm,
+            device_id=device_id,
+            stage=stage,
+            epochs=epochs,
+            max_steps=max_steps,
+            global_batch_size=global_batch_size,
+            per_device_batch_size=per_device_batch_size,
+            learning_rate=learning_rate,
+            weight_decay=weight_decay,
+            max_grad_norm=max_grad_norm,
+            lr_scheduler_type=lr_scheduler_type,
+            warmup_ratio=warmup_ratio,
+            enable_gradient_checkpointing=enable_gradient_checkpointing,
+            enable_mixed_precision_training=enable_mixed_precision_training,
+            reduce_in_full_precision=reduce_in_full_precision,
+            mixed_precision_dtype=mixed_precision_dtype,
+            worker_init_fn=worker_init_fn,
+        )
+
+        # FSDP-Specific Parameters
+        if sharding_strategy == "shard-grad-op":
+            self.fsdp_sharding_strategy = ShardingStrategy._HYBRID_SHARD_ZERO2
+        elif sharding_strategy == "full-shard":
+            self.fsdp_sharding_strategy = ShardingStrategy.HYBRID_SHARD
+        else:
+            raise ValueError(f"FSDP Sharding Strategy {sharding_strategy} is not supported!")
+
+        assert state_dict_type == StateDictType.FULL_STATE_DICT, "Sharded state saving is not yet implemented!"
+        self.fsdp_state_dict_type = state_dict_type
+        self.fsdp_save_policy = FullStateDictConfig(offload_to_cpu=True, rank0_only=True)
+
+    def save_checkpoint(
+        self,
+        run_dir: Path,
+        global_step: int,
+        epoch: int,
+        train_loss: Optional[float] = None,
+        only_trainable: bool = True,
+    ) -> None:
+        """Save a checkpoint to the `run_dir` only containing the state_dicts for trainable parameters by default."""
+        assert isinstance(self.vlm, FSDP), "FSDPStrategy.save_checkpoint assumes VLM is already wrapped in FSDP!"
+
+        # Summon Full State Dictionary =>> Reconstitute from Shards
+        with FSDP.state_dict_type(self.vlm, self.fsdp_state_dict_type, self.fsdp_save_policy):
+            full_vlm_state_dict = self.vlm.state_dict()
+            model_state_dicts = {
+                mkey: OrderedDict() for mkey in (self.trainable_module_keys if only_trainable else self.all_module_keys)
+            }
+
+            # Iterate through `full_vlm_state_dict` and split `mkey.{full_dotted_path}` -> `mkey: {full_dotted_path}`
+            for key, param in full_vlm_state_dict.items():
+                for mkey in model_state_dicts:
+                    if key.startswith(mprefix := f"{mkey}."):
+                        model_state_dicts[mkey][key.removeprefix(mprefix)] = param
+
+            # Save on rank zero *only*
+            if overwatch.is_rank_zero():
+                checkpoint_dir = run_dir / "checkpoints"
+                if train_loss is None:
+                    checkpoint_path = checkpoint_dir / f"step-{global_step:06d}-epoch-{epoch:02d}-loss=inf.pt"
+                else:
+                    checkpoint_path = (
+                        checkpoint_dir / f"step-{global_step:06d}-epoch-{epoch:02d}-loss={train_loss:.4f}.pt"
+                    )
+
+                # Save Checkpoint & Copy Latest to `latest-checkpoint.pt`
+                torch.save({"model": model_state_dicts}, checkpoint_path)
+
+                # TODO (siddk) :: This breaks w/ Sagemaker default permissions (root vs. <user>)... skip?
+                # shutil.copy(checkpoint_path, checkpoint_dir / "latest-checkpoint.pt")
+
+    def run_setup(self, run_dir: Path, n_train_examples: int) -> None:
+        # Iteratively Assemble FSDP Wrapping Policy by fetching the wrapping policies for each backbone/constituent
+        vlm_fsdp_wrapping_policy = self.vlm.get_fsdp_wrapping_policy()
+
+        # Assemble the Default FSDP Mixed Precision Policy
+        if self.enable_mixed_precision_training and self.mixed_precision_dtype == torch.bfloat16:
+            # MixedPrecision `param_dtype` specifies *compute* dtype (for forward/backward only)
+            #   => Reference: https://pytorch.org/docs/stable/fsdp.html#torch.distributed.fsdp.MixedPrecision
+            reduce_buffer_dtype = torch.bfloat16 if not self.reduce_in_full_precision else torch.float32
+            fsdp_precision_policy = MixedPrecision(
+                param_dtype=torch.bfloat16, reduce_dtype=reduce_buffer_dtype, buffer_dtype=reduce_buffer_dtype
+            )
+
+            # When running FSDP with a frozen vision backbone --> move to half precision!
+            if self.stage not in {"full-finetune", "vla-full-train", "vla-sandwich-train"}:
+                overwatch.info("Casting Vision Backbone to *Half Precision* via `.to(dtype=...)`")
+                self.vlm.vision_backbone.to(dtype=self.vlm.vision_backbone.half_precision_dtype)
+
+        else:
+            # If we're not using mixed precision, everything is in default full precision!
+            fsdp_precision_policy = MixedPrecision(
+                param_dtype=torch.float32, reduce_dtype=torch.float32, buffer_dtype=torch.float32
+            )
+
+        # <FSDP> => note that FSDP will automatically take care of device placement (similar to `autocast`)
+        self.vlm = FSDP(
+            self.vlm,
+            auto_wrap_policy=vlm_fsdp_wrapping_policy,
+            mixed_precision=fsdp_precision_policy,
+            sharding_strategy=self.fsdp_sharding_strategy,
+            device_id=torch.cuda.current_device(),
+            limit_all_gathers=True,
+            use_orig_params=True,
+        )
+
+        # Gradient Checkpoint Setup
+        if self.enable_gradient_checkpointing:
+            # For Gradient Checkpointing under FSDP --> we make the same assumption as in the DDP/other strategies; the
+            #   bulk of activation memory is taken up by the LLM activations. However, unlike other strategies, we
+            #   cannot rely on the HF Transformers default `gradient_checkpointing_enable()` --> FSDP breaks semantics!
+            #
+            # Instead, we need to write our own *NO-REENTRANT* wrapper, and apply it to the LLM's Transformer Layer.
+            non_reentrant_wrapper = partial(checkpoint_wrapper, checkpoint_impl=CheckpointImpl.NO_REENTRANT)
+
+            def check_fn(submodule: nn.Module) -> bool:
+                return isinstance(submodule, self.llm_transformer_layer_cls)
+
+            # Note that the terms "activation checkpointing" and "gradient checkpointing" are synonymous!
+            apply_activation_checkpointing(self.vlm, checkpoint_wrapper_fn=non_reentrant_wrapper, check_fn=check_fn)
+
+        # Barrier =>> Sharding takes a minute?
+        dist.barrier()
+
+        # Create Optimizer and LR Scheduler =>> note that most of the LR Schedulers we use require `max_steps/epochs`
+        #   => Optimizer should only operate on parameters that are *unfrozen* / trainable!
+        n_train_examples = math.ceil(n_train_examples / self.global_batch_size) * self.global_batch_size
+        if self.max_steps is None:
+            num_training_steps = (n_train_examples * self.epochs) // self.global_batch_size
+        else:
+            num_training_steps = self.max_steps
+
+        if self.lr_scheduler_type == "linear-warmup+cosine-decay":
+            # Set warmup steps (floor) based on `warmup_ratio` (should be 0.03 - 0.05)
+            num_warmup_steps = int(num_training_steps * self.warmup_ratio)
+
+            # Default AdamW w/ specified LR & Linear Warmup / Cosine Decay & Weight Decay
+            #   => Create Parameter Groups --> bias terms, normalization layer parameters shouldn't be decayed!
+            decay, no_decay = [], []
+            for name, param in self.vlm.named_parameters():
+                if not param.requires_grad:
+                    continue
+
+                # Check on any parameters with fewer than 2 dimensions or with "bias" in the name
+                if param.ndim <= 1 or name.endswith(".bias"):
+                    no_decay.append(param)
+                else:
+                    decay.append(param)
+
+            # Build Parameter Groups
+            groups = [{"params": decay, "weight_decay": self.weight_decay}, {"params": no_decay, "weight_decay": 0.0}]
+
+            # Create Optimizer & LR Scheduler
+            self.optimizer = AdamW(groups, lr=self.learning_rate)
+            self.lr_scheduler = get_cosine_schedule_with_warmup(self.optimizer, num_warmup_steps, num_training_steps)
+            for param_group in self.optimizer.param_groups:
+                param_group["lr"] = 0.0
+
+        elif self.lr_scheduler_type == "constant":
+            num_warmup_steps = 0
+
+            # Default AdamW w/ specified LR & Linear Warmup / Cosine Decay & Weight Decay
+            #   => Create Parameter Groups --> bias terms, normalization layer parameters shouldn't be decayed!
+            decay, no_decay = [], []
+            for name, param in self.vlm.named_parameters():
+                if not param.requires_grad:
+                    continue
+
+                # Check on any parameters with fewer than 2 dimensions or with "bias" in the name
+                if param.ndim <= 1 or name.endswith(".bias"):
+                    no_decay.append(param)
+                else:
+                    decay.append(param)
+
+            # Build Parameter Groups
+            groups = [{"params": decay, "weight_decay": self.weight_decay}, {"params": no_decay, "weight_decay": 0.0}]
+
+            # Create Optimizer & LR Scheduler
+            self.optimizer = AdamW(groups, lr=self.learning_rate)
+            self.lr_scheduler = get_constant_schedule(self.optimizer)
+
+        else:
+            raise ValueError(f"Learning Rate Schedule with type `{self.lr_scheduler_type}` is not supported!")
+
+        # Finalize Setup =>> Log!
+        overwatch.info(
+            "FSDP Full-Shard Strategy =>> Finalized Training Setup:\n"
+            f"         |-> Global (Effective) Batch Size = {self.global_batch_size}\n"
+            f"         |-> Per-Device Batch Size = {self.per_device_batch_size}\n"
+            f"         |-> Distributed World Size = {overwatch.world_size()}\n"
+            f"         |-> Gradient Accumulation Steps = {self.grad_accumulation_steps}\n\n"
+            f"         |-> LLM Backbone FSDP Gradient Checkpointing = {self.enable_gradient_checkpointing}\n"
+            f"         |-> Use FSDP Mixed Precision = {self.enable_mixed_precision_training}\n"
+            f"                 |-> Parameter Precision = {fsdp_precision_policy.param_dtype}\n"
+            f"                 |-> Reduction Precision = {fsdp_precision_policy.reduce_dtype}\n"
+            f"                 |-> Buffer Precision = {fsdp_precision_policy.buffer_dtype}\n\n"
+            f"         |-> Default AdamW LR = {self.learning_rate}\n"
+            f"         |-> AdamW Weight Decay = {self.weight_decay}\n"
+            f"         |-> LR Scheduler Type = {self.lr_scheduler_type}\n"
+            f"         |-> LR Scheduler Warmup Steps (Ratio) = {num_warmup_steps} ({self.warmup_ratio})\n"
+            f"         |-> Dataset Size = {n_train_examples} Examples\n"
+            f"         |-> Max Steps = {num_training_steps}\n"
+        )
+
+    def clip_grad_norm(self) -> None:
+        # Note =>> FSDP uses a custom `clip_grad_norm_` function; requires *uniform grad dtype*
+        self.vlm.clip_grad_norm_(max_norm=self.max_grad_norm)

capvector-oft/prismatic/training/train_utils.py

@@ -0,0 +1,56 @@
+"""Utils for training/fine-tuning scripts."""
+
+import torch
+
+from prismatic.vla.constants import ACTION_DIM, ACTION_TOKEN_BEGIN_IDX, IGNORE_INDEX
+
+
+def get_current_action_mask(token_ids):
+    # Create a tensor marking positions of IGNORE_INDEX
+    newline_positions = token_ids != IGNORE_INDEX
+
+    # Calculate cumulative sum to identify regions between newlines
+    cumsum = torch.cumsum(newline_positions, dim=1)
+
+    # Create the mask
+    mask = (1 <= cumsum) & (cumsum <= ACTION_DIM)
+
+    # Extract the action part only
+    action_tokens_only_mask = token_ids > ACTION_TOKEN_BEGIN_IDX
+    mask = action_tokens_only_mask * mask
+
+    return mask
+
+
+def get_next_actions_mask(token_ids):
+    # Create a tensor marking positions of IGNORE_INDEX
+    newline_positions = token_ids != IGNORE_INDEX
+
+    # Calculate cumulative sum to identify regions between newlines
+    cumsum = torch.cumsum(newline_positions, dim=1)
+
+    # Create the mask
+    mask = cumsum > ACTION_DIM
+
+    # Extract the action part only
+    action_tokens_only_mask = token_ids > ACTION_TOKEN_BEGIN_IDX
+    mask = action_tokens_only_mask * mask
+
+    return mask
+
+
+def compute_token_accuracy(predicted_token_ids, ground_truth_token_ids, mask):
+    correct_preds = (predicted_token_ids == ground_truth_token_ids) & mask
+    accuracy = correct_preds.sum().float() / mask.sum().float()
+    return accuracy
+
+
+def compute_actions_l1_loss(action_tokenizer, predicted_token_ids, ground_truth_token_ids, mask):
+    pred_continuous_actions = torch.tensor(
+        action_tokenizer.decode_token_ids_to_actions(predicted_token_ids[mask].cpu().numpy())
+    )
+    true_continuous_actions = torch.tensor(
+        action_tokenizer.decode_token_ids_to_actions(ground_truth_token_ids[mask].cpu().numpy())
+    )
+    l1_loss = torch.nn.functional.l1_loss(pred_continuous_actions, true_continuous_actions)
+    return l1_loss

capvector-oft/prismatic/util/__init__.py

@@ -0,0 +1 @@
+from .torch_utils import check_bloat16_supported, set_global_seed

capvector-oft/prismatic/util/batching_utils.py

@@ -0,0 +1,212 @@
+"""
+batching_utils.py
+
+Core definitions of (Distributed) Samplers for VLM finetuning; provides functionality for construction and allocating
+"split-modality" batches as described in the LLaVa paper; this makes sure that a given device/batch is either entirely
+(vision, language) or (language-only) data, which leads to sizeable efficiency gains.
+"""
+
+import math
+from typing import Iterator, List, Optional, Tuple
+
+import numpy as np
+import torch
+import torch.distributed as dist
+from torch.utils.data import Dataset, Sampler
+
+
+# High-Fidelity Bitwise Reproduction of the LLaVa Codebase Sampler Strategy + Per-Rank Allocation Scheme (following
+#   the default batching behavior of HF's Trainer Class --> derived from `accelerate`).
+#
+#   =>> Reference: https://github.com/haotian-liu/LLaVA/blob/main/llava/train/llava_trainer.py#L60
+#   =>> Reference: https://github.com/huggingface/transformers/blob/main/src/transformers/trainer_pt_utils.py#L603
+class SplitModalitySampler(Sampler):
+    def __init__(
+        self,
+        dataset: Dataset,
+        modality_lengths: List[Tuple[bool, int]],
+        global_batch_size: int,
+        num_replicas: Optional[int] = None,
+        rank: Optional[int] = None,
+        seed: int = 0,
+        drop_last: bool = False,
+    ) -> None:
+        super().__init__()
+        self.num_replicas = num_replicas if num_replicas is not None else dist.get_world_size()
+        self.rank = rank if rank is not None else dist.get_rank()
+        self.seed, self.epoch = seed, 0
+
+        # Custom Parameters
+        self.dataset, self.modality_lengths, self.drop_last = dataset, modality_lengths, drop_last
+        self.global_batch_size = global_batch_size
+
+        # For our purposes, `drop_last` is always False!
+        assert not self.drop_last, "SplitModalitySampler must set `drop_last = False`!"
+        self.total_size = math.ceil(len(self.dataset) / self.global_batch_size) * self.global_batch_size
+        self.num_samples = self.total_size // self.num_replicas
+
+    @staticmethod
+    def reindex_batch(batch_idxs: List[int], idx2lengths: List[int], n_buckets: int) -> List[List[int]]:
+        """Re-indexes a batch in a way that is conducive to DistributedSampler + grouping by seqlen per rank."""
+        assert len(batch_idxs) % n_buckets == 0, "Batch length is not divisible by `num_replicas`!"
+
+        # Establish initial buckets, capacities, and max number of elements per bucket
+        n_examples_per_bucket = len(batch_idxs) // n_buckets
+        bucket_indices = [[] for _ in range(n_buckets)]
+        bucket_lengths = [0 for _ in range(n_buckets)]
+
+        # Note that `batch_idxs` is already sorted by corresponding length (in descending order)
+        for idx in batch_idxs:
+            shortest_bucket_idx = bucket_lengths.index(min(bucket_lengths))
+            bucket_indices[shortest_bucket_idx].append(idx)
+
+            # Update `bucket_lengths` --> set length to infinity if at capacity!
+            bucket_lengths[shortest_bucket_idx] += idx2lengths[idx]
+            if len(bucket_indices[shortest_bucket_idx]) == n_examples_per_bucket:
+                bucket_lengths[shortest_bucket_idx] = float("inf")
+
+        return bucket_indices
+
+    def get_modality_and_length_grouped_indices(self, generator: torch.Generator) -> List[int]:
+        """
+        Returns a list of indices so that each slice of `global_batch_size` consecutive indices corresponds to elements
+        of the same modality with each sub-sequence of `per_replica_batch_size` (the batch size each unique device sees
+        during distributed training) is roughly grouped by sequence length (for training efficiency).
+        """
+        multimodal_indices, multimodal_lengths = zip(
+            *[(idx, length) for idx, (is_multimodal, length) in enumerate(self.modality_lengths) if is_multimodal]
+        )
+
+        # Handle Special Case --> no "unimodal" inputs
+        unimodal_split = [
+            (idx, length) for idx, (is_multimodal, length) in enumerate(self.modality_lengths) if not is_multimodal
+        ]
+        if len(unimodal_split) == 0:
+            unimodal_indices, unimodal_lengths = [], []
+        else:
+            unimodal_indices, unimodal_lengths = zip(*unimodal_split)
+
+        # Create a permutation of indices for each of the multimodal and unimodal data
+        mm_shuffled_idxs = torch.randperm(len(multimodal_indices), generator=generator)
+        uni_shuffled_idxs = torch.randperm(len(unimodal_indices), generator=generator)
+
+        # We're going to be running sorting/grouping relative to `self.global_batch_size` and `self.num_replicas`
+        g_bsz = self.global_batch_size
+
+        # Break each of the permutations into batches of length `global_batch_size`
+        mm_batch_idxs = [mm_shuffled_idxs[i : i + g_bsz].tolist() for i in range(0, len(mm_shuffled_idxs), g_bsz)]
+        uni_batch_idxs = [uni_shuffled_idxs[i : i + g_bsz].tolist() for i in range(0, len(uni_shuffled_idxs), g_bsz)]
+
+        # If "last" batch is not of length `g_bsz` --> PAD by stealing indices from the first batch!
+        if len(mm_batch_idxs[-1]) < g_bsz:
+            n_missing = g_bsz - len(mm_batch_idxs[-1])
+            mm_batch_idxs[-1].extend(mm_batch_idxs[0][:n_missing])
+
+        if len(uni_batch_idxs) > 0 and len(uni_batch_idxs[-1]) < g_bsz:
+            n_missing = g_bsz - len(uni_batch_idxs[-1])
+            uni_batch_idxs[-1].extend(uni_batch_idxs[0][:n_missing])
+
+        # Now we're going to sort each batch by length --> this will aid in grouping by length by rank (efficiency!)
+        mm_sorted_batch_idxs = [sorted(b, key=lambda i: multimodal_lengths[i], reverse=True) for b in mm_batch_idxs]
+        uni_sorted_batch_idxs = [sorted(b, key=lambda i: unimodal_lengths[i], reverse=True) for b in uni_batch_idxs]
+
+        # IMPORTANT :: At this point, for each modality, we have a list of "batches" (made up of indices) where indices
+        # are sorted by example sequence length *within* each batch. To make this more concrete, consider the following:
+        #   => World Size (`num_replicas`) = 2
+        #   => Global Batch Size (`g_bsz`) = 4
+        #   => `multimodal_indices` = [0,  1,  2,  3,  4,  5,  6,  7,  8,  9,  10, 11]
+        #      `multimodal_lengths` = [20, 90, 21, 22, 91, 18, 89, 19, 93, 88, 92, 17]
+        #
+        # At this point in the code, `mm_sorted_batch_idxs` might then look like the following (length in parenthesis):
+        #   => `mm_sorted_batch_idxs`: [
+        #       [4  (91), 3  (21), 0  (20), 5  (18)]    => Batch 1
+        #       [6  (89), 9  (88), 7  (19), 11 (17)]    => Batch 2
+        #       [8  (93), 10 (92), 1  (90), 2  (21)]    => Batch 3
+        #   ]
+        #
+        # In practice: `g_bsz` is large (= 128), and for contiguous mini-batch "slices", length variance is low.
+
+        # PROBLEM :: We want to split these "global batches" into equal-sized pieces, so that each "replica" (GPU)
+        # sees a "mini-batch" of roughly the same sequence lengths; this is super useful for efficient training.
+
+        # HOWEVER :: The default "access pattern" for splitting a large batch into mini-batches by a DistributedSampler
+        # is akin to a "take every k" where `k` is equal to the number of replicas (GPUs) you're training on. Or, in
+        # Python notation --> `rank_k_indices = flatten(mm_sorted_batch_idxs)[k::num_replicas].
+        #
+        # Naively translating this our example means each GPU (in our world of 2 total) sees the following indices
+        # (grouped by "mini-batch" = `g_bsz / num_replicas` = 2 for convenience):
+        #   => `rank_0_indices`: [ [4 (91), 0 (20)] =>> [6 (89), 7  (19)] =>> [8  (93), 1 (90)] ]
+        #   => `rank_1_indices`: [ [3 (21), 5 (18)] =>> [9 (88), 11 (17)] =>> [10 (92), 2 (21)] ]
+        #
+        # We get lucky sometimes, but for the most part, each "mini-batch" has VASTLY DIFFERENT lengths! Bad!
+
+        # FIX :: If we "undo" the access pattern with the following code and re-arrange the way we allocate batches
+        # inside the __iter__ method below, we can allocate indices appropriately. Running the following code gives us
+        # the following indices (grouped by "mini-batch" again for convenience):
+        #   => `rank_0_indices`: [ [4 (91), 3 (21)] =>> [6  (89), 9 (88)] =>> [8 (93), 10 (92)] ]
+        #   => `rank_1_indices`: [ [5 (18), 0 (20)] =>> [11 (17), 7 (19)] =>> [2 (21),  1 (90)] ]
+        #
+        # Much better! As `g_bsz` and `dataset` grow, we're more often than not getting *decent* groupings!
+        mm_length_bucketed_idxs = [
+            self.reindex_batch(batch, multimodal_lengths, self.num_replicas) for batch in mm_sorted_batch_idxs
+        ]
+        uni_length_bucketed_idxs = [
+            self.reindex_batch(batch, unimodal_lengths, self.num_replicas) for batch in uni_sorted_batch_idxs
+        ]
+
+        # Note :: Because of the initial `randperm` --> we're indexing both sets from 0 (we're clobbering the range)
+        #   => Flatten indices --> index into original `{modality}_indices` then re-batch!
+        mm_output_idxs = [idx for batch in mm_length_bucketed_idxs for bucket in batch for idx in bucket]
+        mm_reindexed = [multimodal_indices[idx] for idx in mm_output_idxs]
+        mm_batches = [mm_reindexed[i : i + g_bsz] for i in range(0, len(mm_reindexed), g_bsz)]
+
+        uni_output_idxs = [idx for batch in uni_length_bucketed_idxs for bucket in batch for idx in bucket]
+        uni_reindexed = [unimodal_indices[idx] for idx in uni_output_idxs]
+        uni_batches = [uni_reindexed[i : i + g_bsz] for i in range(0, len(uni_reindexed), g_bsz)]
+
+        # Finally, randomly permute the multimodal & unimodal batches, merging into a single stream of indices
+        merged_batches = mm_batches + uni_batches
+        merge_idxs = torch.randperm(len(merged_batches), generator=generator)
+        all_batches = [merged_batches[idx] for idx in merge_idxs]
+
+        # [Quality of Life] Shift "max length" batch to index 0 --> if we OOM, it happens immediately!
+        all_lengths = [length + ((_n_patches := 24 * 24) if is_mm else 0) for is_mm, length in self.modality_lengths]
+        all_batches_max_lengths = []
+        for batch in all_batches:
+            all_batches_max_lengths.append(max([all_lengths[idx] for idx in batch]))
+
+        # Identify Batch with "max length" --> Swap into Index 0
+        longest_batch_idx = np.argmax(all_batches_max_lengths)
+        all_batches[0], all_batches[longest_batch_idx] = all_batches[longest_batch_idx], all_batches[0]
+
+        # Flatten & Return all Indices
+        indices = [idx for batch in all_batches for idx in batch]
+        return indices
+
+    def __iter__(self) -> Iterator:
+        """Deterministically shuffle, then split indices by modality and length."""
+        g = torch.Generator()
+        g.manual_seed(self.seed + self.epoch)
+        indices = self.get_modality_and_length_grouped_indices(g)
+        assert len(set(indices)) == len(self.modality_lengths) == len(self.dataset), "Oops!"
+        assert (len(indices) % self.global_batch_size == 0) and (len(indices) % self.num_replicas) == 0, "Oops"
+
+        # Note :: We compute per-replica batch size as a function of `global_batch` and `num_replicas` to ensure that
+        # gradient accumulation doesn't affect what indices are assigned a given rank.
+        per_replica_batch_size = self.global_batch_size // self.num_replicas
+
+        # Tensorize & Unravel --> rather than yielding via a `take_every` --> we want to partition a global batch
+        # across replicas by assigning each a contiguous sub-sequence.
+        indices_t = torch.as_tensor(indices)
+        per_replica_batch_indices_t = indices_t.reshape(-1, per_replica_batch_size)
+        replica_indices_t = per_replica_batch_indices_t[self.rank :: self.num_replicas]
+
+        replica_indices = replica_indices_t.flatten().tolist()
+        return iter(replica_indices)
+
+    def __len__(self) -> int:
+        return self.num_samples
+
+    def set_epoch(self, epoch: int) -> None:
+        """To be called *between* epochs, prior to DataLoader instantiation; ensures random order across epochs."""
+        self.epoch = epoch

capvector-oft/prismatic/util/data_utils.py

@@ -0,0 +1,156 @@
+"""
+data_utils.py
+
+General utilities and classes for facilitating data loading and collation.
+"""
+
+from dataclasses import dataclass
+from typing import Callable, Dict, Sequence, Tuple
+
+import numpy as np
+import torch
+from torch.nn.utils.rnn import pad_sequence
+
+# HuggingFace Default / LLaMa-2 IGNORE_INDEX (for labels)
+IGNORE_INDEX = -100
+
+
+def tree_map(fn: Callable, tree: dict) -> dict:
+    """Maps a function over a nested dictionary."""
+    return {k: tree_map(fn, v) if isinstance(v, dict) else fn(v) for k, v in tree.items()}
+
+
+def tree_map_with_key(fn: Callable, tree: dict, keys: Sequence = ()) -> dict:
+    """Maps a function over a nested dictionary."""
+    return {
+        k: tree_map_with_key(fn, v, (*keys, k)) if isinstance(v, dict) else fn((*keys, k), v) for k, v in tree.items()
+    }
+
+
+@dataclass
+class PaddedCollatorForLanguageModeling:
+    model_max_length: int
+    pad_token_id: int
+    default_image_resolution: Tuple[int, int, int]
+    padding_side: str = "right"
+    pixel_values_dtype: torch.dtype = torch.float32
+
+    def __post_init__(self) -> None:
+        self.dummy_pixel_values = torch.zeros(self.default_image_resolution, dtype=self.pixel_values_dtype)
+
+    def __call__(self, instances: Sequence[Dict[str, torch.Tensor]]) -> Dict[str, torch.Tensor]:
+        input_ids, labels = tuple([instance[key] for instance in instances] for key in ("input_ids", "labels"))
+        pixel_values = [instance["pixel_values"] for instance in instances]
+
+        # For now, we only support Tokenizers with `padding_side = "right"` during Training (but plan to extend!)
+        #   => Handle padding via RNN Utils => `pad_sequence`
+        input_ids = pad_sequence(input_ids, batch_first=True, padding_value=self.pad_token_id)
+        labels = pad_sequence(labels, batch_first=True, padding_value=IGNORE_INDEX)
+
+        # Truncate (if necessary)
+        input_ids, labels = input_ids[:, : self.model_max_length], labels[:, : self.model_max_length]
+
+        # Get `attention_mask` by checking for `pad_token_id`
+        attention_mask = input_ids.ne(self.pad_token_id)
+
+        # === Handle "unimodal" (language-only) vs. "multimodal" ===
+
+        # Some examples are "language-only" --> build a Tensor of `multimodal_indices` that we can slice into easily
+        multimodal_indices = torch.tensor(
+            [idx for idx in range(len(pixel_values)) if pixel_values[idx] is not None], dtype=torch.long
+        )
+
+        # Stack all `pixel_values` --> depending on type (torch.Tensor, or Dict[str, torch.Tensor]) & presence of None
+        if len(multimodal_indices) == 0:
+            pixel_values = torch.stack([self.dummy_pixel_values for _ in range(len(input_ids))])
+        elif isinstance(pv_example := pixel_values[multimodal_indices[0]], torch.Tensor):
+            pixel_values = torch.stack(
+                [
+                    pixel_values[idx] if idx in multimodal_indices else self.dummy_pixel_values
+                    for idx in range(len(input_ids))
+                ]
+            )
+        elif isinstance(pv_example, dict):
+            pixel_values = {
+                k: torch.stack(
+                    [
+                        pixel_values[idx][k] if idx in multimodal_indices else self.dummy_pixel_values
+                        for idx in range(len(input_ids))
+                    ]
+                )
+                for k in pv_example
+            }
+        else:
+            raise ValueError(f"Unsupported `pixel_values` type = {type(pixel_values)}")
+
+        return dict(
+            pixel_values=pixel_values,
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            labels=labels,
+            multimodal_indices=multimodal_indices,
+        )
+
+
+@dataclass
+class PaddedCollatorForActionPrediction:
+    model_max_length: int
+    pad_token_id: int
+    padding_side: str = "right"
+    pixel_values_dtype: torch.dtype = torch.float32
+
+    def __call__(self, instances: Sequence[Dict[str, torch.Tensor]]) -> Dict[str, torch.Tensor]:
+        input_ids, labels = tuple([instance[key] for instance in instances] for key in ("input_ids", "labels"))
+        pixel_values = [instance["pixel_values"] for instance in instances]
+        if "dataset_name" in instances[0]:
+            dataset_names = [instance["dataset_name"] for instance in instances]
+        else:
+            dataset_names = None
+
+        # For now, we only support Tokenizers with `padding_side = "right"` during training
+        #   => Handle padding via RNN Utils => `pad_sequence`
+        assert self.padding_side == "right", f"Invalid Tokenizer `{self.padding_side = }`"
+        input_ids = pad_sequence(input_ids, batch_first=True, padding_value=self.pad_token_id)
+        labels = pad_sequence(labels, batch_first=True, padding_value=IGNORE_INDEX)
+
+        # Truncate (if necessary)
+        input_ids, labels = input_ids[:, : self.model_max_length], labels[:, : self.model_max_length]
+
+        # Get `attention_mask` by checking for `pad_token_id`
+        attention_mask = input_ids.ne(self.pad_token_id)
+
+        # [Contract] For VLA Training =>> No "Unimodal" Data!
+        assert all([pv is not None for pv in pixel_values]), "Invalid VLA Example with `pixel_values = None`!"
+
+        # Stack all `pixel_values` --> depending on type is torch.Tensor or Dict[str, torch.Tensor]
+        if isinstance(pixel_values[0], torch.Tensor):
+            if "pixel_values_wrist" in instances[0]:
+                pixel_values_wrist = [instance["pixel_values_wrist"] for instance in instances]
+                pixel_values = torch.cat((torch.stack(pixel_values), torch.stack(pixel_values_wrist)), dim=1)
+            else:
+                pixel_values = torch.stack(pixel_values)
+        else:
+            raise ValueError(f"Unsupported `pixel_values` type = {type(pixel_values)}")
+
+        # Stack all actions
+        actions = [torch.from_numpy(np.copy(instance["actions"])) for instance in instances]
+        actions = torch.stack(actions)
+
+        # Stack proprio
+        if "proprio" in instances[0]:
+            proprio = [instance["proprio"] for instance in instances]
+            proprio = torch.Tensor(np.squeeze(np.stack(proprio)))
+        else:
+            proprio = None
+
+        output = dict(
+            pixel_values=pixel_values,
+            proprio=proprio,
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            labels=labels,
+            actions=actions,
+        )
+        if dataset_names is not None:
+            output["dataset_names"] = dataset_names
+        return output

capvector-oft/prismatic/util/nn_utils.py

@@ -0,0 +1,53 @@
+"""
+nn_utils.py
+
+Utility functions and PyTorch submodule definitions.
+"""
+
+import torch
+import torch.nn as nn
+
+
+# === Definitions for Various Projection Modules, with Signature :: [..., in_dim] --> [..., out_dim] ===
+class LinearProjector(nn.Module):
+    def __init__(self, vision_dim: int, llm_dim: int) -> None:
+        super().__init__()
+        self.projector = nn.Linear(vision_dim, llm_dim, bias=True)
+
+    def forward(self, img_patches: torch.Tensor) -> torch.Tensor:
+        return self.projector(img_patches)
+
+
+class MLPProjector(nn.Module):
+    def __init__(self, vision_dim: int, llm_dim: int, mlp_type: str = "gelu-mlp") -> None:
+        super().__init__()
+        if mlp_type == "gelu-mlp":
+            self.projector = nn.Sequential(
+                nn.Linear(vision_dim, llm_dim, bias=True),
+                nn.GELU(),
+                nn.Linear(llm_dim, llm_dim, bias=True),
+            )
+        else:
+            raise ValueError(f"Projector with `{mlp_type = }` is not supported!")
+
+    def forward(self, img_patches: torch.Tensor) -> torch.Tensor:
+        return self.projector(img_patches)
+
+
+class FusedMLPProjector(nn.Module):
+    def __init__(self, fused_vision_dim: int, llm_dim: int, mlp_type: str = "fused-gelu-mlp") -> None:
+        super().__init__()
+        self.initial_projection_dim = fused_vision_dim * 4
+        if mlp_type == "fused-gelu-mlp":
+            self.projector = nn.Sequential(
+                nn.Linear(fused_vision_dim, self.initial_projection_dim, bias=True),
+                nn.GELU(),
+                nn.Linear(self.initial_projection_dim, llm_dim, bias=True),
+                nn.GELU(),
+                nn.Linear(llm_dim, llm_dim, bias=True),
+            )
+        else:
+            raise ValueError(f"Fused Projector with `{mlp_type = }` is not supported!")
+
+    def forward(self, fused_img_patches: torch.Tensor) -> torch.Tensor:
+        return self.projector(fused_img_patches)

capvector-oft/prismatic/util/torch_utils.py

@@ -0,0 +1,95 @@
+"""
+torch_utils.py
+
+General utilities for randomness, mixed precision training, and miscellaneous checks in PyTorch.
+
+Random `set_global_seed` functionality is taken directly from PyTorch-Lighting:
+    > Ref: https://github.com/PyTorchLightning/pytorch-lightning/blob/master/pytorch_lightning/utilities/seed.py
+
+This is pretty important to get right if we're every randomly generating our masks (or prefix dropout) inside our
+Dataset __getitem__() with multiple workers... if not handled properly, we will get repeated augmentations anytime
+we inject randomness from non-PyTorch sources (e.g., numpy, random)!
+    > Ref: https://tanelp.github.io/posts/a-bug-that-plagues-thousands-of-open-source-ml-projects/
+
+Terminology
+    -> World Size :: Total number of processes distributed over (# nodes x # devices) -- assumed homogenous!
+    -> Rank :: Integer index of current process in the total world size
+    -> Local Rank :: Local index on given node in [0, Devices per Node]
+"""
+
+import os
+import random
+from typing import Callable, Optional
+
+import numpy as np
+import torch
+
+# === Randomness ===
+
+
+def set_global_seed(seed: int, get_worker_init_fn: bool = False) -> Optional[Callable[[int], None]]:
+    """Sets seed for all randomness libraries (mostly random, numpy, torch) and produces a `worker_init_fn`"""
+    assert np.iinfo(np.uint32).min < seed < np.iinfo(np.uint32).max, "Seed outside the np.uint32 bounds!"
+
+    # Set Seed as an Environment Variable
+    os.environ["EXPERIMENT_GLOBAL_SEED"] = str(seed)
+    random.seed(seed)
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+
+    return worker_init_function if get_worker_init_fn else None
+
+
+def worker_init_function(worker_id: int) -> None:
+    """
+    Borrowed directly from PyTorch-Lightning; inspired by this issue comment in the PyTorch repo:
+        > Ref: https://github.com/pytorch/pytorch/issues/5059#issuecomment-817392562
+
+    Intuition: You can think of the seed sequence spawn function as a "janky" torch.Generator() or jax.PRNGKey that
+    you can run iterative splitting on to get new (predictable) randomness.
+
+    :param worker_id: Identifier for the given worker [0, num_workers) for the Dataloader in question.
+    """
+    # Get current `rank` (if running distributed) and `process_seed`
+    global_rank, process_seed = int(os.environ["LOCAL_RANK"]), torch.initial_seed()
+
+    # Back out the "base" (original) seed - the per-worker seed is set in PyTorch:
+    #   > https://pytorch.org/docs/stable/data.html#data-loading-randomness
+    base_seed = process_seed - worker_id
+
+    # "Magic" code --> basically creates a seed sequence that mixes different "sources" and seeds every library...
+    seed_seq = np.random.SeedSequence([base_seed, worker_id, global_rank])
+
+    # Use 128 bits (4 x 32-bit words) to represent seed --> generate_state(k) produces a `k` element array!
+    np.random.seed(seed_seq.generate_state(4))
+
+    # Spawn distinct child sequences for PyTorch (reseed) and stdlib random
+    torch_seed_seq, random_seed_seq = seed_seq.spawn(2)
+
+    # Torch Manual seed takes 64 bits (so just specify a dtype of uint64
+    torch.manual_seed(torch_seed_seq.generate_state(1, dtype=np.uint64)[0])
+
+    # Use 128 Bits for `random`, but express as integer instead of as an array
+    random_seed = (random_seed_seq.generate_state(2, dtype=np.uint64).astype(list) * [1 << 64, 1]).sum()
+    random.seed(random_seed)
+
+
+# === BFloat16 Support ===
+
+
+def check_bloat16_supported() -> bool:
+    try:
+        import packaging.version
+        import torch.cuda.nccl as nccl
+        import torch.distributed as dist
+
+        return (
+            (torch.version.cuda is not None)
+            and torch.cuda.is_bf16_supported()
+            and (packaging.version.parse(torch.version.cuda).release >= (11, 0))
+            and dist.is_nccl_available()
+            and (nccl.version() >= (2, 10))
+        )
+
+    except Exception:
+        return False