From 2c3d0ebad85c98f38e40bd52b7d7bb430354fbb6 Mon Sep 17 00:00:00 2001 From: Simran Kaur Date: Thu, 8 Oct 2026 09:38:16 +0000 Subject: [PATCH 1/9] feat(chs): support single-host TPU diagnostics on multi-host slices --- Ironwood/src/benchmark_gemm.py | 55 +++++++++++++++++++++------ Ironwood/src/benchmark_hbm.py | 6 ++- Ironwood/src/benchmark_host_device.py | 2 + Ironwood/src/benchmark_utils.py | 47 +++++++++++++++++++++++ Ironwood/src/run_benchmark.py | 9 +++++ requirements.txt | 6 +-- 6 files changed, 109 insertions(+), 16 deletions(-) diff --git a/Ironwood/src/benchmark_gemm.py b/Ironwood/src/benchmark_gemm.py index f7fbe91d..15deb76a 100644 --- a/Ironwood/src/benchmark_gemm.py +++ b/Ironwood/src/benchmark_gemm.py @@ -102,22 +102,53 @@ def f(x, y): lhs_dtype = dtype rhs_dtype = dtype - key = jax.random.key(SEED) + if run_on_local_node: + key = jax.device_put(jax.random.key(SEED), jax.devices("cpu")[0]) + mesh_devices = list(mesh.devices.flat) + + def data_generator(): + """Generates random data independently on each local device.""" + nonlocal key + key, key_lhs, key_rhs = jax.random.split(key, 3) + lhs_device = jax.make_array_from_single_device_arrays( + lhs_shape, + lhs_sharding, + [ + jax.random.normal( + jax.device_put(key_lhs, d), lhs_shape, dtype=lhs_dtype + ) + for d in mesh_devices + ], + ) + rhs_device = jax.make_array_from_single_device_arrays( + rhs_shape, + rhs_sharding, + [ + jax.random.normal( + jax.device_put(key_rhs, d), rhs_shape, dtype=rhs_dtype + ) + for d in mesh_devices + ], + ) + return (lhs_device, rhs_device) - def data_generator(): - """Creates new random data on host and puts it on device.""" - nonlocal key # Use and update the outer 'key' - key, key_lhs, key_rhs = jax.random.split(key, 3) + else: + key = jax.random.key(SEED) - # Create random data on host - lhs_host = jax.random.normal(key_lhs, lhs_shape).astype(lhs_dtype) - rhs_host = jax.random.normal(key_rhs, rhs_shape).astype(rhs_dtype) + def data_generator(): + """Creates new random data on host and puts it on device.""" + nonlocal key # Use and update the outer 'key' + key, key_lhs, key_rhs = jax.random.split(key, 3) - # Put on device (HBM) - lhs_device = jax.device_put(lhs_host, lhs_sharding) - rhs_device = jax.device_put(rhs_host, rhs_sharding) + # Create random data on host + lhs_host = jax.random.normal(key_lhs, lhs_shape).astype(lhs_dtype) + rhs_host = jax.random.normal(key_rhs, rhs_shape).astype(rhs_dtype) - return (lhs_device, rhs_device) + # Put on device (HBM) + lhs_device = jax.device_put(lhs_host, lhs_sharding) + rhs_device = jax.device_put(rhs_host, rhs_sharding) + + return (lhs_device, rhs_device) # Run the benchmark diff --git a/Ironwood/src/benchmark_hbm.py b/Ironwood/src/benchmark_hbm.py index f72f7684..7f9aae51 100644 --- a/Ironwood/src/benchmark_hbm.py +++ b/Ironwood/src/benchmark_hbm.py @@ -38,6 +38,7 @@ def single_device_hbm_copy( dtype: jnp.dtype, num_runs: int = 1, trace_dir: str = None, + run_on_local_node: bool = False, # pylint: disable=unused-argument ) -> Dict[str, Any]: """Benchmarks HBM with copy(read and write) on a single device.""" @@ -67,7 +68,10 @@ def f(a): def single_device_hbm_copy_calculate_metrics( - num_elements: int, dtype: jnp.dtype, time_ms_list: list + num_elements: int, + dtype: jnp.dtype, + time_ms_list: list, + run_on_local_node: bool = False, # pylint: disable=unused-argument ) -> Dict[str, Any]: """Calculates the metrics for the single device hbm copy benchmark.""" # Build dictionary of all the parameters in the function diff --git a/Ironwood/src/benchmark_host_device.py b/Ironwood/src/benchmark_host_device.py index 6d47dab2..3632063c 100644 --- a/Ironwood/src/benchmark_host_device.py +++ b/Ironwood/src/benchmark_host_device.py @@ -25,6 +25,7 @@ def benchmark_host_device( data_size_mib: int, num_runs: int = 100, trace_dir: str = None, + run_on_local_node: bool = False, # pylint: disable=unused-argument ) -> Dict[str, Any]: """Benchmarks H2D/D2H transfer using simple device_put/device_get.""" @@ -105,6 +106,7 @@ def benchmark_host_device_calculate_metrics( data_size_mib: int, h2d_bandwidth_ms: List[float], d2h_bandwidth_ms: List[float], + run_on_local_node: bool = False, # pylint: disable=unused-argument ) -> Tuple[Dict[str, Any], Dict[str, Any]]: """Calculates metrics for Host-Device transfer.""" params = locals().items() diff --git a/Ironwood/src/benchmark_utils.py b/Ironwood/src/benchmark_utils.py index 22a5356a..9546a4ec 100644 --- a/Ironwood/src/benchmark_utils.py +++ b/Ironwood/src/benchmark_utils.py @@ -29,6 +29,36 @@ from tensorflow.tsl.profiler.protobuf import xplane_pb2 +def configure_local_node_tpu(): + """Configures local node TPU execution on a single host without slicebuilder.""" + os.environ["TPU_HOST_BOUNDS"] = "1,1,1" + os.environ["TPU_SKIP_MDS_QUERY"] = "true" + os.environ["TPU_WORKER_ID"] = "0" + os.environ["TPU_WORKER_HOSTNAMES"] = "localhost" + local_flags = [ + "--deepsea_hal_test_skip_slicebuilder=true", + "--deepsea_hal_test_allow_multichip_skip_slicebuilder=true", + "--xla_tpu_enable_async_collective_fusion=false", + "--xla_tpu_use_enhanced_launch_barrier=false", + ] + existing_args = ( + os.environ.get("LIBTPU_INIT_ARGS", "") + .replace( + "--xla_tpu_enable_async_collective_fusion=true", + "--xla_tpu_enable_async_collective_fusion=false", + ) + .replace( + "--xla_tpu_use_enhanced_launch_barrier=true", + "--xla_tpu_use_enhanced_launch_barrier=false", + ) + ) + for flag in local_flags: + flag_name = flag.split("=")[0] + if flag_name not in existing_args: + existing_args = f"{existing_args} {flag}".strip() + os.environ["LIBTPU_INIT_ARGS"] = existing_args + + def get_real_dtype_bytes(dtype) -> float: """Returns the real byte size of a dtype, handling sub-byte types.""" try: @@ -1024,6 +1054,9 @@ def rename_xla_dump( ) continue + if os.path.isdir(original_filepath): + continue + # Copy the renamed files to desired location if is_local_directory_path(dest_xla_dump_dir): try: @@ -1108,6 +1141,9 @@ def extract_hlo_features_from_file( rg_match = re.search( r"replica_groups=({{[0-9,]+(?:},{[0-9,]+)*}})", content, re.DOTALL ) + mesh_rg_match = re.search( + r"replica_groups=(mesh\[([^\]]+)\]\s*\{([^}]+)\})", content + ) if rg_match: replica_groups_str = rg_match.group(1) try: @@ -1117,6 +1153,17 @@ def extract_hlo_features_from_file( except ValueError as e: print(f"Could not parse replica_groups in hlo_text: {e}") first_replica_group = None + elif mesh_rg_match: + replica_groups_str = mesh_rg_match.group(1) + dims = { + k: int(v) + for k, v in re.findall(r"'([^']+)'=(\d+)", mesh_rg_match.group(2)) + } + axes = re.findall(r"'([^']+)'", mesh_rg_match.group(3)) + group_size = 1 + for axis in axes: + group_size *= dims.get(axis, 1) + first_replica_group = list(range(group_size)) else: print(f"Could not find replica_groups in {hlo_file_path}.") diff --git a/Ironwood/src/run_benchmark.py b/Ironwood/src/run_benchmark.py index 2dd12b75..b5c3d1ff 100644 --- a/Ironwood/src/run_benchmark.py +++ b/Ironwood/src/run_benchmark.py @@ -13,6 +13,7 @@ import string from typing import Any, Callable, Dict, List, Tuple from benchmark_utils import ( + configure_local_node_tpu, maybe_write_metrics_file, rename_xla_dump, MetricsStatistics, @@ -375,6 +376,12 @@ def run_single_benchmark(benchmark_config: Dict[str, Any], output_path: str): for param in benchmark_params: if "topology" not in param: param["topology"] = global_topology + # Inject run_on_local_node from config if present in benchmark_config + global_run_on_local_node = benchmark_config.get("run_on_local_node") + if global_run_on_local_node is not None: + for param in benchmark_params: + if "run_on_local_node" not in param: + param["run_on_local_node"] = global_run_on_local_node if not benchmark_name: raise ValueError("Each benchmark must have a benchmark_name.") @@ -384,6 +391,8 @@ def run_single_benchmark(benchmark_config: Dict[str, Any], output_path: str): benchmark_func, calculate_metrics_func = get_benchmark_functions( benchmark_name ) + if benchmark_config.get("run_on_local_node", False): + configure_local_node_tpu() print(f"\n{'=' * 30}Starting benchmark '{benchmark_name}'{'=' * 30}\n") diff --git a/requirements.txt b/requirements.txt index b37bec99..77049aa1 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,8 +1,8 @@ -jax[tpu]==0.10.1 +jax[tpu]==0.11.2 absl-py flax==0.12.7 -jaxlib==0.10.1 -libtpu==0.0.41 +jaxlib==0.11.2 +libtpu==0.0.48 numpy pandas jsonlines From 40b1cd326fbc07dd24553c4027b53daa9fcb77eb Mon Sep 17 00:00:00 2001 From: Simran Kaur Date: Thu, 8 Oct 2026 09:44:45 +0000 Subject: [PATCH 2/9] refactor(chs): gate libtpu init on separate benchmark-level single_host_tpu_init flag - Add 'single_host_tpu_init' as a benchmark-level config key that triggers configure_single_host_tpu() (renamed from configure_local_node_tpu). It acts at the libtpu init layer, once per process. - 'run_on_local_node' stays a sweep-param-only knob (JAX device selection) and is no longer injected from benchmark_config. - Revert the unused run_on_local_node kwargs in benchmark_hbm.py and benchmark_host_device.py since params are no longer injected. --- Ironwood/src/benchmark_hbm.py | 6 +----- Ironwood/src/benchmark_host_device.py | 2 -- Ironwood/src/benchmark_utils.py | 10 ++++++++-- Ironwood/src/run_benchmark.py | 16 +++++++--------- 4 files changed, 16 insertions(+), 18 deletions(-) diff --git a/Ironwood/src/benchmark_hbm.py b/Ironwood/src/benchmark_hbm.py index 7f9aae51..f72f7684 100644 --- a/Ironwood/src/benchmark_hbm.py +++ b/Ironwood/src/benchmark_hbm.py @@ -38,7 +38,6 @@ def single_device_hbm_copy( dtype: jnp.dtype, num_runs: int = 1, trace_dir: str = None, - run_on_local_node: bool = False, # pylint: disable=unused-argument ) -> Dict[str, Any]: """Benchmarks HBM with copy(read and write) on a single device.""" @@ -68,10 +67,7 @@ def f(a): def single_device_hbm_copy_calculate_metrics( - num_elements: int, - dtype: jnp.dtype, - time_ms_list: list, - run_on_local_node: bool = False, # pylint: disable=unused-argument + num_elements: int, dtype: jnp.dtype, time_ms_list: list ) -> Dict[str, Any]: """Calculates the metrics for the single device hbm copy benchmark.""" # Build dictionary of all the parameters in the function diff --git a/Ironwood/src/benchmark_host_device.py b/Ironwood/src/benchmark_host_device.py index 3632063c..6d47dab2 100644 --- a/Ironwood/src/benchmark_host_device.py +++ b/Ironwood/src/benchmark_host_device.py @@ -25,7 +25,6 @@ def benchmark_host_device( data_size_mib: int, num_runs: int = 100, trace_dir: str = None, - run_on_local_node: bool = False, # pylint: disable=unused-argument ) -> Dict[str, Any]: """Benchmarks H2D/D2H transfer using simple device_put/device_get.""" @@ -106,7 +105,6 @@ def benchmark_host_device_calculate_metrics( data_size_mib: int, h2d_bandwidth_ms: List[float], d2h_bandwidth_ms: List[float], - run_on_local_node: bool = False, # pylint: disable=unused-argument ) -> Tuple[Dict[str, Any], Dict[str, Any]]: """Calculates metrics for Host-Device transfer.""" params = locals().items() diff --git a/Ironwood/src/benchmark_utils.py b/Ironwood/src/benchmark_utils.py index 9546a4ec..aba83b81 100644 --- a/Ironwood/src/benchmark_utils.py +++ b/Ironwood/src/benchmark_utils.py @@ -29,8 +29,14 @@ from tensorflow.tsl.profiler.protobuf import xplane_pb2 -def configure_local_node_tpu(): - """Configures local node TPU execution on a single host without slicebuilder.""" +def configure_single_host_tpu(): + """Initializes libtpu as a standalone single host inside a multi-host slice. + + Process-wide and must be called before the first JAX device call. Enabled + via the benchmark-level `single_host_tpu_init` config key. Forces a 1-host + topology, skips SliceBuilder/tpunetd (leaving slice ICI routing untouched), + and disables cross-chip XLA launch barriers and continuation fusion. + """ os.environ["TPU_HOST_BOUNDS"] = "1,1,1" os.environ["TPU_SKIP_MDS_QUERY"] = "true" os.environ["TPU_WORKER_ID"] = "0" diff --git a/Ironwood/src/run_benchmark.py b/Ironwood/src/run_benchmark.py index b5c3d1ff..ae06d5dc 100644 --- a/Ironwood/src/run_benchmark.py +++ b/Ironwood/src/run_benchmark.py @@ -13,7 +13,7 @@ import string from typing import Any, Callable, Dict, List, Tuple from benchmark_utils import ( - configure_local_node_tpu, + configure_single_host_tpu, maybe_write_metrics_file, rename_xla_dump, MetricsStatistics, @@ -376,12 +376,6 @@ def run_single_benchmark(benchmark_config: Dict[str, Any], output_path: str): for param in benchmark_params: if "topology" not in param: param["topology"] = global_topology - # Inject run_on_local_node from config if present in benchmark_config - global_run_on_local_node = benchmark_config.get("run_on_local_node") - if global_run_on_local_node is not None: - for param in benchmark_params: - if "run_on_local_node" not in param: - param["run_on_local_node"] = global_run_on_local_node if not benchmark_name: raise ValueError("Each benchmark must have a benchmark_name.") @@ -391,8 +385,12 @@ def run_single_benchmark(benchmark_config: Dict[str, Any], output_path: str): benchmark_func, calculate_metrics_func = get_benchmark_functions( benchmark_name ) - if benchmark_config.get("run_on_local_node", False): - configure_local_node_tpu() + # `single_host_tpu_init` is a benchmark-level flag (not a sweep param) because + # libtpu is initialized once per process. It must run after the benchmark + # module import above (which sets LIBTPU_INIT_ARGS) and before the first + # JAX device call. + if benchmark_config.get("single_host_tpu_init", False): + configure_single_host_tpu() print(f"\n{'=' * 30}Starting benchmark '{benchmark_name}'{'=' * 30}\n") From 30aa86ea84f669c67c1c45810621c38f33cd8e9e Mon Sep 17 00:00:00 2001 From: Simran Kaur Date: Thu, 8 Oct 2026 12:29:07 +0000 Subject: [PATCH 3/9] chore(chs): log effective single-host TPU init config Prints the env vars and final LIBTPU_INIT_ARGS so a run's logs show that single_host_tpu_init actually took effect before libtpu initialized. --- Ironwood/src/benchmark_utils.py | 12 ++++++++++++ 1 file changed, 12 insertions(+) diff --git a/Ironwood/src/benchmark_utils.py b/Ironwood/src/benchmark_utils.py index aba83b81..222f142c 100644 --- a/Ironwood/src/benchmark_utils.py +++ b/Ironwood/src/benchmark_utils.py @@ -63,6 +63,18 @@ def configure_single_host_tpu(): if flag_name not in existing_args: existing_args = f"{existing_args} {flag}".strip() os.environ["LIBTPU_INIT_ARGS"] = existing_args + print( + "single_host_tpu_init: TPU_HOST_BOUNDS=%s TPU_SKIP_MDS_QUERY=%s " + "TPU_WORKER_ID=%s TPU_WORKER_HOSTNAMES=%s LIBTPU_INIT_ARGS=%r" + % ( + os.environ["TPU_HOST_BOUNDS"], + os.environ["TPU_SKIP_MDS_QUERY"], + os.environ["TPU_WORKER_ID"], + os.environ["TPU_WORKER_HOSTNAMES"], + os.environ["LIBTPU_INIT_ARGS"], + ), + flush=True, + ) def get_real_dtype_bytes(dtype) -> float: From c19a1bf0e528e68d53e955d58388977e899a4963 Mon Sep 17 00:00:00 2001 From: Simran Kaur Date: Thu, 8 Oct 2026 13:55:53 +0000 Subject: [PATCH 4/9] fix: disable async collective fusion in LIBTPU_INIT_ARGS libtpu 0.0.48 rejects --xla_tpu_enable_async_collective_fusion=true at backend init on v6e (ghostlite) and 7x with FAILED_PRECONDITION ("Continuation fusion for AllGather ... not supported on platforms other than Viperlite"), which broke gemm_multiple_run regardless of the single_host_tpu_init flag. These benchmarks are per-device matmuls with no all-gathers, so the fusion flags have no effect on what is measured. Turn them off unconditionally in all modules that set them, and drop the now-redundant rewrite from configure_single_host_tpu(). --- Ironwood/src/benchmark_compute.py | 6 +++--- Ironwood/src/benchmark_gemm.py | 6 +++--- Ironwood/src/benchmark_gemm_numerics.py | 6 +++--- Ironwood/src/benchmark_gemm_throttling.py | 6 +++--- Ironwood/src/benchmark_inference_compute.py | 6 +++--- Ironwood/src/benchmark_utils.py | 16 ++++------------ 6 files changed, 19 insertions(+), 27 deletions(-) diff --git a/Ironwood/src/benchmark_compute.py b/Ironwood/src/benchmark_compute.py index 2c42d19c..48b596ad 100644 --- a/Ironwood/src/benchmark_compute.py +++ b/Ironwood/src/benchmark_compute.py @@ -34,9 +34,9 @@ # Set the environment variable for TPU initialization arguments to optimize # collective matmul. Setting the flags to false will disable the optimization. os.environ["LIBTPU_INIT_ARGS"] = ( - "--xla_tpu_enable_async_collective_fusion=true " - "--xla_tpu_enable_async_collective_fusion_fuse_all_gather=true " - "--xla_tpu_enable_async_collective_fusion_multiple_steps=true " + "--xla_tpu_enable_async_collective_fusion=false " + "--xla_tpu_enable_async_collective_fusion_fuse_all_gather=false " + "--xla_tpu_enable_async_collective_fusion_multiple_steps=false " "--xla_tpu_overlap_compute_collective_tc=true " "--xla_enable_async_all_gather=true " "--xla_enable_async_collective_permute=true " diff --git a/Ironwood/src/benchmark_gemm.py b/Ironwood/src/benchmark_gemm.py index 15deb76a..81f47fd8 100644 --- a/Ironwood/src/benchmark_gemm.py +++ b/Ironwood/src/benchmark_gemm.py @@ -30,9 +30,9 @@ os.environ["LIBTPU_INIT_ARGS"] = ( - "--xla_tpu_enable_async_collective_fusion=true " - "--xla_tpu_enable_async_collective_fusion_fuse_all_gather=true " - "--xla_tpu_enable_async_collective_fusion_multiple_steps=true " + "--xla_tpu_enable_async_collective_fusion=false " + "--xla_tpu_enable_async_collective_fusion_fuse_all_gather=false " + "--xla_tpu_enable_async_collective_fusion_multiple_steps=false " "--xla_tpu_overlap_compute_collective_tc=true " "--xla_enable_async_all_gather=true " "--xla_enable_async_collective_permute=true " diff --git a/Ironwood/src/benchmark_gemm_numerics.py b/Ironwood/src/benchmark_gemm_numerics.py index 0e1d54ff..72b6df5a 100644 --- a/Ironwood/src/benchmark_gemm_numerics.py +++ b/Ironwood/src/benchmark_gemm_numerics.py @@ -30,9 +30,9 @@ # Set the environment variable for TPU initialization arguments to optimize # collective matmul. Setting the flags to false will disable the optimization. os.environ["LIBTPU_INIT_ARGS"] = ( - "--xla_tpu_enable_async_collective_fusion=true " - "--xla_tpu_enable_async_collective_fusion_fuse_all_gather=true " - "--xla_tpu_enable_async_collective_fusion_multiple_steps=true " + "--xla_tpu_enable_async_collective_fusion=false " + "--xla_tpu_enable_async_collective_fusion_fuse_all_gather=false " + "--xla_tpu_enable_async_collective_fusion_multiple_steps=false " "--xla_tpu_overlap_compute_collective_tc=true " "--xla_enable_async_all_gather=true " "--xla_enable_async_collective_permute=true " diff --git a/Ironwood/src/benchmark_gemm_throttling.py b/Ironwood/src/benchmark_gemm_throttling.py index c98d5131..b62c8eb9 100644 --- a/Ironwood/src/benchmark_gemm_throttling.py +++ b/Ironwood/src/benchmark_gemm_throttling.py @@ -18,9 +18,9 @@ import jax.numpy as jnp os.environ["LIBTPU_INIT_ARGS"] = ( - "--xla_tpu_enable_async_collective_fusion=true " - "--xla_tpu_enable_async_collective_fusion_fuse_all_gather=true " - "--xla_tpu_enable_async_collective_fusion_multiple_steps=true " + "--xla_tpu_enable_async_collective_fusion=false " + "--xla_tpu_enable_async_collective_fusion_fuse_all_gather=false " + "--xla_tpu_enable_async_collective_fusion_multiple_steps=false " "--xla_tpu_overlap_compute_collective_tc=true " "--xla_enable_async_all_gather=true " "--xla_enable_async_collective_permute=true " diff --git a/Ironwood/src/benchmark_inference_compute.py b/Ironwood/src/benchmark_inference_compute.py index f2ae12f0..6b4adc09 100644 --- a/Ironwood/src/benchmark_inference_compute.py +++ b/Ironwood/src/benchmark_inference_compute.py @@ -24,9 +24,9 @@ # Set the environment variable for TPU initialization arguments to optimize # collective matmul. Setting the flags to false will disable the optimization. os.environ["LIBTPU_INIT_ARGS"] = ( - "--xla_tpu_enable_async_collective_fusion=true " - "--xla_tpu_enable_async_collective_fusion_fuse_all_gather=true " - "--xla_tpu_enable_async_collective_fusion_multiple_steps=true " + "--xla_tpu_enable_async_collective_fusion=false " + "--xla_tpu_enable_async_collective_fusion_fuse_all_gather=false " + "--xla_tpu_enable_async_collective_fusion_multiple_steps=false " "--xla_tpu_overlap_compute_collective_tc=true " "--xla_enable_async_all_gather=true " "--xla_enable_async_collective_permute=true " diff --git a/Ironwood/src/benchmark_utils.py b/Ironwood/src/benchmark_utils.py index 222f142c..25bccf92 100644 --- a/Ironwood/src/benchmark_utils.py +++ b/Ironwood/src/benchmark_utils.py @@ -35,7 +35,7 @@ def configure_single_host_tpu(): Process-wide and must be called before the first JAX device call. Enabled via the benchmark-level `single_host_tpu_init` config key. Forces a 1-host topology, skips SliceBuilder/tpunetd (leaving slice ICI routing untouched), - and disables cross-chip XLA launch barriers and continuation fusion. + and disables the cross-host XLA launch barrier. """ os.environ["TPU_HOST_BOUNDS"] = "1,1,1" os.environ["TPU_SKIP_MDS_QUERY"] = "true" @@ -44,19 +44,11 @@ def configure_single_host_tpu(): local_flags = [ "--deepsea_hal_test_skip_slicebuilder=true", "--deepsea_hal_test_allow_multichip_skip_slicebuilder=true", - "--xla_tpu_enable_async_collective_fusion=false", "--xla_tpu_use_enhanced_launch_barrier=false", ] - existing_args = ( - os.environ.get("LIBTPU_INIT_ARGS", "") - .replace( - "--xla_tpu_enable_async_collective_fusion=true", - "--xla_tpu_enable_async_collective_fusion=false", - ) - .replace( - "--xla_tpu_use_enhanced_launch_barrier=true", - "--xla_tpu_use_enhanced_launch_barrier=false", - ) + existing_args = os.environ.get("LIBTPU_INIT_ARGS", "").replace( + "--xla_tpu_use_enhanced_launch_barrier=true", + "--xla_tpu_use_enhanced_launch_barrier=false", ) for flag in local_flags: flag_name = flag.split("=")[0] From 6e142604cbdb01ba3d0cc408d18f7eda1a85bff2 Mon Sep 17 00:00:00 2001 From: Simran Kaur Date: Thu, 8 Oct 2026 14:01:27 +0000 Subject: [PATCH 5/9] chore: align async collective fusion flags with upstream op_flags.yaml Keep --xla_tpu_enable_async_collective_fusion=false (required by libtpu 0.0.48 on v6e/7x), restore _multiple_steps=true and drop _fuse_all_gather, matching the gemm_flags set in AI-Hypercomputer/accelerator-microbenchmarks main. No behavioral change while the master switch is off. --- Ironwood/src/benchmark_compute.py | 3 +-- Ironwood/src/benchmark_gemm.py | 3 +-- Ironwood/src/benchmark_gemm_numerics.py | 3 +-- Ironwood/src/benchmark_gemm_throttling.py | 3 +-- Ironwood/src/benchmark_inference_compute.py | 3 +-- 5 files changed, 5 insertions(+), 10 deletions(-) diff --git a/Ironwood/src/benchmark_compute.py b/Ironwood/src/benchmark_compute.py index 48b596ad..18572227 100644 --- a/Ironwood/src/benchmark_compute.py +++ b/Ironwood/src/benchmark_compute.py @@ -35,8 +35,7 @@ # collective matmul. Setting the flags to false will disable the optimization. os.environ["LIBTPU_INIT_ARGS"] = ( "--xla_tpu_enable_async_collective_fusion=false " - "--xla_tpu_enable_async_collective_fusion_fuse_all_gather=false " - "--xla_tpu_enable_async_collective_fusion_multiple_steps=false " + "--xla_tpu_enable_async_collective_fusion_multiple_steps=true " "--xla_tpu_overlap_compute_collective_tc=true " "--xla_enable_async_all_gather=true " "--xla_enable_async_collective_permute=true " diff --git a/Ironwood/src/benchmark_gemm.py b/Ironwood/src/benchmark_gemm.py index 81f47fd8..46bb0ceb 100644 --- a/Ironwood/src/benchmark_gemm.py +++ b/Ironwood/src/benchmark_gemm.py @@ -31,8 +31,7 @@ os.environ["LIBTPU_INIT_ARGS"] = ( "--xla_tpu_enable_async_collective_fusion=false " - "--xla_tpu_enable_async_collective_fusion_fuse_all_gather=false " - "--xla_tpu_enable_async_collective_fusion_multiple_steps=false " + "--xla_tpu_enable_async_collective_fusion_multiple_steps=true " "--xla_tpu_overlap_compute_collective_tc=true " "--xla_enable_async_all_gather=true " "--xla_enable_async_collective_permute=true " diff --git a/Ironwood/src/benchmark_gemm_numerics.py b/Ironwood/src/benchmark_gemm_numerics.py index 72b6df5a..6818bd3b 100644 --- a/Ironwood/src/benchmark_gemm_numerics.py +++ b/Ironwood/src/benchmark_gemm_numerics.py @@ -31,8 +31,7 @@ # collective matmul. Setting the flags to false will disable the optimization. os.environ["LIBTPU_INIT_ARGS"] = ( "--xla_tpu_enable_async_collective_fusion=false " - "--xla_tpu_enable_async_collective_fusion_fuse_all_gather=false " - "--xla_tpu_enable_async_collective_fusion_multiple_steps=false " + "--xla_tpu_enable_async_collective_fusion_multiple_steps=true " "--xla_tpu_overlap_compute_collective_tc=true " "--xla_enable_async_all_gather=true " "--xla_enable_async_collective_permute=true " diff --git a/Ironwood/src/benchmark_gemm_throttling.py b/Ironwood/src/benchmark_gemm_throttling.py index b62c8eb9..5ceb955e 100644 --- a/Ironwood/src/benchmark_gemm_throttling.py +++ b/Ironwood/src/benchmark_gemm_throttling.py @@ -19,8 +19,7 @@ os.environ["LIBTPU_INIT_ARGS"] = ( "--xla_tpu_enable_async_collective_fusion=false " - "--xla_tpu_enable_async_collective_fusion_fuse_all_gather=false " - "--xla_tpu_enable_async_collective_fusion_multiple_steps=false " + "--xla_tpu_enable_async_collective_fusion_multiple_steps=true " "--xla_tpu_overlap_compute_collective_tc=true " "--xla_enable_async_all_gather=true " "--xla_enable_async_collective_permute=true " diff --git a/Ironwood/src/benchmark_inference_compute.py b/Ironwood/src/benchmark_inference_compute.py index 6b4adc09..f62dc18c 100644 --- a/Ironwood/src/benchmark_inference_compute.py +++ b/Ironwood/src/benchmark_inference_compute.py @@ -25,8 +25,7 @@ # collective matmul. Setting the flags to false will disable the optimization. os.environ["LIBTPU_INIT_ARGS"] = ( "--xla_tpu_enable_async_collective_fusion=false " - "--xla_tpu_enable_async_collective_fusion_fuse_all_gather=false " - "--xla_tpu_enable_async_collective_fusion_multiple_steps=false " + "--xla_tpu_enable_async_collective_fusion_multiple_steps=true " "--xla_tpu_overlap_compute_collective_tc=true " "--xla_enable_async_all_gather=true " "--xla_enable_async_collective_permute=true " From 58e342a5920c1a0291ae49b1fd0efca959b2a89e Mon Sep 17 00:00:00 2001 From: Simran Kaur Date: Thu, 8 Oct 2026 14:04:14 +0000 Subject: [PATCH 6/9] chore: limit async collective fusion change to benchmark_gemm.py Keep the PR scoped to the benchmarks CHS runs. Other modules that set --xla_tpu_enable_async_collective_fusion=true will need the same change to run on libtpu 0.0.48; see op_flags.yaml on main. --- Ironwood/src/benchmark_compute.py | 3 ++- Ironwood/src/benchmark_gemm_numerics.py | 3 ++- Ironwood/src/benchmark_gemm_throttling.py | 3 ++- Ironwood/src/benchmark_inference_compute.py | 3 ++- 4 files changed, 8 insertions(+), 4 deletions(-) diff --git a/Ironwood/src/benchmark_compute.py b/Ironwood/src/benchmark_compute.py index 18572227..2c42d19c 100644 --- a/Ironwood/src/benchmark_compute.py +++ b/Ironwood/src/benchmark_compute.py @@ -34,7 +34,8 @@ # Set the environment variable for TPU initialization arguments to optimize # collective matmul. Setting the flags to false will disable the optimization. os.environ["LIBTPU_INIT_ARGS"] = ( - "--xla_tpu_enable_async_collective_fusion=false " + "--xla_tpu_enable_async_collective_fusion=true " + "--xla_tpu_enable_async_collective_fusion_fuse_all_gather=true " "--xla_tpu_enable_async_collective_fusion_multiple_steps=true " "--xla_tpu_overlap_compute_collective_tc=true " "--xla_enable_async_all_gather=true " diff --git a/Ironwood/src/benchmark_gemm_numerics.py b/Ironwood/src/benchmark_gemm_numerics.py index 6818bd3b..0e1d54ff 100644 --- a/Ironwood/src/benchmark_gemm_numerics.py +++ b/Ironwood/src/benchmark_gemm_numerics.py @@ -30,7 +30,8 @@ # Set the environment variable for TPU initialization arguments to optimize # collective matmul. Setting the flags to false will disable the optimization. os.environ["LIBTPU_INIT_ARGS"] = ( - "--xla_tpu_enable_async_collective_fusion=false " + "--xla_tpu_enable_async_collective_fusion=true " + "--xla_tpu_enable_async_collective_fusion_fuse_all_gather=true " "--xla_tpu_enable_async_collective_fusion_multiple_steps=true " "--xla_tpu_overlap_compute_collective_tc=true " "--xla_enable_async_all_gather=true " diff --git a/Ironwood/src/benchmark_gemm_throttling.py b/Ironwood/src/benchmark_gemm_throttling.py index 5ceb955e..c98d5131 100644 --- a/Ironwood/src/benchmark_gemm_throttling.py +++ b/Ironwood/src/benchmark_gemm_throttling.py @@ -18,7 +18,8 @@ import jax.numpy as jnp os.environ["LIBTPU_INIT_ARGS"] = ( - "--xla_tpu_enable_async_collective_fusion=false " + "--xla_tpu_enable_async_collective_fusion=true " + "--xla_tpu_enable_async_collective_fusion_fuse_all_gather=true " "--xla_tpu_enable_async_collective_fusion_multiple_steps=true " "--xla_tpu_overlap_compute_collective_tc=true " "--xla_enable_async_all_gather=true " diff --git a/Ironwood/src/benchmark_inference_compute.py b/Ironwood/src/benchmark_inference_compute.py index f62dc18c..f2ae12f0 100644 --- a/Ironwood/src/benchmark_inference_compute.py +++ b/Ironwood/src/benchmark_inference_compute.py @@ -24,7 +24,8 @@ # Set the environment variable for TPU initialization arguments to optimize # collective matmul. Setting the flags to false will disable the optimization. os.environ["LIBTPU_INIT_ARGS"] = ( - "--xla_tpu_enable_async_collective_fusion=false " + "--xla_tpu_enable_async_collective_fusion=true " + "--xla_tpu_enable_async_collective_fusion_fuse_all_gather=true " "--xla_tpu_enable_async_collective_fusion_multiple_steps=true " "--xla_tpu_overlap_compute_collective_tc=true " "--xla_enable_async_all_gather=true " From 9d68aa1703196bcd43cb43a192454b6fc3c8a771 Mon Sep 17 00:00:00 2001 From: Simran Kaur Date: Thu, 8 Oct 2026 14:33:57 +0000 Subject: [PATCH 7/9] refactor(chs): gate per-device GEMM input generation on single_host_tpu_init The per-device data path exists because SliceBuilder is skipped under single_host_tpu_init, leaving intra-host ICI unprogrammed so a device_put from device 0 would be a cross-chip DMA. Tie the branch to that condition (via is_single_host_tpu_init()) instead of to the mesh choice (run_on_local_node), which keeps the default path unchanged for all non-single-host runs. --- Ironwood/src/benchmark_gemm.py | 3 ++- Ironwood/src/benchmark_utils.py | 22 ++++++++++------------ Ironwood/src/run_benchmark.py | 4 ---- 3 files changed, 12 insertions(+), 17 deletions(-) diff --git a/Ironwood/src/benchmark_gemm.py b/Ironwood/src/benchmark_gemm.py index 46bb0ceb..1edec7c9 100644 --- a/Ironwood/src/benchmark_gemm.py +++ b/Ironwood/src/benchmark_gemm.py @@ -17,6 +17,7 @@ from benchmark_utils import get_peak_flops_multiplier from benchmark_utils import get_rhs_named_shading from benchmark_utils import handle_based_on_sharding +from benchmark_utils import is_single_host_tpu_init from benchmark_utils import iteration_timeit from benchmark_utils import multiple_iteration_timeit_from_trace from benchmark_utils import ShardingStrategy @@ -101,7 +102,7 @@ def f(x, y): lhs_dtype = dtype rhs_dtype = dtype - if run_on_local_node: + if is_single_host_tpu_init(): key = jax.device_put(jax.random.key(SEED), jax.devices("cpu")[0]) mesh_devices = list(mesh.devices.flat) diff --git a/Ironwood/src/benchmark_utils.py b/Ironwood/src/benchmark_utils.py index 25bccf92..f3dbc662 100644 --- a/Ironwood/src/benchmark_utils.py +++ b/Ironwood/src/benchmark_utils.py @@ -29,6 +29,14 @@ from tensorflow.tsl.profiler.protobuf import xplane_pb2 +_single_host_tpu_init = False + + +def is_single_host_tpu_init() -> bool: + """Returns True if configure_single_host_tpu() was applied in this process.""" + return _single_host_tpu_init + + def configure_single_host_tpu(): """Initializes libtpu as a standalone single host inside a multi-host slice. @@ -37,6 +45,8 @@ def configure_single_host_tpu(): topology, skips SliceBuilder/tpunetd (leaving slice ICI routing untouched), and disables the cross-host XLA launch barrier. """ + global _single_host_tpu_init + _single_host_tpu_init = True os.environ["TPU_HOST_BOUNDS"] = "1,1,1" os.environ["TPU_SKIP_MDS_QUERY"] = "true" os.environ["TPU_WORKER_ID"] = "0" @@ -55,18 +65,6 @@ def configure_single_host_tpu(): if flag_name not in existing_args: existing_args = f"{existing_args} {flag}".strip() os.environ["LIBTPU_INIT_ARGS"] = existing_args - print( - "single_host_tpu_init: TPU_HOST_BOUNDS=%s TPU_SKIP_MDS_QUERY=%s " - "TPU_WORKER_ID=%s TPU_WORKER_HOSTNAMES=%s LIBTPU_INIT_ARGS=%r" - % ( - os.environ["TPU_HOST_BOUNDS"], - os.environ["TPU_SKIP_MDS_QUERY"], - os.environ["TPU_WORKER_ID"], - os.environ["TPU_WORKER_HOSTNAMES"], - os.environ["LIBTPU_INIT_ARGS"], - ), - flush=True, - ) def get_real_dtype_bytes(dtype) -> float: diff --git a/Ironwood/src/run_benchmark.py b/Ironwood/src/run_benchmark.py index ae06d5dc..2092bcfb 100644 --- a/Ironwood/src/run_benchmark.py +++ b/Ironwood/src/run_benchmark.py @@ -385,10 +385,6 @@ def run_single_benchmark(benchmark_config: Dict[str, Any], output_path: str): benchmark_func, calculate_metrics_func = get_benchmark_functions( benchmark_name ) - # `single_host_tpu_init` is a benchmark-level flag (not a sweep param) because - # libtpu is initialized once per process. It must run after the benchmark - # module import above (which sets LIBTPU_INIT_ARGS) and before the first - # JAX device call. if benchmark_config.get("single_host_tpu_init", False): configure_single_host_tpu() From 64984172b94cb2531774e8753dcf719b5040a61c Mon Sep 17 00:00:00 2001 From: Simran Kaur Date: Thu, 8 Oct 2026 17:26:02 +0000 Subject: [PATCH 8/9] fix(chs): match existing LIBTPU_INIT_ARGS flags by exact name in configure_single_host_tpu Substring matching could skip a flag whose name is a prefix of an already-present flag. Compare against the set of exact flag names instead. --- Ironwood/src/benchmark_utils.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/Ironwood/src/benchmark_utils.py b/Ironwood/src/benchmark_utils.py index f3dbc662..e7a4d924 100644 --- a/Ironwood/src/benchmark_utils.py +++ b/Ironwood/src/benchmark_utils.py @@ -60,9 +60,10 @@ def configure_single_host_tpu(): "--xla_tpu_use_enhanced_launch_barrier=true", "--xla_tpu_use_enhanced_launch_barrier=false", ) + existing_flag_names = {f.split("=")[0] for f in existing_args.split()} for flag in local_flags: flag_name = flag.split("=")[0] - if flag_name not in existing_args: + if flag_name not in existing_flag_names: existing_args = f"{existing_args} {flag}".strip() os.environ["LIBTPU_INIT_ARGS"] = existing_args From cfc5876964954660d0665b881b36c931fe9f0ae1 Mon Sep 17 00:00:00 2001 From: Simran Kaur Date: Thu, 8 Oct 2026 17:32:42 +0000 Subject: [PATCH 9/9] style(chs): fix pre-commit findings - benchmark_utils.py: wrap docstring to 80 cols, use maxsplit=1 in flag name splits, strip trailing whitespace - requirements.txt: add missing trailing newline --- Ironwood/src/benchmark_utils.py | 10 ++++++---- requirements.txt | 2 +- 2 files changed, 7 insertions(+), 5 deletions(-) diff --git a/Ironwood/src/benchmark_utils.py b/Ironwood/src/benchmark_utils.py index e7a4d924..880d4dd0 100644 --- a/Ironwood/src/benchmark_utils.py +++ b/Ironwood/src/benchmark_utils.py @@ -33,7 +33,7 @@ def is_single_host_tpu_init() -> bool: - """Returns True if configure_single_host_tpu() was applied in this process.""" + """Returns True if configure_single_host_tpu() ran in this process.""" return _single_host_tpu_init @@ -60,9 +60,11 @@ def configure_single_host_tpu(): "--xla_tpu_use_enhanced_launch_barrier=true", "--xla_tpu_use_enhanced_launch_barrier=false", ) - existing_flag_names = {f.split("=")[0] for f in existing_args.split()} + existing_flag_names = { + f.split("=", maxsplit=1)[0] for f in existing_args.split() + } for flag in local_flags: - flag_name = flag.split("=")[0] + flag_name = flag.split("=", maxsplit=1)[0] if flag_name not in existing_flag_names: existing_args = f"{existing_args} {flag}".strip() os.environ["LIBTPU_INIT_ARGS"] = existing_args @@ -1296,7 +1298,7 @@ def handle_based_on_sharding( def create_mesh(strategy: ShardingStrategy, local_mesh: bool = False) -> Mesh: """Creates a mesh. - + Args: strategy: The sharding strategy to apply. local_mesh: If True, restricts the mesh to local devices. diff --git a/requirements.txt b/requirements.txt index 77049aa1..88cfd122 100644 --- a/requirements.txt +++ b/requirements.txt @@ -14,4 +14,4 @@ tensorflow-cpu qwix@git+https://github.com/google/qwix.git tokamax tune-jax==0.7.2 -immutabledict \ No newline at end of file +immutabledict