Lightning Talk: Bringing Google’s Colossus to PyTorch: Rapid Stor... Ankita Luthra & Trinadh Kotturu

This talk discusses the integration of Google Cloud Storage file system with PyTorch, addressing the challenges of increasing model sizes and the associated storage costs, known as 'storage tax'. The speaker highlights the performance gains achieved by using Rapid Storage, which offers 15 TiB/s throughput and supports 20 million queries per second at sub-millisecond latency. They explain how the underlying architecture, Google Colossus, enhances storage efficiency through optimizations in metadata distribution and intelligent data placement. The integration of Rapid Storage into GCSFS provides significant improvements in random and sequential read/write operations, while requiring no changes to existing PyTorch code. The session concludes with an invitation to explore these advancements further at the Google booth.

Good morning everyone. Thank you for attending this session where we are going to talk about how we are bringing Google Cloud Storage file system to PyTorch. And for anyone who is not aware of Google Cloud Storage, it's the same infrastructure that powers the loud apps like YouTube, Google Search, and Google Photos. Myself I'm a senior product manager in Google Cloud Storage joined by Ankita Luthra. She's a

senior software engineer in Google Cloud Storage. It's incredible to see the amount of innovation that is happening in PyTorch. And as your models evolve to support various business challenges and multi-modalities, they're paying a tax a hidden tax called storage tax. What's happening is like I let me explain that storage tax with a specific example. Let me take Llama 4 Scout which is 17 billion active parameters even

though it's like 109 billion parameters and mixture of experts model. It's trained on 40 trillion tokens which is three times the training tokens for Llama 3.3 70 billion. So your training data size is increasing almost multiplying. And if you consider 12 bytes per the checkpoint size can go up to 1.4 terabytes. Even though it's 17 billion active parameters, the total number of parameters are like 109 billion.

So that the checkpoint size can be 1.4 with the GPU cluster reliability you might want to save the checkpoint as frequently as 5 minutes. So you need to store a checkpoint of 1.4 terabyte size in 5 minutes. And it supports a context length of 10 million tokens which is great because now you can put in your entire repo and then get your answers. But that also means

that the KV cache can run into And you wanted it to store the KV cache into remote storage mechanism like an object storage because the clusters cannot hold 1.4 or like the terabyte size of a KV cache in its cluster size, and you wanted to store it to object storage, and you wanted to retrieve it. Which means your latency should be very optimized in order to serve

the KV cache from remote storage object storage mechanism to your GPE clusters. So, you have two things that you need to solve simultaneously. One, the massive throughput because you wanted to run things at scale, and you wanted to get data faster. So, that is where we are Sorry. So, that's where like we are excited to talk about rapid storage, which solves this exact same problem. Where it

is able to support 15 TiB/s. Please note it's a capital B. And it is able to support 20 million QPS at sub-millisecond latency. sub-millisecond latency for random reads. If anybody of you worked on the storage systems, you understand like optimizing the random reads is a very hard problem. So, it's able to support a and up and writes. And it's the storage system that is optimized for your

AML workload. And when we thought about how do you bring this innovation, FSSpec was the obvious choice because it is the Pythonic file interface that is supported by the many frameworks that you are already using. Be it data training or checkpointing or inference. And it supports various file system implementations, including the GCS, the Google Cloud Storage one, GCSFS. by integrating rapid storage into GCSFS, we observed a

4.8 x improvement in both random reads and sequential reads. And then we are able to observe 2.8 x faster writes with rapid storage integration into GCS FS. And to talk about how we are able to achieve, I'm going to hand it out to Ankita. Uh thanks, Srinath. So, I'm going to talk about how Rapid is able to achieve such a high performance. The core system which is

driving this whole performance is Google Colossus. Colossus is the same system on which GCS stores all of its data, including the services like YouTube, Drive, etc. How Google Colossus is able to maintain such a such a high speed at such a massive level is due to its sophisticated architecture, as well as its optimization techniques. Some of those techniques include distributed metadata. So, the metadata, just like data,

is also distributed across server in Colossus. The storage is also intelligent inside Colossus. We use ML-based algorithms to identify the usage patterns and place the data according to those usage patterns. Uh Colossus also uses stateful protocol instead of HTTP stateless protocol, which is much higher performance than usual which is much higher performance giving the that performance to our reads Now, we are bringing all those Colossus APIs

directly to GCS in the form of Rapid buckets or Rapid storage. So, let's take a step back and see how the data is traveling between your GPUs and the storage. So, we load tensors or write tensors using PyTorch from GPUs, and PyTorch internally uses FS spec. FS spec is a popular Python library which is used to expose cloud storage systems as file systems as local file systems.

GCS FS is an implementation of FS spec, and it is used to to talk to Google Cloud Storage in PyTorch. So, how we optimized this entire communication? We use direct connectivity in rapid storage. Direct connectivity means we don't need any front-end proxy between FS spec and Google Cloud Storage. The The protocol is also now stateful and GRPC based instead of HTTP based, and it is much Uh

the data is also co-located in the same zone as your compute. So, the data late So, the network latency between data and the compute is also reduced using rapid buckets. We are also We have also optimized the on-disk format on the rapid buckets, which means now the appends are native and more easier to do in the in the cloud storage buckets. Rapid buckets also support hierarchical namespaces,

which means folder is now a first-class resource inside rapid buckets. That means the layout the data layout inside the rapid buckets is more closer to local file system. So, to bring all of these features and high performance of read and writes to PyTorch, we have integrated these things directly into GCSFS. Now, what these GCSFS integration includes is GCSFS now supports Google supported Google supported SDKs for rapid

and HNS buckets. We also introduced stateful read handles, which means the once the a connection is established between your GPUs between the PyTorch and the cloud storage, we can use the subsequent calls and we can use the subse- we can use the metadata from the subsequent call we can use the metadata and the subsequent calls can actually use that metadata, uh making the writes and reads much

faster, especially the random writes. This would help improve your training time by quite a lot. We also introduced a bucket type auto detection. So, now what type of bucket you're using, whether it is rapid or it is standard, is automatically detected inside GCSFS. All the folder APIs are also now supported as including and especially atomic rename folders, which will help reduce the checkpointing time. The appends are

now native to GCS. So, what does that essentially mean is till now appends were either not supported or to append we have to download the complete object, add the data, and upload the complete object again. But now with rapid buckets, the appends are automatically supported in the buckets and you can directly add the data directly to object in the storage. To intro- to add all these features

into the GCS, we contributed to open source repository of GCSFS. We contributed almost 30K lines to add all these features. A huge shout out to Martin, who is the main maintainer for FS spec from Anaconda, and he has been very helpful and collaborative with us during this whole effort. You can check out all the features, including all the benchmarks comparing standard bucket as well as rapid bucket,

on the GitHub repository. The numbers are there, the documentation, all the features are added in the open source repository. So, as a PyTorch developer, what changes you would have to do in your code to get all these features and to give your GPUs more efficiency? And the answer is no changes. So, the rapid buckets are integrated the all the heavy lifting for the rapid buckets is done

inside GCSFS and our SDKs. So, you don't need to make any code changes. Just create a bucket which is of type rapid, add the add the data in this bucket, and point your workload to use these buckets. Your data access will be much faster and your GPUs can be used and and your GPUs can be used more efficiently. They can use that saved time in processing. And

that's it from us. Please use the rapid buckets and give your GPUs a performance boost. If you have any questions or any discussions you want to do, we'll be available at the Google booth as well. Thank you so much.