autoquant using aqt

[HDCharles]

Summary:

changing autoquant to use aqt instead of the old subclass subtensors changed aqt to first dispatch to a static _quantized_linear_op which then dispatches to the normal function. This way autoquant has an extention point to modify the kernel functions for various quantization modes without editing the main kernel function of all the classes. linear_activation_quantized_tensor got the same treatment.

there were some transposes found in the aqt kernels not present in the subclass kernels, however they do not seen to affect performance (see benchmark_results.txt for an autoquant perf run)

Test Plan:

sh benchmarks.sh

python test_integration.py

Reviewers:

Subscribers:

Tasks:

Tags:

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/ao/609

• 📄 Preview Python docs built from this PR

Note: Links to docs will display an error until the docs builds have been completed.

✅ No Failures

As of commit f5ac4bf with merge base 87869f2 ():
💚 Looks good so far! There are no failures yet. 💚

This comment was automatically generated by Dr. CI and updates every 15 minutes.

[jerryzh168]

can we reuse the code from quant_api.py? also we need to refactor the input_quant_func to be a normal function (not lambda) in order for serialization to work I think, might help to do that refactor at the same time

[HDCharles]

the issue is that we need to call from_float with super in order to get this to work correctly. The code in quant_api would generate an aqt, not an autoquant class inheriting from an aqt. (I tried this approach initially since that's how it worked with subclass)

If we want to reuse the code, it may make sense to make a function that like prepares all the variables needed to go into from_float

def int8_weight_only_kwargs(weight):
       ...do the code up to to_affine_quantized and put it all together
       return a_bunch_of_kwargs
       
then you could have the quant_api code be like

def int8_weight_only():
    def apply_int8wo_quant(weight):
        kwargs = int8_weight_only_kwargs(weight)   
        return to_affine_quantized(**kwargs)
    return _get_linear_subclass_inserter(apply_int8wo_quant)

then in autoquant we could do similar

def from_float(weight):
      kwargs = int8_weight_only_kwargs(weight)
      super().from_float(**kwargs)
      

[jerryzh168]

does move apply_int8wo_quant function to top level help? I'm doing it here: https://github.com/pytorch/ao/pull/630/files, seems like you call the super().from_float after the aqt is produced right

[jerryzh168]

we can move the apply quant to weight Tensor function to top level as well I think

[jerryzh168]

thanks, the changes looks good to me, please fix CI before landing

[jerryzh168]

what are these? should these be enabled or removed?

[HDCharles]

added a todo