Sweep across KV cache layouts

MaxText/inference_microbenchmark.py

@@ -66,9 +66,9 @@ def prefill_benchmark(
       f"\tPrefill TFLOPs/sec/device: {tflops_per_sec_per_device:.3f}\n\n\n\n"
   )
   result_dict = {
-      "prefill_time_in_ms": prefill_average_ms,
-      "prefill_total_tflops_per_device": prefill_tflops_per_device,
-      "prefill_tflops_per_sec_per_device": tflops_per_sec_per_device,
+      "time_in_ms": prefill_average_ms,
+      "total_tflops_per_device": prefill_tflops_per_device,
+      "tflops_per_sec_per_device": tflops_per_sec_per_device,
   }
   return result_dict
 
@@ -109,7 +109,7 @@ def prefill_insert_benchmark(
       f"\tPrefill + Insert step average time: {prefill_insert_average_ms:.3f} ms\n\n\n\n"
   )
   result_dict = {
-      "prefill_insert_time_in_ms": prefill_insert_average_ms
+      "insert_time_in_ms": prefill_insert_average_ms
   }
   return result_dict, decode_state
 
@@ -150,20 +150,20 @@ def ar_benchmark(config, engine, params, decode_state, global_batch_size, cache_
   )
 
   result_dict = {
-      "ar_step_in_ms": ar_average_ms,
-      "ar_step_in_ms_per_seq": ar_average_ms / global_batch_size,
-      "ar_global_batch_size": global_batch_size,
-      "ar_total_throughput_tokens_per_second": total_throughput,
-      "ar_device_bandwidth_GB_per_second": bw_per_device,
+      "step_in_ms": ar_average_ms,
+      "step_in_ms_per_seq": ar_average_ms / global_batch_size,
+      "global_batch_size": global_batch_size,
+      "total_throughput_tokens_per_second": total_throughput,
+      "device_bandwidth_GB_per_second": bw_per_device,
   }
   return result_dict, decode_state
 
 
 def collate_results(config, results, model_size, cache_size, num_model_params, incl_config=False):
   """Adds model/cache size info and optionally config info to results."""
   results["sizes"] = {
-      "Model_size_in_GB": model_size / 1e9,
-      "cache_size_in_GB": cache_size / 1e9,
+      "model_size_in_gb": model_size / 1e9,
+      "cache_size_in_gb": cache_size / 1e9,
       "model_params_in_billions": num_model_params / 1e9,
   }
   if incl_config:
@@ -198,20 +198,20 @@ def print_results_for_analyze(results):
   """Print results."""
   print("\nFor usage in analyze_sharegpt.py :")
 
-  if "Prefill" in results:
+  if "prefill" in results:
     prefill_bucket_size_to_ms = {}
-    for k, v in results["Prefill"].items():
-      prefill_bucket_size_to_ms[int(k)] = round(v["prefill_time_in_ms"], 3)
+    for k, v in results["prefill"].items():
+      prefill_bucket_size_to_ms[int(k)] = round(v["time_in_ms"], 3)
     print(f"PREFILL_BUCKET_SIZE_TO_MS = {prefill_bucket_size_to_ms}")
 
-  if "Prefill_Insert" in results:
+  if "prefill-insert" in results:
     insert_bucket_size_to_ms = {}
-    for k, v in results["Prefill_Insert"].items():
-      insert_bucket_size_to_ms[int(k)] = round(v["prefill_insert_time_in_ms"], 3)
+    for k, v in results["prefill-insert"].items():
+      insert_bucket_size_to_ms[int(k)] = round(v["insert_time_in_ms"], 3)
     print(f"PREFILL_INSERT_BUCKET_SIZE_TO_MS = {insert_bucket_size_to_ms}")
 
-  if "AutoRegressive" in results:
-    print(f"SYSTEM_TIME_PER_DECODE_TOKEN_MS = {results['AutoRegressive']['ar_step_in_ms_per_seq']}")
+  if "autoregressive" in results:
+    print(f"SYSTEM_TIME_PER_DECODE_TOKEN_MS = {results['autoregressive']['step_in_ms_per_seq']}")
 
 
 def summarize_prefill_result(engine, params, tokens, true_length):
