How Deepgram optimized voice AI at scale with AWS Trainium

Deepgram如何利用AWS Trainium大规模优化语音AI | 人工智能

Deepgram如何利用AWS Trainium大规模优化语音AI

随着语音成为人工智能领域的一等模态，Deepgram的模型覆盖语音转文本、文本转语音和语音代理。目前，这些模型每年处理数千万亿个音频标记，服务于超过20万名开发者和1400家组织。客户可通过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）团队将Nova-3模型通过PyTorch重构方式移植到AWS Trainium2，并针对新硬件进行调优。最终配置在单个Amazon Elastic Compute Cloud（Amazon EC2）trn2.3xlarge实例上运行八个模型副本。在Deepgram面向离线、吞吐量导向的批量工作负载场景中，该配置相比此前使用的计算方案，每美元可处理11%-70%更多的音频。我们通过评估时的实例定价进行了测量。这些优化成果帮助Deepgram实现高效扩展，持续推出更优质的模型，并以更低价格服务客户。这些优势同样适用于在Amazon SageMaker AI上自托管Deepgram模型的客户。在该场景中，Deepgram与AWS继续改进部署便捷性、系统可靠性和可观测性。

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

通过原生PyTorch将Nova-3优化到Trainium2

从标准PyTorch代码开始，我们使用了TorchNeuron——Neuron编译器的原生PyTorch集成工具，支持eager execution和torch.compile两种模式。我们首先在Trainium2上以eager模式运行Nova-3，检查哪些操作回退到了CPU，然后通过torch.compile进行编译。编译后的路径在Trainium2上执行所有模型操作，未检测到任何CPU回退。

我们通过torch.compile将编码器和解码器分别编译为独立的静态形状图。在整个过程中，我们在两个层面验证了数值准确性：一是与CPU的逐层输出一致性，二是端到端词错误率（WER），该指标必须与Deepgram的参考值和目标保持一致。

基于这一可行的基准，我们启动了优化工作，将多个层替换为来自公共NKI库（NKILib）的融合内核。该库包含使用Neuron内核接口（NKI）编写的预构建参考内核。这些内核覆盖了融合QKV投影、融合前馈网络（FFN）层和深度可分离1D卷积等操作。

我们针对吞吐量和由此带来的成本效率进行了优化，目标是 Deepgram 的离线推理工作负载。吞吐量以相对于基线配置的相对增益进行报告，而非以每秒处理的绝对音频量来衡量，我们通过调整批量大小 1–128 来找到最大化吞吐量的批量大小。

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

启用内核只是优化飞轮的第一步：启用 NKI 内核（自定义或来自 NKILib），对工作负载进行分析，然后根据分析结果调整现有内核或用新内核替换。这个循环会持续到满足性能目标为止。在每个步骤中，Neuron 栈的不同部分都会加速工作：NKILib 提供即插即用的替代内核，Neuron Explorer 是分析 Neuron 工作负载的首选工具，而 Neuron Agentic Development（NAD）工具包则通过其分析和内核编写能力加速分析和内核调优。

在我们的情况下，NKILib 内核已经使我们取得了大部分进展，而分析显示在某些地方手动编写的内核表现会更好。每次分析中，NAD 的分析能力都能从 Neuron Explorer 的分析结果中找到根本原因，而 NAD 的内核编写能力则帮助生成了替代内核。

| 层 | 发现和修改内容 | 增益 | |--------------|------------------------------------------------------------------------------|-----------------------------------------| | FFN | 参考内核的分块计划在每个 token 分块时都从片外内存重新加载完整的投影权重，导致 token 数量增加时内存带宽限制了吞吐量。我们的手动调优内核每次加载每个权重分块一次，跨多个 token 分块重用，并将权重读取与片上计算流水线化。 | 在批量大小 96 时相比参考内核延迟降低 2.0 倍 | | Depthwise conv1d | 实验性 NKILib 内核将操作表示为隐式的通用矩阵乘法（GEMM），导致 Tensor 引擎利用率较低。我们将其重写为 Vector 引擎上的广播乘加循环，实现了引擎的完全利用，足以弥补 Vector 引擎较低的峰值吞吐量。 | 在批量大小 96 时相比参考内核速度提升 1.5 倍 | | Conv2d | 模型的 conv2d 层遇到了编译瓶颈：批量大小超过 32 时，这些层会阻止模型作为单个图进行编译。由于 NKILib 没有 conv2d 内核，我们用 NKI 从头编写了一个。该内核经过多次由分析结果指导的迭代，按照上述飞轮流程进行优化，从而解锁了批量大小超过 32 的能力。 | 在批量大小 ≤32 时卷积模块速度提升最高达 50 倍 |

Rust 中的数据并行运行时

每个 Trainium2 设备配备八个物理 NeuronCores，默认情况下暴露为四个逻辑 NeuronCores（LNC2）。由于每个 Nova-3 模型实例都可以舒适地运行在单个物理核心上，我们将每个逻辑核心对应的物理核心数减少到一个（LNC1）。这种配置使我们可以在每个 Trainium2 设备上运行八个模型副本（DP8）。输入在最大化设备吞吐量的同时并行处理。数据并行运行时由一个工作窃取队列驱动，每个副本对应一个线程。

Deepgram 的生产服务堆栈使用 Rust 编写，因此将 Nova-3 移植到 Trainium2 上不仅需要编译模型，还必须确保服务路径中完全不依赖 Python 或 PyTorch。为此我们构建了一个轻量级运行时，通过共享的工作窃取队列由 Rust 协调模型副本。C 层直接通过 Neuron 运行时 API（libnrt）加载并执行预先编译的 Neuron 可执行文件格式（NEFFs）。该运行时直接集成到 Deepgram 的服务架构中，在所有测试批次大小下都能精确复现参考标记序列。

我们将其与多线程（MT）和多进程（MP）Python 数据并行运行时进行了基准测试。与 MP 基准相比，Rust 运行时完全消除了 Python 解释器的开销。与 MT 基准相比，它进一步缓解了全局解释器锁（GIL）序列化问题，因此并发工作线程的调度调用不再在单个进程中发生竞争。这些改进至关重要，因为 Nova-3 为每个音频块执行一系列小规模解码器调用，使得主机端调度在低批次大小时成为瓶颈。因此吞吐量提升在低批次大小时最为显著，即使在吞吐量最大化的批次大小下，提升依然具有实际意义。在该场景下，Rust 运行时相比多进程 Python 基准实现了 7.5% 的吞吐量提升。

图 1：AWS Trainium2 上的 Deepgram Nova-3，预先编译并由 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 运行时 API。我们的方案在离线、吞吐量导向场景中验证了成本节约效果，证明 AWS Trainium 是 Deepgram 推理工作负载的可行且高效选项。

要尝试在您自己的模型上使用此方法，请从 AWS Neuron 文档 开始，该文档涵盖编译器、运行时和受支持框架。Neuron 内核接口（NKI）指南解释了如何编写本文所述的内核，NKI 示例仓库包含我们最初使用的参考内核。Neuron 开发者工具文档涵盖了我们在此使用的性能分析和智能工具，包括 Neuron Explorer 和 Neuron Agentic Development。有关实例规格和定价信息，请参阅 Amazon EC2 Trn2 实例。如需帮助规划类似的移植工作，请联系 AWS Forward Deployed Engineering。

作者简介

'"`