TensorFlow - TPU のトラブルシューティング

このガイドでは、Cloud TPU で TensorFlow モデルをトレーニングしているユーザー向けに、トラブルシューティングのヘルプをよくある質問とともに提供します。Pytorch ま��は JAX のトレーニングのトラブルシューティングを行う場合は、以下のフレームワークのトラブルシューティングに関するドキュメントをご覧ください。

PyTorch - TPU のトラブルシューティング
JAX - TPU のトラブルシューティング

重要: Cloud TPU は、TPU VM と TPU ノードという 2 つの TPU アーキテクチャをサポートしています。この 2 つの VM アーキテクチャについては、システムアーキテクチャをご覧ください。このドキュメントのトラブルシューティング手順は両方の構成に適用されますが、TPU VM 構成でのみ発生する問題のトラブルシューティングのための、専用の別の TPU VM セクションもあります。

Cloud TPU の使用方法に関する全般的なガイドについては、以下をご覧ください。

概要

Cloud TPU で発生する一般的な問題は、次のカテゴリに分類されます。

TPU サーバーへの接続の問題

このセクションでは、TPU への接続時に TensorFlow がレスポ��スを停止、またはエラーを出力する状況でのトラブルシューティング方法について説明します。大規模なモデルの場合は TPU グラフのコンパイルステップに時間がかかることがあるため、スクリプトを少なくとも 5 分間実行するまでは、レスポンスが停止していると判断しないでください。

最初のステップは、問題がサーバー自体にあるのか、TensorFlow トレーニングパイプラインにあるのかを確認することです。これを行うには、TPU サーバー URL を使用して MNIST チュートリアルを実行し、これが正しく動作することを確認します。MNIST チュートリアルとの接続の問題がまだある場合は、TPU サーバーの問題であることが確認されます。この場合、次の手順に従います。

次のコマンドを実行して、使用可能な TPU を一覧表示します。zone と project-id は、ゾーンとプロジェクト ID に置き換えます。

(vm)$ gcloud compute tpus list --zone zone --project project-id

これにより、次のような出力が表示されます。

NAME       ZONE           ACCELERATOR_TYPE  NETWORK_ENDPOINT   NETWORK  RANGE          STATUS
demo-tpu   us-central1-b  v2-8              10.240.1.2:8470    default  10.240.1.0  READY

--tpu に正しい値を渡していること（上記の例では demo-tpu）と、この TPU が READY として表示されていることを確認します。
TPU が READY として表示されていない場合や、接続の問題が続いている場合には、以下のコマンドを使用して手動でサーバーを再起動します。
```
(vm)$ gcloud compute tpus stop $TPU_SERVER_NAME && gcloud compute tpus start $TPU_SERVER_NAME
```
上記の例では、$TPU_SERVER_NAME は demo-tpu です。完了までに数分かかることがあります。
上記の ... tpus list コマンドを再実行し、TPU が READY 状態になるのを待ちます。この処理には数分かかることがあります。
MNIST チュートリアルを再度実行してみます。
この手順に従っても MNIST チュートリアルを正常に実行できない場合は、サポートの利用で説明しているメカニズムのいずれかを使用してサポートをご依頼ください。

MNIST の例が正しく実行されるがモデルがまだ反応しない場合、問題はトレーニングパイプラインにある可能性があります。これをデバッグするには、まず、コード内の TPUStrategy をデフォルトの戦略に置き換えます。デフォルトの戦略を使用する場合、strategy.scope() または strategy.run() を使用すると、モデルは TPU ではなく CPU（GPU が存在する場合は GPU）で実行されます。モデルが TPU ではなく CPU で実行されている場合は、TPU 固有の問題があるはずです。それでも実行されない場合は、CPU の問題をデバッグすることをおすすめします。

トレーニング中の `ssh` 接続の喪失

Cloud TPU への ssh 接続は、長時間実行されるトレーニング中にタイムアウトする可能性があります（特に Cloud Shell を使用している場合）。この時点で、TPU コンソールへの出力はなく、TPU がトレーニングを停止したかのように表示されることがあります。これを回避するには、ターミナルマルチプレクサやセッション管理ツール（tmux、screen など）を使用してトレーニングセッションを実行します。これにより、トレーニングの長さに関係なく、ssh 接続がアクティブになります。

