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

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

@@ -23,20 +23,16 @@
 
 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 AutoModelForCausalLM, GPT2LMHeadModel, PretrainedConfig
 from transformers.modeling_outputs import CausalLMOutputWithCrossAttentions, Seq2SeqLMOutput
 
-from ....modeling_base import RBLNBaseModel
+from ....modeling_base import RBLNModel
 from ....modeling_config import DEFAULT_COMPILED_MODEL_NAME, RBLNConfig, RBLNRuntimeConfig
 from ....utils.runtime_utils import RBLNPytorchRuntime
-from ....utils.save_utils import maybe_save_preprocessors
 from ...generation.utils import RBLNGenerationMixin
 from .gpt2_architecture import GPT2LMHeadModelWrapper
 
@@ -59,12 +55,12 @@ class RBLNRuntimeDecoder(RBLNPytorchRuntime):
         return Seq2SeqLMOutput(logits=logits)
 
 
-class RBLNGPT2LMHeadModel(RBLNBaseModel, RBLNGenerationMixin):
+class RBLNGPT2LMHeadModel(RBLNModel, RBLNGenerationMixin):
     """
     The GPT2 Model transformer with a language modeling head on top (linear layer with weights tied to the input
     embeddings).
 
-    This model inherits from [`RBLNBaseModel`]. Check the superclass documentation for the generic methods the
+    This model inherits from [`RBLNMultiModel`]. Check the superclass documentation for the generic methods the
     library implements for all its model.
 
     It implements the methods to convert a pre-trained transformers GPT2 model into a RBLN transformer model by:
@@ -89,8 +85,8 @@ class RBLNGPT2LMHeadModel(RBLNBaseModel, RBLNGenerationMixin):
             torch.ones(batch_size, 1, self.prefill_chunk_size, self.prefill_chunk_size), diagonal=1
         )
 
-        self.prefill_decoder = RBLNRuntimeDecoder(runtime=self.runtimes[0])
-        self.decoder = RBLNRuntimeDecoder(runtime=self.runtimes[1])
+        self.prefill_decoder = RBLNRuntimeDecoder(runtime=self.model[0])
+        self.decoder = RBLNRuntimeDecoder(runtime=self.model[1])
         self.pad_token_id = self.rbln_config.meta["rbln_pad_token_id"]
         self.past_cached_length = 0
 
@@ -117,38 +113,7 @@ class RBLNGPT2LMHeadModel(RBLNBaseModel, RBLNGenerationMixin):
         raise NotImplementedError
 
     @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,
-    ) -> "RBLNGPT2LMHeadModel":
-        """
-        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,
@@ -156,82 +121,45 @@ class RBLNGPT2LMHeadModel(RBLNBaseModel, RBLNGenerationMixin):
                 "use_cache": True,
             }
         )
+        return kwargs
 
-        rbln_config_kwargs, rbln_constructor_kwargs = cls.pop_rbln_kwargs_from_kwargs(kwargs)
-
-        model: GPT2LMHeadModel = 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
-            )
-
-        def compile_gpt2():
-            wrapped_decoder = GPT2LMHeadModelWrapper(model).eval()
-
-            prefill_rbln_runtime_config = rbln_config[DEFAULT_COMPILED_MODEL_NAME][0]
-            dec_rbln_runtime_config = rbln_config[DEFAULT_COMPILED_MODEL_NAME][1]
+    @classmethod
+    @torch.inference_mode()
+    def get_compiled_model(cls, model: GPT2LMHeadModel, rbln_config: RBLNConfig):
+        wrapped_decoder = GPT2LMHeadModelWrapper(model).eval()
 
-            prefill_example_inputs = prefill_rbln_runtime_config.get_dummy_inputs(fill=0)
-            dec_example_inputs = dec_rbln_runtime_config.get_dummy_inputs(fill=0)
+        prefill_rbln_runtime_config = rbln_config[DEFAULT_COMPILED_MODEL_NAME][0]
+        dec_rbln_runtime_config = rbln_config[DEFAULT_COMPILED_MODEL_NAME][1]
 
