optimum-rbln 0.8.2a4__py3-none-any.whl → 0.9.3__py3-none-any.whl

optimum/rbln/transformers/models/clip/modeling_clip.py

@@ -51,8 +51,10 @@ class RBLNCLIPTextModel(RBLNModel):
     on RBLN devices, supporting text encoding for multimodal tasks.
     """
 
+    _tp_support = False
+
     @classmethod
-    def wrap_model_if_needed(cls, model: torch.nn.Module, rbln_config: RBLNCLIPTextModelConfig) -> torch.nn.Module:
+    def _wrap_model_if_needed(cls, model: torch.nn.Module, rbln_config: RBLNCLIPTextModelConfig) -> torch.nn.Module:
         return _TextEncoder(model).eval()
 
     @classmethod
@@ -83,7 +85,18 @@ class RBLNCLIPTextModel(RBLNModel):
         rbln_config.set_compile_cfgs([RBLNCompileConfig(input_info=input_info)])
         return rbln_config
 
-    def forward(self, input_ids: torch.LongTensor, return_dict: bool = None, **kwargs) -> torch.FloatTensor:
+    def forward(self, input_ids: torch.LongTensor, return_dict: Optional[bool] = None, **kwargs) -> torch.FloatTensor:
+        """
+        Forward pass for the RBLN-optimized CLIP text encoder model.
+
+        Args:
+            input_ids (torch.LongTensor): The input ids to the model.
+            return_dict (Optional[bool]): Whether to return a dictionary of outputs.
+
+        Returns:
+            The model outputs. If return_dict=False is passed, returns a tuple of tensors. Otherwise, returns a CLIPTextModelOutput object.
+        """
+
         # To ignore using attention_mask, we override forward method.
         output = super().forward(input_ids, return_dict=return_dict)
         return output
@@ -144,8 +157,10 @@ class RBLNCLIPVisionModel(RBLNModel):
     on RBLN devices, supporting image encoding for multimodal tasks.
     """
 
+    _tp_support = False
+
     @classmethod
