TensorFlow 在推薦系統中的分佈式訓練優化實踐

美團內部深度定製的 TensorFlow 版本，基於原生 TensorFlow 1.x 架構與接口，從大規模稀疏參數的支持、訓練模式、分佈式通信優化、流水線優化、算子優化融合等多維度進行了深度優化。在推薦系統場景中，分佈式擴展性提升 10 倍以上，單位算力性能也有顯著提升，並在美團內部業務中大量使用，本文介紹了相關的優化與實踐工作。

1 背景

TensorFlow（下文簡稱 TF）是谷歌推出的一個開源深度學習框架，在美團推薦系統場景中得到了廣泛的使用。但 TensorFlow 官方版本對工業級場景的支持，目前做得並不是特別的完善。美團在大規模生產落地的過程中，遇到了以下幾方面的挑戰：

所有參數都是用 Variable 表達，對於百億以上的稀疏參數開闢了大量的內存，造成了資源的浪費；
只支持百級別 Worker 的分佈式擴展，對上千 Worker 的擴展性較差；
由於不支持大規模稀疏參數動態添加、刪除，增量導出，導致無法支持 Online Learning；
大規模集羣運行時，會遇到慢機和宕機；由於框架層不能處理，導會致任務運行異常。

以上這些問題，並不是 TensorFlow 設計的問題，更多是底層實現的問題。考慮到美團大量業務的使用習慣以及社區的兼容性，我們基於原生 TensorFlow 1.x 架構與接口，從大規模稀疏參數的支持、訓練模式、分佈式通信優化、流水線優化、算子優化融合等多維度進行了深度定製，從而解決了該場景的核心痛點問題。

首先新系統在支持能力層面，目前可以做到千億參數模型，上千 Worker 分佈式訓練的近線性加速，全年樣本數據能夠 1 天內完成訓練，並支持 Online Learning 的能力。同時，新系統的各種架構和接口更加友好，美團內部包括美團外賣、美團優選、美團搜索、廣告平臺、大衆點評 Feeds 等業務部門都在使用。本文將重點介紹大規模分佈式訓練優化的工作，希望對大家能夠有所幫助或啓發。

2 大規模訓練優化挑戰

2.1 業務迭代帶來的挑戰

隨着美團業務的發展，推薦系統模型的規模和複雜度也在快速增長，具體表現如下：

訓練數據：訓練樣本從到百億增長到千億，增長了近 10 倍。
稀疏參數：個數從幾百到幾千，也增長了近 10 倍；總參數量從幾億增長到百億，增長了 10~20 倍。
模型複雜度：越來越複雜，模型單步計算時間增長 10 倍以上。

對於大流量業務，一次訓練實驗，從幾個小時增長到了幾天，而此場景一次實驗保持在 1 天之內是基本的需求。

2.2 系統負載分析

2.2.1 問題分析工具鏈

TensorFlow 是一個非常龐大的開源項目，代碼有幾百萬行之多，原生系統的監控指標太粗，且不支持全局的監控，如果要定位一些複雜的性能瓶頸點，就比較困難。我們基於美團已經開源的監控系統 CAT[2]，構建了 TensorFlow 的細粒度監控鏈路（如下圖 1 所示），可以精準定位到性能的瓶頸問題。

圖 1 TensorFlow PS 架構全鏈路監控

同時，在性能優化的過程中，會涉及到大量的性能測試和結果分析，這也是一個非常耗費人力的工作。我們抽象了一套自動化的實驗框架（如下圖 2 所示），可以自動化、多輪次地進行實驗，並自動採集各類監控指標，然後生成報告。

圖 2 自動化實驗框架

2.2.2 業務視角的負載分析

在推薦系統場景中，我們使用了 TensorFlow Parameter Server[3]（簡稱 PS）異步訓練模式來支持業務分佈式訓練需求。對於這套架構，上述的業務變化會帶來什麼樣的負載變化？如下圖 3 所示：

圖 3 TensorFlow PS 架構大規模訓練負載分析

總結來看，主要包括通信壓力、PS 併發壓力、Worker 計算壓力。對於分佈式系統來說，通常是通過橫向擴展來解決負載問題。雖然看來起可以解決問題，但從實驗結果來看，當 PS 擴展到一定數量後，單步訓練時間反而會增加，如下圖 4 所示：

圖 4 擴展 PS 提升訓練性能實驗

導致這種結果的核心原因是：Worker 單步訓練需要和所有的 PS 通信同步完成，每增加 1 個 PS 要增加 N 條通信鏈路，這大大增加了鏈路延遲（如下圖 5 所示）。而一次訓練要執行上百萬、上千萬步訓練。最終導致鏈路延遲超過了加 PS 算力併發的收益。

