大規模エンベディングモデル（LEM）用の SparseCore の詳細

SparseCore は、不規則でスパースなメモリアクセスと計算を伴うワークロード（特に高帯域幅メモリ（HBM）に保存された大規模なデータセット）の高性能アクセラレーション用に設計された特殊なタイルプロセッサです。エンベディングルックアップなどのタスクに優れていますが、さまざまな動的でスパースなワークロードの高速化にも対応できます。

1. SparseCore の概要

主なアーキテクチャ機能:

タイルアーキテクチャ: 複数のコンピューティングタイル（各タイルは独自のローカルメモリと処理ユニットを備えた完全なデータフローユニット）で構成され、並列処理が可能です。
動的実行: スパースデータに不可欠な、データ依存の制御フローとメモリアクセスをネイティブにサポートします。
ベクトル処理: 小さなベクトルタスク（ハードウェアバージョンに応じて 8 要素または 16 要素）を利用して、効率的な計算を行います。
一元管理: 単一の SparseCore シーケンサーがすべてのタイルでタスクをオーケストレートし、オペレーションの同期を保証します。
データ要約のサポート: 並べ替え、フィルタリング、プレフィックス合計などのタスクに役立つ特殊なクロスレーンオペレーションが含まれています。
メモリ階層: 大規模なデータセットの保存に HBM を戦略的に活用し、頻繁にアクセスされるデータのステージングにローカルスクラッチパッドメモリ（SPMEM）を活用することで、HBM レイテンシを大幅に短縮します。

仕様の概要:

属性	TPU v4	TPU v5p	Trillium
SparseCores/チップ	4	4	2
Tiles/SparseCore	16	16	16
SIMD 幅	8	8	8（F32）、16（BF16）
HBM 容量	32 GiB	96 GiB	32 GiB

2. SparseCore ホストの前処理

SparseCore のパフォーマンスには効果的なデータ準備が不可欠であり、ホストの前処理が重要な役割を果たします。これには、次のような主要な機能が含まれます。

データ変換:
- 未加工の入力データに必要な変換を適用します。
- ID 変換を管理します。これは、特徴やテーブルのスタッキングを扱う場合に特に重要です。
- 入力データを、次のセクションで説明する Coordinate（COO）スパース形式に変換します。
- チップで使用可能なさまざまな SparseCore に効率的に分散できるようにデータをパーティショニングします。
上限の検証:
- 入力データの特性（ID の数など）が、max_ids_per_partition や max_unique_ids_per_partition などの SparseCore の事前定義された運用上の上限に準拠していることを確認します。
- 入力データがこれらの上限を超えると、ホストのプリプロセスレイヤは、制約内に収まる小さなミニバッチにデータをセグメント化しようとします。
データ転送:
- 処理および検証されたデータを TPU の高帯域幅メモリ（HBM）に効率的にコピーし、SparseCore の実行の準備を整えます。

テーブルのスタッキングについて:

テーブルスタッキングは、複数のエンベディングテーブルを論理的に結合してエンベディングルックアップの効率を高める重要な最適化手法です。このプロセスは通常、基盤となる ML フレームワークによって自動的に処理されます。

特徴量のスタッキング: 複数の異なる特徴量が同じ基盤となる埋め込みテーブルを共有している場合に発生します。一般的な例としては、さまざまなコンテキストの郵便番号などのさまざまなカテゴリ特徴に単一のエンベディング辞書を使用することがあります。
テーブルのスタッキング: このシナリオでは、複数の個別のエンベディングテーブルが積み重ねられます。同じエンベディングディメンションとオプティマイザー構成を共有するテーブルは、多くの場合グループ化されます。

テーブルスタッキングの主な利点は、これらのスタックされたテーブルに対するオペレーションの有効なバッチサイズを大きくできることです。これにより、計算オーバーヘッドが減少し、チップ間通信（ICI）レイテンシを隠すのに効果的です。最適なパフォーマンスを得るには、適度な数のスタックテーブル（通常は 5 ～ 100 個）をおすすめします。

3. COO テンソルへの変換

SparseCore でデータを処理する前に、通常は座標（COO）スパーステンソル形式に変換されます。COO 形式は、通常 3 つの配列を使用してスパース行列を効率的に表現する方法です。

row_ids: ゼロ以外の各要素の行インデックスを含む配列。バッチ処理のコンテキストでは、これは多くの場合、バッチディメンションに対応します。
col_ids: ゼロ以外の各要素の列インデックスを含む配列。エンベディングの場合、これらは特徴値または ID 値であることがよくあります。
values（省略可）: 対応する（row、col）座標にある非ゼロ要素の実際の値を保持する配列。ID 数に関連する制限計算（後述）では、これらの値（ゲイン）は考慮されないことがよくあります。

説明のための例:

ID のバッチを表す入力スパース行列を考えてみましょう。

[
    [id_A],                 // Sample 0
    [id_A, id_B, id_C],     // Sample 1
    [id_B, id_B, id_D],     // Sample 2 (note duplicate id_B)
]

COO 形式に変換した後（および、場合によっては同じサンプル内の ID を重複除去した後）:

row_ids = [0, 1, 1, 1, 2, 2]
col_ids = [id_A, id_A, id_B, id_C, id_B, id_D]

