apex降级保护能用几次,apex降落时的音乐

2022-07-23 04:42:54 创业研究

　　机器之心专栏　　

　　机器之心编辑部 　　

　　Transformer已经成为很多NLP任务和一些CV任务的主流模型。但由于硬件资源不足，很多大学实验室或公司无法训练大模型，缩小批量等措施会导致训练时间翻倍。针对这一痛点，字节跳动推出了LightSeq训练加速引擎，优化了变形金刚训练的整个计算过程，最多可以实现3倍以上的加速。　　

　　现在NLP和CV领域的大部分任务都使用Transformer作为基本模型。早在2019年12月，字节跳动就已经开放了一个类似变形金刚的模型推理加速引擎——light seq 1。作为业界首款支持多种模型和解码方式的推理加速引擎，LightSeq的推理速度比其他同类软件都要快，并且远超TensorFlow和PyTorch。　　

　　今天，LightSeq 的最新版本发布了！ 　　

　　此版本引入了一个主要的引擎update-支持了 Transformer 全流程训练加速，在不同的批处理大小下相比主流训练库最高可加速 3 倍多！至此从训练到推理部署的整个过程都已被 LightSeq 打通。. 　　

　　那么它采用了哪些技术呢？作者将根据LightSeq发布的信息为您一一揭秘。　　

　　LightSeq 能做什么？ 　　

　　图1:1:Transformer模型结构图(以机器翻译为例) 　　

　　自2017年由Google提出以来，Transformer Model 2已经成为许多NLP任务和一些CV任务的主流模型，尤其是机器翻译、文本生成、文本摘要、时间序列预测等任务。图1是一个使用Transformer进行机器翻译任务训练的例子。　　

　　但由于硬件资源不足，很多大学实验室或公司无法训练大模型，缩小批量等措施会导致训练时间翻倍。因此，如何利用有限的资源尽快训练出模型，成为一个亟待解决的问题。　　

　　针对这一痛点，字节跳动推出了LightSeq训练加速引擎，优化了变形金刚训练的整个计算过程。从字嵌入层、编码层、解码层，到最后的损失函数层，从正向传播、反向传播、梯度同步，到最后的参数更新，LightSeq都做了详细的性能分析和优化。　　

　　以机器翻译任务为例。在这个任务中，我们只需要修改几行代码就可以启动LightSeq的训练加速，可以达到3倍的加速。　　

　　一般来说，LightSeq有以下优点: 　　

　　1. 支持Transformer 的完整加速。 　　

　　SEQ是业内首个全面支持整个Transformer模型加速的训练引擎，包括单词嵌入层、编码层、解码层、损失函数层等高效定制层。相比之下，另一个知名深度学习优化引擎DeepSpeed 3只支持编码层加速，所以只能在BERT等模型中使用，局限性很大。　　

　　2. 训练速度快。 　　

　　Seq训练速度很快。比如在WMT14 4英德机器翻译任务中，使用NVIDIA最新的训练显卡A100，与主流的序列生成库相比，LightSeq最快只需要三分之一的训练时间。　　

　　3. 功能全面，简单易用。 　　

　　Seq提供了高效的TensorFlow和PyTorch自定义层供用户灵活使用，可以自由插入拥抱脸等主流训练库。此外，它还与当前流行的训练库(如Fairseq 5和NeurST 6)进行了深度集成。安装LightSeq后，用户只需修改几个命令行参数，就可以在这些培训库上使用LightSeq。　　

　　4. 提供丰富的二次开发工具。 　　

　　Seq提供了完整的CUDA内核和Transformer自定义层的单元测试功能，可以测试自定义操作符的正确性，分析运行时间和加速比，帮助开发者更快的验证功能的正确性和有效性。　　

　　下表1列出了不同训练加速引擎之间支持的功能的比较：　　

　　表1:不同训练加速引擎的支持功能比较　　

　　如何快速上手 LightSeq？ 　　

　　SEQ提供了多种便捷的操作模式。点击文末链接，快速体验。以下是快速访问LightSeq进行培训的两种方法的简要说明。　　

g>使用自定义层

LightSeq 提供了许多自定义层供用户灵活使用。

例如，如果你想将 Transformer 模型中的编码层替换为 LightSeq 的编码层，只需要提供一个编码层参数，用来初始化 LightSeq 编码层，然后就可以用它来替换原始的编码层，加速模型训练。详细代码如下（这里省略了部分配置参数）：

