optimum-rbln 0.1.15__py3-none-any.whl → 0.2.1a0__py3-none-any.whl

optimum/rbln/transformers/models/whisper/modeling_whisper.py

@@ -27,6 +27,7 @@ from typing import TYPE_CHECKING, Any, Callable, Dict, List, Optional, Union
 
 import rebel
 import torch
+from rebel.compile_context import CompileContext
 from transformers import (
     AutoModelForSpeechSeq2Seq,
     AutoProcessor,
@@ -37,23 +38,16 @@ from transformers import (
 from transformers.modeling_outputs import BaseModelOutput, Seq2SeqLMOutput
 
 from ....modeling import RBLNModel
-from ....modeling_config import DEFAULT_COMPILED_MODEL_NAME, RBLNCompileConfig, RBLNConfig
+from ....modeling_config import RBLNCompileConfig, RBLNConfig
 from ....utils.runtime_utils import RBLNPytorchRuntime
 from .generation_whisper import RBLNWhisperGenerationMixin
-from .whisper_architecture import (
-    _WhisperDecoderWrapper,
-    _WhisperEncoderWrapper,
-)
+from .whisper_architecture import WhisperWrapper
 
 
 logger = logging.getLogger(__name__)
 
 if TYPE_CHECKING:
-    from transformers import (
-        AutoFeatureExtractor,
-        AutoProcessor,
-        PretrainedConfig,
-    )
+    from transformers import AutoFeatureExtractor, AutoProcessor, PretrainedConfig, PreTrainedModel
 
 
 class RBLNRuntimeEncoder(RBLNPytorchRuntime):
@@ -164,47 +158,51 @@ class RBLNWhisperForConditionalGeneration(RBLNModel, RBLNWhisperGenerationMixin)
         # TODO(jongho): implement
         raise NotImplementedError
 
+    @classmethod
+    def wrap_model_if_needed(self, model: "PreTrainedModel", rbln_config: "RBLNConfig"):
+        rbln_token_timestamps = rbln_config.model_cfg["token_timestamps"]
+        return WhisperWrapper(model, rbln_token_timestamps)
+
     @classmethod
     @torch.inference_mode()
     def get_compiled_model(cls, model, rbln_config: RBLNConfig):
-        rbln_token_timestamps = rbln_config.model_cfg["token_timestamps"]
-        wrapped_encoder = _WhisperEncoderWrapper(model).eval()
-        wrapped_decoder = _WhisperDecoderWrapper(model, output_attentions=rbln_token_timestamps).eval()
+        wrapped_model = cls.wrap_model_if_needed(model, rbln_config)
 
-        enc_rbln_compile_config = rbln_config.compile_cfgs[0]
-        dec_rbln_compile_config = rbln_config.compile_cfgs[1]
+        enc_compile_config = rbln_config.compile_cfgs[0]
+        dec_compile_config = rbln_config.compile_cfgs[1]
 
-        enc_example_inputs = enc_rbln_compile_config.get_dummy_inputs(fill=1)
-        dec_example_inputs = dec_rbln_compile_config.get_dummy_inputs(fill=1)
+        context = CompileContext(use_weight_sharing=False)
 
-        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_example_inputs = enc_compile_config.get_dummy_inputs(fill=0)
 
-        enc_ir = rebel.torchscript_to_ir(
-            enc_scripted_model,
-            input_names=[v[0] for v in enc_rbln_compile_config.input_info],
-            name=enc_rbln_compile_config.mod_name,
+        # Mark encoder's static tensors (cross kv states)
+        static_tensors = {}
+        for (name, _, _), tensor in zip(enc_compile_config.input_info, enc_example_inputs):
+            if "key_value_states" in name:
+                static_tensors[name] = tensor
+                context.mark_static_address(tensor)
+
+        dec_example_inputs = dec_compile_config.get_dummy_inputs(fill=0, static_tensors=static_tensors)
+
+        # Mark decoder's static tensors (self kv states)
+        for (name, _, _), tensor in zip(dec_compile_config.input_info, dec_example_inputs):
+            if "key_value_states" in name:
+                context.mark_static_address(tensor)
+
+        compiled_encoder = super().compile(
+            wrapped_model.encoder,
+            enc_compile_config,
+            example_inputs=enc_example_inputs,
+            compile_context=context,
         )
-        dec_ir = rebel.torchscript_to_ir(
-            dec_scripted_model,
-            input_names=[v[0] for v in dec_rbln_compile_config.input_info],
-            name=dec_rbln_compile_config.mod_name,
+        compiled_decoder = super().compile(
+            wrapped_model.decoder,
+            dec_compile_config,
+            example_inputs=dec_example_inputs,
+            compile_context=context,
         )
 
-        # Caching encoder/decoder I/O
-        connections = [
-            (enc_ir.outputs[0], dec_ir.inputs[4]),
-            (dec_ir.outputs[1], dec_ir.inputs[3]),
-        ]
-        compiled_model = rebel.compile(
-            enc_ir,
-            dec_ir,
-            connections=connections,
-            fusion=enc_rbln_compile_config.fusion,
-            npu=enc_rbln_compile_config.npu,
-            tensor_parallel_size=enc_rbln_compile_config.tensor_parallel_size,
-        )
-        return compiled_model
+        return {"encoder": compiled_encoder, "decoder": compiled_decoder}
 
     @classmethod
     def _get_rbln_config(
@@ -228,6 +226,22 @@ class RBLNWhisperForConditionalGeneration(RBLNModel, RBLNWhisperGenerationMixin)
 
         # model input info
         enc_input_info = [("input_features", [rbln_batch_size, num_mel_bins, expected_seq_len], "float32")]
+        enc_input_info.extend(
+            [
+                (
+                    "cross_key_value_states",
+                    [
+                        model_config.decoder_layers * 2,
+                        rbln_batch_size,
+                        model_config.decoder_attention_heads,
+                        enc_max_seq_len,
+                        model_config.d_model // model_config.decoder_attention_heads,
+                    ],
+                    "float32",
+                )
+            ]
+        )
+
         dec_input_info = [
             ("decoder_input_ids", [rbln_batch_size, 1], "int64"),
             ("decoder_attention_mask", [rbln_batch_size, rbln_dec_max_seq_len], "int64"),
@@ -236,13 +250,13 @@ class RBLNWhisperForConditionalGeneration(RBLNModel, RBLNWhisperGenerationMixin)
         dec_input_info.extend(
             [
                 (
-                    "self_key_value_states",
+                    "cross_key_value_states",
                     [
                         model_config.decoder_layers * 2,
                         rbln_batch_size,
                         model_config.decoder_attention_heads,
-                        rbln_dec_max_seq_len,
-                        model_config.d_model // model_config.encoder_attention_heads,
+                        enc_max_seq_len,
+                        model_config.d_model // model_config.decoder_attention_heads,
                     ],
                     "float32",
                 )
@@ -251,25 +265,25 @@ class RBLNWhisperForConditionalGeneration(RBLNModel, RBLNWhisperGenerationMixin)
         dec_input_info.extend(
             [
                 (
-                    "cross_key_value_states",
+                    f"self_key_value_states_{i}",
                     [
-                        model_config.decoder_layers * 2,
                         rbln_batch_size,
                         model_config.decoder_attention_heads,
-                        enc_max_seq_len,
+                        rbln_dec_max_seq_len,
                         model_config.d_model // model_config.encoder_attention_heads,
                     ],
                     "float32",
                 )
+                for i in range(model_config.decoder_layers * 2)
             ]
         )
 
-        enc_rbln_compile_config = RBLNCompileConfig(mod_name="encoder", input_info=enc_input_info)
-        dec_rbln_compile_config = RBLNCompileConfig(mod_name="decoder", input_info=dec_input_info)
+        enc_compile_config = RBLNCompileConfig(compiled_model_name="encoder", input_info=enc_input_info)
+        dec_compile_config = RBLNCompileConfig(compiled_model_name="decoder", input_info=dec_input_info)
 
         rbln_config = RBLNConfig(
             rbln_cls=cls.__name__,
-            compile_cfgs=[enc_rbln_compile_config, dec_rbln_compile_config],
+            compile_cfgs=[enc_compile_config, dec_compile_config],
             rbln_kwargs=rbln_kwargs,
         )
 
@@ -285,12 +299,21 @@ class RBLNWhisperForConditionalGeneration(RBLNModel, RBLNWhisperGenerationMixin)
 
     @classmethod
     def _create_runtimes(
-        cls, compiled_models: List[rebel.RBLNCompiledModel], rbln_device_map: Dict[str, int]
+        cls,
+        compiled_models: List[rebel.RBLNCompiledModel],
+        rbln_device_map: Dict[str, int],
+        activate_profiler: Optional[bool] = None,
     ) -> List[rebel.Runtime]:
-        device_val = rbln_device_map[DEFAULT_COMPILED_MODEL_NAME]
+        if any(model_name not in rbln_device_map for model_name in ["encoder", "decoder"]):
+            cls._raise_missing_compiled_file_error(["encoder", "decoder"])
+
         return [
-            compiled_models[0].create_runtime("encoder", tensor_type="pt", device=device_val),
-            compiled_models[0].create_runtime("decoder", tensor_type="pt", device=device_val),
+            compiled_models[0].create_runtime(
+                tensor_type="pt", device=rbln_device_map["encoder"], activate_profiler=activate_profiler
+            ),
+            compiled_models[1].create_runtime(
+                tensor_type="pt", device=rbln_device_map["decoder"], activate_profiler=activate_profiler
+            ),
         ]
 
     def prepare_inputs_for_generation(