Repository navigation
Add vocabulary-parallel fused linear cross entropy - #8792
Conversation
Add deepspeed.sequence.linear_cross_entropy with vocab_parallel_linear_cross_entropy, which fuses the vocabulary-parallel LM-head projection into the cross-entropy loss. Tokens are processed in row chunks so only one [chunk, vocab / tp] logits block is alive at a time, and the input, weight and bias gradients are accumulated during forward. The hidden gradient is all-reduced over the TP group in backward, matching the column-parallel head. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Jin, Youzhi <youzhi.jin@intel.com>
…umulation in FLCE - Derive the default chunk size from the global vocab size so ranks with uneven shards issue the same collectives. - Scale gradients in fp32 before casting, so fp16 loss scales do not overflow. - Run the matmuls in the autocast dtype and return gradients in input dtypes. - Accumulate dW/db in fp32 so results do not depend on the chunk count. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Jin, Youzhi <youzhi.jin@intel.com>
Each chunk reads and writes the fp32 weight-gradient buffer. With large vocab shards the default chunk was small enough for that traffic to dominate, making 256K vocab / TP2 up to 2x slower than the unfused baseline. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Jin, Youzhi <youzhi.jin@intel.com>
Computing gradients during forward over token chunks needed an fp32 weight-gradient accumulator plus a cast copy, about 3x the weight shard. Forward now keeps only per-token softmax statistics, and backward recomputes the logits one vocabulary block at a time, so each weight-gradient block is a single GEMM over all tokens. This also gives an ordinary backward and supports reduction="none". Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Jin, Youzhi <youzhi.jin@intel.com>
Without an fp32-output GEMM, the low-precision product could overflow before a cancelling bias was added, giving NaN where F.linear stays finite. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Jin, Youzhi <youzhi.jin@intel.com>
jinyouzhi
changed the title
Upcast the fallback hidden-gradient block GEMM to fp32: a single fp16 block product can overflow even when the sum over all vocabulary blocks is finite, which turned cancelling contributions into NaN. Look up the autocast dtype through the per-device getters when torch.get_autocast_dtype is missing (PyTorch < 2.4), since DeepSpeed still supports torch >= 2.0. 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: 74ae27cb45
ℹ️ 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".
| _accumulate_matmul(grad_hidden, grad_logits, weight_block) | ||
|
|
||
| if grad_hidden is not None: | ||
| grad_hidden = grad_hidden.to(hidden.dtype) |
There was a problem hiding this comment.
When FP16 is used with TP and loss scaling makes the per-shard input-gradient partials exceed 65504, this narrows each rank's FP32 accumulator before the SUM. For example, with a 65536 loss scale, two equal-probability classes on different TP ranks, and weights of 8, the ranks produce opposing FP32 partials whose correct sum is zero, but they become +inf and -inf here and the all-reduce returns NaN. Sum the FP32 partials before converting to hidden.dtype so cancellation across vocabulary shards remains finite.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Good catch, fixed in 1228ea0. The hidden-gradient partials are now all-reduced in FP32 over the TP group and cast to hidden.dtype afterwards, so partials that overflow FP16 but cancel across vocabulary shards stay finite. TestVocabParallelLinearCrossEntropyTP::test_fp16_hidden_gradient_partials_that_cancel_across_ranks_stay_finite reproduces your example (65536 loss scale, two classes on different TP ranks, weights of 8): it returns NaN before this change and matches the full-vocabulary reference after it.
Each TP rank's partial input gradient can overflow fp16 under loss scaling even when the partials cancel across vocabulary shards, so narrowing before the all-reduce turned a zero gradient into NaN. Reduce the fp32 accumulator first and cast afterwards. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Jin, Youzhi <youzhi.jin@intel.com>
| return torch.get_autocast_dtype(device_type) | ||
| return None | ||
| # PyTorch < 2.4 only has per-device getters, and CUDA uses the unsuffixed ones. | ||
| if device_type == "cuda": |
There was a problem hiding this comment.
why cuda device have to be treated differently?
There was a problem hiding this comment.
I see. I think this piece of code can be here since its for old version of pytorch.
| else: | ||
| return torch.mm(hidden, weight.t(), out_dtype=torch.float32) | ||
| return logits if row_offset is None else logits.add_(row_offset) | ||
| except _MIXED_GEMM_ERRORS: |
There was a problem hiding this comment.
This looks a little bit unstable. If torch.addmm or torch.mm failed for some random reason (i.e. OOM), it might be recognized as dtype not supported. Could device capability be declared in accelerator interface?
There was a problem hiding this comment.
May I introduce new interface into abstract_accelerator?
There was a problem hiding this comment.
Addressed in 8582e19 without a new accelerator interface. Instead of attempting a GEMM and treating any exception as "unsupported" (which could misclassify an OOM), _has_mixed_gemm(device_type) now asks the PyTorch dispatcher whether the device registers kernels for aten::mm.dtype, aten::addmm.dtype and aten::addmm.dtype_out. No computation or allocation happens, and the result is cached per device type.
- Upstream registers these ops for CUDA (incl. ROCm) and XPU; CPU returns False and takes the upcast path.
- PyTorch < 2.8 (ops absent) or a build without the
torch._Cquery helpers falls back to the upcast path, so the result stays correct. - The fast path is now only taken for fp16/bf16 inputs with an fp32 output, matching the kernels' dtype contract. This also removes the old behavior where an fp64 input could permanently disable the fast path for a device.
- No try/except around the GEMM anymore, so real failures (e.g. OOM) propagate.
There was a problem hiding this comment.
Will approve. Use pytorch internal API might be fragile but this is very local so should be okay.
…ed GEMM Catching an exception from the first mixed-dtype GEMM also caught failures unrelated to dtype support (e.g. out of memory) and then disabled the fast path for the device for the rest of the process. Ask the dispatcher whether the device registers mm.dtype / addmm.dtype kernels instead, and only use them for the fp16/bf16 -> fp32 case they implement. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Jin, Youzhi <youzhi.jin@intel.com>
2 participants
Follow-up to #8173.
Motivation
After #8173, a vocabulary-parallel LM head still materializes
[tokens, vocab / tp]logits on every rank. Its CE backward also keeps FP32 softmax buffers of the same size. For long sequences and large vocabularies, this is the largest activation in the step. For example, at 256K vocab, TP2, and 16K tokens, the loss path alone peaks at about 28 GiB per rank.Existing fused linear CE kernels do not cover this case. Liger's FLCE is not TP-aware (linkedin/Liger-Kernel#826, linkedin/Liger-Kernel#868). PyTorch's
F.linear_cross_entropychunks only over the batch dimension and takes no process group.What
deepspeed.sequence.linear_cross_entropy.vocab_parallel_linear_cross_entropy(hidden, weight, target, bias=None, tp_group=None, sp_group=None, vocab_start_index=None, vocab_end_index=None, ignore_index=-100, reduction="mean", chunk_size=None)(reductionis"none","sum"or"mean")head(hidden)plus the vocab-parallel CE, takingweight/biasfromVocabParallelLinear. Uneven shards are supported.h @ W_localᵀstraight into FP32 and runs one MAX all-reduce plus one SUM all-reduce (sum-exp and target logit packed together). Only the per-token max and sum-exp are kept.dWblock then comes from one GEMM over all tokens.dhis accumulated in FP32 and all-reduced over TP, matching the column-parallel head's backward. Backward needs no other collectives.dWis held, and backward is an ordinary autograd backward. The cost is one extra logits GEMM.hidden.chunk_sizetrades memory for speed.Validation
tests/unit/v1/sequence_parallelism/test_vocab_parallel_linear_cross_entropy.pyhas 24 tests. Values and gradients are checked againstF.linear+F.cross_entropyon the full vocabulary. Coverage:none/sum/meanreduction with a non-uniform upstream gradient, with and without bias, and forced chunking;F.linearhandles it;Ran on Intel Arc Pro B70 (XPU, XCCL): 30 passed.
test_vocab_parallel_cross_entropy.pystill passes (18 passed).Memory and time
Measured on 1× Intel Arc Pro B70, single-rank shard, H=4096, BF16. Peak is fwd+bwd and excludes inputs, but includes
dW/dh. The baseline ishead.weightlogits followed byvocab_parallel_cross_entropy.linear_cross_entropy(no TP, reference)Peak memory drops by 3–15×, and time is 4–18% higher because backward recomputes the logits. The first design computed gradients during forward and accumulated
dWin FP32 over token chunks. That needed an FP32dWbuffer plus a cast copy, which came to about 3× the size ofdW: 3.6 GiB peak at 256K / TP2.Follow-ups
loss_function, so the existing loss hook cannot intercept them. The options are to patch the CausalLM forward behindvocab_parallel_ce_backend: "fused_linear", or to document the functional API only.VocabParallelLinearandSimpleModel, plus multi-rank e2e memory numbers.torch.linear_cross_entropy, e.g. by fusing the softmax-gradient elementwise ops.Limitations
out_dtype=torch.float32) are used where the backend supports them, such as XPU. Other backends, such as CPU, fall back to unfused ops.torch.autocast, the GEMMs run in the autocast dtype, asF.linearwould.cc @delock @PKUWZP