一般的なエラーのデバッグ

TPU を作成できない

Cloud TPU を作成する際に、次のエラーが発生することがあります。

googleapiclient.errors.HttpError: < HttpError 403 when requesting https://content-tpu.googleapis.com/v1/projects/{PROJECT}/locations/{ZONE}/nodes/{TPU_NAME}?alt=json returned "Request had insufficient authentication scopes."

これは権限の問題であり、次のコマンドを実行して解決できます。

gcloud auth login --update-adc

このコマンドは、アプリケーションのデフォルト認証情報（ADC）を更新し、問題を解決します。詳細については、gcloud auth login をご覧ください。

ローカルファイルシステムを使用できない

エラーメッセージ

InvalidArgumentError: Unimplemented: File system scheme '[local]' not implemented

影響を受けるフレームワークと構成

このメッセージは、TPU ノードアーキテクチャを使用して TensorFlow でトレーニングを行うと出力されることがあります。

詳細

すべての入力ファイルとモデルディレクトリで Cloud Storage バケットパス（gs://bucket-name/...）が使用されている必要があり、このバケットに TPU サーバーからアクセスできる必要があります。すべてのデータ処理とモデルチェックポインティングは、ローカルマシンではなく TPU サーバー上で実行されることにご注意ください。Cloud Storage で TPU との使用を適切に構成する方法については、Cloud Storage バケットへの接続のガイドをご覧ください。

サポートされないデータ型

エラーメッセージ

TypeError: DataType is not a supported TPU infeed type.

影響を受けるフレームワークと構成

このメッセージは、TPU ノードアーキテクチャを使用して TensorFlow でトレーニングを行うと出力されることがあります。

詳細

現在、tf.float32、tf.int32、tf.bfloat16、tf.bool データ型のみが TPU でサポートされています。その他の一般的なデータ型（tf.uint8、tf.string、tf.int64 など）は、データの前処理中に（つまり、tf.data.Dataset パイプラインで）サポートされているデータ型のいずれかに変換する必要があります。

MNIST トレーニングで使用する decode_image 関数での変換の例をご覧ください。

動的形状がサポートされない

エラーメッセージ

ValueError: shape [Shape] must have a fixed size for dimension
d that is known at graph construction time.

影響を受けるフレームワークと構成

このメッセージは、TensorFlow を使用した XLA コンパイル中にのみ出力されます。

詳細

TPU でモデルを実行するために、TensorFlow はXLA コンパイラを使用してモデルをコンパイルします。このコンパイルステップにより、トレーニング速度とメモリ使用量は大幅に改善されますが、グラフ内のすべてのテンソルの形状（ディメンションサイズ）は、グラフをコンパイルする時点で既知である必要があります。コンパイルする時点で形状を判別できない場合、上記のようなエラーで TPU コンパイルは失敗します。

動的形状が返される一般的な演算の 1 つは、dataset.batch(batch_size) です。これは、ストリームに残っているサンプル数がバッチサイズよりも小さい可能性があるためです。したがって、TPU でのトレーニングでは、dataset.batch に drop remainder=True を設定します。この場合、すべてのバッチが batch_size という静的形状になるように、最後のいくつかのサンプルがファイルから削除される可能性があります。次に例を示します。

dataset = tf.data.Dataset.range(8)
dataset = dataset.batch(3, drop_remainder=True)

利用できない TensorFlow 演算

エラーメッセージ

NotFoundError: No registered 'OpName' OpKernel for XLA_TPU_JIT
devices compatible with node

影響を受けるフレームワークと構成

このメッセージは、TensorFlow を使用してトレーニングするときに出力されることがあります。

詳細

TPU で現在は利用できない TensorFlow 演算がモデルで使用されています。

TPU で利用可能な演算のリストと、今後のサポート計画や回避策の提案については、利用可能な TensorFlow 演算のガイドをご覧ください。

メモリ不足のエラーメッセージ

エラーメッセージ

ResourceExhaustedError: Ran out of memory in memory space hbm; used:
YYY; limit: 7.48G.

影響を受けるフレームワークと構成