圖 5 增加 PS 帶來的鏈路開銷

而對於這個系統，優化的核心難點在於：如何在有限的 PS 實例下，進行分佈式計算的優化。

3 優化實踐

3.1 大規模稀疏參數介紹

對於推薦系統模型，絕大多數參數都是稀疏參數，而對稀疏參數來說有一個非常重要的操作是 Embedding，這個操作通常也是負載最重的，也是後續優化的重點。由於我們對稀疏參數進行了重新定義，後續的優化也基於此之上，所以我們先介紹一下這部分的工作。

在原生的 TensorFlow 中構建 Embedding 模塊，用戶需要首先創建一個足夠裝得下所有稀疏參數的 Variable，然後在這個 Variable 上進行 Embedding 的學習。然而，使用 Variable 來進行 Embedding 訓練存在很多弊端：

Variable 的大小必須提前設定好，對於百億千億的場景，該設定會帶來巨大的空間浪費；
訓練速度慢，無法針對稀疏模型進行定製優化。

我們首先解決了有無的問題，使用 HashTable 來替代 Variable，將稀疏特徵 ID 作爲 Key，Embedding 向量作爲 Value。相比原生使用 Variable 進行 Embedding 的方式，具備以下的優勢：

HashTable 的大小可以在訓練過程中自動伸縮，避免了開闢冗餘的存儲空間，同時用戶無需關注申請大小，從而降低了使用成本。
針對 HashTable 方案實施了一系列定製優化，訓練速度相比 Variable 有了很大的提高，可以進行千億規模模型的訓練，擴展性較好。
得益於稀疏參數的動態伸縮，我們在此基礎上支持了 Online Learning。
API 設計上保持與社區版本兼容，在使用上幾乎與原生 Variable 一致，對接成本極低。

簡化版的基於 PS 架構的實現示意如下圖 6 所示：

圖 6 支撐大規模稀疏參數的 HashTable 方案

核心流程大致可以分爲以下幾步：

稀疏特徵 ID（通常我們會提前完成統一編碼的工作）進入 Embedding 模塊，藉助 TensorFlow 搭建的 Send-Recv 機制，這些稀疏特徵 ID 被拉取到 PS 端，PS 端上的 Lookup 等算子會實際從底層 HashTable 中查詢並組裝 Embedding 向量。
上述 Embedding 向量被 Worker 拉回進行後續訓練，並通過反向傳播計算出這部分參數的梯度，這些梯度進一步被位於 PS 端的優化器拉回。
PS 端的優化器首先調用 Find 算子，從 HashTable 獲取到梯度對應的原始稀疏參數向量和相應的優化器參數，最終通過優化算法，完成對 Embedding 向量和優化器參數的更新計算，再通過 Insert 算子插入 HashTable 中。

3.2 分佈式負載均衡優化

這部分優化，是分佈式計算的經典優化方向。PS 架構是一個典型的 “水桶模型”，爲了完成一步訓練，Worker 端需要和所有 PS 完成交互，因此 PS 之間的平衡就顯得非常重要。但是在實踐中，我們發現多個 PS 的耗時並不均衡，其中的原因，既包括 TensorFlow PS 架構簡單的切圖邏輯（Round-Robin）帶來的負載不均衡，也有異構機器導致的不均衡。

對於推薦模型來說，我們的主要優化策略是，把所有稀疏參數和大的稠密參數自動、均勻的切分到每個 PS 上，可以解決大多數這類問題。而在實踐過程中，我們也發現一個比較難排查的問題：原生 Adam 優化器，實現導致 PS 負載不均衡。下面會詳細介紹一下。

在 Adam 優化器中，它的參數優化過程需要兩個β參與計算，在原生 TensorFlow 的實現中，這兩個β是所有需要此優化器進行優化的 Variabl（或 HashTable）所共享的，並且會與第一個 Variable（名字字典序）落在同一個 PS 上面，這會帶來一個問題：每個優化器只擁有一個β和一個β，且僅位於某個 PS 上。因此，在參數優化的過程中，該 PS 會承受遠高於其他 PS 的請求，從而導致該 PS 成爲性能瓶頸。

圖 7 Adam 優化算法

但是通過觀察 Adam 的優化算法，我們可以看到β和β都是常量，且藍色高亮的部分都是相對獨立的計算過程，各個 PS 之間可以獨立完成。基於這樣的發現，優化的方法也就非常直觀了，我們爲每一個 PS 上的 Adam 優化器冗餘創建了β參數，並在本地計算 t 和 alpha 值，去除了因此負載不均導致的 PS 熱點問題。

