使用 Pax 在单主机 TPU 上训练

创建 TPU 虚拟机

借助 Cloud TPU 虚拟机，您的模型和代码直接在 TPU 虚拟机上运行。您通过 SSH 直接连接到 TPU 虚拟机。您可以直接在 TPU 虚拟机上运行任意代码、安装软件包、查看日志和调试代码。

从安装了 Google Cloud CLI 的 Cloud Shell 或计算机终端运行以下命令，创建 TPU 虚拟机。

根据合同中的可用性设置 zone，如果需要，请参阅 TPU 区域和可用区。

将 accelerator-type 变量设置为 v2-8、v3-8 或 v4-8。

对于 v2 和 v3 TPU 版本，将 version 变量设置为 tpu-vm-base；对于 v4 TPU，将变量设置为 tpu-vm-v4-base。

$ gcloud compute tpus tpu-vm create tpu-name \
--zone zone \
--accelerator-type accelerator-type \
--version version

连接到您的 Google Cloud TPU 虚拟机

使用以下命令通过 SSH 连接到您的 TPU 虚拟机：

$ gcloud compute tpus tpu-vm ssh tpu-name --zone zone

登录虚拟机后，shell 提示符会从 username@projectname 更改为 username@vm-name：

在 Google Cloud TPU 虚拟机上安装 Pax

使用以下命令在 TPU 虚拟机上安装 Pax、JAX 和 libtpu：

(vm)$ python3 -m pip install -U pip \
python3 -m pip install paxml jax[tpu]
-f https://storage.googleapis.com/jax-releases/libtpu_releases.html

系统检查

通过检查 JAX 是否看到 TPU 核心来测试所有内容是否已正确安装：

(vm)$ python3 -c "import jax; print(jax.device_count())"

此时系统会显示 TPU 核心的数量，如果您使用的是 v2-8 或 v3-8，该值应为 8；如果您使用的是 v4-8，该值应为 4。

在 TPU 虚拟机上运行 Pax 代码

您现在可以运行所需的任何 Pax 代码。如需开始在 Pax 中运行模型，可参阅 lm_cloud 示例。例如，以下命令会根据合成数据训练一个基于 2B 参数转换器的 SPMD 语言模型。

以下命令显示了 SPMD 语言模型的训练输出。它可以在大约 20 分钟内训练 300 步。

(vm)$ python3 .local/lib/python3.8/site-packages/paxml/main.py  --exp=tasks.lm.params.lm_cloud.LmCloudSpmd2BLimitSteps --job_log_dir=job_log_dir

在 v4-8 切片上，输出应包含：

损失和单步用时

summary 张量，位于 step=step_# loss = loss
summary 张量，位于 step=step_# 步/秒 x