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>
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.