In 2020 what is the optimal way to train a model in Pytorch on more than one GPU on one computer?

user_5256048 0
Viewed 1003

What are the best practices for training one neural net on more than one GPU on one machine?

I'm a little confused by the different options from nn.DataParallel vs putting different layers on different GPUs with .to('cuda:0') and .to('cuda:1'). I see in the Pytorch docs the latter method the date was 2017. Is there a standard or does it depend on preference or the type of model?

Method 1

class ToyModel(nn.Module):
    def __init__(self):
        super(ToyModel, self).__init__()
        self.net1 = torch.nn.Linear(10, 10)
        self.relu = torch.nn.ReLU()
        self.net2 = torch.nn.Linear(10, 5)

    def forward(self, x):
        x = self.relu(self.net1(x))
        return self.net2(x)

model = ToyModel().to('cuda')
model = nn.DataParallel(model)

Method 2

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

I'm not sure there aren't more ways Pytorch provides to train on more than one GPU. Both of these methods seem to cause my system to freeze depending on what model I use them. In Jupyter the cell stays at a [*] and if I don't restart the kernel the screen freezes and I have to do a hard reset. A few tutorials on multi-gpu cause my system to hang and freeze like this.

pytorch neural-network parallel-processing gpu
2 Answers

If you cannot fit all the layers of your model on a single GPU, then you can use model parallel (that article describes model parallel on a single machine, with layer0.to('cuda:0') and layer1.to('cuda:1') like you mentioned).

If you can, then you can try distributed data parallel - each worker will hold its own copy of the entire model (all layers), and will work on a small portion of the data in each batch. DDP is recommended instead of DP, even if you only use a single machine.

Do you have some examples that can reproduce the issues you're having? Have you tried running your code with tiny inputs, and adding print statements to see whether progress is being made?

user_85826610

Since you said 2020 and optimal, the answer is definitely DDP, Distributed Data Parallel. DataParallel is much less efficient and is only good for quick-and-easy training in 2020. I wouldn't consider DataParallel to be close to optimal. The switch from DP to DDP was painful but was definitely worth it for me.

As of today, the official tutorial on DDP is still a pile of mess in my opinion. I personally went with this tutorial.

I will highlight some of the points that had confused me in the learning process. It is essential to get these points, not just to make the code runs, but to make sure it trains correctly. There can be subtle mistakes that corrupt training without runtime errors.

  1. In DDP, we spawn one process for each GPU. If your code starts from main.py and you are using 2 GPUs, imagine someone will call python main.py --local_rank=0 and python main.py --local_rank=1 to spawn two python processes.
  2. Each process receives its local rank and the global pool size. Use them to distinguish the processes.
  3. Each process initializes its Dataset objects independently. They collect their own samples from their own dataloaders. How do we ensure that they don't get the same piece of data point then? Pass a DistributedSampler to the dataloader.
  4. Each process performs forward, backward, and loss computation independently. They communicate the gradient in lockstep.
  5. Warp your model in nn.parallel.DistributedDataParallel. Then PyTorch will do (4) for you.
user_32374380
Related