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

Do not materalize weights on ostris quantizer when getting state dict. Require dequantization of weight like other quantization methods.

Browse Source
This commit is contained in:
Jaret Burkett
2026-07-16 08:55:51 -06:00
parent 31c45cf37d
commit bbaef7852a
2 changed files with 53 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files

6
toolkit/util/quantize.py
View File

@@ -82,8 +82,10 @@ def get_qtype(qtype: Union[str, qtype]) -> qtype:
def is_quantized_tensor(t) -> bool:
# torchao stores quantized weights as tensor subclasses (e.g. AffineQuantizedTensor) under torchao.*
# that still report as nn.Parameter and expose .dequantize(). (quanto is handled separately.)
# OstrisLinear.weight returns an already-dequantized tensor tagged with _is_ostris_weight
# (its .dequantize() is a no-op) so the merge paths route through requantize_module_weight.
# _is_ostris_weight tags two OstrisLinear tensors: the .weight property's eager tensor
# (already dequantized; .dequantize() is a no-op) so the merge paths route through
# requantize_module_weight, and the lazy OstrisLazyWeight emitted by state_dict()
# (holds no data; .dequantize() materializes) so save loops dequantize it per key.
if getattr(t, '_is_ostris_weight', False):
return True
return 'torchao' in type(t).__module__ and hasattr(t, 'dequantize')