一、為什么 GPU 推薦模型訓(xùn)練框架是剛需

1. PCG 算力集群缺點(diǎn)

最開始的時(shí)候，騰訊 PCG 所有的推薦模型訓(xùn)練都是使用 CPU。但隨著業(yè)務(wù)的深入，以及深度學(xué)習(xí)模型的發(fā)展，PCG 算力集群在做下一代推薦模型時(shí)會(huì)遇到各種問題：

• 首先，系統(tǒng)網(wǎng)絡(luò)帶寬小，不穩(wěn)定。
• 另外，很多推薦模型都很大，我們要考慮用多機(jī)多卡還是單機(jī)多卡，這就涉及到硬件的選型。
• 第三，云上分配到的 CPU 型號(hào)不能保證，有時(shí)會(huì)有一些 AMD 的 CPU，有時(shí)也會(huì)是一些英特爾的 CPU，這對(duì)于參數(shù)服務(wù)器架構(gòu)也是非常不利的，如果 CPU 型號(hào)老舊，就會(huì)導(dǎo)致性能瓶頸，影響整體訓(xùn)練框架的性能。
• 第四，云容器非獨(dú)占，整個(gè)機(jī)器的 IO 網(wǎng)絡(luò)都是共享的，因此可能導(dǎo)致整體訓(xùn)練框架不穩(wěn)定。

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

我們遇到的問題主要包括，傳統(tǒng)網(wǎng)絡(luò)帶寬小，不穩(wěn)定，InfiniBand 價(jià)格昂貴，而且要改造機(jī)房。另外，如果在設(shè)計(jì) GPU 訓(xùn)練的時(shí)候采用多機(jī)多卡，就會(huì)涉及到要把哪些機(jī)器換下來，加上一些支持 InfiniBand 網(wǎng)卡。也需要對(duì)機(jī)房進(jìn)行 switch 的改造，所以也會(huì)導(dǎo)致 GPU 的機(jī)器改造成本非常的昂貴，這就決定了整個(gè)技術(shù)架構(gòu)要采用單機(jī)多卡這樣一個(gè)硬件選型。

另外一個(gè)問題是 CPU 型號(hào)老，如果用多機(jī)多卡也會(huì)導(dǎo)致訓(xùn)練集群每個(gè)節(jié)點(diǎn)的性能都不穩(wěn)定。這是因?yàn)?GPU 在容器中是共享的，因此不能確保每一個(gè) GPU affiliate 到特定的 CPU。如果多機(jī)多卡，CPU 調(diào)度分配會(huì)給上層的 K8s 的調(diào)度集群帶來非常大的壓力，也會(huì)導(dǎo)致出現(xiàn)很多 GPU 碎片。

2. 業(yè)務(wù)，上游生態(tài)

接下來從硬件現(xiàn)狀、業(yè)務(wù)以及上游生態(tài)的角度來看一下我們的技術(shù)目標(biāo)。首先確定一個(gè)大目標(biāo)，就是單機(jī)多卡去訓(xùn)練一個(gè)非常大的模型，這個(gè)大模型的量級(jí)是幾 TB 到 10 TB。

另一個(gè)目標(biāo)是盡量利用硬件的一些特性，利用硬件的整個(gè)機(jī)器不單只是 SSD、host memory，還有網(wǎng)卡，來提高模型訓(xùn)練的大小上限。

整個(gè)架構(gòu)要滿足離線訓(xùn)練以及在線訓(xùn)練。這是因?yàn)殡x線訓(xùn)練是用來追模型的，而在線訓(xùn)練是實(shí)時(shí)更新模型。所以我們的架構(gòu)也要考慮如何快速上線、快速部署。

關(guān)于 GPU 硬件加速器的選型，要考慮 GPU 的供貨問題，所以我們?cè)谡麄€(gè)架構(gòu)設(shè)計(jì)上也要考慮到不能僅支持 GPU，還要支持 XPU。

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

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

以上就是我們根據(jù)目標(biāo)進(jìn)行的技術(shù)選型。

二、GPU 推薦模型訓(xùn)練框架怎么做才最高效

1. GPU 訓(xùn)練的數(shù)據(jù)結(jié)構(gòu)

