单卡 RTX 3090 实现 LLaMA-3 8B 低比特量化全参微调
背景与挑战
自 2010 年起,AI 技术历经多个重大发展阶段,深度学习的崛起显著推动了技术前进。尤其是 2022 年底推出的 ChatGPT,彰显了大语言模型(LLM)的能力达到了前所未有的水平。生成式 AI 大模型迅速进入高速发展期,被誉为第四次工业革命的驱动力。然而,在商业化过程中,大模型因成本过高成为企业的一大负担。持续的技术突破虽然令人鼓舞,但如果落地阶段的成本无法控制,便难以持续资助研发。
开源大模型的兴起正逐步改变这一局面。普通的消费级 GPU 就能够支持 7B/8B 规模模型的全参数微调操作,可能比采用高成本闭源模型成本低几个数量级。在这种去中心化的 AI 范式下,开源模型的应用在保证质量的前提下,可以显著降低边际成本。此外,经过量化压缩的较大模型在性能上往往优于同等大小的预训练小模型,说明量化压缩后的模型仍然保持了优秀的能力。
在云端大模型不断探索技术极限的同时,市场对于能够快速落地和支撑高速成长的智能应用有着迫切需求。边缘计算中的大模型——特别是中小型模型如 7B 和 13B 的模型——因其高性价比和良好的可调性而受到青睐。企业更倾向于自行微调这些模型,以确保应用的稳定运行和数据质量的持续控制。尽管云端模型在处理复杂任务时精度高,但它们面临关键挑战:
- 推理服务的基础设施成本:支持 AI 推理的高性能硬件稀缺且昂贵,集中式商业运营带来的边际成本递增问题成为 AI 业务从 1 到 10 必须翻越的障碍。
- 推理延迟:生产环境中模型必须快速响应,任何延迟都会直接影响用户体验。
- 隐私和数据保护:涉及敏感信息的商业应用场景中,使用第三方云服务处理敏感数据可能会引发隐私和安全问题。
考虑到这些挑战,边缘计算提供了一个有吸引力的替代方案。在边缘设备上直接运行中小模型不仅能降低数据传输的延迟,提高响应速度,而且有助于在本地处理敏感数据。
现有方案的不足与我们的出发点
在当前的开源模型和工具生态中,仍面临一系列不足之处。首先,这些模型和工具往往并未针对本地部署场景进行优化,导致在本地运用时常常受限于算力和内存资源。例如,即便是相对较小的 7B 规模模型,也可能需要高达 60GB 的 GPU 显存来进行全参数微调。此外,现有的量化技术虽然在模型推理部署中表现良好,但其主要用途是减少模型部署时的内存占用。量化后的模型权重在微调过程中无法进行优化,这限制了开发者在资源有限的情况下使用较大模型的能力。
我们的出发点在于解决上述痛点。基于 Neural Architecture Search (NAS) 以及相匹配的 Post-Training Quantization (PTQ) 量化方案,我们提供了超过 200 个从不同规模开源大模型序列压缩而来的低比特量化小模型。同时,我们推出了 Bitorch Engine 开源框架以及专为低比特模型训练设计的 DiodeMix 优化器。开发者可以直接对低比特量化模型在量化空间进行全参数监督微调与继续训练,实现了训练与推理表征的对齐,大幅度压缩模型开发与部署的中间环节。
通过结合低比特权重训练技术和低秩梯度技术,我们就能实现在单卡 RTX 3090 GPU 上对 LLaMA-3 8B 模型进行全参数微调。上述解决方案简洁有效,不仅节省资源,而且有效地解决了量化模型精度损失的问题。