-    def wrap_model_if_needed(cls, model: torch.nn.Module, rbln_config: RBLNCLIPVisionModelConfig) -> torch.nn.Module:
+    def _wrap_model_if_needed(cls, model: torch.nn.Module, rbln_config: RBLNCLIPVisionModelConfig) -> torch.nn.Module:
         wrapper_cfg = {
             "interpolate_pos_encoding": rbln_config.interpolate_pos_encoding,
             "output_hidden_states": rbln_config.output_hidden_states,
@@ -202,13 +217,27 @@ class RBLNCLIPVisionModel(RBLNModel):
 
     def forward(
         self,
-        pixel_values: Optional[torch.FloatTensor] = None,
+        pixel_values: torch.FloatTensor,
         return_dict: bool = True,
-        output_attentions: bool = None,
-        output_hidden_states: bool = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
         interpolate_pos_encoding: bool = False,
         **kwargs,
     ) -> Union[Tuple, BaseModelOutputWithPooling]:
+        """
+        Forward pass for the RBLN-optimized CLIP vision encoder model.
+
+        Args:
+            pixel_values (torch.Tensor): The pixel values to the model.
+            return_dict (bool): Whether to return a dictionary of outputs.
+            output_attentions (Optional[bool]): Whether to return attentions.
+            output_hidden_states (Optional[bool]): Whether to return hidden states.
+            interpolate_pos_encoding (bool): Whether to interpolate position encoding.
+
+        Returns:
+            The model outputs. If return_dict=False is passed, returns a tuple of tensors. Otherwise, returns a BaseModelOutputWithPooling object.
+        """
+
         if len(kwargs) > 0 and any(value is not None for value in kwargs.values()):
             logger.warning(
                 f"Currently, optimum-rbln does not support kwargs {kwargs.keys()} for {self.__class__.__name__}."
@@ -284,6 +313,38 @@ class RBLNCLIPVisionModelWithProjection(RBLNCLIPVisionModel):
     multimodal embedding alignment tasks.
     """
 
+    def forward(
+        self,
+        pixel_values: torch.FloatTensor,
+        return_dict: bool = True,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        interpolate_pos_encoding: bool = False,
+        **kwargs,
+    ) -> Union[Tuple, CLIPVisionModelOutput]:
+        """
+        Forward pass for the RBLN-optimized CLIP vision encoder model with projection.
+
+        Args:
+            pixel_values (torch.Tensor): The pixel values to the model.
+            return_dict (bool): Whether to return a dictionary of outputs.
+            output_attentions (Optional[bool]): Whether to return attentions.
+            output_hidden_states (Optional[bool]): Whether to return hidden states.
+            interpolate_pos_encoding (bool): Whether to interpolate position encoding.
+
+        Returns:
+            The model outputs. If return_dict=False is passed, returns a tuple of tensors. Otherwise, returns a CLIPVisionModelOutput object.
+        """
+
+        return super().forward(
+            pixel_values=pixel_values,
+            return_dict=return_dict,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            interpolate_pos_encoding=interpolate_pos_encoding,
+            **kwargs,
+        )
+
     def _prepare_output(self, output, return_dict):
         # Prepare model output based on return_dict flag.
         # This method can be overridden by subclasses to provide task-specific output handling.

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

@@ -4,10 +4,7 @@ import torch
 from torch import nn
 from transformers import GemmaForCausalLM, GemmaModel
 
-from ..decoderonly.decoderonly_architecture import (
-    RotaryEmbedding,
-    apply_rotary_pos_emb,
-)
+from ..decoderonly.decoderonly_architecture import RotaryEmbedding, apply_rotary_pos_emb
 
 
 def slice_and_unsqueeze_cos_sin(cos, sin, position_ids):
@@ -27,11 +24,11 @@ class RBLNColPaliForRetrievalWrapper(nn.Module):
         output_hidden_states: bool = False,
     ):
         super().__init__()
-        self.text_config = causal_lm.config
+        self.text_config = causal_lm.config.text_config
         self.rotary_emb = self.get_rotary_emb(max_seq_len=max_seq_len)
 
         self.output_hidden_states = output_hidden_states
-        self.language_model = self.convert_to_rbln_language_model(causal_lm.model, max_seq_len)
+        self.language_model = self.convert_to_rbln_language_model(causal_lm.model.language_model, max_seq_len)
 
         self.num_hidden_layers = getattr(self.text_config, "num_hidden_layers", None)
         self.embedding_proj_layer = embedding_proj_layer

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

@@ -11,9 +11,13 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-from typing import Any, Dict, List, Optional, Union
+from typing import Any, List, Optional, Union
 
 from ....configuration_utils import RBLNModelConfig
+from ....utils.logging import get_logger
+
+
+logger = get_logger(__name__)
 
 
 class RBLNColPaliForRetrievalConfig(RBLNModelConfig):
@@ -24,45 +28,57 @@ class RBLNColPaliForRetrievalConfig(RBLNModelConfig):
     including vision tower settings and multi-sequence length support.
 
     Example usage:
-    ```python
-    from optimum.rbln import RBLNColPaliForRetrieval, RBLNColPaliForRetrievalConfig
-
-    # Create a configuration object
-    config = RBLNColPaliForRetrievalConfig(
-        max_seq_lens=1152,
-        output_hidden_states=False,
-        tensor_parallel_size=4
-    )
-
-    # Use the configuration with from_pretrained
-    model = RBLNColPaliForRetrieval.from_pretrained(
-        "vidore/colpali-v1.3-hf",
-        export=True,
-        rbln_config=config
-    )
-    ```
+        ```python
+        from optimum.rbln import RBLNColPaliForRetrieval, RBLNColPaliForRetrievalConfig
+
+        # Create a configuration object
+        config = RBLNColPaliForRetrievalConfig(
+            max_seq_lens=1152,
+            output_hidden_states=False,
+            tensor_parallel_size=4
+        )
+
+        # Use the configuration with from_pretrained
+        model = RBLNColPaliForRetrieval.from_pretrained(
+            "vidore/colpali-v1.3-hf",
+            export=True,
+            rbln_config=config
+        )
+        ```
     """
 
     submodules = ["vision_tower"]
 
     def __init__(
         self,
+        batch_size: Optional[int] = None,
         max_seq_lens: Union[int, List[int]] = None,
         output_hidden_states: Optional[bool] = None,
         vision_tower: Optional[RBLNModelConfig] = None,
-        **kwargs: Dict[str, Any],
+        **kwargs: Any,
     ):
         """
         Args:
+            batch_size (Optional[int]): The batch size for the model.
             vision_tower (Optional[RBLNModelConfig]): Configuration for the vision encoder component.
             max_seq_lens (Union[int, List[int]]): The maximum sequence lengths for the language model.
                 This can be multiple values, and the model will be compiled for each max_seq_len, allowing selection of the most appropriate max_seq_len at inference time.
             output_hidden_states (Optional[bool]): Whether to output the hidden states of the language model.
-            **kwargs: Additional arguments passed to the parent RBLNModelConfig.
+            vision_tower (Optional[RBLNModelConfig]): Configuration for the vision encoder component.
+            kwargs: Additional arguments passed to the parent RBLNModelConfig.
         Raises:
             ValueError: If batch_size is not a positive integer.
         """
         super().__init__(**kwargs)
-        self.vision_tower = vision_tower
+        self.batch_size = batch_size or 1
+        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}")
+
+        if self.batch_size != 1:
+            logger.warning("Ignore batch_size for ColPali vision tower. It will be set to 1.")
+
+        self.vision_tower = self.initialize_submodule_config(
+            submodule_config=vision_tower, batch_size=1, force_kwargs=True
+        )
         self.max_seq_lens = max_seq_lens
         self.output_hidden_states = output_hidden_states

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

@@ -14,13 +14,10 @@
 
 import bisect
 from pathlib import Path
-from typing import TYPE_CHECKING, Any, Optional, Union
+from typing import TYPE_CHECKING, Optional, Tuple, Union
 
 import torch
-from transformers import (
-    PretrainedConfig,
-    PreTrainedModel,
-)
+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
@@ -28,105 +25,72 @@ from transformers.models.paligemma.modeling_paligemma import PaliGemmaMultiModal
 
 from ....configuration_utils import RBLNCompileConfig, 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,
-    )
+    from transformers import AutoFeatureExtractor, AutoProcessor, AutoTokenizer, PretrainedConfig
 
 
-class LoopVisionTower:
-    def __init__(self, vision_tower: RBLNModel) -> None:
-        self.vision_tower = vision_tower
+class LoopVisionTower(LoopProcessor):
+    def __init__(self, vision_tower: "RBLNModel"):
+        super().__init__(model=vision_tower.model[0])
 
