Merlin HugeCTR 分級參數(shù)服務(wù)器系列之三——集成到 TensorFlow

前兩期中我們介紹了 HugeCTR 分級參數(shù)服務(wù)器 （HPS）的三級存儲結(jié)構(gòu)的設(shè)計(jì)，相關(guān)配置使用，數(shù)據(jù)后端以及流式在線模型更新方案。本期將為大家介紹如何將 HPS 集成到 TensorFlow 中，從而實(shí)現(xiàn)在 TensorFlow 中利用分級存儲來部署包含龐大 Embedding Tables 的模型。

引言

當(dāng)需要基于 TensorFlow 來部署包含龐大 Embedding Tables 的深度學(xué)習(xí)模型時(shí)，數(shù)據(jù)科學(xué)家和機(jī)器學(xué)習(xí)工程師需要面對以下挑戰(zhàn)：

• 龐大的 Embedding Tables：訓(xùn)練好的 Embedding Tables 的大小往往達(dá)到幾百GB，使用 TensorFlow 原生的 Layers 或 Variable 無法放入 GPU 內(nèi)存；

• 低延時(shí)需求: 在線推理要求 Embedding 查詢的延時(shí)要足夠低（幾毫秒級），以保證體驗(yàn)質(zhì)量和維持用戶粘度；

• 多 GPU 擴(kuò)展能力：分布式推理框架需要將多個(gè)模型部署在多個(gè) GPU上，每個(gè)模型包含一個(gè)或多個(gè) Embedding Tables；

• 支持加載為 Pre-trained Embeddings：對于遷移學(xué)習(xí)等任務(wù)，需要支持以 Pre-trained Embeddings 的形式加載龐大的 Embedding Tables。

針對以上挑戰(zhàn)，我們?yōu)?HPS 提供了一個(gè)面向 TensorFlow 的 Python 定制化插件，以方便用戶將 HPS 集成到 TensorFlow 模型圖中，實(shí)現(xiàn)包含龐大 Embedding Tables 的模型的高效部署：

• HPS 通過使用集群中可用的存儲資源來擴(kuò)展 GPU 內(nèi)存，包含 CPU RAM 以及非易失性存儲如 HDD 和 SSD，從而實(shí)現(xiàn)龐大 Embedding Tables 的分級存儲，如圖 1 所示；

• HPS 通過使用 GPU Embedding Cache 來利用 embedding key 的長尾特性，當(dāng)查詢請求持續(xù)不斷涌入時(shí)，緩存機(jī)制保證 GPU 內(nèi)存可以自動(dòng)存儲熱門（高頻訪問）key 的 Embeddings，從而可以提供低延遲的查詢服務(wù)；

• 在 GPU 內(nèi)存，CPU 內(nèi)存以及 SSD 組成的存儲層級中，HPS 以層級結(jié)構(gòu)化的方式來管理多個(gè)模型的 Embedding Tables，實(shí)現(xiàn)參數(shù)服務(wù)器的功能；

• HPS 的查詢服務(wù)通過 Custom TensorFlow Layers 來接入，無論是推理還是類似遷移學(xué)習(xí)的任務(wù)，都可以有效支持。

圖 1：HPS 的分級存儲架構(gòu)

TensorFlow 用戶可以使用我們提供的 Python APIs，輕松利用 HPS 的上述特性，我們將在下文中進(jìn)一步介紹。

工作流程

利用 HPS 來部署包含龐大 Embedding Tables 的 TensorFlow 模型的工作流程如圖 2 所示：

圖 2：利用 HPS 部署 TensorFlow 模型的工作流程

流程中的步驟可以總結(jié)如下：

• 訓(xùn)練階段：用戶可以用原生的 TensorFlow Embedding Layers（例如 tf.nn.embedding_lookup_sparse）或者支持模型并行的 SOK^[1] Embedding Layers（例如 sok.DistributedEmbedding）來搭建模型圖并進(jìn)行訓(xùn)練。只要模型可以用 TensorFlow 進(jìn)行訓(xùn)練，則無論密集層以及模型圖的拓?fù)浣Y(jié)構(gòu)如何，HPS 都可以在推理階段集成進(jìn)來。

• 分解訓(xùn)練的模型圖：用戶需要從訓(xùn)練的模型圖中提取由密集層組成的子圖，并將其單獨(dú)保存下來。至于訓(xùn)練好的 Embedding Weights，如果使用的是原生 TensorFlow Embedding Layers，則需要提取 Embedding Weights 并將其轉(zhuǎn)換成 HPS 支持的格式；如果使用的是 SOK Embedding Layers，可以利用 sok.Saver.dump_to_file 來直接得到所需的格式。HPS 的格式要求為：每個(gè) Embedding Table 都被保存在一個(gè)文件夾中，包含兩個(gè)二進(jìn)制文件，key （int64）和 emb_vector（float32）。舉例來說，如果一共有 1000 個(gè)訓(xùn)練好的鍵值對，并且 embedding vector 的長度是 16，那么 key 文件和 emb_vector 文件的大小分別為 1000*8 bytes 和 1000*16*4 bytes。

