Hey everyone, it's great to meet you all. It's really great to be here today. My name is Rhythm, and this is my co-founder Lyndon. Our third co-founder, Yash, couldn't make it today, but we're all very excited to be here. The three of us were previously researchers at OpenAI, and now we're bringing Frontier AI inside of enterprise at Applied Compute. Today, we're going to be talking about efficient reinforcement learning.

As some context on Applied Compute, we help enterprises build their own intelligence to power real work in their company. We think a lot about how to push AI beyond productivity into real automations that deliver ROI that is quantitative for the company. Once we build a system that is specialized to the way a company operates for a particular use case, we deploy it with a data flywheel so that it gets better over time the more you use it. Picture an in-house expert at a company who is always at the forefront of their field.

RL mechanically is the tool that we use in order to bring these out-of-distribution data sets in distribution for the models today. Yash, Lyndon, and I all worked on the RL effort at OpenAI in its early days, and we saw firsthand the power of RL in going and maximizing these public benchmarks. Now we're taking that a step further and helping enterprises go solve the problems they care the most about, sort of their private benchmarks.

So here's a very high-level overview of how Applied Compute RL helps LM acquire these reasoning and intelligence capabilities. Let's say that you have a data set of math problems and we pick four of them for an RL training step.

Then we'll take an open source model, say one of the GPOSS models or one of the Llama models, and we have the model attempt each of those four problems 100 times. So each of these 100 attempts is the model thinking through how it would get to the final answer and then ending off with the final answer itself. And these are many, many reasoning tokens in its thinking trajectory.

We can grade all of these answers, and when the model is correct, we can bias the model's weights to reinforce its thinking trace in that attempt. When it's incorrect, we can discourage the model from having that kind of behavior again. So in this fashion, as we do more and more training steps with batches of four problems, 100 attempts each, the model learns to reason and solve math problems, and it becomes really, really good at math. Of course, at Applied Compute, we're not really helping enterprises solve math problems, but this is kind of the mechanism by which we're able to teach the models to get really, really good at tasks that they care about.

So, as we mentioned, the type of RL work that we do at Applied Compute is actually quite different from the lab. So, these are some real life photos from the labs and a photo we took at the Applied Compute office the other day. They, you know, the labs do these big training runs over several weeks. We do more specialized runs. And you know, there's a couple of aspects of RL training that are particularly important to us.

We need our runs to be fast so that we can train a model and deliver it to a customer very quickly, on the order of days.

They have to be cheap so that our unit costs work and we're able to scale the business sustainably.

And importantly, and this is a point that I think, you know, it's easy to miss, we need our estimates for how long these training jobs will be to be very low variance because we don't want to just be generally fast. We want to be reliably fast when we work with customers.

And so the research problem for us that is very business critical is can we build an RL stack that is so efficient so that in conjunction with our agent building platform, we are really able to scale up this use case specific training motion.

So let's start with an inefficient form of RL, which is synchronous RL. In synchronous RL, sampling and training happen in lock step. So there's some simplifications here, but let's say that we want to train on batches of eight samples. That means we're going to wait for all eight samples to finish and basically finish completion before we start training. And then we're going to repeat this process again. As a result, we have a lot of idle GPUs that are waiting on that third straggler sample to complete.

So in other words, in synchronous RL, our step times are dictated by whatever sample takes the longest time in order to complete.

To illustrate why this is bad, we took 40 arithmetic problems, requested 32 samples each for each of them with Quen 30B, and we measured how long it would take for these samples to complete.

It turns out that 99% of the samples completed in about 40 seconds. It took another 80 seconds to get that last percent of samples to complete. It really has a long tail.

So, as you'd expect, if you look at the throughput chart, the GPUs are doing a lot of work at the beginning when all of the sampling requests are launched, but by the end, they're very, very underutilized because they're waiting on those last samples to complete. The technical term we use at Applied Compute is the GPUs are slacking. So synchronous RL is not an efficient way to use these GPUs.

