Taking Reinforcement Learning Cross Datacenter

Talk transcript

1. 0:00

   [upbeat music] All right. Cool. Hi, everyone. Uh, hope you all have a good time at the conference.

2. 0:17

   Uh, I'm Nan from Modal. Uh, at Modal, 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.

3. 0:28

   Today, I want to talk about one place where everything became, like, gets really interesting, the RL post-training. A lot of our discussion right now is about algorithm and the environments.

4. 0:37

   Sandbox, PPO, GRPO, like tool 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?

5. 0:50

   Are they in the same region? Uh, do they have fast fabric? Uh, can we get them right now? Maybe the default shape of RL compute is too restrictive. Maybe some of the work we usually force in one cluster, uh, like can actually run on scattered auto-scaled capacity.

6. 1:06

   So can we do RL across the globe? So this is my talk about, mainly about. So to make this more concrete, let's start with the RL loop itself. So in the standard, uh, RL post-training loop, we can see there's one trainer, and the trainer updates the policy.

7. 1:21

   Rollout worker, or maybe people call sampler, use those policy to generate tra- tra- trajectories. The environment will be returning the reward and observations. Those trajectory will go back to trainer for the next updates.

8. 1:34

   The important error here is the wait sync. In the default setup, trainer under the rollout will be living in the same cluster, and the, the wait sync will be super fast with RDMA.

9. 1:44

   But that also coupled the rollout fleet to the, to the trainer cluster. If the rollout needs to more ca- to have more capacity, maybe no, no more nodes during runs, um, you are normally limited by the fixed size during your trainer, maybe your trainer cluster.

10. 2:00

    So the ne- next question is, what kind of compute shape did we actually force everything into?

11. 2:06

    On the left side is the cathedral. Uh, one region, one fast interconnect. Uh, many GPUs wired to- together. This is the right shape for the trainer. On the right is the sprawl.

12. 2:17

    This is where a lot of like usable, useful, usable compute actually lives. Different providers, different regions, different price, and different availabilities. There's still a lot of capacity out there, but it's not one perfect RDMA island.

13. 2:32

    This is the mismatch. Available compute is distributed, but the default RL loop asks one tightly coupled cluster, and that cluster is exactly the hard part to get.

14. 2:43

    RL 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, they are pretty hard to get at the same time.

15. 2:56

    RDMA capacity is not as elastic in the way that inference capacity is elastic. You cannot assume you can grow the trainer cluster, uh, halfway through a run just because rollout wants more nodes for a trainer.

16. 3:08

    So if the whole RL loop has no live, like, has to live inside the, the one cluster, rollout inherits the hardest part, capacity constraint in here. So that leads to the key question: Does the whole RL loop actually need that, like this kind of shape?

17. 3:25

    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, it actually wants one fast fabric RDMA-connected.

18. 3:37

    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 across, uh, rollout jobs, there's no global all-reduce.

19. 3:49

    So the thing I want to move here is not backpropagation. Backpropagation 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.

20. 4:09

    Inside the island, a large model may still be having, like local parallelism. They can do PD data aggregation. They can have, like local serving constraints. So across islands, the dependence is much lighter right now.

21. 4:20

    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.

22. 4:29

    Once we define the movable unit, the architecture is very straightforward in this case. We just have trainer staying in RDMA cluster, and that's where the backprop and the collective go.

23. 4:38

    The rollout side will be fanning out across the sprawl. Each rollout island will be, can be single engine or maybe a local serving group, depending on the model and the serving topology there.

24. 4:48

    Across island, there is no global all-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.

25. 5:01

    At this point, the architecture depends on one remaining link, the weight update. So if we want to s- if we want to send the full parameter, like full checkpoint from disk, or maybe through the network, then everything is, like minimalist and, like it will be breaking immediately.

26. 5:16

    So, uh, after this aggregation, the things that we will be discussing about, like the size of like go th- the size of the full parameters go through the disk.