• 搭建并保存推理圖：推理圖由 HPS Layers（如 hps.SparseLookupLayer）和保存好的密集層子圖搭建而成。只需將訓(xùn)練圖中的 Embedding Layer 用 HPS Layers 加以替換，便可以得到推理圖。該推理圖保存后便可在生產(chǎn)環(huán)境中部署。

• 部署推理圖：為了利用 HPS，用戶需要提供一個(gè) JSON 文件，來指定待部署模型的配置信息，用以啟動(dòng) HPS 查詢服務(wù)。接下來便可以部署保存好的推理圖來執(zhí)行在線推理任務(wù)，在此過程中有效地利用 HPS Embedding 查詢的優(yōu)化設(shè)計(jì)。關(guān)于配置信息的更多細(xì)節(jié)，請參考 HPS Configuration^[2]。

API

HPS 提供了簡潔易用的 Python API，可以與 TensorFlow 無縫地銜接。用戶只需幾行代碼，便可以啟動(dòng) HPS 查詢服務(wù)以及將 HPS 集成到 TensorFlow 模型圖中。

• hierarchical_parameter_server.Init：該方法用來針對待部署的模型啟動(dòng) HPS 查詢服務(wù)，需要在執(zhí)行推理任務(wù)前被調(diào)用一次。必須提供的參數(shù)為：

• global_batch_size：整型，待部署模型的全局批大小。例如模型部署在 4 個(gè) GPUs 上，每個(gè) GPU 上批大小為 4096，則全局批大小為 16384

• ps_config_file：字符串，HPS 初始化所需的 JSON 配置文件

該方法支持顯式調(diào)用或隱式調(diào)用。顯式調(diào)用用于基于 Python 腳本的測試工作；隱式調(diào)用則用于在生產(chǎn)環(huán)境中部署模型，要求待部署的推理模型中的 hps.SparseLookupLayer 或 hps.LookupLayer 指定好 global_batch_size 和 ps_config_file，當(dāng)模型首次接收到推理請求時(shí)，會以 call_once 且線程安全的方式觸發(fā) HPS 的初始化

• hierarchical_parameter_server.SparseLookupLayer：繼承自 tf.keras.layers.Layer，通過指定模型名和 table id 訂閱到 HPS 查詢服務(wù)。該層執(zhí)行與 tf.nn.embedding_lookup_sparse 基本相同的功能。構(gòu)造時(shí)必須提供的參數(shù)為：

• model_name：字符串，HPS 部署的模型名

• table_id：整型，指定的 model_name 的 Embedding Tables 的索引

• emb_vec_size：整型，指定的 model_name 和 table_id 的 Embedding Vector 的長度

• emb_vec_dtype：返回的 Embedding Vector 的數(shù)據(jù)類型，目前只支持為 tf.float32

• ps_config_file：字符串，HPS 隱式初始化所需的 JSON 配置文件

• global_batch_size：整型，待部署模型的全局批大小

執(zhí)行時(shí)的輸入和返回值為：

• sp_ids：輸入，int64 類型的 id 的 N x M SparseTensor，其中 N 通常是批次大小，M 是任意的

• sp_weights：輸入，可以是具有 float/double weight 的 SparseTensor，或者是 None 以表示所有 weight 應(yīng)為 1。如果指定，則 sp_weights 必須具有與 sp_ids 完全相同的 shape 和 indice

• combiner：輸入，指定 reduction 操作的字符串。目前支持“mean”，“sqrtn”和“sum”

• max_norm：輸入，如果提供，則在 combine 之前將每個(gè) embedding 規(guī)范化為具有等于 max_norm 的 l2 范數(shù)

• emb_vector: 返回值，表示 combined embedding 的密集張量。對于由 sp_ids 表示的密集張量中的每一行，通過 HPS 來查找該行中所有 id 的 embedding，將它們乘以相應(yīng)的 weight，并按指定的方式組合這些 embedding

• hierarchical_parameter_server.LookupLayer：繼承自 tf.keras.layers.Layer，通過指定模型名和 table id 訂閱到 HPS 查詢服務(wù)。該層執(zhí)行與 tf.nn.embedding_lookup 基本相同的功能。構(gòu)造時(shí)的參數(shù)與 hierarchical_parameter_server.SparseLookupLayer 相同。執(zhí)行時(shí)的輸入和返回值為：

