共计 3695 个字符，预计需要花费 10 分钟才能阅读完成。
简介：最近，阿里云 PAI 团队和达摩院智能计算实验室一起公布“低碳版”巨模型 M6，大幅升高万亿参数超大模型训练能耗。借助咱们自研的 Whale 框架仅应用 480 卡 GPU，即训练出了规模达人类神经元 10 倍的万亿参数多模态大模型 M6，与传统海内公司实现万亿参数规模相比，能耗升高超八成、效率晋升近 11 倍。
作者 | 王林
起源 | 阿里技术公众号
最近，阿里云 PAI 团队和达摩院智能计算实验室一起公布“低碳版”巨模型 M6，大幅升高万亿参数超大模型训练能耗。借助咱们自研的 Whale 框架仅应用 480 卡 GPU，即训练出了规模达人类神经元 10 倍的万亿参数多模态大模型 M6，与传统海内公司实现万亿参数规模相比，能耗升高超八成、效率晋升近 11 倍。
M6 是国内首个实现商业化落地的多模态大模型。M6 领有超过传统 AI 的认知和发明能力，善于绘画、写作、问答，在电商、制造业、文学艺术等诸多畛域领有广泛应用前景。
这里来为大家介绍反对万亿参数模型训练的 Whale 框架设计。
一 模型发展趋势和挑战
1 模型发展趋势
随着深度学习的火爆，模型的参数规模也增长迅速，OpenAI 数据显示：
- 2012 年以前，模型计算耗时每 2 年增长一倍，和摩尔定律保持一致；
- 2012 年后，模型计算耗时每 3.4 个月翻一倍，远超硬件倒退速度；
近一年模型参数规模飞速增长，谷歌、英伟达、阿里、智源研究院都公布了万亿参数模型，有大厂也公布了百亿、千亿参数模型。同时，随着模型参数规模增大，模型成果也在逐步提高，Nvidia 测试 Bert 模型不同参数规模，发现模型困惑度随模型参数规模减少而升高。
Google 在 GShard paper 中也发现 MoETransformer 模型参数规模越大，翻译品质越高。
2 大模型训练的挑战
大模型带来模型成果晋升的同时，也为训练框架带来更大的挑战，例如当咱们要训练一个万亿规模的模型时会面临如下挑战：
训练难：
- GPU 显存曾经不够寄存模型正本，数据并行曾经不能满足需要；
- 须要框架提供新的并行策略，协同多 GPU 能力来寄存和训练模型；
- 如何给用户提供简洁、易用的接口，让用户能很容易实现分布式版模型；
- 超大规模模型对计算效率、通信效率都带来很大挑战，如何进步计算和通信效率；
- 上游工作如何对接，如何反对批量预测和在线推理需要；
老本高：
- 以万亿模型为例，模型参数有 4TB 大小、梯度也有 4TB，加上 optimizer states 和 active tensor，显存需要微小；
- 业界训练等同规模模型须要的资源：英伟达 3072 A100、谷歌 2048 TPU v3，老本太高很难落地；
- 如何降本增效，应用更少的资源，更快的训练收敛；
以后曾经有一些分布式训练框架，例如：Horovod、Tensorflow Estimator、PyTorch DDP 等反对数据并行，Gpipe、PipeDream、PipeMare 等反对流水并行，Mesh Tensorflow、FlexFlow、OneFlow、MindSpore 等反对算子拆分，但这些框架还有一些有余：
- 模式繁多：很多框架只反对局部并行策略，不能齐全反对各种混合并行；
- 接入门槛高：用户实现模型分布式版本难度大、老本高，须要有领域专家教训能力实现高效的分布式并行策略；
- 迁徙代价大：不同分布式框架并行化实现割裂，不同框架有各自定义的 DSL，当用户要切换并行策略时，须要学习各种接口，从新改写模型；
- 性能不现实：局部框架实现未思考集群物理环境；
为了应答以后分布式训练的挑战，咱们研发了分布式训练框架 Whale，次要指标是：
- 对立多种并行策略：在一个框架中反对各种并行策略以及这些策略的各种组合；
- 简洁易用的接口：用户只需增加几行 annotation 即可实现并行策略的配置，模型代码不须要改变；
- 高效的训练框架：联合硬件资源、网络拓扑和模型进行协同优化，打造高效分布式训练框架；
二 PAI 自研 Whale 框架
1 Whale 架构
咱们推出对立多种并行策略的高性能分布式训练框架 Whale，从如下角度来应答分布式训练的挑战：
- 将不同并行化策略进行对立形象、封装，在一套分布式训练框架中反对多种并行策略；
- 基于 Tensorflow 设计一套分布式并行接口，齐全兼容 Tensorflow，用户仅仅只需增加几行 annotation 就能够实现丰盛的分布式并行策略；
- 联合模型构造和网络拓扑进行调度和通信优化，提供高效的分布式训练能力。
Whale 框架如下图所示，次要分 4 个模块：
- API：提供简洁易用接口，让用户组合应用各种混合并行策略；
- Whale IR：将并行策略转成外部表白，通过 TaskGraph、Multi-Dimension、VirtualDevices 形象来表白各种并行策略；
- Whale Engine：基于 WhaleIR，通过图编辑工具来构建分布式执行图；
- Runtime：将分布式执行图转成 TFGraph，再调用 TF 的 Runtime 来执行；
2 Whale 简介易用接口
Whale 提供简洁易用的接口来形容各种并行策略，次要的原语：
- cluster：配置 Virtual Device 的划分办法
- replica：数据并行
- stage：划分 TaskGraph
- pipeline：流水并行
- split：算子拆分
用这些接口能够组合各种并行策略，例如：
数据并行：
流水并行：
流水并行 + 数据并行：
更多并行策略示例：
3 Whale 训练流程
应用 Whale 进行分布式训练流程：
并行策略配置：
- 应用 Whale API 来为模型配置并行策略，只需增加几行 annotation，无需批改模型代码，办法如 2.2 节 所示；
- 能够将模型划分为多个 TaskGraph，TaskGraph 反对配置多个并行策略，每个 TaskGraph 能够配置不同的并行策略；
虚构资源划分：
- 按并行策略来划分 Virtual Device，每个 TaskGraph 对应一个 Virtual Device；
- 按 GPU 资源和网络 topo 来为 Virtual Device 抉择 Physical Device；
分布式执行图：
- 基于并行策略和资源分配信息，应用图编辑工具来编辑执行图（图拷贝、拆分、插入通信节点等），生成最终的分布式执行图；
- 调用 TF 的 runtime 来执行分布式 Graph；
三 万亿 M6 模型预训练
万亿模型的算力需要十分大，为了升高算力需要，Whale 中实现了 MoE(Mixture-of-Experts)构造，MoE 的次要特点是稠密激活，应用 Gating(Router)来为输出抉择 Top k 的 expert 进行计算（k 罕用取值 1、2），从而大大减少算力需要。
Whale 中实现了 MoE(Mixture-of-Experts) layer，并反对专家并行，将 experts 拆分到多个 Devices 上，升高单个 Device 的显存和算力需要。同时数据并行有利于晋升训练的并发度，因而采纳数据并行 + 专家并行组合的混合并行策略来训练 M6 模型：MoElayer 采纳专家并行，其余 layer 采纳数据并行。
Whale 中提供简洁易用的接口来进行模型的混合并行训练，只须要减少几行 annotation 来配置并行策略，模型自身不须要任何批改。M6 模型采纳数据并行 + 专家并行的策略，只须要减少如下图的 annotation：
同时为了节约训练资源，进步训练效率，Whale 中提供各种优化技术：
显存优化：
- Auto Gradient Checkpoint，主动抉择最优 checkpoint 节点，节约 activation 的显存；
- Group-wise Apply，优化 Optimizer Apply 阶段的显存；
- CPU Offload 技术，优化 Optimizer status 和 Weight 的显存；
- 通信池化，管制通信的数据块大小和并发，节约通信的显存；
计算、通信减速：
- 采纳 DP+EP 混合并行策略，升高算力需要；
- 采纳分组交融通信、半精度通信、拓扑感知的 All2All 通信算子等技术来进步通信效率；
- 联合混合精度、编译优化等技术进步训练效率；
借助 Whale 框架，首次在 480 V100 上，3 天内实现万亿 M6 模型的预训练。相比此前英伟达应用 3072 A100 GPU 实现万亿参数、谷歌应用 2048 TPU 实现 1.6 万亿参数大模型，此次达摩院仅应用 480 卡 V100 32G GPU 就实现了万亿模型 M6，节俭算力资源超 80%，且训练效率晋升近 11 倍。
四 结语
模型参数规模已越来越大，大模型已成为发展趋势，为解决超大模型训练的挑战，咱们自研 Whale 框架，将不同并行化策略进行对立形象、封装，在一套分布式训练框架中反对多种并行策略。Whale 提供简洁易用的接口，用户只需增加几行 annotation 即可实现各种并行策略，不须要对模型自身进行批改。同时咱们联合硬件资源、网络 topo、模型进行软硬件协同优化，提供高效分布式训练框架。
通过 Whale 框架，咱们用 480 V100 GPU 卡训练万亿规模模型，并在 3 天内实现模型训练收敛，为超大规模模型训练落地提供了可能，后续咱们会进一步欠缺 Whale 框架，从更大规模、更快速度、更高性价比 3 个维度去扩大 Whale 框架的能力。同时也会推动 Whale 能力在更多业务场景落地，让技术能力到产品能力的转变。
原文链接
本文为阿里云原创内容，未经容许不得转载。