[DO NOT REVIEW][SMOKE TEST] Skip prims shape in grad transform #1704
Conversation
for more information, see https://pre-commit.ci
…backward_transform_dependency_fix
…backward_transform_dependency_fix
Co-authored-by: Masaki Kozuki <[email protected]>
for more information, see https://pre-commit.ci
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
| @@ -123,7 +123,7 @@ def keep_or_swap(p): | |||
| if not isinstance(p, NumberProxyInterface): | |||
| return p | |||
| if p.name in seen: | |||
| return p.value # don't make it a duplicate | |||
| return None # don't make it a duplicate | |||
There was a problem hiding this comment.
note for myself, I needed this to avoid an error in rematerialization pass. I'll work on a separate PR with a repro when I get that.
There was a problem hiding this comment.
errrrr... this one doesn't seem to be working any more... I'm seeing the assert in rematerialization.py again.
There was a problem hiding this comment.
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_()
jfoo = thunder.jit(foo, cache="symbolic values")
out = jfoo(a, b)
took me a while to get a repro here.
I think the issue is coming from saving for backward not properly identifying which gradient path isn't required. So after dce kicks in, the saved_for_backward is inconsistent... now I feel that I just missed a dce somewhere.
There was a problem hiding this comment.
trying my luck with #1725
| else: | ||
| assert isinstance(new, ProxyInterface), (old, new) | ||
| swap_map[variableify(new)] = old | ||
| if variableify(old) != variableify(new): | ||
| swap_map[variableify(new.primal)] = old |
There was a problem hiding this comment.
note for myself, this is a separate thing. break this into a separate PR.
No description provided.