optimum-rbln 0.1.0__py3-none-any.whl → 0.1.1__py3-none-any.whl

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

@@ -23,7 +23,6 @@
 
 import inspect
 import logging
-import warnings
 from pathlib import Path
 from tempfile import TemporaryDirectory
 from typing import TYPE_CHECKING, Any, Callable, Dict, List, Optional, Tuple, Union
@@ -32,19 +31,13 @@ import rebel
 import torch
 from optimum.exporters import TasksManager
 from transformers import AutoModelForCausalLM, GPT2LMHeadModel, PretrainedConfig
-from transformers.generation.logits_process import LogitsProcessorList
-from transformers.generation.stopping_criteria import (
-    StoppingCriteriaList,
-    validate_stopping_criteria,
-)
-from transformers.generation.streamers import BaseStreamer
-from transformers.generation.utils import SampleDecoderOnlyOutput
 from transformers.modeling_outputs import CausalLMOutputWithCrossAttentions, Seq2SeqLMOutput
 
 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
+from ...generation.utils import RBLNGenerationMixin
 from .gpt2_architecture import GPT2LMHeadModelWrapper
 
 
@@ -66,7 +59,7 @@ class RBLNRuntimeDecoder(RBLNPytorchRuntime):
         return Seq2SeqLMOutput(logits=logits)
 
 
-class RBLNGPT2LMHeadModel(RBLNBaseModel):
+class RBLNGPT2LMHeadModel(RBLNBaseModel, RBLNGenerationMixin):
     """
     The GPT2 Model transformer with a language modeling head on top (linear layer with weights tied to the input
     embeddings).
@@ -135,6 +128,7 @@ class RBLNGPT2LMHeadModel(RBLNBaseModel):
         subfolder: str = "",
         local_files_only: bool = False,
         trust_remote_code: bool = False,
+        model_save_dir: Optional[Union[str, Path, TemporaryDirectory]] = None,
         **kwargs,
     ) -> "RBLNGPT2LMHeadModel":
         """
@@ -144,8 +138,16 @@ class RBLNGPT2LMHeadModel(RBLNBaseModel):
         if task is None:
             task = TasksManager.infer_task_from_model(cls.auto_model_class)
 
