jupyter笔记本用TPU pod v2-32

Question

我用下面的代码来训练我的模型在v3-8 TPU中，它是一个单一的设备TPU，它工作的很好，但是相同的代码不工作在TPU v2-32上。据我所知，v2-32是通过专用高速网络相互连接的TPU设备集群，那么如何调整代码使其在v2-32上工作呢？

tpu = tf.distribute.cluster_resolver.TPUClusterResolver(tpu="tpu-name", zone="us-central1-a", project="myproject")
tf.config.experimental_connect_to_cluster(tpu)
tf.tpu.experimental.initialize_tpu_system(tpu)
strategy = tf.distribute.TPUStrategy(tpu)

....... code to load data and define model ......

with strategy.scope(): 
    myModel = createModel()
    ....
    ....

它在v3-8单TPU上工作得很好，但在TPU v2-32上却给出了错误:失效参数.找不到相关的张量remote_handle .

如何在jupyter笔记本中配置v2-32作为TPU集群的策略？

Answer 1

要连接到TPU，需要将、TPU实例的名称作为TPUClusterResolver的tpu参数传递。我在2021年12月偶然发现了类似的错误，当时我使用了以下代码(它使用TPUClusterResolver.connect())：

resolver = tf.distribute.cluster_resolver.TPUClusterResolver.connect(
    tpu="grpc://[ip-address-of-the-node]:8470", 
    zone="europe-west4-a", 
    project="my-awesome-project")

并发现传递节点的gRPC地址(grpc://[ip-address-of-the-node]:8470)会导致InvalidArgumentError。

我找不到与这个问题有关的确切信息，但我发现如下：

1. 以前对Google文档的修订警告用户这种行为(尽管错误信息不同)：

对于单一设备培训，您可以指定

名称或IP地址，例如：grpc://1.2.3.4:8470。

对于TPU，您必须使用TPU名称，以便TensorFlow能够发现所有可用于培训分发的主机的IP地址。

为TPU Pod指定IP地址会产生以下错误：

ValueError: TPUConfig.num_shards is not set correctly ....

来自Google文档(rev. . 2020-03-01)，存档于： (最后访问: 2022-01-22)

(https://github.com/google-research/bert/issues/885#issuecomment-546218170)上报道了类似的问题

请注意，这只发生在TPU上；使用TPUv2-8或v3-8，我可以使用gRPC地址或节点名称访问节点。