近日，蚂蚁集团宣布开源大模型分布式训练加速扩展库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全图编译优化；更多分布式优化策略细节连载等。
ATorch开源地址： https://github.com/intelligent-machine-learning/dlrover
更多内容，点击下方关注：
继续阅读
阅读原文
关键词
大模型
模型训练
分布式训练
大模型训练
深度学习