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

警告

分布式 RPC 框架是實(shí)驗(yàn)性的，隨時(shí)可能更改。

本說明將介紹分布式自動(dòng)分級(jí)的詳細(xì)設(shè)計(jì)，并逐步介紹其內(nèi)部。 在繼續(xù)之前，請(qǐng)確保您熟悉 Autograd 機(jī)械手和分布式 RPC 框架。

背景

假設(shè)您有兩個(gè)節(jié)點(diǎn)，并且在兩個(gè)節(jié)點(diǎn)之間劃分了一個(gè)非常簡(jiǎn)單的模型。 可以使用 torch.distributed.rpc 如下實(shí)現(xiàn)：

import torch
import torch.distributed.rpc as rpc


def my_add(t1, t2):
  return torch.add(t1, t2)


## On worker 0:
t1 = torch.rand((3, 3), requires_grad=True)
t2 = torch.rand((3, 3), requires_grad=True)


## Perform some computation remotely.
t3 = rpc.rpc_sync("worker1", my_add, args=(t1, t2))


## Perform some computation locally based on remote result.
t4 = torch.rand((3, 3), requires_grad=True)
t5 = torch.mul(t3, t4)


## Compute some loss.
loss = t5.sum()

分布式 autograd 背后的主要?jiǎng)訖C(jī)是使用我們已經(jīng)計(jì)算并記錄所有需要梯度的張量的合適梯度的loss在這樣的分布式模型上運(yùn)行向后傳遞。

正向通過過程中的自動(dòng)分級(jí)記錄

PyTorch 在正向傳遞過程中會(huì)構(gòu)建自動(dòng)分級(jí)圖，該圖用于執(zhí)行向后傳遞。 有關(guān)更多詳細(xì)信息，請(qǐng)參見 autograd 如何編碼歷史記錄。

對(duì)于分布式 autograd，我們需要在正向傳遞過程中跟蹤所有 RPC，以確保正確執(zhí)行向后傳遞。 為此，我們?cè)趫?zhí)行 RPC 時(shí)將send和recv函數(shù)附加到自動(dòng)縮放圖。

• send函數(shù)附加到 RPC 的源，并且其輸出邊指向 RPC 輸入張量的 autograd 函數(shù)。 從目的地接收反向傳遞期間此功能的輸入，作為適當(dāng)?shù)?code>recv功能的輸出。
• recv函數(shù)附加到 RPC 的目標(biāo)，并且使用輸入張量從在目標(biāo)上執(zhí)行的運(yùn)算符檢索其輸入。 在向后傳遞過程中，此函數(shù)的輸出梯度將發(fā)送到源節(jié)點(diǎn)并發(fā)送到適當(dāng)?shù)?code>send函數(shù)。
• 每個(gè)send-recv對(duì)都分配有一個(gè)全局唯一的autograd_message_id，以唯一地標(biāo)識(shí)該對(duì)。 這對(duì)于在反向傳遞期間在遠(yuǎn)程節(jié)點(diǎn)上查找對(duì)應(yīng)的功能很有用。
• 對(duì)于 RRef ，每當(dāng)我們調(diào)用 torch.distributed.rpc.RRef.to_here() 時(shí)，我們都會(huì)為所涉及的張量附加一個(gè)適當(dāng)?shù)膕end-recv對(duì)。

