Training Framework Ledger: Megatron-LM vs DeepSpeed Parallelism

PyTorch provides primitives (DDP/FSDP); large-scale training frameworks orchestrate them into systems — Megatron-LM excels at tensor/pipeline parallelism, DeepSpeed at ZeRO and kernel optimization. This article compares the two frameworks' parallel abstractions and selection logic, linking to Distributed Parallelism and PyTorch Distributed Engine.

1. Four-Layer Abstraction

Layer	Responsibility	Representative
Scheduling	data loading, steps, checkpoint, logging	Trainer/scripts
Parallelism	orchestrate DP/TP/PP/SP/ZeRO	Megatron, DeepSpeed
Communication	wrap All-Reduce/All-Gather/All-to-All	above NCCL
Kernel	op fusion, mixed precision, memory opt	FlashAttention, ZeRO-Offload

Takeaway: framework competition is mostly the "parallelism + kernel" abstraction — underneath are the same NCCL and op set.

2. Megatron-LM: Tensor & Pipeline Parallelism

Capability	Approach	Fits
Tensor parallel (TP)	split single-layer params column/row across cards	single layer doesn't fit one card
Pipeline parallel (PP)	slice by layers, microbatch pipeline	deep layers, big models
Sequence parallel (SP)	slice attention by sequence	long sequences
Virtual pipeline (VPP)	fine-grained slicing to reduce bubbles	PP bubble optimization

Key: TP is comm-dense (All-Reduce every layer) — must stay within one node/NVLink domain; PP is comm-sparse, can cross nodes (see Topology Ledger).

3. DeepSpeed: ZeRO & Kernel Optimization

Capability	Approach	Fits
ZeRO-1/2/3	shard states, memory-for-comm	memory-bound (see ZeRO Ledger)
Pipeline parallelism	layer slicing + load balance	big models
Offload	optimizer/weights to CPU/NVMe	extreme memory limits
Kernel optimization	op fusion, FlashAttention	throughput gains

Key: DeepSpeed's differentiation is ZeRO & Offload — complementary to Megatron's TP/PP, not competing.

4. Selection Matrix

Scenario	Recommendation	Why
≤32 cards, tight VRAM	DeepSpeed ZeRO-2/3	simple deployment, direct memory gains
70B+ won't fit one node	Megatron TP+PP	TP can't cross nodes; TP first, then PP
1000-card giant models	Megatron+DeepSpeed combo	TP×PP×DP×ZeRO four dimensions
Long-sequence training	Megatron SP + FlashAttention	sequence slicing + attention optimization

Takeaway: not either/or — production big models usually let Megatron handle TP/PP and DeepSpeed handle ZeRO; framework combination is the norm (e.g. Megatron-DeepSpeed fusion).

5. Reproducible Python: TP/PP Layout Constraints

def tp_pp_layout(gpus, tp, pp):
    """TP/PP legality: TP must divide card count; PP groups cross nodes"""
    if gpus % tp != 0:
        return "TP doesn't divide cards — illegal"
    n_nodes = gpus // tp // pp
    return f"{gpus} cards → TP{tp}×PP{pp}×DP{n_nodes}, TP group = {tp} cards"

tp_pp_layout(64, 8, 4)
# 64 cards → TP8×PP4×DP2: TP8 needs same node (8 cards/node), PP4 pipelines across nodes

6. Checklist

1. ≤32 cards: start with DeepSpeed ZeRO; don't jump to TP/PP;
2. 70B+: Megatron TP+PP — pin TP groups to one node, PP across nodes;
3. 1000-card: TP×PP×DP×ZeRO combo, validate efficiency per dimension;
4. Long sequences: add Megatron SP + FlashAttention;
5. Run Acceptance & Baselines after each layout change; watch throughput & comm ratio.

---

Organized and written by this site for learning purposes; formulas and examples are the site's own; specific versions and measurements are subject to on-site conditions.