随着语音成为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）团队将Nova-3移植到 AWS Trainium2 的PyTorch重构版本中，并针对新硬件进行了调优。最终配置在单台Amazon Elastic Compute Cloud（Amazon EC2）trn2.3xlarge实例上运行八个模型副本。在Deepgram面向吞吐量的离线批量工作负载上，该设置比Deepgram此前使用的计算方案每美元多处理11–70%的音频。我们使用评估时的实例定价进行了这一测量。这些收益帮助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集成，它同时支持eager执行和 torch.compile。我们首先在Trainium2上以eager模式运行Nova-3，检查哪些操作回退到了CPU，然后使用 torch.compile进行编译。编译后的路径使每个模型操作都运行在Trainium2上，未检测到任何CPU回退。
我们使用 torch.compile将编码器和解码器编译为独立的静态形状图。在此过程中，我们在两个层级检查数值精度：与CPU对比的层输出一致性，以及端到端的词错误率（WER），后者必须与Deepgram的参考值和目标保持一致。
在这个可用的基线之上，我们通过用公共NKI Library（NKILib）中的融合内核替换多个层来启动优化工作。NKILib是用Neuron Kernel Interface（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 | 参考内核的平铺调度（tiling schedule）对每个token tile都从片外存储器重新加载整套投影权重，因此随着token数量增长，内存带宽限制了吞吐量。我们手工调优的内核每个权重tile只加载一次，在多个token tile间复用它，并将权重读取与片上计算流水线化。 | 在批量大小96时，延迟比参考内核低2.0倍 |
| Depthwise conv1d | 实验性的NKILib内核将该操作视为隐式通用矩阵乘（GEMM），导致Tensor Engine利用率较低。我们将其重写为Vector Engine上的广播乘加循环，达到了引擎的完全利用率，足以弥补Vector Engine峰值吞吐量较低的问题。 | 与参考内核相比,在批大小为 96 时实现 1.5 倍加速 |
| Conv2d | 该模型的 conv2d 层遇到了编译瓶颈:当批大小超过 32 时,这些层导致模型无法编译为单一计算图。由于 NKILib 没有 conv2d 内核,我们用 NKI 从零编写了一个。该内核按照前文所述的飞轮流程经过多次基于性能分析器信息的迭代,最终解锁了 32 以上的批大小。 | 在批大小 ≤32 时,卷积模块上最高达 50 倍的提升 |
Rust 编写的数据并行运行时
每个 Trainium2 设备配备八个物理 NeuronCore,默认暴露为四个逻辑 NeuronCore(LNC2)。由于每个 Nova-3 模型实例可以轻松容纳在单个物理核心上,我们将每个逻辑核心对应的物理核心数量减少为一个(LNC1)。这种配置使我们在每个 Trainium2 设备上托管八个模型副本(DP8)。输入在最大化每设备吞吐量的同时被并发处理。数据并行运行时由一个工作窃取队列驱动,每个副本一个线程。
Deepgram 的生产服务栈是用 Rust 编写的,因此将 Nova-3 移植到 Trainium2 不仅仅是编译模型:它还必须在服务路径中不依赖 Python 或 PyTorch 的情况下运行。因此,我们构建了一个轻量级运行时,其中 Rust 通过共享的工作窃取队列协调模型副本。一个 C 层通过 Neuron Runtime API (libnrt) 直接加载并执行预先编译的 Neuron 可执行文件格式(NEFFs)。该运行时与 Deepgram 的服务架构直接集成,并在每个测试批大小下都精确复现了参考 token 序列。
我们将其与多线程(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 Runtime API。我们的配置在离线、面向吞吐量的场景中展示了成本节省,验证了 AWS Trainium 是 Deepgram 推理工作负载的一个可行且高效的选择。
要在您自己的模型上尝试这种方法,请从 AWS Neuron 文档开始,它涵盖了编译器、运行时和支持的框架。 Neuron Kernel Interface(NKI)指南 解释了如何编写本文所述的内核,而 NKI 示例代码库 则包含我们起步时使用的参考内核。 Neuron 开发者工具文档 涵盖了我们在这里使用的性能分析和 agentic 工具,包括 Neuron Explorer 和 Neuron Agentic Development。有关实例规格和定价,请参见 Amazon EC2 Trn2 Instances。如需帮助规划类似的移植,请联系 AWS Forward Deployed Engineering。