@@ -227,12 +227,12 @@ def summarize_prefill_result(engine, params, tokens, true_length):
   )
   max_utils.delete_pytree(prefill_result)
   return {
-    "num_prefill_logits_params": num_prefill_logits_params,
-    "total_prefill_logits_size": total_prefill_logits_size,
-    "avg_prefill_logits_param_size": avg_prefill_logits_param_size,
-    "num_prefill_cache_params": num_prefill_cache_params,
-    "total_prefill_cache_size": total_prefill_cache_size,
-    "avg_prefill_cache_param_size": avg_prefill_cache_param_size,
+    "num_logits_params": num_prefill_logits_params,
+    "total_logits_size": total_prefill_logits_size,
+    "avg_logits_param_size": avg_prefill_logits_param_size,
+    "num_cache_params": num_prefill_cache_params,
+    "total_cache_size": total_prefill_cache_size,
+    "avg_cache_param_size": avg_prefill_cache_param_size,
   }
 
 
@@ -254,22 +254,22 @@ def main(config, inference_metadata: Optional[Dict[str, Any]] = None):
   benchmark_results = {}
   if "prefill" in stages_to_benchmark:
 
-    benchmark_results["Prefill_Result"] = {}
-    benchmark_results["Prefill"] = {}
-    benchmark_results["Prefill_Insert"] = {}
+    benchmark_results["prefill-result-sizes"] = {}
+    benchmark_results["prefill"] = {}
+    benchmark_results["prefill-insert"] = {}
     prefill_tokens = {}
     prefill_true_lengths = {}
 
     for prefill_length in prefill_lengths:
       prefill_tokens[prefill_length], prefill_true_lengths[prefill_length] = token_utils.tokenize_and_pad(
         text, vocab, is_bos=True, prefill_lengths=[prefill_length]
       )
