optimum-rbln 0.1.11__py3-none-any.whl → 0.1.12__py3-none-any.whl

optimum/rbln/{modeling_seq2seq.py → transformers/models/seq2seq/modeling_seq2seq.py}

@@ -23,24 +23,17 @@
 
 import inspect
 import logging
-from typing import TYPE_CHECKING, Any, Callable, Dict, List, Optional, Tuple, Union
-
-import rebel
-import torch
-from transformers import (
-    AutoModelForSeq2SeqLM,
-    BartConfig,
-    BartForConditionalGeneration,
-    PretrainedConfig,
-    T5ForConditionalGeneration,
-)
+from abc import ABC
+from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Union
+
+import rebel  # noqa: F401
+import torch  # noqa: F401
+from transformers import AutoModelForSeq2SeqLM, GenerationConfig, PretrainedConfig, PreTrainedModel
 from transformers.modeling_outputs import BaseModelOutput, Seq2SeqLMOutput
 
-from .modeling_base import RBLNModel
-from .modeling_config import DEFAULT_COMPILED_MODEL_NAME, RBLNCompileConfig, RBLNConfig
-from .transformers.models.bart import BartDecoderWrapper, BartEncoderWrapper
-from .transformers.models.t5 import T5DecoderWrapper, T5EncoderWrapper
-from .utils.runtime_utils import RBLNPytorchRuntime
+from ....modeling_base import RBLNModel
+from ....modeling_config import DEFAULT_COMPILED_MODEL_NAME, RBLNCompileConfig, RBLNConfig
+from ....utils.runtime_utils import RBLNPytorchRuntime
 
 
 logger = logging.getLogger(__name__)
@@ -59,7 +52,6 @@ class RBLNRuntimeEncoder(RBLNPytorchRuntime):
 
     def forward(self, *args: List[torch.Tensor], **kwargs: Dict[str, torch.Tensor]):
         _ = super().forward(*args, **kwargs)
-        # Just indicates that it is not None
         return BaseModelOutput(last_hidden_state=torch.tensor([1.0]))
 
 
@@ -71,7 +63,7 @@ class RBLNRuntimeDecoder(RBLNPytorchRuntime):
         return Seq2SeqLMOutput(logits=outputs)
 
 
-class RBLNModelForSeq2SeqLM(RBLNModel):
+class RBLNModelForSeq2SeqLM(RBLNModel, ABC):
     """
     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.
@@ -84,91 +76,35 @@ class RBLNModelForSeq2SeqLM(RBLNModel):
     Currently, this model class only supports the 'bart' and 't5' models from the transformers library. Future updates may include support for additional model types.
     """
 
+    main_input_name = "input_ids"
     auto_model_class = AutoModelForSeq2SeqLM
 
     def __post_init__(self, **kwargs):
-        self.model_dim = self.config.d_model
-        self.batch_size = self.rbln_config.model_cfg["batch_size"]
-        self.enc_max_seq_len = self.rbln_config.model_cfg["enc_max_seq_len"]
-        self.dec_max_seq_len = self.rbln_config.model_cfg["dec_max_seq_len"]
-        self.pad_token_id = self.rbln_config.model_cfg["pad_token_id"]
         self.encoder = RBLNRuntimeEncoder(runtime=self.model[0], main_input_name="input_ids")
         self.decoder = RBLNRuntimeDecoder(runtime=self.model[1], main_input_name="input_ids")
-        self.enc_attention_mask = torch.zeros(1, self.enc_max_seq_len, dtype=torch.float32)
-        self.dec_enc_attention_mask = torch.zeros(self.batch_size, self.enc_max_seq_len, dtype=torch.float32)
-
-    def can_generate(self):
-        return True
-
-    def get_encoder(self):
-        return self.encoder
-
-    def get_decoder(self):
-        return self.decoder
-
-    def __getattr__(self, __name: str) -> Any:
-        def redirect(func):
-            return lambda *pargs, **kwargs: func(self, *pargs, **kwargs)
-
-        if "T5ForConditionalGeneration" == self.config.architectures:
-            val = getattr(T5ForConditionalGeneration, __name)
-        else:
-            val = getattr(BartForConditionalGeneration, __name)
-
-        if isinstance(val, Callable) and "self" in set(inspect.signature(val).parameters):
-            return redirect(val)
-        return val
-
-    @classmethod
-    def update_kwargs(cls, kwargs):
-        kwargs.update(
-            {
-                "torchscript": True,
-                "return_dict": False,
-                "use_cache": True,
-            }
-        )
-        return kwargs
 
     @classmethod
