Optimizing Reinforcement Learning at Trillion-Parameter Scale - Songlin Jiang

This talk covers the application of reinforcement learning on trillion parameter large language models, utilizing Varro, Megatron, and VOM as key components of the system. The speaker, Sony John, discusses the challenges related to speed, memory, and consistency in distributed systems and highlights the importance of applying reinforcement learning to improve reasoning behaviors in models that were previously pretrained. The session delves into using low-rank adapters (LoRA) for parameter-efficient updates, allowing the model's strong prior knowledge to be retained while addressing the computational demands of reinforcement learning. Additionally, the talk addresses the critical issue of training inference mismatch in mixture of experts (MoE) models, introducing the concept of router replay to enhance stability during reinforcement training. The presentation concludes by introducing Mint, a managed service that simplifies the operational aspects of running large-scale reinforcement learning systems.

Hi everyone. I'm Sony John from Alt University and Mind Lab. Today I want to talk about what is on reinforcement learning on trillion parameter scale uh large language models using Varro, the most popular reinforcement learning network so far. Together with training engine Megatron and also inference engine VOM. As you can see here, we are targeting at distance distributed across multiple engines. So, the hardest problem are not

only speed and memory, but also consistency. In the first half, I will talk about the system and show how we make it work. In the second half, I will talk about uh training inference mismatch for improving If you are building large-scale reinforcement learning systems, I hope you leave with a few practical design principles you can apply immediately. A quick self-introduction. I'm a doctoral researcher working on AI

infrastructure, especially systems for large-scale reinforcement learning and I'm also a free software developer and then uh mirror maintainer, and I contribute to many projects such as VOM and Megatron Bridge, which bridge the gap between hugging face format and also the Megatron form format. Now, let's start with uh basic Do you need reinforcement learning after pre-training? Pre-training teaches a model to predict the next token. That gives you

a very strong prior, broad knowledge, and general capability. However, many of the behaviors we want from the reasoning models are hard to obtain through imitation alone. We don't want just want the model to produce likely text. We want it to favor trajectories that leads to better outcomes. That is where reinforcement learning helps. Reinforcement Reinforcement let us optimize for multi-step reasoning quality. It also helps with addressing task

specific and shapes long horizon behavior. The model generates uh rollout trajectory. We evaluate it with a reward signal. And then we update the policy. So, you can think of reward learning as the technique to let turns a strong pre-trained model into a model that has useful reasoning behaviors. Now, if you haven't already know, uh here we are targeting at mixture of experts or MoE, not the standard

dense model. In a dense model, every token goes through every layer. And uh model, there is uh there there's a router that selects only a subset of the uh experts for each token. That gives you a very attractive combination, much higher total capacity, uh and to continue continue the scaling law, compute per token. Um to make this concrete, the target model in this store here is Chimera

2 thinking. This is a trillion scale trillion parameter scale MoE reasoning But the active parameters per token are much smaller, 32.6 billion, because only a subset of the total uh 384 experts is active. If you try to do the full parameter reinforce learning at this scale, the resource requirement becomes very large, even with a relatively small context lens, such as 4K, which is barely enough actually for

reason model even for a single term. Uh, not to mention about agent reward setting that involves the multi-term. Uh, you already need a very expensive training and the low-out setup, which will need around 300 views. So, the immediate question becomes, can we still improve a model like this with reinforcement learning but uh, getting more cost effective while still efficient? Since we start from a huge pre-trained MoE

model, which already encodes a very paradigm, we ask whether a small trainable parameter efficient adapters are enough for updating the policy during reinforcement learning. This leads to one question uh, behind the first half of the talk. Can we improve K K 2 using reinforcement learning without updating the whole model? The hope is that if this gives us something that is affordable, effective, and operationally much easier to

run. That leads leads us to low-rank adapters, which is LoRA. So, what's the difference between the full parameter training and LoRA-based Parameter training, every weight is updated. That gives maximum flexibility, but it also means much larger optimizer state, and much more memory pressure. Besides, we will also need to keep a separate frozen weight in memory if we want to get value, which is needed in some reinforcement

learning algorithms, and measure how different the updated weight is compared to the original pre-trained one. With LoRA, we freeze the base model and only train a low-rank update instead of updating the entire pre-trained tensors. a small delta on top of it. According to our experience, which ranked equal to 32 for the Kimi K2 model that attach Laura to all linear layers. The trainable parameter size is similar

to the 32 billion dense model. For a model at That is a huge reduction in trainable state. And we don't need to keep a separate for calculating the K value. As you can temporarily disable the Laura adapters for it. The important point is that Laura is not only a cost saving trick. When the base very strong, a small trainable delta on top of the last strong prior

