エラーコード: E3001

カテゴリ: CompileTime: SparseCore No Viable Logical Replica Count

このエラーは、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_rowmax_unique_nz_per_row に基づいてメモリを割り当てます。これらの値が控えめに推定されている場合(つまり、実際のデータに必要な値よりもはるかに高い値が設定されている場合)、コンパイラは不要なスペースを予約するため、このエラーが発生します。これらのパラメータが、データセットの実際の ID 分布を正確に反映していることを確認します。

フィードバック主導型最適化(FDO)を適用して、これらのパラメータの最適な値を決定することを検討できます。

2. バッチサイズを小さくする

sample_count はグローバル バッチサイズから直接導出されます。バッチサイズを減らすと、各 SparseCore がステップごとに処理する必要があるデータの量が減り、必要なスクラッチパッド バッファのサイズが小さくなります。