From 15760d784752185ebdd3b3991d8e644661e43c11 Mon Sep 17 00:00:00 2001
From: "Jin, Youzhi" <youzhi.jin@intel.com>
Date: Fri, 9 Oct 2026 17:02:21 +0000
Subject: [PATCH 1/8] Add vocabulary-parallel fused linear cross entropy

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>
---
 deepspeed/sequence/linear_cross_entropy.py    | 165 ++++++++++++++++++
 docs/_tutorials/autotp-training.md            |  29 +++
 ...est_vocab_parallel_linear_cross_entropy.py | 128 ++++++++++++++
 3 files changed, 322 insertions(+)
 create mode 100644 deepspeed/sequence/linear_cross_entropy.py
 create mode 100644 tests/unit/v1/sequence_parallelism/test_vocab_parallel_linear_cross_entropy.py

diff --git a/deepspeed/sequence/linear_cross_entropy.py b/deepspeed/sequence/linear_cross_entropy.py
new file mode 100644
index 000000000000..3a51aa6a27df
--- /dev/null
+++ b/deepspeed/sequence/linear_cross_entropy.py
@@ -0,0 +1,165 @@
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+"""Fused LM-head projection and cross entropy over a vocabulary-parallel weight.
+
+A vocabulary-parallel LM head still materializes ``[tokens, vocab / tp]`` logits per rank.
+This module fuses the projection into the loss: tokens are processed in row chunks, so only
+one ``[chunk, vocab / tp]`` logits block is alive at a time, and the input and weight
+gradients are accumulated during forward.
+"""
+
+import math
+
+import torch
+from torch.autograd.function import once_differentiable
+
+import deepspeed.comm as dist
+from deepspeed.sequence.cross_entropy import _global_sp_sum, _resolve_vocab_metadata
+
+
+def _tp_active(tp_group):
+    return tp_group is not None and dist.get_world_size(tp_group) > 1
+
+
+def _default_chunk_rows(num_tokens, local_vocab_size, hidden_size):
+    # Keep each fp32 logits block about the size of the hidden states, which bounds the
+    # extra peak memory independently of the vocabulary size.
+    num_chunks = max(1, math.ceil(local_vocab_size / hidden_size))
+    return max(1, math.ceil(num_tokens / num_chunks))
+
+
+class _VocabParallelLinearCrossEntropy(torch.autograd.Function):
+
+    @staticmethod
+    def forward(ctx, hidden, weight, bias, target, tp_group, vocab_start_index, vocab_end_index, ignore_index,
+                grad_scale, chunk_rows):
+        num_tokens = hidden.shape[0]
+        local_vocab_size = weight.shape[0]
+        tp_active = _tp_active(tp_group)
+
+        loss = torch.zeros(num_tokens, dtype=torch.float32, device=hidden.device)
+        grad_hidden = torch.empty_like(hidden)
+        grad_weight = torch.zeros_like(weight) if weight.requires_grad else None
+        grad_bias = torch.zeros_like(bias) if bias is not None and bias.requires_grad else None
+
+        for start in range(0, num_tokens, chunk_rows):
+            end = min(start + chunk_rows, num_tokens)
+            hidden_chunk = hidden[start:end]
+            target_chunk = target[start:end]
+
+            logits = hidden_chunk @ weight.t()
+            if bias is not None:
+                logits = logits + bias
+            logits = logits.float()
+
+            valid = target_chunk != ignore_index
+            in_shard = valid & (target_chunk >= vocab_start_index) & (target_chunk < vocab_end_index)
+            local_target = (target_chunk - vocab_start_index).clamp(0, local_vocab_size - 1)
+
+            row_max = logits.amax(dim=-1)
+            if tp_active:
+                dist.all_reduce(row_max, op=dist.ReduceOp.MAX, group=tp_group)
+            logits.sub_(row_max.unsqueeze(-1))
+            target_logit = logits.gather(-1, local_target.unsqueeze(-1)).squeeze(-1)
+            target_logit = torch.where(in_shard, target_logit, torch.zeros_like(target_logit))
+            # Reuse the logits buffer for the softmax so only one fp32 block is alive per chunk.
+            probs = logits.exp_()
+            stats = torch.stack([probs.sum(dim=-1), target_logit], dim=-1)
+            if tp_active:
+                dist.all_reduce(stats, op=dist.ReduceOp.SUM, group=tp_group)
+            sum_exp, target_logit = stats.unbind(-1)
+
+            loss[start:end] = torch.where(valid, sum_exp.log() - target_logit, torch.zeros_like(sum_exp))
+
+            # d(loss) / d(logits) = softmax - onehot, zeroed for ignored tokens and pre-scaled by
+            # the reduction factor so backward only has to multiply by the incoming gradient.
+            row_scale = valid.to(probs.dtype) * grad_scale / sum_exp
+            grad_logits = probs.mul_(row_scale.unsqueeze(-1))
+            grad_logits.scatter_add_(-1, local_target.unsqueeze(-1),
+                                     -(in_shard.to(probs.dtype) * grad_scale).unsqueeze(-1))
+            grad_logits = grad_logits.to(hidden.dtype)
+
+            torch.matmul(grad_logits, weight, out=grad_hidden[start:end])
+            if grad_weight is not None:
+                grad_weight.addmm_(grad_logits.t(), hidden_chunk)
+            if grad_bias is not None:
+                grad_bias.add_(grad_logits.sum(dim=0).to(grad_bias.dtype))
+
+        ctx.save_for_backward(grad_hidden, grad_weight, grad_bias)
+        ctx.tp_group = tp_group
+        return loss.sum() * grad_scale
+
+    @staticmethod
+    @once_differentiable
+    def backward(ctx, grad_output):
+        grad_hidden, grad_weight, grad_bias = ctx.saved_tensors
+        # The precomputed gradients assume a unit upstream gradient; rescale only when needed
+        # so the common ``loss.backward()`` path does not allocate another copy.
+        if not torch.equal(grad_output, torch.ones_like(grad_output)):
+            grad_hidden = grad_hidden * grad_output.to(grad_hidden.dtype)
+            grad_weight = grad_weight * grad_output.to(grad_weight.dtype) if grad_weight is not None else None
+            grad_bias = grad_bias * grad_output.to(grad_bias.dtype) if grad_bias is not None else None
+        # Each rank only saw its vocabulary shard, so the input gradient is a partial sum, the
+        # same reduction the column-parallel LM head performs in its own backward.
+        if _tp_active(ctx.tp_group):
+            dist.all_reduce(grad_hidden, op=dist.ReduceOp.SUM, group=ctx.tp_group)
+        return grad_hidden, grad_weight, grad_bias, None, None, None, None, None, None, None
+
+
+def 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):
+    """Cross entropy of ``hidden @ weight.T (+ bias)`` without materializing the logits.
+
+    ``weight`` (and ``bias``) hold this rank's rows of a vocabulary-parallel LM head, for
+    example ``VocabParallelLinear.weight``. ``hidden`` must be identical on every TP rank, as
+    for a column-parallel layer; its gradient is all-reduced over ``tp_group`` in backward.
+
+    Gradients are computed during forward, so only a single first-order backward is
+    supported. ``chunk_size`` sets the number of tokens per chunk; by default each fp32
+    logits chunk is about as large as ``hidden``.
+    """
+    if reduction not in ("sum", "mean"):
+        raise ValueError(f"Unsupported reduction: {reduction}")
+    if hidden.shape[:-1] != target.shape:
+        raise ValueError("hidden and target must have matching non-hidden dimensions")
+    if weight.dim() != 2 or weight.shape[1] != hidden.shape[-1]:
+        raise ValueError("weight must have shape [local_vocab, hidden]")
+    if (vocab_start_index is None) != (vocab_end_index is None):
+        raise ValueError("vocab_start_index and vocab_end_index must be provided together")
+    if chunk_size is not None and chunk_size <= 0:
+        raise ValueError("chunk_size must be positive")
+
+    vocab_start_index, vocab_end_index, global_vocab_size = _resolve_vocab_metadata(
+        weight.shape[0], vocab_start_index, vocab_end_index, tp_group, hidden.device)
+    target = target.reshape(-1).to(dtype=torch.long)
+    invalid_target = (target != ignore_index) & ((target < 0) | (target >= global_vocab_size))
+    if invalid_target.any().item():
+        raise ValueError(f"Target is out of range for vocabulary size {global_vocab_size}")
+
+    hidden = hidden.reshape(-1, hidden.shape[-1]).contiguous()
+    weight = weight.contiguous()
+    num_tokens = hidden.shape[0]
+
+    if reduction == "mean":
+        valid_tokens = (target != ignore_index).sum().to(torch.float32)
+        if sp_group is not None and dist.get_world_size(sp_group) > 1:
+            dist.all_reduce(valid_tokens, op=dist.ReduceOp.SUM, group=sp_group)
+        grad_scale = 1.0 / max(valid_tokens.item(), 1.0)
+    else:
+        grad_scale = 1.0
+    if chunk_size is None:
+        chunk_size = _default_chunk_rows(num_tokens, weight.shape[0], hidden.shape[1])
+
+    loss = _VocabParallelLinearCrossEntropy.apply(hidden, weight, bias, target, tp_group, vocab_start_index,
+                                                  vocab_end_index, ignore_index, grad_scale, chunk_size)
+    return _global_sp_sum(loss, sp_group)
diff --git a/docs/_tutorials/autotp-training.md b/docs/_tutorials/autotp-training.md
index 297f47a87376..4406fe4518a4 100644
--- a/docs/_tutorials/autotp-training.md
+++ b/docs/_tutorials/autotp-training.md
@@ -271,6 +271,35 @@ each shard's mean itself, weighted by the shard's valid-token count, so an
 additional SP reduction would double-count tokens. Pass an `sp_group` only when
 this loss is the sole aggregation over a manually constructed TP x SP mesh.
 
