feat: add opt-in single_host_tpu_init for node-local diagnostics on multi-host slices

README.md

@@ -257,6 +257,7 @@ tpums benchmark run-config configs/sample_configs/parameter_sweep.yaml \
 - **Report Output & Compiler Flags**:
     - **`--output_dir <path>`**: Directory to save `summary.csv` and `detailed.json` (defaults to `results/`).
     - **`--xla_flags_file_path <path>`**: Optional path to a custom YAML file overriding default per-benchmark XLA / compiler runtime flags.
+- **Standalone Single-Host Init (`--single_host_tpu_init` / `single_host_tpu_init: true`, optional)**: Initializes this host as a standalone 1-host topology inside a multi-host slice (skips SliceBuilder, disables the cross-host launch barrier) so that a single-host benchmark can be used for node-local diagnostics without the other hosts of the slice participating. The slice's ICI routing is left untouched, so device-to-device transfers hang in this mode; only `gemm`, `gemm_throttling` and `hbm` are known to avoid them (a warning is printed for other benchmarks). Off by default; has no effect on normal single- or multi-host runs.
 
 ---
 
@@ -273,6 +274,7 @@ Configuration files define a top-level `benchmark:` mapping containing:
 - **`name:` (`<benchmark_name>`)** — Target benchmark name (e.g., `gemm`, `hbm`, `all_reduce`; see [Benchmark Catalog](#benchmark-catalog)).
 - **`xprof_timing:`** *(optional)* — Boolean (`true` / `false`) to enable XProf hardware trace collection and device timing analysis.
 - **`xprof_device_timing_aggregation:`** *(optional)* — Device timing aggregation mode when `xprof_timing: true` (`measured_device_or_first` [default], `local_slowest_device`, or `global_slowest_per_step`).
+- **`single_host_tpu_init:`** *(optional)* — Boolean (`true` / `false`, default `false`) to initialize this host standalone inside a multi-host slice for node-local diagnostics (see [Common Execution Flags](#4-common-execution-flags)).
 - **`params:`** *(optional)* — Baseline execution parameters shared across all generated runs.
 - **`cases:` / `cases_from_csv:` / `sweep:`** *(optional)* — Case override and parameter sweep generators.
 

src/accelerator_microbenchmarks/benchmarks/matmul.py

@@ -82,6 +82,11 @@ class GemmParams(base.BaseBenchmarkParams):
 TGemmConfig = TypeVar("TGemmConfig", bound=base.BaseBenchmarkParams)
 
 
+def _sample_normal(key: jax.Array, shape: tuple[int, ...], dtype: Any):
+  """Module-level so `jax.jit` reuses its compile cache across calls."""
+  return jax.random.normal(key, shape).astype(dtype)
+
+
 class BaseGemmBenchmark(base.BaseBenchmark[TGemmConfig]):
   """Base class for GEMM benchmarks.
 
@@ -198,6 +203,27 @@ def get_run_identifier(self) -> str:
       run_identifier += "_sf"
     return run_identifier
 
+  def _replicated_normal(
+      self, key: jax.Array, shape: tuple[int, ...], dtype: Any
+  ) -> jax.Array:
+    """Returns standard-normal samples replicated across the mesh.
+
+    Generated on every mesh device under `jit` with a replicated
+    `out_shardings` (as in `hbm.py`), so no device-to-device broadcast is
+    needed; required for `single_host_tpu_init`, where such transfers hang.
+    """
+    if self.mesh is None:
+      raise ValueError("Mesh not initialized.")
+    replicated_sharding = jax.sharding.NamedSharding(
+        self.mesh, jax.sharding.PartitionSpec(*([None] * len(shape)))
+    )
+    generate = jax.jit(
+        _sample_normal,
+        static_argnums=(1, 2),
+        out_shardings=replicated_sharding,
+    )
+    return generate(key, shape, dtype)
+
   def generate_inputs(self) -> tuple[Any, ...]:
     (
         m,
@@ -219,57 +245,34 @@ def generate_inputs(self) -> tuple[Any, ...]:
     key = jax.random.PRNGKey(self.config.seed)
     k1, k2, k3, k4, k5 = jax.random.split(key, 5)
 
-    # Data generation in HBM
+    # Data generation in HBM, replicated across the mesh to avoid
+    # discrepancies between devices.
     # Note: JAX might require intermediate conversion for random.normal if
     # dtypes aren't supported
     a_shape = (k, m) if transpose_a else (m, k)
     b_shape = (n, k) if transpose_b else (k, n)
-    a = jax.random.normal(k1, a_shape).astype(in_dtype)
-    b = jax.random.normal(k2, b_shape).astype(in_dtype)
-
-    # Optional scaling factors
-    use_sf = self.config.use_scaling_factors
-    sf0 = jax.random.normal(k3, (m, 1)).astype(jnp.float32) if use_sf else None
-    sf1 = jax.random.normal(k4, (1, n)).astype(jnp.float32) if use_sf else None
-
-    use_accumulator_matrix = True if self.config.beta != 0.0 else False
-    c = (
-        jax.random.normal(k5, (m, n)).astype(in_dtype)
-        if use_accumulator_matrix
-        else None
-    )
-
-    # Replicated sharding for computation ops to avoid discrepancies
-    if self.mesh is None:
-      raise ValueError("Mesh not initialized.")
-    replicated_sharding = jax.sharding.NamedSharding(
-        self.mesh, jax.sharding.PartitionSpec(None, None)
-    )
+    a = self._replicated_normal(k1, a_shape, in_dtype)
+    b = self._replicated_normal(k2, b_shape, in_dtype)
 
     # Row and column scaling factor vectors (sf0: m x 1, sf1: 1 x n)
     # Unpacked as sf0, sf1 args into gemm_fn(a, b, sf0, sf1, c) if enabled,
     # else (None, None).
-    if use_sf:
+    if self.config.use_scaling_factors:
       sf_args = [
-          jax.device_put(sf0, replicated_sharding),
-          jax.device_put(sf1, replicated_sharding),
+          self._replicated_normal(k3, (m, 1), jnp.float32),
+          self._replicated_normal(k4, (1, n), jnp.float32),
       ]
     else:
       sf_args = [None, None]
 
     # Accumulator matrix C (M x N) for fused GEMM addition when beta != 0.0.
     # Unpacked as c arg into gemm_fn(a, b, sf0, sf1, c) if enabled, else [].
-    if use_accumulator_matrix:
-      acc_args = [jax.device_put(c, replicated_sharding)]
+    if self.config.beta != 0.0:
+      acc_args = [self._replicated_normal(k5, (m, n), in_dtype)]
     else:
       acc_args = []
 
-    return (
-        jax.device_put(a, replicated_sharding),
-        jax.device_put(b, replicated_sharding),
-        *sf_args,
-        *acc_args,
-    )
+    return (a, b, *sf_args, *acc_args)
 
   def run_op(self, *args) -> jnp.ndarray:
     if self._jit_fn is None:

src/accelerator_microbenchmarks/cli.py

@@ -39,6 +39,17 @@ def _add_common_execution_args(parser: argparse.ArgumentParser) -> None:
       default=None,
       help="Optional path to custom op_flags YAML file.",
   )
+  parser.add_argument(
+      "--single_host_tpu_init",
+      action="store_true",
+      default=False,
+      help=(
+          "Initialize this host standalone inside a multi-host slice for"
+          " node-local diagnostics (forces a 1-host topology and skips"
+          " SliceBuilder). Device-to-device transfers hang in this mode;"
+          " supported benchmarks: gemm, gemm_throttling, hbm."
+      ),
+  )
 
 
 def _str_to_bool(value: str) -> bool:
@@ -305,12 +316,18 @@ def run(argv: Sequence[str]) -> None:
           "xprof_device_timing_aggregation",
           constants.XprofDeviceTimingAggregation.MEASURED_DEVICE_OR_FIRST,
       )
+    effective_single_host_tpu_init = args.single_host_tpu_init or (
+        bool(raw_configs[0].get("single_host_tpu_init", False))
+        if raw_configs
+        else False
+    )
 
     tasks = []
     for raw in raw_configs:
       name = raw.pop("name")
       raw.pop("xprof_timing", None)
       raw.pop("xprof_device_timing_aggregation", None)
+      raw.pop("single_host_tpu_init", None)
       bench_cls = registry.benchmark_registry.get_benchmark(name)
       tasks.append((name, bench_cls.Config(**raw)))
 
@@ -324,6 +341,7 @@ def run(argv: Sequence[str]) -> None:
         xprof_dir=args.xprof_dir,
         config_path=args.config_path,
         xla_flags_file_path=args.xla_flags_file_path,
+        single_host_tpu_init=effective_single_host_tpu_init,
     )
     return
 
@@ -338,6 +356,7 @@ def run(argv: Sequence[str]) -> None:
         xprof_device_timing_aggregation=args.xprof_device_timing_aggregation,
         xprof_dir=args.xprof_dir,
         xla_flags_file_path=args.xla_flags_file_path,
+        single_host_tpu_init=args.single_host_tpu_init,
     )
     return
 

src/accelerator_microbenchmarks/core/config.py

@@ -20,6 +20,7 @@ class BenchmarkKey(str, enum.Enum):
   SWEEP = "sweep"
   XPROF_TIMING = "xprof_timing"
   XPROF_DEVICE_TIMING_AGGREGATION = "xprof_device_timing_aggregation"
+  SINGLE_HOST_TPU_INIT = "single_host_tpu_init"
 
 
 _ALLOWED_BENCHMARK_KEYS = frozenset(BenchmarkKey)
@@ -304,6 +305,16 @@ def load_config(path: str) -> list[dict[str, Any]]:
         xprof_device_timing_aggregation
     )
 
+  single_host_tpu_init = benchmark_spec.get(BenchmarkKey.SINGLE_HOST_TPU_INIT)
+  if single_host_tpu_init is not None:
+    if not isinstance(single_host_tpu_init, bool):
+      raise ValueError(
+          f"Expected '{BenchmarkKey.SINGLE_HOST_TPU_INIT.value}' in benchmark"
+          f" '{benchmark_name}' to be a bool, but got"
+          f" {type(single_host_tpu_init).__name__}."
+      )
+    base_params["single_host_tpu_init"] = single_host_tpu_init
+
   # Stage 3: Cases (per-case parameter overrides over base_params)
   case_dicts = (
       [{**base_params, **case} for case in cases_list]

src/accelerator_microbenchmarks/core/runner.py

@@ -14,6 +14,7 @@
 from accelerator_microbenchmarks.core import platform
 from accelerator_microbenchmarks.core import registry
 from accelerator_microbenchmarks.core import report
+from accelerator_microbenchmarks.core import single_host
 from accelerator_microbenchmarks.core import system
 import jax
 import yaml
@@ -147,6 +148,7 @@ def run_benchmarks(
     xprof_device_timing_aggregation: str = (
         constants.XprofDeviceTimingAggregation.MEASURED_DEVICE_OR_FIRST
     ),
+    single_host_tpu_init: bool = False,
 ) -> List[base.BenchmarkResult]:
   """Core execution engine for typed benchmark task configurations."""
   xprof_config = base.XprofConfig(
@@ -162,6 +164,20 @@ def run_benchmarks(
       [{"name": task_name} for task_name, _ in tasks],
       xla_flags_file_path=xla_flags_file_path,
   )
+  # Optional standalone single-host init. Must follow set_xla_flags (which
+  # replaces LIBTPU_INIT_ARGS) and precede any JAX device call.
+  if single_host_tpu_init:
+    unsupported = sorted(
+        {name for name, _ in tasks} - single_host.SUPPORTED_BENCHMARKS
+    )
+    if unsupported:
+      print(
+          f"Warning: single_host_tpu_init is not supported for {unsupported}:"
+          " device-to-device transfers hang in this mode, so collectives and"
+          " device_to_device cannot work and other benchmarks may hang."
+          f" Supported: {sorted(single_host.SUPPORTED_BENCHMARKS)}."
+      )
+    single_host.configure_single_host_tpu()
 
   # 2. Ensure JAX is initialized
   init_jax_distributed()

src/accelerator_microbenchmarks/core/single_host.py

@@ -0,0 +1,68 @@
+"""Standalone single-host TPU initialization inside a multi-host slice.
+
+Opt-in (`single_host_tpu_init` config key / CLI flag) for node-local
+diagnostics. `configure_single_host_tpu()` forces a 1-host topology, skips
+SliceBuilder/tpunetd and disables the cross-host XLA launch barrier. It is
+process-wide and must run before the first JAX device call. The slice's ICI
+routing is left untouched, so device-to-device transfers hang in this mode;
+only benchmarks listed in `SUPPORTED_BENCHMARKS` are known to avoid them.
+"""
+
+import os
+
+# Benchmarks known to perform no device-to-device transfers in this mode
+# (gemm_throttling shares GeneralizedGemmBenchmark.generate_inputs). Lists
+# primary names and registered aliases; keep in sync with
+# `@benchmark_registry.register(..., aliases=...)`.
+SUPPORTED_BENCHMARKS = frozenset({
+    "gemm",
+    "gemm_generalized",
+    "gemm_throttling",
+    "hbm",
+    "hbm_bandwidth",
+})
+
+_SINGLE_HOST_ENV = {
+    "TPU_HOST_BOUNDS": "1,1,1",
+    "TPU_SKIP_MDS_QUERY": "true",
+    "TPU_WORKER_ID": "0",
+    "TPU_WORKER_HOSTNAMES": "localhost",
+}
+
+_SINGLE_HOST_LIBTPU_FLAGS = {
+    "--deepsea_hal_test_skip_slicebuilder": "true",
+    "--deepsea_hal_test_allow_multichip_skip_slicebuilder": "true",
+    "--xla_tpu_use_enhanced_launch_barrier": "false",
+}
+
+
+def _set_env(name: str, value: str) -> None:
+  """Sets `os.environ[name]`, logging the previous value when it changes."""
+  previous = os.environ.get(name)
+  if previous == value:
+    return
+  os.environ[name] = value
+  if previous is None:
+    print(f"single_host_tpu_init: set {name}={value}")
+  else:
+    print(
+        f"single_host_tpu_init: override {name}={value} (was {previous!r})"
+    )
+
+
+def configure_single_host_tpu() -> None:
+  """Configures libtpu to initialize this host standalone. Call before JAX init."""
+  for name, value in _SINGLE_HOST_ENV.items():
+    _set_env(name, value)
+
+  # Merge into LIBTPU_INIT_ARGS by flag name so pre-existing values (e.g. from
+  # op_flags.yaml) are overridden rather than duplicated or left in place.
+  flags: dict[str, str | None] = {}
+  for token in os.environ.get("LIBTPU_INIT_ARGS", "").split():
+    name, sep, value = token.partition("=")
+    flags[name] = value if sep else None
+  flags.update(_SINGLE_HOST_LIBTPU_FLAGS)
+  _set_env(
+      "LIBTPU_INIT_ARGS",
+      " ".join(k if v is None else f"{k}={v}" for k, v in flags.items()),
+  )

tests/test_cli.py

@@ -551,6 +551,57 @@ def test_benchmark_run_without_xprof_timing(self, mock_run_benchmarks):
     self.assertFalse(kwargs["xprof_timing"])
     self.assertEqual(kwargs["xprof_dir"], "/tmp/tensorboard")
 
+  @parameterized.named_parameters(
+      ("flag_set", ["--single_host_tpu_init"], True),
+      ("flag_omitted", [], False),
+  )
+  @mock.patch.object(runner, "run_benchmarks")
+  def test_benchmark_run_single_host_tpu_init(
+      self, extra_flags, expected, mock_run_benchmarks
+  ):
+    """Verifies that `benchmark run` forwards --single_host_tpu_init."""
+    cli.run(
+        ["benchmark", "run", "gemm", "-m", "128", "-n", "128", "-k", "128"]
+        + extra_flags
+    )
+    mock_run_benchmarks.assert_called_once()
+    _, kwargs = mock_run_benchmarks.call_args
+    self.assertIs(kwargs["single_host_tpu_init"], expected)
+
+  @parameterized.named_parameters(
+      ("yaml_true", "  single_host_tpu_init: true\n", [], True),
+      ("yaml_false", "  single_host_tpu_init: false\n", [], False),
+      ("omitted", "", [], False),
+      ("cli_flag", "", ["--single_host_tpu_init"], True),
+      (
+          "cli_flag_overrides_yaml_false",
+          "  single_host_tpu_init: false\n",
+          ["--single_host_tpu_init"],
+          True,
+      ),
+  )
+  @mock.patch.object(runner, "run_benchmarks")
+  def test_benchmark_run_config_single_host_tpu_init(
+      self, yaml_snippet, extra_flags, expected, mock_run_benchmarks
+  ):
+    """Verifies single_host_tpu_init resolution for benchmark run-config."""
+    config_yaml = (
+        "benchmark:\n"
+        "  name: gemm\n"
+        f"{yaml_snippet}"
+        "  params:\n"
+        "    warmup_tries: 1\n"
+        "    num_runs: 1\n"
+    )
+    fake_config = self.create_tempfile(content=config_yaml)
+    cli.run(["benchmark", "run-config", fake_config.full_path] + extra_flags)
+    mock_run_benchmarks.assert_called_once()
+    _, kwargs = mock_run_benchmarks.call_args
+    self.assertIs(kwargs["single_host_tpu_init"], expected)
+    # The root-level key must not leak into the benchmark Config dataclass.
+    _, task_config = kwargs["tasks"][0]
+    self.assertFalse(hasattr(task_config, "single_host_tpu_init"))
+
   @parameterized.named_parameters(
       (
           "yaml_xprof_true",