-    def get_compiled_model(cls, model, rbln_config: RBLNConfig):
-        def optimized_models(model):
-            if isinstance(model, T5ForConditionalGeneration):
-                encoder_model = T5EncoderWrapper(model).eval()
-                decoder_model = T5DecoderWrapper(model).eval()
-            elif isinstance(model, BartForConditionalGeneration):
-                encoder_model = BartEncoderWrapper(model).eval()
-                decoder_model = BartDecoderWrapper(model).eval()
-            else:
-                raise ValueError(f"{model.__class__.__name__} is not supported yet.")
-
-            return encoder_model, decoder_model
-
-        wrapped_encoder, wrapped_decoder = optimized_models(model)
+    @torch.inference_mode()
+    def get_compiled_model(cls, model: PreTrainedModel, rbln_config: RBLNConfig):
+        wrapped_model = cls.wrap_model_if_needed(model, rbln_config)
 
-        wrapped_encoder.encoder_max_length = rbln_config.model_cfg["enc_max_seq_len"]
-        wrapped_encoder.decoder_max_length = rbln_config.model_cfg["dec_max_seq_len"]
-        wrapped_encoder.decoder_batch_size = rbln_config.model_cfg["batch_size"]
+        wrapped_model.encoder.encoder_max_length = rbln_config.model_cfg["enc_max_seq_len"]
+        wrapped_model.encoder.decoder_max_length = rbln_config.model_cfg["dec_max_seq_len"]
 
-        wrapped_decoder.encoder_max_length = rbln_config.model_cfg["enc_max_seq_len"]
-        wrapped_decoder.decoder_max_length = rbln_config.model_cfg["dec_max_seq_len"]
-        wrapped_decoder.decoder_batch_size = rbln_config.model_cfg["batch_size"]
+        wrapped_model.decoder.encoder_max_length = rbln_config.model_cfg["enc_max_seq_len"]
+        wrapped_model.decoder.decoder_max_length = rbln_config.model_cfg["dec_max_seq_len"]
 
         enc_rbln_compile_config = rbln_config.compile_cfgs[0]
         dec_rbln_compile_config = rbln_config.compile_cfgs[1]
 
-        if isinstance(model, T5ForConditionalGeneration):
-            enc_example_inputs = enc_rbln_compile_config.get_dummy_inputs(fill=1)
-            dec_example_inputs = dec_rbln_compile_config.get_dummy_inputs(fill=1)
-        else:
-            enc_example_inputs = enc_rbln_compile_config.get_dummy_inputs(fill=0)
-            dec_example_inputs = dec_rbln_compile_config.get_dummy_inputs(fill=0)
+        enc_example_inputs = enc_rbln_compile_config.get_dummy_inputs(fill=0)
+        dec_example_inputs = dec_rbln_compile_config.get_dummy_inputs(fill=0)
 
         enc_example_inputs[3].fill_(0)
         dec_example_inputs[4].fill_(-1)
 
-        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_scripted_model = torch.jit.trace(wrapped_model.encoder, enc_example_inputs, check_trace=False)
+        dec_scripted_model = torch.jit.trace(wrapped_model.decoder, dec_example_inputs, check_trace=False)
 
         enc_ir = rebel.torchscript_to_ir(
             enc_scripted_model,
@@ -180,13 +116,12 @@ class RBLNModelForSeq2SeqLM(RBLNModel):
             input_names=[v[0] for v in dec_rbln_compile_config.input_info],
             name=dec_rbln_compile_config.mod_name,
         )
-        dec_ir.decoder_batch_size = rbln_config.model_cfg["batch_size"]
 
         connections = [
             (enc_ir.outputs[0], enc_ir.inputs[2], dec_ir.inputs[6]),
-            # (enc_ir.outputs[0], enc_ir.inputs[2]),
             (dec_ir.outputs[1], dec_ir.inputs[5]),
         ]
