optimum-rbln 0.7.4a1__tar.gz → 0.7.4a2__tar.gz

{optimum_rbln-0.7.4a1 → optimum_rbln-0.7.4a2}/.github/scripts/auto_code_review.py

@@ -22,8 +22,7 @@ from github import Github
 
 model_name = os.environ["GOOGLE_MODEL_ID"]
 genai.configure(api_key=os.environ["GOOGLE_API_KEY"])
-max_context_token = 100000
-force_review = False
+max_context_token = 500000
 
 
 def get_pr_diff():
@@ -58,19 +57,6 @@ Review the following code changes(GIT DIFF) along with the pull request (PR) det
     return system_prompt, prompt
 
 
-def translate_review(review):
-    model = genai.GenerativeModel(
-        model_name,
-        system_instruction="You are a professional translator specializing in technical and software-related content. Keep the technical words in English, but understand the whole sentence and rephrase it in Korean.",
-    )
-    prompt = f"""Translate the following English text into Korean, maintaining technical accuracy and clarity. Include ONLY the translation, NO OTHER EXPLANATIONS or RESPONSES as a chatbot. :
-
-{review}"""
-    response = model.generate_content(prompt)
-
-    return response.text
-
-
 def review_code(system_prompt, prompt):
     model = genai.GenerativeModel(model_name, system_instruction=system_prompt)
     response = model.generate_content(prompt)
@@ -125,7 +111,7 @@ def main():
     system_prompt, prompt = get_prompt(diff, pr)
     model = genai.GenerativeModel(model_name=model_name, system_instruction=system_prompt)
     num_tokens = model.count_tokens(prompt).total_tokens
-    if num_tokens > max_context_token and not force_review:
+    if num_tokens > max_context_token:
         msg = f"Diff ({len(diff)}) exceeds maximum allowed tokens ({max_context_token}) > ({num_tokens}). Skipping review."
         print(msg)
         pr.create_issue_comment(msg)
@@ -133,14 +119,10 @@ def main():
 
     # Get Auto review
     review = review_code(system_prompt, prompt)
-    translation = translate_review(review)
 
     # Post comment on PR
-    pr.create_issue_comment(f"""# Auto Code Review
-
-- [참고] Auto Code Review를 invoke하려면, commit message의 시작을 [autoreview]로 시작하거나, "/autoreview" 를 comment로 작성한 후,
-해당 commit의 github action에서 code review를 re-run 하시면 됩니다.
-\n\n{review}\n\n{translation}""")
+    pr.create_issue_comment(f"""# Auto Code Review by {model_name}
+\n\n{review}""")
 
 
 if __name__ == "__main__":

{optimum_rbln-0.7.4a1 → optimum_rbln-0.7.4a2}/.github/workflows/auto_code_review.yml

@@ -14,7 +14,7 @@ env:
 
 jobs:
   auto-review:
-    runs-on: ubuntu-latest
+    runs-on: ubuntu-latest-rbln
     steps:
       - name: Checkout repository
         uses: actions/checkout@v3

{optimum_rbln-0.7.4a1 → optimum_rbln-0.7.4a2}/.github/workflows/check_code_quality.yml

@@ -25,7 +25,7 @@ jobs:
       fail-fast: false
       matrix:
         python-version: [3.9]
-        os: [ubuntu-latest]
+        os: [ubuntu-latest-rbln]
 
     runs-on: ${{ matrix.os }}
     steps:

{optimum_rbln-0.7.4a1 → optimum_rbln-0.7.4a2}/.github/workflows/pr-title-check.yaml

@@ -20,7 +20,7 @@ permissions:
 jobs:
   title-checker:
     name: Check PR title
-    runs-on: ubuntu-latest
+    runs-on: ubuntu-latest-rbln
     steps:
       - uses: amannn/action-semantic-pull-request@v5
         id: lint_pr_title

{optimum_rbln-0.7.4a1 → optimum_rbln-0.7.4a2}/.github/workflows/pr_checklist_validator.yml

@@ -10,7 +10,7 @@ on:
 
 jobs:
   validate-pr-checklist:
-    runs-on: ubuntu-latest
+    runs-on: ubuntu-latest-rbln
     steps:
       - name: Checkout repository
         uses: actions/checkout@v2

