TorchNeuronのプロファイリング(パート1):AWS Trainium でのシステムレベルのプロファイリング。

多くのML実務者はモデルを高速に動作させる方法についてある程度の直観を持っていますが、モデルがなぜ遅いのかを分析する方法を知っている人は少なくなっています。このシリーズのPart 1では、推論パフォーマンスを体系的に分析し、時間がどこに費やされているのかを特定する方法を学びます。

この投稿はゼロから始め、プロファイリングの経験を前提としません。このウォークスルーでは、深層学習のトレーニングおよび推論ワークロード向けに特別に構築されたAIチップであるAWS Trainiumを使用し、 TorchNeuron、そのオープンソースのネイティブPyTorchバックエンドです。TorchNeuronは、eager executionを含む標準的なPyTorchワークフローをTrainium上でサポートします。 torch.compile「、および標準的なPyTorch APIと連携し、最小限のコード変更で利用できます。さらに、TorchNeuronは以下と統合します。」 Neuron Kernel Interface(NKI) Trainium ハードウェア向けのカスタムオペレータの有効化とパフォーマンス最適化を可能にする方法です。読み終える頃には、以下のことを理解できるようになっています:

• 設定方法 torch.profiler Neuronバックエンドを使用して、クリーンなトレースを取得します。
• 読み方： key_averages() Neuronデバイス時間を消費する操作を特定し、そのSelf CPU時間を比較するためのテーブルです。
• Neuron Explorer を使って、フレームワークからのディスパッチからランタイムのスケジューリング、そして NeuronCore 上でのハードウェア実行まで、1つのオペレーションを追いかける方法。
• 同じ操作の3つの実装の間で、これらのシグナルがどのように変化するか。

PyTorchバックエンドとしてのNeuron

ワークロードをプロファイリングする前に、その仕組みを理解しておくと役立ちます。 ニューロン PyTorch と統合しています。

NeuronバックエンドはPyTorchの PrivateUse1 「拡張ポイントです。この仕組みにより、アクセラレータは以下のような専用のディスパッチキーを追加することなく、フレームワークと統合できます。 cuda。この torch_neuronx このパッケージは、アップストリームのPyTorchと連携して、NeuronをPyTorchコードのネイティブバックエンドとして登録します。この拡張メカニズムの詳細については、次を参照してください PrivateUse1による新しいバックエンド統合の促進.

プロファイラーのセットアップ

Neuron ワークロードをプロファイリングするには、2 つのオブジェクトを使用します。

• NeuronConfig キャプチャする Neuron データストリームと NeuronCores、および出力の書き込み先を制御します。これを以下に渡します。 torch.profiler.profile として experimental_config.
• NeuronProfiler ~を取る NeuronConfig また、 export_trace コールバックを提供します。このコールバックはフレームワークトレースを書き出し、 trace.jsonをNeuronシステムトレースと併せて出力します。

nc = NeuronConfig(
    modes=[ProfileMode.DEVICE, ProfileMode.RUNTIME],
    profile_output_dir=OUTPUT_DIR,
    capture_enabled_for_nc='0',
    correlate_device_time=True,
)
exporter = NeuronProfiler(nc)

ザ modes パラメータは、デバイス命令やランタイムトレースなど、キャプチャするデータストリームを選択します。この capture_enabled_for_nc パラメータはトレースする NeuronCore を選択します。設定すると correlate_device_time=True トレースからデバイス時刻を伝播して key_averages() テーブルでは、CPUとNeuronデバイスの時間を比較できます。パラメータの完全なリファレンスおよび利用可能な ProfileMode 値については、NeuronConfig を参照してください。 プロファイルモードのドキュメント.

ツール: Neuron Explorer

Neuron Explorer は、システムトレース、デバイスタイムライン、エンジンごとの内訳を表示するブラウザベースのビューアです。SSH トンネルまたは Visual Studio Code 拡張機能を通じてアクセスできます。この投稿では、Explorer を使用して次のことを確認します。 key_averages() 数値的にレポートし、ディスパッチ間で何が起こるかを調べてください。

プロファイリング後、出力ディレクトリにはシステムトレーが含まれ、デバイスデータをキャプチャした場合は、一致した .neff および .ntff ファイル。Neuron Executable File Format(NEFF)にはコンパイル済みの Neuron 実行ファイルが含まれ、対応する NTFF にはそのハードウェア命令トレースが含まれます。これらのファイルは、ペアであることを示す数値ハッシュを共有しています。Neuron Explorer でプロファイルを開くには、システムプロファイルまたは結合プロファイルの場合は出力ディレクトリ全体をアップロードしてください。デバイスのみのプロファイルの場合は、対応する .neff および .ntff ファイルを個別に。詳細は ステップ5：プロフィールをアップロードする UI およびコマンドラインの説明用。