該優化所帶來的提升具備普適性且效果明顯，在美團內部某業務模型上，通過β熱點去除可以帶來 9% 左右的性能提升。此外，由於擺脫了對β的全局依賴，該優化還能提高 PS 架構的可擴展性，在擴增 Worker 數量的時候相比之前會帶來更好的加速比。

3.3 通信優化

通過 2.2 章節的分析可知，系統的通信壓力也非常大，我們主要基於 RDMA 做了通信優化的工作。首先簡單介紹一下 RDMA，相比較於傳統基於套接字 TCP/IP 協議棧的通信過程，RDMA 具有零拷貝、內核旁路的優勢，不僅降低了網絡的延遲，同時也降低了 CPU 的佔用率，RDMA 更適合深度學習模型的相關通信過程。

RDMA 主要包括三種協議 Infiniband、RoCE(V1, V2)、iWARP。在美團內部的深度學習場景中，RDMA 通信協議使用的是 RoCE V2 協議。目前在深度學習訓練領域，尤其是在稠密模型訓練場景（NLP、CV 等），RDMA 已經是大規模分佈式訓練的標配。然而，在大規模稀疏模型的訓練中，開源系統對於 RDMA 的支持非常有限，TensorFlow Verbs[4] 通信模塊已經很長時間沒有更新了，通信效果也並不理想，我們基於此之上進行了很多的改進工作。

經過優化後的版本，在 1TB Click Logs[5] 公開數據集、DLRM[6] 模型、100 個 Worker 以上的訓練，性能提升了 20%~40%。在美團的多個業務模型上，對比 TensorFlow Seastar[7] 改造的通信層實現也有 10%~60% 的速度提升。同時也把我們的工作回饋給了社區。

3.3.1 Memory Registration 優化

RDMA 有三種數據傳輸的方式 SEND/RECV、WRITE、READ，其中 WRITE、READ 類似於數據發送方直接在遠程 Memory 進行讀寫，Receiver 無法感知，WRITE 和 READ 適用於批量數據傳輸。在 TensorFlow 內部，基於 RDMA 的數據傳輸方式使用的是 WRITE 單邊通信模式。

圖 8 RDMA 傳輸方式

在 RDMA 傳輸數據時，需要提前開闢內存空間並將其註冊到網卡設備上（Memory Registration 過程，下稱 MR），使得這片空間可以被網卡直接操作。開闢新的內存並註冊到設備上，整個過程是比較耗時的。下圖 9 展示了不同大小的內存綁定到網卡設備上的耗時，可以看到隨着註冊內存的增大，綁定 MR 的耗時迅速增加。

圖 9 MR 過程開銷

社區版 Tensorflow RDMA 實現，Tensor 創建依舊沿用了統一的 BFC Allocator，並將所有創建的 Tensor 都註冊到 MR 上。正如上面所提到的，MR 的註冊綁定具有性能開銷，高頻、大空間的 MR 註冊會帶來顯著的性能下降。而訓練過程中的 Tensor，只有那些涉及到跨節點通信的 Tensor 有必要進行 MR，其餘 Tensor 並不需要註冊到 MR。因此，優化的方法也就比較直接了，我們識別並管理那些通信 Tensor，僅對這些跨節點通信的 Tensor 進行 MR 註冊就好了。

3.3.2 RDMA 靜態分配器

RDMA 靜態分配器是上一個 MR 註冊優化的延伸。通過 Memory Registration 優化，去除非傳輸 Tensor 的 MR 註冊，我們降低了 MR 註冊數量。但是在稀疏場景大規模的訓練下，並行訓練的 Worker 常有幾百上千個，這會帶來新的問題：

PS 架構中的 PS 和 Worker 互爲 Client-Server，這裏以 PS 端爲例，當 Worker 數目增加到上千個時，Worker 數目的增多，造成 PS 端 MR 註冊頻次還是非常高，增加了內存分配註冊的耗時。
由於稀疏場景不同 Step 之間同一個算子輸出 Tensor 的形狀可能發生變化，導致了創建的 MR 可複用性較差，帶來了較高的內存碎片和重複註冊 MR 開銷。

針對上面的問題，我們引入了 MR 靜態分配器的策略。

圖 10 MR 靜態分配器

這裏核心的設計思路爲：

