LHS 費用模型

tldr;

本頁說明 Latency Hiding Scheduler 使用的成本模型內部運作方式。如要調整模型，請直接前往「調整」部分。

延遲隱藏排程器 (LHS) 是一種編譯器傳遞，可排程 HLO DAG，盡量縮短實際時間。

這項服務的決策依據是統一的成本模型，該模型會混合使用成效表和分析模型。具體來說，XLA 會嵌入 GEMM 和快速互連集合的效能表，並針對其他情況使用分析網路和融合成本模型。本文件其餘部分會從高層次說明這些內容的內部運作方式。

成效表 - ICI 集體投資計畫

效能表包含兩個主要元件：收集器和插補器。

收集器

收集器是 C++ 工具，負責產生集體作業的效能表格。這項指標會評估個別 HLO 作業的效能 (例如：all-gather、all-reduce) 跨靜態定義的參數空間。

運作方式

這項工具會掃描特定叢集的一系列集體作業、轉移大小和轉移配置。這項工具會使用現有的多主機 HLO 執行器基礎架構和 ExecutionProfile 資料，執行產生的 HLO 並收集效能指標。

資料收集參數

系統會針對下列參數的交叉乘積收集延遲時間表：

  • 集體類型：
    • all-reduce
    • all-gather
    • reduce-scatter
  • 轉移大小：
    • 對數比例，從 1024B 到 2GiB (例如 1024B、2048B、4096B 等)
  • 轉移計畫：
    • rail-aligned
    • non-rail-aligned

這項掃描作業適用於具有 2、4 和 8 部裝置的節點內叢集。

輸出

集合執行結果是 .pbtxt 格式的延遲時間表 (每個平台約 116 KB)。

內插器

「插補器」是編譯器元件，會取用產生的效能表，在編譯期間提供執行階段估算值。

內部資料結構

初始化時，插補器會將效能表處理成對應。這個對應會使用 (collective_type, transfer_scheme) 元組做為鍵。

每個鍵的相關值都是 2D 歐幾里得平面。這個平面會根據兩個軸，為網路輸送量 (由收集器測量) 編製索引：

  1. 傳輸大小。
  2. 涉及的裝置數量。

查閱和插補

編譯器遇到集合運算時，插補器會執行下列步驟：

  1. 它會使用作業的 (collective_type, transfer_scheme) 做為對應鍵，找出正確的 2D 輸送量平面。
  2. 然後，系統會使用作業的 (transfer_size, num_devices) 做為查詢點，在該 2D 平面中進行加權平均擷取 (以歐幾里得距離為準)。
  3. 這項查閱作業的結果是單一且不重複的網路總處理量值。

依據：處理量和外插

系統的設計目的是儲存網路輸送量，而非原始延遲時間。這項設計選擇大幅簡化了表格中未明確顯示的傳輸大小的成效外推作業。

如果延遲時間表顯示網路頻寬在集體大小 S 時達到飽和，此時的總處理量 T 即為最大值。對於任何大小為 S' > S 的新集合，執行階段可估算為：

\[\text{EstimatedTime}(S') = \frac{S'}{T_{\text{saturated} } }\]

因此模型可以估算任何大小的集合體效能，即使大於 Collector 測量的 2 GiB 上限也沒問題。

  • 低估最大處理量。
  • 因此，請高估大型轉移作業的執行階段時間。

一般來說，XLA:GPU 團隊會維護效能表，但如果使用者決定自行提供，則由使用者負責產生表格，並確保表格具有代表性，且包含目標硬體的頻寬飽和區域測量結果。

成效表 - GEMM

與集合系統類似，GEMM 延遲時間表由兩個元件支援：收集器和插補器。

收集器

收集器是 C++ 工具，可計算一般矩陣乘法 (GEMM) 的效能表。這項指標會評估 HLO dot op 層級的矩陣乘法運算效能。

運作方式

這項工具會掃描 GEMM 維度的靜態空間 (批次、兩個非收縮和一個收縮維度) 和資料類型。

  • 預設資料型別： LHS = bf16,f32、RHS = bf16,f32、OUT = bf16,f32。
  • 基礎架構：重複使用 HLO 作業剖析器。

收款參數

系統會針對下列維度的跨產品收集延遲時間表：

  • batch： {1, 2, 4}
  • m (非合約)： {256, 512, ..., 4096}
  • n (非收縮)： {256, 512, ..., 4096}
  • k (收縮)： {256, 512, ..., 4096}

輸出和儲存

完整掃描會產生 .pbtxt 延遲時間表，可供插補器使用。

內插器

插補器是編譯器元件，會使用產生的表格估算 GEMM 效能。

原因：FLOPS 飽和度

收集到的延遲時間表可讓插補器重建每個項目的 FLOPS：

\[\text{FLOPS} = \frac{2 \times b \times m \times n \times k}{\text{runtime} }\]

重要洞察資訊是 FLOPS 會在某個時間點飽和，也就是硬體在超過特定矩陣形狀時達到尖峰 FLOPS。這種飽和度可讓您使用與集合體相同的推斷方法。

查閱和插補

插補器會從表格資料建構 4D 歐幾里得空間。為提供成效估計值，系統會在 4D 空間內執行加權平均插補。如果特定資料類型沒有資料表，系統會將每個維度正規化為位元組數，做為啟發式方法。

分析成本模型 - DCN

S 曲線集體成本模型

S 曲線模型是全分析型網路屋頂線模型。

總覽

這個模型會根據一組固定的網路屬性，估算集體作業的效能。

模型輸入

模型需要兩類輸入內容：

  1. 固定網路屬性 (使用者定義)：

    • 集體啟動額外負荷
    • NIC 速度
    • RTT (封包往返時間)

    根據預設，XLA 會自動偵測平台，並使用最常見架構的值。使用者可以設定這些屬性。詳情請參閱「微調」一節。

  2. 每個集合的輸入內容：

    • 集合類型 (例如AllGather、ReduceScatter)
    • 傳輸大小
    • 參與通訊的節點數量

整合

S 曲線模型已整合至 XLA:GPU，並用於 Hopper 和 Blackwell。

分析成本模型 - 融合

對於其他核心，我們會依據 GPU 效能成本模型估算正確的執行階段。詳情請參閱這篇文章。

調整

發出正確的 XLA 旗標，即可調整 S 曲線模型。在大多數情況下，預設設定就已足夠，但模型控制項也會在其他情況下顯示。

export NIC_SPEED_GBPS=... # NIC speed per GPU in Gigabytes
export GPUS_PER_NODE=... # Num of GPUs per cluster interconnected with fast network (e.g. NVLINK)
export XLA_FLAGS=--xla_gpu_analytical_latency_estimator_options="nic_speed_gbps=$NIC_SPEED_GBPS,gpus_per_node=$GPUS_PER_NODE"