optimum-rbln 0.1.4__py3-none-any.whl → 0.1.8__py3-none-any.whl

optimum/rbln/modeling_seq2seq.py

@@ -23,13 +23,10 @@
 
 import inspect
 import logging
-from pathlib import Path
-from tempfile import TemporaryDirectory
 from typing import TYPE_CHECKING, Any, Callable, Dict, List, Optional, Tuple, Union
 
 import rebel
 import torch
-from optimum.exporters import TasksManager
 from transformers import (
     AutoModelForSeq2SeqLM,
     BartConfig,
@@ -39,12 +36,11 @@ from transformers import (
 )
 from transformers.modeling_outputs import BaseModelOutput, Seq2SeqLMOutput
 
-from .modeling_base import RBLNBaseModel
+from .modeling_base import RBLNModel
 from .modeling_config import DEFAULT_COMPILED_MODEL_NAME, RBLNConfig, RBLNRuntimeConfig
 from .transformers.models.bart import BartDecoderWrapper, BartEncoderWrapper
 from .transformers.models.t5 import T5DecoderWrapper, T5EncoderWrapper
 from .utils.runtime_utils import RBLNPytorchRuntime
-from .utils.save_utils import maybe_save_preprocessors
 
 
 logger = logging.getLogger(__name__)
@@ -75,7 +71,7 @@ class RBLNRuntimeDecoder(RBLNPytorchRuntime):
         return Seq2SeqLMOutput(logits=outputs)
 
 
-class RBLNModelForSeq2SeqLM(RBLNBaseModel):
+class RBLNModelForSeq2SeqLM(RBLNModel):
     """
     This is a generic model class that will be instantiated as one of the model classes of the library (with a sequence-to-sequence language modeling head) when created with the from_pretrained() class method.
     This model inherits from [`RBLNBaseModel`]. Check the superclass documentation for the generic methods the library implements for all its models.
@@ -88,7 +84,6 @@ class RBLNModelForSeq2SeqLM(RBLNBaseModel):
     Currently, this model class only supports the 'bart' and 't5' models from the transformers library. Future updates may include support for additional model types.
     """
 
-    model_type = "rbln_model"
     auto_model_class = AutoModelForSeq2SeqLM
 
     def __post_init__(self, **kwargs):
@@ -97,8 +92,8 @@ class RBLNModelForSeq2SeqLM(RBLNBaseModel):
         self.enc_max_seq_len = self.rbln_config.meta["rbln_enc_max_seq_len"]
         self.dec_max_seq_len = self.rbln_config.meta["rbln_dec_max_seq_len"]
         self.pad_token_id = self.rbln_config.meta["rbln_pad_token_id"]
-        self.encoder = RBLNRuntimeEncoder(runtime=self.runtimes[0], main_input_name="input_ids")
-        self.decoder = RBLNRuntimeDecoder(runtime=self.runtimes[1], main_input_name="input_ids")
+        self.encoder = RBLNRuntimeEncoder(runtime=self.model[0], main_input_name="input_ids")
+        self.decoder = RBLNRuntimeDecoder(runtime=self.model[1], main_input_name="input_ids")
 
     def can_generate(self):
         return True
@@ -149,74 +144,18 @@ class RBLNModelForSeq2SeqLM(RBLNBaseModel):
         }
 
     @classmethod
