Calling backward(retain_graph=True) multiple times with TE Layer does not work

[kshitij12345]

import torch
from transformer_engine.pytorch import Linear as TELinear, fp8_autocast

# m = torch.nn.Linear(16, 16).to("cuda") # This works
m = TELinear(16, 16)
x = torch.randn(16, 16, device='cuda')

with fp8_autocast(True):
 o = m(x).sum()

o.backward(retain_graph=True)

# this fails with
# AssertionError: FP8 execution requires 2D input matrices with height divisible by 8 and width divisible by 16, but got tensor with dims=[0]
# looks like TELinear.backward mutates the context object such that it is not reusable.
o.backward()

This would be useful to support benchmarking just the backward pass.

[timmoon10]

In the module-based API, we deliberately clear the intermediate tensors to minimize memory usage (see #509). For example, in _Linear.backward:

TransformerEngine/transformer_engine/pytorch/module/linear.py

Lines 587 to 588 in 56e0b35

	clear_tensor_data(inputmat_total)
	clear_tensor_data(inputmat_t_total)

We should mark these autograd backwards with once_differentiable for clarity.

Note that the functional API in #707 does not manually deallocate intermediate tensors (the operators themselves do deallocate to achieve the same memory savings as the module-based API). Once that is merged, it should be better for this use-case.

[IvanYashchuk]

With TE version 2.1.0+f2b09d2 the error is

File /usr/local/lib/python3.12/dist-packages/transformer_engine/pytorch/module/linear.py:355, in _Linear.backward(ctx, grad_output)
 349 raise NotImplementedError(
 350 "Comm+GEMM overlap is only supported with FP8 delayed scaling"
 351 )
 353 saved_tensors = ctx.saved_tensors
 354 inputmat, weight_fp8, weight, bias = ( # pylint: disable=unbalanced-tuple-unpacking
--> 355 restore_from_saved(ctx.tensor_objects, saved_tensors)
 356 )
 357 # Delete the references to tensor objects once they've been consumed
 358 # by the `restore_from_saved` method to construct back the actual tensors.
 359 ctx.tensor_objects = None

File /usr/local/lib/python3.12/dist-packages/transformer_engine/pytorch/tensor/quantized_tensor.py:46, in restore_from_saved(tensors, saved_tensors)
 44 """Recombine the tensor data and metadata during backward pass."""
 45 tensor_objects = []
---> 46 for tensor in tensors:
 47 if tensor is None:
 48 tensor_objects.append(saved_tensors[0])

TypeError: 'NoneType' object is not iterable

[IvanYashchuk]

We should mark these autograd backwards with once_differentiable for clarity.

In PyTorch, once_differentiable functions are allowed and expected to work with multiple backward calls with retain_grad=True.

retain_grad=True is a request to PyTorch's Autograd Engine not to deallocate saved tensors for backward computation.
These saved tensors should be deallocated with Python's del operator. clear_tensor_data is a great workaround but it's better to use ctx.maybe_clear_saved_tensors() (example in Thunder).