optimum-rbln 0.9.4a2__py3-none-any.whl → 0.9.5a4__py3-none-any.whl

optimum/rbln/transformers/models/colpali/modeling_colpali.py

@@ -14,23 +14,16 @@
 
 import bisect
 from pathlib import Path
-from typing import TYPE_CHECKING, Optional, Tuple, Union
+from typing import Optional, Tuple, Union
 
 import torch
-from transformers import PretrainedConfig, PreTrainedModel
 from transformers.modeling_outputs import BaseModelOutputWithPooling
 from transformers.modeling_utils import no_init_weights
-from transformers.models.colpali.modeling_colpali import ColPaliForRetrievalOutput
-from transformers.models.paligemma.modeling_paligemma import PaliGemmaMultiModalProjector
+from transformers.models.colpali.modeling_colpali import ColPaliForRetrieval, ColPaliForRetrievalOutput
 
-from ....configuration_utils import RBLNCompileConfig, RBLNModelConfig
+from ....configuration_utils import RBLNModelConfig
 from ....modeling import RBLNModel
 from ...utils.rbln_runtime_wrapper import LoopProcessor
-from .colpali_architecture import RBLNColPaliForRetrievalWrapper
-
-
-if TYPE_CHECKING:
-    from transformers import AutoFeatureExtractor, AutoProcessor, AutoTokenizer, PretrainedConfig
 
 
 class LoopVisionTower(LoopProcessor):
@@ -115,17 +108,25 @@ class RBLNColPaliForRetrieval(RBLNModel):
         from optimum.rbln import RBLNColPaliForRetrieval
 
         # Simple usage using rbln_* arguments
