字节跳动与英伟达, 加州大学河岸分校联合发表的论文《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 (arxiv.org/abs/2210.03… )

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

code： github.com/bytedance/B…

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 以进一步进步功能。

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

图一: ByteTransformer 与其他工作 feature 比照

Remove padding 算法

这个算法源自字节跳动 AML 团队之前的工作 “effective transformer”，在 NVIDIA 开源 FasterTransformer 中也有集成。ByteTransformer 同样运用该算法去除对 attention 外矩阵乘的额定核算。

算法过程：

1）核算 attention mask 的前缀和，作为 offsets

2）依据 offsets 把输入张量从[batch_size, seqlen, hidden_size]重摆放为valid_seqlen, hidden_size]，再参加后续的矩阵乘核算，完结 padding free

图二：Remove padding 算法进程

FMHA (Fused Multi-Head Attention)

为了优化 attention 部分的功能，ByteTransformer 中完结了 fused multi-head attention 算子。关于 seqlen 长度, 以 384 为界划分为两种完结方式:

关于短 seqlen, 由于能够把 QK 整行放在共享内存进行 softmax 操作, 经过手写 kernel的方式完结, 矩阵乘经过调用 wmma 接口运用 TensorCore 确保高功能。
关于长 seqlen, 由于共享内存巨细约束，不能在一个手写 kernel 中完结一切操作。基于高功能的CUTLASS[5] groupedGEMM, 分成两个 gemm kernel 完结, 并把 add_bias, softmax 等操作 fused 到 GEMM kernel 中。

CUTLASS groupedGEMM

NVIDIA 开发的 grouped GEMM 能够在一个 kernel 中完结多个独立矩阵乘问题的核算，使用这个性质能够完结 Attention 中的 padding free。

Attention 中的两次矩阵乘操作，都能够拆解为batch_size x head_num个独立的矩阵乘子问题。
每个矩阵乘子问题，把问题巨细传入到 grouped GEMM，其间 seqlen 传递真实的 valid seqlen 即可。

grouped GEMM 原理：kernel 中每个 threadblock (CTA) 固定 tiling size，每个矩阵乘子问题依据 problem size 和 tiling size，拆解为不同数量的待核算块，再把这些块均匀分配到每个 threadblock 中进行核算。

图三：grouped GEMM 原理示意图。每个子问题拆解为不同数量的块，再对这些块均匀分配，高效地完结单个 kernel 核算多个独立 GEMM 问题。

运用 grouped GEMM 完结 attention 时，由于子问题的数量batch_size x head_num一般较大，读取子问题参数会有不小的开支，由于从线程角度看，每个线程都需求遍历读取一切的子问题巨细。

为了解决这个问题，ByteTransformer 对 grouped GEMM 中读取子问题参数进行了功能优化，使其能够忽略不计：

1）共享子问题参数。 对同一个输入，不同 head 的 valid seqlen 相同，problem size 也相同，经过共享使参数存储量从batch_size x head_num削减到batch_size

2）warp prefetch. 原始完结中，每个 CUDA thread 依次读取一切的子问题 problem size，效率很低。改为一个 warp 内线程读取连续的 32 个子问题参数，然后经过 warp 内线程通讯交流数据，每个线程的读取次数降低到 1/32

图四：warp prefetch 示意图。每个 iteration 一个 warp 读取 32 个子问题 size

softmaxfusion

为了进一步进步功能，把 Q x K 之后的 softmax 也 fuse 到矩阵乘算子中，比较单独的 softmax kernel 节省了中心矩阵的访存操作。

由于 softmax 需求对整行数据做归约，但由于共享内存巨细的约束，一个 threadblock 内不能包容整行数据，一起 threadblock 间的通讯很低效，所以不能仅在 Q x K 的 epilogue 中完结整个 softmax 的操作。把 softmax 拆分成三步核算，别离 fuse 到 Q x K 的 epilogue 中， QK x V 的 prologue 中，以及中心再增加一个轻量的 kernel 做规约。

图五：softmax fusion 流程示意图。分为三步核算，大部分核算 fuse 到前后的 GEMM kernel 中

算法过程：

1）partial reduction：Q x K 的 epilogue 中，每个 threadblock 内部规约，核算出 max 和 sum 两个值

2）full reduction：一个轻量级的 kernel，把每一行的 partial reduction 成果持续规约到整行的成果