このメッセージは、TensorFlow、PyTorch、JAX を使用してトレーニングするときに出力されることがあります。

詳細

各 Cloud TPU は 8 個の TPU コアで構成され、v2 TPU には 8 GB、v3 TPU には 16 GB の RAM（または高帯域幅メモリ（HBM））が搭載されています。このメモリは、勾配計算に必要な中間結果のテンソルだけでなく、重み（変数）テンソルを保存するためにも使用されます。モデルが大きすぎて TPU RAM に収まらない場合、初期化は失敗し、上記のエラーメッセージが表示されます。詳細については、メモリ使用量の削減をご覧ください。

メモリ使用量��減らすためのヒント:

過剰なテンソルパディングを確認する
bfloat16 形式を使用する
入力サイズ（モデル）が大きすぎる場合は、TensorFlow の試験運用版モデルの並列処理を使用すると、モデルのサイズに対応できる可能性があります。

実行の停止を起こす問題

TPU の実行中に TensorFlow でエラーが発生した場合、スクリプトはシェルを終了するのではなくレスポンスが停止しているように見えることがあります。このような場合は、キーボードの CTRL+\ キーを押して SIGQUIT をトリガーします。これにより、Python はすぐに終了します。

同様に、TPU の実行中に CTRL+C キーを押すと、TensorFlow はすぐにはシャットダウンされず、現在のイテレーションループの終了を待って正常に終了します。

この方法で終了した後、TPU に再接続しようとして新しいエラーが発生した場合は、次のコマンドを使用して TPU サーバーを手動でリセットします。

gcloud compute tpus stop tpu-name --zone=zone
gcloud compute tpus start tpu-name --zone=zone

tpu-name は、gcloud compute tpus list コマンドで表示される最初の列から選択し、zone は 2 番目の列に表示されるゾーンです。

テンソルパディングが過剰である

メモリの問題の考えられる原因

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

メモリ使用量を減らす方法

TPU ソフトウェアは、計算効率を最大にし、パディングを最小限に抑えるように、メモリ内にテンソルを��イアウトしようとします。ただし、このメモリレイアウトプロセスは複雑であり、最適な結果を得るにはモデルが以下の経験則に従っている必要があります。メモリのオーバーヘッドを最小限に抑え、計算効率を最大化するには、次のいずれかに該当する必要があります。

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

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

バッチサイズ 1,024 および 128 の倍数である特徴ディメンションを使用すると、最大の効率性が得られます。ただし、これをすべてのモデルで実現できるとは限りません。明確には、「特徴ディメンション」は、全結合層の隠れサイズまたは畳み込みにおける出力チャネルの数を指します。すべての層がこのルールに準拠できるわけではありません（特にネットワークの最初と最後の層）。これは問題なく、ほとんどのモデルでは一定量のパディングが必要になると想定されます。

メモリ使用量の削減

TPU でモデルを実行しているときにメモリ不足エラーが発生した場合は、モデルのメモリ使用量を減らすための手順を取る必要があります。

メモリ使用量を減らす最も効果的な方法は次のとおりです。

過剰なテンソルパディングを減らす
バッチサイズを小さくする

バッチサイズまたはモデルが大きすぎる

メモリの問題の考えられる原因

CPU、GPU、TPU でニューラルネットワークをトレーニングする場合、次の 2 か所でメモリの使用が発生します。

メモリ使用量は、モデル内の重みの数に比例します。
バックワードパスの計算に必要なフォワードパスからの中間アクティベーションの保存。メモリ使用量は、バッチサイズ、層のサイズ、層の数に正比例します。

したがって、モデルの所要メモリはバッチサイズに大きく依存します。

モデルに必要なメモリは、ネットワーク内のレイヤの数によって異なります。

TPU ランタイムは、モデルをメモリに収まるように演算子を最適化しようとしますが（再実体化と呼ばれ、勾配チェックポインティングと類似しています）、常に変換できるとは限りません。

メモリ使用量を減らす方法

メモリ内に収まるまでバッチサイズを徐々に減らし、バッチサイズの合計が 64 の倍数になるようにします（コアあたりのバッチサイズは 8 の倍数である必要があります）。TPU ではバッチサイズが大きければ大きいほど効率的になります。バッチサイズの合計を 1,024（コアあたり 128）から始めることが一般的には適切です。

