[论文解读] Colossal-AI: A Unified Deep Learning System For Large-Scale Parallel Training

Li, Shenggui, Hongxin Liu+6 more|arXiv (Cornell University)|Oct 28, 2021

Topic Modeling参考文献 15被引用 21

一句话总结

Colossal-AI 是一个统一的深度学习系统，通过集成数据并行、流水线并行、张量并行和序列并行，以及异构训练和 ZeRO 优化，实现了高效的大规模分布式训练。它通过动态管理内存并支持灵活、用户友好的序列模型并行化，实现了相较于基线系统的最高 2.76 倍训练加速。

ABSTRACT

The success of Transformer models has pushed the deep learning model scale to billions of parameters. Due to the limited memory resource of a single GPU, However, the best practice for choosing the optimal parallel strategy is still lacking, since it requires domain expertise in both deep learning and parallel computing. The Colossal-AI system addressed the above challenge by introducing a unified interface to scale your sequential code of model training to distributed environments. It supports parallel training methods such as data, pipeline, tensor, and sequence parallelism, as well as heterogeneous training methods integrated with zero redundancy optimizer. Compared to the baseline system, Colossal-AI can achieve up to 2.76 times training speedup on large-scale models.

研究动机与目标

• 解决大规模千亿参数模型分布式训练中缺乏统一、用户友好系统的问题。
• 降低研究人员和工程师在选择最优并行策略时的技术门槛。
• 支持包括序列并行和多维张量并行在内的全面并行方法集合。
• 通过动态分片和卸载技术，提升内存效率和训练吞吐量。
• 在无需架构重写的情况下，实现序列模型代码在多 GPU 集群上的无缝扩展。

提出的方法

• Colossal-AI 提供统一的 API，使用户能够以极少的代码修改将序列模型代码扩展到分布式环境。
• 它集成了数据并行、流水线并行、张量并行和序列并行，支持灵活组合以实现最佳性能。
• 该系统采用多维张量并行技术，减少激活内存瓶颈，提升可扩展性。
• 序列并行将输入张量和激活张量沿序列维度拆分，降低通信和内存开销。
• Colossal-AI 实现了动态张量放置，仅在必要时进行卸载，以提升 GPU 内存利用率。
• 它支持 ZeRO-3 和异构训练，最小化冗余优化器状态，并在混合硬件上加速训练。

实验结果

研究问题

• RQ1 如何在统一的深度学习系统中有效结合多种并行技术以实现最优训练性能？
• RQ2 与一维张量并行相比，序列并行对内存效率和可扩展性有何影响？
• RQ3 在低小批量场景下，动态分片和卸载能否提升训练吞吐量和内存利用率？
• RQ4 与 Megatron-LM 和 DeepSpeed 等最先进系统相比，Colossal-AI 的训练加速表现如何？
• RQ5 在增加流水线阶段数和 GPU 数量时，Colossal-AI 能在多大程度上提升训练吞吐量？

主要发现

• 在大规模模型上，Colossal-AI 相较于基线系统实现了最高 2.76 倍的训练加速。
• 在 12 块 GPU 上，序列并行相比一维张量并行可支持 4.44 倍更大的批量大小和 1.18 倍更长的序列长度。
• 在 4 个流水线阶段下，序列并行相比一维张量并行实现了 1.55 倍的更快训练速度。
• 在低小批量场景下，Colossal-AI 的动态卸载通过避免不必要的 CPU 传输提升了吞吐量，优于 DeepSpeed 的静态策略。
• 在 8 块 GPU 上，使用 13B 参数的 OPT 模型和批量大小 32 时，Colossal-AI 相较于 DeepSpeed 实现了 1.33 倍的加速。
• 当使用线性复杂度注意力机制时，Colossal-AI 支持最大序列长度随 GPU 数量线性扩展，从而实现文档级文本理解。