一、為什么 GPU 推薦模型訓練框架是剛需

1. PCG 算力集群缺點

最開始的時候，騰訊 PCG 所有的推薦模型訓練都是使用 CPU。但隨著業務的深入，以及深度學習模型的發展，PCG 算力集群在做下一代推薦模型時會遇到各種問題：

• 首先，系統網絡帶寬小，不穩定。
• 另外，很多推薦模型都很大，我們要考慮用多機多卡還是單機多卡，這就涉及到硬件的選型。
• 第三，云上分配到的 CPU 型號不能保證，有時會有一些 AMD 的 CPU，有時也會是一些英特爾的 CPU，這對于參數服務器架構也是非常不利的，如果 CPU 型號老舊，就會導致性能瓶頸，影響整體訓練框架的性能。
• 第四，云容器非獨占，整個機器的 IO 網絡都是共享的，因此可能導致整體訓練框架不穩定。

上述問題表明整個參數服務器已難以支持大參數量的推薦模型。因此需要改用 GPU，且只能用單機多卡的方式訓練非常大的模型。

我們遇到的問題主要包括，傳統網絡帶寬小，不穩定，InfiniBand 價格昂貴，而且要改造機房。另外，如果在設計 GPU 訓練的時候采用多機多卡，就會涉及到要把哪些機器換下來，加上一些支持 InfiniBand 網卡。也需要對機房進行 switch 的改造，所以也會導致 GPU 的機器改造成本非常的昂貴，這就決定了整個技術架構要采用單機多卡這樣一個硬件選型。

另外一個問題是 CPU 型號老，如果用多機多卡也會導致訓練集群每個節點的性能都不穩定。這是因為 GPU 在容器中是共享的，因此不能確保每一個 GPU affiliate 到特定的 CPU。如果多機多卡，CPU 調度分配會給上層的 K8s 的調度集群帶來非常大的壓力，也會導致出現很多 GPU 碎片。

2. 業務，上游生態

接下來從硬件現狀、業務以及上游生態的角度來看一下我們的技術目標。首先確定一個大目標，就是單機多卡去訓練一個非常大的模型，這個大模型的量級是幾 TB 到 10 TB。

另一個目標是盡量利用硬件的一些特性，利用硬件的整個機器不單只是 SSD、host memory，還有網卡，來提高模型訓練的大小上限。

整個架構要滿足離線訓練以及在線訓練。這是因為離線訓練是用來追模型的，而在線訓練是實時更新模型。所以我們的架構也要考慮如何快速上線、快速部署。

關于 GPU 硬件加速器的選型，要考慮 GPU 的供貨問題，所以我們在整個架構設計上也要考慮到不能僅支持 GPU，還要支持 XPU。

MLP 的架構選型，我們選擇了 TensorFlow，但是整個框架可以快速切換到 PyTorch，目前也正在做這樣的事情。

同時，還要考慮兼容老的參數服務器技術方案，能夠從 CPU 的訓練和推理快速遷移到 GPU，降低遷移成本。出于這樣的考慮，要兼容參數服務器。

以上就是我們根據目標進行的技術選型。

二、GPU 推薦模型訓練框架怎么做才最高效

1. GPU 訓練的數據結構

接下來介紹 GPU 推薦模型訓練框架的設計。首先要明確數據結構。我們的大目標是單機多卡，因此肯定要有 SSD，然后是 host memory 和 GPU，這幾個層級與傳統的 GPU 訓練框架非常相似。我們也采用了三級緩存的設計，但是考慮到最下層是對象存儲，是用來存模型的，還有樣本的輸入，存儲空間非常大，可以有幾 TB，甚至上百 TB 。因為其存在遠端，讀寫速度只有幾 GB 到幾百 MB 的量級。我們需要考慮將讀寫的過程跟訓練的過程進行分離。

我們盡可能將讀寫分離，因此要把樣本以及模型落盤到 SSD，SSD 的性能只有幾到十幾 GB 每秒，所以也需要對 SSD 讀寫以及它的訓練過程進行分離。

另外是 host memory。Host memory 跟 GPU 的 HBM 讀寫速度已經有幾十 GB 每秒，所以在此處設計上會有一個快速的快交換。

GPU 的 HBM 容量有幾十 GB，可以以塊的形式刷進 HBM，然后 HBM 再以 batch 的形式進行訓練。

