diff --git a/deepspeed/sequence/linear_cross_entropy.py b/deepspeed/sequence/linear_cross_entropy.py
new file mode 100644
index 000000000000..8aad59e3c173
--- /dev/null
+++ b/deepspeed/sequence/linear_cross_entropy.py
@@ -0,0 +1,278 @@
+# 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. 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 functools
+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, 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
+    # 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_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):
+    """``accumulator += mat1 @ mat2`` with ``accumulator`` kept in its own (fp32) dtype.
+
+    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
+    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))
+
+
+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.
+    """
+    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()
+    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
+    def forward(ctx, hidden, weight, bias, target, tp_group, vocab_start_index, vocab_end_index, ignore_index,
+                chunk_rows, compute_dtype):
+        num_tokens = hidden.shape[0]
+        local_vocab_size = weight.shape[0]
+        tp_active = _tp_active(tp_group)
+        # 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).
+        hidden_compute = hidden.to(compute_dtype)
+        weight_compute = weight.to(compute_dtype)
+        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)
+        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)
+            logits = _fp32_logits(hidden_compute[start:end], weight_compute, bias_compute)
+
+            chunk_max = logits.amax(dim=-1)
+            if tp_active:
+                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)
+
+            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))
+
+        # 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
+        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_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)
+            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)
+            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:
+            # 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. 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
+
+
+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.
+
+    ``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 ("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")
+    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_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():
+        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 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])
+    device_type = hidden.device.type
+    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):
+        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 297f47a87376..92bc686fe621 100644
--- a/docs/_tutorials/autotp-training.md
+++ b/docs/_tutorials/autotp-training.md
@@ -271,6 +271,36 @@ 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. 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
+
+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 `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.
+
 
 ## 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..6cdb25a8d121
--- /dev/null
+++ b/tests/unit/v1/sequence_parallelism/test_vocab_parallel_linear_cross_entropy.py
@@ -0,0 +1,267 @@
+# 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 import linear_cross_entropy
+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", ["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):
+    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)
+    if reduction == "none":
+        expected = expected.view_as(target)
+    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)
+
+    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)
+
+
+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_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_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))
+    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)
+    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_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)
+
+
+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)
+    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"):
+        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"):
+        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
+
+    # 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, chunk_size):
+        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=chunk_size)
+        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)
+
+    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)