verl Integration: Top-K Forward-KL Distillation via the VeOmni Engine#

VeOmni’s chunk_topk_distill_function computes verl’s top-k forward-KL distillation loss without materializing the [T, V] student logits tensor — the kernel rides the same chunked fused-linear pattern as chunk_logprobs.py, so use_fused_kernels=True + remove-padding + Ulysses SP all carry over.

VeOmni-side contract#

Add two kwargs to the model forward; the wrapper routes them through to the kernel:

teacher_topk_ids: torch.Tensor          # [B, L, K] int64 (dense)
teacher_topk_log_probs: torch.Tensor    # [B, L, K] fp32
log_prob_min_clamp: float | None = None # optional, matches DistillationLossConfig

When return_log_probs=True, the output dataclass exposes a single fused_linear_aux: Optional[FusedLinearAuxOutput] field carrying the per-token tensors:

`output.fused_linear_aux.*`	Shape	dtype	Sign / range	Grad
`log_probs`	`[B, L]`	fp32	`<= 0`	flows
`entropy`	`[B, L]`	fp32	`>= 0`	flows
`distillation_losses`	`[B, L]`	fp32	`>= 0`	flows
`student_mass`	`[B, L]`	fp32	`[0, 1]`	detached
`teacher_mass`	`[B, L]`	fp32	`[0, 1]`	detached

fused_linear_aux is None on the plain loss path (no return_log_probs). All five tensors are 0 at IGNORE_INDEX positions and the trailing pad slot. KL formula is verbatim verl’s: Σ_k exp(log_p_t,k) · (log_p_t,k - log_q_s,k) on the top-k support (with optional clamp on both terms). The nested-payload shape (rather than flat output.log_probs / output.entropy / output.distillation_losses as top-level fields) means future per-token metrics extend FusedLinearAuxOutput only, not every *WithLogProbs subclass.

verl-side changes (two patches, ~30 lines total)#

Both are pure additive — symmetric with how the existing fused-kernel reads work, no impact on the non-fused path or any non-distillation flow. Note the read path moves one level down: output.log_probs → output.fused_linear_aux.log_probs etc. (the fused-linear payload nests under one field to avoid growing every *WithLogProbs subclass when new per-token tensors land later).

1. Thread teacher tensors into `model_inputs`#

verl/workers/engine/veomni/transformer_impl.py — inside the existing use_fused_kernels and use_remove_padding block of VeOmniEngineWithLMHead.prepare_model_inputs:

distillation_use_topk = tu.get_non_tensor_data(
    data=micro_batch, key="distillation_use_topk", default=False
)
if distillation_use_topk:
    teacher_logprobs = micro_batch["teacher_logprobs"]
    teacher_ids = micro_batch["teacher_ids"]
    if teacher_logprobs.is_nested:
        teacher_logprobs = teacher_logprobs.values().unsqueeze(0)
        teacher_ids = teacher_ids.values().unsqueeze(0)
    model_inputs["teacher_topk_log_probs"] = teacher_logprobs
    model_inputs["teacher_topk_ids"] = teacher_ids
    clamp = tu.get_non_tensor_data(data=micro_batch, key="log_prob_min_clamp", default=None)
    if clamp is not None:
        model_inputs["log_prob_min_clamp"] = clamp

2. Read the three new fields off `model_output`#

verl/workers/engine/fsdp/transformer_impl.py — inside the existing if use_fused_kernels: branch of FSDPEngineWithLMHead.prepare_model_outputs (inherited by VeOmniEngineWithLMHead), right after the existing log_probs / entropy_rmpad reads:

# Replace existing top-level reads:
#   log_probs = output.log_probs.squeeze(0)
#   entropy_rmpad = output.entropy.squeeze(0)
# with the nested fused_linear_aux reads:
aux = output.fused_linear_aux
log_probs = aux.log_probs.squeeze(0)
entropy_rmpad = aux.entropy.squeeze(0)

if distillation_use_topk:
    cu_seqlens = input_ids.offsets()
    for k, src in (
        ("distillation_losses", aux.distillation_losses),
        ("student_mass", aux.student_mass),
        ("teacher_mass", aux.teacher_mass),
    ):
        v = src.squeeze(0)
        assert v.shape == log_probs.shape
        if self.use_ulysses_sp:
            pad_size = output_args["pad_size"]
            v = gather_outputs_and_unpad(v, gather_dim=0, unpad_dim=0, padding_size=pad_size)
        model_output[k] = torch.nested.nested_tensor_from_jagged(v, cu_seqlens)

What stays unchanged#

compute_forward_kl_topk, distillation_ppo_loss, compute_topk_loss, DistillationConfig / DistillationLossConfig, and the teacher_logprobs / teacher_ids / distillation_use_topk micro-batch fields — all reused as-is.
The non-fused logits_processor_func branch — untouched.
Estimator-only loss modes (k1, k2, k3, kl, low_var_kl, abs, mse) — they consume only log_probs, which is still populated identically.

Verification on VeOmni’s side#

Test	What it pins
`tests/ops/test_chunk_topk_distill.py` (10 tests)	Forward / backward numerics vs dense reference (encodes verl’s KL formula), chunk-size invariance, IGN masking, clamp consistency between forward & backward, temperature path
`tests/models/test_return_log_probs_e2e.py::test_return_log_probs_with_topk_distill_populates_three_fields` (qwen3-text, qwen3_vl-vlm)	The full `model.forward(..., teacher_topk_=...)` → `outputs.fused_linear_aux.distillation_losses` path on a real toy model. Asserts: fields populated with right shape/range, mass tensors detached, bitwise-equal to the kernel called directly on `model.lm_head.weight`*, backward reaches `lm_head.weight.grad`
`tests/ops/test_chunk_topk_distill.py::test_chunk_topk_distill_saves_memory_vs_eager` (CUDA)	Peak GPU memory of the fused kernel vs the eager `h @ w.T → log_softmax → gather` path on a large-V case (V=64k). Pins that we never materialize `[T, V]` student logits — the whole point of the kernel for verl.

verl Integration: Top-K Forward-KL Distillation via the VeOmni Engine

Contents