NVIDIA AI 技术助力 vivo 文本预训练大模型性能提升

简介

本案例中，vivo AI 团队与 NVIDIA 团队合作，通过算子优化，提升 vivo 文本预训练大模型的训练速度。在实际应用中，训练提速 60%，满足了下游业务应用对模型训练速度的要求。
本案例主要做的优化包括：通过 NVIDIA Nsight Systems (nsys) 性能分析工具进行性能瓶颈分析，并在此基础上，针对 gather、dropout、softmax、scale、layernorm 等算子进行优化。

客户简介及应用背景

vivo 是一家以设计驱动创造伟大产品，以智能终端和智慧服务为核心的科技公司。自 2017 年开始，vivo 不断地思考着如何通过 AI 技术能力，为全球超过 4 亿的用户提供更好的智能服务。基于此愿景，vivo 打造了针对消费互联网场景的 1001 个 AI 便利。其中，vivo AI 团队研发了面向自然语言理解任务的文本预训练模型 3MP-Text。在中文语言理解测评基准 CLUE 榜单上，3MP-Text 1 亿参数模型效果排名同规模第一，7 亿参数模型排名总榜第十（不包括人类）；在 vivo 内部的多个应用场景如内容理解、舆情分析、语音助手上测试，3MP-Text 1 亿模型效果明显优于同规模开源模型，展现出优秀的中文语言理解能力，具有良好的应用价值。

*此图片来源于 vivo*
*如果您有任何疑问或需要使用此图片，请联系 vivo

客户挑战

为提升预训练模型的效果，往往需要对模型的结构做一定修改，（比如改变位置编码的实现方式，改变模型的宽度和深度等）而这些修改，可能造成模型训练速度的下降。

3MP-Text 模型，采用 Deberta V2^[2]的模型结构，该结构使用相对位置编码，相对于绝对位置编码，效果更好，但其相对位置编码的实现过程，增加了模型在注意力机制部分的计算量，从而降低了模型的训练速度。如图 1 所示，在 NVIDIA GPU 单卡测试中，含有相对位置编码的注意力机制的计算耗时占了单次迭代耗时的 71.5%。

另一方面，已有的研究和实践验证显示，相同参数规模下，减小模型隐层维度，增加模型层数，能提升效果（效果对比见图 2），因此，3MP-Text 模型采用了这种 DeepNarrow 的结构。

图1. Deberta V2 xlarge 模型在 NVIDIA GPU 单卡上，batch size = 20 时一次迭代的 nsys timeline。单次迭代耗时 965ms，含有相对位置编码的注意力机制（DisentangledSelfAttention）前后向计算耗时 690ms，占比 71.5%。

Param	Bert	Bert-DeepNarrow
Hidden size	768	512
intermediate_size	3072	2048
num_attention_heads	12	8
attention_head_size	64	64
num_hidden_layers	12	24
CLUE dev score	67.48	68.22（+0.74）

图2. 实际验证的数据表格

以上两点修改，使得 3MP-Text 模型，相比同参数规模的 BERT 模型，训练时间多 60%，训练成本相应增加，对模型在实际业务场景的应用，造成一定障碍。比如，采用领域预训练的方法，提升 3MP-Text 模型在手机舆情领域任务上的整体表现，由于模型训练时间比 BERT 长 60%，采用该模型会使业务功能的上线时间明显延迟，从而影响了正常迭代优化。

应用方案

本案例将以 NVIDIA GPU 单卡训练情况为例，展开介绍 NVIDIA 所进行的算子优化。

如上文提到，含有相对位置编码的注意力机制计算耗时占比达 71.5%，因此，NVIDIA 团队优先对该模块进行了优化，其中包括 gather 算子、dropout 算子、softmax 算子和 scale 算子的优化。

Gather 算子优化：

对于 gather 操作本身，在 cuda kernel 实现方面，采用了 float4/half4 等数据类型进行向量化读写（一次读写 4 个 float 或 4个 half 元素），并且利用 shared memory 确保合并访问，从而优化 gather（前向）/scatter（反向）cuda kernels。