舉例來說，這就是我們上面示例中的 autograd 圖的樣子(為簡(jiǎn)單起見，排除了 t5.sum(））：

分布式 Autograd 上下文

每個(gè)使用分布式 autograd 的正向和反向傳遞都分配有唯一的 torch.distributed.autograd.context ，并且此上下文具有全局唯一的autograd_context_id。 根據(jù)需要在每個(gè)節(jié)點(diǎn)上創(chuàng)建此上下文。

此上下文具有以下目的：

1. 運(yùn)行分布式后向遍歷的多個(gè)節(jié)點(diǎn)可能會(huì)在同一張量上累積梯度，因此，在我們有機(jī)會(huì)運(yùn)行優(yōu)化器之前，張量的.grad字段將具有來自各種分布式后向遍歷的梯度。 這類似于在本地多次調(diào)用 torch.autograd.backward() 。 為了提供一種為每個(gè)后退通道分離梯度的方法，對(duì)于每個(gè)后退通道，梯度會(huì)累積在  torch.distributed.autograd.context 中。
2. 在前向傳遞過程中，我們?cè)谶@種情況下為每個(gè)自動(dòng)分級(jí)傳遞存儲(chǔ)send和recv函數(shù)。 這樣可以確保我們保留對(duì) autograd 圖中適當(dāng)節(jié)點(diǎn)的引用，以使其保持活動(dòng)狀態(tài)。 除此之外，在向后傳遞過程中很容易查找適當(dāng)?shù)?code>send和recv功能。
3. 通常，我們還使用此上下文為每個(gè)分布式 autograd pass 存儲(chǔ)一些元數(shù)據(jù)。

從用戶的角度來看，自動(dòng)分級(jí)上下文的設(shè)置如下：

import torch.distributed.autograd as dist_autograd
with dist_autograd.context() as context_id:
  loss = model.forward()
  dist_autograd.backward(loss)

分布式后向通行證

