harsh040506/distributed-training

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.