LoCo：用于大规模模型训练的低位通信适配器
机器学习
2024-12-02 v2 最优化与控制
摘要
高效训练大规模模型需要将本地GPU节点的全精度梯度压缩为低精度梯度，以提高GPU节点之间的梯度同步效率。然而，这通常会因压缩信息损失而降低训练质量。为此，我们提出低位通信适配器（LoCo），在压缩前对本地GPU节点进行梯度补偿，以确保高效同步而不牺牲训练质量。具体而言，LoCo设计了一个历史补偿误差的移动平均，以稳定估计当前压缩误差，然后采用其对当前梯度压缩进行补偿，从而实现更少的信息损失。该机制使其能够与通用优化器（如Adam）以及分片策略（如FSDP）兼容。理论分析表明，将LoCo集成到像Adam和SGD这样的全精度优化器中，不会影响在非凸问题上的收敛速度。实验结果表明，在像Megatron-LM和PyTorch FSDP这样的大规模模型训练框架中，LoCo显著提高了通信效率，例如在LLAMAs和MoE等大型语言模型上，LoCo可将Adam的训练速度提高14%至40%，且不影响性能。
引用
@article{arxiv.2407.04480,
title = {LoCo: Low-Bit Communication Adaptor for Large-scale Model Training},
author = {Xingyu Xie and Zhijie Lin and Kim-Chuan Toh and Pan Zhou},
journal= {arXiv preprint arXiv:2407.04480},
year = {2024}
}