[simrankaurb]

Summary

Adds an opt‑in single_host_tpu_init (YAML key single_host_tpu_init: true / CLI --single_host_tpu_init, default off) that initializes the current host standalone inside a multi‑host slice, so single‑host benchmarks (gemm, hbm, …) can be used for node‑local diagnostics — e.g. by Cluster Health Scanner — without the other hosts of the slice participating. No existing behaviour changes when the flag is off (the matmul.py diff is pure additions). This brings the chs branch's single‑host init path (#142) to main.

What the flag does

core/single_host.py::configure_single_host_tpu(), called by the runner after op_flags.yaml is applied and before the first JAX device call:

• env: TPU_HOST_BOUNDS=1,1,1, TPU_WORKER_ID=0, TPU_WORKER_HOSTNAMES=localhost, TPU_SKIP_MDS_QUERY=true
• appends to LIBTPU_INIT_ARGS: --deepsea_hal_test_skip_slicebuilder=true --deepsea_hal_test_allow_multichip_skip_slicebuilder=true --xla_tpu_use_enhanced_launch_barrier=false (and flips an existing enhanced_launch_barrier=true)

Because the slice's ICI routing is left untouched, device‑to‑device transfers hang in this mode. GeneralizedGemmBenchmark.generate_inputs samples the operands on the default device and then calls jax.device_put(x, replicated_sharding), which broadcasts over ICI — so gemm hung for any matrix size, with and without xprof (Slow PjRt TPU operation detected on all local devices). When the flag is on, the operands (and optional scaling factors / accumulator) are instead generated independently on each local mesh device from a CPU‑pinned key and assembled with jax.make_array_from_single_device_arrays; the only transfers are host→device. Same sharding (P(None, None)) and same contents (same key everywhere). A key left on TPU:0 was verified to still hang — the per‑device device_put of the key is itself a D2D copy.

Plumbing mirrors xprof_timing: BenchmarkKey.SINGLE_HOST_TPU_INIT + bool validation in config.py, --single_host_tpu_init on benchmark run, pop/forward in run-config, runner.run_benchmarks(single_host_tpu_init=…). README updated.

Verification

tpu7x‑standard‑4t host of a 2x2x4 slice, jax 0.12.0 / libtpu 0.0.50, no manual environment set up in the pod:

case	result
YAML single_host_tpu_init: true, gemm 16384³ bf16 ×20, xprof_device_timing_aggregation: local_slowest_device	✅ 78 s · xprof p50 8.27 ms · 2127 TFLOP/s per chip (1064 / device) · 20/20 xprof runs
--single_host_tpu_init gemm 1024³	✅ 20 s
--single_host_tpu_init hbm copy, 1 G elems, device 0	✅ 2790 GB/s (single‑device path unaffected)
Control: same pod, no flag	❌ init fails: Invalid --deepsea_slice_builder_worker_addresses … Expected 4 worker addresses, got 1

CPU (4‑device host mesh): flagged and default input paths produce identical replicated arrays and bit‑identical GEMM output for the same seed (incl. scaling‑factor + accumulator inputs); single_host_tpu_init: "yes" is rejected with a clear ValueError; the CLI flag parses and is not forwarded into GemmParams.

Note: upstream tests/ import accelerator_microbenchmarks.tests, which is not resolvable in the OSS checkout, so they were not run here.

[gemini-code-assist]

Code Review

This pull request introduces a helper method _replicated_normal to generate replicated arrays of normal samples locally on each device, preventing device-to-device transfers that can hang when ICI routing is not configured. The generate_inputs method is updated to utilize this helper for input and scaling factor generation. The review feedback suggests several improvements: simplifying the creation of replicated_sharding with an empty PartitionSpec(), using the idiomatic self.mesh.local_devices property instead of manual filtering, and catching both RuntimeError and IndexError when accessing the CPU device to prevent potential crashes.

[gemini-code-assist]

We can simplify the creation of replicated_sharding by using an empty PartitionSpec(). In JAX, an empty PartitionSpec() represents a fully replicated sharding across all dimensions of the mesh, which is cleaner and matches the usage elsewhere in this file (e.g., line 524).

Suggested change

	replicated_sharding = jax.sharding.NamedSharding(
	self.mesh, jax.sharding.PartitionSpec(*([None] * len(shape)))
	)
	replicated_sharding = jax.sharding.NamedSharding(
	self.mesh, jax.sharding.PartitionSpec()
	)

[gemini-code-assist]

Instead of manually filtering self.mesh.devices.flat by process_index, we can use the built-in self.mesh.local_devices property. This is more idiomatic, cleaner, and less error-prone.

    local_devices = self.mesh.local_devices

[gemini-code-assist]

If the CPU backend is registered but has no devices, jax.devices("cpu") would return an empty list, and accessing [0] would raise an IndexError. It is safer to catch both RuntimeError and IndexError to ensure robust fallback behavior.

Suggested change

	try:
	key = jax.device_put(key, jax.devices("cpu")[0])
	except RuntimeError:
	pass # No CPU backend (JAX_PLATFORMS=tpu); fall back to device key.
	try:
	key = jax.device_put(key, jax.devices("cpu")[0])
	except (RuntimeError, IndexError):
	pass # No CPU backend (JAX_PLATFORMS=tpu); fall back to device key.

[gemini-code-assist]

Similarly, we should catch both RuntimeError and IndexError here to handle cases where the CPU backend is registered but has no devices.

Suggested change

	try:
	key = jax.device_put(key, jax.devices("cpu")[0])
	except RuntimeError:
	pass # No CPU backend (JAX_PLATFORMS=tpu); fall back to device key.
	try:
	key = jax.device_put(key, jax.devices("cpu")[0])
	except (RuntimeError, IndexError):
	pass # No CPU backend (JAX_PLATFORMS=tpu); fall back to device key.