Process Group Initialization

The number of processes is not hard-coded.

Programming

Why a Single-GPU Script Needs New Pieces

Distributed Data Parallel, usually abbreviated DDP, is a technique for parallelizing model training across multiple GPUs. Rather than starting from theory, this lesson reads a script that is designed to be run as a whole. The script trains the NeuralNetwork model that was built earlier, adapted so that several GPUs can train it at the same time. The original training loop stays in place. What changes is the code around it.

The adaptation is not a single change. It touches four areas: how processes are started, how the process group that connects them is created and closed, how the dataset is split between them, and how the model is wrapped so that gradients are synchronized. Each area corresponds to a different utility, so every new line in the script has a specific job.

What do you think happens?

The main function takes a rank argument, but the call that starts it passes only world size and the number of epochs. Where does the rank value come from?

You type a different rank number into the mp.spawn call for each GPU
mp.spawn supplies a rank to each process automatically
The NCCL backend reads the rank from the GPU hardware

Reveal answer

Answer:

mp.spawn supplies a rank to each process automatically.

The call passes only world size and epochs as arguments. Each process that mp.spawn starts receives its own rank automatically, and the script uses that rank as the GPU ID.

Five Utilities and the Job Each One Does

The first block of imports brings in everything the DDP script needs. Each utility owns one part of the setup. The table below maps each one to the area of the adaptation it belongs to. Read the table as a map for the rest of the article: every section that follows explains one row.

Utility	Area it handles	Job in the script
mp.spawn	Starting processes	Starts several processes and applies the main function to inputs in parallel, one process per GPU
init_process_group	Connecting processes	Creates the process group the processes use to communicate
destroy_process_group	Connecting processes	Releases the process group when training is finished
DistributedSampler	Splitting the dataset	Decides which part of the dataset each process sees
DistributedDataParallel	Synchronizing the model	Wraps the model so its gradients are synchronized between GPUs during training

Each DDP utility corresponds to one of the four areas of the adaptation.

The script does not define its own process group or sampler. It calls the utilities from the imports. The main work for a learner is to recognize which line calls which utility and why that line exists. The sections below follow the order in which the script runs: first the processes are started, then they connect, then the data is divided, and finally the model is wrapped.

Starting One Process per GPU

The entry point of the script sits inside a __name__ == "__main__" clause. Its code runs when the file is executed as a Python script, not when it is imported as a module. Inside that clause the script first prints the number of available GPUs using torch.cuda.device_count(), sets a random seed for reproducibility, and then calls mp.spawn. The number of processes is not hard-coded. It is set to the number of GPUs the machine reports, and that value is stored as world_size.

python

import torch
import torch.multiprocessing as mp

def main(rank, world_size, epochs):
    # mp.spawn supplies rank; it is not passed by the caller
    ...

if __name__ == "__main__":
    print(torch.cuda.device_count())
    torch.manual_seed(0)
    world_size = torch.cuda.device_count()
    epochs = 10  # illustrative value
    mp.spawn(main, args=(world_size, epochs), nprocs=world_size)

import torch
import torch.multiprocessing as mp

def main(rank, world_size, epochs):
    # mp.spawn supplies rank; it is not passed by the caller
    ...

if __name__ == "__main__":
    print(torch.cuda.device_count())
    torch.manual_seed(0)
    world_size = torch.cuda.device_count()
    epochs = 10  # illustrative value
    mp.spawn(main, args=(world_size, epochs), nprocs=world_size)

Notes on the code

The entry point runs only when the file is executed directly.
torch.cuda.device_count() reports how many GPUs are available; the source stores this as world_size.
mp.spawn starts one process per GPU and passes world_size and epochs as the arguments after rank.
The rank parameter is filled in by mp.spawn for each process, not by the caller.

What happens, in order, from the moment mp.spawn is called until each GPU has its own process with its own rank? The main script calls mp.spawn, which starts one process per GPU, and each process receives a different rank.

The diagram shows two processes because a machine with two GPUs reports a world size of two. The same main function runs in every process. Each process receives a different rank, and the source uses that rank as the process identifier and as the GPU ID. Whenever the code writes rank, it is saying which GPU this particular process is responsible for. This single idea explains most of the later lines in the script.

This lesson uses only rank and world size. The source does not introduce a separate local rank, so the GPU-selection logic described below is tied to the rank value that mp.spawn supplies.

Tracing the Startup Sequence

Trace one run on a machine that reports two GPUs. Follow the state of the program at each step, noting what each process knows at that moment.