在本節(jié)中，我們概述了在分布式后向傳遞過程中準(zhǔn)確計(jì)算依賴項(xiàng)的挑戰(zhàn)，并描述了一些關(guān)于如何執(zhí)行分布式后向傳遞的算法(需要權(quán)衡）。

計(jì)算依賴

考慮以下代碼在單臺(tái)計(jì)算機(jī)上運(yùn)行

import torch
a = torch.rand((3, 3), requires_grad=True)
b = torch.rand((3, 3), requires_grad=True)
c = torch.rand((3, 3), requires_grad=True)
d = a + b
e = b * c
d.sum.().backward()

這就是上面代碼的 autograd 圖形：

autograd 引擎作為向后傳遞的一部分執(zhí)行的第一步是計(jì)算 autograd 圖中每個(gè)節(jié)點(diǎn)的依賴項(xiàng)數(shù)量。 這有助于 autograd 引擎知道何時(shí)可以執(zhí)行圖中的節(jié)點(diǎn)。 add(1)和mul(0)括號(hào)中的數(shù)字表示依賴項(xiàng)的數(shù)量。 如您所見，這意味著在向后傳遞期間，add節(jié)點(diǎn)需要 1 個(gè)輸入，mul節(jié)點(diǎn)不需要任何輸入(換句話說，不需要執(zhí)行）。 本地 autograd 引擎通過遍歷根節(jié)點(diǎn)中的圖來計(jì)算這些依賴性(在這種情況下為d）。

autograd 圖中的某些節(jié)點(diǎn)可能無法在向后傳遞中執(zhí)行的事實(shí)對(duì)分布式 autograd 提出了挑戰(zhàn)。 考慮使用 RPC 的這段代碼。

import torch
import torch.distributed.rpc as rpc


a = torch.rand((3, 3), requires_grad=True)
b = torch.rand((3, 3), requires_grad=True)
c = torch.rand((3, 3), requires_grad=True)


d = rpc.rpc_sync("worker1", torch.add, args=(a, b))
e = rpc.rpc_sync("worker1", torch.mul, args=(b, c))
loss = d.sum()

上面的代碼的相關(guān)自動(dòng)分級(jí)圖為：

計(jì)算此分布式 autograd 圖的依賴項(xiàng)更具挑戰(zhàn)性，并且需要一些開銷(無論是在計(jì)算還是在網(wǎng)絡(luò)通信方面）。

對(duì)于性能敏感的應(yīng)用程序，我們可以通過假設(shè)每個(gè)send和recv函數(shù)在反向傳遞中都是有效的(大多數(shù)應(yīng)用程序不執(zhí)行未使用的 RPC）來避免很多開銷。 這簡(jiǎn)化了分布式 autograd 算法，并且效率更高，但代價(jià)是應(yīng)用程序需要意識(shí)到這些限制。 該算法稱為 FAST 模式算法，下面將對(duì)其進(jìn)行詳細(xì)說明。

在一般情況下，可能不需要每個(gè)send和recv函數(shù)都有效作為反向傳遞的一部分。 為了解決這個(gè)問題，我們還有一個(gè) SMART 模式算法，將在后面的部分中進(jìn)行介紹。

快速模式算法

該算法的關(guān)鍵假設(shè)是，當(dāng)我們運(yùn)行向后傳遞時(shí)，每個(gè)send函數(shù)的相關(guān)性均為 1。 換句話說，我們假設(shè)將從另一個(gè)節(jié)點(diǎn)接收到 RPC 上的漸變。

算法如下：

1. 我們從具有向后遍歷的根的工作程序開始(所有根必須是本地的）。
2. 查找當(dāng)前分布式 Autograd 上下文的所有send功能。
3. 從提供的根目錄和我們檢索到的所有send函數(shù)開始，本地計(jì)算依賴項(xiàng)。
4. 計(jì)算依賴關(guān)系后，使用提供的根啟動(dòng)本地 autograd 引擎。
5. 當(dāng) autograd 引擎執(zhí)行recv功能時(shí)，recv功能會(huì)通過 RPC 將輸入梯度發(fā)送到適當(dāng)?shù)墓ぷ鞒绦颉?每個(gè)recv函數(shù)都知道目標(biāo)工作者 ID，因?yàn)樗挥涗洖檎騻鬟f的一部分。 recv功能還將autograd_context_id和autograd_message_id發(fā)送到遠(yuǎn)程主機(jī)。
6. 當(dāng)在遠(yuǎn)程主機(jī)上收到此請(qǐng)求時(shí)，我們使用autograd_context_id和autograd_message_id查找適當(dāng)?shù)?code>send功能。
7. 如果這是工作人員第一次收到對(duì)給定autograd_context_id的請(qǐng)求，則它將如上面的第 1-3 點(diǎn)所述在本地計(jì)算依賴性。
8. 然后，將在 6 中檢索到的send函數(shù)排隊(duì)以便在該工作者的本地 autograd 引擎上執(zhí)行。
9. 最后，我們不是在張量的.grad字段上累積梯度，而是根據(jù)分布式自學(xué)背景分別累積梯度。 梯度存儲(chǔ)在Dict[Tensor, Tensor]中，基本上是從 Tensor 到其相關(guān)梯度的映射，可以使用 get_gradients() API 檢索此映射。

例如，具有分布式 autograd 的完整代碼如下：

import torch
import torch.distributed.autograd as dist_autograd
import torch.distributed.rpc as rpc


def my_add(t1, t2):
  return torch.add(t1, t2)


## On worker 0:


## Setup the autograd context.
with dist_autograd.context() as context_id:
  t1 = torch.rand((3, 3), requires_grad=True)
  t2 = torch.rand((3, 3), requires_grad=True)


  # Perform some computation remotely.
  t3 = rpc.rpc_sync("worker1", my_add, args=(t1, t2))


  # Perform some computation locally based on remote result.
  t4 = torch.rand((3, 3), requires_grad=True)
  t5 = torch.mul(t3, t4)


  # Compute some loss.
  loss = t5.sum()


  # Run the backward pass.
  dist_autograd.backward([loss])


  # Retrieve the gradients from the context.
  dist_autograd.get_gradients(context_id)

具有依賴關(guān)系的分布式 autograd 圖如下所示：

應(yīng)用于以上示例的 FAST 模式算法如下：

1. 在Worker 0上，我們從根loss和send1開始計(jì)算依賴關(guān)系。 結(jié)果，send1的依賴性為 1，mul對(duì)Worker 0的依賴性為 1。
2. 現(xiàn)在，我們?cè)?code>Worker 0上啟動(dòng)本地 autograd 引擎。 我們首先執(zhí)行mul函數(shù)，將其輸出在 autograd 上下文中累積為t4的梯度。 然后，我們執(zhí)行recv2，它將梯度發(fā)送到Worker 1。
3. 由于這是Worker 1第一次聽到有關(guān)此反向傳遞的信息，因此它將開始依賴性計(jì)算并適當(dāng)?shù)貥?biāo)記send2，add和recv1的依賴性。
4. 接下來，將send2排隊(duì)在Worker 1的本地 autograd 引擎上，該引擎依次執(zhí)行add和recv1。
5. 當(dāng)執(zhí)行recv1時(shí)，它將梯度發(fā)送到Worker 0。
6. 由于Worker 0已經(jīng)計(jì)算了此向后傳遞的依賴性，因此它僅排隊(duì)并在本地執(zhí)行send1。
7. 最后，t1，t2和t4的梯度會(huì)累積在分布式 Autograd 上下文中。