-            prefill_scripted_model = torch.jit.trace(wrapped_decoder, prefill_example_inputs)
-            dec_scripted_model = torch.jit.trace(wrapped_decoder, dec_example_inputs)
+        prefill_example_inputs = prefill_rbln_runtime_config.get_dummy_inputs(fill=0)
+        dec_example_inputs = dec_rbln_runtime_config.get_dummy_inputs(fill=0)
 
-            prefill_ir = rebel.torchscript_to_ir(
-                prefill_scripted_model,
-                input_names=[v[0] for v in prefill_rbln_runtime_config.input_info],
-            )
-            dec_ir = rebel.torchscript_to_ir(
-                dec_scripted_model,
-                input_names=[v[0] for v in dec_rbln_runtime_config.input_info],
-            )
-
-            connections = [
-                (prefill_ir.outputs[1], prefill_ir.inputs[1]),
-            ]
+        prefill_scripted_model = torch.jit.trace(wrapped_decoder, prefill_example_inputs, check_trace=False)
+        dec_scripted_model = torch.jit.trace(wrapped_decoder, dec_example_inputs, check_trace=False)
 
-            compiled_model = rebel.compile(
-                prefill_ir,
-                dec_ir,
-                connections=connections,
-                fusion=prefill_rbln_runtime_config.fusion,
-                npu=prefill_rbln_runtime_config.npu,
-                tensor_parallel_size=prefill_rbln_runtime_config.tensor_parallel_size,
-                use_weight_sharing=True,
-            )
-            compiled_model.save(save_dir_path / f"{DEFAULT_COMPILED_MODEL_NAME}.rbln")
+        prefill_ir = rebel.torchscript_to_ir(
+            prefill_scripted_model,
+            input_names=[v[0] for v in prefill_rbln_runtime_config.input_info],
+        )
+        dec_ir = rebel.torchscript_to_ir(
+            dec_scripted_model,
+            input_names=[v[0] for v in dec_rbln_runtime_config.input_info],
+        )
 
-        compile_gpt2()
-        rbln_config.save(save_dir_path)
+        connections = [(prefill_ir.outputs[1 + i], prefill_ir.inputs[3 + i]) for i in range(model.config.n_layer * 2)]
 
-        return cls._from_pretrained(
-            model_id=save_dir_path,
-            config=config,
-            model_save_dir=save_dir,
-            **rbln_constructor_kwargs,
-            **kwargs,
+        compiled_model = rebel.compile(
+            prefill_ir,
+            dec_ir,
+            connections=connections,
+            fusion=prefill_rbln_runtime_config.fusion,
+            npu=prefill_rbln_runtime_config.npu,
+            tensor_parallel_size=prefill_rbln_runtime_config.tensor_parallel_size,
+            use_weight_sharing=True,
         )
 
+        return compiled_model
+
     @classmethod
     def _get_rbln_config(
         cls,
@@ -271,24 +199,24 @@ class RBLNGPT2LMHeadModel(RBLNBaseModel, RBLNGenerationMixin):
         def get_input_info(query_length):
             return [
                 ("input_ids", [rbln_batch_size, query_length], "int64"),
+                ("attention_mask", [rbln_batch_size, 1, query_length, rbln_max_seq_len], "int64"),
+                (
+                    "cache_position",
+                    [],
+                    "int32",
+                ),
+            ] + [
                 (
-                    "past_key_values",
+                    f"past_key_values_{i}",
                     [
-                        model_config.n_layer,
-                        2,
                         rbln_batch_size,
                         model_config.n_head,
                         rbln_max_seq_len,
                         model_config.hidden_size // model_config.n_head,
                     ],
                     "float32",
-                ),
-                ("attention_mask", [rbln_batch_size, 1, query_length, rbln_max_seq_len], "int64"),
-                (
-                    "cache_position",
-                    [],
-                    "int32",
-                ),
+                )
+                for i in range(model_config.n_layer * 2)
             ]
 
         # model input info
@@ -305,11 +233,14 @@ class RBLNGPT2LMHeadModel(RBLNBaseModel, RBLNGenerationMixin):
 
         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(input_info_index=0, tensor_type="pt", device=device_val),
-            self.compiled_models[0].create_runtime(input_info_index=1, tensor_type="pt", device=device_val),
+            compiled_models[0].create_runtime(input_info_index=0, tensor_type="pt", device=device_val),
+            compiled_models[0].create_runtime(input_info_index=1, tensor_type="pt", device=device_val),
         ]
 
     def prepare_inputs_for_generation(self, input_ids, past_key_values=0, attention_mask=None, **kwargs):
