PaperlayerView exact revision on arXiv

Source-linked AI summary

Lightning OPD: Efficient Post-Training for Large Reasoning Models with Offline On-Policy Distillation

Yecheng Wu, Song Han, Hai Cai

arXiv:2604.13010v2cs.LGcs.AI

TL;DR

Standard OPD offers effective post-training supervision but requires a live teacher server throughout training. Lightning OPD enforces teacher consistency and precomputes teacher log-probabilities offline, achieving comparable benchmark performance with 4.0× higher training efficiency.

  • Problem

    Standard OPD requires a live teacher server throughout training, creating substantial infrastructure overhead and making large-scale experiments costly and difficult to reproduce.

  • Method

    Lightning OPD uses the same teacher for SFT and OPD, precomputes teacher log-probabilities over SFT-reference rollouts, and reuses them during offline training.

  • Results

    Lightning OPD achieves comparable performance to standard OPD across math and code reasoning tasks while delivering 4.0× higher training efficiency.

  • Takeaways & Limitations

    Lightning OPD reduces OPD infrastructure to a single standard training job while preserving dense per-token supervision.

  • Takeaways & Limitations

    Lightning OPD requires SFT data to be generated by the same teacher used during OPD.

Abstract

from arXiv · show

On-policy distillation (OPD) is an effective post-training paradigm for large language models but requires a live teacher server throughout training, resulting in substantial infrastructure overhead. We investigate whether OPD can be performed offline by precomputing teacher log-probabilities once over SFT rollouts and reusing them during training. We find that naively doing so fails to reliably match standard OPD, and trace the root cause to a previously overlooked condition we term teacher consistency, requiring that the same teacher be used for both supervised fine-tuning and OPD. Violating this condition introduces a gradient bias that degrades performance for both offline and online OPD. Building on this insight, we propose Lightning OPD, an offline on-policy distillation framework that enforces teacher consistency and eliminates the need for a live teacher server entirely. We prove 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. Experiments on math reasoning and code generation show that Lightning OPD achieves comparable performance to standard OPD while delivering 4.0x higher training efficiency. Starting from an SFT-initialized Qwen3-8B-Base model, Lightning OPD reaches 69.9% on AIME 2024 in just 30 GPU hours. Lightning OPD further scales to MoE architectures, training Qwen3-30B-A3B to 71.0% on AIME 2024 on a single 8xH100 node, substantially lowering the barrier for academic research on LLM post-training. Our code is released at https://github.com/jet-ai-projects/Lightning-OPD.

1. Introduction

Lightning OPD addresses the infrastructure burden of standard OPD by precomputing teacher supervision offline, while identifying teacher consistency as necessary for reliable distillation. Under this condition, it matches standard OPD performance with substantially greater efficiency across reasoning tasks and model scales.

  • Motivation: Standard OPD requires a continuously running teacher server to score student rollouts, creating substantial compute and reproducibility costs.This bottleneck is especially burdensome for academic researchers lacking extensive serving infrastructure.
  • Problem Diagnosis: Naive offline precomputation does not reliably match standard OPD because SFT and OPD may use different teachers, introducing gradient bias.Teacher consistency requires the same teacher model to generate SFT trajectories and provide OPD supervision.
  • Proposed Method: Lightning OPD precomputes teacher log-probabilities over SFT-reference rollouts and reuses them during training, eliminating the live teacher server.The method preserves dense per-token advantage supervision while reducing OPD infrastructure to a standard training job.
  • Evaluation: Lightning OPD achieves comparable performance to standard OPD across math and code reasoning benchmarks while removing parallel teacher-serving infrastructure.Experiments cover diverse student–teacher pairs, including dense and MoE architectures.
  • Evaluation: 69.9% on AIME 2024 is reached in 30 GPU hours for an SFT-initialized Qwen3-8B-Base model, delivering 4.0× higher training efficiency than standard OPD.On a Qwen3-30B-A3B MoE model, Lightning OPD reaches 71.0% on AIME 2024 on a single 8×H100 node.

2. Related Work

Prior work positions OPD as a stronger, denser-supervision alternative to standard offline distillation and related RL methods. Lightning OPD instead focuses on making on-policy distillation operationally offline while addressing teacher consistency rather than sparse-reward or value-estimation challenges.

  • LLM Post-Training: LLM post-training commonly combines supervised fine-tuning with reinforcement learning, using either sparse outcome rewards or dense process supervision.
  • On-Policy Distillation: On-policy distillation aligns students with teachers on student-generated rollouts and often yields stronger post-training gains than standard offline knowledge distillation.
  • Off-Policy Reinforcement Learning: Unlike off-policy reinforcement learning, Lightning OPD uses a fixed dataset with dense token-level teacher supervision, making teacher consistency its central challenge.

