近日，阿里云PAI团队与达摩院智能计算实验室联合发布巨型模型M6“低碳版”，大大降低万亿参数超大模型训练能耗。借助我们自研的Whale框架，我们仅用480块GPU就训练出了万亿参数多模态大模型M6，其规模是人类神经元的10倍。与传统海外公司实现的万亿参数规模相比，能耗降低。80%以上，效率提升近11倍。M6是国内首款实现商业化的多模态大型车型。M6拥有超越传统AI的认知和创造能力，擅长绘画、写作和问答。在电子商务、制造、文艺等诸多领域具有广阔的应用前景。这里介绍一下支持万亿参数模型训练的Whale框架设计。一、模型发展趋势与挑战1、模型发展趋势随着深度学习的普及，模型的参数规模也迅速增加。OpenAI数据显示：2012年之前，模型计算时间每2年翻一番，与摩尔定律一致；2012年后，模型计算时间每3.4个月翻一番，远超硬件发展速度。过去一年，模型参数规模增长迅速。谷歌、英伟达、阿里、致远研究院都发布了万亿参数模型，一些大厂也发布了百亿、千亿参数模型。同时，随着模型参数规模的增大，模型的效果也在逐步提升。Nvidia测试了Bert模型的不同参数尺度，发现模型perplexity随着模型参数尺度的增加而降低。Google在GShard论文中也发现，MoETransformer模型的参数越大，翻译质量越高。2.大模型训练的挑战大模型在提升模型效果的同时，也给训练框架带来了更大的挑战。例如，当我们要训练一个万亿规模的模型时，我们将面临以下挑战：训练难度大：GPU内存不再足够存储模型副本，数据并行性不再满足需求；该框架需要提供一种新的并行策略来配合多GPU能力来存储和训练模型；如何为用户提供简单易用的界面，使用户可以轻松实现分布式版本模型；超大规模模型对计算效率和通信效率带来巨大挑战，如何提高计算和通信效率；如何连接下游任务，如何支持批量预测和在线推理需求；成本高：以万亿模型为例，模型参数大小4TB，梯度也4TB，加上optimizerstates和activetensor，内存需求巨大；业界训练同等规模模型所需的资源：Nvidia3072A100、Google2048TPUv3，成本太高，难以实现；如何降低成本提高效率，使用更少的资源，实现更快的训练收敛；目前有一些分布式训练框架，如：Horovod、TensorflowEstimator、PyTorchDDP等支持数据并行，Gpipe、PipeDream、PipeMare等支持流水线并行，MeshTensorflow。版本控制困难且成本高，需要领域专家经验才能实现高效的分布式并行策略；迁移成本高：不同的分布式框架并行拆分，不同的框架有自己定制的DSL。当用户想要切换并行策略时，需要学习各种接口并重写模型；性能不理想：部分框架的实现没有考虑集群的物理环境；为了应对当前分布式训练的挑战，我们开发了分布式训练框架Whale，主要目标是：UnifymultipleparallelismStrategies：在一个框架中支持各种并行策略和这些策略的各种组合；简单易用的界面：用户只需添加几行注解即可完成并行策略的配置，无需更改模型代码；高效的训练框架：结合硬件资源和网络拓扑和模型协同优化，打造高效的分布式训练框架；2.PAI自研Whale框架1.Whale架构我们推出了统一多种并行策略的高性能分布式训练框架Whale，从以下几个角度来应对分布式训练挑战：对不同并行化策略的统一抽象和封装，以支持分布式训练框架中的多个并行策略；基于Tensorflow设计一套分布式并行接口，完全兼容Tensorflow，用户只需添加几行注解即可实现丰富的分布式并行策略；结合模型结构和网络拓扑进行调度和通信优化，提供高效的分布式训练能力Whale框架如下图所示，主要分为四个模块：API：提供简单易用的接口，允许用户组合使用各种混合并行策略；WhaleIR：将并行策略转化为内部表达式，通过TaskGraph、Multi-Dimension、VirtualDevicesAbstract来表达各种并行策略；WhaleEngine：基于WhaleIR，通过图编辑工具构建分布式执行图；Runtime：将分布式执行图转化为TFGraph，然后调用TFRuntime执行；2.Whale易于使用的界面Whale提供了简洁易用的界面来描述各种并行策略。主要原语有：cluster：配置VirtualDevicereplica的划分方式：dataparallelstage：划分TaskGraphpipeline：pipelineparallelsplit：operatorsplitting。这些接口可以结合各种并行策略，例如：数据并行：管道并行：管道并行+数据并行：更多并行策略示例：3.Whale训练过程使用Whale进行分布式训练过程：并行策略配置：使用WhaleAPI来配置模型Parallel策略，只需添加几行注解即可，无需修改模型代码，方法见2.2节；模型可以划分为多个TaskGraph，TaskGraph支持配置多种并行策略，每个TaskGraph可以配置不同的并行策略；VirtualResource划分：按照并行策略划分VirtualDevice，每个TaskGraph对应一个VirtualDevice；根据GPU资源和网络拓扑选择VirtualDevice的PhysicalDevice；分布式执行图：根据并行策略和资源分配信息，使用图编辑工具对执行图进行编辑（图复制、拆分、插入通信节点等），生成最终的分布式执行图；调用TF的runtime来执行分布式图；3.TrillionM6模型预训练万亿模型算力需求非常大，为了降低对算力的需求，在Whale中实现了MoE（Mixture-of-Experts）结构。MoE的主要特点是稀疏激活，使用Gating(Router)选择Topk个专家进行输入计算（k通常取值1和2），从而大大降低计算能力需求。Whale实现了MoE（Mixture-of-Experts）层，并行支持专家。专家被拆分到多个设备中，以减少单个设备的内存和计算能力需求。同时，数据并行有利于提高训练的并发性。因此，采用数据并行+专家并行的混合并行策略来训练M6模型：MoElayer采用专家并行，其他层采用数据并行。Whale为模型的混合并行训练提供了一个简单易用的界面。只需要添加几行注解来配置并行策略，模型本身不需要做任何修改。M6模型采用数据并行+专家并行的策略，只需要添加如下图所示的注解即可：同时，为了节省训练资源，提高训练效率，提供了各种优化技术inWhale:memoryoptimization:AutoGradientCheckpoint，自动选择最优的checkpoint节点，节省激活内存；Group-wiseApply，优化OptimizerApply阶段的内存；CPUOffload技术，优化Optimizerstatus和Weight的内存；通信池化，控制通信数据块的大小和并发，节省通信内存；计算、通信加速：采用DP+EP混合并行策略，降低算力需求；采用分组融合通信、半精度通信、拓扑感知All2All通信算子等技术，提高通信效率；结合混合精度、编译优化等技术，提高训练效率；借助Whale框架，首次在480V100上3天完成万亿M6模型的预训练。对比英伟达3072A100GPU实现万亿参数，谷歌2048TPU实现1.6万亿参数大模型，本次达摩院仅用480卡V10032GGPU实现万亿模型M6，节省超过1000亿计算资源.80%，训练效率提升近11倍。4结语模型参数规模越来越大，大模型成为发展趋势。为了解决超大模型训练的挑战，我们开发了Whale框架，将不同的并行化策略统一抽象和封装，并集成到一套分布式训练框架中，支持多种并行策略。Whale提供简洁易用的界面。用户只需添加几行注解即可实现各种并行策略，无需修改模型本身。同时结合硬件资源、网络拓扑、模型进行软硬件协同优化，提供高效的分布式训练框架。通过Whale框架，我们使用480V100GPU卡训练万亿规模模型，并在3天内完成模型训练收敛，为超大规模模型训练的实施提供了可能。未来，我们将进一步完善Whale框架，从更大规模、更快的速度和更高性价比三个维度扩展Whale框架的能力。同时，也将推动鲸鱼能力在更多业务场景落地，将技术能力转化为产品能力。