can be an excellent fit to spend the reinforced learning compute budget if your training budget is fixed. In other words, sometimes it's better to do reinforced learning on the strongest model, you can rather than the full reinforced learning on the weaker base model. So, how to make that Laura pass uh actually work at trillion scale? Let's first have an overview of the system. This is the high-level

system architecture. engine. It generates samples efficiently. Varo acts as a controller and also the orchestration layer. Megatron is the training engine. And Megatron Bridge is the Laura engine. Uh sorry, Megatron Bridge is what helps connect the training side representation with representation. In the reinforced learning training loop, Varo runs samples trajectories, then Varo aggregates, and compute rewards and also advantages. Then Megatron updates the adapters. The key operation

here is repeat, or you can call the sync. After training, we need to move the fresh weights uh back into the raw engine so that the policy used for generation stays updated with the policy we are trying to optimize. So, at trillion scale, there is no single person strategy that's Different parts of the model needs starting rules. Tensor parallelism splits large tensors into smaller ones across devices.

Pipeline parallelism splits the layers across stages. Expert parallelism distributed the experts across Context parallelism splits the sequence to reduce the activation memory. The important engineering point is that Laura has to coexist with all of these. The adapter path has to remain correct under the same hybrid parallelism strategy than the base model already needs. Making Laura work at scale, we really mean making it compatible with the full

distributed layout and the trillion MoE model. So, the next video is how we make all for the MoE model to support. It's not enough to attach layers only to the few dense attention projections and call it out. For MoE reinforcement learning, we need to support for both the dense linear layers and expert side linear layers. Uh that includes the expert not just the common dense parts. So,

over the tensor parallel side, the adapter has to respect how the base model how the base model weight is sharded. So, LoRA A should follow the the base partitioning, while LoRA B can shared in the column parallel side, and style, and on the expert parallel to avoid the separate and heavy O2O communications. Uh for adapters, we make all expert shards in an expert on adapter. Now, we

get into one of the most important practical part, repeat. After the training engine updates the adapters, the LoRA the LoRA engine must use those new adapters for the next generation step. Several ways to do this. The first is what I will call the merge pass. You disable LoRA on the inference engine and merge the LoRA LoRA delta back into the full model weights, and move the merge

results. The advantage is the serving side does not need native LoRA support, but the downside is obvious. You move much more data data, and the communication equals is high, and merge operations can introduce additional precision loss. The second is what I call the bridge Instead of merge weights, you move only the adapter deltas. This is much cheaper in communication, and it keeps the representation closer to how

training produce it. But now the LoRA engine needs the LoRA serving support for your model, and depending on the kernel, inference performance can degrade a lot. So, this is really a really a system trade-off. Merge can simplify the serving, but it's communication heavy. Bridge is communication efficient, but it makes serving integration more demanding. have degraded performance on latency. So far, I have focused on the making the

system to run, and now let's focus on making it stable. Training inference mismatch means that the raw engine and also the training engine are not perfectly Even when they start from the same model, they may use different corners, and different fusions, different numerical execution passes, and different implementations of the same operations. The result is that for the same prompt, the raw engine can produce a slightly different

from the logits implied from by the training engine. If that difference is small, PPO style methods can still tolerate it. But once the difference gets larger, you are no longer updating the exact policy you are sampled from. And MoE makes this issue much worse because more numerical difference can change the actual computation path. In a dense model, a small numeric drift usually changes the logits a little.

In an MoE model, a small drift in the can flip the selecting expert entirely. For example, when raw engine sees the expert seven a little bit higher than the expert nine, and while training sees expert nine a little bit higher than the expert seven, so the expert path flips, which results picks expert seven, while the training engine picks expert 9. And once that happens, the downstream activations

are no longer just slightly different. The model is literally taking a different path through the neural network. That means the behavior policy used during rollout and the assumed during training can can diverge much more sharply than in the dense case. So, uh natural reaction to this problem is, can we just uh correct for the mismatch as the loss level? There are two common families of medication. Important

sampling weighs the gradient to account for the policy mismatch. And rejection sampling drops tokens or sequence that are too far off policy. These methods can of course help. And they are useful tools, but they have uh limitation. They correct or filters after the mismatch has already happened. They do not make the two engines took the same route through the MOE. In other words, they operate on the

consequences, not on the source of the divergence. And that's why in severe MOE these methods may still leave a lot of uh instability or throw away too much that data. And you may correctly optimize but you are not restoring the computational path consistently. And this leads to the router replay R3. The idea is simple. During rollouts, we can save the routing decisions. During the training, replay those