3. Methodology

Lightning OPD converts OPD into a fixed-data training procedure by sampling rollouts from the SFT reference policy and storing teacher log-probabilities once. Its analysis characterizes the offline–online gap, shared optima, regularization effect, and the bias caused by teacher mismatch.

  • Preliminaries: OPD trains a student to match a teacher’s token-level distribution using per-token advantages, with online and offline objectives differing in their response distributions.The online objective uses the current student distribution, whereas Lightning OPD fixes rollouts to the reference policy.
  • Stage 1: Supervised Fine-Tuning: Lightning OPD first fine-tunes the base model on trajectories generated by teacher πT, producing reference policy πref, and requires the same teacher for OPD.This teacher-consistency condition is identified as a prerequisite for a sound offline approximation.
  • Stage 2: Offline On-Policy Distillation: During preprocessing, rollouts are sampled from πref and teacher log-probabilities are stored once; training then computes advantages from stored teacher terms and online student terms.No teacher server is required during optimization.
  • Assumptions: Lightning OPD assumes bounded advantages, support coverage, and bounded score functions, with advantage clipping automatically satisfying the first condition.The support condition is natural when the student is initialized from πref.
  • Theoretical Analysis: At initialization, online and offline gradients coincide exactly; their discrepancy grows with policy drift but remains controlled under the stated assumptions.Theorem 3.5 provides the gradient discrepancy bound, while the chi-squared divergence measures drift from πref.
  • Theoretical Analysis: When the teacher is representable, both methods share a common fixed point at the student minimizing KL divergence to the teacher.At that point, the advantage is zero almost surely and both updates vanish.
  • Theoretical Analysis: The offline gradient equals the online gradient minus a covariance correction that vanishes at initialization and acts empirically as a trust-region effect.This correction stabilizes training without requiring an explicit KL penalty.
  • Teacher Consistency: Teacher mismatch biases both offline and standard OPD gradients, whereas teacher consistency restores the conditions under which Lightning OPD matches standard OPD.The mismatch analysis quantifies persistent bias through the teacher discrepancy term.

4. Experiments

Experiments show that Lightning OPD preserves standard OPD performance without a live teacher server, improves training efficiency, and scales to a 30B MoE model on one 8×H100 node. Teacher consistency is essential, with mismatched teachers degrading both methods, especially Lightning OPD.

  • Main Results: Lightning OPD matches standard OPD across five math and code benchmarks at both 4B and 8B scales without a live teacher server.At 4B, it also outperforms ExOPD on AIME 2024 and LCB v6.
  • Main Results: 68.1% versus 61.0% on AIME 2024 and 40.3% versus 29.0% on LCB v6: Lightning OPD substantially outperforms ExOPD at 4B.
  • Training Cost: 4.0× speedup at 8B reduces Lightning OPD training to 30 GPU hours, while 3.6× speedup at 4B reduces it to 20 GPU hours.Offline rollout collection and teacher log-probability precomputation are one-time operations requiring no specialized infrastructure.
  • Scaling to Mixture-of-Experts: 71.0% on AIME 2024 and 60.8% on LiveCodeBench v5: Lightning OPD trains Qwen3-30B-A3B on a single 8×H100 node.Offline precomputation avoids the memory bottleneck of co-hosting 30B student and teacher models.
  • Ablation Study: Teacher-consistent settings achieve the best performance, while mismatches degrade both standard OPD and Lightning OPD.At 8B, a mismatch drops Lightning OPD by 6.8 points versus 3.7 points for standard OPD because fixed rollouts compound the bias.

5. Conclusion

The conclusion presents Lightning OPD as an efficient offline framework that retains OPD’s dense supervision while removing the persistent teacher-server requirement. It identifies teacher consistency as necessary for avoiding bias and preserving the standard OPD optimum.

  • Lightning OPD reduces OPD infrastructure to a single standard training job by precomputing teacher log-probabilities over SFT rollouts.
  • Mismatched SFT and OPD teachers introduce an irreducible gradient bias that drives both paradigms toward a suboptimal fixed point.
  • Under teacher consistency, Lightning OPD provably shares the same optimum as standard OPD.
  • 69.9% on AIME 2024 in 30 GPU hours: Lightning OPD achieves a 4.0× speedup while matching standard OPD across benchmarks.
  • The authors suggest Lightning OPD can retain practical OPD benefits where maintaining a live teacher throughout training is impractical.

A.2. Proof of Theorem 3.5

