Failed to save the static quantized model

[yiliu30]

Hi, I am following this example and want to save the INT8 static quantization result, but it’s failing.
Could you take a look, thanks!

    ... 
    # quantized linear represented as an nn.Linear with modified tensor subclass weights
    # for both activation and weight quantization
    quantize_(m, ApplyStaticQuantConfig(target_dtype), is_observed_linear)
    print("quantized model (applying tensor subclass to weight):", m)
    after_quant = m(*example_inputs)  # <---- Line 310
    torch.save(m.state_dict(), f"qmodel_m.pt")  # <---- My change
    assert compute_error(before_quant, after_quant) > 25
    print("test passed")

• Log

Testing torch.uint8 static quantization:
example inputs shape: torch.Size([1, 64])
quantized model (applying tensor subclass to weight): ToyLinearModel(
  (linear1): Linear(in_features=64, out_features=64, bias=False)
  (linear2): Linear(in_features=64, out_features=32, bias=False)
)
Traceback (most recent call last):
  File "/home/user/workspace/torchao/tutorials/calibration_flow/static_quant.py", line 325, in <module>
    test_static_quant(torch.uint8, MappingType.ASYMMETRIC)
  File "/home/user/workspace/torchao/tutorials/calibration_flow/static_quant.py", line 311, in test_static_quant
    torch.save(m.state_dict(), f"qmodel_m.pt")
  File "/home/user/miniforge3/envs/ao/lib/python3.11/site-packages/torch/serialization.py", line 965, in save
    _save(
  File "/home/user/miniforge3/envs/ao/lib/python3.11/site-packages/torch/serialization.py", line 1211, in _save
    pickler.dump(obj)
AttributeError: Can't pickle local object '_apply_static_quant_transform.<locals>.<lambda>'

• Env info
• torchao 0.10.0+gita99598d8
• torch 2.8.0.dev20250324+cu126

cc @jerryzh168

[supriyar]

@jerryzh168 can you help with this?

[jerryzh168]

yeah I'll take a look, sorry I didn't see this one

[jerryzh168]

Hi @yiliu30 sorry for the late reply, I took a look, and I think the reason is because there are a few local functions defined in _apply_static_quant_transform, and pickle would not work for that, we have an example for float8 static here:

ao/torchao/quantization/quant_api.py

Line 1612 in 34421b1

def _float8_static_activation_float8_weight_transform(

that should be serializable, the example didn't include this for simplicity, please feel free to follow that to support serialization for your use case

specifically:

ao/tutorials/calibration_flow/static_quant.py

Line 108 in 34421b1

def weight_quant_func(weight):

,

ao/tutorials/calibration_flow/static_quant.py

Line 143 in 34421b1

input_quant_func = lambda x: to_affine_quantized_intx_static(

and

ao/tutorials/calibration_flow/static_quant.py

Line 147 in 34421b1

input_quant_func = lambda x: to_affine_quantized_floatx_static(

last two lines might be a bit harder, you can follow float8 example to do it