在 Cloud TPU 上切换软件版本

概览

在 TPU 上运行的框架的软件版本必须与在本地虚拟机上运行的版本相匹配。您现在可以在运行中的 Cloud TPU 上切换此软件版本，而无需删除和重新创建 TPU。您还还可以将 Cloud TPU 配置为使用特定每夜版的软件框架。我们仍然建议您选择这些框架的受支持版本。

使用情况

切换版本的推荐方法是使用 cloud-tpu-client Python 库。

TensorFlow 的示例用法。

from cloud_tpu_client import Client
import argparse

parser = argparse.ArgumentParser()
parser.add_argument('--tpu-name',
                    type=str,
                    required=True,
                    help='Name of the TPU Instance')
parser.add_argument('--target-version',
                    type=str,
                    required=True,
                    help='target TPU Runtime version')

args = parser.parse_args()

c = Client(args.tpu_name)
c.configure_tpu_version(args.target_version, restart_type='ifNeeded')
c.wait_for_healthy()

这会将 Cloud TPU 配置为与在本地虚拟机上运行的 TensorFlow 版本匹配，包括官方版本以及过时的每夜版。

configure_tpu_version API 的 restart_type 参数定义切换版本时的 TPU 重启行为。选项包括“always”（默认）和“ifNeeded”。

• “always”可用于修复具有诸如 UNHEALTHY_TENSORFLOW, 状态或因之前运行中的资源泄露而返回内存不足 (OOM) 错误的 TPU。设置此选项后，即使未安装新的框架版本，TPU 也会重启。

• “ifNeeded”非常有用，因为如果运行时已经具有正确的版本时不会重新启动运行时，因此它不会给训练脚本增加任何显著的启动时间。设置此选项后，仅当 TPU 未安装正确的框架版本时，TPU 才会重启。

该库直接与 Cloud TPU 通信，因此该代码需要在同一网络中的虚拟机中运行。建议您在模型的其余部分的代码中运行此代码。

其他软件选项

TensorFlow 包含一个 tf.__version__ 字符串，这是配置正确版本的最简单方式。其他软件选项包括：

• PyTorch - pytorch-1.13、pytorch-nightly-dev20220930、pytorch-nightly
• Jax - tpu_driver、tpu_driver0.1-dev20200320、tpu_driver_nightly

例如，将 TPU 配置为使用 PyTorch 的最新每夜版运行。

from cloud_tpu_client import Client
c = Client()
c.configure_tpu_version('pytorch-nightly', restart_type='ifNeeded')
c.wait_for_healthy()