微信基于 PyTorch 的大規(guī)模推薦系統(tǒng)訓(xùn)練實(shí)踐

作者：朱子霖 2023-04-04 07:32:35

TorchRec 作為一個(gè)官方的推薦框架，在 2022 年 1 月份推出之時(shí)，Meta就已經(jīng)利用它在 Instagram Reels 業(yè)務(wù)上順利訓(xùn)練并上線了一個(gè) 1250 億參數(shù)的模型，成為了一個(gè)經(jīng)過業(yè)務(wù)測試的 PyTorch 框架。有了 Instagram 這樣一個(gè)大業(yè)務(wù)的支撐，讓我們有了更多信心，終于可以去理性地考量一個(gè)基于 PyTorch 的推薦框架有什么樣的優(yōu)勢了。

本文將介紹微信基于 PyTorch 進(jìn)行的大規(guī)模推薦系統(tǒng)訓(xùn)練。推薦系統(tǒng)和其它一些深度學(xué)習(xí)領(lǐng)域不同，仍在使用 Tensorflow 作為訓(xùn)練框架，被廣大開發(fā)者詬病。雖然也有使用 PyTorch 進(jìn)行推薦訓(xùn)練的一些實(shí)踐，但規(guī)模較小，也沒有實(shí)際的業(yè)務(wù)驗(yàn)證，很難推動(dòng)業(yè)務(wù)嘗鮮。

2022 年 2 月，PyTorch 團(tuán)隊(duì)推出了官方推薦庫 TorchRec。我們團(tuán)隊(duì)在 5 月開始在內(nèi)部業(yè)務(wù)上嘗試 TorchRec，并且與 TorchRec 團(tuán)隊(duì)展開了一系列的合作。在幾個(gè)月的試用過程中，我們體會(huì)到 TorchRec 非常多的優(yōu)點(diǎn)，也感受到 TorchRec 在超大規(guī)模模型上仍存在一些不足。針對這些不足，我們設(shè)計(jì)了擴(kuò)展功能來填補(bǔ)它的問題。在 2022 年 9 月，我們設(shè)計(jì)的擴(kuò)展功能 dynamic embedding 已經(jīng)正式合進(jìn)了 TorchRec 的主分支，目前仍在與官方團(tuán)隊(duì)持續(xù)優(yōu)化。

一、TorchRec 可以為我們帶來什么

我們先來聊聊 TorchRec 可以給我們帶來什么？我們都知道推薦系統(tǒng)往往和公司的現(xiàn)金流直接掛鉤，試錯(cuò)成本非常高，所以大家需要的是一個(gè)經(jīng)過了業(yè)務(wù)測試的框架。這也是為什么之前的一些基于 PyTorch 的推薦框架都未曾被廣泛應(yīng)用過。而 TorchRec 作為一個(gè)官方的推薦框架，在 2022 年 1 月份推出之時(shí)，Meta就已經(jīng)利用它在 Instagram Reels 業(yè)務(wù)上順利訓(xùn)練并上線了一個(gè) 1250 億參數(shù)的模型，成為了一個(gè)經(jīng)過業(yè)務(wù)測試的 PyTorch 框架。有了 Instagram 這樣一個(gè)大業(yè)務(wù)的支撐，讓我們有了更多信心，終于可以去理性地考量一個(gè)基于 PyTorch 的推薦框架有什么樣的優(yōu)勢了。

對于團(tuán)隊(duì)中的不同成員，TorchRec 有不同的好處。首先，對于團(tuán)隊(duì)中占絕大多數(shù)的算法工程師而言，PyTorch 推薦框架讓大家終于可以享受到像 CV、NLP 工程師體會(huì)到的那種更人性化的動(dòng)態(tài)圖和調(diào)試的體驗(yàn)。

另外，PyTorch 極好的兼容性——一個(gè)基于 PyTorch1.8 做的模型，不需要改一行代碼就可以在最新版本 1.13 上運(yùn)行——讓算法工程師終于可以放心地升級(jí)框架，從而享受到最新的框架功能和更優(yōu)秀的性能。而反觀一些基于 TensorFlow 的推薦框架，往往被卡在 TensorFlow 的某一個(gè)版本上，例如很多團(tuán)隊(duì)可能還在使用基于 TensorFlow 1.x 的內(nèi)部框架。TensorFlow 1.x 在 2021 年 1 月份就已經(jīng)停止維護(hù)了，這就意味著在近兩年的時(shí)間內(nèi)，所有新出的 bug、新出的特性都無法得到很好的支持。使用過程中遇到的問題，也只能靠內(nèi)部維護(hù)團(tuán)隊(duì)去修復(fù)，增加了額外的成本。及時(shí)的框架升級(jí)還可以帶來免費(fèi)的速度提升，高版本的 PyTorch 往往匹配更高版本的 CUDA，以及像 CUDA graph 等的一些新特性，可以進(jìn)一步提升訓(xùn)練速度，提升訓(xùn)練效率。

