Lightning-AI / Lightning-AI/lightning-thunder
Grad Transform generates inconsistent saved_for_backward between forward and backward trace.
Nobody has claimed this yet.
- Dominant language
- Python
- Stars
- 1.5k
- Forks
- 121
- PR merge metrics
- No merged PRs in 30d
Description
In branch #1712
```py
import thunder
def foo(a, b):
return a * b
import torch
a = torch.randn(1, 32, 232, 232)
a.requires_grad_()
b = torch.randn(1, 1, 232, 232)
#b.requires_grad_() # adding this fixes the issue. So does this small code change [here](https://github.com/Lightning-AI/lightning-thunder/pull/1725/files#diff-dffa7275f3b675c42b33a9f91f1259dd4167631b3b6318ea03981cfbb1f87073R78-R79)
jfoo = thunder.jit(foo, cache="symbolic values")
out = jfoo(a, b)
out.sum().backward()
```
Hitting the issue
```py
Traceback (most recent call last):
File "/volume/thunder_bwd6.py", line 21, in
out = jfoo(a, b)
^^^^^^^^^^
File "/opt/pytorch/lightning-thunder/thunder/__init__.py", line 741, in wrapped
return fn(*args, **kwargs)
^^^^^^^^^^^^^^^^^^^
File "/opt/pytorch/lightning-thunder/thunder/__init__.py", line 777, in fn_
cache_entry, inps, pro_to_epi = get_computation_and_inputs(*args, **kwargs)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/opt/pytorch/lightning-thunder/thunder/__init__.py", line 723, in wrapped
cache_entry, inps, pro_to_epi = get_computation_and_inputs_fn(*args, **kwargs)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/opt/pytorch/lightning-thunder/thunder/core/langctxs.py", line 136, in _fn
result = fn(*args, **kwargs)
^^^^^^^^^^^^^^^^^^^
File "/opt/pytorch/lightning-thunder/thunder/__init__.py", line 235, in cache_info_wrapper
res = fn(*args, **kwargs)
^^^^^^^^^^^^^^^^^^^
File "/opt/pytorch/lightning-thunder/thunder/__init__.py", line 629, in get_computation_and_inputs
computation_trc, backward_trc = split_forward_backward(computation_trc, cd, cs, *inps)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/opt/pytorch/lightning-thunder/thunder/executors/torch_autograd.py", line 362, in split_forward_backward
fw_extrace, bw_extrace = rematerialize_forward_and_backward(fw_extrace, bw_extrace)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/opt/pytorch/lightning-thunder/thunder/core/rematerialization.py", line 738, in rematerialize_forward_and_backward
_, new_required_for_backward_bw_to_fw_map = match_fw_and_bw_saved_for_bw_proxies(fw_trace, bw_trace)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/opt/pytorch/lightning-thunder/thunder/core/rematerialization.py", line 553, in match_fw_and_bw_saved_for_bw_proxies
assert len(old_saved_for_backward_fw) == len(old_saved_for_backward_bw)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
AssertionError
```
through my debugging, I think the inconsistency is coming in that, when we generate `saved_for_backward`, we look at the backward trace without considering whether certain portion requires_grad or not.
Afterwards, the backward trace trims portion that doesn't require grad and dce the ops away, including arguments used by those, which might be coming from `saved_for_backward`. Hence triggering the issue above.
My mental model is that, requires_grad is supposed to be propagated across the trace and when we run grad transform, we only do put_grad on TensorProxies with requires_grad == True . I think this would avoid such problem so I started it in #1725 , but that seems to be triggering quite a lot of issues, so I'm now questioning if I'm taking the right approach.
Contributor guide
No contributing guide indexed for this repository
First steps
- Read the whole issue, then the project's contributing guide.
- Comment on the issue to say you are picking it up — it saves two people doing the same work.
- Fork the repository and make your change on a branch.
- Open a pull request that references the issue number.
Research direction
Start by running the reproducer in the issue with thunder.jit and symbolic caching, then inspect split_forward_backward in thunder/executors/torch_autograd.py and rematerialize_forward_and_backward in thunder/core/rematerialization.py. Trace how saved_for_backward is produced and trimmed for the forward and backward traces; done means the reproducer no longer hits the assertion and the two traces remain consistent.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python, pytorch
- Domain
- compilers
- Issue type
- Bug
- Difficulty
- 4/5
- Estimated time
- 3-5 days
- Activity status
- Stale
- Clarity
- Needs clarification
- Newbie friendliness
- 35/100