微信基於 PyTorch 的大規模推薦系統訓練實踐

本文將介紹微信基於 PyTorch 進行的大規模推薦系統訓練。推薦系​​統和其它一些深度學習領域不同，仍在使用 Tensorflow 作為訓練框架，被廣大開發者詬病。雖然也有使用 PyTorch 進行推薦訓練的一些實踐，但規模較小，也沒有實際的業務驗證，很難推動業務嚐鮮。

2022 年 2 月，PyTorch 團隊推出了官方推薦函式庫 TorchRec。我們團隊在 5 月開始在內部業務上嘗試 TorchRec，並且與 TorchRec 團隊展開了一系列的合作。在幾個月的試用過程中，我們體會到 TorchRec 非常多的優點，也感受到 TorchRec 在超大規模模型上仍有一些不足之處。針對這些不足，我們設計了擴充功能來填補它的問題。在 2022 年 9 月，我們設計的擴充功能 dynamic embedding 已經正式合進了 TorchRec 的主分支，目前仍在與官方團隊持續優化。

#一、TorchRec 可以為我們帶來什麼

我們先來聊聊TorchRec 可以帶給我們什麼？我們都知道推薦系統往往和公司的現金流直接掛鉤，試誤成本非常高，所以大家需要的是經過了業務測試的框架。這也是為什麼之前的一些基於 PyTorch 的推薦框架都未曾被廣泛應用過。而 TorchRec 作為一個官方的推薦框架，在 2022 年 1 月推出之時，Meta就已經利用它在 Instagram Reels 業務上順利訓練並上線了一個 1250 億參數的模型，成為了一個經過業務測試的 PyTorch 框架。有了 Instagram 這樣一個大業務的支撐，讓我們有了更多信心，終於可以去理性地考慮一個基於 PyTorch 的推薦框架有什麼樣的優勢了。

對於團隊中的不同成員，TorchRec 有不同的好處。首先，對於團隊中佔絕大多數的演算法工程師而言，PyTorch 推薦框架讓大家終於可以享受像 CV、NLP 工程師體會到的那種更人性化的動態圖和調試的體驗。

另外，PyTorch 極佳的相容性－一個基於PyTorch1.8 所做的模型，不需要改一行程式碼就可以在最新版本1.13 上執行——讓演算法工程師終於可以放心地升級框架，從而享受最新的框架功能和更優秀的效能。而反觀一些基於 TensorFlow 的建議框架，往往被卡在 TensorFlow 的某一個版本上，例如許多團隊可能還在使用基於 TensorFlow 1.x 的內部框架。 TensorFlow 1.x 在 2021 年 1 月就已經停止維護了，這意味著在近兩年的時間內，所有新出的 bug、新出的特性都無法得到很好的支援。使用過程中遇到的問題，也只能靠內部維護團隊來修復，增加了額外的成本。及時的框架升級還可以帶來免費的速度提升，高版本的 PyTorch 往往匹配更高版本的  CUDA，以及像 CUDA graph 等的一些新特性，可以進一步提升訓練速度，提升訓練效率。

除了演算法工程師，框架團隊也是推薦團隊的重要組成部分。公司中的框架團隊會在選取開源框架之後基於內部需求進行二次開發。對他們來說，一個 PyTorch 的推薦框架會帶來更簡化的開發體驗。許多傳統的TensorFlow 推薦框架會模仿TF serving 來做一個基於C session 的擴充功能——這樣的設計方案在當時算是非常先進的方案——但這使得只改一行程式碼也需要完整地編譯整個TensorFlow，耗時很長，甚至還要在解決在內網下載外部的依賴之類的瑣碎問題，開發體驗不太好。

使用 PyTorch 不會遇到這樣的問題，因為 PyTorch 以 Python 哲學為核心，它希望大家能夠自如地擴展。我們在進行二次開發的時候，只需要用 pybind11 這樣比較成熟的 Python 函式庫封裝一下，把我們的函式庫打包成一個動態連結函式庫，就可以載入了。這樣自然整體編譯速度會快很多，同時學習成本也會低不少。

前面提到PyTorch 是一個向後相容性非常好的框架，這讓維護團隊不需要去維護多個版本，很多共通性的問題都可以得到官方的解決，大家就可以專注於特化需求，團隊人員效率就會明顯提升。

上面介紹的都是TorchRec 作為一個PyTorch 推薦框架的優勢，讓我們感到非常開心的是，TorchRec 團隊沒有止步於做一個PyTorch 推薦框架。他們觀察了現有推薦模型以及硬體的特點，在框架中加入了許多的新特性，使得 TorchRec 相比於傳統的推薦框架有明顯的效能優勢。接下來我會選擇其中的幾個來進行介紹，分別是 GPU embedding，TorchRec 裡面優秀的 GPU kernel，還有 TorchRec 能夠根據網路通訊進行的 embedding 劃分。

