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

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

@@ -21,61 +21,23 @@
 # copied, modified, or distributed without prior written permission
 # from Rebellions Inc.
 
-import inspect  # noqa: I001
+import inspect
 import logging
-from pathlib import Path
-from tempfile import TemporaryDirectory
-from typing import TYPE_CHECKING, Any, Callable, Dict, List, Optional, Tuple, Union
+from typing import Any, Callable
 
-import torch  # noqa: F401
-import rebel  # noqa: F401
+from transformers import LlamaForCausalLM, PreTrainedModel
 
-from optimum.exporters import TasksManager
-from transformers import AutoModelForCausalLM, LlamaForCausalLM, PretrainedConfig, AutoConfig
-from transformers.modeling_outputs import CausalLMOutputWithPast
-
-from ...generation.utils import RBLNGenerationMixin
-from ....modeling_base import RBLNBaseModel
-from ....modeling_config import DEFAULT_COMPILED_MODEL_NAME, RBLNConfig, RBLNRuntimeConfig
-from ....utils.runtime_utils import RBLNPytorchRuntime
-from ....utils.save_utils import maybe_save_preprocessors
-
-
-# FIXME:: Merge Two architecture Codes
-from .llama_architecture import (
-    LlamaWrapper,
-    wrap_llama,
-    unwrap_llama,
-)
-
-from .llama_architecture_cb import (
-    LlamaDynamicBatchWrapper as LlamaWrapper_cb,
-    wrap_llama as wrap_llama_cb,
-)
+from ...models.decoderonly import RBLNDecoderOnlyModelForCausalLM
+from .llama_architecture import LlamaWrapper
 
 
 logger = logging.getLogger(__name__)
 
-if TYPE_CHECKING:
-    from transformers import (
-        AutoFeatureExtractor,
-        AutoProcessor,
-        AutoTokenizer,
-        PretrainedConfig,
-    )
-
 
-SUPPORTED_BATCHING_MODES = ["static", "vllm"]
-
-
-class RBLNRuntimeModel(RBLNPytorchRuntime):
-    mandatory_members = ["main_input_name"]
-
-
-class RBLNLlamaForCausalLM(RBLNBaseModel, RBLNGenerationMixin):
+class RBLNLlamaForCausalLM(RBLNDecoderOnlyModelForCausalLM):
     """
     The Llama Model transformer with a language modeling head (linear layer) on top.
-    This model inherits from [`RBLNBaseModel`]. Check the superclass documentation for the generic methods the library implements for all its models.
+    This model inherits from [`RBLNMultiModel`]. Check the superclass documentation for the generic methods the library implements for all its models.
 
     A class to convert and run pre-trained transformers based LlamaForCausalLM model on RBLN devices.
     It implements the methods to convert a pre-trained transformers LlamaForCausalLM model into a RBLN transformer model by:
@@ -83,273 +45,9 @@ class RBLNLlamaForCausalLM(RBLNBaseModel, RBLNGenerationMixin):
     - compiling the resulting graph using the RBLN compiler.
     """
 
