From d9b0681809140a041211803e52e70594a18d60d9 Mon Sep 17 00:00:00 2001 From: Siddhartha Raman S Date: Tue, 27 Oct 2026 12:36:47 -0700 Subject: [PATCH 1/5] [PyTorch] Cross entropy: MXFP8 gradient output and faster single-rank launch Add a grad_quantizer argument to parallel_cross_entropy. With an MXFP8Quantizer (E4M3 or E5M2, compact scales), a BF16 input on a single rank, and row count and vocabulary that are multiples of 128, one Triton kernel writes the derivative as an MXFP8 tensor with row-wise and column-wise data, bit-identical to quantizing the BF16 derivative with the quantizer. A following MXFP8 linear layer (an LM head) then skips its own grad-output quantization pass. Other cases compute the derivative as before and quantize it. Launch the single-rank forward with 4096-element blocks and 4 warps and the backward with 4096-element blocks and 8 warps instead of one 32-warp program with up to 32K elements per row, and find the argmax with a min-reduction over the block maximum instead of tl.max(return_indices=True). Only the accumulation order of the softmax statistics changes. Signed-off-by: Siddhartha Raman S Co-Authored-By: Claude Opus 5.5 --- tests/pytorch/test_cross_entropy.py | 205 +++++++++++++++++- .../common/triton/cross_entropy.py | 139 ++++++++++-- transformer_engine/pytorch/cross_entropy.py | 27 ++- .../pytorch/triton/cross_entropy.py | 100 ++++++++- 4 files changed, 446 insertions(+), 25 deletions(-) diff --git a/tests/pytorch/test_cross_entropy.py b/tests/pytorch/test_cross_entropy.py index f34e6666d4b..30d86e13210 100644 --- a/tests/pytorch/test_cross_entropy.py +++ b/tests/pytorch/test_cross_entropy.py @@ -2,10 +2,15 @@ # # See LICENSE for license information. +import math import random import pytest import torch -from transformer_engine.pytorch import parallel_cross_entropy +import transformer_engine.pytorch as te +import transformer_engine_torch as tex +from transformer_engine.common.recipe import Format, MXFP8BlockScaling +from transformer_engine.pytorch import is_mxfp8_available, parallel_cross_entropy +from transformer_engine.pytorch.tensor.mxfp8_tensor import MXFP8Quantizer, MXFP8Tensor from utils import dtype_tols @@ -523,3 +528,201 @@ def test_parallel_cross_entropy_return_argmax_nan(): target = torch.zeros(logits.shape[:-1], dtype=torch.long, device="cuda") _, argmax = parallel_cross_entropy(logits, target, return_argmax=True) torch.testing.assert_close(argmax, torch.argmax(logits, dim=-1), rtol=0, atol=0) + + +mxfp8_available, reason_for_no_mxfp8 = is_mxfp8_available(return_reason=True) +_MXFP8_DTYPES = {"e4m3": tex.DType.kFloat8E4M3, "e5m2": tex.DType.kFloat8E5M2} +_MXFP8_BUFFERS = ( + "_rowwise_data", + "_rowwise_scale_inv", + "_columnwise_data", + "_columnwise_scale_inv", +) + + +def _heavy_tailed_logits(shape, seed): + """BF16 logits whose MXFP8 block scales vary across the row.""" + generator = torch.Generator(device="cuda").manual_seed(seed) + logits = torch.randn(shape, device="cuda", generator=generator) * 3 + logits[..., :64] *= 8 + return logits.to(torch.bfloat16) + + +def _assert_same_mxfp8_bytes(grad, ref): + """Same data and scale inverses. Only the valid part of the scale inverses is compared, + since the quantizer leaves the padding of its scale layout uninitialized.""" + assert isinstance(grad, MXFP8Tensor), type(grad) + rows = math.prod(grad.shape[:-1]) + cols = grad.shape[-1] + valid = { + "_rowwise_scale_inv": (rows, math.ceil(cols / 32)), + "_columnwise_scale_inv": (math.ceil(rows / 32), cols), + } + for name in _MXFP8_BUFFERS: + a, b = getattr(grad, name), getattr(ref, name) + assert (a is None) == (b is None), name + if a is None: + continue + assert a.shape == b.shape, (name, a.shape, b.shape) + if name in valid: + a, b = a[: valid[name][0], : valid[name][1]], b[: valid[name][0], : valid[name][1]] + torch.testing.assert_close(a.view(torch.uint8), b.view(torch.uint8), rtol=0, atol=0) + + +@pytest.mark.skipif(not mxfp8_available, reason=reason_for_no_mxfp8) +@pytest.mark.parametrize( + "shape", + [(2, 64, 4096), (512, 1, 129280), (128, 2, 32000), (96, 1, 4096), (128, 1, 4000)], + ids=["v4096", "v129280", "v32000", "rows96_fallback", "v4000_fallback"], +) +@pytest.mark.parametrize("fp8_dtype", ["e4m3", "e5m2"]) +@pytest.mark.parametrize("usage", ["both", "rowwise", "columnwise"]) +@pytest.mark.parametrize("reduce_loss", [False, True], ids=["none", "mean"]) +@pytest.mark.parametrize("label_smoothing", [0.0, 0.1], ids=["plain", "smoothed"]) +@pytest.mark.parametrize("overwrite_input", [False, True], ids=["safe", "destructive"]) +def test_parallel_cross_entropy_mxfp8_grad( + shape, fp8_dtype, usage, reduce_loss, label_smoothing, overwrite_input +): + """The MXFP8 gradient has the bytes of the BF16 gradient quantized by the quantizer, and + the loss is unchanged.""" + + values = _heavy_tailed_logits(shape, seed=sum(shape)) + target = torch.randint(0, shape[-1], shape[:-1], device="cuda") + target.view(-1)[:3] = -100 + + def run(grad_quantizer): + logits = values.clone().requires_grad_() + loss = parallel_cross_entropy( + logits, + target, + label_smoothing, + reduce_loss, + overwrite_input=overwrite_input, + grad_quantizer=grad_quantizer, + ) + external_grad = ( + torch.full_like(loss, 0.37) + if reduce_loss + else torch.linspace(0.5, 1.5, loss.numel(), device="cuda").reshape(loss.shape) + ) + loss.backward(external_grad) + return loss.detach(), logits.grad + + quantizer = MXFP8Quantizer( + _MXFP8_DTYPES[fp8_dtype], + rowwise=usage != "columnwise", + columnwise=usage != "rowwise", + ) + ref_loss, ref_grad = run(None) + loss, grad = run(quantizer) + + torch.testing.assert_close(loss, ref_loss, rtol=0, atol=0) + _assert_same_mxfp8_bytes(grad, quantizer(ref_grad)) + # The dequantized gradient is the BF16 gradient up to the MXFP8 rounding error: relative + # for normal FP8 values, and absolute (below 2^-16 of the largest magnitude) for values + # that are subnormal or flushed to zero in FP8. + rel_err = 2.0**-3 if fp8_dtype == "e4m3" else 2.0**-2 + abs_err = ref_grad.abs().max().item() * 2.0**-16 + torch.testing.assert_close(grad.dequantize(), ref_grad, rtol=rel_err, atol=abs_err) + + +@pytest.mark.skipif(not mxfp8_available, reason=reason_for_no_mxfp8) +def test_parallel_cross_entropy_mxfp8_grad_unfused_cases(): + """Inputs the fused kernel does not write still return the quantized gradient.""" + + values = torch.randn(2, 64, 4096, device="cuda") + target = torch.randint(0, 4096, (2, 64), device="cuda") + quantizer = MXFP8Quantizer(tex.DType.kFloat8E4M3) + # FP32 input, and a non-contiguous BF16 input. torch.autograd.grad returns the gradient as + # the loss returns it (AccumulateGrad would re-layout it for the non-contiguous leaf). + for logits_values in (values, values.to(torch.bfloat16).transpose(0, 1)): + logits_target = target.reshape(logits_values.shape[:-1]) + logits = logits_values.clone().requires_grad_() + (ref_grad,) = torch.autograd.grad( + parallel_cross_entropy(logits, logits_target).sum(), logits + ) + (grad,) = torch.autograd.grad( + parallel_cross_entropy(logits, logits_target, grad_quantizer=quantizer).sum(), logits + ) + _assert_same_mxfp8_bytes(grad, quantizer(ref_grad)) + + swizzled = MXFP8Quantizer(tex.DType.kFloat8E4M3) + swizzled.optimize_for_gemm = True + logits = values.to(torch.bfloat16).requires_grad_() + parallel_cross_entropy(logits, target, grad_quantizer=swizzled).sum().backward() + assert isinstance(logits.grad, MXFP8Tensor) and logits.grad._with_gemm_swizzled_scales + + internal = MXFP8Quantizer(tex.DType.kFloat8E4M3) + internal.internal = True + with pytest.raises(ValueError, match="internal"): + parallel_cross_entropy(values, target, grad_quantizer=internal) + + +@pytest.mark.skipif(not mxfp8_available, reason=reason_for_no_mxfp8) +@pytest.mark.parametrize("overwrite_input", [False, True], ids=["safe", "destructive"]) +def test_parallel_cross_entropy_mxfp8_grad_lm_head(overwrite_input): + """An MXFP8 linear layer followed by the loss gets the same input and weight gradients when + the loss returns the MXFP8 gradient of the logits.""" + + tokens, hidden, vocab = 512, 1024, 32768 + torch.manual_seed(0) + linear = te.Linear(hidden, vocab, bias=False, params_dtype=torch.bfloat16, device="cuda") + inp = torch.randn(tokens, 1, hidden, device="cuda", dtype=torch.bfloat16) + target = torch.randint(0, vocab, (tokens, 1), device="cuda") + recipe = MXFP8BlockScaling(fp8_format=Format.E4M3) + + def run(grad_quantizer): + linear.weight.grad = None + x = inp.clone().requires_grad_() + with te.autocast(enabled=True, recipe=recipe): + logits = linear(x) + loss = parallel_cross_entropy( + logits, + target, + overwrite_input=overwrite_input, + grad_quantizer=grad_quantizer, + ) + loss.sum().backward() + return loss.detach(), x.grad, linear.weight.grad.clone() + + ref = run(None) + out = run(MXFP8Quantizer(tex.DType.kFloat8E4M3)) + for a, b in zip(out, ref): + torch.testing.assert_close(a, b, rtol=0, atol=0) + + +@pytest.mark.skipif(not mxfp8_available, reason=reason_for_no_mxfp8) +def test_parallel_cross_entropy_mxfp8_grad_cuda_graph(): + """The MXFP8 backward is CUDA graph capturable.""" + + shape = (256, 1, 8192) + static_values = _heavy_tailed_logits(shape, seed=4) + static_target = torch.randint(0, shape[-1], shape[:-1], device="cuda") + quantizer = MXFP8Quantizer(tex.DType.kFloat8E4M3) + buffer = torch.empty_like(static_values) + outputs = {} + + def step(): + buffer.copy_(static_values) + logits = buffer.detach().requires_grad_() + loss = parallel_cross_entropy( + logits, static_target, overwrite_input=True, grad_quantizer=quantizer + ) + outputs["grad"] = torch.autograd.grad(loss.sum(), logits)[0] + + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(stream): + for _ in range(2): + step() + torch.cuda.current_stream().wait_stream(stream) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + step() + + static_values.copy_(_heavy_tailed_logits(shape, seed=5)) + static_target.copy_(torch.randint(0, shape[-1], shape[:-1], device="cuda")) + graph.replay() + logits = static_values.clone().requires_grad_() + parallel_cross_entropy(logits, static_target).sum().backward() + _assert_same_mxfp8_bytes(outputs["grad"], quantizer(logits.grad)) diff --git a/transformer_engine/common/triton/cross_entropy.py b/transformer_engine/common/triton/cross_entropy.py index 3d1fd531782..8a87e854e3d 100644 --- a/transformer_engine/common/triton/cross_entropy.py +++ b/transformer_engine/common/triton/cross_entropy.py @@ -15,30 +15,25 @@ def _is_nan(x): @triton.jit -def _fold_block_argmax(x, i, offsets, n_cols, run_max, run_idx, run_nan): +def _fold_block_argmax(x, offsets, n_cols, run_max, run_idx, run_nan): """Fold one block of a row into a running argmax with torch.argmax's semantics: NaN counts as the maximum, and the first maximum (or the first NaN) wins. - tl.max(..., return_indices=True) is not NaN-safe (its comparisons are false for NaN, so a - block with a NaN can give a pair that is neither the maximum nor the NaN). NaN is therefore - found separately and the maximum is taken over the other values. run_max is kept apart from - the loss's running maximum, which is unchanged.""" + The maximum is taken over the non-NaN values and its first index is found with one + min-reduction, which is cheaper than tl.max(..., return_indices=True) and, unlike it, safe + with NaN. The position of a NaN is only searched for in the first block that has one. + run_max is kept apart from the loss's running maximum, which is unchanged.""" is_nan = _is_nan(x) - block_has_nan = tl.max(is_nan.to(tl.int32)) > 0 - first_nan = tl.min(tl.where(is_nan, offsets, n_cols)) - block_max, block_idx = tl.max( - tl.where(is_nan, float("-inf"), x), - axis=0, - return_indices=True, - return_indices_tie_break_left=True, - ) + x_no_nan = tl.where(is_nan, float("-inf"), x) + block_max = tl.max(x_no_nan) if run_nan == 0: + block_has_nan = tl.max(is_nan.to(tl.int32)) > 0 if block_has_nan: - run_idx = first_nan + run_idx = tl.min(tl.where(is_nan, offsets, n_cols)) run_nan = 1 elif block_max > run_max: # Only a strictly larger maximum replaces an earlier one: first index wins. - run_idx = i + block_idx + run_idx = tl.min(tl.where(x_no_nan == block_max, offsets, n_cols)) run_max = block_max return run_max, run_idx, run_nan @@ -93,7 +88,7 @@ def cross_entropy_forward_kernel( x = x.to(tl.float32) if COMPUTE_ARGMAX: argmax_max, argmax, argmax_nan = _fold_block_argmax( - x, i, offsets, n_cols, argmax_max, argmax, argmax_nan + x, offsets, n_cols, argmax_max, argmax, argmax_nan ) block_max = tl.max(x) m_new = tl.maximum(m, block_max) @@ -179,7 +174,7 @@ def cross_entropy_tp_pre_kernel( x = x.to(tl.float32) if COMPUTE_ARGMAX: argmax_max, argmax, argmax_nan = _fold_block_argmax( - x, i, offsets, n_cols, argmax_max, argmax, argmax_nan + x, offsets, n_cols, argmax_max, argmax, argmax_nan ) block_max = tl.max(x) m_new = tl.maximum(m, block_max) @@ -318,3 +313,113 @@ def cross_entropy_backward_kernel( is_target = target_is_local & (offsets == target_col) grad -= tl.where(is_target, 1 - label_smoothing, 0.0) tl.store(saved_input_ptr + offsets, grad * grad_output, mask=mask) + + +@triton.jit +def _mxfp8_scale_exponent(amax, MAX_NORM: tl.constexpr): + """E8M0 block scale of TE's MXFP8 cast: amax / MAX_NORM in FP32, rounded up to a power of + two (saturating to 0xFE), with 0xFF for NaN.""" + val = amax * (1.0 / MAX_NORM) + bits = val.to(tl.int32, bitcast=True) + exponent = (bits >> 23) & 0xFF + mantissa = bits & 0x7FFFFF + round_up = (mantissa > 0) & (exponent < 0xFE) & ~((exponent == 0) & (mantissa <= 0x400000)) + exponent = tl.where(round_up, exponent + 1, exponent) + exponent = tl.where(exponent == 0xFF, tl.where(mantissa == 0, 0xFE, 0xFF), exponent) + return exponent + + +@triton.jit +def _mxfp8_scale_reciprocal(exponent): + """2^-(exponent - 127) as FP32 (NaN for 0xFF), as TE's exp2f_rcp.""" + bits = tl.where(exponent == 254, 0x00400000, (254 - exponent) << 23) + bits = tl.where(exponent == 255, 0x7FFFFFFF, bits) + return bits.to(tl.float32, bitcast=True) + + +@triton.jit +def cross_entropy_backward_mxfp8_kernel( + saved_input_ptr, + Y_ptr, + stats_ptr, + n_non_ignore_ptr, + grad_output_ptr, + grad_output_stride, + rowwise_data_ptr, + rowwise_scale_ptr, + rowwise_scale_stride, + columnwise_data_ptr, + columnwise_scale_ptr, + columnwise_scale_stride, + n_cols, + ignore_idx, + reduce_loss: tl.constexpr, + label_smoothing: tl.constexpr, + ROWWISE: tl.constexpr, + COLUMNWISE: tl.constexpr, + FP8_E5M2: tl.constexpr, + BLOCK_COLS: tl.constexpr, +): + """Single-rank derivative of a 32-row x BLOCK_COLS tile, written as MXFP8. + + The derivative is computed as cross_entropy_backward_kernel computes it and rounded to the + input dtype (BF16), then quantized as TE's MXFP8 cast quantizes it: row-wise with 1x32 + blocks along the vocabulary and column-wise with 32x1 blocks along the tokens.""" + pid_r = tl.program_id(0) + pid_c = tl.program_id(1) + rows = (pid_r * 32 + tl.arange(0, 32)).to(tl.int64) + cols = pid_c * BLOCK_COLS + tl.arange(0, BLOCK_COLS) + offsets = rows[:, None] * n_cols + cols[None, :] + + y = tl.load(Y_ptr + rows) + m = tl.load(stats_ptr + rows * 2) + d = tl.load(stats_ptr + rows * 2 + 1) + grad_output = tl.load(grad_output_ptr + rows * grad_output_stride).to(tl.float32) + if reduce_loss: + grad_output = grad_output / tl.load(n_non_ignore_ptr) + + eps = label_smoothing / n_cols + x = tl.load(saved_input_ptr + offsets).to(tl.float32) + grad = tl.exp(x - m[:, None]) / d[:, None] - eps + grad -= tl.where(cols[None, :] == y[:, None], 1 - label_smoothing, 0.0) + grad = grad * grad_output[:, None] + grad = tl.where((y == ignore_idx)[:, None], 0.0, grad) + grad = grad.to(saved_input_ptr.dtype.element_ty).to(tl.float32) + grad_abs = tl.abs(grad) + + if FP8_E5M2: + MAX_NORM: tl.constexpr = 57344.0 + else: + MAX_NORM: tl.constexpr = 448.0 + + if ROWWISE: + N_BLOCKS: tl.constexpr = BLOCK_COLS // 32 + amax = tl.max(tl.reshape(grad_abs, (32, N_BLOCKS, 32)), axis=2) + exponent = _mxfp8_scale_exponent(amax, MAX_NORM) + scaled = tl.reshape( + tl.reshape(grad, (32, N_BLOCKS, 32)) * _mxfp8_scale_reciprocal(exponent)[:, :, None], + (32, BLOCK_COLS), + ) + if FP8_E5M2: + data = scaled.to(tl.float8e5) + else: + data = scaled.to(tl.float8e4nv) + tl.store(rowwise_data_ptr + offsets, data.to(tl.uint8, bitcast=True)) + scale_cols = pid_c * N_BLOCKS + tl.arange(0, N_BLOCKS) + tl.store( + rowwise_scale_ptr + rows[:, None] * rowwise_scale_stride + scale_cols[None, :], + exponent.to(tl.uint8), + ) + + if COLUMNWISE: + exponent = _mxfp8_scale_exponent(tl.max(grad_abs, axis=0), MAX_NORM) + scaled = grad * _mxfp8_scale_reciprocal(exponent)[None, :] + if FP8_E5M2: + data = scaled.to(tl.float8e5) + else: + data = scaled.to(tl.float8e4nv) + tl.store(columnwise_data_ptr + offsets, data.to(tl.uint8, bitcast=True)) + tl.store( + columnwise_scale_ptr + pid_r.to(tl.int64) * columnwise_scale_stride + cols, + exponent.to(tl.uint8), + ) diff --git a/transformer_engine/pytorch/cross_entropy.py b/transformer_engine/pytorch/cross_entropy.py index 8ae3b853801..feee673d5c1 100644 --- a/transformer_engine/pytorch/cross_entropy.py +++ b/transformer_engine/pytorch/cross_entropy.py @@ -11,6 +11,7 @@ import transformer_engine.pytorch.triton.cross_entropy as triton_cross_entropy from transformer_engine.pytorch.jit import no_torch_dynamo +from transformer_engine.pytorch.quantized_tensor import Quantizer __all__ = ["parallel_cross_entropy"] @@ -33,6 +34,7 @@ def forward( ignore_idx=-100, overwrite_input=False, return_argmax=False, + grad_quantizer=None, ): """Compute the loss and save the input and softmax statistics for backward.""" @@ -65,6 +67,7 @@ def forward( ctx.world_size = world_size ctx.ignore_idx = ignore_idx ctx.overwrite_input = overwrite_input + ctx.grad_quantizer = grad_quantizer ctx.did_backward = False if return_argmax: ctx.mark_non_differentiable(argmax) @@ -94,10 +97,11 @@ def backward(ctx, grad_output, *_): ctx.rank, ctx.world_size, ctx.ignore_idx, + ctx.grad_quantizer, ) if ctx.overwrite_input: torch.autograd.graph.increment_version(saved_input) - return grad_input, None, None, None, None, None, None, None + return grad_input, None, None, None, None, None, None, None, None def _validate_inputs( @@ -105,6 +109,7 @@ def _validate_inputs( target: torch.Tensor, label_smoothing: float, overwrite_input: bool, + grad_quantizer: Optional[Quantizer] = None, ) -> None: """Validate assumptions made by the cross entropy Triton kernels.""" @@ -129,6 +134,11 @@ def _validate_inputs( raise ValueError(f"label_smoothing must be in [0, 1], but got {label_smoothing}") if overwrite_input and not inp.is_contiguous(): raise ValueError("overwrite_input=True requires a contiguous input") + if grad_quantizer is not None: + if not isinstance(grad_quantizer, Quantizer): + raise TypeError(f"grad_quantizer must be a Quantizer, but got {type(grad_quantizer)}") + if grad_quantizer.internal: + raise ValueError("grad_quantizer must create QuantizedTensors (internal=False)") @no_torch_dynamo() @@ -140,6 +150,7 @@ def _parallel_cross_entropy_overwrite_input( dist_process_group: Optional[torch.distributed.ProcessGroup], ignore_idx: int, return_argmax: bool, + grad_quantizer: Optional[Quantizer], ) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]: """Run destructive cross entropy outside Torch Dynamo's compiled graph.""" @@ -152,6 +163,7 @@ def _parallel_cross_entropy_overwrite_input( ignore_idx, True, return_argmax, + grad_quantizer, ) @@ -166,6 +178,7 @@ def parallel_cross_entropy( *, overwrite_input: bool = False, return_argmax: bool = False, + grad_quantizer: Optional[Quantizer] = None, _input: Optional[torch.Tensor] = None, ) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]: """Cross entropy loss with optional distributed reduction. @@ -211,6 +224,14 @@ def parallel_cross_entropy( if several logits are equal. Rows with an ignored target get one too. As with ``torch.argmax``, a NaN logit counts as the maximum: a row with NaN logits gets the index of its first NaN (and a NaN loss). + grad_quantizer : Quantizer, default = None + Return the gradient of ``inp`` quantized with this quantizer, for a following + quantized linear layer (for example an MXFP8 LM head) that would otherwise quantize + it itself. With an ``MXFP8Quantizer`` (E4M3 or E5M2, compact scales), a BF16 ``inp`` + without ``dist_process_group`` (or with a single rank), a row count that is a multiple + of 128 and a vocabulary that is a multiple of 128, backward writes the MXFP8 gradient + directly, bit-identical to quantizing the BF16 gradient with ``grad_quantizer``. In + all other cases the gradient is computed as usual and then quantized. Returns ------- @@ -233,7 +254,7 @@ def parallel_cross_entropy( FutureWarning, ) - _validate_inputs(inp, target, label_smoothing, overwrite_input) + _validate_inputs(inp, target, label_smoothing, overwrite_input, grad_quantizer) if overwrite_input: return _parallel_cross_entropy_overwrite_input( inp, @@ -243,6 +264,7 @@ def parallel_cross_entropy( dist_process_group, ignore_idx, return_argmax, + grad_quantizer, ) return CrossEntropyFunction.apply( inp, @@ -253,4 +275,5 @@ def parallel_cross_entropy( ignore_idx, overwrite_input, return_argmax, + grad_quantizer, ) diff --git a/transformer_engine/pytorch/triton/cross_entropy.py b/transformer_engine/pytorch/triton/cross_entropy.py index 3bc586c4ece..235fb455d28 100644 --- a/transformer_engine/pytorch/triton/cross_entropy.py +++ b/transformer_engine/pytorch/triton/cross_entropy.py @@ -15,11 +15,26 @@ cross_entropy_tp_pre_kernel, cross_entropy_tp_post_kernel, cross_entropy_backward_kernel, + cross_entropy_backward_mxfp8_kernel, ) # The optimal maximum block size depends on your hardware, your kernel, and your dtype MAX_FUSED_SIZE = 65536 // 2 +# Launch configuration of the single-rank forward kernel and of the backward kernel, which run +# one program per row. Small blocks with few warps keep more rows in flight per SM: one +# 1024-thread program per row reaches well under half of the HBM bandwidth at a 129k +# vocabulary. Only the order in which the softmax statistics are accumulated changes. The +# forward uses one block size with and without the argmax so that the loss does not depend on +# return_argmax. +_FWD_BLOCK_SIZE = 4096 +_FWD_NUM_WARPS = 4 +_BWD_BLOCK_SIZE = 4096 +_BWD_NUM_WARPS = 8 +# Tile width and warps of the backward that writes an MXFP8 gradient (32-row tiles). +_MXFP8_BWD_BLOCK_COLS = 256 +_MXFP8_BWD_NUM_WARPS = 4 + def cross_entropy_forward( _input: torch.Tensor, @@ -76,8 +91,8 @@ def cross_entropy_forward( COUNT_NON_IGNORE=reduce_loss, COPY_INPUT=not overwrite_input, COMPUTE_ARGMAX=compute_argmax, - BLOCK_SIZE=BLOCK_SIZE, - num_warps=32, + BLOCK_SIZE=min(_FWD_BLOCK_SIZE, BLOCK_SIZE), + num_warps=_FWD_NUM_WARPS, ) else: # Per row and rank: max, sum of exponentials, target logit, logit sum, and with @@ -137,6 +152,33 @@ def cross_entropy_forward( return loss, saved_input, stats, target, n_non_ignore, rank, world_size, argmax +def _mxfp8_backward_block_cols( + saved_input: torch.Tensor, world_size: int, grad_quantizer +) -> Optional[int]: + """Tile width of the fused MXFP8 backward, or None if it cannot write this gradient.""" + # pylint: disable=import-outside-toplevel + from transformer_engine.pytorch.tensor.mxfp8_tensor import MXFP8Quantizer + import transformer_engine_torch as tex + + if not isinstance(grad_quantizer, MXFP8Quantizer): + return None + if grad_quantizer.dtype not in (tex.DType.kFloat8E4M3, tex.DType.kFloat8E5M2): + return None + # 2D blocks and GEMM-swizzled scales are left to the quantizer. + if grad_quantizer.with_2d_quantization or grad_quantizer.optimize_for_gemm: + return None + if not (grad_quantizer.rowwise_usage or grad_quantizer.columnwise_usage): + return None + if world_size != 1 or saved_input.dtype != torch.bfloat16 or not saved_input.is_contiguous(): + return None + B, SQ, V = saved_input.shape + # Unpadded scale-inverse layouts: 128 rows x 4 blocks of the row-wise scales and + # 4 x 128 columns of the column-wise scales. + if (B * SQ) % 128 != 0 or V % 128 != 0: + return None + return _MXFP8_BWD_BLOCK_COLS if V % _MXFP8_BWD_BLOCK_COLS == 0 else 128 + + def cross_entropy_backward( saved_input: torch.Tensor, stats: torch.Tensor, @@ -148,14 +190,60 @@ def cross_entropy_backward( rank: int, world_size: int, ignore_idx: int, + grad_quantizer=None, ): - """Reconstruct the derivative in FP32 and overwrite the saved input buffer.""" + """Reconstruct the derivative in FP32. + + Without ``grad_quantizer``, the derivative overwrites the saved input buffer, which is + returned. With it, the derivative is returned quantized by ``grad_quantizer``: a BF16 input + with an MXFP8 quantizer is written as MXFP8 directly by one kernel, and anything else is + computed in the saved input buffer and then quantized. + """ B, SQ, V = saved_input.shape n_rows = B * SQ BLOCK_SIZE = min(MAX_FUSED_SIZE, triton.next_power_of_2(V)) grad_output = grad_output.contiguous() + if grad_quantizer is not None: + block_cols = _mxfp8_backward_block_cols(saved_input, world_size, grad_quantizer) + if block_cols is not None: + import transformer_engine_torch as tex # pylint: disable=import-outside-toplevel + + grad_input = grad_quantizer.make_empty( + saved_input.shape, dtype=saved_input.dtype, device=saved_input.device + ) + rowwise, columnwise = ( + grad_quantizer.rowwise_usage, + grad_quantizer.columnwise_usage, + ) + rowwise_scale = grad_input._rowwise_scale_inv if rowwise else None + columnwise_scale = grad_input._columnwise_scale_inv if columnwise else None + cross_entropy_backward_mxfp8_kernel[(n_rows // 32, V // block_cols)]( + saved_input_ptr=saved_input, + Y_ptr=target, + stats_ptr=stats, + n_non_ignore_ptr=n_non_ignore if n_non_ignore is not None else stats, + grad_output_ptr=grad_output, + grad_output_stride=0 if reduce_loss else 1, + rowwise_data_ptr=grad_input._rowwise_data if rowwise else saved_input, + rowwise_scale_ptr=rowwise_scale if rowwise else saved_input, + rowwise_scale_stride=rowwise_scale.shape[-1] if rowwise else 0, + columnwise_data_ptr=(grad_input._columnwise_data if columnwise else saved_input), + columnwise_scale_ptr=columnwise_scale if columnwise else saved_input, + columnwise_scale_stride=columnwise_scale.shape[-1] if columnwise else 0, + n_cols=V, + ignore_idx=ignore_idx, + reduce_loss=reduce_loss, + label_smoothing=label_smoothing, + ROWWISE=rowwise, + COLUMNWISE=columnwise, + FP8_E5M2=grad_quantizer.dtype == tex.DType.kFloat8E5M2, + BLOCK_COLS=block_cols, + num_warps=_MXFP8_BWD_NUM_WARPS, + ) + return grad_input + cross_entropy_backward_kernel[(n_rows,)]( saved_input_ptr=saved_input, Y_ptr=target, @@ -169,7 +257,9 @@ def cross_entropy_backward( ignore_idx=ignore_idx, reduce_loss=reduce_loss, label_smoothing=label_smoothing, - BLOCK_SIZE=BLOCK_SIZE, - num_warps=32, + BLOCK_SIZE=min(_BWD_BLOCK_SIZE, BLOCK_SIZE), + num_warps=_BWD_NUM_WARPS, ) + if grad_quantizer is not None: + return grad_quantizer(saved_input) return saved_input From 4c6363ecbc02af8fde428855a335f4e3f07a62f3 Mon Sep 17 00:00:00 2001 From: Siddhartha Raman S Date: Tue, 27 Oct 2026 13:04:36 -0700 Subject: [PATCH 2/5] [PyTorch] Cross entropy: drop non-finite special cases from the MXFP8 scale The NaN/Inf selects in the E8M0 exponent and reciprocal helpers slowed the fused MXFP8 backward from about 340 us to 420 us at 4096 x 129280 on a Rubin GPU. Finite gradients, the only ones a training step uses, still get the bytes of TE's MXFP8 cast; this is now stated in the docstring. Signed-off-by: Siddhartha Raman S Co-Authored-By: Claude Opus 5.5 --- .../common/triton/cross_entropy.py | 25 ++++++++----------- transformer_engine/pytorch/cross_entropy.py | 5 ++-- 2 files changed, 14 insertions(+), 16 deletions(-) diff --git a/transformer_engine/common/triton/cross_entropy.py b/transformer_engine/common/triton/cross_entropy.py index 8a87e854e3d..964c0402b44 100644 --- a/transformer_engine/common/triton/cross_entropy.py +++ b/transformer_engine/common/triton/cross_entropy.py @@ -316,24 +316,21 @@ def cross_entropy_backward_kernel( @triton.jit -def _mxfp8_scale_exponent(amax, MAX_NORM: tl.constexpr): - """E8M0 block scale of TE's MXFP8 cast: amax / MAX_NORM in FP32, rounded up to a power of - two (saturating to 0xFE), with 0xFF for NaN.""" - val = amax * (1.0 / MAX_NORM) +def _mxfp8_scale_exponent(amax, MAX_NORM_RCP: tl.constexpr): + """E8M0 block scale of TE's MXFP8 cast: amax * (1 / max_norm) in FP32, rounded up to a power + of two. Finite gradients give the bytes of TE's cast; non-finite ones are not special-cased.""" + val = amax * MAX_NORM_RCP bits = val.to(tl.int32, bitcast=True) exponent = (bits >> 23) & 0xFF mantissa = bits & 0x7FFFFF - round_up = (mantissa > 0) & (exponent < 0xFE) & ~((exponent == 0) & (mantissa <= 0x400000)) - exponent = tl.where(round_up, exponent + 1, exponent) - exponent = tl.where(exponent == 0xFF, tl.where(mantissa == 0, 0xFE, 0xFF), exponent) - return exponent + round_up = (mantissa > 0) & (exponent != 0xFE) & ~((exponent == 0) & (mantissa <= 0x400000)) + return tl.where(round_up, exponent + 1, exponent) @triton.jit def _mxfp8_scale_reciprocal(exponent): - """2^-(exponent - 127) as FP32 (NaN for 0xFF), as TE's exp2f_rcp.""" + """2^-(exponent - 127) as FP32, as TE's exp2f_rcp for finite scales.""" bits = tl.where(exponent == 254, 0x00400000, (254 - exponent) << 23) - bits = tl.where(exponent == 255, 0x7FFFFFFF, bits) return bits.to(tl.float32, bitcast=True) @@ -388,14 +385,14 @@ def cross_entropy_backward_mxfp8_kernel( grad_abs = tl.abs(grad) if FP8_E5M2: - MAX_NORM: tl.constexpr = 57344.0 + MAX_NORM_RCP: tl.constexpr = 1.0 / 57344.0 else: - MAX_NORM: tl.constexpr = 448.0 + MAX_NORM_RCP: tl.constexpr = 1.0 / 448.0 if ROWWISE: N_BLOCKS: tl.constexpr = BLOCK_COLS // 32 amax = tl.max(tl.reshape(grad_abs, (32, N_BLOCKS, 32)), axis=2) - exponent = _mxfp8_scale_exponent(amax, MAX_NORM) + exponent = _mxfp8_scale_exponent(amax, MAX_NORM_RCP) scaled = tl.reshape( tl.reshape(grad, (32, N_BLOCKS, 32)) * _mxfp8_scale_reciprocal(exponent)[:, :, None], (32, BLOCK_COLS), @@ -412,7 +409,7 @@ def cross_entropy_backward_mxfp8_kernel( ) if COLUMNWISE: - exponent = _mxfp8_scale_exponent(tl.max(grad_abs, axis=0), MAX_NORM) + exponent = _mxfp8_scale_exponent(tl.max(grad_abs, axis=0), MAX_NORM_RCP) scaled = grad * _mxfp8_scale_reciprocal(exponent)[None, :] if FP8_E5M2: data = scaled.to(tl.float8e5) diff --git a/transformer_engine/pytorch/cross_entropy.py b/transformer_engine/pytorch/cross_entropy.py index feee673d5c1..2505a317043 100644 --- a/transformer_engine/pytorch/cross_entropy.py +++ b/transformer_engine/pytorch/cross_entropy.py @@ -230,8 +230,9 @@ def parallel_cross_entropy( it itself. With an ``MXFP8Quantizer`` (E4M3 or E5M2, compact scales), a BF16 ``inp`` without ``dist_process_group`` (or with a single rank), a row count that is a multiple of 128 and a vocabulary that is a multiple of 128, backward writes the MXFP8 gradient - directly, bit-identical to quantizing the BF16 gradient with ``grad_quantizer``. In - all other cases the gradient is computed as usual and then quantized. + directly, bit-identical to quantizing the BF16 gradient with ``grad_quantizer`` + (for finite gradients). In all other cases the gradient is computed as usual and + then quantized. Returns ------- From e41b2bea6a8691d4a699e1c1694e6fa3e61eefe0 Mon Sep 17 00:00:00 2001 From: Siddhartha Raman S Date: Wed, 11 Nov 2026 11:04:41 -0800 Subject: [PATCH 3/5] [PyTorch] Cross entropy: keep masked logits from turning the softmax sum into NaN When every logit seen so far in a row is -inf, the online softmax shifted the exponentials by a running maximum of -inf, so exp(-inf - (-inf)) made the sum of exponentials NaN for the rest of the row. With 4096-wide forward blocks this happens for any row whose first block is masked. Shift by 0 until a finite maximum is seen, in the single-rank and tensor-parallel forward kernels and in the cross-rank merge. Rows with a finite maximum compute exactly as before; rows whose logits are all -inf still give NaN, as in PyTorch. Add tests with masked logits against PyTorch, including the MXFP8 gradient, CUDA graph replay and tensor parallelism. Signed-off-by: Siddhartha Raman S Co-Authored-By: Claude Opus 5.5 --- .../test_parallel_cross_entropy.py | 87 +++++++++ tests/pytorch/test_cross_entropy.py | 171 +++++++++++++++++- .../common/triton/cross_entropy.py | 18 +- 3 files changed, 269 insertions(+), 7 deletions(-) diff --git a/tests/pytorch/distributed/test_parallel_cross_entropy.py b/tests/pytorch/distributed/test_parallel_cross_entropy.py index c4e3510c174..eb776fb07d8 100644 --- a/tests/pytorch/distributed/test_parallel_cross_entropy.py +++ b/tests/pytorch/distributed/test_parallel_cross_entropy.py @@ -153,6 +153,76 @@ def _run_tensor_parallel_argmax(rank, world_size, init_file, label_smoothing): dist.destroy_process_group() +def _run_tensor_parallel_masked(rank, world_size, init_file, label_smoothing): + """Two-rank worker: masked (-inf) logits, including shards and blocks masked entirely.""" + + torch.cuda.set_device(rank) + device = torch.device("cuda", rank) + dist.init_process_group( + backend="nccl", + init_method=f"file://{init_file}", + rank=rank, + world_size=world_size, + device_id=device, + ) + try: + generator = torch.Generator().manual_seed(2027) + for shape, block in (((1, 5, 22), 8), ((1, 5, 80000), 32768)): + vocab = shape[-1] + local_vocab = vocab // world_size + values = torch.randn(shape, generator=generator) + values[0, 0, :local_vocab] = -float("inf") # the shard of rank 0 is masked + values[0, 1, local_vocab:] = -float("inf") # the shard of rank 1 is masked + for start in range(0, vocab, local_vocab): + values[0, 2, start : start + block] = -float("inf") # each first block + values[0, 3, :-1] = -float("inf") # only the last logit is unmasked + target = torch.multinomial( + values.isfinite().float().reshape(-1, vocab), 1, generator=generator + ).reshape(shape[:-1]) + values, target = values.to(device), target.to(device) + vocab_start = rank * local_vocab + for dtype in (torch.float32, torch.bfloat16): + global_values = values.to(dtype) + local_values = global_values[..., vocab_start : vocab_start + local_vocab] + for reduce_loss in (False, True): + for overwrite_input in (False, True): + local_logits = local_values.clone().requires_grad_() + ref_logits = global_values.double().clone().requires_grad_() + loss, argmax = parallel_cross_entropy( + local_logits, + target, + label_smoothing=label_smoothing, + reduce_loss=reduce_loss, + dist_process_group=dist.group.WORLD, + overwrite_input=overwrite_input, + return_argmax=True, + ) + ref_loss = torch.nn.functional.cross_entropy( + ref_logits.reshape(-1, vocab), + target.reshape(-1), + label_smoothing=label_smoothing, + reduction="mean" if reduce_loss else "none", + ).reshape_as(loss) + loss_grad = torch.full_like(loss, 0.37) + loss.backward(loss_grad) + ref_loss.backward(loss_grad.double()) + torch.testing.assert_close( + loss, ref_loss.float(), **dtype_tols(torch.float32) + ) + expected_grad = ref_logits.grad[ + ..., vocab_start : vocab_start + local_vocab + ].to(dtype) + assert local_logits.grad.isfinite().all() + torch.testing.assert_close( + local_logits.grad, expected_grad, **dtype_tols(dtype) + ) + torch.testing.assert_close( + argmax, torch.argmax(global_values, dim=-1), rtol=0, atol=0 + ) + finally: + dist.destroy_process_group() + + @pytest.mark.parametrize("label_smoothing", [0.0, 0.1], ids=["plain", "smoothed"]) def test_parallel_cross_entropy_tensor_parallel(label_smoothing): """Validate tensor-parallel loss and gradients on two ranks.""" @@ -185,3 +255,20 @@ def test_parallel_cross_entropy_tensor_parallel_argmax(label_smoothing): nprocs=world_size, join=True, ) + + +@pytest.mark.parametrize("label_smoothing", [0.0, 0.1], ids=["plain", "smoothed"]) +def test_parallel_cross_entropy_tensor_parallel_masked(label_smoothing): + """Validate tensor-parallel loss and gradients with masked logits on two ranks.""" + + if torch.cuda.device_count() < 2: + pytest.skip("tensor-parallel cross entropy test requires two CUDA devices") + world_size = 2 + with tempfile.TemporaryDirectory() as temp_dir: + init_file = os.path.join(temp_dir, "distributed_init") + mp.spawn( + _run_tensor_parallel_masked, + args=(world_size, init_file, label_smoothing), + nprocs=world_size, + join=True, + ) diff --git a/tests/pytorch/test_cross_entropy.py b/tests/pytorch/test_cross_entropy.py index d75e50e8c1e..b81949cf5ad 100644 --- a/tests/pytorch/test_cross_entropy.py +++ b/tests/pytorch/test_cross_entropy.py @@ -530,6 +530,112 @@ def test_parallel_cross_entropy_return_argmax_nan(): torch.testing.assert_close(argmax, torch.argmax(logits, dim=-1), rtol=0, atol=0) +def _masked_logits(n_rows, vocab, dtype, seed): + """Logits with masked (-inf) entries and a finite maximum in every row, and targets that are + unmasked. Rows start with masked blocks of the forward kernels (4096 wide on a single rank, + 32K wide before), keep only one logit, or are masked elsewhere. Row patterns repeat every + eight rows.""" + + generator = torch.Generator(device="cuda").manual_seed(seed) + values = torch.randn((n_rows, vocab), device="cuda", generator=generator) + values[0::8, :4096] = -math.inf # the first block is masked + values[1::8, : vocab // 2] = -math.inf # the first half is masked + values[2::8, :-1] = -math.inf # only the last logit is unmasked + values[3::8, :] = -math.inf + values[3::8, 4100] = 0.0 # only one logit in the second block is unmasked + values[4::8, ::2] = -math.inf # no block is masked entirely + values[5::8, -4096:] = -math.inf # the last block is masked + values[6::8, :-100] = -math.inf # only the last 100 logits are unmasked + values = values.to(dtype) + target = torch.multinomial(values.isfinite().float(), 1, generator=generator).squeeze(-1) + return values, target + + +def _masked_reference(values, target, label_smoothing, reduce_loss, loss_grad): + """Loss and gradient of torch.nn.functional.cross_entropy in FP64.""" + logits = values.double().reshape(-1, values.shape[-1]).requires_grad_() + loss = torch.nn.functional.cross_entropy( + logits, + target.reshape(-1), + label_smoothing=label_smoothing, + reduction="mean" if reduce_loss else "none", + ).reshape(loss_grad.shape) + loss.backward(loss_grad.double()) + return loss.detach(), logits.grad.reshape(values.shape) + + +@pytest.mark.parametrize("vocab", [8192, 12288, 70000, 129280]) +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16], ids=["fp32", "bf16"]) +@pytest.mark.parametrize("reduce_loss", [False, True], ids=["none", "mean"]) +@pytest.mark.parametrize("label_smoothing", [0.0, 0.1], ids=["plain", "smoothed"]) +@pytest.mark.parametrize("overwrite_input", [False, True], ids=["safe", "destructive"]) +def test_parallel_cross_entropy_masked_logits( + vocab, dtype, reduce_loss, label_smoothing, overwrite_input +): + """Masked (-inf) logits give the loss and gradient of PyTorch, including rows whose first + blocks are masked entirely. With label smoothing the loss of a row with masked logits is + +inf, as in PyTorch, and its gradient is finite.""" + + values, target = _masked_logits(16, vocab, dtype, seed=vocab) + values, target = values.reshape(2, 8, vocab), target.reshape(2, 8) + target[1, 7] = -100 + loss_grad = ( + torch.full((1,), 0.37, device="cuda") + if reduce_loss + else torch.linspace(0.5, 1.5, target.numel(), device="cuda").reshape(target.shape) + ) + ref_loss, ref_grad = _masked_reference(values, target, label_smoothing, reduce_loss, loss_grad) + + results = [] + for return_argmax in (False, True): + logits = values.clone().requires_grad_() + out = parallel_cross_entropy( + logits, + target, + label_smoothing, + reduce_loss, + overwrite_input=overwrite_input, + return_argmax=return_argmax, + ) + loss, argmax = out if return_argmax else (out, None) + loss.backward(loss_grad) + results.append((loss.detach(), logits.grad, argmax)) + (loss, grad, _), (argmax_loss, argmax_grad, argmax) = results + + assert grad.isfinite().all() + torch.testing.assert_close(loss, ref_loss.float(), **dtype_tols(torch.float32)) + torch.testing.assert_close(grad, ref_grad.to(dtype), **dtype_tols(dtype)) + torch.testing.assert_close(argmax, torch.argmax(values, dim=-1), rtol=0, atol=0) + torch.testing.assert_close(argmax_loss, loss, rtol=0, atol=0) + torch.testing.assert_close(argmax_grad, grad, rtol=0, atol=0) + + +@pytest.mark.parametrize("vocab", [4096, 8192, 70000]) +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16], ids=["fp32", "bf16"]) +def test_parallel_cross_entropy_fully_masked_row(vocab, dtype): + """A row whose logits are all -inf gets a NaN loss and gradient, as in PyTorch, without + affecting the other rows. Its argmax is 0, and an ignored one gets a zero loss and + gradient.""" + + values = torch.randn((1, 4, vocab), device="cuda").to(dtype) + values[0, 1:3] = -math.inf + target = torch.tensor([[1, 2, -100, 3]], device="cuda") + for return_argmax in (False, True): + logits = values.clone().requires_grad_() + out = parallel_cross_entropy(logits, target, return_argmax=return_argmax) + loss, argmax = out if return_argmax else (out, None) + loss.backward(torch.ones_like(loss)) + ref_loss = torch.nn.functional.cross_entropy( + values[0].double(), target[0], reduction="none" + ) + assert ref_loss[1].isnan() + assert loss[0, 1].isnan() and logits.grad[0, 1].isnan().all() + assert loss[0, 2] == 0 and (logits.grad[0, 2] == 0).all() + assert loss[0, [0, 3]].isfinite().all() and logits.grad[0, [0, 3]].isfinite().all() + if return_argmax: + torch.testing.assert_close(argmax, torch.argmax(values, dim=-1), rtol=0, atol=0) + + mxfp8_available, reason_for_no_mxfp8 = is_mxfp8_available(return_reason=True) _MXFP8_DTYPES = {"e4m3": tex.DType.kFloat8E4M3, "e5m2": tex.DType.kFloat8E5M2} _MXFP8_BUFFERS = ( @@ -626,6 +732,56 @@ def run(grad_quantizer): torch.testing.assert_close(grad.dequantize(), ref_grad, rtol=rel_err, atol=abs_err) +@pytest.mark.skipif(not mxfp8_available, reason=reason_for_no_mxfp8) +@pytest.mark.parametrize("vocab", [8192, 129280]) +@pytest.mark.parametrize("fp8_dtype", ["e4m3", "e5m2"]) +@pytest.mark.parametrize("reduce_loss", [False, True], ids=["none", "mean"]) +@pytest.mark.parametrize("label_smoothing", [0.0, 0.1], ids=["plain", "smoothed"]) +@pytest.mark.parametrize("overwrite_input", [False, True], ids=["safe", "destructive"]) +def test_parallel_cross_entropy_mxfp8_grad_masked_logits( + vocab, fp8_dtype, reduce_loss, label_smoothing, overwrite_input +): + """With masked (-inf) logits, the MXFP8 gradient has the bytes of the BF16 gradient + quantized by the quantizer, and both match PyTorch.""" + + values, target = _masked_logits(128, vocab, torch.bfloat16, seed=vocab + 1) + values, target = values.reshape(128, 1, vocab), target.reshape(128, 1) + target[:3] = -100 + loss_grad = ( + torch.full((1,), 0.37, device="cuda") + if reduce_loss + else torch.linspace(0.5, 1.5, target.numel(), device="cuda").reshape(target.shape) + ) + + def run(grad_quantizer): + logits = values.clone().requires_grad_() + loss = parallel_cross_entropy( + logits, + target, + label_smoothing, + reduce_loss, + overwrite_input=overwrite_input, + grad_quantizer=grad_quantizer, + ) + loss.backward(loss_grad) + return loss.detach(), logits.grad + + quantizer = MXFP8Quantizer(_MXFP8_DTYPES[fp8_dtype]) + ref_loss, ref_grad = run(None) + loss, grad = run(quantizer) + + torch.testing.assert_close(loss, ref_loss, rtol=0, atol=0) + _assert_same_mxfp8_bytes(grad, quantizer(ref_grad)) + assert grad.dequantize().isfinite().all() + torch_loss, torch_grad = _masked_reference( + values, target, label_smoothing, reduce_loss, loss_grad + ) + torch.testing.assert_close(loss, torch_loss.float(), **dtype_tols(torch.float32)) + torch.testing.assert_close( + ref_grad, torch_grad.to(torch.bfloat16), **dtype_tols(torch.bfloat16) + ) + + @pytest.mark.skipif(not mxfp8_available, reason=reason_for_no_mxfp8) def test_parallel_cross_entropy_mxfp8_grad_unfused_cases(): """Inputs the fused kernel does not write still return the quantized gradient.""" @@ -692,8 +848,9 @@ def run(grad_quantizer): @pytest.mark.skipif(not mxfp8_available, reason=reason_for_no_mxfp8) -def test_parallel_cross_entropy_mxfp8_grad_cuda_graph(): - """The MXFP8 backward is CUDA graph capturable.""" +@pytest.mark.parametrize("masked", [False, True], ids=["dense", "masked"]) +def test_parallel_cross_entropy_mxfp8_grad_cuda_graph(masked): + """The MXFP8 backward is CUDA graph capturable, also when the replayed logits are masked.""" shape = (256, 1, 8192) static_values = _heavy_tailed_logits(shape, seed=4) @@ -720,9 +877,15 @@ def step(): with torch.cuda.graph(graph): step() - static_values.copy_(_heavy_tailed_logits(shape, seed=5)) - static_target.copy_(torch.randint(0, shape[-1], shape[:-1], device="cuda")) + if masked: + values, target = _masked_logits(shape[0], shape[-1], torch.bfloat16, seed=5) + static_values.copy_(values.reshape(shape)) + static_target.copy_(target.reshape(shape[:-1])) + else: + static_values.copy_(_heavy_tailed_logits(shape, seed=5)) + static_target.copy_(torch.randint(0, shape[-1], shape[:-1], device="cuda")) graph.replay() logits = static_values.clone().requires_grad_() parallel_cross_entropy(logits, static_target).sum().backward() _assert_same_mxfp8_bytes(outputs["grad"], quantizer(logits.grad)) + assert outputs["grad"].dequantize().isfinite().all() diff --git a/transformer_engine/common/triton/cross_entropy.py b/transformer_engine/common/triton/cross_entropy.py index 964c0402b44..60a2bf6b601 100644 --- a/transformer_engine/common/triton/cross_entropy.py +++ b/transformer_engine/common/triton/cross_entropy.py @@ -14,6 +14,15 @@ def _is_nan(x): return x != x # pylint: disable=comparison-with-itself +@triton.jit +def _safe_shift(m): + """Shift of the exponentials of an online softmax: the running maximum m, or 0 while m is + -inf, that is while every logit seen so far is -inf (masked). Shifting by m = -inf would + compute exp(-inf - (-inf)) = NaN and make the sum of exponentials NaN for good; shifting + by 0 keeps the sum at 0 until a finite logit is seen. With a finite m nothing changes.""" + return tl.where(m == float("-inf"), 0.0, m) + + @triton.jit def _fold_block_argmax(x, offsets, n_cols, run_max, run_idx, run_nan): """Fold one block of a row into a running argmax with torch.argmax's semantics: NaN counts @@ -92,7 +101,8 @@ def cross_entropy_forward_kernel( ) block_max = tl.max(x) m_new = tl.maximum(m, block_max) - d = d * tl.exp(m - m_new) + tl.sum(tl.exp(x - m_new)) + shift = _safe_shift(m_new) + d = d * tl.exp(m - shift) + tl.sum(tl.exp(x - shift)) m = m_new if label_smoothing > 0: x_sum += tl.sum(tl.where(mask, x, 0.0)) @@ -178,7 +188,8 @@ def cross_entropy_tp_pre_kernel( ) block_max = tl.max(x) m_new = tl.maximum(m, block_max) - d = d * tl.exp(m - m_new) + tl.sum(tl.exp(x - m_new)) + shift = _safe_shift(m_new) + d = d * tl.exp(m - shift) + tl.sum(tl.exp(x - shift)) m = m_new if COMPUTE_X_SUM: x_sum += tl.sum(tl.where(mask, x, 0.0)) @@ -241,7 +252,8 @@ def cross_entropy_tp_post_kernel( argmax = rank_idx.to(tl.int64) * n_cols + local_argmax.to(tl.int64) argmax_max = rank_argmax_max global_m = tl.maximum(m, m_new) - d = d * tl.exp(m - global_m) + d_new * tl.exp(m_new - global_m) + shift = _safe_shift(global_m) + d = d * tl.exp(m - shift) + d_new * tl.exp(m_new - shift) m = global_m x_y = tl.maximum(x_y, tl.load(rank_data_ptr + 2)) if label_smoothing > 0: From a41ae4b2259b6972502f6192bc3fe20fd6f1e440 Mon Sep 17 00:00:00 2001 From: Kirthi Shankar Sivamani Date: Fri, 13 Nov 2026 12:22:17 -0800 Subject: [PATCH 4/5] NaN Signed-off-by: Kirthi Shankar Sivamani --- .../test_parallel_cross_entropy.py | 114 +++++++++++++++++- tests/pytorch/test_cross_entropy.py | 102 ++++++++++++++-- .../common/triton/cross_entropy.py | 29 +++-- transformer_engine/pytorch/cross_entropy.py | 6 +- .../pytorch/triton/cross_entropy.py | 15 ++- 5 files changed, 236 insertions(+), 30 deletions(-) diff --git a/tests/pytorch/distributed/test_parallel_cross_entropy.py b/tests/pytorch/distributed/test_parallel_cross_entropy.py index c4e3510c174..c352b588ca3 100644 --- a/tests/pytorch/distributed/test_parallel_cross_entropy.py +++ b/tests/pytorch/distributed/test_parallel_cross_entropy.py @@ -12,7 +12,10 @@ import torch.distributed as dist import torch.multiprocessing as mp -from transformer_engine.pytorch import parallel_cross_entropy +import transformer_engine_torch as tex +from transformer_engine.pytorch import is_mxfp8_available, parallel_cross_entropy +from transformer_engine.pytorch.tensor.mxfp8_tensor import MXFP8Quantizer +import transformer_engine.pytorch.triton.cross_entropy as triton_cross_entropy # Prepend so installed packages with a top-level utils module cannot shadow the test helpers. sys.path = [str(pathlib.Path(__file__).resolve().parent.parent)] + sys.path @@ -42,10 +45,16 @@ def _run_tensor_parallel(rank, world_size, init_file, label_smoothing): generator=generator, ).to(device) target[0, 2] = -100 + mask_rank_0 = label_smoothing == 0.0 # with smoothing, -inf logits give an infinite loss + if mask_rank_0: + target[0, 1] = local_vocab + 1 external_grad = torch.randn(shape[:-1], generator=generator).to(device) for dtype in (torch.float32, torch.bfloat16): global_values = torch.randn(shape, generator=generator).to(device=device, dtype=dtype) + if mask_rank_0: + # Rank 0's whole shard is -inf in one row: it must contribute nothing, not NaN. + global_values[0, 1, :local_vocab] = float("-inf") vocab_start = rank * local_vocab local_values = global_values[..., vocab_start : vocab_start + local_vocab] for reduce_loss in (False, True): @@ -153,6 +162,89 @@ def _run_tensor_parallel_argmax(rank, world_size, init_file, label_smoothing): dist.destroy_process_group() +class _KernelLaunchCounter: + """Stands in for a Triton kernel, counting its launches and forwarding them.""" + + def __init__(self, kernel): + self.kernel = kernel + self.launches = 0 + + def __getitem__(self, grid): + launch = self.kernel[grid] + + def counted_launch(*args, **kwargs): + self.launches += 1 + return launch(*args, **kwargs) + + return counted_launch + + +def _run_tensor_parallel_mxfp8_grad(rank, world_size, init_file, label_smoothing): + """Two-rank worker: each rank's MXFP8 gradient is written by the fused kernel and has the + bytes of its BF16 gradient shard quantized by the quantizer.""" + + torch.cuda.set_device(rank) + device = torch.device("cuda", rank) + dist.init_process_group( + backend="nccl", + init_method=f"file://{init_file}", + rank=rank, + world_size=world_size, + device_id=device, + ) + counter = _KernelLaunchCounter(triton_cross_entropy.cross_entropy_backward_mxfp8_kernel) + triton_cross_entropy.cross_entropy_backward_mxfp8_kernel = counter + try: + generator = torch.Generator().manual_seed(2027) + # Shards of 384 and 4096 columns: 128- and 256-column tiles. + for local_vocab in (384, 4096): + vocab = local_vocab * world_size + shape = (2, 64, vocab) + target = torch.randint(0, vocab, shape[:-1], generator=generator).to(device) + target[0, :3] = -100 + values = (torch.randn(shape, generator=generator) * 3).to( + device=device, dtype=torch.bfloat16 + ) + vocab_start = rank * local_vocab + local_values = values[..., vocab_start : vocab_start + local_vocab] + for fp8_dtype in (tex.DType.kFloat8E4M3, tex.DType.kFloat8E5M2): + for reduce_loss in (False, True): + quantizer = MXFP8Quantizer(fp8_dtype) + grads = [] + for grad_quantizer in (None, quantizer): + logits = local_values.clone().requires_grad_() + loss = parallel_cross_entropy( + logits, + target, + label_smoothing=label_smoothing, + reduce_loss=reduce_loss, + dist_process_group=dist.group.WORLD, + grad_quantizer=grad_quantizer, + ) + launches = counter.launches + (grad,) = torch.autograd.grad(loss, logits, torch.full_like(loss, 0.37)) + assert counter.launches - launches == int(grad_quantizer is not None) + grads.append(grad) + ref_grad, grad = grads + ref = quantizer(ref_grad) + # 128 rows and a shard that is a multiple of 128: no scale padding. + for name in ( + "_rowwise_data", + "_rowwise_scale_inv", + "_columnwise_data", + "_columnwise_scale_inv", + ): + torch.testing.assert_close( + getattr(grad, name).view(torch.uint8), + getattr(ref, name).view(torch.uint8), + rtol=0, + atol=0, + ) + finally: + triton_cross_entropy.cross_entropy_backward_mxfp8_kernel = counter.kernel + dist.destroy_process_group() + + @pytest.mark.parametrize("label_smoothing", [0.0, 0.1], ids=["plain", "smoothed"]) def test_parallel_cross_entropy_tensor_parallel(label_smoothing): """Validate tensor-parallel loss and gradients on two ranks.""" @@ -185,3 +277,23 @@ def test_parallel_cross_entropy_tensor_parallel_argmax(label_smoothing): nprocs=world_size, join=True, ) + + +@pytest.mark.parametrize("label_smoothing", [0.0, 0.1], ids=["plain", "smoothed"]) +def test_parallel_cross_entropy_tensor_parallel_mxfp8_grad(label_smoothing): + """Validate the fused MXFP8 gradient of each vocabulary shard on two ranks.""" + + if torch.cuda.device_count() < 2: + pytest.skip("tensor-parallel cross entropy test requires two CUDA devices") + mxfp8_available, reason_for_no_mxfp8 = is_mxfp8_available(return_reason=True) + if not mxfp8_available: + pytest.skip(reason_for_no_mxfp8) + world_size = 2 + with tempfile.TemporaryDirectory() as temp_dir: + init_file = os.path.join(temp_dir, "distributed_init") + mp.spawn( + _run_tensor_parallel_mxfp8_grad, + args=(world_size, init_file, label_smoothing), + nprocs=world_size, + join=True, + ) diff --git a/tests/pytorch/test_cross_entropy.py b/tests/pytorch/test_cross_entropy.py index 30d86e13210..b11ad554788 100644 --- a/tests/pytorch/test_cross_entropy.py +++ b/tests/pytorch/test_cross_entropy.py @@ -11,6 +11,7 @@ from transformer_engine.common.recipe import Format, MXFP8BlockScaling from transformer_engine.pytorch import is_mxfp8_available, parallel_cross_entropy from transformer_engine.pytorch.tensor.mxfp8_tensor import MXFP8Quantizer, MXFP8Tensor +import transformer_engine.pytorch.triton.cross_entropy as triton_cross_entropy from utils import dtype_tols @@ -530,6 +531,39 @@ def test_parallel_cross_entropy_return_argmax_nan(): torch.testing.assert_close(argmax, torch.argmax(logits, dim=-1), rtol=0, atol=0) +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) +@pytest.mark.parametrize("return_argmax", [False, True], ids=["loss", "argmax"]) +def test_parallel_cross_entropy_masked_leading_logits(dtype, return_argmax): + """Rows whose first blocks of logits are all -inf (a masked part of the vocabulary) get a + finite loss and gradient: the online softmax must not compute -inf - -inf.""" + + vocab = 65536 # two 32K blocks + inf = float("inf") + values = torch.randn(1, 4, vocab, device="cuda") + values[0, 0, :32768] = -inf # the whole first block + values[0, 1, :40000] = -inf # the first block and part of the second + values[0, 2, :-1] = -inf # all but the last logit + values[0, 3, :4096] = -inf # part of the first block + target = torch.tensor([[40000, 50000, vocab - 1, 100]], device="cuda") + values = values.to(dtype) + + logits = values.clone().requires_grad_() + ref_logits = values.float().clone().requires_grad_() + out = parallel_cross_entropy(logits, target, return_argmax=return_argmax) + loss, argmax = out if return_argmax else (out, None) + ref_loss = torch.nn.functional.cross_entropy( + ref_logits.reshape(-1, vocab), target.reshape(-1), reduction="none" + ).reshape_as(loss) + loss.backward(torch.ones_like(loss)) + ref_loss.backward(torch.ones_like(ref_loss)) + + assert torch.isfinite(loss).all() + torch.testing.assert_close(loss, ref_loss, **dtype_tols(torch.float32)) + torch.testing.assert_close(logits.grad, ref_logits.grad.to(dtype), **dtype_tols(dtype)) + if return_argmax: + torch.testing.assert_close(argmax, torch.argmax(values, dim=-1), rtol=0, atol=0) + + mxfp8_available, reason_for_no_mxfp8 = is_mxfp8_available(return_reason=True) _MXFP8_DTYPES = {"e4m3": tex.DType.kFloat8E4M3, "e5m2": tex.DType.kFloat8E5M2} _MXFP8_BUFFERS = ( @@ -540,6 +574,32 @@ def test_parallel_cross_entropy_return_argmax_nan(): ) +class _KernelLaunchCounter: + """Stands in for a Triton kernel, counting its launches and forwarding them.""" + + def __init__(self, kernel): + self.kernel = kernel + self.launches = 0 + + def __getitem__(self, grid): + launch = self.kernel[grid] + + def counted_launch(*args, **kwargs): + self.launches += 1 + return launch(*args, **kwargs) + + return counted_launch + + +@pytest.fixture +def mxfp8_kernel_launches(monkeypatch): + """Counts the launches of the fused MXFP8 backward kernel, so that tests can tell the fused + path from the fallback, which writes the same bytes.""" + counter = _KernelLaunchCounter(triton_cross_entropy.cross_entropy_backward_mxfp8_kernel) + monkeypatch.setattr(triton_cross_entropy, "cross_entropy_backward_mxfp8_kernel", counter) + return counter + + def _heavy_tailed_logits(shape, seed): """BF16 logits whose MXFP8 block scales vary across the row.""" generator = torch.Generator(device="cuda").manual_seed(seed) @@ -572,8 +632,16 @@ def _assert_same_mxfp8_bytes(grad, ref): @pytest.mark.skipif(not mxfp8_available, reason=reason_for_no_mxfp8) @pytest.mark.parametrize( "shape", - [(2, 64, 4096), (512, 1, 129280), (128, 2, 32000), (96, 1, 4096), (128, 1, 4000)], - ids=["v4096", "v129280", "v32000", "rows96_fallback", "v4000_fallback"], + [ + (2, 64, 4096), + (512, 1, 129280), + (128, 2, 32000), + (128, 1, 50304), + (96, 1, 4096), + (128, 1, 4000), + ], + # 50304 is a multiple of 128 but not of 256: 128-column tiles. + ids=["v4096", "v129280", "v32000", "v50304_cols128", "rows96_fallback", "v4000_fallback"], ) @pytest.mark.parametrize("fp8_dtype", ["e4m3", "e5m2"]) @pytest.mark.parametrize("usage", ["both", "rowwise", "columnwise"]) @@ -581,10 +649,10 @@ def _assert_same_mxfp8_bytes(grad, ref): @pytest.mark.parametrize("label_smoothing", [0.0, 0.1], ids=["plain", "smoothed"]) @pytest.mark.parametrize("overwrite_input", [False, True], ids=["safe", "destructive"]) def test_parallel_cross_entropy_mxfp8_grad( - shape, fp8_dtype, usage, reduce_loss, label_smoothing, overwrite_input + shape, fp8_dtype, usage, reduce_loss, label_smoothing, overwrite_input, mxfp8_kernel_launches ): """The MXFP8 gradient has the bytes of the BF16 gradient quantized by the quantizer, and - the loss is unchanged.""" + the loss is unchanged. Shapes the fused kernel supports are written by it.""" values = _heavy_tailed_logits(shape, seed=sum(shape)) target = torch.randint(0, shape[-1], shape[:-1], device="cuda") @@ -616,6 +684,8 @@ def run(grad_quantizer): ref_loss, ref_grad = run(None) loss, grad = run(quantizer) + fused = math.prod(shape[:-1]) % 128 == 0 and shape[-1] % 128 == 0 + assert mxfp8_kernel_launches.launches == int(fused) torch.testing.assert_close(loss, ref_loss, rtol=0, atol=0) _assert_same_mxfp8_bytes(grad, quantizer(ref_grad)) # The dequantized gradient is the BF16 gradient up to the MXFP8 rounding error: relative @@ -627,29 +697,40 @@ def run(grad_quantizer): @pytest.mark.skipif(not mxfp8_available, reason=reason_for_no_mxfp8) -def test_parallel_cross_entropy_mxfp8_grad_unfused_cases(): - """Inputs the fused kernel does not write still return the quantized gradient.""" +def test_parallel_cross_entropy_mxfp8_grad_dispatch(mxfp8_kernel_launches): + """Inputs the fused kernel does not write still return the quantized gradient, and a + non-contiguous BF16 input is fused, since forward saves a contiguous copy of it.""" values = torch.randn(2, 64, 4096, device="cuda") target = torch.randint(0, 4096, (2, 64), device="cuda") quantizer = MXFP8Quantizer(tex.DType.kFloat8E4M3) - # FP32 input, and a non-contiguous BF16 input. torch.autograd.grad returns the gradient as - # the loss returns it (AccumulateGrad would re-layout it for the non-contiguous leaf). - for logits_values in (values, values.to(torch.bfloat16).transpose(0, 1)): + # An FP32 input falls back; a non-contiguous BF16 input is fused. torch.autograd.grad + # returns the gradient as the loss returns it (AccumulateGrad would re-layout it for the + # non-contiguous leaf). + for logits_values, fused in ( + (values, False), + (values.to(torch.bfloat16).transpose(0, 1), True), + ): + assert logits_values.is_contiguous() == (not fused) logits_target = target.reshape(logits_values.shape[:-1]) logits = logits_values.clone().requires_grad_() (ref_grad,) = torch.autograd.grad( parallel_cross_entropy(logits, logits_target).sum(), logits ) + launches = mxfp8_kernel_launches.launches (grad,) = torch.autograd.grad( parallel_cross_entropy(logits, logits_target, grad_quantizer=quantizer).sum(), logits ) + assert mxfp8_kernel_launches.launches - launches == int(fused) _assert_same_mxfp8_bytes(grad, quantizer(ref_grad)) + # GEMM-swizzled scales fall back to the quantizer. swizzled = MXFP8Quantizer(tex.DType.kFloat8E4M3) swizzled.optimize_for_gemm = True logits = values.to(torch.bfloat16).requires_grad_() + launches = mxfp8_kernel_launches.launches parallel_cross_entropy(logits, target, grad_quantizer=swizzled).sum().backward() + assert mxfp8_kernel_launches.launches == launches assert isinstance(logits.grad, MXFP8Tensor) and logits.grad._with_gemm_swizzled_scales internal = MXFP8Quantizer(tex.DType.kFloat8E4M3) @@ -660,7 +741,7 @@ def test_parallel_cross_entropy_mxfp8_grad_unfused_cases(): @pytest.mark.skipif(not mxfp8_available, reason=reason_for_no_mxfp8) @pytest.mark.parametrize("overwrite_input", [False, True], ids=["safe", "destructive"]) -def test_parallel_cross_entropy_mxfp8_grad_lm_head(overwrite_input): +def test_parallel_cross_entropy_mxfp8_grad_lm_head(overwrite_input, mxfp8_kernel_launches): """An MXFP8 linear layer followed by the loss gets the same input and weight gradients when the loss returns the MXFP8 gradient of the logits.""" @@ -687,6 +768,7 @@ def run(grad_quantizer): ref = run(None) out = run(MXFP8Quantizer(tex.DType.kFloat8E4M3)) + assert mxfp8_kernel_launches.launches == 1 for a, b in zip(out, ref): torch.testing.assert_close(a, b, rtol=0, atol=0) diff --git a/transformer_engine/common/triton/cross_entropy.py b/transformer_engine/common/triton/cross_entropy.py index 964c0402b44..7d2da780f74 100644 --- a/transformer_engine/common/triton/cross_entropy.py +++ b/transformer_engine/common/triton/cross_entropy.py @@ -90,9 +90,10 @@ def cross_entropy_forward_kernel( argmax_max, argmax, argmax_nan = _fold_block_argmax( x, offsets, n_cols, argmax_max, argmax, argmax_nan ) - block_max = tl.max(x) - m_new = tl.maximum(m, block_max) - d = d * tl.exp(m - m_new) + tl.sum(tl.exp(x - m_new)) + m_new = tl.maximum(m, tl.max(x)) + # While every logit so far is -inf, m_new is -inf too, and -inf - -inf would be NaN. + m_safe = tl.where(m_new == float("-inf"), 0.0, m_new) + d = d * tl.exp(m - m_safe) + tl.sum(tl.exp(x - m_safe)) m = m_new if label_smoothing > 0: x_sum += tl.sum(tl.where(mask, x, 0.0)) @@ -176,9 +177,10 @@ def cross_entropy_tp_pre_kernel( argmax_max, argmax, argmax_nan = _fold_block_argmax( x, offsets, n_cols, argmax_max, argmax, argmax_nan ) - block_max = tl.max(x) - m_new = tl.maximum(m, block_max) - d = d * tl.exp(m - m_new) + tl.sum(tl.exp(x - m_new)) + m_new = tl.maximum(m, tl.max(x)) + # While every logit so far is -inf, m_new is -inf too, and -inf - -inf would be NaN. + m_safe = tl.where(m_new == float("-inf"), 0.0, m_new) + d = d * tl.exp(m - m_safe) + tl.sum(tl.exp(x - m_safe)) m = m_new if COMPUTE_X_SUM: x_sum += tl.sum(tl.where(mask, x, 0.0)) @@ -241,7 +243,9 @@ def cross_entropy_tp_post_kernel( argmax = rank_idx.to(tl.int64) * n_cols + local_argmax.to(tl.int64) argmax_max = rank_argmax_max global_m = tl.maximum(m, m_new) - d = d * tl.exp(m - global_m) + d_new * tl.exp(m_new - global_m) + # Shards whose logits are all -inf have m = -inf and d = 0 and must not give NaN. + safe_m = tl.where(global_m == float("-inf"), 0.0, global_m) + d = d * tl.exp(m - safe_m) + d_new * tl.exp(m_new - safe_m) m = global_m x_y = tl.maximum(x_y, tl.load(rank_data_ptr + 2)) if label_smoothing > 0: @@ -348,6 +352,8 @@ def cross_entropy_backward_mxfp8_kernel( columnwise_data_ptr, columnwise_scale_ptr, columnwise_scale_stride, + rank, + world_size, n_cols, ignore_idx, reduce_loss: tl.constexpr, @@ -357,7 +363,8 @@ def cross_entropy_backward_mxfp8_kernel( FP8_E5M2: tl.constexpr, BLOCK_COLS: tl.constexpr, ): - """Single-rank derivative of a 32-row x BLOCK_COLS tile, written as MXFP8. + """Derivative of a 32-row x BLOCK_COLS tile of this rank's vocabulary shard, written as + MXFP8. The derivative is computed as cross_entropy_backward_kernel computes it and rounded to the input dtype (BF16), then quantized as TE's MXFP8 cast quantizes it: row-wise with 1x32 @@ -375,10 +382,12 @@ def cross_entropy_backward_mxfp8_kernel( if reduce_loss: grad_output = grad_output / tl.load(n_non_ignore_ptr) - eps = label_smoothing / n_cols + eps = label_smoothing / (n_cols * world_size) + # Local column of the target; outside [0, n_cols) if another rank holds it. + target_col = y - rank * n_cols x = tl.load(saved_input_ptr + offsets).to(tl.float32) grad = tl.exp(x - m[:, None]) / d[:, None] - eps - grad -= tl.where(cols[None, :] == y[:, None], 1 - label_smoothing, 0.0) + grad -= tl.where(cols[None, :] == target_col[:, None], 1 - label_smoothing, 0.0) grad = grad * grad_output[:, None] grad = tl.where((y == ignore_idx)[:, None], 0.0, grad) grad = grad.to(saved_input_ptr.dtype.element_ty).to(tl.float32) diff --git a/transformer_engine/pytorch/cross_entropy.py b/transformer_engine/pytorch/cross_entropy.py index 2505a317043..c1f71c77781 100644 --- a/transformer_engine/pytorch/cross_entropy.py +++ b/transformer_engine/pytorch/cross_entropy.py @@ -227,9 +227,9 @@ def parallel_cross_entropy( grad_quantizer : Quantizer, default = None Return the gradient of ``inp`` quantized with this quantizer, for a following quantized linear layer (for example an MXFP8 LM head) that would otherwise quantize - it itself. With an ``MXFP8Quantizer`` (E4M3 or E5M2, compact scales), a BF16 ``inp`` - without ``dist_process_group`` (or with a single rank), a row count that is a multiple - of 128 and a vocabulary that is a multiple of 128, backward writes the MXFP8 gradient + it itself. With an ``MXFP8Quantizer`` (E4M3 or E5M2, compact scales), a BF16 ``inp``, + a row count that is a multiple of 128 and a vocabulary (per rank, with + ``dist_process_group``) that is a multiple of 128, backward writes the MXFP8 gradient directly, bit-identical to quantizing the BF16 gradient with ``grad_quantizer`` (for finite gradients). In all other cases the gradient is computed as usual and then quantized. diff --git a/transformer_engine/pytorch/triton/cross_entropy.py b/transformer_engine/pytorch/triton/cross_entropy.py index 235fb455d28..db6a026eaef 100644 --- a/transformer_engine/pytorch/triton/cross_entropy.py +++ b/transformer_engine/pytorch/triton/cross_entropy.py @@ -152,9 +152,7 @@ def cross_entropy_forward( return loss, saved_input, stats, target, n_non_ignore, rank, world_size, argmax -def _mxfp8_backward_block_cols( - saved_input: torch.Tensor, world_size: int, grad_quantizer -) -> Optional[int]: +def _mxfp8_backward_block_cols(saved_input: torch.Tensor, grad_quantizer) -> Optional[int]: """Tile width of the fused MXFP8 backward, or None if it cannot write this gradient.""" # pylint: disable=import-outside-toplevel from transformer_engine.pytorch.tensor.mxfp8_tensor import MXFP8Quantizer @@ -169,11 +167,14 @@ def _mxfp8_backward_block_cols( return None if not (grad_quantizer.rowwise_usage or grad_quantizer.columnwise_usage): return None - if world_size != 1 or saved_input.dtype != torch.bfloat16 or not saved_input.is_contiguous(): + # saved_input is always contiguous: overwrite_input requires a contiguous input, and + # otherwise forward copies the input into a new buffer. + if saved_input.dtype != torch.bfloat16: return None B, SQ, V = saved_input.shape # Unpadded scale-inverse layouts: 128 rows x 4 blocks of the row-wise scales and - # 4 x 128 columns of the column-wise scales. + # 4 x 128 columns of the column-wise scales. With tensor parallelism, V is this rank's + # vocabulary shard. if (B * SQ) % 128 != 0 or V % 128 != 0: return None return _MXFP8_BWD_BLOCK_COLS if V % _MXFP8_BWD_BLOCK_COLS == 0 else 128 @@ -206,7 +207,7 @@ def cross_entropy_backward( grad_output = grad_output.contiguous() if grad_quantizer is not None: - block_cols = _mxfp8_backward_block_cols(saved_input, world_size, grad_quantizer) + block_cols = _mxfp8_backward_block_cols(saved_input, grad_quantizer) if block_cols is not None: import transformer_engine_torch as tex # pylint: disable=import-outside-toplevel @@ -232,6 +233,8 @@ def cross_entropy_backward( columnwise_data_ptr=(grad_input._columnwise_data if columnwise else saved_input), columnwise_scale_ptr=columnwise_scale if columnwise else saved_input, columnwise_scale_stride=columnwise_scale.shape[-1] if columnwise else 0, + rank=rank, + world_size=world_size, n_cols=V, ignore_idx=ignore_idx, reduce_loss=reduce_loss, From c64df2244807ec14377315a05715e07deb17a6a8 Mon Sep 17 00:00:00 2001 From: Kirthi Shankar Sivamani Date: Fri, 13 Nov 2026 13:41:08 -0800 Subject: [PATCH 5/5] Small testing refactors Signed-off-by: Kirthi Shankar Sivamani --- .../test_parallel_cross_entropy.py | 31 +++++---------- tests/pytorch/test_cross_entropy.py | 39 +++++++------------ tests/pytorch/utils.py | 21 ++++++++++ 3 files changed, 44 insertions(+), 47 deletions(-) diff --git a/tests/pytorch/distributed/test_parallel_cross_entropy.py b/tests/pytorch/distributed/test_parallel_cross_entropy.py index e3c9797b8e6..dac5af6a164 100644 --- a/tests/pytorch/distributed/test_parallel_cross_entropy.py +++ b/tests/pytorch/distributed/test_parallel_cross_entropy.py @@ -20,7 +20,7 @@ # Prepend so installed packages with a top-level utils module cannot shadow the test helpers. sys.path = [str(pathlib.Path(__file__).resolve().parent.parent)] + sys.path -from utils import dtype_tols +from utils import KernelLaunchCounter, dtype_tols def _run_tensor_parallel(rank, world_size, init_file, label_smoothing): @@ -171,14 +171,20 @@ def _run_tensor_parallel_masked(rank, world_size, init_file, label_smoothing): ) try: generator = torch.Generator().manual_seed(2027) - for shape, block in (((1, 5, 22), 8), ((1, 5, 80000), 32768)): + # Row 2 masks the first masked_prefix logits of each shard: part of the only block of the + # tensor-parallel forward kernel for 11-wide shards, and the whole first block for + # 40000-wide shards. + for shape, masked_prefix in ( + ((1, 5, 22), 8), + ((1, 5, 80000), triton_cross_entropy.MAX_FUSED_SIZE), + ): vocab = shape[-1] local_vocab = vocab // world_size values = torch.randn(shape, generator=generator) values[0, 0, :local_vocab] = -float("inf") # the shard of rank 0 is masked values[0, 1, local_vocab:] = -float("inf") # the shard of rank 1 is masked for start in range(0, vocab, local_vocab): - values[0, 2, start : start + block] = -float("inf") # each first block + values[0, 2, start : start + masked_prefix] = -float("inf") values[0, 3, :-1] = -float("inf") # only the last logit is unmasked target = torch.multinomial( values.isfinite().float().reshape(-1, vocab), 1, generator=generator @@ -227,23 +233,6 @@ def _run_tensor_parallel_masked(rank, world_size, init_file, label_smoothing): dist.destroy_process_group() -class _KernelLaunchCounter: - """Stands in for a Triton kernel, counting its launches and forwarding them.""" - - def __init__(self, kernel): - self.kernel = kernel - self.launches = 0 - - def __getitem__(self, grid): - launch = self.kernel[grid] - - def counted_launch(*args, **kwargs): - self.launches += 1 - return launch(*args, **kwargs) - - return counted_launch - - def _run_tensor_parallel_mxfp8_grad(rank, world_size, init_file, label_smoothing): """Two-rank worker: each rank's MXFP8 gradient is written by the fused kernel and has the bytes of its BF16 gradient shard quantized by the quantizer.""" @@ -257,7 +246,7 @@ def _run_tensor_parallel_mxfp8_grad(rank, world_size, init_file, label_smoothing world_size=world_size, device_id=device, ) - counter = _KernelLaunchCounter(triton_cross_entropy.cross_entropy_backward_mxfp8_kernel) + counter = KernelLaunchCounter(triton_cross_entropy.cross_entropy_backward_mxfp8_kernel) triton_cross_entropy.cross_entropy_backward_mxfp8_kernel = counter try: generator = torch.Generator().manual_seed(2027) diff --git a/tests/pytorch/test_cross_entropy.py b/tests/pytorch/test_cross_entropy.py index b26215837d0..ba144c6682f 100644 --- a/tests/pytorch/test_cross_entropy.py +++ b/tests/pytorch/test_cross_entropy.py @@ -13,7 +13,7 @@ from transformer_engine.pytorch.tensor.mxfp8_tensor import MXFP8Quantizer, MXFP8Tensor import transformer_engine.pytorch.triton.cross_entropy as triton_cross_entropy -from utils import dtype_tols +from utils import KernelLaunchCounter, dtype_tols class TestParallelCrossEntropy: @@ -439,17 +439,20 @@ def test_parallel_cross_entropy_deprecated_input_alias(): def _logits_with_argmax_cases(shape, dtype): - """BF16-representable logits whose rows exercise ties within and across 32K-wide blocks.""" + """BF16-representable logits whose rows exercise ties within and across the blocks of the + single-rank forward kernel.""" torch.manual_seed(7) values = torch.randn(shape, dtype=torch.bfloat16, device="cuda").to(dtype) vocab = shape[-1] + block = triton_cross_entropy._FWD_BLOCK_SIZE peak = 8.0 values[0, 0, [5, vocab - 30000]] = peak # equal maxima in different blocks values[0, 1, [vocab - 30000, vocab - 4]] = peak # equal maxima in later blocks values[0, 2, [100, 7]] = peak # equal maxima in the same block values[0, 3, vocab - 1] = peak # a strictly larger maximum in the last block values[0, 4] = 0.5 # a constant row + values[1, 0, [block - 1, block]] = peak # equal maxima on both sides of a block boundary return values @@ -533,19 +536,20 @@ def test_parallel_cross_entropy_return_argmax_nan(): def _masked_logits(n_rows, vocab, dtype, seed): """Logits with masked (-inf) entries and a finite maximum in every row, and targets that are - unmasked. Rows start with masked blocks of the forward kernels (4096 wide on a single rank, - 32K wide before), keep only one logit, or are masked elsewhere. Row patterns repeat every - eight rows.""" + unmasked. The masks are laid out against the blocks of the single-rank forward kernel: rows + whose first or last block is masked entirely, rows that keep only one or a few logits, and + rows masked elsewhere. Row patterns repeat every eight rows.""" + block = triton_cross_entropy._FWD_BLOCK_SIZE generator = torch.Generator(device="cuda").manual_seed(seed) values = torch.randn((n_rows, vocab), device="cuda", generator=generator) - values[0::8, :4096] = -math.inf # the first block is masked + values[0::8, :block] = -math.inf # the first block is masked values[1::8, : vocab // 2] = -math.inf # the first half is masked values[2::8, :-1] = -math.inf # only the last logit is unmasked values[3::8, :] = -math.inf - values[3::8, 4100] = 0.0 # only one logit in the second block is unmasked + values[3::8, block + 4] = 0.0 # only one logit in the second block is unmasked values[4::8, ::2] = -math.inf # no block is masked entirely - values[5::8, -4096:] = -math.inf # the last block is masked + values[5::8, -block:] = -math.inf # the last block is masked values[6::8, :-100] = -math.inf # only the last 100 logits are unmasked values = values.to(dtype) target = torch.multinomial(values.isfinite().float(), 1, generator=generator).squeeze(-1) @@ -647,28 +651,11 @@ def test_parallel_cross_entropy_fully_masked_row(vocab, dtype): ) -class _KernelLaunchCounter: - """Stands in for a Triton kernel, counting its launches and forwarding them.""" - - def __init__(self, kernel): - self.kernel = kernel - self.launches = 0 - - def __getitem__(self, grid): - launch = self.kernel[grid] - - def counted_launch(*args, **kwargs): - self.launches += 1 - return launch(*args, **kwargs) - - return counted_launch - - @pytest.fixture def mxfp8_kernel_launches(monkeypatch): """Counts the launches of the fused MXFP8 backward kernel, so that tests can tell the fused path from the fallback, which writes the same bytes.""" - counter = _KernelLaunchCounter(triton_cross_entropy.cross_entropy_backward_mxfp8_kernel) + counter = KernelLaunchCounter(triton_cross_entropy.cross_entropy_backward_mxfp8_kernel) monkeypatch.setattr(triton_cross_entropy, "cross_entropy_backward_mxfp8_kernel", counter) return counter diff --git a/tests/pytorch/utils.py b/tests/pytorch/utils.py index 2a848842b52..ee90f73aad5 100644 --- a/tests/pytorch/utils.py +++ b/tests/pytorch/utils.py @@ -592,6 +592,27 @@ def assert_close_rms( ), f"RMS error {error_rms:.4g} exceeds {rtol} * RMS of expected ({expected_rms:.4g})" +class KernelLaunchCounter: + """Stands in for a Triton kernel, counting its launches and forwarding them. + + Replace the kernel in its module with an instance (for example with + ``monkeypatch.setattr``) to check whether a code path launches it. + """ + + def __init__(self, kernel): + self.kernel = kernel + self.launches = 0 + + def __getitem__(self, grid): + launch = self.kernel[grid] + + def counted_launch(*args, **kwargs): + self.launches += 1 + return launch(*args, **kwargs) + + return counted_launch + + def run_distributed( args: Sequence[str], *,