新加坡国大联合国内企业发布Colossal-AI，为在多个GPU或机器上进行分布式训练模型提供统一并行化系统

在深度学习领域，尺寸不断增长的大型模型已经为计算机视觉和自然语言处理等各个领域带来了巨大且飞速的性能提升。例如，BERT 和 Vision Transformer 等神经网络模型凭借着巨量的数据，已经可以掌握高维特征、甚至可以达到人类无法实现的预测水准。

这些模型不仅架构变得更加丰富，近年来其参数数量更是突飞猛进，而且更新速度不断加快。2020 年，GPT-2 超越此前规模最大的模型 BERT-Large 只用了 3 个月的时间，并且前者的参数数量是后者的约 5 倍。

而 GPT-2 并没有停止进化的脚步，就在同一年便升级为参数量达 1750 亿的 GPT-3。但是，GPT-3 很快就被其他模型追赶上来，2021 年发布的 GLM 拥有 1.75 万亿参数，将 GPT-3 远远甩在后面。

尺寸越来越大的模型带来了愈来愈佳的性能，此外也带来了对 GPU 等硬件愈来愈高的挑战。

来自英国阿伯丁大学（University of Aberdeen）研究员詹姆斯·塞维利亚（Jaime Sevilla）等人曾将机器学习的发展分为三个阶段：前深度学习时代、深度学习时代和大规模训练时代，随之而来的则是不断增长的计算需求。

图｜不同年代发布的机器学习的算力需求（纵轴）与发布时间（横轴）（来源：arXiv）

随着大型模型对算力要求的不断提高，仅仅靠单个 GPU 或单个机器来对模型进行训练已经逐渐达到瓶颈。为应对该问题，使用多个 GPU 或机器来进行模型训练的分布式训练，已被公认为是“未来的必然趋势之一”。

不过，在多个硬件上进行并行式训练的分布式训练仍然存在许多挑战，尤其是大型模型通常十分复杂，且不同的模型适合使用不同的并行方法。

目前，深度学习的分布式训练有几种主流的并行方式：数据并行方法在深度学习中最为常用，其中先将数据在批处理维度中进行拆分、再将其模型副本输入到不同硬件中进行处理。

不过，数据并行方法的一个缺点是处理过程中的内存并未被充分利用。之后，有其他方法通过引入管道并行和张量并行方法，可以有效地检查内存冗余，并且在提高速度的同时实现更高参数量的模型训练。

例如，微软的 DeepSpeed 曾提供一种可以同时组合使用这三种并行方法的训练工具，可训练的模型的参数量高达十亿至万亿量级。

（来源：GitHub）

而英伟达的 Megatron-LM 更提供了开源代码，让开发者跨多个 GPU 训练数十亿参数的模型更为简单。不过，Megatron-LM 由于引入了更多的复杂性，对于系统研究经验较少的开发人员来说，由于必须要考虑通信、一致性和扩展效率等因素而不易上手。

为了让 AI 社区开发者们更容易地使用分布式训练以及各种并行方法，潞晨科技与新加坡国立大学团队发布了 Colossal-AI 系统。相关论文以《机器学习的三个时代的算力需求趋势》（Compute Trends Across Three Eras of Machine Learning）为题发表在 arXiv[1]。

该系统基于 PyTorch 并且开源，它提供了统一的并行训练系统，目的是为了开发者们可以更容易地实现数据并行、管道并行、张量并行和序列并行等将不同种类并行化技术的无缝集成。

对于社区开发者来说，通过使用 Colossal-AI 系统，他们在开发分布式模型时可以使用之前编写模型相同的方法，从而可以将精力更多地集中在模型架构的开发上。并且，在开发模型时可以将开发过程与分布式训练需要关注的内容分开考虑。

具体来说，社区开发者在通过 Colossal-AI 系统进行分布式训练时，通过简单易懂的 API，就可以与之前构建单 GPU 中的模型一样，运用并行方法和构建分布式模型。

此外，为了使该系统实现最高限度的鲁棒性和通用性，Colossal-AI 系统还结合了目前的其他系统中的方法，比如微软 DeepSpeed 的零冗余优化器和卸载、英伟达 Megatron-LM 中的的一维张量并行性等。

图｜1D、2D、2.5D 和 3D 张量并行（来源：arXiv）

Colossal-AI 系统的一大特点是其提供的完备的并行技术库，除了常用的数据和管道并行方法之外，该系统还提供了一系列张量并行训练方法。

尤其是针对单纯一维张量并行方法存在的内存冗余问题，该系统提供了 2D、2.5D 和 3D 等不同种类的张量并行方法，甚至可以完全消除内存冗余。

自从 Colossal-AI 开发以来，该团队就坚持模块化的设计原则。这也使得 Colossal-AI 中的不同模块可以彼此解耦，从而使得其在未来可以更快地适应新趋势，通过最少的代码改变就能针对未来所谓的新功能进行扩展。

图｜将 Colossal-AI 在 CIFAR10 数据集上进行收敛性验证的结果（来源：arXiv）

为了验证 Colossal-AI 系统的并行方法是否正确及有效，研究人员将该系统在 CIFAR10 数据集上进行了实验。结果显示，通过该算法可以通过更低成本得到有效的效果。

参考资料：

1.https://doi.org/10.48550/arXiv.2202.05924