RMSNormのプロファイル

この最初の分析では、一連の明確な算術演算とリダクション演算からなる小規模な操作であるRMSNormをプロファイリングします。

同じ操作の3つの実装を比較します：

1. 個別のPyTorch演算(乗算、平均、加算、逆平方根、乗算)を用いた素朴なRMSNorm。
2. torch.rms_norm()、PyTorch の組み込み機能は単一の融合 Neuron カーネルへ lowered されました。
3. 素朴なRMSNormと torch.compile、最初の実装と同じ操作を使用します。

これらのスクリプトは同じ入力、プロファイラ設定、スケジュールを使用するため、結果を直接比較できます。

本研究成果を再現するためのすべてのスクリプトは、以下から入手できます。 torch-neuronx examples ディレクトリ.

プロファイル1: 素朴なPyTorchオペレーション

まず、個別のPyTorch操作を使ってRMSNormを書くことから始めます。各行は1つ以上のATen操作に対応し、中間テンソルを生成することがあります。

device = torch.device('neuron')

x = torch.randn(4, 512, dtype=torch.bfloat16, device=device)
weight = torch.ones(512, dtype=torch.bfloat16, device=device)

def rms_norm_fn(t):
    x_sq = t * t                                    # square
    mean_sq = x_sq.mean(dim=-1, keepdim=True)       # mean of squares
    rms = torch.rsqrt(mean_sq + 1e-6)               # 1 / sqrt(mean_sq + eps)
    return t * rms * weight                         # normalize + scale

Neuronバックエンドに必要な2つの引数を指定してプロファイラーを呼び出します。

1. 含める ProfilerActivity.PrivateUse1 Neuronデバイスからアクティビティを収集すること。
2. 合格 NeuronConfig として experimental_config Neuron 固有の Kineto 機能を有効にするため。

スケジュールでは wait=2, warmup=1、および active=1 4つにわたる prof.step() 呼び出しです。これにより、初期セットアップのオーバーヘッドを含めずに、1回のウォームイテレーションが記録されます。

電話 torch.neuron.synchronize() プロファイラのループ内部。この同期がないと、デバイスの処理完了前にプロファイラがエクスポートを行い、Neuron の時間列が空のままになってしまう可能性があります。

with torch.profiler.profile(
    activities=[ProfilerActivity.CPU, ProfilerActivity.PrivateUse1],
    experimental_config=nc,
    record_shapes=False,
    on_trace_ready=exporter.export_trace,
    schedule=torch.profiler.schedule(wait=2, warmup=1, active=1),
) as prof:
    for step in range(4):
        with torch.no_grad():
            _ = rms_norm_fn(x)
        torch.neuron.synchronize()
        prof.step()

プロファイリングが完了したら、次を確認します: key_averages() 表。並べ替え: self_privateuse1_time_total Neuronデバイスコストが最も高い操作を先頭に配置するためです。ここに示す表は、Neuronデバイス時間を報告するRMSNorm操作にフィルタリングされています。全文 key_averages() 出力には、割り当てやホストからデバイスへの転送など、CPU側のイベントも含まれます。

print(prof.key_averages().table(sort_by="self_privateuse1_time_total", row_limit=4))

主要な平均値の表を見ると、3つの aten::mul 呼び出しが Neuron デバイス時間の最大の割合を占めていることが分かります。

この表は、どのATen操作が実行されたか、実行頻度、およびCPUとNeuronデバイス上でどれほどの時間を費やしたかを示しています。 aten::mul "は、この関数が3回の乗算を実行するため、デバイス時間の合計が最も高くなっています。

Self NEURON 合計: 2.355 ms

各行は関数に直接対応しています。 aten::mul 実行時間は x * x, x * rmsであり、最後の次元に対する最終的な乗算は weight. aten::mean によって行われます。 aten::add はイプシロン値を加算し、スカラーが CPU 上で開始されるためホストからデバイスへのコピーを発生させます。 aten::rsqrt は逆平方根を計算します。

Neuron Explorer は aten::mul 操作を、フレームワークのディスパッチからランタイムのスケジューリングを経て、Neuron ハードウェアレーンでの実行まで追跡します。

Explorer でトレースを読む

Explorerでは、時間は左から右へ、イベントのライフサイクルは上から下へ流れます。各ATen操作は同じパターンに従います。フレームワークレーンでの作成、続いてランタイムレーンでのスケジューリングが行われ、 nrta_execute_schedule、ハードウェアレーンでのデバイス実行を nc_exec_running.

