Cleanup

Signed-off-by: Rafael Vasquez <rafvasq21@gmail.com>

Author: rafvasq

aiu_fms_testing_utils/scripts/drive_paged_programs.py

@@ -377,11 +377,8 @@ def __load_validation_info(
         **distributed_kwargs,
     )
 
-dprint("model.eval()")
 model.eval()
-dprint("finished eval")
 
-######################################### Building valid prompts
 @dataclass
 class ProgramInfo:
     program_id: str
@@ -406,6 +403,7 @@ def parse_program_limit(limit_str: str) -> tuple[int, str]:
         limit_val = int(match.group(2))
     return limit_val, limit_type
 
+# TODO: Add a check or logic for case that prog criteria json must exist if saving CPU outputs
 with open(args.program_criteria_json_path, "r") as f:
     program_criteria_json_list = json.load(f)["programs"]
     program_criteria_list = []
@@ -578,12 +576,6 @@ def parse_program_limit(limit_str: str) -> tuple[int, str]:
                 f"no valid prompt shape was found which would result in program {program_id} that satisfied batch{batch_size_limit_type}{batch_size_limit} and prompt_length{prompt_length_limit_type}{prompt_length_limit}"
             )
 
-dprint("valid prompts are prepared.")
-################################################################
-
-
-############## saving CPU validation ###########################
-
 if not args.skip_validation:
     with stagger_region(args.stagger_load):
         validation_model = get_model(
@@ -611,7 +603,6 @@ def parse_program_limit(limit_str: str) -> tuple[int, str]:
         )
         # if the cpu validation info is not yet computed, compute it
         if cpu_validation_info is None:
-            dprint(f"Let's start to compute it.")
             cpu_validation_info = extract_validation_information(
                 validation_model,
                 input_ids,
@@ -620,7 +611,6 @@ def parse_program_limit(limit_str: str) -> tuple[int, str]:
                 attn_algorithm="math",
                 **extra_kwargs,
             )
-            dprint(f"Save that cpu info.")
             # save the cpu validation info for later consumption
             if save_validation_info_outputs:
                 cpu_validation_info.save(
@@ -638,12 +628,11 @@ def parse_program_limit(limit_str: str) -> tuple[int, str]:
                 )
 
 if args.stop_after_info_outputs:
-    dprint("CPU validation info saved. Exiting as requested.")
-    # sys.exit(0)
+    dprint("CPU validation outputs saved. Exiting as requested.")
+    exit(0)
 
 ################################################################
 
-dprint("onto model compilation related stuff ...")
 fx_config.backed_size_oblivious = True
 model.compile(backend="sendnn", options={"sendnn.dynamic": True})
 __maybe_prepare_fp8_weights(model, is_fp8)