The proof bounds the discrepancy between online and offline surrogate gradients using importance-sampling divergence. The gap is zero at initialization and grows at a controlled rate under the stated assumptions.

  • The proof uses bounded score-function gradients and advantage assumptions to control the importance-sampling-based discrepancy.
  • At initialization, when πθ = πref, the online and offline gradients are identical because the χ2 divergence is zero.
  • After k gradient steps of size η, the online–offline gradient gap grows at rate O(ηk) under standard smoothness conditions.
  • A KL-based bound is vacuous when advantages can diverge, so the theorem instead uses an L2 advantage constant and χ2 divergence.Advantage clipping and KL regularization keep χ2 small during training.

A.3. Proof of Theorem 3.6

Under teacher consistency, standard and offline OPD share the same optimum, while bounded policy drift controls their gradient discrepancy. The offline method also gains a restoring effect that discourages drift from the reference policy.

  • The on-policy objective equals negative KL divergence from the student to the teacher, so its global optimum minimizes KL(πθ‖πT).
  • When the teacher is representable, both online and offline OPD updates vanish at the teacher-matching solution.At this solution, the advantage is zero for every token, making the solution a shared fixed point.
  • When teacher capacity is limited, both methods retain the same irreducible approximation error, while their per-step gradient discrepancy is controlled by policy drift.Empirically, the drift remains small, keeping offline and online updates close in practice.
  • The final divergence includes an offline-online distribution gap, but switching to online rollouts reduces only that term, not the capacity-limited approximation error.The decomposition is heuristic rather than a formal bound for surrogate gradient dynamics.
  • The offline-online gradient gap decomposes into a covariance term induced by importance weighting.The offline gradient equals the online gradient minus Cov_πref[w,f].
  • The covariance term is zero at the reference policy and increasingly opposes movement away from it as drift grows.This creates an implicit restoring regularization effect that helps prevent policy drift.

A.5. Proof of Theorem 3.8

Teacher mismatch introduces an additive gradient bias that persists independently of policy drift, so inconsistent offline OPD updates remain corrupted even at initialization.

  • Teacher mismatch creates an irreducible additive bias in the OPD gradient.
  • The gradient discrepancy separates into a consistent component and a mismatch component associated with the teacher-difference advantage.
  • At the reference policy, the online-offline gap is exactly zero because the importance weight is one, regardless of teacher mismatch.
  • Unlike the drift-dependent online-offline gap, mismatched offline gradients carry a persistent bias that can remain nonzero throughout training.This bias corrupts the update direction even at initialization.

A.6. Proof of Theorem 3.9

Teacher inconsistency biases standard OPD toward a different effective convergence point, while bounded policy drift keeps the mismatch bias on a similar order beyond initialization.

  • For reachable student iterates, the theorem compares gradients under the OPD teacher and the SFT teacher through their teacher-induced advantage difference.
  • At initialization, the importance weight is one, reducing the gradient difference exactly to the expected teacher-mismatch term.
  • With bounded policy drift, the mismatch bias is expected to remain of similar order because the teacher difference does not depend on the student parameters.The resulting term biases standard OPD whenever σΔ > 0, degrading its effective convergence point.

B. Implementation Details

The implementation uses fixed SFT and OPD-stage configurations, with OPD capped at 4,096 response tokens for improved efficiency.

  • OPD stage: 4,096 tokens is the maximum response length used during OPD training.Evaluation generations can be much longer, but the shorter training limit achieves optimal performance while improving efficiency.
  • SFT stage: Table 5 specifies the hyperparameter configuration for the SFT stage.
  • OPD stage: Table 6 specifies the OPD hyperparameters shared by standard OPD and Lightning OPD across math and code domains.

C.1. Training Dynamics

Lightning OPD stays close to its reference policy during training, gains performance rapidly, and remains stable after early optimization steps.

  • Importance weights: The mean importance weight drops from 1 to approximately 0.94 within 20 steps and then plateaus.Its standard deviation rises sharply early and subsequently stabilizes, indicating that the student remains near the reference distribution.
  • Performance dynamics: AIME 2024 pass@1 captures nearly all performance gains within the first 50 OPD steps and remains stable thereafter.This supports using 150 steps as a sufficient training budget.
  • Checkpoint quality: Both standard OPD and Lightning OPD provide large, stable gains over the SFT baseline across SFT checkpoints.All three curves improve consistently with additional SFT steps.

E. Limitations

The evaluation focuses on mathematical reasoning and code generation, while broader interactive post-training settings remain outside the demonstrated scope.

  • Scope: Experiments focus on mathematical reasoning and code generation with well-defined verifiable evaluation metrics.
  • Open directions: Extending Lightning OPD to multi-turn agents, tool use, and open-ended instruction following remains an open direction.These settings introduce multi-turn distribution shift and challenges in precomputing teacher signals over long interaction trajectories.
  • Teacher consistency: Changing the teacher requires regenerating the SFT dataset, creating a resource-intensive one-time cost.This partially offsets training-time savings, although the cost can be amortized across multiple experiments.
Loading 2604.13010v2…