FSDP, DeepSpeed, Megatron: Choosing the Right Stack
Three frameworks dominate large-scale training. They overlap but specialise differently. The wrong choice means months of fighting your training stack instead of training your model.
The three options
All three implement parallelism for training large transformers. They differ in which parallelisms they support, how they expose them, and what the operational complexity looks like.
FSDP (PyTorch native)
Fully Sharded Data Parallel. Shards model parameters, gradients, and optimiser state across data-parallel ranks. The PyTorch-native ZeRO equivalent.
Strengths: in PyTorch core (no extra dependencies), simple API, good for medium-scale (8-128 GPUs). FSDP2 (2024+) is significantly faster than FSDP1.
Weaknesses: limited tensor and pipeline parallelism support. For models > 30B you usually need to combine FSDP with another framework or move to Megatron.
DeepSpeed
Microsoft’s training framework. Originally famous for ZeRO (now matched by FSDP), continues to lead on offloading (CPU, NVMe), MoE training, and ergonomics for medium-scale teams.
Strengths: best CPU/NVMe offloading (train models that don’t fit in GPU memory at all, slowly), strong MoE support, friendly config-driven setup, decent multi-node experience.
Weaknesses: API surface is wider and sometimes confusing, slower than Megatron at the largest scales, less “production NVIDIA-blessed”.
Megatron-LM
NVIDIA’s framework for very-large-scale transformer training. The reference implementation for tensor and pipeline parallelism.
Strengths: fastest absolute throughput at the largest scales, deeply optimised CUDA kernels, the choice for training-from-scratch frontier models.
Weaknesses: opinionated codebase, steeper learning curve, NVIDIA-only, more research-grade ergonomics than enterprise polish.
Picking one
- Fine-tuning under 30B parameters, <= 64 GPUs: FSDP2.
- Fine-tuning with offloading or MoE, mixed cluster: DeepSpeed.
- Pretraining 30B+ from scratch, large NVIDIA cluster: Megatron-LM, often combined with NeMo.
- Mature production stack at > 100B: NeMo Megatron or a research-derived combo of all three.
Most teams shouldn’t train from scratch. Most fine-tuning fits in FSDP. The hard choices are reserved for the small number of teams pretraining frontier models, where NVIDIA support, hardware topology, and research talent dominate.