Add compilation metrics, script to assert metrics, add conditional import of hugginface_hub and update huggingface hub

Signed-off-by: Kunjan patel <kunjanp@google.com>

Author: coolkp

end_to_end/tpu/eval_assert.py

@@ -46,15 +46,76 @@ def test_final_loss(metrics_file, target_loss):
     print("Final loss test passed.")
 
 
+def test_avg_step_time(metrics_file, max_avg_step_time_str, num_samples_str="10"):
+  """Tests if the average of the last N step times is below a maximum threshold."""
+  max_avg_step_time = float(max_avg_step_time_str)
+  num_samples = int(num_samples_str)
+  metric_key = "perf/step_time_seconds"
+  last_n_step_times = get_last_n_data(metrics_file, metric_key, num_samples)
+
+  if not last_n_step_times:
+    raise ValueError(f"Metric '{metric_key}' not found or no data points in {metrics_file}.")
+
+  avg_last_n_step_time = sum(last_n_step_times) / len(last_n_step_times)
+
+  print(f"Found {len(last_n_step_times)} data points for '{metric_key}'.")
+  print(f"Mean of last {len(last_n_step_times)} step times is {avg_last_n_step_time:.4f} s")
+
+  assert avg_last_n_step_time < max_avg_step_time, \
+      f"Average step time {avg_last_n_step_time:.4f}s is not less than target {max_avg_step_time}s."
+  print("Average step time test passed.")
+
+
+def test_avg_tflops(metrics_file, min_avg_tflops_str, num_samples_str="10"):
+  """Tests if the average of the last N TFLOPs/sec values is above a minimum threshold."""
+  min_avg_tflops = float(min_avg_tflops_str)
+  num_samples = int(num_samples_str)
+  metric_key = "perf/per_device_tflops_per_sec"
+
+  last_n_tflops = get_last_n_data(metrics_file, metric_key, num_samples)
+
+  if not last_n_tflops:
+    raise ValueError(f"Metric '{metric_key}' not found or no data points in {metrics_file}.")
+
+  avg_last_n_tflops = sum(last_n_tflops) / len(last_n_tflops)
+
+  print(f"Found {len(last_n_tflops)} data points for '{metric_key}'.")
+  print(f"Mean of last {len(last_n_tflops)} steps TFLOPs/sec is {avg_last_n_tflops:.2f}")
+
+  assert avg_last_n_tflops > min_avg_tflops, \
+      f"Average TFLOPs/sec {avg_last_n_tflops:.2f} is not greater than target {min_avg_tflops}."
+  print("Average TFLOPs/sec test passed.")
+
+
 def main(argv: Sequence[str]) -> None:
+  if len(argv) < 2:
+    print("Usage: python script.py <test_scenario> [test_vars...]")
+    print("Available scenarios: final_loss, avg_step_time, avg_tflops")
+    raise ValueError("Test scenario not specified.")
 
   _, test_scenario, *test_vars = argv
 
   if test_scenario == "final_loss":
-    test_final_loss(*test_vars)
+    if len(test_vars) < 2:
+      raise ValueError("Usage: final_loss <metrics_file> <target_loss> [num_samples]")
+    metrics_file, target_loss, *num_samples_opt = test_vars
+    num_samples = num_samples_opt[0] if num_samples_opt else "10"
+    test_final_loss(metrics_file, target_loss, num_samples)
+  elif test_scenario == "avg_step_time":
+    if len(test_vars) < 2:
+      raise ValueError("Usage: avg_step_time <metrics_file> <max_avg_step_time> [num_samples]")
+    metrics_file, max_avg_step_time, *num_samples_opt = test_vars
+    num_samples = num_samples_opt[0] if num_samples_opt else "10"
+    test_avg_step_time(metrics_file, max_avg_step_time, num_samples)
+  elif test_scenario == "avg_tflops":
+    if len(test_vars) < 2:
+      raise ValueError("Usage: avg_tflops <metrics_file> <min_avg_tflops> [num_samples]")
+    metrics_file, min_avg_tflops, *num_samples_opt = test_vars
+    num_samples = num_samples_opt[0] if num_samples_opt else "10"
+    test_avg_tflops(metrics_file, min_avg_tflops, num_samples)
   else:
-    raise ValueError(f"Unrecognized test_scenario {test_scenario}")
+    raise ValueError(f"Unrecognized test_scenario '{test_scenario}'. Available: final_loss, avg_step_time, avg_tflops")
 
 
 if __name__ == "__main__":
-  app.run(main)
+  app.run(main)

requirements.txt

@@ -26,7 +26,7 @@ git+https://github.com/mlperf/logging.git
 opencv-python-headless==4.10.0.84
 orbax-checkpoint==0.10.3
 tokenizers==0.21.0
-huggingface_hub==0.24.7
+huggingface_hub==0.30.2
 transformers==4.48.1
 einops==0.8.0
 sentencepiece