雖然稀疏場景同一個算子輸出 Tensor 的 Shape 存在變化的可能，但是整體變化幅度可控，通過監控與分析，是可以找到一個較爲穩定的內存大小，滿足多 Step 間 Tensor 的存儲需求。
基於上面的信息，我們修改了原有逐 Tensor(Request) 的 MR 申請策略，通過一次性預申請一塊較大的空間並註冊到網卡端，後續通過自己維護的分配策略進行空間的分配，大大降低了 MR 申請的頻率，絕大多數情況下，訓練全過程中只需要一次 MR 註冊申請即可。
我們引入了一種簡單的交換協議，將傳輸 Tensor 的 Shape，Data 打包到一起，寫到 Client 端。Client 端根據協議，解析出 Tensor 大小，並最終讀取 Data，避免了原生實現中因 Tensor 的 Shape 變化而產生的多次協商過程。

圖 11 MR 靜態分配器構造流程

具體到實現中，我們引入了 Allocation Analysis 模塊，在訓練開始的一段時間，我們會對分配的歷史數據進行分析，以得到一個實際預開闢 MR 大小以及各個 Tensor 的預留空間大小。然後我們會暫停訓練的進程，啓動 Allocator 的構造過程，包括 MR 的創建以及通信雙端的信息同步。利用相關信息構造 MR Info Map，這個 Map 的 Key 是傳輸 Tensor 的唯一標記（ParsedKey，計算圖切圖時確定），Info 結構體中包含了本地地址指針、offset 大小、ibv_send_wr 相關信息等。然後恢復訓練，後續 Tensor 的傳輸就可以使用靜態開闢好的 MR 進行收發，也免去了因 Shape 變化而產生的多次協商過程。

3.3.3 Multi RequestBuffer 與 CQ 負載均衡

TensorFlow 社區版的 RDMA 通信過程，不僅僅包含上面 Tensor 數據的發送和接收過程，還包括傳輸相關的控制消息的發送和接收過程，控制消息的發送和接收過程同樣是使用了 ibv_post_send 和 ibv_post_recv 原語。原生的控制流實現存在一些瓶頸，在大規模訓練時會限制控制流的吞吐，進而影響數據收發的效率。具體體現在：

請求的發送通過同一片 RequestBuffer 內存進行寫出，多個 Client 的請求均依賴這一片 Buffer，也就導致到控制流信息實際是串行發送的，只有等到對端的 Ack 信息後，纔可以下一個 Request 的寫出，限制了請求的發送吞吐。
在 Client 端需要輪詢 RDMA Completion Queue 來獲得請求的到達，以及相關狀態的變更。原生實現僅有一個 Completion Queue，單線程進行輪詢處理，在大規模分佈式訓練中，限制了應答的效率。

針對上面的問題，我們採用了 Multi RequestBuffer 與 CQ 負載均衡優化，破除了在請求發送和請求應答環節可能存在的吞吐瓶頸。

3.3.4 Send-Driven & Rendezvous-Bypass

對於 Tensorflow PS 架構熟悉的同學會瞭解，一整張計算圖被切割爲 Worker 端和 PS 端後，爲了使兩張計算圖能夠彼此交換數據，建立了基於 Rendezvous（匯合點）機制的異步數據交換模式。如下圖 12 所示：

圖 12 TensoFlow 切圖之 Send-Recv 對添加

基於上圖的切圖邏輯，Recv 算子代表着這一側計算圖有 Tensor 的需求，而 Tensor 的生產者則位於與之配對的另一設備上的 Send 算子背後。

在具體實現上，Tensorflow 實現了 Recv-Driven 的數據交換模式，如上圖所示，位於 DeviceA 和 DeviceB 的兩張計算圖會異步併發的執行，位於 DeviceB 的 Recv 執行時會發起一條 RPC 請求發往 DeviceA，DeviceA 收到請求後，會將請求路由到 Rendezvous 中，如果在當中發現所需要的數據已經生產好，並被 Send 算子註冊了進來，那麼就地獲取數據，返回給 DeviceB；如果此時數據還沒有生產好，則將來自於 DeviceB 的 Recv 請求註冊在 Rendezvous 中，等待後續 DeviceA 生產好後，由 Send 算子發送過來，找到註冊的 Recv，觸發回調，返回數據給 DeviceB。

我們看到，匯合點機制優雅地解決了生產者消費者節奏不同情況下數據交換的問題。不過 Recv-Driven 的模式也引入了兩個潛在的問題：

