GCP TPU Pod v3.32初始化TPUStrategy时核心数不一致错误求助

解决方案 #

修改TPUClusterResolver初始化逻辑 #

问题根源在于TensorFlow同时识别到了本地localhost设备和TPU Pod集群的worker设备，导致策略误判主机核心数不一致。通过明确指定TPU名称和job_name为worker，可以避免设备冲突：

import tensorflow as tf
import os

print("Tensorflow version " + tf.__version__)

# 明确指定TPU名称与集群job角色，过滤本地设备干扰
tpu_name = os.environ['TPU_NAME']
cluster_resolver = tf.distribute.cluster_resolver.TPUClusterResolver(tpu=tpu_name, job_name='worker')
print('Running on TPU ', cluster_resolver.cluster_spec().as_dict()['worker'])

tf.config.experimental_connect_to_cluster(cluster_resolver)
tf.tpu.experimental.initialize_tpu_system(cluster_resolver)
strategy = tf.distribute.TPUStrategy(cluster_resolver)

@tf.function
def add_fn(x,y):
  z = x + y
  return z

x = tf.constant(1.)
y = tf.constant(1.)
z = strategy.run(add_fn, args=(x,y))
print(z)

执行命令保持不变：

TPU_NAME=tpu-test-pod python3 tpu-test.py

错误原因分析 #

从提供的设备列表可见，TensorFlow同时枚举了/job:localhost和/job:worker下的TPU设备。尽管每个主机的TPU核心数均为8，但由于所属job标识不同，TPUStrategy会将其判定为不同主机的核心数不一致，触发报错。指定job_name='worker'后，ClusterResolver会仅连接TPU Pod的集群设备，排除本地设备的干扰。

额外排查建议 #

若上述方案无效，可尝试以下操作：

• 验证所有TPU VM节点的TensorFlow版本一致性（使用tpu-vm-tf-2.11.0-pod镜像默认保证版本统一）
• 通过gcloud compute tpus tpu-vm describe tpu-test-pod确认TPU Pod的运行状态与网络配置正常
• 重启所有TPU VM节点后重新测试

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