TorchStore: What We Learned Building Distributed Storage Sol... Lucas P, Danielle P, Allen W, Amir A

This talk introduces Torch Store, a solution for efficient weight synchronization in reinforcement learning (RL) workflows. The speakers discuss the challenges faced with existing methods of weight syncing, specifically the limitations of using S3 and distributed checkpointing, as well as handcrafted communication protocols. They describe how Torch Store combines easy APIs with the speed of custom approaches, allowing researchers to manage model weights between trainers and generators seamlessly. The session covers the architecture of Torch Store, its integration with other PyTorch components like Torch Monarch and Torch Titan, and the inherent flexibility it offers to researchers when defining their algorithms. Additionally, the speakers explore both synchronous and asynchronous methods of weight synchronization, highlighting the benefits of each and providing insights into their performance optimizations.

Hello, good afternoon. Bonjour. Uh, thank you for welcoming us here today. Thanks for your time. We're going to be talking about Torch Store. My name is Lucas. This is Air. This is Danielle. And I'll get right into it. So, I thought it would help to start in the the beginning. It's kind of a weird spot. I'm just going to assume you know a bunch of things. If

you don't, it's okay. We're going to uh do a little content review. But the story goes, we were hanging out in the office. We were looking at this this RL thing, and we were like, "Okay, this is probably going to be around for a while. We got to think of a solution for weights sync. And if you don't know what weightsync is, basically we have a set

of nodes we call trainers and a set of nodes we call generators. And they need to share a full model very quickly, very often, right? And what we basically saw was that people were doing one of two things. On the one hand here, you were doing the easy thing. You were using something like S3, something like distributed checkpoint, which was sharding aware, and you were saving to

some network file system. This was really easy. You could just save a model and then load it somewhere else. But the problem was is it's very slow. It's generally backed by persistent storage. The infrastructure actually wasn't built for uh persistent streaming or or or it was built for persistence and not for streaming. Um yeah, the other thing that people were doing which was very fast was building

webs of these handcrafted comps. Essentially say, okay, the data that I have over here needs to be over here. And at first people were mostly using collectives and it's kind of a nightmare. You need to order these collectives in order. And every time you make a change to your sharding uh algorithms or how you shard, you have to go back to the drawing board. Later people added

like RPC as like peer-to-peer type connections, RDMA, that sort of thing. And it got a little better, but it was still uh not pluggable. It was very brittle brittle. It wasn't a solution you could pick up and drop in other places and reuse. And we thought, hey, we could do better. So what we did is uh we founded Torch Door. We said we're going to take the

the things we love about the about option one, the slow distributed option, uh the easy APIs, the save anywhere, fetch anywhere semantic, and we're going to combine that with the speed and efficiency of handcrafted coms. We're going to hide it all away so the user never sees it. The researcher can focus on research and we're going to build it and call it torch store. And then a

couple months later, we arrived here today. So today, what we're going to do is we're going to walk you through the basics. If you didn't understand what I was talking about, um, we're going to have an awesome review on RL really quick, an awesome review on why we do weight sync. We're going to motivate the problem for you a little bit, and then we're going to walk

you through our integration and basically give you everything that we learned. And I hope that you all learned something and have something to take home as well. >> All right. Uh, yeah. So, I will talk a bit about reinforcement learning, uh, and how torch store fits in. Thank you. Um yeah, so in recent years RL has become a very critical part of LLM training. So you have

pre-training and SFT which are quite effective but RL really allows you to take the model training to the next level and unlock new capabilities. Um and there are many different flavors of RL uh but in general you always have the following three components. So you have a trainer uh which is taking a data set and running uh forward passes uh through the model and then calculating some

loss u and this can vary according to the algorithm you're using um and then updating the model weights uh after back propagating the loss. Um so this is the trainer and then we have the generator which also has a copy of the model weights. Um and what the generator is doing is sampling prompts from a data set and running inference on the model and uh collecting generations.

It can be one generation per prompt or multiple generations and then we're passing those generations to the greater uh and the greater is going to assign rewards to the generations. Um so this is how we teach the model which generations are good and which are bad. Uh and then we're going to pass all of this data. So the generations, the prompts and the rewards to the trainer

and then the trainer is going to train and the process continues. Uh but as I mentioned both the trainer and the generator are going to have a copy of the model weights and as Lucas mentioned it's very important that these weights stay in sync. Uh otherwise it's like you're trying to teach the model and it's learning but it's forgetting everything each time and you have to teach

