系统日报-20220318（大模型并行训练框架 Colossal-AI）
系统日报-20220318（大模型并行训练框架 Colossal-AI）
《系统日报》持续关注分布式系统、AI System，数据库、存储、大数据等相关领域文章。每天以摘要的形式精选不超过三篇系统文章分享给大家。
分享一个大牛同学公司的的 AI System 项目 Colossal-AI：A Unified Deep Learning System for Large-Scale Parallel Training[2]。
以GPT3为代表的大深度学习模型是现在很火的技术，Colossal-AI 的目标就是解决大模型训练过程遇到的各种分布式难题。它也算是目前该领域前沿研究的集大成者，很适合想了解MLSys却又不知如何下手的萌新去围观。
最近几年的 AI 模型正在急速变大，训练常常需要需要多个 GPU，比如训练 GPT3 需要几千个 GPU。因此，在多个 GPU 上分布式训练前沿 AI 大模型已经成为业界常态。
提起深度学习训练框架，大家可能首先想到就是 PyTorch 和 TensorFlow。但是已有的这些框架只支持数据并行，如果模型太大无法放在单个 GPU 中，数据并行就无用武之地了。业界常用的做法是手写很多模型切分和通信的逻辑，这需要开发者同时掌握 AI 与分布式专业知识，还需花费大量时间精力debug和调优，严重限制了开发效率。
Colossal-AI 的愿景是让用户仅需少量修改，便可将已有 PyTorch/TensorFlow 项目与 Colossal-AI 结合，快速将单机代码自动、高效地扩展为分布式系统。这里简单介绍一些 Colossal 的技术细节。
深度学习训练周期内需要管理两种数据。模型数据，训练中不断更新的模型参数，也是最后拿去线上推理的模型。Activation(a.k.a. Feature Map)：每一层输出的中间结果，训练过程中每个神经网络层的输出。
Colossal-AI 实现的分布式训练技术包括数据并行、张量并行、流水线并行、ZeRO并行和 offload 并行。从对数据切分角度来简单解释一下他们的原理：
数据并行：每个GPU都存储同样的全局的模型数据，将Activation切分到不同GPU。
张量并行：每个GPU有全局的Activation，将模型切分到不同GPU。
流水线并行：将模型数据按层切分为流水段，相应的也按层切分activation，每个GPU执行一段流水线的计算。
ZeRO并行[3]：将模型数据按照层切分到不同机器上，但是计算模式和数据并行相同，每个 GPU 只计算部分 Activation。不过计算中需要取远端的模型碎片。
Offload并行[4]：类似虚拟内存，训练中可以利用CPU内存来存储溢出的数据。
Colossal-AI 允许这些并行策略进行自由组合，可以实现丰富的混合并行方式。对于不同并行策略，只需通过一个 Python 字典来配置，而不影响训练的业务逻辑。这就屏蔽了底层的复杂度，让使用者更关注于模型本身而不是实现细节。
那么，对于某个具体情形，如何针对性地选择最优并行策略呢？这块可以利用一些自动化方式搜索出一个最优并行配置，但在 Colossal-AI 中似乎还没有实现。
AI System 领域演化和分布式数据库越来越相近，也许未来会有一种数据库专门管理大模型训练的模型数据。他山之石可以攻玉，感兴趣小伙伴可以点击 “阅读原文” 去围观一个。
参考资料
[1]任何想法都欢迎来提 issue: https://github.com/DistSysCorp/ArticleListWeekly/issues
[2]Colossal-AI：A Unified Deep Learning System for Large-Scale Parallel Training: https://github.com/hpcaitech/ColossalAI
[3]Rajbhandari S, Rasley J, Ruwase O, et al. Zero: Memory optimizations toward training trillion parameter models[C]//SC20: International Conference for High Performance Computing, Networking, Storage and Analysis. IEEE, 2020: 1-16: .
[4]Ren J, Rajbhandari S, Aminabadi R Y, et al. {ZeRO-Offload}: Democratizing {Billion-Scale} Model Training[C]//2021 USENIX Annual Technical Conference (USENIX ATC 21). 2021: 551-564.: .