语音识别迁移 Trainium2 的优化实战

随着语音成为 AI 的一等模态，Deepgram 的模型已覆盖语音转文本、文本转语音和语音 Agent。如今，这些模型每年为超过 200,000 名开发者和 1,400 家组织处理数千万亿个音频 token。客户既可以通过 Deepgram API 使用这些模型，也可以在自己的环境中运行，包括 Amazon SageMaker AI。

对 Deepgram 来说，推理服务的效率决定了在给定的基础设施预算下，系统能处理多少音频。大规模批处理工作负载给预算带来了很大压力，因此，每个加速器小时的利用效率越高，Deepgram 就能用同一批计算资源处理更多音频。推理效率由此成为一个很自然的优化目标，Deepgram 开始借助 AWS AI 芯片来提升效率。他们与 AWS 合作，将 Nova-3 语音转文本模型迁移到 AWS Trainium2，用于联络中心分析、医疗转录等离线批处理工作负载。

合作从 Deepgram 提供标准 PyTorch 代码形式的 Nova-3 开始。在随后的几个月里，Deepgram 与 AWS Forward Deployed Engineering（FDE）团队通过 PyTorch 重构，将 Nova-3 移植到 AWS Trainium2，并针对新硬件进行调优。最终配置在单个 Amazon Elastic Compute Cloud（Amazon EC2）trn2.3xlarge 实例上运行八个模型副本。对于 Deepgram 面向吞吐量的离线批处理工作负载，与其此前使用的计算方案相比，这套配置每美元可多处理 11%–70% 的音频。我们采用评估时的实例价格进行测算。这些收益帮助 Deepgram 高效扩展，持续推出更好的模型，并以更低的价格提供给客户。对于在 Amazon SageMaker AI 上、自己的虚拟私有云（VPC）中自行托管 Deepgram 模型的客户，这些优势同样适用。在这一场景下，Deepgram 和 AWS 还在持续改善部署便捷性、可靠性和可观测性。

本文将完整介绍 Trainium 上的推理优化过程：将 Nova-3 移植到 AWS Neuron 软件栈、随后开展的 kernel 级调优、让 Deepgram 自有模型服务器能够调用 Trainium 的 Rust 数据并行运行时，以及我们在批量转录中测得的吞吐量和成本结果。

使用原生 PyTorch 将 Nova-3 迁移并优化到 Trainium2

我们从标准 PyTorch 代码出发，使用了 TorchNeuron。它是 Neuron Compiler 与原生 PyTorch 的集成，支持即时执行（eager execution）和 torch.compile。我们先在 Trainium2 上以 eager 模式运行 Nova-3，检查哪些算子回退到了 CPU，然后使用 torch.compile 编译。编译后的执行路径中，所有模型算子都在 Trainium2 上运行，没有检测到 CPU 回退。

我们使用 torch.compile，将编码器和解码器分别编译成独立的静态形状计算图。在整个过程中，我们从两个层面检查数值精度：一是各层输出与 CPU 结果的一致性，二是端到端词错误率（Word Error Rate，WER）。WER 必须与 Deepgram 的参考值和目标保持一致。

在这个能够正常运行的基线之上，我们开始优化：使用公开 NKI Library（NKILib）中的融合 kernel 替换多个层。NKILib 收录了一组使用 Neuron Kernel Interface（NKI）编写的预构建参考 kernel。这些 kernel 覆盖了融合 QKV 投影、融合前馈网络（FFN）层，以及一维逐通道卷积等算子。

我们针对 Deepgram 的离线推理工作负载优化吞吐量，并由此提升成本效率。吞吐量以相对于基线配置的提升幅度呈现，而不是每秒处理音频量的绝对值。我们遍历了 1–128 的 batch size，以找出吞吐量最高的 batch size。

使用 Neuron Explorer 和 NAD 进行 kernel 级调优

启用 kernel 只是优化循环的第一步：启用一个 NKI kernel（自定义或来自 NKILib），分析工作负载的性能，再根据分析结果调优现有 kernel，或用新的 kernel 替换它。这个循环会持续进行，直到达到性能目标。在每一步中，Neuron 软件栈的不同组件都会加快工作进度：NKILib 提供可直接替换的 kernel；Neuron Explorer 是分析 Neuron 工作负载性能的首选工具；Neuron Agentic Development（NAD）工具包则通过其性能分析和 kernel 编写 skills，加快性能数据分析与 kernel 调优。

