TensorFlow 是為了大規(guī)模分布式訓(xùn)練和推理而設(shè)計(jì)的，不過(guò)它在支持新機(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)庫(kù)。圖 1 說(shuō)明了它的一般架構(gòu)。 C API layer 將不同語(yǔ)言的用戶級(jí)代碼與核心運(yùn)行時(shí)分開(kāi)。

圖 1

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

客戶端：

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

通過(guò) 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)來(lái)計(jì)劃圖形表示的計(jì)算部分分配給正確的可用硬件（如 cpu，gpu 等）

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

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

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

圖2 說(shuō)明了這些組件的相互作用。“/ 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)練，這種模式就是很常見(jià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)行通信。

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

客戶端

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

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

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

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

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

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

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

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

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

gRPC over TCP

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

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

圖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)核都是通過(guò)使用 Eigen::Tensor 實(shí)現(xiàn)的，它使用 C ++ 模板為多核 CPU 和 GPU 生成高效的并行代碼；但是，我們可以自由地使用像 cuDNN 這樣的庫(kù)，就可以實(shí)現(xiàn)更高效的內(nèi)核運(yùn)行。我們還實(shí)現(xiàn)了 quantization 量化，可以在移動(dòng)設(shè)備和高吞吐量數(shù)據(jù)中心應(yīng)用等環(huán)境中實(shí)現(xiàn)更快的推理，并使用 gemmlowp 低精度矩陣庫(kù)來(lái)加速量化計(jì)算。

如果用戶發(fā)現(xiàn)很難去將子計(jì)算組合，或者說(shuō)組合后發(fā)現(xiàn)效率很低，則用戶可以通過(guò)注冊(cè)額外的 C++ 編寫(xiě)的內(nèi)核來(lá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