カテゴリ: コンパイル時: SparseCore で有効な論理レプリカ数がありません
このエラーは、XLA:SparseCore コンパイラが、ワークロードを SparseCore のローカル スクラッチパッド メモリ（Tilespmem）に収めることができる有効な論理レプリカ数構成を判別できない場合に発生します。
エラー メッセージの例:
XLA:TPU compile permanent error. Compilation failure: No viable logical replica count for the embedding table with metadata: max_nz_per_row = 141352, max_unique_nz_per_row = 8, feature_width = 8, sample_count = 204800 (last tried split factor for vector splitting = 1, last tried split factor for sample dimension splitting = 1, fixed_size_allocation_bytes = 410880, row_dependent_size_allocation_bytes = 1696224, total_spmem_size_bytes = 524288) ...
XLA バックエンド: TPU
概要
このエラーは、SparseCore のユースケース、特に大規模エンベディング モデル（LEM）に固有のものです。
論理レプリカ数 は、スクラッチパッド割り当ての圧力を管理するために、入力バッチをどのように分割するかを決定する内部コンパイラ パラメータです。コンパイラは、各チャンクに必要な中間バッファが SparseCore の限られたスクラッチパッド メモリ に収まるように、ワークロードを小さなチャンク（レプリカ）に分割しようとします。一般に、論理レプリカ数を増やすと、一度に処理するデータバッチが小さくなるため、割り当ての圧力が軽減されます。
このエラーは、さまざまな分割構成を試しても、必要なバッファが Tilespmem メモリに収まる設定をコンパイラが見つけられなかったことを示します。割り当てサイズは、次の組み合わせによって決まります。
sample_count: 各 SparseCore に割り当てられたエンベディング ルックアップ ID の数（バッチサイズから派生）。feature_width: エンベディング ディメンションのサイズ。max_nz_per_row: すべての SparseCore における一意でないエンベディング ルックアップ ID の最大数。max_unique_nz_per_row: 一意のエンベディング ルックアップ ID の最大数。
デバッグ
このエラーを解決するには、SparseCore スクラッチパッドのメモリ圧力を軽減する必要があります。
1. メタデータの推定を改善する
コンパイラは、max_nz_per_row と max_unique_nz_per_row に基づいてメモリを割り当てます。これらの値が控えめに推定されている場合（つまり、実際のデータに必要な値よりもはるかに高い値に設定されている場合）、コンパイラは不要なスペースを予約し、このエラーが発生します。これらのパラメータが、データセットの実際の ID 分布を正確に反映していることを確認してください。
フィードバック主導型最適化（FDO） を適用して、これらのパラメータの最適な値を決定することを検討してください。
2. バッチサイズを小さくする
sample_count は、グローバル バッチサイズから直接派生します。バッチサイズを小さくすると、各 SparseCore が 1 ステップで処理する必要があるデータの量が減り、必要なスクラッチパッド バッファのサイズが小さくなります。