本文是系列博客的第四部分，专注于如何利用纯原生 PyTorch 加速生成式 AI 模型。若想直接查看代码，请访问我们的 GitHub 仓库（seamless_communication, fairseq2）。我们很高兴能分享一系列新发布的 PyTorch 性能特性，并通过实际案例展示如何将 PyTorch 的原生性能发挥到极致。在第一部分中，我们展示了仅使用纯原生 PyTorch 将 Segment Anything 加速超过 8 倍的方法。在第二部分中，我们展示了仅通过原生 PyTorch 优化将 Llama-7B 加速近 10 倍的方法。在第三部分中，我们展示了仅使用原生 PyTorch 优化将 文本到图像扩散模型加速达 3 倍的方法。
在本篇博客中，我们将重点关注如何加速 FAIR 的 Seamless M4T-v2 模型。通过使用 CUDA Graph 和原生 PyTorch 优化，文本解码器模块实现了 2 倍加速，声码器（vocoder）模块实现了 30 倍加速，从而使端到端推理速度提升了 2.7 倍，且准确率毫无损失。
简介
Seamless M4T 是由 FAIR 开发的开源基础语音/文本翻译和转录技术。Seamless M4T 是一款大规模多语言、多模态机器翻译模型，其最新版本（Seamless M4T-v2）于 2023 年 11 月 30 日发布。图 1 展示了 Seamless M4T-v2 的高层模型架构。
图 1. Seamless M4T-v2 的模型架构。
加速推理延迟对于翻译模型至关重要，它能通过实现跨语言的快速交流来提升用户体验。特别是在聊天机器人、语音翻译和实时字幕等应用中，为了降低延迟，通常会使用 batch_size=1。因此，我们对 batch_size=1 的推理性能进行了分析，如图 2 所示，以了解阿姆达尔定律（Amdahl’s Law）的瓶颈所在。结果表明，文本解码器和声码器是最耗时的模块，分别占推理时间的 61% 和 23%。
图 2. 文本解码器和声码器是最耗时的模块。图示为 A100 GPU 上 batch_size=1 时，英西 S2ST（语音到语音/文本）任务的推理时间模块分解。
为了更深入地分析文本解码器和声码器的性能瓶颈，我们分析了 FLEURS 数据集中英西翻译示例的第 8 个样本的 GPU 追踪（Trace），如图 3 所示。分析揭示了文本解码器和声码器是严重的 CPU 密集型模块。我们观察到，CPU 开销导致了明显的延迟，推迟了 GPU 内核（kernel）的启动，从而导致两个模块的执行时间大幅增加。
(a) 文本解码器的 CPU 和 GPU 追踪
(b) 声码器的 CPU 和 GPU 追踪
图 3. 文本解码器和声码器是严重的 CPU 密集型模块。(a) 文本解码器和 (b) 声码器在 A100 GPU 上针对 FLEURS 数据集中英西翻译示例的第 8 个样本进行 batch_size=1 推理时的 CPU 和 GPU 追踪。
基于“Seamless M4T-v2 中的文本解码器和声码器是严重的 CPU 密集型模块”这一实际系统性能分析结果，我们为这些模块启用了 torch.compile + CUDA Graph。在本篇博文中，我们将分享在 batch_size=1 推理场景下，为各模块启用 torch.compile + CUDA Graph 所需的修改，并讨论 CUDA Graph 的相关内容及后续计划。
使用 CUDA Graph 的 Torch.compile
torch.compile 是一款 PyTorch API，允许用户将 PyTorch 模型编译为独立的可执行文件或脚本。它通常用于通过消除不必要的开销来优化模型性能。
CUDA Graph 是 NVIDIA 提供的一项功能，允许优化 CUDA 应用程序中的内核启动。它创建了一个 CUDA 内核的执行图，可以在 GPU 执行之前由驱动程序进行预处理和优化。使用 CUDA Graph 的主要优点是它减少了与启动单个内核相关的开销，因为图可以作为一个单元启动，从而减少了 API 调用次数以及主机与设备之间的数据传输。这可以带来显著的性能提升，特别是对于拥有大量小内核或重复执行同一组内核的应用。如果您有兴趣深入了解，请参阅我们团队成员 Kim Hazelwood 的这篇论文，它强调了数据对加速计算的重要性：数据在哪里？为什么没有答案就无法讨论 CPU 与 GPU 的性能差异。这是在 NVIDIA 大力投资通用 GPU (GPGPU) 领域，且深度学习尚未彻底改变计算行业之前发表的！
然而，由于 CUDA Graph 运行在编译时记录的 1) 固定内存指针和 2) 固定张量形状上，我们针对 CUDA Graph 引入了以下改进：使其能够跨多种输入尺寸重复使用，从而避免在每次迭代时重新生成 CUDA Graph；并允许 CUDA Graph 内部的数据在不同运行之间重用，从而为多个解码步骤共享 KV Cache。
文本解码器
Seamless 中的文本解码器是来自 NLLB [1] 的一个解码器，用于执行 T2TT（文本到文本翻译）。该模块是一个 CPU 密集型模型，由于自回归生成的特性要求顺序处理 Token，GPU 执行时间不足以掩盖 CPU 开销，这限制了 GPU 上可实现的并行度。基于此观察，我们为文本解码器启用了 torch.compile + CUDA Graph，以减少占主导地位的 CPU 开销，如图 4 所示。
图 4. 启用 torch.compile + CUDA Graph 后文本解码器的 CPU 和 GPU 追踪。
1. 更新和检索 KV Cache
在推理过程中，文本解码器有两个计算阶段：消耗 Prompt 的预填充（prefill）阶段，以及逐个生成输出 Token 的增量生成阶段。在 Batch Size 或输入长度足够高的情况下，预填充阶段会在足够多的 Token 上并行操作——此时 GPU 性能是瓶颈，CPU 开销对性能影响不大。另一方面，增量 Token 生成阶段总是以序列长度 1 执行，且通常使用较小的 Batch Size（甚至为 1），例如在交互式用例中。因此，增量生成可能受限于 CPU 速度，是使用 torch.compile + CUDA Graph 的理想候选对象。
然而，在增量 Token 生成阶段，注意力计算中涉及的 Key 和 Value 的 sequence_length 维度每步增加 1，而 Query 的序列长度始终为 1。具体来说，Key/Value 是通过将新计算出的序列长度为 1 的 Key/Value 附加到 KV Cache 中已存储的 Key/Value 上来生成的。但如前所述，CUDA Graph 在编译期间记录所有张量的形状，并在回放时使用记录的形状。因此，我们参考此处的优秀成果进行了一些修改以解决此问题。
a) 我们修改了 KV Cache 的处理方式，使其使用 CUDA 张量中的索引（即 valid_seq_pos）来写入新值，而不是使用 Python 整数。
图 5. KV Cache append 和 get 的修改
b) 我们还修改了注意力机制，使其能够处理 max_seq_length 下 Key 和 Value 的固定形状。我们仅针对序列中直到当前解码步骤（即 valid_seq_pos）的位置计算 Softmax。为了屏蔽掉大于当前解码步骤（即 valid_seq_pos）的序列位置，我们创建了一个布尔掩码张量（即 mask），其中大于 valid_seq_pos 的位置被设置为 False。
图 6. 用于生成 valid_seq_pos 和 mask 的辅助函数
需要指出的是，这些修改会导致所需的计算量增加，因为我们要对超出必要的序列位置（直到 max_seq_length）进行注意力计算。然而，尽管有此缺点，我们的结果证明，与标准 PyTorch 代码相比，torch.compile + CUDA Graph 仍然提供了显著的性能优势。
c) 由于不同的推理样本具有不同的序列长度，这也会生成不同形状的输入，这些输入需要被投影到 Cross-Attention 层所需的 Key 和 Value。因此，我们将输入填充为静态形状，并生成一个填充掩码（padding mask）来屏蔽掉填充部分的输出。
2. 内存指针管理
由于 CUDA Graph 会连同张量形状一起记录内存指针，因此确保不同的推理样本能够正确引用已记录的内存指针（例如 KV Cache）非常重要，从而避免为每个推理样本都编译一次 CUDA Graph。然而，Seamless 代码库的某些部分使不同的推理样本指向不同的内存地址，因此我们进行了修改以改善内存使用情况。
e) Seamless 采用束搜索（Beam Search）作为文本解码策略。在束搜索过程中，我们需要在每个增量解码步骤中对所有注意力层执行 KV Cache 重排（reordering），以确保每个选定的束（beam）在执行时使用对应的 KV Cache，如下面的代码片段所示。
图 8. 束搜索解码策略的 KV Cache 重排操作。
上述代码分配了新的内存空间，并覆盖了 cache_k 和 cache_v 的原始内存指针。因此，我们修改了 KV Cache 重排逻辑，通过使用 copy_ 算子，保留了编译期间记录的每个 Cache 的内存指针。
图 9. 使用 copy_ 算子进行 KV Cache 的原地（In-place）更新
f) 在通过上述修改为文本解码器启用 torch.compile + CUDA Graph 后，文本解码器的开销转移到了 KV Cache 重排上，如图 10 所示。KV Cache 重排重复调用 index_select 96 次（假设有 24 个解码器层，每层包含两种类型的注意力层，分别带有 Key 和 Value 的 Cache）。
图 10. 启用 torch.compile + CUDA Graph 后文本解码器的 CPU 和 GPU 追踪。
作为加速文本解码器工作的一部分，我们额外对 KV Cache 重排应用了 torch.compile，以从内核融合（kernel fusion）中获益，如图 11 所示。请注意，这里不能使用 CUDA Graph（mode='max-autotune'），因为 copy_ 操作会修改输入，这违反了 torch.compile 中 CUDA Graph 版本对静态输入的要求。
图 11. 将 torch.compile 应用于 KV Cache 重排。
对 KV Cache 重排启用 torch.compile 的结果是，原本单独启动的 GPU 内核（图 12(a)）现在被融合了，因此需要启动的 GPU 内核数量大幅减少（图 12(b)）。
(a) 启用 torch.compile **前** KV Cache 重排的 CPU 和 GPU 追踪
(b) 启用 torch.compile **后** KV Cache 重排的 CPU 和 GPU 追踪
图 12. KV Cache 重排启用 torch.compile (a) 前后 的 CPU 和 GPU 追踪对比
声码器
Seamless 中的声码器是一个 HiFi-GAN 单元声码器（unit-vocoder），它将生成的单元转换为波形输出。其中，“单元”是一种语音表示，结合了音素和音节等不同方面，可用于生成人类可听的声音。声码器是一个相对简单的模块，由 Conv1d 和 ConvTranspose1d 层组成，且如图 3 所示，它是一个 CPU 密集型模块。基于此观察，我们决定为声码器启用 torch.compile + CUDA Graph，以减少如图 10 所示的过高 CPU 开销。但在执行过程中需要进行几处修复。
图 13. 启用 torch.compile + CUDA Graph 后声码器的 CPU 和 GPU 追踪。
a) 声码器的输入张量形状在不同推理样本间各不相同。由于 CUDA Graph 会记录张量形状并进行重放，我们必须将输入填充为零以达到固定大小。由于声码器仅由 Conv1d 层组成，我们不需要额外的填充掩码，用零进行填充即已足够。
b) 声码器包含由 torch.nn.utils.weight_norm 包裹的 conv1d 层（参见此处）。然而，直接将 torch.compile 应用于声码器会导致如下的图中断（graph break），从而导致性能提升不佳。这种图中断发生在 PyTorch 代码的 weight_norm 的 Hook 处理部分。
[1/0_2] torch._dynamo.symbolic_convert.__graph_breaks: [DEBUG] Graph break: setattr(UserDefinedObjectVariable) <function Module.__setattr__ at 0x7fac8f483c10> from user code at:
[1/0_2] torch._dynamo.symbolic_convert.__graph_breaks: [DEBUG] File "/mnt/fsx-home/yejinlee/yejinlee/seamless_communication/src/seamless_communication/models/vocoder/vocoder.py", line 49, in forward
[1/0_2] torch._dynamo.symbolic_convert.__graph_breaks: [DEBUG] return self.code_generator(x, dur_prediction) # type: ignore[no-any-return]1/0_2] torch._dynamo.symbolic_convert.__graph_breaks: [DEBUG] File "/data/home/yejinlee/mambaforge/envs/fairseq2_12.1/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1520, in _call_impl
[1/0_2] torch._dynamo.symbolic_convert.__graph_breaks: [DEBUG] return forward_call(*args, **kwargs)
[2023-12-13 04:26:16,822] [1/0_2] torch._dynamo.symbolic_convert.__graph_breaks: [DEBUG] File "/mnt/fsx-home/yejinlee/yejinlee/seamless_communication/src/seamless_communication/models/vocoder/codehifigan.py", line 101, in forward
[1/0_2] torch._dynamo.symbolic_convert.__graph_breaks: [DEBUG] return super().forward(x)
[1/0_2] torch._dynamo.symbolic_convert.__graph_breaks: [DEBUG] File "/mnt/fsx-home/yejinlee/yejinlee/seamless_communication/src/seamless_communication/models/vocoder/hifigan.py", line 185, in forward
[1/0_2] torch._dynamo.symbolic_convert.__graph_breaks: [DEBUG] x = self.ups[i](x)
[1/0_2] torch._dynamo.symbolic_convert.__graph_breaks: [DEBUG] File "/data/home/yejinlee/mambaforge/envs/fairseq2_12.1/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1550, in _call_impl
[1/0_2] torch._dynamo.symbolic_convert.__graph_breaks: [DEBUG] args_result = hook(self, args)
[1/0_2] torch._dynamo.symbolic_convert.__graph_breaks: [DEBUG] File "/data/home/yejinlee/mambaforge/envs/fairseq2_12.1/lib/python3.8/site-packages/torch/nn/utils/weight_norm.py", line 65, in __call__
[1/0_2] torch._dynamo.symbolic_convert.__graph_breaks: [DEBUG] setattr(module, self.name, self.compute_weight(module))
[1/0_2] torch._dynamo.symbolic_convert.__graph_breaks: [DEBUG]
由于层权重在推理过程中不会改变，我们不需要权重归一化。因此，我们利用 Seamless 代码库中已提供的 remove_weight_norm 函数（此处），简单地移除了声码器的权重归一化，如图 14 所示。
图 14. 为声码器移除 weight_norm
性能评估 + CUDA Graph 的影响
图 15 显示了在文本解码器和声码器上启用 torch.compile(mode=”max-autotune”) + CUDA Graph 后的加速结果。我们实现了文本解码器 2 倍加速，声码器 30 倍加速，从而使端到端推理时间加快了 2.7 倍。
 |  |
