[megatron] feat: add memory-efficient chunked top-k KL computation by HSYZhang · Pull Request #7733 · verl-project/verl

HSYZhang · 2026-09-04T03:02:25Z

What does this PR do?

Adds an opt-in, memory-efficient chunked path for Megatron vocab-parallel top-k forward KL. It extends the FSDP chunking approach from #6593 to tensor-parallel vocabulary shards, avoiding full fp32 [N, V_local] log-softmax/probability buffers while preserving the existing outputs and gradients.

With chunked loss implemented, small student models (e.g., 4B or smaller) can reduce GPU memory via an FSDP-like chunked loss approach, eliminating the need to increase TP to lower peak memory usage. This prevents GPU idleness and significantly improves utilization. We previously used a 4B student model, but the trainer OOM'd during loss computation. The original approach required TP=4 to avoid OOM; with chunked loss enabled, we can now run at TP=1 without OOM, greatly improving GPU utilization and reducing idle time.

Test

Results (b=2, s=96, V=2048, K=32; d_* denotes the maximum absolute error):

Config	Comparison	d_loss	d_grad	d_smass	d_ocount	Verdict
fp32, clamp=-10	baseline vs reference	8.9e-8	1.2e-9	—	—	✅
fp32, clamp=-10	chunked(cs=7/64/4096) vs baseline	≤1.2e-7	≤1.3e-9	1.9e-8	0	✅
fp32, clamp=None	baseline vs reference	8.9e-8	1.2e-9	—	—	✅
fp32, clamp=None	chunked vs baseline	≤1.2e-7	≤1.3e-9	1.9e-8	0	✅
bf16, clamp=-10	chunked vs baseline	≤8.9e-8	1.9e-6	1.9e-8	0	✅

Under fp32, chunked and baseline are nearly bit-for-bit identical (same floating-point operators and reduction order when TP=1);
overlap_count is exactly identical (diff=0);
For bf16 (storing bf16 logits and recomputing p in backward), the additional rounding is only ~1.9e-6;
chunk_size=7 (many odd-sized chunks + a residual chunk) produces results consistent with 4096 → chunk boundary handling is correct.

Actual Qwen3.5-4B distillation dimensions: V=248320, K=128, TP=1, N=4096 tokens, clamp=-10.
Individually compared the maximum absolute/relative differences across 5 forward outputs + gradients.

bf16 logits（actual online dtype）：

chunk_size	per_token_kl	student_mass	teacher_mass	overlap_adv	grad	Conclusion
2048	0.00e+00	0.00e+00	0.00e+00	0.00e+00	0.00e+00	Bit-for-bit identical
4096	0.00e+00	0.00e+00	0.00e+00	0.00e+00	0.00e+00	Bit-for-bit identical
8192	0.00e+00	0.00e+00	0.00e+00	0.00e+00	0.00e+00	Bit-for-bit identical

fp32 logits（stricter rounding observation）：

chunk_size	per_token_kl (abs/rel)	grad (abs/rel)	worst_rel	Conclusion
2048	9.31e-10 / 8.13e-08	8.53e-14 / 2.96e-10	8.22e-07	OK
4096	9.31e-10 / 8.13e-08	8.53e-14 / 2.96e-10	8.22e-07	OK
8192	9.31e-10 / 8.13e-08	8.53e-14 / 2.96e-10	8.22e-07	OK

Peak memory

End-to-end forward/backward peak memory with chunk_size=4,096:

TP=1 (V_local=248,320)

tokens (`N`)	logits size	baseline peak	chunked peak	saved	reduction
4,096	1.90G	13.28 GiB	13.29 GiB	-0.01 GiB	-0.1%
8,192	3.80G	26.56 GiB	18.97 GiB	7.59 GiB	28.6%
16,384	7.60G	53.12 GiB	30.36 GiB	22.76 GiB	42.8%
24,576	11.40G	79.68 GiB	45.54 GiB	34.14 GiB	42.8%

TP=4 (V_local=62,080 per rank, matching the E20 launch)

