Skip to content
BranchesTags

Repository files navigation

Multi-GPU Playground

Educational PyTorch examples for understanding major single-GPU and multi-GPU training patterns used in modern deep learning systems.

The repository uses moderately sized synthetic Transformer models so that differences in GPU memory usage, utilization, communication, and scaling are easier to observe. The scripts remain focused on training mechanics such as device placement, process groups, batch semantics, sharding, collectives, mixed precision, pipeline scheduling, tensor parallelism, and checkpoint formats.

The dataset is random by design. Therefore, the loss is not expected to improve meaningfully. With 10 random classes, the loss should remain near:

log(10) ≈ 2.30

The primary purpose of this repository is to compare execution behavior, not model accuracy.

What Is Included

Script Pattern Typical launch
train_single.py Baseline single-device training loop uv run python train_single.py
train_mixed_precision_single_gpu.py Single-GPU FP32/FP16/BF16 AMP comparison uv run python train_mixed_precision_single_gpu.py --precision bf16
train_dataparallel.py Legacy nn.DataParallel uv run python train_dataparallel.py
train_ddp.py DistributedDataParallel data parallelism uv run torchrun --standalone --nproc_per_node=3 train_ddp.py
train_fsdp.py FSDP1 full sharding and full-state checkpointing uv run torchrun --standalone --nproc_per_node=3 train_fsdp.py
train_fsdp2.py FSDP2 fully_shard and distributed checkpointing uv run torchrun --standalone --nproc_per_node=3 train_fsdp2.py
train_model_parallel.py Manual layer placement across 3 GPUs uv run python train_model_parallel.py
train_pipeline_parallel.py 3-stage pipeline parallelism with GPipe scheduling uv run torchrun --standalone --nproc_per_node=3 train_pipeline_parallel.py
train_tensor_parallel.py Pure tensor-parallel MLP sharding with a 1D mesh uv run torchrun --standalone --nproc_per_node=3 train_tensor_parallel.py
train_2d_dp_tp.py Data-parallel groups combined with tensor-parallel groups uv run torchrun --standalone --nproc_per_node=4 train_2d_dp_tp.py --tp-size 2

The interview study guide is split into focused files under notes/. The original consolidated notes remain in multi_gpu_training_notes.md.

Model Sizes

Most examples use a shared Transformer encoder classifier with approximately 110 million parameters.

Vocabulary size:          32,000
Sequence length:             256
Model dimension:             768
Attention heads:              12
Transformer layers:           12
Feed-forward dimension:     3,072
Output classes:               10

The tensor-parallel and 2D DP + TP examples use a Transformer-style MLP model with approximately 138 million parameters.

Vocabulary size:          32,000
Sequence length:             256
Model dimension:             768
MLP hidden dimension:      6,144
MLP layers:                   12
Output classes:               10

These models are intentionally large enough to make GPU differences visible, while remaining practical on modern multi-GPU systems such as NVIDIA A40-class hardware.

Requirements

  • Python 3.10+
  • NVIDIA GPU or GPUs
  • CUDA-capable PyTorch installation
  • uv
  • NumPy
  • Multiple GPUs for distributed examples
  • Exactly 3 visible GPUs for manual model parallelism and pipeline parallelism

Install all declared dependencies:

uv sync

If NumPy is missing, add it to the project:

uv add numpy

Verify the environment:

uv run python -c "import numpy, torch; print('NumPy:', numpy.__version__); print('PyTorch:', torch.__version__); print('CUDA:', torch.cuda.is_available()); print('GPU count:', torch.cuda.device_count())"

Sanity-check the repository entry point:

uv run python main.py

Monitoring GPU Usage

Open a second terminal while a training command is running.

Recommended Live View

Refresh the standard NVIDIA GPU status display every second:

watch -n 1 nvidia-smi

Press Ctrl+C to stop watching.

This view shows:

  • GPU utilization
  • memory usage
  • temperature
  • power consumption
  • active processes
  • process IDs
  • GPU assignments

A typical workflow is:

Terminal 1 — start monitoring

watch -n 1 nvidia-smi

Terminal 2 — start training

CUDA_VISIBLE_DEVICES=0,1,2 uv run torchrun \
    --standalone \
    --nproc_per_node=3 \
    train_pipeline_parallel.py

