Simple Distributed Training Example

Joseph_Konan: <p>I apologize, as I am having trouble following the official PyTorch tutorials. I have one system with two GPUs and I would like to use both for training.</p>
<p>The following example is a modification of the following:<br>
<a href="https://pytorch.org/tutorials/intermediate/ddp_tutorial.html" class="onebox" target="_blank" rel="noopener nofollow ugc">https://pytorch.org/tutorials/intermediate/ddp_tutorial.html</a></p>
<pre><code class="lang-auto">import os
import torch
import torch.distributed as dist
import torch.nn as nn
import torch.optim as optim

from torch.nn.parallel import DistributedDataParallel as DDP

os.environ['MASTER_ADDR'] = 'localhost'
os.environ['MASTER_PORT'] = '12355'

torch.distributed.init_process_group("nccl", rank=0, world_size=1)
</code></pre>
<pre><code class="lang-auto">class ToyModel(nn.Module):
    def __init__(self):
        super(ToyModel, self).__init__()
        self.net1 = nn.Linear(10, 10)
        self.relu = nn.ReLU()
        self.net2 = nn.Linear(10, 5)

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


rank = dist.get_rank()
print(f"Start running basic DDP example on rank {rank}.")

# create model and move it to GPU with id rank
device_id = rank % torch.cuda.device_count()
model = ToyModel().to(device_id)
ddp_model = DDP(model, device_ids=[device_id])

loss_fn = nn.MSELoss()
optimizer = optim.SGD(ddp_model.parameters(), lr=0.001)

for i in range(10000):
    optimizer.zero_grad()
    outputs = ddp_model(torch.randn(20, 10))
    labels = torch.randn(20, 5).to(device_id)
    loss_fn(outputs, labels).backward()
    optimizer.step()
</code></pre>
<p>However, I observe that only one GPU is being used. When I try increasing <code>world_size</code> to 2, the system hangs on <code>init_process_group</code>. I see some people recommending the use of <code>init_method='env://'</code>, but I am thrown an error unless I specify <code>world_size</code> and <code>rank</code> using <code>os.environ</code>.</p>
<p>What is the correct way to use <code>init_process_group</code>?</p>

sudomaze: <p>Hi Joseph,</p>
<p>You can check this DDP example for how to implement it.</p>
<aside class="onebox githubblob" data-onebox-src="https://github.com/sudomaze/ttorch/blob/main/examples/ddp/run.py">
  <header class="source">

      <a href="https://github.com/sudomaze/ttorch/blob/main/examples/ddp/run.py" target="_blank" rel="noopener nofollow ugc">github.com</a>
  </header>

  <article class="onebox-body">
    <h4><a href="https://github.com/sudomaze/ttorch/blob/main/examples/ddp/run.py" target="_blank" rel="noopener nofollow ugc">sudomaze/ttorch/blob/main/examples/ddp/run.py</a></h4>


      <pre><code class="lang-py"># credits:
# how to use DDP module with DDP sampler: https://gist.github.com/sgraaf/5b0caa3a320f28c27c12b5efeb35aa4c
# how to setup a basic DDP example from scratch: https://pytorch.org/tutorials/intermediate/dist_tuto.html
import os
import torch
import torch.distributed as dist
import torch.multiprocessing as mp
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torchvision import datasets, transforms
import random

from torch.nn.parallel import DistributedDataParallel as DDP

from torch.utils.data import DataLoader, Dataset
from torch.utils.data.distributed import DistributedSampler
import math

</code></pre>



  This file has been truncated. <a href="https://github.com/sudomaze/ttorch/blob/main/examples/ddp/run.py" target="_blank" rel="noopener nofollow ugc">show original</a>

  </article>

  <div class="onebox-metadata">
    
    
  </div>

  <div style="clear: both"></div>
</aside>

<p>I am planning to complete the DDP example soon and add it to <a href="https://github.com/pytorch/examples/" rel="noopener nofollow ugc">PyTorch/examples</a> repo sometime soon. Please let me know if you think there is something that needs to be improved.</p>

Yanli_Zhao: <p>That is correct to set world_size and rank using os.environ if you are using init_method='env://</p>

Yanli_Zhao: <p>see this tutorial for init_process_group() <a href="https://pytorch.org/tutorials/intermediate/dist_tuto.html" class="inline-onebox" rel="noopener nofollow ugc">Writing Distributed Applications with PyTorch — PyTorch Tutorials 1.12.0+cu102 documentation</a></p>