[原始碼解析] PyTorch 分散式(18) --- 使用 RPC 的分散式管道並行

之前的教程分散式 RPC 框架入門展示瞭如何使用torch.distributed.rpc 為 RNN 模型實現分散式模型並行。該教程使用一個 GPU 來託管EmbeddingTable，並且提供的程式碼執行良好。但是，如果模型存在於多個 GPU 上，則需要一些額外的步驟來提高所有 GPU 的攤銷利用率。管道並行就是一種在這種情況下可以提供幫助的正規化。

在本教程中，我們使用ResNet50作為示例模型，單機模型並行最佳實踐教程也使用該模型。類似地，ResNet50模型被分成兩個分片，輸入批次被分成多個分片，並以流水線方式輸入到兩個模型分片中。不同之處在於，本教程不是使用 CUDA 流並行執行，而是呼叫非同步 RPC。因此，本教程中提供的解決方案也適用於跨機器邊界。本教程的其餘部分將分四個步驟介紹實現。

0x02 啟動

下面的程式碼顯示了所有程式的目標函式，在所有節點上都會執行 run_worker，但是其執行程式碼不同。

主要邏輯定義在run_master之中，這是本系統的大腦和實際執行者。
worker 被動地等待來自 master 的命令，因此只執行init_rpcand shutdown。
- init_rpc只是建立分散式環境。
- shutdown預設情況下將阻塞，直到所有 RPC 參與者結束工作。
- 具體業務工作都是master通過RPC直接排程到worker節點上來執行。

def run_worker(rank, world_size, num_split):
    os.environ['MASTER_ADDR'] = 'localhost'
    os.environ['MASTER_PORT'] = '29500'

    # Higher timeout is added to accommodate for kernel compilation time in case of ROCm.
    options = rpc.TensorPipeRpcBackendOptions(num_worker_threads=256, rpc_timeout=300)

    if rank == 0:
        rpc.init_rpc(
            "master",
            rank=rank,
            world_size=world_size,
            rpc_backend_options=options
        )
        run_master(num_split)
    else:
        rpc.init_rpc(
            f"worker{rank}",
            rank=rank,
            world_size=world_size,
            rpc_backend_options=options
        )
        pass

    # block until all rpcs finish
    rpc.shutdown()


if __name__=="__main__":
    world_size = 3
    for num_split in [1, 2, 4, 8]:
        tik = time.time()
        mp.spawn(run_worker, args=(world_size, num_split), nprocs=world_size, join=True)
        tok = time.time()
        print(f"number of splits = {num_split}, execution time = {tok - tik}")

邏輯如下：

           torch.multiprocessing.spawn
                      +
                      |
                      |
         +------------+----------------------------+
         |                                         |
         |                                         |
         v                                         v
+--------+---------------------------+      +------+----------+
| "ps"                               |      | f"worker{rank}" |
|                                    |      |                 |
|     rank = 0                       |      |      rank = 1,2 |
|                                    |      |                 |
|     run_worker +----> run_master   |      |      run_worker |
|                                    |      |                 |
+------------------------------------+      +-----------------+

0x03 定義訓練迴圈

現在我們看看訓練迴圈（training loop）。我們使用專門的 "master " worker 來準備隨機輸入和標籤，並控制分散式反向傳播和分散式優化器step。

它首先建立DistResNet50模組的一個例項，指定了每個批次的微批次數量，還提供了兩個 RPC 工作執行緒的名稱（即“worker1”和“worker2”）。
然後定義了損失函式並使用parameter_rrefs()拿到了一個引數列表RRefs，以此建立了DistributedOptimizer。
最後，主訓練迴圈與常規本地訓練非常相似，不同之處在於它用於dist_autograd啟動後向傳播，併為後向傳播和優化器 step()提供了 context_id。

#########################################################
#                   Run RPC Processes                   #
#########################################################

num_batches = 3
batch_size = 120
image_w = 128
image_h = 128


