カテゴリ: コンパイル時間: SparseCore の割り当てエラー
このエラーは、XLA:SparseCore コンパイラが、現在の SparseCore プログラムに必要な指定されたメモリ空間に連続したメモリブロックを割り当てることができない場合に発生します。
エラー メッセージの例:
INTERNAL:Failed to run pass pipeline. Hlo-Op: result.1:279:1: error: 'memref.alloca' op current allocation offset upper bound (140704 words) exceeds the legitimate user allocatable offset upper bound (131071 words) in memory space 201 when allocating 23440 words. result.1:279:1: note: see current operation: %232 = "memref.alloca"() <{operandSegmentSizes = array<i32: 0, 0>}> : () -> memref<23440xf32, 201>
XLA バックエンド: TPU
概要
SparseCore（SC）は、スパース ワークロード専用のプロセッサです。特定のメモリ階層に依存して、データ移動を効率的に管理します。XLA コンパイラは、ハードウェアの上限とユーザー定義の形状に基づいてバッファのサイズを静的に設定し、割り当てようとします。このエラーは、この割り当てフェーズでメモリ不足（OOM）状態が発生したことを示します。
通常、エラー メッセージにはメモリ空間 ID が指定されます。一般的なメモリ空間とその整数エンコードは次のとおりです。
| メモリ スペース ID | 名前 | 説明 |
|---|---|---|
| 0 | Smem | ローカル スカラー メモリ。スカラー レジスタと制御フローに使用されます。 |
| 201 | TileSpmem | タイル固有のスクラッチパッド メモリ。特定の SC タイルで使用可能な高速ローカル SRAM。 |
| 202 | Spmem | 共有スクラッチパッド メモリ。HBM レイテンシを隠すために、データを（入力、出力、中間）段階的に処理するために使用されます。 |
| 203 | HBM | 高帯域幅メモリ。テーブル、ヒープ、スタックの埋め込みに使用される大きな共有メモリ。 |
| 204 | 同期フラグ | 調整に使用される同期プリミティブ。 |
SC とそのメモリ階層の詳細については、SparseCore のドキュメントをご覧ください。
デバッグ
解決策は、割り当てに失敗したメモリ空間によって異なります。
シナリオ 1. HBM 割り当てエラー
メモリ スペース ID: 203
このエラーは、SparseCore プログラムによってリクエストされた単一の一時割り当てが、使用可能な HBM に収まるほど大きくない場合に発生します。標準のエンベディング ワークロードと SC オフロード コレクティブでは、コアあたりのパーティションが非常に大きい場合や、シャーディングの仕様が正しくない場合、コンパイラが大規模なバッファをリクエストする可能性があります。
推奨される対応:
- シャーディングを確認する: エンベディング テーブルと SC 入力/出力テンソルが正しくパーティション分割/シャーディングされていることを確認します。単一のコアが過剰なデータを処理している場合、割り当てが失敗する可能性があります。
- 制限を調整する:
max_ids_per_partitionとmax_unique_ids_per_partitionを確認します。これらの値を不必要に高く設定すると、コンパイラは必要以上のメモリを予約します。上限がテーブルにどのように変換されるかをご覧ください。
シナリオ 2. 内部メモリの障害
メモリ Space ID: 0、201、202、204
Smem、TileSpmem、Spmem、または同期フラグでの割り当ての失敗は、通常、コンパイラのバグまたは割り当て戦略の制限が原因で発生します。コンパイラがすべてのメモリ要件を考慮できない場合です。
推奨される対応:
- 失敗した XLA オペレーションを分離する: 失敗の原因となっている特定の SC HLO または Mosaic カーネルを特定するには、中間コンパイラ表現を生成します。
- SparseCore MLIR をダンプ: フラグ
--xla_sc_dump_mlir_to=/path/to/dumpを設定します。これにより、SparseCore プログラムの MLIR が生成され、エラー メッセージと一致する割り当てサイズを確認できます。 - Dump Mosaic LLO: カスタム カーネルの場合、
--xla_mosaic_dump_to=/path/to/dumpを使用して Mosaic によって出力されたすべての Low Level Optimizer（LLO）プログラムを検査します。
- SparseCore MLIR をダンプ: フラグ
- スクラッチ サイズを減らす（Pallas ユーザー）: Mosaic カーネル内で障害が発生した場合は、
scratch_shapes構成を確認します。pltpu.SMEMリクエストが特定の TPU 世代のハードウェア仕様に収まるようにします。 - Collective Offload を無効にする: エラーが SC オフロードの Collective オペレーションで発生した場合は、SC オフロード機能を無効にしてみてください。
--xla_tpu_enable_sparse_core_collective_offload_all_gather=false--xla_tpu_enable_sparse_core_collective_offload_all_reduce=false
- バグを報告する: 上記の手順で問題が解決しない場合は、コンパイラのバグである可能性があります。バグレポートを提出してください。