In order to solve this problem, we need to break the condition that sampling and training need to happen in lock step. In other words, we need to allow training while we're sampling. This is called asynchronous RL. And there are many approaches to doing asynchronous RL. One that we particularly like is Pipeline RL from P et al.

We're going to make some simplifications here, but in asynchronous Pipeline RL, we dedicate some GPUs to sampling and some GPUs to training. The sampling workers never stop. They're constantly doing inference with high batch size. As samples complete, they get added to a queue for training, and the training workers pull a batch from the queue to train on. After a batch has been trained on, the training workers propagate the new model weights to all of the sampling workers for what's called an in-flight weight update. And this is really what differentiates Pipeline RL. The sampling workers might be in the middle of a sample, but their weights will still get updated if a training step just completed.

As a result, we end up with samples that had multiple versions of the policy that contributed to the sample in order to generate it. In other words, there are stale tokens in some of these samples. Let's take a look at one sample to make this a bit more clear.

As you can see, there's three versions of the policy at time steps t, t+1, and t+2 that were used to generate this sample since there were two completed train steps and in turn two in-flight weight updates while this sample was being generated.

So when this sample gets trained on in the T+3 to T+4 training batch, we will have some tokens that came from policy three steps behind, some that came from policy two steps behind, and those last two tokens that came from a policy that was one step behind.

Now, let's say that we only tolerate staleness up to two. That means we're not going to allow the in-flight weight update after the T+1 to T+2 training batch completes. And that means the training workers are just going to be idle waiting for this sample to complete before they can propagate that in-flight weight update and start training on the next batch. Because if they were to do the in-flight weight update, that would cause this sample to have staleness 3 as we just saw.

And if we only tolerate staleness one, the training workers are going to be idle for even longer, which is bad. So as you increase how much stale you tolerate, you have less idle GPUs in general. But as we all know, there's no free lunch. This is the standard policy gradient with an importance ratio to adjust for the fact that we're sampling from a policy at time step t and training with the policy at time step t+k given that there's k staleness. The importance ratio is what makes this policy gradient unbiased. But the variance of that ratio increases as you increase staleness. And so this is kind of the big issue here because now with higher variance importance ratio, learning can become unstable and cause divergence.

The concrete trade-off is we want a lot of staleness for fast RL runs, but a lot of staleness makes learning unstable, which then requires innovating on the algorithm and the science. And this is one of the primary research problems that we focus on here at Applied Compute. And as I was talking about earlier, it directly flows back into our core business.

For the purpose of this talk, we're going to focus on a simpler sub problem. Let's assume that we have good science and algorithmic innovations that allow us to tolerate staleness up to some fixed threshold and we have some fixed compute budget as usually exists in the world. What is the highest way for us to do RL in this setting?

Cool. Thanks Rhythm. So we posed this as a modeling problem of our end-to-end system, which you know admittedly is a little bit complicated at first, but we did find that we can get surprisingly far with some first principle systems modeling. And as with any modeling problem, let's figure out the cast of characters that describe the system and then we'll think about how they all fit together to model it. So the first cast member is some proxy of compute budget, in which in this case we have as the number of GPUs. In the synchronous setting, like Rhythm just explained, all the GPUs will either be used for training or sampling since they happen one after the other. But in the asynchronous setting, it's a little bit trickier because we can choose to allocate that pool of GPU compute as much as we want for training or as much as we want from sampling, and that leads to some design decisions.

The next is the training batch size, which is some proxy of the workload that we have on the overall system, and this is kind of an ML decision, but in short, what we have is a batch of problems, which is a subset of our data set. Let's say we have N math problems that we want to train on, and for each of these problems, we're going to sample N problems in parallel. So if the problems are really difficult, we might sample more to encourage some diversity in the samples to encourage the model to learn some potentially divergent strategies.

