カテゴリ: ランタイム: プログラム入力バッファの不一致
このエラーは、コンパイル済みプログラムで想定されるメモリバッファのサイズと、実行時に実際に提供されるバッファのサイズとの間に不一致が XLA ランタイムで検出された場合に発生します。
エラー メッセージの例:
XlaRuntimeError: INVALID_ARGUMENT: Executable(jit_embedding_pipeline_step_fn) expected parameter 2482 of size 5242880 (bf16[16,1280,40]{2,1,0:T(8,128)(2,1)}) but got buffer with incompatible size 1638400 (bf16[16,1280,40]{1,2,0:T(8,128)(2,1)}): while running replica 0 and partition 0 of a replicated computation (other replicas may have failed as well).
XLA バックエンド: TPU
概要
エラー メッセージには、想定されるサイズと実際のサイズ、テンソルの形状とレイアウトの両方が示されます。これらのエラーは、2 つのテンソルの形状が同じでも、物理レイアウト（ハードウェア上でデータがタイル状に配置される方法）が異なる場合、メモリ内のサイズが異なるために発生することがあります。
これらのエラーの主な原因は次のとおりです。
- チェックポイントと XLA 構成の不一致 - モデルがトレーニングされ、チェックポイントが保存されます。そのチェックポイントの重みの物理レイアウトは、その時点の正確な XLA バージョンと構成（XLA フラグなど）によって決まります。後で、このチェックポイントは構成が変更された別の環境に読み込まれます。新しいフラグ、デフォルト値の変更、モデル/XLA コードの変更により、実行時に重みの異なる物理レイアウトが想定されることがあります。チェックポイントの古いバッファが新しいコンパイル済み XLA プログラムに渡されると、ランタイムがエラーをスローします。
- ハードウェア/トポロジ固有のレイアウト - XLA コンパイラは、さまざまなハードウェアでパフォーマンスを最適化するために、テンソルのさまざまな物理レイアウトを自由に選択できます。v4 TPU に最適なレイアウトは、v5 TPU と異なる場合があります。同じチップの異なる Pod スライス（4x4x4 と 4x8）。このエラーは、モデルが 1 つのトポロジのレイアウトに関する想定に基づいてコンパイルされたものの、実行時に別のトポロジでスケジュールされた場合、または特定のハードウェアのコンパイラのレイアウト ロジックにバグがある場合に発生します。
デバッグ
- モデルのエクスポートとチェックポイントからの再実行の間で構成の一貫性を確保します。
- レイアウトに影響する変更が加えられていないことが確実な場合を除き、新しいコードで古いチェックポイントを使用しないでください。
- Saved Model を再エクスポートする: チェックポイントと構成の不一致が疑われる場合は、推論またはファインチューニングに使用しているものとまったく同じ（現在の）コードベースと構成を使用して、保存済みモデルを再エクスポートするのが最も確実な解決策です。
- 2 回の実行間で構成の変更（XLA フラグなど）がないか確認します。
- ハードウェア/トポロジ固有のレイアウト:
- ハードウェアまたはトポロジを切り替える場合は、ハードウェア バージョンとトポロジの不一致を確認します。