AWSの技術者はtorch.profilerとNeuron Explorerで処理を追う方法を紹介しました。RMSNormの3実装を使い、CPUとデバイスの時間を比較します。 遅さの原因を演算、起動、待ち時間へ分ける手法です。測定条件を揃え、コンパイルと定常実行を区別して比較することが焦点です。 Trainium上の遅さを処理の経路から確認 AWSの技術者は10月21日、TorchNeuronのシステム計測をHugging Faceで公開しました。Trainiumを標準PyTorchのバックエンドとして使い、処理のどこに時間がかかるかを調べる内容です。torch.profilerから計測し、Neuron Explorerで実行の経路を追います。記事は第1回として、計測の経験がない開発者も対象にしています。RMSNormという同じ演算を3つの方法で書き、見える差を比較します。単に最終的な処理速度を示すのではなく、フレームワークからランタイム、デバイスまでを結び付けて原因を調べる技術解説です。 演算の表と実行時系列を組み合わせる TorchNeuronはPyTorchのPrivateUse1を使い、neuronというデバイスを登録します。記事はNeuronConfigとNeuronProfilerを使い、取得する情報やNeuronCoreを選びます。key_averagesの表でCPUとデバイスの時間を比較し、Explorerで同じ演算の実行経路を確認します。通常のPyTorch演算を組み合わせる方法、torch.rms_norm、torch.compileを使う方法を比べ、処理の分割や呼び出しの違いを読みます。対象はTrainiumで学習や推論を動かす開発者です。モデル全体が遅いときも、まず個々の演算で時間の使い方を確認する入口になります。デバイスの演算時間だけが短くても、CPU側の処理や呼び出しの間隔が全体へ影響する場合があります。表の数値と時系列を照合することが、改善する場所を選ぶ材料になります。 計測条件と最適化の価値を揃える 実務では、初回のコンパイルや準備の時間と、繰り返し実行した際の時間を区別することが論点です。入力サイズ、型、並列実行の条件が異なると、同じ演算でも見え方が変わる可能性があります。計測するCoreや情報を絞り、必要な範囲のトレースを取得することで、結果を読みやすくする方法が考えられます。最適化を加えた場合は速度だけでなく出力の一致も確認する必要があります。RMSNormの例で分かった傾向が、モデル全体の支配的な要因とは限らないためです。実際の負荷で時間を占める処理を確認し、変更前後の条件を残すことが重要です。トレースに業務上の情報が含まれる可能性も踏まえ、共有範囲を管理する対象になります。測定と仮説、変更後の確認をつなぐことで、計算資源の費用に対して意味のある改善かを判断できます。