-    model_type = "rbln_model"
-    main_input_name = "input_ids"
-    auto_model_class = AutoModelForCausalLM
-
-    def __post_init__(self, **kwargs):
-        self.batch_size = self.rbln_config.meta["rbln_batch_size"]
-        self.max_seq_len = self.rbln_config.meta["rbln_max_seq_len"]
-        self.prefill_chunk_size = self.rbln_config.meta["rbln_prefill_chunk_size"]
-        self.use_continuous_batch = self.rbln_config.meta["rbln_batching"] == "vllm"
-
-        prefill_batch_size = self.batch_size if not self.use_continuous_batch else 1
-        self.prefill_attention_mask = torch.zeros(
-            prefill_batch_size, 1, self.prefill_chunk_size, self.max_seq_len, dtype=torch.int64
-        )
-        self.causal_mask = 1 - torch.triu(
-            torch.ones(prefill_batch_size, 1, self.prefill_chunk_size, self.prefill_chunk_size), diagonal=1
-        )
-        self.decoder_attention_mask = torch.zeros(self.batch_size, 1, 1, self.max_seq_len, dtype=torch.int64)
-
-        self.prefill_decoder = RBLNRuntimeModel(runtime=self.runtimes[0], main_input_name="input_ids")
-        self.decoder = RBLNRuntimeModel(runtime=self.runtimes[1], main_input_name="input_ids")
-        self.past_cached_length = 0
-        self.right_padding = True
-
-    @classmethod
-    @torch.no_grad()
-    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,
-    ) -> "RBLNLlamaForCausalLM":
-        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_configs(kwargs):
-            hf_max_position_embeddings = getattr(AutoConfig.from_pretrained(model_id), "max_position_embeddings", None)
-            max_seq_len = kwargs.get("rbln_max_seq_len", None)
-            if max_seq_len is not None:
-                if max_seq_len <= hf_max_position_embeddings:
-                    kwargs.update({"max_position_embeddings": max_seq_len})
-                else:
-                    raise ValueError("`max_seq_len` should be less or equal than max_position_embeddings!")
-
-            kwargs.update(
-                {
-                    "torchscript": True,
-                    "return_dict": False,
-                    "use_cache": True,
-                    "torch_dtype": torch.float32,
-                    "_attn_implementation": "eager",
-                }
-            )
-
-            return kwargs
-
-        kwargs = update_configs(kwargs)
-
-        rbln_config_kwargs, rbln_constructor_kwargs = cls.pop_rbln_kwargs_from_kwargs(kwargs)
-
-        # FIXME :: This should be moved when wrapping removed.
-        use_continuous_batch = rbln_config_kwargs.get("rbln_batching", "static") == "vllm"
-        origin_mehtods = wrap_llama_cb() if use_continuous_batch else wrap_llama()
-
-        model: LlamaForCausalLM = 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_llama(use_continuous_batch, wrapper_cls):
-            wrapped_model = wrapper_cls(model).eval()
-
-            prefill_rbln_runtime_config = rbln_config[DEFAULT_COMPILED_MODEL_NAME][0]
-            dec_rbln_runtime_config = rbln_config[DEFAULT_COMPILED_MODEL_NAME][1]
-
-            prefill_example_inputs = prefill_rbln_runtime_config.get_dummy_inputs(fill=0)
-            dec_example_inputs = dec_rbln_runtime_config.get_dummy_inputs(fill=4)
-
-            if use_continuous_batch:
-                batch_index_index = 3
-                dec_example_inputs[batch_index_index].fill_(-1)  # fill batch_position -1 to indicate it is decoder.
-
-            prefill_scripted_model = torch.jit.trace(wrapped_model, prefill_example_inputs)
-            dec_scripted_model = torch.jit.trace(wrapped_model, dec_example_inputs)
-
-            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],
-            )
-
-            # Caching prefill_decoder/decoder I/O
-            cache_index_offset = 4 if use_continuous_batch else 3
-            connections = [
-                (prefill_ir.outputs[1 + i], prefill_ir.inputs[cache_index_offset + i])
-                for i in range(model.config.num_hidden_layers * 2)
-            ]
-
-            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")
-
-        wrapper_cls = LlamaWrapper_cb if use_continuous_batch else LlamaWrapper
-        compile_llama(use_continuous_batch=use_continuous_batch, wrapper_cls=wrapper_cls)
-        unwrap_llama(origin_mehtods)
-
-        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,
-        )
-
     @classmethod
