使用 TensorPipe CUDA RPC 進行直接的裝置對裝置通訊

建立於：2021 年 3 月 19 日 | 最後更新：2021 年 3 月 19 日 | 最後驗證：2024 年 11 月 05 日

注意

直接裝置對裝置 RPC (CUDA RPC) 在 PyTorch 1.8 中作為原型功能引入。此 API 可能會變更。

在本食譜中，您將學習

• CUDA RPC 的高階概念。

• 如何使用 CUDA RPC。

需求

• PyTorch 1.8+

• 分散式 RPC 框架入門

什麼是 CUDA RPC？

CUDA RPC 支援將張量從本機 CUDA 記憶體直接傳送到遠端 CUDA 記憶體。 在 v1.8 版本之前，PyTorch RPC 僅接受 CPU 張量。 因此，當應用程式需要透過 RPC 傳送 CUDA 張量時，它必須先將張量移至呼叫端的 CPU，然後透過 RPC 傳送它，然後將其移動到被呼叫端的目的地裝置，這會導致不必要的同步處理以及 D2H 和 H2D 複製。 從 v1.8 開始，RPC 允許使用者使用 set_device_map API 配置每個處理程序的全域裝置映射，指定如何將本機裝置映射到遠端裝置。 更具體地說，如果 worker0 的裝置映射具有一個條目 "worker1" : {"cuda:0" : "cuda:1"}，則 worker0 上 "cuda:0" 的所有 RPC 參數將直接傳送到 worker1 上的 "cuda:1"。 RPC 的回應將使用呼叫端裝置映射的反向，即，如果 worker1 在 "cuda:1" 上傳回一個張量，它將直接傳送到 worker0 上的 "cuda:0"。 所有預期的裝置對裝置直接通訊都必須在每個處理程序的裝置映射中指定。 否則，僅允許使用 CPU 張量。

在底層，PyTorch RPC 依賴於 TensorPipe 作為通訊後端。 PyTorch RPC 從每個請求或回應中提取所有張量到一個清單中，並將所有其他內容打包到二進位酬載中。 然後，TensorPipe 將根據張量裝置類型以及呼叫端和被呼叫端的通道可用性，自動為每個張量選擇一個通訊通道。 現有的 TensorPipe 通道涵蓋 NVLink、InfiniBand、SHM、CMA、TCP 等。

如何使用 CUDA RPC？

下面的程式碼展示了如何使用 CUDA RPC。 該模型包含兩個線性層，並分為兩個分片。 這兩個分片分別放置在 worker0 和 worker1 上，worker0 充當驅動前向和後向傳遞的主節點。 請注意，我們有意跳過了 DistributedOptimizer，以突顯使用 CUDA RPC 時的效能改進。 實驗重複前向和後向傳遞 10 次，並測量總執行時間。 它比較了使用 CUDA RPC 與手動暫存到 CPU 記憶體並使用 CPU RPC 的情況。

import torch
import torch.distributed.autograd as autograd
import torch.distributed.rpc as rpc
import torch.multiprocessing as mp
import torch.nn as nn

import os
import time


class MyModule(nn.Module):
    def __init__(self, device, comm_mode):
        super().__init__()
        self.device = device
        self.linear = nn.Linear(1000, 1000).to(device)
        self.comm_mode = comm_mode

    def forward(self, x):
        # x.to() is a no-op if x is already on self.device
        y = self.linear(x.to(self.device))
        return y.cpu() if self.comm_mode == "cpu" else y

    def parameter_rrefs(self):
        return [rpc.RRef(p) for p in self.parameters()]


def measure(comm_mode):
    # local module on "worker0/cuda:0"
    lm = MyModule("cuda:0", comm_mode)
    # remote module on "worker1/cuda:1"
    rm = rpc.remote("worker1", MyModule, args=("cuda:1", comm_mode))
    # prepare random inputs
    x = torch.randn(1000, 1000).cuda(0)

    tik = time.time()
    for _ in range(10):
        with autograd.context() as ctx:
            y = rm.rpc_sync().forward(lm(x))
            autograd.backward(ctx, [y.sum()])
    # synchronize on "cuda:0" to make sure that all pending CUDA ops are
    # included in the measurements
    torch.cuda.current_stream("cuda:0").synchronize()
    tok = time.time()
    print(f"{comm_mode} RPC total execution time: {tok - tik}")


def run_worker(rank):
    os.environ['MASTER_ADDR'] = 'localhost'
    os.environ['MASTER_PORT'] = '29500'
    options = rpc.TensorPipeRpcBackendOptions(num_worker_threads=128)

    if rank == 0:
        options.set_device_map("worker1", {0: 1})
        rpc.init_rpc(
            f"worker{rank}",
            rank=rank,
            world_size=2,
            rpc_backend_options=options
        )
        measure(comm_mode="cpu")
        measure(comm_mode="cuda")
    else:
        rpc.init_rpc(
            f"worker{rank}",
            rank=rank,
            world_size=2,
            rpc_backend_options=options
        )

    # block until all rpcs finish
    rpc.shutdown()


if __name__=="__main__":
    world_size = 2
    mp.spawn(run_worker, nprocs=world_size, join=True)

下面顯示了輸出，顯示在此實驗中，CUDA RPC 可以幫助實現比 CPU RPC 快 34 倍的速度。

cpu RPC total execution time: 2.3145179748535156 Seconds
cuda RPC total execution time: 0.06867480278015137 Seconds