この変換は、SparseCore がワークを処理して分散する仕組みの基本です。特に col_ids は、ID が属する特定の SparseCore パーティションを特定し、効率的なシャーディングとルックアップを可能にするために不可欠です。

4. SparsecoreConfig: 高レベル API

フレームワーク固有のエンベディング API:

JAX: https://github.com/jax-ml/jax-tpu-embedding
TensorFlow: https://www.tensorflow.org/recommenders/api_docs/python/tfrs/layers/embedding/TPUEmbedding
Keras: https://keras.io/keras_rs/api/embedding_layers/distributed_embedding

SparsecoreConfig（または XLA フラグなどの同等のメカニズム）は、さまざまな SparseCore の動作を制御するための高レベルのインターフェースとして機能します。これらのパラメータを十分に理解することは、パフォーマンスを効果的にチューニングし、モデルの正しい動作を保証するために不可欠です。

disable_table_stacking: bool = False
- 説明: このフラグは、テーブルの自動スタッキングを防止して、オーバーヘッドの増加とチップ間インターコネクト（ICI）レイテンシを隠す機能の低下により、パフォーマンスが低下する可能性のあるフレームワークによるテーブルのスタッキングを防止するかどうかを制御します。
- デフォルト: False（フレームワークがサポートしている場合、テーブルのスタックはデフォルトで一般的に有効になっていることを意味します）。
max_ids_per_chip_per_sample: int = 64
- 説明: このパラメータは、入力バッチ内の 1 つのサンプルから 1 つのチップが処理できるエンベディング ID の合計数について、すべてのテーブルで集計されたグローバルな上限を設定します。これは、より詳細なテーブルごとまたはパーティションごとの上限が考慮される前に、チップレベルでリソースを管理するメカニズムです。この値の微調整は通常、特定のモデルの特性とシステム全体の容量によって異なります。
- デフォルト: 64。
max_ids_per_table: Optional[Dict[str, int]] = None
- 説明: このパラメータは、すべての SparseCore にわたるすべてのパーティションを考慮して、各論理テーブルで処理できるエンベディング ID の最大数（重複を含む）を指定します。これは max_ids_per_partition よりも広い上限です。テーブル T が P 個のパーティションに分割されている場合、この上限はすべての P パーティションに転送される ID の合計に適用されます。通常、これは max_ids_per_partition_per_sample とバッチ全体のサイズに関連しています。
- 設定: 通常は、max_ids_per_partition が定義されている制限ファイル（xla_sparse_core_max_ids_file フラグなど）を使用して構成します。このテーブルレベルのコンセプトは、パーティションレベルの上限（max_ids と max_uniques）を設定する方法です。
- デフォルト: None（明示的に指定されていない場合、値はパーティションごとの上限または他の構成から推測されることがあります）。
max_unique_ids_per_table: Optional[Dict[str, int]] = None
- 説明: max_ids_per_table と同様ですが、このパラメータは各論理テーブルの一意の ID の最大数を指定します。これは、一意の ID 処理と後続のベクトルオペレーションで使用されるオンデバイスバッファのサイズを適切に設定するための重要な設定です。
- 設定: 通常は、上限ファイルで定義されるか、max_unique_ids_per_partition_per_sample から派生します。
- デフォルト: None。
allow_id_dropping: bool = False
- 説明: このブール値フラグは、入力データで検出された ID の数（観測された上限）がコンパイル時に設定された上限（max_ids_per_partition など）を超えた場合に ID を削除するかどうかを制御します。
  - True の場合: 上限を超過する ID は通知なしで削除されます。通常、パーティション内の ID はソートされた順序で処理され、指定されたミニバッチの実行カウントが上限を超える ID は破棄されます。これにより、プログラムの実行は継続されますが、モデルの精度に悪影響が及ぶ可能性があります。
  - False の場合: エラーがトリガーされ、観測された上限がコンパイルされた上限を超えると、プロセスが終了する可能性があります。このアプローチでは、すべてのデータが処理されますが、上限をより保守的に構成する必要があります。
- デフォルト: False（データがサイレントにドロップされるのではなく、オーバーフロー時にエラーが発生します）。
initialize_tables_on_host: bool = True
- 説明: このフラグは、エンベディングテーブルがホスト CPU で初期化されてから、TPU の高帯域幅メモリ（HBM）に転送されるかどうかを決定します。通常、テーブルはホストで初期化されます。この命名規則に従って、True に設定します。False に設定されている場合、デバイス上の初期化メカニズムが示唆されます。これは、パフォーマンスに異なる影響を与えたり、特定の初期化の前提条件を必要としたりする可能性があります。
enable_fast_table_initialization: bool = False
- 説明: TPU でテーブルを直接初期化します。これにより、モデルの起動時間を短縮できます。

5. パフォーマンス向上のためのパイプライン処理

パイプライン処理は、TensorCore（TC）と SparseCore（SC）でのオペレーションの同時実行を可能にするパフォーマンス最適化手法です。これらの計算を重複させることで、全体的なスループットを大幅に改善できます。

