作者：王林、飒洋
导读最近阿里云机器学习PAI平台和达摩院智能计算实验室一起公布“低碳版”巨模型M6-10T，模型参数曾经从万亿跃迁到10万亿，规模远超业界此前公布的万亿级模型，成为以后寰球最大的AI预训练模型。同时做到了业内极致的低碳高效，应用512 GPU在10天内即训练出具备可用程度的10万亿模型。相比之前公布的大模型GPT-3，M6实现等同参数规模，能耗仅为其1%。
M6模型训练应用的正是阿里云机器学习PAI平台自研的分布式训练框架EPL(Easy Parallel Library，原名whale)。EPL通过对不同并行化策略进行对立形象、封装，在一套分布式训练框架中反对多种并行策略，并进行显存、计算、通信等全方位的优化，来提供易用、高效的分布式训练框架。
EPL背地的技术框架是如何设计的？作为开发者咱们能够怎么应用EPL？将来有哪些布局？明天一起来深刻理解。
EPL是什么EPL(Easy Parallel Library)是阿里最近开源的，对立多种并行策略的、灵便易用的自研分布式深度学习训练框架。
1.我的项目背景随着近些年深度学习的火爆，模型的参数规模也增长迅速，OpenAI数据显示：
2012年以前，模型计算耗时每2年增长一倍，和摩尔定律保持一致；2012年后，模型计算耗时每3.4个月翻一倍，远超硬件倒退速度；
特地最近一年模型参数规模飞速增长，谷歌、英伟达、阿里、智源研究院都公布了万亿参数模型，有大厂也公布了百亿、千亿参数模型。随着模型参数规模增大，模型成果也在逐步提高，但同时也为训练框架带来更大的挑战。以后曾经有一些分布式训练框架，例如：Horovod、Tensorflow Estimator、PyTorch DDP等反对数据并行，Gpipe、PipeDream、PipeMare等反对流水并行，Mesh Tensorflow、FlexFlow、OneFlow、MindSpore等反对算子拆分，但当咱们要训练一个超大规模的模型时会面临一些挑战：
* 如何简洁易用：接入门槛高：用户实现模型分布式版本难度大、老本高，须要有领域专家教训能力实现高效的分布式并行策略；最优策略难：随着钻研人员设计出越来越灵便的模型，以及越来越多的并行减速办法，如果没有主动并行策略摸索反对，用户很难找到最适宜本身的并行策略；迁徙代价大：不同模型适宜不同的混合并行策略，但切换并行策略时可能须要切换不同的框架，迁徙老本高；
* 如何进步性价比：业界训练万亿规模模型须要的资源：英伟达 3072 A100、谷歌 2048 TPU v3，资源老本十分高；如何降本增效，组合应用各种技术和办法来缩小须要的资源，进步训练的速度；为了应答以后分布式训练的挑战，咱们研发了分布式训练框架EPL，将不同并行化策略进行对立形象、封装，在一套分布式训练框架中反对多种并行策略。EPL同时提供简洁易用的接口，用户只需增加几行annotation即可实现并行策略的配置，模型代码不须要改变。在用户无感的状况下，EPL通过进行各种显存、计算、通信优化，打造高效的分布式训练框架。
2.次要个性多种并行策略对立：在一套分布式训练框架中反对多种并行策略（数据/流水/算子/专家并行）和其各种组合、嵌套应用；接口灵便易用：用户只需增加几行代码就能够应用EPL丰盛的分布式并行策略，模型代码无需批改；主动并行策略摸索：算子拆分时主动摸索拆分策略，流水并行时主动摸索模型切分策略；分布式性能更优：提供了多维度的显存优化、计算优化，同时联合模型构造和网络拓扑进行调度和通信优化，提供高效的分布式训练；3.开源地址EPL(Easy Parallel Library)的开源地址是：https://github.com/alibaba/Ea...
咱们同时提供了model zoo，欢送大家试用：https://github.com/alibaba/Fa...
EPL次要技术特点EPL通过丰盛并行化策略，简略易用的接口，多维度的显存优化技术，和优化的计算通信减速技术，让每一个算法工程师都能轻松训练分布式大模型工作。
丰盛的并行化策略：EPL提供了多种并行化策略及其组合策略，蕴含数据并行，流水并行，算子拆分并行及并行策略的组合嵌套。丰盛的策略抉择使得不同的模型构造都能找到最适宜本人的分布式训练形式。易用性：用户的模型编程接口和训练接口均基于TensorFlow，用户只需在已有的单机单卡模型上做简略的标记即可实现不同的分布式策略。EPL设计了两种简略的策略接口(replicate/split)来表白分布式策略及混合并行。分布式策略标记的形式让用户无需学习新的模型编程接口，仅需几行代码即可实现和转换分布式策略，极大升高了分布式框架的应用门槛。显存优化：EPL提供了多维度的显存优化技术，蕴含主动重算技术(Gradient Checkpoint)，ZeRO数据并行显存优化技术，CPU Offload技术等，帮忙用户用更少的资源训练更大的模型。通信优化技术：EPL深度优化了分布式通信库，包含硬件拓扑感知，通信线程池，梯度分组交融，混合精度通信、梯度压缩等技术。1.技术架构EPL框架如下图所示，次要分为以下几个模块：
接口层：用户的模型编程接口基于TensorFlow，同时EPL提供了易用的并行化策略表白接口，让用户能够组合应用各种混合并行策略；两头表白层：将用户模型和并行策略转成化外部表白，通过TaskGraph、VirtualDevices和策略形象来表白各种并行策略；并行化引擎层：基于两头表白，EPL会对计算图做策略摸索，进行显存/计算/通信优化，并主动生成分布式计算图。Runtime执行引擎：将分布式执行图转成TFGraph，再调用TF 的Runtime来执行；
2.并行化策略表白EPL通过strategy annotation的形式来划分模型为多个TaskGraph，并在此基础上进行并行化。 EPL有两类strategy：replicate 和 split。通过这两种并行化接口，能够表白出各种不同的并行化策略，例如：
数据并行: 上面这个例子是一个数据并行的例子，每个模型正本用一张卡来计算。如果用户申请了8张卡，就是一个并行度为8的数据并行任务。
流水并行：在上面的例子里，模型被切分成2个 TaskGraph, "stage0"和"stage1"，用户能够通过配置pipeline.num_micro_batch参数来设定pipeline的micro batch数量。 在这个例子里，"stage_0"和"stage_1"组成一个模型正本，共须要2张GPU卡。如果用户申请了8张卡，EPL会主动在pipeline外嵌套一层并行度为4的数据并行（4个pipeline正本并行执行）。
算子拆分并行：在以下例子中，EPL会对split scope下的模型定义做拆分，并搁置在不同的GPU卡上做并行计算。
同时EPL反对对上述并行策略进行组合和嵌套，来组成各种混合并行策略，更多示例能够参考开源代码的文档和示例。
3.显存优化当模型增长，GPU的显存经常成为训练大模型的瓶颈。EPL提供了多维度的显存优化技术，极致优化了训练显存消化。
重算 Recomputation (Gradient Checkpoint)：失常的DNN前向过程中会生成activation，这部分activation会在后向过程中用于梯度的计算。因而在梯度生成之前，前向的activation会始终存留在显存中。activation大小和模型构造以及batch size相干，通常占比都十分高。Gradient Checkpoint (GC) 通过保留前向流传过程中的局部activation，在反向流传中重算被开释的activation，用工夫换空间。GC中比拟重要的一部分是如何抉择适合的checkpoint点，在节俭显存，保障性能的同时，又不会影响收敛性。EPL提供了主动GC性能，让用户能够一键开启GC优化性能。ZeRO：在数据并行的场景下，每个卡上会寄存一个模型正本，optimizer state等，这些信息在每张卡上都是一样，存在很大的冗余量。当模型变大，很容易超出单卡的显存限度。在分布式场景下，咱们能够通过相似DeepSpeed ZeRO的思路，将optimizer state和gradient分片存在不同的卡上，从而缩小单卡的persistent memory占用。显存优化的AMP(Auto Mixed Precision)：在惯例的AMP里，须要保护一个FP16的weight buffer，对于参数量比拟大的模型，也是不小的开销。EPL提供了一个显存优化的AMP版本，FP16只有在用的时候才cast，从而节约显存。Offload: Offload将训练的存储空间从显存扩大到内存甚至磁盘，能够用无限的资源来训练大模型。同时，EPL反对各种显存优化技术的组合应用，达到显存的极致优化。咱们在T5模型上开启GC+ZeRO+显存优化的AMP技术，在性能放弃不变的状况下，显存升高2.6倍。
...