# Multi node PyTorch Distributed Training Guide For People In A Hurry

Multi node training with PyTorch DDP, torch.distributed.launch, torchrun and mpirun - YouTube

[Multi node training with PyTorch DDP, torch.distributed.launch, torchrun and mpirun](https://www.youtube.com/watch?v=LuieZTc-hvU)

[PyTorch](https://pytorch.org/) is designed to be the framework that's both easy to use and delivers performance at scale. Indeed it has become [the most popular deep learning framework](https://paperswithcode.com/trends) by a mile among the research community. However, despite some [lengthy official tutorials](https://pytorch.org/tutorials/beginner/dist_overview.html) and a few helpful [community blogs](https://leimao.github.io/blog/PyTorch-Distributed-Training/), it is not always clear what exactly has to be done to make your PyTorch training to work **across multiple nodes**.

The goal of this tutorial is to give a summary of how to **write** and **launch** PyTorch distributed data parallel jobs. We assume readers have some experience of using PyTorch to train Neural Networks, know the basics of [data parallelization](https://pytorch.org/tutorials/intermediate/ddp_tutorial.html), and want to quickly connect the dots between executing the job on a single workstation and scaling it up **across multiple nodes**. With such a goal in mind, this tutorial will focus on:

- The basic idea of how PyTorch distributed data parallelism works under the hood.
- A few examples that showcase the boilerplate of PyTorch DDP training code.
- Have each example work with `torch.distributed.launch`, `torchrun` and `mpirun` API.

## Table of Content

1. [Distributed PyTorch Under the hood](https://lambda.ai/blog/multi-node-pytorch-distributed-training-guide#distributed-pytorch-underthehood)
2. [Write Multi-node PyTorch Distributed applications](https://lambda.ai/blog/multi-node-pytorch-distributed-training-guide#write-multi-node-pytorch-distributed-applications)

2.1. [Message Passing](https://lambda.ai/blog/multi-node-pytorch-distributed-training-guide#message-passing)

2.2. [ResNet Training](https://lambda.ai/blog/multi-node-pytorch-distributed-training-guide#resnet-training)
3. [Launch Multi-node PyTorch Distributed Applications](https://lambda.ai/blog/multi-node-pytorch-distributed-training-guide#launch-multi-node-pytorch-distributed-applications)

3.1. [torch.distributed.launch](https://lambda.ai/blog/multi-node-pytorch-distributed-training-guide#torchdistributedlaunch)

3.2. [torchrun](https://lambda.ai/blog/multi-node-pytorch-distributed-training-guide#torchrun)

3.3. [mpirun](https://lambda.ai/blog/multi-node-pytorch-distributed-training-guide#mpirun)
4. [Reference Performance on Lambda Cloud](https://lambda.ai/blog/multi-node-pytorch-distributed-training-guide#reference-performance-on-lambda-cloud)

## Distributed PyTorch Under the Hood

The basic idea behind distributed PyTorch starts simple – create a bunch of processes that replicate a single job execution for multiple times. This is very much like any ordinary MPI jobs, e.g. `mpirun -n 3 echo hello world` , which replicates the job of printing "hello world" three times:

```
ubunt@ubuntu-desktop:~$ mpirun -n 3 echo hello world
hello world
hello world
hello world
```

There are two main "tricky" parts that separate a PyTorch distributed (data parallel) training job from the above hello-world `mpirun` job. The PyTorch distributed training has to:

- Assign an accelerator (e.g. a GPU) to each process to maximize the **computation** efficiency of the forward and backward passes for each training step.
- Facilitate the **communication** between different processes. This happens at the end of each training step, so that the gradient can be aggregated, and the model can be updated in a synchronized fashion across all processes.

To reach these two goals, PyTorch creates a group of processes that are "device-aware" (`torch.distributed.init_process_group`), hooks them up with fast communication backends (e.g. `nccl`), and prepares your data pipeline and model implementation to work in this multi-process context (typically via the `torch.utils.data.distributed.DistributedSampler` and the `torch.nn.parallel.DistributedDataParallel` APIs)

As it will become much clearer later, a few environment variables are used to glue everything together: `WORLD_SIZE`, `WORLD_RANK` and `LOCAL_RANK`. Think of them as fancy names for "total number of GPUs in your cluster", "the ID of a GPU at the cluster level", and "the ID of a GPU at a node level".

## Write Multi-node PyTorch Distributed Applications

Next we show a couple of examples of writing distributed PyTorch applications across multiple nodes. We will start with a simple message passing example, and explain how PyTorch DDP leverages environment variables to create processes across multiple nodes. We will then discuss how to generalize the practice to a more complicated use case (ResNet training).

This blog uses two `2xA6000` Lambda Cloud instances for demonstration. But the same procedure generalizes to on-prem clusters with more than two inter-connected baremetal servers, for example, a Lambda Echelon cluster.

The two `2xA6000` nodes, namely `104.171.200.62` and `104.171.200.182`, allows us to have `4` workers in total (`2` on each node). We will use the term "process" and "worker" interchangeably since, as far as this tutorial goes, they mean the same thing.

### Message Passing

We will start with a very simple example that passes tensor between the workers. The example is inspired by an official PyTorch [tutorial](https://pytorch.org/tutorials/intermediate/dist_tuto.html).

```
import os
import argparse

import torch
import torch.distributed as dist

# Environment variables set by torch.distributed.launch
LOCAL_RANK = int(os.environ['LOCAL_RANK'])
WORLD_SIZE = int(os.environ['WORLD_SIZE'])
WORLD_RANK = int(os.environ['RANK'])

def run(backend):
    tensor = torch.zeros(1)

# Need to put tensor on a GPU device for nccl backend
    if backend == 'nccl':
        device = torch.device("cuda:{}.format(LOCAL_RANK)")
        tensor = tensor.to(device)

if WORLD_RANK == 0:
        for rank_recv in range(1, WORLD_SIZE):
            dist.send(tensor=tensor, dst=rank_recv)
            print('worker_{} sent data to Rank {}
'.format(0, rank_recv))
    else:
        dist.recv(tensor=tensor, src=0)
        print('worker_{} has received data from rank {}
'.format(WORLD_RANK, 0))

def init_processes(backend):
    dist.init_process_group(backend, rank=WORLD_RANK, world_size=WORLD_SIZE)
    run(backend)

if __name__ == "__main__":

parser = argparse.ArgumentParser()
    parser.add_argument("--local_rank", type=int, help="Local rank. Necessary for using the torch.distributed.launch utility.")
    parser.add_argument("--backend", type=str, default="nccl", choices=['nccl', 'gloo'])
    args = parser.parse_args()

init_processes(backend=args.backend)
```

The code launches `4` workers, where `worker_0` create a tensor on `GPU 0`  of node `104.171.200.62`, and send the tensor to all other three workers.

```
# output from node 104.171.200.62
worker_0 sent data to Rank 1
worker_0 sent data to Rank 2
worker_0 sent data to Rank 3
worker_1 has received data from rank 0

# output from node 104.171.200.182
worker_2 has received data from rank 0
worker_3 has received data from rank 0
```

### ResNet Training

Now we know the basics of writing a multi-node distributed PyTorch application. Next we will analyze a very popular [ResNet training code written by Lei Mao](https://leimao.github.io/blog/PyTorch-Distributed-Training/). We will not repost his entire code here, instead we will compare the common practices used in his code and the above message pass example, and highlight the extra bells and whistles for training a deep neural networks.

Common practices for distributed PyTorch applications:

- `torch.distributed.init_process_group(backend="nccl")`: The [ResNet script](https://leimao.github.io/blog/PyTorch-Distributed-Training/) uses the same function to create the workers. However, `rank` and `world_size` are not specified here. Indeed, they are optional as specified by the [PyTorch documentation](https://pytorch.org/docs/stable/distributed.html#torch.distributed.init_process_group), with default values being the environment variables `WORLD_SIZE` and `RANK`. So no real difference here.
- `device = torch.device("cuda:{}.format(local_rank)")`: Again, local rank is used for specifying GPU device for a worker. This is the same between the [ResNet script](https://leimao.github.io/blog/PyTorch-Distributed-Training/) and the above message passing script.

Compared to the simple message passing example, there are some extra steps for getting a training job to work:

- `ddp_model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[local_rank], output_device=local_rank)`: The ResNet script uses this common PyTorch practice to "wrap" up the ResNet **model** so it can be used in the DDP context.
- `train_sampler = DistributedSampler(dataset=train_set)`: Another common practice to "wrap" up the **dataset** so it can be used in the DDP context.

## Launch Multi-node PyTorch Distributed Applications

So far we have talked about how to write a PyTorch DDP code so it works across multiple nodes. What we haven't touched upon is how to actually **launch** a multi-node distributed job. Now is the time to do so and finally unfold the mystery of how environment variables like `WORLD_SIZE`, `WORLD_RANK` and `LOCAL_RANK` are set.

There are different ways to "launch" applications in a distributed fashion across multiple nodes, implemented by different vendors, such as HPC veterans (Open MPI), Deep Learning framework leaders (PyTorch), or open source A.I. communities (Horovod).

### torch.distributed.launch

In order to `torch.distributed.launch` a two-node ResNet training job, you need to ssh into the two nodes and run the following commands:

```bash
# On 104.171.200.62 (the master node)
python3 -m torch.distributed.launch \
--nproc_per_node=2 --nnodes=2 --node_rank=0 \
--master_addr=104.171.200.62 --master_port=1234 \
main.py \
--backend=nccl --use_syn --batch_size=8192 --arch=resnet152

# On 104.171.200.182 (the worker node)
python3 -m torch.distributed.launch \
--nproc_per_node=2 --nnodes=2 --node_rank=1 \
--master_addr=104.171.200.62 --master_port=1234 \
main.py \
--backend=nccl --use_syn --batch_size=8192 --arch=resnet152
```

- `nproc_per_node` defines the number of workers on each node. It should equal to the number of GPUs on each node. In this case, `2`.
- `nnodes` defines the number of nodes. In this case, `2`.
- `master_addr` and `master_port` are the IP address and port for the master node. You have to set up these environment variables otherwise PyTorch DDP will complain.
- `node_rank` defines the rank of a node. This has to be set **differently** in the two commands — use `0` for the master node, and `1` for the worker node. Training will freeze if master node is not rank `0`.

### torchrun

Created by the PyTorch team, `torchrun` works similarly to `torch.distributed.launch` but with some extra [functionalities](https://pytorch.org/docs/stable/elastic/run.html#launcher-api) that gracefully handle failed workers and elasticity. In fact, `torchrun` can work with the exact same script as `torch.distributed.launch` does:

```bash
# On 104.171.200.62 (the master node)
torchrun \
--nproc_per_node=2 --nnodes=2 --node_rank=0 \
--master_addr=104.171.200.62 --master_port=1234 \
main.py \
--backend=nccl --use_syn --batch_size=8192 --arch=resnet152

# On 104.171.200.182 (the worker node)
torchrun \
--nproc_per_node=2 --nnodes=2 --node_rank=1 \
--master_addr=104.171.200.62 --master_port=1234 \
main.py \
--backend=nccl --use_syn --batch_size=8192 --arch=resnet152
```

### mpirun

Although the above `torch.distributed.launch` method works "out of the box" as the native PyTorch API, one has to modify and run the launch command on each of the nodes. This is non-scalable and error-prone. Alternatively, we can launch the same job by executing an **`mpirun` command only once from the master node.**

MPI, or the [Message Passing Interface](https://en.wikipedia.org/wiki/Message_Passing_Interface), is a standardized means of exchanging messages between processes belongs to a parallel program distributed across multiple computers. In the context of our distributed PyTorch training, the "message" contains the gradient computed on each GPU, and the "exchange" process makes sure these gradients are gathered from all the GPUs, so that the weights of the networks can be updated in a synchronized way.  
There are two main flavors of MPI implementation, [OpenMPI](https://www.open-mpi.org/) and [MPICH](https://www.mpich.org/). Either of them will work and we will use the OpenMPI implementation in this blog.

```bash
# On 104.171.200.62 (the master node)
mpirun -np 4 \
-H 104.171.200.62:2,104.171.200.182:2 \
-x MASTER_ADDR=104.171.200.62 \
-x MASTER_PORT=1234 \
-x PATH \
-bind-to none -map-by slot \
-mca pml ob1 -mca btl ^openib \
python3 main.py --backend=nccl --use_syn --batch_size=8192 --arch=resnet152
```

### Reference Performance on Lambda Cloud

It is important to consider the scaling efficiency when running a distributed training job across multiple nodes. Ideally the training throughput (measured by number of images processed per second) should scale linearly with the number of GPUs. In practice this is often not the case, due to the inter-GPU communication bottleneck. This is especially the case when GPUs are distributed across different nodes.

The On-Demand Lambda Cloud instance used in this blog supports up to `100 Gbps` inter-node bandwidth. We were actually able to achieve satisfactory scaling efficiency for the above `ResNet152` training test for `CIFAR10` data:

## Summary

This blog gives a quick overview of how multi-node distributed PyTorch works under the hood, and how to write and launch multi-node distributed PyTorch applications.

The code for reproducing the examples can be found in this **[repo](https://github.com/LambdaLabsML/examples/tree/main/pytorch/distributed)**.