requirements_with_jax_stable_stack.txt

@@ -8,7 +8,7 @@ ftfy
 git+https://github.com/mlperf/logging.git
 google-cloud-storage==2.17.0
 grain-nightly==0.0.10
-huggingface_hub==0.24.7
+huggingface_hub==0.30.2
 jax>=0.4.30
 jaxlib>=0.4.30
 Jinja2

src/maxdiffusion/configs/base14.yml

@@ -19,6 +19,10 @@ metrics_file: "" # for testing, local file that stores scalar metrics. If empty,
 # If true save metrics such as loss and TFLOPS to GCS in {base_output_directory}/{run_name}/metrics/
 write_metrics: True
 gcs_metrics: True
+
+timing_metrics_file: "" # for testing, local file that stores function timing metrics such as state creation, compilation. If empty, no metrics are written.
+write_timing_metrics: True 
+
 # If true save config to GCS in {base_output_directory}/{run_name}/
 save_config_to_gcs: False
 log_period: 10000000000  # Flushes Tensorboard

src/maxdiffusion/configs/base21.yml

@@ -19,6 +19,10 @@ metrics_file: "" # for testing, local file that stores scalar metrics. If empty,
 # If true save metrics such as loss and TFLOPS to GCS in {base_output_directory}/{run_name}/metrics/
 write_metrics: True
 gcs_metrics: True
+
+timing_metrics_file: "" # for testing, local file that stores function timing metrics such as state creation, compilation. If empty, no metrics are written.
+write_timing_metrics: True 
+
 # If true save config to GCS in {base_output_directory}/{run_name}/
 save_config_to_gcs: False
 log_period: 10000000000  # Flushes Tensorboard

src/maxdiffusion/configs/base_2_base.yml

@@ -19,6 +19,10 @@ metrics_file: "" # for testing, local file that stores scalar metrics. If empty,
 # If true save metrics such as loss and TFLOPS to GCS in {base_output_directory}/{run_name}/metrics/
 write_metrics: True
 gcs_metrics: False
+
+timing_metrics_file: "" # for testing, local file that stores function timing metrics such as state creation, compilation. If empty, no metrics are written.
+write_timing_metrics: True 
+
 # If true save config to GCS in {base_output_directory}/{run_name}/
 save_config_to_gcs: False
 log_period: 10000000000  # Flushes Tensorboard

src/maxdiffusion/configs/base_flux_dev.yml

@@ -18,6 +18,10 @@ run_name: ''
 metrics_file: "" # for testing, local file that stores scalar metrics. If empty, no metrics are written.
 # If true save metrics such as loss and TFLOPS to GCS in {base_output_directory}/{run_name}/metrics/
 write_metrics: True
+
+timing_metrics_file: "" # for testing, local file that stores function timing metrics such as state creation, compilation. If empty, no metrics are written.
+write_timing_metrics: True 
+
 gcs_metrics: False
 # If true save config to GCS in {base_output_directory}/{run_name}/
 save_config_to_gcs: False

src/maxdiffusion/configs/base_flux_dev_multi_res.yml

@@ -19,6 +19,10 @@ metrics_file: "" # for testing, local file that stores scalar metrics. If empty,
 # If true save metrics such as loss and TFLOPS to GCS in {base_output_directory}/{run_name}/metrics/
 write_metrics: True
 gcs_metrics: False
+
+timing_metrics_file: "" # for testing, local file that stores function timing metrics such as state creation, compilation. If empty, no metrics are written.
+write_timing_metrics: True 
+
 # If true save config to GCS in {base_output_directory}/{run_name}/
 save_config_to_gcs: False
 log_period: 100

src/maxdiffusion/configs/base_flux_schnell.yml

@@ -19,6 +19,10 @@ metrics_file: "" # for testing, local file that stores scalar metrics. If empty,
 # If true save metrics such as loss and TFLOPS to GCS in {base_output_directory}/{run_name}/metrics/
 write_metrics: True
 gcs_metrics: False
+
+timing_metrics_file: "" # for testing, local file that stores function timing metrics such as state creation, compilation. If empty, no metrics are written.
+write_timing_metrics: True 
+
 # If true save config to GCS in {base_output_directory}/{run_name}/
 save_config_to_gcs: False
 log_period: 100

src/maxdiffusion/configs/base_xl.yml

@@ -18,6 +18,10 @@ run_name: ''
 metrics_file: "" # for testing, local file that stores scalar metrics. If empty, no metrics are written.
 # If true save metrics such as loss and TFLOPS to GCS in {base_output_directory}/{run_name}/metrics/
 write_metrics: True
+
+timing_metrics_file: "" # for testing, local file that stores function timing metrics such as state creation, compilation. If empty, no metrics are written.
+write_timing_metrics: True 
+
 gcs_metrics: False
 # If true save config to GCS in {base_output_directory}/{run_name}/
 save_config_to_gcs: False