Lightning-AI / Lightning-AI/lightning-thunder
`thunderfx(fsdp2_model)` with Compiled Autograd fails in `ctx.maybe_clear_saved_tensors()` in backward.
- Dominant language
- Python
- Stars
- 1.5k
- Forks
- 121
- PR merge metrics
- No merged PRs in 30d
Description
## 🐛 Bug
When I tried the combination of `thunderfx`, `fsdp2`, and CompiledAutograd, I see the error:
```
[rank0]:[rank0]: Traceback (most recent call last):
[rank0]:[rank0]: File "/opt/pytorch/lightning-thunder/thunder/benchmarks/benchmark_litgpt.py", line 1004, in
[rank0]:[rank0]: CLI(benchmark_main)
[rank0]:[rank0]: File "/usr/local/lib/python3.12/dist-packages/jsonargparse/_cli.py", line 96, in CLI
[rank0]:[rank0]: return _run_component(components, init)
[rank0]:[rank0]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]:[rank0]: File "/usr/local/lib/python3.12/dist-packages/jsonargparse/_cli.py", line 204, in _run_component
[rank0]:[rank0]: return component(**cfg)
[rank0]:[rank0]: ^^^^^^^^^^^^^^^^
[rank0]:[rank0]: File "/opt/pytorch/lightning-thunder/thunder/benchmarks/benchmark_litgpt.py", line 899, in benchmark_main
[rank0]:[rank0]: benchmark.train()
[rank0]:[rank0]: File "/opt/pytorch/lightning-thunder/thunder/benchmarks/benchmark_litgpt.py", line 783, in train
[rank0]:[rank0]: loss.backward()
[rank0]:[rank0]: File "/usr/local/lib/python3.12/dist-packages/torch/_tensor.py", line 648, in backward
[rank0]:[rank0]: torch.autograd.backward(
[rank0]:[rank0]: File "/usr/local/lib/python3.12/dist-packages/torch/autograd/__init__.py", line 354, in backward
[rank0]:[rank0]: _engine_run_backward(
[rank0]:[rank0]: File "/usr/local/lib/python3.12/dist-packages/torch/autograd/graph.py", line 829, in _engine_run_backward
[rank0]:[rank0]: return Variable._execution_engine.run_backward( # Calls into the C++ engine to run the backward pass
[rank0]:[rank0]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]:[rank0]: File "/usr/local/lib/python3.12/dist-packages/torch/_dynamo/compiled_autograd.py", line 1041, in runtime_wrapper
[rank0]:[rank0]: out = compiled_fn(
[rank0]:[rank0]: ^^^^^^^^^^^^
[rank0]:[rank0]: File "/usr/local/lib/python3.12/dist-packages/torch/_dynamo/eval_frame.py", line 372, in __call__
[rank0]:[rank0]: return super().__call__(*args, **kwargs)
[rank0]:[rank0]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]:[rank0]: File "/usr/local/lib/python3.12/dist-packages/torch/nn/modules/module.py", line 1767, in _wrapped_call_impl
[rank0]:[rank0]: return self._call_impl(*args, **kwargs)
[rank0]:[rank0]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]:[rank0]: File "/usr/local/lib/python3.12/dist-packages/torch/nn/modules/module.py", line 1778, in _call_impl
[rank0]:[rank0]: return forward_call(*args, **kwargs)
[rank0]:[rank0]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]:[rank0]: File "/usr/local/lib/python3.12/dist-packages/torch/_dynamo/eval_frame.py", line 708, in compile_wrapper
[rank0]:[rank0]: raise e.with_traceback(None) from e.__cause__ # User compiler error
[rank0]:[rank0]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]:[rank0]: torch._dynamo.exc.Unsupported: Unsupported autograd.Function context method
[rank0]:[rank0]: Explanation: Dynamo does not support calling the method `maybe_clear_saved_tensors` on `autograd.Function` context objects. Supported methods are `__setattr__`, `save_for_backward` and `mark_non_differentiable`.
[rank0]:[rank0]: Hint: It may be possible to write Dynamo tracing rules for this code. Please report an issue to PyTorch if you encounter this graph break often and it is causing performance issues.
[rank0]:
[rank0]:[rank0]: Developer debug context: call_method AutogradFunctionContextVariable(ThunderFunctionBackward) maybe_clear_saved_tensors
[rank0]:
[rank0]:
[rank0]:[rank0]: from user code:
[rank0]:[rank0]: File ".21", line 204, in forward
[rank0]:[rank0]: call_backward = torch__dynamo_external_utils_call_backward(getitem_149, (getitem_9, getitem_10, getitem_11), call_hook_2); getitem_149 = getitem_10 = getitem_11 = call_hook_2 = None
[rank0]:[rank0]: File "/usr/local/lib/python3.12/dist-packages/torch/_dynamo/external_utils.py", line 133, in call_backward
[rank0]:[rank0]: grads = fake._forward_cls.backward(fake, *args) # type: ignore[attr-defined]
[rank0]:[rank0]: File "/usr/local/lib/python3.12/dist-packages/torch/autograd/function.py", line 604, in wrapper
[rank0]:[rank0]: outputs = fn(ctx, *args)
[rank0]:[rank0]: File "/opt/pytorch/lightning-thunder/thunder/executors/torch_autograd.py", line 156, in backward
[rank0]:[rank0]: ctx.maybe_clear_saved_tensors() # Delete the reference to all saved tensors in the context
[rank0]:
[rank0]:[rank0]: Set TORCHDYNAMO_VERBOSE=1 for the internal stack trace (please do this especially if you're reporting a bug to PyTorch). For even more developer context, set TORCH_LOGS="+dynamo"
```
The function call is here -- https://github.com/Lightning-AI/lightning-thunder/blob/74ac17eb1676c064b48ae48a8935ff062a7bcdad/thunder/executors/torch_autograd.py#L156.
### To Reproduce
Steps to reproduce the behavior:
1. Go to https://github.com/Lightning-AI/lightning-thunder/blob/74ac17eb1676c064b48ae48a8935ff062a7bcdad/thunder/benchmarks/benchmark_litgpt.py
2. Modify the script by calling the forward and backward in the context of `torch._dynamo.utilsmaybe_enable_compiled_autograd` -- https://github.com/pytorch/pytorch/blob/e1f28fe17bc55389fc5a31bad588d22e5ecca722/torch/_dynamo/utils.py#L4291
1. Run the script e.g. `torchrun --role rank --tee 3 --local-ranks-filter 0 --nproc-per-node=8 thunder/benchmarks/benchmark_litgpt.py --compile=thunder_dynamo --distributed_mode=fsdp2 --nsys_enabled=False --micro_batch_size=1 --global_batch_size=8 --model_name=Llama-2-7b-hf --shard_mode=zero3 --bucketing_mode=none --json_path "llma-2-13b-zero3-no-bucketing.json" --return_metrics_as_json=true --warmup_iters 2 --max_iters 10 --n_layers 1`
1. See error
#### Code sample
Modified script is available https://gist.github.com/crcrpar/667bae23d4825fb38a6831331491fb6b.
### Expected behavior
### Environment
pjnl-20250619
- PyTorch Version (e.g., 1.0): 2.8.0a0+34c6371d24
- OS (e.g., Linux):
- How you installed PyTorch (`conda`, `pip`, source): source
- Build command you used (if compiling from source):
- Python version: 3.12
- CUDA/cuDNN version:
- GPU models and configuration:
- Any other relevant information:
### Additional context
The reason why I tried CompiledAutograd is that it would put parameter unsharding communications into `GraphModule`s of TorchDynamo outputs.
IMHO it can give thunder more flexibility than without CompiledAutograd when we use fsdp2.
Contributor guide
No contributing guide indexed for this repository
Research direction
Start with thunder/executors/torch_autograd.py at the ctx.maybe_clear_saved_tensors() call and reproduce the failure using thunder/benchmarks/benchmark_litgpt.py with Compiled Autograd and fsdp2. Compare the behavior with the provided modified script and gist, then verify the benchmark no longer raises the Unsupported autograd.Function context method error under the listed torchrun configuration.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python
- Domain
- compilers, distributed-systems
- Issue type
- Bug
- Difficulty
- 4/5
- Estimated time
- 3-5 days
- Activity status
- Stale
- Clarity
- Mostly clear
- Newbie friendliness
- 35/100