SMART 模式算法

該算法的完整細(xì)節(jié)仍在研究中，但是對(duì)于一般概念，您可以參考 RFC 中的分布式 Autograd Algorithm Smart 模式部分。

分布式優(yōu)化器

DistributedOptimizer 的操作如下：

1. 獲取要優(yōu)化的遠(yuǎn)程參數(shù)列表 (RRef)。 這些也可以是包裝在本地RRef中的本地參數(shù)。
2. 將 Optimizer 類作為本地優(yōu)化器，以在所有不同的RRef所有者上運(yùn)行。
3. 分布式優(yōu)化器在每個(gè)工作程序節(jié)點(diǎn)上創(chuàng)建本地Optimizer的實(shí)例，并將其保存RRef。
4. 當(dāng)調(diào)用 torch.distributed.optim.DistributedOptimizer.step() 時(shí)，分布式優(yōu)化器使用 RPC 在適當(dāng)?shù)倪h(yuǎn)程工作器上遠(yuǎn)程執(zhí)行所有本地優(yōu)化器。
5. 如果多個(gè)并發(fā)的分布式優(yōu)化器正在更新工作器上的相同參數(shù)，則這些更新將通過鎖序列化。

簡(jiǎn)單的端到端示例

綜上所述，以下是使用分布式 autograd 和分布式優(yōu)化器的簡(jiǎn)單的端到端示例。 如果將代碼放入名為“ dist_autograd_simple.py”的文件中，則可以使用命令MASTER_ADDR="localhost" MASTER_PORT=29500 python dist_autograd_simple.py運(yùn)行該代碼：

import multiprocessing as mp
import torch
import torch.distributed.autograd as dist_autograd
from torch.distributed import rpc
from torch import optim
from torch.distributed.optim import DistributedOptimizer


def random_tensor():
    return torch.rand((3, 3), requires_grad=True)


def _run_process(rank, dst_rank, world_size):
    name = "worker{}".format(rank)
    dst_name = "worker{}".format(dst_rank)


    # Initialize RPC.
    rpc.init_rpc(
        name=name,
        rank=rank,
        world_size=world_size
    )


    # Use a distributed autograd context.
    with dist_autograd.context() as context_id:
        # Forward pass (create references on remote nodes).
        rref1 = rpc.remote(dst_name, random_tensor)
        rref2 = rpc.remote(dst_name, random_tensor)
        loss = rref1.to_here() + rref2.to_here()


        # Backward pass (run distributed autograd).
        dist_autograd.backward([loss.sum()])


        # Build DistributedOptimizer.
        dist_optim = DistributedOptimizer(
        optim.SGD,
        [rref1, rref2],
        lr=0.05,
        )


        # Run the distributed optimizer step.
        dist_optim.step()


def run_process(rank, dst_rank, world_size):
    _run_process(rank, dst_rank, world_size)
    rpc.shutdown()


processes = []


## Run world_size workers.
world_size = 2
for i in range(world_size):
    p = mp.Process(target=run_process, args=(i, (i + 1) % 2, world_size))
    p.start()
    processes.append(p)


for p in processes:
    p.join()