it from scratch and it's not iterating on previous learnings. So this weight uh weights sync step is very important and also when we're working with large models we may be dealing with terabytes of data. So it's very important that the weights sync step happens efficiently and this is where torch store comes in. Uh yeah so in PyTorch we have several building blocks for RL. Uh so we

have torch monarch which is a library for orchestration of distributed actors. Uh so this is very useful in RL because you have different actors like a trainer, a generator, a grader that you want running different code on different machines and Torchmon Monarch allows you to orchestrate these all um and have them run on separate nodes. Uh also Torch Monarch offers an abstraction on top of RDMA, remote

direct memory access, which we're going to be talking about more later. Uh oh, there's still more. Uh so yeah then there's torch store which we're going to be talking about. It's a distributed tensor store for weight sync and then we also have torch titan which is a library for pre and post training and we're also building out RL training uh recently. Next slide. Thanks. Uh yeah so

in torch titan we're currently building out RL and this is a good place to see an example of these building blocks working together. Uh so we're using torch monarch where we have a single controller which orchestrates the trainer and the generator. Then we have the trainer using uh components from torch titan for training and the generator is calling into vlm. Uh and then one unique thing here

as well is we have what we call the unified model definition. So we have a single model implementation between training and inference as compared to some other RL frameworks that don't have this. They may have uh different model implementation for training and inference. And this can lead to all sorts of pain points and subtle bugs because the uh model on the generator and the model on the

trainer uh may generate uh may have different uh forward passes and have different outputs uh logits and log props. Um however in the unified model we have we can have bitwise identity between the trainer and generator log props. Additionally, for weight sync, the unified model uh makes it easier because if you have different model implementations, you often have to do a lot of state dict uh remapping

and conversion, but you don't have to do this uh with the unified model. And then of course, it's a good showcase of a Yeah. Uh yeah. So now we'll pass it back over to Lucas to talk about async RL. >> So there's two basic ways of doing sync and async. um sync and async weight sync it's it's hard we're I'm going to work on the workshop the

name of the the slide over here but anyway so we had a lot of questions on like where we would focus and what basically we decided was that it wasn't our job to pick the right solution uh the it wasn't our job to pick the right algorithm for the researcher it's our job to abstract away the infrastructure so the researcher can define the requirements themselves and we

should give them the flexibility to uh to build what they want to build without having to worry about the underlying infrastructure. So let me start with sync weight sync. Um so this is how you would naively approach this problem. If you just had to build it and you told cla to build it, maybe this is the first thing that it would do. Um if you have a

small amount of nodes, essentially what you're going to do is your trainer is going to train like we mentioned and then it's going to initiate a direct GPU to trans transfer from your trainers to your generators, right? And so this works really well. It can even be pretty fast depending on your infrastructure. But as you can see from the slide here, like the big uh cost is

that your entire training stack and potentially your entire cluster is blocked during this weight sync uh time span, right? So this is all blocked simultaneously. And so you might make that decision. You might start there because it's easy, it's simple, you don't have to manage any any other uh additional buffers or anything like that. But you might also come to the conclusion that you can't afford to

do that, right? you don't want a majority of your cluster blocked at one time or maybe you have more flexibility on your algorithm. So on the left hand side here maybe if you're doing synchronous training where the weights have to match exactly at all times between um your generators and your trainers then you don't really have an option. You're going to basically do this anyway. But if

you're allowing yourself to explore things like off policyiness, which is a decision the researcher might make, which allows you basically to say, "Okay, it's okay for my trainer or my generator to fall behind the exact version I'm at." Then you now have this opportunity to unblock a majority of your cluster. So what we do is during training, we're going to do that same save, but we're going

to overlap it with um your forward and your backward pass. Right? So the intuition is that your model weights aren't changing during the forward and backward pass. So if you're very quick about saving and you save to a local CPU buffer, you can actually overlap it completely. So if you're familiar with DCP or distributed checkpointing, you'll you might have recognized the blog post. We called it zero

overhead checkpointing. Uh I really hate that we called it that. I wish that we called it a reasonable amount of overhead checkpointing because there's always something someone's going to find it and it'll be there. But that's basically what we do here. And then on the generation side, we're going to do the reverse, right? So what we're basically going to do is during generation I'm going to fetch

my weights locally so that when I'm ready to update them and move them from the CPU to the GPU they're already local right and the benefit here is that this is completely decoupled from a research perspective you can really push the boundaries of how you might orchestrate your algorithm. You could make a claim such that hey actually what I want is generators that are on different versions

