摘要： 千亿参数量的中文大规模预训练语言模型时代到来。
本文分享自华为云社区《MindSpore 开源框架加持，如何「炼出」首个千亿参数、TB 级内存的中文预训练语言模型？》，原文作者：chengxiaoli。
千亿参数量的中文大规模预训练语言模型时代到来。
近段时间，中文大规模预训练语言模型圈有些冷落。26 亿参数量的「悟道 · 文源」，270 亿参数量的 PLUG，以及昨天华为云公布的千亿级别「盘古」NLP 大模型，预训练语言模型曾经成长到仅加载就须要 TB 级的内存或显存。
咱们能够直观地想到，「盘古」成果理当更好，但计算量需要也更大，训练起来更艰难。
然而「盘古」实际上是这样一次摸索：开源框架 MindSpore，昇腾根底软硬件平台，加上超大规模中文预训练模型，意味着基础设施未然欠缺了。
这项工作由华为以及北京大学相干技术团队联手实现，在昇腾根底软硬件平台，以及 MindSpore 框架主动并行等黑科技的帮忙下，训练出以后最大的中文预训练模型。
那么量级一直拔高的盘古大模型是如何训练进去的？接下来，让咱们粗疏解读下「盘古」背地的关键技术。
千亿参数，TB 级内存的模型
以盘古 2000 亿为例，如果咱们训练时权重都用规范的 FP32 数据格式，那么算下来，权重占的空间就达到了 750GB，训练过程中内存开销还会数倍回升。这 750GB 参数，不是放在硬盘上，也不是加载到内存中，而是须要移到昇腾 Atlas 训练服务器 HBM（High Bandwidth Memory 高带宽存储器）内存中，以利用昇腾 Atlas 训练服务器训练模型。
模型大，意味着数据也大，而且都须要是高质量数据。为了满足数据需要，研发团队从互联网爬取了 80 TB 文本，并最初荡涤为 1TB 的中文数据集。
这样的模型与数据，曾经不是咱们几台服务器能加载上的了，更不用说进行训练。好在研发团队会提供 API，个别算法工程师间接调用接口就能试试成果。
能够说，目前盘古是业界独创的千亿规模中文预训练模型，其中最高参数量达 2000 亿。
超大规模主动并行，算法工程师的福音
先思考一个问题，你想到如何训练这样的大模型了吗？
如果给你足够的计算力，你能想到如何训练这么大的模型吗？咱们最罕用的分布式训练形式数据并行，独自这么做必定是不行的，因为没有哪个计算硬件能放下 800GB 的参数。那么再加上模型并行呢？又产生了新问题，咱们该如何拆分如此微小的「盘古」？硬件产品（如 NPU、GPU 等）之间的梯度流、数据流通信又是什么样的？
显然训练如此宏大的模型，远比咱们设想中的简单，须要大量的工程化操作，并保障这些操作不会或极少影响到模型最终收敛成果。
难道盘古真得靠手动并行优化？
如果手动来写分布式训练逻辑，那么须要综合思考计算量与类型、集群带宽、拓扑构造、样本数量等等一大堆简单的货色，而后再设计出性能比拟优良的并行切分策略，并编写大量并行切分和节点间的通信代码。如果零碎环境变了，还要从新设计并批改算法，想想就感觉头大。
假使咱们用 TensorFlow 或其余相似框架，MirroredStrategy 这一系列自带的分布式策略齐全用不上，看起来自行写并行策略是必不可少的。然而，盘古 真正的训练是一种软硬件协同的形式，MindSpore 计算框架、CANN 异构计算架构、昇腾根底软硬件平台整套基础设施。其中，MindSpore 提供的，就蕴含了至关重要的主动并行能力。
交融 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 等工作，图算交融在优化性能方面都有不错的体现。
总结：训练超大模型的完满体现
即便给咱们足够的算力，超大模型的训练还是异样简单，远比设想中的艰难。对于咱们个别算法工程师来说，针对某个工作，上亿参数量曾经算大的了，然而并不会感到训练上会有什么艰难，因为各个深度学习框架间接调用数据并行接口就能搞定。
然而如果模型持续增大到百亿级、千亿级甚至万亿级，并行与优化策略的复杂度猛然回升，算法工程师一点点地编写与优化代码可太难了。MindSpore 通过主动并行，把计算逻辑和并行逻辑解耦，单卡串行代码主动实现分布式并行，从而使得算法科学家将精力都解放到模型自身上。
为了从预训练获取更多的常识，GPT-3 与盘古这样的模型会越来越大，毕竟到当初咱们还没看到大模型预训练成果的极限在哪。届时，这类模型对基础设施的需要会更大，并行与优化策略也会更加简单。只有领有足够优良的基础设施，大规模预训练的成果才会更好，从而在常识问答、常识检索、常识推理、浏览了解等场景施展更大作用，实现智慧客服、营销、文案生成等商业价值。
大规模计算集群及软硬件协同优化，这次在盘古的训练上失去了充沛的完满体现。正如开发团队所言，「基于 Mindspore 和昇腾根底软硬件平台在千亿参数模型上的实际也是一次摸索，大模型的分布式训练、超参调优、数据集组成、模型构造适应性等都存在太多的未知。当初，盘古模型成果很好，刷新了 clue 版第一，这意味着第一次基于国内软硬件协同优化，以及超大规模分布式训练，后果是令人振奋的，咱们本人也具备了足够强的基础设施。」
当然，也诚如以上所言，盘古只是对超大规模分布式训练、超大规模中文预训练模型的一次摸索，将来还须要更多的钻研工作者投入到通用智能与大规模分布式计算的钻研工作中。
点击关注，第一工夫理解华为云陈腐技术~