SOK在手機(jī)行業(yè)的應(yīng)用案例

介紹

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

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

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

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

特性

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

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

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

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

提供增量導(dǎo)出功能，導(dǎo)出某個(gè)時(shí)間點(diǎn)后 Embedding Table 中更新過(guò)的 Key 和 Value，減少訓(xùn)練權(quán)重推送到推理端的開(kāi)銷。

立即使用 SOK

Embedding Table

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

不同于 TensorFlow 的 Variable 將數(shù)據(jù)儲(chǔ)存到某個(gè) GPU 中，SOK 的 Variable 可以將數(shù)據(jù)均勻儲(chǔ)存在訓(xùn)練的所有 GPU 上，也可以將數(shù)據(jù)儲(chǔ)存到某一個(gè) GPU 中，下面是使用 SOK 的 Variable 創(chuàng)建儲(chǔ)存在每一個(gè) GPU 和 Variable 儲(chǔ)存到 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 的申請(qǐng)中，SOK 會(huì)自動(dòng)將 v1 的數(shù)據(jù)均勻儲(chǔ)存到訓(xùn)練域中每一個(gè) GPU 上，v2 的申請(qǐng)中，在參數(shù)中添加了mode="localized:0" 這個(gè)字符串參數(shù)，SOK 會(huì)將 v2 的數(shù)據(jù)放置在 GPU0 中。

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

Static Embedding Table，對(duì)于 Lookup Indice 存在范圍限制，很多用戶輸入的 Lookup Indice 又是 Hash 過(guò)的 Key，所以 SOK 提出 DynamicVariable 來(lái)解決 Lookup Indice 范圍限制的問(wèn)題。SOK 的 DynamicVariable 是使用 Hash Table 來(lái)儲(chǔ)存 Embedding Table 的，解決了這個(gè)問(wèn)題。

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

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

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

HKV 的 Repo 地址為 HKV repo，提供了 C++ Level 的 Hash Table API，也為推薦系統(tǒng)提供了組合 API，方便推薦系統(tǒng)用戶使用。SOK 中申請(qǐng) 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

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

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

Embedding Lookup

SOK embedding lookup_sparse 提供與 tf.nn.embedding_lookup_sparse 相似的 API，與 tf.nn.embedding_lookup_sparse 不同的是，SOK 的 embedding lookup_sparse 可以同時(shí)將多個(gè) 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 和增量導(dǎo)出 Incremental Dump 的功能。

Dump/Load 可以自動(dòng)并行的進(jìn)行 Embedding Table 的 Key、Value、Optimizer 中狀態(tài)變量（可選）在文件系統(tǒng)中的讀寫(xiě)，將 Embedding Table 的 Key、Value、Optimizer 中的狀態(tài)變量?jī)?chǔ)存成二進(jìn)制文件/從二進(jìn)制文件中讀取，下面是 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 中，用戶通常會(huì)將訓(xùn)練一段時(shí)間后更新了的 Key 和 Value 推送到推理端，這樣可以避免推送整個(gè) Embedding Table 產(chǎn)生的巨大開(kāi)銷，SOK 同樣提供了 incremental_dump 的 API 來(lái)實(shí)現(xiàn)這個(gè)功能。incremental_dump 接受一個(gè) UTC Time Threshold，可以將 Time Threshold 后更新的 Key/Value 導(dǎo)出到 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，如果開(kāi)啟 TensorFlow 的 XLA，需要手動(dòng)將 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ù)團(tuán)隊(duì)開(kāi)展的部分手機(jī)行業(yè)推薦場(chǎng)景高性能優(yōu)化項(xiàng)目實(shí)踐經(jīng)驗(yàn)分享：

案例 1：通過(guò) NVIDIA Merlin

HugeCTR SOK 以及 NVTabular

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

應(yīng)用背景：

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

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

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

實(shí)驗(yàn)性能對(duì)比：

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

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

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

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

案例 2：通過(guò) NVIDIA Merlin

HugeCTR SOK 實(shí)現(xiàn)了

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

應(yīng)用背景：

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

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

通過(guò) HugeCTR Sparse Operation Kit 進(jìn)行了 GPU 加速的推薦系統(tǒng)的實(shí)驗(yàn)，性能加速如下：

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

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

結(jié)束語(yǔ)

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

關(guān)于作者

康暉

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