{optimum_rbln-0.7.4a1 → optimum_rbln-0.7.4a2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: optimum-rbln
-Version: 0.7.4a1
+Version: 0.7.4a2
 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.
 Project-URL: Homepage, https://rebellions.ai
 Project-URL: Documentation, https://docs.rbln.ai

{optimum_rbln-0.7.4a1 → optimum_rbln-0.7.4a2}/src/optimum/rbln/__version__.py

@@ -17,5 +17,5 @@ __version__: str
 __version_tuple__: VERSION_TUPLE
 version_tuple: VERSION_TUPLE
 
-__version__ = version = '0.7.4a1'
+__version__ = version = '0.7.4a2'
 __version_tuple__ = version_tuple = (0, 7, 4)

{optimum_rbln-0.7.4a1 → optimum_rbln-0.7.4a2}/src/optimum/rbln/transformers/models/decoderonly/modeling_decoderonly.py

@@ -578,11 +578,41 @@ class RBLNDecoderOnlyModelForCausalLM(RBLNModel):
         nbits_per_param: int,
         n_model_params: int,
     ) -> int:
+        """
+        We are finding max_n_blocks(x) that satisfies the following equation:
+
+        available_dram - kernel_size - buffer
+            - num_layers * 2 * tensor_parallel_size
+            * align_2MB(
+                x
+                * block_size
+                * align_64(head_dim)
+                * math.ceil(num_key_value_heads / tensor_parallel_size)
+                * 2
+            ) > 0
+
+        This inequality can be rewritten as follows:
+
+        a - c * align_2MB(b * x) > 0
+        where
+           a = available_dram - kernel_size - buffer
+           b = block_size * align_64(head_dim) * math.ceil(num_key_value_heads / tensor_parallel_size) * 2
+           c = num_layers * 2 * tensor_parallel_size
+
+        We can rewrite the inequality as follows:
+        k > align_2MB(b*x)
+        where
+           k = a / c
+
+        After that, we can derive the following equation:
+        x = floor(2**21 / b * floor((k - 1) / 2**21))
+        """
+
         def align(x: int, nbytes: int) -> int:
             return int(math.ceil(x / nbytes) * nbytes)
 
         def align_2MB(x: int) -> int:
-            return align(x, 2 * 1024 * 1024)
+            return align(x, 2**21)
 
         num_attention_heads = getattr(config, "n_head", None) or getattr(config, "num_attention_heads")
         num_layers = getattr(config, "n_layer", None) or getattr(config, "num_hidden_layers")
@@ -612,27 +642,16 @@ class RBLNDecoderOnlyModelForCausalLM(RBLNModel):
         available_dram -= kernel_size
 
         # TODO: Accurate buffer estimation
-        buffer = 2**30  # 1GB Buffer
-        if tensor_parallel_size <= 4:
-            buffer /= 4
-
+        buffer_per_core = 2**29  # 500MB per npu
+        buffer = buffer_per_core * tensor_parallel_size
         available_dram -= buffer
 
-        # Estimate nbytes per a single kvcache block
-        nbytes_per_block = (
-            align_2MB(
-                kvcache_block_size
-                * head_dim
-                * math.ceil(num_key_value_heads / tensor_parallel_size)  # Shard
-                * 2  # (fp16)
-            )
-            * num_layers
-            * 2  # (k, v)
-            * tensor_parallel_size
-        )
-        n_blocks = available_dram // nbytes_per_block
+        b = kvcache_block_size * align(head_dim, 64) * math.ceil(num_key_value_heads / tensor_parallel_size) * 2
+        c = num_layers * 2 * tensor_parallel_size
+        k = available_dram / c
+        max_n_blocks = math.floor(2**21 / b * math.floor((k - 1) / 2**21))
 
