Skip to content

Commit 217d968

Browse files
authored
Make FakeQuantizer expose useful config details (#1717)
**Summary:** Expose useful config details when printing FakeQuantizer, which appears when printing QAT prepared models containing linear layers. Before: ``` >>> print(prepared_model.layers[0].attn.qproj) FakeQuantizedLinear( in_features=4096, out_features=4096, bias=False (activation_fake_quantizer): FakeQuantizer() (weight_fake_quantizer): FakeQuantizer() ) ``` After: ``` >>> print(prepared_model.layers[0].attn.qproj) FakeQuantizedLinear( in_features=4096, out_features=4096, bias=False (activation_fake_quantizer): FakeQuantizer(FakeQuantizeConfig(dtype=torch.int8, granularity=PerToken(), mapping_type=<MappingType.ASYMMETRIC: 3>, scale_precision=torch.float32, zero_point_precision=torch.int32, zero_point_domain=<ZeroPointDomain.INT: 1>, is_dynamic=True, range_learning=False)) (weight_fake_quantizer): FakeQuantizer(FakeQuantizeConfig(dtype=torch.int4, granularity=PerGroup(group_size=32), mapping_type=<MappingType.SYMMETRIC: 1>, scale_precision=torch.float32, zero_point_precision=torch.int32, zero_point_domain=<ZeroPointDomain.INT: 1>, is_dynamic=True, range_learning=False)) ) ``` **Test Plan:** python test/quantization/test_qat.py -k test_fake_quantizer_repr
1 parent ceceea5 commit 217d968

File tree

2 files changed

+24
-0
lines changed

2 files changed

+24
-0
lines changed

test/quantization/test_qat.py

+18
Original file line numberDiff line numberDiff line change
@@ -31,6 +31,9 @@
3131
from torchao.quantization.qat.embedding import (
3232
FakeQuantizedEmbedding,
3333
)
34+
from torchao.quantization.qat.fake_quantizer import (
35+
FakeQuantizer,
36+
)
3437
from torchao.quantization.qat.linear import (
3538
FakeQuantizedLinear,
3639
Int4WeightOnlyQATLinear,
@@ -1348,6 +1351,21 @@ def test_fake_quantize_config_torch_intx(self):
13481351
out2 = linear2(*x2)
13491352
torch.testing.assert_close(out1, out2, atol=0, rtol=0)
13501353

1354+
@unittest.skipIf(
1355+
not TORCH_VERSION_AT_LEAST_2_6, "skipping when torch version is 2.6 or lower"
1356+
)
1357+
def test_fake_quantizer_repr(self):
1358+
"""
1359+
Test that `repr(FakeQuantizer(config))` exposes useful config details.
1360+
"""
1361+
config = FakeQuantizeConfig(torch.int4, group_size=128)
1362+
fake_quantizer = FakeQuantizer(config)
1363+
fake_quantizer_repr = repr(fake_quantizer)
1364+
self.assertTrue("dtype=torch.int4" in fake_quantizer_repr)
1365+
self.assertTrue("group_size=128" in fake_quantizer_repr)
1366+
self.assertTrue("PerGroup" in fake_quantizer_repr)
1367+
self.assertTrue("MappingType.SYMMETRIC" in fake_quantizer_repr)
1368+
13511369

13521370
if __name__ == "__main__":
13531371
unittest.main()

torchao/quantization/qat/fake_quantizer.py

+6
Original file line numberDiff line numberDiff line change
@@ -134,3 +134,9 @@ def _should_compute_qparams(self) -> bool:
134134
Return whether we need to compute new scales and zero points.
135135
"""
136136
return self.config.is_dynamic or self.scale is None or self.zero_point is None
137+
138+
def __repr__(self) -> str:
139+
"""
140+
Return a human readable representation of this `FakeQuantizer` with config details.
141+
"""
142+
return "FakeQuantizer(%s)" % self.config

0 commit comments

Comments
 (0)