字节跳动提出高性能 transformer 推理库，获 IPDPS 2023 最佳论文奖

字节跳动技术团队

于 2023-06-15 12:02:45 发布

阅读量5.6k

点赞数 2

CC 4.0 BY-SA版权

文章标签： transformer 深度学习人工智能

本文链接：https://blog.youkuaiyun.com/ByteDanceTech/article/details/131238326

动手点关注

干货不迷路

字节跳动与英伟达, 加州大学河滨分校联合发表的论文《ByteTransformer: A High-Performance Transformer Boosted for Variable-Length》在第 37 届 IEEE 国际并行和分布式处理大会（IPDPS 2023）中，从 396 篇投稿中脱颖而出，荣获了最佳论文奖。该论文提出了字节跳动的 GPU transformer 推理库——ByteTransformer。针对自然语言处理常见的可变长输入，论文提出了一套优化算法，这些算法在保证运算正确性的前提下，成功避免了传统实现中的冗余运算，实现了端到端的推理过程的大幅优化。另外, 论文中还手动调优了 transformer 中的 multi-head attention, layer normalization, activation 等核心算子，将 ByteTransformer 的推理性提升至业界领先水平。与 PyTorch, TensorFlow, NVIDIA FasterTransformer, Microsoft DeepSpeed-Inference 等知名的深度学习库相比，ByteTransformer 在可变长输入下最高实现131%的加速。论文代码已开源。

ByteTransformer: A High-Performance Transformer Boosted for Variable-Length Inputs ( https://arxiv.org/abs/2210.03052 )

IPDPS: 并行和分布式计算方向计算机系统领域的旗舰会议。该会议专注于分享并讨论并行计算，分布式计算，大规模数据处理以及高性能计算等相关领域的最新研究进展。参与的专家学者来自世界各地的顶尖研究机构和企业，共同探讨该领域的创新发展和前沿技术。

code： https://github.com/bytedance/ByteTransformer

transformer 变长文本 padding free

transformer 在自然语言处理（NLP）中被广泛使用，随着 BERT、GPT-3 等大型模型的出现和发展，transformer 模型的重要性变得越来越突出。这些大型模型通常具有超过一亿个参数，需要大量的计算资源和时间进行训练和推理。因此，优化 transformer 性能变得非常重要。

现有的一些深度学习框架，如 Tensorflow，PyTorch，TVM 以及 NVIDIA TensorRT 等，要求输入序列长度相同，才能利用批处理加速 transformer 计算。然而，在实际场景中，输入序列通常是变长的，而零填充会引入大量的额外计算开销。有一些方法在 kernel launch 前对具有相似 seqlen 的输入分组，以最小化 padding，但无法实现 padding free。字节跳动 AML 团队先前提出的“effective transformer” [4]，通过对输入的重排列，实现了 QKV projection 和 MLP 的 padding free，但 self attention 部分仍然需要 padding。

为了解决这个问题，字节跳动 AML 团队提出了 ByteTransformer，它实现了变长输入的 padding free 计算，并且实现了全面的 kernel fusion 以进一步提高性能。