接下來介紹 GPU 推薦模型訓(xùn)練框架的設(shè)計(jì)。首先要明確數(shù)據(jù)結(jié)構(gòu)。我們的大目標(biāo)是單機(jī)多卡，因此肯定要有 SSD，然后是 host memory 和 GPU，這幾個(gè)層級(jí)與傳統(tǒng)的 GPU 訓(xùn)練框架非常相似。我們也采用了三級(jí)緩存的設(shè)計(jì)，但是考慮到最下層是對(duì)象存儲(chǔ)，是用來存模型的，還有樣本的輸入，存儲(chǔ)空間非常大，可以有幾 TB，甚至上百 TB 。因?yàn)槠浯嬖谶h(yuǎn)端，讀寫速度只有幾 GB 到幾百 MB 的量級(jí)。我們需要考慮將讀寫的過程跟訓(xùn)練的過程進(jìn)行分離。

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

另外是 host memory。Host memory 跟 GPU 的 HBM 讀寫速度已經(jīng)有幾十 GB 每秒，所以在此處設(shè)計(jì)上會(huì)有一個(gè)快速的快交換。

GPU 的 HBM 容量有幾十 GB，可以以塊的形式刷進(jìn) HBM，然后 HBM 再以 batch 的形式進(jìn)行訓(xùn)練。

數(shù)據(jù)結(jié)構(gòu)整體分了四級(jí)，如上圖所示，圖中顏色與左側(cè)存儲(chǔ)的地方對(duì)應(yīng)。首先是對(duì)應(yīng)不同 DataSet，會(huì)有很多數(shù)據(jù)，我們會(huì)把所有的數(shù)據(jù)分成不同的 group，不同的 group 又細(xì)分到不同的 PASS，一個(gè) PASS 又會(huì)包括多個(gè) batch。這樣分塊就剛好滿足了模型的讀寫性能以及不同的硬件特性。

2. GPU 訓(xùn)練整體架構(gòu)

訓(xùn)練流程如上圖所示。因?yàn)橐恍┯布匦缘脑?，我們?huì)把訓(xùn)練過程以及數(shù)據(jù)讀寫過程進(jìn)行分離。當(dāng)開始訓(xùn)練時(shí)，首先會(huì)預(yù)先下載非常多的數(shù)據(jù)，將一個(gè) group 的數(shù)據(jù)下載下來，然后進(jìn)行預(yù)處理，處理完一個(gè) group 之后會(huì)把 group 放到 host memory 里面，然后把單個(gè) PASS 以整塊的方式 flash 到 GPU 的 HBM 里面。因?yàn)?GPU 的 HBM 已經(jīng)有很多個(gè) batch 了，GPU 的流式處理器訪問 global memory 的時(shí)候讀寫速度也會(huì)比較快，所以我們?cè)谡麄€(gè)過程里面就可以在 PASS 放到 HBM 以后，以 batch 進(jìn)行每一輪的同步訓(xùn)練。圖中可以看到，compute 就是 forward-backward 同步訓(xùn)練，處理完全部 group 以后，會(huì)進(jìn)行模型導(dǎo)出。

模型的數(shù)據(jù)下載和預(yù)處理等每個(gè)階段都是可以并發(fā)的，所以我們可以做到多級(jí)流水線的并發(fā)，將硬件資源的帶寬打滿。每一個(gè) compute 的過程，forward 和 backward 都會(huì)涉及到 GPU。因?yàn)橥扑]模型訓(xùn)練的計(jì)算任務(wù)都是 memory bound 的，所以我們會(huì)將每一個(gè)訪存的瓶頸列出來。因?yàn)槊恳淮?forward backward 都是一個(gè) batch，所以 global memory 的訪存也是一個(gè)瓶頸。到預(yù)處理階段，因?yàn)槎嗉?jí)緩存的設(shè)計(jì)，所以涉及到 CPU 的瓶頸以及 CPU IO 讀取的瓶頸。還有一些業(yè)務(wù)上線的考量，沒有顯存訪存的限制，這些限制很大程度上決定著大模型訓(xùn)練框架的性能。后面還會(huì)詳細(xì)介紹相關(guān)優(yōu)化。

