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

Yecheng Wu; Song Han; Han Cai

cs.LGApr 14, 2026

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

Authors: Yecheng Wu, Song Han, Han Cai

Organizations: NVIDIA

Abstract

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.

Figures & tables

Panel of Figure 1 — Figure 1: (Top) Performance (Pass@1, %) and training cost of Lightning OPD compared to standard OPD and the SFT baseline on Qwen3-4B-Base and Qwen3-8B-Base models. Lightning OPD achieves comparable performance to standard OPD on both math and coding benchmarks across both scales, while eliminating the need for a live teacher server during training. At the 8B scale, Lightning OPD achieves a state-of-the-art 69.9% on AIME 2024 in just 30 GPU hours , delivering 4.0 $\times$ higher training efficiency than standard OPD. (Bottom) Intuitive comparison of GPU resource allocation. Standard OPD requires co-hosting the student and teacher, fragmenting GPU resources. Lightning OPD collects rollouts and teacher log-probabilities offline, dedicating all GPUs to student training.

Figure 2: Overview of Lightning OPD. In the SFT stage, the teacher

\pi_{T}

generates trajectories on

\mathcal{Q}_{\text{SFT}}

then the base model

\pi_{\text{base}}

is fine-tuned on these trajectories to obtain the SFT model

\pi_{\text{ref}}

. OPD stage proceeds in two phases. In the preprocessing phase, rollouts are sampled from

\pi_{\text{ref}}

\mathcal{Q}_{\text{OPD}}

and the teacher is queried once to compute and store per-token log-probabilities

\log\pi_{T}(a_{t}\mid s_{t})

, forming the offline dataset

\mathcal{D}_{\text{OPD}}

. In the training phase, the student

\pi_{\theta}

is initialized from

\pi_{\text{ref}}

and trained on

\mathcal{D}_{\text{OPD}}

. At each step, the per-token advantage

A_{t}=\log\pi_{T}(a_{t}\mid s_{t})-\log\pi_{\theta}(a_{t}\mid s_{t})

is computed by reading

\log\pi_{T}

from

\mathcal{D}_{\text{OPD}}

while computing

\log\pi_{\theta}

online. During this process, we do not need a live teacher model. In contrast, standard OPD requires a live teacher server throughout training to compute advantages at every gradient step. By eliminating this requirement, Lightning OPD reduces total training cost of OPD by 4.0

\times

at the 8B scale (from 120 to 30 GPU hours).

Method	Math Reasoning				Code Generation
Method	AIME 2024	AIME 2025	HMMT 2025	Avg.	LCB v5	LCB v6	Avg.
Student: Qwen3-4B-Base Teacher: Qwen3-8B
SFT	56.7	52.1	34.0	47.6	33.8	31.5	32.6
ExOPD Yang et al. (2026a)	61.0	56.0	34.4	50.5	–	29.0	–
OPD	65.4	57.9	39.9	54.4	44.2	39.3	41.8
Lightning OPD	68.1	58.4	39.8	55.4	42.8	40.3	41.5

Table 1: Pass@1 on math and code reasoning benchmarks. Lightning OPD achieves comparable performance to standard OPD across all benchmarks at both model scales, while requiring no live teacher server during training. At the 4B scale, Lightning OPD substantially outperforms ExOPD Yang et al. (2026a) , achieving 68.1% compared to 61.0% on AIME 2024 and 40.3% compared to 29.0% on LCB v6. At the 8B scale, Lightning OPD reaches 69.9% on AIME 2024 and 49.5% on LCB v5. Bold indicates the best result within each model scale.

Method	Qwen3-4B-Base	Qwen3-8B-Base
OPD	72	120
Lightning OPD	20	30
Speedup	3.6 $\times$	4.0 $\times$
Lightning OPD breakdown
$\llcorner$ Rollout collection	10	10
$\llcorner$ Teacher logprob precompute	2	4

Table 2: Training cost (GPU hours) of standard OPD vs. Lightning OPD. Lightning OPD achieves a 3.6

\times

speedup at 4B and a 4.0

\times