-    def _export(
-        cls,
-        model_id: str,
-        config: "PretrainedConfig",
-        use_auth_token: Optional[Union[bool, str]] = None,
-        revision: Optional[str] = None,
-        force_download: bool = False,
-        cache_dir: Optional[str] = None,
-        subfolder: str = "",
-        local_files_only: bool = False,
-        trust_remote_code: bool = False,
-        model_save_dir: Optional[Union[str, Path, TemporaryDirectory]] = None,
-        **kwargs,
-    ) -> "AutoModelForSeq2SeqLM":
-        """
-        Exports a vanilla Transformers model into a rbln-compiled Module.
-        """
-        task = kwargs.pop("task", None)
-        if task is None:
-            task = TasksManager.infer_task_from_model(cls.auto_model_class)
-
-        if model_save_dir is None:
-            save_dir = TemporaryDirectory()
-            save_dir_path = Path(save_dir.name)
-        else:
-            save_dir = model_save_dir
-            if isinstance(save_dir, TemporaryDirectory):
-                save_dir_path = Path(model_save_dir.name)
-            else:
-                save_dir_path = Path(model_save_dir)
-                save_dir_path.mkdir(exist_ok=True)
-
+    def update_kwargs(cls, kwargs):
         kwargs.update(
             {
                 "torchscript": True,
                 "return_dict": False,
-                "use_cache": False,
+                "use_cache": True,
             }
         )
+        return kwargs
 
-        rbln_config_kwargs, rbln_constructor_kwargs = cls.pop_rbln_kwargs_from_kwargs(kwargs)
-
-        model: AutoModelForSeq2SeqLM = TasksManager.get_model_from_task(
-            task=task,
-            model_name_or_path=model_id,
-            subfolder=subfolder,
-            revision=revision,
-            framework="pt",
-            cache_dir=cache_dir,
-            use_auth_token=use_auth_token,
-            local_files_only=local_files_only,
-            force_download=force_download,
-            trust_remote_code=trust_remote_code,
-            **kwargs,
-        )
-
-        if config is None:
-            config = model.config
-
-        config.save_pretrained(save_dir_path)
-        preprocessors = maybe_save_preprocessors(model_id, save_dir_path, src_subfolder=subfolder)
-
-        # Get compilation arguments
-        if rbln_config_kwargs.get("rbln_config", None) is None:
-            rbln_config = cls.get_rbln_config(
-                preprocessors=preprocessors, model_config=model.config, **rbln_config_kwargs
-            )
-
+    @classmethod
+    def get_compiled_model(cls, model, rbln_config: RBLNConfig):
         def optimized_models(model):
             if isinstance(model, T5ForConditionalGeneration):
                 encoder_model = T5EncoderWrapper(model).eval()
@@ -229,67 +168,54 @@ class RBLNModelForSeq2SeqLM(RBLNBaseModel):
 
             return encoder_model, decoder_model
 
-        def compile():
-            wrapped_encoder, wrapped_decoder = optimized_models(model)
+        wrapped_encoder, wrapped_decoder = optimized_models(model)
 
-            wrapped_encoder.encoder_max_length = rbln_config.meta["rbln_enc_max_seq_len"]
-            wrapped_encoder.decoder_max_length = rbln_config.meta["rbln_dec_max_seq_len"]
-            wrapped_encoder.decoder_batch_size = rbln_config.meta["rbln_batch_size"]
+        wrapped_encoder.encoder_max_length = rbln_config.meta["rbln_enc_max_seq_len"]
+        wrapped_encoder.decoder_max_length = rbln_config.meta["rbln_dec_max_seq_len"]
+        wrapped_encoder.decoder_batch_size = rbln_config.meta["rbln_batch_size"]
 
-            wrapped_decoder.encoder_max_length = rbln_config.meta["rbln_enc_max_seq_len"]
-            wrapped_decoder.decoder_max_length = rbln_config.meta["rbln_dec_max_seq_len"]
-            wrapped_decoder.decoder_batch_size = rbln_config.meta["rbln_batch_size"]
+        wrapped_decoder.encoder_max_length = rbln_config.meta["rbln_enc_max_seq_len"]
+        wrapped_decoder.decoder_max_length = rbln_config.meta["rbln_dec_max_seq_len"]
+        wrapped_decoder.decoder_batch_size = rbln_config.meta["rbln_batch_size"]
 
-            enc_rbln_runtime_config = rbln_config[DEFAULT_COMPILED_MODEL_NAME][0]
-            dec_rbln_runtime_config = rbln_config[DEFAULT_COMPILED_MODEL_NAME][1]
+        enc_rbln_runtime_config = rbln_config[DEFAULT_COMPILED_MODEL_NAME][0]
+        dec_rbln_runtime_config = rbln_config[DEFAULT_COMPILED_MODEL_NAME][1]
 
