介紹

通過封裝 NVIDIA Merlin HugeCTR，Sparse Operation Kit（以下簡稱 SOK）使得 TensorFlow 用戶可以借助 HugeCTR 的一些相關(guān)特性和優(yōu)化加速 GPU 上的分布式 Embedding 訓練。

在以往文章中（Merlin HugeCTR Sparse Operation Kit 系列之一、Merlin HugeCTR Sparse Operation Kit 系列之二）我們對 HugeCTR SOK 的基本功能、性能、用法和原理做了詳細的介紹。近期 SOK 又發(fā)布了多個版本迭代，這篇博客對最新 v2.0 版本中的新特性（尤其是動態(tài) Embedding 和在線訓練增量導出）、用法進行了歸納總結(jié)和介紹，并在最后介紹了 SOK 在手機行業(yè)的應(yīng)用案例。

圖 1. SOK 訓練的數(shù)據(jù)并行

- 模型并行-數(shù)據(jù)并行流程

特性

SOK 作為 TensorFlow 的 Plugin，專注于針對 Embedding 的儲存和查詢過程進行優(yōu)化，以下是目前 SOK 提供的主要功能：

提供 Embedding 部分的多機多卡模型并行功能；

為 Embedding Table 提供了 Hash Table 的功能，并且該 Hash Table 可以將 Embedding Vector 存放到 Device Memory 和 Host Memory 上，充分利用設(shè)備儲存；

為 Embedding Lookup 提供了高性能的 Embedding Lookup Kernel，采用 HugeCTR 的 Cuda Kernel 為后端，F(xiàn)use 多個 Lookup Op 到一個 Lookup Op，提升 Embedding Lookup 性能；

提供增量導出功能，導出某個時間點后 Embedding Table 中更新過的 Key 和 Value，減少訓練權(quán)重推送到推理端的開銷。

立即使用 SOK

Embedding Table

SOK 基于 TensorFlow 的 Variable 的基礎(chǔ)上，提供了 Variable 和 DynamicVariable 來儲存 Embedding Table。

不同于 TensorFlow 的 Variable 將數(shù)據(jù)儲存到某個 GPU 中，SOK 的 Variable 可以將數(shù)據(jù)均勻儲存在訓練的所有 GPU 上，也可以將數(shù)據(jù)儲存到某一個 GPU 中，下面是使用 SOK 的 Variable 創(chuàng)建儲存在每一個 GPU 和 Variable 儲存到 GPU0 上的例子：

#Defaultmethodofsok.VariableisDistributedmethod,VariableevenlydistributedtoeachGPU
v1 = sok.Variable(np.arange(15 * 16).reshape(15, 16), dtype=tf.float32)


#If you want to assign a sok.Variable to a specific GPU, add the parameter mode=“l(fā)ocalized:gpu_id” when defining sok.variable, where gpu_id refers to the rank number of a GPU in Horovod
v2=sok.Variable(np.arange(15*16).reshape(15,16),dtype=tf.float32,mode="localized:0")

v1 的申請中，SOK 會自動將 v1 的數(shù)據(jù)均勻儲存到訓練域中每一個 GPU 上，v2 的申請中，在參數(shù)中添加了mode="localized:0" 這個字符串參數(shù)，SOK 會將 v2 的數(shù)據(jù)放置在 GPU0 中。

SOK 中提供了 Variable 和 DynamicVariable 兩種不同的形式儲存 Embedding Table。Variable 可以簡稱為 Static Embedding Table，一個是靜態(tài)的二維數(shù)據(jù)，在申請時需要確定 Embedding Table 的 Shape，申請結(jié)束后，SOK 會相應(yīng)申請好數(shù)據(jù)使用的空間（如果 Static Embedding Table 分布在所有 GPU 上，那么所有 GPU 平分這個空間）。因為 Static Embedding Table 是一個靜態(tài)的二維數(shù)組，這也意味著 Lookup 時查找的 Indice，是無法超出這個二維數(shù)組第一個 Dimension 的長度的，否則會發(fā)生越界問題。

Static Embedding Table，對于 Lookup Indice 存在范圍限制，很多用戶輸入的 Lookup Indice 又是 Hash 過的 Key，所以 SOK 提出 DynamicVariable 來解決 Lookup Indice 范圍限制的問題。SOK 的 DynamicVariable 是使用 Hash Table 來儲存 Embedding Table 的，解決了這個問題。

