在 TPU 虚拟机上运行 AI/机器学习工作负载的存储最佳实践
为了最大限度地提高 TPU 虚拟机上 AI/ML 工作负载的性能和成本效益，请为您的工作负载选择并配置合适的存储解决方案。通过消除 I/O 瓶颈，您可以减少 TPU 加速器的空闲时间，从而缩短训练时间并降低训练费用。
本文档针对 TPU 虚拟机上的训练、设置检查点、提供服务和缓存，提供了特定于工作负载的存储建议和优化最佳实践。在应用这些实践之前，请先查看可用的 TPU 数据存储选项。本文档假定您熟悉 TPU 虚拟机，并且具有配置 Cloud Storage 资源的基本经验。
特定于工作负载的指导
下表按偏好顺序列出了针对不同工作负载的存储建议：
| 工作负载 | 建议 | 优化和工具（如适用） |
|---|---|---|
| 训练数据集，包括数据准备 |
|
|
| 检查点和强化学习权重 |
|
对于 Rapid Bucket，请使用 gcsfusecsi-checkpointing 配置文件。
|
| 模型存储和下载 |
|
对于下载模型，请使用 GKE Run:ai Model Streamer 或 Cloud Storage FUSE，并使用 gcsfusecsi-serving 配置文件通过单独的装载点进行装载。
|
| 键值对 (KV) 缓存分流 |
|
Cloud Storage 优化
以下部分介绍了将 Cloud Storage 与 TPU 虚拟机搭配使用时优化性能的最佳实践。
启用分层命名空间以优化元数据
为了提高元数据性能，请在为 AI/机器学习工作负载创建区域存储分区时启用分层命名空间。元数据性能是指 Cloud Storage 处理涉及查找、列出或修改对象路径和文件夹的操作的速度，而不是读取或写入文件内容本身的速度。
在未启用分层命名空间的存储分区中，文件夹并非实际资源，而是由以正斜杠 (/) 分隔的对象名称前缀表示的模拟文件夹。这会导致列出目录内容或重命名目录等操作变慢，因为系统必须扫描具有相应前缀的所有对象。分层命名空间提供真正的文件系统结构，这对于 AI/机器学习工作负载非常重要，原因如下：
- 原子目录重命名：机器学习框架使用目录重命名来完成检查点。分层命名空间支持原子重命名，可确保快速完成检查点。
- 初始 QPS 较高：与未启用分层命名空间的存储分区相比，分层命名空间支持的初始每秒查询次数 (QPS) 最高可提高 8 倍。这样可以防止多个 TPU 同时访问存储空间时出现瓶颈。
- 高效的文件夹级操作：在特定目录中查找和列出文件的速度更快，从而缩短训练和数据加载期间的响应时间。
通过 Rapid Bucket 提供的可用区级存储分区默认使用分层命名空间。如需了解详情，请参阅分层命名空间概览。
将 Cloud Storage FUSE 与适当的分析结果搭配使用
Cloud Storage FUSE 是一种 FUSE 适配器，可让您将存储分区作为本地文件系统装载。使用 Google Kubernetes Engine 时，我们建议使用 Cloud Storage FUSE CSI 驱动程序和 Cloud Storage FUSE 配置文件来自动进行性能调优。
如需详细了解使用 Cloud Storage FUSE 的最佳实践，请参阅性能调优最佳实践。
自定义 TPU 虚拟机启动磁盘
您可以使用启动脚本或创建自定义映像来自定义 TPU 虚拟机的客机操作系统环境。自定义启动磁盘对于以下情况非常有用：
- 预加载软件和库：安装特定的机器学习框架、依赖项或自定义软件，以缩短虚拟机启动时间并确保环境一致。
- 使用非标准操作系统分发版本：使用未包含在 Google 管理的列表中的操作系统分发版本或版本。
- 应用安全和监控配置：应用自定义安全设置、安装监控代理或设置环境变量。
不过，TPU 虚拟机的启动磁盘恢复功能有限。您无法分离或拍摄启动磁盘的快照以进行离线修复，因此在进行会影响启动过程的更改时请务必小心。遵循这些最佳实践有助于降低自定义 TPU 虚拟机环境时出现启动失败的风险。
自定义启动磁盘时，请谨记以下关键原则：
尽量减少对启动磁盘的修改：尽可能在 Persistent Disk 或 Hyperdisk 卷上安装应用和存储数据，而不是对启动磁盘进行大量修改。
使用 UUID 进行装载：向
/etc/fstab文件添加条目时，请始终使用 UUID 来标识磁盘和分区 (UUID=...)，而不是使用/dev/sdb1等设备名称。无法保证自动生成的设备名称在重新启动后保持稳定。
在进行系统更改时，请遵循以下建议，以降低启动失败的风险：
错误处理：在脚本中实现稳健的错误检查和优雅的失败模式。将详细消息记录到串行控制台和 Cloud Logging，以帮助进行调试。
关键依赖项：修改对启动至关重要的文件（例如
/etc/fstab文件、网络配置或引导加载程序设置）时，请务必格外小心。语法错误或输入错误可能会导致虚拟机无法启动。辅助磁盘：如果脚本依赖于辅助磁盘，请确保脚本能够处理磁盘可能不存在或附加时间比预期长的情况。除非绝对必要，否则请避免使启动过程严重依赖于辅助磁盘装载。
以下示例展示了推荐和不推荐的
/etc/fstab条目（用于挂载辅助磁盘）：- 推荐：
UUID=a1b2c3d4-e5f6-7890-1234-567890abcdef /mnt/mydata ext4 defaults,nofail 0 2 - 不推荐：
/dev/sdb1 /mnt/mydata ext4 defaults 0 2
使用
nofail选项可以防止系统在找不到磁盘时停止运行，但请确保您的应用可以处理装载点不可用的情况。- 推荐：
软件包管理：添加第三方代码库时请务必谨慎。 确保它们可信且与基础操作系统映像兼容。了解您安装的任何软件包的依赖项及其对系统库的潜在影响。
磁盘空间：监控启动磁盘使用情况。过多的日志记录或大型软件安装可能会填满启动磁盘，导致虚拟机无法启动。
日志记录：将应用和脚本配置为详细记录到串行控制台，因为这是诊断 TPU 虚拟机启动问题的主要工具。
规划存储空间容量
请务必规划工作负载所需的存储空间容量，以便充分利用加速器。这包括存储容量和检查点带宽。
估算存储空间
您可以将以下估算值作为存储需求的起点：
| 工作负载类型 | 数据集存储 | 检查点存储 |
|---|---|---|
| LLM 预训练 | 每个 TPU 2 TB | 每个 TPU 200 GB |
| 多模态训练 | 每个 TPU 12 TB | 每个 TPU 1 TB |
| 推理 | 每个 TPU 1 TB | 每个 TPU 1 GB |
估算检查点带宽
您可以使用以下公式估算训练工作负载所需的最低检查点带宽。对于数据读取、多次训练运行或训练与推理，请按比例增加估计的带宽要求。
- 检查点大小：参数数量 × 每个参数的字节数（对于 FP16 + 优化器状态，每个参数大约为 12-16 字节）。为优化器状态和不同精度添加了缓冲区（大约 3 倍）。
- 检查点间隔：保存检查点的频率（例如每 15 分钟一次）。
- 所需带宽：检查点大小 ÷ 检查点间隔。
以下示例展示了如何估算 Qwen3-72B 的最低检查点带宽：
- 检查点大小：720 亿个形参 × 12 字节 ≈ 每个检查点 864 GB。含缓冲区，3 × 864 GB ≈ 2.5 TB。
- 检查点间隔：2 分钟 = 120 秒。
- 所需带宽：2.5 TB ÷ 120 秒 ≈ 20 GBps。
参考方案
如需查看特定硬件和工作负载的存储配置示例，请参阅以下方案：
- TPU7x 推理：
- TPU7x 训练：
配额和带宽限制
Cloud Storage 和 Compute Engine 产品的带宽受默认配额限制。如果您超出配额，输入和输出请求可能会受到限制。
如需了解 Cloud Storage 配额以及如何申请增加配额，请参阅 Cloud Storage 文档中的配额和限制。如需了解 Hyperdisk 和 Persistent Disk 的 Compute Engine 配额，请参阅磁盘配额。