tokens (`N`)	logits size per rank	baseline peak	chunked peak	saved per rank	reduction
16,384	1.92G	13.33 GiB	7.63 GiB	5.71 GiB	42.8%
32,768	3.84G	26.67 GiB	15.25 GiB	11.41 GiB	42.8%
49,152	5.75G	40.00 GiB	22.88 GiB	17.12 GiB	42.8%

At small token counts the chunk workspace makes peak memory effectively unchanged. The benefit grows with sequence length, reaching a stable 42.8% reduction; for the TP=4 E20-scale case at N=49,152, this saves 17.12 GiB per rank.

API and Usage Example

No new API is introduced; this uses the existing distillation configuration added for the FSDP path:

distillation_loss:
  use_chunked_topk: true
  chunked_topk_chunk_size: 4096

The feature remains disabled by default.

Design & Code Changes

Forward processes tokens in chunks, performs TP MAX/SUM reductions to compute the global logZ, and gathers only teacher top-k plus local student top-k statistics.
Backward saves the original low-precision logits and per-token fp32 logZ, then recomputes probabilities chunk by chunk.
This removes the persistent full fp32 probability tensor and limits temporary fp32 vocabulary storage to [chunk_size, V_local], at the cost of additional TP collectives and recomputation.
The baseline path is unchanged when chunking is disabled.

Checklist Before Submitting

Searched for similar PRs; this implementation follows [fsdp] feat: chunked gather-logsumexp for top-K loss to avoid OOM at long context #6593.
PR title follows the required format.
Added targeted numerical validation, including chunk boundaries and realistic model dimensions.
No documentation update is required; the existing FSDP configuration is reused without API changes.

tardis-key · 2026-09-21T02:30:20Z

LGTM. Please check CI to confirm there are no failures caused by this PR. The pre-commit failures need you to update the code accordingly. @HSYZhang

HSYZhang · 2026-09-21T03:02:05Z

LGTM. Please check CI to confirm there are no failures caused by this PR. The pre-commit failures need you to update the code accordingly. @HSYZhang

Thanks, done — I’ve fixed the pre-commit failures and rebased the latest code.

…computation

Assisted-by: OpenAI Codex

