AQLM

当大模型遇上"极致压缩术"：AQLM 如何把 70B 参数塞进一张显卡

一个让工程师睡不着觉的问题

想象一下：你训练出了一个效果惊艳的 700 亿参数大模型，团队欢呼雀跃。但下一秒，噩耗传来——这个模型需要 140GB 显存才能运行。一块顶配 A100 80GB 显卡，勉强够用；如果是 4090 消费级显卡？做梦。模型越大、能力越强，但部署门槛也越高——这不是某家公司的困境，而是整个 AI 行业都在面对的核心矛盾。

AQLM（Additive Quantization for Large Language Models） 正是为解决这个问题而来。这是一个由 Vahe Karagulyan 等研究者开发的 PyTorch 原生量化框架，核心目标只有一个：把大模型压缩到极低位宽（2-bit、4-bit），同时尽量保持模型能力不衰减。2024 年 1 月发布论文以来，该项目已在 GitHub 获得超过 1300 颗星，被 vLLM 官方集成，成为 LLM 极致压缩领域的标杆工作之一。

加性量化：从"四舍五入"到"密码本"

传统量化方法理解起来很简单——就像把高清图片压缩成 JPEG：把所有数值四舍五入到更少的位数（FP16 → INT8），损失一点精度换来巨大的体积缩小。但这种方法在大模型上有个致命问题：粗暴截断会严重破坏模型表达能力。一个 70B 参数的模型，如果量化到 4-bit，误差会累积到让模型完全不可用的程度。

AQLM 的核心创新在于"加性量化"（Additive Quantization）思想。它不把整个权重矩阵当成一个整体来量化，而是将每个权重向量分解为 K 个子向量的和，每个子向量来自一个独立的码本（Codebook）。这样，原来需要存储一个完整的浮点向量，现在只需要存储 K 个码字的索引，压缩率极高。更关键的是，通过束搜索（Beam Search）优化码本组合，可以在离散空间中逼近连续最优解。

项目源码中 src/beam_search_l2.py 和 src/beam_search_xtx.py 分别实现了 L2 距离和激活误差最小化的束搜索算法。src/kmeans.py 负责码本的 K-means 初始化。src/aq.py 是核心——QuantizedWeight 类管理码本结构和量化权重，QuantizedLinear 则是替换原生 Linear 层的量化算子。推理侧，inference_lib/src/aqlm/inference.py 提供了 CUDA 加速的推理内核，inference_kernels/ 目录下是底层实现。

PV-Tuning：绕过梯度估计的难题

量化模型的训练有一个经典难题：量化操作是不可导的。在反向传播时，梯度无法流过量化节点——这叫"STE（Straight-Through Estimation）"问题的根源。传统方法用"直通估计"强行让梯度通过，但这种方法在极端低位宽（2-bit）下表现很差。

AQLM 团队在 2024 年 5 月提出了 PV-Tuning，专门解决这个问题。PV-Tuning 通过优化量化码本来间接训练——不再直接更新量化后的权重，而是优化码本向量本身，使得量化后的结果能更好地拟合目标分布。项目 src/pv_utils.py 和 src/pv_optimizer.py 实现了 PV-Tuning 的核心逻辑，包括针对 Transformers 模型的自适应分片包装（FullyShardedDataParallel）策略。

支持的模型与量化配置

AQLM 支持主流的开源 LLM 架构，包括 Llama、Llava、Mistral、Mixtral、Gemma、Phi3、Qwen2 等（见 convert_to_hf.py 中的 model_type 判断逻辑）。支持的量化位宽从 1-bit 到 8-bit 不等，其中 2-bit 和 4-bit 是主力场景。量化后的模型可通过 convert_to_hf.py 转换为 HuggingFace 格式，零门槛加载到 transformers 生态中使用。

项目还提供了与 vLLM 集成的 Notebook（notebooks/aqlm_vllm.ipynb），可利用 vLLM 的高效推理引擎在量化模型上实现高吞吐量的批量推理。此外还有 CUDA Graph 加速示例（notebooks/aqlm_cuda_graph.ipynb），通过减少 CUDA 内核启动开销进一步提升推理速度。

上手体验：有门槛，但回报丰厚

优点：

• 量化效果拔群：2-bit 量化下模型能力损失控制在极低水平，论文在多个基准上验证了有效性
• PyTorch 原生实现，与 transformers/accelerate 生态无缝衔接
• 提供独立推理库 aqlm（PyPI 分发），pip install 即可获得 CUDA 加速推理
• vLLM 官方集成，生产可用
• 代码结构清晰，aq_engine.py 封装了端到端量化流程

缺点与局限：

• 纯 Python 库，无 Docker 一键部署，部署需要手动配置 Python + CUDA 环境
• 量化过程需要全量显存（原始模型需能加载），对硬件有一定要求
• 不支持 Web UI，本质是底层压缩工具链，面向有 ML 经验的开发者
• 目前仅支持 CausalLM 架构，其他任务类型需扩展

行业意义：压缩即竞争力

大模型压缩不是一个边缘课题，而是决定谁能真正落地大模型应用的关键战场。在端侧部署、边缘计算、降低推理成本等场景下，每压缩一个 bit 都意味着更低的硬件门槛和更高的利润率。AQLM 所代表的"极致量化"路线（2-bit）与 GPTQ、AWQ、GGML 等技术路线并驾齐驱，各有适用场景。AQLM 的差异化在于对加性量化数学模型的深入挖掘，以及 PV-Tuning 对训练信号传递问题的独到解决。

该项目源码结构严谨，src/ 目录完整实现了量化算法，inference_lib/ 提供生产级推理支持，benchmark/ 目录包含详细的性能评测脚本。对于想深入理解 LLM 量化原理、或需要在生产环境中压缩模型的开发者而言，AQLM 是值得深入研究的优秀资源。