『AI原理解读』MindSpore1.2强大并行能力介绍与解读

Posted 2022-11-27 小宋是呢

---

融合 5 大维度，强大的自动并行

 

MindSpore 自动并行提供了 5 维的并行方式：数据并行、算子级模型并行、Pipeline 模型并行、优化器模型并行和重计算，并且在图编译阶段，有机融合了 5 个维度的并行。这 5 维并行方式组合起来构成了盘古的并行策略。

 

a. 数据并行

 

数据并行是最基本，应用最广的并行方式，其将训练数据（mini-batch）切分，每台设备取得其中一份；每台设备拥有完整的模型。在训练时，每台设备经过梯度计算后，需要经过设备间的梯度同步，然后才能进行模型参数的更新。

b. 算子级模型并行

 

算子级模型并行是对模型网络中的每个算子涉及到的张量进行切分。MindSpore 对每个算子都独立建模，每个算子可以拥有不同的切分策略。

以矩阵乘算子 MatMul(x, w)为例，x 是训练数据，w 是模型参数，两者都是二维矩阵。并行策略 ((4, 1), (1, 1)) 表示将 x 按行切 4 份，保持 w 不切，如果一共有 4 台设备，那么每台设备拥有一份 x 的切片，和完整的 w。

c.Pipeline 模型并行

 

Pipeline 模型并行将模型的按层分成多个 stage，再把各个 sage 映射到多台设备上。为了提高设备资源的利用率，又将 mini-batch 划分成多个 micro-batch, 这样就能够使得不同设备在同一时刻处理不同 micro-batch 的数据。

一种 Pipeline 并行方式(Gpipe) 要求反向计算要等所有设备的正向计算完成后才开始，而反向计算可能依赖于正向的输出，导致每个卡正向计算过程中累积的 activation 内存与 micro-batch 数量成正比，从而限制了 micro-batch 的数量。MindSpore 的 Pipeline 并行中，将反向提前，每个 micro-batch 计算完成后，就开始计算反向，有效降低 activation 存储时间，从而提升整体并行效率。

d. 优化器模型并行

 

优化器模型并行将优化器涉及到的参数和梯度切分到多台设备上。以 Adam 优化器为例，其内部可能有多份与权重同等大小的「动量」需要参与计算。在数据并行的情况下，每个卡都拥有完整的「动量」，它们在每个卡上都重复计算，造成了内存及计算的浪费。通过引入优化器并行，每个卡只保存权重及「动量」的切片，能降低每个卡的静态内存及提升计算效率。

e. 重计算

 

重计算 (Rematerialization) 针对正向算子的输出累计保存在内存中，导致内存峰值过大的问题，舍弃了部分正向算子的输出，而是在反向阶段用到时再重新计算一遍。这样做有效地降低了训练过程中的内存使用峰值。如下图所示，第一个内存峰值通过重计算消除，第二个内存峰值可以通过前面讲到的优化器并行消除。

有了这 5 维的并行维度后，如何将其组合起来作用于盘古，并且如何将切分后的模型分片分配到每台设备上仍然是难题。MindSpore 自动并行，把这 5 个维度并行有机组合起来，可以实现非常高效的大模型分布式训练能力

下图 (b) 是一典型的树形的硬件拓扑结构，其带宽随着树深度的增加而降低，并且会产生一些流量冲突。为了利用此特征，MindSpore 的目标是最大化计算通信比，将通信量大的并行方式（算子级并行）放置在服务器内部的多卡之间；将通信量较小（Pipeline 并行）的放置在同一机架内的服务器间；将数据并行（优化器并行）的部分放置在不同机架间，因为该通信可以和计算同时执行(overlap)，对带宽要求较低。

在盘古 2000 亿模型中，MindSpore 将 64 层（layer）划分为 16 个 stage，每个 stage 包含 4 层。在每层中，利用算子级并行的方式对张量进行切分。

如下图中的 Q,K,V 的参数在实际中（按列）被切了 8 份，输入张量（按行）被切了 16 份，输出张量因此被切了 128 份（8*16）。重计算配置是配置在每层内的，也就是重计算引入的多余的计算量不会超过一层的计算量。总计，MindSpore 使用了 2048 块昇腾处理器来训练盘古。

MindSpore 对外屏蔽了复杂并行实现的细节，使得用户像编写单机模型脚本那样简单。用户在单机脚本的基础上，仅通过少了配置就能实现多维度的混合并行。下图是简化版的盘古脚本，其中红色加粗字体表示的在 MindSpore 中的并行策略。将红色加粗字体去掉，则是单机脚本。

 

图算跨层联合优化，发挥硬件极致性能

 

除了跨节点间的大规模自动外，在单卡节点内，MindSpore 通过图层和算子层的跨层协同优化，来进一步发挥昇腾算力。

在传统的 NN 网络中，不同算子承载的计算量和计算复杂度也各不相同。如 LayerNorm 由 11 个基本算子组成，而 Add 则只有 1 个基本算子。这种基于用户角度的算子定义，通常是无法充分发挥硬件资源计算能力的。因为计算量过大、过复杂的算子，通常很难生成切分较好的高性能算子。从而降低设备利用率；而计算量过小的算子，由于计算无法有效隐藏数据搬移开销，也可能会造成计算的空等时延，从而降低设备利用率。

为了提升硬件利用率，MindSpore 使用了图算融合优化技术，通过图层和算子层联合优化，并将「用户使用角度的易用性算子」进行重组融合，然后转换为「硬件执行角度的高性能算子」，从而充分提升硬件资源利用率，进而提升整网执行性能。具体优化流程如下图所示：

以 LayerNorm 算子为例，通过算子拆分和重组，11 个小算子，组成了 1 个单算子和 2 个融合算子。这些重组后的算子可以生成更加高性能的算子，从而大大降低了整体网络运行时间。

在盘古模型中，图算融合帮助整体训练时间减少了 20% 以上。除此之外，对于其它 NLP、CV 等任务，图算融合在优化性能方面都有不错的表现。