Minimize the memory usage of the fused cross entropy kernel

benchmarks/benchmark_parallel_cross_entropy.py

@@ -0,0 +1,192 @@
+# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+#
+# See LICENSE for license information.
+
+"""Benchmark fused cross entropy safe-copy and destructive-reuse modes."""
+
+import argparse
+import gc
+from statistics import mean
+
+import torch
+
+from transformer_engine.pytorch import parallel_cross_entropy
+
+MIB = 1024**2
+
+
+def parse_args():
+    """Parse command-line arguments."""
+
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--batch-size", type=int, default=4)
+    parser.add_argument("--sequence-length", type=int, default=2048)
+    parser.add_argument("--vocab-size", type=int, default=128000)
+    parser.add_argument("--dtype", choices=("bf16", "fp32"), default="bf16")
+    parser.add_argument("--warmup", type=int, default=3)
+    parser.add_argument("--trials", type=int, default=10)
+    parser.add_argument("--label-smoothing", type=float, default=0.0)
+    return parser.parse_args()
+
+
+def _operator(implementation, logits, target, label_smoothing):
+    return parallel_cross_entropy(
+        logits,
+        target,
+        label_smoothing=label_smoothing,
+        reduce_loss=True,
+        overwrite_input=implementation == "destructive",
+    )
+
+
+def _make_inputs(shape, dtype):
+    gc.collect()
+    torch.cuda.empty_cache()
+    caller_input = torch.randn(shape, dtype=dtype, device="cuda", requires_grad=True)
+    target = torch.randint(0, shape[-1], shape[:-1], device="cuda")
+    torch.cuda.synchronize()
+    return caller_input, target
+
+
+def _measure_memory(implementation, shape, dtype, label_smoothing):
+    """Measure memory while keeping the caller-owned input alive throughout."""
+
+    caller_input, target = _make_inputs(shape, dtype)
+    baseline = torch.cuda.memory_allocated()
+    torch.cuda.reset_peak_memory_stats()
+
+    loss = _operator(implementation, caller_input, target, label_smoothing)
+    torch.cuda.synchronize()
+    forward_peak = torch.cuda.max_memory_allocated()
+    forward_end_live = torch.cuda.memory_allocated()
+
+    backward_entry = forward_end_live
+    torch.cuda.reset_peak_memory_stats()
+    loss.backward()
+    torch.cuda.synchronize()
+    backward_peak = torch.cuda.max_memory_allocated()
+    backward_end_live = torch.cuda.memory_allocated()
+
+    # Referencing the input here makes its consistent lifetime explicit. In
+    # destructive mode its storage now contains the derivative.
+    assert caller_input.data_ptr() != 0
+    return {
+        "baseline": baseline,
+        "forward_peak": forward_peak,
+        "forward_end": forward_end_live,
+        "backward_entry": backward_entry,
+        "backward_peak": backward_peak,
+        "backward_end": backward_end_live,
+        "e2e_peak": max(forward_peak, backward_peak),
+    }
+
+
+def _measure_latency(implementation, shape, dtype, label_smoothing):
+    """Measure uninterrupted forward and backward device latency."""
+
+    caller_input, target = _make_inputs(shape, dtype)
+    forward_start = torch.cuda.Event(enable_timing=True)
+    forward_end = torch.cuda.Event(enable_timing=True)
+    backward_end = torch.cuda.Event(enable_timing=True)
+
+    forward_start.record()
+    loss = _operator(implementation, caller_input, target, label_smoothing)
+    forward_end.record()
+    loss.backward()
+    backward_end.record()
+    torch.cuda.synchronize()
+
+    assert caller_input.data_ptr() != 0
+    return {
+        "forward_ms": forward_start.elapsed_time(forward_end),
+        "backward_ms": forward_end.elapsed_time(backward_end),
+        "total_ms": forward_start.elapsed_time(backward_end),
+    }
+
+
+def _format_memory(value, baseline):
+    return f"{value / MIB:9.1f} / +{(value - baseline) / MIB:7.1f}"
+
+
+def _print_results(results):
+    memory_fields = (
+        ("fwd peak", "forward_peak"),
+        ("fwd end", "forward_end"),
+        ("bwd entry", "backward_entry"),
+        ("bwd peak", "backward_peak"),
+        ("e2e peak", "e2e_peak"),
+        ("bwd end", "backward_end"),
+    )
+    header = ["implementation", "baseline MiB"]
+    header.extend(f"{label} abs/+inc MiB" for label, _ in memory_fields)
+    header.extend(("fwd ms", "bwd ms", "total ms"))
+    rows = []
+    for implementation, measurements in results.items():
+        memory_sample = measurements["memory"]
+        timings = measurements["timings"]
+        row = [implementation, f"{memory_sample['baseline'] / MIB:.1f}"]
+        row.extend(
+            _format_memory(memory_sample[field], memory_sample["baseline"])
+            for _, field in memory_fields
+        )
+        row.extend(
+            f"{mean(sample[field] for sample in timings):.3f}"
+            for field in ("forward_ms", "backward_ms", "total_ms")
+        )
+        rows.append(row)
+
+    widths = [max(len(header[idx]), *(len(row[idx]) for row in rows)) for idx in range(len(header))]
+    print("  ".join(value.ljust(widths[idx]) for idx, value in enumerate(header)))
+    print("  ".join("-" * width for width in widths))
+    for row in rows:
+        print("  ".join(value.ljust(widths[idx]) for idx, value in enumerate(row)))
+
+
+def main():
+    """Run the benchmark."""
+
+    args = parse_args()
+    if not torch.cuda.is_available():
+        raise RuntimeError("This benchmark requires CUDA")
+    dtype = torch.bfloat16 if args.dtype == "bf16" else torch.float32
+    shape = (args.batch_size, args.sequence_length, args.vocab_size)
+    n_logits = args.batch_size * args.sequence_length * args.vocab_size
+    n_rows = args.batch_size * args.sequence_length
+    input_bytes = n_logits * dtype.itemsize
+
+    expected_multipliers = {
+        "safe": 2 * dtype.itemsize,
+        "destructive": dtype.itemsize,
+    }
+    print(f"shape={shape}, dtype={dtype}, live input={input_bytes / MIB:.1f} MiB")
+    print("Expected logits-sized absolute forward peak (includes live caller input):")
+    for name, bytes_per_logit in expected_multipliers.items():
+        print(f"  {name:11s}: {bytes_per_logit}N = {bytes_per_logit * n_logits / MIB:.1f} MiB")
+    print(
+        "Saved row metadata: "
+        f"{(n_rows * (2 * 4 + 8) + 8) / MIB:.3f} MiB "
+        "(FP32 max/denominator, int64 targets/count)"
+    )
+    print(f"Forward loss temporary: {n_rows * 4 / MIB:.3f} MiB")
+    print("Tensor-parallel communication buffers: 0 MiB (single-GPU benchmark)")
+    print("Memory cells below are 'absolute / +incremental-from-pre-forward-baseline'.")
+
+    implementations = ("safe", "destructive")
+    for implementation in implementations:
+        for _ in range(args.warmup):
+            _measure_latency(implementation, shape, dtype, args.label_smoothing)
+
+    results = {}
+    for implementation in implementations:
+        results[implementation] = {
+            "memory": _measure_memory(implementation, shape, dtype, args.label_smoothing),
+            "timings": [
+                _measure_latency(implementation, shape, dtype, args.label_smoothing)
+                for _ in range(args.trials)
+            ],
+        }
+    _print_results(results)
+
+
+if __name__ == "__main__":
+    main()