首先是 GPU embedding。我們先來回顧傳統的推薦系統 GPU 訓練流程，我們會把具體的模型放在 GPU worker 上，embedding 存在遠端 PS 上。每個迭代步會先從遠端 PS 拉取參數，之後在 GPU 上進行模型的前向和反向計算，把梯度傳回給 PS，在 PS 上進行參數更新。

圖中綠色的部分是在 GPU 上進行的操作，紅色的部分是網路或 CPU 上進行的。可以看到雖然 GPU 是系統中最昂貴的部分，但很多操作都沒有放在 GPU 上。

傳統流程並沒有充分利用好 GPU。同時，從硬體層面來說，GPU 單卡顯存越來越大，dense 部分模型遠遠沒有充分利用GPU；在英偉達的不斷優化下，NV link 以及GPU direct RDMA 也讓卡間通訊速度越來越快。

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

除去 GPU embedding 以外，TorchRec 也實作了非常優秀的 GPU kernel。這些 kernel 充分利用了最新的硬體特性和 CUDA feature。

舉例來說，假如果要實作一個embedding lookup kernel，也就是要從一個大的embedding 裡面找到一堆ID 對應的embedding vector，那麼普通的實作裡，會給每個GPU thread 分配一個ID，讓他們分別去找對應的embedding。這時候我們要考慮到，GPU 底層是按 warp 進行調度的，一個 warp 裡的 32 個 thread 會一起進行顯存讀寫。這意味著，在上述樣流程裡，雖然在讀取 ID 時連續地存取了顯存，但後續的拷貝變成了一個隨機讀寫的狀態。對於硬體來說，隨機讀寫無法充分利用顯存頻寬，運作效率也不夠高。

TorchRec 則是在每個thread 讀到ID 後，利用shuffle_sync 這樣的warp primitive，將ID 廣播至warp 內的所有thread 上，從而讓一個wrap 裡32 個thread 去同時處理同一個embedding，從而可以進行連續的內存讀寫，使得顯存的頻寬利用效率有明顯的提升，讓kernel 的速度得到數倍提升。

這個表是官方測試的 embedding lookup 效能提升。這裡 Fused EBC 是優化後的kernel，可以看到，不同的設定情況下 TorchRec 相較於原生的 PyTorch 有數十倍的效能提升。在 TorchRec 的基礎之上，我們發現對於 embedding 比較小的情況（小於128），可能有半數甚至更多的 thread 空閒，所以進一步把 warp 內的 thread 分組，讓他們同時去處理多條 embedding。

#在我們的改進下，小embedding dim 上kernel 又有了10% 到30% 的提升。這項優化也已經合入官方 repo。要特別指出的是，TorchRec 的 kernel 放在了 FBGEMM 庫裡，有興趣朋友可以去看一看。

最後想介紹一下 TorchRec 的 embedding 分割機制。前面提到，GPU embedding 就是把 embedding 切分放在卡片上，那麼怎麼分就成了一個需要考慮的問題。傳統來說有兩種劃分思路，Row wise 和 Column wise。 Row wise 是指假如有2 萬個feature， 0 號到第10000 號放在卡1 上，10000 號到20000 號放在卡2 上，這樣我們在訓練的時候，如果ID 對應卡1 上，我們就從卡1 上拿，對應卡2，就從卡2 拿。 Row wise 的問題在於，因為我們不清楚前 10000 號的通訊量和後 10000 號的是不是差距很大，所以通訊都是不均衡的，無法充分利用網路硬體。

Column wise 則是從 embedding 長度角度去分割。例如 embedding 總長是128，可以前 64 維和後 64 維放在不同的位置，這樣通訊會比較均衡，但是在讀取的時候，需要和所有的卡片或 PS 通訊。

分割模式上的差異帶來了選型中的 trade-off。傳統的推薦框架會在設計中固定embedding 的劃分方式，而TorchRec 則在支持了多種劃分方式——比如row wise、column wise，甚至table wise，data parallel——的基礎上，在內部提供瞭如Planner、Estimator、PerfModel 等可以根據使用場景的頻寬、顯存、記憶體、模型大小等等參數自動地去計算劃分的方式的模組。這樣就可以根據我們實際硬體情況去最有效率地劃分 embedding，最有效率地利用硬體。這些功能大都是在 Python 裡面實現的。方便我們針對內部環境進行客製化，從而不費力地建構出一套最適合我們內部環境的推薦系統。

二、在百億模型上的實驗效果

在在我們的實驗中，對於DeepFM、DCN 這樣的在標準模，TorchRec 相對於之前的基準的建議框架會有驚人的10 至15 倍的性能提升。拿到了這樣的效能效益，讓我們有信心把 TorchRec 推到了業務上。

對於微信閱讀精排模型，在對齊精確度的基礎上，我們發現在真實資料上有3 倍左右的效能提升，在假數據上甚至有10 倍左右提升。這裡的差異是因為訓練讀取資料變成瓶頸了，這方面我們還在做進一步的最佳化。

