本文是系列博客的第四部分，该系列专注于如何使用纯原生 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 倍的加速和声码器模块 30 倍的加速，从而使端到端推理速度提升了 2.7 倍，且精度无损。
简介
Seamless M4T 是由 FAIR 开发的一项开源基础语音/文本翻译和转录技术。Seamless M4T 是一个大规模多语言和多模态机器翻译模型，其最新版本（Seamless M4T-v2）于 2023 年 11 月 30 日发布。Seamless M4T-v2 的高级模型架构如图 1 所示。
图 1. Seamless M4T-v2 的模型架构。
加速推理延迟对于翻译模型至关重要，它能通过更快的跨语言沟通来提升用户体验。特别是在 batch_size=1 的情况下，这在聊天机器人、语音翻译和实时字幕等对延迟要求很高的应用中非常常见。因此，我们对 batch_size=1 的推理进行了性能分析，如图 2 所示，以了解阿姆达尔定律的瓶颈。我们的结果表明，文本解码器和声码器是耗时最多的模块，分别占推理时间的 61% 和 23%。
图 2. 文本解码器和声码器是耗时最多的模块。针对 batch_size=1 在 A100 GPU 上执行英-西 S2ST（语音到语音-文本）任务的模块推理时间分解图。
为了更深入地了解文本解码器和声码器的性能瓶颈，我们分析了FLEURS数据集中第 8 个样本的英-西翻译示例的文本解码器和声码器的 GPU 轨迹，如图 3 所示。结果显示，文本解码器和声码器都严重受 CPU 限制。我们观察到由 CPU 开销引起的显著延迟，这推迟了 GPU 内核的启动，从而导致这两个模块的执行时间大幅增加。
(a) 文本解码器的 CPU 和 GPU 轨迹
(b) 声码器的 CPU 和 GPU 轨迹
图 3. 文本解码器和声码器是严重受 CPU 限制的模块。图为FLEURS数据集中第 8 个样本的英-西翻译示例的 (a) 文本解码器和 (b) 声码器的 CPU 和 GPU 轨迹。该轨迹是在 A100 GPU 上以 batch_size=1 运行推理时获得的。
基于实际系统性能分析结果，即 text_decoder 和 vocoder 是 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 调用次数以及主机和设备之间的数据传输。这可以带来显著的性能提升，特别是对于那些有大量小内核或多次重复相同内核集的应用程序。如果您对此感兴趣并想了解更多，可以查看这篇论文，它强调了数据在加速计算中的重要作用：数据在哪里？为什么没有答案就无法辩论 CPU vs. GPU 性能，作者是我们自己的 Kim Hazelwood！这篇论文发表于 NVIDIA 大力投资通用 GPU (GPGPU) 且深度学习尚未彻底改变计算行业的时期！
然而，由于 CUDA Graph 在编译时记录并操作 1) 固定的内存指针和 2) 固定的张量形状，我们为 CUDA Graph 引入了以下改进，以便在不同大小的输入之间复用，以避免为每次迭代生成 CUDA Graph，并让 CUDA Graph 内的数据在不同运行中复用，以在多个解码步骤中共享 KV 缓存。
文本解码器
Seamless 中的文本解码器是来自 NLLB [1] 的解码器，用于执行 T2TT（文本到文本翻译）。此外，该模块是一个受 CPU 限制的模型，由于自回归生成需要对 token 进行顺序处理的特性，GPU 执行时间不足以掩盖 CPU 开销，这限制了在 GPU 上可实现的并行度。基于这一观察，我们为文本解码器启用了 torch.compile + CUDA Graph，以减少占主导地位的 CPU 开销，如图 4 所示。
图 4. 启用 torch.compile + CUDA Graph 后文本解码器的 CPU 和 GPU 轨迹。
1. 更新和检索 KV 缓存
在推理期间，文本解码器有两个计算阶段：一个消耗提示的预填充阶段，以及一个逐个生成输出 token 的增量生成阶段。在足够高的批处理大小或输入长度下，预填充阶段并行处理足够多的 token——GPU 性能是瓶颈，CPU 开销对性能影响不大。另一方面，增量 token 生成始终以序列长度 1 执行，并且通常在小批处理大小（甚至为 1）下执行，例如在交互式用例中。因此，增量生成可能受限于 CPU 速度，因此是 torch.compile + CUDA Graph 的一个很好的应用场景。
然而，在增量 token 生成阶段，参与注意力计算的 key 和 value 的 sequence_length 维度每一步增加一，而 query 的序列长度始终保持为 1。具体来说，key/value 是通过将新计算出的序列长度为 1 的 key/value 附加到目前为止存储在 KV 缓存中的 key/value 来生成的。但如上所述，CUDA Graph 在编译期间记录所有张量的形状，并使用记录的形状进行重放。因此，我们参照这篇优秀文章here中的做法，进行了一些修改来解决这个问题。
a) 我们修改了 KV 缓存的处理方式，使其接受一个 CUDA 张量（即 valid_seq_pos）来指定写入新值的位置，而不是一个 Python 整数。
图 5. 对 KV 缓存 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) 由于不同的推理样本具有不同的序列长度，它们也会生成不同形状的输入，这些输入将被投影到交叉注意力层的 key 和 value。因此，我们将输入填充到静态形状，并生成一个填充掩码来屏蔽掉填充后的输出。
2. 内存指针管理
由于 CUDA Graph 会记录内存指针以及张量的形状，因此确保不同的推理样本能正确引用记录的内存指针（例如 KV 缓存）非常重要，以避免为每个推理样本重新编译 CUDA Graph。然而，Seamless 代码库的某些部分使得不同的推理样本引用了不同的内存地址，因此我们进行了修改以改善内存使用情况。
e) Seamless 采用束搜索（beam search）作为文本解码策略。在束搜索过程中，我们需要在每个增量解码步骤中对所有注意力层进行 KV 缓存重排，以确保每个选定的束都使用相应的 KV 缓存，如下面的代码片段所示。
图 8. 束搜索解码策略的 KV 缓存重排操作。
上述代码为 cache_k 和 cache_v 分配了新的内存空间并覆盖了原始的内存指针。因此，我们修改了 KV 缓存的重排方式，通过使用 copy_ 操作符，保持了每个缓存在编译期间记录的内存指针。
图 9. 使用 copy_ 操作符对 KV 缓存进行原地更新
f) 在通过上述修改为文本解码器启用 torch.compile + CUDA Graph 后，文本解码器的开销转移到了 KV 缓存重排上，如图 10 所示。KV 缓存重排重复调用 index_select 96 次（假设有 24 个解码器层，每个层包含两种类型的注意力层，分别有 key 和 value 的缓存）。
图 10. 启用 torch.compile + CUDA Graph 后文本解码器的 CPU 和 GPU 轨迹。
作为加速文本解码器的一部分，我们额外对 KV 缓存重排应用了 torch.compile，以从内核融合中获益，如图 11 所示。请注意，我们不能在这里使用 CUDA Graph（mode='max-autotune'），因为 copy_ 操作会修改输入，这违反了 torch.compile 中 CUDA Graph 版本的静态输入要求。
图 11. 对 KV 缓存重排应用 torch.compile。
对 KV 缓存重排启用 torch.compile 的结果是，之前分开启动的 GPU 内核（图 12(a)）现在被融合了，因此需要启动的 GPU 内核数量大大减少了（图 12(b)）。
(a) 启用 torch.compile 前 KV 缓存重排的 CPU 和 GPU 轨迹
(b) 启用 torch.compile 后 KV 缓存重排的 CPU 和 GPU 轨迹
图 12. 启用 torch.compile 前 (a) 和后 (b) 的 KV 缓存重排的 CPU 和 GPU 轨迹
声码器
Seamless 中的声码器是一个 HiFi-GAN 单元声码器，它将生成的单元转换为波形输出，其中一个单元是语音的表示，结合了音素和音节等不同方面，可用于生成人类可听的声音。声码器是一个相对简单的模块，由 Conv1d 和 ConvTranspose1d 层组成，并且是一个受 CPU 限制的模块，如图 3 所示。基于这一观察，我们决定为声码器启用 torch.compile + CUDA Graph，以减少不成比例的巨大 CPU 开销，如图 10 所示。但仍有一些问题需要修复。
图 13. 启用 torch.compile + CUDA Graph 后声码器的 CPU 和 GPU 轨迹。
a) 声码器的输入张量形状在不同的推理样本之间是不同的。但是由于 CUDA Graph 会记录张量的形状并进行重放，我们必须用零将输入填充到固定大小。由于声码器只包含 Conv1d 层，我们不需要额外的填充掩码，用零填充就足够了。
b) 声码器由包裹着 torch.nn.utils.weight_norm 的 conv1d 层组成（参见此处）。然而，直接对声码器应用 torch.compile 会导致如下的图中断，从而导致性能提升不佳。这个图中断发生在 weight_norm 的 PyTorch 代码中的钩子处理部分。
[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]
由于在推理过程中层的权重不会改变，我们不需要权重归一化。因此，我们简单地移除了声码器的权重归一化，如图 14 所示，利用了 Seamless 代码库中已提供的 remove_weight_norm 函数（此处）。
图 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 的 API（即 torch.compile(mode="max-autotune-no-cudagraphs")）支持，以确定 CUDA Graph 对性能的影响。在没有 CUDA Graph 的情况下，文本解码器和声码器的加速比分别降至 1.17 倍和 18.4 倍。虽然仍然相当显著，但这表明了 CUDA Graph 的重要作用。我们得出结论，Seamless M4T-v2 在启动 CUDA 内核时耗费了大量时间，尤其是在使用小批量大小（例如 1）的情况下，此时 GPU 内核执行时间不足以分摊 GPU 内核启动时间。
图 16. 逐步应用 torch.compile 和 CUDA Graph 的端到端推理加速。a) “增量解码”：仅对文本解码器应用 torch.compile b) “带 CUDA Graph 的增量解码”：对文本解码器应用 torch.compile + CUDA Graph c) “+KV 缓存重排”：在 b) 的基础上，额外对 KV 缓存重排操作应用 torch.compile d) “+声码器”：在 c) 的基础上，额外对声码器应用 torch.compile e) “+带 CUDA Graph 的声码器”：在 d) 的基础上，额外对声码器应用 torch.compile + CUDA Graph。
图 16 展示了在模块上应用带和不带 CUDA Graph 的 torch.compile 的累积效应。结果表明，端到端推理速度显著提升，证明了这些技术在优化整体延迟方面的有效性。最终，对于 batch_size=1 的 Seamless M4T-v2，我们获得了 2.7 倍的端到端推理加速。
致谢
我们感谢 PyTorch 团队和 Seamless 团队在这项工作中给予的巨大支持。