千卡规模训练算力利用率达 60%，蚂蚁开源分布式训练加速扩展库 ATorch

近日，蚂蚁集团宣布开源大模型分布式训练加速扩展库ATorch。ATorch可针对不同模型和硬件资源，实现深度学习自动资源动态优化和分布式训练稳定性提升，可帮助优化深度学习的智能性，解决大模型训练的提效问题。据了解，ATorch大模型训练算力利用率可达60%，对于千亿模型千卡级训练提效非常友好，相当于为跑车装上了强劲的引擎。

Meta开源的 PyTorch和谷歌开源的TensorFlow，是最受开发者欢迎的两大深度学习框架。凭借其简单易用、功能强大、用途广泛等特点，开发者和研究人员可以轻松的构建和训练模型。随着生成式大模型的爆发，模型训练的数据集和参数规模成指数级增长。要带动如此庞然大物，并且满足模型的快速迭代，分布式训练就成为了解题之道。本次蚂蚁开源的ATorch针对于大模型训练场景，提供了基于 PyTorch 的高性能解决方案。

1

为跑车装上强劲引擎

2023年上半年，蚂蚁集团开源了DLRover项目，这是基于云原生技术打造的智能分布式深度学习系统，相当于为一辆跑车提供了自动驾驶系统。那么如何让跑车跑得更快呢？蚂蚁开发了基于PyTorch的分布式训练加速扩展库ATorch，并且集成到DLRover中，将更大程度上帮助大模型开发能够专注于模型架构的设计，而无需处理工程方面的细节，让训练更高效、智能。这就相当于为跑车装上了一个强劲的引擎。

ATorch设计的初衷是，让深度学习和大模型训练继续拥抱成熟活跃的PyTorch生态，并且优雅组合不同开源方案，任意添加针对自身硬件与业务特点新策略，同时针对不同的模型和硬件资源来配置出最优的策略。

ATorch采用了分层的架构设计，功能清晰、设计全面，可为开发者提供极致精简的开发体验，以及领先的稳定性保障。主要包括统一分布式优化策略配置接口、自动分布式策略搜索、自动弹性容错、高效动态显存管理库、自研优化器加速收敛等核心功能。作为PyTorch框架的高性能扩展加速库，ATorch可实现最少化用户代码侵入，可为千亿参数大模型千卡级提供高性能预训练实践方案。

1）整体设计

（ATorch系统架构）

整体上ATorch采用了分层的架构设计：

• Interface层负责用户使用接口，支持极致精简的one api模式和trainer模式。
• IO和预处理层主要负责存储IO等相关的优化。
• 弹性容错层主要负责稳定性保障，主要配合DLRover使用来实现。
• 核心层是ATorch的主要代码，这里定义了优化策略的统一接入和表达，保障开源方法和自研策略统一管理，同时也包含了自动策略最优策略组合的功能。
• 最下层分别是PyTorch社区生态系统和底层的AIDC硬件。

2）核心功能

• 大模型训练关键关键技术

1.千亿参数大模型千卡级高性能预训练实践方案（算力利用率 > 50%）

2.显存/通信联合优化

3.Flash Attention/Fused ops 等算子优化

4.千亿参数模型快速ckpt save/load的支持

5.支持千亿级模型RLHF高效微调

• 统一分布式优化策略配置接口

1.分布式算子，通信，显存优化

2.数据IO预处理优化

• 自动分布式策略搜索：

给定单机模型定义和硬件资源描述，自动搜索最优配置组合

• ATorch+DLRover自动弹性容错，提升大模型训练稳定性（有效训练时长大于95%）
• GLake高效动态显存管理库，显著提升batch size和端到端吞吐
• 国产硬件生态适配调优
• 自研优化器 AGD（加速收敛）和 WSAM（提升泛化，已被 国际顶会KDD '23 接收）

2

千亿千卡模型训练

算力利用率达60%

近期在以开源模型为对象的大模型训练优化实践中，ATorch取得了亮眼的成绩。

1）主要开源大模型预训练算力吞吐提效

• 清华大学GLM-65b：（参数量：65B，A100*1536）
• Meta LLama2-70b（参数量：70B，H800*1536）
• 多模态大模型Stable Diffusion：（参数量：10B，A100*256）
• 视觉生成模型：（参数量：2B，A100*128）

2）ATorch训练加速比

以antglm10B为实验模型，从1机8卡测试到32机256卡，基线为模型官方提供的Megatron+deepspeed实现

3）ATorch千卡训练稳定性

关键指标：

• 日均纯训练时常占比：40.7% -> 95%
• ckpt save耗时：10min -> 1min
• 训练重启耗时：90 min -> 5 min

后续，ATorch将在功能上进行更多升级。包括ATorch RLHF，支持千亿模型高效指令微调，time/seq相对基线4x提升、ATorch大模型快速容错，ckpt 高速保存恢复方案；国产卡适配及优化；Lynx全图编译优化；更多分布式优化策略细节连载等。