The internal torchao tensor subclasses cause errors with torch.compile

[JohnnyRacer]

When I use torch.compile with certain models that cause the following error:

BackendCompilerFailed: backend='inductor' raised:
AssertionError: torch.compile is currently being used with tensor subclass inputs:
<class 'torch._subclasses.fake_tensor.FakeTensor'>,<class 'torch._subclasses.fake_tensor.FakeTensor'>,<class 'torchao.quantization.linear_activation_quantized_tensor.LinearActivationQuantizedTensor'>,<class 'torch._subclasses.fake_tensor.FakeTensor'>,<class 'torchao.quantization.linear_activation_quantized_tensor.LinearActivationQuantizedTensor'>,<class 'torch._subclasses.fake_tensor.FakeTensor'>,<class 'torch._subclasses.fake_tensor.FakeTensor'>,<class 'torch._subclasses.fake_tensor.FakeTensor'>,<class 'torch._subclasses.fake_tensor.FakeTensor'>,<class 'torch._subclasses.fake_tensor.FakeTensor'>,<class 'torch._subclasses.fake_tensor.FakeTensor'>. We are attempting to a compile a graph with two graph outputs
that alias one another, which is currently unsupported in the subclass use case. If you run into this,
please file a github issue

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

LinearActivationQuantizedTensor and FakeTensor which are tensor subclasses are not supported by torch.compile which raises the errors.

The model compiles correct if the errors are suppressed with the following:

import torch._dynamo
torch._dynamo.config.suppress_errors = True

[supriyar]

cc @jerryzh168 @drisspg

[drisspg]

Do you have a repro?
I added some code a while a go to make sure that the compile specific subclasses work with FP8 tensor that probably needs to be reproduced:

ao/torchao/float8/float8_tensor.py

Line 361 in 52b6f4d

def allowed_subclasses(type):

I have often found this is indicative a graph break though somewhere un expected.

[jerryzh168]

please provide a repro, also what is the version of pytorch you are using? can you run pip show torch and paste the result

[JohnnyRacer]

I can't share my repro directly as the codebase for the model that this happens with is large and complicated (no clue on where it is having problems with torchao). I can't seem to reproduce it with small toy models, as it seems to occur only before Triton begins to compile the kernels in the back. The tensor classes seem to have something to do with testing the quantization optimizations being performed on the model with Triton. I will try to narrow down the modules that are problematic and give a repro if possible.
torchao version : 0.7.0
torch version: 2.5.0a0+b465a5843b.nv24.9 # From Docker NGC PyTorch 24.9 image

[jerryzh168]

no problem, could you try upgrade pytorch to most recent version to see if it is fixed? since there could be tensor subclass + compile issues that's fixed recently:

pip install --pre torch --index-url https://download.pytorch.org/whl/nightly/cu124