fix test decoders

Signed-off-by: Antoni Viros i Martin <aviros@ibm.com>

Author: ani300

aiu_fms_testing_utils/testing/validation.py

@@ -187,10 +187,10 @@ def load_validation_information(validation_path, validation_files_type, batch_si
 
     return ValidationInfo(validation_info)
 
-def extract_validation_information(model, input_ids, max_new_tokens, post_iteration_hook, attn_algorithm=None, eos_token_id = None, only_last_token=False, timing="", **padding_kwargs):
+def extract_validation_information(model, input_ids, max_new_tokens, post_iteration_hook, attn_algorithm=None, eos_token_id = None, only_last_token=False, timing="", **extra_kwargs):
     max_seq_len = model.config.max_expected_seq_len
     attention_specific_kwargs = {}
-    if "paged" in padding_kwargs["attn_name"]:
+    if "paged" in extra_kwargs["attn_name"]:
         from aiu_fms_testing_utils.utils.paged import generate
     else:
         # TODO: Add a unified generation dependent on attn_type
@@ -199,7 +199,7 @@ def extract_validation_information(model, input_ids, max_new_tokens, post_iterat
         attention_specific_kwargs["max_seq_len"] = max_seq_len
 
     # Add only_last_token optimization
-    extra_generation_kwargs = {**padding_kwargs}
+    extra_generation_kwargs = {**extra_kwargs}
     if only_last_token:
         extra_generation_kwargs["only_last_token"] = only_last_token
     if attn_algorithm is not None:

aiu_fms_testing_utils/utils/__init__.py

@@ -9,10 +9,10 @@
 import json
 import random
 
-def warmup_model(model: nn.Module, input_ids: torch.Tensor, max_new_tokens: int, compile_dynamic_sendnn = False, use_cache: bool = True, **padding_kwargs):
+def warmup_model(model: nn.Module, input_ids: torch.Tensor, max_new_tokens: int, compile_dynamic_sendnn = False, use_cache: bool = True, **extra_kwargs):
     import torch_sendnn
     attention_specific_kwargs = {}
-    attn_name = padding_kwargs["attn_name"]
+    attn_name = extra_kwargs["attn_name"]
     if "paged" in attn_name:
         from aiu_fms_testing_utils.utils.paged import generate, adjust_inputs_to_batch
     else:
@@ -25,15 +25,15 @@ def warmup_model(model: nn.Module, input_ids: torch.Tensor, max_new_tokens: int,
 
     # adjust inputs depending on attn_type and dynamic shapes
     _warmup_input_ids = input_ids
-    _padding_kwargs = padding_kwargs
+    _extra_kwargs = extra_kwargs
     _max_new_tokens = max_new_tokens
     if compile_dynamic_sendnn:
         _max_new_tokens = 2
         # always warmup with batch size 2 when using attn_type=paged
         if "paged" in attn_name:
-            _warmup_input_ids, _padding_kwargs = adjust_inputs_to_batch(input_ids, **padding_kwargs)
+            _warmup_input_ids, _extra_kwargs = adjust_inputs_to_batch(input_ids, **extra_kwargs)
 
-    extra_kwargs = {**_padding_kwargs, "only_last_token": "paged" not in attn_name}
+    extra_kwargs = {**_extra_kwargs, "only_last_token": "paged" not in attn_name}
 
     with torch_sendnn.warmup_mode():
         generate(model, _warmup_input_ids, max_new_tokens=_max_new_tokens, do_sample=False, use_cache=use_cache, extra_kwargs=extra_kwargs, **attention_specific_kwargs)

aiu_fms_testing_utils/utils/paged.py

@@ -5,7 +5,7 @@
 import torch
 import fms.utils.spyre.paged
 
-def adjust_inputs_to_batch(input_ids: torch.Tensor, **padding_kwargs):
+def adjust_inputs_to_batch(input_ids: torch.Tensor, **extra_kwargs):
     """
     Adjusts the inputs to a batch. Batch size 1 cannot be handled since we want a symbolic shape for the batch 
     and pytorch automatically sets size 1 dimensions as static
@@ -14,11 +14,11 @@ def adjust_inputs_to_batch(input_ids: torch.Tensor, **padding_kwargs):
     """
     input_ids = input_ids[0].repeat(2, 1)
     # ensure we pass along other kwargs
-    kwargs = {**padding_kwargs}
-    mask = padding_kwargs.get("mask", None)
+    kwargs = {**extra_kwargs}
+    mask = extra_kwargs.get("mask", None)
     if mask is not None:
         kwargs["mask"] = torch.stack((mask[0], mask[0]))
-    position_ids = padding_kwargs.get("position_ids", None)
+    position_ids = extra_kwargs.get("position_ids", None)
     if position_ids is not None:
         kwargs["position_ids"] = position_ids[0].repeat(2, 1)
     return input_ids, kwargs

tests/models/test_decoders.py

@@ -499,6 +499,7 @@ def _metric_calculator(r: torch.Tensor, t: torch.Tensor):
                 input_ids, extra_kwargs = __prepare_inputs(
                     batch_size, seq_length, tokenizer, seed=i
                 )
+                extra_kwargs["attn_name"] = ATTN_NAME
                 cpu_validation_info = __load_validation_info(
                     model_path, batch_size, seq_length, max_new_tokens, tokenizer, i
                 )