optimum-rbln 0.8.1rc1__py3-none-any.whl → 0.8.2__py3-none-any.whl

optimum/rbln/transformers/models/gpt2/configuration_gpt2.py

@@ -12,11 +12,39 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from ..decoderonly.configuration_decoderonly import RBLNDecoderOnlyModelForCausalLMConfig
+from ..decoderonly.configuration_decoderonly import RBLNDecoderOnlyModelConfig, RBLNDecoderOnlyModelForCausalLMConfig
 
 
 class RBLNGPT2LMHeadModelConfig(RBLNDecoderOnlyModelForCausalLMConfig):
     """
-    Configuration class for GPT-2 causal language model.
-    Inherits from RBLNDecoderOnlyModelForCausalLMConfig with no additional parameters.
+    Configuration class for RBLN GPT2 models.
+
+    This class is an alias of RBLNDecoderOnlyModelForCausalLMConfig.
+    """
+
+
+class RBLNGPT2ModelConfig(RBLNDecoderOnlyModelConfig):
+    """
+    Configuration class for RBLN GPT2 models.
+
+    This class is an alias of RBLNDecoderOnlyModelConfig.
+
+    Example usage:
+    ```python
+    from optimum.rbln import RBLNGPT2Model, RBLNGPT2ModelConfig
+
+    # Create a configuration object
+    config = RBLNGPT2ModelConfig(
+        batch_size=1,
+        max_seq_len=1024,
+        tensor_parallel_size=4
+    )
+
+    # Use the configuration with from_pretrained
+    model = RBLNGPT2Model.from_pretrained(
+        "openai/gpt2",
+        export=True,
+        rbln_config=config
+    )
+    ```
     """

optimum/rbln/transformers/models/gpt2/gpt2_architecture.py

@@ -13,14 +13,13 @@
 # limitations under the License.
 
 import math
