PyTorch 深度剖析：如何使用模型并行技術(shù) （Model Parallel）
極市導(dǎo)讀
?本文介紹了模型并行技術(shù) Model Parallel，與 DataParallel 相反，Model Parallel 將一個單一的模型分割到不同的 GPU 上，而不是在每個 GPU 上復(fù)制整個模型。Pipeline Model Parallel 是一種進一步加速模型并行的策略。?>>加入極市CV技術(shù)交流群，走在計算機視覺的最前沿
0 背景
模型并行 (Single Machine Model Parallel) 在分布式訓(xùn)練技術(shù)中廣泛使用。
它和 DataParallel 不同，DataParallel 是在多個GPU上訓(xùn)練神經(jīng)網(wǎng)絡(luò)，將同一個模型復(fù)制到所有 GPU 上，每個 GPU 消耗不同的輸入數(shù)據(jù)分區(qū)。雖然 DataParallel 可以大大加快訓(xùn)練過程，但是，有的時候一些模型太大，沒辦法裝入單個 GPU 的時候，就不適用了。
這篇文章展示了如何通過使用模型并行來解決這個問題，與 DataParallel 相反，Model Parallel 將一個單一的模型分割到不同的 GPU 上，而不是在每個 GPU 上復(fù)制整個模型。具體來說，比如一個模型 m 包含10層：當(dāng)使用 DataParallel 時，每個 GPU 將有這10層的每個副本，而在兩個 GPU 上使用模型并行時，每個 GPU 只需要承載5層。
Model Parallel 的 High-level 的理念是將模型的不同子網(wǎng)絡(luò)放在不同的設(shè)備上，并相應(yīng)地實現(xiàn) forward() 方法，在 GPU 之間移動模型中間輸出。每個模型被拆成了多塊，只有一塊在單獨的設(shè)備上運行，所以，一組設(shè)備可以共同為一個更大的模型服務(wù)。
在這篇文章中，我們不會試圖構(gòu)建巨大的模型，并將其硬擠入數(shù)量有限的 GPU。相反，這篇文章的重點是展示 Model Parallel 的具體操作方法。
1 Model Parallel 基操
比如現(xiàn)在有一個包含2個 Linear layers 的模型，我們想在2塊 GPU 上 run 它，辦法可以是在每塊 GPU 上放置1個 Linear layer，并且把得到的中間結(jié)果在 GPU 之間移動。代碼可以是這樣子：
import torch
import torch.nn as nn
import torch.optim as optim
class ToyModel(nn.Module):
def __init__(self):
super(ToyModel, self).__init__()
self.net1 = torch.nn.Linear(10, 10).to('cuda:0')
self.relu = torch.nn.ReLU()
self.net2 = torch.nn.Linear(10, 5).to('cuda:1')
def forward(self, x):
x = self.relu(self.net1(x.to('cuda:0')))
return self.net2(x.to('cuda:1'))
注意，上述 ToyModel 看起來與在單個 GPU 上的實現(xiàn)方式非常相似，除了四個 to(device) 的調(diào)用，將 Linear layer 和張量放在適當(dāng)?shù)脑O(shè)備上。這是該模型中唯一需要改變的地方。backward() 和 torch.optim 將自動處理梯度問題，就像模型是在一個 GPU 上一樣。
你只需要確保在調(diào)用損失函數(shù)時，標(biāo)簽和輸出是在同一個設(shè)備上。像下面這樣：
model = ToyModel()
loss_fn = nn.MSELoss()
optimizer = optim.SGD(model.parameters(), lr=0.001)
optimizer.zero_grad()
outputs = model(torch.randn(20, 10))
labels = torch.randn(20, 5).to('cuda:1')
loss_fn(outputs, labels).backward()
optimizer.step()
這里應(yīng)該把標(biāo)簽 labels 放在1號 GPU 上面，因為模型的輸出就在1號 GPU 上。
2 對已有的模塊使用 Model Parallel
這段我們介紹如何在多個 GPU 上運行一個現(xiàn)有的單 GPU 模塊，只需做幾行修改即可。
下面的代碼顯示了如何將 torchvision.models.resnet50() 分解到兩個 GPU。這個想法是繼承現(xiàn)有的 ResNet 模塊，并在構(gòu)建過程中將各層分割到兩個 GPU 上面。然后，overwrite forward() 方法，通過相應(yīng)地移動中間輸出來縫合兩個子網(wǎng)絡(luò)。
from torchvision.models.resnet import ResNet, Bottleneck
num_classes = 1000
class ModelParallelResNet50(ResNet):
def __init__(self, *args, **kwargs):
super(ModelParallelResNet50, self).__init__(
Bottleneck, [3, 4, 6, 3], num_classes=num_classes, *args, **kwargs)
self.seq1 = nn.Sequential(
self.conv1,
self.bn1,
self.relu,
self.maxpool,
self.layer1,
self.layer2
).to('cuda:0')
self.seq2 = nn.Sequential(
self.layer3,
self.layer4,
self.avgpool,
).to('cuda:1')
self.fc.to('cuda:1')
def forward(self, x):
x = self.seq2(self.seq1(x).to('cuda:1'))
return self.fc(x.view(x.size(0), -1))
上述實現(xiàn)解決了模型太大，無法裝入單個GPU的情況下的問題。然而，對于運行速度而言，它將比在單個 GPU 上運行的速度要慢。這是因為，在任何時候，兩個 GPU 中只有一個在工作，而另一個則是坐在那里啥也不干。由于中間輸出需要在第二層和第三層之間從 cuda:0復(fù)制到 cuda:1，所以性能會進一步惡化。
下面我們通過一個實驗，看看具體的程序運行時間的量化對比。在這個實驗中，我們通過運行隨機輸入和標(biāo)簽來訓(xùn)練 ModelParallelResNet50 和現(xiàn)有的 torchvision.models.resnet50()。在訓(xùn)練之后，這些模型不會產(chǎn)生任何有用的預(yù)測，但我們可以對執(zhí)行時間有一個合理的了解。
import torchvision.models as models
num_batches = 3
batch_size = 120
image_w = 128
image_h = 128
def train(model):
model.train(True)
loss_fn = nn.MSELoss()
optimizer = optim.SGD(model.parameters(), lr=0.001)
one_hot_indices = torch.LongTensor(batch_size) \
.random_(0, num_classes) \
.view(batch_size, 1)
for _ in range(num_batches):
# 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)
# run forward pass
optimizer.zero_grad()
outputs = model(inputs.to('cuda:0'))
# run backward pass
labels = labels.to(outputs.device)
loss_fn(outputs, labels).backward()
optimizer.step()
上面的 train(model) 方法使用 nn.MSELoss 作為損失函數(shù)，optim.SGD 作為優(yōu)化器。它模擬在128 × 128的圖像上進行訓(xùn)練，這些圖像被組織成3個 batches，每個批次包含120張圖像。然后，我們使用 timeit 運行 train(model) 方法10次，并繪制執(zhí)行時間的標(biāo)準(zhǔn)差，代碼如下：
import matplotlib.pyplot as plt
plt.switch_backend('Agg')
import numpy as np
import timeit
num_repeat = 10
stmt = "train(model)"
setup = "model = ModelParallelResNet50()"
mp_run_times = timeit.repeat(
stmt, setup, number=1, repeat=num_repeat, globals=globals())
mp_mean, mp_std = np.mean(mp_run_times), np.std(mp_run_times)
setup = "import torchvision.models as models;" + \
"model = models.resnet50(num_classes=num_classes).to('cuda:0')"
rn_run_times = timeit.repeat(
stmt, setup, number=1, repeat=num_repeat, globals=globals())
rn_mean, rn_std = np.mean(rn_run_times), np.std(rn_run_times)
def plot(means, stds, labels, fig_name):
fig, ax = plt.subplots()
ax.bar(np.arange(len(means)), means, yerr=stds,
align='center', alpha=0.5, ecolor='red', capsize=10, width=0.6)
ax.set_ylabel('ResNet50 Execution Time (Second)')
ax.set_xticks(np.arange(len(means)))
ax.set_xticklabels(labels)
ax.yaxis.grid(True)
plt.tight_layout()
plt.savefig(fig_name)
plt.close(fig)
plot([mp_mean, rn_mean],
[mp_std, rn_std],
['Model Parallel', 'Single GPU'],
'mp_vs_rn.png')
實驗結(jié)果：
結(jié)果顯示，Model Parallel 實現(xiàn)的執(zhí)行時間比現(xiàn)有的 Single GPU 實現(xiàn)長 4.02/3.75-1=7%。因此，我們可以得出結(jié)論，在 GPU 之間來回復(fù)制張量的開銷大約是7%。還有改進的余地，如何改進？
3 通過 Pipelining Inputs 加速模型并行
因為我們知道兩個 GPU 中的一個在整個執(zhí)行過程中是閑置的。一個選擇是將每個批次的 images 進一步劃分為一個個的 splits，這樣當(dāng)一個 split 到達第二個子網(wǎng)絡(luò)時，下面的 split 可以被送入第一個子網(wǎng)絡(luò)。通過這種方式，兩個連續(xù)的 splits 可以在兩個 GPU 上同時運行。
要理解這波操作，就得首先學(xué)習(xí)一個 torch.split 函數(shù)：
https://pytorch.org/docs/stable/generated/torch.split.html
torch.split(tensor,split_size_or_sections,dim=0)
Parameters
tensor (Tensor) – tensor to split. split_size_or_sections (int) or (list(int)) – size of a single chunk or list of sizes for each chunk dim (int) – dimension along which to split the tensor.
tensor:https://pytorch.org/docs/stable/tensors.html#torch.Tensor
list:https://docs.python.org/3/library/stdtypes.html#list
int:https://docs.python.org/3/library/functions.html#int
它的作用的官方描述是：Splits the tensor into chunks. Each chunk is a view of the original tensor.
如果 split_size_or_sections 是一個整數(shù)類型，那么張量將被分割成同等大小的塊。如果張量沿著給定的維度 dim 的大小不能被 split_size 整除，那么最后一個塊會更小。
如果 split_size_or_sections 是一個列表，那么張量將被分割成 len(split_size_or_sections) 個小塊，其大小與 split_size_or_sections 相同。
舉例：
>>> a = torch.arange(10).reshape(5,2)
>>> a
tensor([[0, 1],
[2, 3],
[4, 5],
[6, 7],
[8, 9]])
>>> torch.split(a, 2)
(tensor([[0, 1],
[2, 3]]),
tensor([[4, 5],
[6, 7]]),
tensor([[8, 9]]))
>>> torch.split(a, [1,4])
(tensor([[0, 1]]),
tensor([[2, 3],
[4, 5],
[6, 7],
[8, 9]]))
接下來，我們回到 PipelineParallelResNet50 模型，進一步將每個 batch 的120張圖片分成20張圖片的 split，這步操作可以通過 splits = iter(x.split(self.split_size, dim=0)) 來完成。
由于PyTorch是異步啟動CUDA操作的，因此該實現(xiàn)不需要催生多個線程來實現(xiàn)并發(fā)。代碼如下，簡單梳理一下代碼的含義：
在 forward() 函數(shù)里面做以下這些事情：
對輸入的 batch=120 的圖片分成相同大小為20的 splits：
splits = iter(x.split(self.split_size, dim=0))
從頭開始，每次取出一個 split：
s_next = next(splits)
把第1個 split 通過第1段模型：
s_prev = self.seq1(s_next).to('cuda:1')
for 循環(huán)可以看做每次循環(huán)做2件事：
A. 前半段模型的輸出傳到后半段并前向傳播：
s_prev = self.seq2(s_prev)
B. 下一個 split 輸入前半段模型：
s_prev = self.seq1(s_next).to('cuda:1')
class PipelineParallelResNet50(ModelParallelResNet50):
def __init__(self, split_size=20, *args, **kwargs):
super(PipelineParallelResNet50, self).__init__(*args, **kwargs)
self.split_size = split_size
def forward(self, x):
splits = iter(x.split(self.split_size, dim=0))
s_next = next(splits)
s_prev = self.seq1(s_next).to('cuda:1')
ret = []
for s_next in splits:
# A. s_prev runs on cuda:1
s_prev = self.seq2(s_prev)
ret.append(self.fc(s_prev.view(s_prev.size(0), -1)))
# B. s_next runs on cuda:0, which can run concurrently with A
s_prev = self.seq1(s_next).to('cuda:1')
s_prev = self.seq2(s_prev)
ret.append(self.fc(s_prev.view(s_prev.size(0), -1)))
return torch.cat(ret)
setup = "model = PipelineParallelResNet50()"
pp_run_times = timeit.repeat(
stmt, setup, number=1, repeat=num_repeat, globals=globals())
pp_mean, pp_std = np.mean(pp_run_times), np.std(pp_run_times)
plot([mp_mean, rn_mean, pp_mean],
[mp_std, rn_std, pp_std],
['Model Parallel', 'Single GPU', 'Pipelining Model Parallel'],
'mp_vs_rn_vs_pp.png')
整個過程可以用下圖2表示，其中數(shù)字1代表進入for循環(huán)前的第1步：s_prev = self.seq1(s_next).to('cuda:1')，2-6代表5次for循環(huán)，數(shù)字7代表最后一步：s_prev = self.seq2(s_prev)。
同一個 batch 的數(shù)據(jù)，有些提前進入后半段模型 Model Part 2，而不用等待全部數(shù)據(jù)走完前半段模型 Model Part 1之后再統(tǒng)一進入后半段模型 Model Part 2。這樣子節(jié)約了運行時間
實驗結(jié)果：
實驗結(jié)果表明，通過流水線輸入并行 ResNet50 模型，訓(xùn)練過程大約加快了 3.75/2.51-1=49%。這與理想的 100% 的速度仍然相差甚遠。由于我們在管道并行實現(xiàn)中引入了一個新的參數(shù) split_sizes，目前還不清楚這個新參數(shù)對整個訓(xùn)練時間有什么影響。直觀地說，使用小的 split_size 會導(dǎo)致許多微小的 CUDA 內(nèi)核啟動，而使用大的 split_size 會導(dǎo)致在第一次和最后一次分割時出現(xiàn)相對較長的空閑時間。兩者都不是最優(yōu)的。對于這個特定的實驗來說，可能會有一個最佳的 split_size 配置。讓我們通過使用幾個不同的 split_size 值進行實驗來找到它。代碼如下：
means = []
stds = []
split_sizes = [1, 3, 5, 8, 10, 12, 20, 40, 60]
for split_size in split_sizes:
setup = "model = PipelineParallelResNet50(split_size=%d)" % split_size
pp_run_times = timeit.repeat(
stmt, setup, number=1, repeat=num_repeat, globals=globals())
means.append(np.mean(pp_run_times))
stds.append(np.std(pp_run_times))
fig, ax = plt.subplots()
ax.plot(split_sizes, means)
ax.errorbar(split_sizes, means, yerr=stds, ecolor='red', fmt='ro')
ax.set_ylabel('ResNet50 Execution Time (Second)')
ax.set_xlabel('Pipeline Split Size')
ax.set_xticks(split_sizes)
ax.yaxis.grid(True)
plt.tight_layout()
plt.savefig("split_size_tradeoff.png")
plt.close(fig)
實驗結(jié)果：
如上圖所示，結(jié)果顯示，將 split_size 設(shè)置為12可以達到最快的訓(xùn)練速度，實現(xiàn)了 3.75/2.43-1=54% 的速度提升。
總結(jié)
本文介紹了模型并行技術(shù) Model Parallel，與 DataParallel 相反，Model Parallel 將一個單一的模型分割到不同的 GPU 上，而不是在每個 GPU 上復(fù)制整個模型。Pipeline Model Parallel 是一種進一步加速模型并行的策略。
如果覺得有用，就請分享到朋友圈吧！
公眾號后臺回復(fù)“transformer”獲取最新Transformer綜述論文下載～
#?極市平臺簽約作者#
科技猛獸
知乎：科技猛獸
清華大學(xué)自動化系19級碩士
研究領(lǐng)域：AI邊緣計算 (Efficient AI with Tiny Resource)：專注模型壓縮，搜索，量化，加速，加法網(wǎng)絡(luò)，以及它們與其他任務(wù)的結(jié)合，更好地服務(wù)于端側(cè)設(shè)備。
作品精選
