Optimum Quanto: 为 PyTorch 模型提供高效量化的强大工具

在深度学习领域,模型量化是一项重要的技术,可以显著减少模型的计算和内存开销,从而提高推理效率并便于部署。Hugging Face 最近推出的 Optimum Quanto 是一个专为 PyTorch 模型设计的量化后端,为研究人员和开发者提供了一套强大而灵活的工具,可以轻松地对各种深度学习模型进行量化。

Optimum Quanto 的主要特性

Optimum Quanto 具有以下几个突出的特点:

易用性强: 所有功能都可在 eager 模式下使用,无需对模型进行跟踪。
设备兼容性好: 量化后的模型可以在任何设备上运行,包括 CUDA 和 MPS。
自动化程度高: 自动插入量化和反量化桩,自动插入量化的函数操作和模块。
工作流程完整: 提供从浮点模型到动态量化再到静态量化模型的无缝工作流。
序列化兼容性: 与 PyTorch 的 weight_only 和 Hugging Face 的 safetensors 兼容。
硬件加速: 在 CUDA 设备上提供加速的矩阵乘法(如 int8-int8, fp16-int4 等)。
支持多种量化精度: 支持 int2、int4、int8 和 float8 权重,以及 int8 和 float8 激活。

Optimum Quanto架构图

Optimum Quanto 的性能表现

Optimum Quanto 在保持模型精度的同时,可以显著降低模型的内存占用和推理延迟。以 meta-llama/Meta-Llama-3.1-8B 模型为例:

精度: 使用 int8/float8 权重和 float8 激活的量化模型,其性能与全精度模型非常接近。
延迟: 在有优化内核的情况下,仅量化模型权重时,量化模型的推理速度与全精度模型相当。
设备内存: 内存占用大约减少了 float 位数 / 整数位数倍。

Meta-Llama-3.1-8B模型性能对比

使用 Optimum Quanto 进行模型量化

Optimum Quanto 提供了简单直观的 API,使用户可以轻松地对 Hugging Face 模型进行量化。以下是一个基本的量化工作流程:

安装 Optimum Quanto:

pip install optimum-quanto

量化模型:

from transformers import AutoModelForCausalLM
from optimum.quanto import QuantizedModelForCausalLM, qint4

model = AutoModelForCausalLM.from_pretrained('meta-llama/Meta-Llama-3-8B')
qmodel = QuantizedModelForCausalLM.quantize(model, weights=qint4, exclude='lm_head')

保存量化后的模型:

qmodel.save_pretrained('./Llama-3-8B-quantized')

加载量化后的模型:

from optimum.quanto import QuantizedModelForCausalLM

qmodel = QuantizedModelForCausalLM.from_pretrained('Llama-3-8B-quantized')

除了对语言模型的支持,Optimum Quanto 还可以用于量化 Diffusers 模型中的子模型,为各种深度学习任务提供全面的量化支持。

Optimum Quanto 的设计理念

Optimum Quanto 的核心是一个 Tensor 子类,它负责:

将源 Tensor 投影到给定目标类型的最佳范围内。
将投影后的值映射到目标类型。

这种设计允许 Quanto 在保持精度的同时,最大限度地减少饱和值和被置零值的数量。对于 int8 和 float8,投影是按张量或按通道对称的;对于更低的位宽,则采用分组仿射(带有偏移或"零点")的方式。

在模块层面,Quanto 提供了一种通用机制,可以将 torch 模块替换为能够处理 quanto 张量的 optimum-quanto 模块。这些模块在模型被冻结之前动态转换其权重,虽然会稍微减慢推理速度,但允许模型进行微调。

结语

Optimum Quanto 为 PyTorch 模型量化提供了一个强大而灵活的解决方案。它不仅简化了量化过程,还在各种设备上提供了出色的性能。随着深度学习模型变得越来越大,像 Optimum Quanto 这样的工具将在提高模型效率和可部署性方面发挥越来越重要的作用。

无论您是研究人员、开发者还是机器学习工程师,Optimum Quanto 都为您提供了一个值得探索的强大工具,可以帮助您更好地优化和部署 PyTorch 模型。随着该项目的不断发展,我们可以期待看到更多令人兴奋的特性和改进,进一步推动深度学习模型的效率和可用性。