Loading…
Unlocking Agentic RL Training for GPT-OSS: A Practical Retrospective
Hugging FaceJason Zhu, Hejian Sang, Arup De, Rohit Jain, Yanning Chen
Summary
Engineers evaluated the GPT-OSS model for agentic reinforcement learning across multi-step tasks such as ReTool and GSM8K using the verl framework. Initial training runs suffered from exploding gradient norms, increasing KL divergence, and stagnating rewards. Investigation revealed that non-deterministic expert routing in the Mixture of Experts architecture caused log-probability mismatches between forward passes, which falsely triggered Proximal Policy Optimization clipping. The team resolved this by substituting detached current log-probabilities for old log-probabilities during on-policy updates. To address further training instabilities and memory bottlenecks, the implementation integrated an attention-sink-aware backward pass into FlashAttention v3 and applied sequence parallelism using all-to-all communication around attention layers.
Context
Validating the open-source GPT-OSS model as a backbone for agentic reinforcement learning faced failures, as initial training runs in the verl framework experienced exploding gradient norms, surging KL divergence, and non-improving rewards.
Approach / What changed
The authors fixed on-policy PPO integrity by overriding old log-probabilities with detached current log-probabilities to avoid MoE routing discrepancies, integrated an attention sink backward pass into FlashAttention v3, and implemented sequence parallelism with all-to-all communication around attention layers.
Takeaways
- MoE gating differences between dual forward passes caused log-probability mismatches, which falsely triggered PPO clipping by pushing the importance sampling ratio away from 1.
- Overriding old log-probabilities with detached current log-probabilities when minibatch equals global batch size restores exact on-policy behavior.
- Sequence parallelism avoids non-attention layer synchronization and uses all-to-all communication at the attention-head level before and after attention layers to reduce per-GPU memory.
Related reading
Mastering Agentic Techniques: AI Agent Reinforcement Learning
Language model agents often fail across long-horizon workflows, repeat tool-call errors, and output invalid schemas when relying solely on prompting or retrieval. Reinforcement learning with verifiable rewards provides a practical training signal for domain-specific tasks by scoring model rollouts against deterministic verifiers such as schema checkers, unit tests, and execution sandboxes. Teams can implement group relative policy optimization to sample multiple trajectories per prompt and update policy weights based on relative performance within the group. Establishing environments with explicit harness, state, and verifier definitions allows agents to learn multi-step tool use, whereas static datasets fail to capture sequential dependencies. Inspecting checkpoint rollouts and testing against held-out tasks prevents reward hacking while continuously turning operational failures into verifiable benchmark environments.
Elizabeth Goodman