403 - Distributed Training

Distributed training is the process of training a machine learning model using multiple computing devices simultaneously.

In traditional single-GPU training, one GPU performs the computations needed to process training data, calculate model predictions, compute gradients, and update model parameters.

Distributed training allows multiple GPUs to participate in this process.

Why Use Distributed Training

Researchers use distributed training to:

• Reduce training time: Process more training data simultaneously.
• Scale machine learning workloads: Train models using more computational resources.
• Support larger models: Use advanced parallelism techniques when models cannot fit on a single GPU.
• Increase throughput: Process more training samples per second.
• Improve resource utilization: Distribute computational work across available GPUs.

However, distributed training introduces additional communication and synchronization overhead.

Important: Using more GPUs does not automatically result in faster training. Performance depends on model size, batch size, GPU capabilities, communication overhead, and data-loading efficiency.

Single-GPU vs. Distributed Training

Understanding Data Parallelism and Model Parallelism

Distributed training generally involves two major approaches: data parallelism and model parallelism.

Data Parallelism

Data parallelism distributes training data across multiple GPUs while maintaining a copy of the same model on each GPU.

Each GPU:

1. Receives a different subset of training data.
2. Performs a forward pass using its local model.
3. Calculates the loss and gradients.
4. Synchronizes gradients with the other GPUs.
5. Updates its model parameters.

Because gradients are synchronized, model replicas remain consistent after optimizer updates.

Model Parallelism

Model parallelism distributes different parts of a neural network across multiple GPUs. Unlike data parallelism, each GPU stores only part of the model. Model parallelism is useful when a model is too large to fit entirely into the memory of a single GPU.

Comparison

Note: This module focuses primarily on data parallelism using PyTorch DDP because it provides a practical introduction to distributed training.

Understanding PyTorch DistributedDataParallel (DDP)

PyTorch is an open-source machine learning framework used to develop and train neural networks.

PyTorch includes DistributedDataParallel (DDP), a tool for training models across multiple GPUs.

DDP creates a model replica for each training process and synchronizes gradients during training.

How DDP Works

A typical DDP training workflow involves:

• Initializing a distributed process group.
• Assigning each process to a GPU.
• Creating and wrapping the model with DDP.
• Distributing training data across processes.
• Performing forward and backward passes.
• Synchronizing gradients.
• Updating model parameters.
• Shutting down the process group.

Important Terminology

Example: Two-GPU Training

Suppose you launch a training job using two GPUs on one compute node.

The world size is 2 because two processes participate in training.

Each process works with its assigned GPU and coordinates gradient updates with the other process.

Preparing a PyTorch Environment on ICE

Before running distributed training, you must verify that your computing environment supports PyTorch and the allocated GPUs.

Step 1: Connect to ICE

Log in using SSH or Open OnDemand. Refer to Module 100 for instructions on accessing ICE.

Step 2: Check Available GPU Resources

Run:

sinfo

You can also inspect Slurm’s reported GPU resources with:

sinfo -o "%P %D %G %T"

Step 3: Set Up Your Python Environment

If you have an approved Python environment, activate it.

For example:

source .venv/bin/activate

Otherwise, after selecting a compatible Python installation, you can create an environment:

python -m venv .venv
source .venv/bin/activate

Step 4: Verify PyTorch

Run:

python -c "import torch; print(torch.__version__)"

To check whether your PyTorch installation was built with CUDA support:

python -c "import torch; print(torch.version.cuda)"

The reported CUDA version refers to the CUDA version used to build PyTorch, not necessarily the version of the system’s installed CUDA toolkit.

If PyTorch is unavailable, consult PyTorch’s official installation instructions and PACE software guidance.

Install a build compatible with your Python version, allocated GPU hardware, and NVIDIA driver.

Step 5: Verify GPU Access

After obtaining a GPU allocation and launching a compute-node shell, run:

nvidia-smi

Then verify GPU access from Python:

python -c "import torch; print(torch.cuda.is_available())"

Expected output on a compatible GPU compute node:

True

You can also check how many GPUs PyTorch can see:

python -c "import torch; print(torch.cuda.device_count())"

Important: These checks must be performed on a GPU compute node. Running them directly on a login node may produce different results.

Creating a Simple Distributed Training Script

In this section, we will create a small PyTorch script that demonstrates distributed training with DDP. The example trains a simple neural network using randomly generated data. It does not require downloading a dataset.

Step 1: Create a Python File

Create a file named:

train_ddp.py

Step 2: Add the Training Code

import os

import torch
import torch.distributed as dist
from torch import nn
from torch.nn.parallel import DistributedDataParallel as DDP
from torch.utils.data import DataLoader, TensorDataset
from torch.utils.data.distributed import DistributedSampler