小さなバッチサイズ（64 など）でもモデルを TPU で実行できない場合は、層の数または層サイズを減らして��てください。

トレーニング速度の向上

モデルを TPU で正常に実行できるが、トレーニング速度が想定ほど速くない場合があります。このセクションでは、速度の向上の実現に役立ついくつかの方法の概略を説明します。トレーニングのパフォーマンスを向上させるためのその他のおすすめの方法については、パフォーマンスガイドをご覧ください。

各トレーニングループの実行ごとのステップ数が少なすぎる

パフォーマンスの問題の説明

引数 steps_per_execution を Model.compile に渡すと、ホストコールバック間で実行されるトレーニングステップの数が制御されます。各ホストコールバックは、TPU サーバーのホスト CPU と TPU デバイス間の重要な通信を必要とするため、steps_per_execution が小さすぎるとトレーニングが遅くなる可能性があります。

モデルに影響しているかどうかを確認する方法

TPU プロファイルで、TPU デバイスのステップ間でホスト CPU コールバックが頻繁に発生していることが明らかになった場合は、steps_per_execution 値を大きくするとトレーニングで効果が出る可能性があります。

緩和する方法

steps_per_execution を大きい値に設定します。steps_per_execution に大きい値を設定できますが、メッセージのロギングとチェックポイントの保存は、指定したステップ数を実行した後にのみ行われることに注意してください。

入力処理のボトルネック

パフォーマンスの問題の説明

TPU が特定のデータチャンクでトレーニングしている間に、入力処理関数によって CPU で次のデータチャンクが準備されます。入力関数がモデル関数よりも時間がかかる場合、入力関数がデータを取得している間、TPU はアイドル状態のままです。

モデルに影響しているかどうかを確認する方法

TensorBoard で入力パイプライン分析を表示するには、Cloud TPU ツール: 入力パイプライン分析ツールの指示に従います。

入力パイプライン分析ページには、モデルで入力処理がボトルネックになっているかどうかを示す明確な概要が表示されます。同じページに演算ごとの実行時間も表示され、問題のある演算を特定できます。

緩和する方法

Dataset API を使用してデータを読み込むときに、いくつかの緩和策を利用できます。

データを tf.train.Example 構造のコレクションとして TFRecord ファイルに保存し、それを TFRecordDataset を使用して読み込みます。例については、Dataset API チュートリアルまたは ResNet チュートリアルをご覧ください。
dataset.cache() か dataset.prefetch()、またはその両方を使用して入力データをバッ��ァリングします。これにより、ファイルアクセスが突発的に遅くなってボトルネックが発生するのを防ぎます。
dataset.map() 関数の num_parallel_calls パラメータを指定して、map() 演算のマルチスレッド化を有効にします。num_parallel_calls の値のシンプルな経験則は、使用可能な CPU コアの数を使用することです。
すべてのトレーニングのすべてのエポックでコストをかけるのではなく、コストのかかるデータ前処理を 1 回限りのコストとしてオフラインで実行します。

すべての入力処理はローカルマシン上ではなく TPU サーバー上にある CPU で実行されるため、ローカルマシンの速度は要因ではありません。

ステップ速度が遅く、MXU 使用率が低い

パフォーマンスの問題の説明

Cloud TPU は、行列乗算と畳み込みを非常に高速に実行できます。他の TensorFlow 演算の大部分は TPU に効率的に実装されていますが、これらは他のハードウェアに対する TPU の主要な強みではありません。したがって、TPU を十分に利用するには、モデルで行列乗算または畳み込みが多くを占めている必要があります。

モデルに影響しているかどうかを確認する方法

この場合の症状としては、ステップ速度が遅いことが、パフォーマンスをプロファイリングしたときに表示される MXU 使用率が低いことと相まって見て取れます。

緩和する方法

行列乗算以外の演算の数を減らすようにしてください。行列乗算の数を減らした後、再度ベンチマークを行って、TPU で許容できるパフォーマンスかどうかを確認します。

テンソルパディングが過剰である

パフォーマンスの問題の説明

