Minimize the memory usage of the fused cross entropy kernel

[ptrendx]

Description

The changes introduced in #3193, while fixing the numerical issue with the fused cross entropy kernel, introduced additional memory requirement. This PR aims to greatly reduce that additional memory usage while preserving the numerics and optimizing runtime performance.

Type of change

• Documentation change (change only to the documentation, either a fix or a new content)
• Bug fix (non-breaking change which fixes an issue)
• New feature (non-breaking change which adds functionality)
• Breaking change (fix or feature that would cause existing functionality to not work as expected)
• Infra/Build change
• Code refactoring

Changes

Please list the changes introduced in this PR:

• The original algorithm computed the full derivative in the forward pass and then passed it to backward pass in order to be multiplied by the incoming gradients. In this PR the thing saved to backward is instead the statistics needed to compute the derivative. This lowers the memory consumption and the number of bytes needed to be read by the backward kernel.
• The new flow contains 2 paths - single-GPU with a single kernel for forward computation and the multi-GPU one with 2 kernels, similar to the existing path.

Checklist:

• I have read and followed the contributing guidelines
• The functionality is complete
• I have commented my code, particularly in hard-to-understand areas
• I have made corresponding changes to the documentation
• My changes generate no new warnings
• I have added tests that prove my fix is effective or that my feature works
• New and existing unit tests pass locally with my changes

[ptrendx]

/te-ci pytorch

[greptile-apps]

Greptile Summary

The PR reduces fused cross-entropy memory usage by saving compact softmax statistics and reconstructing gradients during backward.

• Adds separate single-GPU and tensor-parallel Triton forward paths.
• Adds an opt-in destructive mode that reuses input storage during backward.
• Expands single-GPU and distributed numerical, layout, validation, and buffer-reuse coverage.
• Updates CI routing and adds a memory and latency benchmark.

Confidence Score: 5/5

The PR appears safe to merge because no blocking failure remains.

No blocking failure remains.

Important Files Changed

Filename	Overview
transformer_engine/common/triton/cross_entropy.py	Replaces forward-computed full derivatives with single-rank and tensor-parallel statistics kernels plus backward gradient reconstruction.
transformer_engine/pytorch/triton/cross_entropy.py	Allocates compact saved state, coordinates tensor-parallel statistics gathering, and launches the reconstructed-gradient kernel.
transformer_engine/pytorch/cross_entropy.py	Extends the public API with validated destructive buffer reuse and updates autograd state handling.
tests/pytorch/test_cross_entropy.py	Adds numerical, saved-state, buffer-aliasing, layout, validation, and compatibility coverage.
tests/pytorch/distributed/test_parallel_cross_entropy.py	Adds two-rank tensor-parallel comparisons against PyTorch for supported dtypes and modes.
benchmarks/benchmark_parallel_cross_entropy.py	Adds forward/backward memory and latency comparisons for safe and destructive modes.
qa/L0_pytorch_unittest/test.sh	Updates the single-GPU CI invocation for the renamed cross-entropy test file.
qa/L1_pytorch_distributed_unittest/test.sh	Adds the distributed cross-entropy suite to multi-GPU CI.

Sequence Diagram

sequenceDiagram
    participant Caller
    participant Autograd as CrossEntropyFunction
    participant Forward as Triton forward
    participant Collective as TP all-gather
    participant Backward as Triton backward
    Caller->>Autograd: logits, targets, options
    Autograd->>Forward: launch forward kernel
    alt Single GPU
        Forward-->>Autograd: loss, saved input, max/denominator stats
    else Tensor parallel
        Forward->>Collective: gather local statistics
        Collective-->>Forward: global statistics
        Forward-->>Autograd: loss, saved input, global stats
    end
    Autograd-->>Caller: loss
    Caller->>Autograd: backward(grad_output)
    Autograd->>Backward: saved input, targets, stats, grad_output
    Backward-->>Autograd: reconstructed input gradient
    Autograd-->>Caller: logits gradient

Loading

_{Reviews (2): Last reviewed commit: "Add the guard against torch.compile when..." | Re-trigger Greptile}

[janbernloehr]

Thanks for the fix. We validated commit 41157db on a downstream H100 TP=1 BF16 workload with 4096 × 128256 logits.

• Baseline: 1.96 GiB FP32 CE allocation, OOM.
• overwrite_input=False: 1002 MiB BF16 allocation, still OOM.
• overwrite_input=True: no additional CE allocation; 5/5 iterations passed.

With destructive reuse enabled, peak memory matched the previous baseline exactly, with no throughput regression. The numerical fix from #3193 remained active.
This suggests that the default safe mode can still OOM for large padded vocabularies. Please consider documenting this distinction and supporting the overwrite_input=True path in downstream callers where the input is not reused after backward. A release_v2.18 backport would also be valuable for users still on that release.
These results were obtained against 41157db; since the current PR head has changed, it would be useful to rerun this workload on the latest commit.

[pggPL]

Torch.compile support triton kernels and adding overwrite_input=True can lead to silent incorectness under compile.
May be worth to force no_torch_dynamo in that case.

[pggPL]

Otherwise LGTM

[ptrendx]

/te-ci pytorch

[ptrendx]

@pggPL Addressed the torch.compile comment