-    def _get_rbln_config(
-        cls,
-        preprocessors: Union["AutoFeatureExtractor", "AutoProcessor", "AutoTokenizer"],
-        model_config: "PretrainedConfig",
-        rbln_max_seq_len: Optional[int] = None,
-        rbln_batch_size: Optional[int] = None,
-        rbln_batching: Optional[str] = None,
-    ) -> RBLNConfig:
-        meta = {}
-
-        prefill_chunk_size = 128
-        if rbln_max_seq_len is None:
-            rbln_max_seq_len = getattr(model_config, "max_position_embeddings", None)
-        rbln_batch_size = 1 if rbln_batch_size is None else rbln_batch_size
-        rbln_batching = "static" if rbln_batching is None else rbln_batching
-
-        meta["rbln_max_seq_len"] = rbln_max_seq_len
-        meta["rbln_batch_size"] = rbln_batch_size
-        meta["rbln_prefill_chunk_size"] = prefill_chunk_size
-        meta["rbln_batching"] = rbln_batching
-        use_continuous_batching = meta["rbln_batching"] == "vllm"
-
-        if rbln_batching not in SUPPORTED_BATCHING_MODES:
-            raise ValueError(
-                f'rbln_batching="{rbln_batching}" is not a supported batch mode, '
-                f"Possible: {SUPPORTED_BATCHING_MODES}"
-            )
-
-        def get_input_info(
-            batch_size,  # should be 1 if continous batch prefill
-            query_length,
-            continuous_batch=False,  # determines the shape of `cache position`
-        ):
-            input_info = [
-                ("input_ids", [batch_size, query_length], "int64"),
-                ("attention_mask", [batch_size, 1, query_length, rbln_max_seq_len], "int64"),
-                (
-                    "cache_position",
-                    [batch_size, query_length] if continuous_batch else [],
-                    "int32",
-                ),
-            ]
-
-            if continuous_batch:
-                input_info.append(("batch_position", [], "int16"))
-
-            input_info.extend(
-                [
-                    (
-                        f"past_key_values_{i}",
-                        [
-                            rbln_batch_size,
-                            model_config.num_key_value_heads,
-                            rbln_max_seq_len,
-                            model_config.hidden_size // model_config.num_attention_heads,
-                        ],
-                        "float32",
-                    )
-                    for i in range(model_config.num_hidden_layers * 2)
-                ]
-            )
-
-            return input_info
-
-        prefill_input_info = get_input_info(
-            batch_size=1 if use_continuous_batching else rbln_batch_size,
-            query_length=prefill_chunk_size,
-            continuous_batch=use_continuous_batching,
-        )
-        dec_input_info = get_input_info(
-            batch_size=rbln_batch_size,
-            query_length=1,
-            continuous_batch=use_continuous_batching,
-        )
-
-        prefill_rbln_runtime_config = RBLNRuntimeConfig(input_info=prefill_input_info)
-        dec_rbln_runtime_config = RBLNRuntimeConfig(input_info=dec_input_info)
-
-        dec_rbln_runtime_config.batch_size = rbln_batch_size
-
-        rbln_config = RBLNConfig.from_rbln_runtime_configs(
-            [prefill_rbln_runtime_config, dec_rbln_runtime_config],
-            _rbln_meta=meta,
-        )
-
-        return rbln_config
-
-    def _create_runtimes(self, 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),
-        ]
-
-    def get_decoder(self):
-        return self.decoder
-
-    def can_generate(self):
-        return True
+    def wrapping_torch_model(self, model: "PreTrainedModel", rbln_max_seq_len: int):
+        return LlamaWrapper(model, rbln_max_seq_len).eval()
 
     def __getattr__(self, __name: str) -> Any:
         def redirect(func):
@@ -361,161 +59,3 @@ class RBLNLlamaForCausalLM(RBLNBaseModel, RBLNGenerationMixin):
             return redirect(val)
 
         return val
