一塊GPU訓練TB級推薦模型不是夢,OneEmbedding效能一騎絕塵

語言: CN / TW / HK

來源 | 機器之心

個性化推薦已成為人們獲取資訊的主要形式。以往,人們更多通過主動搜尋自己感興趣的資訊,而現在,基於演算法推薦技術的資訊分發平臺會自動識別使用者興趣,快速篩選資訊,推送使用者所感興趣的資訊。

一方面,推薦系統大幅提升了使用者體驗,另一方面,個性化分發資訊更精準、高效,可以幫助平臺更準確地匹配使用者和資訊,大大提高流量變現效率,基於推薦技術的流量變現引擎甚至成就了萬億市值的龐大商業帝國。

從短視訊資訊流推薦、廣告搜尋到線上購物,這些應用都構築於精準的推薦系統之上,背後的核心功臣就是深度學習模型。

不過,隨著海量資料的積累以及更加頻繁的使用者資料迭代,底層系統可擴充套件性和訓練速度面臨嚴峻的挑戰。人們發現,通用深度學習框架都不能直接滿足工業級推薦系統的需求,而是必須基於通用深度學習框架做深度定製,甚至於要開發專門的系統才行。

針對現代推薦系統的種種痛點,OneFlow 團隊推出了一款高效能、可擴充套件、靈活度高的推薦系統元件 OneEmbedding。它的使用方式和通用深度學習框架一樣簡單,效能卻遠超通用框架,甚至超過了 NVIDIA HugeCTR 這樣為推薦場景定製開發的系統。

具體而言,在 DCN、DeepFM 兩個模型上,無論是 FP32 還是混合精度(automatic mixed-precision, AMP)訓練,OneEmbedding 的效能大幅超過 HugeCTR,而在 HugeCTR 深度優化以至於有點 “過擬合” 的 DLRM 模型上,OneEmbedding 效能與 HugeCTR 基本持平。

(以上測試環境均為:CPU Intel(R) Xeon(R) Platinum 8336C CPU @ 2.30GHz * 2;CPU Memory 1920GB;GPU NVIDIA A100-SXM-80GB * 8;SSD Intel SSD D7P5510 Series 3.84TB * 4)

當用戶使用 OneFlow 搭建推薦模型時,只需使用以下數行程式碼對 Embedding 詞表進行配置即可訓練含有 TB 級別詞表的推薦模型:

```

self.embedding = nn.Embedding(vocab_size, embedding_vec_size)

self.embedding = flow.one_embedding.MultiTableEmbedding( "sparse_embedding", embedding_dim=embedding_vec_size, dtype=flow.float, key_type=flow.int64, tables=tables, store_options=store_options, ) ```

基於 OneEmbedding 搭建的常見搜尋推薦廣告模型案例地址:\ https://github.com/Oneflow-Inc/models/tree/main/RecommenderSystems

1

大規模推薦系統的挑戰

一般而言,推薦系統需要使用類似性別、年齡、行為等方面的離散特徵(sparse feature),在一個 Embedding 詞表中用特徵 ID 進行查表(lookup),取得對應的 Embedding 向量並送到下游使用。

常用的公開資料集 Criteo1T 中大概包含十億個特徵 ID,如果 embedding_dims 配置為 128,那總共需要 512 GB 空間來容納 Embedding 引數,如果使用 Adam 優化器,由於需要儲存額外的兩個狀態變數 m 和 v,所需儲存容量就增加到 1536 GB。實際應用場景中,資料規模比 Criteo 還要高出幾個數量級,模型的容量就更大了。

大規模推薦系統的核心問題就是,如何高效經濟地支援大規模 Embedding 的查詢和更新。 權衡規模、成本和效率,出現瞭如下三種常見的解決方案。

最常見也是最早出現的一種解決方案是將 Embedding 全部部署在 CPU 上,利用 CPU 記憶體容量大、成本低的特點擴充套件引數規模,優點是模型規模幾乎可以無限大。不過,其缺點也很明顯,無論是計算效能還是頻寬,CPU 都遠低於 GPU,導致 Embedding 部分成為顯著瓶頸,往往需要數十乃至上百臺 CPU 伺服器才能支撐一個工業級的推薦系統。

鑑於 GPU 在稠密計算中得天獨厚的優勢,也有人建議用 GPU 來訓練大型 Embedding 模型。問題是,GPU 很貴且視訊記憶體容量有限,如果使用視訊記憶體容量為 40GB 的 A100 來基於 Criteo 資料訓練 128 維嵌入向量,至少需要 13 張顯示卡才能放下 512GB 的 Embedding 詞表。每張卡只有 40GB 視訊記憶體容量,分散式 Embedding 需要使用所謂的模型並行技術,理想情況下,為了解決更大規模的模型只需要增加 GPU 的數量即可。