在我们的案例中，NKILib kernel 已经完成了大部分优化，而性能分析则揭示了哪些地方使用手写 kernel 会表现得更好。每一轮中，NAD 的性能分析 skills 都帮助我们从 Neuron Explorer 的性能分析结果中定位根因，NAD 的 kernel 编写 skills 则帮助我们生成用于替换的 kernel。

使用 Rust 构建数据并行运行时

每个 Trainium2 设备包含八个物理 NeuronCore，默认对外呈现为四个逻辑 NeuronCore（LNC2）。由于每个 Nova-3 模型实例都能轻松放入单个物理核心，我们将每个逻辑核心包含的物理核心数量降为一个（LNC1）。这项配置使我们能够在每个 Trainium2 设备上部署八个模型副本（DP8），并发处理输入，同时最大化单设备吞吐量。数据并行运行时由一个工作窃取队列驱动，每个副本对应一个线程。

Deepgram 的生产推理服务栈使用 Rust 编写，因此，将 Nova-3 移植到 Trainium2 不只是编译模型：模型还必须能够在服务执行路径中不依赖 Python 或 PyTorch 运行。为此，我们构建了一个轻量级运行时，由 Rust 通过共享的工作窃取队列协调模型副本。一个 C 层通过 Neuron Runtime API（libnrt），直接加载和执行提前编译的 Neuron Executable File Format（NEFF）文件。这个运行时直接集成到 Deepgram 的服务架构中，并且在所有测试过的 batch size 下，都精确复现了参考 token 序列。

我们将它与采用多线程（MT）和多进程（MP）的 Python 数据并行运行时进行了基准比较。相较于 MP 基线，Rust 运行时完全消除了 Python 解释器的开销。相较于 MT 基线，它还进一步缓解了全局解释器锁（GIL）带来的串行化问题，使并发工作线程发起的调度调用不再在同一个进程内相互争用。这些收益很关键，因为 Nova-3 会为每个音频分块执行一连串小型解码器调用，在 batch size 较小时，主机端调度会成为瓶颈。因此，吞吐量提升在小 batch size 下最明显，即使在吞吐量最高的 batch size 下，提升仍然可观。在该 batch size 下，Rust 运行时的吞吐量比多进程 Python 基线高出 7.5%。

图 1：AWS Trainium2 上的 Deepgram Nova-3，提前编译并通过 Rust 提供推理服务

我们以每美元处理的音频量衡量性价比。对于 Deepgram 面向吞吐量的离线批处理工作负载，与其此前使用的计算方案相比，单设备 Trainium2 配置（trn2.3xlarge）每美元可多处理 11%–70% 的音频。

本文展示了 Deepgram 与 AWS FDE 如何使用 TorchNeuron，将 Nova-3 语音转文本 PyTorch 模型移植到 Trainium 硬件。TorchNeuron 为 AWS Trainium 提供原生 PyTorch 框架集成。我们直接使用 NKILib 中的预构建 kernel 替换模型层，以提升性能。随后，我们使用 Neuron Explorer 评估性能，并利用 Neuron Agentic Development（NAD）的性能分析 skills 查询性能数据，加快性能问题排查。在这些分析结果的基础上，我们使用 NAD 的 kernel 编写 skills 调整 kernel 源代码。最后，我们为 AWS Trainium 构建了轻量级 Rust 数据并行运行时，从 Rust 调用 Neuron Runtime API。这套配置在面向吞吐量的离线场景中实现了成本节省，验证了 AWS Trainium 是承载 Deepgram 推理工作负载的一种可行且高效的选择。

如果想在自己的模型上尝试这种方法，可以从 AWS Neuron 文档入手，其中涵盖了编译器、运行时和受支持的框架。Neuron Kernel Interface（NKI）指南介绍了如何编写本文所述的 kernel，NKI samples 仓库则收录了我们最初使用的参考 kernel。Neuron 开发者工具文档介绍了本文用到的性能分析与 Agent 工具，包括 Neuron Explorer 和 Neuron Agentic Development。实例规格和定价可参阅 Amazon EC2 Trn2 Instances。如需帮助评估类似移植项目的范围，可以联系 AWS Forward Deployed Engineering。