TensorFlow在美團外賣推薦場景的GPU訓練優化實踐

　　美團機器學習平臺基於內部深度定製的TensorFlow研發了Booster GPU訓練架構。該架構在整體設計上充分考慮了演算法、架構、新硬體的特性，從資料、計算、通訊等多個角度進行了深度的優化，最終其價效比達到CPU任務的2~4倍。本文主要講述Booster架構的設計實現、效能優化及業務落地工作，希望能對從事相關開發的同學有所幫助或者啟發。

　　 1 背景

　　在推薦系統訓練場景中，美團內部深度定製的TenorFlow（簡稱TF）版本[1]，通過CPU算力支撐了美團內部大量的業務。但隨著業務的發展，模型單次訓練的樣本量越來越多，結構也變得越來越複雜。以美團外賣推薦的精排模型為例，單次訓練的樣本量已達百億甚至千億，一次實驗要耗費上千核，且優化後的訓練任務CPU使用率已達90%以上。為了支援業務的高速發展，模型迭代實驗的頻次和併發度都在不斷增加，進一步增加了算力使用需求。在預算有限的前提下，如何以較高的價效比來實現高速的模型訓練，從而保障高效率的模型研發迭代，是我們迫切需要解決的問題。

　　近幾年，GPU伺服器的硬體能力突飛猛進，新一代的NVIDIA A100 80GB SXM GPU伺服器（8卡）[2]，在儲存方面可以做到：視訊記憶體640GB、記憶體1~2TB、SSD10+TB，在通訊方面可以做到：卡間雙向通訊600GB/s、多機通訊800~1000Gbps/s，在算力方面可以做到：GPU 1248TFLOPS（TF32 Tensor Cores），CPU 96~128物理核。如果訓練架構能充分發揮新硬體的優勢，模型訓練的成本將會大大降低。但TensorFlow社群在推薦系統訓練場景中，並沒有高效和成熟的解決方案。我們也嘗試使用優化後的TensorFlow CPU Parameter Server[3]（簡稱PS）+GPU Worker的模式進行訓練，但只對複雜模型有一定的收益。NVIDIA開源的HugeCTR[4]雖然在經典的深度學習模型上效能表現優異，但要在美團的生產環境直接使用起來，還需要做較多的工作。

　　美團基礎研發機器學習平臺訓練引擎團隊，聯合到家搜推技術部演算法效能團隊、NVIDIA DevTech團隊，成立了聯合專案組。在美團內部深度定製的TenorFlow以及NVIDIA HugeCTR的基礎上，研發了推薦系統場景的高效能GPU訓練架構Booster。目前在美團外賣推薦場景中進行了部署，多代模型全面對齊演算法的離線效果，對比之前，優化後的CPU任務，價效比提升了2~4倍。由於Booster對原生TensorFlow介面有較好的相容性，原TensorFlow CPU任務只需要一行程式碼就可完成遷移。這樣讓Booster可以快速在美團多條業務線上進行初步驗證，相比之前的CPU任務，平均價效比都提升到2倍以上。本文將重點介紹Booster架構的設計與優化，以及在美團外賣推薦場景落地的全過程，希望能對大家有所幫助或啟發。

　　 2 GPU訓練優化挑戰

　　GPU訓練在美團內已經廣泛應用到CV、NLP、ASR等場景的深度學習模型，但在推薦系統場景中，卻遲遲沒有得到大規模的應用，這跟場景的模型特點、GPU伺服器的硬體特點都有較強的關係。

　　推薦系統深度學習模型特點

　　讀取樣本量大：訓練樣本在幾十TB~幾百TB，而CV等場景通常在幾百GB以內。

　　模型引數量大：同時有大規模稀疏引數和稠密引數，需要幾百GB甚至上TB儲存，而CV等場景模型主要是稠密引數，通常在幾十GB以內。

　　模型計算複雜度相對低一些：推薦系統模型在GPU上單步執行只需要10~100ms，而CV模型在GPU上單步執行是100~500ms，NLP模型在GPU上單步執行是500ms~1s。

　　GPU伺服器特點

　　GPU卡算力很強，但視訊記憶體仍有限：如果要充分發揮GPU算力，需要把GPU計算用到的各種資料提前放置到視訊記憶體中。而從2016年~2020年，NVIDIA Tesla GPU卡[5]計算能力提升了10倍以上，但視訊記憶體大小隻提升了3倍左右。

　　其它維度資源並不是很充足：相比GPU算力的提升速度，單機的CPU、網路頻寬的增長速度較慢，如果遇到這兩類資源負載較重的模型，將無法充分發揮GPU的能力，GPU伺服器相比CPU伺服器的價效比不會太高。

　　總結來說，CV、NLP等場景的模型訓練屬於計算密集型任務，而且大多模型單張卡的視訊記憶體都可以裝下，這和GPU伺服器的優勢非常好地進行了匹配。但在推薦系統場景中，由於模型相對沒有那麼複雜，遠端讀取的樣本量大，特徵處理耗費CPU多，給單機CPU和網路帶來較大的壓力。同時面對模型引數量大的情況，單機的GPU視訊記憶體是無法放下的。這些GPU伺服器的劣勢，恰恰都被推薦系統場景命中。

　　好在NVIDIA A100 GPU伺服器，在硬體上的升級彌補了視訊記憶體、CPU、頻寬這些短板，但如果系統實現和優化不當，依然不會有太高的價效比收益。在落地Booster架構的過程中，我們主要面臨如下挑戰：

