孤本网
/ 0 阅读
0

Deepgram 将 Nova-3 语音模型移植至 AWS Trainium2,离线批量处理性价比最高提升 70%

一句话结论

Deepgram 优化 Nova-3 模型在 AWS Trainium2 上的离线批量推理,使每美元处理音频量较此前方案最高提升 70%。

关键要点

  • Deepgram Nova-3 语音转文本模型已在 AWS Trainium2 上完成离线批量部署,运行于单台 Amazon EC2 trn2.3xlarge 实例,承载 8 个模型副本。
  • 离线批量工作负载中,每美元处理音频量较 Deepgram 此前计算选项提升 11% 至 70%,基于评估时实例定价。
  • TorchNeuron 编译后,模型全部操作在 Trainium2 芯片内完成,无 CPU 回退;编码器与解码器分别编译为独立静态形状图。
  • 使用 NKI Library(NKILib)手写内核替换部分层后,Depthwise conv1d 在批量大小 96 时延迟比参考内核低 2.0 倍,Conv2d 在批量大小 32 及以下最高提升 50 倍。
  • Rust 运行时在达到最大吞吐量的批量大小下,比多进程 Python 基线吞吐量提升 7.5%。

背景与事实

Deepgram 是一家提供语音人工智能服务的公司,其模型覆盖语音转文本、文本转语音及语音代理三类能力,年处理音频 token 达数千万级,服务超过 20 万名开发者与 1,400 家机构。客户既可通过 Deepgram API 使用其服务,也可通过 Amazon SageMaker AI(亚马逊 SageMaker 人工智能)自行部署模型。

本次项目中,Deepgram 将 Nova-3 语音转文本模型移植到 AWS Trainium2,定位为离线批量工作负载,典型场景包括联络中心通话分析与医疗语音转录。与在线低延迟推理不同,离线批量场景更关注单位成本下的总吞吐量,而非单次请求的响应速度。

项目初期,团队采用标准 PyTorch 代码作为起点。经过数月协作,Deepgram 工程师与 AWS 现场工程(FDE)团队共同完成模型端口与调优。最终配置在单台 Amazon EC2 trn2.3xlarge 实例上以 LNC1 模式运行 8 个 Nova-3 模型副本,即数据并行度为 8(DP8)。模型经 TorchNeuron 编译器编译后,所有算子均在 Trainium2 芯片上执行,不再出现回退到 CPU 的情况;编码器与解码器被编译为彼此独立的静态形状图,以简化运行时调度。

在正确性验证层面,团队执行了两级校验:一是逐层输出一致性比对,将芯片上的中间结果与 CPU 参考实现逐层对比;二是端到端字错率(WER)对齐,确保与参考值一致。优化过程中,团队使用 NKI Library(NKILib)替换了部分网络层为融合内核,涉及查询、键、值(QKV)投影、前馈网络(FFN)层以及深度 1D 卷积等算子。吞吐量指标以相对基线的增益报告,扫描范围覆盖批量大小 1 至 128,以找到最大吞吐量对应的最优点。优化循环包括启用内核、使用 Neuron Explorer 分析瓶颈、调优或替换内核三个环节。

具体内核优化方面,FFN 层被发现因参考内核的分块调度策略导致内存带宽受限。团队手写的新内核通过单次加载权重块、跨块重用以及流水线化权重读取加以改进,在批量大小 96 时延迟比参考内核低 2.0 倍。Depthwise conv1d 内核在批量大小 96 时比参考内核快 1.5 倍。新编写的 Conv2d 内核在批量大小 32 及以下的卷积模块上实现了最高 50 倍的性能提升。

在运行时层面,团队测试了 Rust 运行时与多进程 Python 基线的对比。在达到最大吞吐量的批量大小下,Rust 运行时比多进程 Python 基线吞吐量高 7.5%。综合硬件成本与吞吐量,单台 trn2.3xlarge 设备在离线批量工作负载中的性价比,即每美元处理的音频量,比 Deepgram 此前使用的计算选项高出 11% 至 70%,该数值基于评估时的实例定价。

影响分析

对中文开发者与从业者而言,这一案例提供了语音大模型在非 GPU 加速器上批量推理的完整工程路径。过去语音模型批量处理几乎完全依赖 NVIDIA GPU,而本项目表明,通过 PyTorch 到 TorchNeuron 的编译链路、NKI Library 内核手写以及运行时替换,可以在 AWS Trainium2 上实现有竞争力的成本效率。对于运营联络中心、医疗转录等离线语音管道的团队,每美元音频处理量提升 11% 至 70% 意味着同等预算下可处理的数据量显著增加,或同等数据量下成本降低,直接改善语音人工智能服务的毛利结构。

从技术栈角度看,NKI Library 的手写内核优化、Neuron Explorer 的瓶颈分析以及静态形状图编译等实践,为使用非传统 GPU 加速器的开发者提供了可复用的方法论。Rust 运行时在吞吐量上优于多进程 Python 基线 7.5% 的结果,也提示了在大规模推理服务中,运行时语言选择与进程模型的设计是隐藏但关键的优化维度。

适用边界

上述性价比提升与内核优化收益主要适用于离线批量工作负载,其核心指标是单位成本下的总吞吐量,而非单次请求延迟。对于在线低延迟语音交互场景,如实时语音代理、电话客服等,Trainium2 的推理延迟表现、冷启动时间与静态形状图的灵活性可能成为限制因素,不能直接套用本文结论。此外,11% 至 70% 的成本优势基于评估时的实例定价,云服务商的定价策略、区域可用性折扣以及合同承诺消费可能改变实际成本结构。Nova-3 模型的特定架构与 8 副本配置是本次调优的对象,迁移至其他模型或不同副本数需要重新执行内核优化与批量大小扫描,性能增益不可直接外推。

孤本观察

本文最值得关注的事实是,团队在标准 PyTorch 起点上通过数月端口与调优即实现全芯片推理且无 CPU 回退,这一工程成熟度表明 Trainium2 在语音模型领域已跨越早期适配阶段,进入可生产部署区间。50 倍 Conv2d 内核提升与 2.0 倍 FFN 延迟降低的具体数字,为后续在非 GPU 加速器上优化卷积网络提供了可引用的基准参考。

常见问题

Nova-3 在 AWS Trainium2 上离线批量处理的性价比提升幅度是多少?

每美元处理音频量较此前方案最高提升 70%,具体提升范围为 11% 至 70%,该数值基于评估时的实例定价。

部署 Nova-3 模型时使用了多少台实例和多少个模型副本?

运行于单台 Amazon EC2 trn2.3xlarge 实例,承载 8 个模型副本,数据并行度为 8。

优化后的 Conv2d 和 Depthwise conv1d 内核性能提升了多少?

Conv2d 在批量大小 32 及以下最高提升 50 倍;Depthwise conv1d 在批量大小 96 时延迟比参考内核低 2.0 倍。

Rust 运行时相比多进程 Python 基线在吞吐量上有何表现?

在达到最大吞吐量的批量大小下,Rust 运行时比多进程 Python 基线吞吐量提升 7.5%。

该方案适用于哪些典型场景,不适用于哪些场景?

适用于联络中心通话分析与医疗语音转录等离线批量场景;不适用于实时语音代理、电话客服等在线低延迟场景。

来源:AWS Machine Learning Blog