{
  "@context": "https://schema.org",
  "@id": "https://aggr.aymericbeaumet.com/items/huggingface-co/2026-01-27-unlocking-agentic-rl-training-for-gpt-oss-a-practical/#webpage",
  "@type": [
    "WebPage",
    "ArchiveComponent"
  ],
  "articleBody": "Agentic reinforcement learning (RL) extends traditional LLM training by optimizing not just a single-turn response, but an entire decision-making process learned through direct interaction with an environment during training. Unlike traditional single-turn reinforcement learning or offline preference-based methods that rely on static datasets, agentic RL trains policies by actively collecting on-policy data as the agent plans actions, invokes tools, observes outcomes, and adapts its behavior over multi-step trajectories in either simulated or real environments. This interaction-driven optimization assigns credit across long-horizon decisions, where intermediate choices such as query reformulation, tool selection, and execution order directly influence downstream success. Training follows an iterative closed loop in which the agent interacts with the environment to collect rollout trajectories, computes rewards over these trajectories, updates the policy based on observed outcomes, and then uses the updated policy to drive the next round of interaction and data collection such as GRPO or PPO algorithms..\n\nLinkedIn is an AI-first company that's built agents to help professionals be more successful. In this setting, models must reason over incomplete information, interact with structured services, and adapt to evolving user intent across multiple steps rather than produce a single static response. These capabilities are especially critical for agents that support the goals of recruiters, job and knowledge seekers, and learners end users, such as retrieving information, refining queries, coordinating tools, and executing multi-step workflows. By learning robust decision policies through interaction, agentic RL provides a principled foundation for building scalable, reliable, and adaptable AI systems through end-to-end optimization.\n\nThe GPT-OSS model has shown comparable performance to OpenAI o3-mini and o4-mini \\[[ref](https://openai.com/index/introducing-gpt-oss/)\\], but its suitability for agentic reinforcement learning training has not yet been validated. Most recent work focuses on fine-tuning without tool calling, such as: [Fine-tuning with gpt-oss and Hugging Face Transformers](https://cookbook.openai.com/articles/gpt-oss/fine-tune-transfomers) and [unsloth tutorial: how to fine-tune gpt-oss](https://unsloth.ai/docs/models/gpt-oss-how-to-run-and-fine-tune/tutorial-how-to-fine-tune-gpt-oss). This blog explores the journey to unlock [agentic RL](https://github.com/volcengine/verl/issues/3794) training for GPT-OSS as a potential backbone model for agentic applications.\n\nIn our experiments, we use [verl](https://github.com/volcengine/verl) as our training framework since it is one of the most popular adopted frameworks in the open source community. We use gsm8k, [Retool](https://github.com/verl-project/verl-recipe/tree/21892b9276936efab5375c3f6b8415e472ef7118/retool) task, [verifiable instruction following task](https://arxiv.org/abs/2507.02833), which are commonly used in RL training. We focus on presenting experimental results for the GPT-OSS-20B model, and our attention-sink fix also works for GPT-OSS-120B. The Qwen-2.5-32B model is additionally used to benchmark standard metric trends during RL training.\n\n## Challenges of GPT-OSS RL Training\n\nverl has been an OSS framework used by the team, and the team has previously collaborated and contributed to it to help democratize agentic reinforcement learning training. With the introduction of the new Harmony chat template in GPT-OSS, the first step is to ensure that the training framework fully supports the updated message format and conversation semantics required by Harmony. This step helps rollout generation, trajectory construction, and tool parsing remain consistent and correct under the new template.\n\nThe team uses ReTool as a representative example to verify code correctness. ReTool is an agentic coding task in which the model is asked to solve a math problem with the assistance of a code compiler tool. This setup allows the model to focus on core reasoning and algorithmic logic, while delegating the actual arithmetic and execution to the tool. During an episode, the model interacts with the code tool multiple times, using execution results as feedback to refine its solution. At the end of the trajectory, the model produces a final answer, on which the reward is computed.\n\nDuring the initial training runs, we observed exploding KL divergence and entropy, along with non-increasing rewards, indicating underlying issues in the GPT-OSS training setup, as shown in Figure 1.\n\n| Average Gradient Norm                                                                                                                                                                                                                                      | Average Reward                                                                                                                                                                                                                                      |\n| ---------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | --------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- |\n| [![Average gradient norm in a batch](https://cdn-uploads.huggingface.co/production/uploads/64efbd469e7770db74cb72f5/5eWLHQ-EAWKjPc7T0c6CK.png)](https://cdn-uploads.huggingface.co/production/uploads/64efbd469e7770db74cb72f5/5eWLHQ-EAWKjPc7T0c6CK.png) | [![Average reward in a batch](https://cdn-uploads.huggingface.co/production/uploads/64efbd469e7770db74cb72f5/ZpBvRXLKyAT-PIdsuRi11.png)](https://cdn-uploads.huggingface.co/production/uploads/64efbd469e7770db74cb72f5/ZpBvRXLKyAT-PIdsuRi11.png) |\n\n**Figure 1.** Left: Qwen32b has significantly higher rewards compared to GPT-OSS 20B; Right: The gradient norm exploded as training progressed.\n\n## A Practical Debugging Journey in verl: Restoring PPO On-Policy Integrity\n\n### Restoring PPO On-Policy Integrity: A Fix for MoE Log-Probability Mismatch\n\n![](https://cdn-uploads.huggingface.co/production/uploads/64efbd469e7770db74cb72f5/21xIlXKeAlTK5qKSp-TkX.png)\n\n**Figure 2.** Non-zero importance sampling clip value even for on-policy training.\n\nWe focus on on-policy methods because they provide greater stability and more reliable convergence. The foundation of pure on-policy Proximal Policy Optimization (PPO) mandates that the importance sampling ratio must be exactly 1. The mathematical definition of the importance ratio is:\n\nratio=π(a∣s)πold(a∣s) \\\\text{ratio} = \\\\frac{\\\\pi(a \\\\mid s)}{\\\\pi\\_{\\\\text{old}}(a \\\\mid s)}\n\nThis requirement ensures that the policy update is executed only on the data generated by the current policy π(a | s) = πold(a | s), preventing unintended clipping.\n\nWe have observed the non-zero clipping value in our ReTool training, as shown in Figure 2, stemming from a mismatch between the two log-probabilities:\n\n- Current log-probability `log_prob`: log(π(a | s))\n- Old log-probability `old_log_prob`: log(πold(a | s))\n\n**Root Cause: The Dual Forward Pass and MoE Architecture**\n\nPrior to verl 0.3.0, the implementation relied on two separate forward passes (one to compute the current `log_prob` and one to retrieve the stored `old_log_prob`) for the same state-action pair.\n\nIn a Mixture of Experts (MoE) architecture like GPT-OSS, the gating network routes the input to different experts. Due to implementation factors (e.g., subtle floating-point differences or explicit stochasticity), the expert routing can differ slightly between the two passes. Readers who are interested can further read *[Stabilizing MoE Reinforcement Learning by Aligning Training and Inference Routers](https://arxiv.org/pdf/2510.11370v1)*.\n\nThis difference in routing leads to:\n\nlog⁡(π(a∣s))≠log⁡(πold(a∣s)) \\\\log(\\\\pi(a \\\\mid s)) \\\\neq \\\\log(\\\\pi\\_{\\\\text{old}}(a \\\\mid s))\n\nThe resulting ratio deviates from 1, falsely triggering the PPO clip and violating the core on-policy assumption.\n\n**Solution: Enforcing Ratio = 1 via Log-Probability Substitution**\n\nThe fix resolves the issue by logically overriding the flawed computation when the environment is known to be on-policy (i.e., when the minibatch size equals the global batch size):\n\n```python\nif on_policy:\n    old_log_prob = log_prob.detach()\nelse:\n    old_log_prob = model_inputs[\"old_log_probs\"]\n```\n\nBy setting `old_log_prob` equal to the newly computed `log_prob` (detached to prevent gradient flow through the reference value), the importance ratio is mathematically forced back to 1. This strategy bypasses the instability caused by MoE's non-deterministic routing and guarantees strict on-policy behavior during PPO training.\n\n* * *\n\n### Correcting Training–Inference Mismatch\n\nAlthough fixing the log-probability mismatch reduced the importance-sampling clip ratio to zero, gradient norms continued to explode and rewards failed to improve. To isolate the issue, we simplified training to GSM8K, a single-step task without agentic tool use. The same instability persisted, as shown in the green curves in Figure 3, indicating a **fundamental issue in basic RL training with GPT-OSS under verl.**\n\nWe hypothesize that **training–inference mismatch** could be a potential cause: discrepancies between inference-time execution—where engines such as vLLM and SGLang aggressively optimize for throughput—and training-time execution under FSDP, which prioritizes numerical precision and stability, can effectively turn otherwise **on-policy RL into off-policy optimization.**\n\nThis [blog](https://yingru.notion.site/When-Speed-Kills-Stability-Demystifying-RL-Collapse-from-the-Training-Inference-Mismatch-271211a558b7808d8b12d403fd15edda) details why such mismatches lead to unstable gradients and non-improving rewards. Figure 3 compares training runs with and without rollout correction (see this [verl blog](https://verl.readthedocs.io/en/latest/algo/rollout_corr.html) for details). After applying rollout correction, training dynamics improve significantly, with gradient norms remaining stable rather than exploding.\n\nHowever, as shown in the left plot of Figure 4, the reward increases only modestly, and convergence on the simple GSM8K task remains substantially slower compared to smaller dense model variants.\n\n| Average Entropy                                                                                               | Average Gradient Norm                                                                                         | Average KL Loss                                                                                               |\n| ------------------------------------------------------------------------------------------------------------- | ------------------------------------------------------------------------------------------------------------- | ------------------------------------------------------------------------------------------------------------- |\n| ![](https://cdn-uploads.huggingface.co/production/uploads/64efbd469e7770db74cb72f5/MoBnatEipCi_OScgm5fAJ.png) | ![](https://cdn-uploads.huggingface.co/production/uploads/64efbd469e7770db74cb72f5/k4Su03zqq6Cg-7jM6Il6J.png) | ![](https://cdn-uploads.huggingface.co/production/uploads/64efbd469e7770db74cb72f5/g77U9c8a-FY1rvwru96Ez.png) |\n\n**Figure 3.** Gradient norm behavior under different training configurations. Green: Training without rollout correction, exhibiting unstable gradients. Red: Training with the attention layer frozen to isolate the issue to the attention mechanism, resulting in partial stabilization. Blue: Training with rollout correction enabled (sequence-level importance sampling), yielding stable gradient norms.\n\n| Average Reward                                                                                                                                                                                                                                      | Max Log-Perplexity Difference                                                                                                                                                                                                                                                                                              |\n| --------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | -------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- |\n| [![Average reward in a batch](https://cdn-uploads.huggingface.co/production/uploads/64efbd469e7770db74cb72f5/ihx-XsWH51V0-JM46jODE.png)](https://cdn-uploads.huggingface.co/production/uploads/64efbd469e7770db74cb72f5/ihx-XsWH51V0-JM46jODE.png) | [![Maximum absolute log-perplexity difference in a batch between rollout policy and training policy](https://cdn-uploads.huggingface.co/production/uploads/64efbd469e7770db74cb72f5/fP5KNR2XYY7EH-muYBk_X.png)](https://cdn-uploads.huggingface.co/production/uploads/64efbd469e7770db74cb72f5/fP5KNR2XYY7EH-muYBk_X.png) |\n\n**Figure 4.** Left: Reward improvement on GSM8K remains slow even after applying rollout correction, with performance comparable to runs where the attention layer is frozen during training. Right: A substantial log-ppl mismatch is observed between the inference engine (SGLang with Triton kernels supporting attention-sink forward passes) and the training stack (FSDP with FlashAttention-v2), indicating a large training–inference inconsistency.\n\nTo further isolate the root cause, we freeze the attention layers during training and observe reward dynamics similar to those of runs without freezing (blue curve vs yellow curve in Figure 4). This indicates that learning is primarily driven by the MoE layers, while the attention mechanism contributes less effectively than expected. In addition, we observe a substantial token-level probability mismatch between the inference engine and the distributed training stack which are using different attention kernels. Together, these observations motivate a deeper investigation into the attention mechanism.\n\n* * *\n\n## Attention Sink Support in FlashAttentionV3\n\nAttention sinks used in GPT-OSS are learnable scalar parameters (one per attention head) that act as \"virtual tokens\" in the softmax computation. They allow the model to allocate attention mass to a learned sink rather than forcing all attention to content tokens, which has been shown to improve attention stability in streaming inference and training with sliding-window attention.\n\nAfter a deeper investigation, we identified several major issues:\n\n- verl hard-codes FlashAttention v2 in `fsdp_worker`, which does not support attention sinks.\n- The attention sink backward pass is not supported in FlashAttention v2 and v3, so it does not work as expected even when FlashAttention v3 is enabled.\n- Since the forward pass has not yet been merged into the original FlashAttention v3 repository, we leveraged the forward pass from the vLLM FlashAttention fork ([PR #75](https://github.com/vllm-project/flash-attention/pull/75)) and implemented the backward pass to compute the sink gradient.\n\n### Standard Attention\n\n```python\nscores = QK^T / sqrt(d)               # [B, H, N_q, N_k]\nprobs = softmax(scores, dim=-1)       # Σ_j P_ij = 1\noutput = probs @ V                   # [B, H, N_q, d_v]\n```\n\n### Attention with Sinks (GPT-OSS)\n\n```python\nscores = QK^T / sqrt(d)                               # [B, H, N_q, N_k]\ncombined = concat([scores, sink_param], dim=-1)       # [B, H, N_q, N_k+1]\nprobs = softmax(combined, dim=-1)                     # Σ_j P_ij + P_sink = 1\nprobs_content = probs[..., :-1]                       # Drop sink component\noutput = probs_content @ V                           # [B, H, N_q, d_v]\n```\n\n**Key difference:** The sink participates in softmax normalization but doesn't contribute to the output.\n\n### Mathematical Formulation\n\nThe attention weight for content token j in row i is defined as:\n\nPij=exp⁡(Sij)∑j′=1Nkexp⁡(Sij′)+exp⁡(Sh) P\\_{ij} = \\\\frac{\\\\exp(S\\_{ij})} {\\\\sum\\_{j'=1}^{N\\_k} \\\\exp(S\\_{ij'}) + \\\\exp(S\\_h)}\n\n**Where:**\n\n- Sij = Qi Kj⊤ / √d are the attention scores\n- Pij are the attention weights for the content tokens\n- Sh is the learnable sink parameter for head h\n\n**Sink Probability:**\n\nThe sink probability is computed but not used in the output:\n\nPi,h=exp⁡(Sh)∑j′=1Nkexp⁡(Sij′)+exp⁡(Sh) P\\_{i,h} = \\\\frac{\\\\exp(S\\_h)} {\\\\sum\\_{j'=1}^{N\\_k} \\\\exp(S\\_{ij'}) + \\\\exp(S\\_h)}\n\n### Backward Pass\n\nThe gradient of the loss L with respect to the sink parameter Sh is:\n\n∂L∂Sh=−∑iPi,h(∂L∂Si,h−∑j∈{1,…,Nk}Pij∂L∂Sij) \\\\frac{\\\\partial L}{\\\\partial S\\_h} = - \\\\sum\\_i P\\_{i,h} \\\\left( \\\\frac{\\\\partial L}{\\\\partial S\\_{i,h}} - \\\\sum\\_{j \\\\in \\\\{1,\\\\ldots,N\\_k\\\\}} P\\_{ij} \\\\frac{\\\\partial L}{\\\\partial S\\_{ij}} \\\\right)\n\n**Where:**\n\n- Pi,h is the sink attention probability for row i\n- ∂L/∂Sij is the gradient with respect to the attention scores, including the sink\n\n**Simplified Gradient:**\n\nSince the sink is computed but not used in the output, its gradient ∂L/∂Si,h = 0.\n\nTherefore, the backward equation simplifies to:\n\n∂L∂Sh=−∑iPi,h(∑j∈{1,…,Nk}Pij∂L∂Sij) \\\\frac{\\\\partial L}{\\\\partial S\\_h} = - \\\\sum\\_i P\\_{i,h} \\\\left( \\\\sum\\_{j \\\\in \\\\{1,\\\\ldots,N\\_k\\\\}} P\\_{ij} \\\\frac{\\\\partial L}{\\\\partial S\\_{ij}} \\\\right)\n\nThe forward pass was adapted from vLLM's FlashAttention fork, and we implemented the backward pass to compute gradients for the sink parameters. The implementation will be released following the internal review process.\n\n* * *\n\n### Results\n\nAfter applying the fix in FlashAttention v3, we observe substantially faster convergence for GPT-OSS-20B across a range of reinforcement learning tasks. These include single-turn RL on math reasoning (GSM8K — red curve in Figure 5), instruction following (VerifyIf, evaluated on an out-of-domain multi-if benchmark — Figure 6), and multi-turn agentic RL with tool use (ReTool — Figure 7).\n\nAcross all settings, training becomes stable and exhibits steady reward improvement.\n\n![](https://cdn-uploads.huggingface.co/production/uploads/64efbd469e7770db74cb72f5/6TYGevydK99nQ-I1QTouf.png)\n\n**Figure 5.**. Single Turn GSM8K, the red curve converges much faster than the rest without the fix\n\n| Average Entropy                                                                                                                                                                                                                                      | Average Gradient Norm                                                                                                                                                                                                                                      | Average Reward                                                                                                                                                                                                                                      |\n| ---------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | ---------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | --------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- |\n| [![Average entropy in a batch](https://cdn-uploads.huggingface.co/production/uploads/64efbd469e7770db74cb72f5/ydmBLCSGlD9YKWiIocI1S.png)](https://cdn-uploads.huggingface.co/production/uploads/64efbd469e7770db74cb72f5/ydmBLCSGlD9YKWiIocI1S.png) | [![Average gradient norm in a batch](https://cdn-uploads.huggingface.co/production/uploads/64efbd469e7770db74cb72f5/_kl8mn_CXPsRYJ467IbFs.png)](https://cdn-uploads.huggingface.co/production/uploads/64efbd469e7770db74cb72f5/_kl8mn_CXPsRYJ467IbFs.png) | [![Average reward in a batch](https://cdn-uploads.huggingface.co/production/uploads/64efbd469e7770db74cb72f5/vb2JmmSu-LI5szC_84KsM.png)](https://cdn-uploads.huggingface.co/production/uploads/64efbd469e7770db74cb72f5/vb2JmmSu-LI5szC_84KsM.png) |\n\n**Figure 6**. On verifiable instruction following the task, the run without the fix collapsed (blue), and the run with fix showed steady reward improvement.\n\n| Average Gradient Norm                                                                                                                                                                                                                                      | Average Reward                                                                                                                                                                                                                                      | Validation Accuracy                                                                                                                                                                                                                                                |\n| ---------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | --------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | ------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------ |\n| [![Average gradient norm in a batch](https://cdn-uploads.huggingface.co/production/uploads/64efbd469e7770db74cb72f5/_Xz7_RLhYuYhzGAeMjDXs.png)](https://cdn-uploads.huggingface.co/production/uploads/64efbd469e7770db74cb72f5/_Xz7_RLhYuYhzGAeMjDXs.png) | [![Average reward in a batch](https://cdn-uploads.huggingface.co/production/uploads/64efbd469e7770db74cb72f5/uFX-sTWI6knecIf56uahk.png)](https://cdn-uploads.huggingface.co/production/uploads/64efbd469e7770db74cb72f5/uFX-sTWI6knecIf56uahk.png) | [![val score accuracy mean@30 for aime_2025](https://cdn-uploads.huggingface.co/production/uploads/64efbd469e7770db74cb72f5/5x2mdkpHcdvctZ96yfv58.png)](https://cdn-uploads.huggingface.co/production/uploads/64efbd469e7770db74cb72f5/5x2mdkpHcdvctZ96yfv58.png) |\n\n**Figure 7**. On the Retool task, the run with fix showed steady reward improvement and no gradient exploding (fa2 is the flash attention 2 without the fix while fa3 is the flash attention 3 with the fix). After the fix, the validation accuracy score goes up now.\n\n## Memory-Efficient Training\n\n### Mitigating FSDP Memory Blow-Ups Caused by Repeated MoE Expert Materialization\n\nOne issue we consistently encountered was excessive memory allocation during the FSDP forward pass, which led to repeated out-of-memory (OOM) failures when training GPT-OSS-20B bf16 models on 16 H200 nodes (max response length: 16k, prompt length: 8k). This behavior is highly unexpected for a 20B-parameter MoE model.\n\n```text\n2025-11-27T11:15:27.927Z [36m(TaskRunner pid=32081)[0m File \"/home/jobuser/.local/lib/python3.10/site-packages/transformers/models/gpt_oss/modeling_gpt_oss.py\", line 123, in forward\n2025-11-27T11:15:27.927Z [36m(TaskRunner pid=32081)[0m hidden_states = hidden_states.repeat(num_experts, 1)\n2025-11-27T11:15:27.927Z [36m(TaskRunner pid=32081)[0m torch.OutOfMemoryError: CUDA out of memory. Tried to allocate 180.00 GiB. GPU 0 has a total capacity of 139.72 GiB of which 110.94 GiB is free. Process 685851 has 24.88 GiB memory in use. Process 692458 has 3.87 GiB memory in use. Of the allocated memory 23.28 GiB is allocated by PyTorch, and 84.43 MiB is reserved by PyTorch but unallocated.\n```\n\nWe identified the issue as originating from two different implementations of the MoE forward path in Hugging Face Transformers. This issue has also been reported by other users: [https://github.com/huggingface/transformers/issues/40073](https://github.com/huggingface/transformers/issues/40073); When verl computes log-probabilities under FSDP, the inference forward path is triggered. In the current Hugging Face implementation, this path duplicates hidden states for all experts and performs batched matrix multiplication, materializing extremely large tensors in GPU memory. By contrast, the training forward path uses a for-loop to process each expert sequentially and then combines the results. While slower, this approach is significantly more memory efficient.\n\n```python\n    @GPUMemoryLogger(role=\"dp actor\", logger=logger)\n    def compute_log_prob(self, data: DataProto, calculate_entropy=False) -> torch.Tensor:\n        \"\"\"\n        ....\n        \"\"\"\n        # set to eval, this essentially prioritizes parallelism at the cost of memory efficiency\n        self.actor_module.eval()\n        ...\n```\n\nWe patched the Hugging Face implementation to use a more memory-efficient execution path, avoiding repeated materialization of experts.\n\n### Sequence Parallel with Flash Attention V3\n\nAgentic RL requires the agent to interact with the environment over multiple steps while maintaining an ever-expanding context. Observations and environment feedback from each step are appended to the context and used as input for subsequent decision-making, which introduces significant challenges for memory efficiency and scalability during training.\n\nUnder fully sharded data parallelism (FSDP), model parameters, optimizer states, and gradients are sharded across the entire world size (i.e., all GPUs in the training cluster). Each GPU stores and updates only its assigned parameter shards, while rollout data are replicated across all GPUs—meaning every GPU processes the full agent interaction history for each rollout.\n\nDuring the forward pass, when computation reaches a layer whose parameters are not locally available, an `all_gather` operation is triggered to materialize the full parameters across GPUs. During the backward pass, a corresponding `reduce_scatter` operation aggregates gradients and ensures that each GPU retains only its local shard. This provides a degree of scaling: as the number of GPUs increases, the per-GPU memory footprint decreases.\n\nFSDP provides model-level scaling by sharding model parameters, gradients, and optimizer states across GPUs. Sequence parallelism (or context parallelism) further reduces per-GPU memory consumption by partitioning the input sequence across devices, thereby lowering the peak activation memory on each GPU.\n\nAs the number of sequence-parallel dimensions increases, the maximum activation memory per GPU correspondingly decreases. We have implemented sequence parallelism to be attention-sink-aware and compatible with FlashAttention v3 (Figure 8, right).\n\n[![SP  (2)](https://cdn-uploads.huggingface.co/production/uploads/64efbd469e7770db74cb72f5/ryT_y9BpbFSdMDxNYlVlK.png)](https://cdn-uploads.huggingface.co/production/uploads/64efbd469e7770db74cb72f5/ryT_y9BpbFSdMDxNYlVlK.png)\n\n**Figure 8**. Left: Inference without sequence parallelism. Right: Inference with sequence parallelism, where additional all-to-all communication is performed before and after the attention layer. This partitions the sequence across parallel workers and reduces the peak memory footprint of attention computation by a factor proportional to the sequence-parallelism degree.\n\nSequence parallelism scales along the sequence dimension to reduce the per-GPU activation footprint. Input tokens from all sequences are packed into a single contiguous list by removing padding tokens, while position IDs are used to distinguish tokens belonging to different sequences. This design naturally benefits from FlashAttention’s variable-length support. For sequence parallelism, layers other than the attention layer do not have inter-position dependencies; therefore, they do not require each GPU to hold a complete sequence shard, and no additional communication is needed for these layers.\n\nThe attention layer, however, requires all tokens belonging to the same sequence to be present on the same GPU in order to compute attention weights correctly. To satisfy this constraint, an all-to-all communication is performed to gather sequence elements, with the split performed at the attention-head level. This design avoids communication within the attention computation itself, which would otherwise be prohibitively expensive. After the attention layer, a single all-to-all communication redistributes the outputs back to their original sequence-parallel layout, after which the remaining non-attention layers can proceed without further synchronization.\n\n## Conclusion\n\nOur journey to enable agentic RL training for the GPT-OSS backbone model was a practical retrospective, highlighting that unlocking advanced capabilities in open-source LLMs requires meticulous, deep-dive engineering.\n\nWe made contributions that transformed the viability of GPT-OSS for agentic applications, specifically by:\n\n- **Stabilizing PPO:** We contributed a fix to restore on-policy integrity, overriding the log-probability mismatch caused by the MoE architecture’s non-determinism (Figure 2).\n\n- **Enabling Attention Sink Support:** We successfully implemented and integrated the attention sink backward pass into FlashAttention v3, correcting the catastrophic training–inference mismatch that had previously caused instability and slow convergence (Figures 5, 6, and 7).\n\n- **Scaling Memory Efficiency:** We introduced crucial memory optimizations, including patching the MoE materialization process and integrating sequence parallelism with the new attention sink support, enabling training with the long context windows essential for multi-step agents (Figure 8).\n\nThese engineering efforts validate GPT-OSS as a scalable and high-performance backbone for building the next generation of intelligent, multi-step decision-making agents.\n\n## Acknowledgments\n\nThanks to Deepak Agarwal, Bee-Chung Chen, Animesh Singh, Gungor Polatkan, Balaji Krishnapuram, and Jitendra Agarwal for their leadership support.\n\n## References\n\n1. Feng, Jiazhan, et al. *Retool: Reinforcement Learning for Strategic Tool Use in LLMs.* arXiv preprint arXiv:2504.11536 (2025).\n\n2. Xiao, Guangxuan, et al. *Efficient Streaming Language Models with Attention Sinks.* arXiv preprint arXiv:2309.17453 (2023).\n\n3. When Speed Kills Stability: Demystifying RL Collapse from the Training–Inference Mismatch.\\\n   [https://yingru.notion.site/When-Speed-Kills-Stability-Demystifying-RL-Collapse-from-the-Training-Inference-Mismatch-271211a558b7808d8b12d403fd15edda](https://yingru.notion.site/When-Speed-Kills-Stability-Demystifying-RL-Collapse-from-the-Training-Inference-Mismatch-271211a558b7808d8b12d403fd15edda)\n",
  "authors": [],
  "captured_at": "2026-10-04T12:50:30.804479567+00:00",
  "category": "ai",
  "content_html": "<p>Agentic reinforcement learning (RL) extends traditional LLM training by optimizing not just a single-turn response, but an entire decision-making process learned through direct interaction with an environment during training. Unlike traditional single-turn reinforcement learning or offline preference-based methods that rely on static datasets, agentic RL trains policies by actively collecting on-policy data as the agent plans actions, invokes tools, observes outcomes, and adapts its behavior over multi-step trajectories in either simulated or real environments. This interaction-driven optimization assigns credit across long-horizon decisions, where intermediate choices such as query reformulation, tool selection, and execution order directly influence downstream success. Training follows an iterative closed loop in which the agent interacts with the environment to collect rollout trajectories, computes rewards over these trajectories, updates the policy based on observed outcomes, and then uses the updated policy to drive the next round of interaction and data collection such as GRPO or PPO algorithms..</p>\n<p>LinkedIn is an AI-first company that's built agents to help professionals be more successful. In this setting, models must reason over incomplete information, interact with structured services, and adapt to evolving user intent across multiple steps rather than produce a single static response. These capabilities are especially critical for agents that support the goals of recruiters, job and knowledge seekers, and learners end users, such as retrieving information, refining queries, coordinating tools, and executing multi-step workflows. By learning robust decision policies through interaction, agentic RL provides a principled foundation for building scalable, reliable, and adaptable AI systems through end-to-end optimization.</p>\n<p>The GPT-OSS model has shown comparable performance to OpenAI o3-mini and o4-mini [<a target=\"_blank\" rel=\"noopener noreferrer\" href=\"https://openai.com/index/introducing-gpt-oss/\">ref</a>], but its suitability for agentic reinforcement learning training has not yet been validated. Most recent work focuses on fine-tuning without tool calling, such as: <a target=\"_blank\" rel=\"noopener noreferrer\" href=\"https://cookbook.openai.com/articles/gpt-oss/fine-tune-transfomers\">Fine-tuning with gpt-oss and Hugging Face Transformers</a> and <a target=\"_blank\" rel=\"noopener noreferrer\" href=\"https://unsloth.ai/docs/models/gpt-oss-how-to-run-and-fine-tune/tutorial-how-to-fine-tune-gpt-oss\">unsloth tutorial: how to fine-tune gpt-oss</a>. This blog explores the journey to unlock <a target=\"_blank\" rel=\"noopener noreferrer\" href=\"https://github.com/volcengine/verl/issues/3794\">agentic RL</a> training for GPT-OSS as a potential backbone model for agentic applications.</p>\n<p>In our experiments, we use <a target=\"_blank\" rel=\"noopener noreferrer\" href=\"https://github.com/volcengine/verl\">verl</a> as our training framework since it is one of the most popular adopted frameworks in the open source community. We use gsm8k, <a target=\"_blank\" rel=\"noopener noreferrer\" href=\"https://github.com/verl-project/verl-recipe/tree/21892b9276936efab5375c3f6b8415e472ef7118/retool\">Retool</a> task, <a target=\"_blank\" rel=\"noopener noreferrer\" href=\"https://arxiv.org/abs/2507.02833\">verifiable instruction following task</a>, which are commonly used in RL training. We focus on presenting experimental results for the GPT-OSS-20B model, and our attention-sink fix also works for GPT-OSS-120B. The Qwen-2.5-32B model is additionally used to benchmark standard metric trends during RL training.</p>\n<h2>Challenges of GPT-OSS RL Training</h2>\n<p>verl has been an OSS framework used by the team, and the team has previously collaborated and contributed to it to help democratize agentic reinforcement learning training. With the introduction of the new Harmony chat template in GPT-OSS, the first step is to ensure that the training framework fully supports the updated message format and conversation semantics required by Harmony. This step helps rollout generation, trajectory construction, and tool parsing remain consistent and correct under the new template.</p>\n<p>The team uses ReTool as a representative example to verify code correctness. ReTool is an agentic coding task in which the model is asked to solve a math problem with the assistance of a code compiler tool. This setup allows the model to focus on core reasoning and algorithmic logic, while delegating the actual arithmetic and execution to the tool. During an episode, the model interacts with the code tool multiple times, using execution results as feedback to refine its solution. At the end of the trajectory, the model produces a final answer, on which the reward is computed.</p>\n<p>During the initial training runs, we observed exploding KL divergence and entropy, along with non-increasing rewards, indicating underlying issues in the GPT-OSS training setup, as shown in Figure 1.</p>\n<table>\n<thead>\n<tr>\n<th>Average Gradient Norm</th>\n<th>Average Reward</th>\n</tr>\n</thead>\n<tbody>\n<tr>\n<td><a target=\"_blank\" rel=\"noopener noreferrer\" href=\"https://cdn-uploads.huggingface.co/production/uploads/64efbd469e7770db74cb72f5/5eWLHQ-EAWKjPc7T0c6CK.png\"><picture class=\"article-picture article-picture-fallback\" style=\"--image-ratio:16 / 9\"><img src=\"https://cdn-uploads.huggingface.co/production/uploads/64efbd469e7770db74cb72f5/5eWLHQ-EAWKjPc7T0c6CK.png\" alt=\"Average gradient norm in a batch\" loading=\"eager\" decoding=\"async\" fetchpriority=\"high\" referrerpolicy=\"no-referrer\" class=\"progressive-image\"></picture></a></td>\n<td><a target=\"_blank\" rel=\"noopener noreferrer\" href=\"https://cdn-uploads.huggingface.co/production/uploads/64efbd469e7770db74cb72f5/ZpBvRXLKyAT-PIdsuRi11.png\"><picture class=\"article-picture article-picture-fallback\" style=\"--image-ratio:16 / 9\"><img src=\"https://cdn-uploads.huggingface.co/production/uploads/64efbd469e7770db74cb72f5/ZpBvRXLKyAT-PIdsuRi11.png\" alt=\"Average reward in a batch\" loading=\"lazy\" decoding=\"async\" fetchpriority=\"low\" referrerpolicy=\"no-referrer\" class=\"progressive-image\"></picture></a></td>\n</tr>\n</tbody>\n</table>\n<p><strong>Figure 1.</strong> Left: Qwen32b has significantly higher rewards compared to GPT-OSS 20B; Right: The gradient norm exploded as training progressed.</p>\n<h2>A Practical Debugging Journey in verl: Restoring PPO On-Policy Integrity</h2>\n<h3>Restoring PPO On-Policy Integrity: A Fix for MoE Log-Probability Mismatch</h3>\n<p><picture class=\"article-picture article-picture-fallback\" style=\"--image-ratio:16 / 9\"><img src=\"https://cdn-uploads.huggingface.co/production/uploads/64efbd469e7770db74cb72f5/21xIlXKeAlTK5qKSp-TkX.png\" alt=\"\" loading=\"lazy\" decoding=\"async\" fetchpriority=\"low\" referrerpolicy=\"no-referrer\" class=\"progressive-image\"></picture></p>\n<p><strong>Figure 2.</strong> Non-zero importance sampling clip value even for on-policy training.</p>\n<p>We focus on on-policy methods because they provide greater stability and more reliable convergence. The foundation of pure on-policy Proximal Policy Optimization (PPO) mandates that the importance sampling ratio must be exactly 1. The mathematical definition of the importance ratio is:</p>\n<p>ratio=π(a∣s)πold(a∣s) \\text{ratio} = \\frac{\\pi(a \\mid s)}{\\pi_{\\text{old}}(a \\mid s)}</p>\n<p>This requirement ensures that the policy update is executed only on the data generated by the current policy π(a | s) = πold(a | s), preventing unintended clipping.</p>\n<p>We have observed the non-zero clipping value in our ReTool training, as shown in Figure 2, stemming from a mismatch between the two log-probabilities:</p>\n<ul>\n<li>Current log-probability <code>log_prob</code>: log(π(a | s))</li>\n<li>Old log-probability <code>old_log_prob</code>: log(πold(a | s))</li>\n</ul>\n<p><strong>Root Cause: The Dual Forward Pass and MoE Architecture</strong></p>\n<p>Prior to verl 0.3.0, the implementation relied on two separate forward passes (one to compute the current <code>log_prob</code> and one to retrieve the stored <code>old_log_prob</code>) for the same state-action pair.</p>\n<p>In a Mixture of Experts (MoE) architecture like GPT-OSS, the gating network routes the input to different experts. Due to implementation factors (e.g., subtle floating-point differences or explicit stochasticity), the expert routing can differ slightly between the two passes. Readers who are interested can further read <em><a target=\"_blank\" rel=\"noopener noreferrer\" href=\"https://arxiv.org/pdf/2510.11370v1\">Stabilizing MoE Reinforcement Learning by Aligning Training and Inference Routers</a></em>.</p>\n<p>This difference in routing leads to:</p>\n<p>log⁡(π(a∣s))≠log⁡(πold(a∣s)) \\log(\\pi(a \\mid s)) \\neq \\log(\\pi_{\\text{old}}(a \\mid s))</p>\n<p>The resulting ratio deviates from 1, falsely triggering the PPO clip and violating the core on-policy assumption.</p>\n<p><strong>Solution: Enforcing Ratio = 1 via Log-Probability Substitution</strong></p>\n<p>The fix resolves the issue by logically overriding the flawed computation when the environment is known to be on-policy (i.e., when the minibatch size equals the global batch size):</p>\n<pre><code class=\"language-python\"><span class=\"code-snippet\" data-language=\"Python\"><span class=\"syntax-source syntax-python\"><span class=\"syntax-meta syntax-statement syntax-if syntax-python\"><span class=\"syntax-keyword syntax-control syntax-flow syntax-conditional syntax-python\">if</span> <span class=\"syntax-meta syntax-qualified-name syntax-python\"><span class=\"syntax-meta syntax-generic-name syntax-python\">on_policy</span></span><span class=\"syntax-punctuation syntax-section syntax-block syntax-conditional syntax-python\">:</span></span>\n    <span class=\"syntax-meta syntax-qualified-name syntax-python\"><span class=\"syntax-meta syntax-generic-name syntax-python\">old_log_prob</span></span> <span class=\"syntax-keyword syntax-operator syntax-assignment syntax-python\">=</span> <span class=\"syntax-meta syntax-function-call syntax-python\"><span class=\"syntax-meta syntax-qualified-name syntax-python\"><span class=\"syntax-meta syntax-generic-name syntax-python\">log_prob</span><span class=\"syntax-punctuation syntax-accessor syntax-dot syntax-python\">.</span></span><span class=\"syntax-meta syntax-qualified-name syntax-python\"><span class=\"syntax-variable syntax-function syntax-python\">detach</span></span><span class=\"syntax-punctuation syntax-section syntax-arguments syntax-begin syntax-python\">(</span><span class=\"syntax-punctuation syntax-section syntax-arguments syntax-end syntax-python\">)</span></span>\n<span class=\"syntax-meta syntax-statement syntax-conditional syntax-python\"><span class=\"syntax-keyword syntax-control syntax-flow syntax-conditional syntax-python\">else</span><span class=\"syntax-punctuation syntax-section syntax-block syntax-python\">:</span></span>\n    <span class=\"syntax-meta syntax-qualified-name syntax-python\"><span class=\"syntax-meta syntax-generic-name syntax-python\">old_log_prob</span></span> <span class=\"syntax-keyword syntax-operator syntax-assignment syntax-python\">=</span> <span class=\"syntax-meta syntax-item-access syntax-python\"><span class=\"syntax-meta syntax-qualified-name syntax-python\"><span class=\"syntax-meta syntax-generic-name syntax-python\">model_inputs</span></span></span><span class=\"syntax-meta syntax-item-access syntax-python\"><span class=\"syntax-punctuation syntax-section syntax-brackets syntax-begin syntax-python\">[</span></span><span class=\"syntax-meta syntax-item-access syntax-arguments syntax-python\"><span class=\"syntax-meta syntax-string syntax-python\"><span class=\"syntax-string syntax-quoted syntax-double syntax-python\"><span class=\"syntax-punctuation syntax-definition syntax-string syntax-begin syntax-python\">&quot;</span></span></span><span class=\"syntax-meta syntax-string syntax-python\"><span class=\"syntax-string syntax-quoted syntax-double syntax-python\">old_log_probs<span class=\"syntax-punctuation syntax-definition syntax-string syntax-end syntax-python\">&quot;</span></span></span></span><span class=\"syntax-meta syntax-item-access syntax-python\"><span class=\"syntax-punctuation syntax-section syntax-brackets syntax-end syntax-python\">]</span></span>\n</span></span></code></pre>\n<p>By setting <code>old_log_prob</code> equal to the newly computed <code>log_prob</code> (detached to prevent gradient flow through the reference value), the importance ratio is mathematically forced back to 1. This strategy bypasses the instability caused by MoE's non-deterministic routing and guarantees strict on-policy behavior during PPO training.</p>\n<hr />\n<h3>Correcting Training–Inference Mismatch</h3>\n<p>Although fixing the log-probability mismatch reduced the importance-sampling clip ratio to zero, gradient norms continued to explode and rewards failed to improve. To isolate the issue, we simplified training to GSM8K, a single-step task without agentic tool use. The same instability persisted, as shown in the green curves in Figure 3, indicating a <strong>fundamental issue in basic RL training with GPT-OSS under verl.</strong></p>\n<p>We hypothesize that <strong>training–inference mismatch</strong> could be a potential cause: discrepancies between inference-time execution—where engines such as vLLM and SGLang aggressively optimize for throughput—and training-time execution under FSDP, which prioritizes numerical precision and stability, can effectively turn otherwise <strong>on-policy RL into off-policy optimization.</strong></p>\n<p>This <a target=\"_blank\" rel=\"noopener noreferrer\" href=\"https://yingru.notion.site/When-Speed-Kills-Stability-Demystifying-RL-Collapse-from-the-Training-Inference-Mismatch-271211a558b7808d8b12d403fd15edda\">blog</a> details why such mismatches lead to unstable gradients and non-improving rewards. Figure 3 compares training runs with and without rollout correction (see this <a target=\"_blank\" rel=\"noopener noreferrer\" href=\"https://verl.readthedocs.io/en/latest/algo/rollout_corr.html\">verl blog</a> for details). After applying rollout correction, training dynamics improve significantly, with gradient norms remaining stable rather than exploding.</p>\n<p>However, as shown in the left plot of Figure 4, the reward increases only modestly, and convergence on the simple GSM8K task remains substantially slower compared to smaller dense model variants.</p>\n<table>\n<thead>\n<tr>\n<th>Average Entropy</th>\n<th>Average Gradient Norm</th>\n<th>Average KL Loss</th>\n</tr>\n</thead>\n<tbody>\n<tr>\n<td><picture class=\"article-picture article-picture-fallback\" style=\"--image-ratio:16 / 9\"><img src=\"https://cdn-uploads.huggingface.co/production/uploads/64efbd469e7770db74cb72f5/MoBnatEipCi_OScgm5fAJ.png\" alt=\"\" loading=\"lazy\" decoding=\"async\" fetchpriority=\"low\" referrerpolicy=\"no-referrer\" class=\"progressive-image\"></picture></td>\n<td><picture class=\"article-picture article-picture-fallback\" style=\"--image-ratio:16 / 9\"><img src=\"https://cdn-uploads.huggingface.co/production/uploads/64efbd469e7770db74cb72f5/k4Su03zqq6Cg-7jM6Il6J.png\" alt=\"\" loading=\"lazy\" decoding=\"async\" fetchpriority=\"low\" referrerpolicy=\"no-referrer\" class=\"progressive-image\"></picture></td>\n<td><picture class=\"article-picture article-picture-fallback\" style=\"--image-ratio:16 / 9\"><img src=\"https://cdn-uploads.huggingface.co/production/uploads/64efbd469e7770db74cb72f5/g77U9c8a-FY1rvwru96Ez.png\" alt=\"\" loading=\"lazy\" decoding=\"async\" fetchpriority=\"low\" referrerpolicy=\"no-referrer\" class=\"progressive-image\"></picture></td>\n</tr>\n</tbody>\n</table>\n<p><strong>Figure 3.</strong> Gradient norm behavior under different training configurations. Green: Training without rollout correction, exhibiting unstable gradients. Red: Training with the attention layer frozen to isolate the issue to the attention mechanism, resulting in partial stabilization. Blue: Training with rollout correction enabled (sequence-level importance sampling), yielding stable gradient norms.</p>\n<table>\n<thead>\n<tr>\n<th>Average Reward</th>\n<th>Max Log-Perplexity Difference</th>\n</tr>\n</thead>\n<tbody>\n<tr>\n<td><a target=\"_blank\" rel=\"noopener noreferrer\" href=\"https://cdn-uploads.huggingface.co/production/uploads/64efbd469e7770db74cb72f5/ihx-XsWH51V0-JM46jODE.png\"><picture class=\"article-picture article-picture-fallback\" style=\"--image-ratio:16 / 9\"><img src=\"https://cdn-uploads.huggingface.co/production/uploads/64efbd469e7770db74cb72f5/ihx-XsWH51V0-JM46jODE.png\" alt=\"Average reward in a batch\" loading=\"lazy\" decoding=\"async\" fetchpriority=\"low\" referrerpolicy=\"no-referrer\" class=\"progressive-image\"></picture></a></td>\n<td><a target=\"_blank\" rel=\"noopener noreferrer\" href=\"https://cdn-uploads.huggingface.co/production/uploads/64efbd469e7770db74cb72f5/fP5KNR2XYY7EH-muYBk_X.png\"><picture class=\"article-picture article-picture-fallback\" style=\"--image-ratio:16 / 9\"><img src=\"https://cdn-uploads.huggingface.co/production/uploads/64efbd469e7770db74cb72f5/fP5KNR2XYY7EH-muYBk_X.png\" alt=\"Maximum absolute log-perplexity difference in a batch between rollout policy and training policy\" loading=\"lazy\" decoding=\"async\" fetchpriority=\"low\" referrerpolicy=\"no-referrer\" class=\"progressive-image\"></picture></a></td>\n</tr>\n</tbody>\n</table>\n<p><strong>Figure 4.</strong> Left: Reward improvement on GSM8K remains slow even after applying rollout correction, with performance comparable to runs where the attention layer is frozen during training. Right: A substantial log-ppl mismatch is observed between the inference engine (SGLang with Triton kernels supporting attention-sink forward passes) and the training stack (FSDP with FlashAttention-v2), indicating a large training–inference inconsistency.</p>\n<p>To further isolate the root cause, we freeze the attention layers during training and observe reward dynamics similar to those of runs without freezing (blue curve vs yellow curve in Figure 4). This indicates that learning is primarily driven by the MoE layers, while the attention mechanism contributes less effectively than expected. In addition, we observe a substantial token-level probability mismatch between the inference engine and the distributed training stack which are using different attention kernels. Together, these observations motivate a deeper investigation into the attention mechanism.</p>\n<hr />\n<h2>Attention Sink Support in FlashAttentionV3</h2>\n<p>Attention sinks used in GPT-OSS are learnable scalar parameters (one per attention head) that act as &quot;virtual tokens&quot; in the softmax computation. They allow the model to allocate attention mass to a learned sink rather than forcing all attention to content tokens, which has been shown to improve attention stability in streaming inference and training with sliding-window attention.</p>\n<p>After a deeper investigation, we identified several major issues:</p>\n<ul>\n<li>verl hard-codes FlashAttention v2 in <code>fsdp_worker</code>, which does not support attention sinks.</li>\n<li>The attention sink backward pass is not supported in FlashAttention v2 and v3, so it does not work as expected even when FlashAttention v3 is enabled.</li>\n<li>Since the forward pass has not yet been merged into the original FlashAttention v3 repository, we leveraged the forward pass from the vLLM FlashAttention fork (<a target=\"_blank\" rel=\"noopener noreferrer\" href=\"https://github.com/vllm-project/flash-attention/pull/75\">PR #75</a>) and implemented the backward pass to compute the sink gradient.</li>\n</ul>\n<h3>Standard Attention</h3>\n<pre><code class=\"language-python\"><span class=\"code-snippet\" data-language=\"Python\"><span class=\"syntax-source syntax-python\"><span class=\"syntax-meta syntax-qualified-name syntax-python\"><span class=\"syntax-meta syntax-generic-name syntax-python\">scores</span></span> <span class=\"syntax-keyword syntax-operator syntax-assignment syntax-python\">=</span> <span class=\"syntax-meta syntax-qualified-name syntax-python\"><span class=\"syntax-meta syntax-generic-name syntax-python\">QK</span></span><span class=\"syntax-keyword syntax-operator syntax-arithmetic syntax-python\">^</span><span class=\"syntax-meta syntax-qualified-name syntax-python\"><span class=\"syntax-meta syntax-generic-name syntax-python\">T</span></span> <span class=\"syntax-keyword syntax-operator syntax-arithmetic syntax-python\">/</span> <span class=\"syntax-meta syntax-function-call syntax-python\"><span class=\"syntax-meta syntax-qualified-name syntax-python\"><span class=\"syntax-variable syntax-function syntax-python\">sqrt</span></span><span class=\"syntax-punctuation syntax-section syntax-arguments syntax-begin syntax-python\">(</span><span class=\"syntax-meta syntax-function-call syntax-arguments syntax-python\"><span class=\"syntax-meta syntax-qualified-name syntax-python\"><span class=\"syntax-meta syntax-generic-name syntax-python\">d</span></span></span><span class=\"syntax-punctuation syntax-section syntax-arguments syntax-end syntax-python\">)</span></span>               <span class=\"syntax-comment syntax-line syntax-number-sign syntax-python\"><span class=\"syntax-punctuation syntax-definition syntax-comment syntax-python\">#</span> [B, H, N_q, N_k]\n</span><span class=\"syntax-meta syntax-qualified-name syntax-python\"><span class=\"syntax-meta syntax-generic-name syntax-python\">probs</span></span> <span class=\"syntax-keyword syntax-operator syntax-assignment syntax-python\">=</span> <span class=\"syntax-meta syntax-function-call syntax-python\"><span class=\"syntax-meta syntax-qualified-name syntax-python\"><span class=\"syntax-variable syntax-function syntax-python\">softmax</span></span><span class=\"syntax-punctuation syntax-section syntax-arguments syntax-begin syntax-python\">(</span><span class=\"syntax-meta syntax-function-call syntax-arguments syntax-python\"><span class=\"syntax-meta syntax-qualified-name syntax-python\"><span class=\"syntax-meta syntax-generic-name syntax-python\">scores</span></span><span class=\"syntax-punctuation syntax-separator syntax-arguments syntax-python\">,</span> <span class=\"syntax-variable syntax-parameter syntax-python\">dim</span><span class=\"syntax-keyword syntax-operator syntax-assignment syntax-python\">=</span><span class=\"syntax-keyword syntax-operator syntax-arithmetic syntax-python\">-</span><span class=\"syntax-constant syntax-numeric syntax-integer syntax-decimal syntax-python\">1</span></span><span class=\"syntax-punctuation syntax-section syntax-arguments syntax-end syntax-python\">)</span></span>       <span class=\"syntax-comment syntax-line syntax-number-sign syntax-python\"><span class=\"syntax-punctuation syntax-definition syntax-comment syntax-python\">#</span> Σ_j P_ij = 1\n</span><span class=\"syntax-meta syntax-qualified-name syntax-python\"><span class=\"syntax-meta syntax-generic-name syntax-python\">output</span></span> <span class=\"syntax-keyword syntax-operator syntax-assignment syntax-python\">=</span> <span class=\"syntax-meta syntax-qualified-name syntax-python\"><span class=\"syntax-meta syntax-generic-name syntax-python\">probs</span></span> <span class=\"syntax-keyword syntax-operator syntax-matrix syntax-python\">@</span> <span class=\"syntax-meta syntax-qualified-name syntax-python\"><span class=\"syntax-meta syntax-generic-name syntax-python\">V</span></span>                   <span class=\"syntax-comment syntax-line syntax-number-sign syntax-python\"><span class=\"syntax-punctuation syntax-definition syntax-comment syntax-python\">#</span> [B, H, N_q, d_v]\n</span></span></span></code></pre>\n<h3>Attention with Sinks (GPT-OSS)</h3>\n<pre><code class=\"language-python\"><span class=\"code-snippet\" data-language=\"Python\"><span class=\"syntax-source syntax-python\"><span class=\"syntax-meta syntax-qualified-name syntax-python\"><span class=\"syntax-meta syntax-generic-name syntax-python\">scores</span></span> <span class=\"syntax-keyword syntax-operator syntax-assignment syntax-python\">=</span> <span class=\"syntax-meta syntax-qualified-name syntax-python\"><span class=\"syntax-meta syntax-generic-name syntax-python\">QK</span></span><span class=\"syntax-keyword syntax-operator syntax-arithmetic syntax-python\">^</span><span class=\"syntax-meta syntax-qualified-name syntax-python\"><span class=\"syntax-meta syntax-generic-name syntax-python\">T</span></span> <span class=\"syntax-keyword syntax-operator syntax-arithmetic syntax-python\">/</span> <span class=\"syntax-meta syntax-function-call syntax-python\"><span class=\"syntax-meta syntax-qualified-name syntax-python\"><span class=\"syntax-variable syntax-function syntax-python\">sqrt</span></span><span class=\"syntax-punctuation syntax-section syntax-arguments syntax-begin syntax-python\">(</span><span class=\"syntax-meta syntax-function-call syntax-arguments syntax-python\"><span class=\"syntax-meta syntax-qualified-name syntax-python\"><span class=\"syntax-meta syntax-generic-name syntax-python\">d</span></span></span><span class=\"syntax-punctuation syntax-section syntax-arguments syntax-end syntax-python\">)</span></span>                               <span class=\"syntax-comment syntax-line syntax-number-sign syntax-python\"><span class=\"syntax-punctuation syntax-definition syntax-comment syntax-python\">#</span> [B, H, N_q, N_k]\n</span><span class=\"syntax-meta syntax-qualified-name syntax-python\"><span class=\"syntax-meta syntax-generic-name syntax-python\">combined</span></span> <span class=\"syntax-keyword syntax-operator syntax-assignment syntax-python\">=</span> <span class=\"syntax-meta syntax-function-call syntax-python\"><span class=\"syntax-meta syntax-qualified-name syntax-python\"><span class=\"syntax-variable syntax-function syntax-python\">concat</span></span><span class=\"syntax-punctuation syntax-section syntax-arguments syntax-begin syntax-python\">(</span><span class=\"syntax-meta syntax-function-call syntax-arguments syntax-python\"><span class=\"syntax-meta syntax-structure syntax-list syntax-python\"><span class=\"syntax-punctuation syntax-section syntax-list syntax-begin syntax-python\">[</span><span class=\"syntax-meta syntax-qualified-name syntax-python\"><span class=\"syntax-meta syntax-generic-name syntax-python\">scores</span></span><span class=\"syntax-punctuation syntax-separator syntax-list syntax-python\">,</span> <span class=\"syntax-meta syntax-qualified-name syntax-python\"><span class=\"syntax-meta syntax-generic-name syntax-python\">sink_param</span></span><span class=\"syntax-punctuation syntax-section syntax-list syntax-end syntax-python\">]</span></span><span class=\"syntax-punctuation syntax-separator syntax-arguments syntax-python\">,</span> <span class=\"syntax-variable syntax-parameter syntax-python\">dim</span><span class=\"syntax-keyword syntax-operator syntax-assignment syntax-python\">=</span><span class=\"syntax-keyword syntax-operator syntax-arithmetic syntax-python\">-</span><span class=\"syntax-constant syntax-numeric syntax-integer syntax-decimal syntax-python\">1</span></span><span class=\"syntax-punctuation syntax-section syntax-arguments syntax-end syntax-python\">)</span></span>       <span class=\"syntax-comment syntax-line syntax-number-sign syntax-python\"><span class=\"syntax-punctuation syntax-definition syntax-comment syntax-python\">#</span> [B, H, N_q, N_k+1]\n</span><span class=\"syntax-meta syntax-qualified-name syntax-python\"><span class=\"syntax-meta syntax-generic-name syntax-python\">probs</span></span> <span class=\"syntax-keyword syntax-operator syntax-assignment syntax-python\">=</span> <span class=\"syntax-meta syntax-function-call syntax-python\"><span class=\"syntax-meta syntax-qualified-name syntax-python\"><span class=\"syntax-variable syntax-function syntax-python\">softmax</span></span><span class=\"syntax-punctuation syntax-section syntax-arguments syntax-begin syntax-python\">(</span><span class=\"syntax-meta syntax-function-call syntax-arguments syntax-python\"><span class=\"syntax-meta syntax-qualified-name syntax-python\"><span class=\"syntax-meta syntax-generic-name syntax-python\">combined</span></span><span class=\"syntax-punctuation syntax-separator syntax-arguments syntax-python\">,</span> <span class=\"syntax-variable syntax-parameter syntax-python\">dim</span><span class=\"syntax-keyword syntax-operator syntax-assignment syntax-python\">=</span><span class=\"syntax-keyword syntax-operator syntax-arithmetic syntax-python\">-</span><span class=\"syntax-constant syntax-numeric syntax-integer syntax-decimal syntax-python\">1</span></span><span class=\"syntax-punctuation syntax-section syntax-arguments syntax-end syntax-python\">)</span></span>                     <span class=\"syntax-comment syntax-line syntax-number-sign syntax-python\"><span class=\"syntax-punctuation syntax-definition syntax-comment syntax-python\">#</span> Σ_j P_ij + P_sink = 1\n</span><span class=\"syntax-meta syntax-qualified-name syntax-python\"><span class=\"syntax-meta syntax-generic-name syntax-python\">probs_content</span></span> <span class=\"syntax-keyword syntax-operator syntax-assignment syntax-python\">=</span> <span class=\"syntax-meta syntax-item-access syntax-python\"><span class=\"syntax-meta syntax-qualified-name syntax-python\"><span class=\"syntax-meta syntax-generic-name syntax-python\">probs</span></span></span><span class=\"syntax-meta syntax-item-access syntax-python\"><span class=\"syntax-punctuation syntax-section syntax-brackets syntax-begin syntax-python\">[</span></span><span class=\"syntax-meta syntax-item-access syntax-arguments syntax-python\"><span class=\"syntax-constant syntax-language syntax-python\">...</span>, <span class=\"syntax-punctuation syntax-separator syntax-slice syntax-python\">:</span><span class=\"syntax-keyword syntax-operator syntax-arithmetic syntax-python\">-</span><span class=\"syntax-constant syntax-numeric syntax-integer syntax-decimal syntax-python\">1</span></span><span class=\"syntax-meta syntax-item-access syntax-python\"><span class=\"syntax-punctuation syntax-section syntax-brackets syntax-end syntax-python\">]</span></span>                       <span class=\"syntax-comment syntax-line syntax-number-sign syntax-python\"><span class=\"syntax-punctuation syntax-definition syntax-comment syntax-python\">#</span> Drop sink component\n</span><span class=\"syntax-meta syntax-qualified-name syntax-python\"><span class=\"syntax-meta syntax-generic-name syntax-python\">output</span></span> <span class=\"syntax-keyword syntax-operator syntax-assignment syntax-python\">=</span> <span class=\"syntax-meta syntax-qualified-name syntax-python\"><span class=\"syntax-meta syntax-generic-name syntax-python\">probs_content</span></span> <span class=\"syntax-keyword syntax-operator syntax-matrix syntax-python\">@</span> <span class=\"syntax-meta syntax-qualified-name syntax-python\"><span class=\"syntax-meta syntax-generic-name syntax-python\">V</span></span>                           <span class=\"syntax-comment syntax-line syntax-number-sign syntax-python\"><span class=\"syntax-punctuation syntax-definition syntax-comment syntax-python\">#</span> [B, H, N_q, d_v]\n</span></span></span></code></pre>\n<p><strong>Key difference:</strong> The sink participates in softmax normalization but doesn't contribute to the output.</p>\n<h3>Mathematical Formulation</h3>\n<p>The attention weight for content token j in row i is defined as:</p>\n<p>Pij=exp⁡(Sij)∑j′=1Nkexp⁡(Sij′)+exp⁡(Sh) P_{ij} = \\frac{\\exp(S_{ij})} {\\sum_{j'=1}^{N_k} \\exp(S_{ij'}) + \\exp(S_h)}</p>\n<p><strong>Where:</strong></p>\n<ul>\n<li>Sij = Qi Kj⊤ / √d are the attention scores</li>\n<li>Pij are the attention weights for the content tokens</li>\n<li>Sh is the learnable sink parameter for head h</li>\n</ul>\n<p><strong>Sink Probability:</strong></p>\n<p>The sink probability is computed but not used in the output:</p>\n<p>Pi,h=exp⁡(Sh)∑j′=1Nkexp⁡(Sij′)+exp⁡(Sh) P_{i,h} = \\frac{\\exp(S_h)} {\\sum_{j'=1}^{N_k} \\exp(S_{ij'}) + \\exp(S_h)}</p>\n<h3>Backward Pass</h3>\n<p>The gradient of the loss L with respect to the sink parameter Sh is:</p>\n<p>∂L∂Sh=−∑iPi,h(∂L∂Si,h−∑j∈{1,…,Nk}Pij∂L∂Sij) \\frac{\\partial L}{\\partial S_h} = - \\sum_i P_{i,h} \\left( \\frac{\\partial L}{\\partial S_{i,h}} - \\sum_{j \\in \\{1,\\ldots,N_k\\}} P_{ij} \\frac{\\partial L}{\\partial S_{ij}} \\right)</p>\n<p><strong>Where:</strong></p>\n<ul>\n<li>Pi,h is the sink attention probability for row i</li>\n<li>∂L/∂Sij is the gradient with respect to the attention scores, including the sink</li>\n</ul>\n<p><strong>Simplified Gradient:</strong></p>\n<p>Since the sink is computed but not used in the output, its gradient ∂L/∂Si,h = 0.</p>\n<p>Therefore, the backward equation simplifies to:</p>\n<p>∂L∂Sh=−∑iPi,h(∑j∈{1,…,Nk}Pij∂L∂Sij) \\frac{\\partial L}{\\partial S_h} = - \\sum_i P_{i,h} \\left( \\sum_{j \\in \\{1,\\ldots,N_k\\}} P_{ij} \\frac{\\partial L}{\\partial S_{ij}} \\right)</p>\n<p>The forward pass was adapted from vLLM's FlashAttention fork, and we implemented the backward pass to compute gradients for the sink parameters. The implementation will be released following the internal review process.</p>\n<hr />\n<h3>Results</h3>\n<p>After applying the fix in FlashAttention v3, we observe substantially faster convergence for GPT-OSS-20B across a range of reinforcement learning tasks. These include single-turn RL on math reasoning (GSM8K — red curve in Figure 5), instruction following (VerifyIf, evaluated on an out-of-domain multi-if benchmark — Figure 6), and multi-turn agentic RL with tool use (ReTool — Figure 7).</p>\n<p>Across all settings, training becomes stable and exhibits steady reward improvement.</p>\n<p><picture class=\"article-picture article-picture-fallback\" style=\"--image-ratio:16 / 9\"><img src=\"https://cdn-uploads.huggingface.co/production/uploads/64efbd469e7770db74cb72f5/6TYGevydK99nQ-I1QTouf.png\" alt=\"\" loading=\"lazy\" decoding=\"async\" fetchpriority=\"low\" referrerpolicy=\"no-referrer\" class=\"progressive-image\"></picture></p>\n<p><strong>Figure 5.</strong>. Single Turn GSM8K, the red curve converges much faster than the rest without the fix</p>\n<table>\n<thead>\n<tr>\n<th>Average Entropy</th>\n<th>Average Gradient Norm</th>\n<th>Average Reward</th>\n</tr>\n</thead>\n<tbody>\n<tr>\n<td><a target=\"_blank\" rel=\"noopener noreferrer\" href=\"https://cdn-uploads.huggingface.co/production/uploads/64efbd469e7770db74cb72f5/ydmBLCSGlD9YKWiIocI1S.png\"><picture class=\"article-picture article-picture-fallback\" style=\"--image-ratio:16 / 9\"><img src=\"https://cdn-uploads.huggingface.co/production/uploads/64efbd469e7770db74cb72f5/ydmBLCSGlD9YKWiIocI1S.png\" alt=\"Average entropy in a batch\" loading=\"lazy\" decoding=\"async\" fetchpriority=\"low\" referrerpolicy=\"no-referrer\" class=\"progressive-image\"></picture></a></td>\n<td><a target=\"_blank\" rel=\"noopener noreferrer\" href=\"https://cdn-uploads.huggingface.co/production/uploads/64efbd469e7770db74cb72f5/_kl8mn_CXPsRYJ467IbFs.png\"><picture class=\"article-picture article-picture-fallback\" style=\"--image-ratio:16 / 9\"><img src=\"https://cdn-uploads.huggingface.co/production/uploads/64efbd469e7770db74cb72f5/_kl8mn_CXPsRYJ467IbFs.png\" alt=\"Average gradient norm in a batch\" loading=\"lazy\" decoding=\"async\" fetchpriority=\"low\" referrerpolicy=\"no-referrer\" class=\"progressive-image\"></picture></a></td>\n<td><a target=\"_blank\" rel=\"noopener noreferrer\" href=\"https://cdn-uploads.huggingface.co/production/uploads/64efbd469e7770db74cb72f5/vb2JmmSu-LI5szC_84KsM.png\"><picture class=\"article-picture article-picture-fallback\" style=\"--image-ratio:16 / 9\"><img src=\"https://cdn-uploads.huggingface.co/production/uploads/64efbd469e7770db74cb72f5/vb2JmmSu-LI5szC_84KsM.png\" alt=\"Average reward in a batch\" loading=\"lazy\" decoding=\"async\" fetchpriority=\"low\" referrerpolicy=\"no-referrer\" class=\"progressive-image\"></picture></a></td>\n</tr>\n</tbody>\n</table>\n<p><strong>Figure 6</strong>. On verifiable instruction following the task, the run without the fix collapsed (blue), and the run with fix showed steady reward improvement.</p>\n<table>\n<thead>\n<tr>\n<th>Average Gradient Norm</th>\n<th>Average Reward</th>\n<th>Validation Accuracy</th>\n</tr>\n</thead>\n<tbody>\n<tr>\n<td><a target=\"_blank\" rel=\"noopener noreferrer\" href=\"https://cdn-uploads.huggingface.co/production/uploads/64efbd469e7770db74cb72f5/_Xz7_RLhYuYhzGAeMjDXs.png\"><picture class=\"article-picture article-picture-fallback\" style=\"--image-ratio:16 / 9\"><img src=\"https://cdn-uploads.huggingface.co/production/uploads/64efbd469e7770db74cb72f5/_Xz7_RLhYuYhzGAeMjDXs.png\" alt=\"Average gradient norm in a batch\" loading=\"lazy\" decoding=\"async\" fetchpriority=\"low\" referrerpolicy=\"no-referrer\" class=\"progressive-image\"></picture></a></td>\n<td><a target=\"_blank\" rel=\"noopener noreferrer\" href=\"https://cdn-uploads.huggingface.co/production/uploads/64efbd469e7770db74cb72f5/uFX-sTWI6knecIf56uahk.png\"><picture class=\"article-picture article-picture-fallback\" style=\"--image-ratio:16 / 9\"><img src=\"https://cdn-uploads.huggingface.co/production/uploads/64efbd469e7770db74cb72f5/uFX-sTWI6knecIf56uahk.png\" alt=\"Average reward in a batch\" loading=\"lazy\" decoding=\"async\" fetchpriority=\"low\" referrerpolicy=\"no-referrer\" class=\"progressive-image\"></picture></a></td>\n<td><a target=\"_blank\" rel=\"noopener noreferrer\" href=\"https://cdn-uploads.huggingface.co/production/uploads/64efbd469e7770db74cb72f5/5x2mdkpHcdvctZ96yfv58.png\"><picture class=\"article-picture article-picture-fallback\" style=\"--image-ratio:16 / 9\"><img src=\"https://cdn-uploads.huggingface.co/production/uploads/64efbd469e7770db74cb72f5/5x2mdkpHcdvctZ96yfv58.png\" alt=\"val score accuracy mean@30 for aime_2025\" loading=\"lazy\" decoding=\"async\" fetchpriority=\"low\" referrerpolicy=\"no-referrer\" class=\"progressive-image\"></picture></a></td>\n</tr>\n</tbody>\n</table>\n<p><strong>Figure 7</strong>. On the Retool task, the run with fix showed steady reward improvement and no gradient exploding (fa2 is the flash attention 2 without the fix while fa3 is the flash attention 3 with the fix). After the fix, the validation accuracy score goes up now.</p>\n<h2>Memory-Efficient Training</h2>\n<h3>Mitigating FSDP Memory Blow-Ups Caused by Repeated MoE Expert Materialization</h3>\n<p>One issue we consistently encountered was excessive memory allocation during the FSDP forward pass, which led to repeated out-of-memory (OOM) failures when training GPT-OSS-20B bf16 models on 16 H200 nodes (max response length: 16k, prompt length: 8k). This behavior is highly unexpected for a 20B-parameter MoE model.</p>\n<pre><code class=\"language-text\"><span class=\"code-snippet\">2025-11-27T11:15:27.927Z [36m(TaskRunner pid=32081)[0m File &quot;/home/jobuser/.local/lib/python3.10/site-packages/transformers/models/gpt_oss/modeling_gpt_oss.py&quot;, line 123, in forward\n2025-11-27T11:15:27.927Z [36m(TaskRunner pid=32081)[0m hidden_states = hidden_states.repeat(num_experts, 1)\n2025-11-27T11:15:27.927Z [36m(TaskRunner pid=32081)[0m torch.OutOfMemoryError: CUDA out of memory. Tried to allocate 180.00 GiB. GPU 0 has a total capacity of 139.72 GiB of which 110.94 GiB is free. Process 685851 has 24.88 GiB memory in use. Process 692458 has 3.87 GiB memory in use. Of the allocated memory 23.28 GiB is allocated by PyTorch, and 84.43 MiB is reserved by PyTorch but unallocated.\n</span></code></pre>\n<p>We identified the issue as originating from two different implementations of the MoE forward path in Hugging Face Transformers. This issue has also been reported by other users: <a target=\"_blank\" rel=\"noopener noreferrer\" href=\"https://github.com/huggingface/transformers/issues/40073\">https://github.com/huggingface/transformers/issues/40073</a>; When verl computes log-probabilities under FSDP, the inference forward path is triggered. In the current Hugging Face implementation, this path duplicates hidden states for all experts and performs batched matrix multiplication, materializing extremely large tensors in GPU memory. By contrast, the training forward path uses a for-loop to process each expert sequentially and then combines the results. While slower, this approach is significantly more memory efficient.</p>\n<pre><code class=\"language-python\"><span class=\"code-snippet\" data-language=\"Python\"><span class=\"syntax-source syntax-python\">    <span class=\"syntax-meta syntax-annotation syntax-python\"><span class=\"syntax-punctuation syntax-definition syntax-annotation syntax-python\">@</span><span class=\"syntax-meta syntax-annotation syntax-function syntax-python\"></span></span><span class=\"syntax-meta syntax-annotation syntax-function syntax-python\"><span class=\"syntax-meta syntax-qualified-name syntax-python\"><span class=\"syntax-variable syntax-annotation syntax-function syntax-python\">GPUMemoryLogger</span></span><span class=\"syntax-punctuation syntax-section syntax-arguments syntax-begin syntax-python\">(</span><span class=\"syntax-meta syntax-annotation syntax-arguments syntax-python\"><span class=\"syntax-variable syntax-parameter syntax-python\">role</span><span class=\"syntax-keyword syntax-operator syntax-assignment syntax-python\">=</span><span class=\"syntax-meta syntax-string syntax-python\"><span class=\"syntax-string syntax-quoted syntax-double syntax-python\"><span class=\"syntax-punctuation syntax-definition syntax-string syntax-begin syntax-python\">&quot;</span></span></span><span class=\"syntax-meta syntax-string syntax-python\"><span class=\"syntax-string syntax-quoted syntax-double syntax-python\">dp actor<span class=\"syntax-punctuation syntax-definition syntax-string syntax-end syntax-python\">&quot;</span></span></span><span class=\"syntax-punctuation syntax-separator syntax-arguments syntax-python\">,</span> <span class=\"syntax-variable syntax-parameter syntax-python\">logger</span><span class=\"syntax-keyword syntax-operator syntax-assignment syntax-python\">=</span><span class=\"syntax-meta syntax-qualified-name syntax-python\"><span class=\"syntax-meta syntax-generic-name syntax-python\">logger</span></span></span><span class=\"syntax-punctuation syntax-section syntax-arguments syntax-end syntax-python\">)</span></span>\n<span class=\"syntax-meta syntax-function syntax-python\">    <span class=\"syntax-storage syntax-type syntax-function syntax-python\">def</span> <span class=\"syntax-entity syntax-name syntax-function syntax-python\"><span class=\"syntax-meta syntax-generic-name syntax-python\">compute_log_prob</span></span></span><span class=\"syntax-meta syntax-function syntax-parameters syntax-python\"><span class=\"syntax-punctuation syntax-section syntax-parameters syntax-begin syntax-python\">(</span></span><span class=\"syntax-meta syntax-function syntax-parameters syntax-python\"><span class=\"syntax-variable syntax-parameter syntax-python\">self</span><span class=\"syntax-punctuation syntax-separator syntax-parameters syntax-python\">,</span> <span class=\"syntax-variable syntax-parameter syntax-python\">data</span></span><span class=\"syntax-meta syntax-function syntax-parameters syntax-annotation syntax-python\"><span class=\"syntax-punctuation syntax-separator syntax-annotation syntax-parameter syntax-python\">:</span> <span class=\"syntax-meta syntax-qualified-name syntax-python\"><span class=\"syntax-meta syntax-generic-name syntax-python\">DataProto</span></span></span><span class=\"syntax-meta syntax-function syntax-parameters syntax-python\"><span class=\"syntax-punctuation syntax-separator syntax-parameters syntax-python\">,</span> <span class=\"syntax-variable syntax-parameter syntax-python\">calculate_entropy</span></span><span class=\"syntax-meta syntax-function syntax-parameters syntax-default-value syntax-python\"><span class=\"syntax-keyword syntax-operator syntax-assignment syntax-python\">=</span><span class=\"syntax-constant syntax-language syntax-python\">False</span></span><span class=\"syntax-meta syntax-function syntax-parameters syntax-python\"><span class=\"syntax-punctuation syntax-section syntax-parameters syntax-end syntax-python\">)</span></span><span class=\"syntax-meta syntax-function syntax-python\"> </span><span class=\"syntax-meta syntax-function syntax-annotation syntax-return syntax-python\"><span class=\"syntax-punctuation syntax-separator syntax-annotation syntax-return syntax-python\">-&gt;</span> <span class=\"syntax-meta syntax-qualified-name syntax-python\"><span class=\"syntax-meta syntax-generic-name syntax-python\">torch</span><span class=\"syntax-punctuation syntax-accessor syntax-dot syntax-python\">.</span><span class=\"syntax-meta syntax-generic-name syntax-python\">Tensor</span></span></span><span class=\"syntax-meta syntax-function syntax-python\"><span class=\"syntax-punctuation syntax-section syntax-function syntax-begin syntax-python\">:</span></span>\n        <span class=\"syntax-comment syntax-block syntax-documentation syntax-python\"><span class=\"syntax-punctuation syntax-definition syntax-comment syntax-begin syntax-python\">&quot;&quot;&quot;</span>\n        ....\n        <span class=\"syntax-punctuation syntax-definition syntax-comment syntax-end syntax-python\">&quot;&quot;&quot;</span></span>\n        <span class=\"syntax-comment syntax-line syntax-number-sign syntax-python\"><span class=\"syntax-punctuation syntax-definition syntax-comment syntax-python\">#</span> set to eval, this essentially prioritizes parallelism at the cost of memory efficiency\n</span>        <span class=\"syntax-meta syntax-function-call syntax-python\"><span class=\"syntax-meta syntax-qualified-name syntax-python\"><span class=\"syntax-variable syntax-language syntax-python\">self</span><span class=\"syntax-punctuation syntax-accessor syntax-dot syntax-python\">.</span><span class=\"syntax-meta syntax-generic-name syntax-python\">actor_module</span><span class=\"syntax-punctuation syntax-accessor syntax-dot syntax-python\">.</span></span><span class=\"syntax-meta syntax-qualified-name syntax-python\"><span class=\"syntax-variable syntax-function syntax-python\">eval</span></span><span class=\"syntax-punctuation syntax-section syntax-arguments syntax-begin syntax-python\">(</span><span class=\"syntax-punctuation syntax-section syntax-arguments syntax-end syntax-python\">)</span></span>\n        <span class=\"syntax-constant syntax-language syntax-python\">...</span>\n</span></span></code></pre>\n<p>We patched the Hugging Face implementation to use a more memory-efficient execution path, avoiding repeated materialization of experts.</p>\n<h3>Sequence Parallel with Flash Attention V3</h3>\n<p>Agentic RL requires the agent to interact with the environment over multiple steps while maintaining an ever-expanding context. Observations and environment feedback from each step are appended to the context and used as input for subsequent decision-making, which introduces significant challenges for memory efficiency and scalability during training.</p>\n<p>Under fully sharded data parallelism (FSDP), model parameters, optimizer states, and gradients are sharded across the entire world size (i.e., all GPUs in the training cluster). Each GPU stores and updates only its assigned parameter shards, while rollout data are replicated across all GPUs—meaning every GPU processes the full agent interaction history for each rollout.</p>\n<p>During the forward pass, when computation reaches a layer whose parameters are not locally available, an <code>all_gather</code> operation is triggered to materialize the full parameters across GPUs. During the backward pass, a corresponding <code>reduce_scatter</code> operation aggregates gradients and ensures that each GPU retains only its local shard. This provides a degree of scaling: as the number of GPUs increases, the per-GPU memory footprint decreases.</p>\n<p>FSDP provides model-level scaling by sharding model parameters, gradients, and optimizer states across GPUs. Sequence parallelism (or context parallelism) further reduces per-GPU memory consumption by partitioning the input sequence across devices, thereby lowering the peak activation memory on each GPU.</p>\n<p>As the number of sequence-parallel dimensions increases, the maximum activation memory per GPU correspondingly decreases. We have implemented sequence parallelism to be attention-sink-aware and compatible with FlashAttention v3 (Figure 8, right).</p>\n<p><a target=\"_blank\" rel=\"noopener noreferrer\" href=\"https://cdn-uploads.huggingface.co/production/uploads/64efbd469e7770db74cb72f5/ryT_y9BpbFSdMDxNYlVlK.png\"><picture class=\"article-picture article-picture-fallback\" style=\"--image-ratio:16 / 9\"><img src=\"https://cdn-uploads.huggingface.co/production/uploads/64efbd469e7770db74cb72f5/ryT_y9BpbFSdMDxNYlVlK.png\" alt=\"SP  (2)\" loading=\"lazy\" decoding=\"async\" fetchpriority=\"low\" referrerpolicy=\"no-referrer\" class=\"progressive-image\"></picture></a></p>\n<p><strong>Figure 8</strong>. Left: Inference without sequence parallelism. Right: Inference with sequence parallelism, where additional all-to-all communication is performed before and after the attention layer. This partitions the sequence across parallel workers and reduces the peak memory footprint of attention computation by a factor proportional to the sequence-parallelism degree.</p>\n<p>Sequence parallelism scales along the sequence dimension to reduce the per-GPU activation footprint. Input tokens from all sequences are packed into a single contiguous list by removing padding tokens, while position IDs are used to distinguish tokens belonging to different sequences. This design naturally benefits from FlashAttention’s variable-length support. For sequence parallelism, layers other than the attention layer do not have inter-position dependencies; therefore, they do not require each GPU to hold a complete sequence shard, and no additional communication is needed for these layers.</p>\n<p>The attention layer, however, requires all tokens belonging to the same sequence to be present on the same GPU in order to compute attention weights correctly. To satisfy this constraint, an all-to-all communication is performed to gather sequence elements, with the split performed at the attention-head level. This design avoids communication within the attention computation itself, which would otherwise be prohibitively expensive. After the attention layer, a single all-to-all communication redistributes the outputs back to their original sequence-parallel layout, after which the remaining non-attention layers can proceed without further synchronization.</p>\n<h2>Conclusion</h2>\n<p>Our journey to enable agentic RL training for the GPT-OSS backbone model was a practical retrospective, highlighting that unlocking advanced capabilities in open-source LLMs requires meticulous, deep-dive engineering.</p>\n<p>We made contributions that transformed the viability of GPT-OSS for agentic applications, specifically by:</p>\n<ul>\n<li>\n<p><strong>Stabilizing PPO:</strong> We contributed a fix to restore on-policy integrity, overriding the log-probability mismatch caused by the MoE architecture’s non-determinism (Figure 2).</p>\n</li>\n<li>\n<p><strong>Enabling Attention Sink Support:</strong> We successfully implemented and integrated the attention sink backward pass into FlashAttention v3, correcting the catastrophic training–inference mismatch that had previously caused instability and slow convergence (Figures 5, 6, and 7).</p>\n</li>\n<li>\n<p><strong>Scaling Memory Efficiency:</strong> We introduced crucial memory optimizations, including patching the MoE materialization process and integrating sequence parallelism with the new attention sink support, enabling training with the long context windows essential for multi-step agents (Figure 8).</p>\n</li>\n</ul>\n<p>These engineering efforts validate GPT-OSS as a scalable and high-performance backbone for building the next generation of intelligent, multi-step decision-making agents.</p>\n<h2>Acknowledgments</h2>\n<p>Thanks to Deepak Agarwal, Bee-Chung Chen, Animesh Singh, Gungor Polatkan, Balaji Krishnapuram, and Jitendra Agarwal for their leadership support.</p>\n<h2>References</h2>\n<ol>\n<li>\n<p>Feng, Jiazhan, et al. <em>Retool: Reinforcement Learning for Strategic Tool Use in LLMs.</em> arXiv preprint arXiv:2504.11536 (2025).</p>\n</li>\n<li>\n<p>Xiao, Guangxuan, et al. <em>Efficient Streaming Language Models with Attention Sinks.</em> arXiv preprint arXiv:2309.17453 (2023).</p>\n</li>\n<li>\n<p>When Speed Kills Stability: Demystifying RL Collapse from the Training–Inference Mismatch.<br />\n<a target=\"_blank\" rel=\"noopener noreferrer\" href=\"https://yingru.notion.site/When-Speed-Kills-Stability-Demystifying-RL-Collapse-from-the-Training-Inference-Mismatch-271211a558b7808d8b12d403fd15edda\">https://yingru.notion.site/When-Speed-Kills-Stability-Demystifying-RL-Collapse-from-the-Training-Inference-Mismatch-271211a558b7808d8b12d403fd15edda</a></p>\n</li>\n</ol>\n",
  "content_markdown": "Agentic reinforcement learning (RL) extends traditional LLM training by optimizing not just a single-turn response, but an entire decision-making process learned through direct interaction with an environment during training. Unlike traditional single-turn reinforcement learning or offline preference-based methods that rely on static datasets, agentic RL trains policies by actively collecting on-policy data as the agent plans actions, invokes tools, observes outcomes, and adapts its behavior over multi-step trajectories in either simulated or real environments. This interaction-driven optimization assigns credit across long-horizon decisions, where intermediate choices such as query reformulation, tool selection, and execution order directly influence downstream success. Training follows an iterative closed loop in which the agent interacts with the environment to collect rollout trajectories, computes rewards over these trajectories, updates the policy based on observed outcomes, and then uses the updated policy to drive the next round of interaction and data collection such as GRPO or PPO algorithms..\n\nLinkedIn is an AI-first company that's built agents to help professionals be more successful. In this setting, models must reason over incomplete information, interact with structured services, and adapt to evolving user intent across multiple steps rather than produce a single static response. These capabilities are especially critical for agents that support the goals of recruiters, job and knowledge seekers, and learners end users, such as retrieving information, refining queries, coordinating tools, and executing multi-step workflows. By learning robust decision policies through interaction, agentic RL provides a principled foundation for building scalable, reliable, and adaptable AI systems through end-to-end optimization.\n\nThe GPT-OSS model has shown comparable performance to OpenAI o3-mini and o4-mini \\[[ref](https://openai.com/index/introducing-gpt-oss/)\\], but its suitability for agentic reinforcement learning training has not yet been validated. Most recent work focuses on fine-tuning without tool calling, such as: [Fine-tuning with gpt-oss and Hugging Face Transformers](https://cookbook.openai.com/articles/gpt-oss/fine-tune-transfomers) and [unsloth tutorial: how to fine-tune gpt-oss](https://unsloth.ai/docs/models/gpt-oss-how-to-run-and-fine-tune/tutorial-how-to-fine-tune-gpt-oss). This blog explores the journey to unlock [agentic RL](https://github.com/volcengine/verl/issues/3794) training for GPT-OSS as a potential backbone model for agentic applications.\n\nIn our experiments, we use [verl](https://github.com/volcengine/verl) as our training framework since it is one of the most popular adopted frameworks in the open source community. We use gsm8k, [Retool](https://github.com/verl-project/verl-recipe/tree/21892b9276936efab5375c3f6b8415e472ef7118/retool) task, [verifiable instruction following task](https://arxiv.org/abs/2507.02833), which are commonly used in RL training. We focus on presenting experimental results for the GPT-OSS-20B model, and our attention-sink fix also works for GPT-OSS-120B. The Qwen-2.5-32B model is additionally used to benchmark standard metric trends during RL training.\n\n## Challenges of GPT-OSS RL Training\n\nverl has been an OSS framework used by the team, and the team has previously collaborated and contributed to it to help democratize agentic reinforcement learning training. With the introduction of the new Harmony chat template in GPT-OSS, the first step is to ensure that the training framework fully supports the updated message format and conversation semantics required by Harmony. This step helps rollout generation, trajectory construction, and tool parsing remain consistent and correct under the new template.\n\nThe team uses ReTool as a representative example to verify code correctness. ReTool is an agentic coding task in which the model is asked to solve a math problem with the assistance of a code compiler tool. This setup allows the model to focus on core reasoning and algorithmic logic, while delegating the actual arithmetic and execution to the tool. During an episode, the model interacts with the code tool multiple times, using execution results as feedback to refine its solution. At the end of the trajectory, the model produces a final answer, on which the reward is computed.\n\nDuring the initial training runs, we observed exploding KL divergence and entropy, along with non-increasing rewards, indicating underlying issues in the GPT-OSS training setup, as shown in Figure 1.\n\n| Average Gradient Norm                                                                                                                                                                                                                                      | Average Reward                                                                                                                                                                                                                                      |\n| ---------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | --------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- |\n| [![Average gradient norm in a batch](https://cdn-uploads.huggingface.co/production/uploads/64efbd469e7770db74cb72f5/5eWLHQ-EAWKjPc7T0c6CK.png)](https://cdn-uploads.huggingface.co/production/uploads/64efbd469e7770db74cb72f5/5eWLHQ-EAWKjPc7T0c6CK.png) | [![Average reward in a batch](https://cdn-uploads.huggingface.co/production/uploads/64efbd469e7770db74cb72f5/ZpBvRXLKyAT-PIdsuRi11.png)](https://cdn-uploads.huggingface.co/production/uploads/64efbd469e7770db74cb72f5/ZpBvRXLKyAT-PIdsuRi11.png) |\n\n**Figure 1.** Left: Qwen32b has significantly higher rewards compared to GPT-OSS 20B; Right: The gradient norm exploded as training progressed.\n\n## A Practical Debugging Journey in verl: Restoring PPO On-Policy Integrity\n\n### Restoring PPO On-Policy Integrity: A Fix for MoE Log-Probability Mismatch\n\n![](https://cdn-uploads.huggingface.co/production/uploads/64efbd469e7770db74cb72f5/21xIlXKeAlTK5qKSp-TkX.png)\n\n**Figure 2.** Non-zero importance sampling clip value even for on-policy training.\n\nWe focus on on-policy methods because they provide greater stability and more reliable convergence. The foundation of pure on-policy Proximal Policy Optimization (PPO) mandates that the importance sampling ratio must be exactly 1. The mathematical definition of the importance ratio is:\n\nratio=π(a∣s)πold(a∣s) \\\\text{ratio} = \\\\frac{\\\\pi(a \\\\mid s)}{\\\\pi\\_{\\\\text{old}}(a \\\\mid s)}\n\nThis requirement ensures that the policy update is executed only on the data generated by the current policy π(a | s) = πold(a | s), preventing unintended clipping.\n\nWe have observed the non-zero clipping value in our ReTool training, as shown in Figure 2, stemming from a mismatch between the two log-probabilities:\n\n- Current log-probability `log_prob`: log(π(a | s))\n- Old log-probability `old_log_prob`: log(πold(a | s))\n\n**Root Cause: The Dual Forward Pass and MoE Architecture**\n\nPrior to verl 0.3.0, the implementation relied on two separate forward passes (one to compute the current `log_prob` and one to retrieve the stored `old_log_prob`) for the same state-action pair.\n\nIn a Mixture of Experts (MoE) architecture like GPT-OSS, the gating network routes the input to different experts. Due to implementation factors (e.g., subtle floating-point differences or explicit stochasticity), the expert routing can differ slightly between the two passes. Readers who are interested can further read *[Stabilizing MoE Reinforcement Learning by Aligning Training and Inference Routers](https://arxiv.org/pdf/2510.11370v1)*.\n\nThis difference in routing leads to:\n\nlog⁡(π(a∣s))≠log⁡(πold(a∣s)) \\\\log(\\\\pi(a \\\\mid s)) \\\\neq \\\\log(\\\\pi\\_{\\\\text{old}}(a \\\\mid s))\n\nThe resulting ratio deviates from 1, falsely triggering the PPO clip and violating the core on-policy assumption.\n\n**Solution: Enforcing Ratio = 1 via Log-Probability Substitution**\n\nThe fix resolves the issue by logically overriding the flawed computation when the environment is known to be on-policy (i.e., when the minibatch size equals the global batch size):\n\n```python\nif on_policy:\n    old_log_prob = log_prob.detach()\nelse:\n    old_log_prob = model_inputs[\"old_log_probs\"]\n```\n\nBy setting `old_log_prob` equal to the newly computed `log_prob` (detached to prevent gradient flow through the reference value), the importance ratio is mathematically forced back to 1. This strategy bypasses the instability caused by MoE's non-deterministic routing and guarantees strict on-policy behavior during PPO training.\n\n* * *\n\n### Correcting Training–Inference Mismatch\n\nAlthough fixing the log-probability mismatch reduced the importance-sampling clip ratio to zero, gradient norms continued to explode and rewards failed to improve. To isolate the issue, we simplified training to GSM8K, a single-step task without agentic tool use. The same instability persisted, as shown in the green curves in Figure 3, indicating a **fundamental issue in basic RL training with GPT-OSS under verl.**\n\nWe hypothesize that **training–inference mismatch** could be a potential cause: discrepancies between inference-time execution—where engines such as vLLM and SGLang aggressively optimize for throughput—and training-time execution under FSDP, which prioritizes numerical precision and stability, can effectively turn otherwise **on-policy RL into off-policy optimization.**\n\nThis [blog](https://yingru.notion.site/When-Speed-Kills-Stability-Demystifying-RL-Collapse-from-the-Training-Inference-Mismatch-271211a558b7808d8b12d403fd15edda) details why such mismatches lead to unstable gradients and non-improving rewards. Figure 3 compares training runs with and without rollout correction (see this [verl blog](https://verl.readthedocs.io/en/latest/algo/rollout_corr.html) for details). After applying rollout correction, training dynamics improve significantly, with gradient norms remaining stable rather than exploding.\n\nHowever, as shown in the left plot of Figure 4, the reward increases only modestly, and convergence on the simple GSM8K task remains substantially slower compared to smaller dense model variants.\n\n| Average Entropy                                                                                               | Average Gradient Norm                                                                                         | Average KL Loss                                                                                               |\n| ------------------------------------------------------------------------------------------------------------- | ------------------------------------------------------------------------------------------------------------- | ------------------------------------------------------------------------------------------------------------- |\n| ![](https://cdn-uploads.huggingface.co/production/uploads/64efbd469e7770db74cb72f5/MoBnatEipCi_OScgm5fAJ.png) | ![](https://cdn-uploads.huggingface.co/production/uploads/64efbd469e7770db74cb72f5/k4Su03zqq6Cg-7jM6Il6J.png) | ![](https://cdn-uploads.huggingface.co/production/uploads/64efbd469e7770db74cb72f5/g77U9c8a-FY1rvwru96Ez.png) |\n\n**Figure 3.** Gradient norm behavior under different training configurations. Green: Training without rollout correction, exhibiting unstable gradients. Red: Training with the attention layer frozen to isolate the issue to the attention mechanism, resulting in partial stabilization. Blue: Training with rollout correction enabled (sequence-level importance sampling), yielding stable gradient norms.\n\n| Average Reward                                                                                                                                                                                                                                      | Max Log-Perplexity Difference                                                                                                                                                                                                                                                                                              |\n| --------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | -------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- |\n| [![Average reward in a batch](https://cdn-uploads.huggingface.co/production/uploads/64efbd469e7770db74cb72f5/ihx-XsWH51V0-JM46jODE.png)](https://cdn-uploads.huggingface.co/production/uploads/64efbd469e7770db74cb72f5/ihx-XsWH51V0-JM46jODE.png) | [![Maximum absolute log-perplexity difference in a batch between rollout policy and training policy](https://cdn-uploads.huggingface.co/production/uploads/64efbd469e7770db74cb72f5/fP5KNR2XYY7EH-muYBk_X.png)](https://cdn-uploads.huggingface.co/production/uploads/64efbd469e7770db74cb72f5/fP5KNR2XYY7EH-muYBk_X.png) |\n\n**Figure 4.** Left: Reward improvement on GSM8K remains slow even after applying rollout correction, with performance comparable to runs where the attention layer is frozen during training. Right: A substantial log-ppl mismatch is observed between the inference engine (SGLang with Triton kernels supporting attention-sink forward passes) and the training stack (FSDP with FlashAttention-v2), indicating a large training–inference inconsistency.\n\nTo further isolate the root cause, we freeze the attention layers during training and observe reward dynamics similar to those of runs without freezing (blue curve vs yellow curve in Figure 4). This indicates that learning is primarily driven by the MoE layers, while the attention mechanism contributes less effectively than expected. In addition, we observe a substantial token-level probability mismatch between the inference engine and the distributed training stack which are using different attention kernels. Together, these observations motivate a deeper investigation into the attention mechanism.\n\n* * *\n\n## Attention Sink Support in FlashAttentionV3\n\nAttention sinks used in GPT-OSS are learnable scalar parameters (one per attention head) that act as \"virtual tokens\" in the softmax computation. They allow the model to allocate attention mass to a learned sink rather than forcing all attention to content tokens, which has been shown to improve attention stability in streaming inference and training with sliding-window attention.\n\nAfter a deeper investigation, we identified several major issues:\n\n- verl hard-codes FlashAttention v2 in `fsdp_worker`, which does not support attention sinks.\n- The attention sink backward pass is not supported in FlashAttention v2 and v3, so it does not work as expected even when FlashAttention v3 is enabled.\n- Since the forward pass has not yet been merged into the original FlashAttention v3 repository, we leveraged the forward pass from the vLLM FlashAttention fork ([PR #75](https://github.com/vllm-project/flash-attention/pull/75)) and implemented the backward pass to compute the sink gradient.\n\n### Standard Attention\n\n```python\nscores = QK^T / sqrt(d)               # [B, H, N_q, N_k]\nprobs = softmax(scores, dim=-1)       # Σ_j P_ij = 1\noutput = probs @ V                   # [B, H, N_q, d_v]\n```\n\n### Attention with Sinks (GPT-OSS)\n\n```python\nscores = QK^T / sqrt(d)                               # [B, H, N_q, N_k]\ncombined = concat([scores, sink_param], dim=-1)       # [B, H, N_q, N_k+1]\nprobs = softmax(combined, dim=-1)                     # Σ_j P_ij + P_sink = 1\nprobs_content = probs[..., :-1]                       # Drop sink component\noutput = probs_content @ V                           # [B, H, N_q, d_v]\n```\n\n**Key difference:** The sink participates in softmax normalization but doesn't contribute to the output.\n\n### Mathematical Formulation\n\nThe attention weight for content token j in row i is defined as:\n\nPij=exp⁡(Sij)∑j′=1Nkexp⁡(Sij′)+exp⁡(Sh) P\\_{ij} = \\\\frac{\\\\exp(S\\_{ij})} {\\\\sum\\_{j'=1}^{N\\_k} \\\\exp(S\\_{ij'}) + \\\\exp(S\\_h)}\n\n**Where:**\n\n- Sij = Qi Kj⊤ / √d are the attention scores\n- Pij are the attention weights for the content tokens\n- Sh is the learnable sink parameter for head h\n\n**Sink Probability:**\n\nThe sink probability is computed but not used in the output:\n\nPi,h=exp⁡(Sh)∑j′=1Nkexp⁡(Sij′)+exp⁡(Sh) P\\_{i,h} = \\\\frac{\\\\exp(S\\_h)} {\\\\sum\\_{j'=1}^{N\\_k} \\\\exp(S\\_{ij'}) + \\\\exp(S\\_h)}\n\n### Backward Pass\n\nThe gradient of the loss L with respect to the sink parameter Sh is:\n\n∂L∂Sh=−∑iPi,h(∂L∂Si,h−∑j∈{1,…,Nk}Pij∂L∂Sij) \\\\frac{\\\\partial L}{\\\\partial S\\_h} = - \\\\sum\\_i P\\_{i,h} \\\\left( \\\\frac{\\\\partial L}{\\\\partial S\\_{i,h}} - \\\\sum\\_{j \\\\in \\\\{1,\\\\ldots,N\\_k\\\\}} P\\_{ij} \\\\frac{\\\\partial L}{\\\\partial S\\_{ij}} \\\\right)\n\n**Where:**\n\n- Pi,h is the sink attention probability for row i\n- ∂L/∂Sij is the gradient with respect to the attention scores, including the sink\n\n**Simplified Gradient:**\n\nSince the sink is computed but not used in the output, its gradient ∂L/∂Si,h = 0.\n\nTherefore, the backward equation simplifies to:\n\n∂L∂Sh=−∑iPi,h(∑j∈{1,…,Nk}Pij∂L∂Sij) \\\\frac{\\\\partial L}{\\\\partial S\\_h} = - \\\\sum\\_i P\\_{i,h} \\\\left( \\\\sum\\_{j \\\\in \\\\{1,\\\\ldots,N\\_k\\\\}} P\\_{ij} \\\\frac{\\\\partial L}{\\\\partial S\\_{ij}} \\\\right)\n\nThe forward pass was adapted from vLLM's FlashAttention fork, and we implemented the backward pass to compute gradients for the sink parameters. The implementation will be released following the internal review process.\n\n* * *\n\n### Results\n\nAfter applying the fix in FlashAttention v3, we observe substantially faster convergence for GPT-OSS-20B across a range of reinforcement learning tasks. These include single-turn RL on math reasoning (GSM8K — red curve in Figure 5), instruction following (VerifyIf, evaluated on an out-of-domain multi-if benchmark — Figure 6), and multi-turn agentic RL with tool use (ReTool — Figure 7).\n\nAcross all settings, training becomes stable and exhibits steady reward improvement.\n\n![](https://cdn-uploads.huggingface.co/production/uploads/64efbd469e7770db74cb72f5/6TYGevydK99nQ-I1QTouf.png)\n\n**Figure 5.**. Single Turn GSM8K, the red curve converges much faster than the rest without the fix\n\n| Average Entropy                                                                                                                                                                                                                                      | Average Gradient Norm                                                                                                                                                                                                                                      | Average Reward                                                                                                                                                                                                                                      |\n| ---------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | ---------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | --------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- |\n| [![Average entropy in a batch](https://cdn-uploads.huggingface.co/production/uploads/64efbd469e7770db74cb72f5/ydmBLCSGlD9YKWiIocI1S.png)](https://cdn-uploads.huggingface.co/production/uploads/64efbd469e7770db74cb72f5/ydmBLCSGlD9YKWiIocI1S.png) | [![Average gradient norm in a batch](https://cdn-uploads.huggingface.co/production/uploads/64efbd469e7770db74cb72f5/_kl8mn_CXPsRYJ467IbFs.png)](https://cdn-uploads.huggingface.co/production/uploads/64efbd469e7770db74cb72f5/_kl8mn_CXPsRYJ467IbFs.png) | [![Average reward in a batch](https://cdn-uploads.huggingface.co/production/uploads/64efbd469e7770db74cb72f5/vb2JmmSu-LI5szC_84KsM.png)](https://cdn-uploads.huggingface.co/production/uploads/64efbd469e7770db74cb72f5/vb2JmmSu-LI5szC_84KsM.png) |\n\n**Figure 6**. On verifiable instruction following the task, the run without the fix collapsed (blue), and the run with fix showed steady reward improvement.\n\n| Average Gradient Norm                                                                                                                                                                                                                                      | Average Reward                                                                                                                                                                                                                                      | Validation Accuracy                                                                                                                                                                                                                                                |\n| ---------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | --------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | ------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------ |\n| [![Average gradient norm in a batch](https://cdn-uploads.huggingface.co/production/uploads/64efbd469e7770db74cb72f5/_Xz7_RLhYuYhzGAeMjDXs.png)](https://cdn-uploads.huggingface.co/production/uploads/64efbd469e7770db74cb72f5/_Xz7_RLhYuYhzGAeMjDXs.png) | [![Average reward in a batch](https://cdn-uploads.huggingface.co/production/uploads/64efbd469e7770db74cb72f5/uFX-sTWI6knecIf56uahk.png)](https://cdn-uploads.huggingface.co/production/uploads/64efbd469e7770db74cb72f5/uFX-sTWI6knecIf56uahk.png) | [![val score accuracy mean@30 for aime_2025](https://cdn-uploads.huggingface.co/production/uploads/64efbd469e7770db74cb72f5/5x2mdkpHcdvctZ96yfv58.png)](https://cdn-uploads.huggingface.co/production/uploads/64efbd469e7770db74cb72f5/5x2mdkpHcdvctZ96yfv58.png) |\n\n**Figure 7**. On the Retool task, the run with fix showed steady reward improvement and no gradient exploding (fa2 is the flash attention 2 without the fix while fa3 is the flash attention 3 with the fix). After the fix, the validation accuracy score goes up now.\n\n## Memory-Efficient Training\n\n### Mitigating FSDP Memory Blow-Ups Caused by Repeated MoE Expert Materialization\n\nOne issue we consistently encountered was excessive memory allocation during the FSDP forward pass, which led to repeated out-of-memory (OOM) failures when training GPT-OSS-20B bf16 models on 16 H200 nodes (max response length: 16k, prompt length: 8k). This behavior is highly unexpected for a 20B-parameter MoE model.\n\n```text\n2025-11-27T11:15:27.927Z [36m(TaskRunner pid=32081)[0m File \"/home/jobuser/.local/lib/python3.10/site-packages/transformers/models/gpt_oss/modeling_gpt_oss.py\", line 123, in forward\n2025-11-27T11:15:27.927Z [36m(TaskRunner pid=32081)[0m hidden_states = hidden_states.repeat(num_experts, 1)\n2025-11-27T11:15:27.927Z [36m(TaskRunner pid=32081)[0m torch.OutOfMemoryError: CUDA out of memory. Tried to allocate 180.00 GiB. GPU 0 has a total capacity of 139.72 GiB of which 110.94 GiB is free. Process 685851 has 24.88 GiB memory in use. Process 692458 has 3.87 GiB memory in use. Of the allocated memory 23.28 GiB is allocated by PyTorch, and 84.43 MiB is reserved by PyTorch but unallocated.\n```\n\nWe identified the issue as originating from two different implementations of the MoE forward path in Hugging Face Transformers. This issue has also been reported by other users: [https://github.com/huggingface/transformers/issues/40073](https://github.com/huggingface/transformers/issues/40073); When verl computes log-probabilities under FSDP, the inference forward path is triggered. In the current Hugging Face implementation, this path duplicates hidden states for all experts and performs batched matrix multiplication, materializing extremely large tensors in GPU memory. By contrast, the training forward path uses a for-loop to process each expert sequentially and then combines the results. While slower, this approach is significantly more memory efficient.\n\n```python\n    @GPUMemoryLogger(role=\"dp actor\", logger=logger)\n    def compute_log_prob(self, data: DataProto, calculate_entropy=False) -> torch.Tensor:\n        \"\"\"\n        ....\n        \"\"\"\n        # set to eval, this essentially prioritizes parallelism at the cost of memory efficiency\n        self.actor_module.eval()\n        ...\n```\n\nWe patched the Hugging Face implementation to use a more memory-efficient execution path, avoiding repeated materialization of experts.\n\n### Sequence Parallel with Flash Attention V3\n\nAgentic RL requires the agent to interact with the environment over multiple steps while maintaining an ever-expanding context. Observations and environment feedback from each step are appended to the context and used as input for subsequent decision-making, which introduces significant challenges for memory efficiency and scalability during training.\n\nUnder fully sharded data parallelism (FSDP), model parameters, optimizer states, and gradients are sharded across the entire world size (i.e., all GPUs in the training cluster). Each GPU stores and updates only its assigned parameter shards, while rollout data are replicated across all GPUs—meaning every GPU processes the full agent interaction history for each rollout.\n\nDuring the forward pass, when computation reaches a layer whose parameters are not locally available, an `all_gather` operation is triggered to materialize the full parameters across GPUs. During the backward pass, a corresponding `reduce_scatter` operation aggregates gradients and ensures that each GPU retains only its local shard. This provides a degree of scaling: as the number of GPUs increases, the per-GPU memory footprint decreases.\n\nFSDP provides model-level scaling by sharding model parameters, gradients, and optimizer states across GPUs. Sequence parallelism (or context parallelism) further reduces per-GPU memory consumption by partitioning the input sequence across devices, thereby lowering the peak activation memory on each GPU.\n\nAs the number of sequence-parallel dimensions increases, the maximum activation memory per GPU correspondingly decreases. We have implemented sequence parallelism to be attention-sink-aware and compatible with FlashAttention v3 (Figure 8, right).\n\n[![SP  (2)](https://cdn-uploads.huggingface.co/production/uploads/64efbd469e7770db74cb72f5/ryT_y9BpbFSdMDxNYlVlK.png)](https://cdn-uploads.huggingface.co/production/uploads/64efbd469e7770db74cb72f5/ryT_y9BpbFSdMDxNYlVlK.png)\n\n**Figure 8**. Left: Inference without sequence parallelism. Right: Inference with sequence parallelism, where additional all-to-all communication is performed before and after the attention layer. This partitions the sequence across parallel workers and reduces the peak memory footprint of attention computation by a factor proportional to the sequence-parallelism degree.\n\nSequence parallelism scales along the sequence dimension to reduce the per-GPU activation footprint. Input tokens from all sequences are packed into a single contiguous list by removing padding tokens, while position IDs are used to distinguish tokens belonging to different sequences. This design naturally benefits from FlashAttention’s variable-length support. For sequence parallelism, layers other than the attention layer do not have inter-position dependencies; therefore, they do not require each GPU to hold a complete sequence shard, and no additional communication is needed for these layers.\n\nThe attention layer, however, requires all tokens belonging to the same sequence to be present on the same GPU in order to compute attention weights correctly. To satisfy this constraint, an all-to-all communication is performed to gather sequence elements, with the split performed at the attention-head level. This design avoids communication within the attention computation itself, which would otherwise be prohibitively expensive. After the attention layer, a single all-to-all communication redistributes the outputs back to their original sequence-parallel layout, after which the remaining non-attention layers can proceed without further synchronization.\n\n## Conclusion\n\nOur journey to enable agentic RL training for the GPT-OSS backbone model was a practical retrospective, highlighting that unlocking advanced capabilities in open-source LLMs requires meticulous, deep-dive engineering.\n\nWe made contributions that transformed the viability of GPT-OSS for agentic applications, specifically by:\n\n- **Stabilizing PPO:** We contributed a fix to restore on-policy integrity, overriding the log-probability mismatch caused by the MoE architecture’s non-determinism (Figure 2).\n\n- **Enabling Attention Sink Support:** We successfully implemented and integrated the attention sink backward pass into FlashAttention v3, correcting the catastrophic training–inference mismatch that had previously caused instability and slow convergence (Figures 5, 6, and 7).\n\n- **Scaling Memory Efficiency:** We introduced crucial memory optimizations, including patching the MoE materialization process and integrating sequence parallelism with the new attention sink support, enabling training with the long context windows essential for multi-step agents (Figure 8).\n\nThese engineering efforts validate GPT-OSS as a scalable and high-performance backbone for building the next generation of intelligent, multi-step decision-making agents.\n\n## Acknowledgments\n\nThanks to Deepak Agarwal, Bee-Chung Chen, Animesh Singh, Gungor Polatkan, Balaji Krishnapuram, and Jitendra Agarwal for their leadership support.\n\n## References\n\n1. Feng, Jiazhan, et al. *Retool: Reinforcement Learning for Strategic Tool Use in LLMs.* arXiv preprint arXiv:2504.11536 (2025).\n\n2. Xiao, Guangxuan, et al. *Efficient Streaming Language Models with Attention Sinks.* arXiv preprint arXiv:2309.17453 (2023).\n\n3. When Speed Kills Stability: Demystifying RL Collapse from the Training–Inference Mismatch.\\\n   [https://yingru.notion.site/When-Speed-Kills-Stability-Demystifying-RL-Collapse-from-the-Training-Inference-Mismatch-271211a558b7808d8b12d403fd15edda](https://yingru.notion.site/When-Speed-Kills-Stability-Demystifying-RL-Collapse-from-the-Training-Inference-Mismatch-271211a558b7808d8b12d403fd15edda)\n",
  "dateCreated": "2026-10-04T12:50:30.804479567+00:00",
  "date_created": "2026-01-27T01:53:15+00:00",
  "date_modified": "2026-01-27T01:53:15+00:00",
  "date_published": "2026-01-27T01:53:15+00:00",
  "description": "Agentic reinforcement learning (RL) extends traditional LLM training by optimizing not just a single-turn response, but an entire decision-making process learned through direct interaction with an environment during training. Unlike…",
  "external_url": "https://huggingface.co/blog/LinkedIn/gpt-oss-agentic-rl",
  "git": {
    "commit": "ebb0941c5c06a2ab2377342d607535262a3e969e",
    "path": "items/huggingface-co/2026/01/2026-01-27-unlocking-agentic-rl-training-for-gpt-oss-a-practical.md",
    "permalink": "https://github.com/aymericbeaumet/aggr-instance/blob/ebb0941c5c06a2ab2377342d607535262a3e969e/items/huggingface-co/2026/01/2026-01-27-unlocking-agentic-rl-training-for-gpt-oss-a-practical.md"
  },
  "id": "urn:aggr:item:6c23b017065b9b1d5137072da55fff8f1b5ed005",
  "isBasedOn": {
    "@id": "https://huggingface.co/blog/LinkedIn/gpt-oss-agentic-rl"
  },
  "mainEntity": {
    "@id": "https://huggingface.co/blog/LinkedIn/gpt-oss-agentic-rl",
    "@type": "CreativeWork",
    "archivedAt": {
      "@id": "https://aggr.aymericbeaumet.com/items/huggingface-co/2026-01-27-unlocking-agentic-rl-training-for-gpt-oss-a-practical/#webpage"
    },
    "dateModified": "2026-01-27T01:53:15+00:00",
    "datePublished": "2026-01-27T01:53:15+00:00",
    "headline": "Unlocking Agentic RL Training for GPT-OSS: A Practical Retrospective",
    "timeRequired": "PT15M",
    "url": "https://huggingface.co/blog/LinkedIn/gpt-oss-agentic-rl",
    "wordCount": 3240
  },
  "name": "Archived snapshot: Unlocking Agentic RL Training for GPT-OSS: A Practical Retrospective",
  "preview": null,
  "reading_minutes": 15,
  "replicated_at": null,
  "source": "huggingface-co",
  "source_name": "Hugging Face - Blog",
  "summary": null,
  "tags": [],
  "temporalCoverage": "2026-10-04T12:50:30.804479567+00:00",
  "title": "Unlocking Agentic RL Training for GPT-OSS: A Practical Retrospective",
  "url": "https://aggr.aymericbeaumet.com/items/huggingface-co/2026-01-27-unlocking-agentic-rl-training-for-gpt-oss-a-practical/",
  "word_count": 3240
}