+
         compiled_model = rebel.compile(
             enc_ir,
             dec_ir,
@@ -209,14 +144,13 @@ class RBLNModelForSeq2SeqLM(RBLNModel):
         rbln_batch_size = rbln_kwargs.get("batch_size", None)
         rbln_batch_size = 1 if rbln_batch_size is None else rbln_batch_size
 
-        if isinstance(model_config, BartConfig):
-            n_layer = model_config.decoder_layers
-            n_head = model_config.decoder_attention_heads
-            d_kv = model_config.d_model // model_config.encoder_attention_heads
-        else:
-            n_layer = model_config.num_layers
-            n_head = model_config.num_heads
-            d_kv = model_config.d_kv
+        n_layer = getattr(model_config, "decoder_layers", None) or getattr(model_config, "num_layers")
+        n_head = getattr(model_config, "decoder_attention_heads", None) or getattr(model_config, "num_heads")
+        d_kv = (
+            model_config.d_kv
+            if hasattr(model_config, "d_kv")
+            else model_config.d_model // model_config.encoder_attention_heads
+        )
 
         max_position_embeddings = getattr(model_config, "n_positions", None) or getattr(
             model_config, "max_position_embeddings", None
@@ -270,7 +204,6 @@ class RBLNModelForSeq2SeqLM(RBLNModel):
                 ],
                 "float32",
             ),
-            # int16 available?
             ("batch_idx", [], "int32"),
         ]
 
@@ -281,7 +214,6 @@ class RBLNModelForSeq2SeqLM(RBLNModel):
             (
                 "cache_position",
                 [rbln_batch_size, 1],
-                # [],
                 "int32",
             ),
             ("batch_position", [], "int32"),
@@ -346,35 +278,32 @@ class RBLNModelForSeq2SeqLM(RBLNModel):
             compiled_models[0].create_runtime("decoder", tensor_type="pt", device=device_val),
         ]
 
+    def can_generate(self):
+        return True
+
+    def get_encoder(self):
+        return self.encoder
+
+    def get_decoder(self):
+        return self.decoder
+
     def prepare_inputs_for_generation(
         self,
         input_ids,
-        past_key_values=None,
         attention_mask=None,
         decoder_attention_mask=None,
         **kwargs,
     ):
-        past_cache_length = past_key_values
-        if past_cache_length == 0:
-            cache_pos = []
-            for i in range(input_ids.shape[0]):
-                cache_pos.append([0])
-            cache_position = torch.tensor(cache_pos, dtype=torch.int32)
-
-        max_seq_len = self.dec_max_seq_len
         cur_seq_len = input_ids.shape[-1]
+        cache_position = cur_seq_len - 1
+        max_seq_len = self.rbln_config.model_cfg["dec_max_seq_len"]
         decoder_batch_size = input_ids.shape[0]
         input_ids = input_ids[:, cur_seq_len - 1 : cur_seq_len].contiguous()
-        # In greedy decoding
         decoder_attention_mask = torch.zeros(decoder_batch_size, max_seq_len, dtype=torch.float32)
         decoder_attention_mask[:, :cur_seq_len] = 1
