torch.compile incorrect when imperative autograd APIs are used

[zou3519]

Issue description

torch.compile may be silently incorrect when Tensor.retain_grad or Tensor.register_hook are involved.

Code example

Example with retain_grad:

 # Test 1:
 # Tensor.retain_grad inside function
 # on intermediate, and on output
 # See if the grad can be used :P

 def f(x):
     y = x.clone()
     y.retain_grad()
     z = y.clone()
     z.retain_grad()
     return z, y

 of = torch.compile(backend='aot_eager')(f)

 x = torch.randn([], requires_grad=True)
 z, y = of(x)
 z.clone().backward()
 # inspect z.grad, y.grad
 # Bug: they do not exist. May lead to silent correctness problems

Example with register_hook:

# If you register a hook on an intermediate, it won't work.
def f(x):
    y = x * x
    z = y * x * x * x * 1 * 1 * 1. * 1. * 1.
    print("graph_break")
    y.register_hook(lambda x: x if x is None else 3.14 * x)
    result, = torch.autograd.grad(z, x)
    return result

of = torch.compile(backend='aot_eager')(f)
x = torch.tensor(1., requires_grad=True)

expected = f(x)
result = of(x)
# assertion failed
assert torch.allclose(result, expected)

cc @ezyang @gchanan @kadeng @albanD @gqchen @pearu @nikitaved @soulitzer @lezcano @Varal7 @msaroufim @bdhirsh @anijain2305 @chauhang @wconstab @soumith @ngimel

[ngimel]

cc @Chillee, high priority for silently incorrect results.

[albanD]

Hook handling is a general issue in Dynamo yes...

[ngimel]

Related #91665

[voznesenskym]

Still a bug

[zou3519]

graph-breaking isn't sufficient; the problem is that the autograd.Function we generate for AOTAutograd changes autograd semantics. This isn't an easy fix.

[bdhirsh]

Talking with @zou3519:

Not sure how we can fix this at all in compile, but we might want to consider erroring on this behavior if we don't allow it.

One option (using the register_hook repro from the first post): AOTAutograd create a graph that has y and z both as graph outputs. One option is that AOTAutograd can detect that y is a "special" output, because the other output z was computed off of y. We can then set a bit on y that says you are not allowed to register any hooks to y, since they will not fire properly.

This seems difficult to do though.