　　資料流系統：如何利用好多網路卡、多路CPU，實現高效能的資料流水線，讓資料的供給可以跟上GPU的消費速度。

　　混合引數計算：對於大規模稀疏引數，GPU視訊記憶體直接裝不下的情況，如何充分利用GPU高算力、GPU卡間的高頻寬，實現一套大規模稀疏引數的計算，同時還需要兼顧稠密引數的計算。

　　 3 系統設計與實現

　　面對上面的挑戰，如果純從系統的的角度去設計，難度較大。Booster採用了“演算法+系統”Co-design的設計思路，讓這代系統的設計大大得到簡化。在系統實施路徑上，考慮到業務預期交付時間、實施風險，我們並沒有一步到位落地Booster的多機多卡版本，而是第一版先落地了GPU單機多卡版本，本文重點介紹的也是單機多卡的工作。另外，依託於NVIDIA A100 GPU伺服器強大的計算能力，單機的算力可以滿足美團絕大多數業務的單次實驗需求。

　　3.1 引數規模的合理化

　　大規模稀疏離散特徵的使用，導致深度預估模型的Embedding引數量急劇膨脹，數TB大小的模型一度流行於業界推搜的各大頭部業務場景。但是業界很快意識到，在硬體成本有限的情況下，過於龐大的模型給生產部署運維和實驗迭代創新增添了沉重的負擔。學術研究表明[10-13]，模型效果強依賴於模型的資訊容量，並非引數量。實踐證明，前者可以通過模型結構的優化來進行提升，而後者在保證效果的前提下，尚存有很大的優化空間。Facebook在2020年提出了Compositional Embedding[14]，實現推薦模型引數規模數個量級的壓縮。阿里巴巴也發表了相關工作[15]，將核心業務場景的預估模型由數TB壓縮至幾十GB甚至更小。總的來看，業界的做法主要有以下幾種思路：

　　去交叉特徵：交叉特徵由單特徵間做笛卡爾積產生，這會生成巨大的特徵ID取值空間和對應Embedding參數列。深度預估模型發展至今，已經有大量的方法通過模型結構來建模單特徵間的互動，避免了交叉特徵造成的Embedding規模膨脹，如FM系列[16]、AutoInt[17]、CAN[18]等。

　　精簡特徵：特別是基於NAS的思路，以較低的訓練成本實現深度神經網路自適應特徵選擇，如Dropout Rank[19]和FSCD[20]等工作。

　　壓縮Embedding向量數：對特徵取值進行復合ID編碼和Embedding對映，以遠小於特徵取值空間的Embedding向量數，來實現豐富的特徵Embedding表達，如Compositional Embedding[14]、Binary Code Hash Embedding[21]等工作。

　　壓縮Embedding向量維度：一個特徵Embedding向量的維度決定了其表徵資訊的上限，但是並非所有的特徵取值都有那麼大的資訊量，需要Embedding表達。因此，可以每一個特徵值自適應的學習精簡Embedding維度，從而壓縮引數總量，如AutoDim[22]和AMTL[23]等工作。

　　量化壓縮：使用半精度甚至int8等更激進的方式，對模型引數做量化壓縮，如DPQ[24]和MGQE[25]。

　　美團外賣推薦的模型一度達到100G以上，通過應用以上方案，我們在模型預估精度損失可控的前提下，將模型控制在10GB以下。

　　基於這個演算法基礎假設，我們將第一階段的設計目標定義到支援100G以下的引數規模。這可以比較好的適配A100的視訊記憶體，存放在單機多卡上，GPU卡間雙向頻寬600GB/s，可以充分發揮GPU的處理能力，同時也可以滿足美團大多數模型的需求。

　　3.2 系統架構

　　基於GPU系統的架構設計，要充分考慮硬體的特性才能充分發揮效能的優勢。我們NVIDIA A100伺服器的硬體拓撲和NVIDIA DGX A100[6]比較類似，每臺伺服器包含：2顆CPU，8張GPU，8張網路卡。Booster架構的架構圖如下所示：

　　圖1 系統架構

　　整個系統主要包括三個核心模組：資料模組，計算模組，通訊模組：

　　資料模組：美團自研了一套支援多資料來源、多框架的資料分發系統，在GPU系統上，我們改造資料模組支援了多網路卡資料下載，以及考慮到NUMA Awareness的特性，在每顆CPU上都部署了一個資料分發服務。

　　計算模組：每張GPU卡啟動一個TensorFlow訓練程式執行訓練。

　　通訊模組：我們使用了Horovod[7]來做分散式訓練的卡間通訊，我們在每個節點上啟動一個Horovod程式來執行對應的通訊任務。

　　上述的設計，符合TensorFlow和Horovod原生的設計正規化。幾個核心模組可以相互解耦，獨立迭代，而且如果合併開源社群的最新特性，也不會對系統造成架構性的衝擊。

　　我們再來看一下整個系統的簡要執行流程，每張GPU卡上啟動的TensorFlow程式內部的執行邏輯如下圖：

　　圖2 程式內部執行邏輯

　　整個訓練流程涉及引數儲存、優化器、卡間通訊等幾個關鍵模組。對於樣本的輸入特徵，我們分為稀疏特徵（ID類特徵）和稠密特徵。在實際業務場景中，稀疏特徵通常IDs總量較多，對應的稀疏引數使用HashTable資料結構儲存更合適，而且由於引數量較大，GPU單卡視訊記憶體放不下，我們會通過ID Modulo的方式Partition到多張GPU卡的視訊記憶體中存放。對於IDs總量較少的稀疏特徵，業務通常使用多維矩陣資料結構表達（在TensorFlow裡面的資料結構是Variable），由於引數量不大，GPU單卡視訊記憶體可以放下，我們使用Replica的方式，每張GPU卡的視訊記憶體都放置一份引數。對於稠密引數，通常使用Variable資料結構，以Replica的方式放置到GPU視訊記憶體中。下邊將詳細介紹Booster架構的內部實現。