-        cache_pos = []
-        for i in range(input_ids.shape[0]):
-            cache_pos.append([cur_seq_len - 1])
-        cache_position = torch.tensor(cache_pos, dtype=torch.int32)
+
         return {
             "decoder_input_ids": input_ids,
-            "past_key_values": past_key_values,
             "attention_mask": attention_mask.to(torch.float32),
             "decoder_attention_mask": decoder_attention_mask,
             "cache_position": cache_position,
@@ -383,41 +312,12 @@ class RBLNModelForSeq2SeqLM(RBLNModel):
     def forward(
         self,
         input_ids: torch.LongTensor = None,
-        cache_position: Union[List[torch.Tensor], torch.Tensor] = None,  # vllm keyword argument
-        batch_idx: Optional[torch.LongTensor] = None,
-        enc_lengths: List[int] = None,
+        cache_position: Union[List[torch.Tensor], torch.Tensor] = None,
         **kwargs,
     ) -> Tuple[torch.FloatTensor]:
         # common decoder
-        if enc_lengths is None:
-            output = self._forward_decoder(input_ids=input_ids, cache_position=cache_position, **kwargs)
-            return output
-
-        # vllm & encoder
-        if batch_idx is not None:
-            enc_attention_mask = self.enc_attention_mask.clone()
-            enc_attention_mask[0][: enc_lengths[batch_idx] + 1] = 1
-            padding_need = self.enc_max_seq_len - input_ids.shape[-1]
-            input_ids = torch.nn.functional.pad(input_ids, (0, padding_need))
-            _ = self.encoder(input_ids, enc_attention_mask, batch_idx=batch_idx.to(torch.int32))
-            logits = torch.zeros(1, 1, self.config.vocab_size + 100)
-            logits[0][0][-1] = 1
-        # vllm & decoder
-        else:
-            input_ids[input_ids == (self.config.vocab_size + 99)] = self.config.decoder_start_token_id
-            cache_position[cache_position != 0] = cache_position[cache_position != 0] - 2
-
-            enc_attention_mask = self.dec_enc_attention_mask.clone()
-            dec_attention_mask = torch.zeros(self.batch_size, self.dec_max_seq_len, dtype=torch.float32)
-            for batch_idx in range(self.batch_size):
-                enc_attention_mask[batch_idx, : enc_lengths[batch_idx] + 1] = 1
-
-            logits = self._forward_decoder(
-                attention_mask=enc_attention_mask,
-                decoder_input_ids=input_ids,
-                decoder_attention_mask=dec_attention_mask,
-                cache_position=cache_position,
-            ).logits
+        cache_position = torch.full((self.rbln_config.model_cfg["batch_size"], 1), cache_position, dtype=torch.int32)
+        logits = self._forward_decoder(input_ids=input_ids, cache_position=cache_position, **kwargs).logits
 
         return Seq2SeqLMOutput(
             logits=logits,
@@ -446,25 +346,58 @@ class RBLNModelForSeq2SeqLM(RBLNModel):
 
         return Seq2SeqLMOutput(logits=lm_logits)
 
+    def vllm_forward(
+        self,
+        input_ids: torch.LongTensor = None,
+        cache_position: Union[List[torch.Tensor], torch.Tensor] = None,  # vllm keyword argument
+        batch_idx: Optional[torch.LongTensor] = None,
+        enc_lengths: List[int] = None,  # vllm return current attention_mask length
+        **kwargs,
+    ) -> Tuple[torch.FloatTensor]:
+        # When using vllm, need the output of the encoder (ex. vocab_size + 100) and use that value act as start_token_id in decoder (ex. vocab_size + 99)
+        # encoder
+        if batch_idx is not None:
+            enc_attention_mask = torch.zeros(1, self.rbln_config.model_cfg["enc_max_seq_len"], dtype=torch.float32)
+            enc_attention_mask[0][: enc_lengths[batch_idx] + 1] = 1
+            padding_need = self.rbln_config.model_cfg["enc_max_seq_len"] - input_ids.shape[-1]
+            input_ids = torch.nn.functional.pad(input_ids, (0, padding_need))
+            _ = self.encoder(input_ids, enc_attention_mask, batch_idx=batch_idx.to(torch.int32))
+            logits = torch.zeros(1, 1, self.config.vocab_size + 100)
+            logits[0][0][-1] = 1
+        # decoder
+        else:
+            input_ids[input_ids == (self.config.vocab_size + 99)] = self.config.decoder_start_token_id
+            cache_position[cache_position != 0] = cache_position[cache_position != 0] - 2
+
+            enc_attention_mask = torch.zeros(
+                self.rbln_config.model_cfg["batch_size"],
+                self.rbln_config.model_cfg["enc_max_seq_len"],
+                dtype=torch.float32,
+            )
+            dec_attention_mask = torch.zeros(
+                self.rbln_config.model_cfg["batch_size"],
+                self.rbln_config.model_cfg["dec_max_seq_len"],
+                dtype=torch.float32,
+            )
+            for batch_idx in range(self.rbln_config.model_cfg["batch_size"]):
+                enc_attention_mask[batch_idx, : enc_lengths[batch_idx] + 1] = 1
+
+            logits = self._forward_decoder(
+                attention_mask=enc_attention_mask,
+                decoder_input_ids=input_ids,
+                decoder_attention_mask=dec_attention_mask,
+                cache_position=cache_position,
+            ).logits
+
+        return Seq2SeqLMOutput(logits=logits)
+
     def _prepare_encoder_decoder_kwargs_for_generation(
         self,
         inputs_tensor: torch.Tensor,
         model_kwargs,
         model_input_name: Optional[str] = None,
-        *args,
-        **kwargs,
+        generation_config: Optional[GenerationConfig] = None,
     ) -> Dict[str, Any]:
-        ########## thkim change start ###################
-        # padding input_ids & attention_mask regardless of user's tokenizer usage
-        batch_size, input_len = inputs_tensor.shape
-        inputs_tensor = torch.nn.functional.pad(
-            inputs_tensor, (0, self.enc_max_seq_len - input_len), value=self.pad_token_id
-        )
-        model_kwargs["attention_mask"] = torch.nn.functional.pad(
-            model_kwargs["attention_mask"], (0, self.enc_max_seq_len - input_len), value=0
-        )
-        ########## thkim change end ###################
-
         # 1. get encoder
         encoder = self.get_encoder()
 
@@ -482,18 +415,26 @@ class RBLNModelForSeq2SeqLM(RBLNModel):
                 argument: value for argument, value in encoder_kwargs.items() if argument in encoder_signature
             }
 
+        batch_size, input_len = inputs_tensor.shape
+        inputs_tensor = torch.nn.functional.pad(
+            inputs_tensor,
+            (0, self.rbln_config.model_cfg["enc_max_seq_len"] - input_len),
+            value=self.rbln_config.model_cfg["pad_token_id"],
+        )
+        model_kwargs["attention_mask"] = torch.nn.functional.pad(
+            model_kwargs["attention_mask"], (0, self.rbln_config.model_cfg["enc_max_seq_len"] - input_len)
+        )
+
         # 3. make sure that encoder returns `ModelOutput`
         model_input_name = model_input_name if model_input_name is not None else self.main_input_name
         encoder_kwargs["return_dict"] = True
-        encoder_kwargs[model_input_name] = inputs_tensor
+        encoder_kwargs["output_hidden_states"] = False
+        encoder_kwargs["output_attentions"] = False
+
         for b in range(batch_size):
             batch_idx = torch.tensor(b, dtype=torch.int32)
-            cb_inputs = {}
-            cb_inputs["return_dict"] = True
-            cb_inputs["output_hidden_states"] = False
-            cb_inputs["output_attentions"] = False
-            cb_inputs["input_ids"] = encoder_kwargs["input_ids"][b].unsqueeze(0)
-            cb_inputs["attention_mask"] = encoder_kwargs["attention_mask"][b].unsqueeze(0).to(torch.float32)
-            model_kwargs["encoder_outputs"] = encoder(**cb_inputs, batch_idx=batch_idx)
+            encoder_kwargs["input_ids"] = inputs_tensor[b].unsqueeze(0)
+            encoder_kwargs["attention_mask"] = model_kwargs["attention_mask"][b].unsqueeze(0).to(torch.float32)
+            model_kwargs["encoder_outputs"] = encoder(**encoder_kwargs, batch_idx=batch_idx)
 
         return model_kwargs

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

@@ -21,4 +21,5 @@
 # copied, modified, or distributed without prior written permission
 # from Rebellions Inc.
 
+from .modeling_t5 import RBLNT5ForConditionalGeneration
 from .t5_architecture import T5DecoderWrapper, T5EncoderWrapper

optimum/rbln/transformers/models/t5/modeling_t5.py

@@ -0,0 +1,55 @@
+# 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.
+
+import inspect
+from typing import TYPE_CHECKING, Any, Callable
+
+from transformers import T5ForConditionalGeneration
+
+from ....modeling_config import RBLNConfig
+from ....utils.logging import get_logger
+from ...models.seq2seq import RBLNModelForSeq2SeqLM
+from .t5_architecture import T5Wrapper
+
+
+logger = get_logger()
+
+if TYPE_CHECKING:
+    from transformers import PreTrainedModel
+
+
+class RBLNT5ForConditionalGeneration(RBLNModelForSeq2SeqLM):
+    @classmethod
+    def wrap_model_if_needed(self, model: "PreTrainedModel", rbln_config: "RBLNConfig"):
+        return T5Wrapper(model)
+
+    def __getattr__(self, __name: str) -> Any:
+        def redirect(func):
+            return lambda *pargs, **kwargs: func(self, *pargs, **kwargs)
+
+        val = getattr(T5ForConditionalGeneration, __name)
+
+        if isinstance(val, Callable) and "self" in set(inspect.signature(val).parameters):
+            return redirect(val)
+
+        return val

optimum/rbln/transformers/models/t5/t5_architecture.py

@@ -43,6 +43,12 @@ if TYPE_CHECKING:
     from transformers import T5ForConditionalGeneration
 
 
+class T5Wrapper:
+    def __init__(self, model):
+        self.encoder = T5EncoderWrapper(model)
+        self.decoder = T5DecoderWrapper(model)
+
+
 class T5Encoder(T5Stack):
     def forward(
         self,
@@ -122,19 +128,26 @@ class T5EncoderWrapper(torch.nn.Module):
         )
         self.encoder_max_length = None
         self.decoder_max_length = None
-        self.decoder_batch_size = 1
 
-    def forward(self, input_ids: torch.Tensor, attention_mask: torch.Tensor,
-                cross_key_value: torch.Tensor = None, batch_idx: torch.Tensor = None,
-                ) -> torch.Tensor:
-        encoder_batch_size = input_ids.shape[0]
-        decoder_batch_size = self.decoder_batch_size
+    def forward(
+        self,
+        input_ids: torch.Tensor,
+        attention_mask: torch.Tensor,
+        cross_key_value: torch.Tensor = None,
+        batch_idx: torch.Tensor = None,
+    ) -> torch.Tensor:
         decoder_max_length = self.decoder_max_length or self.default_max_length
         encoder_max_length = self.encoder_max_length or self.default_max_length
 
         attn_layer = self.encoder.block[0].layer[0].SelfAttention
         encoder_position_bias = T5Attention.compute_bias(attn_layer, encoder_max_length, encoder_max_length)
-        encoder_outputs = T5Encoder.forward(self.encoder, input_ids, attention_mask, encoder_position_bias, batch_ids=torch.tensor(0, dtype=torch.int32))
+        encoder_outputs = T5Encoder.forward(
+            self.encoder,
+            input_ids,
+            attention_mask,
+            encoder_position_bias,
+            batch_ids=torch.tensor(0, dtype=torch.int32),
+        )
 
         attn_layer = self.decoder.block[0].layer[0].SelfAttention
         decoder_position_bias = T5Attention.compute_bias(attn_layer, decoder_max_length, decoder_max_length)
@@ -145,22 +158,14 @@ class T5EncoderWrapper(torch.nn.Module):
 
         dummy_past_key_value = []
         for i in range(self.config.num_layers):
-            pkv_self_attn_key = torch.zeros(
-                decoder_batch_size, self.config.num_heads, decoder_max_length, self.config.d_kv
-            )
-            pkv_self_attn_value = torch.zeros(
-                decoder_batch_size, self.config.num_heads, decoder_max_length, self.config.d_kv
-            )
-            pkv_cross_attn_key = torch.zeros(
-                encoder_batch_size, self.config.num_heads, encoder_max_length, self.config.d_kv
-            )
-            pkv_cross_attn_value = torch.zeros(
-                encoder_batch_size, self.config.num_heads, encoder_max_length, self.config.d_kv
-            )
+            pkv_self_attn_key = torch.zeros(1, self.config.num_heads, decoder_max_length, self.config.d_kv)
+            pkv_self_attn_value = torch.zeros(1, self.config.num_heads, decoder_max_length, self.config.d_kv)
+            pkv_cross_attn_key = torch.zeros(1, self.config.num_heads, encoder_max_length, self.config.d_kv)
+            pkv_cross_attn_value = torch.zeros(1, self.config.num_heads, encoder_max_length, self.config.d_kv)
             layer_pkv = (pkv_self_attn_key, pkv_self_attn_value, pkv_cross_attn_key, pkv_cross_attn_value)
             dummy_past_key_value.append(layer_pkv)
 
-        decoder_attention_mask = torch.zeros(decoder_batch_size, decoder_max_length, dtype=torch.float32)
+        decoder_attention_mask = torch.zeros(1, decoder_max_length, dtype=torch.float32)
         decoder_attention_mask[:, :1] = 1
 
         # Since first step of decoder has different graph to further step of it,
@@ -168,7 +173,7 @@ class T5EncoderWrapper(torch.nn.Module):
         # TODO(jongho): Separate first-step-decoder.
         decoder_outputs = T5Decoder.forward(
             self.decoder,
-            input_ids=torch.zeros(decoder_batch_size, 1, dtype=torch.int64),
+            input_ids=torch.zeros(1, 1, dtype=torch.int64),
             attention_mask=decoder_attention_mask,
             position_bias=decoder_position_bias,
             encoder_decoder_position_bias=encoder_decoder_position_bias,
@@ -187,7 +192,7 @@ class T5EncoderWrapper(torch.nn.Module):
             cross_kv_cache.append(past_key_values[i][3])
         cross_kv_cache = torch.stack(cross_kv_cache, dim=0)
 
-        cross_key_value = cross_key_value.slice_scatter(cross_kv_cache, dim=1, start=batch_idx, end=batch_idx+1)
+        cross_key_value = cross_key_value.slice_scatter(cross_kv_cache, dim=1, start=batch_idx, end=batch_idx + 1)
 
         return cross_key_value
 
@@ -240,6 +245,7 @@ class T5DecoderWrapper(torch.nn.Module):
         attn_layer = self.model.decoder.block[0].layer[0].SelfAttention
         _decoder_position_bias = T5Attention.compute_bias(attn_layer, decoder_max_length, decoder_max_length)
 
+        # position_bias need to compute with batch (for cb)
         batch_decoder_position_bias = []
         for i in range(input_ids.shape[0]):
             batch_position_bias = _decoder_position_bias[:, :, cache_position[i][0]].unsqueeze(2)
@@ -259,7 +265,7 @@ class T5DecoderWrapper(torch.nn.Module):
             encoder_decoder_position_bias=encoder_decoder_position_bias,
             past_key_values=kv_cache,
             cache_position=cache_position,
-            batch_ids=rbln_batch_position
+            batch_ids=rbln_batch_position,
         )
 
         past_key_values = decoder_outputs.past_key_values
@@ -312,7 +318,7 @@ class _T5Attention(T5Attention):
             value_states = shape(self.v(hidden_states), batch_size)
         else:
             # cross-attn
-            if cache_position.dim() == 0 :
+            if cache_position.dim() == 0:
                 key_states = shape(self.k(key_value_states), key_value_states.shape[0])
                 value_states = shape(self.v(key_value_states), key_value_states.shape[0])
                 past_key_value = key_states, value_states
@@ -331,18 +337,24 @@ class _T5Attention(T5Attention):
                 batch_value_states = value_states[b].unsqueeze(0)
 
                 if is_self_attn and past_key_value is not None:
-                    batch_key_states = past_key_value[0][b].unsqueeze(0).slice_scatter(
-                        batch_key_states, dim=-2, start=cache_position[b][0], end=cache_position[b][0] + 1
+                    batch_key_states = (
+                        past_key_value[0][b]
+                        .unsqueeze(0)
+                        .slice_scatter(
+                            batch_key_states, dim=-2, start=cache_position[b][0], end=cache_position[b][0] + 1
+                        )
                     )
-                    batch_value_states = past_key_value[1][b].unsqueeze(0).slice_scatter(
-                        batch_value_states, dim=-2, start=cache_position[b][0], end=cache_position[b][0] + 1
+                    batch_value_states = (
+                        past_key_value[1][b]
+                        .unsqueeze(0)
+                        .slice_scatter(
+                            batch_value_states, dim=-2, start=cache_position[b][0], end=cache_position[b][0] + 1
+                        )
                     )
 
                 scores = torch.matmul(batch_query_states, batch_key_states.transpose(3, 2))
                 scores += position_bias[b]
-                attn_weights = nn.functional.softmax(scores.float(), dim=-1).type_as(
-                    scores
-                )
+                attn_weights = nn.functional.softmax(scores.float(), dim=-1).type_as(scores)
                 attn_output = unshape(torch.matmul(attn_weights, batch_value_states), 1)
                 all_key_states.append(batch_key_states)
                 all_value_states.append(batch_value_states)
@@ -371,7 +383,9 @@ class _T5Attention(T5Attention):
                 scores
             )  # (batch_size, n_heads, seq_length, key_length)
 
-            attn_output = unshape(torch.matmul(attn_weights, value_states), batch_size)  # (batch_size, seq_length, dim)
+            attn_output = unshape(
+                torch.matmul(attn_weights, value_states), batch_size
+            )  # (batch_size, seq_length, dim)
 
         attn_output = self.o(attn_output)
         present_key_value = (key_states, value_states)

optimum/rbln/transformers/models/wav2vec2/modeling_wav2vec2.py

@@ -65,7 +65,6 @@ class RBLNWav2Vec2ForCTC(RBLNModel):
     - compiling the resulting graph using the RBLN compiler.
     """
 
-    model_type = "rbln_model"
     main_input_name = "input_values"
     auto_model_class = AutoModelForMaskedLM