跳到正文
原文
上海人工智能实验室 InternLM:原创项目· InternLM·· 2 小时前AI 评分40

InternLM 推出 ResOPD:稀疏在线策略蒸馏的尾部残差化方法

InternLM/ResOPD

AI 导读

上海人工智能实验室 InternLM 项目发布 ResOPD,一种仅用教师模型 Top-k 概率和学生采样 token 概率即可估计全词表 reverse-KL 梯度的稀疏在线策略蒸馏方法,教师传输量和前向次数与稀疏 OPD 持平。

正文

Tail Residualization for Sparse On-Policy Distillation

Penghui Yang, Long Xing, Xuanlang Dai, Ziyu Liu, Kai Chen, and Yuhang Zang

ResOPD estimates the full-vocabulary reverse-KL gradient using only the teacher's Top-k probabilities and the probability of the token sampled by the student. It integrates the observable coarse distribution exactly and samples the remaining within-tail residual. The teacher payload and number of teacher forward passes stay the same as sparse on-policy distillation.

This is the core implementation on top of verl. See UPSTREAM.md for the pinned public base and the scope of the changes, and README.verl.md for the original framework documentation.

Core implementation

File Purpose
resopd.py Standalone PyTorch loss and stopped-coefficient gradient
fsdp/losses.py Packed teacher tensors and sequence-parallel alignment
losses.py Loss registration, dispatch and response-masked metrics
prompt_logprobs.py Preserve both Top-k and observed-token teacher scores
run_resopd.sh Synchronous FSDP2/vLLM training example
tests/resopd Finite-action gradient, teacher payload and verl integration checks

The public loss name is resopd. The experimental name event_cv_reverse_kl_topk is an alias for the same implementation.

For a fixed prefix, let $p$ be the student distribution, $q$ the teacher, $S$ the teacher Top-k support, and $A$ the actual token sampled from $p$. 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 implemented gradient is

$$ \widehat g=\underbrace{\sum_{i\in S}p_i r_i s_i+\tau_p\mu_B}_{\text{exact coarse gradient}} +\mathbf{1}[A\notin S](r_A s_A-\mu_B), \qquad \mu_B=\log(\tau_p/\tau_q)\nabla_\theta\log\tau_p. $$

The code uses an algebraically equivalent centered support-event control. All score coefficients and the event-centering mass are detached. The forward loss is an estimator used for logging; backward supplies the gradient above. Directly differentiating a truncated or renormalized Top-k KL gives a different method.

The unbiasedness statement assumes sampling from the current full student probabilities. The example therefore uses temperature 1, top_p=1, top_k=-1, one PPO epoch and one mini-batch per rollout batch. Truncated sampling, stale rollouts, multiple optimization passes, or logits processors that change the sampling law require separate treatment. No uniform variance improvement is claimed for every possible teacher/student pair.

Installation

Use Python 3.12 and a GPU environment compatible with the pinned verl base. The upstream dependency setup is retained:

git clone https://github.com/InternLM/ResOPD.git
cd ResOPD
uv sync --extra fsdp --extra vllm
source .venv/bin/activate

The separate recipe submodule is not needed for this example. See upstream installation guidance for alternative environments. The example uses vLLM for teacher scoring and eager FSDP2 for the student. The ResOPD adapter also accepts the FSDP/VeOmni eager logits path; Megatron, SGLang teachers, fused student loss kernels and chunked Top-k are not enabled for ResOPD in this release.

Run training

Prepare data in verl's parquet format. For a small public example:

python examples/data_preprocess/gsm8k.py --local_save_dir ./data/gsm8k
export TRAIN_FILE="$PWD/data/gsm8k/train.parquet"
export VAL_FILE="$PWD/data/gsm8k/test.parquet"
bash examples/resopd/run_resopd.sh

Defaults are Qwen3.5-4B as student, Qwen3.5-27B as teacher, teacher Top-k=4, 4 student GPUs plus 4 separate teacher GPUs on one node, and response length 2048. This is an integration example; it does not reproduce the paper's dataset filtering and long-context evaluation protocol. Adjust model IDs, GPU counts, teacher tensor parallelism and sequence lengths for the available memory. The student and teacher must share the tokenizer/token-ID vocabulary.

For example, use smaller models on two GPUs:

STUDENT_MODEL=Qwen/Qwen3-0.6B TEACHER_MODEL=Qwen/Qwen3-4B \
STUDENT_GPUS=1 TEACHER_GPUS=1 TEACHER_TP=1 \
TRAIN_BATCH_SIZE=2 MAX_RESPONSE_LENGTH=256 \
bash examples/resopd/run_resopd.sh trainer.total_training_steps=1

Inspect the command without launching training with DRY_RUN=1. Additional Hydra overrides can be passed as arguments to the script.

The essential settings for an existing verl OPD configuration are:

distillation:
  enabled: true
  distillation_loss:
    loss_mode: resopd
    topk: 4
    use_task_rewards: false
    use_policy_gradient: false
    loss_max_clamp: null
    log_prob_min_clamp: null
actor_rollout_ref:
  model:
    use_fused_kernels: false
  actor:
    use_torch_compile: false
    ppo_epochs: 1
  rollout:
    temperature: 1.0
    top_p: 1.0
    top_k: -1

Keep use_chunked_topk=false (the default) and set ppo_mini_batch_size equal to the training batch size for one update per fresh rollout batch. This uses full student logits; sparsity reduces teacher communication, not the student's vocabulary projection.

Verification

The standalone tests need only PyTorch and pytest, and load the production kernel/parser directly without importing the distributed training stack:

python -m pytest tests/resopd/test_estimator.py tests/resopd/test_teacher_payload.py -q

They enumerate every next-token action, compare each gradient with the paper's coarse-plus-residual expression, check the probability-weighted gradient against autodiff of full reverse KL, and check the variance identity. They also cover teacher-score detachment, FP16/BF16 inputs, and preservation/causal shifting of sampled teacher scores inside and outside Top-k.

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

python -m pytest tests/resopd/test_verl_integration.py -q

These cover loss registration, configuration guards, nested tensors, teacher requests, agent output fields and padding. The GitHub workflow runs the standalone checks. This extracted port has not been rerun through a full GPU training experiment.

Citation

@misc{yang2026resopd,
  title  = {ResOPD: Tail Residualization for Sparse On-Policy Distillation},
  author = {Penghui Yang and Long Xing and Xuanlang Dai and Ziyu Liu and Kai Chen and Yuhang Zang},
  year   = {2026},
  url    = {https://github.com/InternLM/ResOPD}
}

License

Apache-2.0. This project builds on verl and retains its original copyright and license notices.

来源:上海人工智能实验室 InternLM:原创项目 · github.com