カテゴリ: ランタイム: バッファ割り当ての失敗
このエラーは、XLA:TPU ランタイムのメモリ アロケータが、リクエストされた割り当てに対してアクセラレータの HBM 上で適切なメモリブロックを見つけられなかったことを示します。
エラー メッセージの例:
ValueError: RESOURCE_EXHAUSTED: Error allocating device buffer: Attempting to allocate 8.00M. That was not possible. There are 6.43M free.; (0x0x1_HBM0)
XLA バックエンド: TPU
概要
このエラーは、次の条件でスローされます。
jax.device_putを介したユーザー開始バッファ割り当ての失敗
または
- ユーザーがスケジュールしたプログラムの出力割り当ての失敗。
通常、これらの障害は次のいずれかの原因で発生します。
- メモリ不足（OOM）: ユーザーが、TPU の HBM で使用可能な空きメモリの合計量よりも大きいメモリチャンクを割り当てようとしています。
- メモリの断片化: メモリ空間にリクエストされたサイズを満たすのに十分な大きさの単一の連続した空きブロックがないため、割り当てが失敗します。空きメモリの合計量は割り当てに十分ですが、メモリ空間に小さな非連続ブロックとして分散しています。
TPU ランタイムには、割り当ての失敗を再試行するためのメカニズムがいくつか用意されています。
- キューに登録された割り当て解除がある場合、ランタイムは失敗した割り当てを再試行します。
- フラグメンテーションが原因で OOM が発生した場合、ランタイムは自動的にデフラグと再試行をトリガーできます。
- TPU ランタイムは、プログラムの読み込みを維持するよりもバッファ割り当てを優先します。HBM の不足が原因でバッファ割り当てが失敗した場合、システムは、バッファに十分なメモリが使用可能になるまで、読み込まれた TPU プログラムを削除します。
そのため、上記の緩和策を講じた後に発生したエラーには、通常、ユーザーの操作が必要です。
デバッグ
- モデルのメモリ使用量を削減する
- バッチサイズを小さくする: バッチサイズを小さくすると、メモリ使用量が直接的に減少します。
- パラメータ シャーディング: 非常に大規模なモデルの場合は、モデル並列処理やシャーディングなどの手法を使用して、複数の TPU コアまたはホストの HBM にパラメータを分散します。
- シーケンス/コンテキストの長さを短縮する: シーケンス（言語モデルなど）で動作するモデルの場合、入力シーケンス長を短縮すると、メモリ使用量を大幅に削減できます。
- バッファの提供: フレームワーク機能（
jax.jit(..., donate_argnums=...)など）を利用して、特定の入力バッファを上書きして出力に再利用できることを XLA に通知します。詳しくは、バッファの寄付をご覧ください。 - チェックポイント戦略を最適化する: モデルの状態全体を一度に保存するのではなく、モデルの重みのみを保存するか、シャード チェックポイント戦略を使用することを検討してください。
- メモリー レイアウトとパディングのアドレス指定
- TPU メモリはチャンク単位で割り当てられ、パディングによってテンソルの実際のサイズが増加する可能性があります。
- メモリリークが発生しないことを確認する
jax.Arrayオブジェクトへの参照が意図したよりも長く保持されていないことを確認します。jax.Arrayオブジェクトを保持すると、プログラムのコンパイルが完了した後でも自動割り当て解除が妨げられる可能性があります。
各プログラムが使用するメモリ量を減らすために使用できるその他の方法については、エラーコード: E1000 もご覧ください。
ツール
アロケータが OOM 発生時に現在のすべてのアロケーションの詳細レポートをダンプする tpu_log_allocations_on_oom フラグを有効にします。これはデバッグに非常に役立ちます。