ColoSSal-AI是一款大规模并行AI训练系统,旨在作为深度学习框架的核心,帮助用户高效提升AI部署的效率,降低成本。

开源地址:GitHub
自开源以来,ColoSSal-AI便引起广泛关注,连续多日登顶GitHub热榜Python类别第一,与众多备受瞩目的开源项目一同受到全球关注。
经过开发者们的辛勤努力,ColoSSal-AI在经过数月的密集测试后,迎来了正式版发布,此版本包含300多次提交。

此次正式版更新集中于优化分布式训练性能和提升开发者的使用体验,主要亮点包括:
- 重构ZeRO以改善性能和易用性;
- 新增细粒度ProfileR TensoRBOARd监控插件,能够监测训练过程中的内存和网络状态;
- 更灵活的checkpoint策略和可扩展的pipeline模块;
- 开源蛋白质预测方案FAstFold等多种行业解决方案;
- 提供中文教程,包括MOE、BERT等实例,开放用户社群和论坛。
近年来,随着深度学习的快速发展,大模型的规模在短时间内迅速增长,远超硬件的进步速度。前沿的AI大模型不仅超出了单个GPU的处理能力,所需的计算资源往往需要持续数百甚至上千年才能完成。
因此,如何提升单个GPU的容量,以及如何高效利用分布式技术联合多个GPU以低成本实现并行训练加速,已成为AI大模型面临的关键挑战。

针对现有方案中并行维度有限、效率低下、通用性差、部署困难和缺乏维护等问题,ColoSSal-AI通过多维并行、显存优化和大规模优化库等手段,让用户在进行极少量修改的情况下即可快速高效地部署AI大模型训练。
与现有方案仅支持数据并行、一维张量并行和流水并行不同,ColoSSal-AI进一步支持2/2.5/3维张量并行和序列并行,并提供便捷的多维混合并行解决方案。

其中,高维张量并行能够显著减少显存消耗,提升通信效率,实现更高效的计算资源利用。
序列并行则可以帮助BERT提高2倍的训练速度,或支持1.5倍的序列长度。
针对大图片、视频和长文本等数据,序列并行可以突破机器的能力限制,直接处理长序列数据。

ColoSSal-AI结合了多种显存优化技术,包括多维并行、ZeRO冗余内存消除、CPU offload、Gradient Checkpoint和自动混合精度(AMP),最大限度地帮助用户避免显存瓶颈,降低训练硬件需求。

ColoSSal-AI的接口设计与PyTorch风格保持一致,降低了学习和使用成本,只需少量修改即可将现有项目与ColoSSal-AI结合,方便扩展至大规模并行。此外,该系统还保持了良好的扩展性,便于根据需求添加新功能并与现有模块兼容。
细粒度监测方面,ColoSSal-AI的ProfileR TensoRBOARd插件相比于PyTorch的ITeRation监测能力更强,能够在ITeRation内监测网络、通信和内存状态,便于开发者进行精准分析和调试,从而提高开发效率。
大规模优化库方面,ColoSSal-AI提供了大规模并行优化器如LAMB和LARS,首次将训练批量大小扩展到65536,同时与PyTorch自带的各类优化器兼容,并不断探索添加新的前沿优化技术,以满足各种模型需求。
ColoSSal-AI已与自动驾驶、云计算、零售、医药、芯片等多个行业知名企业展开合作,并与AI领域顶尖开源组织建立合作关系。
在蛋白质结构预测方面,基于ColoSSal-AI的加速方案FAstFold,成功将GPU优化和大模型训练技术应用于AlphaFold的训练和推理,显著减少了训练时间和成本。

FAstFold将AlphaFold的训练时间从11天缩短至67小时,并在长序列推理中实现了9.3至11.6倍的速度提升。

对于超大AI模型如GPT-3,ColoSSal-AI所需计算资源仅为英伟达方案的一半;若使用相同计算资源,则能够提升训练速度11%,降低GPT-3的训练成本超过百万美元。
ColoSSal-AI重视开源社区建设,提供中文教程,开放用户社群与论坛,积极响应用户反馈并进行迭代更新,不断增加MoE等前沿应用。
潞晨技术团队的核心成员均来自多所知名高校,拥有在Google、IBM、Intel、Microsoft、NVIDIA等公司的工作经历,且公司成立初期便获得多家顶尖投资机构的支持。

论文地址:arXiv
项目地址:GitHub
文档地址:ColoSSal-AI官网
