import torch
from transformers import TorchAoConfig, AutoModelForCausalLM, AutoTokenizer
from torchao.quantization import Int8DynamicActivationInt8WeightConfig, Int8WeightOnlyConfig
quant_config = Int8WeightOnlyConfig()
quantization_config = TorchAoConfig(quant_type=quant_config)
# Load and quantize the model
quantized_model = AutoModelForCausalLM.from_pretrained(
"meta-llama/Llama-3.1-8B-Instruct",
dtype="auto",
device_map="cpu",
quantization_config=quantization_config
)
quantized_model.save_pretrained("./llama-int8wo", safe_serialization=False)Traceback (most recent call last):
File "/home/jiqing/HuggingFace/tests/workloads/test_ao_2.py", line 15, in <module>
quantized_model.save_pretrained("./llama-int8wo", safe_serialization=False)
File "/home/jiqing/transformers/src/transformers/modeling_utils.py", line 3107, in save_pretrained
state_dict, metadata = hf_quantizer.get_state_dict_and_metadata(self)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/home/jiqing/transformers/src/transformers/quantizers/quantizer_torchao.py", line 160, in get_state_dict_and_metadata
return flatten_tensor_state_dict(model.state_dict()), {}
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/opt/venv/lib/python3.12/site-packages/torchao/prototype/safetensors/safetensors_support.py", line 195, in flatten_tensor_state
_dict
raise ValueError(f"Unsupported tensor type: {type(tensor)}") ValueError: Unsupported tensor type: <class 'torchao.dtypes.affine_quantized_tensor.AffineQuantizedTensor'>
Should save successfully.
System Info
regression PR #42734
Who can help?
@Cyrilvallez
Information
Tasks
examplesfolder (such as GLUE/SQuAD, ...)Reproduction
output:
Expected behavior
Should save successfully.