Cloud TPU 多切片概览

Cloud TPU 多切片是一项全栈性能扩缩技术，可让训练作业在单个切片中或在多个 Pod 中的切片上使用多个 TPU 切片，并采用标准数据并行处理。使用 TPU v4 芯片时，这意味着训练作业可以在单次运行中使用超过 4096 个芯片。对于需要的芯片数少于 4096 个的训练作业，单个切片可以提供最佳性能。不过，多个较小的切片更容易获得，因此当将多切片与较小的切片搭配使用时，启动时间会更短。

多个切片线性扩缩性能

在多切片配置中部署时，每个切片中的 TPU 芯片通过芯片间互连 (ICI) 进行通信。不同切片中的 TPU 芯片通过将数据传输到 CPU（主机）来进行通信，而 CPU 又会通过数据中心网络 (DCN) 传输数据。如需详细了解如何使用多切片进行扩容，请参阅如何使用多切片将 AI 训练扩容到多达数万个 Cloud TPU 芯片。

多切片数据流

开发者无需编写代码即可实现芯片间 DCN 通信。XLA 编译器会为您生成该代码，并将通信与计算重叠，以实现最佳性能。

概念

加速器类型
构成多切片的每个 TPU 切片的形状。多切片请求中的每个切片都采用相同的加速器类型。加速器类型由 TPU 类型（v4 或更高版本）和紧随其后的 TensorCore 数量组成。例如，v5litepod-128 指定一个具有 128 个 TensorCore 的 TPU v5e。
自动修复
当切片遇到维护事件、抢占或硬件故障时，Cloud TPU 会创建新的切片。如果资源不足以创建新的切片，则在有可用硬件之前，创建操作将不会完成。创建新的切片后，多切片环境中的所有其他切片都会重启，以便继续训练。通过正确配置的启动脚本，训练脚本无需用户干预即可自动重新启动，并从最新的检查点加载和恢复。
数据中心网络 (DCN)
延迟时间较长、吞吐量较低的网络（与 ICI 相比），用于在多切片配置中连接 TPU 切片。
Gang 调度
同时预配所有 TPU 切片时，保证所有切片都成功预配或都未成功预配。
芯片间互连
用于在 TPU Pod 内连接 TPU 的高速、低延迟内部链接。
多切片
两个或更多个可通过 DCN 进行通信的 TPU 芯片切片。
节点
在多切片上下文中，节点是指单个 TPU 切片。多切片中的每个 TPU 切片都有一个节点 ID。
启动脚本
每次启动或重新启动虚拟机时运行的标准 Compute Engine 启动脚本。对于多切片，该值在 QR 创建请求中指定。如需详细了解 Cloud TPU 启动脚本，请参阅管理 TPU 资源。
Tensor
一种数据结构，用于在机器学习模型中表示多维数据。
Cloud TPU 容量的类型

您可以使用不同类型的容量创建 TPU（请参阅 TPU 定价方式中的“使用选项”部分）：

  • 预留：如需使用预留，您必须与 Google 签订预留协议。创建资源时，请使用 --reserved 标志。

  • Spot：使用 Spot 虚拟机定位抢占式配额。系统可能会抢占您的资源，以便为更高优先级作业的请求留出空间。创建资源时，请使用 --spot 标志。

  • 按需：定位按需配额，无需预留且不会被抢占。TPU 请求将排入 Cloud TPU 提供的按需配额队列，但无法保证有可用的资源。默认处于选中状态，无需标志。

开始使用

  1. 设置您的 Cloud TPU 环境。

  2. In the Google Cloud console, activate Cloud Shell.

    Activate Cloud Shell

    At the bottom of the Google Cloud console, a Cloud Shell session starts and displays a command-line prompt. Cloud Shell is a shell environment with the Google Cloud CLI already installed and with values already set for your current project. It can take a few seconds for the session to initialize.

    3. 如需使用多切片，您的 TPU 资源必须作为已排队的资源进行管理。

    入门示例

    本教程使用 MaxText GitHub 代码库中的代码。MaxText 是一种高性能、可任意扩缩、开源且经过充分测试的基本 LLM，采用 Python 和 Jax 编写。MaxText 旨在能够在 Cloud TPU 上高效训练。

    shardings.py 中的代码旨在帮助您开始尝试使用不同的并行处理选项。例如，数据并行处理、完全分片数据并行处理 (FSDP) 和张量并行处理。代码可从单切片环境扩容到多切片环境。

    ICI 并行处理

    ICI 是指用于连接单个切片中的 TPU 的高速互连。ICI 分片对应于切片内的分片。shardings.py 提供三个 ICI 并行处理参数：

    • ici_data_parallelism
    • ici_fsdp_parallelism
    • ici_tensor_parallelism

    您为这些参数指定的值决定了每个并行处理方法的分片数。

    这些输入必须受到限制，以便 ici_data_parallelism * ici_fsdp_parallelism * ici_tensor_parallelism 等于切片中的芯片数。

    下表展示了 v4-8 中可用的四个芯片的 ICI 并行处理的示例用户输入：

    ici_data_parallelism ici_fsdp_parallelism