除了 gather 本身的优化外，如图 3 所示的 pytorch 代码中看到，有不少 elementwise 的操作（红框所示）可以通过 kernel 融合（kernel fusion）的优化手段，把它们都融合到一个 cuda kernel（蓝框所示）中，从而提升性能。如图 4 所示，在进行 kernel 融合前，完成相应计算需要 9 个 cuda kernels，kernel 融合后，只需要4 个 cuda kernels。

综合 gather kernel 优化和 kernel 融合优化，该模块性能提升 3.3 倍。

图3. gather 及相关操作的 pytorch 源码。红框为 gather 操作上下游的 elementwise 操作。蓝框示意进行 kernel 融合后，对应 cuda kernel 所执行的全部操作。

图4. gather 及相关操作优化 nsys timeline 对比。(a) 优化前，前向(fw)6 个 cuda kernels 耗时 2.6ms，反向(bw)3 个cuda kernels耗时 3.6 ms；(b)优化后，前向(fw)2 个 cuda kernels 耗时 0.88ms，反向(bw)2 个 cuda kernels 耗时 0.99ms。优化后加速比 3.3x。

Dropout 算子优化：

在 debertaV2 中会使用 StableDropout，如果仔细对比 pytorch 代码，会发现其计算公式绝大部分情况下可以简化为：

Step 1. rand_data = torch.rand_like(input)

Step 2. x.bernoulli_(1 – dropout) == rand_data < (1 – dropout)

Step 3. mask = (1 – torch.empty_like(input).bernoulli_(1 – dropout)).to(torch.bool)

Step 4. input.masked_fill(mask, 0) * (1.0 / (1 – dropout))

显然上述操作涉及大量的 elementwise 的操作，因此把 step 2~4 融合到一个独立的 cuda kernel中，同时再次采用了 float4/half4 等数据类型进行向量化读写来优化 cuda kernel。

如图 5 所示，在进行 kernel 融合前，完成相应计算需要 9 个 cuda kernels，kernel 融合后，只需要 3 个 cuda kernels。

综合 dropout kernel 优化和 kernel 融合优化，该模块性能提升 4.5 倍。

图5. dropout 及相关操作优化 nsys timeline 对比。(a) 优化前，前向(fw)6 个 cuda kernels 耗时 3.4ms，反向(bw)3 个 cuda kernels 耗时 1.9 ms；(b)优化后，前向(fw)2 个 cuda kernels 耗时 0.82ms，反向(bw)1 个 cuda kernels 耗时 0.37ms。优化后加速比 4.5x。

Softmax 算子优化：

与 dropout 类似，根据源码对 Softmax 算子的计算步骤进行如下划分：

Step 1. rmask = ~(mask.to(torch.bool))

Step 2. output = input.masked_fill(rmask, torch.tensor(torch.finfo(input.dtype).min))

Step 3. output = torch.softmax(output, self.dim)

Step 4. output.masked_fill_(rmask, 0)

把step 1~4 融合到一个独立的 cuda kernel 中。由于 softmax 计算中涉及 cuda 线程之间的同步操作，当采用float4/half4等数据类型进行向量化读写时，也减少了参与同步的 cuda 线程数目，从而减少了同步的开销。此外，NVIDIA 团队也利用寄存器数组来缓存数据，避免了多次从全局内存中读取数据。

在 softmax 优化中，只优化了其前向，沿用了原有的反向实现。如图 6 所示，经过优化后，该模块前向性能提升 4 倍。

图6. softmax 及相关操作优化 nsys timeline 对比。(a) 优化前，前向(fw)6 个 cuda kernels 耗时 2.1ms；(b)优化后，前向(fw)1 个 cuda kernels 耗时 0.5ms。优化后加速比 4x。

Scale 算子优化：

如图 7 所示，在 attention 部分，计算 attention score 时会有一个除以 scale 的操作，这个除法操作其实可以很容易通过 cublas 的 API 融合到矩阵乘法之中，因此，直接调用 cublasGemmStridedBatchedEx() API，实现了一个融合 gemm + scale 的 torch op。取得了 1.9x 的加速比（优化前 1.42 ms，优化后 0.75 ms）。