SOK 的 DynamicVariable 封裝了 2 個基于 GPU 的 Hash Table（HierarchicalKV 和 Dynamic Embedding Table），其中 HierarchicalKV（以下簡稱 HKV）是 NVIDIA Merlin 框架下的 Hash Table Repo，它有以下 2 個特性：

HKV 可以利用 GPU Memory 和 Host Memory 儲存 Embedding Table 中的 Embedding Vector，充分利用訓練中的內(nèi)存資源；

HKV 的 Hash Table 擁有 Eviction 功能，Hash Table 滿了后不會繼續(xù)增長，繼續(xù) Insert 會淘汰掉最不常用的 Key/Value，提供 LRU、LFU 等常用淘汰策略，也可以自定義淘汰策略。

HKV 的 Repo 地址為 HKV repo，提供了 C++ Level 的 Hash Table API，也為推薦系統(tǒng)提供了組合 API，方便推薦系統(tǒng)用戶使用。SOK 中申請 HKV Embedding Table 的代碼如下：

#init_capacityandmax_capacityareparametersacceptedbytheHKVtable.ThemeaningsoftheseparameterscanbefoundintheHKVdocumentation.
v2 = sok.DynamicVariable(
  dimension=16, # embedding vector length
  var_type="hybrid", # use HKV backend
  initializer="uniform", # use uniform distribution random initializer
  init_capacity=1024 * 1024, # The number of embedding vectors allocated initially in the hash table
  max_capacity=1024 * 1024, # The number of embedding vectors allocated finally in the hash table
max_hbm_for_vectors=2,#howmanyGPUmemoryshouldthishashtableuse,unitisGB

從上面代碼可見，SOK 申請 HKV Hash Table 時，需要設(shè)置一些參數(shù)，這些參數(shù)如何設(shè)置可以閱讀 SOK 和 HKV 的文檔。

SOK 的 Variable 和 DynamicVariable 均繼承自 TensorFlow 的 ResourceVariable，除了 SOK 的自己定義的參數(shù)外，讀者和用戶可以按照使用 ResourceVariable 的習慣添加 ResourceVariable 的參數(shù)。

Embedding Lookup

SOK embedding lookup_sparse 提供與 tf.nn.embedding_lookup_sparse 相似的 API，與 tf.nn.embedding_lookup_sparse 不同的是，SOK 的 embedding lookup_sparse 可以同時將多個 Lookup Fuse 到一起，代碼如下：

#UseSOKembeddinglookupsparsetodo2embeddinglookups
    emb1, emb2 = sok.lookup_sparse([sok_var1, sok_var2],[keys1, keys2],combiners=["sum", "mean"])
# equals to
    emb1 = tf.nn.embedding_lookup_sparse(v1, indices1, combiner="sum")
emb2=tf.nn.embedding_lookup_sparse(v2,indices2,combiner="mean")

Embedding Table Dump/Load

SOK 關(guān)于 Embedding Table 的權(quán)重提供了 Dump/Load 和增量導出 Incremental Dump 的功能。

Dump/Load 可以自動并行的進行 Embedding Table 的 Key、Value、Optimizer 中狀態(tài)變量（可選）在文件系統(tǒng)中的讀寫，將 Embedding Table 的 Key、Value、Optimizer 中的狀態(tài)變量儲存成二進制文件/從二進制文件中讀取，下面是 SOK 中 Dump/Load 的例子：

#optimizerstatesareoptional.IftheyareunspecifiedincallingtheAPIsabove,onlythekeysandvaluesareloaded.
optimizer = tf.keras.optimizers.SGD(learning_rate=1.0)
sok_optimizer = sok.OptimizerWrapper(optimizer)
path = "./weights"
sok_vars = [sok_var1,sok_var2]
sok.dump(path, sok_vars, sok_optimizer)
sok.load(path,sok_vars,sok_optimizer)

在大部分推薦系統(tǒng)業(yè)務(wù)中，Embedding Table 的內(nèi)存占用非常大，因此，在 Continued Training 中，用戶通常會將訓練一段時間后更新了的 Key 和 Value 推送到推理端，這樣可以避免推送整個 Embedding Table 產(chǎn)生的巨大開銷，SOK 同樣提供了 incremental_dump 的 API 來實現(xiàn)這個功能。incremental_dump 接受一個 UTC Time Threshold，可以將 Time Threshold 后更新的 Key/Value 導出到 Numpy Array 中：

#sokincrementaldump
import pytz
from datetime import datetime


#should convert datatime to utc time
utc_time_threshold = datetime.now(pytz.utc)


sok_vars = [sok_var1,sok_var2]
#keys and values are Numpy array
keys,values=sok.incremental_model_dump(sok_vars,utc_time_threshold)

SOK 與 TensorFlow 的兼容性

SOK 目前兼容 TF2 的靜態(tài)圖，但是不支持 TensorFlow 的 XLA，如果開啟 TensorFlow 的 XLA，需要手動將 SOK 的 Lookup 的 Layer 排除在外，偽代碼如下所示：

@tf.function
def sok_layer(inputs): 
     return sok.lookup_sparse(inputs)


@tf.function(jit_compile=True) 
def xla_layer(inputs): 
     x = xla_layer(inputs) 
returnx

應(yīng)用案例分享

以下是近期 NVIDIA 技術(shù)團隊開展的部分手機行業(yè)推薦場景高性能優(yōu)化項目實踐經(jīng)驗分享：

案例 1：通過 NVIDIA Merlin

HugeCTR SOK 以及 NVTabular

實現(xiàn)了 GPU 加速的推薦系統(tǒng)

應(yīng)用背景：

客戶的推薦系統(tǒng)是針對客戶手機端的廣告和內(nèi)容，該推薦系統(tǒng)之前是使用 CPU PS 架構(gòu)運行，客戶希望使用 GPU 架構(gòu)來進行加速。

應(yīng)用方案/效果以及影響：

通過 HugeCTR Sparse Operation Kit + NVTabular 實現(xiàn)了 GPU 加速的推薦系統(tǒng)，性能加速如下：

實驗性能對比：

SOK 將原有只能單 GPU 訓練的任務(wù)擴展到多 GPU 訓練，并且達到了很好的弱擴展性，支持更大的模型。SOK + tfrecord 3GPU 耗時約 1.3 個小時 vs TF + tfrecord 1GPU 耗時約 3.4 個小時。

業(yè)務(wù)性能對比：

由于業(yè)務(wù)系統(tǒng)輸入數(shù)據(jù)較大，因此客戶采取 Parquet 數(shù)據(jù)格式進行 Input 數(shù)據(jù)壓縮，業(yè)務(wù)性能提升如下，NVTabular Load Parquet 耗時約 1.8 個小時 vs TF Load Parquet 耗時約 7.8 個小時。

在 SOK 加速的基礎(chǔ)上，NVTabular Parguet Datareader 進一步解決了數(shù)據(jù)讀取瓶頸的問題，在實際業(yè)務(wù)測試中，相比原生 TensorFlow Parquet Datareader 達到了 400% 的速度提升。

案例 2：通過 NVIDIA Merlin

HugeCTR SOK 實現(xiàn)了

GPU 加速的推薦系統(tǒng)

應(yīng)用背景：

客戶的推薦系統(tǒng)是針對客戶手機端的廣告和內(nèi)容，該推薦系統(tǒng)使用 CPU + PS 和 GPU + PS 的架構(gòu)運行。

應(yīng)用方案/效果以及影響：

通過 HugeCTR Sparse Operation Kit 進行了 GPU 加速的推薦系統(tǒng)的實驗，性能加速如下：

End to End 時間從 CPU + PS 的 760 毫秒優(yōu)化至 300 毫秒，性能提升約 60%。GPU + PS 的 450 毫秒優(yōu)化至 300 毫秒，性能提升約 33%。

GPU + PS 架構(gòu)的 Embedding 部分耗時為 250 毫秒，SOK 的 Embedding 部分耗時為 85 毫秒，性能提升 67%。

結(jié)束語

SOK v2.0 通過封裝 HKV 和 HugeCTR 的底層代碼提供了模型并行，且功能完善的動態(tài) Embedding Table 和高效的相關(guān)計算。在多個業(yè)內(nèi)實際使用場景中，從速度效率、模型擴展、功能完善等角度，使用戶獲得了不錯的收益。

關(guān)于作者

康暉

GPU 計算專家，2022 年加入 NVIDIA，當前主要從事 SOK 的設(shè)計與開發(fā)。研究領(lǐng)域包括推薦系統(tǒng)工程實現(xiàn)與優(yōu)化，通訊優(yōu)化。