The next thing we need is some proxy of sampling throughput. And to get some intuition of what we should choose here as a modeling decision, let's look at how some modern inference engine surface requests. So in GPU memory, we have the model weights, the activations, and some runtime state called the KV cache in memory. And given this train model, we're going to run the forward pass several times where each forward pass samples the next token and then we'll write to the KV cache. And so what this model shows is that a principal estimate that we should do is we should find some way to measure the latency per GPU of the forward pass. And this ends up being a pretty good choice in practice because from the systems angle, the inference throughput that we choose is largely determined by the batch size that we perform sampling with. So what I've shown here in the red square is a batch of tokens that are all forwarded at the same time. And this sampling forward pass needs to be as large as possible to efficiently utilize the GPUs subject to the runtime constraint that we don't actually run out of memory in the KV cache.

So what we can then do is we can fit a latency curve as a function of batch size, and that latency curve will look something like this. You'll have some regime where it's memory bound, and when it increases, it becomes compute bound, and there's some functional form below. And to explain the details of why we chose this decision, what we have here is an equation that's based in the roofline model from systems. At lower batch sizes, which I've highlighted in yellow here, we don't have that much work to do because there isn't that much compute to do on the processor, and there's so many parameters you need to load in at the same time. And so, as a result, when you add incremental work, it doesn't really add that much latency to the overall system since the processor is so fast at doing math that we're just waiting on memory to stream parameters in from memory to the processor. But as the batch sizes begin to get larger, we then get bottlenecked by the processor. And the more we add to our batch, the slower the forward pass takes. And just for good measure, we have this sigmoid here that just sort of modulates the smooth transition at this hinge point here to show that there's a subtle transition from a memory bound computation to one that's more compute bound and bottlenecked by the processor.

The final cast member is some proxy of training throughput, and we chose to measure this on a per GPU basis. So in this case, the model takes in the training batch size. So the parameter we saw earlier, and we typically do this by fitting a proxy of our empirical workloads. The units here is how many each train, how many tokens per second each training GPU processes. So it needs to do the forward, the backward, and some optimizer steps.

So given these forecast members, we can then begin modeling the system. And the first idea we had, although Rhythm, you know, suggested that this might not be a great idea, we can think about how to use a synchronous setup. And this might be a good idea from first principles because we definitely meet the staleness constraint because we don't train on stale data and we always use the entire GPU fleet for either training or sampling, making sure that we're using efficient use of the hardware. Let's think about how to actually model this. There are two things we need to know. We need to know the batch size at which generation runs. And we also need to know the response length distribution to figure out how our training workload's going to work and also how long the sampling's going to take. And so what I'm showing here in this simulation is a couple of engines. Each square is a request being processed, and they get darker and darker as we make progress throughout the batch. And as they finish samples, they write to the queue. And on the right hand side is a time series metric, maybe something that you'd see in Graphana if you're monitoring production metrics. And what you can see is that the batch size begins very high, but it slowly goes down over time as it eventually goes to zero and all the samples complete. And we can finally run an optimization step. After the step completes, we run this in a loop and we move on to the next step. And so as a result, we can have the following sampling procedure. We do maximum tokens inference forward passes where maximum tokens is the total number of forward passes we do for the longest request. We use the fitted latency estimator to figure out how long that forward pass will take. And then the response length distribution will tell us how many responses to drop. And so what we're showing in this video here is this entire thing of the response length distribution that we feed into the latency estimator. At training time, we can compute the total number of tokens that we just sampled in the batch and divide by the total training throughput, which is just the number of GPUs multiplied by the per GPU training throughput. And so what we have here is a simulation of what this latency curve looks like. So we have the CDF of the response length distribution that tells us how many responses we should drop on the left and the latency curve on the right. And this roughly kind of tracks because as we add more GPUs, we'd expect the latency per step to go down.

