Lightning-AI / Lightning-AI/lightning-thunder

Grad Transform generates inconsistent saved_for_backward between forward and backward trace.

Open
#1,732 5 comments 1 reaction 0 assignees View on GitHub

Nobody has claimed this yet.

autograd
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

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. 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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.