[not for land] towards QAT with exact forward pass

[vkuzo]

Summary:

Exploration for QAT with exact (not emulated) forward pass, WIP and not ready for review.

Test Plan:

(pytorch) [[email protected] ~/local/ao (20250701_qat_exact_fwd_pass)]$ python torchao/prototype/qat_exact/main.py 
# SQNR between high precision and emulated quantization
sqnr_hp_qref tensor(23.3884, device='cuda:0')
# SQNR between high precision and real quantization (the new "kernel" in this PR)
sqnr_hp_q tensor(23.3884, device='cuda:0')
# SQNR between emulated quantization and real quantization (the new "kernel" in this PR)
sqnr_qref_q tensor(136.4019, device='cuda:0')

Reviewers:

Subscribers:

Tasks:

Tags:

[pytorch-bot]

🔗 Helpful Links

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

• 📄 Preview Python docs built from this PR

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

❌ 6 New Failures

As of commit ff8cb2e with merge base c2a6568 ():

NEW FAILURES - The following jobs have failed:

• PR Label Check / Check PR Labels (gh)
  Process completed with exit code 1.
• Run Regression Tests / test (CUDA 2.5.1, linux.g5.12xlarge.nvidia.gpu, torch==2.5.1 --index-url https://download.pytorch... / linux-job (gh)
  test/quantization/test_quant_api.py::TestQuantFlow::test_int8_wo_quant_save_load
• Run Regression Tests / test (CUDA 2.6, linux.g5.12xlarge.nvidia.gpu, torch==2.6.0, cuda, 12.6) / linux-job (gh)
  test/quantization/test_quant_api.py::TestQuantFlow::test_int8_wo_quant_save_load
• Run Regression Tests / test (CUDA 2.7, linux.g5.12xlarge.nvidia.gpu, torch==2.7.0, cuda, 12.6) / linux-job (gh)
  test/quantization/test_quant_api.py::TestQuantFlow::test_int8_wo_quant_save_load
• Run Regression Tests / test-nightly (CUDA Nightly, linux.g5.12xlarge.nvidia.gpu, --pre torch --index-url https://downloa... / linux-job (gh)
  test/quantization/test_quant_api.py::TestQuantFlow::test_int8_wo_quant_save_load
• Run TorchAO Experimental Tests / test-cpu-ops (linux.arm64.2xlarge) (gh)
  test_replace_q_dq_patterns_with_quantized_linear_ops_pass

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