PyTorch DDP

Environment Setup Before We Begin

Before we start, you need to install PyTorch and CUDA. You can do this by following the instructions on the official website. You can install PyTorch as usual, but I have noticed that some people try to install CUDA manually. Personally, I do not recommend this. There are many settings to configure, especially on Windows. Therefore, I strongly recommend creating a virtual environment using miniconda or venv. Below, I will briefly outline the setup method I usually use.

I prefer using miniconda over anaconda. It is much lighter because it installs only the essential components needed to create a virtual environment.

1. Installing Miniconda:

Linux:

1. Go to this link and download Miniconda3-latest-Linux-x86_64.sh.
2. Run bash Miniconda3-latest-Linux-x86_64.sh.

Windows:

1. Go to this link and download and run Miniconda3-latest-Windows-x86_64.exe.

2. Creating a Miniconda Virtual Environment:

Linux:

1. In the terminal: conda create -n your_own_env_name python=3.9

Windows:

1. Open Anaconda Prompt from the Start menu.
2. In the terminal: conda create -n your_own_env_name python=3.9

You can set your_own_env_name to any name you prefer.

3. Installing Packages within Miniconda (Same for Linux & Windows)

1. Run conda activate your_own_env_name in the terminal.
2. Select your desired version and OS on this link and run the command in the terminal. (Whether you use pip or conda does not matter.)
3. You must run the command that includes CUDA (e.g., conda install pytorch==2.2.1 torchvision==0.17.1 torchaudio==2.2.1 pytorch-cuda=11.8 -c pytorch -c nvidia).

Additionally, recently, NumPy version 2 is sometimes installed by default. If that happens, you can reinstall version 1. (For example, pip install numpy==1.26.*.)

Now, PyTorch and CUDA are automatically installed within your virtual environment. If you wish to install additional CUDA-related packages such as cuDNN, you can run conda install -c anaconda cudatoolkit==[desired version] and conda install -c anaconda cudnn. This will install the desired cudatoolkit version and the matching cudnn version.

Applying DDP

Now, let’s get straight to the point. We will look at applying DDP in two parts. The first is the terminal and script input method, and the second is the setting within the Python code.

Terminal and Script Input Method

First, assuming that the DDP setup is complete within the Python code, you can enter the following:

CUDA_VISIBLE_DEVICES=0,1,2,3 torchrun --nproc_per_node=4 --master_port 56789 main.py

Let’s look at this step by step.

• CUDA_VISIBLE_DEVICES=0,1,2,3: This specifies which GPUs you will use from those available on your machine. For example, if your local machine or server has 8 GPUs, this means you will use only GPUs 0 through 3. If you do not specify this separately, all available GPUs will be used.
• torchrun: Before PyTorch version 2.0, we used python -m torch.distributed.launch, but since version 2.0, torchrun has allowed us to execute DDP without typing python -m.
• --nproc_per_node=4: This specifies how many processes to run on that node. In this case, each process corresponds to one GPU, so you should match this number to the number of GPUs.
  (In the example above, since we use 4 GPUs, it is set to 4; the node is the local machine or server.)
• --master_port 56789: This part is actually optional, but it sets the port number used when running DDP. Sometimes, when running multiple DDP instances on a local machine or server, execution might fail if the port numbers overlap. In that case, you can enter any available port number.
• main.py: The name of the file you want to execute. (Obviously…)

There are more arguments you can use. For instance, --master_addr is used when utilizing multiple nodes, i.e., multiple servers. However, since this post is intended for people who are new to DDP and most readers will not need multiple nodes, I will skip this for now.

Settings Within Python Code

Before we dive in, there is one short thing to mention. As you can see from arguments like --master_addr above, DDP creates a separate process for each GPU. In other words, it may be easier to understand if you imagine that each GPU has its own process for running the Python file. During this process, variables are set within Python for each GPU (os.environ['variable_name']). Let’s briefly touch on two of these variables.

