PyTorch 單機模型并行最佳實踐

原文： PyTorch 單機模型并行最佳實踐

注意

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

作者：申力

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

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

Note

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

基本用法

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

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'))

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

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

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))

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

讓我們進行實驗以更定量地了解執(zhí)行時間。 在此實驗中，我們通過運行隨機輸入和標簽來訓(xùn)練ModelParallelResNet50和現(xiàn)有的torchvision.models.reset50()。 訓(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 X 128圖像的訓(xùn)練，這些圖像分為 3 批，每批包含 120 張圖像。 然后，我們使用timeit來運行train(model)方法 10 次，并繪制帶有標準偏差的執(zhí)行時間。

import matplotlib.pyplot as plt
plt.switch_backend('Agg')
import numpy as np
import timeit


num_repeat = 10


stmt = "train(model)"


setup = "model = ModelParallelResNet50()"
## globals arg is only available in Python 3\. In Python 2, use the following
## import __builtin__
## __builtin__.__dict__.update(locals())
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é)果表明，模型并行實現(xiàn)的執(zhí)行時間比現(xiàn)有的單 GPU 實現(xiàn)長4.02/3.75-1=7%。 因此，我們可以得出結(jié)論，在 GPU 之間來回復(fù)制張量大約有 7％的開銷。 有改進的余地，因為我們知道兩個 GPU 之一在整個執(zhí)行過程中處于空閑狀態(tài)。 一種選擇是將每個批次進一步劃分為拆分管道，這樣，當一個拆分到達第二個子網(wǎng)時，可以將下一個拆分饋入第一個子網(wǎng)。 這樣，兩個連續(xù)的拆分可以在兩個 GPU 上同時運行。

通過流水線輸入加速

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

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')

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

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

注意：

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

腳本的總運行時間：(5 分鐘 53.174 秒）

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