Refactor tensor subclass API to also use paramterization

[jerryzh168]

Summary:
Also added tests for tensor subclass api + AOTI compilation

Test Plan:
python test/integration/test_integration.py -k test_aoti

Two issues right now:

1. AOTI test: change_linear_weights_to_int8_dqtensors + cuda device doesn't work
2. another AOTI test also failed with cuda device (only in CI), will create another PR to repro later

Reviewers:

Subscribers:

Tasks:

Tags:

[cpuhrsch]

Why is this needed?

[jerryzh168]

there is some issue with this in AOT Inductor I think. cc @desertfire

[desertfire]

I think @cpuhrsch 's question is why rewriting "+=". I am not aware any AOTI restriction that needs this rewrite.

[jerryzh168]

@desertfire the error I'm getting with "+=" is this: https://gist.github.com/jerryzh168/d4ea2fb8138376cff903c38aaef8f5ef, is this expected?

just menat to review the yaml change

[cpuhrsch]

Will this also show what dtype was expected? It seems like an opportunity for a dtype guard decorator or somesuch

def guard_dtype_size(tensor_arg, arg_name, dtype=None, size=None):
    if dtype is not None and tensor_arg.dtype != dtype:
        raise ValueError("Expected Tensor argument {arg_name} to have dtype {dtype}, but got {tensor_arg.dtype} instead.")
    if size is not None and tensor_arg.size() != size:
        raise ValueError("Expected Tensor argument {arg_name} to have dtype {dtype}, but got {tensor_arg.dtype} instead.")

guard_dtype_size(scales, "scales", torch.bfloat16, zeros.size())
guard_dtype_size(zeros, "zeros", torch.bfloat16)

See ValueError reference manual for why I chose ValueError here.

[cpuhrsch]

Can you use cls.__tensor_flatten__(*args) for this?

[jerryzh168]

you mean tensor unflatten? we can't use cls in forward because of pytorch/pytorch#124735 right now

[cpuhrsch]

If you wrap do

def create_parameterization_module(cls):
    class SubclassParameterization:
        [...]

    def forward(self, args):
        cls.[...](args)

    return SubclassParameterization

then cls is given as an argument to create_parameterization_module and you return an instance of SubclassParameterization where cls is that argument. Essentially a module factory function.

These methods also shouldn't be static.

[jerryzh168]

isn't this using cls in forward? I tried this before, and with @torch._dynamo.allow_in_graph for the constructor function and it fails because we can't use class variable in dynamo right now I think.

are you suggesting something like this: 25abb31#diff-bf4d50867e3d649de2d89146592bf47d2f258c4c19126c8acf0e120ee904b726R134 (but using cls instead of hardcoding the class?

[cpuhrsch]

Yes exactly and using __tensor_unflatten__ instead of from_qtensor_components

[jerryzh168]

for reference using cls in forward is not supported until pytorch/pytorch#123350 is landed, according to Brain

[cpuhrsch]

Can you use return self.__tensor_flatten__ for this?

[jerryzh168]

this works, thanks. I'll create a parent class to host init and right_inverse

[jcaip]

noob question - why do we want to enable this parameterization support?

[jerryzh168]

this is for supporting exporting the tensor subclass model, needed by aot_compile and also torch.export.export

[cpuhrsch]

Tensor subclasses don't work with AOTI

[cpuhrsch]

Can this made generic for all tensor subclasses?

[jerryzh168]

can't do it now because of pytorch/pytorch#124735, should be able to do it after this is fixed

[jerryzh168]

btw @cpuhrsch we need to use checks + skip test here I think, otherwise this test would fail:
FAIL: test_aoti (main.TestAOTI)

[albanD]

I would expect we use parametrization only for AOTI? As some kind of "pre-processing" there.
Especially given that my understanding of the long term plan is that AOTI will do this pre-processing themselves and we wll be able to remove it from there.

[jerryzh168]

we also need this for torch.export (used by executorch), I'll add a test in next PR, also we want to have a consistent code path for all backends/runtimes I think. is there any problems with enabling this for all use cases?

[bdhirsh]

@albanD do you think that long term we want export to do the pre-processing?

I think if that's the case, then we might just want to figure out that story now (it might be less work than getting dynamo to handle parametrizations).

The main contentious bit is probably just where this pre-processing should live. One possible answer is that it should happen transparently as part of torch.export.export(): automatically search the created state dict for subclasses and flatten them (although this might be a problem if the user expects the state dict of the ExportedProgram to alias the original model's state dict)