-    def forward(self, pixel_values, **kwargs):
-        batch_size = pixel_values.shape[0]
-        outputs = []
-        for i in range(batch_size):
-            outputs.append(self.vision_tower(pixel_values[i : i + 1]))
+    def _get_batch_size(self, pixel_values, **kwargs):
+        return pixel_values.shape[0]
 
-        last_hidden_states = [output.last_hidden_state for output in outputs]
-        last_hidden_states = torch.cat(last_hidden_states, dim=0)
+    def _prepare_inputs_for_iteration(self, index, common_inputs, pixel_values, **kwargs):
+        pixel_values_item = pixel_values[index : index + 1]
+        out_buffer = kwargs["out"][index : index + 1]
+        return ([pixel_values_item], {"out": out_buffer})
 
+    def _process_outputs(self, outputs: list, **kwargs) -> "BaseModelOutputWithPooling":
         return BaseModelOutputWithPooling(
-            last_hidden_state=last_hidden_states,
+            last_hidden_state=kwargs["out"],
         )
 
-    def __call__(self, *args: Any, **kwds: Any) -> Any:
-        return self.forward(*args, **kwds)
-
-    def __repr__(self) -> str:
-        return repr(self.vision_tower)
-
 
-class LoopLanguageModel:
-    def __init__(self, language_model: RBLNModel, rbln_config: RBLNModelConfig) -> None:
-        self.language_model = language_model
+class LoopLanguageModel(LoopProcessor):
+    def __init__(self, language_model: RBLNModel, rbln_config: RBLNModelConfig):
+        super().__init__(model=language_model)
         self.rbln_config = rbln_config
 