　　3.3 關鍵實現

　　3.3.1 引數儲存

　　早在CPU場景的PS架構下，我們就實現了大規模稀疏引數的整套邏輯，現在要把這套邏輯搬到GPU上，首先要實現的就是GPU版本的HashTable。我們調研了業界多種GPU HashTable的實現，如cuDF、cuDPP、cuCollections、WarpCore等，最終選擇了基於cuCollections實現TensorFlow版本的GPUHashTable。究其原因，主要是因為實際業務場景中，大規模稀疏特徵的總量通常是未知的，並且隨時可能出現特徵交叉，從而致使稀疏特徵的總量變化很大，這就導致“動態擴容”能力將成為我們GPU HashTable的必備功能，能夠做到動態擴容的只有cuCollections的實現。我們在cuCollections的GPU HashTable基礎上實現了特殊介面（find_or_insert），對大規模讀寫效能進行了優化，然後封裝到了TensorFlow中，並在其上實現了低頻過濾的功能，能力上對齊CPU版本的稀疏引數儲存模組。

　　3.3.2 優化器

　　目前，稀疏引數的優化器與稠密引數的優化器並不相容，我們在GPU HashTable的基礎上，實現了多種稀疏優化器，並且都做了優化器動量Fusion等功能，主要實現了Adam、Adagrad、FTRL、Momentum等優化器。對實際業務場景來說，這些優化器已經能夠覆蓋到絕大多數業務的使用。稠密部分引數可以直接使用TensorFlow原生支援的稀疏/稠密優化器。

　　3.3.2 卡間通訊

　　實際訓練期間，對於不同型別的特徵，我們的處理流程也有所不同：

　　稀疏特徵（ID類特徵，規模較大，使用HashTable儲存）：由於每張卡的輸入樣本資料不同，因此輸入的稀疏特徵對應的特徵向量，可能存放在其他GPU卡上。具體流程上，訓練的前向我們通過卡間AllToAll通訊，將每張卡的ID特徵以Modulo的方式Partition到其他卡中，每張卡再去卡內的GPUHashTable查詢稀疏特徵向量，然後再通過卡間AllToAll通訊，將第一次AllToAll從其他卡上拿到的ID特徵以及對應的特徵向量原路返回，通過兩次卡間AllToAll通訊，每張卡樣本輸入的ID特徵都拿到對應的特徵向量。訓練的反向則會再次通過卡間AllToAll通訊，將稀疏引數的梯度以Modulo的方式Partition到其他卡中，每張卡拿到自己的稀疏梯度後再執行稀疏優化器，完成大規模稀疏特徵的優化。詳細流程如下圖所示：

　　圖3 稀疏特徵處理流程

　　稀疏特徵（規模較小，使用Variable儲存）：相比使用HashTable的區別，由於每張GPU卡都有全量的引數，直接在卡內查詢模型引數即可。在反向聚合梯度的時候，會通過卡間AllGather獲取所有卡上的梯度求平均，然後交給優化器執行引數優化。

　　稠密特徵：稠密引數也是每張卡都有全量的引數，卡內可以直接獲取引數執行訓練，最後通過卡間AllReduce聚合多卡的稠密梯度，執行稠密優化器。

　　在整個的執行過程中，稀疏引數和稠密引數全部放置在GPU視訊記憶體中，模型計算也全部在GPU上處理，GPU卡間通訊頻寬也足夠快，能夠充分發揮了GPU的強大算力。

　　這裡小結一下，Booster訓練架構，與CPU場景PS架構的核心區別在於：

　　訓練模式：PS架構是非同步訓練模式，Booster架構是同步訓練模式。

　　引數分佈：PS架構下模型引數都存放在PS記憶體中，Booster架構下稀疏引數（HashTable）是Partition方式分佈在單機八卡中，稠密引數（Variable）是Replica方式存放在每張卡中，因此Booster架構下的Worker角色兼顧了PS架構下PS/Worker角色的功能。

　　通訊方式：PS架構下PS/Worker間通訊走的是TCP（Grpc/Seastar），Booster架構下Worker間通訊走的是NVSwitch（NCCL），任意兩卡間雙向頻寬600GB/s，這也是Booster架構的訓練速度取得較大提升的原因之一。

　　由於每張卡的輸入資料不同，並且模型引數既有在卡間Partition儲存的，也有在卡間Replica儲存的，因此Booster架構同時存在模型並行、資料並行。此外，由於NVIDIA A100要求CUDA版本>=11.0，而TensorFlow 1.x版本只有NV1.15.4才支援CUDA11.0。美團絕大多數業務場景都還在使用TensorFlow 1.x，因此我們所有改造都是在NV1.15.4版本基礎上開發的。

　　以上就是Booster整體系統架構及內部執行流程的介紹。下文主要介紹在初步實現的Booster架構的基礎上，我們所做的一些效能優化工作。

　　 4 系統效能優化

　　基於上述的設計實現完第一版系統後，我們發現端到端效能並不是很符合預期，GPU的SM利用率（SM Activity指標）只有10%~20%，相比CPU並沒有太大的優勢。為了分析架構的效能瓶頸，我們使用NVIDIA Nsight Systems（以下簡稱nsys）、Perf、uPerf等工具，通過模組化壓測、模擬分析等多種分析手段，最終定位到資料層、計算層、通訊層等幾方面的效能瓶頸，並分別做了相應的效能優化。以下我們將以美團外賣某推薦模型為例，分別從GPU架構的資料層、計算層、通訊層，逐個介紹我們所做的效能優化工作。

