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

MindSpore开源框架加持,如何「炼出」首个千亿参数、TB级内存的中文预训练语言模型?


融合 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 等任务,图算融合在优化性能方面都有不错的表现。

猜你喜欢

转载自juejin.im/post/7000407145971712007