-        save_dir = TemporaryDirectory()
-        save_dir_path = Path(save_dir.name)
+        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)
 
         kwargs.update(
             {
@@ -264,8 +266,7 @@ class RBLNGPT2LMHeadModel(RBLNBaseModel):
         meta["rbln_max_seq_len"] = rbln_max_seq_len
         meta["rbln_pad_token_id"] = rbln_pad_token_id
 
-        if rbln_batch_size is None:
-            rbln_batch_size = 1
+        rbln_batch_size = 1 if rbln_batch_size is None else rbln_batch_size
 
         def get_input_info(query_length):
             return [
@@ -320,6 +321,7 @@ class RBLNGPT2LMHeadModel(RBLNBaseModel):
             self.prompt_ids = input_ids
             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
 
             if cur_len % self.prefill_chunk_size == 0:
                 pad_len = 0
@@ -329,12 +331,12 @@ class RBLNGPT2LMHeadModel(RBLNBaseModel):
             attention_mask = self.prefill_attention_mask.clone()
             cache_position = torch.tensor(past_cached_length, dtype=torch.int32)
 
-            query_length = prompt_min_len
+            query_length = prompt_min_len.item()
         else:
             cache_position = torch.tensor(past_cached_length, dtype=torch.int32)
             attention_mask = torch.zeros(batch_size, 1, 1, self.max_seq_len, dtype=torch.int64)
             attention_mask[:, :, :, : cache_position + 1] = 1
-            input_ids = input_ids[:, -1:].contiguous()
+            input_ids = input_ids[:, cache_position : cache_position + 1].contiguous()
             query_length = 1
 
         model_inputs = {
@@ -357,25 +359,23 @@ class RBLNGPT2LMHeadModel(RBLNBaseModel):
         query_length: Optional[torch.Tensor] = None,
         **kwargs,
     ) -> Union[Tuple, CausalLMOutputWithCrossAttentions]:
-
         if past_key_values is not None:
             past_key_values += query_length
 
         if cache_position == 0:
-            for _ in range(0, query_length, self.prefill_chunk_size):
-                sliced_input_ids = input_ids[:, cache_position : cache_position + self.prefill_chunk_size]
-                attention_mask[:, :, :, :cache_position] = 1
-                attention_mask[:, :, :, cache_position : cache_position + self.prefill_chunk_size] = self.causal_mask
+            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
 
                 output = self.prefill_decoder(
                     input_ids=sliced_input_ids.contiguous(),
                     attention_mask=attention_mask.contiguous(),
-                    cache_position=cache_position,
+                    cache_position=cache_position + step,
                 )
-                query_length -= self.prefill_chunk_size
-                cache_position += self.prefill_chunk_size
 
-            output = output.logits[:, query_length - 1].unsqueeze(1)
+            idx = query_length % self.prefill_chunk_size - 1
+            output = output.logits[:, idx].unsqueeze(1)
 
         else:
             output = self.decoder(
@@ -389,312 +389,3 @@ class RBLNGPT2LMHeadModel(RBLNBaseModel):
 
     def __repr__(self):
         return repr(self.runtimes[0]) + "\n" + repr(self.runtimes[1])
-
-    # call 'greedy_search` directly is deprecated and removed in v4.41.
-    def greedy_search(self, *args, **kwargs):
-        return self._greedy_search(*args, **kwargs)
-
-    def _greedy_search(
-        self,
-        input_ids: torch.LongTensor,
-        logits_processor: Optional[LogitsProcessorList] = None,
-        stopping_criteria: Optional[StoppingCriteriaList] = None,
-        max_length: Optional[int] = None,
-        pad_token_id: Optional[int] = None,
-        eos_token_id: Optional[Union[int, List[int]]] = None,
-        output_logits: Optional[bool] = None,
-        return_dict_in_generate: Optional[bool] = None,
-        streamer: Optional["BaseStreamer"] = None,
-        **model_kwargs,
-    ) -> Union[SampleDecoderOnlyOutput, torch.LongTensor]:
-
-        # init values
-        logits_processor = logits_processor if logits_processor is not None else LogitsProcessorList()
-        stopping_criteria = stopping_criteria if stopping_criteria is not None else StoppingCriteriaList()
-
-        if max_length is not None:
-            warnings.warn(
-                "`max_length` is deprecated in this function, use"
-                " `stopping_criteria=StoppingCriteriaList([MaxLengthCriteria(max_length=max_length)])` instead.",
-                UserWarning,
-            )
-            stopping_criteria = validate_stopping_criteria(stopping_criteria, max_length)
-
-        pad_token_id = pad_token_id if pad_token_id is not None else self.generation_config.pad_token_id
-        eos_token_id = eos_token_id if eos_token_id is not None else self.generation_config.eos_token_id
-        if isinstance(eos_token_id, int):
-            eos_token_id = [eos_token_id]
-        eos_token_id_tensor = torch.tensor(eos_token_id).to(input_ids.device) if eos_token_id is not None else None
-
-        return_dict_in_generate = (
-            return_dict_in_generate
-            if return_dict_in_generate is not None
-            else self.generation_config.return_dict_in_generate
-        )
-
-        # init attention / hidden states / scores tuples
-        raw_logits = () if (return_dict_in_generate and output_logits) else None
-
-        # keep track of which sequences are already finished
-        unfinished_sequences = torch.ones(input_ids.shape[0], dtype=torch.long, device=input_ids.device)
-
-        this_peer_finished = False  # used by synced_gpus only
-
-        while True:
-            # prepare model inputs
-            model_inputs = self.prepare_inputs_for_generation(input_ids, **model_kwargs)
-
-            # forward pass to get next token
-            outputs = self(
-                **model_inputs,
-                return_dict=True,
-            )
-            next_token_logits = outputs.logits[:, -1, :]
-
-            # pre-process distribution
-            next_tokens_scores = logits_processor(input_ids, next_token_logits)
-
-            # Store scores, attentions and hidden_states when required
-            if return_dict_in_generate:
-                if output_logits:
-                    raw_logits += (next_token_logits,)
-
-            # argmax
-            next_tokens = torch.argmax(next_tokens_scores, dim=-1)
-
-            # finished sentences should have their next token be a padding token
-            if eos_token_id is not None:
-                if pad_token_id is None:
-                    raise ValueError("If `eos_token_id` is defined, make sure that `pad_token_id` is defined.")
-                next_tokens = next_tokens * unfinished_sequences + pad_token_id * (1 - unfinished_sequences)
-
-            ########################################################################################################
-            # thkim change for right-padding batch
-            # if min_input_len <= update_idx < max_input_len
-            #   update validate input_ids[:,update_idx]
-            # TODO : raw_logits contains dummy next_token's logits
-            update_idx = model_inputs["cache_position"] + model_inputs["query_length"]
-            if update_idx < self.rightpad_max_len:
-                # update exist input_ids rather than concat
-                valid_indices = model_kwargs["attention_mask"][:, update_idx] == 0
-                input_ids[valid_indices, update_idx] = next_tokens[valid_indices]
-                model_kwargs["attention_mask"][valid_indices, update_idx] = 1
-
-                # dummy next_token -> pad_token_id for streamer
-                # in order to skip by 'skip_special_tokens = True"
-                dummy_indices = ~valid_indices
-                next_tokens[dummy_indices] = pad_token_id
-            else:
-                ############################################END#########################################################
-                # update generated ids, model inputs, and length for next step
-                input_ids = torch.cat([input_ids, next_tokens[:, None]], dim=-1)
-
-                model_kwargs = self._update_model_kwargs_for_generation(
-                    outputs,
-                    model_kwargs,
-                    is_encoder_decoder=self.config.is_encoder_decoder,
-                )
-
-            if streamer is not None:
-                streamer.put(next_tokens.cpu())
-
-            # if eos_token was found in one sentence, set sentence to finished
-            if eos_token_id_tensor is not None:
-                ####################################################################
-                # thkim : to do not finish sequence of dummy_decoder of right_padding
-                if hasattr(self, "rightpad_max_len"):
-                    update_idx = model_inputs["cache_position"] + model_inputs["query_length"]
-                    if update_idx < self.rightpad_max_len:
-                        next_tokens += model_kwargs["attention_mask"][:, update_idx] * eos_token_id_tensor
-                ######################################################################
-                unfinished_sequences = unfinished_sequences.mul(
-                    next_tokens.tile(eos_token_id_tensor.shape[0], 1).ne(eos_token_id_tensor.unsqueeze(1)).prod(dim=0)
-                )
-
-                # stop when each sentence is finished
-                if unfinished_sequences.max() == 0:
-                    this_peer_finished = True
-
-            # stop if we exceed the maximum length
-            # thkim : backward compatibility bool vs torch.BoolTensor
-            is_stop = stopping_criteria(input_ids, None)
-            if isinstance(is_stop, torch.BoolTensor):
-                is_stop = torch.all(is_stop)
-            if is_stop:
-                this_peer_finished = True
-
-            if this_peer_finished:
-                break
-
-        if streamer is not None:
-            streamer.end()
-
-        if return_dict_in_generate:
-            return SampleDecoderOnlyOutput(
-                sequences=input_ids,
-                logits=raw_logits,
-            )
-        else:
-            return input_ids
-
-    # call 'sample` directly is deprecated and removed in v4.41.
-    def sample(self, *args, **kwargs):
-        return self._sample(*args, **kwargs)
-
-    def _sample(
-        self,
-        input_ids: torch.LongTensor,
-        logits_processor: Optional[LogitsProcessorList] = None,
-        stopping_criteria: Optional[StoppingCriteriaList] = None,
-        logits_warper: Optional[LogitsProcessorList] = None,
-        max_length: Optional[int] = None,
-        pad_token_id: Optional[int] = None,
-        eos_token_id: Optional[Union[int, List[int]]] = None,
-        output_attentions: Optional[bool] = None,
-        output_hidden_states: Optional[bool] = None,
-        output_scores: Optional[bool] = None,
-        output_logits: Optional[bool] = None,
-        return_dict_in_generate: Optional[bool] = None,
-        synced_gpus: bool = False,
-        streamer: Optional["BaseStreamer"] = None,
-        **model_kwargs,
-    ) -> Union[SampleDecoderOnlyOutput, torch.LongTensor]:
-        # init values
-        logits_processor = logits_processor if logits_processor is not None else LogitsProcessorList()
-        stopping_criteria = stopping_criteria if stopping_criteria is not None else StoppingCriteriaList()
-
-        if max_length is not None:
-            warnings.warn(
-                "`max_length` is deprecated in this function, use"
-                " `stopping_criteria=StoppingCriteriaList([MaxLengthCriteria(max_length=max_length)])` instead.",
-                UserWarning,
-            )
-            stopping_criteria = validate_stopping_criteria(stopping_criteria, max_length)
-
-        logits_warper = logits_warper if logits_warper is not None else LogitsProcessorList()
-        pad_token_id = pad_token_id if pad_token_id is not None else self.generation_config.pad_token_id
-        eos_token_id = eos_token_id if eos_token_id is not None else self.generation_config.eos_token_id
-
-        if isinstance(eos_token_id, int):
-            eos_token_id = [eos_token_id]
-        eos_token_id_tensor = torch.tensor(eos_token_id).to(input_ids.device) if eos_token_id is not None else None
-
-        output_scores = output_scores if output_scores is not None else self.generation_config.output_scores
-        output_logits = output_logits if output_logits is not None else False
-
-        # init attention / hidden states / scores tuples
-        scores = () if (return_dict_in_generate and output_scores) else None
-        raw_logits = () if (return_dict_in_generate and output_logits) else None
-
-        # keep track of which sequences are already finished
-        batch_size, cur_len = input_ids.shape
-        unfinished_sequences = torch.ones(batch_size, dtype=torch.long, device=input_ids.device)
-        this_peer_finished = False
-
-        # model_kwargs["cache_position"] = torch.arange(cur_len, device=input_ids.device)
-
-        while True:
-            # prepare model inputs
-            model_inputs = self.prepare_inputs_for_generation(input_ids, **model_kwargs)
-
-            # forward pass to get next token
-            outputs = self(
-                **model_inputs,
-                return_dict=True,
-                output_attentions=output_attentions,
-                output_hidden_states=output_hidden_states,
-            )
-
-            next_token_logits = outputs.logits[:, -1, :]
-
-            # pre-process distribution
-            next_token_scores = logits_processor(input_ids, next_token_logits)
-            next_token_scores = logits_warper(input_ids, next_token_scores)
-
-            # Store scores, attentions and hidden_states when required
-            if return_dict_in_generate:
-                if output_scores:
-                    scores += (next_token_scores,)
-                if output_logits:
-                    raw_logits += (next_token_logits,)
-
-            # sample
-            probs = torch.nn.functional.softmax(next_token_scores, dim=-1)
-            next_tokens = torch.multinomial(probs, num_samples=1).squeeze(1)
-
-            # finished sentences should have their next token be a padding token
-            if eos_token_id is not None:
-                if pad_token_id is None:
-                    raise ValueError("If `eos_token_id` is defined, make sure that `pad_token_id` is defined.")
-                next_tokens = next_tokens * unfinished_sequences + pad_token_id * (1 - unfinished_sequences)
-
-            ########################################################################################################
-            # thkim change for right-padding batch
-            # if min_input_len <= update_idx < max_input_len
-            #   update validate input_ids[:,update_idx]
-            # TODO : raw_logits contains dummy next_token's logits
-            update_idx = model_inputs["cache_position"] + model_inputs["query_length"]
-            if update_idx < self.rightpad_max_len:
-                # update exist input_ids rather than concat
-                valid_indices = model_kwargs["attention_mask"][:, update_idx] == 0
-                input_ids[valid_indices, update_idx] = next_tokens[valid_indices]
-                model_kwargs["attention_mask"][valid_indices, update_idx] = 1
-
-                # dummy next_token -> pad_token_id for streamer
-                # in order to skip by 'skip_special_tokens = True"
-                dummy_indices = ~valid_indices
-                next_tokens[dummy_indices] = pad_token_id
-            else:
-                ############################################END#########################################################
-
-                # update generated ids, model inputs, and length for next step
-                input_ids = torch.cat([input_ids, next_tokens[:, None]], dim=-1)
-
-                model_kwargs = self._update_model_kwargs_for_generation(
-                    outputs,
-                    model_kwargs,
-                    is_encoder_decoder=self.config.is_encoder_decoder,
-                )
-
-            if streamer is not None:
-                streamer.put(next_tokens.cpu())
-
-            # if eos_token was found in one sentence, set sentence to finished
-            if eos_token_id_tensor is not None:
-                ####################################################################
-                # thkim : to do not finish sequence of dummy_decoder of right_padding
-                if hasattr(self, "rightpad_max_len"):
-                    update_idx = model_inputs["cache_position"] + model_inputs["query_length"]
-                    if update_idx < self.rightpad_max_len:
-                        next_tokens += model_kwargs["attention_mask"][:, update_idx] * eos_token_id_tensor
-                ######################################################################
-                unfinished_sequences = unfinished_sequences.mul(
-                    next_tokens.tile(eos_token_id_tensor.shape[0], 1).ne(eos_token_id_tensor.unsqueeze(1)).prod(dim=0)
-                )
-
-                # stop when each sentence is finished
-                if unfinished_sequences.max() == 0:
-                    this_peer_finished = True
-
-            # stop if we exceed the maximum length
-            # thkim : backward compatibility bool vs list[bool]
-            is_stop = stopping_criteria(input_ids, None)
-            if isinstance(is_stop, torch.BoolTensor):
-                is_stop = torch.all(is_stop)
-            if is_stop:
-                this_peer_finished = True
-
-            if this_peer_finished:
-                break
-
-        if streamer is not None:
-            streamer.end()
-
-        if return_dict_in_generate:
-            return SampleDecoderOnlyOutput(
-                sequences=input_ids,
-                scores=scores,
-                logits=raw_logits,
-            )
-        else:
-            return input_ids

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

@@ -34,6 +34,7 @@ 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
@@ -75,7 +76,7 @@ class RBLNRuntimeModel(RBLNPytorchRuntime):
         return logits
 
 
-class RBLNLlamaForCausalLM(RBLNBaseModel):
+class RBLNLlamaForCausalLM(RBLNBaseModel, RBLNGenerationMixin):
     """
     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.
@@ -91,7 +92,6 @@ class RBLNLlamaForCausalLM(RBLNBaseModel):
     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"]
@@ -106,6 +106,7 @@ class RBLNLlamaForCausalLM(RBLNBaseModel):
         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()
@@ -120,14 +121,23 @@ class RBLNLlamaForCausalLM(RBLNBaseModel):
         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)
 
-        save_dir = TemporaryDirectory()
-        save_dir_path = Path(save_dir.name)
+        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)
@@ -245,6 +255,7 @@ class RBLNLlamaForCausalLM(RBLNBaseModel):
         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
 
         meta["rbln_max_seq_len"] = rbln_max_seq_len
         meta["rbln_batch_size"] = rbln_batch_size
@@ -321,23 +332,47 @@ class RBLNLlamaForCausalLM(RBLNBaseModel):
 
     # 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, hf_input_length = input_ids.shape
+        batch_size, cur_len = input_ids.shape
         past_cached_length = past_key_values
-        query_length = hf_input_length - past_cached_length
 
         # In greedy decoding
-        if past_key_values == 0:
-            self.prompt_length = query_length
-            self.prompt_ids = input_ids
-            self.prompt_attn_mask = attention_mask.unsqueeze(1).unsqueeze(1).contiguous()
-
-            attention_mask = torch.zeros(batch_size, 1, self.prefill_chunk_size, self.max_seq_len, dtype=torch.int64)
+        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:
-            attention_mask = torch.nn.functional.pad(attention_mask, (0, self.max_seq_len - hf_input_length))
-            attention_mask = attention_mask.reshape(batch_size, 1, 1, -1).contiguous()
+            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)
-            input_ids = input_ids[:, -1:]
+            query_length = 1
 
         model_inputs = {
             "input_ids": input_ids,
@@ -358,43 +393,26 @@ class RBLNLlamaForCausalLM(RBLNBaseModel):
         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:
-            while query_length > self.prefill_chunk_size:
-                # prepare input_ids & attention_mask
-                sliced_input_ids = input_ids[:, cache_position : cache_position + self.prefill_chunk_size].contiguous()
-                attention_mask[:, :, :, :cache_position] = 1
-                attention_mask[:, :, :, cache_position : cache_position + self.prefill_chunk_size] = self.causal_mask
-                attention_mask[:, :, :, : self.prompt_length] *= self.prompt_attn_mask[:, :, :, :]
-
-                _ = self.prefill_decoder(
-                    sliced_input_ids,
-                    attention_mask,
-                    cache_position,
+            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,
                 )
-                # update query_length & cache_position
-                query_length -= self.prefill_chunk_size
-                cache_position += self.prefill_chunk_size
-
-            # prepare input_ids & attention_mask
-            last_input_ids = input_ids[:, cache_position : cache_position + query_length]
-            last_input_ids = torch.nn.functional.pad(last_input_ids, (0, self.prefill_chunk_size - query_length))
-
-            attention_mask[:, :, :, :cache_position] = 1
-            mask_slice = self.causal_mask[:, :, :query_length, :query_length]
-            attention_mask[:, :, :query_length, cache_position : cache_position + query_length] = mask_slice
-            attention_mask[:, :, :, : self.prompt_length] *= self.prompt_attn_mask[:, :, :, :]
-
-            outputs = self.prefill_decoder(
-                last_input_ids.contiguous(),
-                attention_mask.contiguous(),
-                cache_position,
-            )
+            outputs = outputs[:, query_length % self.prefill_chunk_size - 1].unsqueeze(1)
 
-            outputs = outputs[:, query_length - 1].unsqueeze(1)
         # decoder
         else:
             outputs = self.decoder(

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

@@ -163,6 +163,7 @@ class RBLNWhisperForConditionalGeneration(RBLNBaseModel, GenerationMixin):
         subfolder: str = "",
         local_files_only: bool = False,
         trust_remote_code: bool = False,
+        model_save_dir: Optional[Union[str, Path, TemporaryDirectory]] = None,
         **kwargs,
     ) -> "RBLNWhisperForConditionalGeneration":
         """
@@ -172,8 +173,16 @@ class RBLNWhisperForConditionalGeneration(RBLNBaseModel, GenerationMixin):
         if task is None:
             task = TasksManager.infer_task_from_model(cls.auto_model_class)
 
-        save_dir = TemporaryDirectory()
-        save_dir_path = Path(save_dir.name)
+        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)
 
         kwargs.update(
             {
@@ -266,7 +275,7 @@ class RBLNWhisperForConditionalGeneration(RBLNBaseModel, GenerationMixin):
         cls,
         preprocessors: Union["AutoFeatureExtractor", "AutoProcessor"],
         model_config: "PretrainedConfig",
-        rbln_batch_size: Optional[int] = 1,
+        rbln_batch_size: Optional[int] = None,
     ) -> RBLNConfig:
         meta = {}
 
@@ -287,6 +296,7 @@ class RBLNWhisperForConditionalGeneration(RBLNBaseModel, GenerationMixin):
         if rbln_dec_max_seq_len is None:
             raise ValueError("`rbln_dec_max_seq_len` should be specified!")
 
+        rbln_batch_size = 1 if rbln_batch_size is None else rbln_batch_size
         decoder_batch_size = rbln_batch_size
 
         meta["rbln_dec_max_seq_len"] = rbln_dec_max_seq_len

{optimum_rbln-0.1.0.dist-info → optimum_rbln-0.1.1.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: optimum-rbln
-Version: 0.1.0
+Version: 0.1.1
 Summary: Optimum RBLN is the interface between the Hugging Face Transformers and Diffusers libraries and RBLN accelerators.
         It provides a set of tools enabling easy model loading and inference on single and multiple rbln device settings for different downstream tasks.
 Keywords: transformers,diffusers,inference,rbln,atom,rebel