of my model. Maybe I want some generators that are focused on multicycle long prompt iteration very challenging problems. Maybe they they're executing code environments and then coming back and other groups of generations that are doing other things and and whatever, right? So again, so the the idea here is that the core algorithm should be something that the researcher decides and we want to basically allow them to

make any decision they want based on the science and not whether or not they have the infrastructure and the time to go build it. Um so in review, right, why would you do async RL? Well, maybe it's a scaling issue, right? Maybe you just can just you're looking at your cluster efficiency numbers and you're like, ah, we actually can't afford to have the entire cluster go down.

Maybe I have a lot more generators than I have trainers and it just doesn't make sense every time I need to update that everybody needs to slow down and and sync up, right? You don't want that synchron to pay that synchronization cost. And then you might want something more decoupled and you might want something more flexible. Um, so we kind of mentioned uh, you know, sync async.

So I want to walk you through some of the primitives and how we basically arrive there. I'm going to fly through this slide because I think to really understand this, you know, you should look at the tutorial, run some jobs. But the basic takeaway from this slide is you can initialize a global store, right? TS.initialize. Um I'm going to talk a little bit about what mesh is

in the following slide, but the key idea here is you can place the data anywhere and fetch it from anywhere, right? And so when I'm placing the data here, if your model happens to be detensor, right? Then what we'll what we know uh what torch store does is it understands what shard of data you have locally. It saves that data and then when you fetch, even if

you're on a different sharding pattern, it also understands the data you have locally, the window of the tensors you have, and it fetches only that. And that always works across any number of nodes, across any sharding pattern. Um, right? So now that you know this is the part I kind of skipped over. So during initialization these are two of our most import primitives that you that it's

up to the user to define. We have uh a concept called a storage volume and then a strategy. And I'll I'll run through a quick tangent here. I think I have a minute. I remember being at work working on checkpointing and I I talked to an old mentor of mine and I was like, "Hey, I'm so bored. Like actually all we do is like there's some bits

over here and some bites over here and we move them over there." And he was like, "Yeah, that's like kind of the entire job." And so he was basically right. And these primitives, I think, are like basically a direct evolution of like that joke or me being tired of solving the same problem over and over again. So what we believe is if you define a storage volume,

which is just a process, but really what it represents is some space in memory somewhere, right? So it could be on GPU, it could be on CPU, it could be local, it could be remote. If you define a storage volume and you give me a basic rule for for where I should save data when I'm saving it. So you notice in the previous slide I just did

torch.putstate dictic and I didn't say putstate dictic here on these nodes right? So if you give me these two basic rules then torch store will figure out the rest for you. So we'll automatically detect if your data is local we'll use shim. If it's remote we'll use RDMA. If you don't have RDMA we'll fall back to TCP. Um and so those are the the basic building blocks.

Um the next most important thing is you could have the the most beautiful primitives in the world but if they're not fast no one's going to use them. So I'm going to pass it off to Amir here who's going to talk about how we made these things performant. >> Okay. So as Lucas described, Torch store is giving you the right primitives to design a data flow that

best suits your application. Um in the case of async RL, we worked a bit closer on the endto-end data design because the weight transfer path is essential to the overall application throughput. So I'm going to walk through one potential async RL flow so we can see what these design choices look like in practice. Uh so we're starting off here and the first step is the weight push.

So this sequence starts off with the trainer ranks issuing a torch store put state dicticked call. Uh and they do this with whatever shards they own. Uh because the training stack is detensor native. We can operate on the local shards directly rather than having to reshape them into some type of intermediary format to work with. Uh the design that you see here is intentionally similar to something

that you'd see in maybe like model checkpointing. And the goal is to avoid stopping training just to publish Uh so after a new version of the weight is committed and we publish it. Uh each GPU is going to be doing a non-blocking copy into CPU memory and that transfer will overlap with your training work. And uh you know we can also organize the memory here as circular

segments. Uh so this would allow us to pay the allocation cost once up front and then reuse those segments going forward. This is really important in something like RL where you're continuously doing uh you know updated model states and you don't want to pay that registration or allocation cost on every run. Uh that was actually one of the bigger bottlenecks that we faced with our early torch

forge uh integration and we've been iterating on that and it's a lot better now. And another idea that we see here is the storage volume. So Lucas talked about this but basically every rank will publish into their own storage volume and this is a torch store abstraction. It's basically an independent actor process that uh serves as an endpoint for other peers to interact with and retrieve uh