图 15. 应用 torch.compile 和 torch.compile + CUDA Graph 后，文本解码器和声码器的推理时间加速情况。
我们还报告了使用 torch.compile 但不带 CUDA Graph（即使用 torch.compile(mode="max-autotune-no-cudagraphs")）时文本解码器和声码器的加速比，以便识别 CUDA Graph 对性能的影响。在没有 CUDA Graph 的情况下，文本解码器和声码器的加速比分别降至 1.17 倍和 18.4 倍。虽然效果依然显著，但这也凸显了 CUDA Graph 的重要作用。我们的结论是，Seamless M4T-v2 在启动 CUDA 内核上耗费了大量时间，特别是在使用较小 Batch Size（如 1）时，GPU 内核的执行时间不足以摊销内核启动的开销。
图 16. 逐步应用 torch.compile 和 CUDA Graph 的端到端推理加速。a) “Inc. Decoding”：仅对文本解码器应用 torch.compile。b) “Inc. Decoding w/ CUDA Graph”：对文本解码器应用 torch.compile + CUDA Graph。c) “+KV Cache Reordering”：在 b) 的基础上，额外对 KV Cache 重排操作应用 torch.compile。d) “+Vocoder”：在 c) 的基础上，额外对声码器应用 torch.compile。e) “+Vocoder w/ CUDA Graph”：在 d) 的基础上，额外对声码器应用 torch.compile + CUDA Graph。
图 16 展示了对各模块应用 torch.compile（带或不带 CUDA Graph）的累加效应。结果表明，端到端推理加速有显著改善，证明了这些技术在优化整体延迟方面的有效性。最终，我们使 Seamless M4T-v2 在 batch_size=1 时获得了 2.7 倍的端到端推理加速。
致谢
感谢 PyTorch 团队和 Seamless 团队对本项目提供的巨大支持。