显存占用最高直降52%！哈工大&华为LLM分布式训练量化方案，激活梯度一起压缩

随着大语言模型参数规模飙升至数十亿甚至上百亿级别，分布式训练早已成为行业标配。但当下 LLM 训练始终绕不开一大核心痛点：GPU 显存资源极度紧张，模型参数、梯度、优化器状态与海量中间激活值挤占大量显存，动辄出现显存溢出问题。

现有量化方案虽能一定程度压缩开销，却陷入两难困境：4-bit 激活量化极易造成精度暴跌、收敛变慢，8-bit 梯度量化又存在梯度累加溢出、通信时延居高不下的问题。同时 Megatron-LM、DeepSpeed、ZeRO 等主流框架，也无法兼顾显存压缩、训练速度与模型效果三者平衡。

针对这一行业瓶颈，哈尔滨工业大学联合华为团队提出全新AGoQ 量化训练方案，融合激活量化与梯度量化两大核心思路，适配 LLM 分布式训练全流程。在 8B 至 34B 量级模型训练中，AGoQ 最高可降低 52% 训练显存占用，训练速度较主流框架提升至 1.34 倍，全程不损失预训练收敛效果与下游任务精度，为大模型高效低成本训练提供了全新解法。

• 论文标题：

  AGoQ: Activation and Gradient Quantization for Memory-Efficient Distributed Training of LLMs

• 论文链接：

  https://arxiv.org/html/2605.00539v2

AGoQ 基于 Megatron-LM 框架搭建，兼容 8-bit 优化器量化，从激活值、梯度两大显存占用大户入手，搭配精细化工程优化，打破传统量化局限。

（1）层感知激活量化策略

研究团队通过严谨的梯度误差理论分析发现，Transformer 架构中不同模块对量化的耐受度天差地别：注意力模块的 QKV 层量化后会产生极大梯度误差，直接影响训练稳定性，因此保留高精度不做压缩；而 RMSNorm、SiLU 激活函数、MLP 前馈网络等模块，经过中间值重计算设计，可安全压缩至近 4-bit 精度。

同时针对流水线并行训练中不同 GPU 设备激活存储量不均衡的问题，设计动态位宽补偿机制，给显存富余的设备分配更高量化位宽，在不增加峰值显存的前提下，弥补量化带来的精度损耗。

此外，AGoQ 还通过量化/反量化与 GEMM 算子内核融合，把逐元素量化运算和矩阵计算尽量合并到同一个 GPU kernel 中执行，从而减少额外调度和访存开销。

（2）保精度 8-bit 梯度量化算法

传统梯度量化仅用于通信传输，仍以高精度格式存储，无法节省显存。AGoQ 采用 8-bit FP8 格式存储梯度，在本地小批次梯度累加时，先将低比特梯度反量化为高精度进行求和计算，再重新量化回 8-bit，这样既保留了低比特存储的显存优势，又有效避免了直接在 FP8 上累加带来的溢出与精度损失风险。

同时重构分布式 All-Reduce 通信逻辑，将原本的归约散射全聚合流程，拆解为逐设备通信、本地高精度归约再全局聚合的模式，既压缩了梯度通信数据量，又保证了多卡梯度聚合的精度。

本次实验搭建最高 64 卡 NVIDIA A6000 GPU 集群，搭载 200Gb/s 高速 InfiniBand 网络，同时兼容华为昇腾 910 NPU 硬件平台；软件环境适配主流深度学习栈，选用 OLMo-1B、LLaMA2-7B/13B、LLaMA3-8B、CodeLLaMA-34B 等主流模型，序列长度覆盖 16K 至 80K 长序列场景。实验选取 Megatron-LM、DeepSpeed、ZeRO 系列、COAT 等当下主流训练系统作为基线，从显存占用、训练速度、收敛能力、下游任务精度四大维度展开对比测试。

显存层面，相较于全精度 BF16 训练，AGoQ 最高削减 52% 显存占用，对比 COAT 方案额外降低 30% 激活显存与 75% 梯度显存，在 32K、64K 超长序列训练场景下，大幅减少显存溢出故障。

速度层面，在 64 卡分布式环境中，AGoQ 平均训练提速 1.23 倍，80K 超长序列下更是达到 1.34 倍加速；即便在 10Gbps 低带宽商用网络环境中，梯度通信时延仍能实现 3.7 倍大幅缩减，适配各类集群部署场景。

在模型效果上，AGoQ 预训练损失曲线与全精度基线高度重合，远优于微软 FP8 梯度量化方案的收敛表现。在零样本评测任务中，AGoQ 在 ARC、HellaSwag、PIQA、SciQ、Winogrande 等任务上的整体表现接近基线，部分任务小幅提升，个别任务也存在轻微下降，但整体未观察到系统性精度退化。

消融实验也进一步证实如下图，层感知量化避免了统一 4-bit 量化的收敛失效问题，动态位宽补偿与内核融合设计，共同保障了精度与训练效率的双赢。

AGoQ 构建了层感知 4-bit 激活量化 + 保精度 8-bit 梯度量化的一体化方案，精准解决了大模型分布式训练显存紧张、量化易失精度、通信开销过大三大痛点。通过模块级误差适配、流水线动态优化、算子内核重构等多重设计，无需牺牲模型收敛能力与下游任务精度，就能实现显存大幅压缩与训练吞吐量显著提升。

该方案可无缝适配 LLaMA 系列主流大模型，兼容 GPU 与国产 NPU 硬件，能直接集成到 Megatron-LM 等主流训练框架中，非常适合企业与科研机构开展大模型预训练、长序列训练等任务。

未来可进一步拓展至更大万亿级参数模型、微调场景，同时优化多机多卡超大规模集群的通信调度策略，让低位宽量化技术真正落地普及，进一步降低大模型训练的硬件门槛与成本。