optimum-rbln 0.8.4a8__py3-none-any.whl → 0.9.2__py3-none-any.whl

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

@@ -26,6 +26,7 @@ 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
 
 
@@ -33,93 +34,64 @@ if TYPE_CHECKING:
     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):
@@ -212,7 +184,7 @@ class RBLNColPaliForRetrieval(RBLNModel):
     @classmethod
     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,
@@ -252,9 +224,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)
 
@@ -298,7 +270,7 @@ class RBLNColPaliForRetrieval(RBLNModel):
         """
         if not hasattr(model, "vision_tower"):
             model.vision_tower = model.vlm.vision_tower
-            del model.vlm.vision_tower
+            del model.vlm.model.vision_tower
         model = super().from_model(model, config, rbln_config, model_save_dir, subfolder, **kwargs)
         return model
 
@@ -306,8 +278,7 @@ class RBLNColPaliForRetrieval(RBLNModel):
     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
 
     def get_image_features(self, pixel_values: torch.Tensor):
@@ -318,8 +289,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
 
@@ -385,11 +362,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

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

@@ -0,0 +1,233 @@
+# Copyright 2025 Rebellions Inc. All rights reserved.
+
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at:
+
+#     http://www.apache.org/licenses/LICENSE-2.0
+
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# 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 List, Optional, Tuple, Union
+
+import torch
+import torch.nn as nn
+from transformers import PreTrainedModel
+
+from optimum.rbln.transformers.models.decoderonly.decoderonly_architecture import (
+    DecoderOnlyLayer,
+    DecoderOnlyModel,
+    DecoderOnlyWrapper,
+)
+
+from .configuration_colqwen2 import (
+    RBLNColQwen2ForRetrievalConfig,
+)
+
+
+def slice_and_unsqueeze_cos_sin(cos, sin, position_ids):
+    """Slice cos[cache_position], sin[cache_position] vector for the query."""
+    cos = cos[position_ids[0]][None, None, None, :, :]
+    sin = sin[position_ids[0]][None, None, None, :, :]
+
+    return cos, sin
+
+
+class ColQwen2LanguageModelWrapper(DecoderOnlyWrapper):
+    def __init__(
+        self, model: PreTrainedModel, rbln_config: "RBLNColQwen2ForRetrievalConfig", use_rotary_emb: bool = True
+    ):
+        model.config = (
+            model.config.vlm_config.text_config if hasattr(model.config, "vlm_config") else model.config.text_config
+        )
+        super().__init__(model, rbln_config, use_rotary_emb)
+
+    def get_decoder_layers(self, model: PreTrainedModel):
+        return model.language_model.layers
+
+    def convert_to_rbln_class(self, model: PreTrainedModel, max_seq_len: int):
+        new_layers = []
+        for layer_idx, layer in enumerate(self.get_decoder_layers(model)):
+            is_sliding = layer_idx in self.rbln_config.sliding_window_layers
+            new_self_attn = self.get_rbln_attn_class()(
+                self.get_attn_layer(layer),
+                self.rbln_config,
+                is_sliding=is_sliding,
+            )
+            new_layer = self.get_rbln_layer_class()(layer, new_self_attn)
+            new_layers.append(new_layer)
+
+        new_model = self.get_rbln_model_class()(
+            model.language_model,
+            new_layers,
+            self.rbln_config,
+            use_learned_pos_emb=self.__class__._use_learned_pos_emb,
+        )
+
+        # text_projection layer from model
+        self.embedding_proj_layer = (
+            model.embedding_proj_layer if hasattr(model, "embedding_proj_layer") else model.custom_text_proj
+        )
+        return new_model
+
+    def get_rbln_model_class(self):
+        return RBLNColQwen2LanguageModel
+
+    def prepare_forward_args(self, *args):
+        args = list(args)
+        input_ids = None if self.rbln_config.use_inputs_embeds else args.pop(0)
+        inputs_embeds = args.pop(0) if self.rbln_config.use_inputs_embeds else None
+        cache_position = args.pop(0)
+        global_block_tables = args.pop(0)
+        local_block_tables = None
+        position_embeds = args.pop(0)
+        position_ids = None
+        attention_mask = args.pop(0) if self.rbln_config.use_attention_mask else None
+        past_key_values = args
+
+        if len(past_key_values) != 2 * self.num_hidden_layers:
+            raise ValueError(
+                f"Different past_key_values to model's config. {len(past_key_values)} != {2 * self.num_hidden_layers}"
+            )
+
+        _past_key_values = []
+        for i in range(self.config.num_hidden_layers):
+            key_states = past_key_values[i * 2]
+            value_states = past_key_values[i * 2 + 1]
+            past_key_value = [key_states, value_states]
+            _past_key_values.append(past_key_value)
+        past_key_values = _past_key_values
+
+        return (
+            input_ids,
+            inputs_embeds,
+            cache_position,
+            global_block_tables,
+            local_block_tables,
+            attention_mask,
+            position_ids,
+            past_key_values,
+            position_embeds,
+        )
+
+    def forward(self, *args):
+        (
+            input_ids,
+            inputs_embeds,
+            cache_position,
+            global_block_tables,
+            local_block_tables,
+            attention_mask,
+            position_ids,
+            past_key_values,
+            rotary_emb,
+        ) = self.prepare_forward_args(*args)
+
+        last_hidden_states = self.model(
+            input_ids=input_ids,
+            inputs_embeds=inputs_embeds,
+            attention_mask=attention_mask,
+            cache_position=cache_position,
+            position_ids=position_ids,
+            past_key_values=past_key_values,
+            rotary_emb=rotary_emb,
+            global_block_tables=global_block_tables,
+            local_block_tables=local_block_tables,
+        )
+
+        proj = self.embedding_proj_layer(last_hidden_states[0])
+        all_hidden_states = last_hidden_states[1] if self.rbln_config.output_hidden_states else None
+
+        if self.rbln_config.output_hidden_states:
+            return proj, all_hidden_states
+        else:
+            return proj
+
+
+class RBLNColQwen2LanguageModel(DecoderOnlyModel):
+    def __init__(
+        self,
+        model,
+        layers: List["DecoderOnlyLayer"],
+        rbln_config: "RBLNColQwen2ForRetrievalConfig",
+        use_learned_pos_emb=None,
+    ):
+        super().__init__(model, layers, rbln_config, use_learned_pos_emb)
+
+        self.output_hidden_states = rbln_config.output_hidden_states
+
+    def forward(
+        self,
+        input_ids: torch.Tensor = None,
+        inputs_embeds: Optional[torch.Tensor] = None,
+        attention_mask: torch.Tensor = None,
+        cache_position: torch.Tensor = None,
+        position_ids: torch.Tensor = None,
+        query_position: torch.Tensor = None,
+        past_key_values: Tuple[Tuple[torch.Tensor]] = None,
+        rotary_emb: Optional[Union[nn.Module, torch.Tensor]] = None,
+        global_block_tables: Optional[torch.Tensor] = None,
+        local_block_tables: Optional[torch.Tensor] = None,
+        lora_int_id: Optional[torch.Tensor] = None,
+    ):
+        # retrieve input_ids and inputs_embeds
+        if (input_ids is None) ^ (inputs_embeds is not None):
+            raise ValueError(
+                "You cannot specify both input_ids and inputs_embeds at the same time, and must specify either one"
+            )
+
+        # embed positions
+        if inputs_embeds is None:
+            inputs_embeds = self.get_embedding()(input_ids)
+
+        hidden_states = inputs_embeds * self.hidden_multiplier
+
+        # get cos,sin vector if needed
+        position_ids = position_ids if position_ids is not None else cache_position
+        if rotary_emb is not None:
+            if isinstance(rotary_emb, torch.Tensor):
+                cos = rotary_emb[0]
+                sin = rotary_emb[1]
+            else:
+                cos, sin = rotary_emb(hidden_states, self.max_seq_len)  # dtype carrier, max_seq_len
+                cos, sin = slice_and_unsqueeze_cos_sin(cos, sin, position_ids)
+
+        # Get sequence positions for flash attention
+        if self.attn_impl == "flash_attn":
+            seq_positions = cache_position[:, 0]
+            seq_positions = self.convert_sequence_positions_for_flash_attn(
+                seq_positions=seq_positions, max_seq_len=self.max_seq_len
+            )
+        else:
+            seq_positions = cache_position[:, :1]
+
+        # Get local cache positions for sliding window layers
+        if len(self.sliding_window_layers) > 0:
+            sliding_cache_pos = self.get_local_cache_positions(position_ids, query_position)
+
+        all_hidden_states = () if self.output_hidden_states else None
+        for layer_idx, layer in enumerate(self.layers):
+            if self.output_hidden_states:
+                all_hidden_states += (hidden_states,)
+
+            is_sliding = True if layer_idx in self.sliding_window_layers else False
+            hidden_states = layer(
+                hidden_states=hidden_states,
+                attention_mask=attention_mask,
+                seq_positions=sliding_cache_pos if is_sliding else seq_positions,
+                past_key_values=past_key_values,
+                cos=cos,
+                sin=sin,
+                block_tables=local_block_tables if is_sliding else global_block_tables,
+                lora_int_id=lora_int_id,
+            )
+
+        hidden_states = self.get_last_layernorm()(hidden_states)
+        if self.output_hidden_states:
+            all_hidden_states += (hidden_states,)
+
+        return hidden_states, all_hidden_states

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

@@ -0,0 +1,74 @@
+# Copyright 2025 Rebellions Inc. All rights reserved.
+
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at:
+
+#     http://www.apache.org/licenses/LICENSE-2.0
+
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# 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 Optional
+
+from optimum.rbln.configuration_utils import RBLNModelConfig
+
+from ..decoderonly.configuration_decoderonly import RBLNDecoderOnlyModelConfig
+
+
+class RBLNColQwen2ForRetrievalConfig(RBLNDecoderOnlyModelConfig):
+    """
+    Configuration class for RBLN ColQwen2 models for document retrieval.
+
+    This class extends RBLNModelConfig with specific configurations for ColQwen2 models,
+    including vision tower settings and multi-sequence length support.
+
+    Example usage:
+        ```python
+        from optimum.rbln import RBLNColQwen2ForRetrievalConfig, RBLNColQwen2ForRetrievalConfig
+
+        # 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,
+        )
+
+        # Use the configuration with from_pretrained
+        model = RBLNColQwen2ForRetrieval.from_pretrained(
+            "vidore/colqwen2-v1.0-hf",
+            export=True,
+            rbln_config=config
+        )
+        ```
+    """
+
+    submodules = ["visual"]
+
+    def __init__(
+        self,
+        visual: Optional[RBLNModelConfig] = None,
+        batch_size: Optional[int] = None,
+        use_inputs_embeds: bool = True,
+        output_hidden_states: Optional[bool] = False,
+        **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
+        self.output_hidden_states = output_hidden_states