TPU7x (Ironwood) 效能最佳化

本指南說明幾種方法，可有效管理 TPU7x (Ironwood) 多層記憶體系統之間的資料移動，進而提升效能。包括低精確度訓練、sharding、通訊最佳化、活化重新具體化、範圍向量記憶體 (VMEM) 調整，以及自訂加速器核心等技術。

如要充分發揮 TPU7x 的效能，請先熟悉 Ironwood 架構，尤其是記憶體階層和互連拓撲。詳情請參閱「TPU7x (Ironwood)」。

使用 FP8 進行低精確度訓練

FP8 (8 位元浮點) 是一種高效的數值型資料格式，主要用於加速模型訓練和推論。TPU 使用 8 位元表示數字，而非標準的 16 位元格式 (FP16 或 BF16) 和 32 位元 (FP32)，因此能大幅加快資料處理速度，並減少記憶體用量。

TPU7x 支援 FP8 資料類型的內建硬體加速功能，每個晶片的理論峰值效能為 4614 TFLOPS。這項功能可大幅縮短端對端訓練時間。對於相容的運算，尤其是 AI 工作負載常見的密集矩陣乘法，使用 FP8 可將效能提升至標準 BF16 訓練的 1.3 倍。與 BF16 相比，FP8 的尖峰 FLOPs 增加一倍，權重和活化函式的記憶體用量則減半。對於運算受限的工作負載，以及受記憶體容量或頻寬限制的情境，FP8 應是主要的調整槓桿。

使用 FP8 可帶來下列效能優勢：

  • 降低高頻寬記憶體 (HBM) 壓力：記憶體用量較小，因此較大的模型或推論期間 KV 快取較大的模型，可完全容納在 192 GiB 的 HBM 中。這可避免將資料卸載至速度較慢的主機記憶體，進而降低成本。
  • 提高有效批次大小：FP8 可減少活化所需的記憶體，因此能使用更大的批次。這可提升資料平行處理能力，進而提高處理量，並更有效地運用運算單元。
  • 降低記憶體頻寬需求：每次運算移動的資料量減半，可降低 HBM 到 MXU 資料路徑的需求。在資料移動經常成為瓶頸的系統中，這有助於讓 MXU 保持滿載。

如要使用 FP8，且效能不降低或僅略微降低，必須謹慎選擇量化技術。以下是 FP8 訓練的最佳做法：

  • 縮放精細程度：以每個張量縮放為基準。如果發生品質或效能問題，請改用逐軸縮放。可能不需要調整子管道。
  • 縮放模式：動態縮放會在執行階段計算縮放比例係數，是維持品質的良好預設值。雖然靜態縮放可消除運算，大幅提升效能，但需要仔細分析才能判斷正確的縮放比例係數，而且可能不適用於所有用途，尤其是模型設定變更時。相反地，部分強大的模型和設定可以將權重或活化函式的縮放比例固定在 FP8 限制，讓您減少量化負擔，同時維持準確度並提升效能。
  • FP8 格式 (E4M3 和 E5M2)：常見且有效的方法是混合使用 FP8 格式。舉例來說，在正向傳遞中，權重和活化函式可使用 E4M3，充分發揮 E4M3 的高精確度優勢；在反向傳遞中，梯度可使用 E5M2，因應梯度較大的動態範圍。
  • 四捨五入：使用「四捨五入至最接近的偶數」(RNE)，而非隨機四捨五入梯度，可維持品質，同時提供更佳的效能和重現性。
  • 在 MaxText 中啟用 FP8：MaxText 透過 QWIX 量化程式庫支援 FP8 訓練。如要啟用量化功能，請在設定中設定下列旗標：use_qwix_quantization=true。

資料分割和平行處理

Sharding 是將大型模型或其訓練資料切成較小的片段，並分配到多個 TPU 晶片或核心的程序。選擇合適的分區策略，對於在 TPU7x 上達到高效能至關重要。

如果只是一味追求平行處理程度，往往會導致通訊受限，成效不彰。最佳做法通常是選取符合記憶體限制的最簡單分區策略，這樣可以盡量減少通訊負擔，並有效率地運用運算單元。

在選取 shard 策略之前，任何效能調整作業的第一步都應該是算術強度分析。這項分析會判斷特定運算是否受限於運算、記憶體頻寬或互連頻寬。計算方式是將浮點運算次數除以必須移動的資料位元組數。

算術強度高，代表工作負載受限於運算能力。如果算術強度偏低，代表工作負載受限於記憶體或通訊，效能會受限於資料從 HBM 或透過 ICI 網路移動的速度。這項分析有助於找出理想的批次大小和資料分割策略。舉例來說，如果工作負載受限於通訊，則不適合採用會增加通訊量的分區策略，例如高度張量平行處理。