-    def prepare_inputs(self, inputs_embeds: torch.Tensor, attention_mask: torch.Tensor):
+    def _get_batch_size(self, inputs_embeds, **kwargs):
+        return inputs_embeds.shape[0]
+
+    def _prepare_inputs_before_loop(self, *, inputs_embeds: torch.Tensor, attention_mask: torch.Tensor, **kwargs):
         input_len = inputs_embeds.shape[1]
         idx = bisect.bisect_left(self.rbln_config.max_seq_lens, input_len)
         if idx == len(self.rbln_config.max_seq_lens):
             raise ValueError(
                 f"Required seq_len({input_len}) is larger than available max_seq_lens({self.rbln_config.max_seq_lens})."
             )
-        else:
-            max_seq_len = self.rbln_config.max_seq_lens[idx]
-
-        inputs_embed = torch.nn.functional.pad(inputs_embeds, (0, 0, 0, max_seq_len - input_len))
-        attn_mask = torch.nn.functional.pad(attention_mask, (0, max_seq_len - input_len)).to(torch.float32)
-        position_ids = torch.arange(max_seq_len, dtype=torch.int32).view(1, -1)
-
-        return inputs_embed, attn_mask, position_ids
-
-    def forward(self, inputs_embeds: torch.Tensor, attention_mask: torch.Tensor, **kwargs):
-        padded_inputs_embed, padded_attn_mask, padded_position_ids = self.prepare_inputs(inputs_embeds, attention_mask)
-        input_batch_size = inputs_embeds.shape[0]
-        input_seq_len = inputs_embeds.shape[1]
-
-        all_embeddings = []
-        all_hidden_states = []
-        for i in range(input_batch_size):
-            outputs = self.language_model(
-                inputs_embeds=padded_inputs_embed[i : i + 1],
-                attention_mask=padded_attn_mask[i : i + 1],
-                position_ids=padded_position_ids,
-            )
-
-            if self.rbln_config.output_hidden_states:
-                embedding = outputs[0]
-                hidden_states = outputs[1:]
-            else:
-                embedding = outputs
-                hidden_states = None
+        max_seq_len = self.rbln_config.max_seq_lens[idx]
+        padded_inputs_embed = torch.nn.functional.pad(inputs_embeds, (0, 0, 0, max_seq_len - input_len))
+        padded_attn_mask = torch.nn.functional.pad(attention_mask, (0, max_seq_len - input_len)).to(torch.float32)
+        padded_position_ids = torch.arange(max_seq_len, dtype=torch.int32).view(1, -1)
+
+        return {
+            "padded_inputs_embed": padded_inputs_embed,
+            "padded_attn_mask": padded_attn_mask,
+            "padded_position_ids": padded_position_ids,
+        }
 
-            all_embeddings.append(embedding)
-            all_hidden_states.append(hidden_states)
+    def _prepare_inputs_for_iteration(self, index: int, common_inputs, *args, **kwargs):
+        item_kwargs = {
+            "inputs_embeds": common_inputs["padded_inputs_embed"][index : index + 1],
+            "attention_mask": common_inputs["padded_attn_mask"][index : index + 1],
+            "position_ids": common_inputs["padded_position_ids"],
+            "out": [tensor[index : index + 1] for tensor in kwargs["out"]],
+        }
+        return ([], item_kwargs)
 
-        embeddings = torch.cat(all_embeddings, dim=0)[:, :input_seq_len]
+    def _process_outputs(self, outputs: list, **kwargs):
         if self.rbln_config.output_hidden_states:
-            hidden_states = [
-                torch.cat(
-                    [batch_hidden_states[layer_idx][:, :input_seq_len] for batch_hidden_states in all_hidden_states],
-                    dim=0,
-                )
-                for layer_idx in range(len(all_hidden_states[0]))
-            ]
-            return embeddings, tuple(hidden_states)
+            return kwargs["out"][0], tuple(kwargs["out"][1:])
         else:
-            return embeddings
-
-    def __call__(self, *args: Any, **kwds: Any) -> Any:
-        return self.forward(*args, **kwds)
-
-    def __repr__(self) -> str:
-        return repr(self.language_model)
+            return kwargs["out"]
 
 
 class RBLNColPaliForRetrieval(RBLNModel):
@@ -134,8 +98,8 @@ class RBLNColPaliForRetrieval(RBLNModel):
     The ColPali Model transformer for document retrieval using vision-language models.
     This model inherits from [`RBLNModel`]. Check the superclass documentation for the generic methods the library implements for all its models.
 
