Meta 內部都在用的 FX 工具大起底:利用 Graph Transformation 最佳化 PyTorch 模型

超神經HyperAI 發表於 2022-11-30
PyTorch
PyTorch 中的 graph mode 在效能方面表示更為出色,本文介紹 Torch.FX 這個強大工具,可以捕捉和最佳化 PyTorch 程式 graph。

一、簡介

PyTorch 支援兩種執行模式:eager mode 和 graph mode。

eager mode 中,模型中的運算子在讀取時會立即執行,它易於使用,對機器學習從業者更友好,因此被設定為預設的執行模式。

graph mode 中,運算子先被合成一個 graph,然後作為一個整體進行編譯和執行,它的效能更高,因此在實際生產中大量使用。

具體來說,graph mode 支援運算元融合,兩個運算元透過合併,可以降低或本地化記憶體讀取以及核心啟動總開銷。

融合可以是橫向 (horizontal) 的:採取應用於多個 operand 的單一操作(如 BatchNorm),並將這些 operand 合併到一個陣列中。

融合也可以是縱向 (vertical) 的:將一個核心與另一個核心合併,後者需要使用第一個核心的輸出(如 ReLU 後接卷積)。

Torch.FX(縮寫為 FX)是一個公開可用的工具包,作為 PyTorch 軟體包的一部分,支援 graph mode 的執行。它可以:

  1. 從 PyTorch 程式中獲取 graph
  2. 允許開發者在獲取的 graph 上編寫 transformation**

Meta 內部先前已經在用 FX 來最佳化生產模型 (production model) 的訓練吞吐量 (training throughput)。本文將透過介紹 Meta 開發的基於 FX 的最佳化,來展示利用圖結構轉換 (graph transformation) 最佳化 PyTorch 部署模型效能的方法

二、背景

embedding table 廣泛存在於推薦系統中,本節將介紹 FX 和 embedding table 的背景知識。

2.1. FX

圖 1 是一個簡單示例,演示瞭如何用 FX 轉換 PyTorch 程式,它包含三個步驟:

  • 從程式中獲取 graph
  • 修改 graph(在本例中,我們用 GELU 代替 RELU)
  • 從修改後的 graph 中生成一個新程式

Meta 內部都在用的 FX 工具大起底:利用 Graph Transformation 最佳化 PyTorch 模型

圖1:在 PyTorch 模組中用 GELU 取代 RELU 的 FX

FX API 為檢查和轉換 PyTorch 程式 graph 還提供了許多其他功能。

2.2. embedding table

Meta 內部都在用的 FX 工具大起底:利用 Graph Transformation 最佳化 PyTorch 模型

圖2:批尺寸=1 的稀疏特徵 embedding table 示意圖

在推薦系統中,稀疏特徵(例如,User ID,Story ID)由 embedding table 表示。

embedding table E 是一個 HxD 矩陣,其中 H 是雜湊大小,D 是嵌入向量維度。E 的每一行都是一個浮點數向量。

feature hashing 的作用是將一個稀疏特徵對映到 E的索引列表中,例如 [S1,S2,...,Sk],其中 0≤Si<H。它的輸出值計算為 f(E[S1], E[S2],...,E[Sk]),其中 E[Si] 是 Si 行的向量,f 是池化函式,通常是 sum,average,max 三個函式之一。

為了充分利用 GPU,稀疏特徵通常為批處理。批處理中的每個實體都有自己的索引列表。如果一個批次有 B 個實體,可以簡單理解為一個表徵有 B 個索引列表。

更為嚴謹的表示方法是將 B 個索引列表合併成一個索引列表,並新增一個索引長度的列表(該批中的每個實體都有一個長度 length)。

例如,如果一批包含 3 個實體,其索引列表如下:

  • Entity 1: indices = [10, 20]
  • Entity 2: indices = [5, 9, 77, 81]
  • Entity 3: indices = [15, 20, 45]

則完整批尺寸的 indice 和 length 將是:

  • Indices = [10, 20, 5, 9, 77, 81, 15, 20, 45]
  • Lengths = [2, 4, 3]

而整個 batch 的 embedding table 查詢,輸出為是一個 BxD 矩陣。

三、3 種 FX Transformation

PyTorch 更新了 3 個 FX transformation,以加速對 embedding table 的訪問,本節將逐一介紹。

下文 3.1 關於將多個小輸入張量結合成一個大張量的轉換;3.2 關於將多個平行計算鏈融合成一個計算鏈的轉換;3.3 關於將通訊與計算重疊的轉換。

3.1 結合輸入稀疏特徵

batch 中的每個輸入稀疏特徵,都可以表示為兩個列表:一個索引列表和一個 B length 列表,其中 B 表示批尺寸。

在 PyTorch 中,這兩個列表都可以以張量的形式存在。當 PyTorch 模型在 GPU 上執行時,embedding table 通常儲存在 GPU 記憶體中(它更接近 GPU,讀寫頻寬比 CPU 記憶體更高)。

需要使用輸入稀疏特徵時,兩個張量都要先從 CPU 複製到 GPU。然而每個主機到裝置的記憶體複製都需要啟動核心,這對於實際的資料傳輸來說,會更加耗費時間。

如果一個模型使用了多個輸入稀疏特徵,這種複製可能成為效能瓶頸(例如,1000 個輸入稀疏特徵將需要從主機到裝置複製 2000 個張量)。

一個減少主機到裝置 memcpy 數量的最佳化方法,就是在多個輸入稀疏特徵傳送到裝置之前,先將其進行組合。

