AWS Trainium、RMSNormを約6倍高速化

2026年10月21日、Hugging Face上でAWS TrainiumのPyTorch推論をTorchNeuronで解析する手順が公開されました。RMSNormを題材に、torch.profilerの集計表とNeuron Explorerのトレースで処理時間の所在を特定します。単純実装と2つの演算融合方式を同条件で比べ、デバイス時間が約6分の1になる結果を示しました。

計測ではProfilerActivity.PrivateUse1とNeuronConfigを指定し、CPUとNeuronの活動を同時に収集します。correlate_device_time=Trueでデバイス時間をkey_averages()へ反映し、torch.neuron.synchronize()で計測終了前の未完了処理を待ちます。表をデバイス時間順に並べ、Explorerでフレームワーク、ランタイム、ハードウェアを横断して確認します。

RMSNormの単純実装は乗算、平均、加算、逆平方根を個別に実行し、6回の演算呼び出しで2.355ミリ秒を要しました。3回のaten::mulがデバイス時間の49.5％を占め、演算ごとのディスパッチと中間テンソルの生成が負担になります。

組み込みのtorch.rms_normは処理を1つのカーネルに融合し、デバイス時間を403.3マイクロ秒まで短縮しました。元のPython関数をtorch.compileでコンパイルした方式も単一の実行単位となり、375.5マイクロ秒を記録しました。

組み込み演算とコンパイルは異なる経路を使いますが、どちらも単純実装のおよそ6分の1のデバイス時間でした。実務ではkey_averages()でボトルネックを絞り、Explorerで実行間の空白や転送を確かめたうえで、融合演算かtorch.compileの採用を検討できます。