@@ -386,6 +317,3 @@ class RBLNGPT2LMHeadModel(RBLNBaseModel, RBLNGenerationMixin):
             output = output.logits
 
         return CausalLMOutputWithCrossAttentions(logits=output, past_key_values=past_key_values)
-
-    def __repr__(self):
-        return repr(self.runtimes[0]) + "\n" + repr(self.runtimes[1])

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

@@ -107,7 +107,6 @@ class _LlamaAttention(LlamaAttention):
         use_cache: bool = False,
         **kwargs,
     ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:
-
         bsz, q_len, _ = hidden_states.size()
 
         if self.config.pretraining_tp > 1:
@@ -227,7 +226,6 @@ class _LlamaDecoderLayer(LlamaDecoderLayer):
         use_cache: Optional[bool] = False,
         **kwargs,
     ) -> Tuple[torch.FloatTensor, Optional[Tuple[torch.FloatTensor, torch.FloatTensor]]]:
-
         residual = hidden_states
 
         hidden_states = self.input_layernorm(hidden_states)
@@ -414,7 +412,6 @@ class _LlamaForCausalLM(LlamaForCausalLM):
         output_hidden_states: Optional[bool] = None,
         return_dict: Optional[bool] = None,
     ) -> Union[Tuple, CausalLMOutputWithPast]:
-
         output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
         output_hidden_states = (
             output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
@@ -617,23 +614,23 @@ def apply_rotary_pos_emb(q, k, cos, sin, position_ids, unsqueeze_dim=1):
     return q_embed, k_embed
 
 
-def wrap_llama():
-    origin_mehtods = {}
-    origin_mehtods["LlamaRotaryEmbedding_INIT"] = LlamaRotaryEmbedding.__init__
-    origin_mehtods["LlamaRotaryEmbedding_forward"] = LlamaRotaryEmbedding.forward
-    origin_mehtods["LlamaModel_forward"] = LlamaModel.forward
-    origin_mehtods["LlamaForCausalLM_forward"] = LlamaForCausalLM.forward
+origin_methods = {}
+origin_methods["LlamaRotaryEmbedding_INIT"] = LlamaRotaryEmbedding.__init__
+origin_methods["LlamaRotaryEmbedding_forward"] = LlamaRotaryEmbedding.forward
+origin_methods["LlamaModel_forward"] = LlamaModel.forward
+origin_methods["LlamaForCausalLM_forward"] = LlamaForCausalLM.forward
 
+
+def wrap_llama():
     LlamaRotaryEmbedding.__init__ = _LlamaRotaryEmbedding.__init__
     LlamaRotaryEmbedding.forward = _LlamaRotaryEmbedding.forward
     LlamaModel.forward = _LlamaModel.forward
     LlamaForCausalLM.forward = _LlamaForCausalLM.forward
 
-    return origin_mehtods
-
 
-def unwrap_llama(origin_mehtods):
-    LlamaRotaryEmbedding.__init__ = origin_mehtods["LlamaRotaryEmbedding_INIT"]
-    LlamaRotaryEmbedding.forward = origin_mehtods["LlamaRotaryEmbedding_forward"]
-    LlamaModel.forward = origin_mehtods["LlamaModel_forward"]
-    LlamaForCausalLM.forward = origin_mehtods["LlamaForCausalLM_forward"]
+def unwrap_llama():
+    global origin_methods
+    LlamaRotaryEmbedding.__init__ = origin_methods["LlamaRotaryEmbedding_INIT"]
+    LlamaRotaryEmbedding.forward = origin_methods["LlamaRotaryEmbedding_forward"]
+    LlamaModel.forward = origin_methods["LlamaModel_forward"]
+    LlamaForCausalLM.forward = origin_methods["LlamaForCausalLM_forward"]

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

@@ -118,6 +118,9 @@ class _LlamaAttention(LlamaAttention):
         batch_index: Optional[int] = None,
         output_attentions: bool = False,
         use_cache: bool = False,
+        cos: Optional[torch.Tensor] = None,
+        sin: Optional[torch.Tensor] = None,
+        layer_id: int = 0,
         **kwargs,
     ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:
         bsz, q_len, _ = hidden_states.size()
@@ -156,8 +159,11 @@ class _LlamaAttention(LlamaAttention):
                     "with a layer index."
                 )
             kv_seq_len += past_key_value.get_usable_length(kv_seq_len, self.layer_idx)
