Why Distribution Is Required
A 70B parameter model in FP16 needs roughly 140 GB just for weights, before optimizer states, gradients, and activations. No single consumer GPU holds that. Even smaller models benefit from distribution to shorten training time by processing larger effective batch sizes across many devices.
Distributed training splits work across GPUs, nodes, and datacenters while preserving mathematical equivalence (or controlled approximation) to single-device training.
Parallelism Dimensions
Data parallelism (DP)
Each GPU holds a full model copy and processes different micro-batches. Gradients are all-reduced and averaged before the optimizer step. Simple and widely used, but memory per GPU still holds the full model.
Tensor parallelism (TP)
Individual layers are sharded across GPUs. Attention and MLP matrix multiplies split along hidden or head dimensions. Requires fast NVLink-style interconnect because activations cross devices every layer.
Pipeline parallelism (PP)
Layers are grouped into stages on different GPUs. Micro-batches flow through the pipeline. Bubble time (idle GPUs) is reduced with techniques like 1F1B scheduling. Useful when model depth exceeds what TP alone can shard efficiently.
Expert parallelism (EP)
For MoE models, different expert FFNs live on different GPUs. Tokens route to remote experts with all-to-all communication. Load imbalance across experts is a first-class ops problem.
Large pretraining jobs combine 3D parallelism: DP + TP + PP, tuned to cluster topology.
Refer: Megatron-LM · DeepSpeed ZeRO
Memory Optimizations
ZeRO (Zero Redundancy Optimizer)
DeepSpeed ZeRO partitions optimizer states (Stage 1), gradients (Stage 2), and parameters (Stage 3) across data parallel ranks instead of replicating everything. Stage 3 enables training models far larger than per-GPU memory at the cost of more communication.
FSDP (Fully Sharded Data Parallel)
PyTorch-native sharding similar to ZeRO-3. Parameters are sharded across ranks and all-gathered just-in-time for each forward/backward layer. Popular for fine-tuning large models on moderate clusters.
from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
model = FSDP(
model,
auto_wrap_policy=transformer_auto_wrap_policy,
mixed_precision=MixedPrecision(param_dtype=torch.bfloat16),
device_id=torch.cuda.current_device(),
)Gradient checkpointing
Recomputes activations during backward instead of storing them, trading compute for memory. Essential for long sequences or large batches on memory-constrained GPUs.
Communication and Hardware
All-reduce bandwidth limits data parallel scaling. Tensor parallel groups should sit on the same node with NVLink. Cross-node training depends on InfiniBand or RoCE latency and topology-aware job schedulers.
Fault tolerance matters at scale: checkpoint to distributed storage frequently, use elastic training where supported, and design restarts so a single straggler node does not waste an entire week's compute.
Fine-Tuning vs Pretraining Distribution
Fine-tuning often uses LoRA/QLoRA on fewer GPUs: base weights frozen or quantized, only adapter matrices trained. QLoRA keeps the base model in 4-bit NF4 while adapters train in higher precision, democratizing 70B fine-tunes on single high-end GPUs or small pods.
Full fine-tuning of large models still needs FSDP multi-GPU setups similar to pretraining but for shorter durations and smaller datasets.
Frameworks and Operations
Common stack: PyTorch + DeepSpeed or FSDP + Megatron-LM patterns for TP/PP + NCCL for collective ops. Kubernetes, Slurm, or cloud ML platforms schedule multi-node jobs. Observability includes per-GPU utilization, gradient norm spikes, throughput (tokens/sec), and communication overhead ratios.
Even if you only fine-tune, understanding distributed training explains why certain open models ship only with specific parallel configs, why context length extensions are expensive, and how serving-time tensor parallel in inference mirrors training-time sharding.