InternLM 开源 ResOPD：基于 verl 的稀疏在线策略蒸馏尾部残差方法

InternLM; ResOPD

Repository navigation

NotificationsYou must be signed in to change notification settings; Fork0 (0); Star65 (65)You must be signed in to star a repository

ResOPD: Tail Residualization for Sparse On-Policy Distillation. Core implementation built on verl.

Repository files navigation

README; Contributing; Apache-2.0 license

Integrate the coarse distribution exactly. Sample only the fine-tail residual.

3 University of Science and Technology of China 4 Fudan University

ResOPD estimates the full-vocabulary reverse-KL gradient from the teacher's Top-k probabilities and the sampled-token score. It computes the observable coarse gradient exactly and samples only the unresolved within-tail residual, using the same sparse teacher payload and no additional teacher forward passes.

This repository provides the core implementation and its integration with verl. The paper contains the derivations and experiments. UPSTREAM.md records the pinned framework base and release scope.

Paper highlights

These are results from the paper's experimental protocols. The launcher below uses the released DeepMath training and held-out sets, with a shorter response budget for the integration example.

Quick start

1. Install

Use Python 3.12 and a GPU environment compatible with the pinned verl base.

The separate recipe submodule is optional for this example. See upstream installation guidance for alternative environments.

2. Prepare data

Download the paper's processed DeepMath subset from ygyjrc/ResOPD-deepmath. The files are already in verl format and can be used directly:

The launcher defaults to these files under data/resopd-deepmath. To use another local directory, set DATA_DIR, or override TRAIN_FILE and VAL_FILE individually. The held-out split comes from the same screened DeepMath subset; the paper also reports separate mathematical reasoning benchmarks. Evaluation uses verl's boxed-answer math scorer; task rewards remain disabled for distillation.

3. Inspect and launch

Adjust model IDs, GPU allocation, tensor parallelism, and sequence lengths for the available hardware. The student and teacher must share the token-ID vocabulary. Extra Hydra overrides can be passed directly to the launcher.

Keep use_chunked_topk=false and set ppo_mini_batch_size equal to the training batch size. Set loss_mode: resopd to select the estimator.

Method

ResOPD separates an exact coarse update from a sampled fine-tail correction:

Integrate the head. Evaluate the teacher Top-k contributions analytically; Account for the tail bucket. Use the residual student and teacher masses to compute the observable aggregate tail gradient; Sample the residual. On tail hits, add only the unresolved fine-tail correction

At a fixed prefix, let $p$ and $q$ be the student and teacher distributions, $S$ the teacher Top-k support, and $A\sim p$ the actual student action. Define $r_i=\log(p_i/q_i)$, $s_i=\nabla_\theta\log p_i$, $\tau_p=1-\sum_{i\in S}p_i$, and $\tau_q=1-\sum_{i\in S}q_i$.

The code uses an algebraically equivalent centered support-event control. Score coefficients and the event-centering mass are detached: the forward value is used for logging, while backward supplies the estimator above.

Sampling matters. Unbiasedness assumes draws from the current full student distribution. The launcher uses full-support sampling and a single update per rollout batch. Truncated sampling, stale rollouts, repeated optimization passes, and sampling-altering logits processors require separate treatment, as discussed in the paper. No uniform variance ordering is claimed for every distribution.

Implementation map

Supported path: vLLM teacher scoring and eager FSDP/FSDP2/VeOmni student logits; the launcher uses FSDP2. Megatron students, SGLang teachers, fused student loss kernels, and chunked Top-k are not enabled for ResOPD in this release. The student still performs full-vocabulary normalization; sparse feedback reduces teacher communication.

For the underlying framework documentation, see README.verl.md.

Verification

Standalone estimator and teacher-payload checks require only PyTorch and pytest:

With verl's runtime dependencies installed, run the CPU integration checks:

The checks cover exact finite-action gradients, the reverse-KL expectation, variance identities, numerical precision, teacher-score alignment, and the verl data path. Release validation records 22 standalone checks and 13 integration checks, plus comparison with the original research kernel. The GitHub workflow runs the standalone checks. A full GPU training run of this extracted port is outside the recorded release validation.

Citation

If you use ResOPD, please cite the paper:

GitHub's Cite this repository entry also points to the paper through CITATION.cff.

License and acknowledgments

ResOPD: Tail Residualization for Sparse On-Policy Distillation. Core implementation built on verl.