-            if isinstance(model, T5ForConditionalGeneration):
-                enc_example_inputs = enc_rbln_runtime_config.get_dummy_inputs(fill=1)
-                dec_example_inputs = dec_rbln_runtime_config.get_dummy_inputs(fill=1)
-            else:
-                enc_example_inputs = enc_rbln_runtime_config.get_dummy_inputs(fill=0)
-                dec_example_inputs = dec_rbln_runtime_config.get_dummy_inputs(fill=0)
-
-            enc_scripted_model = torch.jit.trace(wrapped_encoder, enc_example_inputs)
-            dec_scripted_model = torch.jit.trace(wrapped_decoder, dec_example_inputs)
-
-            enc_ir = rebel.torchscript_to_ir(
-                enc_scripted_model,
-                input_names=[v[0] for v in enc_rbln_runtime_config.input_info],
-                name=enc_rbln_runtime_config.rbln_mod_name,
-            )
-            dec_ir = rebel.torchscript_to_ir(
-                dec_scripted_model,
-                input_names=[v[0] for v in dec_rbln_runtime_config.input_info],
-                name=dec_rbln_runtime_config.rbln_mod_name,
-            )
-            dec_ir.decoder_batch_size = rbln_config.meta["rbln_batch_size"]
-
-            connections = [
-                (enc_ir.outputs[0], dec_ir.inputs[5]),
-                (dec_ir.outputs[1], dec_ir.inputs[4]),
-            ]
-            compiled_model = rebel.compile(
-                enc_ir,
-                dec_ir,
-                connections=connections,
-                fusion=enc_rbln_runtime_config.fusion,
-                npu=enc_rbln_runtime_config.npu,
-                tensor_parallel_size=enc_rbln_runtime_config.tensor_parallel_size,
-            )
-            compiled_model.save(save_dir_path / f"{DEFAULT_COMPILED_MODEL_NAME}.rbln")
-
-        compile()
-
-        rbln_config.save(save_dir_path)
-
-        return cls._from_pretrained(
-            model_id=save_dir_path,
-            config=config,
-            model_save_dir=save_dir,
-            **rbln_constructor_kwargs,
-            **kwargs,
+        if isinstance(model, T5ForConditionalGeneration):
+            enc_example_inputs = enc_rbln_runtime_config.get_dummy_inputs(fill=1)
+            dec_example_inputs = dec_rbln_runtime_config.get_dummy_inputs(fill=1)
+        else:
+            enc_example_inputs = enc_rbln_runtime_config.get_dummy_inputs(fill=0)
+            dec_example_inputs = dec_rbln_runtime_config.get_dummy_inputs(fill=0)
+
+        enc_scripted_model = torch.jit.trace(wrapped_encoder, enc_example_inputs, check_trace=False)
+        dec_scripted_model = torch.jit.trace(wrapped_decoder, dec_example_inputs, check_trace=False)
+
+        enc_ir = rebel.torchscript_to_ir(
+            enc_scripted_model,
+            input_names=[v[0] for v in enc_rbln_runtime_config.input_info],
+            name=enc_rbln_runtime_config.rbln_mod_name,
+        )
+        dec_ir = rebel.torchscript_to_ir(
+            dec_scripted_model,
+            input_names=[v[0] for v in dec_rbln_runtime_config.input_info],
+            name=dec_rbln_runtime_config.rbln_mod_name,
         )
+        dec_ir.decoder_batch_size = rbln_config.meta["rbln_batch_size"]
+
+        connections = [
+            (enc_ir.outputs[0], dec_ir.inputs[5]),
+            (dec_ir.outputs[1], dec_ir.inputs[4]),
+        ]
+        compiled_model = rebel.compile(
+            enc_ir,
+            dec_ir,
+            connections=connections,
+            fusion=enc_rbln_runtime_config.fusion,
+            npu=enc_rbln_runtime_config.npu,
+            tensor_parallel_size=enc_rbln_runtime_config.tensor_parallel_size,
+        )
+        return compiled_model
 
     @classmethod
     def _get_rbln_config(
@@ -411,11 +337,14 @@ class RBLNModelForSeq2SeqLM(RBLNBaseModel):
 
         return rbln_config
 
-    def _create_runtimes(self, rbln_device_map: Dict[str, int]) -> List[rebel.Runtime]:
+    @classmethod
+    def _create_runtimes(
+        cls, compiled_models: List[rebel.RBLNCompiledModel], rbln_device_map: Dict[str, int]
+    ) -> List[rebel.Runtime]:
         device_val = rbln_device_map[DEFAULT_COMPILED_MODEL_NAME]
         return [
-            self.compiled_models[0].create_runtime("encoder", tensor_type="pt", device=device_val),
-            self.compiled_models[0].create_runtime("decoder", tensor_type="pt", device=device_val),
+            compiled_models[0].create_runtime("encoder", tensor_type="pt", device=device_val),
+            compiled_models[0].create_runtime("decoder", tensor_type="pt", device=device_val),
         ]
 
     def forward(
@@ -436,9 +365,6 @@ class RBLNModelForSeq2SeqLM(RBLNBaseModel):
 
         return Seq2SeqLMOutput(logits=lm_logits)
 
-    def __repr__(self):
-        return repr(self.runtimes[0]) + "\n" + repr(self.runtimes[1])
-
     def _prepare_encoder_decoder_kwargs_for_generation(
         self,
         inputs_tensor: torch.Tensor,

optimum/rbln/transformers/__init__.py

@@ -27,28 +27,36 @@ from transformers.utils import _LazyModule
 
 
 _import_structure = {
+    "cache_utils": ["RebelDynamicCache"],
     "generation": ["BatchTextIteratorStreamer"],
     "models": [
         "RBLNCLIPTextModel",
         "RBLNCLIPTextModelWithProjection",
+        "RBLNDPTForDepthEstimation",
+        "RBLNGemmaForCausalLM",
         "RBLNGPT2LMHeadModel",
         "RBLNWav2Vec2ForCTC",
         "RBLNWhisperForConditionalGeneration",
         "RBLNLlamaForCausalLM",
         "RBLNMidmLMHeadModel",
+        "RBLNXLMRobertaModel"
     ],
 }
 
 if TYPE_CHECKING:
+    from .cache_utils import RebelDynamicCache
     from .generation import BatchTextIteratorStreamer
     from .models import (
         RBLNCLIPTextModel,
         RBLNCLIPTextModelWithProjection,
+        RBLNDPTForDepthEstimation,
+        RBLNGemmaForCausalLM,
         RBLNGPT2LMHeadModel,
         RBLNLlamaForCausalLM,
         RBLNMidmLMHeadModel,
         RBLNWav2Vec2ForCTC,
         RBLNWhisperForConditionalGeneration,
+        RBLNXLMRobertaModel,
     )
 else:
     import sys

optimum/rbln/transformers/cache_utils.py

@@ -0,0 +1,111 @@
+from typing import Optional, Tuple
+
+import torch
+from transformers.cache_utils import DynamicCache
+
+
+class RebelDynamicCache(DynamicCache):
+    """
+    A cache that grows dynamically as more tokens are generated. This is the default for generative models.
+
+    It stores the Key and Value states as a list of tensors, one for each layer. The expected shape for each tensor is
+    `[batch_size, num_heads, seq_len, head_dim]`.
+    """
+
+    def __init__(self, current_steps) -> None:
+        super().__init__()
+        self.current_steps = current_steps
+
+    def assign(
+        self,
+        key_states: torch.Tensor,
+        value_states: torch.Tensor,
+        layer_idx: int,
+    ) -> None:
+        self.key_cache[layer_idx] = key_states.squeeze(2)
+        self.value_cache[layer_idx] = value_states.squeeze(2)
+
+    def update(
+        self,
+        key_states: torch.Tensor,
+        value_states: torch.Tensor,
+        layer_idx: int,
+        batch_idx: int,
+        read_first_step: Optional[bool] = False,
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        """
+        Updates the cache with the new `key_states` and `value_states` for the layer `layer_idx` and the batch 'batch_inx'
+        based on self.current_step,
+        """
+        current_step = self.current_steps[0 if read_first_step else batch_idx]
+        kend = current_step + key_states.shape[-2]
+        vend = current_step + value_states.shape[-2]
+        update_key_states = (
+            self.key_cache[layer_idx][batch_idx]
+            .unsqueeze(0)
+            .unsqueeze(2)
+            .slice_scatter(key_states, dim=-2, start=current_step, end=kend)
+        )
+        update_value_states = (
+            self.value_cache[layer_idx][batch_idx]
+            .unsqueeze(0)
+            .unsqueeze(2)
+            .slice_scatter(value_states, dim=-2, start=current_step, end=vend)
+        )
+
+        return update_key_states, update_value_states
+
+    @classmethod
+    def from_input_format(cls, position_ids, num_hidden_layer, *past_key_values) -> "DynamicCache":
+        """Converts a cache in the rbln cache format (list of past_kv) into an equivalent `DynamicCache`."""
+
+        batch, _ = position_ids.shape
+        current_steps = [position_ids[b][0] for b in range(batch)]
+
+        assert len(current_steps) == batch
+        cache = cls(current_steps)
+
+        for layer_idx in range(num_hidden_layer):
+            key_states = past_key_values[layer_idx * 2]
+            value_states = past_key_values[layer_idx * 2 + 1]
+            cache.key_cache.append(key_states)
+            cache.value_cache.append(value_states)
+
+        return cache
+
+
+class RebelDynamicCache_4D(RebelDynamicCache):
+    def assign(
+        self,
+        keys: torch.Tensor,
+        values: torch.Tensor,
+        layer_idx: int,
+    ) -> None:
+        self.key_cache[layer_idx] = keys
+        self.value_cache[layer_idx] = values
+
+    def update(
+        self,
+        keys: torch.Tensor,
+        values: torch.Tensor,
+        layer_idx: int,
+        batch_idx: int,
+        read_first_step: Optional[bool] = False,
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        """
+        Updates the cache with the new `keys` and `values` for the layer `layer_idx` and the batch 'batch_inx'
+        based on self.current_step,
+        """
+        current_step = self.current_steps[0 if read_first_step else batch_idx]
+        kend = current_step + keys.shape[-2]
+        vend = current_step + values.shape[-2]
+        update_keys = (
+            self.key_cache[layer_idx][batch_idx].unsqueeze(0).slice_scatter(keys, dim=-2, start=current_step, end=kend)
+        )
+        update_values = (
+            self.value_cache[layer_idx][batch_idx]
+            .unsqueeze(0)
+            .slice_scatter(values, dim=-2, start=current_step, end=vend)
+        )
+
+        return update_keys, update_values

optimum/rbln/transformers/generation/utils.py

@@ -32,7 +32,6 @@ class RBLNGenerationMixin:
         generation_config: Optional[GenerationConfig] = None,  # thkim change for 4.41.0
         **model_kwargs,
     ) -> Union[SampleDecoderOnlyOutput, torch.LongTensor]:
-
         ###################### thkim change for 4.41.0 ############################
         if generation_config is not None:
             pad_token_id = generation_config.pad_token_id
@@ -216,7 +215,6 @@ class RBLNGenerationMixin:
         do_sample: Optional[bool] = True,
         **model_kwargs,
     ) -> Union[SampleDecoderOnlyOutput, torch.LongTensor]:
-
         ###################### thkim change for 4.41.0 ############################
         if generation_config is not None:
             pad_token_id = generation_config.pad_token_id

optimum/rbln/transformers/models/__init__.py

@@ -22,8 +22,11 @@
 # from Rebellions Inc.
 
 from .clip import RBLNCLIPTextModel, RBLNCLIPTextModelWithProjection
+from .dpt import RBLNDPTForDepthEstimation
+from .gemma import RBLNGemmaForCausalLM
 from .gpt2 import RBLNGPT2LMHeadModel
 from .llama import RBLNLlamaForCausalLM
 from .midm import RBLNMidmLMHeadModel
 from .wav2vec2 import RBLNWav2Vec2ForCTC
 from .whisper import RBLNWhisperForConditionalGeneration
+from .xlm_roberta import RBLNXLMRobertaModel

optimum/rbln/transformers/models/bart/bart_architecture.py

@@ -56,7 +56,6 @@ class _BartAttention(BartAttention):
         cache_position: torch.Tensor,
         key_value_states: Optional[torch.Tensor] = None,
     ) -> Tuple[torch.Tensor, Tuple[torch.Tensor]]:
-
         bsz, tgt_len, _ = hidden_states.size()
         is_cross_attention = key_value_states is not None
 
@@ -111,7 +110,6 @@ class _BartSdpaAttention(BartSdpaAttention):
         cache_position: torch.Tensor,
         key_value_states: Optional[torch.Tensor] = None,
     ) -> Tuple[torch.Tensor, Tuple[torch.Tensor]]:
-
         bsz, tgt_len, _ = hidden_states.size()
         is_cross_attention = key_value_states is not None
 
@@ -166,7 +164,6 @@ class _BartDecoderLayer(BartDecoderLayer):
         cache_position: torch.Tensor,
         attn_impl: str = "eager",
     ) -> Tuple[torch.Tensor, Tuple[torch.Tensor]]:
-
         # Self Attention Block
         residual = hidden_states
         self_attn_past_key_value = past_key_value[:2]
@@ -218,7 +215,6 @@ class _BartDecoder(BartDecoder):
         cache_position: torch.Tensor,
         attn_impl: str = "eager",
     ):