資料分割策略決策框架

MaxText 提供多種 sharding 策略。最佳選擇取決於模型架構、序列長度，以及平衡運算負載與通訊負擔的需求。

  • 完全分區資料平行處理 (FSDP)：這是資料平行處理的首選預設策略。FSDP 會在資料平行裝置上，將模型權重、梯度和最佳化器狀態進行 shard。在運算期間，每部裝置都會執行 All-Gather 作業，擷取本機微批次所需的完整權重。只要每部裝置的批次大小夠大，足以隱藏這項 All-Gather 通訊的延遲時間，FSDP 就會非常有效。如果是混合專家 (MoE) 模型，計算算術強度時必須考量稀疏性。
  • Tensor 平行處理 (TP)：TP 會在裝置間分割個別 Tensor。一般來說，張量是多層感知器 (MLP) 和注意力區塊中的權重矩陣。硬體的高算術強度 (11.5k) 對模型的維度要求非常高，才能讓 TP 在 ICI 上可行，而嘗試使用 TP 可能會導致系統受限於通訊。
  • 專家平行處理 (EP)：這是訓練 MoE 模型的標準和必要策略。EP 會在一組裝置上對「專家」層進行 shard，並使用 All-to-All 通訊集合，將權杖轉送至指定的專家裝置。如果模型的 MLP 維度夠大，接近效能上限，EP 就能發揮效率。
  • 情境平行處理 (CP)：CP 是一項專門策略，對於訓練序列長度極長的模型至關重要。主要功能是管理活化函數的記憶體消耗量，這會隨著序列長度呈二次方成長，並可能超過 HBM 容量。CP 會將活化張量的序列維度進行 shard，因此可使用分數形式的每部裝置批次大小。由於 CP 引入的通訊量比 FSDP 多，一般規則是使用滿足記憶體限制所需的最小 CP 程度，並確保批次軸 shard 保持為整數。

下表列出常見工作負載類型與最佳資料分割策略的對應關係：

工作負載類型 建議的主要 sharding 次要 sharding 主要瓶頸 理由
密集模型 - 短序列 FSDP 不適用 重新實體化、FF Matmuls FSDP 可提供最佳平衡。如果序列較短，啟用記憶體可能不是主要問題。關鍵在於全域批次大小要夠大，才能隱藏 FSDP 的權重 All-Gather。隨著批次大小增加，活化大小也會增加，因此需要適當的重新具體化政策，確保這項設定不會記憶體不足。
密集模型 - 長序列 FSDP CP Flash Attention、活化記憶體 活化記憶體會成為主要限制。必須啟用 CP，才能啟用分數形式的個別裝置批次大小，並避免記憶體不足 (OOM) 問題。Flash Attention 是運算和時間浪費的主要來源。
MoE 模型 - 短序列 FSDP + EP 不適用 All-to-All (專家轉送)、重新具體化 MoE 模型需要 EP 來將專家模型進行 shard。權杖轉送的全對全通訊是主要瓶頸，必須重疊。重新具體化也是浪費資源的重要原因。
MoE 模型 - 超大規模 FSDP + EP + PP 模型平行處理 (MP) 先前所述的所有瓶頸，以及管線氣泡 如果模型超出單一 Pod 的記憶體，則需要 PP 將層級分片至多個 Pod。這會產生 DCN 通訊和管道泡泡負擔。這項設定非常複雜，需要仔細調整。

通訊最佳化

在 TPU7x 上，重疊通訊和運算的主要機制稱為 SparseCore 集合卸載。Ironwood 架構包含專屬的 SparseCore 單元，可做為獨立的控制執行緒，管理 ICI 網狀架構上的資料移動。這樣一來，集體通訊作業 (例如 All-Gather 或 Reduce-Scatter) 就能與 TensorCore 上的主要運算作業平行執行。建議在 TPU7x 上使用這個方法執行非同步集體作業。使用建議的旗標，為最常見的集合啟用卸載功能。

啟用重新具體化

活化值重新具體化 (又稱梯度檢查點) 是減少模型 HBM 占用空間的基本技術。這項技術不會將前向傳遞的所有中間活化值儲存在 HBM 中，供反向傳遞使用，而是只儲存少數關鍵活化值 (檢查點)，並在反向傳遞期間隨選重新計算其他活化值。這會大幅節省記憶體，但會增加運算量 (標準 Transformer 區塊的 FLOPs 約增加 25-30%)。