27. 5:24

    Naively, that means shipping all f- full checkpoints every time rollout needs a new weight version. At this scale, the checkpoint is very huge. So a Kimi-scale NVFP4 checkpoint, you have like five hundred gigabytes.

28. 5:36

    Normally, it take a minute or maybe i- normally it take mult- multiple minutes to hours to just do the wait sync. So moving that over commodity links might not be the smartest choice 'cause like, uh, when you're doing async, maybe even, even fully async training, you still want the weight updates latency to be as low as possible,

29. 5:53

    like within seconds. So the problem here is not whether rollout can leave the cluster. The problem is like the full checkpoint is the wrong unit of synchronization. So the next question is, can we keep the exact same serve version there but send a much smaller object?

30. 6:09

    So this is the bet. What if less than 1% of the rollout visible weights got changed from one version to another one? By rollout visible weights, I mean the weights in the served rollout ver- uh, checkpoint, uh, not the F- FP32 optimizer state, not the Adam, like, moments.

31. 6:25

    Like, the weights are the rollout's engine, would ac- which we'll actually use to serve, maybe let's say, the FPA or maybe NVFP4 format. If that's true, we do not need to ship the entire full parameter over the network.

32. 6:37

    We just need to ship the change to serv- server view, the precision delta, got, got difference. The important part here is like still bitwise re- reconstruction. The rollout engine gets the same served, ser- uh, served version you would have gotten to as synced to the full checkpoint there.

33. 6:53

    So if this works, then the link shrinks from hundreds of gigabytes to maybe hundreds of megabytes, and this is something small enough we can just send it across, like, the network.

34. 7:03

    So right now we need to justify the less than 1%, like, claim. Why would this rollout visible, like, weights barely change?

35. 7:10

    Uh, now we get, we will get into this mechanism. So it's kind of small Adam, like, step meets finite precision. We need, we need a two,

36. 7:19

    uh, prerequisite. The first one... Two gra- two ingredients. The ingredient one is the precision. The optimizer may keep very high precision ma- uh, master weights, but the next forward pass read the BF16 visible view.

37. 7:31

    That view has finite resolution. Around a value of magnitude Theta, uh, BF16 spacing is roughly, uh, Theta over 128. Uh, that spacing is, people call it ULP, basically the unit in the last place.

38. 7:45

    Basically, it's the distance between the adjacent representable BF16 value. But an update only needs to cross the nearest surrounding boundary to be v- viable, for it to be visible.

39. 7:53

    That boundary is about half of the ULP, so roughly, like, this Theta over 256. For weight around one, the BF16 ULP is around, uh, .0078, and the nearest surrounding boundary is about .0039.

40. 8:08

    If the optimizer nudged the master weight by something, like, smaller than that, uh, the BF16's invisible value will run back. So you will not see any change from the weight perspective, rollout weight perspective.

41. 8:20

    So that's the floor. The second primitive there is, we call push. So for Adam, or maybe AdaLevel here, we just, like, we ignore the weight decay term. The per, per...

42. 8:30

    Sorry, the, the per- parameter update is the learning rate times the normalized direction. The raw gradient can be dense and can have very different magnitudes across parameters. An Adam, like, divides by running, uh, gradients statistics.

43. 8:45

    So the per-weight push is usually on the order of learning rates. The paper Paused, I said there, uh, they have, they prove a bound. The Adam step is at most b times the learning rates.

44. 8:56

    So you do not need to actually remember the exact, like, bound there. Like, the important notes here is the Adam makes the push small, and they're very controlled. So at RL post-training learning rates, the push is very, very tiny.

45. 9:09

    So that's the push. Combining these two pr- primitives, now we have to, we have a whole, like, better picture. A served va- value changed only if the push cleared the floor.

46. 9:20

    The push is the Adam step, roughly the learning rates. The floor is the nearest BF16 rounding boundary, uh, roughly Theta over 256. Take a Theta equal one, the BF16 boundary is about .0039.