-from typing import TYPE_CHECKING, Tuple
+from typing import TYPE_CHECKING, Tuple, Union
 
 import torch
 import torch.nn as nn
 
 from ..decoderonly.decoderonly_architecture import (
     DecoderOnlyAttention,
-    DecoderOnlyForCausalLM,
     DecoderOnlyLayer,
     DecoderOnlyModel,
     DecoderOnlyWrapper,
@@ -28,31 +27,27 @@ from ..decoderonly.decoderonly_architecture import (
 
 
 if TYPE_CHECKING:
-    from transformers import GPT2LMHeadModel
+    from transformers import GPT2LMHeadModel, GPT2Model
 
 
 class GPT2Wrapper(DecoderOnlyWrapper):
-    def convert_to_rbln_causal_lm(self, causal_lm: "GPT2LMHeadModel", max_seq_len: int):
-        if self.attn_impl != "eager":
-            raise NotImplementedError(f"flash attention ({self.attn_impl}) is not implemented for {self.__class__}")
-        new_layers = []
-        for layer in causal_lm.transformer.h:
-            new_self_attn = GPT2Attention(
-                layer.attn,
-                self.use_attention_mask,
-                kvcache_block_size=self.kvcache_block_size,
-                use_position_ids=self.use_position_ids,
-            )
-            new_layer = GPT2Layer(layer, new_self_attn)
-            new_layers.append(new_layer)
-        new_model = GPT2Model(
-            causal_lm.transformer,
-            new_layers,
-            max_seq_len=max_seq_len,
-            sliding_window_layers=self.sliding_window_layers,
-        )
-        new_causal_lm = DecoderOnlyForCausalLM(causal_lm, new_model)
-        return new_causal_lm
+    def get_rbln_attn_class(self):
+        return GPT2Attention
+
+    def get_rbln_layer_class(self):
+        return GPT2Layer
+
+    def get_rbln_model_class(self):
+        return GPT2Model
+
+    def get_attn_layer(self, layer: nn.Module):
+        return layer.attn
+
+    def get_model_layer(self, model: Union["GPT2LMHeadModel", "GPT2Model"]):
+        return model.transformer if self.is_causal_lm else model
+
+    def get_decoder_layers(self, model: Union["GPT2LMHeadModel", "GPT2Model"]):
+        return model.transformer.h if self.is_causal_lm else model.h
 
 
 class GPT2Model(DecoderOnlyModel):

optimum/rbln/transformers/models/gpt2/modeling_gpt2.py

@@ -13,7 +13,7 @@
 # limitations under the License.
 
 from ....utils import logging
-from ...models.decoderonly import RBLNDecoderOnlyModelForCausalLM
+from ...models.decoderonly import RBLNDecoderOnlyModel, RBLNDecoderOnlyModelForCausalLM
 from .gpt2_architecture import GPT2Wrapper
 
 
@@ -36,3 +36,20 @@ class RBLNGPT2LMHeadModel(RBLNDecoderOnlyModelForCausalLM):
 
     _decoder_wrapper_cls = GPT2Wrapper
     _use_rotary_emb = False
+
+
+class RBLNGPT2Model(RBLNDecoderOnlyModel):
+    """
+    The GPT2 Model transformer without a language modeling head.
+
+    This model inherits from [`RBLNDecoderOnlyModel`]. Check the superclass documentation for the generic methods the
+    library implements for all its model.
+
+    A class to convert and run pre-trained transformers based GPT2Model model on RBLN devices.
+    It implements the methods to convert a pre-trained transformers GPT2Model 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.
+    """
+
+    _decoder_wrapper_cls = GPT2Wrapper
+    _use_rotary_emb = False

optimum/rbln/transformers/models/idefics3/configuration_idefics3.py

@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import Any, Dict, Optional
+from typing import Any, Optional
 
 from ....configuration_utils import RBLNModelConfig
 
@@ -39,7 +39,7 @@ class RBLNIdefics3ForConditionalGenerationConfig(RBLNModelConfig):
         batch_size: Optional[int] = None,
         vision_model: Optional[RBLNModelConfig] = None,
         text_model: Optional[RBLNModelConfig] = None,
-        **kwargs: Dict[str, Any],
+        **kwargs: Any,
     ):
         """
         Args:

optimum/rbln/transformers/models/idefics3/modeling_idefics3.py

@@ -34,17 +34,11 @@ from transformers.models.idefics3.modeling_idefics3 import Idefics3CausalLMOutpu
 from ....configuration_utils import RBLNCompileConfig, RBLNModelConfig
 from ....modeling import RBLNModel
 from ....utils.runtime_utils import RBLNPytorchRuntime
-from ..decoderonly.modeling_decoderonly import (
-    RBLNDecoderOnlyOutput,
-)
+from ..decoderonly.modeling_decoderonly import RBLNDecoderOnlyForCausalLMOutput
 
 
 if TYPE_CHECKING:
-    from transformers import (
-        AutoFeatureExtractor,
-        AutoProcessor,
-        AutoTokenizer,
-    )
+    from transformers import AutoFeatureExtractor, AutoProcessor, AutoTokenizer
 
 
 class RBLNRuntimeVisionModel(RBLNPytorchRuntime):
@@ -494,7 +488,7 @@ class RBLNIdefics3ForConditionalGeneration(RBLNModel):
         if not return_dict:
             return logits, generate_idx
         else:
-            return RBLNDecoderOnlyOutput(
+            return RBLNDecoderOnlyForCausalLMOutput(
                 logits=logits,
                 generate_idx=generate_idx,
             )

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

@@ -12,5 +12,5 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from .configuration_llama import RBLNLlamaForCausalLMConfig
-from .modeling_llama import RBLNLlamaForCausalLM
+from .configuration_llama import RBLNLlamaForCausalLMConfig, RBLNLlamaModelConfig
+from .modeling_llama import RBLNLlamaForCausalLM, RBLNLlamaModel

optimum/rbln/transformers/models/llama/configuration_llama.py

@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from ..decoderonly.configuration_decoderonly import RBLNDecoderOnlyModelForCausalLMConfig
+from ..decoderonly.configuration_decoderonly import RBLNDecoderOnlyModelConfig, RBLNDecoderOnlyModelForCausalLMConfig
 
 
 class RBLNLlamaForCausalLMConfig(RBLNDecoderOnlyModelForCausalLMConfig):
@@ -40,3 +40,11 @@ class RBLNLlamaForCausalLMConfig(RBLNDecoderOnlyModelForCausalLMConfig):
     )
     ```
     """
+
+
+class RBLNLlamaModelConfig(RBLNDecoderOnlyModelConfig):
+    """
+    Configuration class for RBLN Llama models.
+
+    This class is an alias of RBLNDecoderOnlyModelConfig.
+    """

optimum/rbln/transformers/models/llama/modeling_llama.py

@@ -13,7 +13,7 @@
 # limitations under the License.
 
 from ....utils import logging
-from ...models.decoderonly import RBLNDecoderOnlyModelForCausalLM
+from ...models.decoderonly import RBLNDecoderOnlyModel, RBLNDecoderOnlyModelForCausalLM
 from .llama_architecture import LlamaWrapper
 
 
@@ -81,3 +81,24 @@ class RBLNLlamaForCausalLM(RBLNDecoderOnlyModelForCausalLM):
     """
 
     _decoder_wrapper_cls = LlamaWrapper
+
+
+class RBLNLlamaModel(RBLNDecoderOnlyModel):
+    """
+    The Llama Model transformer outputting raw hidden-states without any specific head on top.
+    This model inherits from [`RBLNDecoderOnlyModel`]. 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 LlamaModel on RBLN devices.
+    It implements the methods to convert a pre-trained transformers LlamaModel 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.
+
+    **Configuration:**
+    This model uses [`RBLNLlamaModelConfig`] for configuration. When calling methods like `from_pretrained` or `from_model`,
+    the `rbln_config` parameter should be an instance of [`RBLNLlamaModelConfig`] or a dictionary conforming to its structure.
+
+    See the [`RBLNLlamaModelConfig`] class for all available configuration options.
+    """
+
+    _decoder_wrapper_cls = LlamaWrapper

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

@@ -0,0 +1,16 @@
+# 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 .configuration_llava import RBLNLlavaForConditionalGenerationConfig
+from .modeling_llava import RBLNLlavaForConditionalGeneration

optimum/rbln/transformers/models/llava/configuration_llava.py

@@ -0,0 +1,54 @@
+# 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 Any, Optional
+
+from ....configuration_utils import RBLNModelConfig
+
+
+class RBLNLlavaForConditionalGenerationConfig(RBLNModelConfig):
+    """
+    Configuration class for RBLNLlavaForConditionalGenerationConfig.
+
+    This configuration class stores the configuration parameters specific to
+    RBLN-optimized LLaVA models for multimodal conditional generation tasks
+    that combine vision and language processing capabilities.
+    """
+
+    submodules = ["vision_tower", "language_model"]
+
+    def __init__(
+        self,
+        batch_size: Optional[int] = None,
+        vision_tower: Optional[RBLNModelConfig] = None,
+        language_model: Optional[RBLNModelConfig] = None,
+        **kwargs: Any,
+    ):
+        """
+        Args:
+            batch_size (Optional[int]): The batch size for inference. Defaults to 1.
+            vision_tower (Optional[RBLNModelConfig]): Configuration for the vision encoder component.
+            language_model (Optional[RBLNModelConfig]): Configuration for the language model 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
+        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.vision_tower = vision_tower
+        self.language_model = language_model