數據結構整體分了四級，如上圖所示，圖中顏色與左側存儲的地方對應。首先是對應不同 DataSet，會有很多數據，我們會把所有的數據分成不同的 group，不同的 group 又細分到不同的 PASS，一個 PASS 又會包括多個 batch。這樣分塊就剛好滿足了模型的讀寫性能以及不同的硬件特性。

2. GPU 訓練整體架構

訓練流程如上圖所示。因為一些硬件特性的原因，我們會把訓練過程以及數據讀寫過程進行分離。當開始訓練時，首先會預先下載非常多的數據，將一個 group 的數據下載下來，然后進行預處理，處理完一個 group 之后會把 group 放到 host memory 里面，然后把單個 PASS 以整塊的方式 flash 到 GPU 的 HBM 里面。因為 GPU 的 HBM 已經有很多個 batch 了，GPU 的流式處理器訪問 global memory 的時候讀寫速度也會比較快，所以我們在整個過程里面就可以在 PASS 放到 HBM 以后，以 batch 進行每一輪的同步訓練。圖中可以看到，compute 就是 forward-backward 同步訓練，處理完全部 group 以后，會進行模型導出。

模型的數據下載和預處理等每個階段都是可以并發的，所以我們可以做到多級流水線的并發，將硬件資源的帶寬打滿。每一個 compute 的過程，forward 和 backward 都會涉及到 GPU。因為推薦模型訓練的計算任務都是 memory bound 的，所以我們會將每一個訪存的瓶頸列出來。因為每一次 forward backward 都是一個 batch，所以 global memory 的訪存也是一個瓶頸。到預處理階段，因為多級緩存的設計，所以涉及到 CPU 的瓶頸以及 CPU IO 讀取的瓶頸。還有一些業務上線的考量，沒有顯存訪存的限制，這些限制很大程度上決定著大模型訓練框架的性能。后面還會詳細介紹相關優化。

下載數據過程中的優化手段比較雜，難成體系。首先因為下載會涉及到網絡，因此會用到 DPDK 這樣的網絡優化庫。在網絡訪問中，可能要通過 TCP 的手段去訪問遠端的一些東西，所以也會用到 Google 的 BBR 擁塞控制算法去優化網絡窗口，從而優化下載速度。第三個是 DMA，網卡每下載一批數據，就可以用更少的 CPU 去完成這一下載動作。第四是落盤，在計算過程中，以及 group 交換的過程中需要用到這樣的一個優化手段，因為每一個 SSD 盤的帶寬是有限的，假如一臺機器配備多個 SSD 盤，然后每多個 SSD 盤可以通過 LVM 虛擬出來，這樣多個盤就可以并行訪問了，即可加快數據下載的速度。第五是 direct IO，因為我們不需要做太多的 cache IO，所以可以直接用 direct IO 去完成 group 落盤的操作。最后是數據結構的一些優化，例如我們采用了 Parquet 列存，這樣既可減少 IO 的一些 linux kernel 內核的終端操作，也可以把數據變成一個平整的大塊，更利于對數據進行并發處理。

3. 各階段優化詳細介紹

下面介紹預處理階段的優化。首先我們會把每一次每一個 group 所用到的去重之后的 key 先撈出來，將這些 key 里的 feature 變成 CSR 格式，這兩點非常重要，因為 CSR 格式可以減少訪存的 cache miss rate。另外，可以通過去重的 keys，對所需要的數據進行預處理，這樣就可以實現整塊 flash 到 GPU memory 的高效操作。

在前文對數據結構的介紹中提到，整體是一個三塊的結構，其大小是從下往上遞減的，那么其內部數據結構是怎樣的呢？一個 group 會先預處理，然后落盤，與訓練過程是分離的，所以它會把每一個 key 要用到的 embedding Vector 以及它的 optimizer state 先提取出來落盤，然后我們會把它變成一個 sample，一個 sample 變成一個 CSR 的數據，然后多個 CSR 的 sample 會變成一個 CSR 的 batch，放到 PASS 里。每一個 CSR 會存在 PASS 里面，由多個 PASS 組成一個 group。因為整個過程中要盡量節省 host memory，所以會先把它落盤到 SSD。

