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