PyTorch 單機(jī)模型并行最佳實(shí)踐

2020-09-16 13:44 更新

原文： PyTorch 單機(jī)模型并行最佳實(shí)踐

注意

單擊此處的下載完整的示例代碼

作者：申力

模型并行在分布式訓(xùn)練技術(shù)中被廣泛使用。 先前的帖子已經(jīng)解釋了如何使用 DataParallel 在多個(gè) GPU 上訓(xùn)練神經(jīng)網(wǎng)絡(luò)； 此功能將相同的模型復(fù)制到所有 GPU，其中每個(gè) GPU 消耗輸入數(shù)據(jù)的不同分區(qū)。 盡管它可以極大地加快訓(xùn)練過程，但不適用于某些模型太大而無法放入單個(gè) GPU 的用例。 這篇文章展示了如何通過使用模型并行解決該問題，與DataParallel相比，該模型將單個(gè)模型拆分到不同的 GPU 上，而不是在每個(gè) GPU 上復(fù)制整個(gè)模型(具體來說， 假設(shè)模型m包含 10 層：使用DataParallel時(shí)，每個(gè) GPU 都具有這 10 層中每個(gè)層的副本，而當(dāng)在兩個(gè) GPU 上并行使用模型時(shí)，每個(gè) GPU 可以承載 5 層）。

模型并行化的高級思想是將模型的不同子網(wǎng)放置在不同的設(shè)備上，并相應(yīng)地實(shí)現(xiàn)forward方法以在設(shè)備之間移動(dòng)中間輸出。 由于模型的一部分僅在任何單個(gè)設(shè)備上運(yùn)行，因此一組設(shè)備可以共同為更大的模型服務(wù)。 在本文中，我們不會嘗試構(gòu)建龐大的模型并將其壓縮到有限數(shù)量的 GPU 中。 相反，本文著重展示并行模型的概念。 讀者可以將這些想法應(yīng)用到實(shí)際應(yīng)用中。

Note

對于模型跨越多個(gè)服務(wù)器的分布式模型并行訓(xùn)練，請參見分布式 RPC 框架入門，以獲取示例和詳細(xì)信息。

基本用法

讓我們從包含兩個(gè)線性層的玩具模型開始。 要在兩個(gè) GPU 上運(yùn)行此模型，只需將每個(gè)線性層放在不同的 GPU 上，然后移動(dòng)輸入和中間輸出以匹配層設(shè)備。

  1. import torch
  2. import torch.nn as nn
  3. import torch.optim as optim
  6. class ToyModel(nn.Module):
  7.     def __init__(self):
  8.         super(ToyModel, self).__init__()
  9.         self.net1 = torch.nn.Linear(10, 10).to('cuda:0')
  10.         self.relu = torch.nn.ReLU()
  11.         self.net2 = torch.nn.Linear(10, 5).to('cuda:1')
  14.     def forward(self, x):
  15.         x = self.relu(self.net1(x.to('cuda:0')))
  16.         return self.net2(x.to('cuda:1'))

請注意，除了五個(gè)to(device)調(diào)用將線性層和張量放置在適當(dāng)?shù)脑O(shè)備上之外，上述ToyModel看起來非常類似于在單個(gè) GPU 上實(shí)現(xiàn)它的方式。 那是模型中唯一需要更改的地方。 backward()和torch.optim將自動(dòng)處理漸變，就像模型在一個(gè) GPU 上一樣。 調(diào)用損失函數(shù)時(shí)，只需確保標(biāo)簽與輸出位于同一設(shè)備上。

  1. model = ToyModel()
  2. loss_fn = nn.MSELoss()
  3. optimizer = optim.SGD(model.parameters(), lr=0.001)
  6. optimizer.zero_grad()
  7. outputs = model(torch.randn(20, 10))
  8. labels = torch.randn(20, 5).to('cuda:1')
  9. loss_fn(outputs, labels).backward()
  10. optimizer.step()

將模型并行應(yīng)用于現(xiàn)有模塊

