开发者学堂课程【PAL 平台学习路线：机器学习入门到应用：Whale 基于 Tensorflow 深度学习分布式训练框架】学习笔记，与课程紧密联系，让用户快速学习知识。
课程地址：https://developer.aliyun.com/learning/course/855/detail/14097
Whale 基于 Tensorflow 深度学习分布式训练框架
内容介绍：
一、Whale 的特点
二、用户接口介绍
一、Whale 的特点
工业级深度学习分布式训练框架
1.支持多种不同分布式策略。
Whale 不仅支持最基础的数据并行，同时也支持模型并行、算子拆分、流水并行及各种不同并行化的混合并行化组合。
2. 支持用户自定义并行化策略，也支持自动并行化。
3. 高效的分布式训练效率。
4.接口灵活易用，完全兼容 Tensorflow。
用户定义好的单机单卡的 Tensorflow 代码只需要利用 Whale 简单的借口，就可以完成非常丰富的并行化策略。
二、用户接口介绍
Whale 提供了两组简单的用户接口。
1.cluster Virtual Devices 定义
Cluster 将物理设备分组成虚拟设备，最终生成的虚拟设备即 Virtual Devices ,其会将具体的模型部分映射到虚拟 Virtual Devices 上进行执行。
2.Scope 并用化原语
（1）用于描述并行化类型:
①replica（数据并行）
②stage（模型并行）
③split（算子拆分）
④pipeline（流水并行）
⑤auto-parallel（自动并行）
3.接口使用实例
实例介绍接口如何使用，具体主要有数据并行、算子拆分和流水并行的实例。
（1）数据并行实例
数据并行在每一个具体的计算设备上都会放置一个完整的模型副本，在完成前项后项之后，每个 Devices 上的模型副本都会生成出一个对应的 local batch 的梯度。经过全局的AllReduce 后，每张卡会拿到 AllReduce 之后的梯度，再进行本地 ApplyUpdate 的权重。这个过程用户本身定义单机单卡代码，即如下①所展示，只要用户定义好单机单卡的代码，在其前面加上 with wh.cluster():，里面不传参数即代表所有的 device 都会映射成一个虚拟 device，然后再通过 with wh.replica(): 接口来告诉 whale框架。下面定义的这些模型，在 replica 下面所有的代码都会以数据并行的方式 Place 到各个设备上，具体在背后如何进行梯度的聚合，全部由 whale 框架自动完成。
①数据并行代码
with wh.cluster():
with wh.replica():
out =Model()
②数据并行逻辑图
（2）算子拆分实例：大规模分类任务
·ResNet-50 with 100,000 classess
该模型 ResNet-50 部分权重大约有90 M左右，而全连接层权重大约有782 M。在实际应用中如果用数据并行训练时，可以看到 Bottleneck 和 performing 得到的信息，整个全连接层的 allreduce 通信是整个训练中的一个瓶颈，这主要原因还是因为 FC 全连接层的参数量是 ResNet-50 的参数量的8.7倍，它的参数量太大，导致整个训练过程中全连接层的梯度同步是一个瓶颈。那针对这种场景下，如何进行更优的分布式训练呢？这里提供的策略就是通过算子拆分。
瓶颈的主要原因：
①FC 层参数量=8.7X ResTet50参数量
②数据并行场景下 FC 层的同时是瓶颈
对于 ResNet-50 大规模分类任务，算子拆分实际上就是将模型分成两部分，一是 backbone 部分，也就是 Stage0 ，其划分完后剩下的一部分全连接层和 Softmax 层，会将其划分到 Stage1 里。这样模型就会被划分为两个部分，一是 Stage0 ，ResNet-50 部分，二是 Stage1 ，全连接层和 Softmax 层部分。
这里又通过 class 定义了两个 Virtual Device,其实际上是可以通过 Whale 定义出两个相同的 Virtual Device，也可以定义出两个不同的 Virtual Device。这里举的例子是第一个 Virtual Device 其背后映射的具体物理资源是相同的，映射完后，其背后的混合映射执行图如下图所示，数据并行部分会全部以副本的形式放在 GPU0到 GPU5上。对于全连接层和Softmax 部分会通过分片的方式放在 GPU0到 GPU5上，中间会通过 Features Allgather50输出的 Feature，通过 Allgather 的操作再分发到具体的全连接层分片上，完成前项，后项也是同理。这样一个复杂的逻辑，通过 Whale 如何进行分布式的实现呢？如下代码。
import whale as wh
cluster = wh.cluster(layout = {"a1l"})
with cluster:
with wh. replica():
features = ResNet50(inputs)
with wh. split():
logits = FC(features)
predictions = Softmax(logits)
对于简单的单机单卡代码来说，用户只需要定义 feature 、ResNet-50的模型部分、FC 和 Softmax 部分，这里是按照前面的介绍需要对 ResNet-50 部分采用数据并行的逻辑操作，只需要通过对 ResNet-50添加 whih.replica 操作就可以完成 ResNet-50数据并行部分的实现。FC 层和 Softmax 层，按照前面的逻辑来说是需要进行 soft 拆分的，只需要在它的前面加上 with wh.split，这个 scope 接口就让 whale 在底层框架层自动的将 FC 层和 Softmax 层进行分片的放置在所有的物理 Devices 上，在最外层套的是 class，按照刚才的描述，需要生成一个相同的 class 即可。
这里就可以用 which class logits 来进行完成，通过这样一个简单的接口，就可以实现非常复杂的算子拆分便易化，这样的接口非常简单，而且通过这样的接口实现，可以将模型的分类扩展到亿级分类，这种场景实际上对数据并行无法实现，但是通过算子拆分，可以提高模型的规模。这里也对10万分类进行了一个测试，10万分类的测试的 Bottleneck 是 whale 的数据运行，可以看到，在10万分类上，64卡训练性能是数据运行的14.8倍，这个也得益于通过算子拆分。这里只需要去同步通梯度，梯度的通信只需要去通信 ResNet-50部分。算子拆分的 FC 部分，就不需要进行梯度通信了，只需要用一些的feature map通信来替代算子拆分 FC 部分的梯度通信，这样可以大大降低通信的成本。
（3）流水并存实例：BertLarge 任务
介绍一下流水运行，这里是以 BertLarge 任务为例。单机单卡的任务，这是 bert 的 encoders embedding加一个24层的 encoder 最后加一个 pooler 层。
由于这个模型比较大，通过数据并行的时候，通梯度的通信开销比较高，那首先想到需要把这个模型进行分片，放到不同的卡上，所以第一步，把分片放到不同卡上，就是这里的意思是将 embedding 的1到8层的 encoder 放在了 GPU0（stage0）上，把第9层和第16层的 encoder放在 GPU1（stage1），把第17到24层 encoder 放在 GPU2（stage2）上。
这样的模式运行还具有一些问题，与下面的持续图一样，当训练的过程中，在同一时刻，只有一张卡会正在使用，其他卡都处于空闲状态，所以也就是 bubble 的时间会非常长。为了解决这个问题，需要在此基础上再加上流水并行的策略。流水并行就是如下图所示，先加一个模型，按照模型并行的方式拆分完之后，再输入的时候喂进去多个不同的macbatch，这样可以将整个时序上的 gpu 使用率拉的更满，可以明显的看到右侧的时序图和比左侧的时序图空白的部分要少很多，只在一个模型的模型并行的基础上进行流水并行是远远不够的，因为在大规模训练时，常常具体计算硬件资源，有时到达几十上百张卡，这种情况很难将一个模型把它拆分成上百份，这个性能会非常差，而且模型不一定能够拆得这么多。因此，应该先将一个模型在机内进行拆分，拆分成 N 份之后，在不同的机器间做数据并行，第一个副本放在 work0上，第二个副本也同样的以模型并行加流水并行的方式放在 work1上，work2上也有同样的副本。这是一种并行化的方式。另一种更优的变化方式，可以先做跨级流水，也就是 work0、work1、work2上的 GPU0，这是三张卡。然后分别放一个模型副本的stage0、stage1、stage2，像 work0 只保存三个 stage0 的副本。
编写这样一个模型并行、流水并行、数据并行的 Bertlarge 的训练通过 word 接口怎么写？这里举一个编写代码的例子。
首先是单机单卡的代码，先定义好 embedding 以及 encoder layer，然后最后加 pooler。
output = embedding （inputs )
output = encoder_layer 0_24(output)
output = pooler(output )
首先要做的是把此模型变成一个模型并行的代码，模型并行代码通过stage接口，这里有三个 with stage。用户只需要将三个 with stage 接口放在相应的模型定义的地方。举个例子：像 embedding 和 encoder layer 的0到8被当作第一个 stage，放在第一个 with stage 下面，那 encoder layer 的8到16层，就用另外一个 with stage 来进行进行嵌套 with。同样的道理定义第三个stage，定义好这三个 stage 之后，这就是一个标准的模型并行的代码了，需要对这个模型并行的代码再增加一个流水并行的逻辑，所以只需要在它到外面再嵌套一个 with pipeline ，Michael 的number 数填成4，这样就完成了模型并行加流水并行的代码的撰写。刚才也提到，当卡数特别多的时候，只用流水并行是完全不够的，那还需要套一个 replica 接口，通过 with replica会自动的将模型放置到不同的 work 上的 Devices 上，自动完成数据并行的 placement。另外最外层还需要加一个 with class，其是有默认的 Layout 来进行生成 class ，自己也可以按照自己需要指定 class 的 layout 类型。通过这样一个简单的代码撰写，就可以完成复杂的数据并行、流水并行以及模型并行的混合并行的 BertLarge 任务的训练。
import whale as wh
with wh.cluster():
with wh. replica():
with wh. pipeline(micro_ batch=4):
with wh. stage():
output . embedding inputs )
output . encoder layer _0 8(output)
with wh.stage():
output . encoder layer 8 16(output)
with wh. stage():
output . encoder layer 16 24(output)
output . pooler( output)
通过这样的混合并行的效果，这里也做了测试。在64卡的任务上，当 Batch Size=6时，流水变形的性能比 whale 数据并行有1.34倍的提升。另外比 horovod 数据并行有2.32倍的提升。Whale 数据并行之所以比 horovod 数据并行更优秀，是因为 whale 数据并行的在逻辑上有大量的优化,以及通项的优化。