speedup at 8B, requiring only 20 and 30 GPU hours to train a reasoning model respectively. The lower panel breaks down Lightning OPD costs by stage, showing that the actual OPD training consumes only a moderate fraction of the total budget, highlighting how the minimal infrastructure requirement of Lightning OPD makes the training highly efficient.

Method	Math Reasoning				Code Generation
Method	AIME 2024	AIME 2025	HMMT 2025	Avg.	LCB v5	LCB v6	Avg.
Student: Qwen3-30B-A3B-Base Teacher: Qwen3-30B-A3B-Thinking-2507
SFT	66.8	63.2	44.6	58.2	39.4	33.0	36.2
OPD (OOM)	✗	✗	✗	✗	✗	✗	✗
Lightning OPD	71.0	66.3	48.3	61.9	60.8	54.4	57.6

Table 3: Lightning OPD applied to a Mixture-of-Experts architecture (Qwen3-30B-A3B). By precomputing teacher log-probabilities offline, Lightning OPD enables on-policy distillation of a 30B MoE model on a single 8

\times

H100 node, achieving state-of-the-art performance on math and code reasoning tasks. In contrast, OPD runs out of memory as it requires co-hosting both a 30B teacher and a 30B student. Bold indicates improvement over the SFT baseline.

Table 4: Teacher consistency ablation (AIME 2024 pass@1, %). Rows indicate the SFT-stage teacher and columns indicate the OPD-stage teacher. The consistent setting ( diagonal, bold ) always achieves the best performance for both standard OPD and Lightning OPD. Teacher inconsistency degrades Lightning OPD more severely than standard OPD (up to 7 points at 8B), because the fixed rollout distribution compounds the gradient bias when the SFT and OPD teachers differ.

Appendix figures & tables3 assets

Supplementary material from the paper’s appendix.

Appendix

Hyperparameter	4B Scale	8B Scale
Training steps	3000	3000
Global batch size	256	128
Max sequence length	16384	16384
Learning rate	$8\times 10^{-5}$	$8\times 10^{-5}$
LR schedule	cosine	cosine
Warmup ratio	0.1	0.1

Appendix

Table 5: SFT stage hyperparameters.

Hyperparameter	4B Scale	8B Scale
Training steps	150	150
Global batch size	256	256
Max response length	4096	4096
Learning rate	$2\times 10^{-6}$	$2\times 10^{-6}$
LR schedule	constant	constant
Weight decay	0.1	0.1

Appendix

Table 6: OPD stage hyperparameters for standard OPD and Lightning OPD. The same configuration is used for both the math and code domains.

Panel of Figure 3 — Figure 3: Training dynamics of Lightning OPD (Qwen3-4B-Base student). (a) The mean importance weight $w_{t}=\pi_{\theta}/\pi_{\text{ref}}$ drops rapidly to $\approx 0.94$ within the first 20 steps and plateaus, with its standard deviation stabilizing concurrently, validating the implicit regularization in Theorem 3.7 . (b) AIME 2024 pass@1 rises steeply in the first 50 OPD steps and saturates thereafter, confirming that 150 steps is a sufficient training budget. (c) AIME 2024 pass@1 of the SFT model, standard OPD, and Lightning OPD as a function of SFT checkpoint quality. All three improve consistently with more SFT steps, and both OPD variants provide a large, stable gain on top of the SFT baseline at every checkpoint.

Explore similar work

Jul 30, 2026

Lightning OPD 2.0: Mitigating Style Bias in Cross-Teacher On-Policy Distillation for Large Reasoning Models

