AMSP：降低 ZeRO 通信开销以实现高效 LLM 训练
分布式、并行与集群计算
2024-03-14 v2
摘要
训练大语言模型（LLMs）因模型状态的高内存需求而面临 GPU 内存消耗的挑战。广泛使用的 Zero Redundancy Optimizer（ZeRO）通过策略性分片解决了这一问题，但在大规模下引入了通信挑战。为应对该问题，我们提出 AMSP，一个旨在优化 ZeRO 以实现可扩展 LLM 训练的系统。AMSP 包含三种灵活的分片策略：Full-Replica、Full-Sharding 和 Partial-Sharding，并允许模型状态（参数、梯度、优化器状态）内的每个组件独立选择分片策略以及设备网格。我们对通信成本进行了深入分析，将发现问题形式化为一个优化问题以寻找最优分片策略。此外，AMSP 通过高效重叠通信与计算来优化分布式 LLM 训练。评估表明，在 1024 块 GPU 上训练基于 LLaMA 的模型时，Model FLOPs Utilization（MFU）最高可达 52%，相较于 MiCS 和 ZeRO++ 等 newly proposed systems，训练吞吐量提升了 1.56 倍。
引用
@article{arxiv.2311.00257,
title = {AMSP: Reducing Communication Overhead of ZeRO for Efficient LLM Training},
author = {Qiaoling Chen and Qinghao Hu and Guoteng Wang and Yingtong Xiong and Ting Huang and Xun Chen and Yang Gao and Hang Yan and Yonggang Wen and Tianwei Zhang and Peng Sun},
journal= {arXiv preprint arXiv:2311.00257},
year = {2024}
}