-        cos, sin = self.rotary_emb(value_states, seq_len=kv_seq_len)
-        query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin, position_ids)
+        if layer_id == 0:
+            cos, sin = self.rotary_emb(value_states, seq_len=kv_seq_len)
+        query_states, key_states, cos, sin = apply_rotary_pos_emb(
+            query_states, key_states, cos, sin, position_ids, layer_id
+        )
         if past_key_value is not None:
             cache_kwargs = {"sin": sin, "cos": cos}  # Specific to RoPE models
         if (batch_index is None or batch_index == -1) and bsz > 1:
@@ -261,7 +267,7 @@ class _LlamaAttention(LlamaAttention):
         if not output_attentions:
             attn_weights = None
 
-        return attn_output, attn_weights, key_states, value_states
+        return attn_output, attn_weights, key_states, value_states, cos, sin
 
 
 class _LlamaDecoderLayer(LlamaDecoderLayer):
@@ -275,6 +281,9 @@ class _LlamaDecoderLayer(LlamaDecoderLayer):
         output_attentions: Optional[bool] = False,
         use_cache: Optional[bool] = False,
         batch_ids: Optional[torch.LongTensor] = None,
+        cos: Optional[torch.Tensor] = None,
+        sin: Optional[torch.Tensor] = None,
+        layer_id: int = 0,
         **kwargs,
     ) -> Tuple[torch.FloatTensor, Optional[Tuple[torch.FloatTensor, torch.FloatTensor]]]:
         residual = hidden_states
@@ -282,7 +291,7 @@ class _LlamaDecoderLayer(LlamaDecoderLayer):
         hidden_states = self.input_layernorm(hidden_states)
         bsz, _, _ = hidden_states.size()
 
