PyTorch 遠(yuǎn)程參考協(xié)議

原文： https://pytorch.org/docs/stable/notes/rref.html

警告

RRef API 是實(shí)驗(yàn)性的，隨時(shí)可能更改。

本說(shuō)明描述了遠(yuǎn)程引用協(xié)議的設(shè)計(jì)細(xì)節(jié)，并逐步介紹了不同情況下的消息流。 在繼續(xù)之前，請(qǐng)確保您熟悉分布式 RPC 框架。

背景

RRef 代表遠(yuǎn)程參考。 它是位于本地或遠(yuǎn)程工作人員上的對(duì)象的引用，并且透明地在內(nèi)部進(jìn)行引用計(jì)數(shù)。 從概念上講，它可以視為分布式共享指針。 應(yīng)用程序可以通過(guò)調(diào)用 remote() 創(chuàng)建 RRef。 每個(gè) RRef 都由 remote() 呼叫的被調(diào)用方工作者(即所有者）擁有，并且可以由多個(gè)用戶使用。 所有者存儲(chǔ)實(shí)際數(shù)據(jù)，并跟蹤全局參考計(jì)數(shù)。 每個(gè) RRef 可以由全局RRefId唯一標(biāo)識(shí)，該全局RRefId在創(chuàng)建時(shí)在 remote() 調(diào)用的調(diào)用方上分配。

在所有者工作程序中，只有一個(gè)OwnerRRef實(shí)例包含真實(shí)數(shù)據(jù)，而在用戶工作程序中，可以根據(jù)需要包含任意數(shù)量的UserRRefs，而UserRRef不保存數(shù)據(jù)。 所有者上的所有用法都將使用全局唯一的RRefId來(lái)檢索唯一的OwnerRRef實(shí)例。 在 rpc_sync() ， rpc_async() 或  remote() 調(diào)用中將UserRRef用作參數(shù)或返回值時(shí)，將創(chuàng)建該UserRRef 將會(huì)根據(jù)更新的參考計(jì)數(shù)通知所有者。 如果全局沒有UserRRef實(shí)例，并且所有者上也沒有對(duì)OwnerRRef的引用，則OwnerRRef及其數(shù)據(jù)將被刪除。

假設(shè)條件

RRef 協(xié)議的設(shè)計(jì)基于以下假設(shè)。

• 瞬態(tài)網(wǎng)絡(luò)故障：RRef 設(shè)計(jì)旨在通過(guò)重試消息來(lái)處理瞬態(tài)網(wǎng)絡(luò)故障。 節(jié)點(diǎn)崩潰或永久性網(wǎng)絡(luò)分區(qū)超出了范圍。 當(dāng)這些事件發(fā)生時(shí)，該應(yīng)用程序可能會(huì)關(guān)閉所有工作人員，還原到先前的檢查點(diǎn)，然后恢復(fù)訓(xùn)練。
• 非冪等 UDF ：我們假設(shè)提供給 rpc_sync() ， rpc_async() 或  remote() 的用戶功能(UDF） 不是冪等的，因此無(wú)法重試。 但是，內(nèi)部 RRef 控制消息將成為冪等且可重試。
• 消息傳遞無(wú)序：我們不假定任何一對(duì)節(jié)點(diǎn)之間的消息傳遞順序，因?yàn)榘l(fā)送者和接收者都使用多個(gè)線程。 無(wú)法保證首先處理哪個(gè)消息。

RRef 生命周期

該協(xié)議的目標(biāo)是在適當(dāng)?shù)臅r(shí)候刪除OwnerRRef。 刪除OwnerRRef的正確時(shí)機(jī)是在沒有活動(dòng)UserRRef實(shí)例且用戶代碼也沒有保存對(duì)OwnerRRef的引用的情況下。 棘手的部分是確定是否存在任何活動(dòng)的UserRRef實(shí)例。

設(shè)計(jì)推理

用戶可以在以下三種情況下獲得UserRRef：

1. 從所有者那里收到了UserRRef。
2. 接收到另一個(gè)用戶的UserRRef。
3. 創(chuàng)建另一個(gè)工人擁有的新UserRRef。

情況 1 是最簡(jiǎn)單的，所有者將其 RRef 傳遞給用戶，所有者調(diào)用 rpc_sync() ， rpc_async() 或  remote() 使用其 RRef 作為參數(shù)。 在這種情況下，將在用戶上創(chuàng)建一個(gè)新的UserRRef。 由于所有者是調(diào)用者，因此可以輕松地在OwnerRRef上更新其本地引用計(jì)數(shù)。

唯一的要求是任何UserRRef必須在銷毀時(shí)通知所有者。 因此，我們需要第一個(gè)保證：

G1。 刪除任何“ UserRRef”時(shí)，都會(huì)通知所有者。

