實(shí)現(xiàn) TensorFlow 架構(gòu)的規(guī)模性和靈活性

TensorFlow 是為了大規(guī)模分布式訓(xùn)練和推理而設(shè)計(jì)的，不過它在支持新機(jī)器學(xué)習(xí)模型和系統(tǒng)級(jí)優(yōu)化的實(shí)驗(yàn)中的表現(xiàn)也足夠靈活。

本文對(duì)能夠同時(shí)兼具規(guī)模性和靈活性的系統(tǒng)架構(gòu)進(jìn)行了闡述。設(shè)定的人群是已基本熟悉 TensorFlow 編程概念，例如 computation graph, operations, and sessions。有關(guān)這些主題的介紹，請(qǐng)參閱

https://tensorflow.google.cn/guide/low_level_intro?hl=zh-CN。如已熟悉Distributed TensorFlow，本文對(duì)您也很有幫助。行至文尾，您應(yīng)該能了解 TensorFlow 架構(gòu)，足以閱讀和修改核心 TensorFlow 代碼了。

概覽

TensorFlow runtime 是一個(gè)跨平臺(tái)庫。圖 1 說明了它的一般架構(gòu)。 C API layer 將不同語言的用戶級(jí)代碼與核心運(yùn)行時(shí)分開。

圖 1

本文重點(diǎn)介紹以下幾個(gè)方面：

客戶端：

將整個(gè)計(jì)算過程轉(zhuǎn)義成一個(gè)數(shù)據(jù)流圖

通過 session，啟動(dòng)圖形執(zhí)行

分布式主節(jié)點(diǎn)

基于用戶傳遞給 Session.run() 中的參數(shù)對(duì)整個(gè)完整的圖形進(jìn)行修剪，提取其中特定子圖

將上述子圖劃分成不同片段，并將其對(duì)應(yīng)不同的進(jìn)程和設(shè)備當(dāng)中

將上述劃分的片段分布到 worker services 工作節(jié)點(diǎn)服務(wù)上

每個(gè) worker services 工作節(jié)點(diǎn)服務(wù)上執(zhí)行其收到的圖形片段

工作節(jié)點(diǎn)服務(wù) Worker Services(每一任務(wù)一個(gè))

使用內(nèi)核實(shí)現(xiàn)來計(jì)劃圖形表示的計(jì)算部分分配給正確的可用硬件（如 cpu，gpu 等）

與其他工作節(jié)點(diǎn)服務(wù) work services 相互發(fā)送和接收計(jì)算結(jié)果

內(nèi)核實(shí)現(xiàn)

執(zhí)行單個(gè)圖形操作的計(jì)算部分

圖2 說明了這些組件的相互作用?！? job：worker / task：0” 和 “/ job：ps / task：0” 都是工作節(jié)點(diǎn)服務(wù) worker services 上執(zhí)行的任務(wù)。“PS” 表示“參數(shù)服務(wù)器”：負(fù)責(zé)存儲(chǔ)和更新模型參數(shù)。其他任務(wù)在迭代優(yōu)化參數(shù)時(shí)會(huì)對(duì)這些參數(shù)發(fā)送更新。如果在單機(jī)環(huán)境下，上述 PS 和 worker 不是必須的，不需要在任務(wù)之間進(jìn)行這種特定的分工，但是對(duì)于分布式訓(xùn)練，這種模式就是很常見的。

圖 2

請(qǐng)注意，分布式主節(jié)點(diǎn) Distributed Master 和工作節(jié)點(diǎn)服務(wù) Worker Service 僅存在于分布式 TensorFlow 中。TensorFlow 的單進(jìn)程版本包含一個(gè)特殊的 Session 實(shí)現(xiàn)，它可以執(zhí)行分布式主服務(wù)器執(zhí)行的所有操作，但只與本地進(jìn)程中的設(shè)備進(jìn)行通信。

下面，通過逐步處理示例圖來詳細(xì)介紹一下 TensorFlow 核心模塊。

客戶端

用戶在客戶端編寫 TensorFlow 程序來構(gòu)建計(jì)算圖。該程序可以直接組成單獨(dú)的操作，也可以使用 Estimators API 之類的便利庫來組合神經(jīng)網(wǎng)絡(luò)層和其他更高級(jí)別的抽象概念。TensorFlow 支持多種客戶端語言，我們優(yōu)先考慮 Python 和 C ++，因?yàn)槲覀兊膬?nèi)部用戶最熟悉這些語言。隨著功能的日趨完善，我們一般會(huì)將它們移植到 C ++，以便用戶可以從所有客戶端語言優(yōu)化訪問。盡管大多數(shù)訓(xùn)練庫仍然只支持 Python，但 C ++ 確實(shí)能夠支持有效的推理。

客戶端創(chuàng)建會(huì)話，該會(huì)話將圖形定義作為tf.GraphDef 協(xié)議緩沖區(qū)發(fā)送到分布式主節(jié)點(diǎn)。當(dāng)客戶端評(píng)估圖中的一個(gè)或多個(gè)節(jié)點(diǎn)時(shí)，評(píng)估會(huì)觸發(fā)對(duì)分布式主節(jié)點(diǎn)的調(diào)用以啟動(dòng)計(jì)算。