same And so the training path follows the same expert path that the sample actually used during the generation. That gives us two benefits. It keeps the sample instead of discarding. Second, it reduces root-induced drift by aligning the communication path between routing and training, which is often the most damaging part of mismatch in MoE Here I want to emphasize that this is the system fix that we can

supplement on top of an algorithm fix. Just as I mentioned previously, uh the to make router replay work, the capture, serialization, and replay path has to be implemented consis- across VoM, Varo, and Megatron. So, the important point is not only the concept of router replay, but also making that implementation real across the full stack. Now, let's look at how router replay is implemented inside VoM. the MoE

router already decides which each token goes to. a routed expert capture turned on, VoM can record those top K experts IDs token by token and also layer by layer as the model runs. I replace about path consistency, not only about the probability correction. That means we are not replaying a routing score later. We are replaying the actual sparse communication path uh that router used. Now, here comes

the harder part, which is make capture IDs are semantically correct and aligned with the training side. So, there were two major greatness issues that we found when we initially tried router replay on Varo. To get And the practical symptom when this is wrong is very clear. And the training appears to be replaying router information. But in reality, it's falling completely different experts from the ones used during

rollout, which makes what produced from raw engine really different from training when with router replay disabled. So after investigation, the first bug we found is expert layer mismatch. On the real end side, the reported expert router information can be indexed across all transformer layers, including the While the K make A2 X MOE model, it has like the first three layers as a dense layer and the later

parts as the MOE layer. And on the MOE on the Megatron training uh only makes sense for the actual MOE router. So if those two views are not aligned, and then if the IDs are numeric, even if the IDs are like numeric values, they are attached to the wrong layers replay becomes meaningless. ID can be wrong under the EBLP. EPLB, which stands for the expert partition load

balancer. EPLB duplicates whole experts and remaps logical experts to physical placements across the GPUs. This is good for but it introduces semantically problems for replay. Training wants the logical expert IDs. Because those correspond to the model's actual routing decisions. Um if what we export from inference is already the remapped physical ID, training will replay the wrong path. So, the feature is conceptually simple, absolutely critical. We need

to capture the logical top K routing decisions before EPLB mapping happens. And both of the bugs has already been submitted and accepted by VOA, and two bugs are no longer the problem. uh that's all key points for today's Here are the three design principles for you to take away. First, preserve the strongest prior you have. For trillion parameter models, it's more cheaper. And uh a small trend

will build up on top of a very strong model can be a better use of reinforcement uh than for reinforcement learning on a weaker base. Second, treat Laura merge as the and Laura merge bridge as it's a system trade-off. Merge can simplify serving, but is Bridge is communication light, but requires stronger serving engine integration and careful performance engineering. Third, treat routing as part of the In MOE

reinforcement learning, hybridizing the positional load is often not enough if the expert path changes. Router replay should be thought of as a first-class stability mechanism, not just an information trick. And uh build on top of this principles, Here we introduce Mint. The basic idea is simple. Instead of hand-wiring rollout training, checkpoint, and evaluations, users should define the loop, what to train, and what data to experience to

learn from, and how to learn, and how to Then Mint handles the infrastructure underneath. Scheduling compute, running distributed jobs, managing model states, and recovering from failures. Why does that matter? Because large-scale reinforcement learning today carries still carries a big operational burden. What goes into running together the systems, rather than improving the learning loop itself. The goal here is to abstract away that operation burden, keep the training

to evaluation to deploy loop tight, and turn one-off systems work into something repeatable. And this is meant to be practical, not just conceptual. Mint is a managed service for post-training and It comes with a Python SDK, so it can support both quick starts and scale jobs. It is designed to work across a wide range of model scale from the small QN3 models up to 100K to And

if you are interested to learn more about it, then please visit the link, or you can scan the QR code to join our Discord to learn Discord community. that's basically all, since I didn't dive too deep into the technical implementation details in this talk because of the time limit. If you want to learn more, this blog post here goes in deeper into both parts And the first

one is done by the Thinking Machine Lab. Which is a blog post that talks why we should choose lower and basically you can just like with a full parameter reinforcement learning you can increase the learning rate to up to like 10 times you can basically achieve the same result as a full parameter. building trillion parameter reasoning reinforcement learning with 100% GPUs is uh I would that covers

the first part of this talk. And the root cause of layer off 3 Y it failed and how we fix it is the second part of the talk. And if you want to learn more about training inference mismatch can happen and uh and like if you want to learn more why it's like very like in engineering if intensive to fix those like to cover the both the

lower side and also Uh so that they exactly match uh you can go into this blog post. This talks about like they use the batch environment kernel. and basically it will like uh take more time to do the inference so it will have some efficiency loss. And thank you very happy to answer if you have any questions. Maybe Yeah.