Deepgram 如何用 AWS Trainium2 优化语音 AI 推理

随着语音成为 AI 的一等模态，Deepgram 的模型涵盖语音转文本、文本转语音和语音代理。目前，这些模型每年为 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 面向离线、吞吐量优化的批处理工作负载中，该配置每美元处理的音频量，比 Deepgram 之前使用的计算选项高出 11–70 percent。我们根据评估时的实例定价进行了测算。这些提升有助于 Deepgram 高效扩展、持续推出更好的模型，并以更低的价格将模型提供给客户。对于在 Amazon SageMaker AI 上的自有虚拟私有云 (VPC) 中自行托管 Deepgram 模型的客户，这些优势同样适用。Deepgram 与 AWS 还在持续改善设置的便捷性、可靠性和可观测性。

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

使用原生 PyTorch 将 Nova-3 优化到 Trainium2 上

我们从标准 PyTorch 代码开始，使用 TorchNeuron（Neuron Compiler 的原生 PyTorch 集成），它同时支持即时执行和 torch.compile。我们首先在 Trainium2 上以即时模式运行 Nova-3，检查哪些操作回退到了 CPU，然后使用 torch.compile 进行编译。编译后的路径在 Trainium2 上运行了所有模型操作，未检测到 CPU 回退。

我们使用 torch.compile 将编码器和解码器编译为形状固定的独立计算图。整个过程中，我们从两个层面检查数值准确性：与 CPU 结果进行层级输出一致性对比，以及端到端词错误率 (WER)；后者必须与 Deepgram 的参考值和目标保持一致。

在此基础上，我们启动了优化工作，使用来自公开 NKI Library (NKILib) 的融合内核替换了若干层。NKILib 是一个预构建参考内核集合，这些内核使用 Neuron Kernel Interface (NKI) 编写。内核涵盖融合 QKV 投影、融合前馈网络 (FFN) 层和深度可分离一维卷积等操作。

我们针对 Deepgram 的离线推理工作负载优化了吞吐量，并进一步提升成本效率。吞吐量以相对于基线配置的增益来报告，而不是以每秒处理的绝对音频量来报告；我们测试了 1–128 的批量大小，以找出吞吐量最大的批量大小。

使用 Neuron Explorer 和 NAD 进行内核级调优

启用内核只是优化飞轮的第一步：启用一个 NKI 内核（自定义内核或来自 NKILib 的内核），分析工作负载，然后根据分析结果调优现有内核，或用新的内核替换它。这个循环会不断重复，直到达到性能目标。在每一步中，Neuron 堆栈的不同部分都会加速相关工作：NKILib 提供可直接替换的内核，Neuron Explorer 是分析 Neuron 工作负载的首选工具，而 Neuron Agentic Development (NAD) 工具包则通过其分析和内核编写技能，加快分析结果解读与内核调优。

在我们的案例中，NKILib 内核帮助我们完成了大部分工作，而分析结果显示了手写内核能够改进的地方。每一次，NAD 分析技能都帮助我们从 Neuron Explorer 分析结果定位到根本原因，而 NAD 内核编写技能则帮助生成替换内核。

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 的服务架构中，并且在所有测试过的批次大小下，都能精确复现参考 token 序列。

我们将其与多线程（MT）和多进程（MP）Python 数据并行运行时进行了基准测试。与 MP 基线相比，Rust 运行时彻底消除了 Python 解释器的开销。与 MT 基线相比，它进一步缓解了全局解释器锁（GIL）带来的串行化，因此并发工作线程发出的调用不再在单个进程内相互争用。这些提升之所以重要，是因为 Nova-3 会针对每个音频块执行一系列小型解码器调用，导致低批次大小下的瓶颈出现在主机端调度。因此，吞吐量提升在低批次大小时最大；即使在吞吐量最高的批次大小下，提升仍然显著。在该批次大小下，Rust 运行时相比多进程 Python 基线实现了 7.5% 的吞吐量提升。

图 1：Deepgram Nova-3 在 AWS Trainium2 上提前编译并由 Rust 提供服务

成本节省效果

我们以每美元处理的音频量衡量性价比。在 Deepgram 面向离线、注重吞吐量的批处理工作负载中，单设备 Trainium2 配置（trn2.3xlarge）每美元处理的音频量比 Deepgram 之前使用的计算选项高出 11–70%。

结论

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

若要在自己的模型上尝试这种方法，请从 AWS Neuron 文档入手，其中介绍了编译器、运行时和支持的框架。Neuron Kernel Interface（NKI）指南说明了如何编写本文介绍的内核，而 NKI 示例代码仓库则提供了我们采用的参考内核。Neuron 开发者工具文档介绍了我们在此使用的性能分析和智能体工具，包括 Neuron Explorer 和 Neuron Agentic Development。有关实例规格和定价，请参阅 Amazon EC2 Trn2 实例。如需帮助评估类似移植项目的范围，请联系 AWS Forward Deployed Engineering。

作者简介

Pierre Lienhart

Pierre 是 AWS 伦敦前沿部署工程团队的高级机器学习工程师。他与模型提供商和 AI 公司合作，充分发挥其推理工作负载在各种计算类型上的潜力，帮助他们在生产环境中以更快的速度和更低的成本提供模型服务。

Victor Wang

Victor 是 Deepgram 的资深软件工程师，也是工程副总裁的技术顾问。加入 Deepgram 之前，Victor 曾在 AWS 担任多个职务，包括高级解决方案架构师、技术项目经理、专业服务顾问和软件开发人员。他热衷于学习新技术和周游世界，累计飞行里程超过 200 万英里。

Andre Gomes

Andre 是 AWS Annapurna Labs ML 的解决方案架构师，拥有帮助初创公司运行大规模训练和实时推理的经验。他曾是初创公司的联合创始人兼 CEO，如今与创始人合作，帮助他们在 Trainium 和 Inferentia 上扩展工作负载。Andre 在 UFMG 获得博士学位，期间曾在剑桥大学开展研究，将大型语言模型和神经主题建模应用于大规模数据分析与决策。

Elyse Zhang

Elyse 是 AWS 的高级生成式 AI 策略师，负责帮助企业和初创公司客户推进模型定制之旅。她专注于将业务需求转化为适用于各行业的有效定制策略。

Haotian An

Haotian 是 AWS 前沿部署工程团队的机器学习工程师，专长于定制基础模型和大规模优化推理。他与客户紧密合作，根据特定使用场景调整生成式 AI，帮助客户发掘新能力并取得可衡量的业务成果。

Niraj Chaudhari

Niraj 是 AWS 前沿部署工程团队的深度学习架构师。他与企业客户合作，通过监督微调和强化微调定制基础模型、构建智能体系统并优化 AI 工作负载。他还帮助客户在 Amazon SageMaker AI 和 Amazon Bedrock 上将生成式 AI 应用从概念验证推向生产环境。

Rohit Thekkanal

Rohit 是 AWS 前沿部署工程团队模型定制部门的机器学习工程高级经理，负责领导专注于模型优化的可扩展生成式 AI 应用开发。在 Amazon 任职近十年间，他参与了对 Amazon 零售商品目录产生重大影响的机器学习项目。Rohit 获得了芝加哥大学布斯商学院 MBA 学位和卡内基梅隆大学硕士学位。