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

描述

 

 

vivo AI 团队与 NVIDIA 团队合作,通过算子优化,提升 vivo 文本预训练大模型的训练速度。在实际应用中,训练提速 60%,满足了下游业务应用对模型训练速度的要求。通过 NVIDIA Nsight Systems 性能分析工具进行性能瓶颈分析,并在此基础上,针对 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 的模型结构,该结构使用相对位置编码,相对于绝对位置编码,效果更好,但其相对位置编码的实现过程,增加了模型在注意力机制部分的计算量,从而降低了模型的训练速度。如图 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%。

 

英伟达

图 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)。

 

英伟达

图 7. Attention 部分,scale 操作相关源码。

 

 Layernorm 算子优化:

 

除了上述提到算子外,还通过改造 apex 中的 layer_norm 模块(https://github.com/NVIDIA/apex/tree/master/apex/contrib/csrc/layer_norm),以便在 hidden dim=512 情况下,优化 layernorm 算子,取得了 2.4 倍的加速比(优化前 0.53 ms,优化后 0.22 ms)。

 

使用效果及影响

 

使用 NVIDIA 做的算子优化,vivo 3MP-Text 模型的训练速度提升 60%,达到了和同规模 BERT 模型相同的速度,下游业务应用时,模型的训练速度不再成为瓶颈,训练成本进一步降低。另外,这些算子优化,也可以应用到其他使用 Deberta V2 模型的场景中。

 

未来,vivo AI 团队和 NVIDIA 将在大模型分布式训练、推理等方面持续合作,共同推进生成式 AI 在手机场景行业的应用落地(如语音助手、智能创作、智能办公等)和性能提升。

 

       点击 “阅读原文”,或扫描下方海报二维码,在 5 月 29 日观看 NVIDIA 创始人兼 CEO 黄仁勋为 COMPUTEX 2023 带来的主题演讲直播,了解 AI、图形及其他领域的最新进展!  

 


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

文章出处:【微信公众号:NVIDIA英伟达】欢迎添加关注!文章转载请注明出处。


打开APP阅读更多精彩内容
声明:本文内容及配图由入驻作者撰写或者入驻合作网站授权转载。文章观点仅代表作者本人,不代表电子发烧友网立场。文章及其配图仅供工程师学习之用,如有内容侵权或者其他违规问题,请联系本站处理。 举报投诉

全部0条评论

快来发表一下你的评论吧 !

×
20
完善资料,
赚取积分