Lightning-AI / Lightning-AI/lightning-thunder

`thunderfx(fsdp2_model)` with Compiled Autograd fails in `ctx.maybe_clear_saved_tensors()` in backward.

Open
#2,263 2 comments 0 reactions 0 assignees View on GitHub
thunderfx
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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.