Distributed training basics in MLOps - Commands & Configuration

Introduction

Training machine learning models on large datasets can take a long time on a single computer. Distributed training splits the work across multiple machines or processors to finish faster and handle bigger data.

When your dataset is too large to fit into one machine's memory.

When training a deep learning model takes hours or days on a single GPU.

When you want to speed up model training by using multiple GPUs or machines.

When you need to scale your training to handle more complex models.

When you want to improve resource usage by distributing workload efficiently.

Commands

This command starts distributed training using PyTorch's built-in launcher with 2 processes on one machine. Each process handles one GPU to train the model in parallel.

Terminal

python -m torch.distributed.launch --nproc_per_node=2 train.py

Expected Output

INFO: Distributed training initialized with 2 processes

Epoch 1/10: loss=0.45

Epoch 2/10: loss=0.38

... (training logs continue)

→

--nproc_per_node - Number of processes to launch on the current node, usually set to the number of GPUs.

This command runs one process of a distributed training job using the NCCL backend for GPU communication. It specifies the total number of processes (world_size), this process's rank, and the master node's address and port.

Terminal

python train.py --backend nccl --world_size 4 --rank 0 --master_addr 127.0.0.1 --master_port 29500

Expected Output

INFO: Process 0 initialized for distributed training

Epoch 1/10: loss=0.46

Epoch 2/10: loss=0.39

... (training logs continue)

→

--backend - Communication backend for distributed training, NCCL is optimized for GPUs.

→

--world_size - Total number of processes participating in training.

→

--rank - Unique ID of this process among all processes.

→

--master_addr - IP address of the master node coordinating training.

→

--master_port - Port on the master node for communication.

Key Concept

If you remember nothing else from distributed training, remember: splitting training across multiple processes or machines lets you handle bigger data and train faster by working together.

Code Example

MLOps

import torch
import torch.distributed as dist
import torch.nn as nn
import torch.optim as optim
from torch.nn.parallel import DistributedDataParallel as DDP
import os

def setup():
    dist.init_process_group(backend='nccl', init_method='env://')
    torch.cuda.set_device(int(os.environ['LOCAL_RANK']))

def cleanup():
    dist.destroy_process_group()

class SimpleModel(nn.Module):
    def __init__(self):
        super(SimpleModel, self).__init__()
        self.linear = nn.Linear(10, 1)
    def forward(self, x):
        return self.linear(x)


def train():
    setup()
    model = SimpleModel().cuda()
    ddp_model = DDP(model, device_ids=[int(os.environ['LOCAL_RANK'])])
    loss_fn = nn.MSELoss()
    optimizer = optim.SGD(ddp_model.parameters(), lr=0.001)

    for epoch in range(2):
        inputs = torch.randn(20, 10).cuda()
        labels = torch.randn(20, 1).cuda()
        optimizer.zero_grad()
        outputs = ddp_model(inputs)
        loss = loss_fn(outputs, labels)
        loss.backward()
        optimizer.step()
        print(f"Epoch {epoch+1}: loss={loss.item():.4f}")

    cleanup()

if __name__ == '__main__':
    train()

Common Mistakes

Not setting the correct world_size or rank for each process.

Processes won't coordinate properly, causing errors or hanging training.

Ensure each process has a unique rank and the total world_size matches the number of processes.

Using the wrong backend for the hardware (e.g., using NCCL on CPUs).

Communication will fail or be very slow because the backend is not compatible.

Use NCCL for GPUs and Gloo for CPUs.

Not launching the right number of processes matching available GPUs.

Some GPUs remain idle or processes compete for the same GPU, reducing efficiency.

Set nproc_per_node to the number of GPUs on the machine.

Summary

Use torch.distributed.launch or environment variables to start multiple training processes for distributed training.

Set backend, world_size, rank, master address, and port correctly to enable communication between processes.

DistributedDataParallel wraps your model to synchronize gradients and speed up training across GPUs or machines.

Practice

1. What is the main purpose of distributed training in machine learning?

easy

A. To avoid using GPUs during training

B. To split the training workload across multiple machines or GPUs

C. To increase the learning rate automatically

D. To reduce the size of the training dataset

Solution

1. Step 1: Understand distributed training goal

   Distributed training is designed to share the training task among several machines or GPUs to speed up the process.

2. Step 2: Analyze options

   Only To split the training workload across multiple machines or GPUs correctly describes this purpose. Options A, B, and C do not relate to workload distribution.

3. Final Answer:

   To split the training workload across multiple machines or GPUs -> Option B

4. Quick Check:

   Distributed training = workload split [OK]

Hint: Distributed training means sharing work across machines [OK]

Common Mistakes:

• Thinking distributed training reduces dataset size
• Confusing distributed training with hyperparameter tuning
• Believing distributed training avoids GPU use