　　4.1 資料層

　　如前文所述，推薦系統的深度學習模型，樣本量大，模型相對不復雜，資料I/O本身就是瓶頸點。如果幾十臺CPU伺服器上的資料I/O操作，都要在單臺GPU伺服器上完成，那麼資料I/O的壓力會變得更大。我們先看一下在當前系統下的樣本資料流程，如下圖所示：

　　圖4 樣本資料流程及核心優化點

　　核心流程：資料分發程式通過網路讀取HDFS樣本資料（TFRecord格式）到記憶體中，然後通過共享記憶體（Shared Memory）的方式把樣本資料傳輸給TensorFlow訓練程式。TensrFlow訓練程式收到樣本資料後，走原生的TensrFlow特徵解析邏輯，拿到特徵資料後通過GPU MemcpyH2D到GPU視訊記憶體中。我們通過模組化壓測分析發現，資料分發層的樣本拉取、TensrFlow層的特徵解析以及特徵資料MemcpyH2D到GPU等幾個流程，都存在較大的效能問題（圖中黃色流程所示），以下詳細介紹我們在這幾塊所做的效能優化工作。

　　4.1.1 樣本拉取優化

　　樣本拉取、組裝Batch是由資料分發程式完成的，我們在這裡所做的主要優化工作是，首先將資料分發程式通過numactl獨立到NUMA內部執行，避免了NUMA間的資料傳輸；其次，資料下載從單網路卡擴充到了多網路卡，增大資料下載頻寬；最後，資料分發程式與TensrFlow程式之間的傳輸通道，從單個Shared Memory擴充套件到每張GPU卡有獨立的Shared Memory，避免了單Shared Memory所帶來的記憶體頻寬問題，並在TensrFlow內部實現了特徵解析時對輸入資料零拷貝的能力。

　　4.1.2 特徵解析優化

　　目前，美團內部絕大多數業務的樣本資料都還是TFRecord格式，TFRecord實際上是ProtoBuf（簡稱PB）格式。PB反序列化非常耗費CPU，其中ReadVarint64Fallback方法CPU佔用較為突出，實際profiling結果如下圖：

　　圖5 樣本解析profiling結果

　　究其原因，CTR場景的訓練樣本通常包含了大量的int64型別的特徵，int64在PB中是以Varint64型別資料儲存的，ReadVarint64Fallback方法就是用來解析int64型別的特徵。普通的int64資料型別需要佔用8個位元組，而Varint64針對不同的資料範圍，使用了變長的儲存長度。PB在解析Varint型別資料時，首先要確定當前資料的長度，Varint用7bit儲存資料，高位1bit儲存標記位，該標記位表示下一個位元組是否有效，如果當前位元組最高位為0，則說明當前Varint資料在該位元組處結束。我們實際業務場景的ID特徵大多是經過Hash後的值，用Varint64型別表達會比較長，這也就導致在特徵解析過程中要多次判斷資料是否結束，以及多次位移和拼接來生成最終資料，這使得CPU在解析過程中存在大量的分支預測和臨時變數，非常影響效能。以下是4位元組Varint的解析流程圖：

　　圖6 ProtoBuf Varint解析流程圖

　　這個處理流程，非常適合用SIMD指令集批處理優化。以4位元組的Varint型別為例，我們的優化流程主要包括兩步：

　　SIMD尋找最高位：通過SIMD指令將Varint型別資料的每個位元組與0xF0做與運算，找到第一個結果等於0的位元組，這個位元組就是當前Varint資料的結束位置。

　　SIMD處理Varint：按理來說，通過SIMD指令將Varint資料高位清零後的每個位元組依次右移3/2/1/0位元組，就可得到最終的int型別資料，但SIMD沒有這樣的指令。因此，我們通過SIMD指令分別處理每個位元組的高4bit、低4bit，完成了這個功能。我們將Varint資料的高低4bit分別處理成int_h4與int_l4，再做或運算，就得到了最終的int型別資料。具體優化流程如下圖所示（4位元組資料）：

　　圖7 ProtoBuf Varint解析優化後流程圖

　　對於Varint64型別資料的處理，我們直接分成了兩個Varint型別資料來處理。通過這兩步的SIMD指令集優化，樣本解析速度得到大大提升，在GPU端到端訓練速度提升的同時，CPU使用率下降了15%。這裡我們主要使用了SSE指令集優化，期間也嘗試了AVX等更大長度的指令集，但效果不是很明顯，最終並沒有使用。此外，SIMD指令集在老的機器上會導致CPU嚴重降頻，因此官方社群並沒有引入這個優化，而我們GPU機器的CPU都比較新，完全可以使用SIMD指令集進行優化。

　　4.1.3 MemcpyH2D流水線

　　解析完樣本得到特徵資料後，需要將特徵資料拉到GPU中才能執行模型計算，這裡需要通過CUDA的MemcpyH2D操作。我們通過nsys分析這塊的效能，發現GPU在執行期間有較多的停頓時間，GPU需要等待特徵資料Memcpy到GPU上之後才能執行模型訓練，如下圖所示：

　　圖8 nsys profiling結果