據我們的觀察，在實際業務模型中，在 Rendezvous 中 Recv 算子等待 Send 算子的比例和 Send 算子等待 Recv 算子的比例相當，也就是說對於 Send 等到 Recv 的數據，在 Send 準備好的那一剎那就可以發給對端，但是由於機制實現問題，還是等待 Recv 算子過來，纔將數據拉取回去，通信過程耗時較長。
Rendezvous 作爲一個數據交換的熱點，它內部的邏輯開銷並不低。

針對上面提到的問題，我們在 RDMA 上實現了另外一種數據交換的模式，叫做 Send-Driven 模式。與 Recv-Driven 模式相對，顧名思義就是有 Send 算子直接將數據寫到 Recv 端，Recv 端接收數據並註冊到本地 Rendezvous 中，Recv 算子直接從本地的 Rendezvous 中獲取數據。具體流程如下圖 13 所示：

圖 13 原生的 Recv-Driven 與補充的 Send-Driven 機制

從圖中可以看到，相較於 Recv-Driven 模式，Send-Driven 模式的通信流程得到了比較大的簡化，另外在數據 ready 後立即發送的特性，跳過了一側的 Rendezvous，並且對於生產者先於消費者的情況，可以加快消費端數據獲取的速度。

3.4 延遲優化

這部分優化，也是分佈式計算的經典優化方向。整個流程鏈路上那些可以精簡、合併、重疊需要不斷去挖掘。對於機器學習系統來說，相比其它的系統，還可以用一些近似的算法來做這部分工作，從而獲得較大的性能提升。下面介紹我們在兩個這方面做的一些優化實踐。

3.4.1 稀疏域參數聚合

在啓用 HashTable 存儲稀疏參數後，對應的，一些配套參數也需要替換爲 HashTable 實現，這樣整個計算圖中會出現多張 HashTable 以及大量的相關算子。在實踐中，我們發現需要儘量降低 Lookup/Insert 等算子的個數，一方面降低 PS 的負載，一方面降低 RPC QPS。因此，針對稀疏模型的常見用法，我們進行了相關的聚合工作。

以 Adam 優化器爲例，需要創建兩個 slot，以保存優化中的動量信息，它的 Shape 與 Embedding 相同。在原生優化器中，這兩個 Variable 是單獨創建的，並在反向梯度更新的時候會去讀寫。同理，使用 HashTable 方案時，我們需要同時創建兩張單獨的 HashTable 用來訓練 m、v 參數。那麼在前向，反向中需要分別對 Embedding、 m、v 進行一次 Lookup 和一次 Insert，總共需要三次 Lookup 和三次 Insert。

這裏一個優化點就是將 Embedding、 m、v，以及低頻過濾的計數器（見下圖 14 的 Counting HashTable）聚合到一起，作爲 HashTable 的 Value，這樣對稀疏參數的相關操作就可以聚合執行，大大減少了稀疏參數操作頻次，降低了 PS 的壓力。

圖 14 基於 HashTable 的參數融合策略

該特性屬於一個普適型優化，開啓聚合功能後，訓練速度有了顯著的提高，性能提升幅度隨着模型和 Worker 規模的變化，效果總是正向的。在美團內部真實業務模型上，聚合之後性能相比非聚合方式能提升了 45% 左右。

3.4.2 Embedding 流水線優化

流水線，在工業生產中，指每一個生產單位只專注處理某個片段的工作，以提高工作效率及產量的一種生產方式。在計算機領域內，更爲大家熟知的是，流水線代表一種多任務之間 Overlap 執行的並行化技術。例如在典型的 RISC 處理器中，用戶的程序由大量指令構成，而一條指令的執行又可以大致分爲：取指、譯碼、執行、訪存、寫回等環節。這些環節會利用到指令 Cache、數據 Cache、寄存器、ALU 等多種不同的硬件單元，在每一個指令週期內，這 5 個環節的硬件單元會並行執行，得以更加充分的利用硬件能力，以此提高整個處理器的指令吞吐性能。處理器的指令流水線是一套複雜而系統的底層技術，但其中的思想在分佈式深度學習框架中也被大量的使用，例如：

如果將分佈式訓練簡單的抽象爲計算和通信兩個過程，絕大多數主流的深度學習框架都支持在執行計算圖 DAG 時，通信和計算的 Overlap。
如果將深度模型訓練簡單的分爲前向和反向，在單步內，由於兩者的強依賴性，無法做到有效並行，字節 BytePS[8] 中引入的通信調度打破了 step iteration 間的屏障，上一輪的部分參數更新完畢後，即可提前開始下輪的前向計算，增強了整體視角下前反向的 Overlap。
百度 AIBox[9] 爲了解決 CTR 場景 GPU 訓練時，參數位於主存，但計算位於 GPU 的問題，巧妙調度不同硬件設備，搭建起了主要利用 CPU / 主存 / 網卡的參數預準備階段和主要利用 GPU/NVLink 的網絡計算階段，通過兩個階段的 Overlap 達到更高的訓練吞吐。

