[PyTorch]Distributed Data Parallel - Practice

Reference: https://pytorch.org/tutorials/intermediate/model_parallel_tutorial.html

모델 병렬화는 분산 훈련 기술에서 아주 잘 활용된다. [Pytorch]Distributed Data Prarallel - DataParallel 에서 여러 대의 GPU를 사용한 훈련을 보여줬다; 이것의 특징은 같은 모델을 모든 GPU에 똑같이 복사한다는 점인데, 각각의 GPU는 input data의 서로 다른 부분들을 입력으로 받는다. 훈련에서 이런 방법이 가속화에 도움이 안되는건 아니지만, 모델이 매우 큰 경우 별로 좋은 방법은 아님.

따라서 이러한 문제를 해결하기 위해서 어떻게 model parallel을 통해서 하나의 모델을 서로 다른 GPU에 할당할 수 있는가에 대해 얘기하고자 한다.

The high-level idea of model parallel is to place different sub-networks of a model onto different devices, and implement the forward method accordingly to move intermediate outputs across devices.

Basic Usage

모델이 두개의 linear layer를 가지고 있다고 하자. 우리는 이 모델을 두개의 GPU에서 실행할 것이고, 간단하게 각각의 layer를 서로 다른 GPU에 할당할 것이다. 그리고 나서 input and intermediate outputs를 각 레이어 device에 맞게 줄 거임

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

backward() and torch.optim은 자동으로 맞춰줌

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

Apply Model Parallel to Existing Modules

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

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