Repository navigation
Conversation
…nfig (verl#6676) Introduce three opt-in fields on DistillationConfig that scaffold the upcoming student-top-K reverse-KL pipeline (verl issue verl-project#6676): - teacher_backend ("vllm" | "fsdp"): selects whether the teacher is served by the existing vLLM/SGLang rollout (default, unchanged) or a frozen FSDP worker that gathers logprobs at student-top-K IDs each step. - teacher_chunk_size: token-axis chunk size for the FSDP gather kernel. - teacher_fsdp_config: reuses FSDPEngineConfig so the existing engine builder can construct the teacher model. When teacher_backend='fsdp' the post-init now skips the vLLM-only prompt_logprobs / max_logprobs / resource-pool-sizing checks (those assume a rollout deployment) and pins loss_mode to 'reverse_kl_topk' so misuse is caught early. The default vllm path is unchanged. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Implement the mode-seeking objective requested in verl issue verl-project#6676: L_t = sum_{i in TopK_student} q_i * (log q_i - log p_i) with the support taken on the *student*'s own top-K and the expectation under the student distribution. The teacher log-probabilities at those student-top-K IDs are precomputed by the upcoming FSDP teacher worker (P3), so this loss only consumes the gathered (B, T, K) tensor and does not need to materialize a full [B, T, V] log-softmax. Changes: - verl/trainer/distillation/fsdp/losses.py: add reverse_kl_divergence helper and compute_reverse_kl_topk, mirroring compute_forward_kl_topk's nested-tensor / ulysses-SP unwrap path. Returns distillation_losses, student_mass (≈ 1, sanity), and teacher_mass (the key mode-seeking diagnostic — should rise during training). - verl/trainer/distillation/losses.py: dispatch compute_topk_loss by loss_mode; route 'reverse_kl_topk' to the fsdp/veomni implementation (megatron raises NotImplementedError pending a vocab-parallel port); register the top-level wrapper that re-pads outputs, aggregates the mass diagnostics over the response mask, and clamps minor negative drift introduced by log_prob_min_clamp. - tests/workers/test_reverse_kl_topk_on_cpu.py: CPU-only unit tests for non-negativity, zero at distribution match, gradient flow into the student logits, agreement with the manual reference formula, and registry/metric-aggregation contracts. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Introduce the worker-side scaffold for the FSDP teacher backend: - verl/workers/teacher/utils.py: chunked_gather_logprobs — kernel that computes log_softmax(logits).gather(idx) without materializing the full log_softmax tensor, by using log_softmax(x).gather(idx) == x.gather(idx) - logsumexp(x) in fp32 over chunks of the leading axis. Peak memory becomes chunk_size * V * 4 bytes per side instead of N * V * 4. - verl/workers/teacher/fsdp_teacher_worker.py: TeacherFSDPWorker — subclasses TrainingWorker and pins engine_config.forward_only=True so FSDPEngine builds without optimizer/grad. Adds compute_logprobs_at_ids which reuses FSDPEngineWithLMHead.prepare_model_inputs (matching the actor's rmpad / nested-tensor layout), runs a single forward, and delegates to chunked_gather_logprobs at trainer-supplied topk_ids. The trainer is responsible for the roll-by-1 alignment; Ulysses SP on the teacher is gated off until a follow-up. - verl/experimental/teacher_loop/teacher_fsdp_manager.py: FSDPTeacherManager and MultiTeacherFSDPManager — mirror the public surface of TeacherModelManager / MultiTeacherModelManager so the trainer's resource-pool plumbing can switch on teacher_backend without diverging downstream consumers. Worker-group spawning and main-loop dispatch are explicitly NotImplementedError-marked for P4. - tests/workers/teacher/test_chunked_gather_on_cpu.py: CPU-only unit tests covering chunk_size sweep including > N rows, 2D rmpad input, shape/dtype preservation, validation errors, and fp32-stable behavior on large-magnitude logits. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
…676) Connect the post-rollout student-top-K -> teacher-gather pipeline into the trainer's RayPPOTrainer fit loop. With teacher_backend='fsdp': - agent_loop.py: gate the per-trajectory vLLM teacher logprob hook on use_vllm_teacher (False under the FSDP backend) so rollouts no longer attempt to reach a non-existent LLMServerClient. - workers/engine/fsdp/transformer_impl.py: relax the per-key shape assertion in prepare_model_outputs from `v.shape == log_probs.shape` to `v.shape[0] == log_probs.shape[0]`, allowing the logits processor to emit per-token, top-K outputs (e.g. (total_nnz, K)) which the reverse-KL pipeline needs for student_topk_log_probs / student_topk_ids. - trainer/distillation/losses.py: add a `compute_student_topk_only` data-flag fast path on compute_topk_loss that returns the student's own top-K log-probs and IDs, used by the actor's first-pass forward before the teacher is queried. - workers/engine_workers.py: add ActorRolloutRefWorker.compute_student_topk which stamps the data flag above and runs a no-grad infer_batch on the actor engine, returning the per-token student top-K as a TensorDict the trainer can hand to the teacher worker. - experimental/teacher_loop/teacher_fsdp_manager.py: rewrite as a thin manager (single-teacher only for now) that holds a Ray worker group handle installed by the trainer after the unified spawn. Multi-teacher raises NotImplementedError until the per-teacher resource-pool split is finalized. - trainer/ppo/ray_trainer.py: branch teacher manager creation by teacher_backend; pre-spawn, register TeacherFSDPWorker on the teacher resource pool with a TrainingWorkerConfig assembled from distillation.teacher_fsdp_config; post-spawn, install the spawned handle on MultiTeacherFSDPManager and reset the worker. Add helpers _build_fsdp_teacher_training_config and _inject_fsdp_teacher_data, and call the latter inside the fit loop's update_actor block (under a dedicated marked_timer) so update_actor sees student_topk_ids / teacher_on_student_logp on every step. The vLLM path is unchanged. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
…l#6676) Single-teacher OPD recipe that runs the verl#6676 student-top-K reverse-KL pipeline end-to-end: - distillation.teacher_backend=fsdp routes the teacher through the TeacherFSDPWorker instead of a vLLM rollout replica. - distillation.distillation_loss.loss_mode=reverse_kl_topk drives the student-top-K reverse-KL loss, with topk repurposed as the *student*-top-K width. - teacher_fsdp_config exposes the FSDPEngineConfig defaults the trainer now wires through (param_offload + bf16 + fsdp_size=-1). Defaults are sized for a single 8-GPU node with a 1.7B student and an 8B FSDP teacher; the script's user-adjustable block exposes the knobs needed to retarget at larger pairs. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
…#6676) - losses.py: short-circuit policy loss on the FSDP-teacher first pass (compute_student_topk_only) so the student top-K emit step does not trigger the final loss path. - ray_trainer.py: wrap TeacherFSDPWorker with ray.remote, materialize the non-blocking compute_logprobs_at_ids future, and pass temperature=1.0 so the FSDP engine emits the teacher's raw distribution. - engine.py: allow runtime override of EngineRouterReplayConfig.mode and type max/micro batch fields as Optional[int]. - fsdp_teacher_worker.py: move cu_seqlens to the values device before nested_tensor_from_jagged. - teacher/utils.py: move topk_ids to logits.device for torch.gather. - distillation.yaml: add teacher_backend / teacher_chunk_size / teacher_fsdp_config knobs for the FSDP teacher path; regenerate the trainer yaml variants. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
… padding (verl#6676) - ray_trainer.py: Phase 3 now scatters student_topk_ids / teacher_on_student_logp back to left-right padded dense (B, max_seq_len, K) using the Phase-1 ``indices`` / ``max_seq_len``, instead of no_padding_2_padding (which sliced response-only and left-shifted, breaking input_ids alignment). - padding.py: left_right_2_no_padding now also re-nests student_topk_ids and teacher_on_student_logp aligned to input_ids, matching the teacher_logprobs / teacher_ids handling so the reverse_kl_topk logits processor sees aligned jagged tensors.
…erl#6676) Track the student/teacher top-K intersection during reverse-KL OPD, following "Rethinking On-Policy Distillation of Large Language Models". The FSDP teacher now also emits its own top-K ids, carried through the left_right padding glue, so compute_reverse_kl_topk can report overlap_ratio and overlap_token_advantage. Also default use_policy_gradient to False. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
… (verl#6676) The student-top-K reverse-KL loss summed q_i*(log q_i - log p_i) over the student's own top-K without renormalizing. Once the student starts to de-peak, that truncated sum can be minimized by leaking probability mass *outside* the top-K (every q_i shrinks, so every term shrinks) instead of by matching the teacher. This drives a self-reinforcing collapse: entropy explodes, student_mass falls from ~0.97 to ~0.4 and overlap_ratio drops sharply, while the reported loss keeps decreasing. The onset is governed by the deterministic optimization drift, not batch noise, so it reproduces at the same step regardless of train batch size. Fix: renormalize both student and teacher log-probs onto the shared student-top-K support before the divergence, so the per-token objective is a proper KL(Q_K || P_K). student_mass / teacher_mass remain un-normalized as de-peaking diagnostics. Add a regression test asserting the loss is invariant to mass leaked outside the top-K, and update the reference-formula test.
…ture (verl#6676) The reverse-KL OPD loss compares the student and teacher distributions, so both must be taken at the same temperature. The FSDP teacher was queried at a fixed T=1.0 while the student rollout uses rollout.temperature, so the divergence was computed between two distributions at mismatched temperatures. Pass rollout.temperature through to the teacher forward to align them.
# Conflicts: # verl/workers/config/distillation.py # verl/workers/engine/fsdp/transformer_impl.py
The top-K distillation path may produce tensors of shape (total_nnz,) or (total_nnz, K) for student_topk_log_probs / student_topk_ids. The old strict assert required v.shape == (total_nnz,), which would fail for 2D outputs. Change it to only compare the first dimension.
YangxSong
requested review from
ArronHZG,
PeterSH6,
eric-haibin-lin,
tongyx361,
vermouth1992 and
wuxibin89
as code owners
July 2, 2026 16:46
|
|
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Code Review
This pull request introduces support for student-top-K reverse-KL on-policy distillation (OPD) using a frozen FSDP-served teacher backend. Key additions include the TeacherFSDPWorker and a memory-efficient chunked_gather_logprobs utility to compute teacher log-probabilities at student-selected token IDs without materializing full vocabulary-sized tensors. The PR also adds the reverse_kl_topk loss mode, a dedicated FSDPTeacherManager, configuration extensions, padding utilities, integration tests, and example bash scripts for running reverse-KL distillation. I have no feedback to provide as there are no review comments to assess.
Important
The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.
7 of 8 tasks
1 of 3 tasks
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
What does this PR do?
Adds reverse-KL on-policy distillation with student-top-K support (verl issue #6676), following the formulation in Rethinking On-Policy Distillation: the loss is
KL(Q_student_topk || P_teacher_topk), which the teacher log-probs are evaluated at the student's own top-K tokens rather than the teacher's. Because the query IDs are student-specific and change every step, the teacher is served as a frozen forward-only FSDP worker (teacher_backend='fsdp') that gathers log-probs at arbitrary IDs via a chunked gather kernel. The existing vLLM teacher path is unchanged.Checklist Before Starting
[{modules}] {type}: {description}(This will be checked by the CI)...Test
Unit tests
The following unit tests have been verified:
python -m pytest tests/workers/test_reverse_kl_topk_on_cpu.py \ tests/workers/teacher/test_chunked_gather_on_cpu.py -vEnd-to-end test
Run the script:
examples/on_policy_distillation_trainer/run_dsr1_1.5b_reverse_kl_fsdp.sh. Use reverse-kl-topk to distill knowledge from JustRL-DeepSeek-1.5B into DeepSeek-R1-Distill-Qwen-1.5B on DAPO-17k.API and Usage Example
New opt-in
DistillationConfigfields (defaults preserve current behavior):distillation.teacher_backend("vllm"default /"fsdp"),teacher_chunk_size,teacher_fsdp_config. Under the FSDP backenddistillation_loss.topkrefers to the student top-K width.# Full recipe: examples/on_policy_distillation_trainer/run_dsr1_1.5b_reverse_kl_fsdp.sh python3 -m verl.trainer.main_ppo ... \ distillation.enabled=True \ distillation.teacher_backend=fsdp \ distillation.teacher_models.teacher_model.model_path=hbx/JustRL-DeepSeek-1.5B \ +distillation.teacher_fsdp_config.strategy=fsdp \ +distillation.teacher_fsdp_config.dtype=bfloat16 \ +distillation.teacher_fsdp_config.param_offload=True \ distillation.distillation_loss.loss_mode=reverse_kl_topk \ distillation.distillation_loss.topk=16Design & Code Changes
New files
verl/workers/teacher/fsdp_teacher_worker.py:TeacherFSDPWorker— a forward-only FSDP worker that loads and hosts the frozen teacher model, withget_teacher_logprobs_at_ids(student_topk_ids)as the main entry point.verl/workers/teacher/utils.py:chunked_gather_logprobs— streams log-prob gathering in chunks along(B*T), usinggather(logits, ids) - logsumexp(logits)to avoid materializing[B, T, V].verl/experimental/teacher_loop/teacher_fsdp_manager.py:TeacherManagerFSDP— thin single-teacher manager mirroring the vLLM teacher-manager interface.verl/workers/teacher/__init__.py: package init for the teacher module.examples/on_policy_distillation_trainer/run_dsr1_1.5b_reverse_kl_fsdp.sh: end-to-end recipe for DSR1-1.5B distillation.examples/on_policy_distillation_trainer/run_qwen3_reverse_kl_fsdp.sh: Qwen3 variant recipe.tests/workers/test_reverse_kl_topk_on_cpu.py: CPU unit tests for the reverse-KL loss computation and renormalization.tests/workers/teacher/test_chunked_gather_on_cpu.py: CPU unit tests for the chunked gather kernel.Modified files
verl/trainer/distillation/fsdp/losses.py: addcompute_reverse_kl_topk— gathers student log-probs atstudent_topk_ids, renormalizes both sides to the shared student-top-K support, computesKL(Q_K || P_K), and emits raw mass + overlap diagnostics.verl/trainer/distillation/losses.py: registerreverse_kl_topkloss mode; add top-level wrapper that re-pads per-micro-batch tensors, collects diagnostics, and clamps ≥0.verl/workers/config/distillation.py: addteacher_backend,teacher_chunk_size,teacher_fsdp_configfields toDistillationConfigwith post-init validation.verl/workers/config/engine.py: add FSDP engine config entries needed by the teacher worker.verl/trainer/ppo/ray_trainer.py: branch teacher creation onteacher_backend; in the FSDP path, run a lightweight student forward (compute_student_topk_only=True) to collectstudent_topk_ids, query the teacher, and stamp the result into the batch.verl/workers/engine_workers.py: expose teacher worker integration on the actor rollout worker group.verl/workers/engine/fsdp/transformer_impl.py: supportcompute_student_topk_onlyflag — short-circuits after computing student top-K without running policy loss.verl/workers/utils/padding.py: add pad/unpad helpers for nested tensor handling in the logits processor pipeline.verl/experimental/agent_loop/agent_loop.py: wire FSDP teacher manager into the agent loop.verl/trainer/config/distillation/distillation.yaml& generated trainer configs: addteacher_backend,teacher_chunk_size,teacher_fsdp_configfields.Checklist Before Submitting
pre-commit install && pre-commit run --all-files --show-diff-on-failure --color=always