[2210.03052] ByteTransformer: A High-Performance Transformer Boosted for Variable-Length Inputs

ByteTransformer源码解析

“ByteTransformer: A High-Performance Transformer Boosted for Variable-Length Inputs”由Yujia Zhai等人撰写。文章提出了ByteTransformer，这是一种针对可变长度输入进行优化的高性能Transformer框架，在自然语言处理任务中展现出显著的性能优势。

  1. 研究背景

    • Transformer的广泛应用与挑战：Transformer在自然语言处理等领域应用广泛，但模型参数规模的不断扩大对计算性能提出了更高要求。处理可变长度输入时，现有深度学习框架通过填充序列的方式会带来显著的内存和计算开销。
    • 相关工作与不足：传统深度学习框架为固定长度输入设计，处理可变长度输入时存在冗余计算。部分框架虽支持可变长度输入，但在关键算法优化和内核融合方面存在不足。

  2. ByteTransformer的设计与优化

    • BERT Transformer编码器的数学表达与性能分析：对BERT Transformer编码器的架构和计算流程进行数学表达，通过性能分析确定计算密集型和内存密集型操作的占比，发现GEMM操作和注意力模块是性能瓶颈。
    • 融合内存密集型操作：将层归一化和激活等内存密集型操作与相邻内核融合，减少全局内存访问。融合层归一化和加偏置操作，平均提升单层BERT性能3.2%；融合加偏置和激活操作，进一步提升3.8%。
    • 可变长度输入的零填充算法：提出零填充算法，通过计算掩码矩阵的前缀和来确定有效令牌的位置，打包输入张量，避免对填充令牌的计算。该算法在平均序列长度为最大长度60%时，可使BERT加速24.7%。
    • 优化多头注意力机制：针对短序列，将中间矩阵存储在共享内存和寄存器中，避免冗余计算；针对长序列，采用分组GEMM的融合多头注意力机制，并对软max操作进行融合优化。优化后的MHA比标准PyTorch MHA快6.13倍，使单层BERT性能提升19%。

  3. 实验评估

    • 实验设置：在NVIDIA A100 GPU上进行实验，与TensorFlow、PyTorch等先进Transformer框架对比，使用半精度浮点数格式（FP16），变量序列长度随机生成。
    • 实验结果：内核融合优化在层归一化和加偏置、GEMM和加偏置与激活操作上分别实现69%和24%的性能提升。优化后的MHA在短序列和长序列上均优于其他变体。在单层BERT和端到端BERT性能测试中，ByteTransformer分别比基线版本提升60%和87%-138%。此外，ByteTransformer的优化策略对其他BERT类模型同样有效。

  4. 研究结论：ByteTransformer通过算法创新和架构感知优化，显著提升了可变长度输入下Transformer的性能，在多个模型和任务上超越了现有先进框架。未来将致力于完全开源ByteTransformer，并进一步扩展优化策略，加速更多BERT类模型的推理和训练。