Skip to content

feat: replace soon-deprecated torch quantization functions - #211

Merged
andrea-fasoli merged 2 commits into
mainfrom
afasoli/replace_quantize
Sep 8, 2026
Merged

feat: replace soon-deprecated torch quantization functions#211
andrea-fasoli merged 2 commits into
mainfrom
afasoli/replace_quantize

Conversation

@andrea-fasoli

Copy link
Copy Markdown
Collaborator

Description of the change

Replace deprecated torch.quantize_per_tensor/torch.quantize_per_channel calls with equivalent plain arithmetic to unblock the torch 2.13 upgrade (#209).

PyTorch 2.13 emits a UserWarning from these ops (pytorch/pytorch#184982), which CI promotes to an error via -W error::UserWarning in tox.ini, failing test_save_aiu.py::test_large_outlier_bert and test_lsq.py::test_lsq_single_sided.

Converted 13 call sites across quant/quantizers.py, quant_refactor/ (torch_quantizer.py, per_tensor_ste.py, per_channel_ste.py, pactplussym_rc.py, quantizers_new.py), modules/bmm.py, modules/linear.py, and custom_ext_kernels/utils.py.

Also added shared tables to quant_refactor/base_quant.py:

  • _INT_REPR_DTYPES: maps each quantized dtype to the storage dtype int_repr() used to return (qint8→int8, quint8→uint8, qint32→int32), so the arithmetic replacements cast to the right integer type instead of leaving a float.
  • _DTYPE_RANGES: gives each dtype's storage bounds so the replacements can reproduce the ops' saturation before the caller's narrower quant_min/quant_max clamp (which matters for int4-in-int8), and supplies the explicit bounds needed to clamp qint32 in float64, since 2**31-1 isn't representable in fp32 and wraps on cast.

The replacements mirror the deprecated kernels bit-for-bit — reciprocal multiply in fp32 (not division, which is marginally more accurate but drifts 1 LSB on exact rounding ties), round half-to-even, and saturation to the storage dtype.

Note: one site (modules/linear.py:439) is unchanged, it still calls torch.quantize_per_tensor. It's the only site that needs the quantized tensor object itself (passed to set_weight_bias for FBGEMM), not int_repr(). Plain arithmetic can't produce that, and converting it means rewriting the FBGEMM interop. It's in an untested from_reference classmethod and doesn't trigger under -W error today, but it will break when torch actually removes these ops.

Related issues or PRs

Closes #210 , enables #209

Was the PR tested

  • I have ensured all unit tests pass

Signed-off-by: Andrea Fasoli <andrea.fasoli@ibm.com>
@andrea-fasoli andrea-fasoli changed the title replace soon-deprecated torch quantization functions feat: replace soon-deprecated torch quantization functions Sep 8, 2026
@github-actions github-actions Bot added the feat label Sep 8, 2026
Signed-off-by: Andrea Fasoli <andrea.fasoli@ibm.com>
@chichun-charlie-liu

Copy link
Copy Markdown
Collaborator

Here's the feedback by claude. Looks ready for merge once CI is done. Please also double check Minor Findings 2.

Review Summary

Replaces 13 call sites of torch.quantize_per_tensor/torch.quantize_per_channel (deprecated in pytorch/pytorch#184982) with equivalent plain arithmetic: round(x * scale.reciprocal()) + zero_point, then saturate to the storage dtype's range, then cast. Adds _INT_REPR_DTYPES/_DTYPE_RANGES lookup tables to reproduce the deprecated ops' storage-dtype and saturation-range behavior exactly.

Correctness — verified by hand-tracing the arithmetic:

  • The core substitution (reciprocal-multiply in fp32, round-half-to-even via torch.round, saturate, cast) faithfully reproduces the deprecated kernels internally, including multiplying by the reciprocal rather than dividing — called out explicitly in the torch_quantizer.py comment, and it matters because this is meant to match bit-for-bit, not just approximately.
  • Per-channel broadcast reshaping (bcast[axis] = -1, the new _broadcast_qparams helper) correctly handles the axis parameter across per_channel_ste.py, quantizers_new.py, and torch_quantizer.py.
  • The float64 clamp step before casting to qint32/int32 is a real correctness detail (fp32 can't exactly represent 2**31-1) that's handled correctly and commented.
  • Skipping linear.py:439 (FBGEMM interop) is reasonable since that path needs an actual quantized tensor object, not just its int representation.
  • Confirmed torch.int is torch.int32, so no key mismatch between zero_point's dtype and the lookup tables.

Two minor findings:

  1. Duplication: _INT_REPR_DTYPES is defined identically in both quant/quantizers.py and quant_refactor/base_quant.py instead of one importing from the other. Low risk today, but a future dtype-mapping change could update one copy and miss the other.

  2. Possible precision asymmetry (worth confirming): In torch_quantizer.py's per-tensor branch, tensor * scale.reciprocal() has no explicit .float() upcast on tensor, while the equivalent code in linear.py/bmm.py/custom_ext_kernels/utils.py does cast to float first. Since TorchQuantizer is explicitly the reference class the tests compare against, any precision drift here would undercut the "bit-for-bit" goal. Could you confirm tensor is always fp32 by the time it reaches this path, or would it be safer to add the explicit cast for parity with the other call sites?

Neither issue blocks merge — overall this is a solid, consistent conversion across all 13 sites.

@chichun-charlie-liu chichun-charlie-liu left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

lgtm

@andrea-fasoli
andrea-fasoli merged commit cd8a062 into main Sep 8, 2026
14 checks passed
@andrea-fasoli
andrea-fasoli deleted the afasoli/replace_quantize branch September 8, 2026 22:37
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

Projects

None yet

Development

Successfully merging this pull request may close these issues.

torch.quantize_per_channel triggers deprecation warning in torch 2.13

2 participants