一键启动

LightSeq 还和几个当前流行的训练库（例如 Fairseq 和 NeurST）做了深度集成，仅需修改很少的代码就能开启 LightSeq 加速。

Fairseq

LightSeq 为 Fairseq 提供了一套完整便捷的 Transformer 训练样例。

首先，你需要安装 Fairseq 以及必要的第三方库，然后用如下命令安装 LightSeq 加速库：

接着就可以通过 LightSeq 提供的启动器，灵活地指定使用 LightSeq 优化版本的 Transformer 模型、参数优化器和损失函数。下面是启动参数样例（这里省略了和 LightSeq 无关的参数）：

NeurST

NeurST 是一款同时支持 TensorFlow 和 PyTorch 的开源序列生成库，可以用来做文本生成、机器翻译和语音翻译。LightSeq 已经与其进行了深度融合，无需用户修改代码和启动参数即可直接进行加速训练。

首先需要安装 NeurST，官方提供了详细的安装教程 <6>。

然后安装 TensorFlow 版本的 LightSeq，命令如下：

这样，NeurST 就会自动识别 LightSeq 已经安装成功，调用 lightseq 库进行模型构建来加速训练，无需修改启动参数。运行命令详见 NeurST 提供的机器翻译样例 <7>。

性能测试

在 WMT14 标准的英德翻译任务上，LightSeq 做了评测实验。以当前流行的 Fairseq 训练库（基于 PyTorch）和被广泛使用 Apex 工具库 <8> 为基准，测试了 LightSeq 的训练性能。实验在 NVIDIA Tesla V100 和 NVIDIA Ampere A100 上进行，采用单机八卡数据并行训练和 16 位浮点数混合精度。

在不同模型大小和批处理大小下，LightSeq 对单步训练速度的提升结果如图 2 所示：

图 2：A100 和 V100 显卡下不同层数模型加速比

这里所有模型的词表大小为 40k，编码层和解码层隐层维度是 1024，注意力头数为 16。e 和 d 分别表示编码器和解码器的层数。加速比的计算采用了每秒训练有效单词数（real word per second）这一常见指标。

可以发现：

1. 使用了 LightSeq 后，单步训练速度有 45%-250% 的提升。作为对比，Apex 仅有 5%-30% 的提升。LightSeq 能取得这种大幅度的性能提升，源自其对模型进行了全流程的细致优化。

2. Apex 显存使用量对比原生 Fairseq 略有提升，例如 V100 上，Apex 导致 6e6d 的模型在 15000 批处理大小上出现了显存溢出，而 LightSeq 仍然可以正常训练。这证明了 LightSeq 在保证高效计算的同时，也做到了高效使用显存，这个特征在显存有限或者训练大模型时非常关键。

3. 随着批处理大小的增加，LightSeq 加速比逐渐降低。其原因是，经过 LightSeq 优化后，单步训练中矩阵乘法占比提高，显卡的计算吞吐成为训练速度的瓶颈。这表明 LightSeq 已经对计算资源做到了充分利用。也解释了为什么计算吞吐更高的 A100，平均会取得比 V100 高 15% 左右的加速比。

最后在 WMT14 英德翻译数据集上测试了图 1 中 Transformer 模型训练至收敛的时间，结果如图 3 所示：

图 3：A100 和 V100 显卡下不同模型训练总时间

由于 LightSeq 的计算优化是无损的，不会影响模型训练至收敛的训练步数，所以收敛时间的提升和单步训练时间的提升趋于一致。观察可以发现，LightSeq 最多可将模型训练时间由 8.5 小时降低到 3.8 小时。

可视化分析

为了更清楚地看出 LightSeq 优化前后模型的运算情况，我们用 Nsight Systems <9> 可视化模型训练过程中单步的算子调用情况来展示。

首先是 Fairseq+Apex 的可视化，结果如图 4 所示。总耗时在 288ms 左右，三个红色框分别表示前向传播、反向传播、梯度同步与参数更新。可以看出前向传播的算子排列比较稀疏，存在很大的优化空间。

图 4：Fairseq+Apex 单步训练过程可视化

然后是 Fairseq+LightSeq 的可视化，结果如图 5 所示，总耗时降到了 185ms 左右。而且 LightSeq 的算子排列更加紧密，大大增加了显卡的利用率。

图 5：Fairseq+LightSeq 单步训练过程可视化

CUDA kernel性能