-        hidden_states, self_attn_weights, k, v = _LlamaAttention.forward(
+        hidden_states, self_attn_weights, k, v, cos, sin = _LlamaAttention.forward(
             self.self_attn,
             hidden_states=hidden_states,
             attention_mask=attention_mask,
@@ -291,6 +300,9 @@ class _LlamaDecoderLayer(LlamaDecoderLayer):
             output_attentions=output_attentions,
             batch_index=batch_ids,
             use_cache=use_cache,
+            cos=cos,
+            sin=sin,
+            layer_id=layer_id,
             **kwargs,
         )
         past_key_value.assign(k, v, layer_idx)
@@ -313,7 +325,7 @@ class _LlamaDecoderLayer(LlamaDecoderLayer):
         if use_cache:
             outputs += (present_key_value,)
 
-        return outputs
+        return outputs, cos, sin
 
 
 class _LlamaModel(LlamaModel):
@@ -415,10 +427,11 @@ class _LlamaModel(LlamaModel):
         all_self_attns = () if output_attentions else None
         next_decoder_cache = () if use_cache else None
 
+        cos = None
+        sin = None
         for layer_idx, decoder_layer in enumerate(self.layers):
             if output_hidden_states:
                 all_hidden_states += (hidden_states,)
-
             layer_outputs = _LlamaDecoderLayer.forward(
                 decoder_layer,
                 hidden_states,
@@ -429,7 +442,13 @@ class _LlamaModel(LlamaModel):
                 output_attentions=output_attentions,
                 use_cache=use_cache,
                 batch_ids=batch_ids,
+                cos=cos,
+                sin=sin,
+                layer_id=layer_idx,
             )
+            cos = layer_outputs[-2]
+            sin = layer_outputs[-1]
+            layer_outputs = layer_outputs[0]
 
             hidden_states = layer_outputs[0]
 
@@ -697,7 +716,7 @@ def rotate_half(x):
     return torch.cat((-x2, x1), dim=-1)
 
 
-def apply_rotary_pos_emb(q, k, cos, sin, position_ids, unsqueeze_dim=1):
+def apply_rotary_pos_emb(q, k, cos, sin, position_ids, layer_id, unsqueeze_dim=1):
     """Applies Rotary Position Embedding to the query and key tensors.
 
     Args:
@@ -718,42 +737,28 @@ def apply_rotary_pos_emb(q, k, cos, sin, position_ids, unsqueeze_dim=1):
     Returns:
         `tuple(torch.Tensor)` comprising of the query and key tensors rotated using the Rotary Position Embedding.
     """
-    if position_ids.shape[0] > 1:
-        cos_all = []
-        sin_all = []
-        for i in range(position_ids.shape[0]):
-            cos_all.append(cos[position_ids[i : i + 1]].unsqueeze(unsqueeze_dim))
-            sin_all.append(sin[position_ids[i : i + 1]].unsqueeze(unsqueeze_dim))
-        cos = torch.cat(cos_all, dim=0)
-        sin = torch.cat(sin_all, dim=0)
-    else:
-        cos = cos[position_ids].unsqueeze(unsqueeze_dim)
-        sin = sin[position_ids].unsqueeze(unsqueeze_dim)
-    # cos = cos[position_ids].unsqueeze(unsqueeze_dim)
-    # sin = sin[position_ids].unsqueeze(unsqueeze_dim)
+    if layer_id == 0:
+        if position_ids.shape[0] > 1:
+            cos_all = []
+            sin_all = []
+            for i in range(position_ids.shape[0]):
+                cos_all.append(cos[position_ids[i : i + 1]].unsqueeze(unsqueeze_dim))
+                sin_all.append(sin[position_ids[i : i + 1]].unsqueeze(unsqueeze_dim))
+            cos = torch.cat(cos_all, dim=0)
+            sin = torch.cat(sin_all, dim=0)
+        else:
+            cos = cos[position_ids].unsqueeze(unsqueeze_dim)
+            sin = sin[position_ids].unsqueeze(unsqueeze_dim)
+        # cos = cos[position_ids].unsqueeze(unsqueeze_dim)
+        # sin = sin[position_ids].unsqueeze(unsqueeze_dim)
 
     q_embed = (q * cos) + (rotate_half(q) * sin)
     k_embed = (k * cos) + (rotate_half(k) * sin)
-    return q_embed, k_embed
+    return q_embed, k_embed, cos, sin
 
 
 def wrap_llama():
-    origin_mehtods = {}
-    origin_mehtods["LlamaRotaryEmbedding_INIT"] = LlamaRotaryEmbedding.__init__
-    origin_mehtods["LlamaRotaryEmbedding_forward"] = LlamaRotaryEmbedding.forward
-    origin_mehtods["LlamaModel_forward"] = LlamaModel.forward
-    origin_mehtods["LlamaForCausalLM_forward"] = LlamaForCausalLM.forward
-
     LlamaRotaryEmbedding.__init__ = _LlamaRotaryEmbedding.__init__
     LlamaRotaryEmbedding.forward = _LlamaRotaryEmbedding.forward
     LlamaModel.forward = _LlamaModel.forward
     LlamaForCausalLM.forward = _LlamaForCausalLM.forward
-
-    return origin_mehtods
-
-
-def unwrap_llama(origin_mehtods):
-    LlamaRotaryEmbedding.__init__ = origin_mehtods["LlamaRotaryEmbedding_INIT"]
-    LlamaRotaryEmbedding.forward = origin_mehtods["LlamaRotaryEmbedding_forward"]
-    LlamaModel.forward = origin_mehtods["LlamaModel_forward"]
-    LlamaForCausalLM.forward = origin_mehtods["LlamaForCausalLM_forward"]