メカニズム: スパースエンベディングルックアップ（SC によって処理）と密レイヤ計算（TC によって処理）を含む標準のトレーニングステップでは、パイプライン処理により、TC が同じステップ i の別の部分、または i-1 や i+1 などの隣接するステップの部分を同時に処理している間に、SC がステップ i の一部（フォワードパスやバックワードパスなど）を処理できます。
勾配への影響: SparseCore が「古い」勾配で動作する可能性があります。たとえば、ステップ i のバックプロパゲーションフェーズで計算されたグラデーションは、ステップ i+2 まで完全に更新されず、SC に表示されない場合があります。
パフォーマンスと数値のトレードオフ: この重複実行により、大幅な高速化が実現し、デバイスのステップ時間が 2 倍に短縮される可能性があります。ただし、古いグラデーションの使用によって生じる数値（embedding_weights）のわずかな変化が、モデルの収束動作や最終的な精度に影響する可能性があります。このトレードオフの許容度はモデルに大きく依存し、多くの場合、実証的な検証が必要になります。
制御フラグ: パイプラインは tf_xla_disable_full_embedding_pipelining で制御できます。このフラグを true に設定すると、完全なパイプライン処理（TensorCore と SparseCore の計算の重複）が無効になります。一方、false に設定すると（または、フラグのセマンティクスが false のときに有効になることを意味する場合）、有効になります。

概念的なパイプラインフロー:

パイプライン処理なし（簡略化された順次フロー）:

Loop: SC/F_i -> TC/F_i -> TC/B_i -> SC/B_i
パイプライン処理あり（簡略化されたオーバーラップフロー）:
```
Time ->
Step i:   SC/F_i | TC/F_i | TC/B_i | SC/B_i
Step i+1:          SC/F_i+1| TC/F_i+1| TC/B_i+1| SC/B_i+1
```
注: ハードウェアとコンパイラに実装される実際のパイプラインステージは、より複雑になる可能性があります。多くの場合、データ依存関係を管理して正確性を確保するために、プリループ、メイン実行ループ、ポストループが含まれます。

6. XLA の役割

XLA（Accelerated Linear Algebra）は、通常は TensorFlow などのフレームワークから、高度に最適化された TPU 向けのマシンコードに、高レベルの計算グラフを変換するドメイン固有のコンパイラです。これには、SparseCore 宛てのオペレーションの命令の生成が含まれます。

SparseCore コンテキストの主な関数:

スパース演算のコンパイル: XLA は、エンベディングルックアップ演算（SparseDenseMatmulOp など）やその他のスパース計算を低レベルの実行可能な SparseCore プログラムにコンパイルします。
制限の統合: 構成された運用上の制限（max_ids_per_partition、max_unique_ids_per_partition など。多くの場合、xla_sparse_core_max_ids_file などのフラグで指定された制限ファイルを通じて提供されます）を利用して、オンデバイスのメモリバッファのサイズを静的に決定し、割り当てます（特に SPMEM 内）。
ターゲット最適化: XLA は、SparseCore アーキテクチャ用に特別に設計された一連の最適化を実行します。これには、効率を最大化するための命令スケジューリング、メモリレイアウト変換、オペレーションの融合などが含まれます。
フラグを使用した制御: SparseCore の動作、チューニングパラメータ、最適化戦略の多くの側面が公開され、XLA フラグ（制限の見積もり用の xla_sparse_core_estimate_max_ids やデバッグ用の xla_sc_detect_nan など）を通じて制御されます。

オープンソースステータス:

現在、Sparsecore 実装は内部専用で、libtpu.so を使用して提供されています。

エラー報告と診断:

SparseCore の構成またはリソース制約に関連するコンパイルエラーは、XLA:TPU コンパイル時エラーとして現れることがよくあります。これらのエラーメッセージは、使用可能な SPMEM に対して上限が高すぎる、サポートされていない構成が使用されているなどの問題に関する貴重な分析情報を提供します。

7. 制限が SparseCore のテーブルにどのように反映されるか

SparseCore では、「上限」は、主に、使用可能な SparseCore にシャーディング（分散）される各テーブルのパーティションごとの 2 つの設定を指す基本的な構成パラメータです。

max_ids_per_partition: 単一の SparseCore が単一の計算ステップ内で特定のテーブルの特定のパーティションに送信または処理することが想定される ID の合計数（重複を含む）の最大数を定義します。
max_unique_ids_per_partition: これは、単一の SparseCore が送信または処理することが想定される一意の ID の最大数を定義します。

物理テーブルレイアウトと処理への変換:

テーブルシャーディング戦略: エンベディングテーブルは通常、システム内のすべての SparseCore で「mod-sharded」されます。つまり、各 SparseCore は、各テーブルの語彙（行）の個別のサブセットを担当します。通常、ID j は k = j % num_total_sparse_cores のような式に基づいて SparseCore_k に割り当てられます。
「パーティション」の定義: このコンテキストでは、「パーティション」とは、単一の SparseCore がルックアップを処理するエンベディングテーブルの特定のセグメントを指します。
SPMEM バッファ割り当て: これらの上限は、XLA コンパイラがデバイス上のスクラッチパッドメモリ（SPMEM）内のバッファのサイズを静的に設定して割り当てるために使用されます。バッファは、特定のパーティションの ID に関連する必要なすべてのデータ（指定された max_ids と max_unique_ids の上限まで）を処理のために SPMEM に読み込めるようにサイズ設定されます。これは、パーティション内の重複 ID の削減（圧縮スパース行（CSR）表現の作成時など）など、要素ごとの計算以外の場合に特に重要です。この場合、そのパーティションの ID に関連するデータセット全体を高速メモリですぐに利用できるようにする必要があります。
コンパイルされた上限と観測された上限:
- 観測された上限: 処理される入力データに基づいて、実行時に各パーティションで検出された実際の ID の数。
- 観測された上限がコンパイルされた上限を超えると、ID のドロップ（allow_id_dropping が有効になっている場合）またはエラーが発生する可能性があります。
制限の計算: 適切な制限を決定するプロセスでは、入力データの分布を慎重に分析する必要があります。任意のテーブル（T1 とします。これは、より大きなスタックテーブル T の一部である可能性があります）の場合:
1. 入力バッチ（形状 [BatchSize, MaxSequenceLength] の 2D SparseTensor など）は、最初に使用可能な SparseCore に分割されます。たとえば、TensorCore が 2 つの SparseCore とペアになっている場合、各 SparseCore は [BatchSize/2, MaxSequenceLength] の形状のサブバッチを受け取ることがあります。
2. このサブバッチは COO 形式に変換され、row_ids と col_ids が生成されます。
3. 同じサンプル内の重複する ID（同じ row_id と col_id を持つエントリ）は削除されます。
4. 残りの一意の col_id（サンプル内）ごとに、この ID を担当するターゲット SparseCore は、mod-sharding ルール target_sc_id = col_id % num_total_sparse_cores を使用して決定されます。
5. 各 target_sc_id に送信される ID の総数（ids_per_sparse_core[target_sc_id]++）と一意の ID の数（unique_ids_per_sparse_core[target_sc_id]++。特定の target_sc_id の一意性を確認した後）のカウントが維持されます。
6. テーブル T1 の max_ids_per_partition は max(ids_per_sparse_core_array) に設定されます。
7. 同様に、テーブル T1 の max_unique_ids_per_partition は max(unique_ids_per_sparse_core_array) に設定されます。
8. テーブル T1 が積み重ねテーブルのコンポーネントである場合、すべての構成テーブルの統計情報を合計する前に、回転やシフトなどの追加の変換が ID 分布に適用されることがあります。これにより、チップ間の負荷を分散できます。

これらの上限を正しく設定することはバランスを取る行為です。上限を低くすると、パフォーマンスが向上する可能性があります（ステップごとに処理する必要があるデータが減り、SPMEM の負荷が軽減されるため）。ただし、上限を低くしすぎると、ミニバッチ処理が過剰になったり、望ましくない ID の削除が発生したりする可能性があります。

8. 各 SparseCore の通信方法

SparseCore 通信（特にエンベディングルックアップの ID リストの処理のコンテキスト）は、いくつかの連携メカニズムに依存しています。

Mod シャーディングと暗黙的ルーティング:
- エンベディングテーブルは、システム内のすべての SparseCore にわたって mod シャーディングされます。
- ホストが入力データのバッチ（後で col_ids を含む COO 形式に前処理される）を提供すると、col_id 値を使用して、その特定の ID（target_sc_id = col_id % num_total_sparse_cores）を担当する SparseCore が決定されます。
- 各 SparseCore は、割り当てられた語彙パーティションにマッピングされる ID のサブセットのみを効果的に受信して処理します。ホストの前処理ステージは、各 SparseCore が関連する ID を簡単に識別して操作できるようにデータを準備するうえで重要です。
ホスト別のデータ分布:
- ホストのプリプロセスロジックは、入力バッチ全体をパーティショニングし、row_ids と col_ids の関連部分（関連する特徴や重みがある場合はそれらも含む）を、各 SparseCore が直接アクセスできるメモリ（HBM）または SparseCore が必要なデータを取得する共有 HBM に分散します。
Intra-SparseCore 処理:
- SparseCore は、特定のテーブルパーティションに指定された ID のセットを受け取ると、これらの ID の重複除去や対応するエンベディングベクトルの収集などのオペレーションを実行します。これらは主に、SparseCore 独自のタイル内で実行され、ローカル SPMEM を利用するローカルコンピューティングです。
Inter-SparseCore 通信（All-to-All）:
- 初期処理フェーズ（エンベディングルックアップなど）の後、「all-to-all」通信パターンを使用して、SparseCore 全体で結果を結合または再配布できます（たとえば、すべての元のサンプル位置に対応する入力を想定する TensorCore レイヤにアクティベーションをフィードする前など）。これは、元の入力バッチが並列処理用に分散された場合に、アクティベーションの完全なセットを再構築するために不可欠です。
TensorCore との通信:
- SparseCore は TensorCore と通信して、エンベディングアクティベーション（フォワードパス中）を送信し、グラデーション（バックワードパス中）を受信します。このインタラクションは XLA コンパイルプログラムによってオーケストレートされ、多くの場合、HBM が仲介バッファとして使用されます。パイプライン戦略（前述）は、この SC-TC 通信のタイミングと同期に大きく影響します。

基本的に、適切な SparseCore への ID の最初の「分散」は、シャーディングスキームとホストの事前処理ステップによって処理されます。後続の通信では、SparseCore がローカルデータで動作します。データが TensorCore によるさらなる処理の前に SparseCore 間でグローバルに交換または並べ替えられる必要がある場合は、all-to-all などの集団通信オペレーションが続く可能性があります。