def run_master(split_size):

    # put the two model parts on worker1 and worker2 respectively
    model = DistResNet50(split_size, ["worker1", "worker2"])
    loss_fn = nn.MSELoss()
    opt = DistributedOptimizer( # 分散式優化器
        optim.SGD,
        model.parameter_rrefs(),
        lr=0.05,
    )

    one_hot_indices = torch.LongTensor(batch_size) \
                           .random_(0, num_classes) \
                           .view(batch_size, 1)

    for i in range(num_batches):
        print(f"Processing batch {i}")
        # generate random inputs and labels
        inputs = torch.randn(batch_size, 3, image_w, image_h)
        labels = torch.zeros(batch_size, num_classes) \
                      .scatter_(1, one_hot_indices, 1)

        # The distributed autograd context is the dedicated scope for the
        # distributed backward pass to store gradients, which can later be
        # retrieved using the context_id by the distributed optimizer.
        with dist_autograd.context() as context_id:
            outputs = model(inputs)
            dist_autograd.backward(context_id, [loss_fn(outputs, labels)])  # 分散式梯度
            opt.step(context_id)

我們先按照單機的思路來畫圖，下一節會再擴充套件。從單機角度看，好像沒啥稀奇的地方。

           torch.multiprocessing.spawn
                      +
                      |
                      |
         +------------+-------------------------------------------+
         |                                                        |
         |                                                        |
         v                                                        v
+--------+----------------------------------------------+  +------+----------+
| "ps"                                                  |  | f"worker{rank}" |
|                                                       |  |                 |
|     rank = 0                                          |  |      rank = 1,2 |
|                                                       |  |                 |
|     run_worker +----> run_master                      |  |      run_worker |
|                           +                           |  |                 |
|                           |                           |  |                 |
|                           |                           |  +-----------------+
|                           v                           |
| +-------------------------+-------------------------+ |
| |                                                   | |
| |                                                   | |
| |  model = DistResNet50(split_size,                 | |
| |                      ["worker1", "worker2"])      | |
| |  loss_fn = nn.MSELoss()                           | |
| |  opt = DistributedOptimizer(                      | |
| |      optim.SGD,                                   | |
| |      model.parameter_rrefs(),                     | |
| |      lr=0.05,                                     | |
| |  )                                                | |
| |  for i in range(num_batches):                     | |
| |      with dist_autograd.context() as context_id:  | |
| |          outputs = model(inputs)                  | |
| |          dist_autograd.backward(context_id,       | |
| |                    [loss_fn(outputs, labels)])    | |
| |          opt.step(context_id)                     | |
| |                                                   | |
| |                                                   | |
| +---------------------------------------------------+ |
|                                                       |
+-------------------------------------------------------+

0x04 將 ResNet50 模型分片拼接成一個模組

我們這裡先假定分片是個黑盒子。

首先，我們建立一個DistResNet50模組來組裝兩個分片並實現流水線並行邏輯。在建構函式中，我們使用兩次 rpc.remote呼叫將兩個分片分別放在兩個不同的 RPC 工作執行緒上，並保持RRef指向到兩個模型部分，以便在前向傳遞中引用它們。
forward函式將輸入批次拆分為多個微批次，並以流水線方式將這些微批次提供給兩個模型部件。
- 首先使用 rpc.remote呼叫將第一個分片應用於微批次，然後將中間輸出RRef轉發到第二個模型分片。
- 之後收集所有微輸出（micro-outputs）的Future ，並在迴圈後等待所有微輸出。
- 請注意，remote()和rpc_async()都立即返回並非同步執行。因此，整個迴圈是非阻塞的，並且會同時啟動多個 RPC。
一個 micro-batch 在兩個模型部分上的執行順序由一箇中間輸出y_rref變數來維護。微批次之間的執行順序無關緊要。
最後，前向函式將所有微批次的輸出連線成一個單一的輸出張量並返回。該parameter_rrefs函式可以讓我們簡化分散式優化器構建，後面會用到。parameter_rrefs 的作用是：從 worker 1，worker 2 取出每個分片需要優化的引數。最後這些引數會傳遞給DistributedOptimizer。

