Transformers 中的 AQLM 量化Additive Quantization of Language Models 原理、配置与实战【免费下载链接】transformers Transformers: the model-definition framework for state-of-the-art machine learning models in text, vision, audio, and multimodal models, for both inference and training.项目地址: https://gitcode.com/GitHub_Trending/tra/transformersAQLMAdditive Quantization of Language Models语言模型加法量化是 Transformers 中继 AWQ、GPTQ 之后又一重量级权重量化方案其核心思路是对多个权重联合量化利用权重之间的相互依赖关系换取更低的比特率与更高的精度。本文以官方指南 aqlm.md 为骨架结合仓库内量化器实现与集成测试源码系统讲解 AQLM 的量化原理、主流配置形态、加载方式、LoRA 微调与torch.compile加速以及可扩展的配置参数细节。读完本文你将掌握如何在 Transformers 中正确加载、训练、保存并深度调优一个 AQLM 量化模型。AQLM 量化原理与定位AQLMAdditive Quantization of Language Models源自论文Additive Quantization for Language Models见 docs/source/en/quantization/aqlm.md 首段引用的研究条目。与逐权重独立量化的方案不同AQLM将 8~16 个权重作为一组并把这组权重表示成多个向量码本vector code的加和AQLM represents groups of 8-16 weights as a sum of multiple vector codes.换言之传统标量量化scalar quantization只关注单个权重本身的值域而 AQLM 采用向量量化vector quantization思想让一组权重共享一组可学习的码本向量。由于它显式捕捉并利用了组内权重之间的统计相关性interdependenciesAQLM 可以在极低比特如 2-bit下依然保持较好的模型精度这一特性使它尤其适合大模型如 Llama-2-7B、Mixtral-8x7B在单卡上的部署。在 Transformers 中AQLM 属于预量化权重加载型集成模型权重在 Hub 上已由 AQLM 工具库完成量化并发布为 Hugging Face 格式from_pretrained时由仓库自动识别quantization_config并调用 AQLM 的QuantizedLinear层替换原始nn.Linear见下文源码分析。此外它具备两个重要的扩展能力可配合 PEFT 库以LoRA方式对量化后的模型做微调官方文档明确支持与torch.compile完全兼容可获得更快的推理与训练速度。安装与环境要求AQLM 官方指南给出的安装命令非常精简但其背后有三个依赖需要齐备。运行pip install aqlm[gpu,cpu]其中aqlm[gpu,cpu]是带 GPU 与 CPU 推理/训练内核kernel支持的完整安装官方明确AQLM 仅支持 Python 3.10结合仓库源码实际运行还需要 Transformers 侧的Accelerate。关于第三点量化器实现 quantizer_aqlm.py 在validate_environment中做了两项强制校验任一缺失都会直接抛出ImportErrorif not is_accelerate_available(): raise ImportError(Using aqlm quantization requires Accelerate: pip install accelerate) if not is_aqlm_available(): raise ImportError(Using aqlm quantization requires AQLM: pip install aqlm[gpu,cpu])因此最小可用环境的完整安装命令实际应为pip install aqlm[gpu,cpu] accelerate另外需注意 AQLM 库版本与能力绑定关系源自 quantizer_aqlm.py 的is_trainable属性与集成测试训练支持aqlm版本须 1.0.2否则即使加载成功也会给出pip install aqlm1.0.2的告警并判定模型不可训练torch.compile支持集成测试 test_aqlm.py 明确要求aqlm1.0.3才允许执行torch.compile用例。Transformers 侧的可用性探测函数为is_aqlm_available()定义于 import_utils.py通过_is_package_available(aqlm)检测。加载一个已量化的 AQLM 模型与加载普通模型没有任何心智负担差异直接使用from_pretrained即可。官方示例以 2-bit 量化的 Mixtral 8x7B 模型为例from transformers import AutoTokenizer, AutoModelForCausalLM quantized_model AutoModelForCausalLM.from_pretrained( ISTA-DASLab/Mixtral-8x7b-AQLM-2Bit-1x16-hf, dtypeauto, device_mapauto, )代码中两点值得展开device_mapauto表示模型按显存自动切分到可用设备仓库集成测试 test_aqlm.py 还验证了device_mapauto下的多 GPU 加载场景检查hf_device_map覆盖全部 GPU 编号。dtypeauto表示沿用 checkpoint 存储时的 dtype。加载链路从源码角度可以还原为三条主线配置解析AutoModelForCausalLM.from_pretrained读取模型config.json中的quantization_config自动映射到 AqlmConfig量化器调度自动量化分发表 auto.py 与 auto.py 分别把方法名aqlm绑定到AqlmHfQuantizer与AqlmConfig层替换权重加载前AqlmHfQuantizer._process_model_before_weight_loading调用replace_with_aqlm_linear完成层替换见下节。仓库还提供了一个更小的端到端验证模型BlackSamorez/Llama-2-7b-AQLM-2Bit-1x16-hf集成测试即用它做生成与保存/加载回归。加载后即可正常使用generatetokenizer AutoTokenizer.from_pretrained(BlackSamorez/Llama-2-7b-AQLM-2Bit-1x16-hf) input_ids tokenizer(Hello my name is, return_tensorspt) output quantized_model.generate(**input_ids, max_new_tokens32) print(tokenizer.decode(output[0], skip_special_tokensTrue))加载背后的层替换机制核心替换函数 replace_with_aqlm_linear 是所有 AQLM 加载行为的落点。其执行逻辑为遍历模型所有命名模块对每个未被排除且类型为torch.nn.Linear的模块在meta设备上实例化对应的aqlm.QuantizedLinear依次传入量化配置中的in_group_size、out_group_size、num_codebooks、nbits_per_codebook记录原模块类型到source_clsrequires_grad_(False)并通过set_submodule原地替换。值得注意的是第 2 步中模块是否替换由should_convert_module(module_name, modules_to_not_convert)判定凡是命中linear_weights_not_to_quantize列表如lm_head.weight的模块会保留原精度不被量化。替换后若整个模型没有任何 Linear 被转换会打印告警提示检查模型结构。集成测试 test_aqlm.py 对这一点做了精确的数值验证对facebook/opt-350m全量替换后模型中QuantizedLinear数量与原始nn.Linear数量严格相等传入modules_to_not_convert[lm_head]后QuantizedLinear数量恰好比原 Linear 数少 1证明排除逻辑按预期生效。AQLM 配置与内核支持矩阵AQLM 设置的核心自由度官方文档指出不同 AQLM 量化设置主要差异集中在两点码本数量number of codebooks记作 K码本大小codebook size以比特计记作 N。配置的记号约定为KxN。例如1x16表示 1 个码本、每个码本 16-bit码本规模为2^16个向量2x8表示 2 个码本、每个码本 8-bitKx8表示 K 个 8-bit 码本。这与AqlmConfig中的字段一一对应num_codebooks默认 1即码本数量 Knbits_per_codebook默认 16即单码本向量编码比特数 N实际码本大小为2**nbits_per_codebook见 quantization_config.py 的 docstring。内核支持矩阵继承官方表格不同内核Kernel对不同配置的支持情况、精度与加速比表现如下Kernel码本数量码本大小(bits)记法精度加速比快速 GPU 推理快速 CPU 推理TritonKNKxN-最高约 0.7x✅❌CUDA1161x16最佳最高约 1.3x✅❌CUDA282x8尚可最高约 3.0x✅❌NumbaK8Kx8良好最高约 4.0x❌✅从表中可以提炼出几条实用决策准则追求最佳精度选择 CUDA 内核 1x16配置官方给出的示例模型Mixtral-8x7b-AQLM-2Bit-1x16-hf与Llama-2-7b-AQLM-2Bit-1x16-hf即属于此类追求速度且为 GPUCUDA 2x8最高约 3.0x精度尚可OKCPU 上的快速推理Numba Kx8最高约 4.0x这也是唯一支持 CPU 快速推理的组合Triton 内核组合通用但加速不显著最高约 0.7x且不支持 CPU。需要说明以上加速比为官方指南对 AQLM 各内核相对原始模型推理的标称数据具体收益与模型规模、批次、算子及硬件环境相关请以实测为准。AqlmConfig 参数详解若需要手工构造量化配置例如将未量化模型以 AQLM 配置加载或多卡实验可导入AqlmConfig。其完整参数及默认值整理如下源码见 quantization_config.py参数类型默认值说明in_group_sizeint8沿输入维度的分组大小out_group_sizeint1沿输出维度的分组大小官方建议恒为 1num_codebooksint1加法量化过程使用的码本数量即 Knbits_per_codebookint16单个码本向量的编码比特数码本大小为2**nbits_per_codebook即 Nlinear_weights_not_to_quantizeOptional[list[str]]None不参与量化的nn.Linear权重完整路径列表如lm_head.weightpost_init阶段会自动执行类型与取值校验四个核心参数必须是int否则抛出TypeErrorlinear_weights_not_to_quantize若非列表会抛出ValueError为None时会被规范化为空列表。AqlmConfig使用范例from transformers import AqlmConfig quantization_config AqlmConfig( in_group_size8, out_group_size1, num_codebooks1, nbits_per_codebook16, linear_weights_not_to_quantize[lm_head.weight], )仓库测试覆盖了其两种序列化路径见 test_aqlm.pyto_dict()与配置对象字段完全等价from_dict()可接受包含in_group_size32, num_codebooks8, nbits_per_codebook8等自定义取值的字典并正确还原。附加说明不要对未量化模型误用 AQLM 配置AQLM 是加载预量化权重的方案其本身并不在from_pretrained时对随机权重执行端到端量化校准AqlmHfQuantizer.requires_calibration True量化过程需要外部校准工具完成。集成测试中的test_raise_if_non_quantized用例test_aqlm.py证实若试图用AqlmConfig(bits4)等配置加载一个没有对应量化权重的普通模型如facebook/opt-125m会触发ValueError。因此把AqlmConfig传给from_pretrained前务必确认目标 checkpoint 确实是以 AQLM 格式量化的模型。LoRA 微调与 torch.compileAQLM 的优势不止于极低比特推理还在于它可以继续参与训练和编译优化。LoRA 微调AQLM 与 PEFT 库的 LoRA 方法配套使用文档原话 AQLM also supports fine-tuning with LoRA with the PEFT library。你可以在量化后的权重之上插入低秩适配器进行参数高效微调而无需更新庞大的量化权重本体。使用前提是把aqlm库升级到 1.0.2见前文is_trainable逻辑。torch.compile 加速量化模型在加载后可直接与torch.compile组合以获得更快的推理与训练。集成测试 test_aqlm.py 给出了完整的编译推理范式核心步骤包括使用StaticCache构造静态 KV 缓存预分配max_cache_len 前缀长度 max_new_tokens 1的缓存空间先跑一次完整前向以填充前缀 KV 缓存并触发内核编译用torch.compile(decode_one_tokens, modereduce-overhead, fullgraphTrue)编译逐 token 解码函数把每步解码编译为 CUDA Graph 以削减 Python 与 kernel launch 开销在torch.no_grad()内逐 token 生成并校验输出。注意该用例声明了aqlm1.0.3的前置条件同时测试中对含torch.Any类型的某些近期 PyTorch 版本做了显式跳过标记skip注释说明等待 AQLM 侧修复意味着实际执行torch.compile时请留意本机 PyTorch 与 aqlm 的版本组合兼容性。保存与再次加载量化模型AqlmHfQuantizer.is_serializable()返回True因此 AQLM 量化模型可以被保存并重新加载集成测试 test_aqlm.py 覆盖了完整的保存-加载-生成回归。操作方式与普通模型一致with tempfile.TemporaryDirectory() as tmpdirname: quantized_model.save_pretrained(tmpdirname) reloaded AutoModelForCausalLM.from_pretrained(tmpdirname, device_mapauto) # 重新加载后可继续 generate 并应得到与保存前一致的输出从量化到发布模型获取与再生产物说明官方指南同时指出AQLM 的完整工作流量化你的模型、把量化模型推送到 Hub 等可以参考 AQLM 官方仓库及其实例 notebook。需要说明的是量化压缩权重本身由 AQLM 训练工具完成Transformers 在本文所述链路中承担的是已量化 checkpoint 的加载、训练LoRA、推理与保存角色。如果你希望亲手量化模型并发布...-AQLM-2Bit-1x16-hf这类 Hub 格式的 checkpoint应使用 AQLM 上游工具的校准流程生成quantization_config再经由本文的from_pretrained路径验证加载。小结AQLM 在 Transformers 中的正确打开方式梳理 AQLM 在 Transformers 中的完整技术闭环可以收敛为以下要点环节关键结论证据位置原理组级向量加法量化8~16 权重一组、码本加和表示aqlm.md环境pip install aqlm[gpu,cpu] acceleratePython 3.10训练需aqlm1.0.2quantizer_aqlm.py加载直接from_pretrained自动触发QuantizedLinear层替换integrations/aqlm.py配置AqlmConfig的 K/N 组合决定精度与内核选择quantization_config.py训练与加速支持 PEFT LoRA 微调与torch.compile需新版本 aqlmtest_aqlm.py保存可save_pretrained并原样重新加载test_aqlm.py总体而言AQLM 是 Transformers 生态中面向超低比特2-bit 级别 精度保持诉求的主力方案在追求极致压缩比的场景如 Mixtral 8x7B 等超大规模 MoE 模型落地单卡下1x16配置提供最佳精度2x8/Kx8提供更强吞吐配合 LoRA 与torch.compile还可无缝衔接下游训练与部署优化。相关主文档位于 docs/source/en/quantization/aqlm.md你可以在仓库docs/source/en/quantization/目录下对比阅读其他量化方案指南以选择最适配自身任务的路线。【免费下载链接】transformers Transformers: the model-definition framework for state-of-the-art machine learning models in text, vision, audio, and multimodal models, for both inference and training.项目地址: https://gitcode.com/GitHub_Trending/tra/transformers创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考