distributed-training

name	distributed-training
description	This skill should be used when the user asks about "distributed training", "multi-GPU training", "data parallelism", "model parallelism", "pipeline parallelism", "tensor parallelism", "DDP", "FSDP", "ZeRO", "DeepSpeed", "Megatron-LM", "GPU utilization", "NCCL", "torchrun", "gradient communication", "checkpoint recovery", "spot instance preemption", "NVLink", "InfiniBand", "training throughput", or when a model is too large for a single GPU or training speed needs to be scaled.
version	1.0.0

Distributed Training — Parallelism Strategies and Cluster Management

Provides systematic guidance for scaling ML training across multiple GPUs, multiple nodes, and heterogeneous hardware configurations. Covers parallelism taxonomy, communication efficiency, fault tolerance, and GPU utilization optimization.

Parallelism Strategy Selection

Strategy	Model Fits on 1 GPU?	Reduces Memory?	Communication Overhead
Data Parallel (DDP)	Yes	No	Low (gradient sync)
FSDP	No	Yes (shards all states)	Medium (all-gather)
Tensor Parallel (TP)	No	Yes (shards weight matrices)	High (within-layer)
Pipeline Parallel (PP)	No	Yes (splits layers)	Medium (activation transfer)
3D Hybrid (TP+PP+DP)	No	Maximum	Highest

Quick heuristics:

• < 7B parameters, fits in VRAM → DDP with automatic mixed precision
• 7B–70B parameters → FSDP with FULL_SHARD strategy

• 70B parameters → 3D Hybrid parallelism (Megatron-LM / DeepSpeed)

Data Parallel (DDP) — Standard Multi-GPU

# Launch: torchrun --nproc_per_node=8 train.py
model = DistributedDataParallel(model, device_ids=[local_rank])
sampler = DistributedSampler(dataset)  # ensures no overlap across GPUs

All-reduce gradient synchronization happens automatically via NCCL ring-allreduce.

Bucket tuning: bucket_cap_mb=100 (default: 25MB) — larger buckets reduce communication rounds but increase memory. Tune based on model size and interconnect bandwidth.

FSDP — Fully Sharded Data Parallel

Shards model parameters, gradients, and optimizer states across all ranks:

• FULL_SHARD: maximum memory saving — parameter gathered only when needed
• SHARD_GRAD_OP: shard gradients/optimizer states only — faster but more memory
• Use transformer_auto_wrap_policy to wrap Transformer blocks for efficient sharding

FSDP enables training of 70B+ parameter models on 8× A100s (80GB) without offloading.

See references/parallelism-strategies.md for TP, PP, and 3D hybrid implementation patterns.

Communication Efficiency

Gradient compression (for bandwidth-limited clusters):

• PowerSGD: low-rank gradient compression, 10–100× compression ratio
• 1-bit Adam: quantize gradients to 1-bit; apply error feedback accumulation
• Use only when NCCL all-reduce is the bottleneck (profile first)

DeepSpeed ZeRO Optimizer:

• Stage 1: shard optimizer states only (4× memory reduction)
• Stage 2: + shard gradients (8× reduction)
• Stage 3: + shard parameters (linear reduction in N GPUs)
• CPU offload: ZeRO-Infinity offloads to CPU/NVMe (enabling trillion-parameter models)

Memory Optimization Techniques

1. Gradient checkpointing: discard activations during forward pass; recompute during backward. Saves O(√L) activations (L = number of layers). ~33% compute overhead, 60–80% VRAM savings.
2. Activation offloading: swap activations to CPU RAM when VRAM is full.
3. Flash Attention: I/O-aware attention algorithm — O(n) memory vs. O(n²) for standard attention.
4. torch.compile(): Dynamo + Inductor compilation, 1.5–2.5× throughput without code changes.

Fault Tolerance

• Checkpointing cadence: every 500 steps to shared storage (NFS or S3); retain 3 most recent.
• Checkpoint contents: weights, optimizer state, scheduler state, step count, RNG state, Python/library versions.
• Elastic training: torchrun --max_restarts=3 for automatic fault recovery.
• Spot instance resilience: catch SIGTERM (15s warning), flush checkpoint, exit cleanly, allow elasticity layer to restart.

See references/cluster-management.md for Kubernetes ML cluster setup, autoscaling, and multi-cloud training patterns.

GPU Utilization Targets

Utilization	Status	Primary Cause
≥ 85%	Excellent	—
70–85%	Acceptable	Minor data loading or scheduling overhead
50–70%	Poor	CPU/IO bottleneck or small batch size
< 50%	Critical	Synchronous data loading, GPU idle waiting

Profile with torch.profiler.profile() for the first 100 steps; identify the bottleneck category before optimizing.