[RFC] Tail-aware top-k KL for on-policy distillation

Feature request

A new on-policy distillation loss mode, forward_kl_topk_tail, that makes the truncated teacher-top-k forward KL a valid coarse-grained KL by aggregating the out-of-top-k mass into a single tail bucket.

The objective keeps the existing teacher top-k payload ([S, K] log-probs, i.e. vLLM prompt_logprobs=k), so teacher compute, schema and network traffic are unchanged. Implementation PR: #7855 (FSDP, Megatron TP, VeOmni paths + tests + benchmark + docs).

Feedback is welcome on the objective itself and on the two config constraints noted below before this lands.

Motivation

verl's top-k OPD loss retains only the teacher's top-k terms:

L_truncated = Σ_{i∈T} p_i (log p_i − log q_i)

T is truncated, so P = Σ_{i∈T} p_i and Q = Σ_{i∈T} q_i are both below 1. Two consequences:

1. The quantity is not a KL. It is negative whenever Q > P, so the current implementation applies clamp_min(0) and silently drops the gradient information carried by those tokens. In our runs the per-token minimum was exactly 0.0 (clamping active) in 415 / 467 training steps (89 %).
2. Tail-mass mismatch is unpenalized. The loss is blind to how the student's mass outside the teacher's top-k compares with the teacher's own tail mass, even though that is exactly the part of the distribution the truncation threw away.

The kl/k1/k3 single-sample reverse-KL estimators are unbiased but mode-seeking, so they do not cover this either.

Proposed objective

Treat every token outside T as one aggregate category, giving two K+1-category distributions [p_1..p_K, p_tail] and [q_1..q_K, q_tail]:

P = Σ_{i∈T} p_i,  Q = Σ_{i∈T} q_i,  p_tail = 1 − P,  q_tail = 1 − Q

L_tail = Σ_{i∈T} p_i log(p_i / q_i) + p_tail log(p_tail / q_tail)

This is the exact KL between the two coarse-grained distributions, so it is:

• non-negative by construction — no clamp_min(0) and no discarded gradient;
• a lower bound on the full forward KL (coarse-graining only removes information);
• free in teacher-side cost, because it reuses the existing top-k payload.

The analytic gradient with respect to a student logit z_j is

j ∈ teacher top-k:  dL/dz_j = q_j − p_j
j ∉ teacher top-k:  dL/dz_j = q_j (P − Q) / (1 − Q)

The second form is what makes the vocab-parallel (Megatron TP) path tractable: each rank applies it to its local shard without gathering logits. It also matches the derivative of the q_tail.clamp_min(tail_mass_eps) numerical guard when that guard is active.

Why not just use the full-vocabulary KL (#7239 / #7375)?

They solve a different constraint and the two are complementary:

• Full-vocabulary KL is exact but requires materializing or gathering the full [B, L, V] logits/distributions, which is why #7239 and #7375 pursue it (and why #7733 / #7809 optimize the memory of the existing top-k path).
• forward_kl_topk_tail is for the regime where the teacher payload is capped at K by construction — the standard vLLM prompt_logprobs=k OPD setup where the full distribution is never transmitted. In that regime the truncated objective is not a KL at all, and this change repairs it at zero extra teacher cost rather than requiring a new payload.
• If a user later moves to full-vocabulary OPD, forward_kl_topk_tail simply becomes unnecessary — it is not an alternative route to the same endpoint, it is the correct coarse-grained objective for the K-capped payload.

Validation so far (details in #7855)

• CPU tests: non-negativity, equality against an independent coarse-distribution oracle, autograd gradient equality, analytic sharded-gradient equality with cross-shard teacher IDs, the lower-bound property across K, equality with the full KL when K spans the vocabulary, regular + chunked FSDP paths, invalid clamp/epsilon configuration, fused-mass fail-closed behaviour — 41 passed.
• Megatron TP parity (2 GPUs): loss, student_mass, teacher_mass and tail_loss match the FSDP reference; every vocabulary-shard gradient matches within atol=rtol=1e-4; existing truncated cases do not regress.
• Synthetic full-KL benchmark: gradient cosine similarity to the true full-KL gradient improves at every K (0.303→0.455 at K=8; 0.488→0.793 at K=64). Scalar error versus the full KL is not uniformly better — expected, since this is a lower bound, and we do not claim otherwise.
• Real-model training efficiency (Qwen3-1.7B student ← Qwen3-32B teacher, vLLM TP=2, 3×A6000, GSM8K, full 1319-example test set): at equal step budgets the tail-aware arm leads in every comparison. At 175 steps it already matches what the baseline reaches at 350 steps (0.3472 ≥ 0.3457), i.e. ≤50 % of the steps for the same accuracy, with unchanged throughput and peak memory (142.6 vs 143.0 tok/s, 40.58 vs 40.61 GB).

The claim is training efficiency (fewer steps to a given accuracy), not a higher final ceiling: the two curves share an asymptote and the gap narrows with longer training. Current evidence is single-seed per arm; the measured run-to-run noise for the same arm at the same step is ≈3.3 pp, so multi-seed error bars are still running.

Two config constraints worth review

1. forward_kl_topk_tail requires log_prob_min_clamp=null. Per-entry log-prob clamping changes the selected top-k mass and therefore invalidates the coarse-grained KL definition.
2. A new tail_mass_eps (default 1e-6, strictly in (0,1)) guards q_tail against underflow when the student's top-k mass rounds to one in finite precision.
3. On the VeOmni backend, use_fused_kernels=True is currently unsupported for this mode: the fused kernel's top-k mass outputs are degenerate on our host (teacher_mass ≈ 2.11e-4, effectively constant per token). The same values appear with the pre-existing forward_kl_topk mode, so this is a property of the fused path rather than of this change; we document the constraint and leave the fused-kernel issue as a follow-up.

Your contribution

Yes — PR #7855 implements the loss, the config and validation rules, the FSDP/Megatron/VeOmni paths, the metrics, the CPU and GPU tests, a synthetic benchmark and the docs. Happy to adjust the design based on RFC feedback, and to add multi-seed results to the PR as they complete.