03

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

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

而且對於大模型的場景，演算法團隊往往會提出 embedding 的動態增刪需求，例如刪除一週沒有訪問過的 ID。 TorchRec 的方案是不支援這樣特性的。還有，超大模型的業務通常會涉及許多團隊，遷移基層框架會遇到很大的阻力。我們需要的支援逐步地漸進遷移，而不能讓大家一起放下手邊的工作，那樣的成本過高，風險太大。

根據上述的需求，我們考慮如何去修改 TorchRec，使得它能夠適應超大規模模型的場景。我們認為在超大規模訓練中，仍然需要支援連接遠端的 PS，因為遠端 CPU PS 已經非常成熟了，非常容易支援 embedding 的動態增添。同時，對於跨團隊的合作，可以用 PS 來隔離開訓練和推理，達到漸進的遷移。

那麼接下來一個問題就是該如何引入 PS。如果把 PS 直接連到 GPU embedding 上，每個迭代步還是要去存取遠端的 PS，會重新讓網路和 CPU 整體操作的佔比提升，GPU 使用率又被拉下來了。

#

04

#微信團隊的dynamic embedding 如何解決問題

#這時候我們發現單位時間內資料中的新ID 實際上只佔總資料中很少的一部分， HugeCTR 發表論文中也提到相似的結論：只有一小部分的ID 會被頻繁存取。由此，我們想到先正常使用 GPU embedding 進行訓練，在顯存放滿時再將 ID 批次驅逐至 PS。

根據這樣的一個思路，假如GPU embedding 裡面只能存下n 個ID，而總ID 有N個，甚至無限多個。可以將全域的 ID 依序對應到 0、1、2、3…，並且把把映射關係存在一個叫做 ID transform 的結構中，讓 GPU embedding 利用映射的結果進行正常的訓練。當 GPU embedding 放滿了，也就是 ID transformer 中 n 對映射的時候，再批次驅逐 ID 至 PS。

#在這個設計下，可以讓PS 很少介入，只有在驅逐時才需要GPU worker 和PS 通訊。

除此之外，這樣的設計中PS 只需要作為KV，不需要支援參數更新，也就不需要實現優化器相關的操作，讓PS 團隊專注於儲存相關的工作。我們也支援實作了任意 KV 儲存的插件，在開源版本中更是內建了 Redis 插件，讓 Redis 也可以當作一個 PS 來使用。

下面介紹一些 dynamic embedding 中的設計細節。我們實現的最簡基礎的 ID Transformer，其實也就是用一個哈希表，使用的是 PyTorch 裡高效能的 ska::flat_hash_map。

 ID Transformer 作為流程中僅有的CPU 操作，對效能需求可能會比較高，所以我們也實現了一個高效能的版本，以L1 cacheline 為單位存儲，從而進一步提升記憶體的存取效率。

另外，對於驅逐方案，我們希望在不增加記憶體快取壓力的情況下，高效地融合LRU 和LFU。受到 Redis 的 LFU 方案的啟發，我們設計了一個機率的演算法：只儲存 ID 存取頻數的指數。例如訪問了 32 次即儲存 5。在更新頻數時，如果又訪問到這個 ID，就產生 5 位的隨機數，如果在 5 位全為 0，也就是發生了機率為 1/ 32 的事件，我們就增加頻數指數為 6。透過這樣的機率演算法，就可以把 LRU 和LFU 的頻數放到 uint32 裡面，在不提高訪存壓力的情況下融合了 LRU 和 LFU。

#

最後來簡單介紹一下我們的多卡方案。我們目前是將所有卡片的資料都先 gather 到卡一的 ID Transformer 上，之後再 broadcast 回去。因為我們實現的 ID Transformer 的效能非常高，而且可以和 GPU 計算 Pipeline 起來，不會成為具體的效能瓶頸。

以上就是 dynamic embedding 在設計上一些想法。在我們內部的一個萬億級的業務上，在對齊精度情況下，dynamic embedding 方案相對於我們內部原始的 GPU Tensorflow 框架有 3 倍左右的效能提升。相較於 TF 優化版也仍有 50% 以上的效能優勢。

最後推薦大家去試用 Torchrec。對於相對較小的業務，例如百億下的業務，推薦大家直接使用原生的 TorchRec：即插即用，不需要任何的二次開發，性能可以得到成倍的提升。對於極大的業務，則推薦大家嘗試配合我們合進 TorchRec 的 dynamic embedding，一方面方便連接內部的 PS，另一方面也支援 embedding 的擴展和漸進遷移，同時還是可以獲得一定的性能提升。

這裡是我們已經對齊的一些精確度的模型和已有的應用場景，有興趣的朋友可以去試試。

#

以上是微信基於 PyTorch 的大規模推薦系統訓練實踐的詳細內容。更多資訊請關注PHP中文網其他相關文章！