• os.environ['WORLD_SIZE']: This refers to the total number of processes launched by torchrun. In a single-node setting, where each GPU runs one process, this is usually the number of GPUs. For example, when using 4 GPUs, os.environ['WORLD_SIZE']=4.
• os.environ['LOCAL_RANK']: This refers to the local index of each process launched by torchrun. In other words, you can understand this as the GPU number (0 to 3 in the current example).

DDP Initialization

Based on this, let’s see how to set it up within the Python code. Everyone has their own coding style, but I usually write it as follows:

1
2
3
4
5
6
7
8

args.device = 'cuda:0'
args.world_size = 1
args.rank = 0
args.local_rank = int(os.environ.get("LOCAL_RANK", 0))
torch.cuda.set_device(args.local_rank)
torch.distributed.init_process_group(backend='nccl', init_method='env://')
args.world_size = torch.distributed.get_world_size()
args.local_rank = torch.distributed.get_rank()

This is just my coding style; I often use args.xxx. (You could define the variables separately, but keeping them inside args makes them convenient to use anywhere in the code.)

Let’s examine the code line by line.

• args.device='cuda:0' ~ args.rank=0: You can understand this as the step where we initialize the variables needed for DDP.
• args.local_rank = int(os.environ.get("LOCAL_RANK", 0)): As mentioned above, os.environ['LOCAL_RANK'] refers to the GPU number on which the code is running. Therefore, this line retrieves the GPU number through the get() function of os.environ and stores it in args.local_rank.
  • Example 1: If it corresponds to GPU 0, args.local_rank = 0
  • Example 2: If it corresponds to GPU 3, args.local_rank = 3
• torch.cuda.set_device(args.local_rank): Usually, there are two ways to move tensors or parameters to a GPU in PyTorch: params.to('cuda:#') and params.cuda(). I’ll skip explaining the first method because it is widely used. The second method moves them to the current default GPU. The default GPU is usually the first GPU (index 0), but set_device can be used to change this. In other words, this line designates the default GPU using the args.local_rank specified above.
• torch.distributed.init_process_group(backend='nccl', init_method='env://'): Understand this as the part that initializes the process group so that each GPU can communicate and operate correctly. Here, backend=’nccl’ specifies that CUDA will handle the operations, and init_method=’env://’ tells PyTorch to read the initialization information from environment variables. Since most people use this form, there is usually no need to change it.
• args.world_size, args.local_rank: I typically use these two to store the final WORLD_SIZE and LOCAL_RANK values. For args.local_rank, you can understand this as initializing the value once more for verification purposes.

Applying DDP to the Model

Next, let’s look at how to apply DDP to the model and train it. (It’s very simple.)

1
2
3
4
5
6

from torch.nn.parallel import DistributedDataParallel as DDP
model = DDP(model,device_ids=[args.local_rank])
...
logits = model(x)
loss = loss_fn(logits, labels)
loss.backward() 

As you can see, the model is wrapped with DistributedDataParallel, which is provided by PyTorch. In this process, the model is assigned to each GPU using the args.local_rank designated earlier. (Note: brackets are required for device_ids=[args.local_rank].)

Additional Tip - export NCCL_P2P_DISABLE=1

Unfortunately, sometimes it does not work smoothly… It might fail for various reasons, such as misconfigured settings or conflicts. Still, nowadays, thanks to LLMs (ChatGPT, Claude, Gemini, etc.), debugging has become much easier if you capture the error message clearly. I highly recommend using them properly when the problem is not caused by an internal code issue.

However… there are times when no error message appears, and you get stuck in an infinite loading state. This happened to me; right when entering the DDP function, the program would suddenly freeze and then hang indefinitely. I do not know the exact cause, but it seems to happen when the GPUs are processing the model during the internal DDP operations… (I’ve even tried debugging by printing everything inside the PyTorch framework… )

In such cases, you can try export NCCL_P2P_DISABLE=1. If you run this once in the terminal and then run the code, it often works smoothly. Haha.

Conclusion…

This was my first time posting about coding, and I hope this information is helpful to those who are new to PyTorch and DDP. I have experienced this myself, and from what I have seen, many people struggle for quite a long time when trying DDP for the first time. Sometimes searching through blogs does not help, and LLMs can be unkind, so I hope this post is a big help in those situations!