class DistResNet50(nn.Module):
    """
    Assemble two parts as an nn.Module and define pipelining logic
    """
    def __init__(self, split_size, workers, *args, **kwargs):
        super(DistResNet50, self).__init__()

        self.split_size = split_size

        # Put the first part of the ResNet50 on workers[0]
        self.p1_rref = rpc.remote(
            workers[0], # 放到第一個worker之上
            ResNetShard1,
            args = ("cuda:0",) + args,
            kwargs = kwargs
        )

        # Put the second part of the ResNet50 on workers[1]
        self.p2_rref = rpc.remote(
            workers[1], # 放到第二個worker之上
            ResNetShard2,
            args = ("cuda:1",) + args,
            kwargs = kwargs
        )

    def forward(self, xs):
        # Split the input batch xs into micro-batches, and collect async RPC
        # futures into a list
        out_futures = []
        for x in iter(xs.split(self.split_size, dim=0)): # 將輸入批次拆分為多個微批次
            x_rref = RRef(x) # 封裝成RRef
            y_rref = self.p1_rref.remote().forward(x_rref) # 第一個worker處理微批次
            z_fut = self.p2_rref.rpc_async().forward(y_rref) # 第二個worker繼續處理
            out_futures.append(z_fut)

        # collect and cat all output tensors into one tensor.
        return torch.cat(torch.futures.wait_all(out_futures))

    def parameter_rrefs(self):
        remote_params = []
        remote_params.extend(self.p1_rref.remote().parameter_rrefs().to_here())
        remote_params.extend(self.p2_rref.remote().parameter_rrefs().to_here())
        return remote_params

為了演示，我們這裡只畫出了一個worker 1 的內部細節，請大家記住，worker 1 和 worker 2是一樣的。同時，對 run_master 也進行了簡化。流水線是在master的forward方法之中完成，具體在圖上的1，2兩個數字代表的箭頭上體現。

         torch.multiprocessing.spawn
                    +
                    |
                    |
       +------------+---------------------------------------------------+
       |                                                                |
       |                                                                |
       v                                                                v
+------+----------------------------------------------+         +-------+--------------+
|  "ps"                                     rank = 0  |         |"worker 1"   rank = 1 |
|                                                     |         |                      |
|  run_worker   DistributedOptimizer(p1_rref,p2_rref) |         | run_worker           |
|      +                                              |         |                      |
|      |                                              |         |                      |
|      |        DistResNet50                          |         |      +-------------+ |
|      |                                              |    +-------->  |ResNetShard1 | |
|      v                                              |    |    |      |             | |
|  run_master   p1_rref  +------------------------------------------>  |             | |
|      +                                              |    |    |      +-------+-----+ |
|      |                                              |    |    |              |       |
|      |        p2_rref +-------------------------------+  |    +----------------------+
|      |                                              | |  |                   |
|      |                                              | |  |                   |
|      v                                              | |  |                   |
| +----+--------------------------------------------+ | |  |                   |
| | model = DistResNet50(split_size,                | | |  |                   |
| |                     ["worker1", "worker2"])     | | |  |1                  |2
| | loss_fn = nn.MSELoss()                          | | |  |                   |
| | opt = DistributedOptimizer(                     | | |  |                   |
| |     optim.SGD,                                  | | |  |                   |
| |     model.parameter_rrefs(),                    | | |  |                   |
| | )                                               | | |  |                   v
| | for i in range(num_batches):                    | | |  |    +--------------+--------+
| |     with dist_autograd.context() as context_id: | | |  |    | "worker 2"  rank = 2  |
| |         outputs = model(inputs) +----------------------+    |                       |
| |         dist_autograd.backward(context_id,      | | |       |      +--------------+ |
| |                   [loss_fn(outputs, labels)])   | | +------------> |ResNetShard2  | |
| |         opt.step(context_id)                    | |         |      |              | |
| +-------------------------------------------------+ |         |      +--------------+ |
+-----------------------------------------------------+         +-----------------------+

0x05 對 ResNet50 模型進行分割槽

這是ResNet50在兩個模型分片中實現的準備步驟。下面的程式碼是從torchvision 中的 ResNet 實現中借用的。該ResNetBase模組包含兩個 ResNet 分片（shards）的通用構建塊和屬性。

現在，我們已準備好定義兩個模型分片。在建構函式之中，我們簡單地將所有 ResNet50 層分成兩部分，並將每個部分移動到提供的裝置中。兩個分片的forward功能如下：

