加速AI：通过GCSFS和Rapid Bucket将Google Colossus引入PyTorch
阅读原文· developers.googleblog.comGoogle 把自家 Colossus 存储架构直接接进 PyTorch，号称零代码改动提速 23%，做大规模训练的团队值得评估一下，但对大多数人来说这更像 GCP 的护城河加固。
Google Cloud推出了一项高性能集成方案，通过fsspec接口将Rapid Storage直接连接至PyTorch，以消除AI训练瓶颈。该方案利用Google的Colossus架构和双向gRPC流技术，可提供高达15 TiB/s的聚合吞吐量，并显著降低延迟。开发者仅需更新存储桶类型而无需修改代码，即可将总训练时间缩短23%。
Speeding Up AI: Bringing Google Colossus to PyTorch via GCSFS and Rapid Bucket
Today, we are announcing a major performance boost for AI/ML workloads using the PyTorch ecosystem on Google Cloud. By integrating Rapid Storage, powered by Google’s Colossus storage architecture, directly with PyTorch via the industry-standard fsspec interface, we are enabling researchers and developers to keep their GPUs busier than ever before.
The challenge: Keeping GPUs fed
As model sizes grow, data loading and checkpointing often become the primary bottlenecks in training. Data preparation activities to train models involve fetching and processing terabytes and petabytes of data from remote storage mechanisms like object storage. Standard REST-based storage access can struggle to meet the extreme throughput and low-latency requirements of modern distributed training, wasting valuable GPU resources.
Rapid Bucket: Rapid Storage via bi-di gRPC
Our new Rapid Bucket solution provides high-performance object storage in dedicated zonal buckets. By bypassing legacy REST APIs and utilizing persistent gRPC bidirectional streams, we’ve brought the power of Colossus, filesystem stateful protocols that power YouTube and Google Search, directly to the PyTorch ecosystem.
Key performance metrics of Rapid Storage
- Extreme Throughput: 15+ TiB/s aggregate throughput.
- Ultra-Low Latency: <1ms for random reads and append writes.
- High QPS: Rapid Bucket provides 20M+ QPS.
Fsspec - PyTorch’s Pythonic file interface
fsspec is the pervasive Pythonic interface for file systems in the PyTorch ecosystem. It is already used for:
- Data preparation: Dask, Pandas, Hugging Face Datasets, Ray Data
- Checkpoints: PyTorch Lightning, Torch.dist, Weights & Biases
- Inference: vLLM
There are various backend implementations of fsspec for many different storage systems, which can all be integrated under a single layer, eliminating the need to write specific code for each backend. By integrating Rapid Storage with gcsfs (the Google Cloud Storage implementation of fsspec), developers can leverage speed gains provided by Rapid with a simple fsspec.open() call — no complex code rewrites required.
Under the hood: Leveraging Colossus
To achieve a performance boost with Rapid Buckets, we optimized the entire data path:
- Stateful grpc-based streaming: gRPC bi-directional streaming keeps the connection alive, minimizing per-operation overhead like connection setup, auth, metadata etc., and enabling efficient, stateful data exchange for multiple reads or appends within a single object.
- Direct path: Google Cloud Storage(GCS) Rapid Bucket uses direct connectivity for its gRPC bi-directional streaming APIs (BidiReadObject, BidiWriteObject) to achieve maximum performance by connecting clients directly to underlying Colossus files. Non-Rapid traffic to GCS would typically have more network hops than direct paths, making read/write latencies over Rapid significantly lower. For more details, see Rapid storage internal working.
- Zonal co-location: By placing storage in the same zone as your compute (e.g.,
us-central1-a), we eliminate cross-zone latency. Prior to Rapid buckets, data in a regional bucket and compute(accelerators) can be in different zones and access the data induced latency. - No-Op User Migration: Preserved the existing
fsspecAPI while entirely upgrading internal traffic from HTTP to BiDi-gRPC for Rapid buckets. By adding bucket-type auto-detection to gcsfs, PyTorch and otherfsspecclients transparently utilize Rapid with zero manual configuration.
Results
A dataset of 134M rows totaling around 451GB was loaded onto 16 GKE nodes, each containing eight A4 GPUs. Training was conducted in 100 steps, with a checkpoint after every 25 steps using PyTorch Lightning. We benchmarked the performance of total training time, including the data load times, and we observed a performance gain of 23% using Rapid Bucket compared with Standard regional bucket.
Microbenchmarking — that is, measuring the performance of a building block like I/O or resource usage — confirms these gains. Throughput improved by 4.8x for reads (both sequential and random) and 2.8x for writes. These tests used 16MB IO sizes across 48 processes. You can find more details at GCSFS-performance-benchmarks.
Getting started with GCSFS on Rapid Bucket is easy. Your existing code and scripts remain the same. You just need to change the bucket to a Rapid Bucket to take advantage of the performance boost.
To install:
Rapid Bucket integration is available from version 2026.3.0.
pip install gcsfsCode sample to read/write from 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...")