The next idea, given that the synchronous setup might not be the most principled choice, as Rhythm showed, is an asynchronous setup. But it's not just as easy as just sort of provisioning the compute between training and inference because if we don't do this carefully, we might actually run into the idle GPU problem again. And to show this, let's illustrate two extremes of what this allocation problem looks like. Let's first look at one end of the spectrum where we provision way too many inference GPUs and not that many samplers. In this case, we're consuming from a queue much faster than we're actually producing from it because the sampling workers are producing work significantly faster than significantly slower than we can actually consume them. When the red square grays out, it shows that they're idle. And what this diagram should hopefully illustrate is that for a lot of the time, we're actually not using that, and that has the same problem of low GPU utilization in the synchronous case as shown earlier.

On the other end of the extreme, we can provision way too many sampling GPUs, in which case our production rate is way faster than the rate that we actually consume them in. So here we've doubled the number of overall sampling GPUs and have the number of training GPUs. As you can see, they produce samples at much more rapid of a rate. But this index here in each yellow square, which is the staleness count of each sample, goes up. And as time moves on, we get more and more stale. And so the samples get more and more kind of less, more and more transparent as a result. And we learn less from each individual sample. So let's think about how we can actually model this workload then to determine an optimal async workload. In this case, the picture looks a little bit different because in steady state, the batch size is relatively consistent compared to the synchronous setup where it kind of goes down over time. So on the right hand side here, we have the same time series metrics. But in this case, it's a little bit different because the yellow squares are always full because every time we complete a sample, a new sample goes in, and we can continue writing to the queue. And so that batch size with a little bit of wiggles just for good measure is like a is pretty consistent over the course of a run. Now obviously the caveat here is that this batch size will certainly go down as we, you know, as response lengths go up because we run out of cache, uh, KV cache, but that's kind of a separate story and actually our model accommodates for that because we're actually accommodating for a response length distribution.

We can then begin to figure out the optimal layout, and there's two kind of constraints that we have to satisfy now that we know that the generation batch size is roughly consistent throughout the course of a run. The first invariant that we need to have is that the production consumption rate are roughly equal. So on the left hand side of this equality, we have the training throughput, which is the number of training GPUs multiplied by the per GPU throughput, and then also we have the number of sampling GPUs multiplied by the sampling throughput, which is just the batch size multiplied by the latency to actually do a forward pass on that batch size. And the next thing is that given that Rhythm you indicated that if we have too much staleness, that can be bad from an ML perspective, we want to make sure that our max theoretical staleness or simulated staleness doesn't exceed what our ML can handle. And so here we have the max staleness on the left, which is equal to on the top how much time the longest request took in the batch, which is just the maximum number of tokens multiplied by the number of by the amount of time each token forward pass takes. And on the bottom here we have the length of a training step, which is the training batch size multiplied by the mean sequence length.

So the simulation here then will sweep through multiple different values of the number of training GPUs. And since we have a fixed pool of compute, that then implies a certain number of GPUs used for sampling. And for this number of sampling GPUs, we can compute the minimum steady-state generation batch size to make sure that we don't blow out of memory subject to our KV cache memory constraints and also such that we have maximum throughput on the sampling side. And the final thing is we want to prune out all simulations where the sampling throughput brings us over the maximum possible staleness. When we look at that simulation, we can run an end-to-end similarly parameterized by the response length. We see that this kind of roughly simulates a 60% speed up relative to our synchronous baseline, assuming that the GPU compute is optimally allocated between training and sampling.

As a result, when we sweep layouts within these constraints, this allows us to limit staleness, but also make sure that we have our runs running at maximal throughput without actually doing the run itself. And so this gives us insight to simulate different workloads before actually running them on the GPU because these runs can actually be fairly expensive. And so this allows us to ask answer scientific questions from first principles like what is the optimal configuration that we should have of our GPU compute if we made response lengths very long because often times when models learn via reinforcement learning, they begin to think for much longer and also what empirical throughputs we should target during our performance optimization. So this has been a really useful piece of technology for simulation has informed a lot of the systems and research design decisions that we make. Cool. Thanks for your time and find us afterwards to jam on some more RL research engineering together later. Thank you.