此外还测试了 LightSeq 单卡情况下所有 CUDA kernel 的性能，对比了 PyTorch、TensorFlow（XLA 编译优化）、DeepSpeed 和 LightSeq 四种实现方式。由于 kernel 太多，这里只列举了部分实验结果。

首先对比了最常用的 dropout，图 6 是 V100 显卡上 16 位和 32 位浮点数 dropout 不同实现的加速对比结果：

图 6：dropout 加速对比

从图 6 中可以看出，LightSeq 的实现要远远快于 PyTorch 和 DeepSpeed。DeepSpeed 在元素数量过百万之后逐渐不如 PyTorch，而 LightSeq 始终比 PyTorch 快。随着元素数量的增加，LightSeq 和 DeepSpeed 的速度都会有明显下降。TensorFlow 在开启了 XLA 之后速度依然落后于 PyTorch，且随着元素数量的增加差距逐渐缩小。

然后对比了注意力机制中的 softmax 函数，测试了实际训练场景中批处理大小为 8192 情况下的加速比。图 7 是 V100 显卡上 16 位和 32 位浮点数 softmax 不同实现的加速对比结果，因为 DeepSpeed 只支持句子长度为 8 的整数倍，所以这里只测试了长度为 32 的整数倍的句子计算速度：

图 7：softmax 加速对比

可以看出，LightSeq 几乎在所有情况下都远远快于 DeepSpeed。且随着序列长度增加，LightSeq 加速比逐渐增大。而 DeepSpeed 在 16 位浮点数情况下加速比会逐渐减小，在 32 位浮点数情况下甚至会出现比 PyTorch 还要慢的情况。TensorFlow 即便使用 XLA 融合算子，运算速度也远远落后于其它三种实现。

在其它多数 kernel 测试中，LightSeq 都要快于其它三种实现。

关键技术

图 8：模型训练过程

如图 8 所示，以 2 卡为例，一个完整的 Transformer 模型训练过程主要包括四个步骤：前向传播、后向传播、梯度同步和参数更新。其中前向传播和后向传播占模型训练总时间的 70% 多，包含了众多计算密集型和 I/O 密集型操作，因此是优化的重点。而梯度同步虽然耗时相对较少，但是也可以通过和反向传播并行化等方法隐藏掉大部分延时。最后优化器更新参数方面也大有文章可做，可以从计算和 I/O 两个方面降低延时。

下面详细介绍优化这四个步骤的几种技术：算子多运算融合、输入输出层融合、动态显存复用和参数连续化，最后介绍一下 LightSeq 提供的单元测试功能。

算子多运算融合

在大多数深度学习框架（例如 TensorFlow 和 PyTorch）中，一个简单的运算通常都需要很多细粒度的核函数来实现。例如在 TensorFlow 中，一次层归一化（Layer Normalization）操作需要调用三次核函数以及两次中间显存读写，非常耗时。而基于 CUDA，LightSeq 定制化了一个层归一化专用的核函数，将两次中间结果的写入寄存器。从而实现一次核函数调用，同时没有中间结果显存读写，因此大大节省了计算开销。

基于这个思路，LightSeq 利用 CUDA 矩阵运算库 cuBLAS <10> 提供的矩阵乘法和自定义核函数实现了 Transformer 的编码器和解码器。以编码层为例，具体结构如图 9 所示：

图 9：编码层计算过程

蓝色部分是自定义核函数，黄色部分是矩阵乘法。可以发现，矩阵乘法之间的运算全部都用一个定制化核函数实现了，因此大大减少了核函数调用和显存读写，最终提升了运算速度。

LightSeq 还优化了核函数的实现，采用 float4 数据类型来进行数据读写，大大增加了数据的吞吐量，减小了核函数执行的延时。LightSeq 支持任意长度的序列输入，根据不同的序列长度选择最合适的核函数来调用执行。相比之下，DeepSpeed 多数核函数只支持句子长度为 8 的整数倍，因此速度较慢，局限性较大。

输入输出层融合

此外 LightSeq 还实现了词嵌入层和损失函数层的算子融合。对于词嵌入层，LightSeq 将词表查找与放缩、位置向量融合以及 dropout 操作都写成了一个核函数。对于损失函数层，将交叉熵损失融合成一个核函数。通过输入输出层的融合，进一步减小了模型训练的时间，增加了显卡利用率。

以词嵌入层为例，图 10 展示了输入句子中单词 id 之后，词嵌入层进行的计算过程：