我們看到，在深度學習框架設計上，通過分析場景，可以從不同的視角發掘可並行的階段，來提高整體的訓練吞吐。

對於大規模稀疏模型訓練時，核心模型流程是：先執行稀疏參數的 Embedding，然後執行稠密部分子網絡。其中稀疏參數 Embedding 在遠端 PS 上執行，主要耗費網絡資源，而稠密部分子網絡在本地 Worker 執行，主要耗費計算資源。這兩部分佔了整個流程的大部分時間，在美團某實際業務模型上分別耗時佔比：40%+、50%+。

那我們是否可以提前執行稀疏參數的 Embedding，來做到通信和計算的 Overlap，隱藏掉這部分時間呢？從系統實現上肯定是可行的，但從算法上講，這樣做會引入參數 Staleness 的問題，可能會導致模型精度受到影響。但在實際的生產場景中，大規模異步訓練時本身就會帶來幾十到幾百個步的滯後性問題。經過我們測試，提前獲取一兩步的稀疏參數，模型精度並未受到影響。

在具體實現上，我們把整個計算圖拆分爲 Embedding Graph（EG）和 Main Graph（MG）兩張子圖，兩者異步獨立執行，做到拆分流程的 Overlap（整個拆分過程，可以做到對用戶透明）。EG 主要覆蓋從樣本中抽取 Embedding Key，查詢組裝 Embedding 向量，Embedding 向量更新等環節。MG 主要包含稠密部分子網絡計算、梯度計算、稠密參數部分更新等環節。

圖 15 Embedding 流水線模塊交互關係

兩張子圖的交互關係爲：EG 向 MG 傳遞 Embeding 向量（從 MG 的視角看，是從一個稠密 Variable 讀取數值）；MG 向 EG 傳遞 Embedding 參數對應的梯度。上述兩個過程的表達都是 TensorFlow 的計算圖，我們利用兩個線程，兩個 Session 併發的執行兩張計算圖，使得兩個階段 Overlap 起來，以此到達了更大的訓練吞吐。

圖 16 Embedding 流水線架構流程圖

上圖是 Embedding 流水線的架構流程圖。直觀來看分爲左側的樣本分發模塊，頂部的跨 Session 數據交換模塊，以及自動圖切分得到的 Embedding Graph 和 Main Graph，藍色的圓圈代表新增算子，橙色箭頭代表 EG 重點流程，藍色箭頭代表 MG 重點流程，紅色箭頭代表樣本數據重點流程。

以對用戶透明的形式引入了一層名爲 Pipeline Dataset 的抽象層，這一層的產生是爲了滿足 EG/MG 兩張計算圖以不同節奏運行的需求，支持自定義配置。另外，爲了使得整個流水線中的數據做到彼此的配套，這裏還會負責進行一個全局 Batch ID 的生成及註冊工作。Pipeline Dataset 對外暴露兩種 Iterator，一個供 EG 使用，一個供 MG 使用。Pipeline Dataset 底部共享 TensorFlow 原生的各層 Dataset。
頂部的 ExchangeManager 是一個靜態的，跨 Session 的數據交換媒介，對外暴露數據註冊和數據拉取的能力。抽象這個模塊的原因是，EG 和 MG 原本歸屬於一張計算圖，因爲流水線的原因拆解爲拆爲兩張圖，這樣我們需要建立一種跨 Session 的數據交換機制，並準確進行配套。它內部以全局 Batch ID 做 Key，後面管理了樣本數據、Embeding 向量、Embedding 梯度、Unique 後的 Index 等數據，並負責這些數據的生命週期管理。
中間的 Embedding Graph 由獨立的 TF Session 運行於一個獨立的線程中，通過 a 算子獲得樣本數據後，進行特徵 ID 的抽取等動作，並進行基於 HashTable 方法的稀疏參數查詢，查詢結果通過 c 算子放置到 ExchangeManager 中。EG 中還包含用於反向更新的 f 算子，它會從 ExchangeManager 中獲取 Embedding 梯度和與其配套的前向參數，然後執行梯度更新參數邏輯。
下面的 Main Graph 負責實際稠密子網絡的計算，我們繼承並實現一種可訓練的 EmbeddingVariable，它的構建過程（d 算子）會從 ExchangeManager 查找與自己配套的 Embedding 向量封裝成 EmbeddingVariable，給稠密子網絡。此外，在 EmbeddingVariable 註冊的反向方法中，我們添加了 e 算子使得 Embedding 梯度得以添加到 ExchangeManager 中，供 EG 中的 f 算子消費。

