[TransformerEngine] Support `backward(retain_graph=True)`

@kshitij12345 is already working on this.

Since Feb 7, 2025.

```python
from thunder.benchmarks.targets import LitGPTConfig, LitGPTBenchmark, backward_only
import torch
import thunder
from thunder.executors.transformer_engineex import transformer_engine_ex

# With bigger model
# cfg: LitGPTConfig = LitGPTConfig.from_name("Llama-2-7b-hf")
# cfg.n_layer = 3
# b = LitGPTBenchmark(cfg, batchdims=(2,), device="cuda:0", dtype=torch.bfloat16, requires_grad=True)
# args, kwargs = b.make_batch()
# fn = thunder.jit(b.fn(), executors=[transformer_engine_ex,])

# With smaller model
def foo(x, w):
return torch.nn.functional.linear(x, w)
x = torch.randn(16, 16, requires_grad=True, device='cuda')
w = torch.randn(16, 16, requires_grad=True, device='cuda')
args = (x, w)
kwargs = {}
fn = thunder.jit(foo, executors=[transformer_engine_ex,])

# backward_only creates a graph and calls `torch.autograd.backward` with `retain_graph=True`.
backward_fn, backward_setup = backward_only(fn, *args, **kwargs)
backward_args = backward_setup()
backward_fn(*backward_args)

# Second usage of `backward_fn` fails with
# File "/home/kkalambarkar/lightning-thunder/thunder/executors/transformer_engineex.py", line 352, in _te_functional_linear_backward_impl
# with enable_grad(ctx.saved_tensors[2]):
# IndexError: tuple index out of range
backward_fn(*backward_args)

```

cc: @IvanYashchuk