-      benchmark_results["Prefill_Result"]["prefill_length"] = summarize_prefill_result(
+      benchmark_results["prefill-result-sizes"][prefill_length] = summarize_prefill_result(
         engine, params, prefill_tokens[prefill_length], prefill_true_lengths[prefill_length]
       )
 
     for prefill_length in prefill_lengths:
-      benchmark_results["Prefill"][prefill_length] = prefill_benchmark(
+      benchmark_results["prefill"][prefill_length] = prefill_benchmark(
         config,
         engine,
         params,
@@ -279,7 +279,7 @@ def main(config, inference_metadata: Optional[Dict[str, Any]] = None):
         benchmark_loop_iters
       )
 
-      benchmark_results["Prefill_Insert"][prefill_length], decode_state = prefill_insert_benchmark(
+      benchmark_results["prefill-insert"][prefill_length], decode_state = prefill_insert_benchmark(
         config,
         engine,
         decode_state,
@@ -291,7 +291,7 @@ def main(config, inference_metadata: Optional[Dict[str, Any]] = None):
       )
 
   if "generate" in stages_to_benchmark:
-    benchmark_results["AutoRegressive"], decode_state = ar_benchmark(
+    benchmark_results["autoregressive"], decode_state = ar_benchmark(
       config, engine, params, decode_state, engine.max_concurrent_decodes, cache_size, model_size, benchmark_loop_iters)
 
   results = collate_results(config, benchmark_results, model_size, cache_size, num_model_params)

MaxText/inference_microbenchmark_sweep.py

@@ -43,6 +43,7 @@ def main():
   """
   pyconfig.initialize(sys.argv)
   config = pyconfig.config
+  base_run_name = config.run_name
 
   with open(config.inference_metadata_file, encoding='utf-8') as json_file:
     inference_metadata = json.load(json_file)
@@ -54,6 +55,9 @@ def main():
   ar_key_axis_order_list = inference_metadata['ar_key_axis_order_list'].split(':')
   ar_value_axis_order_list = inference_metadata['ar_value_axis_order_list'].split(':')
 
+  start_key_value_axis_order_product_id = key_value_axis_order_product_id_list[0]
+  end_key_value_axis_order_product_id = key_value_axis_order_product_id_list[-1]
+
   results = []
   for (
     key_value_axis_order_product_id,
@@ -74,25 +78,17 @@ def main():
     print(f"ar_key_axis_order {ar_key_axis_order}")
     print(f"ar_value_axis_order {ar_value_axis_order}")
 
-    # Manually update the config
-    # Don't set key_value_axis_order_product_id; otherwise it will recompute
-    # ar_key_axis_order and ar_value_axis_order
-    quant = 'bf16' if not config.quantization else config.quantization
-    run_name = (
-      f"{inference_metadata['accelerator']}-{config.model_name}-"
-      f"{quant}-{key_value_axis_order_product_id}-{prefill_key_axis_order}-"
-      f"{ar_key_axis_order}"
+    run_tag = (
+      f"{key_value_axis_order_product_id}-{prefill_key_axis_order.replace(',','')}-{ar_key_axis_order.replace(',','')}"
     )
+    run_name = f"{base_run_name}/{run_tag}"
+
     tensorboard_dir = os.path.join(config.base_output_directory, run_name, "tensorboard", "")
-    checkpoint_dir = os.path.join(config.base_output_directory, run_name, "checkpoint", "")
-    metrics_dir = os.path.join(config.base_output_directory, run_name, "metrics", "")
     pyconfig._config.keys['prefill_key_axis_order'] = prefill_key_axis_order # pylint: disable=protected-access
     pyconfig._config.keys['prefill_value_axis_order'] = prefill_value_axis_order # pylint: disable=protected-access
     pyconfig._config.keys['ar_key_axis_order'] = ar_key_axis_order # pylint: disable=protected-access
     pyconfig._config.keys['ar_value_axis_order'] = ar_value_axis_order # pylint: disable=protected-access
     pyconfig._config.keys['tensorboard_dir'] = tensorboard_dir # pylint: disable=protected-access
-    pyconfig._config.keys['checkpoint_dir'] = checkpoint_dir # pylint: disable=protected-access
-    pyconfig._config.keys['metrics_dir'] = metrics_dir # pylint: disable=protected-access
     pyconfig._config.keys['run_name'] = run_name # pylint: disable=protected-access
     max_utils.write_config_raw_keys_for_gcs(pyconfig._config.keys) # pylint: disable=protected-access
 
@@ -121,6 +117,8 @@ def main():
       "prefill_value_axis_order": f"{prefill_value_axis_order}",
       "ar_key_axis_order": f"{ar_key_axis_order}",
       "ar_value_axis_order": f"{ar_value_axis_order}",
+      "run_name": f"{run_name}",
+      "run_tag": f"{run_tag}",
       "config_json_string": json.dumps(
           pyconfig._config.keys, # pylint: disable=protected-access
           default=lambda x: f"<<non-serializable: {type(x).__qualname__}>>"
@@ -135,10 +133,14 @@ def main():
       metrics = microbenchmark_results['flattened_results']
       metrics = {k.lower(): v for k, v in metrics.items()}
       dimensions_json['oom'] = 'False'
+      print(f"Completed run {key_value_axis_order_product_id} out of: "
+            f"{start_key_value_axis_order_product_id} to {end_key_value_axis_order_product_id}")
     except xla_extension.XlaRuntimeError:
       # OOM
       metrics = {}
       dimensions_json['oom'] = 'True'
+      print(f"Failed at run {key_value_axis_order_product_id} out of: "
+            f"{start_key_value_axis_order_product_id} to {end_key_value_axis_order_product_id}")
 
     final = {'metrics': metrics, 'dimensions': dimensions_json}
     print(f"Result: {final}")
@@ -151,4 +153,4 @@ def main():
 
 
 if __name__ == "__main__":
-  main()
+  main()