github/ai-toolkit
1
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

Move uintx quantization to ostris quant with bit identical matching. Now we are not bound to an older version of torch ao.

Browse Source
This commit is contained in:
Jaret Burkett
2026-07-27 12:59:42 -06:00
parent fb204b7677
commit b677cdb026
3 changed files with 221 additions and 10 deletions
Download Patch File Download Diff File Expand all files Collapse all files

11
toolkit/util/quantize.py
View File

@@ -7,7 +7,6 @@ from optimum.quanto.tensor import Optimizer, qtype, qtypes
from torchao.quantization.quant_api import (
quantize_ as torchao_quantize_,
Float8WeightOnlyConfig,
UIntXWeightOnlyConfig,
Int8WeightOnlyConfig
)
from optimum.quanto import freeze
@@ -42,13 +41,9 @@ Q_MODULES = [
torchao_qtypes = {
# "int4": Int4WeightOnlyConfig(),
"uint2": UIntXWeightOnlyConfig(torch.uint2),
"uint3": UIntXWeightOnlyConfig(torch.uint3),
"uint4": UIntXWeightOnlyConfig(torch.uint4),
"uint5": UIntXWeightOnlyConfig(torch.uint5),
"uint6": UIntXWeightOnlyConfig(torch.uint6),
"uint7": UIntXWeightOnlyConfig(torch.uint7),
"uint8": UIntXWeightOnlyConfig(torch.uint8),
# uint2..uint8 are handled by the UIntXQuantizer ostris backend
# (toolkit/util/uintx_quant.py), a bit-exact reproduction of torchao 0.10.0's
# UIntXWeightOnlyConfig, so ARAs stay byte-identical after torchao upgrades
"int8": Int8WeightOnlyConfig(),
"float8": Float8WeightOnlyConfig(),
}