使用 Pax 在单主机 TPU 上训练
准备工作
运行以下命令将 gcloud 配置为使用 您的 Cloud TPU 项目，并安装训练所需的组件 在单主机 TPU 上运行 Pax 的模型。
安装 Google Cloud CLI
Google Cloud CLI 包含用于交互 Google Cloud CLI 产品和服务。如果您尚未安装 请立即按照安装 Google Cloud CLI 中的说明进行安装。
配置 gcloud 命令
（运行 gcloud auth list 以查看可用的账号）。
$ gcloud config set account account$ gcloud config set project project-id
启用 Cloud TPU API
使用以下 gcloud 命令启用 Cloud TPU API 在 Cloud Shell 中运行。 （您也可以通过 Google Cloud 控制台。
$ gcloud services enable tpu.googleapis.com
运行以下命令以创建服务身份（服务账号）。
$ gcloud beta services identity create --service tpu.googleapis.com
创建 TPU 虚拟机
借助 Cloud TPU 虚拟机，您的模型和代码可以直接在 TPU 虚拟机上运行。您可以直接通过 SSH 连接到 TPU 虚拟机。你可以 代码、安装软件包、查看日志和调试代码。
从以下位置运行以下命令来创建 TPU 虚拟机： Cloud Shell 或您的计算机终端 安装了 Google Cloud CLI。
请根据合同中的空闲情况设置zone， 参考 TPU 区域和可用区 （如果需要的话）。
将 accelerator-type 变量设置为 v2-8、v3-8 或 v4-8。
将 version 变量设置为 tpu-vm-base 或 tpu-vm-v4-base（对于 v4 TPU）。
$ 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；如果您使用的是 v4-8，则为 4。
在 TPU 虚拟机上运行 Pax 代码
现在，您可以根据需要运行任何 Pax 代码。lm_cloud 样本 是开始在 Pax 中运行模型的绝佳起点。例如， 来训练 基于 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 切片上，输出应包括：
落败次数和单步用时
摘要张量，位于 step=* step_# * loss = loss
摘要张量，位于 step=* step_# * Steps per second x
文章来源：google cloud
推荐阅读
- 使用 Cloud TPU 进行 BERT 微调：句子和句对分类任务 (TF 2.x)
- 在 Cloud TPU Pod 上训练 PyTorch 模型
- 在 Cloud TPU 上训练 NCF (TF 2.x)
更多芯擎AI开发板干货请关注芯擎AI开发板专栏。欢迎添加极术小姐姐微信（id:aijishu20)加入技术交流群，请备注研究方向。