Vollständiges Transkript anzeigen (3.649 Wörter)
[00:12] I'm Jason. I'm from Modal, and today I'm talking about scaling agent RL across the globe.
[00:21] But first, I think a lot of AI engineers view models as kind of like static backends that um, that they kind of swap out. Like as the new models come out, they'll maybe test the models, run evals, uh do AB AB tests, and kind of like rebuild their harness around whatever model is cheapest, but also highest quality for for their use cases.
[00:50] And the claim I want to make for this talk is that this is actually a losing strategy when you end up having to compete on like, not only the agent economics, but also the experience of the agent, and like these two competing priorities of like offering differentiated value, uh.
[01:08] And I think this is what we hear from our customers who do fine-tuning. Um, it's like the first the first area where I think, um, fine-tuning took off is is people who had agent deployments, who um, need reliable SLOs. Like, like they they need reliability SLOs. If they if they need those, then they need to own their own inference, and uh the open-source models out of the box were not always capable of supporting their deployments uh without, you know, uh degrading in in quality.
[01:40] And I think fine-tuning, especially, you know, after these deployments became more mature, I think was the only path to owning your own inference. Uh, but more importantly, kind of like pushing yourself up this curve in terms of latency, quality, uh and cost all together. Like if you if you do fine-tuning, you will like own all these yourself, kind of. Um, there were other use cases though. I think one we heard which was interesting was uh proprietary datasets that their agents actually needed to reason about and understand and take action on, but uh these these datasets were like not present in any pre-training set. And so there- like out of the box, sometimes they would build agents that had to deal with this type of data, and Claude, GPT, just like would not work at all uh for these tasks.
[02:30] And so, they actually decided no other alternative but to fine-tune open-source models.
[02:37] And lastly, I think there's an emerging attitude around fine-tuning, which is um, in order to build differentiated value with agents, uh, like more and more people are deciding to vertically integrate across all the different levels of the stack. Um, these are the most common in industries where uh people are already competing on like some metric of task intelligence per dollar, like coding. But I think um what we should see as the industry matures outside of just coding, is like we should see this trend kind of continue elsewhere. Um, so I think that's- that's why I think a lot of AI engineers have not really been fully bought in because the the latest model comes out, and it's better than any model they could fine-tune from the last generation, but I think there's like a lot of different reasons people do fine-tuning, and in the long haul I think uh, yeah, this- this is only going to become more common.
[03:34] So you may have heard of fine-tuning in the context of RL or SFT, and the question is why is this talk about RL and not SFT? Um, I think the idea that I want to communicate here is that uh, if you're if you're doing SFT, you're basically training on traces recorded uh during production. Um, but the problem is that agents in particular have this like branching structure, uh based on the number of turns that they run. Um, so in production, if you train on this particular trace, uh later on when you deploy that model, it may see like a completely different branch in the kind of like session- like the- the tree of possible sessions, and uh this is quite bad. SFT is quite brittle in in this regard. So, kind of like why RL has emerged as this like dominant paradigm of fine-tuning, is that um, in RL the model actually explores all these branches uh during training time. At least that's kind of the hope. And when people talk about RL like generalizing better, that's often what they mean.
[04:45] I think one reason why this tradeoff has been hard for people is that they view like post-training and fine-tuning and RL as difficult. Um, but I think my claim here is actually that if you have an agent deployment in production, you probably already have the ingredients you need. Um, so if you're in production, you have an agent harness that that's deployed somewhere, and hopefully you have evals of some kind. Maybe calling it evals is not necessarily required. It could be AB tests, it could be um basically any way you do change management for this harness, for this deployment. Um, but those two together give you the tools you need to basically build out something that you can consider an environment, an RL environment.
[05:31] And why production traces are important for this is that you need to kind of get high-quality tasks to train on in this environment from somewhere, and the best place often to do that is to synthetically generate these tasks from actual live production traces of things that your users are- are doing, um with with your product, and yeah, kind of forming- finding examples where it's failing, mining those, and creating tasks around them so that fine-tuning actually gives you some chance of improving your product in those areas.
[06:05] And this is basically all you need to do RL.
[06:09] But then the next part about it being hard is that like, okay, just because you have an agent deployment doesn't mean that your team may be capable, or maybe that's how you feel. Um, and probably because RL has a reputation of just being difficult generally. Um, but I think over the last couple of- at least over the last year, over the last several months, uh like we have kind of figured the science of of RL, and that's what I'm going to basically talk about today is that uh, a particular RL recipe that we've had a lot of success with, um that I think simplifies and stabilizes a lot of what you know, was difficult before.
[06:49] But first, okay, just to kind of recap how- how does reinforcement learning work? Uh, it's it's this loop, mechanically this is- this is what's happening. You have a bunch of uh inference GPUs that are running your sampler, aka the model that you're training. Um, and those rollouts you- you basically generate rollouts in your environment, uh with your harness. Those rollouts get collected and fed to a trainer, which is consuming all this data, and updating the model uh with, you know, gradient descent. And then uh that updated model, you- you broadcast those new weights to all the samplers, and um over time you- the rollouts you're generating are from this like continuously improving model, uh so you're- you're- you're kind of like going back and forth between them, and iterating, and and that's how the model improves.
[07:40] Uh so it turns out that like the math of RL prefers this to be done sequentially. So first you do your rollouts, then you train and update your model, then you- you take that model update, send it back to your rollout side so that the new rollouts you generate are from that new version of the model. Uh but the problem is that this wastes a ton of idle compute. Um, there are ways to like share the GPUs between the two, but I think uh especially with agents, um what's much better is that if you actually are continuously doing both in parallel, asynchronously. Um, this is better because you can like pack your GPUs, get high utilization from them, and uh not waste money. But the downside is that um, your rollouts that you generated were very likely from an old copy of the model, and this is what people refer to as staleness in asynchronous RL. Basically, uh rollouts are generated- generated by V1, uh the current model that you're training is V3, are those rollouts still useful for V3 if the model's gotten better since then?
[08:44] And this is like one of the reasons why it was unstable. Another reason why it was unstable is actually purely at the systems level. There's a lot of details that can show up as differences in between what the trainer thinks it's training and what the samplers think we're training. Uh basically, at all these different levels of the stack, there can be these like small differences in behavior. And that means that like in- in the math of RL, I think the- the idea is that both- both sides of this think they're doing the same thing. They think they're learning from the same data, training like the same distribution, um but actually it's like slightly different.
[09:21] And so the- the solution that I'll present here is accounting for these differences, um broadly with like kind of the simplest set of fixes that will uh correct for this mismatch without drowning out the signal of the- the training loop.
[09:37] So yeah, what's the recipe? Here it is. Uh we start from GRPO, this is uh what DeepSeek introduced for DeepSeek R1 and since then has become kind of a cornerstone of of doing RL on language models. We add something called Masked Importance Sampling. This is correcting for token-wise mismatch. And this mismatch is not just the mismatch from all the systems-level stuff that I pointed out, uh it's also from the staleness as well. So, if an old version of the model would have sampled a different token here, uh the ratio between the two can be can be very different, and then uh this is a mismatch that we want to correct for somehow. And and uh I won't go too much into detail about how these mechanisms work, but uh I will, you know, reference them.
[10:24] And then, the second is score centering. So, even though we're correcting for each token-wise mismatch, what ends up happening is over the course of training, whatever mismatch is still remaining, accumulates, and it kind of accumulates in the- the weights themselves, like the latest version of the model is slightly different than uh what it should be, even though we've like corrected for the most obvious differences. Um, and so this score centering approach is it basically removes this drift over the course of training, uh not just at each token or at each step. Um, it it kind of like builds up a term that accumulates with the drift, and removes it.
[11:02] And then the last thing is top-p sampling, uh with mask replay. So, top-p is- is one of the parameters you use when you sample from a language model, basically says only consider the top-p of the of the distribution, P being like 0.98, 0.95, basically, only the tokens that have probability mask greater than 95- uh or greater than 5%. Um, and so it just removes a bunch of very rare, likely bad tokens that can just never be sampled from. And so if your sampler decides to only sample from these tokens, but your trainer decides, unluckily, to sample one of these error tokens, this mask- mismatch can be basically catastrophic if you don't correct for it.
[11:42] So we correct for it. Um, and it turns out actually, in the math of score centering, that if you're doing top-p sampling like this, that actually the- the technique we we use to do score centering uh is exact instead of approximate. Because if you don't do top-p, it's actually only an approximate correction for drift. So, all this is to say that uh with this recipe, uh things get much, much better, much more stable.
[11:58] And this is in- this basically just an ablation that that we did on Qwen 3.6 35B uh the MoE. Um, it's trained on the Mimo code environments with evals on Swe-bench Pro, and the Hard 51 subset. Um, basically what this shows is an ablation of these different techniques, uh putting them all together showing like much more stable training uh than any of uh the alternatives that we've tested here.
[12:31] So yeah, I think the- that's all well and good, but I guess the question is like how far can we push this recipe? Does it actually scale? Does it- does it work for broad problems? Uh, before I talk about that, I think I want to talk briefly about uh why we care about RL at Modal, and like some things about doing RL on Modal are different. I kind of want to just like get into that a little bit first.
[12:52] So yeah, first why Modal for agent RL in particular? Well, we do training. We just released uh clusters uh as GA, uh last week I think, or week before. Um, we also do inference really well. We have a bunch of endpoints, uh we've- we've done a lo- a lot of the hard work around doing LLM inference at scale.
[13:10] And the environments you need, they run in Modal sandboxes. So, like naively, we have all of the kind of infra primitive layers like- or infra primitives that you'd want to be able to compose into like a full production agent RL training system.
[13:25] The problem- the only problem, actually, uh because all that sounds great, but the only problem is that our capacity model is very different compared to other people. The reason for that is we aggregate tons of capacity across many different providers. Here there's three pictured, but in reality it's like 15 or 20 at this point. Um, so we- we run on all the hyperscalers, but we also run on a bunch of Neoclouds, um and we handle all of the kind of like GPU health and reliability concerns.
[13:56] But this is great, but uh it's great for reliability, it's great for kind of like fast and flexible scheduling, it gets us around compute shortages, uh but it also fragments our compute pool. So, uh if you want to schedule large training clusters, it's actually quite difficult in this like multi-provider, multi-region fleet. Uh we've engineered it so that we can schedule clusters pretty quickly up to like 4 nodes or 8 nodes, uh but once you get past that to 16 nodes or 32 nodes, uh it gets pretty painful, it takes a really long time to find like a contiguous block of all these nodes.
[14:30] And in RL, this is particularly bad because, in RL we generally have, especially for agent RL, we have like 2, 2 inference GPUs for every training GPU. So 2/3, ish, of your cluster is actually just stuck doing inference. Um, so if you're scheduling an 8-node cluster, max, you're probably only using 2 or 3 of those nodes for the trainer side and most of it's just spent doing inference.
[14:53] It'd be really great if we could move these inference GPUs out of the cluster.
[15:00] And one reason why we think we can do this is because uh during RL, the weight updates, the actual like gradients that you're applying to your model, end up being incredibly sparse, and there's uh two reasons for that. But um, yeah, essentially if you- if you take the model at step T and subtract the- the new model T minus- or sorry, take the model at T plus one, subtract T, uh you get this delta, and it turns out that even though the full model weights at each step are very dense and very large, uh the updates themselves are incredibly sparse, 100x or 1000x uh as sparse.
[15:37] So you can actually just feed this to off-the-shelf compression, and it turns out that the amount of data you need to communicate from the- the trainer- uh yeah, from the trainer to the rollout engines, goes down by a factor of 1000.
[15:50] This means that you can actually move your inference GPUs out of the cluster.
[15:55] And so actually for the- for the run that I mentioned, uh of Qwen 3.6, we ran the training cluster, I think it it landed somewhere in GCP, but the inference GPUs actually ran all across our fleet and all the different hyper-scalers and Neoclouds that we run across. Um and it turns out with this sparse delta compression technique, you can do this, uh it works well and it does not introduce like bad performance.
[16:21] Not only that, we- we actually fragment our rollout fleet across different GPU types. So, again, because the- we have fragmentation due to all the different providers we run on. Uh but we also have fragmentation because we have a lot of different GPUs that different customers want to use. Uh it'd be really great if we could leverage as many of them as possible for RL runs and that's- that's what we've done.
[16:43] Each of these GPU types, they- they increase the mismatch potentially. So if the trainer is running in B200, uh but your rollouts are running in H100 or A100, there are hardware-level differences in the- the kernels, the math you can run, that implement the model and so this like increases more and more mismatch which traditionally should cause your training runs to destabilize.
[17:04] This is just kind of like a graph to show uh how many requests are being processed by each GPU type in the run we did, over the course of training.
[17:14] So actually, the recipe that I mentioned uh is capable of handling this amount of mismatch that we throw at it, which uh yeah, I think is like kind of unprecedented. I think this just goes to show how much we've figured out the science of RL over the past year or two. Um basically this shows that you have up to like 4, 5 times as much mismatch, uh but the training reward and gradient norm, everything is- is training as stably as it was uh in the homogeneous case, which is basically running training on B200s and inference on B200s.
[17:55] And kind of a bonus, like a side effect of this is that uh the rollouts that we run for this full training run are actually 15% cheaper. And that's because the the demand for the high-end GPUs like B200s uh means that the price is just generally higher. But if you can actually run your rollouts cheaply on on cheaper hardware that there's less demand for, like the overall cost of doing this run is is significantly lower. This is actually without any inference tuning for each GPU type. We're just doing like the basic off-the-shelf inference recipe for each GPU. If you actually decided to uh you know, optimize each individual deployment for each GPU type with for example, different levels of quantization in the model, um you could actually probably make this significantly more cheap. We just uh this is not the focus of this work, but uh but yeah.
[18:44] Oh, I did also- maybe I forgot to mention, but in this case, there are cases where like, for example, H200, we have two different pools of H200s. That's because we run two different uh quantizations of the model uh on H200s. So we have one H200 pool doing it at BF16, which is like the the safest way to do it, but we have another running at FP8, uh and again like the the correction that we do can like drown out the mismatch that's introduced by even excessive quantization of the model, and still learn from these quantized rollouts.
[19:16] So yeah, I think the- the point I'm trying to make here is that uh like we've released this, there's a blog post that where we go into like all the technical detail. We'll be releasing the source code for it very soon. It's based on Mials and SG Learn, the the reinforcement learning framework and the inference engine. And um, I think the- the takeaway is that um, I do believe that if you want to be competitive in a world where a bunch of different competitors are running agent stacks, you do need to to some degree own the full stack yourself, or at least like build competency in every area of the stack.
[19:51] I think a lot of work traditionally ends at this agent harness, but actually, um the vertical integration across all of this is just extremely important in a- in a hyper-competitive uh industry. We see this again in certain industries like coding, um but I think it we- we also see it in other industries and it's only going to keep growing, I think.
[20:13] Um and of course like, if you're building agent deployments, you own everything up the stack as well, and vertically integrating across all of this is extremely important uh for maintaining kind of competency and- and competitive edge.
[20:27] Um yeah, you should own your own agent RL, and it's actually possible now because the science has been figured out, and there are very good tools and very good infrastructure that lets you do that.
[20:38] That's all I got, so thanks. We're working on hard problems. Come visit our booth and- and talk to us.
Video
Video
