Refactor QAT to use common fake_quantize_affine primitive

[andrewor14]

Summary: Currently there are two QAT quantizers, 8da4w and 4w. Today, these use different autograd functions to represent their fake quantization numerics, but this is not scalable because new QAT quantizers may introduce yet another divergent code path. To address this, this commit refactors both quantizers to use the common fake_quantize_affine QAT primitive.

Test Plan:
python test/quantization/test_qat.py

Reviewers: jerryzh168

Subscribers: jerryzh168, supriyar, msaroufim

[pytorch-bot]

🔗 Helpful Links

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

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

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

[jerryzh168]

if it's helpful we could have a general util like:

def get_block_size(granularity, **kw_params) -> Callable:
    if granularity == Granularity.PER_BLOCK:
        ...
    elif type == Granularity.PER_TOKEN:
        ...
     ...


block_size = get_block_size(Granularity.PER_TOKEN)(x)

[andrewor14]

Sounds good, let's do that separately