DistributedDataParallel : what did I wrong? with simple code

Hyeonuk_Woo: <p>Hi, I’m new to DistributedDataParallel(DDP). I have lots of problem with this module…</p>
<p>I tried to run code like below.(from <a href="https://pytorch.org/tutorials/intermediate/ddp_tutorial.html" class="inline-onebox" rel="noopener nofollow ugc">Getting Started with Distributed Data Parallel — PyTorch Tutorials 2.2.0+cu121 documentation</a>)</p>
<pre><code class="lang-auto">import os
import tempfile
import torch
import torch.distributed as dist
import torch.multiprocessing as mp
import torch.nn as nn
import torch.optim as optim

from torch.nn.parallel import DistributedDataParallel as DDP


def setup(rank, world_size):
    os.environ['MASTER_ADDR'] = 'localhost'
    os.environ['MASTER_PORT'] = '12355'

    # initialize the process group
    dist.init_process_group("gloo", rank=rank, world_size=world_size)


def cleanup():
    dist.destroy_process_group()


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


def demo_basic(rank, world_size):
    print(f"Running basic DDP example on rank {rank}.")
    setup(rank, world_size)

    # create model and move it to GPU with id rank
    for x in range( 1, 10000):
        model = ToyModel().to(rank)
        ddp_model = DDP(model, device_ids=[rank])

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

        optimizer.zero_grad()
        outputs = ddp_model(torch.randn(20, 10))
        labels = torch.randn(20, 5).to(rank)
        loss_fn(outputs, labels).backward()
        optimizer.step()



def run_demo(demo_fn, world_size):
    mp.spawn(demo_fn,
             args=(world_size,),
             nprocs=world_size,
             join=True)


if __name__ == "__main__":
    n_gpus = torch.cuda.device_count()
    if n_gpus &lt; 4:
        print(f"Requires at least 4 GPUs to run, but got {n_gpus}.")
    else:
        run_demo(demo_basic, 4)

</code></pre>
<p>run with</p>
<pre><code class="lang-auto">python3 new.py
</code></pre>
<p>However, when I monitor ‘nvidia-smi’. There are 4 process in GPU 0. Why not one process for each GPU device?<br>
<div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/c/5/c5f313eaea86c3069a3278f0ab9dcf267502f6ff.png" data-download-href="https://discuss.pytorch.org/uploads/default/c5f313eaea86c3069a3278f0ab9dcf267502f6ff" title="Screenshot from 2021-09-06 17-13-06-1"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/c/5/c5f313eaea86c3069a3278f0ab9dcf267502f6ff_2_534x500.png" alt="Screenshot from 2021-09-06 17-13-06-1" data-base62-sha1="sf8QUJPvj14dopAjYJ6xSiz6jsH" width="534" height="500" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/c/5/c5f313eaea86c3069a3278f0ab9dcf267502f6ff_2_534x500.png, https://discuss.pytorch.org/uploads/default/original/3X/c/5/c5f313eaea86c3069a3278f0ab9dcf267502f6ff.png 1.5x, https://discuss.pytorch.org/uploads/default/original/3X/c/5/c5f313eaea86c3069a3278f0ab9dcf267502f6ff.png 2x" data-dominant-color="3D1A32"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">Screenshot from 2021-09-06 17-13-06-1</span><span class="informations">708×662 80.9 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div></p>

huahuanZ: <p>Add this line right after the <code>setup(rank, world_size)</code></p>
<pre><code class="lang-auto">setup(rank, world_size)
torch.cuda.set_device(rank)
for x in range(1, 10000):
...
</code></pre>

Hyeonuk_Woo: <p>Thanks! Everything become good!</p>