Autodiff does not place the correct producer for recomputation(checkpointing)

@riccardofelluga is already working on this.

Since Jun 3, 2025.

## 🐛 Bug

I've noticed that with the new autodiff, if an executor emits more than one symbol for a given primitive, not all producer symbols are placed in the backward trace but only the last one.

In my mind the correct behavior would be to find all the producers up to the input to the checkpointing function and place all the producer symbols in the backward trace for recomputation.

### To Reproduce

Unfortunately the repro is not super small but it is minimal to show the behavior:

```python
from functools import partial

import torch

from thunder.core.prims import get_grad, put_grad
from thunder.core.proxies import TensorProxy
from thunder.extend import OperatorExecutor, register_executor
import thunder.torch as ltorch

repro_ex = OperatorExecutor("repro_ex")
register_executor(repro_ex)

def _linear_fwd_meta(a, w, bias):
out_shape = (*a.shape[:-1], w.shape[0])
return TensorProxy(like=a, shape=out_shape)

def _linear_fwd_impl(a, w, bias):
return torch.nn.functional.linear(a, w, bias=bias)

_linear_fwd = repro_ex.register_operator("dummy_linear_fwd", meta=_linear_fwd_meta, fn=_linear_fwd_impl)

def _pass_token_meta(tokens):
return (*(TensorProxy(like=token) for token in tokens), )

def _pass_token_impl(tokens):
return tokens

_dummy_pass_token = repro_ex.register_operator("dummy_pass_token", meta=_pass_token_meta, fn=_pass_token_impl)

def _linear_bwd_meta(grad, a, w):
return TensorProxy(like=a), TensorProxy(like=w)

def _linear_bwd_impl(grad, a, w):
return a, w

_linear_bwd = repro_ex.register_operator("dummy_linear_bwd", meta=_linear_bwd_meta, fn=_linear_bwd_impl)

def exec_transform(a, w, bias):
primal = _linear_fwd(a, w, bias)
(primal, ) = _dummy_pass_token((primal, ))

def grad_transform(a, w, bias):
primal = _linear_fwd(a, w, bias)
(primal, ) = _dummy_pass_token((primal, ))

grad_out = get_grad(primal)

grad_a, grad_w = _linear_bwd(grad_out, a, w)
(grad_a, grad_w) = _dummy_pass_token((grad_a, grad_w))

put_grad(a, grad_a)
put_grad(w, grad_w)

return primal

repro_ex.register_implementation(
ltorch.linear,
checker=lambda a, w, bias: True,
execution_transform=exec_transform,
grad_transform=grad_transform,
)

def fn_to_checkpoint(x, y):
o = torch.nn.functional.linear(x, y, bias=None)
return o.sin()

checkpoint_fn = partial(torch.utils.checkpoint.checkpoint, use_reentrant=False)

def f(x, y):
return checkpoint_fn(fn_to_checkpoint, x, y)

x = torch.randn((16, 16), device="cuda", dtype=torch.float32, requires_grad=True)
y = torch.randn((16, 16), device="cuda", dtype=torch.float32, requires_grad=True)

from thunder.dynamo import thunderfx

jf = thunderfx(f, executors=[repro_ex])

out = jf(x, y)

print(jf.last_backward_traces[-1])
```

This will print the backward trace where there are no `dummy_linear_fwd` even tho they should.

```python
import torch
from thunder.executors.torchex import no_autocast

@torch.no_grad()
@no_autocast
def backward_fn(saved_for_backward, cotangents):
# saved_for_backward: "Collection"
# cotangents: "Collection"
C0, _, = saved_for_backward
# C0: "Collection"
# None
clear_mutable_collection(saved_for_backward)
del saved_for_backward
t12, = cotangents
# t12: "cuda:0 f32[16, 16]"
clear_mutable_collection(cotangents)
del cotangents
l_x_, l_y_, t4, = C0
# l_x_: "cuda:0 f32[16, 16]"
# l_y_: "cuda:0 f32[16, 16]"
# t4: "cuda:0 f32[16, 16]"
clear_mutable_collection(C0)
del C0

# .7:6: tag_activation_checkpoint = torch.ops.higher_order.tag_activation_checkpoint(wrap_body_0, l_x_, l_y_, use_reentrant = False); wrap_body_0 = l_x_ = l_y_ = None
(bw_t5,) = dummy_pass_token((t4,))
del t4
bw_t16 = torch.cos(bw_t5) # bw_t16: "cuda:0 f32[16, 16]"
# bw_t16 = ltorch.cos(bw_t5) # bw_t16: "cuda:0 f32[16, 16]"
# bw_t16 = prims.cos(bw_t5) # bw_t16: "cuda:0 f32[16, 16]"
del bw_t5
bw_t17 = torch.mul(t12, bw_t16) # bw_t17: "cuda:0 f32[16, 16]"
# bw_t17 = ltorch.mul(t12, bw_t16) # bw_t17: "cuda:0 f32[16, 16]"
# bw_t17 = prims.mul(t12, bw_t16) # bw_t17: "cuda:0 f32[16, 16]"
del t12, bw_t16
(bw_t7, bw_t8) = dummy_linear_bwd(bw_t17, l_x_, l_y_)
del bw_t17, l_x_, l_y_
(grad_for_l_x_, grad_for_l_y_) = dummy_pass_token((bw_t7, bw_t8))
del bw_t7, bw_t8
return (grad_for_l_x_, grad_for_l_y_)
```

cc. @t-vi @beverlylytle

cc @mruberry @lantiga