在開始訓練 PASS 全部數據之前，先將 PASS 需要的全部 embedding vector 整塊 flash 到 GPU 顯存上的 embedding 里面。如上圖中紅框的部分，就是已經準備好數據了，綠色的 PASS 部分剛好對應上三級緩存中綠色的一塊。整塊 flash 到 GPU 的 global memory 里面，在 global memory 里面存的一個 PASS 的結構是 CSR 的 batch 以及這個 batch 所要用到的 embedding Vector。這個 embedding 綠色這一塊其實是一個 GPU 的 hashtable。這里我們做了一個特異化的處理，傳統的 GPU hashtable 有增刪查改之類的過程，而我們直接將插入功能去除了，因為我們已經把整塊數據處理好了，不再需要 hashtable 的 insert 過程，所以將 insert 去掉，然后把整塊的 hashtable 的數據結構準備好，直接 flash 到 global memory 里面。這樣性能也會非常好，而且每一次 embedding 只存一個 group 所需要的訓練數據，GPU 的顯存就只關注這一個 group 數據的 embedding vector 就夠了。

只要 SSD 能存得下，那么本地單機多卡就能訓練得了多大的模型。例如這樣一個架構，雖然單機多卡，總的顯存只有 320 GB，但可以訓練十幾 TB 的模型，這也是我們的框架的優勢之一。正如前面介紹的，每一個計算過程里面，我們會把 PASS flash 到 GPU memory，在每一次 forward backward 的過程中會撈一個 batch 進行 forward，然后到 TensorFlow 的一個計算圖，再 backward 回來。這中間有一個黃色的箭頭，代表 TensorFlow 在 forward backward 時可以用 tape。因為每一個 TensorFlow 的計算都是 batch 數據并行的，所以每一個 TF 的計算圖都會虛擬一個小的塊，這個塊會在 forward backward 之后返回一個梯度，梯度會在一個多卡的 buffer 里面做 reduce 這樣一個操作。然后 scatter 到不同的卡上，因為 key 已經對于每一個卡做了哈希，所以在每一個 batch 里面，哪一個 key 在哪一個卡上也是固定。

接下來介紹一下 Compute 過程中的優化手段。空間優化方面，首先 embedding 部分，我們對其空間復雜度進行了優化，主要手段是 Dynamic embedding，就是 key value 指向一個 embedding vector，并且這個 embedding vector 不是定長的，embedding vector Lens 可以從 0 到最大的 Lens 過渡，這樣就很容易實現 Dynamic embedding。

第二個優化是 Multiple Hash，即一個 hashtable 可以分解成兩個 hashtable，其空間大小也可以從一個 hashtable 變成多個小的 hashtable，這樣空間也可以得到優化。

第三是混合精度訓練，這個做法比較大眾化，但確實是一個比較好的空間優化手段。

計算優化方面，首先是 Unified Hash Table，因為不同的特征會放在不同的 embedding table 里面。通過對 key 進行 rehash，因為 key 的數據特征是 int64，我們用低 48 位去存真實的 key，高 16 位存 hashtable 的 ID，這樣就可以把所有特征打散到一個非常大的 hashtable，查找的過程就只要一個 kernel，backward 時候也只要一個 kernel 就可以進行 update，從而節省了 kernel launch 的開銷。

第二是 kernel merge，這也是我們最近做的一個比較新的功能。

第三是混合精度，其實不僅是優化空間，因為在 GPU 上面我們有 FP16，FP16 可以用向量化訪存，也可以用 half2 這樣的子數據結構去做多數據的指令融合。所以這個其實也是有計算優化的。

最后一個是非對稱結構的 hash table。

上圖是對 unified hash table 的詳細介紹。

再詳細介紹一下 kernel merge，把相同的一些參數合到一起，這樣 pooling 的kernel 開銷就可以降低了。另外這些都涉及到 GPU 的一些優化，在此不做詳細展開。

我們的業務上需要增量模型、全量模型的分塊操作，每天更新一個全量模型，每小時更新一個增量模型，這樣就可以快速地進行上線。

模型更新、模型上線的優化手段主要包括，INT 量化、異步 IO 以及騰訊云的 COS 協程隊列等。

三、未來展望

最后是對未來工作的展望。首先，因為 GPU 卡供應比較緊張，將來可能無法使用 A100 去做訓練，因此需要考慮如何在 A10/T4/P4 上進行訓練，或在非英偉達 GPU 上進行訓練。第二是推薦大模型與 GPT 的結合。第三是更靈活的架構，例如 PS 以及 embedding 是不是一定要在 GPU 上，甚至 GPU 的訓練能不能結合 PS 的架構去做更大。第四是更大規模的訓練，我們目前只支持單機多卡，模型級別上限只達到 TB 級，之后要考慮能否支持 PB 級。最后是希望使用更低的硬件配置來訓練更大的模型。