Match QAT prepare and convert numerics exactly

[andrewor14]

Summary: Previously, Int8DynActInt4QATQuantizer had slightly diverging numerics between the prepare and convert steps. This is because the prepare step uses quantization primitives shared with AQT (specifically quantize_affine and dequantize_affine), while the convert step relies on old ops from the torch.ops.quantized_decomposed namespace. The diverging numerics is negligible for small models, but the quantization errors begin to compound for larger models with many linear layers.

More specifically, there are three different places where the divergence occurs during activation quantization:

1. Choose qparams. The prepare step casts the qparams to torch.float32, whereas the convert step casts the scales to torch.float64 and zero points to torch.int64.

2. Quantize. The prepare step performs round before adding zero points and uses torch functions, while the convert step adds before rounding and uses torch tensor methods.

# Prepare
x = torch.clamp(
    torch.round(x * (1.0 / scale)) + zero_point, qmin, qmax,
)

# Convert
x = ( 
    x.mul(1.0 / scale)
    .add(zero_point)
    .round()
    .clamp(qmin, qmax)
    .to(quantize_dtype)
)

3. Dequantize. The prepare step casts to torch.int32 before adding the zero points, and casts back to the original dtype before multiplying the scale. The convert step only casts at the very end.

# Prepare
x = x.to(torch.int32) - zero_point.to(torch.int32)
x = x.to(orig_dtype)
x = x * scale

# Convert
x = x - zero_point
x = x * scale 
x = x.to(orig_dtype)

This commit makes the convert path use the same torchao quantization primitives as the prepare path, thereby resolving the 3 above differences. Now, the prepare and convert steps match exactly in terms of numerics over many trials.

Test Plan:

python test/quantization/test_qat.py -k test_fake_quantize_per_token_vs_convert
python test/quantization/test_qat.py -k test_qat_8da4w_prepare_vs_convert

[pytorch-bot]

🔗 Helpful Links

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

• 📄 Preview Python docs built from this PR

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

❌ 2 New Failures

As of commit 890438a with merge base dfbd681 ():

NEW FAILURES - The following jobs have failed:

• .github/workflows/float8nocompile_test.yaml (gh)
• Run TorchAO Experimental Tests / test-mps-ops (macos-m1-stable) (gh)
  test_linear_117

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

[facebook-github-bot]

@andrewor14 has imported this pull request. If you are a Meta employee, you can view this diff on Phabricator.

[jerryzh168]

Thanks for the updates, will be good to setup a deprecation plan for quantized_decomposed.choose_qparams op then

[facebook-github-bot]

@andrewor14 has imported this pull request. If you are a Meta employee, you can view this diff on Phabricator.