下載數(shù)據(jù)過程中的優(yōu)化手段比較雜，難成體系。首先因?yàn)橄螺d會(huì)涉及到網(wǎng)絡(luò)，因此會(huì)用到 DPDK 這樣的網(wǎng)絡(luò)優(yōu)化庫。在網(wǎng)絡(luò)訪問中，可能要通過 TCP 的手段去訪問遠(yuǎn)端的一些東西，所以也會(huì)用到 Google 的 BBR 擁塞控制算法去優(yōu)化網(wǎng)絡(luò)窗口，從而優(yōu)化下載速度。第三個(gè)是 DMA，網(wǎng)卡每下載一批數(shù)據(jù)，就可以用更少的 CPU 去完成這一下載動(dòng)作。第四是落盤，在計(jì)算過程中，以及 group 交換的過程中需要用到這樣的一個(gè)優(yōu)化手段，因?yàn)槊恳粋€(gè) SSD 盤的帶寬是有限的，假如一臺(tái)機(jī)器配備多個(gè) SSD 盤，然后每多個(gè) SSD 盤可以通過 LVM 虛擬出來，這樣多個(gè)盤就可以并行訪問了，即可加快數(shù)據(jù)下載的速度。第五是 direct IO，因?yàn)槲覀儾恍枰鎏嗟?cache IO，所以可以直接用 direct IO 去完成 group 落盤的操作。最后是數(shù)據(jù)結(jié)構(gòu)的一些優(yōu)化，例如我們采用了 Parquet 列存，這樣既可減少 IO 的一些 linux kernel 內(nèi)核的終端操作，也可以把數(shù)據(jù)變成一個(gè)平整的大塊，更利于對(duì)數(shù)據(jù)進(jìn)行并發(fā)處理。

3. 各階段優(yōu)化詳細(xì)介紹

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

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

在開始訓(xùn)練 PASS 全部數(shù)據(jù)之前，先將 PASS 需要的全部 embedding vector 整塊 flash 到 GPU 顯存上的 embedding 里面。如上圖中紅框的部分，就是已經(jīng)準(zhǔn)備好數(shù)據(jù)了，綠色的 PASS 部分剛好對(duì)應(yīng)上三級(jí)緩存中綠色的一塊。整塊 flash 到 GPU 的 global memory 里面，在 global memory 里面存的一個(gè) PASS 的結(jié)構(gòu)是 CSR 的 batch 以及這個(gè) batch 所要用到的 embedding Vector。這個(gè) embedding 綠色這一塊其實(shí)是一個(gè) GPU 的 hashtable。這里我們做了一個(gè)特異化的處理，傳統(tǒng)的 GPU hashtable 有增刪查改之類的過程，而我們直接將插入功能去除了，因?yàn)槲覀円呀?jīng)把整塊數(shù)據(jù)處理好了，不再需要 hashtable 的 insert 過程，所以將 insert 去掉，然后把整塊的 hashtable 的數(shù)據(jù)結(jié)構(gòu)準(zhǔn)備好，直接 flash 到 global memory 里面。這樣性能也會(huì)非常好，而且每一次 embedding 只存一個(gè) group 所需要的訓(xùn)練數(shù)據(jù)，GPU 的顯存就只關(guān)注這一個(gè) group 數(shù)據(jù)的 embedding vector 就夠了。

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

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

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

第三是混合精度訓(xùn)練，這個(gè)做法比較大眾化，但確實(shí)是一個(gè)比較好的空間優(yōu)化手段。

計(jì)算優(yōu)化方面，首先是 Unified Hash Table，因?yàn)椴煌奶卣鲿?huì)放在不同的 embedding table 里面。通過對(duì) key 進(jìn)行 rehash，因?yàn)?key 的數(shù)據(jù)特征是 int64，我們用低 48 位去存真實(shí)的 key，高 16 位存 hashtable 的 ID，這樣就可以把所有特征打散到一個(gè)非常大的 hashtable，查找的過程就只要一個(gè) kernel，backward 時(shí)候也只要一個(gè) kernel 就可以進(jìn)行 update，從而節(jié)省了 kernel launch 的開銷。

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

第三是混合精度，其實(shí)不僅是優(yōu)化空間，因?yàn)樵?GPU 上面我們有 FP16，F(xiàn)P16 可以用向量化訪存，也可以用 half2 這樣的子數(shù)據(jù)結(jié)構(gòu)去做多數(shù)據(jù)的指令融合。所以這個(gè)其實(shí)也是有計(jì)算優(yōu)化的。

最后一個(gè)是非對(duì)稱結(jié)構(gòu)的 hash table。

上圖是對(duì) unified hash table 的詳細(xì)介紹。

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

我們的業(yè)務(wù)上需要增量模型、全量模型的分塊操作，每天更新一個(gè)全量模型，每小時(shí)更新一個(gè)增量模型，這樣就可以快速地進(jìn)行上線。

模型更新、模型上線的優(yōu)化手段主要包括，INT 量化、異步 IO 以及騰訊云的 COS 協(xié)程隊(duì)列等。

三、未來展望

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