这篇文章是关于如何使用纯原生 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（语音到语音文本）任务时，按模块划分的推理时间明细。
为了更深入地了解文本解码器和声码器的性能瓶颈，我们分析了图 3 所示的 FLEURS 数据集英语-西班牙语翻译示例中第 8 个样本的文本解码器和声码器的 GPU 轨迹。结果显示，文本解码器和声码器是严重的 CPU 瓶颈模块。我们观察到由于 CPU 开销导致启动 GPU 内核延迟而产生的显著差距，这导致这两个模块的执行时间大幅增加。
(a) 文本解码器的 CPU 和 GPU 轨迹
(b) 声码器的 CPU 和 GPU 轨迹
图 3. 文本解码器和声码器是严重的 CPU 瓶颈模块。FLEURS 数据集英语-西班牙语翻译示例中第 8 个样本的 (a) 文本解码器和 (b) 声码器的 CPU 和 GPU 轨迹。轨迹是在 A100 GPU 上使用 batch_size=1 运行推理时获得的。
基于对 Seamless M4T-v2 中文本解码器和声码器存在严重 CPU 瓶颈的真实系统性能分析结果，我们对这些模块启用了 torch.compile + CUDA Graph。在这篇文章中，我们将分享在 batch_size=1 推理场景下，为每个模块启用 torch.compile + CUDA Graph 所需的修改、关于 CUDA Graph 的讨论以及下一步计划。
Torch.compile 与 CUDA Graph
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 的预填充阶段，以及一个逐个生成输出 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 Tensor 中写入新值的索引（即 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）上计算注意力。然而，尽管存在这个缺点，我们的结果表明 torch.compile + CUDA Graph 与标准 PyTorch 代码相比仍然提供了显著的性能优势。
c) 由于不同的推理样本具有不同的序列长度，这也会产生需要投影到交叉注意力层的 key 和 value 的不同形状的输入。因此，我们对输入进行填充以获得静态形状，并生成一个填充掩码来屏蔽填充的输出。
2. 内存指针管理
由于 CUDA Graph 记录内存指针以及张量形状，因此重要的是让不同的推理样本正确引用记录的内存指针（例如 KV cache），以避免为每个推理样本重新编译 CUDA Graph。然而，Seamless codebase 中的某些部分导致不同的推理样本引用不同的内存地址，因此我们进行了修改以改善内存影响。
e) Seamless 采用束搜索作为文本解码策略。在束搜索过程中，我们需要对每个增量解码步骤的所有注意力层执行 KV cache 重排，以确保每个选定的 beam 使用相应的 KV cache 进行操作，如下面的代码片段所示。
图 8. 束搜索解码策略的 KV cache 重排操作。
上面的代码分配了新的内存空间，并覆盖了 cache_k 和 cache_v 的原始内存指针。因此，我们修改了 KV cache 重排操作，通过使用 copy_ 操作符来保留编译期间记录的每个 cache 的内存指针。
图 9. 使用 copy_ 操作符对 KV cache 进行就地更新
f) 在按照上述修改对文本解码器启用 torch.compile + CUDA Graph 后，文本解码器的开销转移到了 KV cache 重排操作，如图 10 所示。KV cache 重排操作重复调用 index_select 96 次（假设有 24 个解码器层，每个层包含两种类型的注意力层，分别缓存 key 和 value）。
图 10. 启用 torch.compile + CUDA Graph 后文本解码器的 CPU 和 GPU 轨迹。
作为加速文本解码器的一部分，我们还额外对 KV cache 重排操作应用了 torch.compile，以从内核融合中获益，如图 11 所示。请注意，此处我们无法使用 CUDA Graph (mode='max-autotune')，因为 copy_ 操作会修改输入，这违反了 torch.compile 中 CUDA graph 版本的静态输入要求。
图 11. 将 torch.compile 应用于 KV Cache 重排。
将 torch.compile 应用于 KV cache 重排后，原来单独启动的 GPU 内核（图 12(a)）现在被融合，从而需要启动的 GPU 内核数量大大减少（图 12(b)）。
(a) 启用 torch.compile 之前 KV cache 重排的 CPU 和 GPU 轨迹
(b) 启用 torch.compile 之后 KV cache 重排的 CPU 和 GPU 轨迹
图 12. KV cache 重排的 CPU 和 GPU 轨迹 (a) 启用 torch.compile 之前和 (b) 启用 torch.compile 之后
声码器
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），从而导致性能提升不理想。这种图中断发生在 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 codebase 中已提供的 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 内核上花费了大量时间，特别是当我们使用小 batch size（例如 1）时，GPU 内核执行时间不足以分摊 GPU 内核启动时间。
图 16. 逐步应用 torch.compile 和 CUDA Graph 的端到端推理加速比。a) “增量解码”: 仅对文本解码器应用 torch.compile b) “增量解码 (含 CUDA Graph)”: 对文本解码器应用 torch.compile + CUDA Graph c) “+KV Cache 重排”: 在 b) 的基础上额外对 KV cache 重排操作应用 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 团队在这项工作中给予的巨大支持。
