加速 AI：通过 GCSFS 和 Rapid Bucket 将 Google Colossus 引入 PyTorch

Source: Google Developers Blog

请提供您希望翻译的具体文本内容，我将为您翻译成简体中文。

挑战：保持 GPU 供给

随着模型规模的增长，数据加载和检查点保存往往成为训练的主要瓶颈。

训练模型的数据准备工作涉及从远程存储系统（如对象存储）获取并处理 TB 到 PB 级别的数据。

标准的基于 REST 的存储访问往往难以满足现代分布式训练对极高吞吐量和低延迟的要求，导致 GPU 资源被浪费。

Rapid Bucket：通过双向 gRPC 实现快速存储

我们的全新 Rapid Bucket 解决方案在专用区域桶中提供高性能对象存储。通过绕过传统的 REST API 并使用持久化的 gRPC 双向流，我们将支撑 YouTube 和 Google Search 的 Colossus——文件系统状态协议的强大能力，直接带入 PyTorch 生态系统。

关键性能指标

• 极高吞吐量： 15 + TiB/s 总吞吐量
• 超低延迟： 随机读取和追加写入 < 1 ms
• 高 QPS： 每秒 20 M+ 查询

Fsspec – PyTorch 的 Pythonic 文件接口

fsspec 是 PyTorch 生态系统中用于文件系统的通用 Pythonic 接口。它已经被用于：

• 数据准备： Dask、Pandas、Hugging Face Datasets、Ray Data
• 检查点： PyTorch Lightning、torch.distributed、Weights & Biases
• 推理： vLLM

fsspec 为许多不同的存储系统提供了各种后端实现，所有这些都可以在同一层中集成。这消除了为每个后端编写特定代码的需求。通过将 Rapid Storage 与 gcsfs（fsspec 的 Google Cloud Storage 实现）集成，开发者只需一次简单的 fsspec.open() 调用即可利用 Rapid 提供的速度提升——无需进行复杂的代码重写。

Under the Hood: Leveraging Colossus

为了在 Rapid Buckets 上实现性能提升，我们对整个数据路径进行了优化：

• 基于 Stateful gRPC 的流式传输 – 双向 gRPC 保持连接活跃，消除每次操作的开销，例如连接建立、身份验证和元数据交换。这使得在单个对象内进行多次读取或追加时能够实现高效的有状态数据交换。
• 直连路径 – Google Cloud Storage（GCS）Rapid Buckets 使用直接连接的 gRPC 双向流式 API（BidiReadObject、BidiWriteObject）。通过将客户端直接连接到底层的 Colossus 文件，我们实现了最大性能。非 Rapid 流量通常会经过额外的网络跳数，导致更高的读写延迟。更多细节请参阅博客文章 How the Colossus stateful protocol benefits Rapid storage。
• 同区共置 – 将数据存储在与计算资源相同的区域（例如 us-central1-a）可消除跨区延迟。在 Rapid Buckets 出现之前，区域性存储桶和计算（或加速器）可能位于不同的可用区，从而产生额外的延迟。
• 无感用户迁移 – 保持现有的 fsspec API 不变，同时将 Rapid Buckets 的所有内部流量从 HTTP 升级为双向 gRPC。通过在 gcsfs 中添加存储桶类型自动检测，PyTorch 及其他 fsspec 客户端能够自动使用 Rapid Buckets，零手动配置。

结果

一个包含 134 M 行（约 451 GB）的数据集被加载到 16 台 GKE 节点上，每台节点配备八块 A4 GPU。训练运行了 100 步，每 25 步保存一次检查点，使用 PyTorch Lightning。

我们对 总训练时间（包括数据加载延迟）进行了基准测试，观察到使用 Rapid Bucket 相比标准区域存储桶可获得 23 % 的性能提升。

微基准测试

微基准测试——测量 I/O 或资源使用等单个构建块的性能——证实了这些提升：

测试使用 16 MB 的 I/O 大小，在 48 个进程上运行。
更多细节请参见 GCSFS 性能基准测试。

示例

入门

在 Rapid Bucket 上开始使用 GCSFS 非常简单。您现有的代码和脚本保持不变——只需指向一个 Rapid Bucket 即可利用性能提升。

安装

Rapid Bucket 集成自版本 2026.3.0 起可用。

pip install gcsfs

代码示例：使用 GCS Rapid 读写

import gcsfs

# Initialize the filesystem
fs = gcsfs.GCSFileSystem()

# Writing to a Rapid bucket
with fs.open('my-zonal-rapid-bucket/data/checkpoint.pt', 'wb') as f:
    f.write(b"model data...")

# Appending to an existing object (Native Rapid feature)
with fs.open('my-zonal-rapid-bucket/data/checkpoint.pt', 'ab') as f:
    f.write(b"appended data...")