除了算法工程師，框架團(tuán)隊(duì)也是推薦團(tuán)隊(duì)的重要組成部分。公司中的框架團(tuán)隊(duì)會(huì)在選取開源框架之后基于內(nèi)部需求進(jìn)行二次開發(fā)。對于他們來說，一個(gè) PyTorch 的推薦框架會(huì)帶來更簡化的開發(fā)體驗(yàn)。很多傳統(tǒng)的 TensorFlow 推薦框架會(huì)模仿 TF serving 來做一個(gè)基于 C++ session 的擴(kuò)展——這樣的設(shè)計(jì)方案在當(dāng)時(shí)算是非常先進(jìn)的方案——但這使得只改一行代碼也需要完整地編譯整個(gè) TensorFlow，耗時(shí)很長，甚至還要在解決在內(nèi)網(wǎng)下載外部的依賴之類的瑣碎問題，開發(fā)體驗(yàn)不太好。

使用 PyTorch 不會(huì)遇到這樣的問題，因?yàn)?PyTorch 以 Python 哲學(xué)為核心，它希望大家能夠自如地進(jìn)行擴(kuò)展。我們在進(jìn)行二次開發(fā)的時(shí)候，只需要用 pybind11 這樣比較成熟的 Python 庫封裝一下，把我們的庫打包成一個(gè)動(dòng)態(tài)鏈接庫，就可以加載了。這樣自然整體編譯速度會(huì)快很多，同時(shí)學(xué)習(xí)成本也會(huì)低不少。

前面提到 PyTorch 是一個(gè)向后兼容性非常好的框架，這讓維護(hù)團(tuán)隊(duì)不需要去維護(hù)多個(gè)版本，很多共性的問題都可以得到官方的解決，大家就可以專注于特化需求，團(tuán)隊(duì)人員效率就會(huì)有明顯提升。

上面介紹的都是 TorchRec 作為一個(gè) PyTorch 推薦框架的優(yōu)勢，讓我們感到非常開心的是，TorchRec 團(tuán)隊(duì)沒有止步于做一個(gè) PyTorch 推薦框架。他們觀察了現(xiàn)有推薦模型以及硬件的特點(diǎn)，在框架中加入了許多的新特性，使得 TorchRec 相比于傳統(tǒng)的推薦框架有明顯的性能優(yōu)勢。接下來我會(huì)選擇其中的幾個(gè)來進(jìn)行介紹，分別是 GPU embedding，TorchRec 里面優(yōu)秀的 GPU kernel，還有 TorchRec 能夠根據(jù)網(wǎng)絡(luò)通信進(jìn)行的 embedding 劃分。

?首先是 GPU embedding。我們先來回顧一下傳統(tǒng)的推薦系統(tǒng) GPU 訓(xùn)練流程，我們會(huì)把具體的模型放在 GPU worker 上，embedding 存在遠(yuǎn)端 PS 上。每個(gè)迭代步會(huì)先從遠(yuǎn)端 PS 拉取參數(shù)，之后在 GPU 上進(jìn)行模型的前向和反向計(jì)算，把梯度傳回給 PS，在 PS 上進(jìn)行參數(shù)更新。

圖中綠色的部分是在 GPU 上進(jìn)行的操作，紅色的部分是網(wǎng)絡(luò)或者 CPU 上進(jìn)行的。可以看到雖然 GPU 是系統(tǒng)中最昂貴的部分，很多操作卻都沒有放在 GPU 上。?

傳統(tǒng)流程并沒有充分利用好 GPU。同時(shí)，從硬件層面來說，GPU 單卡顯存越來越大，dense 部分模型遠(yuǎn)遠(yuǎn)沒有充分利用 GPU；在英偉達(dá)的不斷優(yōu)化下，NV link 以及 GPU direct RDMA 還讓卡間通信速度越來越快。