-
         # embedding
         positions_idx = cache_position + self.embed_positions.offset
         positions = self.embed_positions.weight[positions_idx]
@@ -284,7 +280,6 @@ class BartDecoderWrapper(torch.nn.Module):
         self_kv_cache: torch.Tensor,
         cross_kv_cache: torch.Tensor,
     ) -> Tuple[torch.FloatTensor, Tuple[torch.FloatTensor]]:
-
         # prepare past_key_values
         kv_cache = ()
         for i in range(0, self.num_layers * 2, 2):

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

@@ -51,7 +51,6 @@ class _TextEncoder(torch.nn.Module):
 
 
 class RBLNCLIPTextModel(RBLNModel):
-    model_type = "rbln_clip"
     auto_model_class = AutoModel  # feature extraction
     original_model_class = CLIPTextModel
     original_config_class = CLIPTextConfig

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

@@ -0,0 +1,36 @@
+# Copyright 2024 Rebellions Inc.
+
+# 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.
+
+# Portions of this software are licensed under the Apache License,
+# Version 2.0. See the NOTICE file distributed with this work for
+# additional information regarding copyright ownership.
+
+# All other portions of this software, including proprietary code,
+# are the intellectual property of Rebellions Inc. and may not be
+# copied, modified, or distributed without prior written permission
+# from Rebellions Inc.
+
+from .decoderonly_architecture import (
+    DecoderOnlyAttention,
+    DecoderOnlyDecoderLayer,
+    DecoderOnlyModel,
+    DecoderOnlyWrapper,
+    DynamicNTKScalingRotaryEmbedding,
+    LinearScalingRotaryEmbedding,
+    RotaryEmbedding,
+    apply_rotary_pos_emb,
+    rotate_half,
+    slice_and_unsqueeze_cos_sin,
+)
+from .modeling_decoderonly import RBLNDecoderOnlyModelForCausalLM