-    A class to convert and run pre-trained transformers based ColPaliForRetrieval model on RBLN devices.
-    It implements the methods to convert a pre-trained transformers ColPaliForRetrieval model into a RBLN transformer model by:
+    A class to convert and run pre-trained transformers based `ColPaliForRetrieval` model on RBLN devices.
+    It implements the methods to convert a pre-trained transformers `ColPaliForRetrieval` model into a RBLN transformer model by:
 
     - transferring the checkpoint weights of the original into an optimized RBLN graph,
     - compiling the resulting graph using the RBLN compiler.
@@ -217,9 +181,9 @@ class RBLNColPaliForRetrieval(RBLNModel):
         return multi_modal_projector
 
     @classmethod
-    def wrap_model_if_needed(cls, model: "PreTrainedModel", rbln_config: RBLNModelConfig):
+    def _wrap_model_if_needed(cls, model: "PreTrainedModel", rbln_config: RBLNModelConfig):
         return RBLNColPaliForRetrievalWrapper(
-            causal_lm=model.vlm.language_model,
+            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,
@@ -259,9 +223,9 @@ class RBLNColPaliForRetrieval(RBLNModel):
         input_infos = []
         for max_seq_len in rbln_config.max_seq_lens:
             input_info = [
-                ("inputs_embeds", [1, max_seq_len, hidden_size], "float32"),
-                ("attention_mask", [1, max_seq_len], "float32"),
-                ("position_ids", [1, max_seq_len], "int32"),
+                ("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)
 
@@ -271,19 +235,11 @@ class RBLNColPaliForRetrieval(RBLNModel):
         return rbln_config
 
     @classmethod
-    def from_model(cls, model: "PreTrainedModel", *args, **kwargs):
-        if not hasattr(model, "vision_tower"):
+    def _reconstruct_model_if_needed(cls, model: "PreTrainedModel"):
+        if hasattr(model, "vlm"):
             model.vision_tower = model.vlm.vision_tower
-            del model.vlm.vision_tower
-        model = super().from_model(model, *args, **kwargs)
-        return model
-
-    @classmethod
-    def get_pytorch_model(cls, *args, **kwargs):
-        model = super().get_pytorch_model(*args, **kwargs)
-        model.vision_tower = model.vlm.vision_tower
-        del model.vlm.vision_tower
-
+            del model.vlm.model.vision_tower
+            return model
         return model
 
     def get_image_features(self, pixel_values: torch.Tensor):
@@ -294,8 +250,14 @@ class RBLNColPaliForRetrieval(RBLNModel):
         # Returns:
         #     image_features (`torch.Tensor`): Image feature tensor of shape `(num_images, image_length, embed_dim)`).
 
-        vision_outputs = self.vision_tower(pixel_values).last_hidden_state
-        image_features = self.multi_modal_projector(vision_outputs)
+        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
 
@@ -342,7 +304,7 @@ class RBLNColPaliForRetrieval(RBLNModel):
         output_hidden_states: Optional[bool] = None,
         return_dict: Optional[bool] = None,
         **kwargs,
-    ) -> ColPaliForRetrievalOutput:
+    ) -> Union[Tuple, ColPaliForRetrievalOutput]:
         if pixel_values is not None:
             pixel_values = pixel_values.to(dtype=self.dtype)
 
@@ -361,11 +323,27 @@ class RBLNColPaliForRetrieval(RBLNModel):
             input_ids=input_ids, inputs_embeds=inputs_embeds, pixel_values=pixel_values
         )
 
+        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 i 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.
-        outputs = self.language_model(inputs_embeds=inputs_embeds, attention_mask=attention_mask)
+        self.language_model(inputs_embeds=inputs_embeds, attention_mask=attention_mask, out=outputs)
 
-        embeddings = outputs if not self.rbln_config.output_hidden_states else outputs[0]
-        hidden_states = None if not self.rbln_config.output_hidden_states else outputs[1]
+        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)

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

@@ -0,0 +1,2 @@
+from .configuration_colqwen2 import RBLNColQwen2ForRetrievalConfig
+from .modeling_colqwen2 import RBLNColQwen2ForRetrieval