…olicy distillation (#7855) ### What does this PR do? Adds a `forward_kl_topk_tail` loss mode for on-policy distillation. The existing `forward_kl_topk` objective retains only the teacher's top-k terms: ``` L_truncated = Σ_{i∈T} p_i (log p_i − log q_i) ``` Because `T` is truncated, the retained masses `P = Σ_{i∈T} p_i` and `Q = Σ_{i∈T} q_i` are both below 1. The quantity is therefore not a KL and goes negative whenever `Q > P`; the current implementation clamps it with `clamp_min(0)`, which silently discards the gradient carried by those tokens. In our runs the per-token minimum was exactly `0.0` (clamping active) in **415 / 467 training steps (89 %)**. This PR treats every token outside `T` as one aggregate category: ``` 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) ``` which is the exact KL between the two `K+1`-category distributions: non-negative by construction (no clamp), a lower bound on the full forward KL, and it reuses the existing teacher top-k payload, so teacher compute, schema and network traffic are unchanged. **The claim of this PR is training efficiency**: at an equal training-step budget, `forward_kl_topk_tail` reaches a higher GSM8K accuracy than `forward_kl_topk`; equivalently, it reaches a given accuracy in fewer steps. It is not a claim about a higher final ceiling — the two curves approach the same asymptote and the gap narrows with longer training, which is the expected shape of faster convergence. Related in verl: #7239 / #7375 (full-vocabulary OPD), #6676 / #7590 / #7806 (reverse-KL on the student top-k support), #7733 / #7809 (memory/perf work on the existing top-k path) — none of them change the truncated forward-KL objective. **The objective itself is prior art, and this PR does not claim it as new.** Section 3.3 of the [LFM2 Technical Report](https://arxiv.org/abs/2511.23404v1) optimizes `KL(Bern(P_T) || Bern(P_S)) + P_T * KL(P_T(.|T) || P_S(.|T))`; at temperature 1 that is algebraically the coarse-grained KL written above, and the report states the same lower bound. The aggregate tail bucket is the *ghost token* of [Sparse Logit Sampling](https://aclanthology.org/2025.acl-long.885.pdf). What this PR adds is the integration into the OPD path: the analytic gradient on the out-of-top-k tokens for vocab-parallel (Megatron TP) shards so no rank gathers logits, reuse of the existing `[S, K]` teacher payload with no schema or communication change, the `tail_mass_eps` guard semantics, tests + TP parity, and the efficiency measurement below. ### Checklist Before Starting - [x] Search for similar PRs. Query links: - https://github.com/verl-project/verl/pulls?q=is%3Apr+is%3Aopen+distillation+kl - https://github.com/verl-project/verl/pulls?q=is%3Apr+is%3Aopen+topk - https://github.com/verl-project/verl/issues?q=is%3Aissue+is%3Aopen+distillation+kl - [x] Format the PR title as `[{modules}] {type}: {description}` ### Test Unit tests (CPU), on Linux with `torch 2.11.0+cu129`: ```bash python -m pytest tests/workers/test_distillation_topk_symmetry_on_cpu.py \ tests/workers/test_chunked_topk_log_probs_on_cpu.py \ tests/test_base_config_on_cpu.py tests/utils/test_config_on_cpu.py -q # 41 passed ``` The new coverage asserts: 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 and chunked FSDP paths; invalid clamp/epsilon configuration; and the fused-mass fallback / fail-closed behaviour. Megatron tensor-parallel parity (2 GPUs): ```bash torchrun --standalone --nnodes=1 --nproc-per-node=2 tests/utils/test_special_megatron_kl_loss_tp.py # [PASS] VP KL divergence correctness verified for test case 0..3 # [PASS] VP tail-aware KL correctness verified for test case 0..3 ``` Loss, `student_mass`, `teacher_mass` and `tail_loss` match the FSDP reference, and every vocabulary-shard gradient matches FSDP within `atol=rtol=1e-4`; existing truncated top-k cases do not regress. Synthetic full-KL benchmark (`examples/on_policy_distillation_trainer/benchmark_tail_kl.py`) — gradient cosine similarity to the true full-KL gradient, and the frequency of negative (clamped) truncated losses: | K | truncated | tail-aware | negative truncated loss frequency | |---|---|---|---| | 8 | 0.303 | 0.455 | 6.25 % → 0 % | | 16 | 0.327 | 0.561 | 0 % | | 32 | 0.393 | 0.653 | 0 % | | 64 | 0.488 | 0.793 | 0 % | Scalar error relative to the full KL is *not* uniformly better (it is slightly worse), which is expected: tail-aware KL is a lower bound, not a better point estimate. We do not claim otherwise. Real-model efficiency measurement — student Qwen3-1.7B (FSDP + colocated vLLM rollout) and teacher Qwen3-32B (vLLM TP=2) on 3×RTX A6000, GSM8K, batch 16, identical seed/data/batch/optimizer for both arms, validation on the full 1319-example test set (standard error ≈ ±1.3 pp). The only difference between arms is `loss_mode`; both arms run with `log_prob_min_clamp=null`. Experiment A (same run, same step grid): | steps | `forward_kl_topk` | `forward_kl_topk_tail` | |---|---|---| | 175 | 0.3177 | **0.3472** | | 350 | 0.3457 | **0.3616** | At 175 steps the tail-aware arm already matches what the baseline reaches at 350 steps (0.3472 ≥ 0.3457) — i.e. **≤ 50 % of the training steps for the same accuracy**. Experiment B (independent seed, earliest phase): | steps | `forward_kl_topk` | `forward_kl_topk_tail` | |---|---|---| | 50 | 0.3412 | **0.4246** | The leading margin shrinks with training (50 steps +8.34 pp, 175 steps +2.96 pp, 350 steps +1.59 pp), consistent with two curves sharing an asymptote but converging at different rates. Cost is unchanged: 142.6 vs 143.0 tokens/s, 40.58 vs 40.61 GB peak HBM, and `grad_norm` of the same order, so the efficiency difference is not bought with extra compute. Mechanistically, `loss_min` is exactly `0.0` (clamped) in 415/467 steps for the baseline versus `≈ −1e−7` (finite-precision noise, never clamped) for the tail-aware arm, while the top-k overlap ratio is essentially identical (0.7153 vs 0.7119), i.e. the difference does not come from different token selection. Caveats: both arms are single-seed, and the measured run-to-run noise for the same arm at the same step is ≈3.3 pp, so Experiment B's margin is ≈2.5× the noise while Experiment A's is of the same order. We report three step budgets rather than only the most favourable one, and multi-seed error bars are in progress. ### API and Usage Example No API change; it is a new value for an existing config field. ```bash distillation.distillation_loss.loss_mode=forward_kl_topk_tail distillation.distillation_loss.topk=16 distillation.distillation_loss.log_prob_min_clamp=null # required by this mode distillation.distillation_loss.use_policy_gradient=false # intended setting ``` `tail_mass_eps` (new, default `1e-6`) is the numerical guard on `q_tail`; it must lie strictly between 0 and 1. ### Design & Code Changes - `verl/trainer/distillation/tail_kl.py` (new): stable tail-bucket value computation in FP32, plus an analytic, Megatron-independent logit-gradient helper. For `j` inside the teacher top-k `dL/dz_j = q_j − p_j`; for `j` outside, `dL/dz_j = q_j (P − Q) / (1 − Q)`, which every vocabulary shard can apply locally without gathering logits. The helper also matches the derivative of the `q_tail.clamp_min(tail_mass_eps)` guard when it is active. - `verl/trainer/distillation/losses.py`: registers the mode with `use_topk=True` (so the teacher path requests top-k log-probs as before), dispatches to the FSDP/VeOmni and Megatron backends, keeps `clamp_min(0)` only for the legacy mode, and adds `head_loss`, `tail_loss`, `teacher_tail_mass`, `student_tail_mass` metrics. If a fused engine returns a detached `student_mass` it fails closed instead of silently training with an incomplete gradient. - `verl/trainer/distillation/fsdp/losses.py`: refactors the existing top-k path without changing legacy behaviour, adds `compute_forward_kl_topk_tail`, and supports both the regular and `use_chunked_topk=True` paths. - `verl/trainer/distillation/megatron/losses.py`: extends `_VocabParallelKLDivergence` with `include_tail` / `tail_mass_eps`, all-reduces the student top-k mass across TP ranks, and uses the shard-local analytic gradient in backward. - `verl/workers/config/distillation.py` + config YAMLs: `tail_mass_eps`, validation rules (strictly between 0 and 1; `forward_kl_topk_tail` requires `log_prob_min_clamp=null`), and a warning when combined with `use_policy_gradient=true`. Per-entry log-prob clamping is rejected for this mode because it would change the selected top-k mass and invalidate the coarse-grained KL. - Tests, benchmark, docs and generated configs as listed in the diff. **VeOmni backend.** Two notes, neither of which is a defect in VeOmni: - *Fused-kernel masses — our earlier claim was wrong.* Our first measurement with `model_engine=veomni` + `use_fused_kernels=True` showed `teacher_mass ≈ 2.11e-4`, which we described as the fused kernel's mass outputs being degenerate. Instrumenting the payload showed instead that every teacher top-k log-prob was exactly `-log V` — the *teacher payload* was uniform — and `2.11e-4 = K / V` is the correct mass of that payload. Serving the same teacher directly through vLLM returns a real distribution, and re-running the example with the default `distillation.teacher_models.teacher_model.*` entry instead of the named `+distillation.teacher_models.gsm8k.*` entry delivers a real payload (`Σ exp = 0.0377`). The artefact was in our own configuration, not in the fused kernel. - *Design limitation (documented).* `chunk_topk_distill` returns `student_mass` / `teacher_mass` with `requires_grad=False` by design, so the tail bucket cannot receive gradient through the fully-fused path. `forward_kl_topk_tail` therefore requires `use_fused_kernels=False` on the VeOmni backend; the fail-closed guard rejects the fused combination rather than silently training with an incomplete gradient. Two environment-level blockers were also worked around for the measurement: veomni 0.1.11's generated Qwen3 patch passes `cache_position` to `transformers.masking_utils.create_causal_mask`, which the pinned transformers 5.9.0 no longer accepts; and single-GPU VeOmni runs need `actor_rollout_ref.actor.veomni.init_device=cuda` because `fsdp_enabled` is false when `fsdp_size == 1` while the default is `meta`. ### Checklist Before Submitting - [x] Read the [Contribute Guide](https://github.com/verl-project/verl/blob/main/CONTRIBUTING.md). - [x] Apply pre-commit checks: `pre-commit run --all-files` (ruff, ruff-format, mypy 1.17.0, generated-config verification, docstring/license/device-API/DataProto/structure/example-naming/uv-gpu-only checks and `compileall` with warnings-as-errors all pass on Linux). - [x] Add / Update the documentation (`docs/algo/opd.md`, `examples/on_policy_distillation_trainer/README.md`). - [ ] Add unit or end-to-end test(s) to the CI workflow. The new CPU tests extend the existing `tests/workers/test_distillation_topk_symmetry_on_cpu.py` and the Megatron TP test covers the GPU path; wiring them into `.github/workflows` may need a maintainer's guidance on where they best fit. - [ ] CI request sent in the `ci-request` channel. **AI assistance disclosure:** the implementation, tests and experiment harness were produced with AI assistance (Claude Code). The submitting human reviews every changed line and ran the tests and experiments reported above. --------- Co-authored-by: lipengyu <eric.lpy@antgroup.com> Co-authored-by: OpenAI Codex <noreply@openai.com> Co-authored-by: Claude Code <noreply@anthropic.com>

HSYZhang requested review from PeterSH6, eric-haibin-lin, tongyx361 and vermouth1992 as code owners September 4, 2026 03:02

HSYZhang changed the title ~~[megatron] feat: add memory-efficient chunked top-k KL computation~~ [distillation, megatron] feat: add memory-efficient chunked top-k KL computation Sep 4, 2026

HSYZhang mentioned this pull request Sep 9, 2026

[megatron] feat: fuse linear top-k log probabilities #7809

Open

7 of 8 tasks

This was referenced Sep 13, 2026

[algo, fsdp, megatron, veomni] feat: add tail-aware top-k KL for on-policy distillation #7855

Merged

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

Open

tardis-key changed the title ~~[distillation, megatron] feat: add memory-efficient chunked top-k KL computation~~ [megatron] feat: add memory-efficient chunked top-k KL computation Sep 21, 2026

HSYZhang force-pushed the megatron-chunked-topk-kl branch from 3ddcf08 to 34f14a0 Compare September 21, 2026 02:51

HSYZhang added 2 commits September 21, 2026 14:27

[distillation, megatron] feat: add memory-efficient chunked top-k KL …

4f35432

…computation

[megatron] style: apply Ruff formatting to chunked top-k KL

7c52182

Assisted-by: OpenAI Codex

HSYZhang force-pushed the megatron-chunked-topk-kl branch from 6c89d36 to 7c52182 Compare September 21, 2026 07:01

lindicaphxag-tech mentioned this pull request Oct 2, 2026

[OPD] [megatron/fsdp] use forward_kl_topk method occurs oom #6810

Open

2 of 4 tasks

taking-lying-flat mentioned this pull request Nov 5, 2026

[fsdp, algo] fix: preserve fp32 log probabilities for tail KL #8548

Open

5 of 7 tasks

Navigation Menu

Repository navigation

[megatron] feat: add memory-efficient chunked top-k KL computation - #7733

[megatron] feat: add memory-efficient chunked top-k KL computation#7733
HSYZhang wants to merge 2 commits into
verl-project:mainfrom
HSYZhang:megatron-chunked-topk-kl

HSYZhang commented Sep 4, 2026 •

edited

Loading

Uh oh!

tardis-key commented Sep 21, 2026

Uh oh!

HSYZhang commented Sep 21, 2026

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

Conversation

HSYZhang commented Sep 4, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

What does this PR do?

Test

Peak memory

API and Usage Example

Design & Code Changes

Checklist Before Submitting

Uh oh!

tardis-key commented Sep 21, 2026

Uh oh!

HSYZhang commented Sep 21, 2026

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

HSYZhang commented Sep 4, 2026 •

edited

Loading