3）element-wise op：修改了 CUTLASS 的代码，使其支撑 prologue fusion，即在加载输入矩阵后，fuse 一些 element-wise 的操作。在 QK x V 的 prologue 中，读取当前行的规约成果，核算出 softmax 的终究成果，再参加后续的矩阵乘核算

功能数据

短 seqlen 手写 kernel 的功能

在 <= 384的短 seqlen 情况下，cuBLAS batch GEMM 比较 PyTorch MHA 功能进步 5 倍，而启用 zero padding 算法优化 softmax 后进一步进步 9% 的功能，ByteTransformer 的 MHA 把两个 batched GEMM 和中心的 softmax 彻底 fuse 成一个 kernel，比较三种变体完结，均匀加快别离为 617%、42% 和 30% 。

图六：手写 attention kernel 功能比照。注：cuBLAS + zero padding 指对 softmax 的 zero padding

长 seqlen CUTLASS kernel 的功能

在 448~1024 seqlen 下， cuBLAS batched GEMM 比 PyTorch 的 MHA 功能进步了3倍，一起对 softmax 的 zero padding，进一步进步了 17% 的功能，经过引进高功能的 CUTLASS grouped GEMM 以及 softmax fusion，ByteTransformer 的 fused MHA 比变体 MHA 完结的功能进步了451%，110% 和 79% 。

图七：长 seqlen CUTLASS FMHA 功能比照

全面的 kernel fusion

除矩阵乘和 attention 的优化外，ByteTransformer 还对一些小的操作进行了全面的 kernel fusion，经过削减显存拜访和 kernel launch 的开支，能够获得更极致的功能。

add-bias & LayerNorm fusion

矩阵乘之后的 add-bias 和 LayerNorm 操作，经过手写 kernel 的方式做 fusion，这部分操作在 seqlen 为 256 和 1024 的情况下别离占 10% 和 6% 的推迟，fused kernel 能够优化 61% 的功能，对单层 BERT transformer 的功能进步 3.2% （均匀 seqlen 128 – 1024 的情况）。

GEMM & add-bias & GELU fusion

经过 CUTLASS fuse epilogue 的方式，把矩阵乘后的 add-bias 操作和 GELU activation 操作 fuse 到矩阵乘 kernel 中。add-bias 和 GELU 在 seqlen 为 256 和 1024 的情况下占比耗时别离为 7% 和 5%。把 add-bias 和 GELU 融合 GEMM 能够完美躲藏这部分的访存推迟，进一步使单层 transformer功能进步 3.8% 。

变种 transformer 支撑

现在，字节跳动 AML 团队已经在 GitHub 上开源了 ByteTransformer 的规范 BERT 完结（github.com/bytedance/B… 。除此之外，字节内部版别还支撑了许多 transformer 变种，比如 Deberta, Roformer，T5 等等。代码完结易于拓宽，而且上述各种优化手法也能够方便地使用到变种 transformer 中。

定论

ByteTransformer 是一种高效的 transformer 完结，它经过一系列优化手法，完结了在 BERT transformer 上的高功能体现。关于变长文本输入，比较其他 transformer 完结，ByteTransformer 具有显着的优势，试验中均匀加快可达 50% 以上。适用于加快自然语言处理使命，进步模型练习与推理的效率。一起，ByteTransformer 也为其他研究者提供了一种高效的 transformer 完结方式。其优化手法和功能体现关于实践使用具有重要意义。

引证：

[1] J. Fang, Y. Yu, C. Zhao, and J. Zhou, “Turbotransformers: an efficient gpu serving system for transformer models,” in Proceedings of the 26th ACM SIGPLAN Symposium on Principles and Practice of Parallel Programming, 2021, pp. 389–402.

[2] R. Y. Aminabadi, S. Rajbhandari, M. Zhang, A. A. Awan, C. Li, D. Li, E. Zheng, J. Rasley, S. Smith, O. Ruwase et al., “Deepspeed inference: Enabling efficient inference of transformer models at unprecedented scale,” arXiv preprint arXiv:2207.00032, 2022.

[3] NVIDIA, github.com/NVIDIA/Fast…

[4] ByteDance github.com/bytedance/e…

[5] NVIDIA github.com/NVIDIA/cutl…

[6] T. Dao, D. Y. Fu, S. Ermon, A. Rudra, and C. Re, “Flashattention: Fast and memory-efficient exact attention with io-awareness,” arXiv preprint arXiv:2205.14135, 2022.

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