+### Fused Projection and Loss
+
+The vocabulary-parallel loss above still materializes `[tokens, vocab / tp]`
+logits on every rank. `vocab_parallel_linear_cross_entropy` fuses the LM-head
+projection into the loss and processes tokens in chunks, so only one
+`[chunk, vocab / tp]` block is alive at a time:
+
+```python
+from deepspeed.sequence.linear_cross_entropy import vocab_parallel_linear_cross_entropy
+
+head = model.lm_head  # VocabParallelLinear
+hidden = model.model(input_ids).last_hidden_state[:, :-1]
+loss = vocab_parallel_linear_cross_entropy(hidden,
+                                           head.weight,
+                                           labels[:, 1:],
+                                           bias=head.bias,
+                                           tp_group=head.mp_group,
+                                           vocab_start_index=head.vocab_start_index,
+                                           vocab_end_index=head.vocab_end_index)
+```
+
+The call replaces both `head(hidden)` and the loss: it all-reduces the hidden
+gradient over the TP group in backward, as the column-parallel head would. It
+supports `sum` and `mean` reductions and computes gradients during forward, so
+it allows only first-order backward. `chunk_size` (tokens per chunk) trades
+peak memory for speed; by default each FP32 logits chunk is about the size of
+the hidden states. This API is not yet wired into the HuggingFace forward or
+the `vocab_parallel_lm_head` configuration.
+
 
 ## Custom Patterns
 
diff --git a/tests/unit/v1/sequence_parallelism/test_vocab_parallel_linear_cross_entropy.py b/tests/unit/v1/sequence_parallelism/test_vocab_parallel_linear_cross_entropy.py
new file mode 100644
index 000000000000..a3eb4115b009
--- /dev/null
+++ b/tests/unit/v1/sequence_parallelism/test_vocab_parallel_linear_cross_entropy.py
@@ -0,0 +1,128 @@
+# SPDX-License-Identifier: Apache-2.0
+# DeepSpeed Team
+
+import pytest
+import torch
+import torch.nn.functional as F
+
+import deepspeed.comm as dist
+from deepspeed.accelerator import get_accelerator
+from deepspeed.module_inject.tp_shard import AutoTPMeta, get_shard_size_list
+from deepspeed.sequence.linear_cross_entropy import vocab_parallel_linear_cross_entropy
+from unit.common import DistributedTest
+
+
+def _inputs(vocab_size=17, hidden_size=8, device="cpu", dtype=torch.float32, bias=True):
+    torch.manual_seed(7)
+    hidden = torch.randn(2, 5, hidden_size, device=device, dtype=dtype)
+    weight = torch.randn(vocab_size, hidden_size, device=device, dtype=dtype)
+    bias_tensor = torch.randn(vocab_size, device=device, dtype=dtype) if bias else None
+    target = torch.randint(0, vocab_size, (2, 5), device=device)
+    target[0, 1] = -100
+    target[1, 4] = -100
+    return hidden, weight, bias_tensor, target
+
+
+def _leaf(tensor):
+    return None if tensor is None else tensor.detach().clone().requires_grad_(True)
+
+
+def _reference(hidden, weight, bias, target, reduction):
+    logits = F.linear(hidden.float(), weight.float(), None if bias is None else bias.float())
+    return F.cross_entropy(logits.view(-1, weight.shape[0]), target.view(-1), reduction=reduction)
+
+
+@pytest.mark.parametrize("reduction", ["sum", "mean"])
+@pytest.mark.parametrize("use_bias", [True, False])
+@pytest.mark.parametrize("chunk_size", [None, 1, 3])
+def test_matches_linear_then_cross_entropy(reduction, use_bias, chunk_size):
+    hidden, weight, bias, target = _inputs(bias=use_bias)
+    ref_args = [_leaf(hidden), _leaf(weight), _leaf(bias)]
+    args = [_leaf(hidden), _leaf(weight), _leaf(bias)]
+
+    expected = _reference(*ref_args, target, reduction)
+    actual = vocab_parallel_linear_cross_entropy(args[0],
+                                                 args[1],
+                                                 target,
+                                                 bias=args[2],
+                                                 reduction=reduction,
+                                                 chunk_size=chunk_size)
+    torch.testing.assert_close(actual, expected)
+
+    # A non-unit upstream gradient must rescale the gradients precomputed in forward.
+    (actual * 3).backward()
+    (expected * 3).backward()
+    for got, want in zip(args, ref_args):
+        if want is not None:
+            torch.testing.assert_close(got.grad, want.grad)
+
+
+def test_all_ignored_is_zero():
+    hidden, weight, _, target = _inputs(bias=False)
+    hidden, weight = _leaf(hidden), _leaf(weight)
+    loss = vocab_parallel_linear_cross_entropy(hidden, weight, torch.full_like(target, -100))
+
+    torch.testing.assert_close(loss, torch.zeros_like(loss))
+    loss.backward()
+    torch.testing.assert_close(hidden.grad, torch.zeros_like(hidden))
+    torch.testing.assert_close(weight.grad, torch.zeros_like(weight))
+
+
+def test_frozen_weight_only_produces_hidden_grad():
+    hidden, weight, _, target = _inputs(bias=False)
+    hidden = _leaf(hidden)
+    vocab_parallel_linear_cross_entropy(hidden, weight, target).backward()
+
+    reference_hidden = _leaf(hidden)
+    _reference(reference_hidden, weight, None, target, "mean").backward()
+    torch.testing.assert_close(hidden.grad, reference_hidden.grad)
+    assert weight.grad is None
+
+
+def test_validates_inputs():
+    hidden, weight, _, target = _inputs(bias=False)
+    with pytest.raises(ValueError, match="Unsupported reduction"):
+        vocab_parallel_linear_cross_entropy(hidden, weight, target, reduction="none")
+    with pytest.raises(ValueError, match="matching non-hidden dimensions"):
+        vocab_parallel_linear_cross_entropy(hidden, weight, target[:, :-1])
+    with pytest.raises(ValueError, match="local_vocab, hidden"):
+        vocab_parallel_linear_cross_entropy(hidden, weight[:, :-1], target)
+    with pytest.raises(ValueError, match="chunk_size must be positive"):
+        vocab_parallel_linear_cross_entropy(hidden, weight, target, chunk_size=0)
+    with pytest.raises(ValueError, match="out of range"):
+        vocab_parallel_linear_cross_entropy(hidden, weight, torch.full_like(target, 17))
+
+
+class TestVocabParallelLinearCrossEntropyTP(DistributedTest):
+    world_size = 2
+
+    @pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16])
+    def test_uneven_vocab_matches_full_vocab(self, dtype):
+        device = torch.device(get_accelerator().current_device_name())
+        rank = dist.get_rank()
+        vocab_size = 17
+        partition_sizes = get_shard_size_list(vocab_size, self.world_size, AutoTPMeta(), "lm_head")
+        start = sum(partition_sizes[:rank])
+        end = start + partition_sizes[rank]
+
+        hidden, weight, bias, target = _inputs(vocab_size, device=device, dtype=dtype)
+        ref_args = [_leaf(hidden), _leaf(weight), _leaf(bias)]
+        args = [_leaf(hidden), _leaf(weight[start:end]), _leaf(bias[start:end])]
+
+        expected = _reference(*ref_args, target, "mean")
+        actual = vocab_parallel_linear_cross_entropy(args[0],
+                                                     args[1],
+                                                     target,
+                                                     bias=args[2],
+                                                     tp_group=dist.get_world_group(),
+                                                     vocab_start_index=start,
+                                                     vocab_end_index=end,
+                                                     chunk_size=3)
+        actual.backward()
+        expected.backward()
+
+        tolerance = {} if dtype == torch.float32 else dict(atol=2e-2, rtol=2e-2)
+        torch.testing.assert_close(actual, expected, **tolerance)
+        torch.testing.assert_close(args[0].grad, ref_args[0].grad, **tolerance)
+        torch.testing.assert_close(args[1].grad, ref_args[1].grad[start:end], **tolerance)
+        torch.testing.assert_close(args[2].grad, ref_args[2].grad[start:end], **tolerance)

From 9428dbf32ce3b14f14a48397a74eab9429b9a330 Mon Sep 17 00:00:00 2001
From: "Jin, Youzhi" <youzhi.jin@intel.com>
Date: Sat, 10 Oct 2026 09:20:47 +0000
Subject: [PATCH 2/8] Fix uneven-shard chunking, loss-scale overflow, autocast
 and bf16 accumulation 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>
---
 deepspeed/sequence/linear_cross_entropy.py    | 94 +++++++++++++++----
 ...est_vocab_parallel_linear_cross_entropy.py | 57 ++++++++++-
 2 files changed, 129 insertions(+), 22 deletions(-)

diff --git a/deepspeed/sequence/linear_cross_entropy.py b/deepspeed/sequence/linear_cross_entropy.py
index 3a51aa6a27df..81086e7091a2 100644
--- a/deepspeed/sequence/linear_cross_entropy.py
+++ b/deepspeed/sequence/linear_cross_entropy.py
@@ -22,26 +22,62 @@ def _tp_active(tp_group):
     return tp_group is not None and dist.get_world_size(tp_group) > 1
 
 
-def _default_chunk_rows(num_tokens, local_vocab_size, hidden_size):
+def _default_chunk_rows(num_tokens, global_vocab_size, tp_world_size, hidden_size):
     # Keep each fp32 logits block about the size of the hidden states, which bounds the
-    # extra peak memory independently of the vocabulary size.
-    num_chunks = max(1, math.ceil(local_vocab_size / hidden_size))
+    # extra peak memory independently of the vocabulary size. Derived from the global shard
+    # layout rather than the local shard so that every TP rank issues the same collectives
+    # even when the vocabulary is split unevenly.
+    max_local_vocab_size = math.ceil(global_vocab_size / tp_world_size)
+    num_chunks = max(1, math.ceil(max_local_vocab_size / hidden_size))
     return max(1, math.ceil(num_tokens / num_chunks))
 
 
+_MIXED_ADDMM_UNSUPPORTED = set()
+
+
+def _accumulate_matmul(accumulator, mat1, mat2):
+    """``accumulator += mat1 @ mat2`` with ``accumulator`` kept in its own (fp32) dtype.
+
+    Accumulating chunk gradients in a low-precision buffer drops later contributions once the
+    sum grows, so the gradient would depend on the chunk size.
+    """
+    if mat1.dtype == accumulator.dtype:
+        accumulator.addmm_(mat1, mat2)
+        return
+    device_type = accumulator.device.type
+    if device_type not in _MIXED_ADDMM_UNSUPPORTED:
+        try:
+            torch.addmm(accumulator, mat1, mat2, out_dtype=accumulator.dtype, out=accumulator)
+            return
+        except (RuntimeError, TypeError, NotImplementedError):
+            # Older PyTorch or backends without a mixed-precision GEMM (e.g. CPU).
+            _MIXED_ADDMM_UNSUPPORTED.add(device_type)
+    accumulator.add_(torch.mm(mat1, mat2))
+
+
 class _VocabParallelLinearCrossEntropy(torch.autograd.Function):
 
     @staticmethod
     def forward(ctx, hidden, weight, bias, target, tp_group, vocab_start_index, vocab_end_index, ignore_index,
-                grad_scale, chunk_rows):
+                grad_scale, chunk_rows, compute_dtype):
         num_tokens = hidden.shape[0]
         local_vocab_size = weight.shape[0]
         tp_active = _tp_active(tp_group)