-
-    def _reorder_cache(self, past_key_values, beam_idx):
-        raise NotImplementedError
-
-    # args input_ids, past_key_values and attention_mask are updated by _update_model_kwargs_for_generation() in _greedy_search() in GenerationMixin
-    def prepare_inputs_for_generation(self, input_ids, past_key_values=0, attention_mask=None, **kwargs):
-        batch_size, cur_len = input_ids.shape
-        past_cached_length = past_key_values
-
-        # In greedy decoding
-        if past_cached_length == 0:
-            # padding with prefill_chunk_size
-            # TODO left padding + left padding has issue on stoppingcriteria(max_len)
-            if cur_len % self.prefill_chunk_size != 0:
-                pad_len = self.prefill_chunk_size - cur_len % self.prefill_chunk_size
-                input_ids = torch.nn.functional.pad(input_ids, (0, pad_len))
-
-            # padding_side
-            if batch_size > 1 and torch.all(attention_mask[..., -1] == 1):
-                self.right_padding = False
-
-            if self.right_padding:
-                self.rightpad_max_len = cur_len
-                prompt_min_len = torch.min(torch.sum(attention_mask, dim=-1))
-                self.dummy_len = torch.sum(attention_mask, dim=-1) - prompt_min_len  # dummy_decoder generation length
-                query_length = prompt_min_len.item()
-            else:
-                query_length = cur_len - past_cached_length
-                self.prompt_length = query_length
-                self.prompt_attn_mask = attention_mask.unsqueeze(1).unsqueeze(1).contiguous()
-
-            attention_mask = self.prefill_attention_mask.clone()
-            cache_position = torch.tensor(0, dtype=torch.int32)
-
-        else:
-            if self.right_padding:
-                attention_mask = torch.zeros(batch_size, 1, 1, self.max_seq_len, dtype=torch.int64)
-                attention_mask[:, :, :, : past_cached_length + 1] = 1
-                input_ids = input_ids[:, past_cached_length : past_cached_length + 1].contiguous()
-            else:
-                attention_mask = torch.nn.functional.pad(attention_mask, (0, self.max_seq_len - cur_len))
-                attention_mask = attention_mask.reshape(batch_size, 1, 1, -1).contiguous()
-                input_ids = input_ids[:, -1:]
-
-            cache_position = torch.tensor(past_cached_length, dtype=torch.int32)
-            query_length = 1
-
-        model_inputs = {
-            "input_ids": input_ids,
-            "past_key_values": past_key_values,
-            "attention_mask": attention_mask,
-            "cache_position": cache_position,
-            "query_length": query_length,
-        }
-
-        return model_inputs
-
-    def forward(self, *args, **kwargs):
-        if self.use_continuous_batch:
-            return self.forward_cb(*args, **kwargs)
-        else:
-            return self.forward_static(*args, **kwargs)
-
-    def forward_static(
-        self,
-        input_ids: torch.LongTensor = None,
-        attention_mask: Optional[torch.Tensor] = None,
-        past_key_values: int = None,
-        cache_position: Optional[torch.Tensor] = None,
-        query_length: Optional[torch.Tensor] = None,
-        **kwargs,
-    ) -> Tuple[torch.FloatTensor]:
-        if past_key_values is not None:
-            past_key_values += query_length
-
-        # prefill_decoder
-        if cache_position == 0:
-            for step in range(0, query_length, self.prefill_chunk_size):
-                sliced_input_ids = input_ids[:, step : step + self.prefill_chunk_size]
-                attention_mask[:, :, :, :step] = 1
-                attention_mask[:, :, :, step : step + self.prefill_chunk_size] = self.causal_mask
-                if not self.right_padding:
-                    attention_mask[:, :, :, : self.prompt_length] &= self.prompt_attn_mask[:, :, :, :]
-
-                outputs = self.prefill_decoder(
-                    input_ids=sliced_input_ids.contiguous(),
-                    attention_mask=attention_mask.contiguous(),
-                    cache_position=cache_position + step,
-                )
-            outputs = outputs[:, query_length % self.prefill_chunk_size - 1].unsqueeze(1)
-
-        # decoder
-        else:
-            outputs = self.decoder(
-                input_ids.contiguous(),
-                attention_mask.contiguous(),
-                cache_position=cache_position,
-            )
-
-        return CausalLMOutputWithPast(
-            logits=outputs,
-            past_key_values=past_key_values,
-        )
-
-    def forward_cb(
-        self,
-        input_ids: torch.LongTensor = None,
-        cache_position: Optional[torch.Tensor] = None,  # torch.tensor(,dtype=int32) (1,64) // (4,1)
-        batch_idx: int = None,
-        **kwargs,
-    ) -> Tuple[torch.FloatTensor]:
-        # prefill_decoder
-        if cache_position.shape[1] > 1:
-            query_length = input_ids.shape[1]
-            attention_mask = self.prefill_attention_mask.clone()
-            for step in range(0, query_length, self.prefill_chunk_size):
-                if step + self.prefill_chunk_size > query_length:
-                    input_ids = torch.nn.functional.pad(input_ids, (0, step + self.prefill_chunk_size - query_length))
-                    cache_position = torch.cat(
-                        [
-                            cache_position,
-                            torch.arange(
-                                query_length,
-                                step + self.prefill_chunk_size,
-                                dtype=torch.int32,
-                            ).unsqueeze(0),
-                        ],
-                        dim=-1,
-                    )
-
-                sliced_input_ids = input_ids[:, step : step + self.prefill_chunk_size]
-                sliced_cache_positions = cache_position[:, step : step + self.prefill_chunk_size]
-                attention_mask[:, :, :, :step] = 1
-                attention_mask[:, :, :, step : step + self.prefill_chunk_size] = self.causal_mask
-
-                outputs, _ = self.prefill_decoder(
-                    sliced_input_ids.contiguous(),
-                    attention_mask.contiguous(),
-                    sliced_cache_positions.contiguous(),
-                    torch.tensor(batch_idx, dtype=torch.int16),
-                )
-            outputs = outputs[:, query_length % self.prefill_chunk_size - 1].unsqueeze(1)
-        # decoder
-        else:
-            attention_mask = self.decoder_attention_mask.clone()
-            for b_idx in range(self.batch_size):
-                attention_mask[b_idx, :, :, : cache_position[b_idx].item() + 1] = 1
-
-            outputs = self.decoder(
-                input_ids.contiguous(),
-                attention_mask.contiguous(),
-                cache_position.contiguous(),
-                torch.tensor(0, dtype=torch.int16),
-            )[0]
-
-        return CausalLMOutputWithPast(
-            logits=outputs,
-        )