TPU では、その計算ユニットを効率的に使用できるように、メモリ内のテンソルはパディングされます。パディングによって、メモリとメモリ帯域幅の両方の使用量が増加する可能性があります。テンソルパディングの問題を理解して解決するには、テンソルパディングをご覧ください。

スループットが低く、メモリ使用量が少ない

パフォーマンスの問題の説明

一般的に、使用するバッチサイズが大きいほど、サンプル数 / 秒に関して TPU でのトレーニング速度は向上します。

モデルに影響しているかどうかを確認する方法

モデルのバッチサイズは、常に少なくとも 64（TPU コアあたり 8）である必要があります。これは、TPU ではテンソルは常にこのサイズにパディングされるためです。TPU でのトレーニングの理想的なバッチサイズは、1024（TPU コアあたり 128）です。これにより、メモリ転送およびパディングに関する非効率性が排除されるためです。

緩和する方法

メモリに収まり 64 の倍数である最大のバッチサイズを使用することをおすすめします。これを実現する最も簡単な方法は 1024 から開始することであり、メモリ不足エラーが発生した場合は、モデルが正常に実行されるまでバッチサイズを減らしてみます。モデルのバッチサイズを変更するには、同じモデル精度（学習率など）を実現するために他のハイパーパラメータを調整する必要がある場合がありますが、これはケースバイケースで評価する必要があります。

層サイズが小さすぎる

パフォーマンスの問題の説明

モデルで行列乗算または畳み込みが多くを占めていても、入力テンソルが小さい場合、TPU は最も効率的に動作しないことがあります。他のハードウェアと比較すると、TPU は��ッチサイズ��層サイズの両方が大きい場合に最も効率的に動作します（たとえば、ディメンション >= 512）。

モデルに影響しているかどうかを確認する方法

一般に、層サイズが 128 よりも小さいと、128 が TPU 行列乗算ユニットの固有のディメンションであるため、TPU の効率は悪くなります。全結合層の場合、高い効率を実現するには 512 の最小隠れサイズをおすすめします。畳み込み層は、通常、等しい効率レベルを実現するために、全結合層ほどの大きさである必要はないことに注意してください。

緩和する方法

モデル内の層サイズを小さくする主な動機がトレーニング速度である場合は、より大きい層を使用して TPU でモデルを再ベンチマークします。たとえば、層の出力サイズを 256 から 512 に増やすと、モデルが 2 倍の計算を実行してもトレーニング時間は 20% しか増加しません。

演算レベルのモデルプロファイリング

パフォーマンスのボトルネックを特定するには、多くの場合、演算レベルの実行時間とメモリ使用量を測定することが有益です。これを行う方法については、
Cloud TPU ツール: トレースビューアのガイドをご覧ください。

モデルの精度の低下のデバッグ

Cloud TPU エコシステムの目標の 1 つは、現在 CPU または GPU でトレーニングしているモデルを TPU でトレーニングしたときに、バッチサイズや学習率などのハイパーパラメータを多少調整するだけで、同様の精度が実現されるようにすることです。ただし、TPU でモデルをトレーニングしたときに、精度の低下が確認されることがあります。このような問題のデバッグは、ニューラルネットワークトレーニングのランダムな性質のために、非常にもどかしい場合があります。このセクションでは、モデルを TPU に移植する際に、モデルの精度の低下の根本原因を特定する方法について説明します。

データシャーディング（データ並列処理）について

TensorFlow の主要な目標の 1 つは、CPU、GPU、TPU のいずれで実行しても、各演算からほぼ同じ結果が生成されるようにすることです。これには、ランダム演算などいくつかの例外があります。一般に、TPU と CPU で非ランダム演算の出力に大きな違いがある場合は、バグとして報告してください。

ただし、トレーニングパイプライン全体としては、CPU / GPU と TPU でのトレーニングには大きな違いがあります。TPU でのトレーニングでは、TensorFlow がデータシャーディングを実行します。各 Cloud TPU には 8 個の TPU コアがあり、それらは独立したプロセッシングユニットとして動作します。トレーニングのステップごとに、各 TPU コアにデータのバッチが受け渡され、重みの勾配が計算され、勾配が他の TPU コアと交換され、重み更新が計算されます。デフォルトでは、損失はコア全体で平均化されますが、CrossShardOptimizer ��パラメータを変更することで合計できます。