9. SparseCore のメモリ管理

各 SparseCore は、複数の異なるタイプのメモリを効率的に管理して計算を実行します。

スクラッチパッドメモリ（SPMEM）:
- Nature: 比較的小さいながらも非常に高速なローカル SRAM で、各 SparseCore 専用です。SPMEM はキャッシュではありません。その使用は XLA コンパイラによって明示的に管理およびオーケストレーションされます。
- 目的: SPMEM は「機会的データステージング」に使用されます。これには、継続的な SC 計算に必要な入力、出力、中間結果が含まれます。SPMEM にデータをステージングすると、通常 HBM へのアクセスに関連する高レイテンシが大幅に短縮されます。
- サイズ設定: 「制限」のセクションで説明したように、SPMEM バッファはコンパイル時に静的にサイズ設定されます。このサイズ設定は、max_ids_per_partition や max_unique_ids_per_partition などのパラメータに基づいています。この静的割り当てにより、テーブルパーティションに対する特定のオペレーション（CSR 削減など）で、そのパーティションの ID に必要なすべてのデータ（定義された上限まで）が SPMEM に収まるようになります。
- コンパイラの最適化: XLA コンパイラは、高度な最適化を組み込んで、HBM レイテンシを効果的に隠蔽し、パフォーマンスを最大化するために、SPMEM にステージングする必要があるデータの量と、どの特定のデータ要素をステージングする必要があるかを正確に判断します。
- 動的割り当ての制約: SparseCore コンパイラは現在、動的スクラッチパッド割り当てをサポートしていません。これは、上限を慎重に構成することで、静的サイジングが非常に重要であることを示しています。
高帯域幅メモリ（HBM）:
- Nature: すべての SparseCore、TensorCore、ホストシステムからアクセス可能な大規模な共有メモリリソース。プライマリエンベディングテーブルは HBM に保存されます。
- スタック使用量: SparseCore オペレーションでは、中間結果を HBM に一時的に保存することがよくあります。これは、限られた SPMEM に収まらないか、処理パイプラインの大きなステージ間で渡す必要があるためです。フォワードパスとバックワードパスの両方での HBM スタックの使用量は、次のように見積もることができます。
  - フォワードパス HBM スタック（単一テーブル）≈（2 * feature_width + 1）* max_unique_nz_per_row * logical_replica_count * 4 バイト
  - Backward Pass HBM Stack（単一テーブル）≈ 3 * feature_width * max_unique_nz_per_row * logical_replica_count * 4 バイト
- ヒープ使用量: HBM は、ホストによって管理されるヒープにも対応しています。ヒープには、密なレイヤの重み、モデルで使用される定数、プリフェッチされた入力データなどのデータが保存されます。ヒープ使用量は、ホストがデータをプリフェッチするステップ数（maximum_parallel_iterations フラグで制御）とともに増加する傾向があります。プリフェッチを増やすと、ホストからデバイスへの転送とデバイスの計算をオーバーラップさせることでパフォーマンスを向上させることができますが、HBM の消費量も増えます。
- HBM 最適化のためのシリアル化: フラグ xla_sc_num_serialized_tables_to_optimize_hbm は、特定の時点で HBM スタックメモリに「ライブ」で保持されるテーブルの数を制御するメカニズムを提供します。この数を増やすと、より多くのテーブルの処理が効果的にシリアル化されるため、HBM スタックのピーク使用量を削減できますが、並列処理の減少によりパフォーマンスが低下する可能性があります。
ベクトルメモリ（VMEM）:
- VMEM は、TC（TensorCore）でのみ使用されるローカルスクラッチパッドメモリです。VMEM は SparseCore によって直接管理されませんが、SC が主に TensorCore を介してやり取りするメモリエコシステムの不可欠な部分です。

全体的なメモリ管理戦略:

SparseCore のコアメモリ管理戦略は、SparseCore タイルでアクティブに処理されている「ホット」データに高速な小容量の SPMEM を使用し、低速な HBM へのアクセスを最小限に抑えることを中心としています。構成された上限は、SPMEM がオーバーフローしないようにするための主要なメカニズムです。HBM は、SPMEM 容量を超えるか、異なる処理ユニットまたはパイプラインステージ間で共有する必要がある大規模なエンベディングテーブルと一時データの保存に使用されます。XLA コンパイラは、これらのアーキテクチャ原則とユーザーが構成した上限に基づいて、すべてのデータ移動とバッファ割り当てを調整します。

10. パフォーマンスとメモリのボトルネック

SparseCore で最適なパフォーマンスを実現するには、潜在的なボトルネックとその対処方法を明確に理解する必要があります。これらは、ホスト、SparseCore 自体、または TensorCore とのやり取りで発生する可能性があります。

一般的なパフォーマンスのボトルネック:

ホストのボトルネック:
- 問題: ホスト CPU がデータを前処理して TPU に十分な速度でフィードできないため、SparseCore と TensorCore の使用率が低下する可能性があります。これはパフォーマンスの制限要因となることがよくあります。
- 軽減策: ホストの CPU 使用率と入力パイプラインの指標をモニタリングします。ホスト側のデータ読み込みと前処理ルーティンを最適化します（COO 変換のヒントを参照）。maximum_parallel_iterations フラグを調整して、データプリフェッチをファインチューニングします。
最適でない TC/SC 同期（パイプライン処理の欠如）:
- 問題: TensorCore と SparseCore 間のパイプライン処理が無効になっているか、効率的に動作していない場合、一方のユニットが他方のユニットを待つ時間が長くなり、システム全体のスループットが低下する可能性があります。
- 軽減策: パイプライン処理が有効になっていることを確認します（tf_xla_disable_full_embedding_pipelining = false など）。
上限によるボトルネック:
- 問題:
  - 上限が低すぎる: 過剰なミニバッチ処理（厳しい上限を満たすために、入力バッチを多数の小さなサブバッチに分割する）がトリガーされる可能性があります。これにより正確性は維持されますが、各ミニバッチで処理オーバーヘッドが発生し、全体的な実行速度が低下する可能性があります。allow_id_dropping が true の場合、上限が低すぎると ID のドロップにつながり、モデルの精度に影響する可能性があります。
  - 上限が高すぎる（ただし適合している）: 上限が高すぎると、ミニバッチ処理が妨げられる可能性があります。また、実際のデータ特性がこれらのピーク値に達することがほとんどない場合、SPMEM の負荷が不必要に増加する可能性があります。また、厳密には必要のない HBM スタックの使用量が増加する可能性もあります。
  - コンパイルの失敗: 構成された上限で、使用可能な物理メモリよりも多くの SPMEM または HBM スタックが必要な場合、コンパイルは失敗します。
- 軽減策: 上限が正しく設定されていることを確認します。
データ分布の偏り:
- 問題: 特定の SparseCore パーティションが、他のパーティションと比較して不均衡に多くの ID を一貫して受け取っている場合（ID の分布が悪いことを示します）、過負荷の SparseCore がパフォーマンスのボトルネックになります。
- 軽減策: ミニバッチ処理中の ID シャッフルは、特に「ホット」なユーザーテーブルを含むスタックテーブルで、この問題を軽減するのに役立ちます。ID の分布を慎重に分析して、テーブルごとに適切な上限値を設定します。
表の積み重ねに関する問題:
- 問題:
  - テーブルのスタックが少なすぎる: ICI レイテンシを効果的に隠蔽したり、処理オーバーヘッドを十分に削減したりするには不十分な可能性があります。
  - テーブルのスタックが多すぎる: 管理が困難になるほど大きな論理テーブルが作成されたり、使用可能なリソースの上限を超えたりする可能性があります。
- 軽減策:
  - スタッキングに最適なテーブル数を確保します。一般的なガイドラインでは、スタックするテーブルの「スイートスポット」は 5 ～ 100 個とされています。
非効率的な数値/量子化:
- 問題: BF16 や量子化整数などの低精度形式で十分な場合（計算速度も向上する場合）に、FP32 の完全な精度を使用すると、パフォーマンスのボトルネックになる可能性があります。
- 緩和策: 低精度のオプションを検討します。ただし、量子化自体にオーバーヘッドがあり、モデルの精度を維持するには量子化パラメータの慎重なチューニングが必要になる場合があります。
HBM 帯域幅の飽和:
- 問題: HBM との間で過剰なデータ移動が発生する。これは、特徴の幅が非常に小さい（パディングのオーバーヘッドが大きい）、メモリアクセスパターンが非効率的である、ルックアップの数が非常に多いことが原因である可能性があり、使用可能な HBM 帯域幅が飽和する可能性があります。
- 軽減策: TPU の数をスケーリングすると、HBM 帯域幅の飽和を軽減できます。

一般的なメモリのボトルネック:

SPMEM オーバーフロー（コンパイルの失敗）:
- 問題: max_ids_per_partition と max_unique_ids_per_partition の設定が高すぎると、XLA コンパイラが十分な SPMEM を割り当てることができず、"Fixed size allocations (...) do not fit in TileSpmem (...)" などのコンパイルエラーが発生する可能性があります。また、(sample_count * feature_width) / kNumTiles（kNumTiles は SC あたりのタイルの数）が大きすぎて、タイル SPMEM 内のステージング収集オペランドに収まらない場合、"Gather operand too large..." などのエラーが発生することがあります。
- 緩和策: バッチサイズを小さくするか、処理に使用するチップの数を増やします。
HBM スタックオーバーフロー（ランタイムまたはコンパイル）:
- 問題: feature_width、max_unique_nz_per_row、logical_replica_count の組み合わせにより、HBM スタックのメモリ要件が使用可能な HBM を超える場合、実行時またはコンパイル時にメモリ不足（OOM）エラーが発生する可能性があります。
- 軽減策: xla_sc_num_serialized_tables_to_optimize_hbm フラグを調整して、テーブルの処理をシリアル化することで HBM スタックの使用量を減らします（通常、パフォーマンスの低下を伴います）。
HBM ヒープの枯渇:
- 問題: 主に、非常に大きな密なレイヤの重み、メモリに保存された多数の定数、過度にアグレッシブな入力プリフェッチ（高い maximum_parallel_iterations）が原因で発生します。
- 軽減策: XProf Memory Viewer などのツールを使用してヒープ使用量をモニタリングします。
パディングのオーバーヘッド:
- 問題: エンベディングテーブルは、特徴ディメンションで 32B 境界に揃うように（8 個の浮動小数点数に相当）パディングされます。そのため、特徴の幅が小さい（たとえば 1 つの浮動小数点）と、パディングのオーバーヘッドが大きくなり（たとえば、割り当てられたバッファスペースの 7/8 がパディング）、HBM が無駄になります。テーブルの語彙ディメンションも、システムの SparseCore の数の倍数になるようにパディングされます。ただし、語彙サイズが十分に大きいテーブルの場合、この影響は通常無視できます。