GPU embedding 是一個(gè)非常簡單的方案。他直接把 embedding 切分放在 GPU 上——比如單機(jī)上有 8 張卡，我們把 embedding 直接切分為 8 份，每份放在一張卡上——從而保證所有的操作全都留在卡上。GPU 的利用效率就會(huì)有明顯提升，訓(xùn)練速度也會(huì)有質(zhì)的飛躍。如果擔(dān)心 GPU 上面的顯存空間不足，TorchRec 還做了 UVM 的支持，可以提前劃分一部分主機(jī)上的內(nèi)存作為顯存的補(bǔ)充，從而提升單機(jī)內(nèi)部能放下的 embedding 大小。

除去 GPU embedding 以外，TorchRec 還實(shí)現(xiàn)了非常優(yōu)秀的 GPU kernel。這些 kernel 充分利用了最新的硬件特性和 CUDA feature。

舉例來說，假如果要實(shí)現(xiàn)一個(gè) embedding lookup kernel，也就是要從一個(gè)大的 embedding 里面找到一堆 ID 對應(yīng)的 embedding vector，那么普通的實(shí)現(xiàn)里，會(huì)給每個(gè) GPU thread 分配一個(gè) ID，讓他們分別去找對應(yīng)的 embedding。這個(gè)時(shí)候我們要考慮到，GPU 底層是按 warp 進(jìn)行調(diào)度的，一個(gè) warp 里的 32 個(gè) thread 會(huì)一起進(jìn)行顯存讀寫。這意味著，在上述樣流程里，雖然在讀取 ID 時(shí)連續(xù)地訪問了顯存，但后續(xù)的拷貝變成了一個(gè)隨機(jī)讀寫的狀態(tài)。對于硬件來說，隨機(jī)讀寫無法充分利用顯存帶寬，運(yùn)行效率也就不夠高。

TorchRec 則是在每個(gè) thread 讀到 ID 后，利用 shuffle_sync 這樣的 warp primitive，將 ID 廣播至 warp 內(nèi)的所有thread 上，從而讓一個(gè) wrap 里 32 個(gè) thread 去同時(shí)處理同一個(gè) embedding，從而可以進(jìn)行連續(xù)的內(nèi)存讀寫，使得顯存的帶寬利用效率有明顯的提升，讓 kernel 的速度得到數(shù)倍提升。

這個(gè)表是官方測試的 embedding lookup 性能提升。這里 Fused EBC 是優(yōu)化后的kernel，可以看到，不同的設(shè)置情況下 TorchRec 相較于原生的 PyTorch 有數(shù)十倍的性能提升。在 TorchRec 的基礎(chǔ)之上，我們發(fā)現(xiàn)對于 embedding 比較小的情況（小于128），可能有半數(shù)甚至更多的 thread 空閑，所以進(jìn)一步把 warp 內(nèi)的 thread 分組，讓他們同時(shí)去處理多條 embedding。

在我們的改進(jìn)下，小 embedding dim 上 kernel 又有了 10% 到 30% 的提升。這一優(yōu)化也已經(jīng)合入官方 repo。要特別指出的是，TorchRec 的 kernel 放在了 FBGEMM 庫里，有興趣朋友可以去看一看。

最后想介紹一下 TorchRec 的 embedding 劃分機(jī)制。前面提到，GPU embedding 就是把 embedding 切分一下放在卡上，那么怎么分就成了一個(gè)需要考慮的問題。傳統(tǒng)來說有兩種劃分思路，Row wise 和 Column wise。Row wise 是指假如有 2 萬個(gè) feature， 0 號(hào)到第 10000 號(hào)放在卡 1 上，10000 號(hào)到 20000 號(hào)放在卡 2 上，這樣我們在訓(xùn)練的時(shí)候，如果 ID 對應(yīng)卡 1，我們就從卡 1 上拿，對應(yīng)卡 2，就從卡 2 上拿。Row wise 的問題在于，因?yàn)槲覀儾磺宄?10000 號(hào)的通信量和后 10000 號(hào)的是不是差距很大，通信都是不均衡的，無法充分利用網(wǎng)絡(luò)硬件。

Column wise 則是從 embedding 長度角度去劃分。例如 embedding 總長是128，可以前 64 維和后 64 維放在不同的位置，這樣通信會(huì)比較均衡，但是在讀取的時(shí)候，需要和所有的卡或者 PS 通信。