optimum/rbln/transformers/models/midm/hf_hub_cached/midm_bitext_tokenization.py

@@ -10,7 +10,8 @@
 # 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.
-""" Tokenization class for model Midm_bitext_tonkenizer."""
+"""Tokenization class for model Midm_bitext_tonkenizer."""
+
 import os
 import re
 import warnings

optimum/rbln/transformers/models/midm/hf_hub_cached/modeling_midm.py

@@ -817,7 +817,6 @@ class MidmModel(MidmPreTrainedModel):
         all_cross_attentions = () if output_attentions and self.config.add_cross_attention else None
         all_hidden_states = () if output_hidden_states else None
         for i, (block, layer_past) in enumerate(zip(self.h, past_key_values)):
-
             # Model parallel
             if self.model_parallel:
                 torch.cuda.set_device(hidden_states.device)
@@ -833,7 +832,6 @@ class MidmModel(MidmPreTrainedModel):
                 all_hidden_states = all_hidden_states + (hidden_states,)
 
             if self.gradient_checkpointing and self.training:
-
                 if use_cache:
                     logger.warning(
                         "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`..."
@@ -1174,7 +1172,6 @@ class MidmDoubleHeadsModel(MidmPreTrainedModel):
         return_dict=None,
         **kwargs,
     ):
-
         return_dict = return_dict if return_dict is not None else self.config.use_return_dict
 
         transformer_outputs = self.transformer(
@@ -1445,7 +1442,6 @@ def get_submodule(module, target: str):  # -> "Module":
     mod: torch.nn.Module = module
 
     for item in atoms:
-
         if not hasattr(mod, item):
             raise AttributeError(mod._get_name() + " has no " "attribute `" + item + "`")