天天看點

[源碼解析] PyTorch 分布式 Autograd (2) ---- RPC基礎

前文我們給出了分布式autograd的設計思路,本文開始,我們進行具體源碼分析。因為無論是前向傳播還是反向傳播,都需要依賴 RPC 來完成,是以我們先看看封裝于 RPC 之上的一些基本功能,比如初始化,代理(RPC 相關功能都是基于代理完成),消息接受,發送等等。

目錄

[源碼解析] PyTorch 分布式 Autograd (2) ---- RPC基礎

0x00 摘要

0x01 示例

0x02 RPC 基礎

2.1 初始化

2.1.1 初始化後端

2.1.2 生成代理

2.1.3 設定代理

2.1.4 靜态類變量

2.2 RPC 代理

2.2.1 RpcAgent

2.2.2 ProcessGroupAgent

2.2.3 TensorPipeAgent

2.2.4 回調函數

2.2.4.1 RequestCallback

2.2.4.2 RequestCallbackNoPython

0x03 發送邏輯

3.1 Python

3.2 C++

3.2.1 pyRpcBuiltin

3.2.2 sendMessageWithAutograd

0x04 接受邏輯

4.1 回調

4.2 operator()

4.3 RequestCallbackImpl

0xFF 參考

通過本文,大家可以了解:如何初始化RPC後端,如何生成 RPC 代理,如何使用RPC代理進行發送和接受消息,如何連接配接遠端 dist.autograd 自動微分引擎。

PyTorch分布式其他文章如下:

深度學習利器之自動微分(1)

深度學習利器之自動微分(2)

[源碼解析]深度學習利器之自動微分(3) --- 示例解讀

[源碼解析]PyTorch如何實作前向傳播(1) --- 基礎類(上)

[源碼解析]PyTorch如何實作前向傳播(2) --- 基礎類(下)

[源碼解析] PyTorch如何實作前向傳播(3) --- 具體實作

[源碼解析] Pytorch 如何實作後向傳播 (1)---- 調用引擎

[源碼解析] Pytorch 如何實作後向傳播 (2)---- 引擎靜态結構

[源碼解析] Pytorch 如何實作後向傳播 (3)---- 引擎動态邏輯

[源碼解析] PyTorch 如何實作後向傳播 (4)---- 具體算法

[源碼解析] PyTorch 分布式(1)------曆史和概述

[源碼解析] PyTorch 分布式(2) ----- DataParallel(上)

[源碼解析] PyTorch 分布式(3) ----- DataParallel(下)

[源碼解析] PyTorch 分布式(4)------分布式應用基礎概念

[源碼解析] PyTorch分布式(5) ------ DistributedDataParallel 總述&如何使用

[源碼解析] PyTorch分布式(6) ---DistributedDataParallel -- 初始化&store

[源碼解析] PyTorch 分布式(7) ----- DistributedDataParallel 之程序組

[源碼解析] PyTorch 分布式(8) -------- DistributedDataParallel之論文篇

[源碼解析] PyTorch 分布式(9) ----- DistributedDataParallel 之初始化

[源碼解析] PyTorch 分布式(10)------DistributedDataParallel 之 Reducer靜态架構

[源碼解析] PyTorch 分布式(11) ----- DistributedDataParallel 之 建構Reducer和Join操作

[源碼解析] PyTorch 分布式(12) ----- DistributedDataParallel 之 前向傳播

[源碼解析] PyTorch 分布式(13) ----- DistributedDataParallel 之 反向傳播

[源碼解析] PyTorch 分布式 Autograd (1) ---- 設計

為了更好的說明,本文代碼會依據具體情況來進行相應精簡。

我們從 PyTorch 示例部分之中摘錄示例代碼并且修改了一些,代碼目的是讓兩個 worker 之間就通過 RPC 進行協作。示例 worker 具體分為兩部分:

RPC操作,建構依賴基礎。

執行後向傳播。

可以用如下辦法來啟動了兩個 worker,其中使用了 rpc.init_rpc 來初始化 rpc。worker0 會啟動,然後利用 RPC 在 worker 1 之上也進行了一些操作。

