近日，阿里巴巴云(AlibabaCloud ) PAI团队与达摩医院智能计算实验室一起推出了“低碳版”巨型机型M6，大大降低了万亿参数超大机型的训练能耗。 利用我们自研的Whale框架只使用480卡的GPU，训练了规模是人类神经元10倍的万亿参数多模大模型M6。 与传统的海外公司实现万亿参数规模相比，能源消费量超过8成，效率提高了近1.1倍。
M6是国内首个实现商业化落地的多模态大机型。 M6拥有超越传统AI的认知和创造力，擅长绘画、写作、答疑，在电子商务、制造业、文艺等多个领域具有广泛的应用前景。
本节介绍了支持万亿参数模型培训的Whale框架的设计。
模特的趋势和挑战1 模型发展趋势
随着深度学习的爆发，模型的参数规模也迅速增加。 根据OpenAI数据，如下所示。
2012年以前，模型计算的时间每两年翻一倍，与摩尔定律一致； 从2012年开始，模型计算每3.4个月需要两倍的时间，远远超过硬件的发展速度
近一年模型参数规模快速增长，谷歌、NVIDIA、阿里、智源研究院均发布万亿参数模型，一些大工厂也发布百亿、千亿参数模型。 同时，随着模型参数规模的增大，模型效果也逐渐提高，Nvidia测试了Bert模型的不同参数规模，发现模型的困惑度随着模型参数规模的增加而降低。
谷歌还在GShard paper上发现，MoETransformer模型参数越大，翻译质量越高。
2 大模型训练的挑战
大模型在带来模型效果提高的同时，也给培训框架带来巨大挑战。 例如，在培训规模为1万亿美元的模型时，您将面临以下挑战：
培训困难： GPU图形存储没有足够的模型副本，数据并行处理无法满足需要； 需要框架提供新的并行策略，并结合多个GPU能力存储和训练模型； 如何为用户提供简洁易用的界面，使用户轻松实现分布式模型； 超大规模模型给计算效率、通信效率带来巨大挑战，如何提高计算和通信效率； 下游任务如何对接，如何支持批量预测和在线推理需求； 成本高：以万亿元模型为例，模型参数有4TB大小，梯度也有4TB，加上优化器状态和活动传感器，培训内存需求明显较大的行业同等规模模型所需资源： 如何降低这个插件，使用更少的资源，更快地收敛训练； 当前支持数据并行处理，如Horovod、Tensorflow Estimator、PyTorch DDP；支持流水并行处理，如Gpipe、PipeDream、pipe mare；mesh tensor flow、FlexFlow
单模：许多框架只支持部分并行策略，不能完全支持各种混合并行； 高访问阈值：用户实现模型的分布式版本难度大、成本高，需要专家经验才能实现高效的分布式并行策略； 迁移成本大：不同分布式框架的并行化实现分割，不同框架有各自定义的DSL，用户切换并行策略时，需要学习各种界面，重写模型； 性能不理想：一些框架的实现没有考虑群集的物理环境； 为了应对当前分布式培训的挑战，开发了分布式培训框架Whale。 主要目标是：
统一多种并行策略：在一个框架内支持各种并行策略和这些策略的不同组合； 简单易用的界面：用户只需添加几行annotation即可完成并行策略配置，无需修改模型代码； 高效培训框架：结合硬件资源、网络拓扑和模型共同优化，构建高效的分布式培训框架； 二PAI自研Whale框架1 Whale架构
我们推出了一个统一多种并行策略的高性能分布式培训框架Whale，从以下角度应对分布式培训挑战：
统一抽象和封装不同的并行化策略，在一系列分布式培训框架中支持多种并行化策略； 基于Tensorflow设计了一组分布式并行接口，与Tensorflow完全兼容，用户只需添加几行annotation即可实现丰富的分布式并行策略； 结合模型结构和网络拓扑进行调度和通信优化，提供高效的分布式训练能力。 Whale框架主要分为四个模块，如下图所示。
API :提供简单易用的界面，允许用户组合使用各种混合并行策略； Whale IR :将并行策略转换为内部表示，通过任务图、多维、虚拟设备抽象来表示各种并行策略； 基于whaleengine:whaleIR，用图编辑工具构建分布式执行图Runtime :将分布式执行图转换为TFGraph，调用TF的Runtime执行；
2 Whale简介易用接口
Whale提供了一个简单易用的界面，用于描述各种并行策略。 主要原文如下。
构成集群：虚拟设备的分割方法replica :分割数据并行stage :任务组
hpipeline：流水并行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能力在更多业务场景落地，让技术能力到产品能力的转变。
作者 | 精明的大象
原文链接
本文为阿里云原创内容，未经允许不得转载。