On-policy distillation (OPD) provides dense token-level supervision from a teacher, but its effectiveness can depend on teacher consistency, meaning that the model providing OPD supervision should also have generated the demonstrations used to train the supervised fine-tuning (SFT) reference. However, this condition is frequently violated in practice when SFT data have mixed or unknown provenance or when different models are preferred for SFT data generation and subsequent distillation. In such cross-teacher settings, even a stronger OPD teacher can yield little improvement over the SFT reference. We find that raw teacher--reference disagreement contains potentially useful context-specific teacher evidence as well as a recurring component associated with differences in wording, formatting, and reasoning cadence. We introduce Lightning OPD 2.0 with cross-fitted style residualization, which uses rollout-level cross-fitting to estimate this recurring component as an operational proxy for style-token bias and subtracts it before constructing the token-level OPD update. Across mathematical reasoning and code generation benchmarks, Lightning OPD 2.0 consistently outperforms Lightning OPD in cross-teacher settings. Starting from Klear-Reasoner-8B-SFT, Lightning OPD 2.0 reaches 82.4% on AIME 2024 and 63.0% on LiveCodeBench v5. Together, these results establish Lightning OPD 2.0 as a practical approach to cross-teacher OPD, relaxing teacher consistency as a prerequisite and allowing the SFT data generator and distillation teacher to be selected independently. Code will be released soon.

Yecheng Wu, Song Han, Han Cai

NVIDIA

Language Model Distillation On-Policy Distillation Knowledge Distillation LLM Reasoning

Ask PDF

Jun 15, 2026

PowerOPD: Stabilizing On-Policy Distillation with Bounded Power Transformation

Standard on-policy distillation (OPD) for large language models estimates the reverse-KL objective using student-sampled tokens, yielding an unbiased single-sample Monte Carlo estimator that avoids vocabulary-wide computation. However, we show that this estimator suffers from severe training pathologies in practice: sample inefficiency, unstable generation dynamics, and a substantial performance gap compared to exact full-vocabulary OPD. Reward-level diagnosis traces these pathologies to the log-ratio reward, which is unbounded by construction, producing extremely high-variance gradients concentrated at early positions and persisting throughout training; standard post-hoc scaling fail as they operate only after this distortion occurs. To solve this problem, we propose PowerOPD: a family of natively bounded, sign-consistent rewards from the Box-Cox power transformation, parameterized by alpha > 0, of which the log-ratio is the degenerate alpha -> 0 limit. Across six mathematical reasoning benchmarks and four Qwen3 teacher-student pairs, PowerOPD achieves benchmark-averaged Avg@8/Pass@8 gains of up to +6.37/+5.71 over vanilla OPD, +3.01/+3.54 over post-hoc stabilization, and +2.59/+8.90 over full-vocabulary OPD, while reducing wall-clock time by 59.2% and peak GPU memory by 23.1%. Larger alpha generally improves accuracy, consistently shortens responses, and keeps gradient norms more than 3,000x smaller than vanilla OPD.

Anhao Zhao, Junlong Tong, Yingqi Fan +3

Eastern Institute of Technology, Ningbo · Shanghai Jiao Tong University · University of Waterloo +1

On-Policy Distillation Teacher-Student Learning

Ask PDF

May 12, 2026

Learning to Foresee: Unveiling the Unlocking Efficiency of On-Policy Distillation

On-policy distillation (OPD) has emerged as an efficient post-training paradigm for large language models. However, existing studies largely attribute this advantage to denser and more stable supervision, while the parameter-level mechanisms underlying OPD's efficiency remain poorly understood. In this work, we argue that OPD's efficiency stems from a form of ``foresight'': it establishes a stable update trajectory toward the final model early in training. This foresight manifests in two aspects. First, at the Module-Allocation Level, OPD identifies regions with low marginal utility and concentrates updates on modules that are more critical to reasoning. Second, at the Update-Direction Level, OPD exhibits stronger low-rank concentration, with its dominant subspaces aligning closely with the final update subspace early in training. Building on these findings, we propose EffOPD, a plug-and-play acceleration method that speeds up OPD by adaptively selecting an extrapolation step size and moving along the current update direction. EffOPD requires no additional trainable modules or complex hyperparameter tuning, and achieves an average training acceleration of

3\times

while maintaining comparable final performance. Overall, our findings provide a parameter-dynamics perspective for understanding the efficiency of OPD and offer practical insights for designing more efficient post-training methods for large language models.

Yuchen Cai, Ding Cao, Liang Lin +9

1USTC · 2Tencent · 4HKUST(GZ) +3

Efficient Language Model Training On-Policy Distillation Language Model Post-Training

Ask PDF