如何積極套用重新具體化，是關鍵的調整參數，完全取決於主要瓶頸，而瓶頸通常會隨著序列長度而變化。

長序列工作負載 (例如 128k)：在這種情況下，活化張量的大小是 HBM 的主要消費者。工作負載通常受記憶體限制。因此，套用積極的重新具體化政策非常有利。節省記憶體後，訓練作業就能順利進行，不會發生記憶體不足錯誤，也能使用較大的批次大小，而重新計算的運算負擔是值得的取捨。

短序列工作負載 (例如 8k)：在這些情況下，活化記憶體較不重要，工作負載較可能受限於運算能力。重新具體化的運算負擔可能是導致效率低下的最大原因。

調整 MaxText 的重新具體化政策

MaxText 提供一組預設和自訂政策，可讓您精細控管重新具體化，並使用 remat_policy 旗標設定。

預設政策

MaxText 提供下列內建政策：

  • full：最激進的政策，幾乎會重新具體化所有內容。這會盡可能減少 HBM 用量，但會盡可能提高重新運算負擔。非常適合記憶體極度受限、長序列的情境。
  • minimal：最不積極的政策，會儲存大部分的活化。這可盡量提高 HBM 使用率，同時盡量減少重新運算。最適合短序列、受運算限制的工作負載，且不需考量記憶體。
  • 中繼政策：save_dot_with_context_except_mlp、save_qkv_proj 和 save_out_proj 等選項提供各種折衷方案，可選擇性地檢查昂貴的點積運算輸出，同時重新具體化較便宜的元素運算。

自訂政策

如要進一步控管，可以將 remat_policy 設為 custom。這可讓您指定模型解碼模組中個別層的行為。每個圖層都可以指派下列三種行為之一：

  • device：啟用項目會儲存在 TPU 裝置的高頻寬記憶體 (HBM) 中。
  • remat：捨棄活化值，並在反向傳遞期間重新具體化。
  • offload：啟用作業會從 HBM 移至 CPU 主機的記憶體，釋出 HBM，但會產生 PCIe 傳輸延遲。

範圍 VMEM 調整

核心效能 (例如 FlashAttention) 取決於核心中選取的圖塊大小，而圖塊大小受限於可用的向量記憶體 (VMEM)。TPU7x 晶片中的兩個 TensorCore 各有 64 MiB 的向量記憶體 (VMEM)。這項 VMEM 容量可分配給目前範圍 (範圍 VMEM) 和未來權重預先擷取。增加範圍 VMEM 可增加核心中的圖塊大小，進而減少記憶體停滯並提高核心效能。您可以設定 xla_tpu_scoped_vmem_limit_kib (以 LIBTPU_INIT_ARGS 為單位) 來變更範圍內的 VMEM 大小，這可用於探索核心效能以及端對端效能限制。最佳化限定範圍的 VMEM 大小，可能會間接影響自訂 Pallas 核心效能，因為增加限定範圍的 VMEM，可為核心內圖塊大小解鎖更大的超參數搜尋空間。

Tokamax 核心

Tokamax 是高效能 JAX 核心程式庫，內含許多高度最佳化的 TPU 核心，可解決多項常見的硬體特定瓶頸：

  • Splash attention：Splash attention 是主要的注意力機制實作方式，可消除標準注意力機制的 HBM 瓶頸，並在 TPU 上使用最有效率的注意力機制實作方式。
  • Megablox 分組矩陣乘法 (GMM)：針對 MoE 工作負載，Megablox 會透過不規則活化表示法進行運算，有效率地處理分組矩陣乘法。這項作業會有效率地對應不規則維度，計算 LHS 中不規則資料列群組與對應專家矩陣之間的矩陣乘法，因此不必將批次填補至固定大小。
  • 使用 tune-jax 進行經驗調整：tune-jax 程式庫提供公用程式，可對最佳區塊大小進行經驗搜尋。預設核心大小通常不是最佳選擇，調整後可選擇適合硬體的 VMEM 圖塊大小，充分運用硬體。
  • 最大 logits 估計值：設定 max_logit_const 的值，可進一步最佳化 Tokamax Splash 注意力核心。設定後，系統會在注意力機制的 softmax 運算 (softmax(Q * KT)) 期間，取代最大 logit 的縮減計算，減少部分運算和同步處理負擔。在 MaxText 中，這項功能是透過設定 use_max_logits_estimate 實作，可設為 None (停用) 或浮點值。請確認特定模型的 logit 範圍與預估值相容，以免發生數值溢位。如果設定這個值，建議進行收斂測試。