[论文解读] MiCS: Near-linear Scaling for Training Gigantic Model on Public Cloud

Zhen Zhang, Shuai Zheng+6 more|arXiv (Cornell University)|Apr 30, 2022

Stochastic Gradient Optimization Techniques被引用 6

一句话总结

MiCS 通过基于规模感知的模型划分、分层通信和两跳梯度同步，在公共云上实现了对超大深度学习模型的近线性可扩展训练，最大限度地减少了通信规模。在 512 个 GPU 上训练 1000 亿参数模型时，其吞吐量最高提升 2.89 倍，弱扩展效率达 99.4%，在异构、低带宽的云环境中优于 DeepSpeed ZeRO。

ABSTRACT

Existing general purpose frameworks for gigantic model training, i.e., dense models with billions of parameters, cannot scale efficiently on cloud environment with various networking conditions due to large communication overheads. In this paper, we propose MiCS, which Minimizes the Communication Scale to bring down communication overhead. Specifically, by decreasing the number of participants in a communication collective, MiCS can utilize heterogeneous network bandwidth, reduce network traffic over slower links, reduce the latency of communications for maintaining high network bandwidth utilization, and amortize expensive global gradient synchronization overhead. Our evaluation on AWS shows that the system throughput of MiCS is up to 2.89$ imes$ that of the state-of-the-art large model training systems. MiCS achieves near-linear scaling efficiency, which is up to 1.27$ imes$ that of DeepSpeed. MiCS allows us to train a proprietary model with 100 billion parameters on 512 GPUs with 99.4% weak-scaling efficiency, and it is able to saturate over 54.5% theoretical computation power of each GPU on a public cloud with less GPU memory and more restricted networks than DGX-A100 clusters.

研究动机与目标

• 解决现有大规模模型训练框架在异构、低带宽网络的公共云环境中部署时面临的高通信开销问题。
• 提升在缺乏高速 InfiniBand 互连的公共云集群上的训练吞吐量和可扩展性，与 DGX 集群形成对比。
• 在无需专用硬件或复杂模型并行化的标准公共云 GPU 实例上，实现对超大模型（如 1000 亿参数）的高效训练。
• 在大规模异构集群中显著降低通信成本的同时，保持数据并行的简洁性和通用性。
• 即使在网络受限条件下，也能实现与高端集群相当的近线性扩展效率。

提出的方法

• 提出规模感知的模型划分方法，将 GPU 划分为更小的、隔离的组，每组持有模型状态的完整副本，从而减少频繁参数同步中的参与方数量。
• 采用分层通信策略，将集体通信操作限制在组内，并行化组间通信，降低在较慢节点间链路上的流量。
• 设计两跳梯度同步机制，通过先在组内聚合梯度，再在组间同步组级别汇总结果，避免全局 all-reduce 操作。
• 根据 GPU 内存容量和网络拓扑结构配置分组，平衡内存使用与通信效率，将组大小作为可调参数。
• 重用现有数据并行框架（如 ZeRO），但重构通信模式，以利用节点内高带宽链路，减少节点间延迟和数据量。
• 使用启发式方法确定模型状态复制的最优组大小，确保内存效率并最小化通信瓶颈。

实验结果

研究问题

• RQ1 数据并行训练系统是否能在异构、低带宽网络的公共云集群上实现近线性扩展？
• RQ2 在不依赖 InfiniBand 等专用硬件的情况下，如何减少大规模模型训练中的通信开销？
• RQ3 通过最小化通信参与方数量，能在多大程度上提升大规模 GPU 集群中的训练吞吐量和可扩展性？
• RQ4 系统能否在避免复杂模型并行或流水线并行的同时，保持高效率和低通信成本？
• RQ5 在真实公共云环境约束下（如 GPU 内存有限、网络带宽波动），系统表现如何？

主要发现

• MiCS 在 AWS 上的系统吞吐量最高比 DeepSpeed ZeRO 提升 2.89 倍，显著提升了大规模训练性能。
• 在 512 个 GPU 上训练 1000 亿参数模型时，系统弱扩展效率达 99.4%，表明具备近线性可扩展性。
• MiCS 在公共云实例上使每个 GPU 的理论计算能力利用率超过 54.5%，且单 GPU 内存使用率低于 32%，在相同条件下优于 ZeRO。
• 分层通信策略降低了节点间网络流量和延迟，提升了在较慢节点间链路上的带宽利用率。
• 两跳梯度同步机制通过将 all-reduce 操作限制在组级别汇总结果，减少了全局同步开销，提升了大规模集群中的效率。
• MiCS 即使在 GPU 内存有限和网络带宽受限的条件下仍保持高性能，适用于标准公共云部署。