Cloud TPU パフォーマンスガイド

TPU のパフォーマンスに関するトラブルシューティングの最初の手順は、モデルのプロファイリングを行うことです。パフォーマンスプロファイルのキャプチャの詳細については、Cloud TPU でのモデルのプロファイリングをご覧ください。

TPU モデルのパフォーマンス

このセクションでは、モデルのパフォーマンスを低下させる可能性がある一般的な問題と、その処理方法について説明します。

モデルが入力バウンド

TPU は非常に高速に計算を実行します。TPU がアイドル状態になっていないことを確認する場合、TPU に読み込まれているデータの安定したストリームがあることを確認することが重要です。これを行う方法は、データセットを読み込んで前処理する方法によって異なります。たとえば、tf.data.TFRecordset() と num_parallel_reads パラメータを使用してデータファイルを並列で読み取ることができます。
シャーディング（コア間でのバッチの分割）のためにバッチサイズが小さすぎる

TPU ランタイムは、TPU デバイスの 8 つのコアすべて（v2-8 や v3-8 など）にバッチを分割します。グローバルバッチサイズを 128 に指定すると、各コアは 16（128 / 8）のバッチサイズを受け取ります。

最適なメモリ使用量を実現するには、TPU メモリに収まる最も大きいバッチサイズを使用します。各 TPU コアは、行列乗算の処理に 2 次元 8 X 128 ベクトルレジスタを使用します。一般的に、バッチサイズは 8 または 128 で割り切れる必要があります。

XLA コンパイラの最適化

XLA は、TPU、CPU、GPU などのプラットフォーム用のバイナリを生成できる機械学習用のコンパイラです。XLA は標準の TensorFlow コードベースの一部ですが、PyTorch モデルと JAX モデルでも使用できます。Cloud TPU のモデルは XLA グラフに変換され、XLA は TPU 実行可能ファイルにコンパイルします。XLA の詳細については、XLA: 機械学習のためのコンパイラの最適化をご覧ください。

パディング

TPU のメモリを効率的に使用するには、128 x 8 のチャンクにタイリングできるようにデータを構造化します。行列計算のデータが 128 x 8 チャンク全体を満たさない場合、XLA コンパイラはテンソルをパディングします。パディングには次の 2 つの欠点があります。

パディングされたテンソルでは、TPU コアの利用が不十分になります。
パディングによって、テンソルに必要なオンチップメモリストレージ量が増加し、メモリ不足エラーにつながる可能性があります。

パディングは必要に応じて XLA コンパイラによって自動的に行われますが、メモリビューアツールを使用して行われるパディングの量を決定できます。TPU に適したテンソル次元を選択することで、パディングを回避できます。

テンソルディメンション

計算をより効率的に実行するために、XLA コンパイラは TPU HBM メモリに保存されたテンソルのサイズを切り上げます。このパディングはハードウェアレベルで透過的に行われ、結果には影響しません。ただし、場合によっては、パディングによってメモリ使用量と実行時間が大幅に増加する可能性があります。

TPU ランタイムは、計算効率を最大にし、パディングを最小限に抑えるように、メモリ内にテンソルをレイアウトします。メモリのオーバーヘッドを最小限に抑え、計算効率を最大化するには、次のいずれかに該当する必要があります。

バッチサイズの合計は 64 の倍数（TPU コアあたり 8）であり、特徴ディメンションサイズは 128 の倍数である。
バッチサイズの合計は 1,024 の倍数（TPU コアあたり 128）であり、特徴ディメンションサイズは 8 の倍数である。

バッチサイズ 1,024 および 128 の倍数である特徴ディメンションを使用すると、最大の効率性が得られます。ただし、これをすべてのモデルで実現できるとは限りません。

融合

融合は、XLA コンパイラがプログラムの最適化に使用する一般的な手法です。融合演算は、組み合わせて実行する、複数の構成演算の組み合わせです。

たとえば、次の一連の演算について考えてみます。

    tmp = tf.add(x, y)
    result = tf.multiply(tmp, z)

このコードは、次の擬似コードとほぼ同じです。

    for (i = 0; i < element_count; i++) {
      tmp[i] = x[i] + y[i];
    }

    for (i = 0; i < element_count; i++) {
      result = tmp[i] * z[i];
    }

融合により、配列アクセスが同時に発生します。

    for (i = 0; i < element_count; i++) {
      result = (x[i] + y[i]) * z[i];
    }

この例では、メモリのラウンドトリップ数が削減され、XLA は「tmp」にスペースを割り当てる必要はありません。

融合は重要な最適化であり、次のいくつかの点で Cloud TPU にメリットがあります。

低速のメインメモリに中間結果を保存する必要をなくして、メモリ転送を低減します。
通常は使用されないようなハードウェアユニットを有効に活用できます。
同時にアクティブ化されるバッファの数が少ないため、モデルのメモリ使用率を低減できます。

ブロードキャスト

ブロードキャストは、異なるが互換性のあるシェイプを持つ 2 つのテンソルが組み合わされたときに暗黙的に発生します。

たとえば、tf.add(vector, matrix) は、マトリックスのシェイプにベクトルをブロードキャストすることを要求します。演算の結果は、マトリックスと同じシェイプになります。詳細については、配列のブロードキャストのガイドをご覧ください。

ブロードキャストは多くの場合コンシューマと融合できますが、ブロードキャストを強制するとパフォーマンスが低下し、メモリ使用量が増加する可能性があります。

次の例では、ベクトルとマトリックスを追加する際の暗黙のブロードキャストを、ブロードキャストを実体化する argmax と融合することはできません。

`tf.argmax(tf.add(vector, zero_matrix), axis=0)`