美團機器學習平臺基於內部深度定製的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計算用到的各種數據提前放置到顯存中。而從16年~20年，NVIDIA Tesla GPU卡[5]計算能力提升了10倍以上，但顯存大小隻提升了3倍左右。
- 其它維度資源並不是很充足：相比GPU算力的提升速度，單機的CPU、網絡帶寬的增長速度較慢，如果遇到這兩類資源workload較重的模型，將無法充分發揮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架構的架構圖如下所示：
整個系統主要包括三個核心模塊：數據模塊，計算模塊，通信模塊：
- 數據模塊：美團自研了一套支持多數據源、多框架的數據分發系統，在GPU系統上，我們改造數據模塊支持了多網卡數據下載，以及考慮到NUMA Awareness的特性，在每顆CPU上都部署了一個數據分發服務。
- 計算模塊：每張GPU卡啓動一個TensorFlow訓練進程執行訓練。
- 通信模塊：我們使用了Horovod[7]來做分佈式訓練的卡間通信，我們在每個節點上啓動一個Horovod進程來執行對應的通信任務。
上述的設計，符合TensorFlow和Horovod原生的設計範式。幾個核心模塊可以相互解耦，獨立迭代，而且如果合併開源社區的最新特性，也不會對系統造成架構性的衝擊。
我們再來看一下整個系統的簡要執行流程，每張GPU卡上啓動的TensorFlow進程內部的執行邏輯如下圖：
整個訓練流程涉及參數存儲、優化器、卡間通信等幾個關鍵模塊。對於樣本的輸入特徵，我們分爲稀疏特徵（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到其他卡中，每張卡拿到自己的稀疏梯度後再執行稀疏優化器，完成大規模稀疏特徵的優化。詳細流程如下圖所示：
- 稀疏特徵（規模較小，使用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的壓力會變得更大。我們先看一下在當前系統下的樣本數據流程，如下圖所示：
核心流程：數據分發進程通過網絡讀取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結果如下圖：
究其原因，CTR場景的訓練樣本通常包含了大量的int64類型的特徵，int64在PB中是以Varint64類型數據存儲的，ReadVarint64Fallback方法就是用來解析int64類型的特徵。普通的int64數據類型需要佔用8個字節，而Varint64針對不同的數據範圍，使用了變長的存儲長度。PB在解析Varint類型數據時，首先要確定當前數據的長度，Varint用7bit存儲數據，高位1bit存儲標記位，該標記位表示下一個字節是否有效，如果當前字節最高位爲0，則說明當前Varint數據在該字節處結束。我們實際業務場景的ID特徵大多是經過Hash後的值，用Varint64類型表達會比較長，這也就導致在特徵解析過程中要多次判斷數據是否結束，以及多次位移和拼接來生成最終數據，這使得CPU在解析過程中存在大量的分支預測和臨時變量，非常影響性能。以下是4字節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字節數據）：
對於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上之後才能執行模型訓練，如下圖所示：
對於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主要包含稠密部分子網絡計算、梯度計算、稠密參數部分更新等環節。
兩張子圖的交互關係爲：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使用率也非常低，如下圖所示：
從圖中可以看出，訓練期間卡間通信耗時比較長，同時在通信期間GPU使用率也非常低，卡間通信是影響訓練性能提升的關鍵瓶頸點。我們對通信過程進行拆解打點後發現，卡間通信（AllToAll、AllReduce、AllGather等）協商的時間遠遠高於數據傳輸的時間：
分析具體原因，以負責大規模稀疏參數通信的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的訓練性能，具體數據如下：
可以看到，我們優化後的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] https://tech.meituan.com/2021/12/09/meituan-tensorflow-in-recommender-systems.html
- [2] https://images.nvidia.cn/aem-dam/en-zz/Solutions/data-center/nvidia-ampere-architecture-whitepaper.pdf
- [3] https://www.usenix.org/system/files/conference/osdi14/osdi14-paper-li_mu.pdf
- [4] https://github.com/NVIDIA-Merlin/HugeCTR
- [5] https://en.wikipedia.org/wiki/Nvidia_Tesla
- [6] https://www.nvidia.com/en-us/data-center/dgx-a100
- [7] https://github.com/horovod/horovod
- [8] https://github.com/NVIDIA/nccl
- [9] https://www.tensorflow.org/xla
- [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.
閱讀美團技術團隊更多技術文章合集
前端 | 算法 | 後端 | 數據 | 安全 | 運維 | iOS | Android | 測試
| 在公衆號菜單欄對話框回覆【2021年貨】、【2020年貨】、【2019年貨】、【2018年貨】、【2017年貨】等關鍵詞，可查看美團技術團隊歷年技術文章合集。
