Colossal-AI：一个大规模并行AI训练系统的正式版发布

Colossal-AI是一个旨在提升AI部署效率，最小化部署成本的大规模并行AI训练系统。经过数月的密集测试，Colossal-AI迎来正式版，更新重点优化了分布式训练性能及开发者的易用性。

本次正式版更新的主要亮点包括：

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

Colossal-AI的优势主要体现在以下几个方面：

专业助力大模型训练

随着深度学习的兴起及大模型横扫各大性能榜单，前沿AI模型的大小在短短几年内便已增大万倍，远超硬件数倍的缓慢增长。Colossal-AI通过高效多维并行、显存优化、大规模优化库、细粒度监测等方式，让用户仅需极少量修改，即可高效快速部署AI大模型训练。

Colossal-AI相比现有方案中并行维度仅包括数据并行、一维张量并行、流水并行三种方案，进一步提供2/2.5/3维张量并行和序列并行，以及便捷的多维混合并行解决方案。高维张量并行能极大减轻显存消耗，提升通信效率，使得计算资源利用更加高效。序列并行针对大图片、视频、长文本、长时间医疗监测等数据，可以帮助突破原有机器能力限制，直接处理长序列数据。

显存优化

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

灵活易用

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

细粒度监测和大规模优化库

Colossal-AI提供细粒度ProfilerTensorBoard插件，相较于PyTorch仅能以iteration为单位进行记录训练过程，Colossal-AI能够监测iteration内的网络、通信、内存等状态，方便开发者进行精确分析和调试，提高开发效率。Colossal-AI还提供大规模并行优化器LAMB、LARS等，首次将训练batchsize扩展到65536。Colossal-AI还与PyTorch自带各类optimizer兼容，并不断探索添加最新前沿优化技术，满足各类模型需求。

丰富的行业解决方案

Colossal-AI目前已与自动驾驶、云计算、零售、医药、芯片等行业知名厂商达成合作，与AI领域顶级开源组织HuggingFace等建立合作。蛋白质结构预测加速方案：FastFoldAlphaFold因强大的AI预测蛋白质结构能力被Science和Nature评选为2021年十大科学突破之首，但存在训练时间长、成本高等问题。基于Colossal-AI的加速方案FastFold，将GPU优化和大模型训练技术引入AlphaFold的训练和推理，成功超越谷歌和哥伦比亚大学的方案，将AlphaFold训练时间从11天减少到67小时，且总成本更低，在长序列推理中也实现9.3～11.6倍的速度提升。对于超大AI模型，如GPT-3，相比英伟达方案，Colossal-AI仅需一半的计算资源，即可启动训练；若使用相同计算资源，则能提速11%，可降低GPT-3训练成本超百万美元。

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

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