Merge pull request #179 from foundation-model-stack/fix_add_eos_id

pass EOS token id to validate

Author: Ssukriti

aiu_fms_testing_utils/scripts/drive_paged_programs.py

@@ -682,6 +682,7 @@ def __metric_calculator(r: torch.Tensor, t: torch.Tensor):
                 max_new_tokens,
                 LogitsExtractorHook(),
                 attn_algorithm="math",
+                eos_token_id=tokenizer.eos_token_id,
                 **extra_kwargs,
             )
             # save the cpu validation info for later consumption
@@ -706,6 +707,7 @@ def __metric_calculator(r: torch.Tensor, t: torch.Tensor):
                 input_ids,
                 max_new_tokens,
                 GoldenTokenHook(cpu_validation_info.get_info("tokens")),
+                eos_token_id=tokenizer.eos_token_id,
                 last_n_tokens=64,
                 timing=TIMING,
                 prefill_chunk_size=args.prefill_chunk_size,
@@ -751,6 +753,7 @@ def __metric_calculator(r: torch.Tensor, t: torch.Tensor):
                 input_ids,
                 max_new_tokens,
                 None,
+                eos_token_id=tokenizer.eos_token_id,
                 last_n_tokens=64,
                 timing=TIMING,
                 prefill_chunk_size=args.prefill_chunk_size,
@@ -794,6 +797,7 @@ def __metric_calculator(r: torch.Tensor, t: torch.Tensor):
             input_ids,
             max_new_tokens,
             None,
+            eos_token_id=tokenizer.eos_token_id,
             last_n_tokens=64,
             timing=TIMING,
             prefill_chunk_size=args.prefill_chunk_size,

aiu_fms_testing_utils/scripts/generate_metrics.py

@@ -245,6 +245,7 @@ def write_csv(metrics, path, metric_name):
         args.max_new_tokens,
         LogitsExtractorHook(),
         attn_algorithm="math",
+        eos_token_id=tokenizer.eos_token_id,
         **padding_kwargs,
     )
     cpu_static_tokens = cpu_validation_info.get_info("tokens")
@@ -259,6 +260,7 @@ def write_csv(metrics, path, metric_name):
         ids.to("cuda"),
         args.max_new_tokens,
         None,
+        eos_token_id=tokenizer.eos_token_id,
         last_n_tokens=1,
         **{k: v.to("cuda") for k, v in padding_kwargs.items()},
     )
@@ -325,6 +327,7 @@ def write_csv(metrics, path, metric_name):
                 args.max_new_tokens,
                 LogitsExtractorHook(),
                 attn_algorithm="math",
+                eos_token_id=tokenizer.eos_token_id,
                 **padding_kwargs,
             )
 
@@ -334,6 +337,7 @@ def write_csv(metrics, path, metric_name):
             ids.to("cuda"),
             args.max_new_tokens,
             GoldenTokenHook(cpu_validation_info.get_info("tokens"), "cuda"),
+            eos_token_id=tokenizer.eos_token_id,
             last_n_tokens=1,
             **{k: v.to("cuda") for k, v in padding_kwargs.items()},
         )

aiu_fms_testing_utils/scripts/save_cpu_data.py

@@ -97,6 +97,7 @@ def process_row(row):
             max_new_tokens,
             LogitsExtractorHook(),
             attn_algorithm="math",
+            eos_token_id=tokenizer.eos_token_id,
         )
     return {"id": id, "input_ids": input_ids, "validation": cpu_validation_info}
 

aiu_fms_testing_utils/scripts/validation.py

@@ -690,6 +690,7 @@ def print_result(result, result_idx: int = 0, file_prefix: str = ""):
         ids.to(validation_device),
         args.max_new_tokens,
         LogitsExtractorHook(),
+        eos_token_id=None if args.no_early_termination else tokenizer.eos_token_id,
         attn_algorithm="math",
         **padding_kwargs,
     )

tests/models/test_decoders.py

@@ -612,6 +612,7 @@ def _get_device_validation_information(
         max_new_tokens,
         post_iteration_hook,
         timing=TIMING,
+        eos_token_id=tokenizer.eos_token_id,
         **extra_kwargs,
         **device_dependent_kwargs,
     )