カテゴリ: ランタイム: バッファ割り当ての失敗
このエラーは、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オブジェクトを保持すると、プログラムのコンパイルが完了した後でも自動割り当て解除が妨げられる可能性があります。
ツール
アロケータが OOM 発生時に現在のすべてのアロケーションの詳細なレポートをダンプする tpu_log_allocations_on_oom フラグを有効にします。これはデバッグに非常に役立ちます。