在圖 3 中，客戶端構(gòu)建了一個(gè)圖表，將權(quán)重（w）應(yīng)用于特征向量（x），添加偏差項(xiàng)（b）并將結(jié)果保存在變量中。

圖 3

代碼：

tf.Session

分布式主節(jié)點(diǎn) Distributed master

分布式主節(jié)點(diǎn)：

基于客戶端指定的節(jié)點(diǎn)，從完整的圖形中截取所需的子圖

對(duì)圖表進(jìn)一步進(jìn)行劃分，使其可以將每個(gè)圖形片段映射到不同的執(zhí)行設(shè)備上

以及緩存這些劃分好的片段，以便在后續(xù)步驟中再次使用

由于主節(jié)點(diǎn)可以總攬步驟計(jì)算，因此它可以使用標(biāo)準(zhǔn)的優(yōu)化方法去做優(yōu)化，例如公共子表達(dá)式消除和常量的綁定。然后，對(duì)一組任務(wù)中優(yōu)化后的子圖或者片段執(zhí)行協(xié)調(diào)。

圖4

圖 5 顯示了示例圖的可能分區(qū)。分布式主節(jié)點(diǎn)已對(duì)模型參數(shù)進(jìn)行分組，以便將它們放在參數(shù)服務(wù)器上。

圖5

在分區(qū)切割圖形邊緣的情況下，分布式主節(jié)點(diǎn)插入發(fā)送和接收節(jié)點(diǎn)以在分布式任務(wù)之間傳遞信息（圖 6）。

圖 6

然后，分布式主節(jié)點(diǎn)將圖形片段傳送到分布式任務(wù)。

圖 7

代碼：

MasterService API definition

Master interface

工作節(jié)點(diǎn)服務(wù) Worker Service

在每個(gè)任務(wù)中，該部分負(fù)責(zé)：

處理來自 master 發(fā)來的請(qǐng)求

為包含本地子圖規(guī)劃所需要的內(nèi)核執(zhí)行

協(xié)調(diào)與其他任務(wù)之間的直接信息交換

我們優(yōu)化了 Worker Service，以便其以較低的負(fù)載便能夠運(yùn)行大型圖形。當(dāng)前的版本可以執(zhí)行每秒上萬個(gè)子圖，這使得大量副本可以進(jìn)行快速的，細(xì)粒度的訓(xùn)練步驟。Worker service 將內(nèi)核分派給本地設(shè)備并在可能的情況下并行執(zhí)行內(nèi)核，例如通過使用多個(gè) CPU 內(nèi)核或者 GPU 流。

我們還特別針對(duì)每對(duì)源設(shè)備和目標(biāo)設(shè)備類型的 Send 和 Recv 操作進(jìn)行了專攻：

使用 cudaMemcpyAsync() 來進(jìn)行本地 CPU 和 GPU 設(shè)備之間的重疊計(jì)算和數(shù)據(jù)傳輸

兩個(gè)本地 GPU 之間的傳輸使用對(duì)等 DMA，以避免通過主機(jī) CPU 主內(nèi)存進(jìn)行高負(fù)載的復(fù)制

對(duì)于任務(wù)之間的傳輸，TensorFlow 使用多種協(xié)議，包括：

gRPC over TCP

融合以太網(wǎng)上的 RDMA

另外，我們還初步支持 NVIDIA 用于多 GPU 通信的 NCCL 庫

圖8

代碼：

WorkerService API definition

Worker interface

Remote rendezvous (for Send and Recv implementations)

內(nèi)核運(yùn)行

該運(yùn)行時(shí)包含了 200 多個(gè)標(biāo)準(zhǔn)操作，其中涉及了數(shù)學(xué)，數(shù)組，控制流和狀態(tài)管理等操作。每個(gè)操作都有對(duì)應(yīng)各種設(shè)備優(yōu)化后的內(nèi)核運(yùn)行。其中許多操作內(nèi)核都是通過使用 Eigen::Tensor 實(shí)現(xiàn)的，它使用 C ++ 模板為多核 CPU 和 GPU 生成高效的并行代碼；但是，我們可以自由地使用像 cuDNN 這樣的庫，就可以實(shí)現(xiàn)更高效的內(nèi)核運(yùn)行。我們還實(shí)現(xiàn)了 quantization 量化，可以在移動(dòng)設(shè)備和高吞吐量數(shù)據(jù)中心應(yīng)用等環(huán)境中實(shí)現(xiàn)更快的推理，并使用 gemmlowp 低精度矩陣庫來加速量化計(jì)算。

如果用戶發(fā)現(xiàn)很難去將子計(jì)算組合，或者說組合后發(fā)現(xiàn)效率很低，則用戶可以通過注冊(cè)額外的 C++ 編寫的內(nèi)核來提供有效的運(yùn)行。例如，我們建議為一些性能的關(guān)鍵操作注冊(cè)自己的融合內(nèi)核，例如 ReLU 和 Sigmoid 激活函數(shù)及其對(duì)應(yīng)的梯度等。XLA Compiler提供了一個(gè)實(shí)驗(yàn)性質(zhì)的自動(dòng)內(nèi)核融合實(shí)現(xiàn)。

代碼：

OpKernel?interface