我們從頭看看示例代碼,當腳本啟動時候,會調用到 rpc.init_rpc 來初始化 rpc。從 RPC 注釋中可以看到兩個概念,就是大家常見的 rank 和 world_size。

具體初始化代碼是:

其中我們關心的是:_init_rpc_backend 會設定後端。

_init_rpc_backend 這裡會依據配置來看看最後生成什麼 Agent,然後把這個代理設定到目前上下文。RPC有兩種後端,TENSORPIPE 和 PROCESS_GROUP,其中PROCESS_GROUP已經被廢棄,會逐漸遷移到TENSORPIPE。

可以看到,預設會生成 TensorPipeAgent。

我們接下來看看如何生成 TensorPipeAgent,具體是在 torch/csrc/distributed/rpc/init.cpp。當這裡生成 TensorPipeAgent 時候,把 RequestCallbackImpl 配置為回調函數。代理内部就用這個回調函數用來處理接收到的請求。

具體如下:

_init_rpc_states 會把代理設定在PyTorch環境之中,其定義在 torch/distributed/rpc/api.py 之中有。

接下來就要進入了C++世界。在 torch/csrc/distributed/rpc/init.cpp 中有 _set_and_start_rpc_agent,其作用是:

RpcAgent::setCurrentRpcAgent 設定了代理。

調用 rpcAgent->start() 來啟動代理。

setCurrentRpcAgent 定義在 torch/csrc/distributed/rpc/rpc_agent.cpp 之中。

在 RpcAgent 之中,有一個靜态成員變量 currentRpcAgent_。

在 C++ 之中,靜态成員變量有如下特點:

其屬于整個類所有。

其生命期不依賴于任何對象,為程式的生命周期。

可以通過類名直接通路公有靜态成員變量。

可以通過對象名通路一個類的公有靜态成員變量。

類的所有派生對象共享該類的靜态成員變量。

靜态成員變量需要在該類外單獨配置設定空間。

靜态成員變量在程式内部位于全局資料區。

是以,我們可知<code>RpcAgent::currentRpcAgent_</code> 可以認為就是全局變量,rpc 統一使用這個變量進行協調。具體通過 RpcAgent 的一些公有成員函數來完成這些功能。

于是目前拓展如下,以後進行 RPC 操作,都會通過 RpcAgent::currentRpcAgent_ 這個全局變量進行。

dist.autograd 的相關功能都是基于 RPC 代理完成,是以我們需要仔細看看代理。

這是用來傳遞RPC的代理,是收發 RPC消息的代理基類,其:

提供了<code>send</code> API用來處理request 和 response。

也配置了 cb_ 用來處理接收到的請求。

<code>WorkerInfo</code> 是代理執行個體所在 worker 的全局唯一标示,包括<code>name_</code>和<code>id_</code>這兩個成員變量。<code>name_</code>是全局唯一名字,<code>id_</code>是全局唯一ID。

ProcessGroupAgent 是 RpcAgent 的派生類。這是之前使用的,但是 PyTorch 提供了更優秀的 TensorAgent。我們隻選取了部分成員變量。

TensorPipeAgent 定義在 torch/csrc/distributed/rpc/tensorpipe_agent.h,這是目前和未來使用的。TensorPipeAgent利用TensorPipe在可用傳輸或通道之中透明地移動張量和資料。它就像一個混合的RPC傳輸,提供共享記憶體(linux)和TCP(linux&amp;mac)支援。PyTorch 正在開發其支援CUDA版本。

我們隻選取了部分成員變量。

Agent 在收到消息時候,會調用回調函數。而 RequestCallbackImpl 實作了回調邏輯。RequestCallbackImpl 是派生類,我們先來看看基類 RequestCallbackNoPython,結果找到了RequestCallback 這個接口,是以 RequestCallback 才是這個派生體系的基礎。

RequestCallback 是處理 RPC 消息的接口,是一個抽象類。

