pytorch_distributed_data_parallel_ddp_multi_gpu_training_quickstart.py

1import torch
2import torch.nn as nn
3import torch.optim as optim
4import torch.distributed as dist
5import torch.multiprocessing as mp
6from torch.nn.parallel import DistributedDataParallel as DDP
7
8def setup(rank, world_size):
9    # initialize the process group
10    dist.init_process_group("gloo", rank=rank, world_size=world_size)
11
12def cleanup():
13    dist.destroy_process_group()
14
15class ToyModel(nn.Module):
16    def __init__(self):
17        super(ToyModel, self).__init__()
18        self.net1 = nn.Linear(10, 10)
19        self.relu = nn.ReLU()
20        self.net2 = nn.Linear(10, 5)
21
22    def forward(self, x):
23        return self.net2(self.relu(self.net1(x)))
24
25def demo_basic(rank, world_size):
26    print(f"Running basic DDP example on rank {rank}.")
27    setup(rank, world_size)
28
29    # create model and move it to GPU with id rank
30    model = ToyModel().to(rank)
31    ddp_model = DDP(model, device_ids=[rank])
32
33    loss_fn = nn.MSELoss()
34    optimizer = optim.SGD(ddp_model.parameters(), lr=0.001)
35
36    optimizer.zero_grad()
37    outputs = ddp_model(torch.randn(20, 10).to(rank))
38    labels = torch.randn(20, 5).to(rank)
39    loss_fn(outputs, labels).backward()
40    optimizer.step()
41
42    cleanup()
43
44def run_demo(demo_fn, world_size):
45    mp.spawn(demo_fn,
46             args=(world_size,),
47             nprocs=world_size,
48             join=True)
49
50if __name__ == "__main__":
51    n_gpus = torch.cuda.device_count()
52    if n_gpus < 2:
53      print(f"Requires at least 2 GPUs to run, but only found {n_gpus}")
54    else:
55      run_demo(demo_basic, n_gpus)