Autodiff transform prunes incorrectly registered operations (non-prims without subsymbols) without a grad_transform when followed by a stateful nn.Module

@t-vi is already working on this.

Since Jul 29, 2025.

## 🐛 Bug

Thunder's `autodiff` transform appears to unexpectedly prune operations that do not have a registered `grad_transform`. This behavior is observed when such an operation is followed in the computation graph by a stateful `nn.Module` (like `nn.Linear`) that triggers parameter lifting.

The introduction of a stateful module causes its parameters to be "lifted" and treated as inputs to the graph. It seems that in this scenario, the dependency analysis within the `AugmentedForwardProcessor` does not fully trace the gradient requirements from the stateful module's backward pass back to the output of the preceding operation (the one without a `grad_transform`).

Consequently, the processor may incorrectly conclude that the operation is not needed for any gradient computation and prune it from the graph, which can lead to a corrupted trace or a runtime error.

### To Reproduce

Run the following script with `--repro` option.

#### Code sample

```python
import argparse

import torch
import torch.nn as nn

import thunder
from thunder.core.symbol import Symbol
from thunder.core.proxies import TensorProxy
from thunder.executors.torchex import ex
from thunder.executors.torchex import _always_executable

# 1. Define a simple Python function and its meta-function
def my_simple_mul(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor:
# The actual implementation doesn't matter for demonstrating the bug,
# as it's the graph structure that triggers it.
return a * b

def my_simple_mul_meta(a: TensorProxy, b: TensorProxy) -> TensorProxy:
# Meta functions should return empty tensors with the correct metadata
return TensorProxy(like=a)

# 2. Manually create a Thunder Symbol for the function
# This symbol will be "opaque" to the autodiff transform because it has no
# registered grad_transform.
my_simple_mul_symbol = Symbol(
name="my_simple_mul",
meta=my_simple_mul_meta,
is_prim=False,
)

# 3. Register an implementation for the symbol with an executor.
ex.register_implementation(my_simple_mul_symbol, my_simple_mul, checker=_always_executable)

# 4. Make the symbol available in a thunder namespace to be called from the model
# (This is for convenience in the reproducer script)
setattr(thunder.torch, "my_simple_mul", my_simple_mul_symbol)

# 5. Define the two nn.Modules to compare
# This module should work correctly.
class SinModule(nn.Module):
def forward(self, x, y):
intermediate = thunder.torch.my_simple_mul(x, y)
output = torch.sin(intermediate)
return output

# This module will demonstrate the bug.
class LinearModule(nn.Module):
def __init__(self):
super().__init__()
self.linear = nn.Linear(2, 2, bias=False)

def forward(self, x, y):
intermediate = thunder.torch.my_simple_mul(x, y)
output = self.linear(intermediate)
return output

if __name__ == "__main__":
parser = argparse.ArgumentParser(description="Demonstrate a bug in Thunder's autodiff transform.")
parser.add_argument(
"--repro",
action="store_true",
help="Run the LinearModule case to reproduce the bug. Default is to run the working SinModule case.",
)
args = parser.parse_args()

x = torch.randn(4, 2, device="cuda")
y = torch.randn(4, 2, device="cuda")

if args.repro:
print("\n--- Testing LinearModule (This should demonstrate the bug) ---")
model_class = LinearModule
model_name = "LinearModule"
else:
print("--- Testing SinModule (This should work correctly) ---")
model_class = SinModule
model_name = "SinModule"

model = model_class().cuda()
jitted_model = thunder.jit(
model,
disable_torch_autograd=False,
)

jitted_model(x, y)
```

Running this script with `--repro` would result in
```
--- Testing LinearModule (This should demonstrate the bug) ---
Traceback (most recent call last):
File "/home/mkozuki/ghq/github.com/Lightning-AI/lightning-thunder/concise_repro.py", line 92, in
jitted_model(x, y)
File "/home/mkozuki/ghq/github.com/crcrpar/torch-2/torch/nn/modules/module.py", line 1773, in _wrapped_call_impl
return self._call_impl(*args, **kwargs)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/home/mkozuki/ghq/github.com/crcrpar/torch-2/torch/nn/modules/module.py", line 1784, in _call_impl
return forward_call(*args, **kwargs)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/home/mkozuki/ghq/github.com/Lightning-AI/lightning-thunder/thunder/core/module.py", line 80, in forward
res = self._forward_fn(*args, **kwargs)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/home/mkozuki/ghq/github.com/Lightning-AI/lightning-thunder/thunder/__init__.py", line 839, in wrapped
return fn(*args, **kwargs)
^^^^^^^^^^^^^^^^^^^
File "/home/mkozuki/ghq/github.com/Lightning-AI/lightning-thunder/thunder/__init__.py", line 881, in fn_
result = cache_entry.computation_fn(*inps)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/home/mkozuki/ghq/github.com/Lightning-AI/lightning-thunder/thunder/__init__.py", line 800, in wrapped
return fn(*args, **kwargs)
^^^^^^^^^^^^^^^^^^^
File "/home/mkozuki/ghq/github.com/crcrpar/torch-2/torch/utils/_contextlib.py", line 120, in decorate_context
return func(*args, **kwargs)
^^^^^^^^^^^^^^^^^^^^^
File "/home/mkozuki/ghq/github.com/Lightning-AI/lightning-thunder/thunder/executors/torchex.py", line 167, in no_autocast_fn
return fn(*args, **kwargs)
^^^^^^^^^^^^^^^^^^^
File "thunder.computation_3", line 14, in computation
NameError: name 'intermediate' is not defined
```

### Expected behavior

Both path should not behave differently.

### Environment

- PyTorch Version (e.g., 1.0): 2.9.0a0+git476874b
- OS (e.g., Linux): Ubuntu 24.04
- How you installed PyTorch (`conda`, `pip`, source): source
- Build command you used (if compiling from source): N/A
- Python version: 3.12.10
- CUDA/cuDNN version: 12.9/ N/A
- GPU models and configuration: RTX600Ada
- Any other relevant information: N/A