RequestCallbackNoPython 的定義在 torch/csrc/distributed/rpc/request_callback_no_python.h,其實作了一些處理機制,因為其包含太多方法,我們隻能摘錄部分,如果有興趣的朋友請深入研究。

我們會在後續分析接受邏輯時候,看到如何調用到回調函數。

我們先來看看發送邏輯。也就是 rpc.rpc_sync 的作用:建立 root,添加 send等。

我們從 python 部分開始。

首先來到 rpc_sync,發現其調用了_invoke_rpc。

其次來到<code>_invoke_rpc</code>,可以看到此函數依據調用類型不同(内置操作,script,udf這三種),選擇了不同路徑。

從這裡開始就進入到了C++世界,torch/csrc/distributed/rpc/init.cpp。

這裡可以看到<code>_invoke_rpc_builtin</code> 對應了 pyRpcBuiltin,<code>_invoke_rpc_python_udf</code> 對應了 pyRpcPythonUdf。

我們選用 <code>_invoke_rpc_builtin</code> 對應的 pyRpcBuiltin 來看看。

在 torch/csrc/distributed/rpc/python_functions.cpp可以看到,pyRpcBuiltin 會調用到 sendMessageWithAutograd。

在 torch/csrc/distributed/autograd/utils.cpp 這裡利用 agent 來進行發送 FORWARD_AUTOGRAD_REQ。

後面在接收方,我們将會看到處理 FORWARD_AUTOGRAD_REQ 消息,是以發送和接受大緻可以聯系起來。

發送流程如下,其中 sendMessageWithAutograd 會使用 RpcAgent::getCurrentRpcAgent() 得到 RpcAgent::currentRpcAgent_,就是得到了全局設定的代理,然後通過代理進行發送。

當Agent接受到消息之後,會調用到RequestCallback::operator()。就是我們前面所說的回調函數。代碼位于 torch/csrc/distributed/rpc/tensorpipe_agent.cpp。

operator() 之中會調用 processMessage 處理消息。

随後,會調用到 RequestCallbackNoPython::processMessage 之中。

先調用 RequestCallbackImpl 中實作的 deserializePythonRpcCommand 來對 PythonUDF 反序列化。

然後調用 processRpcWithErrors 來處理消息。

然後調用到 processRpcWithErrors。

接下來是 processRpc。這裡能夠看到處理 FORWARD_AUTOGRAD_REQ。

這時候,讀者會有疑問,之前 TensorPipeAgent 明明設定了 RequestCallbackImpl 作為回調函數,怎麼隻調用了其 deserializePythonRpcCommand呢,deserialXXX 看起來是序列化相關的,按說應該調用一些業務處理函數,比如processXXXX 之類的。我們接下來就看看 RequestCallbackImpl。

RequestCallbackImpl 定義在 torch/csrc/distributed/rpc/request_callback_impl.h。

因為最終生成的是 RequestCallbackImpl,是以實際上,上圖中間有一步 processRpcWithErrors 實際調用的是 RequestCallbackImpl 這裡的函數 processRpcWithErrors,其就是增加了一些異常處理邏輯。

邏輯圖修改如下:

如果結合之前的發送,我們拓展圖例如下:

當發送者需要在遠端運作自動梯度計算時候,調用 rpc.rpc_sync。

從 Python 調用到 C++ 世界,函數為 pyRpcBuiltin。

調用 sendMessageWithAutograd,以此通知Receiver。

會調用 RpcAgent::getCurrentRpcAgent() 來得到本地的 Agent。

調用 current Agent 的 send 函數。

send 函數發送 FORWARD_AUTOGRAD_REQ給 Receiver worker。

respond 函數會調用 Receiver 之中 Agent 的回調函數 cb_。

調用到 RequestCallbackImpl 的 processRpcWithErrors。

然後調用 processRpc。

最後調用到 processForwardAutogradReq,完成了基于RPC的分布式autograd的啟動過程。

手機如下:

[源碼解析] PyTorch 分布式 Autograd (2) ---- RPC基礎

至此,RPC介紹完畢,我們下一篇介紹上下文相關等管理類,敬請期待。