孤本网
/ 1 阅读
0
0

低精度多项式近似技术可在 GB200 架构上提升大语言模型训练吞吐量

一句话结论

低精度多项式替代原生激活函数可提升 NVIDIA GB200 架构的训练吞吐量最高达 8.0%。

关键要点

  • 测试采用 Bfloat16 格式下 3 次或 4 次多项式程序,替代原生 Sigmoid、Tanh 和 SiLU 激活函数。
  • 在 L2 缓存驻留工作集中,打包融合乘加程序相比原生 PyTorch 评估实现了 1.19 倍至 2.19 倍加速。
  • 在 HBM 驻留工作集中,多项式程序实现了 1.00 倍至 1.70 倍加速,部分场景下加速比接近 1.0 倍。
  • 门控线性单元替换使完整训练步吞吐量提升 8.0%,Tanh 软上限注意力提升 2.9%,稠密 SiLU 提升 2.7%。
  • 在约 1000 亿 token 训练时程下,多项式与原生实现的平滑训练损失差异介于 -0.107 至 +0.079 之间。

背景与事实

随着 GPU 架构的演进,矩阵运算、特殊函数单元和内存流水线的扩展速率存在不均衡。NVIDIA 的 FlashAttention-4 架构在 Blackwell 平台中暴露了注意力计算内部的这一瓶颈,促使研究者探索替代方案。相关研究于 2026 年 9 月 3 日在计算机科学机器学习领域发表,主要测试了短多项式程序在加速大语言模型中特殊函数单元操作方面的可行性。研究团队在隔离的 IEEE binary16 格式下,对比了原生 PyTorch 评估与打包融合乘加程序的执行效率,工作集覆盖 L2 缓存驻留与 HBM 驻留两种场景。

在具体的集成任务中,研究采用 4 块 GB200 计算芯片,使用 Bfloat16 格式的 3 次或 4 次多项式程序替换了原生 Sigmoid、Tanh 和 SiLU 激活函数。多项式程序结合了解析对称性、目标格式舍入策略以及消费内核内部的打包算术。测试涵盖四个任务:稠密 SiLU、Tanh 软上限注意力、Sigmoid 注意力和路由专家 Swish 门控线性单元。性能数据显示,稠密 SiLU、Tanh 软上限注意力和路由专家替换分别使完整训练步吞吐量提升了 2.7%、2.9% 和 8.0%。对于 Sigmoid 注意力替换,完整注意力前向计算提升了 7.4%,但完整 GPU 步骤仅提升 0.3%。

为了验证模型行为层面的影响,研究进行了同检查点开源权重消融实验,并在每个任务中执行了一次配对预训练比较。在接近 1000 亿 token 的常见训练时程中,四个任务的最终平滑训练损失差异(多项式减去原生)在 -0.107 到 +0.079 之间波动。研究未披露具体的模型参数量、数据集构成或优化器超参数,损失差异的统计显著性水平也未在摘要中明确说明。多项式系数由分析推导得出,舍入策略直接嵌入内核而非在计算后应用。

影响分析

对于中文开发者与从业者而言,这一结果提供了明确的性能优化方向,但需注意其特定的硬件前提。在 NVIDIA GB200 集群上运行的 Bfloat16 训练任务中,替换激活函数为 3 次或 4 次多项式可获得 2.7% 至 8.0% 的吞吐收益,对大规模预训练而言是可观的算力节省。判断认为,在 L2 缓存驻留场景下加速效果显著(最高 2.19 倍),而 HBM 驻留场景收益有限(最高 1.70 倍),这意味着数据重排或算子融合策略需结合内存层级进行调整。

对于模型收敛性而言,损失差异在正负 0.1 以内的波动范围表明多项式近似在特定条件下具有可行性,但该结论仅基于 100 亿 token 时程。若训练目标对损失精度要求极高,或在超过该时程的长周期训练中,多项式近似的误差累积效应缺乏验证,盲目应用可能带来收敛风险。开发者需评估自身任务对吞吐量与损失精度的优先级排序,再决定是否采用该方案。

适用边界

该结论严格限制于 NVIDIA GB200 硬件架构与 Bfloat16 数据格式,不支持其他 GPU 代次或 IEEE binary16 原生评估场景。多项式程序仅替代了 Sigmoid、Tanh 和 SiLU 三类激活函数,其他特殊函数如 Softmax、LayerNorm 或矩阵乘法本身不在该测试范围内。在内存受限的 HBM 驻留场景下,多项式程序部分配置未表现出显著优势,此时原生实现可能更优。

此外,模型行为验证仅覆盖四个特定集成任务,且损失差异统计基于约 100 亿 token 时程,不能推广至其他模型架构或更长训练周期。路由专家任务 8.0% 的吞吐收益不能直接套用至稠密模型或非门控线性单元场景。在 L2 缓存不足或工作集超出 L2 容量的实际部署中,1.0 倍至 1.70 倍的 HBM 加速比是更现实的预期。

孤本观察

研究将 FlashAttention-4 暴露的内存瓶颈直接映射到激活函数替换策略,但未对比 FlashAttention-4 内部多项式近似方案与本文独立内核方案的性能差异。100 亿 token 时程下的损失差异范围较宽,正负波动表明多项式系数对特定任务敏感,缺乏统一的误差界保证。

孤本观察

研究将 FlashAttention-4 暴露的内存瓶颈直接映射到激活函数替换策略,但未对比 FlashAttention-4 内部多项式近似方案与本文独立内核方案的性能差异。100 亿 token 时程下的损失差异范围较宽,正负波动表明多项式系数对特定任务敏感,缺乏统一的误差界保证。

低精度多项式近似技术可在 GB200 架构上提升大语言模型训练吞吐量

低精度多项式近似技术可在 GB200 架构上提升大语言模型训练吞吐量

低精度多项式近似技术可在 GB200 架构上提升大语言模型训练吞吐量

低精度多项式近似技术可在 GB200 架构上提升大语言模型训练吞吐量

低精度多项式近似技术可在 GB200 架构上提升大语言模型训练吞吐量

低精度多项式近似技术可在 GB200 架构上提升大语言模型训练吞吐量

常见问题

在 NVIDIA GB200 架构上,将激活函数替换为多项式后,完整训练步吞吐量最高提升了多少?

门控线性单元替换使完整训练步吞吐量最高提升了 8.0%。

多项式近似替代原生激活函数时,使用的多项式次数和数据格式是什么?

研究采用 Bfloat16 格式下的 3 次或 4 次多项式程序。

在 L2 缓存驻留和 HBM 驻留场景下,多项式程序相比原生评估的加速比分别是多少?

L2 缓存驻留场景加速比为 1.19 至 2.19 倍,HBM 驻留场景加速比为 1.00 至 1.70 倍。

约 1000 亿 token 训练时程下,多项式与原生实现的平滑训练损失差异范围是多少?

差异介于 -0.107 至 +0.079 之间。

来源:arXiv cs.LG


评论