This makes it easy to observe how memory and utilization change across GPUs during initialization, forward passes, backward passes, communication, and optimizer steps.

One-Time GPU Status

nvidia-smi

Watch Selected Physical GPUs

watch -n 1 nvidia-smi -i 0,1,2

Compact Live Metrics

nvidia-smi \
    --query-gpu=index,name,utilization.gpu,memory.used,memory.total,temperature.gpu,power.draw \
    --format=csv \
    -l 1

Process-Focused View

nvidia-smi pmon -s um

Watch GPU Topology

GPU topology is useful when interpreting peer-to-peer communication performance:

nvidia-smi topo -m

Look for connections such as:

NV#    NVLink connection
PIX    same PCIe switch
PXB    multiple PCIe bridges
PHB    PCIe host bridge
SYS    CPU/socket or NUMA interconnect

Pipeline, tensor, and distributed training performance may depend strongly on GPU interconnect topology.

Selecting GPUs

Use CUDA_VISIBLE_DEVICES before the command.

Run on one selected physical GPU:

CUDA_VISIBLE_DEVICES=1 uv run python train_single.py

Use stable PCI bus ordering:

CUDA_DEVICE_ORDER=PCI_BUS_ID \
CUDA_VISIBLE_DEVICES=0,2 \
uv run python train_dataparallel.py

CUDA remaps visible physical GPUs to local process indices.

For example:

CUDA_VISIBLE_DEVICES=2,3

Inside the program:

cuda:0 -> physical GPU 2
cuda:1 -> physical GPU 3

Therefore, scripts with a --device argument should generally use local device 0 after selecting a single physical GPU:

CUDA_VISIBLE_DEVICES=2 \
uv run python train_mixed_precision_single_gpu.py \
    --device 0 \
    --precision bf16

Running the Examples

1. Single-GPU Baseline

CUDA_VISIBLE_DEVICES=0 \
uv run python train_single.py

Use this as the baseline for:

  • iteration time
  • GPU memory usage
  • GPU utilization
  • model replication overhead
  • comparison with multi-GPU methods

2. Mixed Precision

FP32:

CUDA_VISIBLE_DEVICES=0 \
uv run python train_mixed_precision_single_gpu.py \
    --device 0 \
    --precision fp32

FP16:

CUDA_VISIBLE_DEVICES=0 \
uv run python train_mixed_precision_single_gpu.py \
    --device 0 \
    --precision fp16

BF16:

CUDA_VISIBLE_DEVICES=0 \
uv run python train_mixed_precision_single_gpu.py \
    --device 0 \
    --precision bf16

Compare:

  • peak allocated memory
  • reserved memory
  • step time
  • numerical stability
  • GPU utilization

BF16 is generally the most practical reduced-precision mode on A40-class GPUs.

3. nn.DataParallel

CUDA_VISIBLE_DEVICES=0,1 \
uv run python train_dataparallel.py

nn.DataParallel uses one Python process and splits batches across visible GPUs. GPU 0 acts as the primary device and may use more memory than the other GPUs.

Watch for uneven memory usage:

watch -n 1 nvidia-smi

4. DistributedDataParallel

Two GPUs:

CUDA_VISIBLE_DEVICES=0,1 \
uv run torchrun \
    --standalone \
    --nproc_per_node=2 \
    train_ddp.py

Three GPUs:

CUDA_VISIBLE_DEVICES=0,1,2 \
uv run torchrun \
    --standalone \
    --nproc_per_node=3 \
    train_ddp.py

For torchrun, --nproc_per_node must equal the number of visible GPUs.

DDP creates one process per GPU and synchronizes gradients through collectives.

Observe:

  • similar memory usage across GPUs
  • similar utilization across GPUs
  • gradient synchronization overhead
  • global batch-size growth with world size

5. FSDP1

CUDA_VISIBLE_DEVICES=0,1,2 \
uv run torchrun \
    --standalone \
    --nproc_per_node=3 \
    train_fsdp.py

FSDP shards model parameters, gradients, and optimizer-related state across ranks.

Observe:

  • lower persistent model-state memory per GPU
  • temporary memory increases during parameter all-gathers
  • communication during forward and backward passes
  • checkpoint creation at the end of training

6. FSDP2

