本文將介紹微信基于 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/ target=_blank class=infotextkey>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 的擴展和漸進遷移，同時還是可以獲得一定的性能提升。
這里是我們已經對齊的一些精度的模型和已有的應用場景，有興趣的朋友可以去試一下。
