Multi-GPU Playground

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

The repository uses a tiny synthetic Transformer classifier so the scripts stay focused on training mechanics: device placement, process groups, batch semantics, sharding, collectives, mixed precision, and checkpoint formats. The dataset is random by design, so the loss is not expected to meaningfully improve; for 10 random classes it should hover near log(10) ~= 2.30.

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 DCP sharded 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 plus 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.

Requirements

• Python 3.10+
• NVIDIA GPU(s) with a CUDA-capable PyTorch install
• uv for the commands shown above, or an equivalent virtual environment
• Multiple GPUs for the distributed examples; scripts that require a fixed count raise a clear error when the hardware is not available

Install dependencies:

uv sync

Sanity check the repo entrypoint:

uv run python main.py

Run Controls

Watch GPU Utilization

Basic one-shot GPU status:

nvidia-smi

Refresh every second:

watch -n 1 nvidia-smi

Watch only selected physical GPU IDs:

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

Compact live view with GPU utilization, memory, temperature, and power:

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

Select Which GPUs A Run Can See

Use CUDA_VISIBLE_DEVICES before uv run:

CUDA_VISIBLE_DEVICES=1 uv run python train_single.py

For stable numbering that matches PCI bus order:

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

Important: CUDA remaps visible GPUs to local IDs. With:

CUDA_VISIBLE_DEVICES=2,3

the program sees:

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

So scripts with a --device flag should usually use --device 0 after selecting one visible GPU:

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

Single-Process Examples

Run on one selected GPU:

CUDA_VISIBLE_DEVICES=0 uv run python train_single.py
CUDA_VISIBLE_DEVICES=0 uv run python train_mixed_precision_single_gpu.py --device 0 --precision fp32
CUDA_VISIBLE_DEVICES=0 uv run python train_mixed_precision_single_gpu.py --device 0 --precision fp16
CUDA_VISIBLE_DEVICES=0 uv run python train_mixed_precision_single_gpu.py --device 0 --precision bf16

Run nn.DataParallel on a selected subset:

CUDA_VISIBLE_DEVICES=0,1 uv run python train_dataparallel.py

Run manual model parallelism on exactly 3 visible GPUs:

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

Distributed Examples With torchrun

For torchrun, set --nproc_per_node equal to the number of visible GPUs you want to use.

DDP on two selected GPUs:

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

FSDP/FSDP2 on three selected GPUs:

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

Pipeline parallelism expects exactly 3 ranks:

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

Pure tensor parallelism on three selected GPUs:

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

2D DP + TP on four GPUs with two TP ranks per DP replica:

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

Useful Debug Environment Variables

Show distributed setup details:

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

Show NCCL communication diagnostics:

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

Make CUDA operations synchronous for debugging stack traces:

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

Recommended Learning Order

1. Run train_single.py to anchor the baseline loop.
2. Run train_mixed_precision_single_gpu.py with fp32, fp16, and bf16 to compare memory behavior.
3. Compare train_dataparallel.py with train_ddp.py to see why DDP is preferred.
4. Run train_fsdp.py and train_fsdp2.py to study model-state sharding and checkpoint formats.
5. Explore model parallel, pipeline parallel, tensor parallel, and 2D DP+TP once the data-parallel examples are clear. For train_2d_dp_tp.py, pass --tp-size when you want multiple DP replicas on one node; otherwise it defaults to TP-within-node for multi-node layouts.

Project Layout

src/
  data.py       synthetic random token dataset
  model.py      tiny Transformer classifier shared by most demos
  runtime.py    small CUDA/runtime helpers

notes/          split interview notes by training paradigm
train_*.py      focused training examples
main.py         quick index of common commands

Generated Artifacts

FSDP/FSDP2 examples can create .pt files and distributed checkpoint directories such as fsdp2_tiny_transformer_sharded_dcp/. These are intentionally ignored by git because they are generated training outputs, not source files.

Notes For GitHub

This is a learning repo, not a benchmark suite. The examples favor explicitness over abstraction so the distributed-training mechanics remain visible during interview prep.