只需進(jìn)行幾行更改，就可以在多個(gè) GPU 上運(yùn)行現(xiàn)有的單 GPU 模塊。 以下代碼顯示了如何將torchvision.models.reset50()分解為兩個(gè) GPU。 這個(gè)想法是繼承現(xiàn)有的ResNet模塊，并在構(gòu)建過程中將層拆分為兩個(gè) GPU。 然后，通過相應(yīng)地移動(dòng)中間輸出，覆蓋forward方法來縫合兩個(gè)子網(wǎng)。

  1. from torchvision.models.resnet import ResNet, Bottleneck
  4. num_classes = 1000
  7. class ModelParallelResNet50(ResNet):
  8.     def __init__(self, *args, **kwargs):
  9.         super(ModelParallelResNet50, self).__init__(
  10.             Bottleneck, [3, 4, 6, 3], num_classes=num_classes, *args, **kwargs)
  13.         self.seq1 = nn.Sequential(
  14.             self.conv1,
  15.             self.bn1,
  16.             self.relu,
  17.             self.maxpool,
  20.             self.layer1,
  21.             self.layer2
  22.         ).to('cuda:0')
  25.         self.seq2 = nn.Sequential(
  26.             self.layer3,
  27.             self.layer4,
  28.             self.avgpool,
  29.         ).to('cuda:1')
  32.         self.fc.to('cuda:1')
  35.     def forward(self, x):
  36.         x = self.seq2(self.seq1(x).to('cuda:1'))
  37.         return self.fc(x.view(x.size(0), -1))

對于模型太大而無法放入單個(gè) GPU 的情況，上述實(shí)現(xiàn)解決了該問題。 但是，您可能已經(jīng)注意到，如果您的模型合適，它將比在單個(gè) GPU 上運(yùn)行它要慢。 這是因?yàn)樵谌魏螘r(shí)間點(diǎn)，兩個(gè) GPU 中只有一個(gè)在工作，而另一個(gè)在那兒什么也沒做。 由于中間輸出需要在layer2和layer3之間從cuda:0復(fù)制到cuda:1，因此性能進(jìn)一步惡化。

讓我們進(jìn)行實(shí)驗(yàn)以更定量地了解執(zhí)行時(shí)間。 在此實(shí)驗(yàn)中，我們通過運(yùn)行隨機(jī)輸入和標(biāo)簽來訓(xùn)練ModelParallelResNet50和現(xiàn)有的torchvision.models.reset50()。 訓(xùn)練后，模型將不會產(chǎn)生任何有用的預(yù)測，但是我們可以對執(zhí)行時(shí)間有一個(gè)合理的了解。

  1. import torchvision.models as models
  4. num_batches = 3
  5. batch_size = 120
  6. image_w = 128
  7. image_h = 128
  10. def train(model):
  11.     model.train(True)
  12.     loss_fn = nn.MSELoss()
  13.     optimizer = optim.SGD(model.parameters(), lr=0.001)
  16.     one_hot_indices = torch.LongTensor(batch_size) \
  17.                            .random_(0, num_classes) \
  18.                            .view(batch_size, 1)
  21.     for _ in range(num_batches):
  22.         # generate random inputs and labels
  23.         inputs = torch.randn(batch_size, 3, image_w, image_h)
  24.         labels = torch.zeros(batch_size, num_classes) \
  25.                       .scatter_(1, one_hot_indices, 1)
  28.         # run forward pass
  29.         optimizer.zero_grad()
  30.         outputs = model(inputs.to('cuda:0'))
  33.         # run backward pass
  34.         labels = labels.to(outputs.device)
  35.         loss_fn(outputs, labels).backward()
  36.         optimizer.step()