qa/L0_pytorch_unittest/test.sh

@@ -48,7 +48,7 @@ python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_fusible_ops.xml
 python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_distributed_weight.xml $TE_PATH/tests/pytorch/test_distributed_weight.py || test_fail "test_distributed_weight.py"
 python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_backward_override.xml $TE_PATH/tests/pytorch/test_backward_override.py || test_fail "test_backward_override.py"
 python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_permutation.xml $TE_PATH/tests/pytorch/test_permutation.py || test_fail "test_permutation.py"
-python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_parallel_cross_entropy.xml $TE_PATH/tests/pytorch/test_parallel_cross_entropy.py || test_fail "test_parallel_cross_entropy.py"
+python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_cross_entropy.xml $TE_PATH/tests/pytorch/test_cross_entropy.py || test_fail "test_cross_entropy.py"
 python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_cpu_offloading.xml $TE_PATH/tests/pytorch/test_cpu_offloading.py || test_fail "test_cpu_offloading.py"
 NVTE_FLASH_ATTN=0 NVTE_CPU_OFFLOAD_V1=1 python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_cpu_offloading_v1.xml $TE_PATH/tests/pytorch/test_cpu_offloading_v1.py || test_fail "test_cpu_offloading_v1.py"
 NVTE_ALLOW_UNSAFE_PICKLE_EXTRA_STATE=1 python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_attention.xml $TE_PATH/tests/pytorch/attention/test_attention.py || test_fail "test_attention.py"