由于郵件可能會(huì)延遲或出現(xiàn)亂序，因此我們還需要一項(xiàng)保證，以確保刪除郵件不會(huì)過(guò)早處理。 如果 A 向 B 發(fā)送涉及 RRef 的消息，我們將 A 上的 RRef 稱為父 RRef，將 B 上的 RRef 稱為子 RRef。

G2。 在所有者確認(rèn)子 RRef 之前，不會(huì)刪除父 RRef。

在情況 2 和 3 中，所有者可能僅對(duì) RRef 分支圖有部分了解或根本不了解。 例如，可以在用戶上構(gòu)建 RRef，并且在所有者收到任何 RPC 調(diào)用之前，創(chuàng)建者用戶可能已經(jīng)與其他用戶共享了 RRef，并且這些用戶可以進(jìn)一步共享 RRef。 一個(gè)不變性是，任何 RRef 的派生圖始終都是一棵樹，因?yàn)榕缮?RRef 總是在被調(diào)用方上創(chuàng)建一個(gè)新的UserRRef實(shí)例(除非被調(diào)用方是所有者），因此每個(gè) RRef 都有一個(gè)父級(jí)。

所有者對(duì)樹中任何UserRRef的視圖分為三個(gè)階段：

1) unknown -> 2) known -> 3) deleted.

所有者對(duì)整棵樹的看法不斷變化。 擁有者認(rèn)為沒有活動(dòng)的UserRRef實(shí)例時(shí)，即刪除OwnerRRef實(shí)例時(shí)，所有UserRRef實(shí)例都可能確實(shí)被刪除或未知，因此所有者刪除了其OwnerRRef實(shí)例。 危險(xiǎn)的情況是某些分叉未知，而另一些被刪除。

G2 簡(jiǎn)單地保證在擁有者知道其所有子級(jí)UserRRef實(shí)例之前，不能刪除任何父級(jí)UserRRef。 但是，有可能在擁有者知道其父項(xiàng)UserRRef之前刪除了子項(xiàng)UserRRef。

考慮下面的示例，其中OwnerRRef分支到 A，然后 A 分支到 Y，Y 分支到 Z：

OwnerRRef -> A -> Y -> Z

如果所有者在處理所有來(lái)自 Z 的消息(包括刪除消息）之前先處理來(lái)自 Y 的所有消息，那么所有者將在知道 Y 之前就知道 Z 的刪除。但這不會(huì)造成任何問題。 因?yàn)?，Y 的祖先中至少有一個(gè)還活著(在本例中為 A），這將阻止所有者刪除OwnerRRef。 更具體地說(shuō)，如果所有者不知道 Y，則由于 G2 而無(wú)法刪除 A，并且所有者知道 A，因?yàn)樗姓呤?A 的父母。

如果在用戶上創(chuàng)建 RRef，事情會(huì)變得有些棘手：

OwnerRRef
    ^
    |
    A -> Y -> Z

如果 Z 在UserRRef上調(diào)用 to_here() ，則所有者至少知道刪除 Z 時(shí)的 A，否則， to_here() 不會(huì)結(jié)束。 如果 Z 沒有調(diào)用 to_here() ，則所有者可能在從 A 和 Y 發(fā)送任何消息之前就已從 Z 接收了所有消息。在這種情況下，由于尚未獲得OwnerRRef的真實(shí)數(shù)據(jù) 創(chuàng)建后，也沒有要?jiǎng)h除的內(nèi)容。 就像 Z 根本不存在一樣。 因此，仍然可以。

實(shí)作

