深入瞭解大型嵌入模型 (LEM) 的 SparseCore

SparseCore 是專門的平鋪處理器，可針對涉及不規則、稀疏記憶體存取和運算的工作負載，提供高效能加速功能，特別是儲存在高頻寬記憶體 (HBM) 中的大型資料集。雖然這項技術擅長嵌入式查閱等工作，但也能加速處理各種其他動態和稀疏工作負載。

1. SparseCore 簡介

主要架構功能：

分塊架構：包含多個運算分塊 (每個分塊都是完整的資料流單元，具有自己的本機記憶體和處理單元)，可進行平行處理。
動態執行：原生支援資料相關的控制流程和記憶體存取，對於稀疏資料至關重要。
向量處理：使用小型向量工作 (8 個或 16 個元素，視硬體版本而定) 進行有效率的運算。
集中式控制：單一 SparseCore 排序器會協調所有圖塊的工作，確保作業同步。
支援資料摘要：包括專用的跨通道作業，有助於排序、篩選和前置總和等工作。
記憶體階層：策略性地運用 HBM 儲存大型資料集，並運用本機暫存記憶體 (SPMEM) 暫存經常存取的資料，大幅降低 HBM 延遲。

規格一覽：

屬性	TPU v4	TPU v5p	Trillium
SparseCores/Chip	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 轉換，這在處理特徵或表格堆疊時尤其重要。
- 將輸入資料轉換為座標 (COO) 稀疏格式，詳情請參閱下一節。
- 將資料分區，以便在晶片上可用的不同 SparseCore 之間有效分配資料。
限制驗證：
- 請確保輸入資料的特徵 (例如 ID 數量) 符合 SparseCore 的預先定義作業限制，例如 max_ids_per_partition 和 max_unique_ids_per_partition。
- 如果輸入資料超過這些限制，主機預先處理層可以嘗試將資料分割成較小的迷你批次，以符合限制。
資料移轉：
- 將處理及驗證過的資料有效複製到 TPU 的高頻寬記憶體 (HBM)，準備執行 SparseCore。

瞭解表格堆疊：

資料表堆疊是一項重要的最佳化技術，可將多個嵌入資料表邏輯合併，提升嵌入式查詢效率。這項程序通常由基礎機器學習架構自動處理。

特徵堆疊：多個不同的特徵共用同一個基礎嵌入資料表時，就會發生這種情況。常見的例子是針對各種類別特徵 (例如不同環境的郵遞區號) 使用單一嵌入字典。
資料表堆疊：在此情境中，多個不同的嵌入資料表會堆疊在一起。共用相同嵌入維度和最佳化工具設定的資料表通常會分組。

資料表堆疊的主要優點，是為這些堆疊資料表上的作業建立較大的有效批次大小。這可減少運算負擔，並有效隱藏晶片間通訊 (ICI) 的延遲。為求最佳效能，建議使用適量的堆疊表格 (一般介於 5 到 100 個)。

3. 轉換為 COO 張量

資料通常會先轉換為座標 (COO) 稀疏張量格式，再由 SparseCore 處理。COO 格式通常會使用三個陣列，有效率地表示稀疏矩陣：

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
- 說明：這個參數會針對單一晶片可從輸入批次中的一個樣本處理的嵌入 ID 總數，設定全域上限 (匯總所有表格)。這是管理晶片層級資源的機制，之後才會考量更精細的每個資料表或每個資料分割區限制。微調這個值通常取決於特定模型特徵和整體系統容量。
- 預設：64。
max_ids_per_table: Optional[Dict[str, int]] = None
- 說明：這個參數會指定每個邏輯資料表可處理的嵌入 ID 數量上限 (可包含重複項目)，並考量所有 SparseCore 的所有分割區。這比 max_ids_per_partition 的限制更寬鬆。如果資料表 T 分成 P 個分區，這項限制會套用至導向所有 P 分區的 ID 總和。這通常與 max_ids_per_partition_per_sample 和整體批次大小有關。
- 設定：通常是使用限制檔案 (例如使用 xla_sparse_core_max_ids_file 標記) 設定，其中定義了 max_ids_per_partition。這個資料表層級的概念是設定分區層級限制 (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 處理) 的標準訓練步驟中，管道化可讓 SC 處理步驟 i 的一部分 (例如正向或反向傳遞)，而 TC 同時處理相同步驟 i 的不同部分，甚至是相鄰步驟 (如 i-1 或 i+1) 的部分。
對梯度造成的影響：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 (加速線性代數) 是特定領域專用的編譯器，可將高階運算圖 (通常來自 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 中：

max_ids_per_partition：這項參數定義單一 SparseCore 在單一運算步驟中，預計傳送至或處理特定資料表分區的 ID 總數上限 (包括重複 ID)。
max_unique_ids_per_partition：這項設定定義任何單一 SparseCore 預期傳送至或處理的專屬 ID 數量上限。

轉換為實體資料表版面配置和處理：

資料表分片策略：嵌入資料表通常會在系統中的所有 SparseCore 之間「模數分片」。也就是說，每個 SparseCore 都會負責每個資料表的不同詞彙子集 (資料列)。一般來說，系統會根據 k = j % num_total_sparse_cores 等公式，為 SparseCore_k 指派 ID j。
「分區」的定義：在此情境中，「分區」是指嵌入資料表的特定區隔，單一 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 清單以進行嵌入式查閱時) 仰賴多種協調機制：

