[algo, trainer, vllm] feat: add ResOPD sparse on-policy distillation

docs/algo/opd.md

@@ -2,7 +2,7 @@
 
 **Author:** [Jacob Helwig](https://jacobhelwig.github.io/)
 
-Last updated: 10/03/2026.
+Last updated: 10/20/2026.
 
 ## Background
 
@@ -55,7 +55,7 @@ In practice, the sampled rollout is treated as fixed during the student update.
 
 ### Loss Variants
 
-We implement two OPD variants.
+We implement three OPD variants.
 
 #### GKD OPD
 
@@ -127,6 +127,47 @@ $$
 
 The stop-gradient is required because the reward is used inside a policy-gradient objective. Without it, differentiating through the estimator with respect to student parameters would zero the teacher's logprob, leading to a loss that is independent of the teacher signal. 
 
+#### ResOPD
+
+ResOPD [9] targets the full-vocabulary **reverse** KL using the teacher top-$k$
+probabilities plus the teacher score of the actual sampled token. Writing the
+reverse KL over the coarse-graining induced by the teacher top-$k$ support
+$S$ and one aggregate tail bucket, the head and bucket terms are available in
+closed form and only the *within-tail* residual is sampled.
+
+Let $r_v=\log\bigl(\pi_\theta(v \mid s_t)/\nu(v \mid s_t)\bigr)$,
+$s_v=\nabla_\theta\log\pi_\theta(v \mid s_t)$,
+$\tau_\pi=1-\sum_{v\in S}\pi_\theta(v \mid s_t)$ and
+$\tau_\nu=1-\sum_{v\in S}\nu(v \mid s_t)$. For
+$y_t\sim\pi_\theta(\cdot \mid s_t)$ the gradient estimator is
+
+$$
+\widehat{g}(s_t)
+=
+\sum_{v\in S}\pi_\theta(v \mid s_t)\, r_v s_v
++
+\tau_\pi \mu_B
++
+\mathbf{1}[y_t\notin S]\bigl(r_{y_t} s_{y_t}-\mu_B\bigr),
+\qquad
+\mu_B=\log\frac{\tau_\pi}{\tau_\nu}\nabla_\theta\log\tau_\pi .
+$$
+
+When the rollout is sampled from the current student at temperature 1, this is
+an unbiased estimator of the full-vocabulary reverse-KL gradient; the third term
+is the residual that `forward_kl_topk_tail` drops in the forward direction.
+
+Relation to the other variants: `forward_kl_topk_tail` optimizes the coarse
+**forward** KL and discards the within-tail term, yielding a deterministic lower
+bound, whereas ResOPD estimates that term and so needs one extra teacher score,
+namely the log-probability of the sampled token. PG OPD uses a single sampled
+token only and carries no distributional head.
+
+`loss_mode=resopd` is implemented as a stopped-coefficient surrogate whose
+backward pass is $\widehat{g}$. Its logged forward value is a Monte Carlo
+estimate and can be negative, so it must not be clamped. FP16/BF16 logits are
+promoted to FP32 before the log-softmax.
+
 ### Multi-Teacher OPD
 
 Multi-teacher OPD (MOPD) extends OPD to multiple domain-specialized teachers [4,5,6,7]. This is useful when distilling knowledge to a student across multiple domains. In each domain, such as math, coding, or instruction following, different teachers specialized to the domain can be used.
@@ -153,6 +194,8 @@ MOPD consolidates multiple specialized policies into a single student model whil
 
 [8] Li, Yaxuan, et al. "Rethinking On-Policy Distillation of Large Language Models: Phenomenology, Mechanism, and Recipe." arXiv preprint arXiv:2604.13016, 2026.
 
+[9] Yang, Penghui, et al. "ResOPD: Tail Residualization for Sparse On-Policy Distillation." arXiv preprint arXiv:2610.04882, 2026.
+
 ## Configuration Parameters
 
 OPD parameters live under three namespaces:
@@ -281,6 +324,8 @@ Two registered families:
     teacher's top-k support.
   - `forward_kl_topk_tail`: a coarse-grained forward KL that additionally
     collapses every non-top-k token into one aggregate tail bucket.
+  - `resopd`: a reverse-KL estimator that keeps the coarse head and tail terms
+    in closed form and samples the within-tail residual.
 - **Single-sample KL estimators** (`kl`, `k1`, `abs`, `mse`, `k2`,
   `low_var_kl`, `k3`): per-token Monte Carlo estimators of reverse KL
   computed from the student's `log_probs` and the teacher's single
@@ -290,8 +335,8 @@ Two registered families:
 
 `k` for top-k distillation losses. Default: `32`.
 
-Only used when `loss_mode` requires top-$k$ (e.g. `forward_kl_topk` or
-`forward_kl_topk_tail`). Drives both
+Only used when `loss_mode` requires top-$k$ (e.g. `forward_kl_topk`,
+`forward_kl_topk_tail` or `resopd`). Drives both
 the teacher's `prompt_logprobs` request size and (for vLLM) the engine's
 `max_logprobs` cap.
 
@@ -517,6 +562,53 @@ distillation:
 
 Currently, only `policy_loss_mode=vanilla` is supported. Other policy-loss modes, such as `dppo_tv`, require additional parameters and are not implemented for OPD.
 
+### ResOPD
+
+ResOPD requires the teacher log-probability of the sampled token, which vLLM
+already returns, so the teacher is queried no more than for `forward_kl_topk`.
+The observed token is appended as column $K$ of the existing `teacher_ids` /
+`teacher_logprobs` tensors; the first $K$ columns remain the top-$k$ support,
+even when the sampled token is among them.
+
+```yaml
+distillation:
+   distillation_loss:
+      loss_mode: resopd
+      topk: 4
+      use_policy_gradient: false
+      use_task_rewards: false
+      loss_max_clamp: null
+      log_prob_min_clamp: null
+```
+
+Unbiasedness assumes the rollout is drawn from the current student
+distribution, so the sampler must be left untruncated and the update kept
+on-policy:
+
+```yaml
+actor_rollout_ref:
+   model:
+      use_fused_kernels: false
+   actor:
+      use_torch_compile: false
+      ppo_epochs: 1
+      ppo_mini_batch_size: ${data.train_batch_size}
+   rollout:
+      temperature: 1.0
+      top_p: 1.0
+      top_k: -1
+      n: 1
+```
+
+Use a vLLM teacher sharing the student's token-ID vocabulary, and keep
+`use_chunked_topk=false`. Eager FSDP/FSDP2/VeOmni student logits are supported;
+Megatron and the fused loss kernels are not.
+
+Unlike `forward_kl_topk_tail`, ResOPD does not use `tail_mass_eps`: it clamps
+the tail masses at machine epsilon instead. In `forward_kl_topk_tail` that
+epsilon is part of the objective and carries gradient, while in ResOPD it only
+guards `log(0)` inside a term whose coefficient is detached.
+
 ### Task rewards
 
 OPD can be optimized alone or combined with the standard PPO/GRPO task-reward loss.
@@ -669,6 +761,18 @@ These metrics are logged for top-$k$ loss modes such as `forward_kl_topk`.
   Mean probability mass outside the teacher-selected top-$k$ support. Emitted
   by `forward_kl_topk_tail`.
 
+- `actor/distillation/resopd_head_hit`
+  Fraction of tokens whose sampled token lies inside the teacher top-$k$
+  support. Emitted by `resopd`.
+
+- `actor/distillation/resopd_centered_event`
+  Mean of the centered membership indicator used as the control variate.
+  Emitted by `resopd`.
+
+- `actor/distillation/resopd_tail_log_ratio`
+  Mean $\log(\tau_\pi/\tau_\nu)$ between the student and teacher tail masses.
+  Emitted by `resopd`.
+
 - `actor/distillation/overlap_ratio`
   Average fraction of teacher top-$k$ tokens that also appear in the student's
   top-$k$ tokens, computed as

tests/special_sanity/check_license.py

@@ -22,6 +22,7 @@
 # Add custom license headers below
 license_head_prime = "Copyright 2024 PRIME team and/or its affiliates"
 license_head_individual = "Copyright 2025 Individual Contributor:"
+license_head_individual_26 = "Copyright 2026 Individual Contributor:"
 license_head_sglang = "Copyright 2023-2024 SGLang Team"
 license_head_modelbest = "Copyright 2025 ModelBest Inc. and/or its affiliates"
 license_head_amazon = "Copyright 2025 Amazon.com Inc and/or its affiliates"
@@ -40,6 +41,7 @@
     license_head_bytedance_26,
     license_head_prime,
     license_head_individual,
+    license_head_individual_26,
     license_head_sglang,
     license_head_modelbest,
     license_head_amazon,

tests/trainer/test_resopd_integration_on_cpu.py

@@ -0,0 +1,160 @@
+# Copyright 2026 Individual Contributor: Penghui Yang and the ResOPD authors
+# SPDX-License-Identifier: Apache-2.0
+"""CPU checks for the sparse teacher transport and ResOPD loss dispatch."""
+
+import asyncio
+from types import SimpleNamespace
+from unittest.mock import AsyncMock, patch
+
+import pytest
+import torch
+from tensordict import TensorDict
+
+from verl.experimental.teacher_loop.teacher_manager import AsyncTeacherLLMServerManager, _get_teacher_sampling_params
+from verl.trainer.distillation.losses import collect_resopd, compute_topk_loss
+from verl.trainer.distillation.resopd import resopd_loss
+from verl.trainer.ppo.padding_utils import SYNTHETIC_PADDING_SEQ_LEN, construct_minimal_padding_template
+from verl.workers.config import DistillationLossConfig
+from verl.workers.utils.padding import left_right_2_no_padding
+
+
+def loss_config(**overrides):
+    args = dict(loss_mode="resopd", use_policy_gradient=False, loss_max_clamp=None, log_prob_min_clamp=None, topk=2)
+    return DistillationLossConfig(**(args | overrides))
+
+
+@pytest.mark.parametrize(
+    "bad",
+    [
+        dict(use_policy_gradient=True),
+        dict(loss_max_clamp=10),
+        dict(log_prob_min_clamp=-10),
+        dict(use_chunked_topk=True),
+        dict(topk=0),
+        dict(topk=None),
+    ],
+)
+def test_invalid_configuration(bad):
+    with pytest.raises(ValueError, match="ResOPD"):
+        loss_config(**bad)
+
+
+@pytest.mark.parametrize("mode", ["resopd", "k1", "forward_kl_topk"])
+def test_teacher_request_keeps_two_tensor_contract(mode):
+    manager = object.__new__(AsyncTeacherLLMServerManager)
+    manager.distillation_loss_config = loss_config(loss_mode=mode, use_policy_gradient=mode == "k1")
+    manager.teacher_model_configs = {
+        "teacher": SimpleNamespace(inference=SimpleNamespace(name="vllm", temperature=1.0))
+    }
+    width = {"resopd": 3, "forward_kl_topk": 2, "k1": 1}[mode]
+    payload = dict(prompt_ids=[[0] * width] * 3, prompt_logprobs=[[-1.0] * width] * 3)
+    client = SimpleNamespace(generate=AsyncMock(return_value=SimpleNamespace(extra_fields=payload)))
+    manager.teacher_client = {"teacher": client}
+    ids, logprobs = asyncio.run(manager.compute_teacher_logprobs_single([8, 2, 0]))
+    assert ids.shape == logprobs.shape == (3, width)
+    client.generate.assert_awaited_once()
+    params = client.generate.call_args.kwargs["sampling_params"]
+    assert params["prompt_logprobs"] == (0 if mode == "k1" else 2)
+    assert params.get("include_prompt_token", False) == (mode == "resopd")
+
+
+def test_unsupported_teacher_and_student_fail():
+    teacher = SimpleNamespace(inference=SimpleNamespace(name="sglang", temperature=1.0))
+    cfg = SimpleNamespace(distillation_loss=loss_config())
+    with pytest.raises(NotImplementedError, match="vLLM"):
+        _get_teacher_sampling_params(teacher, cfg.distillation_loss)
+    with pytest.raises(NotImplementedError, match="eager"):
+        compute_topk_loss(SimpleNamespace(strategy="megatron"), cfg, {}, None, "thd")
+    with pytest.raises(ValueError, match="disable fused kernels"):
+        collect_resopd(None, cfg, {}, {})
+
+
+@pytest.mark.parametrize("strategy", ["fsdp", "fsdp2", "veomni"])
+def test_packed_teacher_scores_and_response_mask(strategy):
+    # Column 2 is the observed token, including a head hit (duplicated ID).
+    ids = torch.tensor([[[0, 0, 0], [0, 1, 2], [0, 1, 1], [0, 0, 0]], [[0, 1, 2], [0, 1, 0], [0, 0, 0], [0, 0, 0]]])
+    scores = torch.tensor(
+        [
+            [[-1.0, -1.0, -1.0], [-0.7, -1.2, -2.0], [-0.6, -1.3, -1.3], [0.0, 0.0, 0.0]],
+            [[-0.8, -1.1, -2.1], [-0.9, -1.0, -0.9], [0.0, 0.0, 0.0], [0.0, 0.0, 0.0]],
+        ]
+    )
+    data = TensorDict(
+        dict(
+            input_ids=torch.tensor([[0, 8, 2, 1], [7, 2, 0, 0]]),
+            position_ids=torch.tensor([[0, 0, 1, 2], [0, 1, 2, 0]]),
+            attention_mask=torch.tensor([[0, 1, 1, 1], [1, 1, 1, 0]]),
+            prompts=torch.tensor([[0, 8], [7, 2]]),
+            responses=torch.tensor([[2, 1], [0, 0]]),
+            response_mask=torch.tensor([[1, 0], [1, 0]]),
+            teacher_ids=ids,
+            teacher_logprobs=scores,
+        ),
+        batch_size=2,
+    )
+    data = left_right_2_no_padding(data)
+    assert data["teacher_ids"].values()[:, -1].tolist() == [2, 1, 0, 2, 0, 0]
+    z = torch.tensor([[[0.2, 0.4, -0.3]] * 6], requires_grad=True)
+    cfg = SimpleNamespace(distillation_loss=loss_config())
+    output = compute_topk_loss(SimpleNamespace(strategy=strategy), cfg, data, z, "thd")
+    losses, metrics = collect_resopd(None, cfg, {key: value.squeeze(0) for key, value in output.items()}, data)
+    expected = torch.stack(
+        [output["distillation_losses"][0, :2], torch.stack([output["distillation_losses"][0, 4], z.new_zeros(())])]
+    )
+    torch.testing.assert_close(losses, expected)
+    assert metrics["distillation/resopd_head_hit"].aggregate() == 0.5
+    (losses * data["response_mask"]).sum().backward()
+    assert z.grad[0, [1, 2, 3, 5]].eq(0).all()
+    assert torch.isfinite(z.grad).all()
+
+
+@pytest.mark.parametrize("rank", [0, 1])
+def test_sequence_parallel_slicing_and_padding(rank):
+    ids = torch.tensor([[0, 1, 2], [0, 1, 0], [0, 1, 1]])
+    scores = torch.tensor([[-0.7, -1.2, -2.0], [-0.7, -1.2, -0.7], [-0.7, -1.2, -1.2]])
+    data = {
+        key: torch.nested.as_nested_tensor([value], layout=torch.jagged)
+        for key, value in (("teacher_ids", ids), ("teacher_logprobs", scores))
+    }
+    logits = torch.zeros(1, 2, 3, requires_grad=True)
+    with (
+        patch("verl.trainer.distillation.fsdp.losses.get_ulysses_sequence_parallel_world_size", return_value=2),
+        patch("verl.utils.ulysses.dist.get_world_size", return_value=2),
+        patch("verl.utils.ulysses.get_ulysses_sequence_parallel_rank", return_value=rank),
+    ):
+        output = compute_topk_loss(
+            SimpleNamespace(strategy="fsdp"), SimpleNamespace(distillation_loss=loss_config()), data, logits, "thd"
+        )
+    ids = torch.cat([ids, torch.zeros_like(ids[:1])])[2 * rank : 2 * rank + 2].unsqueeze(0)
+    scores = torch.cat([scores, torch.zeros_like(scores[:1])])[2 * rank : 2 * rank + 2].unsqueeze(0)
+    expected = resopd_loss(logits, scores[..., :-1], ids[..., :-1], scores[..., -1:], ids[..., -1:])
+    torch.testing.assert_close(output["distillation_losses"], expected["distillation_losses"])
+
+
+def test_synthetic_padding_preserves_payload_width():
+    source = {"teacher_ids": torch.ones(7, 3, dtype=torch.long), "teacher_logprobs": torch.ones(7, 3)}
+    output, tag = construct_minimal_padding_template(source, {}, eos_token_id=0)
+    assert tag["is_padding"] and output["response_mask"].sum() == 0
+    for key in source:
+        assert output[key].shape == (SYNTHETIC_PADDING_SEQ_LEN, 3)
+
+
+def test_negative_surrogate_is_not_clipped():
+    names = (
+        "distillation_losses",
+        "student_mass",
+        "teacher_mass",
+        "resopd_head_hit",
+        "resopd_centered_event",
+        "resopd_tail_log_ratio",
+    )
+    output = {key: torch.tensor([-2.0, 999.0], requires_grad=key == "distillation_losses") for key in names}
+    data = {
+        "prompts": torch.nested.as_nested_tensor([torch.tensor([7])], layout=torch.jagged),
+        "responses": torch.nested.as_nested_tensor([torch.tensor([1])], layout=torch.jagged),
+        "response_mask": torch.tensor([[1]]),
+    }
+    losses, _ = collect_resopd(None, None, output, data)
+    assert losses.item() == -2.0
+    losses.sum().backward()
+    assert output["distillation_losses"].grad.tolist() == [1.0, 0.0]