functorch doesn't work in debug mode

[zou3519]

It's that autograd assert that we run into often:

import torch
from functorch import make_fx
from functorch.compile import nnc_jit


def f(x, y):
    return torch.broadcast_tensors(x, y)


inp1 = torch.rand(())
inp2 = torch.rand(3)

print(f(inp1, inp2))  # without nnc compile everything works fine

print(make_fx(f)(inp1, inp2))  # fails
print(nnc_jit(f)(inp1, inp2))
# RuntimeError: self__storage_saved.value().is_alias_of(result.storage())INTERNAL ASSERT FAILED at "autograd/generated/VariableType_3.cpp":3899, please report a bug to PyTorch.

cc @albanD @soulitzer what's the chance we can add an option to turn these off? They've been more harmful (e.g. prevent debugging in debug mode) than useful for us.

[albanD]

Why do they fail though?
Things that should be views should also be views in the context of functorch.

[zou3519]

@albanD It fails in AOTAutograd and AOTAutograd uses __torch_dispatch__ and there is no existing easy way for __torch_dispatch__ to set the alias relationship. Because that doesn't exist yet, most people who use __torch_dispatch__ with wrapper tensors in debug mode will run into this assert, right?

[albanD]

I don't think we run any debug build anywhere in CI? EDIT: after checking, there is actually a CI debug build that runs all the testss.
I also stopped building debug locally as 32GB of ram are not enough to make one unless you set a super low MAX_JOBS count.
So I don't know tbh.

[gmagogsfm]

Hit this issue when using functorch, I actually regularly develop in DEBUG mode because I need to step through compiler stuff often, would be great if we could fix it.

[soulitzer]

Maybe we could temporarily update self.has_storage() check to self.has_storage() && self.storage().data() != nullptr if we can assume that is the case for all tensors created by make_wrapper_subclass. Then once we figure out a way to set the alias relationship revert this change.

[zou3519]

@soulitzer another way to go about it is to check self for the Python dispatch key -- if it has the Python dispatch key then we temporarily don't do the check. Is there an issue on the pytorch side for fixing the alias relationship for Tensor Subclasses?

[soulitzer]

Is there an issue on the pytorch side for fixing the alias relationship for Tensor Subclasses?

This one seems related pytorch/pytorch#65339

[gmagogsfm]

@zhxchen17

[albanD]

Quick point following some offline discussion with Jeffrey:

• We don't want to disable these tests in general, they have a good reason for existing.
• We need to improve our view story for subclass as discussed in the issue linked above. We have quite a few possible ways that need to be explored. We should dive into them for sure.
• If there are other DEBUG tests that you want to run for functorch, then I think it is ok to have a way for you to disable these. But they should run by default on DEBUG builds unless specified otherwise.

[Chillee]

Ed's PR resolves this for AOTAutograd, but not for vmap/grad more generally from my understanding.

[zou3519]

vmap/grad don't hit the asserts I think. But BatchedTensor and GradTensor do not have storage with leads to some other fun things...