print statement causes inplace error

[ptrblck]

🐛 Describe the bug

Reported in: https://discuss.pytorch.org/t/error-with-view-no-grad-and-inplace-modify/173082
but I couldn't find the created GitHub issue and the author didn't follow up.

Code to reproduce the issue:

net = nn.Sequential(
    nn.Linear(10, 10),
    nn.ReLU(),
    nn.Linear(10, 10),
    nn.ReLU(),
    nn.Linear(10, 10),
    nn.ReLU(),
)

with torch.no_grad():
  for param in net.parameters():
    for j in param.flatten():
        #print("current j", j)
        j += 1

Comment the print statement in and the code will fail with:

RuntimeError: A view was created in no_grad mode and its base or another view of its base has been modified inplace with grad mode enabled. Given that this use case is ambiguous and error-prone, it is forbidden. You can clarify your code by moving both the view and the inplace either both inside the no_grad block (if you don't want the inplace to be tracked) or both outside (if you want the inplace to be tracked).

I would assume the inplace operation is allowed as it's in a no_grad block and no computation graph was ever created.

Also, maybe related to: https://discuss.pytorch.org/t/old-problem-but-strange-things-trying-to-backward-through-the-graph-a-second-time/178369

but no executable code snippet was posted yet.

Versions

Reproduced in a nightly build: 2.1.0.dev20230407+cu118.

CC @albanD as we talked about this issue before.

cc @ezyang @gchanan @zou3519 @kadeng @albanD @gqchen @pearu @nikitaved @soulitzer @lezcano @Varal7

[albanD]

I'm not sure from the top of my head.
But most likely that we're creating a view during printing, then the inplace happens which invalidates that view and then for some reason we try to recover it.
We would need to investigate this further.

[soulitzer]

Smaller repro:

a = torch.rand(1, requires_grad=True)

with torch.no_grad():
    b = a[:]
    b += 1

# Doing any of the of the following produces an error
b.sin()   # (1) 
b.grad_fn  # (2)
print(b)  # (3) the reason this fails is because it calls into t.grad_fn for printing purposes

The easy fix for (3), is just to special case on views in no-grad in the printing code, but maybe there is a more general fix.

See also #11390

[hjmshi]

Hi, we are running into a similar issue as we are implementing an updated version of Distributed Shampoo and seeking to apply torch.compile on it from PyTorch 2. In our new version of the code, we create many views of the parameters in the optimizer, and seek to apply in-place foreach operators on the parameter views.

When we print the list of parameter views or apply torch.compile on this list, we observe the same error, even though everything is performed under @torch.no_grad. Prior to applying the in-place add for the first time, if we print the list of views, we do observe requires_grad = True, consistent with #11390.

Any suggestions on how to proceed? Thanks in advance!

Interestingly, we found that logging the tensor does not trigger this issue, but printing does.

cc: @tsunghsienlee @shintaro-iwasaki @minddrummer @csmiler @mlazos @bdhirsh @yuchenhao

[albanD]

Note that here we can easily fix printing (try/except around access to grad_fn and print an invalid grad_fn). The other errors are expected: This is undefined behavior to do this so we rather raise an error.

[albanD]

Actionable to fix print.
High pri for user activity