CUDA_VISIBLE_DEVICES=0,1,2 \
uv run torchrun \
    --standalone \
    --nproc_per_node=3 \
    train_fsdp2.py

FSDP2 uses fully_shard and DTensor-backed state.

The script may generate:

fsdp2_tiny_transformer_full_state.pt
fsdp2_tiny_transformer_sharded_dcp/

Observe:

  • sharded parameter memory
  • resharding behavior
  • distributed checkpoint writing
  • differences from FSDP1

7. Manual Model Parallelism

This example expects exactly 3 visible GPUs:

CUDA_VISIBLE_DEVICES=0,1,2 \
uv run python train_model_parallel.py

The model is divided by layers:

GPU 0:
    embeddings
    first 4 Transformer layers

GPU 1:
    middle 4 Transformer layers

GPU 2:
    final 4 Transformer layers
    normalization
    classifier head
    loss

Observe:

  • model memory distributed across GPUs
  • activation transfers between devices
  • sequential execution between model partitions
  • GPU utilization gaps caused by lack of pipeline overlap

8. Pipeline Parallelism

Pipeline parallelism expects exactly 3 ranks:

CUDA_VISIBLE_DEVICES=0,1,2 \
uv run torchrun \
    --standalone \
    --nproc_per_node=3 \
    train_pipeline_parallel.py

Current pipeline configuration:

Pipeline stages:       3
Layers per stage:      4
Total layers:         12
Batch size:           96
Microbatches:          6
Microbatch size:      16

The global batch is divided into microbatches and processed using GPipe scheduling.

Watch GPU activity in another terminal:

watch -n 1 nvidia-smi

Expected startup messages may include NCCL warnings about unbatched point-to-point operations. These warnings do not necessarily indicate a failure.

NumPy must be installed because pipeline stage metadata exchange may internally use NumPy:

uv add numpy

Observe:

  • pipeline fill and drain periods
  • utilization overlap between stages
  • memory differences among stages
  • point-to-point activation communication
  • pipeline bubbles

9. Pure Tensor Parallelism

CUDA_VISIBLE_DEVICES=0,1,2 \
uv run torchrun \
    --standalone \
    --nproc_per_node=3 \
    train_tensor_parallel.py

All tensor-parallel ranks process the same logical batch.

The MLP projection matrices are sharded:

fc1: column parallel
fc2: row parallel

Observe:

  • model tensor shards across GPUs
  • communication inside each MLP block
  • similar input data on all ranks
  • reduced parameter size for sharded linear layers
  • synchronization overhead at every layer

10. 2D Data Parallel + Tensor Parallel

Four GPUs with TP size 2:

CUDA_VISIBLE_DEVICES=0,1,2,3 \
uv run torchrun \
    --standalone \
    --nproc_per_node=4 \
    train_2d_dp_tp.py \
    --tp-size 2

Conceptual layout:

DP replica 0:
    ranks 0 and 1 form TP group 0

DP replica 1:
    ranks 2 and 3 form TP group 1

Within a TP group, ranks process the same batch. Across DP replicas, different batches are processed.

Observe:

  • tensor-parallel communication within each TP group
  • gradient synchronization across DP replicas
  • memory distribution across all GPUs
  • interaction between DP and TP communication

Useful Debug Environment Variables

Detailed Distributed Debugging

TORCH_DISTRIBUTED_DEBUG=DETAIL \
CUDA_VISIBLE_DEVICES=0,1 \
uv run torchrun \
    --standalone \
    --nproc_per_node=2 \
    train_ddp.py

NCCL Diagnostics

NCCL_DEBUG=INFO \
CUDA_VISIBLE_DEVICES=0,1 \
uv run torchrun \
    --standalone \
    --nproc_per_node=2 \
    train_ddp.py

For additional NCCL subsystem details:

NCCL_DEBUG=INFO \
NCCL_DEBUG_SUBSYS=INIT,COLL,P2P \
CUDA_VISIBLE_DEVICES=0,1,2 \
uv run torchrun \
    --standalone \
    --nproc_per_node=3 \
    train_pipeline_parallel.py

Synchronous CUDA Execution

CUDA_LAUNCH_BLOCKING=1 \
CUDA_VISIBLE_DEVICES=0 \
uv run python train_single.py