獲取一個輸入資料的RRef，這樣就可以在本地獲取資料，然後將其移動到預期的裝置之上。
將所有層應用於輸入後，它將輸出移動到 CPU 並返回。這是因為 RPC API 需要張量駐留在 CPU 上，以避免在呼叫方和被呼叫方中的裝置數量不匹配時出現無效裝置錯誤。

import threading
import torch
import torch.nn as nn
from torchvision.models.resnet import Bottleneck

num_classes = 1000

def conv1x1(in_planes, out_planes, stride=1):
    """1x1 convolution"""
    return nn.Conv2d(in_planes, out_planes, kernel_size=1, stride=stride, bias=False)

class ResNetBase(nn.Module):
    def __init__(self, block, inplanes, num_classes=1000,
                 groups=1, width_per_group=64, norm_layer=None):
        super(ResNetBase, self).__init__()

        self._lock = threading.Lock()
        self._block = block
        self._norm_layer = nn.BatchNorm2d
        self.inplanes = inplanes
        self.dilation = 1
        self.groups = groups
        self.base_width = width_per_group

    # 輔助函式，用來構建Sequential    
    def _make_layer(self, planes, blocks, stride=1):
        norm_layer = self._norm_layer
        downsample = None
        previous_dilation = self.dilation
        if stride != 1 or self.inplanes != planes * self._block.expansion:
            downsample = nn.Sequential(
                conv1x1(self.inplanes, planes * self._block.expansion, stride),
                norm_layer(planes * self._block.expansion),
            )

        layers = []
        layers.append(self._block(self.inplanes, planes, stride, downsample, self.groups,
                                  self.base_width, previous_dilation, norm_layer))
        self.inplanes = planes * self._block.expansion
        for _ in range(1, blocks):
            layers.append(self._block(self.inplanes, planes, groups=self.groups,
                                      base_width=self.base_width, dilation=self.dilation,
                                      norm_layer=norm_layer))

        return nn.Sequential(*layers)

    def parameter_rrefs(self):
        r"""
        Create one RRef for each parameter in the given local module, and return a
        list of RRefs.
        """
        return [RRef(p) for p in self.parameters()]


class ResNetShard1(ResNetBase):
    """
    The first part of ResNet.
    """
    def __init__(self, device, *args, **kwargs):
        super(ResNetShard1, self).__init__(
            Bottleneck, 64, num_classes=num_classes, *args, **kwargs)

        self.device = device # 配置裝置
        self.seq = nn.Sequential( # 構建Sequential模組
            nn.Conv2d(3, self.inplanes, kernel_size=7, stride=2, padding=3, bias=False),
            self._norm_layer(self.inplanes),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(kernel_size=3, stride=2, padding=1),
            self._make_layer(64, 3),
            self._make_layer(128, 4, stride=2)
        ).to(self.device) # 放到裝置之上

        for m in self.modules():
            if isinstance(m, nn.Conv2d):
                nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
            elif isinstance(m, nn.BatchNorm2d):
                nn.init.ones_(m.weight)
                nn.init.zeros_(m.bias)

    def forward(self, x_rref):
        x = x_rref.to_here().to(self.device) # 把輸入放到裝置之上
        with self._lock:
            out =  self.seq(x) # 把所有層都應用到輸入之上
        return out.cpu() # 輸出需要移動到CPU之上


class ResNetShard2(ResNetBase):
    """
    The second part of ResNet.
    """
    def __init__(self, device, *args, **kwargs):
        super(ResNetShard2, self).__init__(
            Bottleneck, 512, num_classes=num_classes, *args, **kwargs)

        self.device = device # 配置裝置
        self.seq = nn.Sequential( # 構建Sequential模組
            self._make_layer(256, 6, stride=2),
            self._make_layer(512, 3, stride=2),
            nn.AdaptiveAvgPool2d((1, 1)),
        ).to(self.device) # 放到裝置上

        self.fc =  nn.Linear(512 * self._block.expansion, num_classes).to(self.device)

    def forward(self, x_rref):
        x = x_rref.to_here().to(self.device) # 把輸入放到裝置之上
        with self._lock:
            out =  self.seq(x) # 把所有層都應用到輸入之上
        return out.cpu() # 輸出需要移動到CPU之上

