diff --git a/Ironwood/src/benchmark_gemm.py b/Ironwood/src/benchmark_gemm.py index f7fbe91d..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 @@ -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 diff --git a/Ironwood/src/benchmark_utils.py b/Ironwood/src/benchmark_utils.py index 22a5356a..880d4dd0 100644 --- a/Ironwood/src/benchmark_utils.py +++ b/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. diff --git a/Ironwood/src/run_benchmark.py b/Ironwood/src/run_benchmark.py index 2dd12b75..2092bcfb 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_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") diff --git a/requirements.txt b/requirements.txt index b37bec99..88cfd122 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 @@ -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