模型量化技术细节
大模型时代的显著特征之一便是模型对计算资源需求的大幅度攀升。GPTQ 与 AWQ 等权重 PTQ 压缩方案以可扩展的方式验证了大语言模型在 4-bit 表征上的可靠性。与此同时,QLoRA 巧妙地将 4-bit LLM 表征与 LoRA 技术相结合,将低比特表征推广至监督微调阶段。
相比于 INT4,更低比特的 Round-To-Nearest (RTN) 量化如 INT2 表征通常要求原始模型有着更平滑的连续参数空间才能保持较低的量化损失。近期 layer-Importance 等多项工作则进一步观测到了不同深度的 transformer 模块对模型容量的参与度也表现出非均匀分布的特性。基于这些工作的启发,我们探索了一种搜索与校准结合的 Two-stage LLM 低比特量化方案。
1. 混合精度 NAS 搜索
首先,利用 NAS 相关方法对大语言模型参数空间的量化敏感性进行搜索与排序,利用经典的混合精度表征实现模型参数中的最优比特位分配。为降低模型量化后的大规模硬件部署难度,我们放弃了复杂的矢量量化与 INT3 表征等设计,采用经典 Group-wise MinMax Quantizer,同时仅选择 INT4 (group size 128) 与 INT2 (group size 64) 作为基础量化表征。
我们探索了 Layer-mix 与 Channel-mix 两种排布下的混合精度搜索空间。其中,Channel-mix 量化由于能更好适配 transformer 架构的系统性激活涌现现象,往往能达到更低的量化损失,而 Layer-mix 量化在具备更优的硬件友好度的同时仍然保持了极佳的模型容量。利用高效的混合精度 NAS 算法,我们能在数小时内基于低端 GPU 如 RTX 3090 上完成对 Qwen1.5 110B 大模型的量化排布统计,并基于统计特性在数十秒内完成任意低比特量级模型的最优架构搜索。
2. 离线知识蒸馏校准
搜索得到模型的量化排布后,引入了一种基于离线知识蒸馏的可扩展 PTQ 校准算法,以应对超低比特量化 (如 2 到 3-bit) 带来的累积分布漂移问题。仅需使用不超过 512 个样本的多源校准数据集,即可在数小时内使用单张 A100 GPU 完成 0.5B-110B 大语言模型的 PTQ 校准。随着当前开源社区 100B + 大模型的持续涌现,如何构建高效且可扩展的量化压缩方案将是 LLM 社区系统工程研究的重要组成部分。
性能分析
基于 Two-stage 量化压缩方案,我们提供了超过 200 个从不同规模开源大模型序列压缩而来的低比特量化小模型,涵盖了最新的 Llama3、Phi-3、Qwen1.5 以及 Mistral 等系列。我们利用 EleutherAI 的 lm-evaluation-harness 库等对低比特量化模型的真实性能与产业场景定位进行了探索。
其中我们的 4-bit 量化校准方案基本实现了相对于 FP16 的 lossless 压缩。基于混合 INT4 与 INT2 表征实现的 sub-4 bit 量化校准方案在多项 zero-shot 评测结果表明,搜索与少量数据校准的经典 INT2 量化表征已经足够维持 LLM 在语言模型阅读理解、常识推理以及自然语言推理方面的核心能力。
同时我们试图利用进一步的 few-shot 消融对比实验探索超低比特在产业应用中的定位。有趣的现象是,INT2 表征为主体的超低比特 (bpw: 2.2/2.5) 模型在 5-shot 帮助下即可实现推理能力的大幅度提升。这种对少量示例样本的利用能力表明,低比特压缩技术在构造容量有限但足够'聪明'的语言模型方面已经接近价值兑现期,配合检索增强 (RAG) 等技术适合构造更具备成本效益的模型服务。
考虑到当前少样本 PTQ 校准仅仅引入了有限的计算资源,以我们开源的低比特模型作为初始化进行更充分的全参数量化训练将进一步提升低比特模型在实际任务中的表现。
开源工具链
我们推出了三款实用的工具来辅助这些模型的使用,并计划未来持续优化和扩展。
Bitorch Engine (BIE)
Bitorch Engine 是一款前沿的神经网络计算库,其设计理念旨在为现代 AI 研究与开发找到灵活性与效率的最佳平衡。BIE 基于 PyTorch,为低位量化的神经网络操作定制了一整套优化的网络组件。它是实现低比特量化 LLM 全参数微调的基础。此外,BIE 还提供了基于 CUTLASS 和 CUDA 的 kernel,支持 1-8 bit 量化感知训练。我们还开发了专为低比特组件设计的优化器 DiodeMix,有效解决了量化训练与推理表征的对齐问题。在开发过程中,我们发现 PyTorch 原生不支持低比特张量的梯度计算,为此我们对 PyTorch 进行了少量调整,提供了支持低比特梯度计算的修改版。
green-bit-llm
green-bit-llm 是为 GreenBitAI low-bit LLM 专门开发的工具包。该工具包支持云端和消费级 GPU 上的高性能推理,并与 Bitorch Engine 配合,完全兼容 transformers、PEFT 和 TRL 等主流训练 / 微调框架,支持直接使用量化 LLM 进行全参数微调和 PEFT。
以目前最新的开源大模型 Llama-3 8b base 模型为例,我们选择它的 2.2/2.5/3.0 bit 作为全参数量化监督微调 (Q-SFT) 对象,使用 huggingface 库托管的 tatsu-lab/alpaca 数据集进行 1 个 epoch 的最小指令微调对齐训练测试。模型完全在量化权重空间进行学习而不涉及常规的 LoRA 参数优化与融合等后处理步骤,训练结束后即可直接实现高性能的量化推理部署。
相比于传统的 LoRA 微调 + 后量化的工程组合,Q-SFT 直接在推理模型量化空间进行学习的方式大幅度简化了大模型从开发到部署之间的工程链条,同时实现了更好的模型微调效果。此外,推理模型与训练模型的表征对齐也为更丝滑的端侧学习与应用提供了可能性。
在我们的自研低比特模型以外,green-bit-llm 完全兼容 AutoGPTQ 系列 4-bit 量化压缩模型,这意味着 huggingface 现存的 4-bit GPTQ 模型都可以基于 green-bit-llm 在量化参数空间进行低资源继续学习 / 微调。
gbx-lm
gbx-lm 工具将 GreenBitAI 的低比特模型适配至苹果的 MLX 框架,进而能在苹果芯片上高效的运行大模型。目前已支持模型的加载和生成等基础操作。用户可以在苹果电脑上迅速建立一个本地聊天演示页面。
技术实现细节与扩展
为了进一步确保低比特量化训练的稳定性和效率,我们在工程实现上做了以下关键优化:
1. 梯度截断与归一化
在低比特量化训练中,梯度的数值范围往往非常敏感。DiodeMix 优化器内部集成了动态梯度截断机制,防止在反向传播过程中出现梯度爆炸或消失。同时,针对量化权重的更新步长,我们采用了自适应归一化策略,确保每一步的参数更新都在量化网格的允许范围内,从而避免量化误差的累积。
2. 混合精度训练支持
虽然目标是全参数低比特微调,但在某些特定层(如 Embedding 层或 LayerNorm),保留 FP16 精度有助于维持模型的初始表达能力。我们的框架支持灵活的混合精度配置,允许用户在量化空间学习的同时,冻结部分关键层的浮点参数。这种策略在验证实验中显示能有效提升收敛初期的稳定性。
3. 硬件兼容性扩展
除了 NVIDIA GPU,该方案同样考虑了 Apple Silicon 的支持。通过 gbx-lm 工具,利用 MLX 框架的底层优化,可以在 Mac Studio 等设备上运行低比特模型。这对于没有高端数据中心资源的个人开发者或小团队来说,是一个重要的补充方案。未来我们将探索更多 ARM 架构设备的适配,进一步扩大边缘计算的场景覆盖。
总结与展望
本文介绍了一套完整的低比特量化全参微调解决方案。通过 NAS 搜索与 PTQ 校准的结合,我们构建了高精度的低比特模型;通过 Bitorch Engine 和 DiodeMix 优化器,我们实现了在消费级显卡上的高效训练。这套方案不仅降低了大模型落地的门槛,也为边缘计算设备上的实时智能应用提供了可行的技术路径。
未来的工作将集中在以下几个方面:一是进一步优化量化感知训练的收敛速度,减少校准数据的需求;二是探索更低比特(如 1-bit)的可行性与边界;三是完善自动化工具链,使得从预训练模型到低比特微调模型的转换更加自动化和标准化。随着技术的成熟,低比特量化将成为大模型普及的关键基础设施之一。