AGoQ：LLM 内存高效分布式训练的激活和梯度量化

量化是减少训练 大语言模型 (LLM) 的 GPU 内存需求的关键方法。然而，当前的方法对于 4 位激活和 8 位梯度无效，这很容易导致收敛缓慢或精度损失。为了解决这个问题，我们引入了 AGoQ，结合了两种新技术：1）层感知激活量化算法，根据层的类型和管道阶段为各个层的激活分配适当的位宽，以实现近 4 位激活存储；2）梯度量化算法，通过采用 8 位梯度存储和保留精度的 8 位 All-Reduce 通信来减少内存使用并缩短通信时间。我们在两个 GPU 集群（最多 64 个 GPU）上使用不同大小的 LLM 进行了广泛的实验，实验结果表明，与最先进的训练系统 Megatron-LM（带或不带 ZeRO）、COAT 和 DeepSpeed（具有 8B 至 32B LLaMA 模型）相比，我们的 AGoQ 减少了高达 52\% 的内存，并实现了高达 1.34$\times$ 的训练速度提升，同时使用 LLaMA 架构实现预训练的收敛损失和下游任务的可比精度。