　　對於GPU系統的資料流，需要提前傳輸到離GPU處理器最近的視訊記憶體中，才能發揮GPU的計算能力。我們基於TensorFlow的prefetch功能，實現了GPU版本的PipelineDataset，在計算之前先把資料拷貝到了GPU視訊記憶體中。需要注意的是CPU記憶體拷貝到GPU視訊記憶體這個過程，CPU記憶體需要使用Pinned Memory，而非原生的Paged Memory，可以加速MemcpyH2D流程。

　　4.1.4 硬體調優

　　在資料層的效能優化期間，美團內部基礎研發平臺的伺服器組、網路組、作業系統組也幫助我們做了相關的調優：

　　在網路傳輸方面，為了減少網路協議棧處理開銷，提高資料拷貝的效率，我們通過優化網路卡配置，開啟LRO（Large-Receive-Offload）、TC Flower的硬體解除安裝、Tx-Nocache-Copy等特性，最終網路頻寬提升了17%。

　　在CPU效能優化方面，經過效能profiling分析，發現記憶體延遲和頻寬是瓶頸。於是我們嘗試了3種NPS配置，綜合業務場景和NUMA特性，選擇了NPS2。此外，結合其他BIOS配置（例如APBDIS，P-state等），可以將記憶體延遲降低8%，記憶體頻寬提升6%。

　　通過上述優化，網路極限頻寬提升了80%，在業務需求頻寬下GPU的H2D頻寬提升了86%。最終在資料解析層面也拿到了10%+的效能收益。

　　經過資料層樣本拉取、特徵解析、MemcpyH2D和硬體的優化，Booster架構端到端訓練速度提升了40%，訓練價效比達到了CPU的1.4倍，資料層也不再成為當前架構的效能瓶頸。

　　4.2 計算層

　　4.2.1 Embedding流水線

　　早在CPU場景做TensorFlow訓練效能優化時，我們就已經實現了Embedding Pipeline[1]的功能：我們把整個計算圖拆分為Embedding Graph（EG）和Main Graph（MG）兩張子圖，兩者非同步獨立執行，做到執行上的Overlap（整個拆分過程，可以做到對使用者透明）。EG主要覆蓋從樣本中抽取Embedding Key，查詢組裝Embedding向量，Embedding向量更新等環節；MG主要包含稠密部分子網路計算、梯度計算、稠密引數部分更新等環節。

　　圖9 Embedding流水線模組互動關係

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

　　我們把這個流程在GPU架構下也實現了一遍，並在其中加入了卡間同步流程，大規模稀疏特徵的AllToAll通訊及其反向梯度的AllToAll通訊都在EG中執行，普通稀疏特徵的反向梯度的卡間AllGather同步、稠密引數的反向梯度的卡間AllReduce同步都在MG中執行。需要注意的是，在GPU場景中，EG、MG是在同一個GPU Stream上執行CUDA Kernel的，我們嘗試過EG、MG分別在獨立的GPU Stream上執行，效能會變差，深層原因與CUDA底層實現有關，這個問題本身還在等待解決。

　　4.2.2 運算元優化及XLA

　　相比CPU層面的優化，GPU上的優化更加複雜。首先對於TensorFlow的運算元，還有一些沒有GPU的實現，當模型中使用了這些CPU運算元，會跟上下游的GPU運算元出現記憶體和視訊記憶體之間的資料來回拷貝，影響整體效能，我們在GPU上實現了使用較為頻繁、影響較大的運算元。另外，對於TensorFlow這代框架，運算元粒度是非常細的，可以方便使用者靈活搭建各種複雜的模型，但這對GPU處理器來說卻是一個災難，大量的Kernel Launch以及訪存開銷導致不能充分利用GPU算力。對於GPU上的優化，通常有兩個方向，手工優化和編譯優化。在手工優化方面，我們重新實現了一些常用的運算元和層（Unique、DynamicPartition、Gather等）。

　　以Unique運算元為例，原生TensorFlow的Unique運算元要求輸出元素的順序與輸入元素的順序一致，而在實際場景中，我們並不需要這個限制，我們修改了Unique運算元的GPU實現，減少了因輸出有序導致的額外執行的GPU Kernel。

　　在編譯優化方面，目前我們主要使用TensorFlow社群提供的XLA[9]來做一些自動優化。原生TensorFlow 1.15中的XLA正常開啟可獲得10～20%端到端的效能提升。但XLA對運算元動態shape不能很好地進行支援，而推薦系統場景的模型中這種情況卻非常常見，這就導致XLA加速效能不符合預期，甚至是負優化，因此我們做了如下的緩解工作：

　　區域性優化：對於我們手動引入的動態shape運算元（如Unique），我們進行了子圖示記，不執行XLA編譯，XLA只優化可以穩定加速的子圖。

　　OOM兜底：XLA會根據運算元的type、input type、shape等資訊，快取編譯中間結果，避免重複編譯。然而由於稀疏場景以及GPU架構實現的特殊性，天然存在Unique、DynamicPartition等Output shape是動態的運算元，這就導致這些運算元以及連線在這些運算元之後的運算元，在執行XLA編譯時無法命中XLA快取而重新編譯，新的快取越來越多，而舊的快取不會被釋放，最終導致CPU記憶體OOM。我們在XLA內部實現了LRUCache，主動淘汰掉舊的XLA快取，避免OOM的問題。

　　Const Memcpy消除：XLA在使用TF_HLO重寫TensorFlow運算元時，對一些編譯期已固定的資料會打上Const標記，然而這些Const運算元的Output只能定義在Host端，為了將Host端的Output送給Device端需要再加一次MemcpyH2D，這就佔用了TensorFlow原有的H2D Stream，影響樣本資料提前拷貝到GPU端。由於XLA的Const Output在編譯期已經固化，因此沒有必要每一步都做一次MemcpyH2D，我們將Device端的Output快取下來，後續使用該Output時，直接從快取中讀取，避免多餘的MemcpyH2D。