+        ctx.input_dtypes = (hidden.dtype, weight.dtype, None if bias is None else bias.dtype)
+
+        # Mirror what autocast would do for ``F.linear``: run the GEMMs in one compute dtype even
+        # when activations and parameters differ (e.g. bf16 activations with fp32 weights).
+        weight_requires_grad = weight.requires_grad
+        bias_requires_grad = bias is not None and bias.requires_grad
+        hidden = hidden.to(compute_dtype)
+        weight = weight.to(compute_dtype)
+        bias = None if bias is None else bias.to(compute_dtype)
 
         loss = torch.zeros(num_tokens, dtype=torch.float32, device=hidden.device)
         grad_hidden = torch.empty_like(hidden)
-        grad_weight = torch.zeros_like(weight) if weight.requires_grad else None
-        grad_bias = torch.zeros_like(bias) if bias is not None and bias.requires_grad else None
+        # Parameter gradients sum over every chunk, so they are accumulated in fp32.
+        grad_weight = torch.zeros_like(weight, dtype=torch.float32) if weight_requires_grad else None
+        grad_bias = torch.zeros_like(bias, dtype=torch.float32) if bias_requires_grad else None
 
         for start in range(0, num_tokens, chunk_rows):
             end = min(start + chunk_rows, num_tokens)
@@ -78,13 +114,13 @@ def forward(ctx, hidden, weight, bias, target, tp_group, vocab_start_index, voca
             grad_logits = probs.mul_(row_scale.unsqueeze(-1))
             grad_logits.scatter_add_(-1, local_target.unsqueeze(-1),
                                      -(in_shard.to(probs.dtype) * grad_scale).unsqueeze(-1))
-            grad_logits = grad_logits.to(hidden.dtype)
+            if grad_bias is not None:
+                grad_bias.add_(grad_logits.sum(dim=0))
+            grad_logits = grad_logits.to(compute_dtype)
 
             torch.matmul(grad_logits, weight, out=grad_hidden[start:end])
             if grad_weight is not None:
-                grad_weight.addmm_(grad_logits.t(), hidden_chunk)
-            if grad_bias is not None:
-                grad_bias.add_(grad_logits.sum(dim=0).to(grad_bias.dtype))
+                _accumulate_matmul(grad_weight, grad_logits.t(), hidden_chunk)
 
         ctx.save_for_backward(grad_hidden, grad_weight, grad_bias)
         ctx.tp_group = tp_group
@@ -94,17 +130,27 @@ def forward(ctx, hidden, weight, bias, target, tp_group, vocab_start_index, voca
     @once_differentiable
     def backward(ctx, grad_output):
         grad_hidden, grad_weight, grad_bias = ctx.saved_tensors
-        # The precomputed gradients assume a unit upstream gradient; rescale only when needed
-        # so the common ``loss.backward()`` path does not allocate another copy.
-        if not torch.equal(grad_output, torch.ones_like(grad_output)):
-            grad_hidden = grad_hidden * grad_output.to(grad_hidden.dtype)
-            grad_weight = grad_weight * grad_output.to(grad_weight.dtype) if grad_weight is not None else None
-            grad_bias = grad_bias * grad_output.to(grad_bias.dtype) if grad_bias is not None else None
+        hidden_dtype, weight_dtype, bias_dtype = ctx.input_dtypes
+        # The precomputed gradients assume a unit upstream gradient. Rescale in fp32 so that a
+        # large loss scale (e.g. 65536 with fp16) does not overflow before it is applied.
+        scale = grad_output.float()
+        unit_scale = torch.equal(scale, torch.ones_like(scale))
+
+        def finalize(grad, dtype):
+            if grad is None:
+                return None
+            if unit_scale:
+                return grad.to(dtype)
+            return (grad.float() * scale).to(dtype)
+
+        grad_hidden = finalize(grad_hidden, hidden_dtype)
+        grad_weight = finalize(grad_weight, weight_dtype)
+        grad_bias = finalize(grad_bias, bias_dtype)
         # Each rank only saw its vocabulary shard, so the input gradient is a partial sum, the
         # same reduction the column-parallel LM head performs in its own backward.
         if _tp_active(ctx.tp_group):
             dist.all_reduce(grad_hidden, op=dist.ReduceOp.SUM, group=ctx.tp_group)
-        return grad_hidden, grad_weight, grad_bias, None, None, None, None, None, None, None
+        return grad_hidden, grad_weight, grad_bias, None, None, None, None, None, None, None, None
 
 
 def vocab_parallel_linear_cross_entropy(hidden,
@@ -158,8 +204,16 @@ def vocab_parallel_linear_cross_entropy(hidden,
     else:
         grad_scale = 1.0
     if chunk_size is None:
-        chunk_size = _default_chunk_rows(num_tokens, weight.shape[0], hidden.shape[1])
+        tp_world_size = dist.get_world_size(tp_group) if tp_group is not None else 1
+        chunk_size = _default_chunk_rows(num_tokens, global_vocab_size, tp_world_size, hidden.shape[1])
+    device_type = hidden.device.type
+    if torch.is_autocast_enabled(device_type):
+        compute_dtype = torch.get_autocast_dtype(device_type)
+    else:
+        compute_dtype = torch.promote_types(hidden.dtype, weight.dtype)
 
-    loss = _VocabParallelLinearCrossEntropy.apply(hidden, weight, bias, target, tp_group, vocab_start_index,
-                                                  vocab_end_index, ignore_index, grad_scale, chunk_size)
+    with torch.autocast(device_type, enabled=False):
+        loss = _VocabParallelLinearCrossEntropy.apply(hidden, weight, bias, target, tp_group, vocab_start_index,
+                                                      vocab_end_index, ignore_index, grad_scale, chunk_size,
+                                                      compute_dtype)
     return _global_sp_sum(loss, sp_group)
diff --git a/tests/unit/v1/sequence_parallelism/test_vocab_parallel_linear_cross_entropy.py b/tests/unit/v1/sequence_parallelism/test_vocab_parallel_linear_cross_entropy.py
index a3eb4115b009..ed5c581a3fa9 100644
--- a/tests/unit/v1/sequence_parallelism/test_vocab_parallel_linear_cross_entropy.py
+++ b/tests/unit/v1/sequence_parallelism/test_vocab_parallel_linear_cross_entropy.py
@@ -79,6 +79,56 @@ def test_frozen_weight_only_produces_hidden_grad():
     assert weight.grad is None
 
 
+def test_fp16_loss_scale_does_not_overflow():
+    # A loss scale of 65536 is not representable in fp16, but the scaled gradients are.
+    hidden = _leaf(torch.ones(4, 1, dtype=torch.float16))
+    weight = _leaf(torch.zeros(2, 1, dtype=torch.float16))
+    bias = _leaf(torch.zeros(2, dtype=torch.float16))
+    target = torch.zeros(4, dtype=torch.long)
+    reference = [_leaf(hidden), _leaf(weight), _leaf(bias)]
+
+    (vocab_parallel_linear_cross_entropy(hidden, weight, target, bias=bias) * 65536).backward()
+    (_reference(*reference, target, "mean") * 65536).backward()
+
+    for got, want in zip((hidden, weight, bias), reference):
+        torch.testing.assert_close(got.grad, want.grad)
+
+
+def test_many_low_precision_chunks_accumulate_parameter_gradients_exactly():
+    # Summing 1024 per-chunk contributions in a bf16 buffer drops most of them, which would make
+    # the gradient depend on chunk_size.
+    hidden = torch.ones(1024, 1, dtype=torch.bfloat16)
+    weight = _leaf(torch.zeros(2, 1, dtype=torch.bfloat16))
+    bias = _leaf(torch.zeros(2, dtype=torch.bfloat16))
+    target = torch.zeros(1024, dtype=torch.long)
+    reference = [_leaf(weight), _leaf(bias)]
+
+    vocab_parallel_linear_cross_entropy(hidden, weight, target, bias=bias, chunk_size=1).backward()
+    _reference(hidden, *reference, target, "mean").backward()
+
+    torch.testing.assert_close(weight.grad, reference[0].grad)
+    torch.testing.assert_close(bias.grad, reference[1].grad)
+
+
+def test_autocast_low_precision_activations_with_fp32_parameters():
+    hidden, weight, bias, target = _inputs()
+    hidden = _leaf(hidden.bfloat16())
+    weight, bias = _leaf(weight), _leaf(bias)
+    reference = [_leaf(hidden), _leaf(weight), _leaf(bias)]
+
+    with torch.autocast("cpu", dtype=torch.bfloat16):
+        actual = vocab_parallel_linear_cross_entropy(hidden, weight, target, bias=bias)
+        logits = F.linear(*reference)
+    expected = F.cross_entropy(logits.float().view(-1, weight.shape[0]), target.view(-1))
+    actual.backward()
+    expected.backward()
+
+    torch.testing.assert_close(actual, expected, atol=2e-2, rtol=2e-2)
+    for got, want in zip((hidden, weight, bias), reference):
+        assert got.grad.dtype == want.dtype
+        torch.testing.assert_close(got.grad, want.grad, atol=2e-2, rtol=2e-2)
+
+
 def test_validates_inputs():
     hidden, weight, _, target = _inputs(bias=False)
     with pytest.raises(ValueError, match="Unsupported reduction"):
@@ -96,8 +146,11 @@ def test_validates_inputs():
 class TestVocabParallelLinearCrossEntropyTP(DistributedTest):
     world_size = 2
 
+    # chunk_size=None checks that ranks with different shard sizes still agree on the chunking,
+    # since every chunk issues collectives.
+    @pytest.mark.parametrize("chunk_size", [None, 3])
     @pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16])
-    def test_uneven_vocab_matches_full_vocab(self, dtype):
+    def test_uneven_vocab_matches_full_vocab(self, dtype, chunk_size):
         device = torch.device(get_accelerator().current_device_name())
         rank = dist.get_rank()
         vocab_size = 17
@@ -117,7 +170,7 @@ def test_uneven_vocab_matches_full_vocab(self, dtype):
                                                      tp_group=dist.get_world_group(),
                                                      vocab_start_index=start,
                                                      vocab_end_index=end,
-                                                     chunk_size=3)
+                                                     chunk_size=chunk_size)
         actual.backward()
         expected.backward()
 