モデルの全損失を、独立したサンプルあたりの損失の平均（または合計）としてコンピューティングできる場合、この手順は数学的には単一の大きなバッチでのトレーニングと同等です。

サンプルごとに独立していない最も一般的な演算はバッチ正規化であり、これはコア別の各バッチに対して個別に実行されます。たとえば、バッチサイズの合計が 128 の場合、コアあたりのバッチサイズは 16 であり、8 個のコアのそれぞれが固有の 16 サンプルにわたってバッチ正規化を行います。場合によっては、小さいバッチ（32 未満など）にわたってバッチ正規化を行うことにより、精度が低下することがわかっています。理想的なシナリオでは、バッチサイズの合計を大きくする必要があります（たとえば、256 から 1024）。バッチサイズが大きすぎてメモリに収まらない場合は、シャーディングの影響をケースバイケースで評価する必要があります。

決定論的なトレーニング

モデルの精度の違いをデバッグすることが難しい理由の 1 つは、異なるフレームワーク（TensorFlow、PyTorch、JAX）のトレーニングソフトウェアでモデルがトレーニングされるたびに、異なる重みの初期化とデータシャッフルが使用されることです。決定論的になるようにトレーニング手順を変更して、複数の実行でほぼ同一のモデルが生成されるようにすると有益です。このセクションでは、MNIST チュートリアルを決定論的に実行する方法を示します。

CPU で 1 ステップ実行することにより、初期チェックポイントファイルを生成します。このステップは、決定論的な重みの初期化を達成するために使用されます。また、モデル内のランダム関数に、固定のランダムシードが使用されていることを確認してください。

# Run training for 1 step to create an initial checkpoint.
python mnist_tpu.py \
  --use_tpu=False \
  --data_dir=${STORAGE_BUCKET}/data/ \
  --model_dir=${STORAGE_BUCKET}/init_output \
  --random_seed=12345 \
  --iterations=1
  --train_steps=1

入力関数内のデータシャッフル関数を変更して、ランダムシードが使用されるようにします。これは MNIST チュートリアルですでに行われています。入力データ処理演算は常に CPU で実行されるため、これは入力データ処理演算で機能します。モデル関数でのランダム演算は、TPU と CPU の間で決定論的ではない場合があります。ランダム演算を呼び出すたびに、固定シードを渡すと実行間で同じ結果を得ることができます。次に例を示します。

# In the flag definitions
tf.flags.DEFINE_integer("batch_size", None, "Random seed for training")

# In the input_fn
if FLAGS.random_seed is not None:
dataset = dataset.shuffle(seed=FLAGS.random_seed)

同じモデルを CPU で 2 回実行し、トレーニングが決定論的であることを確認します。トレーニングは妥当なステップ数（たとえば、1,000）で実行する必要がありますが、収束まで実行する必要はありません。

CPU トレーニングはシングルコア TPU トレーニングと比較されるため、単一の TPU コアに収まるバッチサイズを使用します（通常、バッチ全体を 8 で除算したサイズ）。TensorFlow では実行間でビット対ビットの決定論は保証されませんが、損失は非常に近似するはずです。

初期の重みをコピーする