上面的train(model)方法使用nn.MSELoss作為損失函數(shù)，并使用optim.SGD作為優(yōu)化器。 它模擬了對128 X 128圖像的訓(xùn)練，這些圖像分為 3 批，每批包含 120 張圖像。 然后，我們使用timeit來運(yùn)行train(model)方法 10 次，并繪制帶有標(biāo)準(zhǔn)偏差的執(zhí)行時(shí)間。

  1. import matplotlib.pyplot as plt
  2. plt.switch_backend('Agg')
  3. import numpy as np
  4. import timeit
  7. num_repeat = 10
  10. stmt = "train(model)"
  13. setup = "model = ModelParallelResNet50()"
  14. ## globals arg is only available in Python 3\. In Python 2, use the following
  15. ## import __builtin__
  16. ## __builtin__.__dict__.update(locals())
  17. mp_run_times = timeit.repeat(
  18.     stmt, setup, number=1, repeat=num_repeat, globals=globals())
  19. mp_mean, mp_std = np.mean(mp_run_times), np.std(mp_run_times)
  22. setup = "import torchvision.models as models;" + \
  23.         "model = models.resnet50(num_classes=num_classes).to('cuda:0')"
  24. rn_run_times = timeit.repeat(
  25.     stmt, setup, number=1, repeat=num_repeat, globals=globals())
  26. rn_mean, rn_std = np.mean(rn_run_times), np.std(rn_run_times)
  29. def plot(means, stds, labels, fig_name):
  30.     fig, ax = plt.subplots()
  31.     ax.bar(np.arange(len(means)), means, yerr=stds,
  32.            align='center', alpha=0.5, ecolor='red', capsize=10, width=0.6)
  33.     ax.set_ylabel('ResNet50 Execution Time (Second)')
  34.     ax.set_xticks(np.arange(len(means)))
  35.     ax.set_xticklabels(labels)
  36.     ax.yaxis.grid(True)
  37.     plt.tight_layout()
  38.     plt.savefig(fig_name)
  39.     plt.close(fig)
  42. plot([mp_mean, rn_mean],
  43.      [mp_std, rn_std],
  44.      ['Model Parallel', 'Single GPU'],
  45.      'mp_vs_rn.png')

img

結(jié)果表明，模型并行實(shí)現(xiàn)的執(zhí)行時(shí)間比現(xiàn)有的單 GPU 實(shí)現(xiàn)長4.02/3.75-1=7%。 因此，我們可以得出結(jié)論，在 GPU 之間來回復(fù)制張量大約有 7％的開銷。 有改進(jìn)的余地，因?yàn)槲覀冎纼蓚€(gè) GPU 之一在整個(gè)執(zhí)行過程中處于空閑狀態(tài)。 一種選擇是將每個(gè)批次進(jìn)一步劃分為拆分管道，這樣，當(dāng)一個(gè)拆分到達(dá)第二個(gè)子網(wǎng)時(shí)，可以將下一個(gè)拆分饋入第一個(gè)子網(wǎng)。 這樣，兩個(gè)連續(xù)的拆分可以在兩個(gè) GPU 上同時(shí)運(yùn)行。

通過流水線輸入加速

在以下實(shí)驗(yàn)中，我們將每批次120張圖像，進(jìn)一步劃分為 20 張圖像的均分。 當(dāng) PyTorch 異步啟動(dòng) CUDA 操作時(shí)，該實(shí)現(xiàn)無需生成多個(gè)線程即可實(shí)現(xiàn)并發(fā)。

  1. class PipelineParallelResNet50(ModelParallelResNet50):
  2.     def __init__(self, split_size=20, *args, **kwargs):
  3.         super(PipelineParallelResNet50, self).__init__(*args, **kwargs)
  4.         self.split_size = split_size
  7.     def forward(self, x):
  8.         splits = iter(x.split(self.split_size, dim=0))
  9.         s_next = next(splits)
  10.         s_prev = self.seq1(s_next).to('cuda:1')
  11.         ret = []
  14.         for s_next in splits:
  15.             # A. s_prev runs on cuda:1
  16.             s_prev = self.seq2(s_prev)
  17.             ret.append(self.fc(s_prev.view(s_prev.size(0), -1)))
  20.             # B. s_next runs on cuda:0, which can run concurrently with A
  21.             s_prev = self.seq1(s_next).to('cuda:1')
  24.         s_prev = self.seq2(s_prev)
  25.         ret.append(self.fc(s_prev.view(s_prev.size(0), -1)))
  28.         return torch.cat(ret)
  31. setup = "model = PipelineParallelResNet50()"
  32. pp_run_times = timeit.repeat(
  33.     stmt, setup, number=1, repeat=num_repeat, globals=globals())
  34. pp_mean, pp_std = np.mean(pp_run_times), np.std(pp_run_times)
  37. plot([mp_mean, rn_mean, pp_mean],
  38.      [mp_std, rn_std, pp_std],
  39.      ['Model Parallel', 'Single GPU', 'Pipelining Model Parallel'],
  40.      'mp_vs_rn_vs_pp.png')