劃分模式上的差別帶來了選型中的 trade-off。傳統(tǒng)的推薦框架會(huì)在設(shè)計(jì)中固定 embedding 的劃分方式，而 TorchRec 則在支持了多種劃分方式——比如 row wise、column wise，甚至 table wise，data parallel——的基礎(chǔ)上，在內(nèi)部提供了如 Planner、Estimator、PerfModel 等可以根據(jù)使用場景的帶寬、顯存、內(nèi)存、模型大小等等參數(shù)自動(dòng)地去計(jì)算劃分的方式的模塊。這樣就可以根據(jù)我們實(shí)際硬件情況去最高效地劃分 embedding，最高效地利用硬件。這些功能大都是在 Python 里面實(shí)現(xiàn)的。方便我們針對內(nèi)部環(huán)境進(jìn)行客制化，從而不費(fèi)力地構(gòu)建出一套最適合于我們內(nèi)部環(huán)境的推薦系統(tǒng)。

二、在百億模型上的實(shí)驗(yàn)效果

在我們的實(shí)驗(yàn)中，對于 DeepFM、DCN 這樣的在標(biāo)準(zhǔn)模，TorchRec 相對于之前的基準(zhǔn)的推薦框架會(huì)有驚人的 10 至 15 倍的性能提升。拿到了這樣的性能收益，讓我們有信心把 TorchRec 推到了業(yè)務(wù)上。

對于微信讀書精排模型，在對齊精度的基礎(chǔ)上，我們發(fā)現(xiàn)在真實(shí)數(shù)據(jù)上有 3 倍左右的性能提升，在假數(shù)據(jù)上甚至有 10 倍左右提升。這里的差別是因?yàn)橛?xùn)練讀取數(shù)據(jù)變成瓶頸了，這方面我們還在做進(jìn)一步的優(yōu)化。

原始方案在千億及更大模型上的不足

前面介紹的基本是百億級(jí)別或者以下的模型，也就是單機(jī)就可以放得下的模型。在把 TorchRec 推到更大的模型的時(shí)候，我們觀察到 TorchRec 的原生設(shè)計(jì)的一些問題。對于大模型來說，TorchRec 的純 GPU embedding 方案需要更多的卡——可能原本 8 張卡的訓(xùn)練速度就可以吃進(jìn)全部數(shù)據(jù)，但是我們要用 16 張卡放下 embedding，這使得好不容易提升上去的 GPU 硬件利用效率又被拖了下來。

而且對于大模型的場景，算法團(tuán)隊(duì)往往會(huì)提出 embedding 的動(dòng)態(tài)增刪需求，例如刪除一周沒有訪問過的 ID。TorchRec 的方案是不支持這樣特性的。還有，超大模型的業(yè)務(wù)一般都會(huì)涉及諸多團(tuán)隊(duì)，遷移基層框架會(huì)遇到很大的阻力。我們需要的支持逐步地漸進(jìn)遷移，而不能讓大家一起放下手頭的工作，那樣的成本過高，風(fēng)險(xiǎn)太大。

根據(jù)上述的需求，我們考慮如何去修改 TorchRec，使得它能夠適應(yīng)超大規(guī)模模型的場景。我們認(rèn)為在超大規(guī)模訓(xùn)練中，仍然需要支持連接遠(yuǎn)程的 PS，因?yàn)檫h(yuǎn)端 CPU PS 已經(jīng)非常成熟了，非常容易支持 embedding 的動(dòng)態(tài)增添。同時(shí)，對于跨團(tuán)隊(duì)的合作，可以用 PS 來隔離開訓(xùn)練和推理，實(shí)現(xiàn)漸進(jìn)的遷移。

那么接下來一個(gè)問題就是該如何引入 PS。如果把 PS 直接連到 GPU embedding 上，每個(gè)迭代步還是要去訪問遠(yuǎn)端的 PS，會(huì)重新使網(wǎng)絡(luò)和 CPU 整體操作的占比提升，GPU 利用率又被拉下來了。

微信團(tuán)隊(duì)的 dynamic embedding 如何解決問題

這個(gè)時(shí)候我們發(fā)現(xiàn)單位時(shí)間內(nèi)數(shù)據(jù)中的新 ID 實(shí)際上只占總數(shù)據(jù)中很少的一部分， HugeCTR 發(fā)表論文中也提到相似的結(jié)論：只有一小部分的 ID 會(huì)被頻繁訪問。由此，我們想到先正常使用 GPU embedding 進(jìn)行訓(xùn)練，在顯存放滿時(shí)再將 ID 批量驅(qū)逐至 PS。

