マルチスライス環境のプロファイリング

Cloud TPU マルチスライス環境は、データセンター ネットワーク（DCN）上で通信する複数の TPU Pod スライスで構成されます。DCN Collective Stats ツールを使用すると、マルチスライス環境が DCN ネットワークをどの程度効果的に使用しているかに関する情報を表示できます。具体的には、DCN Collective Stats ツールを使用すると、次のことができます。

• 収集したデータに基づいてスライス間のネットワーク パフォーマンスを表示して理解する
• パフォーマンスのボトルネックを特定する
• モデルのパフォーマンスを最適化する

DCN Collective Stats ツールのすべての指標は、TPU ごとに生成されます。

用語

DCN Collective Stats ツールは、マルチスライス環境内の TPU スライス間で発生する通信を説明する指標を表示します。TPU ランタイムがスライス間通信を開始すると、一連のオペレーションが使用されます。

send

ホストを中断してダイレクト メモリ アクセス（DMA）を開始し、データ転送を開始するために、いっぱいになったバッファをホストに提供します。

send-done

データ転送が完了したことをホストに通知します。

recv

ホストが転送されたデータを入力するための空のバッファを提供します。

recv-done

データを受信したことをホストに通知します。

注: 実際のデータ送信は、send オペレーションが完了した後に行われます。send-done オペレーションはデータの送信後に発生します。同様に、データの受信は、recv オペレーションの完了後です。recv-done オペレーションはデータの受信後に実行されます。

send オペレーションが発生するとグループが開始され、一致する recv-done オペレーションが発生すると完了します。

Slack 時間

グループがデータを送受信できる時間の尺度。これには、send、send-done、recv、recv-done オペレーションは含まれません。 たとえば、次のタイムラインがあるとします。

この例では、Slack 時間は次のように計算されます。

Slack 時間 = t₁ + t₂ + t₃

Slack 時間を増やすと、TPU がまとめて停滞する可能性が低くなります。別のシャーディング方法を選択すると、Slack 時間を長くできます。

ストール期間

送信、送信完了、受信、受信完了のオペレーションにグループが費やした平均時間。データの送信に費やした時間は含まれません。たとえば、次のタイムラインがあるとします。

この例では、ストール期間は次のように計算されます。

ストール期間 = t_送信+ t_送信完了+ t_受信+ t_受信完了

観測期間

データの送受信時間を含む、send オペレーションと recv-done オペレーションの間の時間。たとえば、次のタイムラインがあるとします。

観測期間は次のように計算されます。

観測期間 = t_送信+ t_{1 人}+ t_送信完了+ t_{2 個}+ t_受信+ t_{3 個}+ t_受信完了

発生回数

プロファイル期間中にグループが開始および完了した回数。send オペレーションが発生するとグループが開始され、一致する recv-end オペレーションが発生すると完了します。この指標に含まれる send オペレーションとその一致する recv-done オペレーションは、プロファイル期間内に行う必要があります。

集計ストールの合計

プロファイル期間中にグループが TPU を停止した合計時間。集計ストールの合計は次のように計算されます。

集計ストールの合計 = ストール期間 × 発生回数

転送されるデータのサイズ

プロファイル期間中にグループに対してネットワーク経由で送信されたデータの量。

必要な帯域幅

指定された Slack 内でデータを送信するために必要な帯域幅。この指標を使用して、プロファイル期間中にネットワーク帯域幅で競合するグループの数を確認できます。必要な帯域幅は次のように計算されます。

必要な帯域幅 = 送信データサイズ ÷ Slack 時間

ツールのステータス

次の表に、DCN Collective Stats ツールに表示される各指標に必要な TensorFlow または TPU ランタイム バージョンを示します。

DCN Collective Stats	TPU ランタイム バージョンのサポートされている TensorFlow
Slack 時間	TensorFlow 2.15.0、tensorboard 2.15.1、tensorboard-plugin-profile 2.15.0
ストール期間	TensorFlow 2.15.0、tensorboard 2.15.1、tensorboard-plugin-profile 2.15.0
観測期間	TensorFlow 2.15.0、tensorboard 2.15.1、tensorboard-plugin-profile 2.15.0
発生回数	TensorFlow 2.15.0、tensorboard 2.15.1、tensorboard-plugin-profile 2.15.0
集計ストールの合計	tf-nightly、tb-nightly、tbp-nightly
転送されるデータのサイズ	tf-nightly、tb-nightly、tbp-nightly
必要な帯域幅	tf-nightly、tb-nightly、tbp-nightly

DCN Collective Stats ツールの分析方法

1. TensorBoard サーバーを実行し、[プロファイル] タブに移動します。
2. DCN Collective Stats ツールのテーブルを、[集計ストールの合計] の降順で並べ替えます。
3. [集計ストールの合計] が最も高い DCN グループ名を特定します。このグループの集計ストール期間が他のグループと比較して大幅に長い場合は、DCN グループにボトルネックがあることを示している可能性があります。
4. DCN グループに必要な帯域幅にコア数を掛けます。v4 TPU ホストごとに 8 つのコアがあるため、グループに必要な帯域幅は、表示されている値の 8 倍です。必要な帯域幅が TPU の最大ネットワーク帯域幅より大きい場合は、ネットワークが輻輳している可能性があります。必要な帯域幅を減らすには、使用するシャーディング メカニズムを変更してみてください。シャーディング メカニズムの詳細については、Cloud TPU マルチスライスの概要をご覧ください。
5. HLO ダンプを生成して、コンパイラの問題があるかどうかを確認します。重複する HLO 演算のスケジューリングを可能にするには、send オペレーションと recv-done オペレーションをまとめてファンアウトすることをおすすめします。HLO 演算が重なると、TPU のストール時間が短縮されます。
6. Trace Viewer で、集計ストールの合計が最大である DCN グループの recv-done オペレーションの期間を確認します。転送時間が長い場合、データを取得するためにネットワーク上で recv-done オペレーションは通常ブロックされるため、帯域幅のボトルネックが発生する可能性があります。
7. recv-done オペレーションの時間が Slack 時間に対してそれほど長くない場合は、ハードウェアに問題がある可能性があります。