GPT-OSS Agentic RL Training: A Practical Retrospective

TL;DR

Hugging Face and LinkedIn researchers have successfully unlocked agentic reinforcement learning (RL) for the GPT-OSS model by resolving critical instabilities in PPO on-policy integrity, attention-sink backward passes, and MoE memory materialization. These fixes enable GPT-OSS to serve as a stable backbone for multi-step decision-making agents that interact with environments and tools.

Agentic RL and GPT-OSS

Agentic RL differs from traditional single-turn RL by optimizing an entire decision-making process. Instead of producing a static response, the model learns to plan actions, invoke tools, and adapt behavior over multi-step trajectories. This requires a closed-loop system where the agent collects rollout trajectories, computes rewards, and updates its policy iteratively using algorithms like PPO or GRPO.

While GPT-OSS has shown performance comparable to OpenAI o3-mini and o4-mini, its suitability for agentic RL was previously unvalidated. Researchers used the verl training framework and tested the GPT-OSS-20B model on tasks including GSM8K, Retool (an agentic coding task), and verifiable instruction following.

Resolving PPO On-Policy Instability

Initial training runs exhibited exploding KL divergence and entropy with non-increasing rewards. The team identified a failure in Proximal Policy Optimization (PPO) on-policy integrity caused by the Mixture of Experts (MoE) architecture.

The MoE Log-Probability Mismatch

In pure on-policy PPO, the importance sampling ratio must be exactly 1. However, in MoE architectures like GPT-OSS, the gating network may route inputs to different experts between the forward pass used for rollout generation and the forward pass used for training due to floating-point differences or stochasticity. This results in a mismatch between the current log-probability and the old log-probability, falsely triggering the PPO clip and violating on-policy assumptions.

The Fix: The team implemented a log-probability substitution that overrides the computation when the environment is known to be on-policy (minibatch size equals global batch size), forcing the importance ratio back to 1 by setting old_log_prob = log_prob.detach().

Correcting Training-Inference Mismatch via Attention Sinks

Even after fixing PPO integrity, gradient norms continued to explode. The researchers identified a fundamental training-inference mismatch where the inference engine (SGLang) and the training stack (FSDP with FlashAttention-v2) produced different token-level probabilities.

The Role of Attention Sinks

GPT-OSS utilizes attention sinks—learnable scalar parameters that act as "virtual tokens" in softmax computation. These sinks allow the model to allocate attention mass to a learned parameter rather than forcing it onto content tokens, improving stability in streaming inference.

Implementation in FlashAttention v3

The team discovered that verl hard-coded FlashAttention v2, which does not support attention sinks, and that neither v2 nor v3 supported the necessary backward pass for sink gradients. To resolve this, they:

  1. Leveraged the forward pass from the vLLM FlashAttention fork.
  2. Implemented the backward pass to compute the sink gradient $\frac{\partial L}{\partial S_{h}}$.

Result: This fix led to substantially faster convergence and stable reward improvement across GSM8K, VerifyIf, and Retool tasks.

Scaling Memory Efficiency for Long Contexts

Agentic RL requires expanding context windows as environment feedback is appended to the trajectory. The team implemented two primary memory optimizations to prevent Out-of-Memory (OOM) failures.

Mitigating MoE Expert Materialization

Researchers found that the Hugging Face Transformers inference forward path duplicated hidden states for all experts, materializing extremely large tensors in GPU memory (e.g., attempting to allocate 180 GiB for a 20B model). They patched the implementation to use a more memory-efficient execution path that processes experts sequentially.

Sequence Parallelism with FlashAttention v3

To further reduce per-GPU activation memory, the team implemented sequence parallelism (context parallelism). This partitions the input sequence across devices, reducing the peak activation footprint.

Because the attention layer requires all tokens of a sequence to be present on one GPU, the team implemented an all-to-all communication strategy:

  • Pre-attention: Gather sequence elements, splitting at the attention-head level.
  • Post-attention: Redistribute outputs back to the original sequence-parallel layout.

This design is attention-sink-aware and compatible with FlashAttention v3, enabling the model to handle the long context windows necessary for multi-step agents.

Sources