近日,蚂蚁集团宣布开源大模型分布式训练加速扩展库ATorch。
ATorch可以实现深度学习的自动资源动态优化,提高不同模型和硬件资源的分布式训练的稳定性,有助于提高深度学习的智能性。
据了解,在大模型训练中,ATorch千亿模型千卡级别训练的算力利用率可达60%,相当于在跑车上安装了强大的发动机。
随着大型生成模型的爆炸式增长,模型训练的数据集和参数的大小呈指数级增长。
为了驱动如此庞大的模型并满足模型的快速迭代,分布式训练成为了问题的解决方案。
在大型模型的开发中,许多开发人员和研究人员使用 PyTorch 和 TensorFlow 深度学习框架来构建和训练模型。
业界已经开展了大量的工作,为PyTorch等深度学习框架配备更适合大型模型训练的“工具包”。
此次蚂蚁开源的ATorch就是其中之一。
据了解,ATorch采用分层架构设计,功能清晰、设计全面,可以为开发者提供极其精简的开发体验和领先的稳定性保障。
主要包括统一分布式优化策略配置接口、自动分布式策略搜索、自动弹性容错、高效动态内存管理库、自研优化器加速收敛等核心功能。
ATorch作为PyTorch框架的高性能扩展加速库,可以最大程度地减少用户代码入侵,为千亿参数大型模型的千卡级训练提供易用、高性能的解决方案。
在近期针对开源模型的大模型训练优化实践中,ATorch取得了优异的成绩。
例如,清华大学开源的GLM-65b大模型的预训练算力利用率从28.8%提升到62%,清华大学开发的LLama2-70b大模型的预训练算力利用率元从 42% 增加到 60%。
英国AI公司Stability AI开发的大型多模态模型Stable Diffusion的训练算力利用率从21.8%提升至58.7%。
此外,ATorch具有出色的大卡训练稳定性,日均纯训练时间提升至95%,CKPT保存时间控制在1分钟,训练重启时间最快5分钟,达到行业领先水平。
目前,ATorch已经集成到蚂蚁集团的大模型训练开源产品DLRover中,DLRover是蚂蚁集团基于云原生技术打造的智能分布式深度学习系统。
ATorch的加入将有助于大型模型开发更大程度地专注于模型架构的设计,而无需处理工程细节,使训练更加高效和智能。