模組分片和隱含路徑：
- 嵌入資料表會在系統中的所有 SparseCore 之間進行模數分片。
- 主機提供一批輸入資料 (後續會預先處理為 COO 格式，包括 col_ids) 時，系統會使用 col_id 值判斷負責該特定 ID 的 SparseCore：target_sc_id = col_id % num_total_sparse_cores。
- 每個 SparseCore 實際上只會接收及處理對應至其指派字彙分割區的 ID 子集。主機前處理階段非常重要，可準備資料，讓每個 SparseCore 都能輕鬆識別並處理相關 ID。
依主機劃分的資料分布：
- 主機預先處理邏輯會分割整體輸入批次，並將 row_ids 和 col_ids 的相關部分 (以及任何相關聯的特徵或權重，如適用) 分配至每個 SparseCore 可直接存取的記憶體 (HBM)，或是 SparseCore 會從中擷取所需資料的共用 HBM。
SparseCore 內部處理：
- SparseCore 收到指定 ID 集後，會針對特定資料表分割區執行作業，例如簡化這些 ID，以及收集對應的嵌入向量。這些主要是指在 SparseCore 的自有圖塊中執行的本機運算，並使用其本機 SPMEM。
Inter-SparseCore 通訊 (All-to-All)：
- 在初始處理階段 (例如嵌入式查閱) 之後，系統可能會使用「全對全」通訊模式，合併或重新分配 SparseCore 的結果 (例如，在將啟用項目饋送至 TensorCore 層之前，該層會預期輸入內容對應於所有原始樣本位置)。如果原始輸入批次是為了平行處理而分散，這項資訊對於重建完整的啟用集至關重要。
與 TensorCore 通訊：
- SparseCore 會與 TensorCore 通訊，傳送嵌入啟動 (在前向傳遞期間)，並接收梯度 (在反向傳遞期間)。這項互動是由 XLA 編譯的程式協調，且經常會使用 HBM 做為中介緩衝區。管道化策略 (稍早討論過) 會大幅影響這項 SC-TC 通訊的時序和同步。

基本上，將 ID「分配」給適當的 SparseCore，主要是由分片配置和主機預先處理步驟處理。後續通訊會涉及 SparseCore 對本機資料執行作業，如果資料需要在 TensorCore 進一步處理前，於 SparseCore 之間進行全對全交換或重新排序，則可能還會進行集體通訊作業。

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 編譯器會納入精密的最佳化功能，準確判斷需要在 SPMEM 中暫存多少資料，以及哪些特定資料元素，才能有效隱藏 HBM 延遲時間並盡量提升效能。
- 動態分配限制：SparseCore 編譯器目前不支援動態暫存空間分配。這凸顯了透過仔細設定限制來進行靜態大小調整的重要性。
高頻寬記憶體 (HBM)：
- 性質：所有 SparseCore、TensorCore 和主機系統都能存取的大型共用記憶體資源。主要嵌入資料表會儲存在 HBM 中。
- 堆疊用量：SparseCore 作業通常需要 HBM 中的暫時儲存空間，用於存放不適合有限 SPMEM 的中繼結果，或需要在處理管道的較大階段之間傳遞。正向和反向傳遞期間的 HBM 堆疊用量估算方式如下：
  - 轉送 HBM 堆疊 (單一資料表) ≈ (2 * feature_width + 1) * max_unique_nz_per_row * logical_replica_count * 4 個位元組
  - 反向傳遞 HBM 堆疊 (單一資料表) ≈ 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 的核心記憶體管理策略，是使用小型快速的 SPMEM 處理 SparseCore 晶片主動處理的「熱」資料，盡量減少對較慢 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 個資料表，這是「最佳範圍」。
效率不彰的數值/量化：
- 問題：使用完整的 FP32 精確度時，如果使用 BF16 或量化整數等較低精確度格式就已足夠 (且可加快運算速度)，可能會造成效能瓶頸。
- 減輕影響：考慮使用精確度較低的選項。不過請注意，量化本身會產生一些額外負擔，而且可能需要仔細調整量化參數，才能維持模型準確度。
HBM 頻寬飽和：
- 問題：資料過度移轉至 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 (非數字) 值。如果偵測到各種編譯器傳遞的輸入或輸出內容中含有 NaN，這個旗標就會引發錯誤。這類錯誤通常會提供 NaN 的出現位置資訊。
邊界檢查工具 (記憶體存取)：
- 旗標：xla_sc_assert_level=bounds
- 用途：這個標記會啟用 ASAN (AddressSanitizer) 樣式的工具，重新編寫記憶體存取指令 (例如 VMEM 載入/儲存和 DMA 作業)，加入動態檢查。這些檢查會驗證記憶體存取是否在目標記憶體區域的分配界限內。
- 行為：如果偵測到超出範圍的記憶體存取，執行作業就會失敗。
- 注意：這個檢查程式可能會產生誤判，例如因為複雜的步進存取模式未完全被檢查程式理解，這項轉換作業會在後端編譯程序的後期套用。
緩衝區檢查程式 (記憶體損毀)：
- 旗標：
  - 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，並導致運算資源使用率偏低。日後我們會進行改良，以減少特徵維度較小資料表的這種效率不彰問題。

近期改善項目：

在 HBM 中暫存收集運算元：
- 這項最佳化措施可讓部分收集作業的輸入或輸出內容暫存在較大的 HBM 中，有助於減輕共用暫存記憶體 (SPMEM) 的負擔。
減少堆疊記憶體用量：
- 我們已實作強化功能，可減少 SparseCore 作業期間的 HBM 堆疊記憶體耗用量，理想情況下不會對整體效能或輸送量造成負面影響。

這些強化功能著重於提升 SparseCore 的效能、記憶體效率和作業彈性，以支援更多稀疏工作負載。

深入瞭解大型嵌入模型 (LEM) 的 SparseCore 透過集合功能整理內容 你可以依據偏好儲存及分類內容。