半二次量化(HQQ)
这个仓库包含了我们文章中提出的半二次量化(HQQ)的官方实现:
什么是HQQ?
HQQ是一种快速且准确的模型量化器,无需校准数据。在几分钟内就能量化最大的模型,无需校准数据 🚀。
常见问题
为什么我应该使用HQQ而不是其他量化方法?- HQQ量化模型非常快。
- 支持8、4、3、2、1比特。
- 可以用于任何模型(大语言模型、视觉模型等)。
- 反量化步骤是线性操作,这意味着HQQ兼容各种优化的CUDA/Triton内核。
- HQQ兼容peft训练。
- 我们努力使HQQ完全兼容`torch.compile`,以实现更快的推理和训练。
量化模型的质量如何?
我们对语言和视觉模型都进行了详细的基准测试。请参阅我们的博客文章:HQQ,HQQ+。
量化模型的速度如何?
带有axis=1
的4位模型可以使用优化的推理融合内核,如torchao的int4_gemm。这与gpt-fast使用的是同一个内核,根据我们的基准测试,目前它是最快的内核。我们还支持Marlin内核。此外,我们专注于使hqq完全兼容torch.compile
,这可以加速训练和推理。有关更多详细信息,请参阅下面的后端部分。
我应该使用什么量化设置?
你应该从nbits=4, group_size=64, axis=1
开始。这些设置在质量、显存使用和速度之间提供了良好的平衡。如果你想在相同的显存使用下获得更好的结果,切换到axis=0
并使用ATEN后端。如果你想使用更低的比特数,如nbits=2
,你应该通过HQQ+使用axis=0
和较小的组大小,这意味着添加低秩适配器并使用小数据集进行微调。
axis
参数是什么意思?
axis
参数是进行分组的轴。通常,axis=0
比axis=1
给出更好的结果,尤其是在较低比特数时。然而,目前优化的推理运行时仅支持axis=1
。
HQQ和HQQ+有什么区别?
HQQ+是在HQQ的基础上添加了可训练的低秩适配器,以提高低比特数时的量化质量。
安装
首先,确保你有与你的CUDA版本匹配的Pytorch 2版本:https://pytorch.org/
你可以通过 pip install hqq
安装hqq。
要获取最新版本,你可以直接通过 pip install git+https://github.com/mobiusml/hqq.git
安装核心库。
或者,克隆仓库并在当前文件夹中运行 pip install .
。
基本用法
要使用HQQ进行量化,你只需要按如下方式替换线性层(torch.nn.Linear
):
from hqq.core.quantize import *
#量化设置
quant_config = BaseQuantizeConfig(nbits=4, group_size=64)
#替换你的线性层
hqq_layer = HQQLinear(your_linear_layer, #torch.nn.Linear或None
quant_config=quant_config, #量化配置
compute_dtype=torch.float16, #计算数据类型
device='cuda', #cuda设备
initialize=True, #使用False以稍后量化
del_orig=True #如果为True,删除原始层
)
量化参数设置如下:
nbits
(int):支持8、4、3、2、1比特。group_size
(int):只要weight.numel()
可以被group_size
整除,就没有限制。quant_zero
(bool):如果为True,将零点量化为8位,不分组。quant_scale
(bool):如果为True,将缩放因子量化为8位,组大小为128。offload_meta
(bool):如果为True,元数据会被卸载到CPU。view_as_float
(bool):如果为True,量化参数会被视为浮点类型而不是整数类型。
设置offload_meta=True
会大大减少GPU内存需求,但对于较小的组大小会使处理变慢。启用时,你可以仅使用18.8GB和13GB的显存分别运行HQQ 2位量化的Llama2-70B和Mixtral。
后端
原生后端
HQQLinear
模块可以使用以下原生后端:
HQQLinear.set_backend(HQQBackend.PYTORCH) #Pytorch后端
HQQLinear.set_backend(HQQBackend.PYTORCH_COMPILE) #编译后的Pytorch
HQQLinear.set_backend(HQQBackend.ATEN) #Aten/CUDA后端
HQQBackend.ATEN
后端在可用时会自动安装并默认使用。
注意HQQBackend.ATEN
仅支持axis=0
。对于axis=1
,你需要使用HQQBackend.PYTORCH
或HQQBackend.PYTORCH_COMPILE
。
以下是使用各种后端的加速基准测试,以HQQBackend.PYTORCH
为基准:
更快的推理
我们支持使用融合内核的外部后端以实现更快的推理。你可以在模型量化后按如下方式启用其中一个后端:
from hqq.utils.patching import prepare_for_inference
#使模型与fullgraph torch.compile兼容的Pytorch后端:适用于任何设置
#prepare_for_inference(model)
#Torchao的tiny_gemm后端(最快):nbits=4, compute_dtype=bfloat16, axis=1
prepare_for_inference(model, backend="torchao_int4")
#Marlin后端:nbits=4, axis=1, compute_dtype=float16, group_size=None
#prepare_for_inference(model, backend="marlin", allow_merge=True)
#Bitblas后端:nbits=4/2/1, axis=1, compute_dtype=float16, group_size=None
#prepare_for_inference(model, backend="bitblas")
这些后端仅适用于4位量化和axis=1
。此外,对于Marlin,我们仅支持group_size=None
。以下是不同后端之间的比较。torchao内核在4090上达到195 tokens/sec(生成速度)。
与模型一起使用
Transformers 🤗
对于HF的transformers的使用,请参见文档中的以下示例:
from transformers import AutoModelForCausalLM, HqqConfig
# 所有线性层将使用相同的量化配置
quant_config = HqqConfig(nbits=4, group_size=64, quant_zero=False, quant_scale=False, axis=1)
# 加载并量化
model = AutoModelForCausalLM.from_pretrained(
model_id,
torch_dtype=torch.float16,
device_map="cuda",
quantization_config=quant_config
)
注意:使用这种方法,您无法直接通过save_pretrained
保存/加载量化模型。请改用hqq库中的保存/加载调用。
HQQ库
您也可以使用HQQ库来量化transformers模型:
#在CPU上加载模型
from transformers import AutoModelForCausalLM
model = AutoModelForCausalLM.from_pretrained(model_id, torch_dtype=compute_dtype)
#量化
from hqq.models.hf.base import AutoHQQHFModel
quant_config = BaseQuantizeConfig(nbits=4, group_size=64, quant_scale=False, quant_zero=False, axis=1)
AutoHQQHFModel.quantize_model(model, quant_config=quant_config, compute_dtype=compute_dtype, device=device)
保存/加载
您可以按以下方式保存/加载量化模型:
from hqq.models.hf.base import AutoHQQHFModel
#保存:确保在任何修补之前保存模型
AutoHQQHFModel.save_quantized(model, save_dir)
#加载
model = AutoHQQHFModel.from_quantized(save_dir)
设置后端
您可以按以下方式设置原生后端:
HQQLinear.set_backend(HQQBackend.ATEN if axis==0 else HQQBackend.PYTORCH_COMPILE)
您可以按照后端部分中的说明进行修补以实现更快的推理:
from hqq.utils.patching import prepare_for_inference
prepare_for_inference(model, backend="torchao_int4")
自定义HF模型
AutoHQQHFModel
旨在与任何transformers模型兼容。然而,其适应性有一个缺点 - 在处理层时可能会遇到问题或经历缓慢。如果您遇到这样的问题,您可以选择创建一个自定义模型,其中明确定义了修补逻辑来替代AutoHQQHFModel
。以下是您可以使用或扩展的流行模型示例:
from hqq.models.hf.llama import LlamaHQQ #Llama
from hqq.models.hf.mistral import MistralHQQ #Mistral
from hqq.models.hf.mixtral import MixtralHQQ #Mixtral
自定义量化配置 ⚙️
您可以通过为每个层名称指定设置来为不同的层设置各种量化配置:
Transformers 🤗
# 具有相同标签的每个线性层将使用专用的量化配置
q4_config = {'nbits':4, 'group_size':64, 'quant_zero':False, 'quant_scale':False}
q3_config = {'nbits':3, 'group_size':32, 'quant_zero':False, 'quant_scale':False}
quant_config = HqqConfig(dynamic_config={
'self_attn.q_proj':q4_config,
'self_attn.k_proj':q4_config,
'self_attn.v_proj':q4_config,
'self_attn.o_proj':q4_config,
'mlp.gate_proj':q3_config,
'mlp.up_proj' :q3_config,
'mlp.down_proj':q3_config,
})
HQQ库
from hqq.core.quantize import *
q4_config = BaseQuantizeConfig(nbits=4, group_size=64, quant_zero=False, quant_scale=False)
q3_config = BaseQuantizeConfig(nbits=3, group_size=32, quant_zero=False, quant_scale=False)
quant_config = {'self_attn.q_proj':q4_config,
'self_attn.k_proj':q4_config,
'self_attn.v_proj':q4_config,
'self_attn.o_proj':q4_config,
'mlp.gate_proj':q3_config,
'mlp.up_proj' :q3_config,
'mlp.down_proj':q3_config,
}
Peft训练
您可以按以下方式使用HQQ进行LoRA训练:
#首先,量化/加载一个量化的HQQ模型
from hqq.core.peft import PeftUtils
base_lora_params = {'lora_type':'default', 'r':32, 'lora_alpha':64, 'dropout':0.05, 'train_dtype':torch.float32}
lora_params = {'self_attn.q_proj': base_lora_params,
'self_attn.k_proj': base_lora_params,
'self_attn.v_proj': base_lora_params,
'self_attn.o_proj': base_lora_params,
'mlp.gate_proj' : None,
'mlp.up_proj' : None,
'mlp.down_proj' : None}
#将LoRA添加到线性/HQQ模块
PeftUtils.add_lora(model, lora_params)
#可选:设置您的后端
HQQLinear.set_backend(HQQBackend.ATEN if axis==0 else HQQBackend.PYTORCH_COMPILE)
#训练 ....
#将LoRA权重转换为相同的模型dtype以实现更快的推理
model.eval()
PeftUtils.cast_lora_weights(model, dtype=compute_dtype)
#保存LoRA权重
PeftUtils.save_lora_weights(model, filename)
#加载LoRA权重:自动调用add_lora
PeftUtils.load_lora_weights(model, filename)
我们提供了一个完整的示例,用于使用HQQ/LoRA训练模型,您可以在examples/lora/train_hqq_lora_example.py
中找到。
如果您想通过FSDP使用多GPU训练,请查看Answer.AI的这个出色的仓库:https://github.com/AnswerDotAI/fsdp_qlora
示例
我们在examples
目录中提供了各种示例,展示了不同后端的模型量化。
引用 📜
@misc{badri2023hqq,
title = {Half-Quadratic Quantization of Large Machine Learning Models},
url = {https://mobiusml.github.io/hqq_blog/},
author = {Hicham Badri and Appu Shaji},
month = {November},
year = {2023}