カテゴリ: コンパイル時: SparseCore の割り当てエラー
このエラーは、現在の SparseCore プログラムに必要な指定のメモリ空間に、XLA: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. 内部メモリのエラー（メモリ空間 ID: 0、201、202、204）
Smem、TileSpmem、Spmem、またはSync Flags での割り当てエラーは通常、コンパイラのバグまたは割り当て戦略の制限が原因で発生します。 この場合、コンパイラはすべてのメモリ 要件を考慮できません。
おすすめの対処法:
- 失敗した XLA オペレーションを分離する: 失敗の原因となっている特定の SC HLO または Mosaic カーネルを特定するには、コンパイラの中間表現を生成します。
- SparseCore MLIR をダンプする: フラグ
--xla_sc_dump_mlir_to=/path/to/dumpを設定します。これにより、SparseCore プログラムの MLIR が生成され、エラー メッセージに一致する割り当てサイズを確認できます。 - Mosaic LLO をダンプする: カスタム カーネルの場合は、
--xla_mosaic_dump_to=/path/to/dumpを使用して、Mosaic によって出力されたすべての Low Level Optimizer（LLO）プログラムを検査します。
- SparseCore MLIR をダンプする: フラグ
- スクラッチ サイズを減らす（Pallas ユーザー）: Mosaic カーネル内でエラーが発生した場合は、
scratch_shapes構成を確認します。pltpu.SMEMリクエストが、特定の TPU 世代のハードウェア仕様に収まるようにします。 - コレクティブ オフロードを無効にする: エラーが SC オフロード コレクティブ オペレーションから発生している場合は、SC オフロード機能を無効にしてみてください。
--xla_tpu_enable_sparse_core_collective_offload_all_gather=false--xla_tpu_enable_sparse_core_collective_offload_all_reduce=false
- バグを報告する: 上記の手順で問題が解決しない場合は、コンパイラのバグである可能性があります。バグレポートを提出してください。