Lightning OPD: Efficient Post-Training for Large Reasoning Models with Offline On-Policy Distillation
| Authors | Yecheng Wu et al. |
| Year | 2026 |
| HF Upvotes | 11 |
| arXiv | 2604.13010 |
| Download | |
| HF Page | View on Hugging Face |
Abstract
On-policy distillation (OPD) has emerged as an efficient post-training paradigm for large language models. However, standard OPD requires a live teacher inference server throughout training, resulting in substantial infrastructure overhead. In this work, we investigate whether on-policy distillation can be performed offline. A natural approach is to precompute teacher log-probabilities once over SFT rollouts and reuse them during training. In practice, however, this offline variant fails to reliably match the performance of standard OPD. To understand this discrepancy, we identify a previously overlooked condition that is critical for any OPD pipeline, which we term teacher consistency. This condition requires that the same teacher model be used for both supervised fine-tuning and OPD. We show that violating teacher consistency introduces an irreducible gradient bias, causing both offline and online OPD to converge to a suboptimal fixed point regardless of training duration. Building on this insight, we propose Lightning OPD, an offline on-policy distillation framework that enforces teacher consistency by precomputing teacher log-probabilities over SFT rollouts. This design eliminates the need for a live teacher server entirely. We further show that, under teacher consistency, Lightning OPD shares the same optimum as standard OPD, with bounded gradient discrepancy and an implicit regularization effect that helps prevent policy drift. Extensive experiments on mathematical reasoning and code generation demonstrate that Lightning OPD achieves state-of-the-art performance with significantly improved efficiency. Starting from an SFT-initialized Qwen3-8B-Base model, Lightning OPD reaches 69.9% on AIME 2024 in just 30 GPU hours, achieving a 4.0x speedup over standard OPD and substantially lowering the barrier to entry for academic research on LLM post-training.
Engineering Breakdown
Plain English
This paper investigates whether on-policy distillation (OPD)—a post-training technique for large language models—can be performed without a live teacher inference server running throughout training. The authors discovered that a straightforward offline approach, where teacher log-probabilities are precomputed once and reused, fails to match the performance of standard OPD. They identified a critical missing condition called "teacher consistency," which requires that the same teacher model be used for both the supervised fine-tuning phase and the OPD phase. This finding explains why naive offline distillation underperforms and opens a path to dramatically reduce infrastructure costs during LLM training.
Core Technical Contribution
The paper's core contribution is identifying and formalizing the "teacher consistency" condition—a previously overlooked requirement for OPD pipelines that explains why offline distillation fails. Rather than proposing a complex new algorithm, the authors provide a diagnostic framework that clarifies what goes wrong when you precompute teacher outputs and reuse them without alignment to the SFT teacher. This is a foundational insight that shifts how practitioners should think about the dependency between the SFT phase and the OPD phase. The work demonstrates that the relationship between these two training stages is tighter than standard practice assumes, and respecting this constraint is necessary for offline distillation to work reliably.
How It Works
Standard OPD works by running a live teacher model during training: as the student model generates rollouts on-policy, the teacher immediately provides log-probabilities for those exact outputs, and the student learns to match the teacher's distribution. In the offline variant, the authors propose precomputing all teacher log-probabilities over the SFT dataset once, storing them to disk, and then during training the student simply looks up and reuses these precomputed values—eliminating the need for a live teacher server. The failure of this approach stems from teacher inconsistency: if the SFT phase and OPD phase use different teacher model checkpoints or architectures, or if the student's on-policy rollouts diverge too far from the original SFT distribution, the precomputed log-probabilities become stale or misaligned. The fix involves ensuring that both SFT and OPD use the same teacher model (same weights, same training procedure), so that the cached teacher outputs remain valid throughout training. This constraint is tighter than practitioners typically enforce, which explains why the naive offline approach fails in practice.
Production Impact
For production LLM training, this insight enables substantial infrastructure savings: instead of maintaining a live GPU-based teacher inference server (which can consume 30-50% of training compute), teams can precompute and cache teacher outputs once, reducing the infrastructure footprint to just storage and lookups. However, the teacher consistency requirement introduces a strict dependency: if you want to upgrade your SFT teacher mid-training or apply different regularization techniques to the teacher versus the student, offline distillation will break. This means deployment of offline OPD requires careful orchestration—you must commit to a teacher model upfront, run SFT to completion, then run OPD with that exact teacher in read-only mode. The practical trade-off is cost (lower) versus flexibility (higher constraints): teams with stable, fixed-size training pipelines gain immediate wins, but teams with adaptive training schedules or frequent model updates will find the teacher consistency requirement restrictive.
Limitations and When Not to Use This
The paper's analysis is based on the observation that teacher consistency matters, but does not provide detailed quantitative guidance on how much teacher drift is tolerable before performance degrades meaningfully. The work assumes you can afford to precompute and store teacher outputs for your entire SFT dataset, which may be prohibitive for very large datasets (hundreds of billions of tokens) where storage and I/O become bottlenecks. The paper does not address scenarios where you want to use distillation for continual learning or fine-tuning on new data after SFT is complete—in those settings, the offline precomputed outputs would be stale. Additionally, the paper does not explore whether teacher consistency can be relaxed with periodic updates (e.g., refreshing teacher outputs every N steps) or whether auxiliary techniques like knowledge distillation with soft targets can compensate for teacher drift.
Research Context
This work builds directly on the on-policy distillation paradigm, which emerged as a more sample-efficient alternative to RLHF for post-training LLMs by leveraging the student's own on-policy rollouts rather than replay buffers. The paper's key insight—that the teacher model must remain fixed and consistent—connects to broader literature on distribution mismatch in reinforcement learning and off-policy learning, where using stale behavior policies leads to poor performance. The contribution is primarily diagnostic rather than algorithmic: it explains a previously mysterious failure mode and clarifies the hidden assumptions in OPD implementations. This opens the door to future work on relaxed consistency conditions (partial updates, adaptive teachers) and practical strategies for offline distillation at massive scale.
:::tip Subscribe Get weekly breakdowns of papers like this in AI Letters - the newsletter for engineers building production AI systems. :::