fix(Int8DynActInt4WeightQuantizer): use scales_precision for runtime modules

[Anai-Guo]

Summary

Int8DynActInt4WeightQuantizer._convert_for_runtime passed self.precision as the scales_precision argument of replace_linear_8da4w, so the runtime Int8DynActInt4WeightLinear modules were built with the wrong scales dtype. The code even carried a # TODO: this should be self.scales_precision? next to the line.

This diverges from _create_quantized_state_dict, which uses self.scales_precision for the scales. As reported in #2571, that mismatch makes quantize(model).state_dict() inconsistent with _create_quantized_state_dict(model) whenever precision != scales_precision.

Fix

Pass self.scales_precision (instead of self.precision) as the scales_precision argument, matching _create_quantized_state_dict.

 replace_linear_8da4w(
     model,
     self.groupsize,
     self.padding_allowed,
     self.precision,
-    # TODO: this should be self.scales_precision?
-    self.precision,
+    self.scales_precision,
 )

Fixes #2571

Test plan

• Instantiate Int8DynActInt4WeightQuantizer(precision=torch.float32, scales_precision=torch.bfloat16), quantize a small linear model, and confirm the runtime module scales dtype now matches _create_quantized_state_dict output.

🤖 Generated with Claude Code

[pytorch-bot]

🔗 Helpful Links

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

• 📄 Preview Python docs built from this PR

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

❗ 1 Active SEVs

There are 1 currently active SEVs. If your PR is affected, please view them below:

• [ROCm] MI350 CI jobs will have longer queue times due to CI migration

❌ 2 New Failures, 2 Unclassified Failures

As of commit b341436 with merge base 28e6aca ():

NEW FAILURES - The following jobs have failed:

• Run Regression Tests / test-nightly (CPU Nightly, linux.4xlarge, --pre torch torchvision --index-url https://download.py... / linux-job (gh)
  test/quantization/pt2e/test_x86inductor_fusion.py::DynamicShapesCppWrapperCpuTests::test_qconv2d_maxpool2d_linear_dynamic_cpu_dynamic_shapes_cpp_wrapper
• Run Regression Tests / test-nightly (CUDA Nightly, linux.g5.12xlarge.nvidia.gpu, --pre torch torchvision --index-url htt... / linux-job (gh)
  test/quantization/pt2e/test_x86inductor_fusion.py::DynamicShapesCppWrapperCpuTests::test_qconv2d_maxpool2d_linear_dynamic_cpu_dynamic_shapes_cpp_wrapper

UNCLASSIFIED FAILURES - DrCI could not classify the following jobs because the workflow did not run on the merge base. The failures may be pre-existing on trunk or introduced by this PR:

• Run 1xH100 Tests / test (H100, mt-l-x86iamx-22-225-h100, --pre torch torchvision torchaudio mslk --index-url https:/... / linux-job (gh) (this job did not run on the merge base, so DrCI cannot tell whether the failure is pre-existing)
  RuntimeError: Command bash /__w/_temp/exec_script failed with exit code 1
• Run Regression Tests / test-nightly (CPU Nightly (Clang), linux.4xlarge, --pre torch torchvision --index-url https://dow... / linux-job (gh) (this job did not run on the merge base, so DrCI cannot tell whether the failure is pre-existing)
  test/quantization/pt2e/test_x86inductor_fusion.py::DynamicShapesCppWrapperCpuTests::test_qconv2d_maxpool2d_linear_dynamic_cpu_dynamic_shapes_cpp_wrapper

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

[andrewor14]

@claude review

[claude]

Claude finished @andrewor14's task in 56s —— View job

---

Review: fix(Int8DynActInt4WeightQuantizer): use scales_precision for runtime modules

I've reviewed the change against origin/main. The fix is correct and well-targeted. ✅

The bug

_convert_for_runtime builds the runtime Int8DynActInt4WeightLinear modules via replace_linear_8da4w, whose 5th positional argument is scales_precision (linear_quant_modules.py:521-522). That value flows into the scales/zeros buffer dtypes (linear_quant_modules.py:443-456). Previously the quantizer passed self.precision there, so the runtime buffers were allocated with the activation precision instead of the scales precision.

This matters because of how quantize() works:

state_dict = self._create_quantized_state_dict(model)   # scales in self.scales_precision
model = self._convert_for_runtime(model)                # buffers were in self.precision (bug)
model.load_state_dict(state_dict, strict=False)          # copy_ casts to the buffer dtype

load_state_dict copies into the existing buffers, so the destination dtype wins. With precision=fp32, scales_precision=bf16, the bf16 scales produced by _create_quantized_state_dict were silently up-cast to fp32 on load — exactly the inconsistency reported in #2571. The one-line change makes _convert_for_runtime consistent with _create_quantized_state_dict, and the stale # TODO: this should be self.scales_precision? is correctly removed.

Test

test_8da4w_quantizer_scales_precision is a good regression test: it picks precision != scales_precision (the only configuration that exposes the bug), asserts both scales.dtype and zeros.dtype are bf16, and runs a forward pass. It would fail on the pre-fix code. Nice coverage.

Minor (non-blocking) suggestions

• Tighten the assertion to also cover precision. To fully document the contract, you could also assert the non-scales paths still honor precision — e.g. mod.precision == torch.float32 and mod.bias.dtype == torch.float32 (the bias buffer is allocated with precision at linear_quant_modules.py:459). This guards against a future "fix" that over-corrects by routing everything through scales_precision.
• Imports inside the test. The local from ... import statements are fine and match the lazy-import style elsewhere in this file, so no change needed — just noting Int8DynActInt4WeightQuantizer and Int8DynActInt4WeightLinear may already be importable at module top depending on the test file's existing imports.

Neither is required to land. The core fix is minimal, correct, and matches the existing state-dict path.
• branch fix/int8dynact-int4-scales-precision

[Anai-Guo]

Thanks @claude for the review. I took the non-blocking suggestion and tightened the regression test to document the full dtype contract: the test now uses ToyLinearModel(bias=True) and additionally asserts mod.precision == torch.float32 and mod.bias.dtype == torch.float32 alongside the existing bfloat16 scales/zeros checks. This makes it fail not only on the original bug but also on a future change that over-corrects by routing precision through scales_precision.

🤖 Generated with Claude Code

[andrewor14]

Thanks, will merge this for now but I'm actually surprised this code still exists. We meant to deprecate/delete it, will do that separately.

[andrewor14]

Will wait for Ci to pass first

[Anai-Guo]

@andrewor14 The remaining red checks here look like pre-existing nightly/infra breakage rather than anything from this change:

• All the release test jobs pass: test (CPU 2.12), test (CUDA 2.12), test (SM-89), plus the CPU/MPS ops jobs.
• Only the test-nightly (CPU Nightly / CPU Nightly (Clang) / CUDA Nightly) jobs and the ~2-min test (H100, ...) setup fail.
• The exact same test-nightly jobs are failing identically on unrelated docs-only PRs (e.g. docs(finetuning): replace TorchTune QAT section with Unsloth #4312, docs(quant_api): fix RST warnings in FqnToConfig and Int8DynamicActivationIntxWeightConfig docstrings (#3863) #4475), which can't possibly touch CPU/CUDA test behavior — so the failures are environmental (nightly --pre torch / docker-hash), not introduced here.

So the dtype regression test added in this PR is green on the release runners. Happy to rebase if a fresh nightly helps, but it should be safe to merge once you're comfortable.

🤖 Generated with Claude Code