the data that they need. So the next step here is the cross node fetch. So this is where the torch store controller comes in and the controller is basically responsible for the peer discovery and um basically metadata exchange. So for example, if a receiver needs to build some type of uh local tensor, then the controller will tell it which peers have the shards that it relies on

so that it could uh ping them and get what it needs. And the controller is only on the metadata path. So there's never going to be any heavy data going through it or and it's not going to be a bottleneck in the process. Uh throughout this process, we made sure that we you know batch operations, cache metadata to reduce the amount of RPC calls and make sure

that the hardware cues are full and everything is doing work at all times. So in this example, the trainer storage volumes will initiate the cross node direct memory transfer and this will engage with a framework that provides an abstraction of the network transport layer. Uh like Dell said, building in the PyTorch stack, we have the privilege of having a lot of strong communication libraries. So torchcoms, torch

monarch provide the primitives that we need for cross node uh direct memory access and I encourage you to check out those frameworks if you've never worked with them. uh on the receiving side the generators will receive whatever shards they need based on uh you know their own layout whether that's tensor parallelism fstp or something else and you know if the the trainer side and the generator side

layouts differ then this is actually a resharting step it's not just a simple fetch one more point I'll make is that we use numoare placement for the memory so this makes sure that the local memory is close to the respective CPUs and we get nick affinity as well so we get good network saturation and good uh you know hostto device uh throughput So I'll talk about the

final step here which is what happens on the generator side once those weights arrive. Uh here we use a feature that's currently in flight called the pre-fetching volume. Just like the trainer side we're trying to drill in the fact that there's uh no runtime allocations or uh you know we don't want to block on these operations. So basically we can allow the training and inference to continue

while the new version of the model is um being staged. So on the generator we'll do a background fetch of RDMA and this will land in a circular buffer like the trainer side and when the inference actor is ready to update it'll call torch stores get uh state dictic and give it a target version. Uh so if all goes well then this stage right here is the

only time that uh you'll be blocking which is when you're loading the model weights into the inference engine and doing a GPU memory load. Um here we take advantage of the unified memory model so that these inference engines can only retrieve the shards that they need and they don't need to build an entire view of the the model which was a quirk quirk that we faced with

VLM early on. Uh and you can play around with this however you want. So let's say that you want to do a semi-synchronous path where you skip this staging buffer and the weights go directly into the inference GPUs. You can do this and depending on your approach, if you want to do async uh if you want to do sync RL uh then you can do it directly

from trainer GPUs to uh inference GPUs and Danielle will talk about what that looks like through torch store as well. >> Yeah. So uh yeah as air said in the case where you're doing sync RL uh you can take advantage of an additional optimization that torch offers where you can uh do a direct uh GPU toGPU RDMMA read and you can skip the intermediate uh CPU buffer

step. Um so essentially the GPUs on the generator nodes the GPU ranks can directly access the memory of the GPU ranks uh where the trainers model state deck lives and the way we do this is using uh RDMA. So we register the GPU memory of the trainer uh tensors uh with RDMA and then we pass these RDMA handles which are essentially just memory addresses. We pass these

to the generator and once the generator has these memory addresses, it can directly read uh the tensors on the trainer's GPU memory. Um so yeah depending on your exact setup uh if you're doing sync RL and let's say you want to be strictly on policy uh so that's why you're doing sync RL or let's say you have uh limited CPU memory uh then this can be a

good u optimization for you and we do see a speed up with this uh like in the case of Quen 314B doing cross node uh weight sync it will take about 0.5 seconds uh yeah and there are some complications with this that torch store will handle for you. Uh for example, if you're doing sharding and you have non-ontiguous tensors uh like if you're doing columnwise sharding then

this isn't handled out of the box for you with RDMA. So torch store will handle this for you and torch store will also handle uh resharding. So if you have a different sharding layout on your trainer and generator uh this is handled by torch store. And here to talk more about that is Lucas. >> So I think online charting skipped ahead. Um, yeah, I want to give

a a quick review here. As Danielle said, I think re-sharding is one of the things that makes uh this weight sync problem the most challenging. So, the crux of the matter is that you might choose on your trainers to shard differently than on your generator. So here basically we show okay here we've decided TP equals 4 and we're doing this row-wise and basically each rank will hold

this contiguous segment of memory and then oh my god here on the generator side for whatever reason the researcher or somebody has decided that inf inference is better and more optimized when we shard this way right um it's pretty hard to do this efficiently uh I think as Danielle mentioned you have a lot of kind of parts uh you have end up running into a lot of

