一个悄然崛起的国产开源项目-技术圈

转自 GitHubDaily

大家好，我是小 G。

此前，我曾在公众号推荐过 GitHub 上一个性能颇为强悍的国产 AI 开源项目：Colossal-AI，作为一个大规模并行 AI 训练系统、深度学习框架的内核，该项目可帮助用户便捷实现最大化提升 AI 部署效率，同时最小化部署成本。

仅需一半数量的 GPU，便能完成相同效果的 GPT-3 训练工作，极大降低了项目研发成本。

GitHub 地址：https://github.com/hpcaitech/ColossalAI

Colossal-AI 一经开源便受到广泛关注，连续多日登上 GitHub Python 方向世界第一的位置，与众多已有数万 star 的明星开源项目一起受到海内外关注！

世界顶尖的人工智能机构 Hugging Face 也主动与 Colossal-AI 接洽合作，对于一个国产项目来说，受到如此关注，这也从侧面反映出了它的实力。

近日，经过开发者们的不断努力后，Colossal-AI 在数月的密集测试后迎来正式版！此版本由 300 多次 commits 组成。

本次正式版更新重点优化了分布式训练性能及开发者的易用性，主要亮点包括：

重构 ZeRO 以改善性能和易用性；
添加细粒度 Profiler TensorBoard 监控插件，监测训练过程中内存、网络等状态；
更灵活的 checkpoint 策略，可扩展的 pipeline 模块；
开源蛋白质预测 FastFold 等丰富行业解决方案；
添加中文教程，PaLM、MOE、BERT等实例，开放用户社群及论坛。

专业助力大模型训练

近年来，随着深度学习的兴起及大模型横扫各大性能榜单，前沿 AI 模型的大小在短短几年内便已增大万倍，远超硬件数倍的缓慢增长。前沿 AI 大模型不仅远超单个 GPU 的容纳能力，所需算力也往往需要单个 GPU 运行数百甚至上千年。

因此，如何提升单个 GPU 的容纳能力，如何高效利用分布式技术，联合多个 GPU 低成本实现并行训练加速已成为 AI 大模型的关键痛点。

针对现有方案并行维度有限、效率不高、通用性差、部署困难、缺乏维护等痛点，Colossal-AI 通过高效多维并行、显存优化、大规模优化库、细粒度监测等方式，让用户仅需极少量修改，即可高效快速部署 AI 大模型训练。

多维并行

相比现有方案中并行维度仅包括数据并行、一维张量并行、流水并行三种方案，Colossal-AI 进一步提供 2/2.5/3 维张量并行和序列并行，以及便捷的多维混合并行解决方案。

△ViT 张量并行为 64 时，可提升 14 倍批大小和 5 倍训练速度

其中，高维张量并行能极大减轻显存消耗，提升通信效率，使得计算资源利用更加高效。

△序列并行帮助 BERT 提升 2 倍训练速度，或 1.5 倍序列长度

而序列并行针对大图片、视频、长文本、长时间医疗监测等数据，可以帮助突破原有机器能力限制，直接处理长序列数据。

显存优化

Colossal-AI 综合了多重显存优化技术，包含多维并行，ZeRO 冗余内存消除，CPU offload，Gradient Checkpoint，自动混合精度（AMP）等前沿技术，最大限度帮助用户避免显存瓶颈，降低训练的硬件需求。

△GPT-2 使用 Colossal-AI，同样硬件下提升 24 倍可训练模型大小，或 3 倍训练速度

灵活易用

Colossal-AI 接口设计与 PyTorch 风格保持一致，降低学习和使用成本，仅需极少量修改，便可将已有项目与 Colossal-AI 结合，便捷扩展至大规模并行。此外，该系统还保持了优秀的扩展性，便于根据需求添加新功能，与已有功能模块兼容。

细粒度监测：细粒度 Profiler TensorBoard 插件，相较于 PyTorch 仅能以 iteration 为单位进行记录训练过程，Colossal-AI 能够监测 iteration 内的网络、通信、内存等状态，方便开发者进行精确分析和调试，提高开发效率。

大规模优化库：Colossal-AI 提供大规模并行优化器 LAMB、LARS 等，首次将训练 batch size 扩展到 65536。Colossal-AI 还与 PyTorch 自带各类 optimizer 兼容，并不断探索添加最新前沿优化技术，满足各类模型需求。

丰富的行业解决方案

Colossal-AI 目前已与自动驾驶、云计算、零售、医药、芯片等行业知名厂商达成合作，与 AI 领域顶级开源组织 Hugging Face 等建立合作。

蛋白质结构预测加速方案：FastFold

AlphaFold 因强大的 AI 预测蛋白质结构能力被 Science 和 Nature 评选为 2021 年十大科学突破之首，但存在训练时间长、成本高等问题。

△图片来源：https://arxiv.org/pdf/2203.00854.pdf

基于 Colossal-AI 的加速方案 FastFold，将 GPU 优化和大模型训练技术引入 AlphaFold 的训练和推理，成功超越谷歌和哥伦比亚大学的方案，将 AlphaFold 训练时间从 11 天减少到 67 小时，且总成本更低，在长序列推理中也实现 9.3～11.6 倍的速度提升。

△长序列推理性能对比

半数 GPU 训练 GPT-3

对于超大 AI 模型，如 GPT-3，相比英伟达方案，Colossal-AI 仅需一半的计算资源，即可启动训练；若使用相同计算资源，则能提速 11%，可降低 GPT-3 训练成本超百万美元。

Colossal-AI 注重开源社区建设，提供中文教程，开放用户社群及论坛，对于用户反馈进行高效交流与迭代更新，不断添加 PaLM、MOE 等前沿应用。

项目团队

潞晨技术团队的核心成员均来自美国加州大学伯克利分校，斯坦福大学，清华大学，北京大学，新加坡国立大学，新加坡南洋理工大学等国内外知名高校；拥有 Google Brain、IBM、Intel、 Microsoft、NVIDIA 等知名厂商工作经历。公司成立即获得创新工场、真格基金等多家顶尖 VC 机构种子轮投资。

△潞晨科技创始人尤洋教授：加州大学伯克利分校博士、IPDPS/ICPP 最佳论文、ACM/IEEE George Michael HPC Fellowship、福布斯 30 岁以下精英 (亚洲 2021)、IEEE-CS 超算杰出新人奖、UC 伯克利 EECS Lotfi A. Zadeh 优秀毕业生奖

△潞晨 CSO Prof. James Demmel：加州大学伯克利分校杰出教授、ACM/IEEE Fellow，美国科学院、工程院、艺术与科学院三院院士

传送门

论文地址：https://arxiv.org/abs/2110.14883

项目地址：https://github.com/hpcaitech/ColossalAI

文档地址：https://www.colossalai.org/

* 本文观点参考链接：

https://medium.com/@hpcaitech/5-must-follow-features-that-are-seeing-colossal-ais-success-2d5361e27e4b