-        return n_blocks, nbytes_per_block
+        return max_n_blocks
 
     @classmethod
     def _get_rbln_config(
@@ -689,7 +708,7 @@ class RBLNDecoderOnlyModelForCausalLM(RBLNModel):
 
         rbln_kvcache_num_blocks = (rbln_max_seq_len // rbln_kvcache_block_size) * rbln_batch_size
         if rbln_attn_impl == "flash_attn":
-            max_num_blocks, _ = cls.get_maximum_num_blocks(
+            max_num_blocks = cls.get_maximum_num_blocks(
                 config=model_config,
                 tensor_parallel_size=rbln_kwargs.get("tensor_parallel_size", 1),
                 kvcache_block_size=rbln_kvcache_block_size,

{optimum_rbln-0.7.4a1 → optimum_rbln-0.7.4a2}/src/optimum/rbln/transformers/models/t5/modeling_t5.py

@@ -36,19 +36,50 @@ from .t5_architecture import T5Wrapper
 logger = get_logger()
 
 if TYPE_CHECKING:
+    from rebel import Runtime
     from transformers import AutoFeatureExtractor, AutoProcessor, AutoTokenizer, PreTrainedModel
 
 
 class RBLNRuntimeModel(RBLNPytorchRuntime):
+    def __init__(
+        self,
+        runtime: "Runtime",
+        max_seq_len: int,
+        **kwargs: Any,
+    ) -> None:
+        super().__init__(runtime, **kwargs)
+        self.max_seq_len = max_seq_len
+
+    def _prepare_inputs(
+        self,
+        input_ids: torch.LongTensor,
+        attention_mask: torch.LongTensor,
+    ):
+        input_len = input_ids.shape[-1]
+        pad_len = None
+        if input_len > self.max_seq_len:
+            raise ValueError(f"Error input_len({input_len}) exceed max_seq_len({self.max_seq_len}).")
+        elif input_len < self.max_seq_len and input_len > 0:
+            pad_len = self.max_seq_len - input_len
+            logger.warning(
+                f"Warning: The input was padded with {pad_len} tokens to meet the compiled model's requirements. "
+                "For optimal performance, consider recompiling with a shorter 'rbln_max_seq_len'."
+            )
+            input_ids = torch.nn.functional.pad(input_ids, (0, pad_len))
+            attention_mask = torch.nn.functional.pad(attention_mask, (0, pad_len), value=0)
+
+        return input_ids, attention_mask, pad_len
+
     def forward(
         self,
         input_ids: torch.LongTensor,
-        attention_mask: torch.FloatTensor,
+        attention_mask: torch.LongTensor,
         head_mask: torch.FloatTensor,
         inputs_embeds: torch.FloatTensor,
         **kwargs,
     ):
-        return super().forward(
+        input_ids, attention_mask, pad_len = self._prepare_inputs(input_ids, attention_mask)
+        logits = super().forward(
             input_ids,
             attention_mask,
             head_mask,
@@ -56,6 +87,8 @@ class RBLNRuntimeModel(RBLNPytorchRuntime):
             **kwargs,
         )
 
+        return logits[:, :-pad_len, :] if pad_len is not None else logits
+
 
 class T5EncoderWrapper(torch.nn.Module):
     def __init__(self, model: "T5EncoderModel") -> None:
@@ -72,7 +105,8 @@ class RBLNT5EncoderModel(RBLNModel):
     rbln_model_input_names = ["input_ids", "attention_mask"]
 
     def __post_init__(self, **kwargs):
-        self.model = RBLNRuntimeModel(runtime=self.model[0])
+        max_seq_len = self.rbln_config.model_cfg["max_seq_len"]
+        self.model = RBLNRuntimeModel(runtime=self.model[0], max_seq_len=max_seq_len)
 
     @classmethod
     def wrap_model_if_needed(self, model: "PreTrainedModel", rbln_config: "RBLNConfig"):

{optimum_rbln-0.7.4a1 → optimum_rbln-0.7.4a2}/src/optimum/rbln/transformers/models/whisper/whisper_architecture.py

@@ -190,11 +190,11 @@ class WhisperDecoder(nn.Module):
         all_hiddens = []
         for i in range(inputs_embeds.shape[0]):
             position_id = cache_position[i]
-            position = self.embed_positions(input_ids, position_ids=position_id)
+            position = self.embed_positions.weight[position_id]
             batch_hidden = position + inputs_embeds[i]
             all_hiddens.append(batch_hidden)
 
-        hidden_states = torch.stack(all_hiddens, dim=0)
+        hidden_states = torch.cat(all_hiddens, dim=0).unsqueeze(1)
 
         # prepare attn mask (normal attention - masked)
         if attention_mask is not None: