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:

1. 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.
2. 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.
3. 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.
4. No-Op User Migration: Preserved the existing fsspec API while entirely upgrading internal traffic from HTTP to BiDi-gRPC for Rapid buckets. By adding bucket-type auto-detection to gcsfs, PyTorch and other fsspec clients 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.

Get started

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.

Code 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...")

Python