-        # `max_seq_lens` is automatically inferred from the model config
         model = RBLNColPaliForRetrieval.from_pretrained(
             "vidore/colpali-v1.3-hf",
             export=True,
-            rbln_max_seq_lens=1152,
+            rbln_config={
+                "vlm": {
+                    "language_model": {
+                        "prefill_chunk_size": 8192, # same as model's max_position_embeddings (max_seq_len)
+                    }
+                }
+            }
         )
 
         # Using a config dictionary
         rbln_config = {
-            "max_seq_lens": 1152,
-            "output_hidden_states": False,
+            "vlm": {
+                "language_model": {
+                    "prefill_chunk_size": 8192, # same as model's max_position_embeddings (max_seq_len)
+                }
+            }
         }
         model = RBLNColPaliForRetrieval.from_pretrained(
             "vidore/colpali-v1.3-hf",
@@ -137,7 +138,9 @@ class RBLNColPaliForRetrieval(RBLNModel):
         from optimum.rbln import RBLNColPaliForRetrievalConfig
 
         config = RBLNColPaliForRetrievalConfig(
-            max_seq_lens=1152,
+            vlm={
+                "language_model": {"prefill_chunk_size": 8192},
+            },
             output_hidden_states=False,
             tensor_parallel_size=4
         )
@@ -150,212 +153,93 @@ class RBLNColPaliForRetrieval(RBLNModel):
     """
 
     auto_model_class = None
+    _rbln_submodule_postfix = "model"
     _rbln_submodules = [
-        {"name": "vision_tower"},
+        {"name": "vlm"},
     ]
 
     def __post_init__(self, **kwargs):
-        self.vision_tower = LoopVisionTower(self.rbln_submodules[0])
-        self.language_model = LoopLanguageModel(self.model[0], self.rbln_config)
-
+        self.vlm_model = self.rbln_submodules[0]
         artifacts = torch.load(self.model_save_dir / self.subfolder / "torch_artifacts.pth", weights_only=False)
-        self.embed_tokens = self._create_embedding_layer()
-        self.embed_tokens.load_state_dict(artifacts["embed_tokens"])
-        self.multi_modal_projector = self._create_multi_modal_projector()
-        self.multi_modal_projector.load_state_dict(artifacts["multi_modal_projector"])
-
+        self.embedding_proj_layer = self._create_embedding_proj_layer()
+        self.embedding_proj_layer.load_state_dict(artifacts["embedding_proj_layer"])
         return super().__post_init__(**kwargs)
 
-    def _create_embedding_layer(self):
+    def _create_embedding_proj_layer(self):
         with no_init_weights():
-            embed_tokens = torch.nn.Embedding(
-                self.config.text_config.vocab_size,
-                self.config.text_config.hidden_size,
-                self.config.text_config.pad_token_id,
+            embedding_proj_layer = torch.nn.Linear(
+                self.config.vlm_config.text_config.hidden_size, self.config.embedding_dim
             )
-        return embed_tokens
-
-    def _create_multi_modal_projector(self):
-        with no_init_weights():
-            multi_modal_projector = PaliGemmaMultiModalProjector(self.config.vlm_config)
-        return multi_modal_projector
-
-    @classmethod
-    def _wrap_model_if_needed(cls, model: "PreTrainedModel", rbln_config: RBLNModelConfig):
-        return RBLNColPaliForRetrievalWrapper(
-            causal_lm=model.vlm,
-            embedding_proj_layer=model.embedding_proj_layer,
-            max_seq_len=max(rbln_config.max_seq_lens),
-            output_hidden_states=rbln_config.output_hidden_states,
-        )
+        return embedding_proj_layer
 
     @classmethod
     def save_torch_artifacts(
         cls,
-        model: "PreTrainedModel",
+        model: "ColPaliForRetrieval",
         save_dir_path: Path,
         subfolder: str,
         rbln_config: RBLNModelConfig,
     ):
         save_dict = {}
-        save_dict["embed_tokens"] = model.vlm.get_input_embeddings().state_dict()
-        save_dict["multi_modal_projector"] = model.vlm.multi_modal_projector.state_dict()
+        save_dict["embedding_proj_layer"] = model.embedding_proj_layer.state_dict()
         torch.save(save_dict, save_dir_path / subfolder / "torch_artifacts.pth")
 
-    @classmethod
-    def _update_rbln_config(
-        cls,
-        preprocessors: Optional[Union["AutoFeatureExtractor", "AutoProcessor", "AutoTokenizer"]],
-        model: Optional["PreTrainedModel"] = None,
-        model_config: Optional["PretrainedConfig"] = None,
-        rbln_config: Optional[RBLNModelConfig] = None,
-    ) -> RBLNModelConfig:
-        hidden_size = model_config.vlm_config.text_config.hidden_size
-        if rbln_config.max_seq_lens is None:
-            rbln_config.max_seq_lens = [model_config.vlm_config.text_config.max_position_embeddings]
-        if isinstance(rbln_config.max_seq_lens, int):
-            rbln_config.max_seq_lens = [rbln_config.max_seq_lens]
-        rbln_config.max_seq_lens = sorted(set(rbln_config.max_seq_lens))
-
-        if rbln_config.output_hidden_states is None:
-            rbln_config.output_hidden_states = model_config.vlm_config.text_config.output_hidden_states
-
-        input_infos = []
-        for max_seq_len in rbln_config.max_seq_lens:
-            input_info = [
-                ("inputs_embeds", [rbln_config.vision_tower.batch_size, max_seq_len, hidden_size], "float32"),
-                ("attention_mask", [rbln_config.vision_tower.batch_size, max_seq_len], "float32"),
-                ("position_ids", [rbln_config.vision_tower.batch_size, max_seq_len], "int32"),
-            ]
-            input_infos.append(input_info)
-
-        rbln_compile_config = RBLNCompileConfig(input_info=input_infos)
-        rbln_config.set_compile_cfgs([rbln_compile_config])
-
-        return rbln_config
-
-    @classmethod
-    def _reconstruct_model_if_needed(cls, model: "PreTrainedModel"):
-        if hasattr(model, "vlm"):
-            model.vision_tower = model.vlm.vision_tower
-            del model.vlm.model.vision_tower
-            return model
-        return model
-
-    def get_image_features(self, pixel_values: torch.Tensor):
-        # Projects the last hidden state from the vision model into language model space.
-        # Args:
-        #     pixel_values (`torch.FloatTensor]` of shape `(batch_size, channels, height, width)`)
-        #        The tensors corresponding to the input images.
-        # Returns:
-        #     image_features (`torch.Tensor`): Image feature tensor of shape `(num_images, image_length, embed_dim)`).
-
-        vision_output_size = [
-            pixel_values.shape[0],
-            self.config.vlm_config.vision_config.num_image_tokens,
-            self.config.vlm_config.vision_config.hidden_size,
-        ]
-        vision_output = torch.empty(size=vision_output_size, dtype=torch.float32, device="cpu")
-        self.vision_tower(pixel_values, out=vision_output)
-        image_features = self.multi_modal_projector(vision_output)
-        image_features = image_features / (self.config.text_config.hidden_size**0.5)
-        return image_features
-
-    def _preprocess_inputs(
-        self,
-        input_ids: Optional[torch.LongTensor] = None,
-        inputs_embeds: Optional[torch.FloatTensor] = None,
-        pixel_values: Optional[torch.FloatTensor] = None,
-        **kwargs,
-    ):
-        if (input_ids is None) ^ (inputs_embeds is not None):
-            raise ValueError("You must specify exactly one of input_ids or inputs_embeds")
-
-        # Replace image id woth PAD if the image token if OOV, to avoid index-errors
-        if input_ids is not None and self.config.vlm_config.image_token_index >= self.config.text_config.vocab_size:
-            special_image_mask = input_ids == self.config.vlm_config.image_token_index
-            llm_input_ids = input_ids.clone()
-            llm_input_ids[special_image_mask] = 0
-        else:
-            llm_input_ids = input_ids
-
-        if inputs_embeds is None:
-            inputs_embeds = self.embed_tokens(llm_input_ids)
-
-        # Merge text and images
-        image_features = None
-        if pixel_values is not None:
-            image_features = self.get_image_features(pixel_values)
-            special_image_mask = (input_ids == self.config.vlm_config.image_token_index).unsqueeze(-1)
-            special_image_mask = special_image_mask.expand_as(inputs_embeds).to(inputs_embeds.device)
-
-            image_features = image_features.to(inputs_embeds.device, inputs_embeds.dtype)
-            inputs_embeds = inputs_embeds.masked_scatter(special_image_mask, image_features)
-
-        return inputs_embeds, image_features
-
     def forward(
         self,
         input_ids: Optional[torch.LongTensor] = None,
-        inputs_embeds: Optional[torch.FloatTensor] = None,
         pixel_values: Optional[torch.FloatTensor] = None,
         attention_mask: Optional[torch.Tensor] = None,
-        output_attentions: Optional[bool] = None,
         output_hidden_states: Optional[bool] = None,
         return_dict: Optional[bool] = None,
         **kwargs,
     ) -> Union[Tuple, ColPaliForRetrievalOutput]:
+        """
+        Forward pass for the RBLN-optimized ColPaliForRetrieval model.
+
+        Args:
+            input_ids (torch.LongTensor of shape (batch_size, sequence_length)): Indices of input sequence tokens in the vocabulary.
+            pixel_values (torch.Tensor of shape (batch_size, num_channels, image_size, image_size)): The tensors corresponding to the input images.
+            attention_mask (torch.Tensor of shape (batch_size, sequence_length)): Mask to avoid performing attention on padding token indices.
+            output_hidden_states (bool, optional): Whether or not to return the hidden states of all layers. See hidden_states under returned tensors for more detail.
+            return_dict (bool, optional): Whether or not to return a ModelOutput instead of a plain tuple.
+
+        Returns:
+            ColPaliForRetrievalOutput or tuple(torch.FloatTensor)
+        """
         if pixel_values is not None:
             pixel_values = pixel_values.to(dtype=self.dtype)
 
-        if output_attentions:
-            raise ValueError("output_attentions is not supported for RBLNColPaliForRetrieval")
-
-        if output_hidden_states is not None and output_hidden_states != self.rbln_config.output_hidden_states:
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+        output_hidden_states = (
+            output_hidden_states if output_hidden_states is not None else self.rbln_config.output_hidden_states
+        )
+        if output_hidden_states != self.rbln_config.output_hidden_states:
             raise ValueError(
                 f"Variable output_hidden_states {output_hidden_states} is not equal to rbln_config.output_hidden_states {self.rbln_config.output_hidden_states} "
                 f"Please compile again with the correct argument."
             )
 
-        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
-
-        inputs_embeds, image_features = self._preprocess_inputs(
-            input_ids=input_ids, inputs_embeds=inputs_embeds, pixel_values=pixel_values
+        vlm_output = self.vlm_model(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            pixel_values=pixel_values,
+            output_hidden_states=output_hidden_states,
+            return_dict=True,
+            **kwargs,
         )
+        vlm_hidden_states = vlm_output.hidden_states if output_hidden_states else None
+        vlm_image_hidden_states = vlm_output.image_hidden_states if pixel_values is not None else None
 
-        outputs = []
-        language_model_out_size = [inputs_embeds.shape[0], self.rbln_config.max_seq_lens[0], self.config.embedding_dim]
-        language_model_hidden_states_size = [
-            inputs_embeds.shape[0],
-            self.rbln_config.max_seq_lens[0],
-            self.rbln_config.max_seq_lens[0],
-        ]
-        outputs.append(torch.empty(size=language_model_out_size, dtype=torch.float32, device="cpu"))
-        if self.rbln_config.output_hidden_states:
-            for _ in range(self.config.vlm_config.text_config.num_hidden_layers + 1):
-                outputs.append(torch.empty(size=language_model_hidden_states_size, dtype=torch.float32, device="cpu"))
-
-        # Embedding_proj_layer is fused on the bottom of the language model.
-        self.language_model(inputs_embeds=inputs_embeds, attention_mask=attention_mask, out=outputs)
-
-        embeddings = outputs[0][:, : inputs_embeds.shape[1]]
-        hidden_states = (
-            None
-            if not self.rbln_config.output_hidden_states
-            else [tensor[0][:, : inputs_embeds.shape[1]] for tensor in outputs[1:]]
-        )
-
-        # L2 normalization
-        embeddings = embeddings / embeddings.norm(dim=-1, keepdim=True)  # (batch_size, sequence_length, dim)
+        last_hidden_states = vlm_output[0]
+        proj_dtype = self.embedding_proj_layer.weight.dtype
+        embeddings = self.embedding_proj_layer(last_hidden_states.to(proj_dtype))
+        embeddings = embeddings / embeddings.norm(dim=-1, keepdim=True)
 
         if attention_mask is not None:
-            embeddings = embeddings * attention_mask.unsqueeze(-1)  # (batch_size, sequence_length, dim)
+            embeddings = embeddings * attention_mask.unsqueeze(-1)
 
-        if not return_dict:
-            return (embeddings, hidden_states, image_features)
-        else:
-            return ColPaliForRetrievalOutput(
-                embeddings=embeddings,
-                hidden_states=hidden_states,
-                image_hidden_states=image_features,
-            )
+        return ColPaliForRetrievalOutput(
+            embeddings=embeddings,
+            hidden_states=vlm_hidden_states,
+            image_hidden_states=vlm_image_hidden_states,
+        )

optimum/rbln/transformers/models/colqwen2/configuration_colqwen2.py

@@ -32,14 +32,16 @@ class RBLNColQwen2ForRetrievalConfig(RBLNDecoderOnlyModelConfig):
 
         # Create a configuration object
         config = RBLNColQwen2ForRetrievalConfig(
-            visual={
-                "max_seq_lens": 6400,
-                "device": 0,
-            },
-            max_seq_len=32_768,
-            tensor_parallel_size=4,
-            device=[0, 1, 2, 3],
-            output_hidden_states=False,
+            vlm = {
+                "visual": {
+                    "max_seq_lens": 6400,
+                    "device": 0,
+                },
+                "max_seq_len": 32_768,
+                "tensor_parallel_size": 4,
+                "device": [0, 1, 2, 3],
+                "output_hidden_states": False,
+                }
         )
 
         # Use the configuration with from_pretrained
@@ -51,22 +53,37 @@ class RBLNColQwen2ForRetrievalConfig(RBLNDecoderOnlyModelConfig):
         ```
     """
 
-    submodules = ["visual"]
+    submodules = ["vlm"]
+    _allow_no_compile_cfgs = True
 
     def __init__(
         self,
-        visual: Optional[RBLNModelConfig] = None,
         batch_size: Optional[int] = None,
-        use_inputs_embeds: bool = True,
+        output_hidden_states: Optional[bool] = None,
+        vlm: Optional[RBLNModelConfig] = None,
         **kwargs,
     ):
-        super().__init__(use_inputs_embeds=use_inputs_embeds, **kwargs)
-        if not self.use_inputs_embeds:
-            raise ValueError(
-                "RBLNColQwen2ForRetrievalConfig does not allow `use_inputs_embeds` to be set to False, "
-                "as RBLNColQwen2ForRetrieval accepts only `inputs_embeds` as input."
-            )
-        if batch_size is not None and batch_size != 1:
-            raise ValueError("batch_size is not supported for RBLNColQwen2ForRetrievalConfig")
-
-        self.visual = visual
+        """
+        Args:
+            batch_size (Optional[int]): The batch size for the model.
+            output_hidden_states (Optional[bool]): Whether to output the hidden states of the VLM model.
+            vlm (Optional[RBLNModelConfig]): Configuration for the VLM component.
+            kwargs: Additional arguments passed to the parent RBLNModelConfig.
+        Raises:
+            ValueError: If batch_size is not a positive integer.
+        """
+        super().__init__(**kwargs)
+        self.batch_size = batch_size or 1
+        self.output_hidden_states = output_hidden_states or False
+
+        if not isinstance(self.batch_size, int) or self.batch_size < 0:
+            raise ValueError(f"batch_size must be a positive integer, got {self.batch_size}")
+
+        self.vlm = self.initialize_submodule_config(
+            submodule_config=vlm,
+            batch_size=batch_size,
+            output_hidden_states=output_hidden_states,
+            force_kwargs=True,
+            logits_to_keep=0,
+            use_inputs_embeds=True,
+        )