パフォーマンスとメモリに影響する一般的な要因:

トポロジ: 使用可能なチップの数とその相互接続アーキテクチャ。
バッチサイズ: SparseCore あたりの sample_count に直接影響し、メモリ消費量とコンピューティング負荷に影響します。
データ形式: 最適なパフォーマンスを実現するには、効率的なデバイス上のデータレイアウトを確保することが重要です。

11. SparseCore プロファイルを分析する

パフォーマンスプロファイルの分析は、SparseCore ワークロード内のボトルネックを特定し、最適化の機会を見つけるための重要なステップです。

トレースを取得する:
- XProf などのプロファイリングツールを使用して、モデルのトレーニング中または推論の実行中に詳細な実行トレースをキャプチャします。このトレースは、ホスト、TensorCore、SparseCore で発生するオペレーションのタイムラインを提供します。
トレースビューア（XProf や TensorBoard など）を調べる:
- ホストアクティビティ: ホストのアクティビティを精査します。TPU アクティビティに大きなギャップがあるかどうか。このようなギャップは、ホストがボトルネックになっており、データを十分に速くフィードできていないことを示している可能性があります。入力パイプラインのパフォーマンスを分析します。
- TensorCore（TC）と SparseCore（SC）のアクティビティ:
  - TC と SC の両方の実行タイムラインを確認します。並行して動作している場合は、パイプライン処理が効果的に行われていることを示しています。それとも、一方のユニットが他方を待機してアイドル状態になっている期間が長いのでしょうか？
  - SC と TC の両方で、最も時間がかかるオペレーション（実行時間が最も長いオペレーション）を特定します。
  - ビジュアルトレース出力（多くの場合、TPU:0 SparseCore 1 (pid 1005) のように、さまざまなオペレーションを時系列で表す色付きのブロックが表示されます）は、支配的なオペレーションとアイドル期間を視覚的に特定するうえで非常に役立ちます。
- ステップ時間分析: 全体的なステップ時間を観察し、ホスト処理、SC 計算、TC 計算の間でどのように分散または分割されているかを把握します。
メモリ分析（XProf Memory Viewer）:
- ヒープ使用量: XProf の [Memory Viewer] タブなどのツールを使用して、HBM ヒープ使用量を調べます。これにより、大きなモデルの重み、定数、過度にアグレッシブな入力プリフェッチが過剰な量の HBM を消費しているかどうかを判断できます。--vmodule=best_fit_allocator=1 などのフラグを有効にすると、ヒープ使用量のピーク時のログを取得できる場合があります。
- スタック使用量（間接的）: HBM スタックの直接プロファイリングは複雑になる可能性がありますが、メモリ不足エラーが発生し、ヒープ使用量が妥当な範囲内である場合は、HBM スタックの枯渇（多くの場合、制限や特徴の幅が大きすぎるため）が原因である可能性が高いです。HBM スタックの使用率の計算式は、この見積もりに役立ちます。
特定のパターンを探す:
- ミニバッチ処理: 上限を頻繁に超える場合は、トレースでミニバッチ処理の証拠（グローバルバッチサイズに対して想定されるよりも小さい SC オペレーションの数が多いなど）を確認できます。これは、ログから推測したり、特定のオペレーションの呼び出し回数をモニタリングしたりすることで確認できます。
- ID のドロップ: ID のドロップが有効になっていて発生している場合、システムログにその兆候が示されることがあります。これは、構成された上限が入力データに対して厳しすぎることを示す明確な兆候でもあります。
- コンパイル時間: 再コンパイル時間が長くなると、特にフィードバック指向の最適化（FDO）が有効になっていて、上限が頻繁に調整されている場合は、トレーニング全体の時間に大きなオーバーヘッドが追加される可能性があります。
フラグと構成との相関付け:
- プロファイルで観測された動作を SparseCore 構成（制限ファイルの設定、XLA フラグ）に関連付けます。たとえば、xla_sc_num_serialized_tables_to_optimize_hbm を高い値に設定すると、SC のパフォーマンスは低下しますが、HBM スタックの使用量は減少します。
反復プロセス:
- プロファイリングは、多くの場合、反復的な改善プロセスです。特定の変更（上限の調整、機能の有効化または無効化）を行い、新しいプロファイルをキャプチャしてから、以前のプロファイルと比較して、変更の影響を確認します。

12. 一般的なデバッグフラグ

SparseCore の実行に関連する問題のデバッグを支援するために、いくつかのフラグを有効にできます。これらのチェックを有効にすると、パフォーマンスが低下することが多いため、通常は本番環境の実行では無効にする必要があります。

