Repository navigation
Add Triton kernels for vocab-parallel logprobs, entropy and KL - #8838
Conversation
The vocab-parallel ops ran each row chunk as a dozen separate eager fp32 kernels, making them 2-2.4x slower than a dense log_softmax on the same shard and keeping several fp32 copies of each chunk alive. On CUDA/XPU with Triton available, each pass now runs in one kernel per row that loops over the local vocabulary: local max and sum-exp in a single read (rescaled to the global max, still one MAX and one SUM all-reduce), then the fused forward and backward passes. Other devices, non-contiguous inputs, and deepspeed.HAS_TRITON = False keep the existing eager path unchanged. Intel Arc Pro B70, bf16, 75968-column shard, N=8192, fwd+bwd: reverse KL 340.8 -> 15.5 ms, forward KL 313.0 -> 15.5 ms, entropy 205.5 -> 9.1 ms, logprobs 99.5 -> 7.2 ms; extra peak memory 192-510 MiB -> ~0. KL is within ~10% of the 608 GB/s bandwidth floor. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Jin, Youzhi <youzhi.jin@intel.com>
The Triton forward for entropy and KL ran two passes per input: a stats pass for the local max and sum-exp, then a second pass for the entropy or KL terms. On H100 the kernels are at the memory-bandwidth floor of that pass count, so the remaining cost is the extra read. Each shard now accumulates, in the same pass as its max and sum-exp, the probability-weighted sum of logits centred on that max (for KL, the centred log-ratio weighted by p_A, plus a flag for support that B masks out). After the MAX all-reduce the sums are re-centred on the global max, so H = log S - W / S and KL = W / S_A - log S_A + log S_B. Centring on the max keeps a large common offset from cancelling in fp32. With TP the forward needs one MAX and one SUM all-reduce instead of one MAX and two SUMs. Rows are now addressed with up to two strides, so a slice such as logits[:, :-1] is read in place instead of being copied by reshape in both forward and backward. Other layouts are copied first, as before. Gradients are returned contiguous. H100, bf16, 8192 x 75968 shard, fwd+bwd vs the two-pass kernels: reverse KL 5.83 -> 4.48 ms (fp32 input 11.5 -> 8.39 ms), entropy 3.46 -> 2.69 ms; logits[:, :-1] input: reverse KL 13.6 -> 4.5 ms and 2.3 GiB less peak memory. logprobs is unchanged (within 1% at 8192 rows, about 25 us slower at 1024 rows). Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: 0z5a <dezhen.lu@student.uni-tuebingen.de>
|
Thanks for the Triton follow-up. I validated it on H100 and tried a few kernel-level changes on top of it; numbers and a patch are below. #8838 on H100
Suggested change: read each logit once in the entropy / KL forwardSince the kernels are bandwidth bound, the remaining lever is the number of passes. The forward currently reads each input twice: once for the local max and sum-exp, then again for the entropy or KL terms. The first pass can also accumulate the probability-weighted sum of logits centred on the running max: Σ e^(z−m)(z−m) for entropy, and Σ e^(a−m_a)((a−m_a)−(b−m_b)) for KL(A‖B). When the max grows, these sums are rescaled the same way as the sum-exp. After the MAX all-reduce, each shard's sums are re-centred on the global max:
Centring on the max keeps a large common offset from cancelling in fp32; your large-offset test still passes. Support mismatch, meaning B masks a token that A supports, is a flag decided on the logits; it shares the MAX all-reduce. With TP the forward then needs one MAX and one SUM all-reduce instead of one MAX and two SUMs. The backward kernels are unchanged. Separately, rows are addressed with up to two strides, so a slice such as H100, fwd+bwd, 8192 × 75968 shard. The #8838 and patch columns are medians of 3 interleaved runs with nothing else on the node; the eager column comes from a single run.
All outputs and gradients match the eager path within your test tolerances, with NaN/inf in the same places. Validation:
One side effect I noticed, on a 1024 × 75968 KL in both directions: with the patch, bf16 and fp32 inputs give bitwise identical outputs, and the bf16 gradient equals the round-to-nearest-even of the fp32 one. With the current kernels they differ by up to 4e-6. My guess is that the per-thread layout, and so the reduction order, depends on the dtype; I have not verified that. Either way it is harmless, just a note. I have not run anything on XPU, so the 2048-wide statistics blocks are only tuned on H100. I opened the patch as a single-commit PR against your branch: jinyouzhi#5. Feel free to take all or part of it. AI assistance was used for the implementation and analysis; I reviewed the changes and the measurements. |
Read each logit once in the vocab-parallel entropy/KL forward
|
@0z5a Thank you for your optimization patch. I merged it. |
It’s worth checking batch invariance when we move on to RL rollouts. |
The docs and the test skip now follow the dispatch predicate instead of a hard-coded CUDA/XPU list. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Jin, Youzhi <youzhi.jin@intel.com>
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Jin, Youzhi <youzhi.jin@intel.com>
Codex Review SummaryThis comment shows the latest Codex review activity on this pull request.
ℹ️ About Codex in GitHubYour team has set up Codex to review pull requests in this repo. Reviews are triggered when you
Codex reacts with 👀 while any review is running, comments if it has suggestions, and reacts with 👍 once all reviews finish with no findings. |
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 987bd4aad5
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
| if logits[0].device.type not in ("cuda", "xpu"): | ||
| return False | ||
| return deepspeed.HAS_TRITON |
There was a problem hiding this comment.
On CUDA devices that DeepSpeed marks as Triton-unsupported (CudaAccelerator.is_triton_supported() returns false below compute capability 8, including V100/T4), HAS_TRITON only establishes that a Triton backend driver is active, so this predicate still selects these kernels. Any of these formerly eager vocab-parallel ops then reaches an unsupported Triton JIT instead of falling back; check the accelerator capability here, and represent the validated XPU support through that capability rather than hard-coding device types.
AGENTS.md reference: AGENTS.md:L40-L40
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Removed the device specific dispatch and follow upper capacibility status.
accelerator = get_accelerator()
return deepspeed.HAS_TRITON and accelerator.is_triton_supported() and accelerator.on_accelerator(logits[0])| scale = tl.where(seen, tl.exp(run_max - row_max), 0.0) | ||
| run_weighted = tl.where(seen, scale * (run_weighted - (row_max - run_max) * run_sum), 0.0) |
There was a problem hiding this comment.
When a valid fp32/bfloat16 row spans the finite extremes (for example [-3e38, 3e38]), row_max - run_max overflows to inf while scale underflows to zero, so this evaluates 0 * (0 - inf) and writes NaN into the weighted sum. The eager implementation returns entropy 0 for this one-hot distribution; the analogous updates in _kl_stats_kernel and the TP re-centering in _triton_entropy/_triton_kl_div have the same 0×inf or inf−inf failure, even yielding NaN for KL of two identical rows. Mask the weighted contribution when the rescale is zero, including in the distributed re-centering, before multiplying by the shift.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Reproduced (entropy and KL were NaN on XPU for a row spanning ±3e38) and fixed in 6a53407: zero-weight terms are skipped in both kernels (in-loop rescale, per-element term, lane re-centring) and in the TP re-centring. Added single-process (fp32/bf16, peak in first and last block) and TP=2 tests; on a one-hot row entropy and KL(p || p) are exactly 0 with zero gradient. The new tests fail without the fix.
CUDA devices below compute capability 8 now take the eager path: a Triton driver may load there although Triton does not support them. The XPU accelerator still reports no Triton support, to keep DeepSpeed's other Triton ops off, so XPU is enabled here directly. When a row spans finite extremes such as -3e38 and 3e38, the difference of two maxima overflows: the rescale underflows to 0 while the shift is infinite, and entropy and KL became NaN. Terms with zero weight are now skipped in the kernels and in the tensor-parallel re-centring, which leaves every finite result unchanged. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Jin, Youzhi <youzhi.jin@intel.com>
Use the same check as the vocab-parallel cross entropy: Triton installed, the accelerator reports Triton support, and the logits live on it. XPU falls back to the eager path until its accelerator enables Triton. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Jin, Youzhi <youzhi.jin@intel.com>
…deepspeedai#8958) ## Motivation Teacher-student training compares the models' next-token distributions, not just the probability of a single ground-truth token. DeepSpeed already offers vocabulary-parallel forward/reverse KL, but does not support the interior `0 < beta < 1` generalized Jensen-Shannon divergence used by GKD-style distillation. Applying a local softmax independently on each vocabulary shard does not compute the intended full-vocabulary objective. TRL currently has two relevant implementations on its main branch: `experimental.GKDTrainer` computes generalized JSD from materialized logits; `DistillationTrainer` uses completion filtering and tiled LM-head projection to compute KL/JSD without materializing full logits. Those implementations do not call DeepSpeed's vocabulary-parallel divergence ops. In particular, support for DeepSpeed/ZeRO or FSDP parameter sharding is not itself support for a vocabulary-parallel softmax and divergence. This PR supplies the missing Core logits-based primitive, as a building block for teacher-student workloads and a possible subsequent TRL integration. It does not claim to integrate TRL or replace its memory-efficient tiled loss. Refs deepspeedai#8790. ## Changes - Add `vocab_parallel_jsd(student_logits, teacher_logits, beta=0.5, tp_group=None)`, returning FP32 per-token divergence replicated on TP ranks. - Add fused Triton paired softmax statistics, divergence and backward kernels. Reuse global normalization and reduce the two KL-to-mixture expectations together. Recompute probabilities in registers during backward rather than retaining full FP32 probability tensors; keep bounded eager chunks as fallback. - Enable JSD on an active Intel Triton backend without changing XPU's global Triton capability flag or existing-op dispatch. - Support gradients to both inputs, detached teachers, uneven vocabulary shards, empty local shards, and non-contiguous causal-shift views. - Preserve the existing Core convention: beta=0 selects forward KL and beta=1 selects reverse KL. These are explicit endpoint branches, not continuous limits of the weighted interior JSD. - Leave temperature, optional T-squared scaling, masking and normalization to the caller, following the existing vocabulary-parallel per-token ops. ## Not a duplicate - deepspeedai#8791 / deepspeedai#8838 implement logprobs, entropy and forward/reverse KL, not interior-beta generalized JSD. The new API reuses their KL endpoints. - deepspeedai#8947 fixes counting and accumulation in the existing dense rollout loss. This PR does not edit `runtime/rollout/opsd.py` or its tests. - DeepSpeedExamples#1019 integrates vocabulary-parallel KL into OPSD. This PR does not change Examples, cache handling, rollout or trainer integration. - deepspeedai#8945 connects fused-linear CE to HF forward. This PR is independent of it: no LM-head projection changes or loss-function hooks. ## Scope and limitations This is a new opt-in logits-based API; existing ops and training defaults are unchanged. Unlike the existing single-pass KL forward, JSD needs global normalization before evaluating the nonlinear mixture: statistics and divergence are two input passes. It still consumes materialized local logits. Top-k restriction, pointwise clipping, differing student/teacher shard layouts, fused-linear divergence, DeepCompile support and automatic TRL integration are not included. Every TP rank must supply identically ordered student/teacher vocabulary shards and backpropagate the same replicated per-token loss. ## Validation Hardware: up to 4 Intel Arc Pro B70 GPUs, XCCL, driver 1.15.39122+11, torch 2.14.0.dev20260804+xpu, Intel Triton 3.8.0. Final validation: **67 CPU regressions passed, 11 fused XPU cases passed, and 4 distributed cases passed at each of TP1/2/4**. Existing entropy/KL Triton-only tests still skip on XPU because their accelerator-wide flag remains disabled; JSD uses its own checked Intel backend dispatch. The final TP cases cover FP32/BF16 interior JSD versus an FP64 oracle, empty shards/disjoint support, and head SGD update parity. On NVIDIA L20, driver 580.173.02, torch 2.14.1+cu130 (CUDA 13.0), Triton 3.8.0, NCCL: the **full test file passed, 113 passed with no skips**. This includes all JSD TP1/2/4 cases and FP32/FP16/BF16 single-device kernels. The existing `ds` conda environment was used without dependency changes; tested source hashes match the local implementation and tests. ```bash PYTHONPATH=tests:. OMP_NUM_THREADS=4 pytest \ tests/unit/v1/sequence_parallelism/test_vocab_parallel_ops.py -q \ -k 'not TestVocabParallel and not triton' PYTHONPATH=tests:. OMP_NUM_THREADS=4 pytest \ tests/unit/v1/sequence_parallelism/test_vocab_parallel_ops.py -q -k jsd_triton CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 OMP_NUM_THREADS=4 PYTHONPATH=tests:. \ pytest --forked tests/unit/v1/sequence_parallelism/test_vocab_parallel_ops.py -q \ -k 'TestVocabParallelJSDTP and (0.2 or empty or optimizer)' pre-commit run --files deepspeed/ops/triton_ops/vocab_parallel_softmax.py \ deepspeed/sequence/vocab_parallel_ops.py \ tests/unit/v1/sequence_parallelism/test_vocab_parallel_ops.py \ docs/_tutorials/autotp-training.md git diff --check # L20 full-file run: CUDA_VISIBLE_DEVICES=0,1,2,3 PYTHONPATH=tests:. OMP_NUM_THREADS=4 \ pytest --forked tests/unit/v1/sequence_parallelism/test_vocab_parallel_ops.py \ -q --color=no ``` Scoped pre-commit hooks and diff checks pass. Numerical checks use an independent FP64 dense PyTorch reference, nonuniform vocabulary partitions, FP32/BF16, both-input gradients and frozen inputs. Additional cases cover disjoint support, undefined masked rows, empty local shards, empty token batches, causal-shift views, independent large logit shifts and row chunking. Three SGD updates of a sharded head match the dense reference at TP1/2/4. This head-only training test does not exercise a DeepSpeed Engine or ZeRO. Exploratory logits-only forward+backward A/B on B70, N=1024, V=151936, BF16, beta=0.3, frozen teacher: | TP | Implementation | Peak allocated GiB/rank | Time ms, median | |---|---|---:|---:| | 2 | Gather both logits, dense PyTorch JSD | 6.671 | 112.17 | | 2 | Vocabulary-parallel eager JSD | 1.057 | 58.38 | | 2 | Vocabulary-parallel Triton JSD | 0.435 | 3.08 | | 2 | Existing Triton reverse KL | 0.435 | 3.29 | | 4 | Vocabulary-parallel eager JSD | 0.841 | 31.14 | | 4 | Vocabulary-parallel Triton JSD | 0.217 | 3.41 | | 4 | Existing Triton reverse KL | 0.217 | 3.16 | Five warmups and ten measured trials, synchronized wall time; per-trial MAX over TP ranks, independently for time and memory. Resident input allocations are included in peak. Student gradients are included. All arms start with sharded logits, and the dense arm includes their gather; it is not TRL's existing chunked-projection implementation. Losses and local student gradients match before timing. Clocks were not pinned. This benchmark excludes head projection, backbone, optimizer and rollout; it demonstrates only an op-level comparison, not full-training or TRL speedup. Reproduction script and raw JSON/logs are retained locally and should be attached to the PR when submitted. KL's Triton kernels are enabled only in the XPU benchmark for comparison; their production dispatch is unchanged. L20, same workload and protocol: | TP | Implementation | Peak allocated GiB/rank | Time ms, median | |---|---|---:|---:| | 2 | Gather both logits, dense PyTorch JSD | 6.671 | 98.02 | | 2 | Vocabulary-parallel eager JSD | 1.057 | 34.43 | | 2 | Vocabulary-parallel Triton JSD | 0.435 | 1.95 | | 2 | Existing Triton reverse KL | 0.435 | 2.21 | | 4 | Gather both logits, dense PyTorch JSD | 6.526 | 103.41 | | 4 | Vocabulary-parallel eager JSD | 0.841 | 16.78 | | 4 | Vocabulary-parallel Triton JSD | 0.217 | 1.57 | | 4 | Existing Triton reverse KL | 0.217 | 1.82 | L20 clocks were also not pinned; these remain exploratory op-level timings. All 15 CUDA sm89 offline kernel compilations pass (FP32/FP16/BF16 and both/ student-only/teacher-only backward variants), in addition to the actual L20 runtime validation above. Not validated: ZeRO stages, offload, PP/SP, DeepCompile, real-model convergence or TRL integration. No end-to-end speedup or convergence claim is made. ## AI assistance Implementation, test preparation and this description were prepared with GitHub Copilot CLI. Human review and independently run tests are required before submission. --------- Signed-off-by: Jin, Youzhi <youzhi.jin@intel.com> Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Co-authored-by: Ma, Guokai <guokai.ma@gmail.com>
tohtana
added
needs-gpu-test
Follow-up to #8791 (RFC #8790). The vocab-parallel logprobs / entropy / KL ops ran each row chunk as a dozen separate eager fp32 kernels. On the same local shard they were about 2–2.4× slower than a dense
log_softmax, which offsets most of the memory benefit when they are used for OPSD / RL losses (e.g. deepspeedai/DeepSpeedExamples#1019).Changes
deepspeed/ops/triton_ops/vocab_parallel_softmax.py: Triton kernels with one program per row, looping over the local vocabulary in blocks. The module has no communication.deepspeed/sequence/vocab_parallel_ops.pydispatches to these kernels on CUDA/XPU when the inputs are contiguous anddeepspeed.HAS_TRITONis true. Otherwise it runs the existing eager path, which is unchanged. Collectives and the autograd structure are the same.autotp-training.md.Performance
Intel Arc Pro B70, bf16, 75968-column shard (Qwen vocab / TP2), forward + backward:
torch.compileof the eager chunk functions reached 26.9 ms for reverse KL (N=8192), so I chose hand-written kernels.Tests
test_triton_kernels_match_eager_path(fp32 / bf16 × logprobs / entropy / forward KL / reverse KL). On an accelerator it compares the Triton outputs and gradients with the CPU eager path. Cases: multiple kernel blocks, partially masked vocabulary, all--infrows, support mismatch with underflowing probability, large offsets, duplicate / ignored top-k ids, and a fully ignored row.infis removed.Executed on 2× Intel Arc Pro B70 (XPU), torch 2.14.0.dev+xpu, Triton 3.8.0:
tests/unit/v1/sequence_parallelism/test_vocab_parallel_ops.py: 41 passed.Update: single-pass entropy/KL forward (by @0z5a)
The tables above were measured on the first commit. @0z5a's
e8038bed3reads each logit once in the entropy/KL forward, using online weighted statistics. Both commits were re-measured with the same script on 1× B70: bf16, a 75968-column shard, forward + backward, median of 20 runs.99f05419297d4f66dfExtra peak memory, excluding inputs and gradients, is below 1 MiB on both commits.
Numerics vs the first commit. Compared on 1× and 2× B70 (TP=1, TP=2), fp32 and bf16. Cases: 512×32000, 256×151936, and 64×32000 with a +1000 offset and
-infmasks. Each output and gradient was also compared with an fp64 CPU reference.logprobs: outputs and gradients are bitwise identical.entropy/ KL: not bitwise identical, because the summation order changed. Forward outputs differ by at most 35 ULP. The forward is closer to fp64 on almost every case: e.g. fp32 entropy relative error drops from 3e-7 to 1.5e-7, and KL from about 6–9e-7 to 3–8e-7. Gradient errors stay at the same level, and bf16 gradient errors are unchanged (about 2e-3).inf/NaNplacement differences.Contribution checklist
Related work. This is a performance follow-up to #8791 (RFC #8790). A search for vocab-parallel Triton kernels and entropy/KL found no overlapping PR.
Commands and results. Hardware and software: 2× Intel Arc Pro B70 (XPU, XCCL), Level Zero driver 1.15.39122 (compute runtime 26.27.39122.11), torch 2.14.0.dev20260804+xpu, Triton 3.8.0.
At head
8f463213f: 49 passed on 2× B70; withDS_ACCELERATOR=cpu, 33 passed and the 16 Triton tests are skipped.Compatibility matrix.
row_layoutstrides frome8038bed3ZeRO and offload do not interact with these ops.
AI assistance. This PR was developed with AI assistance (GitHub Copilot CLI). I reviewed every changed line and ran the tests above.