Deepgramの音声認識Nova-3をAWS Trainium2へ移植、バッチ処理のドル当たり音声処理量が11〜70%増

Deepgramは音声認識・音声合成・音声エージェントのモデルをAPIや顧客環境で提供する会社で、AWSの技術記事によると年間でquadrillion（10の15乗）級のaudio tokenを処理し、20万人超の開発者と1,400組織に使われている。大量の音声を決まった予算で処理するには、アクセラレータ1台あたりの処理効率が費用を左右します。今回の取り組みは、こうしたバッチ処理の費用効率を上げることを目的に、DeepgramとAWS Forward Deployed Engineeringのチームが数か月かけて進めました。

移植は、Deepgramが標準のPyTorchコードとして提供したNova-3を、Neuron CompilerのPyTorch統合であるTorchNeuronでTrainium2向けに組み直す形で始まりました。まずeagerモードで動かし、CPUへフォールバックする演算を確認します。次にtorch.compileでコンパイルしました。コンパイル後はすべてのモデル演算がTrainium2で動き、CPUフォールバックは検出されていません。エンコーダーとデコーダーは別々の静的shapeグラフにコンパイルし、層ごとの出力と音声認識精度（WER）がDeepgramの参照値に沿うことを確認しています。

カーネル単位の調整

調整は、Neuron Kernel Interface（NKI）で書かれた既存カーネル集NKILibへの差し替えで大部分を済ませ、残る箇所はプロファイリングツールNeuron Explorerで原因を特定し、Neuron Agentic Development（NAD）のプロファイリング・カーネル作成機能を使って手書きカーネルに置き換える流れで進めました。

モデルの実行環境も作り直しました。Trainium2デバイス1台の物理NeuronCore 8個を論理コア1個あたり1物理コアの設定にして、Nova-3のレプリカ8個（DP8）を1台に載せます。Deepgramの本番サーバーはRustで書かれているため、推論経路からPythonとPyTorchを外し、事前コンパイルしたNEFF（Neuron Executable File Format）をNeuron Runtime API経由で実行する軽量ランタイムをRustで構築しました。レプリカごとにスレッドを持つワークスティーリングキューで入力を分散し、テストしたすべてのバッチサイズで参照と一致するトークン列を再現しています。このRustランタイムは、Pythonのマルチプロセス構成比で7.5%のスループット向上を、スループットが最大になるバッチサイズで確認できました。

費用の測定は、評価時点のインスタンス価格を使い、Deepgramのオフラインのスループット志向バッチ負荷で行われました。記事は効果の範囲も明示しています。対象はオフラインバッチ処理に限られ、リアルタイム用途やあらゆるGPUとの比較ではありません。精度面では、移植後のWERがDeepgramの参照値と目標に沿うことを確認したとしています。

記事は、Amazon SageMaker AI上でDeepgramモデルを自社のVPC内で動かす顧客にも、こうした改善が及ぶ可能性に触れています。11〜70%という幅は負荷とバッチサイズで変わる値であり、DeepgramのAPI価格への反映については記事では触れられていません。音声AIの単位コストがGPU以外のアクセラレータでも下げられることを示した事例として、実測手順ごと公開された点が特徴です。