feat(chs): single-host TPU init for node-local diagnostics + libtpu 0.0.48

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
@@ -30,8 +31,7 @@
 
 
 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=false "
     "--xla_tpu_enable_async_collective_fusion_multiple_steps=true "
     "--xla_tpu_overlap_compute_collective_tc=true "
     "--xla_enable_async_all_gather=true "
@@ -102,22 +102,53 @@ def f(x, y):
     lhs_dtype = dtype
     rhs_dtype = dtype
 
-    key = jax.random.key(SEED)
+    if is_single_host_tpu_init():
+        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
 

Ironwood/src/benchmark_utils.py

@@ -29,6 +29,47 @@
 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() ran 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.
+
+    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 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"
+    os.environ["TPU_WORKER_HOSTNAMES"] = "localhost"
+    local_flags = [
+        "--deepsea_hal_test_skip_slicebuilder=true",
+        "--deepsea_hal_test_allow_multichip_skip_slicebuilder=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",
+    )
+    existing_flag_names = {
+        f.split("=", maxsplit=1)[0] for f in existing_args.split()
+    }
+    for flag in local_flags:
+        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
+
+
 def get_real_dtype_bytes(dtype) -> float:
     """Returns the real byte size of a dtype, handling sub-byte types."""
     try:
@@ -1024,6 +1065,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 +1152,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 +1164,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}.")
 
@@ -1240,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.

Ironwood/src/run_benchmark.py

@@ -13,6 +13,7 @@
 import string
 from typing import Any, Callable, Dict, List, Tuple
 from benchmark_utils import (
+    configure_single_host_tpu,
     maybe_write_metrics_file,
     rename_xla_dump,
     MetricsStatistics,
@@ -384,6 +385,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("single_host_tpu_init", False):
+        configure_single_host_tpu()
 
     print(f"\n{'=' * 30}Starting benchmark '{benchmark_name}'{'=' * 30}\n")
 

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
@@ -14,4 +14,4 @@ tensorflow-cpu
 qwix@git+https://github.com/google/qwix.git
 tokamax
 tune-jax==0.7.2
-immutabledict
+immutabledict