Weird interaction with p.grad in dynamo

[janeyx99]

🐛 Describe the bug

I encountered this while trying to enable dynamo on unit tests. Not mentioning requires_grad and then setting the grad of a parameter within a compiled region when the parameter lives outside of that region causes internal dynamo to think that the grad is None.

In my repro below, to_sparse() could be replaced with any function (like .add will produce the same error).

Error logs

InternalTorchDynamoError: 'NoneType' object has no attribute 'to_sparse'

from user code:
   File "<ipython-input-4-48c0ef0ff017>", line 6, in f
    p.grad = p.grad.to_sparse()

Set TORCH_LOGS="+dynamo" and TORCHDYNAMO_VERBOSE=1 for more information


You can suppress this exception and fall back to eager by setting:
    import torch._dynamo
    torch._dynamo.config.suppress_errors = True

Minified repro

import torch
param = torch.rand(2, 3, dtype=torch.float32, device='cuda')  # moving this inside f would work! 

def f(p):
    p.grad = torch.rand_like(p)
    p.grad = p.grad.to_sparse()  # why does dynamo think p.grad at this moment is empty?

compiled_f = torch._dynamo.optimize("eager")(f)
compiled_f(param)

Versions

on main

cc @ezyang @msaroufim @wconstab @bdhirsh @anijain2305 @zou3519 @voznesenskym @penguinwu @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @wenzhe-nrv @jiayisunx @chenyang78 @aakhundov @kadeng

[voznesenskym]

I think I saw this recently, and maybe have a fix, let me go look.