NVIDIA GPU 与算子优化助力 vivo 文本预训练大模型训练提速 60%

一、方案背景与核心挑战

vivo 是一家以智能终端和智慧服务为核心的科技公司，其 AI 团队研发了面向自然语言理解任务的文本预训练模型 3MP-Text。在中文语言理解测评基准 CLUE 榜单上，3MP-Text 1 亿参数模型效果排名同规模第一，7 亿参数模型排名总榜第十（不包括人类）；在内容理解、舆情分析、语音助手等内部场景测试中，3MP-Text 1 亿模型效果明显优于同规模开源模型。核心挑战：为提升预训练模型效果，往往需要对模型结构做修改（如改变位置编码实现方式、调整模型宽度与深度），而这些修改可能造成训练速度下降。3MP-Text 采用 Deberta V2 结构，相对位置编码的实现增加了注意力机制部分的计算量——在 NVIDIA GPU 单卡测试中，含有相对位置编码的注意力机制计算耗时占单次迭代耗时的 71.5%；同时模型采用减小隐层维度、增加层数的 DeepNarrow 结构，进一步提升了效果但加重了训练负担。两方面因素叠加，使 3MP-Text 相比同参数规模 BERT 模型训练时间多 60%，直接影响业务功能上线节奏。

二、解决方案名称与架构

解决方案名称：基于 NVIDIA GPU 与算子级优化的 3MP-Text 训练加速方案

核心技术：vivo AI 团队与 NVIDIA 团队合作，通过 NVIDIA Nsight Systems（nsys）性能分析工具进行性能瓶颈分析，并在此基础上针对 gather、dropout、softmax、scale、layernorm 等算子进行优化。

三、核心技术组件

Gather 与 kernel 融合：在 cuda kernel 实现上采用 float4/half4 等数据类型进行向量化读写，并利用 shared memory 确保合并访问，从而优化 gather（前向）/scatter（反向）kernels；同时通过 kernel 融合把 gather 上下游的 elementwise 操作合并到单个 cuda kernel，融合前完成相应计算需要 9 个 cuda kernels，融合后只需 4 个。综合两项优化，该模块性能提升 3.3 倍。

Dropout 与 Softmax 优化：将 StableDropout 的多步 elementwise 操作融合到一个独立 cuda kernel 中，并配合向量化读写，kernels 数量由 9 个降至 3 个，模块性能提升 4.5 倍。Softmax 方面将掩码填充与归一化步骤融合到单个 cuda kernel，借助向量化读写减少参与同步的线程数目，并利用寄存器数组缓存数据避免多次访问全局内存，前向性能提升 4 倍。

Scale 与 Layernorm 优化：针对注意力计算中的除法 scale 操作，直接调用 cublasGemmStridedBatchedEx() API 实现融合 gemm + scale 的 torch op，取得 1.9 倍加速比；同时通过改造 Apex 中的 layer_norm 模块，适配 hidden dim=512 等特定结构以获得更优性能。

四、实施成果

通过上述算子优化，vivo 文本预训练大模型在实际应用中训练提速 60%，满足了下游业务应用对模型训练速度的要求。此前因模型训练时间比 BERT 长 60% 而可能延迟上线的领域预训练任务，得以恢复正常迭代优化节奏，3MP-Text 模型在手机舆情等场景的应用障碍得到消除。

五、未来展望

面向自然语言理解与多模态场景的持续演进，NVIDIA 与 vivo 将继续在算子级别挖掘性能空间，帮助模型在效果提升的同时保持可接受训练成本。