2. Which of the following is the correct way to initialize a process group for distributed training in PyTorch?

easy

A. torch.distributed.init_process_group(backend='nccl', rank=0, world_size=1)

B. torch.init_process_group(backend='nccl', rank=0, world_size=1)

C. torch.distributed.start_process_group(backend='nccl', rank=0, world_size=1)

D. torch.distributed.init_group(backend='nccl', rank=0, world_size=1)

Solution

1. Step 1: Identify correct function name

   The correct function to initialize communication is torch.distributed.init_process_group.

2. Step 2: Check syntax correctness

   torch.distributed.init_process_group(backend='nccl', rank=0, world_size=1) uses the correct module and function name with proper parameters. Other options use incorrect function names or modules.

3. Final Answer:

   torch.distributed.init_process_group(backend='nccl', rank=0, world_size=1) -> Option A

4. Quick Check:

   Correct init function = torch.distributed.init_process_group(backend='nccl', rank=0, world_size=1) [OK]

Hint: Use torch.distributed.init_process_group to start communication [OK]

Common Mistakes:

• Using wrong function names like start_process_group
• Calling init_process_group from wrong module
• Misspelling function or module names

3. Given the following code snippet for distributed training setup, what is the output of print(rank, world_size)?

import torch.distributed as dist
rank = 2
world_size = 4
print(rank, world_size)

medium

A. 4 2

B. Error: rank and world_size undefined

C. 0 1

D. 2 4

Solution

1. Step 1: Analyze variable assignments

   Variables rank and world_size are assigned values 2 and 4 respectively before the print statement.

2. Step 2: Understand print output

   Printing rank and world_size will output '2 4' exactly as assigned.

3. Final Answer:

   2 4 -> Option D

4. Quick Check:

   Print rank, world_size = 2 4 [OK]

Hint: Print variables as assigned to see output [OK]

Common Mistakes:

• Confusing rank with world_size order
• Assuming variables are undefined
• Expecting automatic values without assignment

4. You wrote this code to initialize distributed training but get an error:

import torch.distributed as dist
dist.init_process_group(backend='nccl', rank=0)

What is missing that causes the error?

medium

A. The rank parameter should be a string

B. The backend parameter is incorrect

C. The world_size parameter is missing

D. The import statement is wrong

Solution

1. Step 1: Check init_process_group parameters

   The function requires both rank and world_size parameters to know the total number of processes.

2. Step 2: Identify missing parameter

   The code misses world_size, which causes the error.

3. Final Answer:

   The world_size parameter is missing -> Option C

4. Quick Check:

   Missing world_size causes error [OK]

Hint: Always provide world_size with rank in init_process_group [OK]

Common Mistakes:

• Omitting world_size parameter
• Using wrong backend names
• Passing rank as string instead of int

5. In a distributed training setup with 4 GPUs, you want each process to know its rank and the total number of processes. Which code snippet correctly sets this up and prints the rank and world size?

hard

A. import torch.distributed as dist
dist.init_process_group(backend='nccl')
rank = dist.get_rank()
world_size = dist.get_world_size()
print(rank, world_size)

B. import torch.distributed as dist
world_size = 4
rank = dist.get_rank()
dist.init_process_group(backend='nccl', rank=rank, world_size=world_size)
print(rank, world_size)

C. import torch.distributed as dist
rank = 0
world_size = 4
dist.init_process_group(backend='nccl', rank=rank, world_size=world_size)
print(rank, world_size)

D. import torch.distributed as dist
rank = dist.get_rank()
world_size = dist.get_world_size()
dist.init_process_group(backend='nccl', rank=rank, world_size=world_size)
print(rank, world_size)

Solution

1. Step 1: Understand correct initialization order

   dist.init_process_group must be called before calling dist.get_rank() or dist.get_world_size() to initialize communication.

2. Step 2: Analyze each option

   import torch.distributed as dist dist.init_process_group(backend='nccl') rank = dist.get_rank() world_size = dist.get_world_size() print(rank, world_size) initializes the process group first, then gets rank and world size, then prints them. Other options either get rank before initialization or pass rank manually without initialization.

3. Final Answer:

   import torch.distributed as dist dist.init_process_group(backend='nccl') rank = dist.get_rank() world_size = dist.get_world_size() print(rank, world_size) -> Option A

4. Quick Check:

   Init first, then get rank/world_size = import torch.distributed as dist dist.init_process_group(backend='nccl') rank = dist.get_rank() world_size = dist.get_world_size() print(rank, world_size) [OK]

Hint: Initialize before getting rank and world size [OK]

Common Mistakes:

• Calling get_rank before init_process_group
• Passing rank manually without init
• Not calling init_process_group at all