通過上面的設計，我們就搭建起了一套可控的 EG/MG 併發流水線訓練模式。總體來看，Embedding 流水線訓練模式的收益來源有：

經過我們對多個業務模型的 Profiling 分析發現，EG 和 MG 在時間的比例上在 3:7 或 4:6 的左右，通過將這兩個階段並行起來，可以有效的隱藏 Embedding 階段，使得 MG 網絡計算部分幾乎總是可以立即開始，大大加速了整體模型的訓練吞吐。
TensorFlow 引擎中當使用多個優化器（稀疏與非稀疏）的時候，會出現重複構建反向計算圖的問題，一定程度增加了額外計算，通過兩張子圖的拆分，恰好避免了這個問題。
在實施過程中的 ExchangeManager 不僅負責了 Embedding 參數和梯度的交換，還承擔了元數據複用管理的職責。例如 Unique 等算子的結果保存，進一步降低了重複計算。

另外，在 API 設計上，我們做到了對用戶透明，僅需一行代碼即可開啓 Embedding 流水線功能，對用戶隱藏了 EG/MG 的切割過程。目前，在美團某業務訓練中，Embedding 流水線功能在 CPU PS 架構下可以帶來 20%~60% 的性能提升（而且 Worker 併發規模越大，性能越好）。

3.5 單實例 PS 併發優化

經過 2.2 章節的分析可知，我們不能通過持續擴 PS 來提升分佈式任務的吞吐，單實例 PS 的併發優化，也是非常重要的優化方向。我們主要的優化工作如下。

3.5.1 高性能的 HashTable

PS 架構下，大規模稀疏模型訓練對於 HashTable 的併發讀寫要求很高，因爲每個 PS 都要承擔成百乃至上千個 Worker 的 Embedding 壓力，這裏我們綜合速度和穩定性考慮，選用了 tbb::concurrent_hash_map[10] 作爲底層 HashTable 表實現，並將其包裝成一個新的 TBBConcurrentHashTable 算子。經過測試，在千億規模下 TBBConcurrentHashTable 比原生 MutableDenseHashTable 訓練速度上快了 3 倍。

3.5.2 HashTable BucketPool

對於大規模稀疏模型訓練來說，Embedding HashTable 會面對大量的併發操作，通過 Profiling 我們發現，頻繁動態的內存申請會帶來了較大性能開銷（即使 TensorFlow 的 Tensor 有專門的內存分配器）。我們基於內存池化的思路優化了 HashTable 的內存管理。

我們在 HashTable 初始化時，會先爲 Key 和 Value 分別創造兩個 BucketPool，每個池子都會先 Malloc 較大一塊內存備用，考慮到可能會有對 HashTable 進行中的 Key 和 Value 進行 Remove 的場景（如 Online Learning 訓練時），需要對從 HashTable 中刪除的 Key 和 Value 所使用的內存進行回收，因此每個 BucketPool 還有一個 ReuseQueue 來負責維護回收的內存。每次向內部的哈希表數據結構中 Insert Key 和 Value 的時候，Key 和 Value 內存和釋放分配都進行池化管理。用這種方式降低了大規模稀疏訓練中遇到稀疏內存分配開銷，整體端到端訓練性能提升了 5% 左右。

圖 17 HashTable 內存優化

3.6 單位算力吞吐優化

經過 2.2 章節的分析，Worker 的計算壓力也非常大，如果不優化 Worker，同時要保持吞吐，需要橫向擴展更多的 Worker，給 PS 帶來更大的壓力。而對於用戶來說，如果能在有限的計算資源下帶來性能提升，對業務價值更高。我們通過 CAT 統計出了一些高頻算子，並進行了專項優化。這裏選取 Unique&DynamicPartition 算子融合案例進行分享。

在 TensorFlow PS 架構中，包括 Embedding 向量在內的共享參數都存儲在 PS 上，並通過網絡與 Worker 交互，在進行 Embedding 查詢過程中，往往會涉及如下兩個環節：

由於稀疏參數的性質，從樣本中抽取得到的待查詢 Embedding ID，它的重複率往往高達 70%~90%，如果不進行去重查詢，不論是對 HashTable 的查詢還是網絡的傳輸，都會帶來不小的壓力。因此，通常會在查詢前進行 Unique 操作。
在大規模稀疏場景中，爲了存儲千億規模的參數，會有多個 PS 機器共同承載。而 Worker 端會負責對查詢請求按照設定的路由規則進行切分，這裏通常會在查詢前進行 DynamicPartition 動作。