請注意，設(shè)備到設(shè)備的張量復(fù)制操作在源設(shè)備和目標(biāo)設(shè)備上的當(dāng)前流上同步。 如果創(chuàng)建多個(gè)流，則必須確保復(fù)制操作正確同步。 在完成復(fù)制操作之前寫入源張量或讀取/寫入目標(biāo)張量可能導(dǎo)致不確定的行為。 上面的實(shí)現(xiàn)僅在源設(shè)備和目標(biāo)設(shè)備上都使用默認(rèn)流，因此沒有必要強(qiáng)制執(zhí)行其他同步。

img

實(shí)驗(yàn)結(jié)果表明，對并行 ResNet50 進(jìn)行建模的流水線輸入可將訓(xùn)練過程大約加快3.75/2.51-1=49%的速度。 距理想的 100％加速仍然相去甚遠(yuǎn)。 由于我們在管道并行實(shí)現(xiàn)中引入了新參數(shù)split_sizes，因此尚不清楚新參數(shù)如何影響整體訓(xùn)練時(shí)間。 直觀地講，使用較小的split_size會導(dǎo)致許多小的 CUDA 內(nèi)核啟動(dòng)，而使用較大的split_size會導(dǎo)致在第一次和最后一次拆分期間出現(xiàn)較長的空閑時(shí)間。 都不是最佳選擇。 對于此特定實(shí)驗(yàn)，可能會有最佳的split_size配置。 讓我們嘗試通過使用幾個(gè)不同的split_size值進(jìn)行實(shí)驗(yàn)來找到它。

  1. means = []
  2. stds = []
  3. split_sizes = [1, 3, 5, 8, 10, 12, 20, 40, 60]
  6. for split_size in split_sizes:
  7.     setup = "model = PipelineParallelResNet50(split_size=%d)" % split_size
  8.     pp_run_times = timeit.repeat(
  9.         stmt, setup, number=1, repeat=num_repeat, globals=globals())
  10.     means.append(np.mean(pp_run_times))
  11.     stds.append(np.std(pp_run_times))
  14. fig, ax = plt.subplots()
  15. ax.plot(split_sizes, means)
  16. ax.errorbar(split_sizes, means, yerr=stds, ecolor='red', fmt='ro')
  17. ax.set_ylabel('ResNet50 Execution Time (Second)')
  18. ax.set_xlabel('Pipeline Split Size')
  19. ax.set_xticks(split_sizes)
  20. ax.yaxis.grid(True)
  21. plt.tight_layout()
  22. plt.savefig("split_size_tradeoff.png")
  23. plt.close(fig)

img

結(jié)果表明，將split_size設(shè)置為 12 可獲得最快的訓(xùn)練速度，從而導(dǎo)致3.75/2.43-1=54%加速。 仍有機(jī)會進(jìn)一步加快訓(xùn)練過程。 例如，對cuda:0的所有操作都放在其默認(rèn)流上。 這意味著下一個(gè)拆分的計(jì)算不能與上一個(gè)拆分的復(fù)制操作重疊。 但是，由于上一個(gè)和下一個(gè)拆分是不同的張量，因此將一個(gè)計(jì)算與另一個(gè)副本重疊是沒有問題的。 實(shí)現(xiàn)需要在兩個(gè) GPU 上使用多個(gè)流，并且不同的子網(wǎng)結(jié)構(gòu)需要不同的流管理策略。 由于沒有通用的多流解決方案適用于所有模型并行用例，因此在本教程中將不再討論。

注意：

這篇文章顯示了幾個(gè)性能指標(biāo)。 當(dāng)您在自己的計(jì)算機(jī)上運(yùn)行相同的代碼時(shí)，您可能會看到不同的數(shù)字，因?yàn)榻Y(jié)果取決于底層的硬件和軟件。 為了使您的環(huán)境獲得最佳性能，一種正確的方法是首先生成曲線以找出最佳分割尺寸，然后將該分割尺寸用于管道輸入。

腳本的總運(yùn)行時(shí)間：(5 分鐘 53.174 秒）

Download Python source code: model_parallel_tutorial.py Download Jupyter notebook: model_parallel_tutorial.ipynb

以上內(nèi)容是否對您有幫助：
在線筆記
App下載
App下載

掃描二維碼

下載編程獅App

公眾號
微信公眾號

編程獅公眾號