　　對於XLA的優化，確切的來說應該是問題修復，目前能夠做到的是GPU場景下可以正常開啟XLA，並獲得10～20%的訓練速度提升。值得一提的是，對於動態shape的運算元編譯問題，美團內部基礎研發機器學習平臺/深度學習編譯器團隊已經有了徹底的解決方案，後續我們會聯合解決這個問題。

　　經過計算層的Embedding流水線、XLA相關優化，Booster架構端到端訓練速度提升了60%，GPU單機八卡訓練價效比達到同等資源下CPU的2.2倍。

　　4.3 通訊層

　　在單機多卡訓練過程中，我們通過Nsight Systems分析發現，卡間通訊耗時佔比非常高，而且在此期間GPU使用率也非常低，如下圖所示：

　　圖10 nsys profiling結果

　　從圖中可以看出，訓練期間卡間通訊耗時比較長，同時在通訊期間GPU使用率也非常低，卡間通訊是影響訓練效能提升的關鍵瓶頸點。我們對通訊過程進行拆解打點後發現，卡間通訊（AllToAll、AllReduce、AllGather等）協商的時間遠遠高於資料傳輸的時間：

　　圖11 Horovod timeline結果

　　分析具體原因，以負責大規模稀疏引數通訊的AllToAll為例，我們通過Nsight Systems工具，觀察到通訊協商時間長主要是由於某張卡上的運算元執行時間比較晚導致的。由於TensorFlow運算元排程並不是嚴格有序，同一個特徵的embedding_lookup運算元，在不同卡上真正執行的時間點也不盡相同，某張卡上第一個執行embedding_lookup運算元在另一張卡上可能是最後一個執行，因此我們懷疑不同卡上運算元排程的不一致性，導致了各張卡發起通訊的時刻不同，並最終導致了通訊協商時間過長。我們通過幾組模擬實驗也論證了確實是由運算元排程導致的。對於這個問題，最直接的想法是改造TensorFlow計算圖的核心排程演算法，但這個問題在學術界也一直是一個複雜的問題。我們換了一種思路，通過融合關鍵的運算元，來緩解這個問題，通過統計，我們選擇了HashTable和Variable相關的運算元。

　　4.3.1 HashTable相關運算元融合

　　我們設計和實現了一個圖優化過程，這個過程會自動地將圖中可以合併的HashTable及對應的embedding_lookup過程進行合併，合併策略上主要將embedding_size相同的HashTable合併到一塊。同時為了避免HashTable合併之後原始特徵之間發生ID衝突，我們引入了自動統一特徵編碼的功能，對不同的原始特徵分別加上不同的偏移量，歸入不同的特徵域，實現了訓練時的統一特徵編碼。

　　我們在某實際業務模型上進行測試，該圖優化將38張HashTable合併成為了2張HashTable，將38次embedding_lookup合併成了2次，這將EmbeddingGraph中的embedding_lookup相關運算元數量減少了90%，卡間同步通訊次數減少了90%。此外，運算元合併之後，embedding_lookup中的GPU運算元也發生了合併，減少了Kernel Launch次數，使得EmbeddingGraph的執行速度變得更快。

　　4.3.2 Variable相關運算元融合

　　類似於HashTable Fusion的優化思路，我們觀察到業務模型中通常包含數十至數百個TensorFlow原生的Variable，這些Variable在訓練期間梯度需要做卡間同步，同樣的，Variable數量太多導致卡間同步的協商時間變長。我們通過Concat/Split運算元，將所有的Trainable Variables自動合併到一起，使得整個MG的反向只產生幾個梯度Tensor，大大減少了卡間同步的次數。同時，做完Variable Fusion之後，優化器中實際執行的運算元數量也大大減少，加快了計算圖本身的執行速度。

　　需要注意的是，TensorFlow的Variable分為兩種，一種是每個Step全部引數值都參與訓練的Dense Variable，如MLP的Weight；另一種是專門用於embedding_lookup的Variable，每個Step只有部分值參與訓練，我們稱之為Sparse Variable。對於前者，做Variable合併不會影響到演算法效果。而對於後者，它反向梯度是IndexedSlices物件，卡間同步預設走的是AllGather通訊，如果業務模型中對於Sparse Variables的優化採用的是Lazy優化器，即每個Step只優化更新Variable中的某些行，此時對Sparse Variables做合併，會導致其反向梯度從IndexedSlices物件轉為Tensor物件，卡間同步變成AllReduce過程，就可能會影響到演算法效果。對於這種情況，我們提供了一個開關，由業務去控制是否合併Sparse Variables。經過我們的實測，在某推薦模型上合併Sparse Variables會提高5～10%的訓練效能，而對實際業務效果的影響在一個千分點以內。

　　這兩種運算元融合的優化，不僅優化了卡間通訊效能，對卡內計算效能也有一定的提升。經過這兩種運算元融合的優化，GPU架構端到端訓練速度提升了85%，同時不影響業務演算法的效果。

　　4.4 效能指標

　　完成了資料層、計算層、通訊層的效能優化後，對比我們的TensorFlow[3] CPU場景，GPU架構取得了2～4倍的價效比收益（不同業務模型收益不同）。我們基於美團外賣某推薦模型，使用單臺GPU節點（A100單機八卡）和同成本的CPU Cluster，分別對比了原生TensorFlow 1.15和我們優化後的TensorFlow 1.15的訓練效能，具體資料如下：