通常這兩個過程會利用 TensorFlow 既有的算子進行搭建，但在實際使用中，我們發現它並不是很高效，主要問題在於：

Unique 算子原生實現，它內部使用的內存分配策略較爲低效。使用了兩倍輸入參數（Embedding ID）的大小進行內存分配，但由於輸入參數較大，而且重複率高，導致 HashTable 創建過大且非常稀疏。幾乎每次插入都會產生一次 minor_page_fault，導致 HashTable 性能下降。我們使用 Intel Vtune 驗證了這一點（參見圖 18）。
Unique 和 Dynamic Partition 算子存在冗餘數據遍歷，這些操作其實可以在一次數據遍歷中全部做完，節省掉算子切換、冗餘數據遍歷的耗時。

圖 18 Unique 算子內部出現 DRAM Bound 問題

總結來說，HashTable 開闢過大會導致大量的 minor_page_fault，導致訪存的時間增加，HashTable 過小又可能會導致擴容。我們採用了基於啓發式算法的內存自適應 Unique 算子實現，通過對訓練歷史重複率的統計，我們可以得到一個相對合理的 HashTable 大小，來提高訪存的性能；另外 Unique 算子內 HashTable 的具體選擇上，經過我們的多種測試，選擇了 Robin HashTable 替換了原生 TF 中的實現。

進一步，我們對圍繞 Embedding ID 的 Unique 和 Partition 環節進行了算子合併，簡化了邏輯實現。經過上述的優化，Unique 單算子可以取得 51% 的加速，在真實模型端到端上可以獲得 10% 左右的性能提升，算子總數量降低了 4%。

在整個關鍵算子優化的過程中，Intel 公司的林立凡、張向澤、高明進行大量的技術支持，我們也複用了他們的部分優化工作，在此深表感謝！

4 大規模稀疏算法建模

大規模稀疏能力在業務落地的過程中，算法層面還需要從特徵和模型結構上進行對應升級，才能拿到非常好的效果。其中外賣廣告從業務特點出發，引入大規模稀疏特徵完成外賣場景下特徵體系的升級，提供了更高維的特徵空間和參數空間，增強了模型的擬合能力。重新設計了面向高維稀疏場景的特徵編碼方案，解決了特徵編碼過程中的特徵衝突問題，同時編碼過程去掉了部分冗餘的特徵哈希操作，一定程度上簡化了特徵處理邏輯，並降低了特徵計算的耗時。

在系統層面，面對百億參數、百億樣本以上量級的大規模稀疏模型的訓練，會帶來訓練迭代效率的大大降低，單次實驗從一天以內，增長到一週左右。美團機器學習平臺訓練引擎團隊，除了上述 TensorFlow 框架層面的優化、還針對業務模型進行了專項優化，整體吞吐優化了 8 到 10 倍（如果投入更多計算資源，可以進一步加速），大大提升業務的迭代效率，助力外賣廣告業務取得了較爲明顯的提升。

5 總結與展望

TensorFlow 在大規模推薦系統中被廣泛使用，但由於缺乏大規模稀疏的大規模分佈式訓練能力，阻礙了業務的發展。美團基於 TensorFlow 原生架構，支持了大規模稀疏能力，並從多個角度進行了深度優化，做到千億參數、千億樣本高效的分佈式訓練，並在美團內部進行了大規模的使用。對於這類關鍵能力的缺失，TensorFlow 社區也引起了共鳴，社區官方在 2020 年創建了 SIG Recommenders[11]，通過社區共建的方式來解決此類問題，美團後續也會積極的參與到社區的貢獻當中去。

美團推薦系統場景的模型訓練，目前主要運行在 CPU 上，但隨着業務的發展，有些模型變得越來越複雜，CPU 上已經很難有優化空間（優化後的 Worker CPU 使用率在 90% 以上）。而近幾年，GPU 的計算能力突飛猛進，新一代的 NVIDIA A100 GPU，算力達到了 156TFLOPS（TF32 Tensor Cores）、80G 顯存、卡間帶寬 600GB/s。對於這類複雜模型的 Workload，我們基於 A100 GPU 架構，設計了下一代的分佈式訓練架構，經過初步優化，在美團某大流量業務推薦模型上也拿到了較好的效果，目前還在進一步優化當中，後續我們會進行分享，敬請期待。