1
The script runs from the __name__ == "__main__" clause. torch.cuda.device_count() reports 2, so world_size is 2.
2
mp.spawn is called with nprocs set to 2 and the arguments world_size and epochs. At this point only the main script exists.
3
mp.spawn starts one process for each GPU. Process 0 receives rank 0 and Process 1 receives rank 1, without any action from the caller.
4
Each process runs the main function with its own rank and the same world size and epochs values.
5
Each process calls ddp_setup, which connects it to the group and selects the GPU that matches its rank. The next section covers this step in detail.

Worked example

Which rank does each process hold?

A machine reports 4 GPUs. Using the script as written, what world_size does the main script store, how many processes does mp.spawn start, and which rank does each one hold?

1
Read the count
torch.cuda.device_count() reports 4, so world_size is set to 4.
2
Read the spawn call
mp.spawn is called with nprocs=world_size, so the number of processes is 4.
3
Assign ranks
Each process receives a different rank, so the ranks are 0, 1, 2, and 3. The rank of each process is also its GPU ID.

Result

Four processes start with ranks 0 through 3, and rank 0 uses GPU 0 through rank 3 using GPU 3.

Connecting Processes with the Process Group

Once the processes exist, they must be able to communicate, because gradients have to be synchronized between them. The ddp_setup function handles this setup. According to the source, it does four things in order. It sets the address and port of the main node so that processes can reach one another. It initializes the process group with the NCCL backend using init_process_group. It records the rank and world size of the current process. Finally, it selects the GPU that matches the current rank.

1
Set the address and port of the main node, so that every process can find the others.
2
Call init_process_group with the NCCL backend to create the group the processes use to communicate.
3
Record the rank and world size of the current process.
4
Select the GPU that matches the current rank, so this process uses its own device.

The group is not left open indefinitely. destroy_process_group is the counterpart to init_process_group: it releases the process group that the processes used to communicate. The source introduces these two functions together because they bracket the life of the group. Between the two calls, the processes can exchange the information that DDP needs. Outside them, no group exists to exchange it.

Dividing the Data with DistributedSampler

If several processes train on the same dataset without any division, they all see the same examples. The source handles this with DistributedSampler, which divides the dataset into distinct, non-overlapping subsets, one for each process. The sampler is attached to the training DataLoader, and each process then receives its own subset. In this setup the sampler also takes over shuffling of the data.

python

from torch.utils.data import DataLoader
from torch.utils.data.distributed import DistributedSampler

train_loader = DataLoader(
    train_dataset,
    batch_size=32,  # illustrative value
    sampler=DistributedSampler(train_dataset),
)

from torch.utils.data import DataLoader
from torch.utils.data.distributed import DistributedSampler

train_loader = DataLoader(
    train_dataset,
    batch_size=32,  # illustrative value
    sampler=DistributedSampler(train_dataset),
)

Notes on the code

The sampler argument is what gives each process its own subset.
The rest of the DataLoader call looks the same as a single-GPU version.

Which dataset indices does each rank receive? This layout is illustrative. The source guarantees only that each process gets a distinct, non-overlapping subset, not this particular assignment.

The diagram lists eight examples and marks which of two ranks receives each one. Every example belongs to exactly one rank, so the subsets are disjoint. Each GPU then trains on only its own slice, which is why the sampler, not the DataLoader body, decides what each process sees.

Wrapping the Model for Gradient Synchronization

Inside main, the model is moved to the device identified by rank with model.to(rank). The features and labels of each batch are moved the same way. The model is then wrapped with DistributedDataParallel, which the script imports under the name DDP. The source describes this wrapping as the step that enables synchronization of gradients between the different GPUs during training.

python

model = model.to(rank)
model = DDP(model, device_ids=[rank])

model = model.to(rank)
model = DDP(model, device_ids=[rank])

Notes on the code

model.to(rank) places the model on the GPU for this process.
DDP(model, device_ids=[rank]) wraps the model so gradients are synchronized across GPUs during training.

During the backward pass, how do gradients move between GPUs so that every process ends up with synchronized gradients? The source does not spell out the step-by-step computation, so this diagram shows the logical pattern implied by the sampler and the DDP wrapper.

The diagram shows that each GPU works on its own batch, and that the gradients are kept consistent across devices before each GPU applies them. The training loop in the source still contains a placeholder where prediction and backpropagation code belongs. The exact per-step computation is therefore left to that code, and this article does not add to it.

Mistakes That Make Processes Diverge

