Tokenless

Taking Reinforcement Learning Cross Datacenter — Nan Jiang, Modal

Nan Jiang introduces "Adam absorption" to revolutionize RL model synchronization. By exploiting finite precision serving and small Adam steps, less than 1% of served model weights actually change, allowing for 500MB patches instead of 500GB checkpoints. This enables a distributed "bulletin board" architecture, decoupling trainers from global rollout fleets and unlocking elastic, cross-region GPU capacity for RL.

6 min read
Taking Reinforcement Learning Cross Datacenter — Nan Jiang, Modal
source: https://www.youtube.com/watch?v=maRzp4kImJ4

This talk addresses the critical challenge of syncing large model checkpoints in Reinforcement Learning (RL) post-training, particularly when dealing with distributed GPU capacity and the need for rapid weight updates. A standard frontier-scale checkpoint can be around 500GB, making synchronization across regions time-consuming (minutes to hours) and hindering the ability to land weight updates in seconds. The speaker, Nan Jiang, proposes a novel solution where instead of shipping the full 500GB checkpoint, only about 500MB is transferred, allowing the rollout engine to reconstruct a bitwise identical weights version.

The Problem: Mismatch in Compute Architecture

The standard RL post-training loop tightly couples the trainer and rollout workers within the same cluster, often relying on fast RDMA interconnects. This architecture, described as a “cathedral” (one region, fast interconnect, many GPUs), is suitable for the trainer which requires collective operations like all_reduce and model-parallel communication. However, it severely limits the scalability and elasticity of the rollout fleet. Modern GPU capacity, akin to a “bazaar,” is distributed across different regions, providers, and price points. The default RL loop’s demand for a tightly coupled cluster (requiring sufficient GPUs, same region, fast fabric, and immediate availability) makes it hard to scale rollout workers elastically, unlike typical inference capacity.

Decoupling the RL Loop

The key insight is to separate what must remain in the tightly coupled training cluster from what can be distributed. Backpropagation and collective communication remain in the trainer’s cluster. The rollout fleet, however, can be decentralized. The “movable unit” is defined as a “rollout serving island” – a coherent endpoint or group of endpoints serving a single policy version. This allows for local parallelism within an island but maintains a lightweight interface across islands: trainer sends policy version, rollout sends trajectories and metadata.

The Flaw in Full Checkpoint Synchronization

When the rollout fleet is decoupled, the remaining link is the weight update. Shipping a full 500GB checkpoint over commodity links for every weight update is impractical for achieving sub-second latencies. This implies that the full checkpoint is the “wrong unit of synchronization.”

The Core Bet: Less than 1% of Weights Change

Nan Jiang’s central claim is that less than 1% of “rollout visible weights” actually change between consecutive versions. “Rollout visible weights” refer to the weights as served by the rollout engine (e.g., in BF16, FP8, or NVFP4 format), not the FP32 optimizer states or Adam moments. If this holds true, only the changed portion needs to be shipped, significantly reducing the data transfer size.

The Mechanism: Adam Absorption

This phenomenon is explained by “Adam absorption,” which arises from the interaction of small Adam steps and finite precision serving:

  1. Ingredient 1: Finite Precision Floor: Optimizers maintain high-precision (e.g., FP32) master weights, but rollout engines serve models in lower precision (e.g., BF16). BF16 has a finite resolution, meaning a “rounding boundary” (or “floor”) for value changes to be visible. For a weight around 1, this boundary is roughly theta / 256 (approximately 0.0039).
  2. Ingredient 2: Small Adam Steps: At typical RL post-training learning rates, the per-parameter Adam step (or “push”) is extremely small, often around 3 millionths. This push is more than a thousand times smaller than the BF16 rounding boundary.

Adam Absorption in Action: When the Adam step is smaller than the rounding boundary, the master FP32 weight moves, but the served BF16 value does not change because the update isn’t large enough to cross the rounding threshold. This makes the served update effectively sparse. It’s crucial to note that this is not due to sparse gradients; gradients are dense (99% non-zero), and the FP32 master update is also dense but simply small.

Lower precision serving exacerbates this effect. An internal run serving GLM 4.7 Air in FP8 showed only 0.15% of weights changing on the first step (when learning rates are higher), settling near 0.05% in subsequent steps.

Lossless Patching and Bitwise Reconstruction

Instead of a full checkpoint, a “lossless patch” is shipped. This patch contains the positions and replacement bits for the changed weights, along with metadata. Various lossless encoding methods (e.g., selective overrides, XOR) ensure that applying the patch reconstructs the exact same bitwise-identical served version that would have been obtained from a full sync. This avoids additive delta drift.

The Bulletin Board Architecture

To operationalize this, a bulletin board architecture is proposed:

  • The trainer, in its dedicated cluster, publishes immutable rollout weight versions to a shared “bulletin board” after each optimization step.
  • Rollout engines, now decentralized across regions and providers, pull these versions from the bulletin board.
  • A “sidecar” component on each rollout engine makes it “version-aware.” If the engine is at the correct version, the sidecar proxies requests. If it’s behind but can catch up via patches, it applies them. If it cannot, it signals “not ready.” This enables elastic rollout, where any idle GPU can be utilized.

This reduces the synchronization size from 500GB to 500MB, shrinking sync latency from minutes/hours to seconds.

Modal’s Implementation: Stitch

Modal’s implementation of this protocol is called “Stitch.” It is framework-agnostic, async-first, and agentic. Stitch allows rollout engines to autoscale globally, self-synchronize weights, serve version-specific requests, and return rollout metadata. This transforms scattered inference capacity into one elastic RL rollout fleet.

Open Questions

The presentation concludes with open questions:

  • Does Adam absorption sparsity still hold for other optimizers like Muon?
  • How scalable is fully asynchronous RL with this global compute model?
  • Can this paradigm generalize to other stages like pre-training, mid-training, or SFT?

Share post

Stay In The Loop! Subscribe to Our Newsletter.

Get updates straight to your inbox. No spam, just useful content.

Recommended for You

Reinforce Adjoint Matching: Scaling Diffusion RL

Reinforce Adjoint Matching: Scaling Diffusion RL

Andreas Bergmeister presents Reinforce Adjoint Matching (RAM), a novel method for scalable reinforcement learning post-training of diffusion models. RAM sidesteps the complexity of existing RL approaches by deriving a consistency loss that directly corrects the pre-training target with a reward signal, eliminating the need for costly SDE rollouts or reward gradients. This approach achieves superior performance and up to 50x faster convergence on challenging image generation tasks like complex composability, text rendering, and human preference alignment.

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

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

At Applied Compute, efficient Reinforcement Learning is critical for delivering business value. This talk explores the transition from inefficient synchronous RL to a high-throughput asynchronous 'Pipeline RL' system. The core challenge is managing 'staleness'—a side effect of in-flight weight updates that can destabilize training. The speakers detail their first-principles systems model, based on the Roofline model, used to simulate and find the optimal allocation of GPU resources between sampling and training, balancing throughput with algorithmic stability and achieving significant speedups.