def main():
    # Initialize distributed communication.
    dist.init_process_group(backend="nccl")

    local_rank = int(os.environ["LOCAL_RANK"])
    rank = dist.get_rank()
    world_size = dist.get_world_size()

    torch.cuda.set_device(local_rank)
    device = torch.device("cuda", local_rank)

    # Build a small neural network.
    model = nn.Sequential(
        nn.Linear(10, 32),
        nn.ReLU(),
        nn.Linear(32, 2)
    ).to(device)

    model = DDP(model, device_ids=[local_rank])

    # Generate an example dataset.
    # A fixed seed makes the dataset identical across ranks.
    generator = torch.Generator().manual_seed(42)

    features = torch.randn(
        1024, 10, generator=generator
    )

    labels = torch.randint(
        0, 2, (1024,), generator=generator
    )

    dataset = TensorDataset(features, labels)

    # Give each process a different data subset.
    sampler = DistributedSampler(
        dataset,
        num_replicas=world_size,
        rank=rank,
        shuffle=True
    )

    loader = DataLoader(
        dataset,
        batch_size=32,
        sampler=sampler
    )

    optimizer = torch.optim.SGD(
        model.parameters(),
        lr=0.01
    )

    loss_fn = nn.CrossEntropyLoss()

    # Train for three epochs.
    for epoch in range(3):
        sampler.set_epoch(epoch)

        total_loss = 0.0

        for inputs, targets in loader:
            inputs = inputs.to(device)
            targets = targets.to(device)

            optimizer.zero_grad()

            outputs = model(inputs)
            loss = loss_fn(outputs, targets)

            loss.backward()
            optimizer.step()

            total_loss += loss.item()

        if rank == 0:
            print(
                f"Epoch {epoch + 1}: "
                f"Rank 0 local loss sum = {total_loss:.4f}",
                flush=True
            )

    dist.destroy_process_group()


if __name__ == "__main__":
    main()

Understand the Code

The example prints the sum of rank 0’s local batch losses, not a globally averaged loss across all GPUs.

Important: This example demonstrates how DDP operates. Because the network is small, training with two GPUs may be slower than training with one GPU due to communication overhead.

Requesting Multiple GPUs Using Slurm

ICE uses Slurm to schedule computing jobs and allocate resources.

To train a model across multiple GPUs, you must request the required resources before launching the training processes.

Step 1: Create a Slurm Submission Script

Create a file named:

train_ddp.sbatch

Add the following:

#!/bin/bash
#SBATCH --job-name=ddp-training
#SBATCH --nodes=1
#SBATCH --ntasks=1
#SBATCH --cpus-per-task=4
#SBATCH --gres=gpu:2
#SBATCH --time=00:30:00
#SBATCH --output=ddp-%j.out
#SBATCH --error=ddp-%j.err

set -euo pipefail

## Activate the prepared Python environment.
source .venv/bin/activate

## Display GPU information.
nvidia-smi

## Launch one training process per GPU.
torchrun \
    --standalone \
    --nnodes=1 \
    --nproc-per-node=2 \
    train_ddp.py

Understanding the Slurm Script

Why Use torchrun?

torchrun is a PyTorch utility that launches distributed training processes.

In this example:

torchrun --standalone --nnodes=1 --nproc-per-node=2 train_ddp.py

• --standalone configures a single-node distributed launch.

• --nnodes=1 specifies one compute node.

• --nproc-per-node=2 launches two training processes.

• train_ddp.py is the Python training script.

Each training process is assigned a local rank and uses a different GPU.

Important: The Slurm submission script is a general single-node example, not a verified ICE configuration. You may need to specify an authorized partition or account and adjust the requested resources. GPU requests also depend on Slurm’s configuration. Confirm that the node exposes two GPUs to the launch task before running the experiment.

Submitting and Monitoring a Distributed Job

Step 1: Submit the Job

Run:

sbatch train_ddp.sbatch

Slurm will return a job ID if the submission is accepted.

Example:

Submitted batch job 123456

Step 2: Monitor Job Status

Run:

squeue --me

Common job states include:

Completed jobs may no longer appear in squeue.

Step 3: Inspect Output Logs

After the job starts, inspect its output file:

cat ddp-123456.out

Replace 123456 with your actual job ID.

If the script runs successfully, you should see messages similar to:

Epoch 1: Rank 0 local loss sum = ...
Epoch 2: Rank 0 local loss sum = ...
Epoch 3: Rank 0 local loss sum = ...

The exact loss values may vary.

Step 4: Check Job History

Run:

sacct -j 123456

This displays accounting information for the job, if available.

For more details:

sacct -j 123456 \
    --format=JobID,JobName,State,Elapsed,ExitCode

Step 5: Cancel a Job

If necessary, cancel a job using:

scancel 123456

Only cancel jobs you are authorized to manage.

Understanding Distributed Training Performance

Using multiple GPUs can improve training performance, but scaling efficiency depends on several factors.

What Affects Performance?

Training Throughput

Training throughput measures how much data is processed within a given amount of time.

For example:

Throughput = Number of Samples Processed / Training Time

A higher throughput generally means the model processes more training samples per second.

Scaling Efficiency

Scaling efficiency measures how well training performance improves when additional GPUs are used.

Speedup = Single-GPU Time / Multi-GPU Time

For example, if an experiment takes 100 seconds on one GPU and 60 seconds on two GPUs:

Speedup = 100 / 60 = 1.67

The experiment achieved approximately a 1.67× speedup.

When comparing configurations, keep the model, dataset, global batch size, and experimental settings consistent unless the purpose is to study how those settings affect scaling.

Important: GPU utilization alone does not determine whether distributed training is efficient. Researchers should also measure throughput, training time, memory usage, and model quality.

Further reading

• PyTorch torchrun documentation
• PyTorch DistributedDataParallel documentation
• PACE documentation