From bfd74595a7d42b2bc9b4ac67a300a2b9eb17d763 Mon Sep 17 00:00:00 2001
From: "Jin, Youzhi" <youzhi.jin@intel.com>
Date: Sat, 10 Oct 2026 09:26:28 +0000
Subject: [PATCH 3/8] Use at least 1024 tokens per default FLCE chunk

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>
---
 deepspeed/sequence/linear_cross_entropy.py | 10 ++++++++--
 docs/_tutorials/autotp-training.md         |  3 ++-
 2 files changed, 10 insertions(+), 3 deletions(-)

diff --git a/deepspeed/sequence/linear_cross_entropy.py b/deepspeed/sequence/linear_cross_entropy.py
index 81086e7091a2..d93e4e0ff6d5 100644
--- a/deepspeed/sequence/linear_cross_entropy.py
+++ b/deepspeed/sequence/linear_cross_entropy.py
@@ -22,6 +22,11 @@ def _tp_active(tp_group):
     return tp_group is not None and dist.get_world_size(tp_group) > 1
 
 
+# Every chunk reads and writes the whole fp32 weight-gradient buffer. Below about this many
+# rows per chunk, that traffic outweighs the chunk's GEMMs and the loss becomes bandwidth-bound.
+_MIN_CHUNK_ROWS = 1024
+
+
 def _default_chunk_rows(num_tokens, global_vocab_size, tp_world_size, hidden_size):
     # Keep each fp32 logits block about the size of the hidden states, which bounds the
     # extra peak memory independently of the vocabulary size. Derived from the global shard
