optimum-rbln 0.1.11__py3-none-any.whl → 0.1.13__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 RBLNT5EncoderModel, RBLNT5ForConditionalGeneration
 from .t5_architecture import T5DecoderWrapper, T5EncoderWrapper

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

@@ -0,0 +1,108 @@
+# 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, Dict, Optional, Union
+
+from transformers import (
+    AutoModelForTextEncoding,
+    PretrainedConfig,
+    T5ForConditionalGeneration,
+)
+
+from ....modeling_base import RBLNModel
+from ....modeling_config import RBLNCompileConfig, 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 AutoFeatureExtractor, AutoProcessor, AutoTokenizer, PreTrainedModel
+
+
+class RBLNT5EncoderModel(RBLNModel):
+    auto_model_class = AutoModelForTextEncoding
+
+    @classmethod
+    def _get_rbln_config(
+        cls,
+        preprocessors: Optional[Union["AutoFeatureExtractor", "AutoProcessor", "AutoTokenizer"]],
+        model_config: Optional["PretrainedConfig"] = None,
+        rbln_kwargs: Dict[str, Any] = {},
+    ) -> RBLNConfig:
+        rbln_max_seq_len = rbln_kwargs.get("max_seq_len", None)
+        rbln_batch_size = rbln_kwargs.get("batch_size", None)
+
+        max_position_embeddings = getattr(model_config, "n_positions", None)
+
+        if rbln_max_seq_len is None:
+            rbln_max_seq_len = max_position_embeddings
+            if rbln_max_seq_len is None:
+                for tokenizer in preprocessors:
+                    if hasattr(tokenizer, "model_max_length"):
+                        rbln_max_seq_len = tokenizer.model_max_length
+                        break
+                if rbln_max_seq_len is None:
+                    raise ValueError("`rbln_max_seq_len` should be specified!")
+
+        if max_position_embeddings is not None and rbln_max_seq_len > max_position_embeddings:
+            raise ValueError("`rbln_max_seq_len` should be less or equal than max_position_embeddings!")
+
+        if rbln_batch_size is None:
+            rbln_batch_size = 1
+
+        input_info = [
+            ("input_ids", [rbln_batch_size, rbln_max_seq_len], "int64"),
+            ("attention_mask", [rbln_batch_size, rbln_max_seq_len], "int64"),
+        ]
+
+        rbln_compile_config = RBLNCompileConfig(input_info=input_info)
+
+        rbln_config = RBLNConfig(
+            rbln_cls=cls.__name__,
+            compile_cfgs=[rbln_compile_config],
+            rbln_kwargs=rbln_kwargs,
+        )
+
+        rbln_config.model_cfg.update({"max_seq_len": rbln_max_seq_len})
+        return rbln_config
+
+
+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