カテゴリ: コンパイル時間: スコープ付き Vmem OOM
このエラーは、プログラムに必要なスコープ付きベクトル メモリ（Vmem）が割り当てられた量よりも多いことを示します。
エラー メッセージの例:
RESOURCE_EXHAUSTED: Ran out of memory in memory space vmem while allocating on stack for %my-custom-kernel = bf16[2048,4096]{1,0:T(8,128)(2,1)} custom-call(...) ...
XLA バックエンド: TPU
概要
TPU には、TensorCore（TC）でのみ使用されるローカル スクラッチパッド メモリであるベクトル メモリ（VMEM）があります。コンパイラは、さまざまなタイプのアロケーションの Vmem を管理します。
- 命令スコープの割り当て: 単一の HLO 命令の実行中に Vmem に一時的に保存されるストレージ。これには、オペランド スパン バッファ（ダブル バッファリングなど）とレジスタ スピルが含まれます。
- プログラム スコープの割り当て: 単一の HLO 命令のスコープを超えて存在する割り当て。これらは通常、HLO 命令の入力または出力である HLO の一時変数と中間結果です。
コンパイル時間スコープの Vmem OOM は、命令スコープの割り当てがその命令の割り当て上限を超えた場合に発生します。この上限は制御されます
- --xla_tpu_scoped_vmem_limit_kib フラグを使用して、プログラム全体でグローバルに設定します。
- vmem_limit_bytes パラメータを使用して、カスタム カーネルごとに設定します。
通常、これらのエラーは内部コンパイラのバグ、または割り当て上限を超えるカスタム カーネルが原因で発生します。
デバッグ
エラー メッセージを慎重に分析して、エラーがカスタム カーネルと標準 HLO のどちらに起因するかを特定します。カスタム カーネルによるエラーには、次のシグネチャが必要です。
Ran out of memory in memory space vmem while allocating on stack for %my-custom-call = <output-shape> custom-call(<params>), custom_call_target="tpu_custom_call" ...
- カスタム カーネル スコープの Vmem OOM: エラーがカスタム カーネルを指している場合 → カーネルを再調整するに移動します。
- 非カーネル Vmem の問題: カスタム カーネル以外の op が原因で Vmem OOM が発生した場合は、コンパイラの内部バグである可能性があります。HLO ダンプを使用して XLA のバグを報告してください。
カーネルを再調整する
エラーがカスタム カーネルに起因する場合は、次の手法を使用してカーネルのメモリ要件を減らします。
- ブロックサイズを調整する: カーネル構成のブロックサイズ（タイルサイズ）を小さくして、Scoped Vmem の使用量を減らします。
- カーネルごとのスコープ付き Vmem 上限を設定する: vmem_limit_bytes パラメータを使用して、特定のカーネルに必要なメモリ量を明示的にリクエストします。
- メモリのカラーリングを変更: pallas.tpu.with_memory_space_constraint を使用して、カーネルの入出力を VMEM に明示的に色付け / 制限します。ただし、Vmem に入力 / 出力を過度に色付けすると、VMEM 全体で OOM が発生する可能性があるため、注意してください。
- カーネル固有の再チューニングが難しい場合や、問題が多くのカーネルに影響する場合は、フラグ --xla_tpu_scoped_vmem_limit_kib を使用してグローバル Vmem 上限を調整できます。