　　圖12 CPU/GPU訓練吞吐對比

　　可以看到，我們優化後的TensorFlow GPU架構訓練吞吐，是原生TensorFlow GPU的3倍以上，是優化後TensorFlow CPU場景的4倍以上。

　　注：原生TensorFlow使用了tf.Variable作為Embedding的引數儲存。

　　 5 業務落地

　　Booster架構要在業務生產中落地，不只是要有一個良好的系統效能，還需要同時關注訓練生態系統的完備性以及訓練產出模型的效果。

　　5.1 完備性

　　一次完整的模型訓練實驗，除了要跑訓練（Train）任務外，往往還需要跑模型的效果評估（Evaluate）或模型的預估（Predict）任務。我們基於TensorFlow Estimator正規化對訓練架構進行封裝，實現使用者側一套程式碼統一支援GPU和CPU場景下的Train、Evaluate和Predict任務，通過開關進行靈活切換，使用者只需要關注模型程式碼本身的開發。我們將架構改動全都封裝到了引擎內部，使用者只需要一行程式碼就能從CPU場景遷移到GPU架構：

　　tf.enable_gpu_booster()

　　實際業務場景，使用者通常會使用train_and_evaluate模式，在跑訓練任務的過程中同時評估模型效果。上了Booster架構後，由於訓練跑的太快，導致Evaluate速度跟不上訓練正常產出Checkpoint的速度。我們在GPU訓練架構的基礎上，支援了Evaluate on GPU的能力，業務可以申請一顆A100 GPU專門用來做Evaluate，單顆GPU做Evaluate的速度是CPU場景下單個Evaluate程式的40倍。同時，我們也支援了Predict on GPU的能力，單機八卡Predict的速度是同等成本下CPU的3倍。

　　此外，我們在任務資源配置上也提供了比較完善的選項。在單機八卡（A100單臺機器至多配置8張卡）的基礎上，我們支援了單機單卡、雙卡、四卡任務，並打通了單機單卡/雙卡/四卡/八卡/CPU PS架構的Checkpoint，使得使用者能夠在這幾種訓練模式間自由切換、斷點續訓，方便使用者選擇合理的資源型別、資源量跑實驗，同時業務也能夠從已有模型的Checkpoint來WarmStart訓練新的模型。

　　5.2 訓練效果

　　相較PS/Worker非同步模式的CPU訓練，單機多卡訓練時卡間是全同步的，因而避免了非同步訓練梯度更新延遲對訓練效果的影響。然而，由於同步模式下每一步迭代的實際Batch Size是每張卡樣本數的總和，並且為了充分利用A100卡的算力，我們會將每張卡的Batch Size（單步迭代的樣本數）儘量調大。這使得實際訓練的Batch Size（1萬~10萬）比PS/Worker非同步模式（1千~1萬）大很多。我們需要面臨大Batch下訓練超參調優的問題[26,27]：在保證Epoch不變的前提下，擴大Batch Size會導致引數有效更新次數減少，可能導致模型訓練的效果變差。

　　我們採用Linear Scaling Rule[28]的原則指導調整學習率。如果訓練Batch Size較PS/Worker模式的Batch Size增大N倍，將學習率也放大N倍即可。這種方式簡單便於操作，實踐效果還不錯。當然需要注意的是，如果原有訓練方式的學習率已經很激進時，大Batch Size訓練學習率的調整幅度則需要適當減小，或者使用學習率Warmup等更復雜的訓練策略[29]。我們會在後續工作中對超參優化模式做更深入的探索。

　　 6 總結與展望

　　在美團推薦系統訓練場景，隨著模型越來越複雜，CPU上優化的邊際效應越來越低。美團基於內部深度定製的TensorFlow、NVIDIA HugeCTR，研發了Booster GPU訓練架構。整體設計充分考慮演算法、架構、新硬體的特性，並從資料、計算、通訊等多個角度深度優化，對比之前CPU的任務，價效比提升到2~4倍。從功能和完備性上支援TensorFlow的各類訓練介面（Train/Evaluate/Rredict等)，支援CPU和GPU模型相互匯入。易用性上TensorFlow CPU任務只需要一行程式碼就可完成GPU架構遷移。目前在美團外賣推薦場景實現了大規模的投產應用，後續我們將會全面推廣到到家搜尋推薦技術部以及美團全業務線。

　　當然，Booster基於NVIDIA A100單機多卡還有不少優化空間，如資料層面的樣本壓縮、序列化、特徵解析，計算層面的多圖運算元排程、動態shape運算元的編譯優化，通訊層面的量化通訊等等。同時為了更廣泛的支援美團內的業務模型，Booster的下一個版本也會支援更大的模型，以及多機多卡的GPU訓練。

　　7 作者簡介

　　家恆、國慶、崢少、曉光、鵬鵬、永宇、俊文、正陽、瑞東、翔宇、秀峰、王慶、封宇、事峰、黃軍等，來自美團基礎研發平臺-機器學習平臺訓練引擎&到家研發平臺-搜尋推薦技術部Booster聯合專案組。

　　8 參考文獻

　　[1]

　　[2] https://images.nvidia.cn/aem-dam/en-zz/Solutions/data-center/nvidia-ampere-architecture-whitepaper.pdf

　　[3]

　　[4]

　　[5]

　　[6]

　　[7]

　　[8]

　　[9]

　　[10] Yann LeCun, John S. Denker, and Sara A. Solla. Optimal brain damage. In NIPS, pp. 598–605. Morgan Kaufmann, 1989.

