TPU切片中核心间互连（ICI）的检测方法问询

启动TPU Pod时ICI是否自动启用？ #

TPU Pod启动时默认自动启用ICI，无需借助Ray Cluster或其他第三方工具。ICI是TPU Pod架构的原生核心组件，Google Cloud在部署TPU Pod时会自动配置好跨主机TPU节点间的ICI通信链路，用户不需要手动进行额外配置或激活操作。

v2-32切片下用Jax脚本检测ICI状态？ #

可以通过Jax的内置API直接检测，以下是实用的检测方法和示例代码：

方法1：对比全局与本地设备数 #

v2-32切片由4个主机组成，每个主机包含8个TPU核心。如果跨主机的ICI正常工作，全局设备数应为32，本地设备数为8。通过以下代码可以验证：

import jax

global_device_count = jax.device_count()
local_device_count = jax.local_device_count()

print(f"全局TPU设备总数: {global_device_count}")
print(f"当前主机TPU设备数: {local_device_count}")

if global_device_count > local_device_count:
    print("ICI跨主机通信已正常启用")
else:
    print("仅本地TPU设备运行，未启用跨主机ICI")

方法2：查看设备拓扑信息 #

通过遍历Jax设备对象，可以查看每个设备所属的主机索引，确认跨主机设备是否存在：

import jax

for idx, device in enumerate(jax.devices()):
    print(f"设备{idx}: ID={device.id}, 所属主机索引={device.process_index}")

如果输出中包含多个不同的process_index值，说明跨主机的ICI连接已生效。

方法3：跨设备通信测试 #

用jax.pmap执行简单的跨设备广播操作，验证通信是否正常：

import jax
import jax.numpy as jnp

@jax.pmap
def broadcast_test(x):
    return jax.lax.broadcast(x, (jax.local_device_count(),))

# 初始化测试数据
test_data = jnp.array([1.0, 2.0])
# 执行跨设备操作
result = broadcast_test(test_data)

print("跨设备广播结果:", result)

如果能正常返回所有设备的广播结果，说明ICI通信链路正常。

内容的提问来源于stack exchange，提问作者Krishna Mohan