Efficient Reinforcement Learning – Rhythm Garg & Linden Li, Applied Compute

Applied Compute focuses on moving AI beyond productivity tools into real, ROI-driven automations for enterprises. A core component of this strategy is Reinforcement Learning (RL), which is used to specialize models for specific business use cases and create a data flywheel where the model improves with use. The RL mechanism involves taking a base model, such as Llama, and having it attempt a task (e.g., solving a math problem) many times. For each attempt, the model generates a full reasoning trace. These traces are graded, and the model’s weights are updated to reinforce correct reasoning paths and discourage incorrect ones. This process allows a general model to become a specialized expert at a task that an enterprise cares about.

The Business Case for Efficient RL

Unlike large lab environments that can afford long-running experiments, enterprise applications demand an RL stack that is:

• Fast: Models must be trained and delivered to customers on the order of days.
• Cheap: Unit economics must be favorable to ensure the business can scale sustainably.
• Reliable: Training time estimates must have low variance so that delivery timelines are predictable and dependable for customers.

This creates a business-critical research problem: building an RL stack efficient enough to scale use-case-specific training.

The Problem with Synchronous RL

A naive approach to RL is synchronous training, where sampling and training happen in lock-step. A batch of samples is generated, and only once all samples are complete does a training step begin.

This method is highly inefficient due to the “straggler” problem. The time for each step is dictated by the single longest sample generation in the batch. Empirical data shows that sample completion times have a long tail; for instance, 99% of samples might complete in 40 seconds, while the final 1% takes another 80 seconds. During this time, the vast majority of GPUs are idle, or “slacking,” leading to poor hardware utilization and slow overall throughput.

Asynchronous RL and the Staleness Trade-off

To overcome these inefficiencies, Applied Compute utilizes asynchronous RL, specifically an approach based on Pipeline RL.

In this paradigm, GPUs are partitioned and dedicated to either sampling or training.

1. Sampling workers continuously perform inference with high batch sizes, adding completed samples to a queue.
2. Training workers pull batches from this queue to perform optimization steps.
3. A key feature is the in-flight weight update. After a training step, the new model weights are immediately propagated to all sampling workers, even if they are in the middle of generating a sample.

This leads to a phenomenon where a single generated sample may have been produced by multiple versions of the model’s policy. These tokens generated by older policies are referred to as “stale.” This introduces a critical trade-off:

• High Staleness: Allowing for more staleness means training workers are rarely idle, maximizing GPU utilization and speeding up wall-clock time.
• Algorithmic Instability: The standard policy gradient algorithm uses an importance ratio to correct for the discrepancy between the sampling policy and the training policy. The variance of this ratio increases with staleness, which can destabilize learning and cause the model to diverge.

Innovating on algorithms that can tolerate high degrees of staleness while maintaining stable learning is a primary research focus.

Modeling the System for Optimal Performance

Assuming a fixed compute budget and a maximum staleness threshold tolerated by the algorithm, the challenge is to find the most efficient system configuration. This is framed as a modeling problem to simulate workloads before running expensive jobs.

The model is built on a few key components:

• Compute Budget: The total number of GPUs available.
• Workload: The training batch size.
• Sampling Throughput: Modeled by the latency of a single forward pass, which is a function of the inference batch size. This relationship is captured by fitting a latency curve inspired by the Roofline Model, which accounts for the system being memory-bound at small batch sizes and compute-bound at large batch sizes.
• Training Throughput: An empirically measured value of tokens per second per GPU.

Finding the Optimal Asynchronous Configuration

In an asynchronous setup, the key decision is how to allocate the GPU pool between sampling and training. An imbalance leads to inefficiency:

• Too many training GPUs: The training workers consume data faster than it’s produced, leading to idle trainers.
• Too many sampling GPUs: The samplers produce data much faster than it’s consumed. Samples become excessively stale in the queue, harming model performance.

The simulation aims to find the optimal balance by enforcing two constraints:

1. Balanced Rates: The token production rate from samplers must equal the token consumption rate of the trainers.
2. Staleness Bound: The maximum theoretical staleness (calculated as the time for the longest sample divided by the time for one training step) must not exceed the predefined algorithmic threshold.

By sweeping through different GPU allocations and pruning configurations that violate these constraints, the model can identify the optimal layout. This simulation predicts a potential 60% speedup over a synchronous baseline. This modeling approach provides critical insights for configuring production runs, guiding performance optimization, and answering scientific questions about system behavior under different workloads without incurring the cost of actual GPU time.

• Tags:
• Reinforcement Learning
• Feature
• LLM Training
• Asynchronous RL
• Systems for ML
• GPU Utilization
• Performance Modeling