In PyTorch, a key component of distributing tasks across multiple processes or machines for parallel computation is setting up Distributed Data Parallel (DDP). A crucial aspect of this setup is determining the "rank" of each process involved in the computation. The rank is essentially an identifier for each process involved in the distributed training, allowing the system to differentiate and coordinate between them.

Understanding Process Rank in PyTorch DDP

Each process in a distributed training scenario needs to know its rank because the rank determines how the data is split among the processes and how they communicate with each other. PyTorch provides several means to set and get the rank of a process either through environment variables or explicitly setting them in the code.

In most implementations of PyTorch's DDP, including tutorials and practical examples, the setup function determines each process's rank in one of the following ways:

1. Environment Variables: PyTorch uses environment variables such as RANK, WORLD_SIZE, and MASTER_ADDR to automatically configure settings for each process. When using tools like torch.distributed.launch or torch.distributed.run, these variables are typically set automatically. Here, RANK refers to the process's unique identifier within the group of processes.
2. Explicit Specification in Code: The rank can also be set manually in the script. This is common in smaller setups or custom distributed training setups where automatic tools are not used. For instance, if using torch.distributed.init_process_group, one can specify the rank explicitly in the function call.

How does setup Obtain the Rank?

The setup function in a PyTorch DDP tutorial generally configures each process for distributed training. The function initializes the process group, which requires knowledge of the rank. Here's a typical flow of how setup might be implemented to determine the rank:

1import os
2import torch.distributed as dist
3
4def setup(rank, world_size):
5    os.environ['MASTER_ADDR'] = 'localhost'
6    os.environ['MASTER_PORT'] = '12345'
7    dist.init_process_group("nccl", rank=rank, world_size=world_size)

In this code:

• The rank and world_size would typically be passed to the setup function when starting the training script.
• Environment variables for MASTER_ADDR and MASTER_PORT are set, which are necessary for initializing the process group.
• dist.init_process_group is called with "nccl" as the backend (suitable for GPU-based training) and the rank and world size are specified.

Execution Example

When executing a PyTorch script using DDP, the setup might be triggered using a command line interface or a launch utility, which internally sets these parameters:

python -m torch.distributed.launch --nproc_per_node=4 --nnodes=1 --node_rank=0 --master_addr="localhost" --master_port=12345 train_script.py

This method automatically sets the environment variables and passes to the training script the necessary rank and world_size.

Summary Table

Conclusion

In PyTorch DDP, the setup function's ability to identify each process's rank is vital for orchestrating distributed training. It ensures that data and tasks are properly divided and that processes communicate effectively throughout the training phase. The specifics can vary depending on the size of the setup and the tools used, but understanding these basics is crucial for effective implementation of distributed deep learning models.