根據(jù)這樣的一個(gè)思路，假如 GPU embedding 里面只能存下 n 個(gè) ID，而總 ID 有 N 個(gè)，甚至無窮多個(gè)。可以將全局的 ID 按順序映射到 0、1、2、3…，并把把映射關(guān)系存在一個(gè)叫 ID transform 的結(jié)構(gòu)中，讓 GPU embedding 利用映射的結(jié)果進(jìn)行正常的訓(xùn)練。當(dāng) GPU embedding 放滿了，也就是 ID transformer 中 n 對映射的時(shí)候，再批量驅(qū)逐 ID 至 PS。

在這種設(shè)計(jì)下，可以使得 PS 很少介入，只有在驅(qū)逐時(shí)才需要 GPU worker 和 PS 通信。

除此之外，這樣的設(shè)計(jì)中 PS 只需要作為 KV，不需要支持參數(shù)更新，也就不需要實(shí)現(xiàn)優(yōu)化器相關(guān)的操作，從而讓 PS 團(tuán)隊(duì)專注于存儲(chǔ)相關(guān)的工作。我們也支持實(shí)現(xiàn)了任意 KV 存儲(chǔ)的插件，在開源版本中更是內(nèi)置了 Redis 插件，讓 Redis 也可以作為一個(gè) PS 來使用。

下面介紹一些 dynamic embedding 中的設(shè)計(jì)細(xì)節(jié)。我們實(shí)現(xiàn)的最簡基礎(chǔ)的 ID Transformer，其實(shí)也就是用一個(gè)哈希表，使用的是 PyTorch 里高性能的 ska::flat_hash_map。

ID Transformer 作為流程中僅有的 CPU 操作，對性能要求可能會(huì)比較高，所以我們還實(shí)現(xiàn)了一個(gè)高性能的版本，以 L1 cacheline 為單位存儲(chǔ)，從而進(jìn)一步提升內(nèi)存的訪存效率。

另外，對于驅(qū)逐方案，我們希望在不增加內(nèi)存緩存壓力的情況下，高效地融合 LRU 和 LFU。受到 Redis 的 LFU 方案的啟發(fā)，我們設(shè)計(jì)了一種概率的算法：只存儲(chǔ) ID 訪問頻數(shù)的指數(shù)。比如訪問了 32 次即存儲(chǔ) 5。在更新頻數(shù)時(shí)，如果又訪問到這個(gè) ID，就生成 5 位的隨機(jī)數(shù)，如果在 5 位全為 0，也就是發(fā)生了概率為 1/ 32 的事件，我們就增加頻數(shù)指數(shù)為 6。通過這樣的概率算法，就可以把 LRU 和LFU 的頻數(shù)放到 uint32 里面，在不提高訪存壓力的情況下融合了 LRU 和 LFU。

最后來簡單介紹一下我們的多卡方案。我們目前是將所有卡的數(shù)據(jù)都先 gather 到卡一的 ID Transformer 上，之后再 broadcast 回去。因?yàn)槲覀儗?shí)現(xiàn)的 ID Transformer 的性能非常高，而且可以和 GPU 計(jì)算 Pipeline 起來，不會(huì)成為具體的性能瓶頸。

以上就是 dynamic embedding 在設(shè)計(jì)上一些想法。在我們內(nèi)部的一個(gè)萬億級(jí)的業(yè)務(wù)上，在對齊精度情況下，dynamic embedding 方案相對于我們內(nèi)部原始的 GPU Tensorflow 框架有 3 倍左右的性能提升。相比于 TF 優(yōu)化版也仍然有 50% 以上的性能優(yōu)勢。

最后推薦大家去試用一下 Torchrec。對于相對較小的業(yè)務(wù)，比如百億下的業(yè)務(wù)，推薦大家直接使用原生的 TorchRec：即插即用，不需要任何的二次開發(fā)，性能可以得到成倍的提升。對于極大的業(yè)務(wù)，則推薦大家嘗試配合我們合進(jìn) TorchRec 的 dynamic embedding，一方面方便連接內(nèi)部的 PS，另一方面也支持 embedding 的擴(kuò)展和漸進(jìn)遷移，同時(shí)還是可以獲得一定的性能提升。