issues where you're trying to write from a contiguous segment of memory where you're doing these all to alls where now to grab this one column I need to go and grab a small amount of data from all of these uh different ranks and then place it in in memory that's non-ontiguous because it's uh sharted columnwise. So I think um so that's pretty challenging and also just computing

I think that this is where a lot of people get mixed up. They're like, "Oh my god, there's context parallelism, there's expert parallelism, there's uh well, there's, you know, different forms of model parallelism, there's pipeline. How do I even think about this?" I think the intuition to take home is u that there's always, you know, if you could, you would probably fit your entire model on one

GPU and never think about this a day in your life. You just have some giant GPU somewhere and it would have the entire model definition, which is what we started with, right? So there is always a the assertion is that there is always a global canonical view of your model and your tensor right is it might never be materialized but the tensor at least like in philosophy

or like you know thinking abstractly like exists as one piece somewhere right and generally this is defined by an FQN a size and a shape um and what torch store does so is builds on that idea right so on save what we do is So we say okay this is the window that I have. If it's denser it does this automatically. We also have our own kind

of primitives that define the same data which is basically an offset and a shape and um an FQN. Maybe I'm forgetting some piece of metadata data type stuff like that. And on save we don't do any collectives which I think is pretty novel in terms of checkpointing. Generally people do a lot of all gathers and we just register our metadata into the controller right and then on

load we we're simply asking the question okay given my current window fetch me the data right and the controller identifies okay the data stored on storage volume zero which might be on this node over here and maybe storage volume one is on a completely different node it doesn't really matter because we know how to talk to all storage volumes across and uh I'm going to fetch that

data locally so the benefit here is that once Once you do this once, your sharding pattern is typically not training or changing during training. So you don't need to have this additional metadata transfers uh like on every step after the first time. You can immortize the cost. That's something that we're exploring. Um a big lesson here for me on top of resharding and and actually this was

a big lesson across the board was to be opinionated about the problem that you're solving. So we knew we had to solve this problem but not about kind of the solution. Don't be too inflexible on the solution you have. So, we actually went through several different versions of defining our primitives of storage volumes and controllers and clients before we got this right. And that was kind of

a learning lesson. Like, you're probably not going to get it right on the first time. So, just throw it away and keep going. Um, and finally, once so that's kind of all of the different pieces. So, we have resharding, we have RDMA, we have these storage volumes. Torch store is going to basically identify everything. And I'm going to pass it off to Amir to tell you how

we're scaling. So it not only has to be fast but it has to work at scale and air amir has some amazing results to share. Okay, I'm going to try to run through this because we're low on time, but basically uh we did some benchmarking to verify towards stores efficiency correctness at scale. And this was done in a synthetic open source environment where we test the weight

transfer primitives uh using DeepSeek V3 across uh two sets of actors that are sending and receiving. And this was done at up to 64 node scale for some sharding patterns like FSTP to FSTP or FSTP to TP where it gets a little bit more complicated. You have row-wise, call-wise, uh replicate sharding. And we can see here with the simpler cases like FSTP even when there's an imbalance

of trainers and generators uh you you maintain a relative happy path for the endto-end weight movement. But when we move to TP where it's a bit more complicated there is a bit of uh room for improvement I guess for the fetch side. But the important thing to note here is that when done correctly for async RL the main blocking segment for the trainers and generators is the

weight update which should always be fast once you have those weights in memory and you can just load them into the inference engine. So obviously we'll try to drill down that fetch time but overall I think the picture looks good. Uh one last thing uh I'll talk about where we see this going. So the the underlying ideas that we've discussed here for weight transfer can be seen

in other parts of the training uh inference stack and we see a natural opportunity to apply tor store for things like inmemory checkpointing or KV cache for uh inference and also if you'd like to extend it for fall tolerance domains like persistent storage because this is all uh involile memory or doing torch FT style uh live weight recovery for trainers that those are definitely applications and we'll

be we're excited to take a look at those as well. Uh we'll end it off there I think. Yeah, I think that's about it. I do want to throw out one last call to action. Um, if you guys res if anything that we said today resonated with you guys. If you think that the primitives we're building can fit your application, give us a star, file an issue,

get involved. We need the community support. I think we're at a place where a lot of these primitives exist, but we want the UX to be really, really solid. And we're not going to get there unless the the community helps us. And uh, yeah, thank you for listening and thank you to my contributors. And let us know if you have any questions. File some issues.