gsutil mkdir ${STORAGE_BUCKET}/cpu_output_1
gsutil cp -f ${STORAGE_BUCKET}/init_output/* ${STORAGE_BUCKET}/cpu_output_1
gsutil mkdir ${STORAGE_BUCKET}/cpu_output_2
gsutil cp -f ${STORAGE_BUCKET}/init_output/* ${STORAGE_BUCKET}/cpu_output_2

実行 1

python mnist_tpu.py \
  --use_tpu=False \
  --data_dir=${STORAGE_BUCKET}/data/ \
  --model_dir=${STORAGE_BUCKET}/cpu_output_1 \
  --batch_size=128 \
  --random_seed=12345 \
  --train_steps=2000 \
  --eval_steps=10

出力 1

accuracy = 0.9910644, global_step = 1000, loss = 0.025323588

実行 2

python mnist_tpu.py \
  --use_tpu=False \
  --data_dir=${STORAGE_BUCKET}/data/ \
  --model_dir=${STORAGE_BUCKET}/cpu_output_1 \
  --batch_size=128 \
  --random_seed=12345 \
  --train_steps=2000 \
  --eval_steps=10

出力 2

accuracy = 0.9910644, global_step = 1000, loss = 0.025323414

シングルコア TPU トレーニング

MNIST チュートリアルを決定論的に実行できたら、次のステップでは、CPU でのトレーニング結果を TPU で複製し、単一の TPU コアを使用して、問題がデータシャーディングに関連するか TPU 実行エンジン自体に関連するかを特定します。

MNIST チュートリアルでシングルコアトレーニングと評価を実行する方法は次のとおりです。

CPU と同じ重みの初期化を使用する

gsutil cp -f ${STORAGE_BUCKET}/init_output/* ${STORAGE_BUCKET}/tpu_output

1,000 ステップのトレーニングを実行する

python mnist.py \
    --use_tpu=True \
    --master=$GRPC_SERVER \
    --train_file=${STORAGE_BUCKET}/data/train.tfrecords \
    --model_dir=${STORAGE_BUCKET}/tpu_output \
    --random_seed=12345 \
    --num_shards=1 \
    --batch_size=128 \
    --train_steps=1000 \
    --eval_steps=10

出力

  accuracy = 0.9910644, global_step = 1000, loss = 0.02514153

損失は CPU トレーニング済みモデルと完全には一致しませんが、近似するはずです。対象のモデルで近似しない場合は、TPU 実行エンジンでバグが見つかったことを示している可能性があります。バグレポートを提出する際は事前に、次のことを再確認してください。

num_shards=1 を TPUConfig に渡している。
モデル関数にランダム演算はなく、入力関数内のランダム演算は正しくシードされている。
CPU と TPU のトレーニングに同じ初期チェックポイントファイルを使用している。

マルチコア TPU トレーニングのデバッグ

モデルにおいて、CPU とシングルコア TPU で同じ損失を達成した場合、問題は次のいずれかである可能性があります。

（a）ニューラルモデルを異なる初期化でトレーニングする際の自然なランダム差異による劣化。

（b）TPU でのデータシャーディングに関する問題��原因��化。

（a）��題であるかどうかを判断するには、上記のように、同じ重みの初期化を使用して CPU / GPU とマルチコア TPU でモデル全体を再トレーニングします。

精度の低下が統計的に有意であると判断した場合、データシャーディングに関して可能性が最も高い問題は次のとおりです。

モデルでバッチ正規化を使用する場合、バッチサイズの合計が 256 未満（たとえば、コアあたり 32 未満）であると、精度が低下する可能性があります。
バッチ式の損失関数はシャーディングの影響を受けます。そのような損失関数は、通常、非常に特殊化されています。たとえば、Karras et al. 2017 は、敵対的生成ネットワークをトレーニングする際に、バッチ弁別子を使用します。

TPU VM のトラブルシューティング

以下の問題と解決策は、TPU VM の構成にのみ適用されます。

gcloud 設定のトラブルシューティング

問題: gcloud components update を実行すると、次のエラーメッセージが表示される。

ERROR: (gcloud.components.update)
You cannot perform this action because the Cloud SDK component manager is
disabled for this installation.

解決策: TPU VM で gcloud を使用するには、パッケージマネージャーで管理されない gcloud インストールを使用する必要があります。ソースコードから gcloud をインストールするには、次の手順を行います。

  sudo apt-get remove google-cloud-sdk
  curl -O https://dl.google.com/dl/cloudsdk/channels/rapid/downloads/google-cloud-sdk-311.0.0-linux-x86_64.tar.gz
  tar -xzf google-cloud-sdk-311.0.0-linux-x86_64.tar.gz
  ./google-cloud-sdk/install.sh
  source ~/.bashrc

問題

gcloud compute tpus tpu-vm ssh ${TPU_NAME} --zone ${ZONE} コマンドを実行すると、次のエラーメッセージが表示される。

Waiting for SSH key to propagate.
ssh: connect to host 34.91.136.59 port 22: Connection timed out
ssh: connect to host 34.91.136.59 port 22: Connection timed out
ssh: connect to host 34.91.136.59 port 22: Connection timed out
ERROR: (gcloud.compute.tpus.tpu-vm.ssh) Could not SSH into the instance.  It is possible that your SSH key has not propagated to the instance yet. Try running this command again.  If you still cannot connect, verify that the firewall and instance are set to accept ssh traffic.

解決策

SSH 認証鍵の伝播に問題がある可能性があります。自動生成された鍵をバックアップの場所に移動して、gcloud で強制的に再作成してみてください。

mv ~/.ssh/google_compute_engine ~/.ssh/old-google_compute_engine
mv ~/.ssh/google_compute_engine.pub ~/.ssh/old-google_compute_engine.pub

デバッグログ

サポートされている Cloud TPU フレームワーク、JAX、PyTorch、TensorFlow は、すべての TPU VM で提供される libtpu という共有ライブラリを介して TPU にアクセスします。このライブラリには、TPU プログラムのコンパイルに使用される XLA コンパイラ、コンパイルされたプログラムの実行に使用される TPU ランタイム、TPU への低レベルアクセスにランタイムで使用される TPU ドライバが含まれます。

libtpu ライブラリは、デバッグに役立つ可能性のある情報をログに記録します。デフォルトでは、これらのログは各 Cloud TPU VM 上の /tmp/tpu_logs に書き込まれます。トレーニングを開始する前に、次の環境変数を設定して、ロギング動作を変更できます。

TPU_LOG_DIR: ログが書き込まれるディレクトリ: ディレクトリの場所はデフォルトで /tmp/tpu_logs になります。ディレクトリがまだ存在しない場合は作成されますが、親ディレクトリは作成されません。指定されたディレクトリの検索または作成時にエラーが発生した場合、メッセージが stderr に出力されますが、プログラムは停止されず、ロギングは無効になります。ディレクトリ名を「disabled」に設定して、ディスクへのロギングを完全に無効にします。
TPU_MIN_LOG_LEVEL: ディスクに記録される最小の重大度: 選択肢は 0（INFO）、1（WARNING）、2（ERROR）、3（FATAL）です。デフォルト値は 0 です。
TPU_STDERR_LOG_LEVEL: 該当する場合にディスクに加えて stderr に記録される最小の重大度: 選択肢は、TPU_MIN_LOG_LEVEL と同じです。デフォルト値は 3 です。
TPU_MAX_LOG_SIZE_MB: 各ログファイルの最大サイズ（メガバイト単位）: ログファイルがほぼこのサイズになると、新しいログファイルが自動的に開始されます。デフォルトは 1,024 です。

TensorFlow - TPU のトラブルシューティング

概要

TPU サーバーへの接続の問題

トレーニング中の ssh 接続の喪失

一般的なエラーのデバッグ

TPU を作成できない

ローカル ファイルシステムを使用できない

サポートされないデータ型

動的形状がサポートされない

利用できない TensorFlow 演算

メモリ不足のエラー メッセージ

実行の停止を起こす問題

テンソル パディングが過剰である

メモリ使用量の削減

バッチサイズまたはモデルが大きすぎる

トレーニング速度の向上

各トレーニング ループの実行ごとのステップ数が少なすぎる

入力処理のボトルネック

ステップ速度が遅く、MXU 使用率が低い

テンソル パディングが過剰である

スループットが低く、メモリ使用量が少ない

層サイズが小さすぎる

演算レベルのモデル プロファイリング

モデルの精度の低下のデバッグ

データ シャーディング（データ並列処理）について

決定論的なトレーニング

初期の重みをコピーする

実行 1

出力 1

実行 2

出力 2

シングルコア TPU トレーニング

CPU と同じ重みの初期化を使用する

1,000 ステップのトレーニングを実行する

出力

マルチコア TPU トレーニングのデバッグ

TPU VM のトラブルシューティング

gcloud 設定のトラブルシューティング

デバッグログ

トレーニング中の `ssh` 接続の喪失

ローカルファイルシステムを使用できない

メモリ不足のエラーメッセージ

テンソルパディングが過剰である

各トレーニングループの実行ごとのステップ数が少なすぎる

テンソルパディングが過剰である

演算レベルのモデルプロファイリング

データシャーディング（データ並列処理）について