diff --git a/tests/pytorch/distributed/test_parallel_cross_entropy.py b/tests/pytorch/distributed/test_parallel_cross_entropy.py index c4e3510c174..dac5af6a164 100644 --- a/tests/pytorch/distributed/test_parallel_cross_entropy.py +++ b/tests/pytorch/distributed/test_parallel_cross_entropy.py @@ -12,11 +12,15 @@ 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 -from utils import dtype_tols +from utils import KernelLaunchCounter, dtype_tols def _run_tensor_parallel(rank, world_size, init_file, label_smoothing): @@ -153,6 +157,148 @@ 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) + # 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 + 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 + ).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() + + +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 +331,40 @@ 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, + ) + + +@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 0cd77f28824..ba144c6682f 100644 --- a/tests/pytorch/test_cross_entropy.py +++ b/tests/pytorch/test_cross_entropy.py @@ -2,12 +2,18 @@ # # 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 +import transformer_engine.pytorch.triton.cross_entropy as triton_cross_entropy -from utils import dtype_tols +from utils import KernelLaunchCounter, dtype_tols class TestParallelCrossEntropy: @@ -433,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 @@ -523,3 +532,397 @@ 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) + + +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. 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, :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, 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, -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) + 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 = ( + "_rowwise_data", + "_rowwise_scale_inv", + "_columnwise_data", + "_columnwise_scale_inv", +) + + +@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) + 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), + (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"]) +@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, mxfp8_kernel_launches +): + """The MXFP8 gradient has the bytes of the BF16 gradient quantized by the quantizer, and + 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") + 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) + + 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 + # 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) +@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, mxfp8_kernel_launches +): + """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) + + assert mxfp8_kernel_launches.launches == 1 + 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_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) + # 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) + 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, 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.""" + + 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)) + assert mxfp8_kernel_launches.launches == 1 + 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) +@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) + 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() + + 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/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], *, diff --git a/transformer_engine/common/triton/cross_entropy.py b/transformer_engine/common/triton/cross_entropy.py index 3d1fd531782..f49c987785f 100644 --- a/transformer_engine/common/triton/cross_entropy.py +++ b/transformer_engine/common/triton/cross_entropy.py @@ -15,30 +15,34 @@ def _is_nan(x): @triton.jit -def _fold_block_argmax(x, i, offsets, n_cols, run_max, run_idx, run_nan): +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 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,11 +97,12 @@ 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) - 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)) @@ -179,11 +184,12 @@ 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) - 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)) @@ -246,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: @@ -318,3 +325,115 @@ 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_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)) + return tl.where(round_up, exponent + 1, exponent) + + +@triton.jit +def _mxfp8_scale_reciprocal(exponent): + """2^-(exponent - 127) as FP32, as TE's exp2f_rcp for finite scales.""" + bits = tl.where(exponent == 254, 0x00400000, (254 - exponent) << 23) + 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, + rank, + world_size, + 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, +): + """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 + 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 * 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, :] == 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) + grad_abs = tl.abs(grad) + + if FP8_E5M2: + MAX_NORM_RCP: tl.constexpr = 1.0 / 57344.0 + else: + 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_RCP) + 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_RCP) + 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 e6ab2b9a10f..738d8686793 100644 --- a/transformer_engine/pytorch/cross_entropy.py +++ b/transformer_engine/pytorch/cross_entropy.py @@ -12,6 +12,7 @@ import transformer_engine.pytorch.triton.cross_entropy as triton_cross_entropy from transformer_engine.pytorch.device import te_device_type from transformer_engine.pytorch.jit import no_torch_dynamo +from transformer_engine.pytorch.quantized_tensor import Quantizer __all__ = ["parallel_cross_entropy"] @@ -34,6 +35,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.""" @@ -66,6 +68,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) @@ -95,10 +98,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( @@ -106,6 +110,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.""" @@ -130,6 +135,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() @@ -141,6 +151,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.""" @@ -153,6 +164,7 @@ def _parallel_cross_entropy_overwrite_input( ignore_idx, True, return_argmax, + grad_quantizer, ) @@ -167,6 +179,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. @@ -212,6 +225,15 @@ 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``, + 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. Returns ------- @@ -234,7 +256,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, @@ -244,6 +266,7 @@ def parallel_cross_entropy( dist_process_group, ignore_idx, return_argmax, + grad_quantizer, ) return CrossEntropyFunction.apply( inp, @@ -254,4 +277,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..db6a026eaef 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,34 @@ 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, 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 + # 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. 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 + + def cross_entropy_backward( saved_input: torch.Tensor, stats: torch.Tensor, @@ -148,14 +191,62 @@ 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, 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, + rank=rank, + world_size=world_size, + 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 +260,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