qa/L1_pytorch_distributed_unittest/test.sh

@@ -41,6 +41,7 @@ else
 fi
 
 python3 -m pytest -v -s --junitxml=$XML_LOG_DIR/pytest_test_sanity.xml $TE_PATH/tests/pytorch/distributed/test_sanity.py || test_fail "test_sanity.py"
+python3 -m pytest -v -s --junitxml=$XML_LOG_DIR/pytest_test_parallel_cross_entropy.xml $TE_PATH/tests/pytorch/distributed/test_parallel_cross_entropy.py || test_fail "test_parallel_cross_entropy.py"
 python3 -m pytest -v -s --junitxml=$XML_LOG_DIR/pytest_test_numerics.xml $TE_PATH/tests/pytorch/distributed/test_numerics.py || test_fail "test_numerics.py"
 python3 -m pytest -v -s --junitxml=$XML_LOG_DIR/pytest_test_numerics_exact.xml $TE_PATH/tests/pytorch/distributed/test_numerics_exact.py || test_fail "test_numerics_exact.py"
 python3 -m pytest -v -s --junitxml=$XML_LOG_DIR/pytest_test_fusible_ops.xml $TE_PATH/tests/pytorch/distributed/test_fusible_ops.py || test_fail "test_fusible_ops.py"

tests/pytorch/distributed/test_parallel_cross_entropy.py

@@ -0,0 +1,105 @@
+# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+#
+# See LICENSE for license information.
+
+import os
+import pathlib
+import sys
+import tempfile
+
+import pytest
+import torch
+import torch.distributed as dist
+import torch.multiprocessing as mp
+
+from transformer_engine.pytorch import parallel_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
+
+
+def _run_tensor_parallel(rank, world_size, init_file, label_smoothing):
+    """Two-rank correctness worker for tensor-parallel cross entropy."""
+
+    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,
+    )
+    try:
+        generator = torch.Generator().manual_seed(2025)
+        shape = (2, 3, 22)
+        local_vocab = shape[-1] // world_size
+        target = torch.randint(
+            0,
+            shape[-1],
+            shape[:-1],
+            generator=generator,
+        ).to(device)
+        target[0, 2] = -100
+        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)
+            vocab_start = rank * local_vocab
+            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_()
+                    local_before = local_logits.detach().clone()
+                    version_before = local_logits._version
+                    ref_logits = global_values.float().clone().requires_grad_()
+
+                    loss = parallel_cross_entropy(
+                        local_logits,
+                        target,
+                        label_smoothing=label_smoothing,
+                        reduce_loss=reduce_loss,
+                        dist_process_group=dist.group.WORLD,
+                        overwrite_input=overwrite_input,
+                    )
+                    ref_loss = torch.nn.functional.cross_entropy(
+                        ref_logits.reshape(-1, shape[-1]),
+                        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) if reduce_loss else external_grad
+                    loss.backward(loss_grad)
+                    ref_loss.backward(loss_grad)
+                    assert local_logits._version == version_before + int(overwrite_input)
+                    if overwrite_input:
+                        assert not torch.equal(local_logits, local_before)
+                    else:
+                        torch.testing.assert_close(local_logits, local_before, rtol=0.0, atol=0.0)
+
+                    torch.testing.assert_close(loss, ref_loss, **dtype_tols(torch.float32))
+                    expected_grad = ref_logits.grad[
+                        ..., vocab_start : vocab_start + local_vocab
+                    ].to(dtype)
+                    torch.testing.assert_close(
+                        local_logits.grad, expected_grad, **dtype_tols(dtype)
+                    )
+    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."""
+
+    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,
+            args=(world_size, init_file, label_smoothing),
+            nprocs=world_size,
+            join=True,
+        )