ID チェック（範囲外）:
- フラグ: xla_sparse_core_enable_id_bound_check = true
- 目的: ホストシステムでチェックを行い、入力データ内のエンベディング ID が、特定のエンベディングテーブルに定義された有効な語彙範囲外にあるかどうかを検出します。これにより、入力データの誤りや破損に関連する問題を検出できます。
NaN チェッカー:
- フラグ: xla_sc_detect_nan = true
- 目的: SparseCore で処理される浮動小数点データ内の NaN（Not a Number）値の検出を有効にします。このフラグは、さまざまなコンパイラパスの入力または出力で NaN が検出された場合にエラーを発生させます。通常、このようなエラーには NaN が検出された場所に関する情報が含まれています。
境界チェッカー（メモリアクセス）:
- フラグ: xla_sc_assert_level=bounds
- 目的: このフラグは、メモリアクセス命令（VMEM の読み込み/保存や DMA 処理など）を書き換えて動的チェックを含める ASAN（AddressSanitizer）スタイルのツールを有効にします。これらのチェックは、メモリアクセスがターゲットメモリ領域の割り当てられた境界内にあるかどうかを検証します。
- 動作: 境界外メモリアクセスが検出されると、実行は失敗します。
- 注意: このチェッカーは、誤検出を生成する可能性があります。たとえば、チェッカーが完全に理解していない複雑なストライドアクセスパターンが原因で誤検出が発生する可能性があります。この変換は、バックエンドコンパイルプロセスの後半で適用されます。
バッファチェッカー（メモリ破損）:
- フラグ:
  - xla_tpu_buffer_contents_sanitizer_config='cores_to_sanitize: [TC, SC_SCS, SC_TILE], sanitizer_mode: LOCAL_ONLY'
  - xla_tpu_verify_launch_id_across_cores=true
- 目的: これらのフラグは、メモリバッファが関係のないオペレーションによって誤って破損または上書きされないようにするのに役立ちます。バッファサニタイザーは、バッファの内容をチェックして、予期しない変更がないことを確認します。

13. 量子化のサポート

SparseCore の SparseDenseMatmulOp は、32 ビット浮動小数点（FP32）と整数データ型の両方を使用して、エンベディングテーブルに対するオペレーションをサポートするように設計されています。通常、モデルトレーニングは埋め込みテーブルに FP32 精度を使用して実行されますが、トレーニング後の量子化（PTQ）を適用できます。PTQ では、推論に低精度のデータ型（8 ビット整数など）を使用できます。これにより、パフォーマンスの向上とメモリフットプリントの削減につながる可能性があります。

シミュレートされた量子化:

SparseDenseMatmulOp は「シミュレートされた量子化」を実行するように構成できます。この動作モードでは、エンベディングベクトルはまず低精度に量子化され、後続の計算で使用される前に高精度（FP32 など）に逆量子化されます。この手法により、量子化ノイズの影響を考慮しながらモデルをトレーニングできます。シミュレートされた量子化でトレーニングすると、推論用に完全に量子化された最終モデルの精度を向上させることができます。

`SparseDenseMatmulOp` の構成属性（量子化の場合）:

quantization_config_num_buckets = 256
- この属性は、32 ビットの浮動小数点数が量子化される個別のバケットまたはレベルの数を指定します。たとえば、8 ビット整数に量子化する場合は、通常 2^8 =256 個のバケットを指定します。
quantization_config_low = -X.X
- この属性は、量子化範囲の最小浮動小数点値を定義します。この指定された最小値を下回る入力値は、量子化中にこの最小値にクリップされます。
quantization_config_high = Y.Y
- この属性は、量子化範囲の最大浮動小数点値を定義します。この指定された最大値を超える入力値は、量子化中にこの最大値にクリップされます。

数値とパイプラインの相互作用:

TensorCore と SparseCore 間のパイプライン処理が有効になっているかどうかによって、モデルの数値動作が変わる可能性があります。パイプライン処理がアクティブな場合、SparseCore によって処理されたグラデーションが「古い」（前のイテレーションのもの）可能性があります。これは量子化プロセスと相互作用し、モデルのトレーニングダイナミクスや最終的な精度に影響を与える可能性があります。

14. 近日提供予定の機能と最近の改善点

SparseCore エコシステムは、継続的な開発と強化の対象となります。

ロードマップ:

サンプルディメンションのミニバッチ処理:
- これは、既存の語彙ディメンションのミニバッチ処理機能を補完する機能として計画されています。
- これにより、サンプルディメンションに沿ってエンベディング入力をさらにパーティショニングできます。これは、一度にサンプルサブセットからルックアップをフィルタして処理できるオンデバイスループを導入することで実現されます。このような機能は、サンプルあたりの ID 数が非常に多い場合や、処理ユニット間のロードバランシングを改善する場合に役立ちます。
行あたりの整数が 8 個未満の埋め込み（特徴の幅が小さい）のサポートを改善:
- 現在の設計では、8 個の浮動小数点数（32 バイトに相当）未満の埋め込み特徴の幅に大きなパディングが使用されることがよくあります。このパディングにより、HBM が無駄になり、コンピューティングリソースが十分に活用されない可能性があります。今後の改善では、特徴量の次元が小さいテーブルのこの非効率性を軽減することを目指しています。

大規模エンベディング モデル（LEM）用の SparseCore の詳細 コレクションでコンテンツを整理 必要に応じて、コンテンツの保存と分類を行います。