システムトレースビューでは、イベントがレーンに整理されます。Default グループ化モードでは、ラベルは instance/PID/subsystem/identifier というパターンに従います:

イーガーモードでは、各操作は独立して実行されます。依存関係チェーンビューアーを使って追跡してください。 aten::mul.Tensor 「フレームワーク層からランタイムスケジューリング、デバイス実行に至るまで。」

依存関係チェーンビューアーは、1つの操作のライフサイクルを、フレームワーク、ランタイム、ハードウェアの各レーンにわたって強調表示します。

Profile 2: torch.rms_norm

個々の演算をPyTorchの組み込み機能に置き換えます。 rms_norm function. Neuronはこの操作を1つの融合カーネルへと低減します:

def rms_norm_fn(t):
    return torch.rms_norm(t, [512], weight)

融合されたRMSNormは、key averagesテーブルで1つのデバイスオペレーションとして表示されます。

Self NEURON 合計: 403.3 μs

融合されたRMSNormは、各算術演算を個別にディスパッチする代わりに、単一のフレームワークからハードウェアへのパスに従います。

Loweringパスの比較

素朴な実装では、各演算が独立してディスパッチ境界を越え、それぞれ独自のフレームワークからランタイム、そしてハードウェアへのパスをたどります。融合実装では、1回のPython呼び出しで計算全体をカバーします。トレースには1回のディスパッチと1つのハードウェアイベントが含まれ、構成演算間に隙間はありません。

プロファイル3: torch.compileを使用した素朴なPyTorch

最後の例として、元の素朴なRMSNorm関数を torch.compile.

With @torch.compile(backend="neuron")でコンパイルすると、TorchDynamoがPythonバイトコードをFXグラフとしてキャプチャし、AOTAutogradがそれをATen演算へと低下（lower）させます。Neuronバックエンドは各演算を個別にディスパッチする代わりに、完全なグラフを受け取ります。バックエンドは演算を融合し、メモリアクセスを最適化し、グラフを単一のNEFFアーティファクトへとコンパイルします。このアーティファクトは1回のハードウェアディスパッチとして実行されます。

@torch.compile(backend='neuron', fullgraph=True)
def naive_compiled(t):
    x_sq = t * t
    mean_sq = x_sq.mean(dim=-1, keepdim=True)
    rms = torch.rsqrt(mean_sq + 1e-6)
    return t * rms * weight

コンパイルされた関数は、主要な平均テーブルに1つのTorch-Compiled Regionとして表示されます。

Self NEURON total: 375.5 μs

コンパイルされたグラフは、ランタイムのスケジューリングからハードウェアの実行まで1つのパスをたどります。

テーブルには、プロファイル2と同様に1回のデバイス実行が示されていますが、コードのたどるパスは異なります。専用に構築された融合演算を呼び出す代わりに、コンパイラが元のPython関数から融合を発見します。この実行では、Neuron時間は375 μsで、403 μsと torch.rms_norm, 375 μs であり、403 μs と比較されます。

並べて見る比較サマリー

両方 torch.rms_norm および torch.compile RMSNormを1つのカーネルディスパッチに融合します。これらの実行では、どちらの手法でも素朴な実装のデバイス時間の約6分の1になります。融合によりハードウェア実行間のギャップがなくなり、のような中間値を具現化する必要がなくなります。 x_sq および mean_sq 個別にディスパッチされる操作の間。

結論

key_averages() どの操作がデバイス時間を消費しているかを示し、その完全な出力にはCPU側のセットアップや転送も明らかにされます。Explorerは対応する視覚的トレースを提供します。1つの操作を、フレームワークのディスパッチからランタイムのスケジューリング、そしてハードウェア実行まで追跡し、融合(fusion)が取り除くギャップを確認できます。

より大きなモデルでも同じワークフローを使えます。表の中でホットスポットを見つけ、トレースでそれを検証し、その後、融合された演算かそれとも torch.compile コストを削減できる。

パート2では、デバイスレベルのプロファイリングに進みます。NEFFを開き、TensorEngine、VectorEngine、およびダイレクトメモリアクセス（DMA）の内訳を確認し、カーネルをルーフライン上に配置して、それがコンピュートバウンドかメモリバウンドかを判定します。

この連載を始めるきっかけを与えてくれたAritra Roy Gosthipaty、そしてレビューをしてくれた仲間のPinak Panigrahi、Samir Araujo、Akhil Azhikodan、Esha Lakhotia、その他の皆さんに心から感謝します!