• inputs：輸入，保存在 Tensor 中的鍵。數(shù)據(jù)類型必須為 tf.int64

• emb_vector: 返回值，查詢到的 Embedding Vector。數(shù)據(jù)類型為 tf.float32

部署方案

按照圖 2 的工作流程得到集成了 HPS 的推理模型后，用戶可以根據(jù)生產(chǎn)環(huán)境選擇多種部署方案：

• Triton TensorFlow backend^[3]：Triton 推理服務(wù)器是開源推理服務(wù)軟件，可簡化 AI 推理流程，支持部署不同深度學(xué)習(xí)框架的推理模型。集成了 HPS 的 TensorFlow 推理模型可以基于 Triton TensorFlow backend 進(jìn)行部署，只需將 HPS 的 embedding lookup 視作 custom op，并在啟動(dòng) tritonserver 前用 LD_PRELOAD 將其 shared library 加載到 Triton 即可

• TensorFlow Serving^[4]：集成了 HPS 的 TensorFlow 推理模型也可方便地部署在 TensorFlow Serving 這一靈活高性能的推理服務(wù)系統(tǒng)上，HPS 的 embedding lookup 同樣可作為 custom op 被加載到 TensorFlow Serving中

如果用戶希望進(jìn)一步對集成了 HPS 的推理模型中的密集網(wǎng)絡(luò)層進(jìn)行優(yōu)化，還可以使用 tensorflow.python.compiler.tensorrt.trt_convert 對 SavedModel 進(jìn)行轉(zhuǎn)換，HPS 的 embedding lookup 可自動(dòng) fallback 到其 TensorFlow plugin 對應(yīng)的 kernels，而可被優(yōu)化的密集網(wǎng)絡(luò)層則會生成 TensorRT engine 來執(zhí)行。轉(zhuǎn)換后的 SavedModel 仍然可以使用 Triton TensorFlow backend 或 TensorFlow Serving 進(jìn)行部署。

除了使用 HPS 的 TensorFlow plugin 外，用戶還可使用 Triton HPS backend^[5]。利用 Triton Ensemble Model^[6]，用戶可以方便地將 HPS backend 和其他 Triton backend 連接起來，搭建出 HPS 用于 embedding lookup、其他 backend 用于密集網(wǎng)絡(luò)層前向傳播的推理服務(wù)流水線。使用該方案來部署模型的工作流程如圖 3 所示：

圖 3：利用 Triton Emsemble Model 部署模型的工作流程

這里密集網(wǎng)絡(luò)層部分除了可以使用 TensorFlow backend 進(jìn)行部署外，還可以使用 TensorRT backend 進(jìn)行部署，此時(shí)需要將密集網(wǎng)絡(luò)層的 SavedModel 通過 TensorFlow->ONNX->TensorRT 的轉(zhuǎn)化，得到性能優(yōu)化的 TensorRT engine。

結(jié)語

在這一期的 HugeCTR 分級參數(shù)服務(wù)器文章中，我們介紹了使用 HPS 部署包含龐大 Embedding Tables 的 TensorFlow 模型的解決方案，工作流程以及 API。更多信息，請參考 HPS 官方文檔：

https://nvidia-merlin.github.io/HugeCTR/master/hierarchical_parameter_server/index.html

在下一期中，我們將著重介紹 HugeCTR 分級參數(shù)服務(wù)器中最關(guān)鍵的組件：Embedding Cache 的設(shè)計(jì)細(xì)節(jié)，敬請期待。

以下是 HugeCTR 的 Github repo 以及其他發(fā)布的文章，歡迎感興趣的朋友閱讀和反饋。Github:

https://github.com/NVIDIA-Merlin/HugeCTR （更多文章詳見 README）

[1]SOK

https://nvidia-merlin.github.io/HugeCTR/sparse_operation_kit/master/index.html

[2]HPS Configuration

https://nvidia-merlin.github.io/HugeCTR/master/hugectr_parameter_server.html#configuration

[3]Triton TensorFlow backend

https://github.com/triton-inference-server/tensorflow_backend

[4]TensorFlow Serving

https://github.com/tensorflow/serving

[5]Triton HPS backend

https://github.com/triton-inference-server/hugectr_backend/tree/main/hps_backend

[6]Triton Ensemble Model

https://github.com/triton-inference-server/server/blob/main/docs/user_guide/architecture.md#ensemble-models

點(diǎn)擊查看關(guān)于 HugeCTR 分級參數(shù)服務(wù)器的更多內(nèi)容

????Merlin HugeCTR 分級參數(shù)服務(wù)器簡介

?

Merlin HugeCTR 分級參數(shù)服務(wù)器簡介之二