Taking Reinforcement Learning Cross Datacenter — Nan Jiang, Modal

[music] All right, cool. Hi everyone. Uh, hope you all have a good time at the conference. Uh, I'm N from Moto. Uh, at Moto, we spend a lot of time thinking about GPU capacity, like where it exists, how do we make it elastic, and what kind of workload can we actually use it. Today I want to talk about one place where everything became like gets really interesting the IO post training. A lot of IO discussion right now is about algorithm and the environments sandbox PO GRPO like to call maybe low precision training maybe deterministic kernels. Um but when you run those experiments at a scale the problem became more physical. Where are the GPUs? Are they in the same region? Uh do they have fast fabric? Uh can we get them right now? Maybe the default shape of IO comput is too restrictive. Maybe some of the work we usually force them one cluster uh like can actually run on scattered autoscaled capacity. So can we do our cross globe? So this is my talk about many about. So to make this more concrete let's start with the IO loop itself. So in the standard uh our post training loop we can see there's a one trainer and the trainer updates the policy rollout worker or maybe people call it sampler use those policy to generate trajectories. The environment will be returning the reward and observations. Those trajectory will go back to trainer for the next updates. The important error here is the way sync in the default setup trainer and the rollout will be living in the same cluster and the the way sync will be super fast with RDMA. But that also couple the rollout fleet to the to the trainer cluster. If the rollout needs to more to have more capacity maybe no more nodes during runs um you are normally limited by the fixed size during your trainer maybe your trainer cluster. So the ne next question is what kind of compute shape did we actually force everything into? On the left side is the cathedral uh one region one faster interconnect uh many GPUs wild together. This is the right shape for the trainer. On the right it is bro. This is where a lot of like usable useful usable compute actually lives. Different providers different regions different price and different availability. There's still a lot of capacity out there, but it's not one perfect RDM in the island. This is the mismatch. Available computer is distributed, but the default IO loop ask one tightly coupled cluster. And that cluster is exactly the hot part hard to get. IO wants all four of these at the same time. Enough GPU, same region, fast fabric, and available now. Any of these like is manageable, but all four of them that are pretty hard to get at the same time. ADMIC capacity is not elastic in the way the inference capacity is elastic. You cannot assume you can grow the trainer cluster uh halfway through a run just because roll out wants more nodes for a trainer. So if the whole out loop has no live like has to live inside the the one cluster rollout inherents the hottest part capacity constraint in here. So that leads to the key question does the whole out loop actually need that like this kind of shape. So let's dive into this training is one tightly coupled job. Every step has collectives all reduced and the model parallel communication that part actually wants one fast fabric a RDM connected rollout is a fleet of serving jobs. It generates trajectories call environments or maybe tools and they will be sending back data back to the trainer. So cross rollout jobs there's no global or reduce. So the thing I want to move here is not back propagation. Back propagation should stay in the cluster. The run the rollout fleet is the one that can leave. More precisely, the movable unit is the rollout serving island a coherent endpoint or maybe a local group of endpoint or that they can be serving one policy version inside island. A large model may still be having like local parallelism. They can do PD segregation. They can have like local serving constraints. So across islands the dependency is much lighter right now. policy version in and the trajectory and the metadata out. So once we define the unit that way the architecture is much much more natural. Once we define the movable unit the architecture is very straightforward. In this case we just have trainer standing audic cluster and that's where the back proper and the collective go the rollout side will be fing out across the SPR. Each rollout island will be can be single engine will be a local serving group depending on the model and the serving topology there. across island there is no global or reduce that's the most important thing there the global interface is very simple the trainer send policy weight version out and the rollout sending trajectory and the metadata back at this point the architecture depends on one remaining link the weight update so if you want to if you want to send the full parameter like full checkpoint from disk or maybe through the network then everything is like minimless and like it will be breaking immediately so uh after this aggregation The things we will be discussing about like the size of like go the size of the full parameters go through the disk. Naively that means shipping all full checkpoints every time rollout needs a new way version. At this scale the checkpoint is very huge. So a Kim scale NVIP for checkpoint you have like 500 gigabytes normally take a minute or maybe normally take multiple minutes to hours to just do the way. So moving that over commodity link is might not be the smartest choice because like uh when you're doing async maybe even even fully async training you still want way updates latency to be as low as possible like within seconds. So the problem here is not whether rollout can leave the cluster the problem is like the full checkpoint is a wrong unit of synchronization. So the next question is can we keep the exact same serve version there but send a much smaller object. So this is the bet. What if less than 1% of rollout visible weights got changed from one version to another one by rollout visible weights I mean the weights in the served rollout checkpoint uh not the FP32 optimizer states not the atom like moments like the weights are the rollouts engine which will actually use to serve maybe let's say the FBA or maybe MVP4 format if that's true we do not need to shift the entire full parameter over the network we just need to shift the change the server view the precision data got different. The important part here is like still a bit wise reconstruction. The rollout engine gets the same served served version. You would have gotten to as sync to the full checkpoint there. So if this works then the link shrinks from hundreds of gigabytes to maybe hundreds of megabytes and this is something smaller enough we can just send it across like the network. So right now we need to justify the less than 1% like claim. Why would this roll of visible like weights barely change? Uh now we get we will get into this mechanism. So it's kind of small atom like step meets finite precision. We need we need a two uh prerequisite. The first one two two ingredients. The ingredient one is the precision. The optimizer may keep very high precision master weights but the next four will pass really a BF6 visible view. That view has final resolution around a value of magnitude theta. uh BF16 spacing is roughly theta over 128. Uh that spacing is people call it oop basically the unit in the last piece. Basically it's the distance between the adjacent representable BF systeming value. But the update only needs to cross the near surrounding boundary to be viable to be visible. That boundary is about half of the oop. So roughly like the over 256 for weight around one. The BF systeming loop is around uh 0.78 and the near surrounding boundary is about 00039. If the optimizer notched the master weight by something like smaller than that uh the BFC visible value will run back so you will not see any change from the weight perspective rollout weight perspective. So that's the floor. The second primitive there is we call push. So for Adam maybe add here we just like we ignore the weight decay turn the per per sorry the the per parameter updates is the learning rate times a normalized direction the raw gradient can be dense and can have very different magnitudes across parameters and add like device by running uh gradients statistics. So the per wise push is usually on the order of learning rates. The paper passed I sight there they prove a bound. The addon step is at most B times the learning rates. So you do not need to actually remember the exact like bound there. Like the important notes is the adden makes the push small and very controlled. So at our post training learning rates the push is very very tiny. So that's the push. Combining these two primitives now we have to we have a whole like better picture assert the value changed only if the push clear the floor the push is the addent step roughly the learning rate the floor is the nearest the BF16 rounding boundary roughly theta over 256 take theta equal one the BF6 boundary is about 039 a typical added step here is around 3 millions so the update is more than a thousand smaller than the boundary so the BF visible value will not change. This is not saying the master weights uh in is frozen forever. It is saying the value that rollout engine would serve does not change on this part. So the whole magnet is pushing is is pushed versus four. Let's visualize this to have better understanding. The x-axis is the weight magnitude and the y-axis is the update magnitude. First we look at the red line. The red diagonal is the bxis invisible boundary. It's like theta over 256. And now we look at the green bound. This is the atom push. Is this roughly around the learning rate and with a conservative upper bound? So now we ask like where the most point fits for most of the weights. The red floor is above the green push. Those updates exists in the master weights but they are not visible in the served BI6 view. This step small weights on the left can move. Large weight on the right. They will just stay the same. They will be absorbed it. This is the addon absorption. This is why the serve update become very sparse. In this case, the object will be shipped is just a diff uh is just a diff not the entire FP32 automated state. We first look at the rollout view. The weight cast or projected to the dype that the rollout engine will be serving. Then we will be comparing the version t minus one and the version t. Uh in this view the patch is the patch is the change the position plus like uh replacement bits and also like some metadata. There are multiple lossless encoding. People can do selective overrides. They can people can also do xor. The important part is like there are bit equivalent bit level equivalent. So it's not a floating point addition. So there's no additive uh delta drift. If a roller engine applied the patch correctly, it reconstruct the sync server version bitwise. So everything so far we explained is about the full parameter which is hot power. So in full parameter reinforcement learning the automizer update the whole model but the roll of view patch is sparse as the thing we just explained lower is a small for a different purpose for for a different reason. The base model is frozen and the adapter is small enough by construction. So you do not need to pu So you do not need to have the push versus flow arguments here. So full parameter delta is small by absorption and the lower updates are small by construction. Let's dive into deeper about the paper itself. So the paper they mentioned more stats I will be showing here. The measurement is not great in sposity. It's not optimized state sposity. They cast weights to BF16 compare consecutive version stness and they compare the the version bitwise and they count what did not change over time across model family the result around 99% of the time is bit identical per step is also it also survives stillness even when the roller lax the changes set remain very small the important part is not only the number it is the patch is lossless change index plus imprint value reconstruct the exact same version. So a common misconception there is like the work it works because all our gradients are sparse. They are not. The paper reports the gradients are dense. About 99% of the parameter gets non-zero gradients. The FP32 master update is also dense. It's just small. The main thing is like the rollout weight change is just 1% from the perspective of rollout engine. So far we mostly talk about BF16 but the rollout of open serving even lower precision such as like MFP4 FA and NVF4 and we can see many many model providers doing this in the roll out. Uh this is not a training precision the training is just like in the normal BFC although people can do QA on that. So for fixed stale flow format the visibility for is roughly theta over two to the mantisa plus one. So as you can see the F4 will be higher and FA also will be between BF60 and F4 which means in even lower precision there will be less weight changed. So plane flows are easy to reason about. So each element has its own rounding sling one value cross the floor one value they just flip and changed group scale such in4 they are they're a bit different. This is the regime where many low precision serving system are moving towards right now. For in four each ways is quantized against a shared group scale and we can apply the same rationale and also we can observe similar thing for NBF4 is hierarchical scales and uh we can see there are different encoding and the displaying mechanism for MPV4. So this is the from one internal run. So here's the model we serve like GM 4.7 air in FP8 and we can see in the beginning there are only like 0.15% of weights got changed in the first step where the learning is high and after we have more training step like when the adden is going relatively stable and you can see the entire curve goes stable when you got only one 05% uh weight change during each step. So we can see this pattern showing more generally. We have a different research uh we saw different research sync have a similar conclusion and we saw a different model providers such as cursor composer 2 mi they all using add in the post training at at this point assume we can produce exact rollout weights version cheaply. The next question is how do we do this in practice? So from sparse delta to to our cross globe. How do we do this with elastic rollout engines and also explicit stness this is the whole shape the trainer saying a cluster after it got updated it publish immutable rollout based version to a shared bulletboard rollout engine leave outside of the training cluster which means they don't need to be RDM connected with the trainer they can be in different region or different providers. This is also request lane. You can see a request does not just say give me the completion. You will also say which version you will be sending request to and which version you will be accepting and the response will come back with the version and also like exact same information as if as if they are in the same cluster as the trainer. They will be returning tokens log prop like router replay information and many more metadata. The trainer writes immutable version to the broad after optimizer state uh after optimizer step engine pool version and the materialize it locally in the checkpoint layout. So they can just like serve directly the artifact defined version the engine choose how to how to load and shot it. It does not change it does it does not choose a different server version since it will be displaying in a hugging HF will be saved tensor format which is accepted widely by many rollout engines such as sjet and VLM. So we can support any compatible back end attention back end back end different parallelism compatible serving these type and any compatible GPUs there. We can talk more about the scar itself. The scikar is basically what makes a normal roller engine version aware. If the version is already at acceptable commit version, the psychar just proxy the request. If the engine is behind but can they can catch up the psychar just apply the missing transition. If they cannot get there the the psychar just simply return not ready. So this will be supporting elasticity rollout and any idle GPU can just be used with design with this design to support this aggregated roll out. So this is more like a system latency analyst incluster way syncing is fast because they have RDMA. Uh a full checkpoint across regions through network is pretty slow and but if we use the delta if we use exactly what we described previously we can decrease the number of like transfer transfer size from like 500 gigabytes to 5 500 megabytes. So it will be like extremely fast in seconds. So everything above was very general protocol. Stitch is one of very concrete implementation from model that we imple implementing everything above. So on the on the trainer side stitch publish uh what defines a rollout weights version and on the contract side you will be pulling out and record everything on the bulletboard and the only rollout side you will be pulling the latest weights and it start doing way sync across different region and different providers. So stitch itself is a very framework agnostic about trainer and engine and also transport. It's very async first and also agent first agentic first. By doing this we can have rollouts uh engines like autoscale globally. Each one self-s sync it weights serve accept version and return rollout metadata. That means scattered inference cap capacity became one elastic rollout fleet. Instead of being limited by the trainer cluster rollout can be the global pool. So inference capacity can now become our capacity. Last section we have some uh ongoing explorations. So we can see a lot of model providers such as moonshot and also deepc4 they have like they're adopting muan in their post training. Um does the spicy still hold for muan because a lot of thing we discussed previously only for Adam. Second question is async RL at the scale right now we can use the compute across the globe then how like how scalable is the fully async RL this is a very open-ended question there and uh third question is like does it generalize pass because like we have pre-training mid training and SFT like do we have can we apply same paradigm there last but not least we are working on some very hot problem and come work with us you can check the link there model Thank you.