現實是,GPU 相對於 CPU 的成本非常高昂,並且推薦系統中模型主體計算部分不大,模型並行在擴充套件過程中只是解決了 Embedding 規模的問題,訓練速度的收益比較有限,甚至會因為多裝置之間引入通訊導致訓練速度下降,因此通常只適用於小規模叢集。

為了緩解 GPU 之間傳輸頻寬的問題,業界發展出比乙太網頻寬更高的 NVSwitch、Infiniband 網路等互聯技術。一方面,這意味著額外的成本,另一方面,很多使用者的基礎設施不具備相應改造、升級的條件。

那麼,有沒有魚和熊掌兼得的方案?

針對上述方案存在的問題,OneFlow 團隊設計了 OneEmbedding,通過分層儲存讓單卡也能支援 TB 級模型的訓練,通過橫向擴充套件讓模型容量沒有天花板,通過 OneFlow 的自動流水線機制、運算元優化和通訊量化壓縮等技術實現極致效能,在用法像 PyTorch 一樣簡單的前提下, OneEmbedding 在 DLRM 模型上效能是 TorchRec 的 3 倍以上,開啟 TorchRec 沒有支援的混合精度後,OneEmbedding 的效能更是 TorchRec 的 7 倍以上。

(TorchRec 效能資料參考 8 卡 A100 測試結果:https://github.com/facebookresearch/dlrm/tree/main/torchrec_dlrm/#preliminary-training-results

2

OneEmbedding 的核心優勢

分層儲存:單卡也能支援 TB 級模型訓練

利用資料的空間區域性性和時間區域性性,多級快取可以很好地實現效能和成本的折衷。OneEmbedding 也基於這個思想實現了多級快取,即使使用者只有一張 GPU 也可以訓練 TB 級別的模型。

使用者既可以把 Embedding 部署到 GPU 視訊記憶體上,也可以把 Embedding 部署到 CPU 記憶體甚至是 SSD 上。這種方案可以發揮 CPU 記憶體或者 SSD 更低的成本優勢,以對 Embedding 引數規模進行擴充套件,還可以利用 GPU 視訊記憶體作為快取記憶體裝置,以實現高效能效果。

OneEmbedding 會動態地將最近頻繁訪問的條目快取到 GPU 視訊記憶體上,同時將最近訪問頻率較低的條目逐出到 CPU 記憶體或者 SSD 等底層儲存中。在資料遵循冪律分佈這一前提下,基於有效的 Cache 管理演算法,OneEmbedding 可以使 GPU 快取的命中率始終維持在較高的水平。

值得強調的是,OneEmbedding 只是將 CPU 記憶體和 SSD 作為儲存裝置,所有計算都在 GPU 上執行。目前,OneEmbedding 提供三種預置儲存方案:

  • 使用 GPU 視訊記憶體儲存模型所有引數
  • 將 CPU 記憶體作為 Embedding 引數儲存裝置,並搭配使用 GPU 作為快取記憶體
  • 將 SSD 作為 Embedding 引數儲存裝置,並搭配使用 GPU 作為快取記憶體

```

使用 SSD 作為儲存裝置,搭配 GPU 作為快取記憶體

store_options = flow.one_embedding.make_cached_ssd_store_options( cache_budget_mb=cache_memory_budget_mb, persistent_path=persistent_path, capacity=vocab_size, ) ```

使用者可以根據實際使用時的硬體裝置情況,只需用短短數行程式碼進行配置,即可一箭三雕實現規模、效率、成本的最優化。

為了掩蓋 CPU 和 SSD 取資料的延遲,OneEmbedding 引入流水線、資料預取等技術手段,使得在以 CPU 記憶體 和 SSD 作為儲存後端的同時,效率依然可以和使用純 GPU 訓練那樣保持在同樣的水平。

分別對三種儲存方案進行測試。其中,測試用例與 MLPerf 的 DLRM 模型一致,引數規模約為 90GB。在使用 SSD 和 CPU 記憶體作為儲存裝置時,我們配置的 GPU 快取大小為每個 GPU 12GB,相比於 90 GB 的總引數量,只能有一部分引數保持在 GPU 視訊記憶體內,其他的引數則是儲存在 CPU 記憶體或者 SSD 上,隨著訓練過程動態的換入到 GPU 快取中來,測試結果如下圖。

(測試環境:CPU Intel(R) Xeon(R) Silver 4214R CPU @ 2.40GHz * 2;CPU Memory 512GB;GPU NVIDIA A100-PCIE-40GB * 4;SSD Intel SSD D7P5510 Series 7.68TB * 4)

從測試結果中可以看到:

(1)純 GPU 視訊記憶體方案效能最佳,但由於 GPU 視訊記憶體只有 4x40GB,理論上最大隻能夠訓練 160 GB 模型;

(2)相比純 GPU 視訊記憶體方案,GPU 快取 + CPU 儲存的方案效能只有微小損失,但可以把引數規模的天花板擴充套件到 CPU 記憶體容量,往往是數百 GB~數 TB;

(3)更進一步,如果能接受更大的效能損失,GPU 快取 + SSD 儲存的方案能將引數規模的天花板擴充套件到 SSD 的容量,模型規模可達數十 TB,甚至更大。

如果我們想在只有一個 NVIDIA A30-24GB GPU 的伺服器上進行上述 DLRM 模型的完整訓練,24G 的視訊記憶體顯然無法直接訓練 90GB 規模的模型。藉助分層儲存,使用 CPU 記憶體作為儲存裝置,GPU 視訊記憶體作為快取記憶體,就可以支援比 90GB 還大的模型。

橫向擴充套件:多卡線性加速,打破模型天花板

使用分層儲存技術,OneEmbedding 提升了單卡情況下的 Embedding 引數規模極限,只要記憶體空間夠大甚至能夠訓練 TB 級別大小的模型。如果模型的容量進一步擴大到甚至大大超過 CPU 記憶體的容量,使用者還可以在多級儲存的基礎上藉助 OneFlow 的並行能力輕鬆地橫向拓展到多機多卡,以訓練更大的模型。

在推薦系統中,模型主體引數相較 Embedding 則小的多。因此我們一般將 Embedding 部分設定為模型並行,模型主體設定為資料並行。通過使用多機多卡,可進一步提升 Embedding 大小。

具體到實現細節,每個 Rank 各自負責一部分 Embedding 的儲存,特徵 ID 進入到各個 Rank,可能存在重複 ID 的情況,首先要進行去重(即下圖的 ID Shuffle);各個 Rank 拿著去重後的 ID 去查詢 Embedding,得到對應的區域性資料,所有 Rank 資料合併後各 Rank 得到完整的 Embedding 資料(即下圖的 Embedding Shuffle);最後,各 Rank 以資料並行的方式完成整個模型訓練過程。

下圖展示了 OneEmbedding 採取純 GPU 視訊記憶體的策略訓練 DLRM 模型時,FP32 和 AMP 配置下,不同 GPU 個數下模型吞吐量。

(測試環境:CPU Intel(R) Xeon(R) Platinum 8336C CPU @ 2.30GHz * 2;CPU Memory 1920GB;GPU NVIDIA A100-SXM-80GB * 8;SSD Intel SSD D7P5510 Series 3.84TB * 4)

可以看到,隨著 GPU 裝置數的增加,模型吞吐量均能顯著增加,在混合精度情況下,單張 GPU 能有 600 萬的吞吐量,當擴充套件到 8 張 GPU 時能有近 4000 萬的吞吐量。

流水線機制:自動重疊計算和資料傳輸

在 DLRM 模型中,Embedding 中的 Dense Feature 會進入到 Bottom MLP 中,而 Sparse Feature 經過 Embedding 查詢得到對應特徵。兩者進入 Interaction 進行特徵交叉,最後進入到 Top MLP。

Embedding 相關操作包含查表(Embedding Lookup)、更新(Embedding Update)。由於 OneEmbedding 使用的是分層儲存的機制,可能會遇到特徵 ID 沒有命中快取記憶體的情況,此時,資料拉取耗時較長,會影響訓練速度。

為避免這個不足,OneEmbedding 加入資料預取(Embedding Prefetch)操作,以保證查表和更新操作均能在 GPU 上執行。由於前後迭代之間的資料預取不存在依賴關係,在當前迭代計算的同時,可以預取下一個迭代需要的 Embedding 資料,實現計算和預取的重疊。

在 Embedding 資料查詢交換的過程中,與 Embedding 操作無關的 Dense Feature 可以進入到 Bottom MLP 進行計算,在時間上進行重疊。完整的重疊執行時序如下圖所示。

如此複雜的資料流水線控制在傳統深度學習框架裡是一個很挑戰的問題。不僅如此,在實際推薦場景中,使用者的資料在不斷變化,這要求流水線機制還要能應對動態資料。

而 OneFlow 的 Actor 機制讓這一切問題變得很簡單,每個 Actor 都通過自己內部的狀態機和訊息機制實現了分散式協同工作。通過為每個 Actor 賦予多份儲存塊,不同的 Actor 可以同時工作,重疊各自的工作時間,從而實現 Actor 之間的流水線。我們只需要將 Embedding 操作分配到單獨的一個 stream 上,即可讓系統自發地形成流水線。

運算元優化:逼近 GPU 極限效能

OneFlow 團隊不僅對通用運算元進行了深度優化,還針對流行的推薦系統模型特點,增加了多個高效能 CUDA 運算元實現。

對於 DLRM、DCN 模型中的特徵交叉部分,OneFlow 分別實現了\ FusedDotFeatureInteraction 和 FusedCrossFeatureInteraction 運算元。

(FusedCrossFeatureInteraction 運算元,圖片出自 《Deep & Cross Network for Ad Click Predictions》)

對於模型中多個全連線層部分,OneFlow 基於 cublasLt 矩陣運算庫實現了 FusedMLP 運算元。

而對於帶 Dropout 操作的全連線層,OneFlow 深度定製了其中的 ReluDropout 操作,使用 bitmask 形式儲存前向產生的 mask,在反向傳播中,通過設定 cublasLt 矩陣乘的引數 alpha=dropout_scale 以實現反向運算元融合。

量化壓縮:壓榨通訊效率

在模型訓練的通訊過程中,近期也有不少工作對資料進行量化壓縮以節省通訊量,提高通訊效率,OneEmbedding 也支援這個特性。

並行訓練中,各個 Rank 之間需要通訊以交換 Embedding 資料,我們先將浮點型別資料量化成 int8 型別,交換完後再反量化恢復。

下圖以 DLRM 模型為例展示了選擇純 GPU 視訊記憶體儲存配置,分別測試在 FP32 和混合精度條件下量化前後模型吞吐量。

量化前後模型精度對比(AUC):

(測試環境:CPU Intel(R) Xeon(R) Silver 4214R CPU @ 2.40GHz * 2;CPU Memory 512GB;GPU NVIDIA A100-PCIE-40GB * 4;SSD Intel SSD D7P5510 Series 7.68TB * 4)

測試結果表明,在不影響模型精度的前提下, 量化通訊在 FP32 的情況下相比預設通訊模式有 64% 的提升, 在混合精度的情況下有 13% 的提升。

易用性:構建大規模推薦模型就像使用 PyTorch 一樣簡單

OneEmbedding 作為 OneFlow 的一個內部擴充套件元件,意味著使用者可以在使用 OneEmbedding 的高階特性同時,還能享受 OneFlow 通用框架的靈活性構建自己的推薦模型。

class DLRMModule(nn.Module): def __init__(self, args): super(DLRMModule, self).__init__() self.bottom_mlp = FusedMLP(...) self.embedding = OneEmbedding(...) self.interaction = FusedDotInteraction(...) self.top_mlp = FusedMLP(...) def forward(self, sparse_feature, dense_feature): dense_fields = self.bottom_mlp(dense_feature) embedding = self.embedding(sparse_feature) features = self.interaction(dense_fields, embedding) return self.top_mlp(features)

最後,值得一提的是,OneEmbedding 通過內建的編碼機制對特徵 ID 進行編碼,支援動態插入新增資料。使用者不需要提前規劃 Embedding 容量,也無需對資料集中的特徵 ID 進行特殊處理。這種動態機制天然地支援增量訓練的場景,同時也減少使用負擔。

目前 OneFlow 的 models 倉庫下提供了基於 OneEmbedding 搭建的一系列模型,如 DLRM, DeepFM, xDeepFM, DCN, PNN, MMoE,後續也會補充更多的推薦模型

https://github.com/Oneflow-Inc/models/tree/main/RecommenderSystems )。

3

結語

OneEmbedding 是應訓練大規模推薦系統模型需求而生的元件,靈活的分層儲存、高度優化的資料流水線以及易於橫向擴充套件的特性,能讓使用者輕鬆訓練 TB 級別的推薦模型。

目前,OneFlow 框架提供了一些模型示例供你一鍵體驗 OneEmbedding。後續,OneFlow 團隊將推出覆蓋業界主流模型的推薦系統模型庫 Flow-Recommender ,它不僅支援分散式訓練,還支援分散式推理。歡迎感興趣的朋友關注。

  • OneEmbedding 地址:
  • https://github.com/Oneflow-Inc/models/tree/main/RecommenderSystems
  • OneEmbedding 文件:
  • https://docs.oneflow.org/master/cookies/one_embedding.html
  • OneEmbedding API 文件:
  • https://oneflow.readthedocs.io/en/master/one_embedding.html

\ \ 歡迎下載體驗 OneFlow v0.8.0 最新版本: https://github.com/Oneflow-Inc/oneflow/