47. 9:32

    A typical, uh, Adam step i- here is around three minutes. So the update is more than a thousand smaller than the boundary, so the BF visible value will not change.

48. 9:42

    This is not saying the master weights are in, is frozen forever. It is saying the value that rollout engine would serve does not change on this part. So the whole mechanism is pushing, uh, is, is push versus floor.

49. 9:56

    Let's visualize this to have better understanding. The x-axis is the weight magnitude, and the y-axis is the u- uh, update magnitude. First, we look at the red line. The red line now is the BF16 visible boundary.

50. 10:07

    It's like Theta over 256. And now we look at the green bound. This is the Adam push. It sits roughly around the learning rate and, uh, with a conservative upper bound.

51. 10:16

    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 exist in the master weights, but they are not visible in the served BF16 view this step.

52. 10:30

    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 Adam absorption. This is why the served update become very sparse.

53. 10:41

    In this case, the object will be shipped as just a diff, uh, it's just a diff, not the entire FP32 optimizer state. We first look at the rollout view.

54. 10:51

    The weight cast, like, all projected to the D type 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.

55. 11:01

    The patch is, the patch is the change to precision, plus, like, re- uh, replacement bits and also, like, some metadata. There are multiple lossless encoding. People can do selective overwrites.

56. 11:12

    They can... People can also do XOR. The important part is, like, they are bit equivalent, bit-level equivalent, so it's not a floating point addition, so there's no additive, uh, delta drift.

57. 11:24

    If a rollout engine apply the patch correctly, it reconstruct the same served version bitwise.

58. 11:30

    So everything so far we explained is about the full parameter, which is hub, hub RL. So in full parameter reinforcement learning, the optimizer update the whole model. But the rollout view patch is sparse, as the thing we just explained.

59. 11:43

    LoRA is a small for a different purpose, uh, for, for a different reason. The base model is frozen, and the adapter is small enough by construction, so do not need to put, so do not need to have the push versus the floor argument here.

60. 11:56

    So full parameter delta is small by absorption, and the LoRA updates are small by construction.

61. 12:03

    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 gradient sparsity. It's not optimizer state sparsity.

62. 12:12

    They cast weights to BF16, compare consecutive versions delta, 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.

63. 12:27

    It's also s- it also survives staleness. Even when the rollout lags, the changes that remain are very small. The important part is not only the number, it is the patch is lossless.

64. 12:38

    Change index plus replacement value reconstructs the exact same version.

65. 12:44

    So a common misconception there is, like, the work- it works because all gradient are sparse. They are not. The paper reports the gradients are dense. About 99% of parameter gives non-zero gradients.

66. 12:56

    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.

67. 13:06

    So far, we mostly talk about BF16, but the rollout often serving even lower precision, such as, like, MF- FP4- F8 and NVFP4, and, uh, we can see many, many model providers doing this in their rollout.

68. 13:18

    Uh, this is not a training precision. The training is ju- just, like, in the normal BF16, although people can do QAT on that. So f- for fixed scale flow format, the visibility for is roughly theta over two to the mantissa plus one.

69. 13:32

    So as you can see, the FP4 will be higher, and FP also will be between BF16 and FP4, which means in even lower precision, there will be less weight change.

70. 13:42

    So plain floats are easy to reason about, so each element has its own rounding ceiling. One value cross the floor, one value they just flip and it changed. Group scale such Int4, they are qui- they are a bit different.

71. 13:54

    This is the regime where many low precision serving system are moving towards right now. For Int4, each weight is quantized against a shared group scale, and we can apply the same rationale, and also we can observe similar thing for NVFP4.

72. 14:07

    It's hierarchical scales, and, uh, we can see there are different encoding and displaying mechanism for NVFP4.

73. 14:16

    So this is the... from one internal run. So here's the model we serve, like, GROM 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 learner is high.

74. 14:29

    And after we have more training step, like, when the Adam is going r- r- relatively stable, and you can see the entire curve goes stable. We got only 1...

