微信基於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 的擴充套件和漸進遷移，同時還是可以獲得一定的效能提升。

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