TPU 配置
TPU v5e 配置
Cloud TPU v5e 是一款训练和推理（服务）产品。如需区分训练和推理环境，请在创建 GKE 节点池时将 AcceleratorType 或 AcceleratorConfig 标志与 TPU API 或 --machine-type 标志结合使用。
训练作业针对吞吐量和可用性进行了优化，而处理作业针对延迟时间进行了优化。因此，在为训练预配的 TPU 上执行的训练作业可能具有较低的可用性；同样，在为训练预配的 TPU 上执行的服务作业可能会具有较长的延迟时间。
您可以使用 AcceleratorType 指定要使用的 TensorCore 数量。使用 gcloud CLI 或 Google Cloud 控制台创建 TPU 时，您可以指定 AcceleratorType。您为 AcceleratorType 指定的值是一个字符串，格式如下：v$VERSION_NUMBER-$CHIP_COUNT。
您还可以使用 AcceleratorConfig 指定要使用的 TensorCore 数量。但是，由于 TPU v5e 没有自定义 2D 拓扑变体，因此使用 AcceleratorConfig 和 AcceleratorType 没有区别。
如需使用 AcceleratorConfig 配置 TPU v5e，请使用 --version 和 --topology 标志。将 --version 设置为您要使用的 TPU 版本，并将 --topology 设置为切片中 TPU 芯片的物理排列方式。您为 AcceleratorConfig 指定的值是一个格式为 AxB 的字符串，其中 A 和 B 是每个方向的条状标签计数。
v5e 支持以下 2D 切片形状：
| 拓扑 | TPU 芯片的数量 | 主机数量 |
| 1x1 | 1 | 1/8 |
| 2x2 | 4 | 1/2 |
| 2x4 | 8 | 1 |
| 4x4 | 16 | 2 |
| 4x8 | 32 | 4 |
| 8x8 | 64 | 8 |
| 8x16 | 128 | 16 |
| 16x16 | 256 | 32 |
v5e TPU 切片中的每个 TPU 虚拟机包含 1、4 或 8 个芯片。在 4 个芯片及更小的切片中，所有 TPU 芯片共享同一个非统一内存访问 (NUMA) 节点。
对于 8 芯片 v5e TPU 虚拟机，在 NUMA 分区内 CPU-TPU 通信更高效。例如，在下图中，CPU0-Chip0 通信比 CPU0-Chip4 通信快。
用于提供服务的 Cloud TPU v5e 类型
最多支持 8 个 v5e 芯片进行单主机传送。支持以下配置：1x1、2x2 和 2x4 切片。每个切片分别有 1、4 和 8 个芯片。
如需为服务作业预配 TPU，请在 CLI 或 API TPU 创建请求中使用以下加速器类型之一：
| AcceleratorType (TPU API) | 机器类型 (GKE API) |
|---|---|
v5litepod-1 |
ct5lp-hightpu-1t |
v5litepod-4 |
ct5lp-hightpu-4t |
v5litepod-8 |
ct5lp-hightpu-8t |
系统支持使用 Sax 在 8 个以上的 v5e 芯片上进行传送（也称为多主机传送）。如需了解详情，请参阅大语言模型服务。
用于训练的 Cloud TPU v5e 类型
最多支持对 256 个芯片进行训练。
如需为 v5e 训练作业预配 TPU，请在 CLI 或 API TPU 创建请求中使用以下加速器类型之一：
| AcceleratorType (TPU API) | 机器类型 (GKE API) | 拓扑 |
|---|---|---|
v5litepod-16 |
ct5lp-hightpu-4t |
4x4 |
v5litepod-32 |
ct5lp-hightpu-4t |
4x8 |
v5litepod-64 |
ct5lp-hightpu-4t |
8x8 |
v5litepod-128 |
ct5lp-hightpu-4t |
8x16 |
v5litepod-256 |
ct5lp-hightpu-4t |
16x16 |
v5e TPU 虚拟机类型比较：
| 虚拟机类型 | n2d-48-24-v5lite-tpu | n2d-192-112-v5lite-tpu | n2d-384-224-v5lite-tpu |
| v5e 芯片数量 | 1 | 4 | 8 |
| vCPU 的数量 | 24 | 112 | 224 |
| RAM (GB) | 48 | 192 | 384 |
| NUMA 个节点的数量 | 1 | 1 | 2 |
| 适用于 | v5litepod-1 | v5litepod-4 | v5litepod-8 |
| 中断 | 高 | 中 | 低 |
为了给需要更多芯片的工作负载腾出空间，调度器可能会抢占更少芯片的虚拟机。因此，8 芯片虚拟机可能会抢占 1 和 4 芯片虚拟机。
TPU v4 配置
TPU v4 Pod 由 4096 个芯片组成，通过可重新配置的高速链路互连。TPU v4 的灵活网络使您能够以多种方式连接相同大小的 Pod 切片中的芯片。创建 TPU Pod 切片时，您可以指定所需的 TPU 版本和 TPU 资源数量。创建 TPU v4 Pod 切片时，您可以通过以下两种方式之一指定其类型和大小：AcceleratorType 和 AccleratorConfig。
使用 AcceleratorType
如果未指定拓扑，请使用 AcceleratorType。如需使用 AcceleratorType 配置 v4 TPU，请在创建 TPU Pod 切片时使用 --accelerator-type 标志。将 --accelerator-type 设置为包含您要使用的 TPU 版本和 TensorCore 数量的字符串。例如，如需创建包含 32 个 TensorCore 的 v4 Pod 切片，您可以使用 --accelerator-type=v4-32。
以下命令使用 --accelerator-type 标志创建一个包含 512 个 TensorCore 的 v4 TPU Pod 切片：
$ gcloud compute tpus tpu-vm create tpu-name
--zone=zone
--accelerator-type=v4-512
--version=tpu-vm-tf-2.15.0-pod-pjrt
TPU 版本后面的数字 (v4) 指定 TensorCore 的数量。v4 TPU 中有两个 TensorCore，因此 TPU 芯片的数量将为 512/2 = 256。
使用 AcceleratorConfig
如果要自定义 TPU 切片的物理拓扑，请使用 AcceleratorConfig。对于超过 256 个芯片的 Pod 切片的性能调整，通常需要这样做。
如需使用 AcceleratorConfig 配置 v4 TPU，请使用 --version 和 --topology 标志。将 --version 设置为您要使用的 TPU 版本，并将 --topology 设置为 Pod 切片中 TPU 芯片的物理排列方式。
您可以使用 3 元组 AxBxC 指定 TPU 拓扑，其中 A<=B<=C 且 A、B、C 全部 <= 4 或全部为 4 的整数倍。值 A、B 和 C 分别是这三个维度中的条状标签计数。例如，如需创建一个包含 16 个芯片的 v4 Pod 切片，您需要设置 --version=v4 和 --topology=2x2x4。
以下命令将创建一个 v4 TPU Pod 切片，其中有 128 个 TPU 芯片以 4x4x8 数组排列：
$ gcloud compute tpus tpu-vm create tpu-name
--zone=zone
--type=v4
--topology=4x4x8
--version=tpu-vm-tf-2.15.0-pod-pjrt
2A=B=C 或 2A=2B=C 的拓扑也具有针对全对通信进行了优化的拓扑变体，例如 4×4×8、8×8×16 和 12×12×24。这类拓扑称为“扭曲鸟居”。
下图显示了一些常见的 TPU v4 拓扑。
较大的 Pod 切片可以通过一个或多个 4x4x4“立方体”芯片构建。
Twist Tori 拓扑
某些 v4 3D 环面切片形状可以选择使用所谓的扭曲环面。例如，可以将两个 v4 立方体排列为 4x4x8 切片或 4x4x8_twisted。扭曲拓扑提供的平分带宽明显更高。增加对分带宽对于使用全球通信模式的工作负载非常有用。扭曲拓扑可以提高大多数模型的性能，其中大型 TPU 嵌入工作负载受益最大。
对于将数据并行用作唯一并行策略的工作负载，扭曲拓扑的表现可能会略有提升。对于 LLM，使用扭曲拓扑的性能可能因并行类型（DP、MP 等）而异。最佳实践是在使用和不使用扭曲拓扑的情况下训练 LLM，以确定哪种拓扑可为模型提供最佳性能。FSDP MaxText 模型的一些实验使用扭曲的拓扑，实现了 1-2 项 MFU 方面的改进。
扭曲拓扑的主要优势在于，它将非对称环形拓扑（例如 4×4×8）转换为密切相关的对称拓扑。对称拓扑具有诸多优势：
- 改进了负载均衡
- 更高的对分带宽
- 缩短数据包路由
这些优势最终会转化为许多全球通信模式的性能的提升。
TPU 软件支持在切片上扭转 tori，其中每个维度的大小是最小维度的大小，或者是最小维度大小的两倍。例如，4x4x8、4x8x8 或 12x12x24。
例如，请考虑以下 4×2 环面拓扑，其中 TPU 在切片中带有 (X,Y) 坐标标签：
为清晰起见，此拓扑图中的边显示为无定向边。实际上，每个边缘都是 TPU 之间的双向连接。我们将此网格的一侧和另一侧之间的边称为环绕边，如图所示。
通过扭转此拓扑，我们最终得到一个完全对称的 4×2 扭曲环面拓扑：
上图与上图之间唯一的变化是 Y 环绕边。这些 TPU 已改为连接到具有 X+2 mod 4 坐标的 TPU，而不是连接到另一个具有相同 X 坐标的 TPU。
这种思路可以泛化到不同的维度大小和不同数量的维度。只要每个维度等于或等于最小维度大小的两倍，生成的网络就是对称的网络。
如需详细了解如何在创建 Cloud TPU 时指定扭曲的 tori 配置，请参阅使用 AcceleratorConfig。
下表显示了支持的扭曲拓扑，以及使用它们与无解拓扑相比在理论上增加了对分带宽。
| 扭曲拓扑 | 与非扭环面相比，带宽在理论上增加了 |
|---|---|
| 4×4×8_扭转 | 约 70% |
| 8x8x16 扭曲 | |
| 12×12×24_扭转 | |
| 4×8×8_扭转 | 约 40% |
| 8×16×16 扭曲 |
TPU v4 拓扑变体
某些包含相同数量芯片的拓扑可以以不同的方式排列。例如，可以使用以下拓扑配置包含 512 个芯片（1024 个 TensorCore）的 TPU Pod 切片：4x4x32、4x8x16 或 8x8x8。包含 2048 个芯片（4096 个 TensorCore）的 TPU Pod 切片可提供更多拓扑选项：4x4x128、4x8x64、4x16x32 和 8x16x16。包含 2048 个芯片（4096 个 TensorCore）的 TPU Pod 切片可提供更多拓扑选项：4x4x128、4x8x64、4x16x32 和 8x16x16。
与给定芯片数量关联的默认拓扑是最类似于立方体的拓扑（请参阅 v4 拓扑）。此形状可能是数据并行机器学习训练的最佳选择。其他拓扑对于具有多种并行处理（例如，模型和数据并行，或模拟的空间分区）的工作负载非常有用。如果拓扑与所用的并行性相匹配，则这些工作负载的性能最佳。例如，如果在 X 维度上并行放置 4 路模型，并在 Y 维度和 Z 维度上并行放置 256 路数据并行处理，则会与 4x16x16 拓扑匹配。
具有多个并行维度的模型在并行维度映射到 TPU 拓扑维度时效果最佳。这些模型通常是数据 + 模型并行大型语言模型 (LLM)。例如，对于拓扑为 8x16x16 的 TPU v4 Pod 切片，TPU 拓扑维度为 8、16 和 16。使用 8 路或 16 路模型并行处理（映射到其中一个物理 TPU 拓扑维度）性能更高。采用此拓扑时，四路模型的并行性不是最优，因为它与任何 TPU 拓扑维度均不对齐，但对于相同数量的芯片，采用 4x16x32 拓扑时效果最佳。
TPU v4 配置由两组组成，一组拓扑小于 64 个芯片（小型拓扑），另一组拓扑大于 64 个芯片（大型拓扑）。
小型 v4 拓扑
Cloud TPU 支持小于 64 个芯片的以下 TPU v4 Pod 切片（4x4x4 立方体）。您可以使用这些小型 v4 拓扑来创建其基于 TensorCore 的名称（例如 v4-32）或其拓扑（例如 2x2x4）：
| 名称（基于 TensorCore 数量） | 条状标签数量 | 拓扑 |
| v4-8 | 4 | 2x2x1 |
| v4-16 | 8 | 2x2x2 |
| v4-32 | 16 | 2x2x4 |
| v4-64 | 32 | 2x4x4 |
大型 v4 拓扑
TPU v4 Pod 切片以 64 个芯片为增量提供，形状在所有三个维度上都是 4 的倍数。这些维度也必须按递增顺序排列。下表显示了几个示例。其中一些拓扑是“自定义”拓扑，只能使用 --type 和 --topology 标志创建，因为排列芯片的方式有多种。
以下命令将创建一个 v4 TPU Pod 切片，其中有 512 个 TPU 芯片以 8x8x8 数组排列：
$ gcloud compute tpus tpu-vm create tpu-name
--zone=zone
--type=v4
--topology=8x8x8
--version=tpu-vm-tf-2.15.0-pod-pjrt
您可以使用 --accelerator-type 创建具有相同数量的 TensorCore 的 v4 TPU Pod 切片：
$ gcloud compute tpus tpu-vm create tpu-name
--zone=zone
--accelerator-type=v4-1024
--version=tpu-vm-tf-2.15.0-pod-pjrt
| 名称（基于 TensorCore 数量） | 条状标签数量 | 拓扑 |
| v4-128 | 64 | 4x4x4 |
| v4-256 | 128 | 4x4x8 |
| v4-512 | 256 | 4x8x8 |
不适用 - 必须使用 --type 和 --topology 标志 |
256 | 4x4x16 |
| v4-1024 | 512 | 8x8x8 |
| v4-1536 | 768 | 8x8x12 |
| v4-2048 | 1024 | 8x8x16 |
不适用 - 必须使用 --type 和 --topology 标志 |
1024 | 4x16x16 |
| v4-4096 | 2048 | 8x16x16 |
| … | … | … |
TPU v3 配置
TPU v3 Pod 由 1024 个芯片组成，通过高速链路互连。如需创建 TPU v3 设备或 Pod 切片，请针对 gcloud compute tpus tpu-vm 命令使用 --accelerator-type 标志。您可以通过指定 TPU 版本和 TPU 核心数来指定加速器类型。对于单个 v3 TPU，请使用 --accelerator-type=v3-8。对于具有 128 个 TensorCore 的 v3 Pod 切片，请使用 --accelerator-type=v3-128。
以下命令展示了如何创建具有 128 个 TensorCore 的 v3 TPU Pod 切片：
$ gcloud compute tpus tpu-vm create tpu-name
--zone=zone
--accelerator-type=v3-128
--version=tpu-vm-tf-2.15.0-pjrt
下表列出了支持的 v3 TPU 类型：
| TPU 版本 | 支持结束 |
|---|---|
| v3-8 | （结束日期尚未确定） |
| v3-32 | （结束日期尚未确定） |
| v3-128 | （结束日期尚未确定） |
| v3-256 | （结束日期尚未确定） |
| v3-512 | （结束日期尚未确定） |
| v3-1024 | （结束日期尚未确定） |
| v3-2048 | （结束日期尚未确定） |
如需详细了解如何管理 TPU，请参阅管理 TPU。如需详细了解不同版本的 Cloud TPU，请参阅系统架构。
TPU v2 配置
TPU v2 Pod 由 512 个芯片组成，通过可重新配置的高速链路互连。如需创建 TPU v2 Pod 切片，请对 gcloud compute tpus tpu-vm 命令使用 --accelerator-type 标志。您可以通过指定 TPU 版本和 TPU 核心数来指定加速器类型。对于单个 v2 TPU，请使用 --accelerator-type=v2-8。对于具有 128 个 TensorCore 的 v2 Pod 切片，请使用 --accelerator-type=v2-128。
以下命令展示了如何创建具有 128 个 TensorCore 的 v2 TPU Pod 切片：
$ gcloud compute tpus tpu-vm create tpu-name
--zone=zone
--accelerator-type=v2-128
--version=tpu-vm-tf-2.15.0-pjrt
如需详细了解如何管理 TPU，请参阅管理 TPU。如需详细了解不同版本的 Cloud TPU，请参阅系统架构。
下表列出了支持的 v2 TPU 类型
| TPU 版本 | 支持结束 |
|---|---|
| v2-8 | （结束日期尚未确定） |
| v2-32 | （结束日期尚未确定） |
| v2-128 | （结束日期尚未确定） |
| v2-256 | （结束日期尚未确定） |
| v2-512 | （结束日期尚未确定） |
TPU 类型兼容性
您可以将 TPU 类型更改为具有相同数量的 TensorCore 或芯片的其他 TPU 类型（例如，v3-128 和 v4-128），并运行您的训练脚本，而无需更改代码。但是，如果更改为具有更多或更少数量的 TensorCore 的 TPU 类型，将需要执行显著的调整和优化。如需了解详情，请参阅在 TPU Pod 上运行训练。
TPU 虚拟机软件版本
本部分介绍了具有 TPU 虚拟机架构的 TPU 应使用的 TPU 软件版本。如需了解 TPU 节点架构，请参阅 TPU 节点软件版本。
TPU 软件版本适用于 TensorFlow、PyTorch 和 JAX 框架。
TensorFlow
对于 TensorFlow 2.15.0 版，请使用与编写模型的 TensorFlow 版本相匹配的 TPU 软件版本。您还必须指定流执行器 (SE) 运行时或 PJRT 运行时。例如，如果您要将 TensorFlow 2.15.0 与 PJRT 运行时搭配使用，请使用 tpu-vm-tf-2.15.0-pjrt TPU 软件版本。
PJRT 具有自动设备内存碎片整理功能，可简化硬件与框架的集成。如需详细了解 PJRT，请参阅 Google 开源博客上的 PJRT：简化机器学习硬件和框架集成。
我们正努力将 TPU v2、v3 和 v4 的所有功能迁移到 PJRT 运行时。下表介绍了 PJRT 或 Steam Executor 当前支持的功能。
| 加速器 | 特征 | PJRT 支持 | 在流执行器上受支持 |
|---|---|---|---|
| TPU v2-v4 | 密集计算（无 TPU 嵌入 API） | 是 | 是 |
| TPU v2-v4 | Dense Compute API + TPU Embedding API | 否 | 是 |
| TPU v2-v4 | tf.summary/tf.print（在软设备上放置） |
否 | 是 |
| TPU v5e | 密集计算（无 TPU 嵌入 API） | 是 | 否 |
| TPU v5e | TPU 嵌入 API | 不适用 - TPU v5e 不支持 TPU embedding API | 不适用 |
TensorFlow 2.14.0 及更低版本仅支持流执行器。使用与编写模型时所用 TensorFlow 版本相匹配的 TPU 软件版本。例如，如果您使用的是 TensorFlow 2.14.0，请使用 tpu-vm-tf-2.14.0 TPU 软件版本。
目前支持的适用于 TPU 的 TensorFlow TPU 虚拟机软件版本为：
- tpu-vm-tf-2.15.0-pjrt
- tpu-vm-tf-2.15.0-se
- tpu-vm-tf-2.14.1
- tpu-vm-tf-2.14.0
- tpu-vm-tf-2.13.1
- tpu-vm-tf-2.13.0
- tpu-vm-tf-2.12.1
- tpu-vm-tf-2.12.0
- tpu-vm-tf-2.11.1
- tpu-vm-tf-2.11.0
- tpu-vm-tf-2.10.1
- tpu-vm-tf-2.10.0
- tpu-vm-tf-2.9.3
- tpu-vm-tf-2.9.1
- tpu-vm-tf-2.8.4
- tpu-vm-tf-2.8.3
- tpu-vm-tf-2.8.0
- tpu-vm-tf-2.7.4
- tpu-vm-tf-2.7.3
如果您使用的是 Pod 切片，请在 TensorFlow 版本号后面附加 -pod。例如 tpu-vm-tf-2.15.0-pod-pjrt。
如需详细了解 TensorFlow 补丁版本，请参阅支持的 TensorFlow 补丁版本。
TPU v4 与 TensorFlow 2.10.0 及更低版本
如果您使用 TensorFlow 在 TPU v4 上训练模型，TensorFlow 2.10.0 及更低版本将使用下表中显示的 v4 专用版本。如果表中未显示您使用的 TensorFlow 版本，请按照 TensorFlow 部分中的指导进行操作。
| TensorFlow 版本 | TPU 软件版本 |
|---|---|
| 2.10.0 | tpu-vm-tf-2.10.0-v4、tpu-vm-tf-2.10.0-pod-v4 |
| 2.9.3 | tpu-vm-tf-2.9.3-v4、tpu-vm-tf-2.9.3-pod-v4 |
| 2.9.2 | tpu-vm-tf-2.9.2-v4、tpu-vm-tf-2.9.2-pod-v4 |
| 2.9.1 | tpu-vm-tf-2.9.1-v4、tpu-vm-tf-2.9.1-pod-v4 |
Libtpu 版本
创建 TPU 虚拟机时预装了 TensorFlow 并预安装了相应的 Libtpu 库。如果您要创建自己的虚拟机映像，请指定以下 TensorFlow TPU 软件版本和相应的 libtpu 版本：
| TensorFlow 版本 | libtpu.so 版本 |
|---|---|
| 2.15.0 | 1.9.0 |
| 2.14.1 | 1.8.1 |
| 2.14.0 | 1.8.0 |
| 2.13.1 | 1.7.1 |
| 2.13.0 | 1.7.0 |
| 2.12.1 | 1.6.1 |
| 2.12.0 | 1.6.0 |
| 2.11.1 | 1.5.1 |
| 2.11.0 | 1.5.0 |
| 2.10.1 | 1.4.1 |
| 2.10.0 | 1.4.0 |
| 2.9.3 | 1.3.2 |
| 2.9.1 | 1.3.0 |
| 2.8.3 | 1.2.3 |
| 2.8.* | 1.2.0 |
| 2.7.3 | 1.1.2 |
PyTorch
使用与编写模型时所用的 PyTorch 版本匹配的 TPU 软件版本。例如，如果您使用的是 PyTorch 1.13 和 TPU v2 或 v3，请使用 tpu-vm-pt-1.13 TPU 软件版本。如果您使用的是 TPU v4，请使用 tpu-vm-v4-pt-1.13 TPU 软件版本。TPU Pod 使用相同的 TPU 软件版本（例如 v2-32、v3-128、v4-32）。当前支持的 TPU 软件版本为：
TPU v2/v3：
- tpu-vm-pt-2.0 (pytorch-2.0)
- tpu-vm-pt-1.13 (pytorch-1.13)
- tpu-vm-pt-1.12 (pytorch-1.12)
- tpu-vm-pt-1.11 (pytorch-1.11)
- tpu-vm-pt-1.10 (pytorch-1.10)
- v2-alpha 版 (pytorch-1.8.1)
TPU v4：
- tpu-vm-v4-pt-2.0 (pytorch-2.0)
- tpu-vm-v4-pt-1.13 (pytorch-1.13)
创建 TPU 虚拟机时，TPU 虚拟机上会预安装最新版本的 PyTorch。当您安装 PyTorch 时，系统会自动安装正确版本的 libtpu.so。
如需更改当前的 PyTorch 软件版本，请参阅更改 PyTorch 版本。
JAX
您必须在 TPU 虚拟机上手动安装 JAX，因为不存在特定于 JAX 的 TPU 软件版本。对于所有 TPU 版本，请使用 tpu-ubuntu2204-base。当您安装 JAX 时，系统会自动安装正确版本的 libtpu.so。
TPU 节点软件版本
本部分介绍了在 TPU 节点架构中应用于 TPU 时应使用的 TPU 软件版本。如需了解 TPU 虚拟机架构，请参阅 TPU 虚拟机软件版本。
TPU 软件版本适用于 TensorFlow、PyTorch 和 JAX 框架。
TensorFlow
使用与编写模型时所用 TensorFlow 版本相匹配的 TPU 软件版本。例如，如果您使用的是 TensorFlow 2.12.0，请使用 2.12.0 TPU 软件版本。TensorFlow 专用 TPU 软件版本如下：
- 2.12.1
- 2.12.0
- 2.11.1
- 2.11.0
- 2.10.1
- 2.10.0
- 2.9.3
- 2.9.1
- 2.8.4
- 2.8.2
- 2.7.3
如需详细了解 TensorFlow 补丁版本，请参阅支持的 TensorFlow 补丁版本。
创建 TPU 节点时，TPU 节点上会预安装最新版本的 TensorFlow。
PyTorch
使用与编写模型的 PyTorch 版本匹配的 TPU 软件版本。例如，如果您使用的是 PyTorch 1.9，请使用 pytorch-1.9 软件版本。
PyTorch 专用 TPU 软件版本如下：
- pytorch-2.0
- pytorch-1.13
- pytorch-1.12
- pytorch-1.11
- pytorch-1.10
- pytorch-1.9
- pytorch-1.8
- pytorch-1.7
pytorch-1.6
pytorch-nightly
创建 TPU 节点时，TPU 节点上会预安装最新版本的 PyTorch。
JAX
您必须在 TPU 虚拟机上手动安装 JAX，因此系统不会预安装特定于 JAX 的 TPU 软件版本。您可以使用列出的适用于 TensorFlow 的任何软件版本。