Optimum-quanto —— PyTorch 多功能量化工具包¶
你可以通过这个 笔记本 来尝试使用 optimum-quanto + transformers!
🤗 optimum-quanto 是一个多功能的 PyTorch 量化工具包。它使用的量化方法是线性量化。Optimum-quanto 提供了以下几个独特功能:
- 权重量化(
float8,int8,int4,int2) - 激活量化(
float8,int8) - 模态无关(例如 CV, LLM)
- 设备无关(例如 CUDA, XPU, MPS, CPU)
- 与
torch.compile兼容 - 容易添加特定设备的自定义内核
- 支持量化感知训练
在开始之前,请确保安装了以下库:
In [ ]:
pip install optimum-quanto accelerate transformers
现在,你可以通过在 from_pretrained() 方法中传递 QuantoConfig 对象来量化模型。只要模型包含 torch.nn.Linear 层,这适用于任何模态的任何模型。
Transformers 集成仅支持权重量化。对于更复杂的情况,如激活量化、校准和量化感知训练,你应该使用 optimum-quanto 库。
默认情况下,无论权重实际存储的数据类型(如 torch.float16),权重都会以全精度(torch.float32)加载。设置 torch_dtype="auto" 可以根据模型的 config.json 文件中定义的数据类型自动加载最节省内存的数据类型。
In [ ]:
from transformers import AutoModelForCausalLM, AutoTokenizer, QuantoConfig
model_id = "facebook/opt-125m"
tokenizer = AutoTokenizer.from_pretrained(model_id)
quantization_config = QuantoConfig(weights="int8")
quantized_model = AutoModelForCausalLM.from_pretrained(model_id, torch_dtype="auto", device_map="cuda:0", quantization_config=quantization_config)
请注意,目前 Transformers 尚不支持序列化,但这即将实现!如果你想要保存模型,可以使用 quanto 库。
Optimum-quanto 库使用线性量化算法进行量化。尽管这是一种基本的量化技术,但我们仍然获得了非常好的结果!请查看以下基准测试(llama-2-7b 在困惑度指标上的表现)。你可以在这里找到更多基准测试:这里

该库足够灵活,可以与大多数 PTQ 优化算法兼容。未来计划以最无缝的方式集成最受欢迎的算法(AWQ, Smoothquant)。