Taking Reinforcement Learning Cross Datacenter — Nan Jiang, Modal

AI Engineer · 19 min · 212 sentences · from YouTube's caption track

Each timecode opens YouTube at the start of that sentence. Line anchors (#s42) are the cue ids in the WebVTT, and every line carries its start and end seconds. All transcripts has every talk, and the whole corpus as one file.

00:01[music]
00:12All right, cool.
00:13Hi everyone.
00:14Uh, hope you all have a good time at the conference.
00:17Uh, I'm N from Moto.
00:19Uh, 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.
00:28Today I want to talk about one place where everything became like gets really interesting the IO post training.
00:34A 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.
00:43Um but when you run those experiments at a scale the problem became more physical.
00:49Where are the GPUs?
00:50Are they in the same region?
00:51Uh do they have fast fabric?
00:53Uh can we get them right now?
00:55Maybe the default shape of IO comput is too restrictive.
00:58Maybe some of the work we usually force them one cluster uh like can actually run on scattered autoscaled capacity.
01:06So can we do our cross globe?
01:08So this is my talk about many about.
01:10So to make this more concrete let's start with the IO loop itself.
01:15So 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
01:24use those policy to generate trajectories.
01:26The environment will be returning the reward and observations.
01:30Those trajectory will go back to trainer for the next updates.
01:34The 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.
01:44But that also couple the rollout fleet to the to the trainer cluster.
01:48If 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.
01:59So the ne next question is what kind of compute shape did we actually force everything into?
02:06On the left side is the cathedral uh one region one faster interconnect uh many GPUs wild together.
02:13This is the right shape for the trainer.
02:15On the right it is bro.
02:17This is where a lot of like usable useful usable compute actually lives.
02:22Different providers different regions different price and different availability.
02:27There's still a lot of capacity out there, but it's not one perfect RDM in the island.
02:32This is the mismatch.
02:33Available computer is distributed, but the default IO loop ask one tightly coupled cluster.
02:40And that cluster is exactly the hot part hard to get.
02:43IO wants all four of these at the same time.
02:46Enough GPU, same region, fast fabric, and available now.
02:50Any of these like is manageable, but all four of them that are pretty hard to get at the same time.
02:56ADMIC capacity is not elastic in the way the inference capacity is elastic.
03:00You cannot assume you can grow the trainer cluster uh halfway through a run just because roll out wants more nodes for a trainer.
03:08So 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.
03:18So that leads to the key question does the whole out loop actually need that like this kind of shape.
03:25So let's dive into this training is one tightly coupled job.
03:29Every 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.
03:39It generates trajectories call environments or maybe tools and they will be sending back data back to the trainer.
03:45So cross rollout jobs there's no global or reduce.
03:49So the thing I want to move here is not back propagation.
03:52Back propagation should stay in the cluster.
03:55The run the rollout fleet is the one that can leave.
03:59More 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.
04:10A large model may still be having like local parallelism.
04:13They can do PD segregation.
04:14They can have like local serving constraints.
04:17So across islands the dependency is much lighter right now.
04:20policy version in and the trajectory and the metadata out.
04:24So once we define the unit that way the architecture is much much more natural.
04:29Once we define the movable unit the architecture is very straightforward.
04:32In 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.
04:41Each rollout island will be can be single engine will be a local serving group depending on the model and the serving topology there.
04:48across 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
05:01at 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
05:16so 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.
05:24Naively that means shipping all full checkpoints every time rollout needs a new way version.
05:30At this scale the checkpoint is very huge.
05:32So 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.
05:42So 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.
05:55So the problem here is not whether rollout can leave the cluster the problem is like the full checkpoint is a wrong unit of synchronization.
06:02So the next question is can we keep the exact same serve version there but send a much smaller object.
06:09So this is the bet.
06:10What 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
06:20uh 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
06:32if 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.
06:42The important part here is like still a bit wise reconstruction.
06:45The rollout engine gets the same served served version.
06:49You would have gotten to as sync to the full checkpoint there.
06:53So 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.
07:03So right now we need to justify the less than 1% like claim.
07:06Why would this roll of visible like weights barely change?
07:10Uh now we get we will get into this mechanism.
07:13So it's kind of small atom like step meets finite precision.
07:17We need we need a two uh prerequisite.
07:20The first one two two ingredients.
07:22The ingredient one is the precision.
07:24The optimizer may keep very high precision master weights but the next four will pass really a BF6 visible view.
07:31That view has final resolution around a value of magnitude theta.
07:36uh BF16 spacing is roughly theta over 128.
07:40Uh that spacing is people call it oop basically the unit in the last piece.
07:45Basically it's the distance between the adjacent representable BF systeming value.
07:49But the update only needs to cross the near surrounding boundary to be viable to be visible.
07:53That boundary is about half of the oop.
07:55So roughly like the over 256 for weight around one.
08:00The BF systeming loop is around uh 0.78 and the near surrounding boundary is about 00039.
08:07If 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.
08:20So that's the floor.
08:21The second primitive there is we call push.
08:24So 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
08:35the raw gradient can be dense and can have very different magnitudes across parameters and add like device by running uh gradients statistics.
08:45So the per wise push is usually on the order of learning rates.
08:49The paper passed I sight there they prove a bound.
08:52The addon step is at most B times the learning rates.
08:56So you do not need to actually remember the exact like bound there.
09:00Like the important notes is the adden makes the push small and very controlled.
09:05So at our post training learning rates the push is very very tiny.
09:09So that's the push.
09:11Combining 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
09:24rounding 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.
09:42This is not saying the master weights uh in is frozen forever.
09:46It is saying the value that rollout engine would serve does not change on this part.
09:51So the whole magnet is pushing is is pushed versus four.
09:56Let's visualize this to have better understanding.
09:58The x-axis is the weight magnitude and the y-axis is the update magnitude.
10:03First we look at the red line.
10:04The red diagonal is the bxis invisible boundary.
10:07It's like theta over 256.
10:09And now we look at the green bound.
10:11This is the atom push.
10:12Is this roughly around the learning rate and with a conservative upper bound?
10:16So now we ask like where the most point fits for most of the weights.
10:20The red floor is above the green push.
10:23Those updates exists in the master weights but they are not visible in the served BI6 view.
10:28This step small weights on the left can move.
10:32Large weight on the right.
10:33They will just stay the same.
10:34They will be absorbed it.
10:36This is the addon absorption.
10:37This is why the serve update become very sparse.
10:41In this case, the object will be shipped is just a diff uh is just a diff not the entire FP32 automated state.
10:49We first look at the rollout view.
10:51The weight cast or projected to the dype that the rollout engine will be serving.
10:56Then 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.
11:08There are multiple lossless encoding.
11:10People can do selective overrides.
11:12They can people can also do xor.
11:14The important part is like there are bit equivalent bit level equivalent.
11:19So it's not a floating point addition.
11:21So there's no additive uh delta drift.
11:24If a roller engine applied the patch correctly, it reconstruct the sync server version bitwise.
11:30So everything so far we explained is about the full parameter which is hot power.
11:35So 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.
11:46The base model is frozen and the adapter is small enough by construction.
11:50So you do not need to pu So you do not need to have the push versus flow arguments here.
11:56So full parameter delta is small by absorption and the lower updates are small by construction.
12:03Let's dive into deeper about the paper itself.
12:05So the paper they mentioned more stats I will be showing here.
12:08The measurement is not great in sposity.
12:10It's not optimized state sposity.
12:12They 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
12:27is 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
12:38change index plus imprint value reconstruct the exact same version.
12:44So a common misconception there is like the work it works because all our gradients are sparse.
12:49They are not.
12:50The paper reports the gradients are dense.
12:52About 99% of the parameter gets non-zero gradients.
12:56The FP32 master update is also dense.
12:59It's just small.
13:00The main thing is like the rollout weight change is just 1% from the perspective of rollout engine.
13:06So 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.
13:18Uh this is not a training precision the training is just like in the normal BFC although people can do QA on that.
13:25So for fixed stale flow format the visibility for is roughly theta over two to the mantisa plus one.
13:31So 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.
13:42So plane flows are easy to reason about.
13:45So 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.
13:54This is the regime where many low precision serving system are moving towards right now.
13:59For 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
14:11encoding and the displaying mechanism for MPV4.
14:16So this is the from one internal run.
14:18So 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
14:29and 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%
14:40uh weight change during each step.
14:43So we can see this pattern showing more generally.
14:46We 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
14:58at at this point assume we can produce exact rollout weights version cheaply.
15:02The next question is how do we do this in practice?
15:07So from sparse delta to to our cross globe.
15:10How 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
15:22rollout 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.
15:34This is also request lane.
15:36You can see a request does not just say give me the completion.
15:39You 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
15:48like exact same information as if as if they are in the same cluster as the trainer.
15:52They will be returning tokens log prop like router replay information and many more metadata.
16:00The 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.
16:10So they can just like serve directly the artifact defined version the engine choose how to how to load and shot it.
16:17It 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.
16:31So we can support any compatible back end attention back end back end different parallelism compatible serving these type and any compatible GPUs there.
16:40We can talk more about the scar itself.
16:43The scikar is basically what makes a normal roller engine version aware.
16:46If the version is already at acceptable commit version, the psychar just proxy the request.
16:52If the engine is behind but can they can catch up the psychar just apply the missing transition.
16:58If they cannot get there the the psychar just simply return not ready.
17:04So 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.
17:13So this is more like a system latency analyst incluster way syncing is fast because they have RDMA.
17:20Uh 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
17:34transfer size from like 500 gigabytes to 5 500 megabytes.
17:37So it will be like extremely fast in seconds.
17:41So everything above was very general protocol.
17:44Stitch is one of very concrete implementation from model that we imple implementing everything above.
17:49So 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
17:59and the only rollout side you will be pulling the latest weights and it start doing way sync across different region and different providers.
18:08So stitch itself is a very framework agnostic about trainer and engine and also transport.
18:13It's very async first and also agent first agentic first.
18:17By doing this we can have rollouts uh engines like autoscale globally.
18:22Each one self-s sync it weights serve accept version and return rollout metadata.
18:27That means scattered inference cap capacity became one elastic rollout fleet.
18:33Instead of being limited by the trainer cluster rollout can be the global pool.
18:38So inference capacity can now become our capacity.
18:43Last section we have some uh ongoing explorations.
18:48So 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.
18:56Um does the spicy still hold for muan because a lot of thing we discussed previously only for Adam.
19:02Second 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
19:15and 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
19:32Thank you.