Refactor QAT to use tensor subclasses

[andrewor14]

Overview

This commit refactors QAT to use tensor subclasses. This is motivated by the general move towards tensor subclasses in torchao for better composability with other subclasses like DTensors. To achieve this, we introduce AffineFakeQuantizedTensor, which is analogous to AffineQuantizedTensor but applies fake quantization instead and requires gradient updates.

How training works

AffineFakeQuantizedTensor wraps the original weight or input activation tensor and applies fake quantize dynamically only when the linear function is called. Gradients only flow to the outer tensor (AffineFakeQuantizedTensor) and never to the inner tensor. For weights, the outer tensor is also a torch.nn.Parameter, and gradient updates received by the outer tensor are then passed to the inner tensor through ops like aten.add_ and aten.mul_.

Input activation

An important difference between the PTQ and the QAT flows is how input activation subclasses are inserted. For QAT, we use the nn.module forward_pre_hook instead of relying on another subclass LinearActivationQuantizedTensor that wraps the weight subclass. The problem with the old PTQ approach is it can create subclasses under __torch_dispatch__, which runs below autograd and so the created subclasses cannot have gradients, so it was difficult to get the gradients to flow correctly in such cases. It's also not super intuitive because quantizing input activation needs to go through the weights. In the new approach used by QAT, we instead register a forward_pre_hook that wraps the input activations before each call to forward. This approach is also motivated by how DTensor wraps their subclasses .

Work items

• Add AffineFakeQuantizedTensor
• Add support for int4 weight only fake quantize
• Add support for int8 dynamic activations + int4 weight fake quantize (8da4w)
• Add prepare and convert path to int4 QAT quantizer
• Add prepare and convert path to 8da4w QAT quantizer
• Support enabling and disabling fake quant dynamically
• Support __repr__ in AffineFakeQuantizedTensor
• Fix backward pass for int4 weight only
• Fix backward pass for int8 dynamic activations + int4 weight

[pytorch-bot]

🔗 Helpful Links

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

• 📄 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 a4148c3 with merge base b523f9f ():
💚 Looks good so far! There are no failures yet. 💚

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

[bdhirsh]

left a few questions but I'll defer to @jerryzh168 for the review

[jerryzh168]

looks good, thanks!

[andrewor14]

For reference, just added some end-to-end benchmarks and evaluation results for tensor subclass QAT here: pytorch/torchtune#1330. Looks pretty good so far