75. 14:39

    0.5%, uh, weight change during each step. So we can see this pattern showing more generally. We have a different research. Uh, we, we saw different research across RL sync have a similar conclusion, and we saw a different model provider such as Cursor, Composer-2, MAI, they all using Adam in their post-training.

76. 14:58

    At, at this point, assume we can produce exact rollout weights version cheaply. The next question is, how do we do this in practice?

77. 15:06

    So from sparse data to, to RL across globe, how do we do this with elastic rollout engines and also explicit staleness?

78. 15:16

    This is the whole shape. The trainer stay in the RDMA cluster. After it got updated, it published immutable rollout weights version to a shared bullet board. Rollout engine live outside of the training cluster, which means they don't need to be RDMA connected with the trainer.

79. 15:31

    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 then which version you will be accepting.

80. 15:44

    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 prob, like, router replay information, and many more me- metadata.

81. 16:00

    The trainer writes immutable version to the broad aft- after optimizer state. Uh, af- after optimizer step. Engine pull version and materialize it locally in the checkpoint layout, so they can just, like, serve directly.

82. 16:13

    The auto-defined version, the engine choose how to, how to load and shard it. It does not change... It does, it does not choose the different server version.

83. 16:22

    Since it will be displaying in a HuggingFa- HDF, it will be saved tensor format, which is accepted widely by many rollout engines such as AstroNet and VLM. So we can support any compatible backend.

84. 16:33

    Attention backend, MOE backend, different parallelism, compatible serving vtype, and any compatible GPUs there.

85. 16:41

    We can talk more about the sidecar itself. The sidecar is basically what makes a normal rollout engine version-aware. If the version is already at acceptable committed version, the sidecar just proxy the re- request.

86. 16:52

    If the engine is behind, but it can, they can catch up, the sidecar just apply the missing transaction. If they cannot get there, the, the sidecar just sim- simply return not ready.

87. 17:04

    So this will be supporting elastic rollouts, and id- any idle GPU can just be used with design, with this design to support this aggregated rollout.

88. 17:14

    So this is more like a system latency analysis. In-cluster way sync is fast because they have RDMA. Uh, a full checkpoint across regions through network is pretty slow.

89. 17:25

    And but if we use the delta c- If, if we use exactly what we described previously, we can decrease the number of, like, transfer, t- transfer size from, like, five hundred gigabytes to fi- five hundred megabytes, so you will be, like, extremely fast, in seconds.

90. 17:42

    So everything above was very general protocol. Stitch is one of very concrete implementation from Modul that we imple- im- 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'll be pulling out and it require everything on the bullet board.

91. 18:00

    And on the rollout side, you'll be pulling the latest weights, and it start doing a way sync across different region and different providers. So Stitch itself is a very framework-agnostic about trainer and engine and also transport.

92. 18:13

    It's very a- async first and also agent first, agentic first.

93. 18:17

    By doing this, we can have rollouts, uh, engines like auto-scale globally. Each one self sync its weights, serve accept version, and return rollout metadata. That means scattered inference capab- uh, capacity became one, uh, elastic rollout fleet.

94. 18:33

    Instead of being limited by the training cluster, rollout can be the global pool. So inference capacity can now become RL capacity.

95. 18:44

    Last section, we have some, uh, ongoing explorations.

96. 18:48

    So w- we can see a lot of model providers such as Moonshot and also DeepSea-V4, they have, like... they are adopting Muon in their post-training. Um, does the sparsity still hold for Muon?

97. 18:59

    'Cause a lot of thing we discussed previously only for Adam. Second question is async RL at scale. Right now, we can use the compute across the globe. Then how, like, how scalable is the fully async RL?

98. 19:12

    This is a very open-ended question there. And, uh, third question is, like, does a generalized PaaS RL, 'cause, like, we have pre-training, mid-training, and SFT. Like, do we have-- can we apply same paradigm there?

99. 19:25

    Last but not least, we are working on some very hard problem, and come work with us. You can check the link there, modul.jobs. Thank you.