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

Cloud TPU マルチスライス環境は、データセンター ネットワーク（DCN）を介して通信する複数の TPU Pod スライスで構成されます。DCN 集計統計ツールを使用すると、マルチスライス環境が 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 オペレーションが発生すると完了します。

休憩時間

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

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

余裕時間 = 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 時間

ツールのステータス

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

DCN の集計統計情報	サポートされている TPU ランタイム バージョンの TensorFlow
休憩時間	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 時間に対してそれほど長くない場合は、ハードウェアに問題がある可能性があります。