[BUG] `CrossEntropyFunction.forward()` modifies input in-place without `ctx.mark_dirty()`, causing GPU memory leak

CrossEntropyFunction in transformer_engine/pytorch/cross_entropy.py modifies the input tensor in-place via Triton kernels (online_softmax_kernel + cross_entropy_kernel) but does not call ctx.mark_dirty(inp). This violates the PyTorch torch.autograd.Function contract and causes Function.apply()'s C++ layer to permanently retain a reference to the input tensor, leaking GPU memory on every call.

Root Cause

In CrossEntropyFunction.forward():

def forward(ctx, inp, target, ...):
    loss, inp = triton_cross_entropy.cross_entropy_forward(inp, target, ...)
    # ↑ Triton kernels overwrite inp data in-place (logits → softmax/gradient data)
    ctx.save_for_backward(inp.detach())
    # ↑ Missing: ctx.mark_dirty(inp)
    return loss

triton_cross_entropy.cross_entropy_forward() passes inp as X_ptr to online_softmax_kernel and cross_entropy_kernel, which overwrite the tensor's data in-place. However, since ctx.mark_dirty(inp) is never called, PyTorch's C++ THPFunction_apply retains a strong reference to the original input tensor for version-counter consistency checking. This reference is held by a C++ object and is invisible to Python's garbage collector.

Impact

• Memory leak: Each call to CrossEntropyFunction.apply() leaks the full logits tensor (shape [seq_len, vocab_size/tp], bf16). In our case, ~200-600 MiB per micro-batch, accumulating to 3-5 GiB per training step with 16 micro-batches.
• OOM: Training crashes after a few steps.
• Affected configurations: Any training using TE's cross entropy via parallel_cross_entropy() or CrossEntropyFunction.apply(), especially noticeable with gradient accumulation (multiple micro-batches per step).

Reproduction

import torch
import transformer_engine.pytorch.cross_entropy as te_ce

# Simulate repeated calls (as in gradient accumulation)
for i in range(16):
    logits = torch.randn(1024, 32000, dtype=torch.bfloat16, device='cuda', requires_grad=True)
    target = torch.randint(0, 32000, (1024,), device='cuda')
    loss = te_ce.parallel_cross_entropy(logits, target)
    loss.sum().backward()
    del loss, logits, target
    
    # Memory keeps growing despite del + backward
    print(f"Step {i}: {torch.cuda.memory_allocated()/1024**3:.2f} GiB")

Environment

• TransformerEngine version: installed via pip (with PyTorch 2.10.0+cu128)
• PyTorch version: 2.10.0+cu128
• GPU: A800-80GB

Workaround

Track the input tensor reference in a global deque during forward(), then storage().resize_(0) after backward completes:

import collections
_leaked_ce_inputs = collections.deque()

def _drain_leaked_inputs():
    while _leaked_ce_inputs:
        t = _leaked_ce_inputs.popleft()
        if t.untyped_storage().size() > 0:
            t.storage().resize_(0)
        del t

# In CrossEntropyFunction.forward():
#   _leaked_ce_inputs.append(inp)  # before triton call

# In training loop, after backward_step():
#   _drain_leaked_inputs()

Suggested Fix

Add ctx.mark_dirty(inp) in CrossEntropyFunction.forward() to properly inform PyTorch's autograd that the input was modified in-place. Note that mark_dirty requires the modified tensor to also be returned as an output, so the function signature may need adjustment. Alternatively, clone the input before passing to Triton kernels so the original input is not modified.

Start in transformer_engine/pytorch/cross_entropy.py at CrossEntropyFunction.forward() and trace triton_cross_entropy.cross_entropy_forward(), including the online_softmax_kernel and cross_entropy_kernel calls. Run the supplied CUDA reproduction and inspect the PyTorch autograd Function contract for ctx.mark_dirty() and its output requirements. Done means repeated calls no longer retain the input logits or grow GPU memory after backward while preserving the cross-entropy behavior.