@@ -29,7 +34,8 @@ def _default_chunk_rows(num_tokens, global_vocab_size, tp_world_size, hidden_siz
     # even when the vocabulary is split unevenly.
     max_local_vocab_size = math.ceil(global_vocab_size / tp_world_size)
     num_chunks = max(1, math.ceil(max_local_vocab_size / hidden_size))
-    return max(1, math.ceil(num_tokens / num_chunks))
+    chunk_rows = max(math.ceil(num_tokens / num_chunks), _MIN_CHUNK_ROWS)
+    return max(1, min(chunk_rows, num_tokens))
 
 
 _MIXED_ADDMM_UNSUPPORTED = set()
@@ -172,7 +178,7 @@ def vocab_parallel_linear_cross_entropy(hidden,
 
     Gradients are computed during forward, so only a single first-order backward is
     supported. ``chunk_size`` sets the number of tokens per chunk; by default each fp32
-    logits chunk is about as large as ``hidden``.
+    logits chunk is about as large as ``hidden``, with at least 1024 tokens per chunk.
     """
     if reduction not in ("sum", "mean"):
         raise ValueError(f"Unsupported reduction: {reduction}")
diff --git a/docs/_tutorials/autotp-training.md b/docs/_tutorials/autotp-training.md
index 4406fe4518a4..1d80b16536e1 100644
--- a/docs/_tutorials/autotp-training.md
+++ b/docs/_tutorials/autotp-training.md
@@ -297,7 +297,8 @@ gradient over the TP group in backward, as the column-parallel head would. It
 supports `sum` and `mean` reductions and computes gradients during forward, so
 it allows only first-order backward. `chunk_size` (tokens per chunk) trades
 peak memory for speed; by default each FP32 logits chunk is about the size of
-the hidden states. This API is not yet wired into the HuggingFace forward or
+the hidden states, with at least 1024 tokens per chunk to amortize the FP32
+weight-gradient accumulation. This API is not yet wired into the HuggingFace forward or
 the `vocab_parallel_lm_head` configuration.
 
 

From 44eb5cba457c07bd203e598dac79fcaba61db4d6 Mon Sep 17 00:00:00 2001
From: "Jin, Youzhi" <youzhi.jin@intel.com>
Date: Sat, 10 Oct 2026 10:27:00 +0000
Subject: [PATCH 4/8] Recompute logits in vocabulary blocks in FLCE backward

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>
---
 deepspeed/sequence/linear_cross_entropy.py    | 230 ++++++++++--------
 docs/_tutorials/autotp-training.md            |  14 +-
 ...est_vocab_parallel_linear_cross_entropy.py |  30 ++-
 3 files changed, 156 insertions(+), 118 deletions(-)

diff --git a/deepspeed/sequence/linear_cross_entropy.py b/deepspeed/sequence/linear_cross_entropy.py
index d93e4e0ff6d5..6039cfed67e7 100644
--- a/deepspeed/sequence/linear_cross_entropy.py
+++ b/deepspeed/sequence/linear_cross_entropy.py
@@ -4,9 +4,9 @@
 """Fused LM-head projection and cross entropy over a vocabulary-parallel weight.
 
 A vocabulary-parallel LM head still materializes ``[tokens, vocab / tp]`` logits per rank.
-This module fuses the projection into the loss: tokens are processed in row chunks, so only
-one ``[chunk, vocab / tp]`` logits block is alive at a time, and the input and weight
-gradients are accumulated during forward.
+This module fuses the projection into the loss. Forward processes tokens in row chunks and
+keeps only per-token softmax statistics. Backward recomputes the logits one vocabulary block
+at a time, so each weight-gradient block comes from a single GEMM over all tokens.
 """
 
 import math
@@ -22,11 +22,6 @@ def _tp_active(tp_group):
     return tp_group is not None and dist.get_world_size(tp_group) > 1
 
 
-# Every chunk reads and writes the whole fp32 weight-gradient buffer. Below about this many
-# rows per chunk, that traffic outweighs the chunk's GEMMs and the loss becomes bandwidth-bound.
-_MIN_CHUNK_ROWS = 1024
-
-
 def _default_chunk_rows(num_tokens, global_vocab_size, tp_world_size, hidden_size):
     # Keep each fp32 logits block about the size of the hidden states, which bounds the
     # extra peak memory independently of the vocabulary size. Derived from the global shard
@@ -34,129 +29,152 @@ def _default_chunk_rows(num_tokens, global_vocab_size, tp_world_size, hidden_siz
     # even when the vocabulary is split unevenly.
     max_local_vocab_size = math.ceil(global_vocab_size / tp_world_size)
     num_chunks = max(1, math.ceil(max_local_vocab_size / hidden_size))
-    chunk_rows = max(math.ceil(num_tokens / num_chunks), _MIN_CHUNK_ROWS)
-    return max(1, min(chunk_rows, num_tokens))
+    return max(1, math.ceil(num_tokens / num_chunks))
 
 
-_MIXED_ADDMM_UNSUPPORTED = set()
+# Devices without GEMMs that take low-precision inputs and produce fp32 outputs (e.g. CPU).
+_MIXED_GEMM_UNSUPPORTED = set()
+_MIXED_GEMM_ERRORS = (RuntimeError, TypeError, NotImplementedError)
 
 
 def _accumulate_matmul(accumulator, mat1, mat2):
     """``accumulator += mat1 @ mat2`` with ``accumulator`` kept in its own (fp32) dtype.
 
-    Accumulating chunk gradients in a low-precision buffer drops later contributions once the
-    sum grows, so the gradient would depend on the chunk size.
+    Accumulating block gradients in a low-precision buffer drops later contributions once the
+    sum grows, so the gradient would depend on the block size.
     """
     if mat1.dtype == accumulator.dtype:
         accumulator.addmm_(mat1, mat2)
         return
     device_type = accumulator.device.type
-    if device_type not in _MIXED_ADDMM_UNSUPPORTED:
+    if device_type not in _MIXED_GEMM_UNSUPPORTED:
         try:
             torch.addmm(accumulator, mat1, mat2, out_dtype=accumulator.dtype, out=accumulator)
             return
-        except (RuntimeError, TypeError, NotImplementedError):
-            # Older PyTorch or backends without a mixed-precision GEMM (e.g. CPU).
-            _MIXED_ADDMM_UNSUPPORTED.add(device_type)
+        except _MIXED_GEMM_ERRORS:
+            _MIXED_GEMM_UNSUPPORTED.add(device_type)
     accumulator.add_(torch.mm(mat1, mat2))
 
 
+def _fp32_logits(hidden, weight, offset=None):
+    """``hidden @ weight.T + offset`` in fp32, with ``offset`` broadcast over the logits.
+
+    Producing fp32 straight from the GEMM, with the offset applied in its epilogue, saves
+    separate upcast and add passes over the ``[tokens, vocab]`` block.
+    """
+    device_type = hidden.device.type
+    if hidden.dtype != torch.float32 and device_type not in _MIXED_GEMM_UNSUPPORTED:
+        try:
+            if offset is None:
+                return torch.mm(hidden, weight.t(), out_dtype=torch.float32)
+            return torch.addmm(offset, hidden, weight.t(), out_dtype=torch.float32)
+        except _MIXED_GEMM_ERRORS:
+            _MIXED_GEMM_UNSUPPORTED.add(device_type)
+    logits = (hidden @ weight.t()).float()
+    if offset is not None:
+        logits.add_(offset)
+    return logits
+
+
 class _VocabParallelLinearCrossEntropy(torch.autograd.Function):
 
     @staticmethod
     def forward(ctx, hidden, weight, bias, target, tp_group, vocab_start_index, vocab_end_index, ignore_index,
-                grad_scale, chunk_rows, compute_dtype):
+                chunk_rows, compute_dtype):
         num_tokens = hidden.shape[0]
         local_vocab_size = weight.shape[0]
         tp_active = _tp_active(tp_group)
-        ctx.input_dtypes = (hidden.dtype, weight.dtype, None if bias is None else bias.dtype)
-
         # Mirror what autocast would do for ``F.linear``: run the GEMMs in one compute dtype even
         # when activations and parameters differ (e.g. bf16 activations with fp32 weights).
-        weight_requires_grad = weight.requires_grad
-        bias_requires_grad = bias is not None and bias.requires_grad
-        hidden = hidden.to(compute_dtype)
-        weight = weight.to(compute_dtype)
-        bias = None if bias is None else bias.to(compute_dtype)
-
-        loss = torch.zeros(num_tokens, dtype=torch.float32, device=hidden.device)
-        grad_hidden = torch.empty_like(hidden)
-        # Parameter gradients sum over every chunk, so they are accumulated in fp32.
-        grad_weight = torch.zeros_like(weight, dtype=torch.float32) if weight_requires_grad else None
-        grad_bias = torch.zeros_like(bias, dtype=torch.float32) if bias_requires_grad else None
+        hidden_compute = hidden.to(compute_dtype)
+        weight_compute = weight.to(compute_dtype)
+        bias_fp32 = None if bias is None else bias.float()
 
+        valid = target != ignore_index
+        in_shard = valid & (target >= vocab_start_index) & (target < vocab_end_index)
+        local_target = (target - vocab_start_index).clamp(0, local_vocab_size - 1)
+
+        loss = torch.empty(num_tokens, dtype=torch.float32, device=hidden.device)
+        row_max = torch.empty_like(loss)
+        row_sum_exp = torch.empty_like(loss)
         for start in range(0, num_tokens, chunk_rows):
             end = min(start + chunk_rows, num_tokens)
-            hidden_chunk = hidden[start:end]
-            target_chunk = target[start:end]
-
-            logits = hidden_chunk @ weight.t()
-            if bias is not None:
-                logits = logits + bias
-            logits = logits.float()
+            logits = _fp32_logits(hidden_compute[start:end], weight_compute, bias_fp32)
 
-            valid = target_chunk != ignore_index
-            in_shard = valid & (target_chunk >= vocab_start_index) & (target_chunk < vocab_end_index)
-            local_target = (target_chunk - vocab_start_index).clamp(0, local_vocab_size - 1)
-
-            row_max = logits.amax(dim=-1)
+            chunk_max = logits.amax(dim=-1)
             if tp_active:
-                dist.all_reduce(row_max, op=dist.ReduceOp.MAX, group=tp_group)
-            logits.sub_(row_max.unsqueeze(-1))
-            target_logit = logits.gather(-1, local_target.unsqueeze(-1)).squeeze(-1)
-            target_logit = torch.where(in_shard, target_logit, torch.zeros_like(target_logit))
-            # Reuse the logits buffer for the softmax so only one fp32 block is alive per chunk.
-            probs = logits.exp_()
-            stats = torch.stack([probs.sum(dim=-1), target_logit], dim=-1)
+                dist.all_reduce(chunk_max, op=dist.ReduceOp.MAX, group=tp_group)
+            logits.sub_(chunk_max.unsqueeze(-1))
+            target_logit = logits.gather(-1, local_target[start:end].unsqueeze(-1)).squeeze(-1)
+            target_logit = torch.where(in_shard[start:end], target_logit, torch.zeros_like(target_logit))
+            stats = torch.stack([logits.exp_().sum(dim=-1), target_logit], dim=-1)
             if tp_active:
                 dist.all_reduce(stats, op=dist.ReduceOp.SUM, group=tp_group)
             sum_exp, target_logit = stats.unbind(-1)
 
-            loss[start:end] = torch.where(valid, sum_exp.log() - target_logit, torch.zeros_like(sum_exp))
-
-            # d(loss) / d(logits) = softmax - onehot, zeroed for ignored tokens and pre-scaled by
-            # the reduction factor so backward only has to multiply by the incoming gradient.
-            row_scale = valid.to(probs.dtype) * grad_scale / sum_exp
-            grad_logits = probs.mul_(row_scale.unsqueeze(-1))
-            grad_logits.scatter_add_(-1, local_target.unsqueeze(-1),
-                                     -(in_shard.to(probs.dtype) * grad_scale).unsqueeze(-1))
-            if grad_bias is not None:
-                grad_bias.add_(grad_logits.sum(dim=0))
-            grad_logits = grad_logits.to(compute_dtype)
-
-            torch.matmul(grad_logits, weight, out=grad_hidden[start:end])
-            if grad_weight is not None:
-                _accumulate_matmul(grad_weight, grad_logits.t(), hidden_chunk)
+            row_max[start:end] = chunk_max
+            row_sum_exp[start:end] = sum_exp
+            loss[start:end] = torch.where(valid[start:end], sum_exp.log() - target_logit, torch.zeros_like(sum_exp))
 
-        ctx.save_for_backward(grad_hidden, grad_weight, grad_bias)
+        # The max and the sum of exponentials are kept apart: folding them into one log-sum-exp
+        # loses the small normalizer to rounding when the logits share a large offset.
+        ctx.save_for_backward(hidden, weight, bias, valid, in_shard, local_target, row_max, row_sum_exp)
         ctx.tp_group = tp_group
-        return loss.sum() * grad_scale
+        ctx.compute_dtype = compute_dtype
+        # Backward recomputes the logits in vocabulary blocks with the same number of elements
+        # as a forward chunk, so both passes have the same peak.
+        num_chunks = math.ceil(num_tokens / chunk_rows) if num_tokens else 1
+        ctx.block_cols = max(1, math.ceil(local_vocab_size / num_chunks))
+        return loss
 
     @staticmethod
     @once_differentiable
-    def backward(ctx, grad_output):
-        grad_hidden, grad_weight, grad_bias = ctx.saved_tensors
-        hidden_dtype, weight_dtype, bias_dtype = ctx.input_dtypes
-        # The precomputed gradients assume a unit upstream gradient. Rescale in fp32 so that a
-        # large loss scale (e.g. 65536 with fp16) does not overflow before it is applied.
-        scale = grad_output.float()
-        unit_scale = torch.equal(scale, torch.ones_like(scale))
-
-        def finalize(grad, dtype):
-            if grad is None:
-                return None
-            if unit_scale:
-                return grad.to(dtype)
-            return (grad.float() * scale).to(dtype)
-
-        grad_hidden = finalize(grad_hidden, hidden_dtype)
-        grad_weight = finalize(grad_weight, weight_dtype)
-        grad_bias = finalize(grad_bias, bias_dtype)
-        # Each rank only saw its vocabulary shard, so the input gradient is a partial sum, the
-        # same reduction the column-parallel LM head performs in its own backward.
-        if _tp_active(ctx.tp_group):
-            dist.all_reduce(grad_hidden, op=dist.ReduceOp.SUM, group=ctx.tp_group)
-        return grad_hidden, grad_weight, grad_bias, None, None, None, None, None, None, None, None
+    def backward(ctx, grad_loss):
+        hidden, weight, bias, valid, in_shard, local_target, row_max, row_sum_exp = ctx.saved_tensors
+        compute_dtype = ctx.compute_dtype
+        needs_hidden, needs_weight, needs_bias = ctx.needs_input_grad[:3]
+        local_vocab_size = weight.shape[0]
+
+        hidden_compute = hidden.to(compute_dtype)
+        token_grad = torch.where(valid, grad_loss.float(), torch.zeros_like(row_max))
+        # softmax * upstream gradient = exp(logits - max) * (upstream gradient / sum_exp).
+        prob_scale = (token_grad / row_sum_exp).unsqueeze(-1)
+        neg_row_max = -row_max.unsqueeze(-1)
+        rows = torch.arange(hidden.shape[0], device=hidden.device)
+        # The hidden gradient sums over every vocabulary block, so it is accumulated in fp32.
+        grad_hidden = torch.zeros_like(hidden, dtype=torch.float32) if needs_hidden else None
+        grad_weight = torch.empty_like(weight) if needs_weight else None
+        grad_bias = torch.empty_like(bias) if needs_bias else None
+
+        for start in range(0, local_vocab_size, ctx.block_cols):
+            end = min(start + ctx.block_cols, local_vocab_size)
+            weight_block = weight[start:end].to(compute_dtype)
+            if bias is None:
+                logits = _fp32_logits(hidden_compute, weight_block, neg_row_max)
+            else:
+                logits = _fp32_logits(hidden_compute, weight_block, bias[start:end].float()).add_(neg_row_max)
+
+            # d(loss) / d(logits) = (softmax - onehot) * upstream gradient, one block at a time.
+            grad_logits = logits.exp_().mul_(prob_scale)
+            in_block = in_shard & (local_target >= start) & (local_target < end)
+            block_rows = rows[in_block]
+            grad_logits[block_rows, local_target[in_block] - start] -= token_grad[in_block]
+
+            if grad_bias is not None:
+                grad_bias[start:end] = grad_logits.sum(dim=0)
+            grad_logits = grad_logits.to(compute_dtype)
+            if grad_weight is not None:
+                grad_weight[start:end] = grad_logits.t() @ hidden_compute
+            if grad_hidden is not None:
+                _accumulate_matmul(grad_hidden, grad_logits, weight_block)
+
+        if grad_hidden is not None:
+            grad_hidden = grad_hidden.to(hidden.dtype)
+            # Each rank only saw its vocabulary shard, so the input gradient is a partial sum, the
+            # same reduction the column-parallel LM head performs in its own backward.
+            if _tp_active(ctx.tp_group):
+                dist.all_reduce(grad_hidden, op=dist.ReduceOp.SUM, group=ctx.tp_group)
+        return grad_hidden, grad_weight, grad_bias, None, None, None, None, None, None, None
 
 
 def vocab_parallel_linear_cross_entropy(hidden,
@@ -176,11 +194,12 @@ def vocab_parallel_linear_cross_entropy(hidden,
     example ``VocabParallelLinear.weight``. ``hidden`` must be identical on every TP rank, as
     for a column-parallel layer; its gradient is all-reduced over ``tp_group`` in backward.
 
-    Gradients are computed during forward, so only a single first-order backward is
-    supported. ``chunk_size`` sets the number of tokens per chunk; by default each fp32
-    logits chunk is about as large as ``hidden``, with at least 1024 tokens per chunk.
+    ``chunk_size`` sets the number of tokens per forward chunk; backward recomputes the
+    logits in vocabulary blocks of the same size. By default each fp32 logits block is about
+    as large as ``hidden``. ``reduction="none"`` returns the per-token losses in the shape of
+    ``target``, with zeros at ignored positions, and does not reduce over ``sp_group``.
     """
-    if reduction not in ("sum", "mean"):
+    if reduction not in ("none", "sum", "mean"):
         raise ValueError(f"Unsupported reduction: {reduction}")
     if hidden.shape[:-1] != target.shape:
         raise ValueError("hidden and target must have matching non-hidden dimensions")
@@ -193,6 +212,7 @@ def vocab_parallel_linear_cross_entropy(hidden,
 
     vocab_start_index, vocab_end_index, global_vocab_size = _resolve_vocab_metadata(
         weight.shape[0], vocab_start_index, vocab_end_index, tp_group, hidden.device)
+    target_shape = target.shape
     target = target.reshape(-1).to(dtype=torch.long)
     invalid_target = (target != ignore_index) & ((target < 0) | (target >= global_vocab_size))
     if invalid_target.any().item():
@@ -202,13 +222,6 @@ def vocab_parallel_linear_cross_entropy(hidden,
     weight = weight.contiguous()
     num_tokens = hidden.shape[0]
 
-    if reduction == "mean":
-        valid_tokens = (target != ignore_index).sum().to(torch.float32)
-        if sp_group is not None and dist.get_world_size(sp_group) > 1:
-            dist.all_reduce(valid_tokens, op=dist.ReduceOp.SUM, group=sp_group)
-        grad_scale = 1.0 / max(valid_tokens.item(), 1.0)
-    else:
-        grad_scale = 1.0
     if chunk_size is None:
         tp_world_size = dist.get_world_size(tp_group) if tp_group is not None else 1
         chunk_size = _default_chunk_rows(num_tokens, global_vocab_size, tp_world_size, hidden.shape[1])
@@ -219,7 +232,16 @@ def vocab_parallel_linear_cross_entropy(hidden,
         compute_dtype = torch.promote_types(hidden.dtype, weight.dtype)
 
     with torch.autocast(device_type, enabled=False):
-        loss = _VocabParallelLinearCrossEntropy.apply(hidden, weight, bias, target, tp_group, vocab_start_index,
-                                                      vocab_end_index, ignore_index, grad_scale, chunk_size,
-                                                      compute_dtype)
+        per_token_loss = _VocabParallelLinearCrossEntropy.apply(hidden, weight, bias, target, tp_group,
+                                                                vocab_start_index, vocab_end_index, ignore_index,
+                                                                chunk_size, compute_dtype)
+    if reduction == "none":
+        return per_token_loss.view(target_shape)
+
+    loss = per_token_loss.sum()
+    if reduction == "mean":
+        valid_tokens = (target != ignore_index).sum().to(torch.float32)
+        if sp_group is not None and dist.get_world_size(sp_group) > 1:
+            dist.all_reduce(valid_tokens, op=dist.ReduceOp.SUM, group=sp_group)
+        loss = loss / valid_tokens.clamp(min=1.0)
     return _global_sp_sum(loss, sp_group)
diff --git a/docs/_tutorials/autotp-training.md b/docs/_tutorials/autotp-training.md
index 1d80b16536e1..92bc686fe621 100644
--- a/docs/_tutorials/autotp-training.md
+++ b/docs/_tutorials/autotp-training.md
@@ -275,8 +275,9 @@ this loss is the sole aggregation over a manually constructed TP x SP mesh.
 
 The vocabulary-parallel loss above still materializes `[tokens, vocab / tp]`
 logits on every rank. `vocab_parallel_linear_cross_entropy` fuses the LM-head
-projection into the loss and processes tokens in chunks, so only one
-`[chunk, vocab / tp]` block is alive at a time:
+projection into the loss. Forward processes tokens in chunks and keeps only
+per-token softmax statistics; backward recomputes the logits one vocabulary block
+at a time. Only one block of logits is alive at a time:
 
 ```python
 from deepspeed.sequence.linear_cross_entropy import vocab_parallel_linear_cross_entropy
@@ -294,11 +295,10 @@ loss = vocab_parallel_linear_cross_entropy(hidden,
 
 The call replaces both `head(hidden)` and the loss: it all-reduces the hidden
 gradient over the TP group in backward, as the column-parallel head would. It
-supports `sum` and `mean` reductions and computes gradients during forward, so
-it allows only first-order backward. `chunk_size` (tokens per chunk) trades
-peak memory for speed; by default each FP32 logits chunk is about the size of
-the hidden states, with at least 1024 tokens per chunk to amortize the FP32
-weight-gradient accumulation. This API is not yet wired into the HuggingFace forward or
+supports `none`, `sum` and `mean` reductions. `chunk_size` (tokens per forward
+chunk; backward uses vocabulary blocks of the same size) trades peak memory for
+speed; by default each FP32 logits block is about the size of the hidden states.
+This API is not yet wired into the HuggingFace forward or
 the `vocab_parallel_lm_head` configuration.
 
 
diff --git a/tests/unit/v1/sequence_parallelism/test_vocab_parallel_linear_cross_entropy.py b/tests/unit/v1/sequence_parallelism/test_vocab_parallel_linear_cross_entropy.py
index ed5c581a3fa9..b1936e9f2a58 100644
--- a/tests/unit/v1/sequence_parallelism/test_vocab_parallel_linear_cross_entropy.py
+++ b/tests/unit/v1/sequence_parallelism/test_vocab_parallel_linear_cross_entropy.py
@@ -32,7 +32,7 @@ def _reference(hidden, weight, bias, target, reduction):
     return F.cross_entropy(logits.view(-1, weight.shape[0]), target.view(-1), reduction=reduction)
 
 
-@pytest.mark.parametrize("reduction", ["sum", "mean"])
+@pytest.mark.parametrize("reduction", ["none", "sum", "mean"])
 @pytest.mark.parametrize("use_bias", [True, False])
 @pytest.mark.parametrize("chunk_size", [None, 1, 3])
 def test_matches_linear_then_cross_entropy(reduction, use_bias, chunk_size):
@@ -41,6 +41,8 @@ def test_matches_linear_then_cross_entropy(reduction, use_bias, chunk_size):
     args = [_leaf(hidden), _leaf(weight), _leaf(bias)]
 
     expected = _reference(*ref_args, target, reduction)
+    if reduction == "none":
+        expected = expected.view_as(target)
     actual = vocab_parallel_linear_cross_entropy(args[0],
                                                  args[1],
                                                  target,
@@ -49,9 +51,9 @@ def test_matches_linear_then_cross_entropy(reduction, use_bias, chunk_size):
                                                  chunk_size=chunk_size)
     torch.testing.assert_close(actual, expected)
 
-    # A non-unit upstream gradient must rescale the gradients precomputed in forward.
-    (actual * 3).backward()
-    (expected * 3).backward()
+    upstream = torch.rand_like(expected) + 0.5
+    actual.backward(upstream)
+    expected.backward(upstream)
     for got, want in zip(args, ref_args):
         if want is not None:
             torch.testing.assert_close(got.grad, want.grad)
@@ -95,8 +97,7 @@ def test_fp16_loss_scale_does_not_overflow():
 
 
 def test_many_low_precision_chunks_accumulate_parameter_gradients_exactly():
-    # Summing 1024 per-chunk contributions in a bf16 buffer drops most of them, which would make
-    # the gradient depend on chunk_size.
+    # The weight gradient sums 1024 per-token contributions; it must not depend on chunk_size.
     hidden = torch.ones(1024, 1, dtype=torch.bfloat16)
     weight = _leaf(torch.zeros(2, 1, dtype=torch.bfloat16))
     bias = _leaf(torch.zeros(2, dtype=torch.bfloat16))
@@ -110,6 +111,21 @@ def test_many_low_precision_chunks_accumulate_parameter_gradients_exactly():
     torch.testing.assert_close(bias.grad, reference[1].grad)
 
 
+def test_hidden_gradient_sums_many_vocab_entries_exactly():
+    # Each token's hidden gradient sums 1023 contributions of 1/1024. A bf16 running sum stalls
+    # near 0.5, so the result must not depend on how the vocabulary is split into blocks.
+    hidden = torch.zeros(1024, 1, dtype=torch.bfloat16, requires_grad=True)
+    weight = torch.ones(1024, 1, dtype=torch.bfloat16)
+    weight[0] = 0
+    target = torch.zeros(1024, dtype=torch.long)
+    reference = _leaf(hidden)
+
+    vocab_parallel_linear_cross_entropy(hidden, weight, target, reduction="sum", chunk_size=1).backward()
+    _reference(reference, weight, None, target, "sum").backward()
+
+    torch.testing.assert_close(hidden.grad, reference.grad)
+
+
 def test_autocast_low_precision_activations_with_fp32_parameters():
     hidden, weight, bias, target = _inputs()
     hidden = _leaf(hidden.bfloat16())
@@ -132,7 +148,7 @@ def test_autocast_low_precision_activations_with_fp32_parameters():
 def test_validates_inputs():
     hidden, weight, _, target = _inputs(bias=False)
     with pytest.raises(ValueError, match="Unsupported reduction"):
-        vocab_parallel_linear_cross_entropy(hidden, weight, target, reduction="none")
+        vocab_parallel_linear_cross_entropy(hidden, weight, target, reduction="avg")
     with pytest.raises(ValueError, match="matching non-hidden dimensions"):
         vocab_parallel_linear_cross_entropy(hidden, weight, target[:, :-1])
     with pytest.raises(ValueError, match="local_vocab, hidden"):

From 138f9c76ec0b844cb1fb378a9e56f38129fabed9 Mon Sep 17 00:00:00 2001
From: "Jin, Youzhi" <youzhi.jin@intel.com>
Date: Sat, 10 Oct 2026 10:45:01 +0000
Subject: [PATCH 5/8] Add the bias inside the GEMM in the FLCE fallback path

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>
---
 deepspeed/sequence/linear_cross_entropy.py    | 30 ++++++++++---------
 ...est_vocab_parallel_linear_cross_entropy.py | 18 +++++++++++
 2 files changed, 34 insertions(+), 14 deletions(-)

diff --git a/deepspeed/sequence/linear_cross_entropy.py b/deepspeed/sequence/linear_cross_entropy.py
index 6039cfed67e7..76da23bb9d5f 100644
--- a/deepspeed/sequence/linear_cross_entropy.py
+++ b/deepspeed/sequence/linear_cross_entropy.py
@@ -56,8 +56,8 @@ def _accumulate_matmul(accumulator, mat1, mat2):
     accumulator.add_(torch.mm(mat1, mat2))
 
 
-def _fp32_logits(hidden, weight, offset=None):
-    """``hidden @ weight.T + offset`` in fp32, with ``offset`` broadcast over the logits.
+def _fp32_logits(hidden, weight, bias=None, row_offset=None):
+    """``hidden @ weight.T + bias + row_offset`` in fp32; ``row_offset`` has shape ``[tokens, 1]``.
 
     Producing fp32 straight from the GEMM, with the offset applied in its epilogue, saves
     separate upcast and add passes over the ``[tokens, vocab]`` block.
@@ -65,15 +65,19 @@ def _fp32_logits(hidden, weight, offset=None):
     device_type = hidden.device.type
     if hidden.dtype != torch.float32 and device_type not in _MIXED_GEMM_UNSUPPORTED:
         try:
-            if offset is None:
+            if bias is not None:
+                logits = torch.addmm(bias.float(), hidden, weight.t(), out_dtype=torch.float32)
+            elif row_offset is not None:
+                return torch.addmm(row_offset, hidden, weight.t(), out_dtype=torch.float32)
+            else:
                 return torch.mm(hidden, weight.t(), out_dtype=torch.float32)
-            return torch.addmm(offset, hidden, weight.t(), out_dtype=torch.float32)
+            return logits if row_offset is None else logits.add_(row_offset)
         except _MIXED_GEMM_ERRORS:
             _MIXED_GEMM_UNSUPPORTED.add(device_type)
-    logits = (hidden @ weight.t()).float()
-    if offset is not None:
-        logits.add_(offset)
-    return logits
+    # Add the bias inside the GEMM as F.linear does, so a large product that the bias cancels
+    # does not overflow a low-precision output first.
+    logits = torch.nn.functional.linear(hidden, weight, bias).float()
+    return logits if row_offset is None else logits.add_(row_offset)
 
 
 class _VocabParallelLinearCrossEntropy(torch.autograd.Function):
@@ -88,7 +92,7 @@ def forward(ctx, hidden, weight, bias, target, tp_group, vocab_start_index, voca
         # when activations and parameters differ (e.g. bf16 activations with fp32 weights).
         hidden_compute = hidden.to(compute_dtype)
         weight_compute = weight.to(compute_dtype)
-        bias_fp32 = None if bias is None else bias.float()
+        bias_compute = None if bias is None else bias.to(compute_dtype)
 
         valid = target != ignore_index
         in_shard = valid & (target >= vocab_start_index) & (target < vocab_end_index)
@@ -99,7 +103,7 @@ def forward(ctx, hidden, weight, bias, target, tp_group, vocab_start_index, voca
         row_sum_exp = torch.empty_like(loss)
         for start in range(0, num_tokens, chunk_rows):
             end = min(start + chunk_rows, num_tokens)
-            logits = _fp32_logits(hidden_compute[start:end], weight_compute, bias_fp32)
+            logits = _fp32_logits(hidden_compute[start:end], weight_compute, bias_compute)
 
             chunk_max = logits.amax(dim=-1)
             if tp_active:
@@ -149,10 +153,8 @@ def backward(ctx, grad_loss):
         for start in range(0, local_vocab_size, ctx.block_cols):
             end = min(start + ctx.block_cols, local_vocab_size)
             weight_block = weight[start:end].to(compute_dtype)
-            if bias is None:
-                logits = _fp32_logits(hidden_compute, weight_block, neg_row_max)
-            else:
-                logits = _fp32_logits(hidden_compute, weight_block, bias[start:end].float()).add_(neg_row_max)
+            bias_block = None if bias is None else bias[start:end].to(compute_dtype)
+            logits = _fp32_logits(hidden_compute, weight_block, bias_block, neg_row_max)
 
             # d(loss) / d(logits) = (softmax - onehot) * upstream gradient, one block at a time.
             grad_logits = logits.exp_().mul_(prob_scale)
diff --git a/tests/unit/v1/sequence_parallelism/test_vocab_parallel_linear_cross_entropy.py b/tests/unit/v1/sequence_parallelism/test_vocab_parallel_linear_cross_entropy.py
index b1936e9f2a58..dd904f46b4cb 100644
--- a/tests/unit/v1/sequence_parallelism/test_vocab_parallel_linear_cross_entropy.py
+++ b/tests/unit/v1/sequence_parallelism/test_vocab_parallel_linear_cross_entropy.py
@@ -96,6 +96,24 @@ def test_fp16_loss_scale_does_not_overflow():
         torch.testing.assert_close(got.grad, want.grad)
 
 
+def test_fp16_product_cancelled_by_bias_stays_finite():
+    # The product 256 * 256 overflows fp16 but the bias brings the logits back to 32, as in F.linear.
+    hidden = _leaf(torch.tensor([[256.0]], dtype=torch.float16))
+    weight = _leaf(torch.tensor([[256.0], [256.0]], dtype=torch.float16))
+    bias = _leaf(torch.tensor([-65504.0, -65504.0], dtype=torch.float16))
+    target = torch.tensor([0])
+    args = [_leaf(hidden), _leaf(weight), _leaf(bias)]
+
+    actual = vocab_parallel_linear_cross_entropy(*args[:2], target, bias=args[2])
+    expected = F.cross_entropy(F.linear(hidden, weight, bias).float(), target)
+    actual.backward()
+    expected.backward()
+
+    torch.testing.assert_close(actual, expected)
+    for got, want in zip(args, (hidden, weight, bias)):
+        torch.testing.assert_close(got.grad, want.grad)
+
+
 def test_many_low_precision_chunks_accumulate_parameter_gradients_exactly():
     # The weight gradient sums 1024 per-token contributions; it must not depend on chunk_size.
     hidden = torch.ones(1024, 1, dtype=torch.bfloat16)

From 74ae27cb45f6511b8996a7f09f2d3c69c9fb40d0 Mon Sep 17 00:00:00 2001
From: "Jin, Youzhi" <youzhi.jin@intel.com>
Date: Mon, 12 Oct 2026 07:56:49 +0000
Subject: [PATCH 6/8] Fix FLCE fp16 fallback overflow and pre-2.4 autocast API

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>
---
 deepspeed/sequence/linear_cross_entropy.py    | 26 ++++++++++++---
 ...est_vocab_parallel_linear_cross_entropy.py | 33 ++++++++++++++++++-
 2 files changed, 54 insertions(+), 5 deletions(-)

diff --git a/deepspeed/sequence/linear_cross_entropy.py b/deepspeed/sequence/linear_cross_entropy.py
index 76da23bb9d5f..064eada15dec 100644
--- a/deepspeed/sequence/linear_cross_entropy.py
+++ b/deepspeed/sequence/linear_cross_entropy.py
@@ -53,7 +53,8 @@ def _accumulate_matmul(accumulator, mat1, mat2):
             return
         except _MIXED_GEMM_ERRORS:
             _MIXED_GEMM_UNSUPPORTED.add(device_type)
-    accumulator.add_(torch.mm(mat1, mat2))
+    # A low-precision block product can overflow even when the sum over all blocks is finite.
+    accumulator.addmm_(mat1.to(accumulator.dtype), mat2.to(accumulator.dtype))
 
 
 def _fp32_logits(hidden, weight, bias=None, row_offset=None):
@@ -80,6 +81,24 @@ def _fp32_logits(hidden, weight, bias=None, row_offset=None):
     return logits if row_offset is None else logits.add_(row_offset)
 
 
+def _autocast_dtype(device_type):
+    """The active autocast dtype for ``device_type``, or ``None`` when autocast is off."""
+    if hasattr(torch, "get_autocast_dtype"):
+        if torch.is_autocast_enabled(device_type):
+            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":
+        enabled, dtype = torch.is_autocast_enabled(), torch.get_autocast_gpu_dtype
+    else:
+        is_enabled = getattr(torch, f"is_autocast_{device_type}_enabled", None)
+        enabled = is_enabled is not None and is_enabled()
+        dtype = getattr(torch, f"get_autocast_{device_type}_dtype", None)
+    if enabled and dtype is not None:
+        return dtype()
+    return None
+
+
 class _VocabParallelLinearCrossEntropy(torch.autograd.Function):
 
     @staticmethod
@@ -228,9 +247,8 @@ def vocab_parallel_linear_cross_entropy(hidden,
         tp_world_size = dist.get_world_size(tp_group) if tp_group is not None else 1
         chunk_size = _default_chunk_rows(num_tokens, global_vocab_size, tp_world_size, hidden.shape[1])
     device_type = hidden.device.type
-    if torch.is_autocast_enabled(device_type):
-        compute_dtype = torch.get_autocast_dtype(device_type)
-    else:
+    compute_dtype = _autocast_dtype(device_type)
+    if compute_dtype is None:
         compute_dtype = torch.promote_types(hidden.dtype, weight.dtype)
 
     with torch.autocast(device_type, enabled=False):
diff --git a/tests/unit/v1/sequence_parallelism/test_vocab_parallel_linear_cross_entropy.py b/tests/unit/v1/sequence_parallelism/test_vocab_parallel_linear_cross_entropy.py
index dd904f46b4cb..53cb09f8ccf8 100644
--- a/tests/unit/v1/sequence_parallelism/test_vocab_parallel_linear_cross_entropy.py
+++ b/tests/unit/v1/sequence_parallelism/test_vocab_parallel_linear_cross_entropy.py
@@ -8,6 +8,7 @@
 import deepspeed.comm as dist
 from deepspeed.accelerator import get_accelerator
 from deepspeed.module_inject.tp_shard import AutoTPMeta, get_shard_size_list
+from deepspeed.sequence import linear_cross_entropy
 from deepspeed.sequence.linear_cross_entropy import vocab_parallel_linear_cross_entropy
 from unit.common import DistributedTest
 
@@ -96,6 +97,20 @@ def test_fp16_loss_scale_does_not_overflow():
         torch.testing.assert_close(got.grad, want.grad)
 
 
+def test_fp16_hidden_gradient_blocks_that_cancel_stay_finite():
+    # Each vocabulary entry contributes +-131072 to the scaled hidden gradient, which overflows
+    # fp16 alone, but the contributions cancel, so the result must not depend on vocab blocking.
+    hidden = _leaf(torch.zeros(4, 1, dtype=torch.float16))
+    weight = torch.full((2, 1), 8.0, dtype=torch.float16)
+    target = torch.zeros(4, dtype=torch.long)
+    reference = _leaf(hidden)
+
+    (vocab_parallel_linear_cross_entropy(hidden, weight, target) * 65536).backward()
+    (F.cross_entropy(F.linear(reference, weight).float(), target) * 65536).backward()
+
+    torch.testing.assert_close(hidden.grad, reference.grad)
+
+
 def test_fp16_product_cancelled_by_bias_stays_finite():
     # The product 256 * 256 overflows fp16 but the bias brings the logits back to 32, as in F.linear.
     hidden = _leaf(torch.tensor([[256.0]], dtype=torch.float16))
@@ -144,12 +159,28 @@ def test_hidden_gradient_sums_many_vocab_entries_exactly():
     torch.testing.assert_close(hidden.grad, reference.grad)
 
 
-def test_autocast_low_precision_activations_with_fp32_parameters():
+class _LegacyAutocastTorch:
+    """``torch`` as seen by code on PyTorch < 2.4, which has no device-aware autocast getters."""
+
+    def __getattr__(self, name):
+        if name == "get_autocast_dtype":
+            raise AttributeError(name)
+        return getattr(torch, name)
+
+    @staticmethod
+    def is_autocast_enabled():
+        return torch.is_autocast_enabled("cuda")
+
+
+@pytest.mark.parametrize("legacy_autocast_api", [False, True])
+def test_autocast_low_precision_activations_with_fp32_parameters(monkeypatch, legacy_autocast_api):
     hidden, weight, bias, target = _inputs()
     hidden = _leaf(hidden.bfloat16())
     weight, bias = _leaf(weight), _leaf(bias)
     reference = [_leaf(hidden), _leaf(weight), _leaf(bias)]
 
+    if legacy_autocast_api:
+        monkeypatch.setattr(linear_cross_entropy, "torch", _LegacyAutocastTorch())
     with torch.autocast("cpu", dtype=torch.bfloat16):
         actual = vocab_parallel_linear_cross_entropy(hidden, weight, target, bias=bias)
         logits = F.linear(*reference)

From 1228ea0d05fac11ab4bcb46eff662db0079e9c29 Mon Sep 17 00:00:00 2001
From: "Jin, Youzhi" <youzhi.jin@intel.com>
Date: Tue, 13 Oct 2026 07:01:51 +0000
Subject: [PATCH 7/8] Sum FLCE hidden-gradient partials over TP in fp32

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>
---
 deepspeed/sequence/linear_cross_entropy.py    |  5 +++--
 ...est_vocab_parallel_linear_cross_entropy.py | 21 +++++++++++++++++++
 2 files changed, 24 insertions(+), 2 deletions(-)

diff --git a/deepspeed/sequence/linear_cross_entropy.py b/deepspeed/sequence/linear_cross_entropy.py
index 064eada15dec..6b32e538d97f 100644
--- a/deepspeed/sequence/linear_cross_entropy.py
+++ b/deepspeed/sequence/linear_cross_entropy.py
@@ -190,11 +190,12 @@ def backward(ctx, grad_loss):
                 _accumulate_matmul(grad_hidden, grad_logits, weight_block)
 
         if grad_hidden is not None:
-            grad_hidden = grad_hidden.to(hidden.dtype)
             # Each rank only saw its vocabulary shard, so the input gradient is a partial sum, the
-            # same reduction the column-parallel LM head performs in its own backward.
+            # same reduction the column-parallel LM head performs in its own backward. The partials
+            # are summed in fp32 because they can overflow fp16 even when their sum does not.
             if _tp_active(ctx.tp_group):
                 dist.all_reduce(grad_hidden, op=dist.ReduceOp.SUM, group=ctx.tp_group)
+            grad_hidden = grad_hidden.to(hidden.dtype)
         return grad_hidden, grad_weight, grad_bias, None, None, None, None, None, None, None
 
 
diff --git a/tests/unit/v1/sequence_parallelism/test_vocab_parallel_linear_cross_entropy.py b/tests/unit/v1/sequence_parallelism/test_vocab_parallel_linear_cross_entropy.py
index 53cb09f8ccf8..6cdb25a8d121 100644
--- a/tests/unit/v1/sequence_parallelism/test_vocab_parallel_linear_cross_entropy.py
+++ b/tests/unit/v1/sequence_parallelism/test_vocab_parallel_linear_cross_entropy.py
@@ -244,3 +244,24 @@ def test_uneven_vocab_matches_full_vocab(self, dtype, chunk_size):
         torch.testing.assert_close(args[0].grad, ref_args[0].grad, **tolerance)
         torch.testing.assert_close(args[1].grad, ref_args[1].grad[start:end], **tolerance)
         torch.testing.assert_close(args[2].grad, ref_args[2].grad[start:end], **tolerance)
+
+    def test_fp16_hidden_gradient_partials_that_cancel_across_ranks_stay_finite(self):
+        # With a 65536 loss scale each rank's hidden-gradient partial is +-65536, which overflows
+        # fp16, but the two classes live on different ranks and their partials cancel to zero.
+        device = torch.device(get_accelerator().current_device_name())
+        rank = dist.get_rank()
+        hidden = _leaf(torch.zeros(4, 1, dtype=torch.float16, device=device))
+        weight = torch.full((2, 1), 8.0, dtype=torch.float16, device=device)
+        target = torch.zeros(4, dtype=torch.long, device=device)
+        reference = _leaf(hidden)
+
+        loss = vocab_parallel_linear_cross_entropy(hidden,
+                                                   weight[rank:rank + 1],
+                                                   target,
+                                                   tp_group=dist.get_world_group(),
+                                                   vocab_start_index=rank,
+                                                   vocab_end_index=rank + 1)
+        (loss * 65536).backward()
+        (F.cross_entropy(F.linear(reference, weight).float(), target) * 65536).backward()
+
+        torch.testing.assert_close(hidden.grad, reference.grad)

From 8582e19c90b57ba9a2d56a74bad2f90b7ae9fa3b Mon Sep 17 00:00:00 2001
From: "Jin, Youzhi" <youzhi.jin@intel.com>
Date: Wed, 14 Oct 2026 05:20:27 +0000
Subject: [PATCH 8/8] Detect fp32-output GEMM support from the dispatcher
 instead of a failed 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>
---
 deepspeed/sequence/linear_cross_entropy.py | 54 +++++++++++++---------
 1 file changed, 32 insertions(+), 22 deletions(-)

diff --git a/deepspeed/sequence/linear_cross_entropy.py b/deepspeed/sequence/linear_cross_entropy.py
index 6b32e538d97f..8aad59e3c173 100644
--- a/deepspeed/sequence/linear_cross_entropy.py
+++ b/deepspeed/sequence/linear_cross_entropy.py
@@ -9,6 +9,7 @@
 at a time, so each weight-gradient block comes from a single GEMM over all tokens.
 """
 
+import functools
 import math
 
 import torch
@@ -32,9 +33,24 @@ def _default_chunk_rows(num_tokens, global_vocab_size, tp_world_size, hidden_siz
     return max(1, math.ceil(num_tokens / num_chunks))
 
 
-# Devices without GEMMs that take low-precision inputs and produce fp32 outputs (e.g. CPU).
-_MIXED_GEMM_UNSUPPORTED = set()
-_MIXED_GEMM_ERRORS = (RuntimeError, TypeError, NotImplementedError)
+_MIXED_GEMM_INPUT_DTYPES = (torch.float16, torch.bfloat16)
+_MIXED_GEMM_OPS = ("aten::mm.dtype", "aten::addmm.dtype", "aten::addmm.dtype_out")
+
+
+@functools.lru_cache(maxsize=None)
+def _has_mixed_gemm(device_type):
+    """Whether ``device_type`` has GEMMs that take fp16/bf16 inputs and produce fp32 outputs.
+
+    The dispatcher is asked whether the device registers these kernels, instead of trying a GEMM
+    and treating any failure as "unsupported": a GEMM can also fail for unrelated reasons such as
+    running out of memory. PyTorch < 2.8 has no such ops, and some devices (e.g. CPU) do not
+    implement them.
+    """
+    try:
+        dispatch_key = torch._C._dispatch_key_for_device(device_type)
+        return all(torch._C._dispatch_has_kernel_for_dispatch_key(op, dispatch_key) for op in _MIXED_GEMM_OPS)
+    except (AttributeError, RuntimeError):
+        return False
 
 
 def _accumulate_matmul(accumulator, mat1, mat2):
@@ -46,13 +62,11 @@ def _accumulate_matmul(accumulator, mat1, mat2):
     if mat1.dtype == accumulator.dtype:
         accumulator.addmm_(mat1, mat2)
         return
-    device_type = accumulator.device.type
-    if device_type not in _MIXED_GEMM_UNSUPPORTED:
-        try:
-            torch.addmm(accumulator, mat1, mat2, out_dtype=accumulator.dtype, out=accumulator)
-            return
-        except _MIXED_GEMM_ERRORS:
-            _MIXED_GEMM_UNSUPPORTED.add(device_type)
+    use_mixed_gemm = (accumulator.dtype == torch.float32 and mat1.dtype in _MIXED_GEMM_INPUT_DTYPES
+                      and _has_mixed_gemm(accumulator.device.type))
+    if use_mixed_gemm:
+        torch.addmm(accumulator, mat1, mat2, out_dtype=accumulator.dtype, out=accumulator)
+        return
     # A low-precision block product can overflow even when the sum over all blocks is finite.
     accumulator.addmm_(mat1.to(accumulator.dtype), mat2.to(accumulator.dtype))
 
@@ -63,18 +77,14 @@ def _fp32_logits(hidden, weight, bias=None, row_offset=None):
     Producing fp32 straight from the GEMM, with the offset applied in its epilogue, saves
     separate upcast and add passes over the ``[tokens, vocab]`` block.
     """
-    device_type = hidden.device.type
-    if hidden.dtype != torch.float32 and device_type not in _MIXED_GEMM_UNSUPPORTED:
-        try:
-            if bias is not None:
-                logits = torch.addmm(bias.float(), hidden, weight.t(), out_dtype=torch.float32)
-            elif row_offset is not None:
-                return torch.addmm(row_offset, hidden, weight.t(), out_dtype=torch.float32)
-            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:
-            _MIXED_GEMM_UNSUPPORTED.add(device_type)
+    if hidden.dtype in _MIXED_GEMM_INPUT_DTYPES and _has_mixed_gemm(hidden.device.type):
+        if bias is not None:
+            logits = torch.addmm(bias.float(), hidden, weight.t(), out_dtype=torch.float32)
+        elif row_offset is not None:
+            return torch.addmm(row_offset, hidden, weight.t(), out_dtype=torch.float32)
+        else:
+            return torch.mm(hidden, weight.t(), out_dtype=torch.float32)
+        return logits if row_offset is None else logits.add_(row_offset)
     # Add the bias inside the GEMM as F.linear does, so a large product that the bias cancels
     # does not overflow a low-precision output first.
     logits = torch.nn.functional.linear(hidden, weight, bias).float()