🚧 wip stuff

Signed-off-by: Prashant Gupta <prashantgupta@us.ibm.com>

Author: prashantgupta24

aiu_fms_testing_utils/scripts/refactored_dpp.py

@@ -189,6 +189,11 @@ def parse_cli_args() -> argparse.Namespace:
         action="store_true",
         help="set to true to save cpu validation outputs for later consumption",
     )
+    parser.add_argument(
+        "--save_validation_info_outputs_only",
+        action="store_true",
+        help="set to true to save cpu validation outputs for later consumption",
+    )
     parser.add_argument(
         "--prioritize_large_batch_sizes",
         action="store_true",
@@ -595,15 +600,11 @@ def get_program_prompt_list(
 def run_validation(
     args: argparse.Namespace,
     model: torch.nn.Module,
-    validation_model: Optional[torch.nn.Module],
     program_id: int,
     valid_prompt,
     input_ids: torch.Tensor,
     extra_kwargs: Dict[str, Any],
-    sample_key: str,
-    attn_name: str,
-    cpu_dtype: str,
-    tokenizer: AutoTokenizer,
+    cpu_validation_info: ValidationInfo
 ):
 
     if local_rank == 0:
@@ -612,46 +613,46 @@ def run_validation(
             f"program id: {program_id}, valid prompt: {valid_prompt}, input shape: {input_ids.shape}"
         )
 
-    cpu_validation_info: Optional[ValidationInfo] = None
-    if not args.skip_validation:
-        # attempt to load the cpu validation info if it is already computed
-        cpu_validation_info = __load_validation_info(
-            model_variant=args.model_variant,
-            batch_size=valid_prompt[0],
-            seq_length=valid_prompt[1],
-            max_new_tokens=args.max_new_tokens,
-            tokenizer=tokenizer,
-            seed=0,
-            cpu_dtype=cpu_dtype,
-            attn_type=attn_name,
-            validation_info_outputs_dir=args.validation_info_outputs_dir,
-            sample_key=sample_key,
-        )
-        # if the cpu validation info is not yet computed, compute it
-        if cpu_validation_info is None and validation_model is not None:
-            cpu_validation_info = extract_validation_information(
-                model=validation_model,
-                input_ids=input_ids,
-                max_new_tokens=args.max_new_tokens,
-                post_iteration_hook=LogitsExtractorHook(),
-                attn_algorithm="math",
-                **extra_kwargs,
-            )
-            # save the cpu validation info if requested
-            if args.save_validation_info_outputs:
-                cpu_validation_info.save(
-                    get_validation_info_path(
-                        validation_info_dir=args.validation_info_outputs_dir,
-                        model_variant=args.model_variant,
-                        batch_size=valid_prompt[0],
-                        seq_length=valid_prompt[1],
-                        max_new_tokens=args.max_new_tokens,
-                        seed=0,
-                        attn_type=attn_name,
-                        dtype=cpu_dtype,
-                        sample_key=sample_key,
-                    )
-                )
+
+    # if not args.skip_validation:
+    #     # attempt to load the cpu validation info if it is already computed
+    #     cpu_validation_info = __load_validation_info(
+    #         model_variant=args.model_variant,
+    #         batch_size=valid_prompt[0],
+    #         seq_length=valid_prompt[1],
+    #         max_new_tokens=args.max_new_tokens,
+    #         tokenizer=tokenizer,
+    #         seed=0,
+    #         cpu_dtype=cpu_dtype,
+    #         attn_type=attn_name,
+    #         validation_info_outputs_dir=args.validation_info_outputs_dir,
+    #         sample_key=sample_key,
+    #     )
+    #     # if the cpu validation info is not yet computed, compute it
+    #     if cpu_validation_info is None and validation_model is not None:
+    #         cpu_validation_info = extract_validation_information(
+    #             model=validation_model,
+    #             input_ids=input_ids,
+    #             max_new_tokens=args.max_new_tokens,
+    #             post_iteration_hook=LogitsExtractorHook(),
+    #             attn_algorithm="math",
+    #             **extra_kwargs,
+    #         )
+    #         # save the cpu validation info if requested
+    #         if args.save_validation_info_outputs:
+    #             cpu_validation_info.save(
+    #                 get_validation_info_path(
+    #                     validation_info_dir=args.validation_info_outputs_dir,
+    #                     model_variant=args.model_variant,
+    #                     batch_size=valid_prompt[0],
+    #                     seq_length=valid_prompt[1],
+    #                     max_new_tokens=args.max_new_tokens,
+    #                     seed=0,
+    #                     attn_type=attn_name,
+    #                     dtype=cpu_dtype,
+    #                     sample_key=sample_key,
+    #                 )
+    #             )
 
     golden_hook = None
     if args.test_type == "metrics":
@@ -800,50 +801,6 @@ def run_tokens_test(
 
 __maybe_prepare_fp8_weights(model, is_fp8)
 
-# Load validation model
-validation_model = None
-if not args.skip_validation:
-    validation_model = load_model(
-        device_type="cpu",
-        model_variant=args.model_variant,
-        is_fp8=is_fp8,
-        distributed_kwargs=distributed_kwargs,
-        stagger_load=args.stagger_load,
-        is_validation=True
-    )
-
-## MODEL WARMUP ##
-
-# warmup with any input so compiler produces criteria json
-# TODO: Swap this with __prepare_inputs once fix for shape_id is available
-# input_ids, extra_kwargs, sample_key = __prepare_inputs(2, max_tkv, tokenizer)
-prompt_list = [torch.arange(0, 64, dtype=torch.int64)]
-# matching vllm warmup to pad to 2 on fp8, and no pad for fp16
-if is_fp8:
-    prompt_list = prompt_list * 2
-input_ids, extra_kwargs = pad_input_ids(prompt_list, min_pad_length=64)
-extra_kwargs["mask"] = extra_kwargs["mask"].to(torch.float16)
-
-extra_kwargs["attn_name"] = ATTN_NAME
-if ( "granite-3.3-8b-instruct" in args.model_variant and args.distributed and dist.get_world_size() == 4):
-    extra_kwargs["_kvcache_num_blocks_hint"] = KVCACHE_NUM_BLOCKS_HINT
-
-warmup_model(
-    model=model,
-    input_ids=input_ids,
-    max_new_tokens=args.max_new_tokens,
-    compile_dynamic_sendnn=True,
-    stagger_update_lazyhandle=args.stagger_update_lazyhandle,
-    prefill_chunk_size=args.prefill_chunk_size,
-    **extra_kwargs,
-)
-
-if args.distributed:
-    # wait for rank0 to be finished as it is the only one generating the criteria json
-    # this is needed since otherwise we may run into a race condition
-    torch.distributed.barrier()
-
-
 ## PREPARE PROGRAM CRITERIA AND PROMPTS ##
 
 with open(args.program_criteria_json_path, "r") as f:
@@ -891,25 +848,134 @@ def run_tokens_test(
     custom_shape=custom_shape,
 )
 
+
+## CPU validation
+
+def get_cpu_validation(
+    args: argparse.Namespace,
+    valid_prompt,
+    input_ids: torch.Tensor,
+    extra_kwargs: Dict[str, Any],
+    sample_key: str,
+    attn_name: str,
+    cpu_dtype: str,
+    tokenizer: AutoTokenizer,
+):
+    if not args.skip_validation:
+        dprint("Generating CPU validation for prompt: {}".format(valid_prompt))
+        # Load validation model
+        validation_model = load_model(
+            device_type="cpu",
+            model_variant=args.model_variant,
+            is_fp8=is_fp8,
+            distributed_kwargs=distributed_kwargs,
+            stagger_load=args.stagger_load,
+            is_validation=True
+        )
+
+        # attempt to load the cpu validation info if it is already computed
+        cpu_validation_info = __load_validation_info(
+            model_variant=args.model_variant,
+            batch_size=valid_prompt[0],
+            seq_length=valid_prompt[1],
+            max_new_tokens=args.max_new_tokens,
+            tokenizer=tokenizer,
+            seed=0,
+            cpu_dtype=cpu_dtype,
+            attn_type=attn_name,
+            validation_info_outputs_dir=args.validation_info_outputs_dir,
+            sample_key=sample_key,
+        )
+
+        if cpu_validation_info is not None:
+            dprint("cpu validation info found, returning it")
+            return cpu_validation_info
+        dprint("cpu validation info not found, computing it now")
+        # if the cpu validation info is not yet computed, compute it
+        if validation_model is not None:
+            dprint("extracting cpu validation info")
+            cpu_validation_info = extract_validation_information(
+                model=validation_model,
+                input_ids=input_ids,
+                max_new_tokens=args.max_new_tokens,
+                post_iteration_hook=LogitsExtractorHook(),
+                attn_algorithm="math",
+                **extra_kwargs,
+            )
+            dprint("cpu validation info extracted")
+            # save the cpu validation info if requested
+            if args.save_validation_info_outputs:
+                dprint("saving cpu validation info")
+                cpu_validation_info.save(
+                    get_validation_info_path(
+                        validation_info_dir=args.validation_info_outputs_dir,
+                        model_variant=args.model_variant,
+                        batch_size=valid_prompt[0],
+                        seq_length=valid_prompt[1],
+                        max_new_tokens=args.max_new_tokens,
+                        seed=0,
+                        attn_type=attn_name,
+                        dtype=cpu_dtype,
+                        sample_key=sample_key,
+                    )
+                )
+                dprint("cpu validation info saved")
+for program_id, valid_prompt, input_ids, extra_kwargs, sample_key in valid_prompts:
+    cpu_validation_info: Optional[ValidationInfo] = None
+    cpu_validation_info = get_cpu_validation(args, valid_prompt, input_ids, extra_kwargs, sample_key, ATTN_NAME, CPU_DTYPE, tokenizer)
+
+if args.save_validation_info_outputs_only:
+    dprint("CPU validation information saved. Exiting.")
+    exit(0)
+
+
+## MODEL WARMUP ##
+
+# warmup with any input so compiler produces criteria json
+# TODO: Swap this with __prepare_inputs once fix for shape_id is available
+# input_ids, extra_kwargs, sample_key = __prepare_inputs(2, max_tkv, tokenizer)
+prompt_list = [torch.arange(0, 64, dtype=torch.int64)]
+# matching vllm warmup to pad to 2 on fp8, and no pad for fp16
+if is_fp8:
+    prompt_list = prompt_list * 2
+input_ids, extra_kwargs = pad_input_ids(prompt_list, min_pad_length=64)
+extra_kwargs["mask"] = extra_kwargs["mask"].to(torch.float16)
+
+extra_kwargs["attn_name"] = ATTN_NAME
+if ( "granite-3.3-8b-instruct" in args.model_variant and args.distributed and dist.get_world_size() == 4):
+    extra_kwargs["_kvcache_num_blocks_hint"] = KVCACHE_NUM_BLOCKS_HINT
+
+warmup_model(
+    model=model,
+    input_ids=input_ids,
+    max_new_tokens=args.max_new_tokens,
+    compile_dynamic_sendnn=True,
+    stagger_update_lazyhandle=args.stagger_update_lazyhandle,
+    prefill_chunk_size=args.prefill_chunk_size,
+    **extra_kwargs,
+)
+
+if args.distributed:
+    # wait for rank0 to be finished as it is the only one generating the criteria json
+    # this is needed since otherwise we may run into a race condition
+    torch.distributed.barrier()
+
+
 ## RUN VALIDATION AND TESTS ##
 
 failed_cases = []
 # for each program and valid prompt (batch size, sequence length)
 for program_id, valid_prompt, input_ids, extra_kwargs, sample_key in valid_prompts:
     extra_kwargs["attn_name"] = ATTN_NAME
 
-    aiu_validation_info, cpu_validation_info = run_validation(
+    aiu_validation_info = run_validation(
                 args=args,
                 model=model, 
-                validation_model=validation_model, 
                 program_id=program_id, 
                 valid_prompt=valid_prompt, 
                 input_ids=input_ids, 
-                extra_kwargs=extra_kwargs, 
-                sample_key=sample_key, 
-                attn_name=ATTN_NAME, 
-                cpu_dtype=CPU_DTYPE, 
-                tokenizer=tokenizer,
+                extra_kwargs=extra_kwargs,
+                cpu_validation_info=cpu_validation_info
             )
 
     if args.test_type == "metrics":