这篇文章是多系列博客的第四部分，重点介绍如何使用纯原生 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. 文本解码器和声码器是最耗时的模块。A100 GPU 上 batch_size=1 的英语-西班牙语 S2ST（语音到语音文本）任务的模块推理时间分解。
为了更仔细地查看文本解码器和声码器的性能瓶颈，我们分析了 FLEURS 数据集的英语-西班牙语翻译示例中第 8 个样本的文本解码器和声码器的 GPU 跟踪，如图 3 所示。结果表明，文本解码器和声码器是严重受 CPU 限制的模块。 我们观察到 CPU 开销造成了显著的差距，延迟了 GPU 内核的启动，导致两个模块的执行时间都大幅增加。
(a) 文本解码器的 CPU 和 GPU 跟踪
(b) 声码器的 CPU 和 GPU 跟踪
图 3. 文本解码器和声码器是严重受 CPU 限制的模块。 (a) 文本解码器 (b) 声码器的 CPU 和 GPU 跟踪，用于 FLEURS 数据集的英语-西班牙语翻译示例的第 8 个样本。跟踪是通过在 A100 gpu 上以 batch_size=1 运行推理获得的。
基于 Seamless M4T-v2 中 text_decoder 和声码器是严重受 CPU 限制的模块的实际系统性能分析结果，我们对这些模块启用了 torch.compile + CUDA Graph。在这篇文章中，我们分享了在 batch_size=1 推理场景中在每个模块上启用 torch.compile + CUDA Graph 所需的修改、关于 CUDA Graph 的讨论以及下一步计划。
Torch.compile 与 CUDA Graph
torch.compile 是一个 PyTorch API，允许用户将 PyTorch 模型编译成独立的exe文件或脚本，通常用于通过消除不必要的开销来优化模型性能。
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 缓存。
文本解码器
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 生成阶段，注意力计算中涉及的键和值的 sequence_length 维度在每个步骤中增加 1，而查询的序列长度始终保持为 1。具体而言，键/值是通过将新计算出的序列长度为 1 的键/值附加到目前存储在 KV 缓存中的键/值来生成的。但如上所述，CUDA Graph 在编译期间记录所有张量的形状，并在重放时使用记录的形状。因此，根据 此处 的出色工作，对解决此问题进行了一些修改。
a) 我们修改了 KV 缓存处理，以获取要在 CUDA 张量中写入新值的索引（即 valid_seq_pos），而不是 Python 整数。
图 5. KV 缓存 append 和 get 的修改
b) 我们还修改了注意力，使其适用于 max_seq_length 上的固定形状的键和值。我们仅计算序列位置直到当前解码步骤（即 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) 由于不同的推理样本具有不同的序列长度，因此它也会生成不同形状的输入，这些输入将投影到交叉注意力层的键和值。因此，我们将输入填充为静态形状，并生成填充掩码以屏蔽填充的输出。
2. 内存指针管理
由于 CUDA Graph 记录内存指针以及张量的形状，因此重要的是使不同的推理样本正确引用记录的内存指针（例如，KV 缓存），以避免为每个推理样本编译 CUDA Graph。但是，Seamless 代码库的某些部分使不同的推理样本引用不同的内存地址，因此我们进行了修改以改善内存影响。
e) Seamless 采用集束搜索作为文本解码策略。在集束搜索过程中，我们需要为每个增量解码步骤对所有注意力层执行 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 个解码器层，其中每个层由两种类型的注意力层组成，每种注意力层都具有键和值的缓存）。
图 10. 启用 torch.compile + CUDA Graph 后文本解码器的 CPU 和 GPU 跟踪。
作为加速文本解码器的一部分，我们还在 KV 缓存重新排序中应用了 torch.compile，以受益于内核融合，如图 11 所示。请注意，我们不能在此处使用 CUDA Graph (mode='max-autotune')，因为 copy_ 操作会修改输入，这违反了 torch.compile 中 CUDA 图版本的静态输入要求。
图 11. 将 torch.compile 应用于 KV 缓存重新排序。
由于对 KV 缓存重新排序启用了 torch.compile，因此单独启动的 gpu 内核（图 12(a)）现在已融合，因此需要启动的 gpu 内核少得多（图 12(b)）。
(a) 在启用 torch.compile 之前 KV 缓存重新排序的 CPU 和 GPU 跟踪
(b) 在启用 torch.compile 之后 KV 缓存重新排序的 CPU 和 GPU 跟踪
图 12. KV 缓存重新排序的 CPU 和 GPU 跟踪（a）启用 torch.compile 之前和（b）之后
声码器
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 代码中的 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]
由于层的权重在推理期间不会改变，因此我们不需要权重归一化。因此，我们只需删除声码器的权重归一化，如图 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 的文本解码器和声码器的推理时间加速
我们还报告了使用不带 CUDA Graph 的 torch.compile 的文本解码器和声码器的加速，这由 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 图的端到端推理加速。a) “Inc. Decoding”：仅将 torch.compile 应用于文本解码器 b) “Inc. Decoding w/ CUDA Graph”：将 torch.compile + CUDA Graph 应用于文本解码器 c) “+KV Cache Reordering”：在 b) 的基础上，另外将 torch.compile 应用于 KV 缓存重新排序操作 d) “+Vocoder”：在 c) 的基础上，另外将 torch.compile 应用于声码器 e) “+Vocoder w/ CUDA Graph”：在 d) 的基础上，另外将 torch.compile + CUDA Graph 应用于声码器。
图 16 表示将带和不带 CUDA Graph 的 torch.compile 应用于模块的累积效果。结果表明，端到端推理加速得到了显著改善，证明了这些技术在优化整体延迟方面的有效性。因此，对于 batch_size=1 的 Seamless M4T-v2，我们获得了 2.7 倍 的端到端推理加速。
致谢
我们感谢 PyTorch 团队和 Seamless 团队为这项工作提供的巨大支持。