例如,給定以下三個輸入特徵:

  • Feature_A: indices = [106, 211, 7], lengths = [2, 1]
  • Feature_B: indices = [52, 498, 616, 870, 1013], lengths = [3, 2]
  • Feature_C: indices = [2011, 19, 351, 790], lengths = [1, 3]

組合後的形式為:

Features_A_B_C: indices = [106, 211, 7, 52, 498, 616, 870, 1013, 2011, 19, 351, 790], lengths = [2, 1, 3, 2, 1, 3]

所以不需要從主機到裝置複製 3x2=6 個張量,只需要複製 2 個張量。

圖 3(b) 描述了這種最佳化的實現,它包含兩個元件:

  • CPU 端:輸入 pipeline 被修改為將所有稀疏特徵的 indices 組合成一個張量,所有 length 組合成另一個張量。然後將這兩個張量複製到 GPU 上。
  • GPU 端:使用 FX,在模型 graph 中插入一個Permute_and_Split 運算元,從合併的張量中恢復單個特徵 indices 和 length 張量,並將其傳送至下游的相應節點。

Meta 內部都在用的 FX 工具大起底:利用 Graph Transformation 最佳化 PyTorch 模型
最佳化前:兩個張量都要從 CPU 複製到 GPU

Meta 內部都在用的 FX 工具大起底:利用 Graph Transformation 最佳化 PyTorch 模型
最佳化後:將輸入稀疏特徵進行組合

3.2 從訪問 embedding table 開始的計算鏈橫向融合

在一個生產模型中,每個 GPU 上有 10 個 embedding table 很常見。出於效能方面的考慮,對這些 table 的查詢被分到一組,這樣它們的輸出就被串聯在一個大張量中(見圖 4(a)中的紅色部分)。

為了對單個特徵輸出進行計算,使用 Split 運算元將大張量分成 N 個小張量(其中 N 為特徵的數量),然後將所需的計算應用於每個張量。

如圖 4(a) 所示,應用於每個特徵輸出 O 的計算是Tanh(LayerNorm(O))。所有的計算結果都被串聯成一個大的張量,然後傳遞給下游的運算元(圖 4(a) 中的 Op1)。

這裡主要的 runtime cost 是 GPU 核心啟動的開銷。例如,圖 4(a) 中的 GPU 核心的啟動次數為 2*N+3(圖中的每個橢圓都表示一個 GPU 核心)。這會影響效能,因為 LayerNorm 和 Tanh 在 GPU 上的執行時間,與它們的核心啟動時間相比很短。

此外,Split 運算元可能會建立一個額外的嵌入向量輸出張量的副本,消耗額外的 GPU 記憶體。

用 FX 來實現一種叫做橫向融合 (horizontal fusion) 的最佳化,可以大大減少 GPU 核心的啟動次數(在這個例子中,最佳化後的 GPU 核心啟動次數為 5,見圖 4(b))。

使用 Add_middle_dim 運算元代替顯式 Split,將 shape 為 (B, NxD) 的 2D 嵌入張量重塑為 shape 為 (B, N, D) 的 3D 張量。接下來將一個單一的 LayerNorm 應用到它的最後一維。對 LayerNorm 的結果應用一個 Tanh。最後,用 Remove_middle_dim 運算元將 Tanh 的結果恢復成 2D 張量。

由於 Add_middle_dim 和 Remove_middle_dim 只是重塑張量,並沒有建立額外的副本,所以也可以減少 GPU 記憶體的消耗。

Meta 內部都在用的 FX 工具大起底:利用 Graph Transformation 最佳化 PyTorch 模型
最佳化前:所有輸出被串聯到一個大張量中

Meta 內部都在用的 FX 工具大起底:利用 Graph Transformation 最佳化 PyTorch 模型
進行橫向融合最佳化後

3.3 計算與通訊間的重疊 (overlap)

面向投產的推薦模型的訓練,通常是在分散式 GPU 系統上完成的。由於每個 GPU 的裝置記憶體容量不足以容納模型中的所有 embedding table,因此需要將其分佈在多個 GPU 上。

在訓練步驟中,GPU 需要從其他 GPU 上的 embedding table 中讀取/寫入特徵值。這被稱為 all-to-all 通訊,可能是影響效能的重要原因。

透過 FX 實現一個 transformation,可以將計算與 all-to-all 通訊重疊。圖 5(a) 顯示了一個具備嵌入向量 table 訪問 (EmbeddingAllToAll) 及其他運算元的模型 graph 例項。如圖 5(b) 所示,在沒有任何最佳化的情況下,它們會在一個 GPU 流上順序執行。

使用FX將 EmbeddingAllToAll 分成 EmbeddingAllToAll_Request和EmbeddingAllToAll_Wait,並在它們之間安排獨立的運算元。
Meta 內部都在用的 FX 工具大起底:利用 Graph Transformation 最佳化 PyTorch 模型

圖5:計算與通訊的重疊

3.4 總結

Meta 內部都在用的 FX 工具大起底:利用 Graph Transformation 最佳化 PyTorch 模型
表1:本節討論的最佳化及解決的相應效能瓶頸

為了發現哪些模型會從這些 transformation 中受益,開發人員對 MAIProf 收集的執行在 Meta 資料中心的模型的效能資料進行分析。得出與 eager mode 相比,這些 transformation 在一組生產模型上實現了 2-3 倍的速度提升。

四、結語

從效能角度考量,PyTorch 中的 graph mode 比生產環境中使用的 eager mode 更受歡迎。FX 是一個強大的工具,可以捕捉和最佳化 PyTorch 程式 graph。本文展示了三種 FX transformation,用於最佳化 Meta 內部的生產推薦模型。

最後希望更多 PyTorch 開發者可以使用 graph transformation 來提升模型的效能。

—— 完 ——