[原始碼解析] PyTorch 分散式(18) --- 使用 RPC 的分散式管道並行
[原始碼解析] PyTorch 分散式(18) --- 使用 RPC 的分散式管道並行
目錄0x00 摘要
在前面的文章之中，我們已經學習了PyTorch 分散式的基本模組，接下來我們通過幾篇文章來看看如何把這些模組應用到實踐之中，順便把PyTorch分散式邏輯整體梳理一下。本文介紹如何使用 RPC 來完成分散式管道並行。
本文以DISTRIBUTED PIPELINE PARALLELISM USING RPC
PyTorch分散式其他文章如下：
[原始碼解析]深度學習利器之自動微分(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 分散式 Autograd (2) ---- RPC基礎
[原始碼解析] PyTorch 分散式 Autograd (3) ---- 上下文相關
[原始碼解析] PyTorch 分散式 Autograd (4) ---- 如何切入引擎
[原始碼解析] PyTorch 分散式 Autograd (5) ---- 引擎(上)
[原始碼解析] PyTorch 分散式 Autograd (6) ---- 引擎(下)
[原始碼解析] PyTorch分散式優化器(1)----基石篇
[原始碼解析] PyTorch分散式優化器(2)----資料並行優化器
[原始碼解析] PyTorch分散式優化器(3)---- 模型並行
[原始碼解析] PyTorch 分散式(14) --使用 Distributed Autograd 和 Distributed Optimizer
[原始碼解析] PyTorch 分散式(15) --- 使用分散式 RPC 框架實現引數伺服器
[原始碼解析] PyTorch 分散式(16) --- 使用非同步執行實現批處理 RPC
[原始碼解析] PyTorch 分散式(17) --- 結合DDP和分散式 RPC 框架
注：本文沒有完全按照原文順序進行翻譯，而是按照自己理解的思路重新組織了文章。原文是從下至上，從細節到整體的順序分析，但是我在理解時候總覺得彆扭，缺乏一個總體的感知，所以我們還是以從上到下的邏輯，配合圖例進行分析。
0x01 綜述
1.1 先決條件
本教程使用 Resnet50 模型來演示使用torch.distributed.rpc API實現分散式管道並行。這可以看作是單機模型並行最佳實踐中討論的多 GPU 流水線並行的分散式對應版本。
本文的先決條件如下：
注意
-
本教程需要 PyTorch v1.6.0 或更高版本。
-
本教程的完整原始碼可以在pytorch/examples找到 。
1.2 基礎知識
之前的教程分散式 RPC 框架入門 展示瞭如何使用torch.distributed.rpc 為 RNN 模型實現分散式模型並行。該教程使用一個 GPU 來託管EmbeddingTable，並且提供的程式碼執行良好。但是，如果模型存在於多個 GPU 上，則需要一些額外的步驟來提高所有 GPU 的攤銷利用率。管道並行就是一種在這種情況下可以提供幫助的正規化。
在本教程中，我們使用ResNet50作為示例模型，單機模型並行最佳實踐 教程也使用該模型。類似地，ResNet50模型被分成兩個分片，輸入批次被分成多個分片，並以流水線方式輸入到兩個模型分片中。不同之處在於，本教程不是使用 CUDA 流並行執行，而是呼叫非同步 RPC。因此，本教程中提供的解決方案也適用於跨機器邊界。本教程的其餘部分將分四個步驟介紹實現。
0x02 啟動
下面的程式碼顯示了所有程序的目標函式，在所有節點上都會執行 run_worker，但是其執行程式碼不同。
- 主要邏輯定義在
run_master之中，這是本系統的大腦和實際執行者。 - worker 被動地等待來自 master 的命令，因此只執行
init_rpcandshutdown。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](