　　[11] Kenji Suzuki, Isao Horiba, and Noboru Sugie. A simple neural network pruning algorithm with application to filter synthesis. Neural Process. Lett., 13(1):43–53, 2001.

　　[12] Suraj Srinivas and R. Venkatesh Babu. Data-free parameter pruning for deep neural networks. In BMVC, pp. 31.1–31.12. BMVA Press, 2015.

　　[13] Jonathan Frankle and Michael Carbin. The lottery ticket hypothesis: Finding sparse, trainable neural networks. In 7th International Conference on Learning Representations, ICLR 2019, New Orleans, LA, USA, May 6-9, 2019. OpenReview.net, 2019.

　　[14] Hao-Jun Michael Shi, Dheevatsa Mudigere, Maxim Naumov, and Jiyan Yang. Compositional embeddings using complementary partitions for memory-efficient recommendation systems. In Proceedings of the 26th ACM SIGKDD International Conference on Knowledge Discovery & Data Mining, pp. 165-175. 2020.

　　[15] https://mp.weixin.qq.com/s/fOA_u3TYeSwAeI6C9QW8Yw

　　[16] Jianxun Lian, Xiaohuan Zhou, Fuzheng Zhang, Zhongxia Chen, Xing Xie, and Guangzhong Sun. 2018. xDeepFM: Combining Explicit and Implicit Feature Interactions for Recommender Systems. arXiv preprint arXiv:1803.05170 (2018).

　　[17] Weiping Song, Chence Shi, Zhiping Xiao, Zhijian Duan, Yewen Xu, Ming Zhang, and Jian Tang. Autoint: Automatic feature interaction learning via self-attentive neural networks. In Proceedings of the 28th ACM International Conference on Information and Knowledge Management, pp. 1161-1170. 2019.

　　[18] Guorui Zhou, Weijie Bian, Kailun Wu, Lejian Ren, Qi Pi, Yujing Zhang, Can Xiao et al. CAN: revisiting feature co-action for click-through rate prediction. arXiv preprint arXiv:2011.05625 (2020).

　　[19] Chun-Hao Chang, Ladislav Rampasek, and Anna Goldenberg. Dropout feature ranking for deep learning models. arXiv preprint arXiv:1712.08645 (2017).

　　[20] Xu Ma, Pengjie Wang, Hui Zhao, Shaoguo Liu, Chuhan Zhao, Wei Lin, Kuang-Chih Lee, Jian Xu, and Bo Zheng. Towards a Better Tradeoff between Effectiveness and Efficiency in Pre-Ranking: A Learnable Feature Selection based Approach. In Proceedings of the 44th International ACM SIGIR Conference on Research and Development in Information Retrieval, pp. 2036-2040. 2021.

　　[21] Bencheng Yan, Pengjie Wang, Jinquan Liu, Wei Lin, Kuang-Chih Lee, Jian Xu, and Bo Zheng. Binary Code based Hash Embedding for Web-scale Applications. In Proceedings of the 30th ACM International Conference on Information & Knowledge Management, pp. 3563-3567. 2021.

　　[22] Xiangyu Zhao, Haochen Liu, Hui Liu, Jiliang Tang, Weiwei Guo, Jun Shi, Sida Wang, Huiji Gao, and Bo Long. Autodim: Field-aware embedding dimension searchin recommender systems. In Proceedings of the Web Conference 2021, pp. 3015-3022. 2021.

　　[23] Bencheng Yan, Pengjie Wang, Kai Zhang, Wei Lin, Kuang-Chih Lee, Jian Xu, and Bo Zheng. Learning Effective and Efficient Embedding via an Adaptively-Masked Twins-based Layer. In Proceedings of the 30th ACM International Conference on Information & Knowledge Management, pp. 3568-3572. 2021.

　　[24] Ting Chen, Lala Li, and Yizhou Sun. Differentiable product quantization for end-to-end embedding compression. In International Conference on Machine Learning, pp. 1617-1626. PMLR, 2020.

　　[25] Wang-Cheng Kang, Derek Zhiyuan Cheng, Ting Chen, Xinyang Yi, Dong Lin, Lichan Hong, and Ed H. Chi. Learning multi-granular quantized embeddings for large-vocab categorical features in recommender systems. In Companion Proceedings of the Web Conference 2020, pp. 562-566. 2020.

　　[26] Nitish Shirish Keskar, Dheevatsa Mudigere, Jorge Nocedal, Mikhail Smelyanskiy, and Ping Tak Peter Tang. On large-batch training for deep learning: Generalization gap and sharp minima. arXiv preprint arXiv:1609.04836 (2016).

　　[27] Elad Hoffer, Itay Hubara, and Daniel Soudry. Train longer, generalize better: closing the generalization gap in large batch training of neural networks. Advances in neural information processing systems 30 (2017).

　　[28] Priya Goyal, Piotr Dollár, Ross Girshick, Pieter Noordhuis, Lukasz Wesolowski, Aapo Kyrola, Andrew Tulloch, Yangqing Jia, and Kaiming He. Accurate, large minibatch sgd: Training imagenet in 1 hour. arXiv preprint arXiv:1706.02677 (2017).

　　[29] Chao Peng, Tete Xiao, Zeming Li, Yuning Jiang, Xiangyu Zhang, Kai Jia, Gang Yu, and Jian Sun. Megdet: A large mini-batch object detector. In Proceedings of the IEEE conference on Computer Vision and Pattern Recognition, pp. 6181-6189. 2018.