This is useful for debugging but may significantly slow execution.

Troubleshooting

NumPy Is Not Available

Error:

RuntimeError: Numpy is not available

Fix:

uv add numpy

Verify:

uv run python -c "import numpy; print(numpy.__version__)"

torchrun Terminates Other Ranks

When one rank fails, torchrun normally sends termination signals to the remaining ranks.

Example:

Sending process ... closing signal SIGTERM

The SIGTERM messages are usually consequences of the first rank failure. Locate the earliest traceback and fix that root cause.

CUDA Out of Memory

Reduce the batch size in the relevant script or through a supported command-line argument.

For mixed precision:

CUDA_VISIBLE_DEVICES=0 \
uv run python train_mixed_precision_single_gpu.py \
    --device 0 \
    --precision bf16 \
    --batch-size 16

For 2D DP + TP:

CUDA_VISIBLE_DEVICES=0,1,2,3 \
uv run torchrun \
    --standalone \
    --nproc_per_node=4 \
    train_2d_dp_tp.py \
    --tp-size 2 \
    --batch-size 16

Also check whether another process is occupying GPU memory:

nvidia-smi

Wrong Number of GPUs

Check visible GPUs:

CUDA_VISIBLE_DEVICES=0,1,2 \
uv run python -c "import torch; print(torch.cuda.device_count())"

For distributed runs, ensure:

number of CUDA_VISIBLE_DEVICES entries
=
--nproc_per_node

Port Already in Use

Use a different rendezvous port:

CUDA_VISIBLE_DEVICES=0,1 \
uv run torchrun \
    --standalone \
    --master-port=29601 \
    --nproc_per_node=2 \
    train_ddp.py

Recommended Learning Order

  1. Run train_single.py as the baseline.
  2. Compare FP32, FP16, and BF16 with train_mixed_precision_single_gpu.py.
  3. Compare train_dataparallel.py with train_ddp.py.
  4. Compare DDP with FSDP1 and FSDP2.
  5. Run manual model parallelism and observe sequential device activity.
  6. Run pipeline parallelism and observe microbatch overlap.
  7. Run pure tensor parallelism.
  8. Run the 2D DP + TP example.

For every experiment, keep this running in a separate terminal:

watch -n 1 nvidia-smi

Record:

  • GPU memory used
  • GPU utilization
  • power draw
  • step time
  • total runtime
  • effective batch size
  • number of GPUs
  • precision mode
  • any communication warnings

Project Layout

src/
  data.py       synthetic random-token dataset
  model.py      shared Transformer classifier
  runtime.py    CUDA and runtime helpers

notes/
  README.md     study-guide index
  ...           topic-specific training notes

train_single.py
train_mixed_precision_single_gpu.py
train_dataparallel.py
train_ddp.py
train_fsdp.py
train_fsdp2.py
train_model_parallel.py
train_pipeline_parallel.py
train_tensor_parallel.py
train_2d_dp_tp.py

main.py
pyproject.toml
uv.lock

Generated Artifacts

FSDP and FSDP2 examples may generate files such as:

fsdp_tiny_transformer_full_state.pt
fsdp2_tiny_transformer_full_state.pt
fsdp2_tiny_transformer_sharded_dcp/

These are training outputs and should normally remain ignored by Git.

Notes for GitHub

This repository is a learning and experimentation project rather than a formal benchmark suite.

The examples favor explicit code over heavy abstraction so that the following behaviors remain visible:

  • process creation
  • device placement
  • batch distribution
  • gradient synchronization
  • parameter sharding
  • activation movement
  • point-to-point communication
  • pipeline scheduling
  • tensor sharding
  • checkpoint formats

When reporting results, include the GPU model, number of GPUs, PyTorch version, CUDA version, precision, batch size, model size, and the command used.

About

A PyTorch-focused playbook for learning multi-GPU neural network training: single-GPU baselines, mixed precision, DataParallel, DDP, FSDP/FSDP2, model parallelism, pipeline parallelism, and tensor parallelism with visual notes, code examples, and communication patterns.

Topics

multi-gpu-trainingparallel-trainingpytorch

Resources

Readme
Activity

Stars

0 stars

Watchers

0 watching

Forks

0 forks
Report repository

Releases

Packages

Contributors

Languages