A faulty DDP setup rarely crashes. Each process keeps running, but the processes stop behaving as one coordinated training run. The table compares a correct setup with three common faults. Use the last column to find the line that owns each fix.

Setup	What each process sees	Gradient behavior	Line that owns the fix
Correct setup	Its own distinct, non-overlapping subset of the data	Synchronized across GPUs by the DDP wrapper	DistributedSampler attached to the loader; DDP(model, device_ids=[rank])
Sampler omitted	The full dataset, so every process sees the same examples	Still wrapped, but the processes train on identical data	The sampler argument on the training DataLoader
Wrapper omitted	A correctly divided subset	Not synchronized between GPUs during training	The DDP(...) call after the model is moved to its device
Wrong rank or device	Data split by rank, but the model sits on a device that ignores the rank	Gradients are computed on the wrong device for the process	model.to(rank) and the device selection in ddp_setup

Three faults that cause processes to diverge from the expected behavior, each tied to one line of the script.

Common mistake
Removing the sampler from the training DataLoader
What it looks like
```
DataLoader(train_dataset, batch_size=32) with no sampler argument
```
Why it fails
Each process iterates over the whole dataset, so all processes see the same examples. The loader still runs without error.
Do this instead
Attach DistributedSampler(train_dataset) through the sampler argument.
Common mistake
Skipping the DDP wrapper
What it looks like
Training with model.to(rank) but never calling DDP(model, device_ids=[rank])
Why it fails
The source ties gradient synchronization between GPUs to the DDP wrapper, so without it the gradients are not synchronized during training.
Do this instead
Wrap the model with DDP after moving it to the device for this rank.
Common mistake
Assigning every process the same device
What it looks like
model.to(0) inside main, executed in every process
Why it fails
The source moves the model to the device identified by rank. A fixed device ignores the rank, so processes that should use different GPUs all use the same one.
Do this instead
Use model.to(rank) so each process uses the device that matches its own rank.
Common mistake
Passing rank yourself in mp.spawn args
What it looks like
```
mp.spawn(main, args=(rank, world_size, epochs), nprocs=world_size)
```
Why it fails
mp.spawn supplies rank to each process automatically, so passing it again conflicts with the rank parameter of main.
Do this instead
Pass only the arguments that follow rank, such as world_size and epochs.
Common mistake
Forgetting destroy_process_group after training, which leaves the process group unreleased

Best practice

When a DDP run produces an unexpected result, ask which stage of the setup could have produced it. The stages are spawning, process group creation, data division, and model wrapping. Each one has a specific line that owns it, so a symptom can often be traced back to one place.

Locating a Fault in a DDP Script

MEDIUM

A teammate runs the DDP script on a machine with four GPUs. The run finishes, but the logs suggest only one process is training. Using the four stages from this article, which stage would you check first, and which line in the script owns that stage?

Hints

Start with the stage that runs first. Which code decides how many processes are started?
The number of processes comes from the value stored as world_size.
Check the spawn call: mp.spawn is given nprocs=world_size.

Your attempt (not saved)

EASY

Name the line that gives each process its data slice, and describe what happens to that process if the line is deleted while the rest of the loader stays the same.

Hints

The sampler argument is the piece that changes between the single-GPU and DDP loaders.
Think about what each process iterates over once the sampler is gone.

Your attempt (not saved)

Key Takeaways

mp.spawn starts one process per GPU, using the number of GPUs the machine reports as world_size, and each process receives its own rank automatically, which the source uses as its GPU ID.
init_process_group creates the NCCL process group that connects the processes, ddp_setup records rank and world size and selects the GPU, and destroy_process_group releases the group when finished.
DistributedSampler attached to the training DataLoader gives each process a distinct, non-overlapping subset of the dataset and takes over shuffling in this setup.
model.to(rank) places the model on the GPU for each process, and wrapping it with DDP enables gradient synchronization between GPUs during training.
When a DDP run misbehaves, trace the symptom to one of the four stages (spawning, process group creation, data division, model wrapping), since each stage has a specific owning line.

Related concepts

Training a Neural Network Model in PyTorch

More concepts starting with P

Why a Single-GPU Script Needs New Pieces

Five Utilities and the Job Each One Does

Starting One Process per GPU

Tracing the Startup Sequence

Connecting Processes with the Process Group

Dividing the Data with DistributedSampler

Wrapping the Model for Gradient Synchronization

Mistakes That Make Processes Diverge

Locating a Fault in a DDP Script

Key Takeaways

Related Topics

Related concepts