Add Warmup Loop for JIT ops (#30)

* Fix num_runs loop for DCN collective tests * Adding warmup logic * Added warmp_up param as example to sample config

Author: Perseus14

configs/sample_benchmark_collectives.yaml

@@ -2,23 +2,27 @@ benchmarks:
 - benchmark_name: ppermute
  benchmark_sweep_params:
  - {matrix_dim_range: {start: 1024, end: 20000, increase_by: 1024}, dtype: "bfloat16", dcn_size_range: 1, ici_size_range: 4}
+ warmup_tries: 10
  trace_dir: "/tmp/microbenchmarks/collectives"
  csv_path: "/tmp/microbenchmarks/collectives"
  xla_dump_dir: "/tmp/microbenchmarks/collective/hlo_graphs"
 - benchmark_name: all_gather
  benchmark_sweep_params:
  - {matrix_dim_range: {start: 1024, end: 20000, increase_by: 1024}, dtype: "bfloat16", dcn_size_range: 1, ici_size_range: 4}
+ warmup_tries: 10
  trace_dir: "/tmp/microbenchmarks/collectives"
  csv_path: "/tmp/microbenchmarks/collectives"
  xla_dump_dir: "/tmp/microbenchmarks/collective/hlo_graphs"
 - benchmark_name: psum
  benchmark_sweep_params:
  - {matrix_dim_range: {start: 1024, end: 20000, increase_by: 1024}, dtype: "bfloat16", dcn_size_range: 1, ici_size_range: 4}
+ warmup_tries: 10
  trace_dir: "/tmp/microbenchmarks/collectives"
  csv_path: "/tmp/microbenchmarks/collectives"
  xla_dump_dir: "/tmp/microbenchmarks/collective/hlo_graphs"
 - benchmark_name: psum_scatter
  benchmark_sweep_params:
  - {matrix_dim_range: {start: 1024, end: 20000, increase_by: 1024}, dtype: "bfloat16", dcn_size_range: 1, ici_size_range: 4}
+ warmup_tries: 10
  trace_dir: "/tmp/microbenchmarks/collectives"
  csv_path: "/tmp/microbenchmarks/collectives"

src/benchmark_collectives.py

@@ -57,6 +57,7 @@ def psum_benchmark(
  ici_size: int,
  num_runs: int = 1,
  trace_dir: str = None,
+ warmup_tries: int = 10,
 ) -> Dict[str, Any]:
  """Benchmarks the psum collective operation.
 
@@ -90,6 +91,7 @@ def f(x):
  jitted_op,
  sharded_matrix,
  matrix_dim=matrix_dim,
+ warmup_tries=warmup_tries,
  tries=num_runs,
  task="psum_dcn_op",
  trace_dir=trace_dir,
@@ -110,6 +112,7 @@ def f(x):
  jitted_op,
  sharded_matrix,
  matrix_dim=matrix_dim,
+ warmup_tries=warmup_tries,
  tries=num_runs,
  task="psum_ici_op",
  trace_dir=trace_dir,
@@ -189,6 +192,7 @@ def psum_scatter_benchmark(
  ici_size: int,
  num_runs: int = 1,
  trace_dir: str = None,
+ warmup_tries: int = 10,
 ) -> Dict[str, Any]:
  """Benchmarks the psum_scatter collective operation.
 
@@ -224,6 +228,7 @@ def f(x):
  jitted_op,
  sharded_matrix,
  matrix_dim=matrix_dim,
+ warmup_tries=warmup_tries,
  tries=num_runs,
  task="psum_scatter_dcn_op",
  trace_dir=trace_dir,
@@ -244,6 +249,7 @@ def f(x):
  jitted_op,
  sharded_matrix,
  matrix_dim=matrix_dim,
+ warmup_tries=warmup_tries,
  tries=num_runs,
  task="psum_scatter_ici_op",
  trace_dir=trace_dir,
@@ -322,6 +328,7 @@ def all_gather_benchmark(
  dtype: jnp.dtype,
  dcn_size: int,
  ici_size: int,
+ warmup_tries: int = 10,
  num_runs: int = 1,
  trace_dir: str = None,
 ) -> Dict[str, Any]:
@@ -364,6 +371,7 @@ def f(x):
  jitted_op,
  sharded_matrix,
  matrix_dim=matrix_dim,
+ warmup_tries=warmup_tries,
  tries=num_runs,
  task="all_gather_dcn_op",
  trace_dir=trace_dir,
@@ -390,6 +398,7 @@ def f(x):
  jitted_op,
  sharded_matrix,
  matrix_dim=matrix_dim,
+ warmup_tries=warmup_tries,
  tries=num_runs,
  task="all_gather_ici_op",
  trace_dir=trace_dir,
@@ -469,6 +478,7 @@ def ppermute_benchmark(
  ici_size: int,
  num_runs: int = 1,
  trace_dir: str = None,
+ warmup_tries: int = 10,
 ) -> Dict[str, Any]:
  """Benchmarks the ppermute collective operation.
 
@@ -506,6 +516,7 @@ def f(x):
  jitted_op,
  sharded_matrix,
  matrix_dim=matrix_dim,
+ warmup_tries=warmup_tries,
  tries=num_runs,
  task="ppermute_dcn_op",
  trace_dir=trace_dir,
@@ -527,6 +538,7 @@ def f(x):
  jitted_op,
  sharded_matrix,
  matrix_dim=matrix_dim,
+ warmup_tries=warmup_tries,
  tries=num_runs,
  task="ppermute_ici_op",
  trace_dir=trace_dir,
@@ -598,6 +610,7 @@ def all_to_all_benchmark(
  ici_size: int,
  num_runs: int = 1,
  trace_dir: str = None,
+ warmup_tries: int = 10,
 ) -> Dict[str, Any]:
  """Benchmarks the all_to_all collective operation.
 
@@ -634,6 +647,7 @@ def f(x):
  jitted_op,
  sharded_matrix,
  matrix_dim=matrix_dim,
+ warmup_tries=warmup_tries,
  tries=num_runs,
  task="all_to_all_dcn_op",
  trace_dir=trace_dir,
@@ -660,6 +674,7 @@ def f(x):
  jitted_op,
  sharded_matrix,
  matrix_dim=matrix_dim,
+ warmup_tries=warmup_tries,
  tries=num_runs,
  task="all_to_all_ici_op",
  trace_dir=trace_dir,

src/benchmark_utils.py

@@ -19,15 +19,19 @@
 import shutil
 
 
-def simple_timeit(f, *args, matrix_dim=None, tries=10, task=None, trace_dir=None) -> float:
+def simple_timeit(f, *args, matrix_dim=None, warmup_tries = 10, tries=10, task=None, trace_dir=None) -> float:
  """Simple utility to time a function for multiple runs."""
  assert task is not None
 
  if trace_dir:
  return timeit_from_trace(f, *args, matrix_dim=matrix_dim, tries=tries, task=task, trace_dir=trace_dir)
 
+ # warmup loop
+ print(f"Running warmup loop with {warmup_tries} tries...")
+ for _ in range(warmup_tries):
+ data = f(*args)
+ jax.block_until_ready(data)
  outcomes_ms = []
- jax.block_until_ready(f(*args)) # warm it up!
  for _ in range(tries):
  jax.devices() # Force synchronization across devices
  s = datetime.datetime.now()
@@ -97,13 +101,17 @@ def is_local_directory_path(dir: str) -> bool:
  return dir.startswith("/") or dir.startswith("./") or dir.startswith("../")
 
 
-def timeit_from_trace(f, *args, matrix_dim=None, tries=10, task=None, trace_dir=None) -> float:
+def timeit_from_trace(f, *args, matrix_dim=None, warmup_tries=10, tries=10, task=None, trace_dir=None) -> float:
  """
  Time a function with jax.profiler and get the run time from the trace.
  """
  LOCAL_TRACE_DIR = "/tmp/microbenchmarks_tmptrace"
 
- jax.block_until_ready(f(*args)) # warm it up!
+ # warmup loop
+ print(f"Running warmup loop with {warmup_tries} tries...")
+ for _ in range(warmup_tries):
+ data = f(*args)
+ jax.block_until_ready(data)
 
  if matrix_dim is not None:
  trace_name = f"{task}_dim_{matrix_dim}"

src/run_benchmark.py

@@ -279,6 +279,8 @@ def run_single_benchmark(benchmark_config: Dict[str, Any]):
  trace_dir = benchmark_config.get("trace_dir")
  xlml_metrics_dir = benchmark_config.get("xlml_metrics_dir")
  xla_dump_dir = benchmark_config.get("xla_dump_dir")
+ warmup_tries = benchmark_config.get("warmup_tries")
+ warmup_tries = warmup_tries if warmup_tries is not None else 1000
 
  if not benchmark_name:
  raise ValueError("Each benchmark must have a 'benchmark_name'.")
@@ -299,7 +301,7 @@ def run_single_benchmark(benchmark_config: Dict[str, Any]):
  test_start_time = (
  datetime.datetime.now(tz=datetime.timezone.utc).isoformat() + "Z"
  ) # "Z" indicates UTC
- benchmark_results = benchmark_func(**benchmark_param)
+ benchmark_results = benchmark_func(**benchmark_param, warmup_tries=warmup_tries)
  test_end_time = (
  datetime.datetime.now(tz=datetime.timezone.utc).isoformat() + "Z"
  )
@@ -400,6 +402,8 @@ def run_benchmark_multithreaded(benchmark_config):
  csv_path = benchmark_config.get("csv_path")
  if not benchmark_name:
  raise ValueError("Each benchmark must have a 'benchmark_name'.")
+ warmup_tries = benchmark_config.get("warmup_tries")
+ warmup_tries = warmup_tries if warmup_tries is not None else 1000
 
  # Get the benchmark function
  benchmark_func, calculate_metrics_func = get_benchmark_functions(benchmark_name)
@@ -427,7 +431,7 @@ def run_benchmark_multithreaded(benchmark_config):
  with ThreadPoolExecutor(max_workers=num_hosts) as executor:
  # Create a mapping of futures to their corresponding parameters
  future_to_param = {
- executor.submit(benchmark_func, **benchmark_param): benchmark_param
+ executor.submit(benchmark_func, **benchmark_param, warmup_tries=warmup_tries): benchmark_param
  for benchmark_param in preprocessed_benchmark_params
  }