G1 通過(guò)在UserRRef析構(gòu)函數(shù)中發(fā)送刪除消息來(lái)實(shí)現(xiàn)。 為了提供 G2 ，無(wú)論何時(shí)將父級(jí)UserRRef派生，都將其置于上下文中，并由新的ForkId對(duì)其進(jìn)行索引。 僅當(dāng)父級(jí)UserRRef從子級(jí)收到確認(rèn)消息(ACK）時(shí)，才會(huì)從上下文中刪除該父級(jí)UserRRef，并且只有當(dāng)所有者確認(rèn)后，該子級(jí)才會(huì)發(fā)出 ACK。

協(xié)議方案

現(xiàn)在，讓我們討論以上設(shè)計(jì)如何在四種情況下轉(zhuǎn)換為協(xié)議。

用戶與所有者共享 RRef 作為返回值

import torch
import torch.distributed.rpc as rpc


## on worker A
rref = rpc.remote('B', torch.add, args=(torch.ones(2), 1))
## say the rref has RRefId 100 and ForkId 1
rref.to_here()

在這種情況下，在用戶工作程序 A 上創(chuàng)建UserRRef，然后將其與遠(yuǎn)程消息一起傳遞給所有者工作程序 B，然后 B 創(chuàng)建OwnerRRef。 方法 remote() 立即返回，這意味著UserRRef可以在所有者了解之前被分叉/使用。

在所有者上，當(dāng)接收到 remote() 調(diào)用時(shí)，它將創(chuàng)建OwnerRRef，并返回一個(gè) ACK 來(lái)確認(rèn){100, 1}(RRefId，ForkId）。 僅在收到此 ACK 后，A 才能刪除其UserRRef。 這涉及 G1 和 G2 。 G1 很明顯。 對(duì)于 G2 而言，OwnerRRef是UserRRef的子級(jí)，并且UserRRef直到收到所有者的 ACK 才被刪除。

上圖顯示了消息流，其中實(shí)心箭頭包含用戶功能，而虛線箭頭是內(nèi)置消息。 請(qǐng)注意，從 A 到 B 的前兩個(gè)消息 (remote() 和 to_here())可以按任何順序到達(dá) B，但最終的刪除消息僅在以下情況下發(fā)送： ：

• B 確認(rèn)UserRRef {100, 1}(G2），并且
• Python GC 同意刪除本地UserRRef實(shí)例。 當(dāng) RRef 不再在范圍內(nèi)并且可以進(jìn)行垃圾回收時(shí)，就會(huì)發(fā)生這種情況。

用戶與所有者共享 RRef 作為參數(shù)

import torch
import torch.distributed.rpc as rpc


## on worker A and worker B
def func(rref):
  pass


## on worker A
rref = rpc.remote('B', torch.add, args=(torch.ones(2), 1))
## say the rref has RRefId 100 and ForkId 1
rpc.rpc_async('B', func, args=(rref, ))

在這種情況下，在 A 上創(chuàng)建UserRRef后，A 會(huì)將其用作對(duì) B 的后續(xù) RPC 調(diào)用中的參數(shù)。A 將使UserRRef {100, 1}保持活動(dòng)狀態(tài)，直到收到 B 的確認(rèn) (G2 ， 而不是 RPC 調(diào)用的返回值）。 這是必要的，因?yàn)樵诮邮盏剿邢惹暗南⒅?，A 不應(yīng)發(fā)出刪除消息，否則，OwnerRRef可以在使用前刪除，因?yàn)槲覀儾荒鼙ＷC消息的傳遞順序。 這是通過(guò)創(chuàng)建 RRef 的子項(xiàng)ForkId并將其保存在地圖中，直到收到所有者確認(rèn)該子項(xiàng)ForkId來(lái)完成的。 下圖顯示了消息流。

注意，UserRRef可以在功能完成甚至啟動(dòng)之前在 B 上刪除。 但是，這是可以的，因?yàn)樵?B 向子ForkId發(fā)送 ACK 時(shí)，它已經(jīng)獲取了OwnerRRef實(shí)例，這將防止它被過(guò)早刪除。

所有者與用戶共享 RRef

所有者對(duì)用戶是最簡(jiǎn)單的情況，所有者可以在本地更新引用計(jì)數(shù)，并且不需要任何其他控制消息即可通知其他人。 關(guān)于 G2 ，因?yàn)楦讣?jí)是所有者，所以它與父級(jí)立即從所有者接收 ACK 相同。

import torch
import torch.distributed.rpc as RRef, rpc


## on worker B and worker C
def func(rref):
  pass


## on worker B, creating a local RRef
rref = RRef("data")
## say the rref has RRefId 100
dist.rpc_async('C', func, args=(rref, ))

上圖顯示了消息流。 請(qǐng)注意，當(dāng)在 rpc_async 調(diào)用之后OwnerRRef退出作用域時(shí)，將不會(huì)刪除它，因?yàn)槿绻嬖谌魏我阎呐缮瑒t內(nèi)部會(huì)存在一個(gè)使它保持活動(dòng)狀態(tài)的映射，在這種情況下為UserRRef {100, 1}。 (G2 )

用戶與用戶共享 RRef

這是最復(fù)雜的情??況，其中調(diào)用者用戶(父級(jí)UserRRef），被調(diào)用者用戶(子級(jí)UserRRef）和所有者都需要參與。

import torch
import torch.distributed.rpc as rpc


## on worker A and worker C
def func(rref):
  pass


## on worker A
rref = rpc.remote('B', torch.add, args=(torch.ones(2), 1))
## say the rref has RRefId 100 and ForkId 1
rpc.rpc_async('C', func, args=(rref, ))

當(dāng) C 從 A 接收到子項(xiàng)UserRRef時(shí)，它向所有者 B 發(fā)送一個(gè)派生請(qǐng)求。稍后，當(dāng) B 確認(rèn) C 上的UserRRef時(shí)，C 將并行執(zhí)行兩個(gè)操作：1）發(fā)送子項(xiàng) ACK 到 A，然后 2）運(yùn)行用戶提供的功能。 在這段時(shí)間中，親本(A）將保持其UserRRef {100, 1}存活以實(shí)現(xiàn) G2 。