Testing(do-not-merge): Isolate GEMM to single chip to remove barrier-cores (v6e only)

[simrankaurb]

No description provided.

[gemini-code-assist]

Code Review

This pull request attempts to restrict the benchmarking execution to a subset of devices by hardcoding device_count = 1 in benchmark_gemm_throttling.py and slicing the devices list to [0:1] in benchmark_utils.py. The reviewer identified two critical issues with these changes: first, slicing to [0:1] selects only a single core, which will trigger an assertion error because a single TPU chip requires two cores (suggesting [0:2] instead); second, hardcoding device_count = 1 will underestimate the total FLOPs and throughput metrics by a factor of two for a single-chip configuration (suggesting device_count = 2 instead).

[gemini-code-assist]

Slicing the devices to [0:1] selects only a single device (core). However, a single TPU chip consists of two cores (as indicated by the assertion num_devices % 2 == 0 and the comment on SHARDING_ON_SINGLE_CHIP_WITH_M). Using [0:1] will cause an AssertionError whenever a single-chip sharding strategy is used because num_devices will be 1, which is not divisible by 2. To isolate to a single chip (2 cores), you should slice [0:2] instead.

Suggested change

	devices = jax.local_devices()[0:1] if local_mesh else jax.devices()[0:1]
	devices = jax.local_devices()[0:2] if local_mesh else jax.devices()[0:2]

[gemini-code-assist]

If the GEMM is isolated to a single chip (which contains 2 devices/cores), the mesh created by create_mesh will contain 2 devices. Hardcoding device_count = 1 here will cause total_flops_all_devices to be underestimated by a factor of 2, leading to incorrect total throughput metrics. Consider setting device_count = 2 to match the single-chip configuration.

Suggested change

	device_count = 1
	device_count = 2