我們把目前邏輯擴充如下，這裡：

DistResNet50 被分成兩個部分，分別在 worker 1，worker 2之上。
兩個部分的引數通過RRef儲存在master之上。
ps就是master，它負責驅動全部業務。
通過分散式優化器和分散式autograd完成後向傳播。
兩個worker就是簡單執行而已：
- 負責搭建分散式環境和等待結束。
- 具體工作是由master通過RPC直接放到worker之上執行。
流水線則是在master的forward 之上顯式配置，由一箇中間輸出來完成，具體在圖上的1，2兩個數字代表的箭頭上體現。

      torch.multiprocessing.spawn
                 +
                 |
                 |
    +------------+---------------------------------------------------+---------------------+
    |                                                                |                     |
    |                                                                |                     |
    v                                                                v                     |
+---+----------------------------------------------+        +--------+------------------+  |
| "ps"                                   rank = 0  |        |"worker 1“       rank = 1  |  |
|                                                  |        |                           |  |
| run_worker DistributedOptimizer(p1_rref,p2_rref) |        | run_worker                |  |
|     +                                            |        |  +----------------------+ |  |
|     |                                            |        |  |  ResNetShard1        | |  |
|     |         DistResNet50 +--------+-------------------> |  |   +----------------+ | |  |
|     |                               |            |        |  |   | ResNetBase     | | |  |
|     v                               |            |        |  |   |                | | |  |
| run_master    p1_rref  +--------------------------------------------> parameters()| | |  |
|     +                               |            |        |  |   |                | | |  |
|     |                               |            |    +-> |  |   |                | | |  |
|     |         p2_rref  +---------------------------+  |   |  |   |                | | |  |
|     |                               |            | |  |   |  |   +----------------+ | |  |
|     |                               |            | |  |   |  +----------------------+ |  |
|     |                               |            | |  |   +---------------------------+  |
|     |                               |            | |  |                |                 |
|     |                               +----------------------------+     |                 |
|     |                                            | |  |          |     |                 |
|     |                                            | |  |          |     |                 |
|     v                                            | |  |          |     |                 |
| +---+------------------------------------------+ | |  | 1        |     |2                |
| | model = DistResNet50(split_size,             | | |  |          |     |                 |
| |                     ["worker1", "worker2"])  | | |  |          |     |                 |
| | loss_fn = nn.MSELoss()                       | | |  |          V     v                 |
| | opt = DistributedOptimizer(                  | | |  |   +------------+--------------+  |
| |   optim.SGD,                                 | | |  |   |"worker 2"       rank = 2  |  |
| |   model.parameter_rrefs(),                   | | |  |   |                           |  |
| | )                                            | | |  |   |    +--------------------+ |  |
| | for i in range(num_batches):                 | | |  |   |    | ResNetShard2       | |  |
| |   with dist_autograd.context() as context_id:| | |  |   |    | +----------------+ | |  |
| |                                              | | |  |   |    | |  ResNetBase    | | +<-+
| |       outputs = model(inputs) +---------------------+   |    | |                | | |
| |                                              | | |      |    | |                | | |
| |       dist_autograd.backward(context_id,     | | +----------------> parameters()| | |
| |                 [loss_fn(outputs, labels)])  | |        |    | |                | | |
| |       opt.step(context_id)                   | |        |    | +----------------+ | |
| +----------------------------------------------+ |        |    +--------------------+ |
+--------------------------------------------------+        +---------------------------+

手機如下：

至此，PyTorch 這幾篇官方示例文章都剖析完畢，從下一篇我們開始介紹彈性訓練，敬請期待。

0xFF 參考

[DISTRIBUTED PIPELINE PARALLELISM USING RPC](

[原始碼解析] PyTorch 分散式(18) --- 使用 RPC 的分散式管道並行

[原始碼解析] PyTorch 分散式(18) --- 使用 RPC 的分散式管道並行

0x00 摘要

0x01 綜述

1.1 先決條件

1.2 基礎知識

0x02 啟動

0x03 定義訓練迴圈

0x04 將 ResNet50 模型分片拼接成一個模組

0x05 對 ResNet50 模型進行分割槽

0xFF 參考

相關文章