Personalize this lesson
Adapt explanations and teaching visuals to your background and preferred voice.
Scaling a language-model training job from eight GPUs to 64 can make each update slower when the global batch stays fixed. Some ranks wait for a layer result; others wait for gradients or a routed expert. More hardware can expose communication and scheduling costs instead of shortening the step, even when nothing runs out of memory.
Which part of the work should each GPU own? The last chapter treated all-reduce, reduce-scatter, all-gather, all-to-all, and send/receive as NCCL contracts: matching ranks, matching dtypes, matching order. This chapter follows that question through a real training mesh.
Megatron-LM is NVIDIA's reference training stack for that job. Megatron Core is the library underneath it: parallel layers, process groups, schedules, optimizers, checkpoint utilities, and model definitions. The current repository splits the two on purpose. Megatron-LM ships recipes a research team can run. Core ships building blocks a framework author can import.[1]
Source baseline: commit 8971d19dda0f2bda87783353dde93a5134422a96, checked on September 2, 2026.[2] The linked latest guides can move independently. Python examples below execute CPU arithmetic, not Megatron kernels or NCCL; hardware experiments are labeled separately. No GPU throughput or convergence result was measured here.
Treat a process rank as a coordinate, not just rank=27. Its axes include data replica, context lane, pipeline stage, and tensor shard. Megatron builds NCCL groups from those coordinates, then each forward and backward operation picks the smallest group that owns its data. Scale comes from factoring one world into several grids, not from making every GPU talk to every other GPU.
Split the dimensions that hurt
Start with the resource that is blocking progress. Parameters, gradients, optimizer state, and activations can exceed one GPU's memory. Long sequences can make attention intermediates too large even when weights fit. A sparse mixture-of-experts (MoE) model stores many experts but activates only a few per token. The same model can therefore hit a memory wall, a link wall, or an uneven-work wall.
Megatron splits different dimensions for different reasons:
| Pressure | Axis or mechanism | What gets partitioned | Typical NCCL op |
|---|---|---|---|
| Wide matrix doesn't fit | Tensor parallelism (TP) | A layer's weight and activation channels | All-reduce or all-gather |
| Stack is too deep | Pipeline parallelism (PP) | Contiguous layer ranges | Point-to-point activation send/receive |
| Batch needs throughput | Data parallelism (DP) | Independent samples or replicas | Gradient reduce or reduce-scatter |
| Prompt is too long | Context parallelism (CP) | Sequence chunks and attention context | Peer-to-peer or collective KV exchange |
| MoE experts are spread out | Expert parallelism (EP) | Expert weights and routed tokens | All-to-all dispatch and combine |
| Activations fill memory | Selective recomputation | Saved intermediates | Extra forward compute, less storage |
Read the table as a diagnosis map, not a menu of interchangeable flags. Each axis owns a different piece of state, so the same extra GPU can remove one bottleneck and add another.
Follow one sample through a candidate mesh before choosing. Data parallelism keeps its computation on one replica. Tensor parallelism makes its layer arithmetic cross tensor shards. Pipeline parallelism sends its activation from one layer range to the next. Context parallelism keeps only part of its sequence local. Expert parallelism may move its routed token to another rank and back. That trace predicts where the next wait can appear.[3]
The first budget is memory and the second is communication. TP lowers per-layer memory but adds collectives. DP raises independent work in flight without making one sample's path shorter. PP can make a model fit across nodes, but it inserts stage-to-stage latency. CP reduces sequence-local activation pressure while exchanging remote context. EP spreads expert weights while paying for token movement. Choose the smallest set of axes that removes the observed wall, then measure the new wait.
Why doesn't adding data replicas directly shorten one sample's forward pass?
Answer
Each sample still traverses one model-parallel replica. DP can improve aggregate training throughput by processing different samples concurrently, but it adds gradient synchronization and doesn't split that sample's layer arithmetic further.
Two projects in one repository, plus a bridge
Once the bottleneck has a name, the repository map tells you where to inspect it. A Megatron-LM clone contains two layers in one tree, plus a third project next door.
Megatron Core
megatron/core/ contains the parts a framework can import: transformer modules, tensor-parallel layers, pipeline schedules, distributed data parallel buffers, optimizers, datasets, model-specific builders, inference utilities, and distributed checkpointing. megatron/core/parallel_state.py owns process-group construction. megatron/core/tensor_parallel/layers.py implements sharded linear and embedding layers. megatron/core/pipeline_parallel/schedules.py sequences microbatches across stages. megatron/core/distributed/param_and_grad_buffer.py manages contiguous gradient and parameter buffers that can overlap communication with compute.
Core is an API surface, not one opinionated training command. Its configuration objects and process groups let a custom trainer choose a pipeline schedule, optimizer, precision recipe, checkpoint format, or model architecture while reusing tested kernels and collectives. Inspect Core when the question is "which ranks own this tensor, buffer, or stage?" The answer depends on the mesh a framework builds.[1]
In an existing Megatron-LM checkout containing the pinned commit, these read-only commands locate the two implementations used below. They don't import Megatron, initialize CUDA, or install dependencies. git show reports an invalid object if that commit hasn't been fetched.
1MEGATRON_REV=8971d19dda0f2bda87783353dde93a5134422a96
2git show "${MEGATRON_REV}:megatron/core/parallel_state.py" | rg -n 'class RankGenerator|def initialize_model_parallel'
3git show "${MEGATRON_REV}:megatron/core/tensor_parallel/layers.py" | rg -n 'class ColumnParallelLinear|class RowParallelLinear'Megatron-LM reference scripts
Top-level training scripts and examples/ turn Core into repeatable experiments. pretrain_gpt.py still lives at the repository root. Model providers configure tokenizers, data loaders, schedules, logging, checkpoint intervals, and command-line arguments. Tests cover rank-group construction, tensor-parallel layers, pipeline schedules, MoE dispatch, precision recipes, and distributed checkpoint round trips. Use these scripts to reproduce a training setup or isolate a mesh change; they aren't a drop-in serving API.
Megatron Bridge
Megatron Bridge is a separate PyTorch-native library in the NVIDIA NeMo family (NVIDIA-NeMo/Megatron-Bridge). It's the conversion, verification, and recipe layer between Hugging Face checkpoints and Megatron Core layouts, including parallelism-aware import and export.[1][4] A training run can start from an ecosystem checkpoint, use Megatron's sharded layout, then export a portable checkpoint for evaluation or serving. Conversion isn't a rename: tensor and pipeline shards must map to global tensors, and MoE expert placement must be preserved.
The split makes the debugging boundary visible. Core owns mesh and kernel behavior, LM owns the runnable recipe, and Bridge owns checkpoint representation. Align versions of Megatron Core, Transformer Engine, CUDA, Bridge, and checkpoint metadata. A source checkout can be ahead of the latest release, so record the commit or package version with each experiment.
Rank groups: factor world size first
Make the grid concrete with a dense run of 64 GPUs using TP=4, PP=2, CP=2, and DP=4. For the pinned source's dense grid, with additional generalized-tensor rematerialization factors left at their default of one:
Thus
Don't multiply EP into that equation while retaining the same dense DP value. Expert layers use a second factorization of the same ranks, introduced below, not another independent copy of the world.[2]
Before looking at any rank IDs, predict the ownership boundary. TP, PP, and CP describe one model copy; DP repeats that model-parallel mesh for independent data. Rank numbering is an implementation detail, while group membership is the contract. initialize_model_parallel builds rank generators from a configurable order. The default is tp-cp-ep-dp-pp, and initialization checks that world size is divisible by the model-parallel product.[3]
Start at rank 3 and predict its frequent peers from that order. TP varies fastest, so tensor peers get consecutive global IDs. PP varies last, so a pipeline peer can be far away in the global numbering. The nested visual reads the same 64-GPU grid outside-in: DP replicas, then CP lanes, then PP stages, then four TP cells.[3]
This standard-library snippet rebuilds that numbering. It doesn't create NCCL communicators. It checks that rank 3's TP group is [0, 1, 2, 3] and its PP group is [3, 35]. The size-one EP coordinate is retained only to match the source's default order.
1ORDER = ("tp", "cp", "ep", "dp", "pp")
2SIZES = {"tp": 4, "cp": 2, "ep": 1, "dp": 4, "pp": 2}
3WORLD = 64
4
5def coords(rank: int) -> dict[str, int]:
6 if type(rank) is not int or not 0 <= rank < WORLD:
7 raise ValueError("rank must be an integer inside this world")
8 remaining = rank
9 out: dict[str, int] = {}
10 for name in ORDER:
11 out[name] = remaining % SIZES[name]
12 remaining //= SIZES[name]
13 if remaining:
14 raise ValueError(f"rank {rank} exceeds world size")
15 return out
16
17def rank_of(c: dict[str, int]) -> int:
18 if set(c) != set(ORDER) or any(type(c[name]) is not int or not 0 <= c[name] < SIZES[name] for name in ORDER):
19 raise ValueError("invalid coordinate")
20 rank = 0
21 stride = 1
22 for name in ORDER:
23 rank += c[name] * stride
24 stride *= SIZES[name]
25 return rank
26
27def group(rank: int, axis: str) -> list[int]:
28 if axis not in ORDER:
29 raise ValueError("unknown axis")
30 base = coords(rank)
31 members = []
32 for index in range(SIZES[axis]):
33 member = dict(base)
34 member[axis] = index
35 members.append(rank_of(member))
36 return members
37
38world = 64
39assert SIZES["tp"] * SIZES["pp"] * SIZES["cp"] * SIZES["ep"] * SIZES["dp"] == world
40assert all(rank_of(coords(rank)) == rank for rank in range(world))
41assert coords(3) == {"tp": 3, "cp": 0, "ep": 0, "dp": 0, "pp": 0}
42assert group(3, "tp") == [0, 1, 2, 3]
43assert group(3, "pp") == [3, 35]
44print("rank 3", coords(3))
45print("TP group", group(3, "tp"))
46print("PP group", group(3, "pp"))
47print("DP group", group(3, "dp"))1rank 3 {'tp': 3, 'cp': 0, 'ep': 0, 'dp': 0, 'pp': 0}
2TP group [0, 1, 2, 3]
3PP group [3, 35]
4DP group [3, 11, 19, 27]The code checks rank arithmetic for one dense mesh, not communication or an MoE launch. Numeric neighbors aren't necessarily physical neighbors: the launcher still maps ranks to devices and nodes. Megatron-FSDP (Fully Sharded Data Parallel) changes state ownership within the configured data-parallel domain; it isn't another factor in . With CP, dense weights also repeat across context peers, so gradient synchronization and optimizer sharding can use the combined DP-with-CP group.[2]
A 64-GPU dense launch sets TP=4, PP=2, and CP=2. What DP size keeps every rank assigned?
Answer
DP must be 4 because 64 / (4 × 2 × 2) = 4. A non-divisible product should fail during process-group initialization rather than silently duplicating or dropping ranks.
What each parallel axis does
Tensor parallelism: split width
Before naming the axis, predict from a tiny MLP. If hidden width H=8, intermediate width I=16, and TP=4, each rank should own four intermediate channels. The first projection can produce those channels locally; the second projection must combine four partial output vectors. That is tensor parallelism (TP): split a layer's matrix so each GPU computes a shard.
TP cuts a matrix along a dimension that keeps local matrix multiplications (GEMMs) large enough to use the GPU. For a column-parallel first MLP projection, A with shape [H, I] is split into TP column shards of shape [H, I/TP]. Many dense GPT-style MLPs use I = 4H. Each rank computes X A_i, applies GeLU locally, and passes its activation shard to a row-parallel second projection. The second projection's rows line up with those local shards, so partial outputs can be summed once at the end.[5]
The running numbers make that split visible: H=8, I=16, TP=4. Each rank stores [8,4] for the up projection and [4,8] for the down projection. Use one token, so X has shape [1,8]. Rank i computes Y_i = GeLU(X A_i) with shape [1,4], then Z_i = Y_i B_i with shape [1,8]. The layer output is Z = Σ_i Z_i. GeLU never needs an all-gather because it acts independently on each channel.
Those are mathematical X @ A shapes. The pinned PyTorch implementation stores linear weights in [out_features, in_features] order, so ColumnParallelLinear.weight is [4,8] and RowParallelLinear.weight is [8,4] here. A transpose in the source isn't a different partition.[6]
The next snippet uses the tanh approximation to GeLU in standard-library Python. Serial and sharded sums agree within a tolerance, not necessarily bit for bit. Megatron's fused kernels can round differently. The example omits biases, dropout, sequence parallelism, and backward.
1import math
2
3H, I, TP = 8, 16, 4
4
5def gelu(x: float) -> float:
6 return 0.5 * x * (1.0 + math.tanh(math.sqrt(2.0 / math.pi) * (x + 0.044715 * x**3)))
7
8def matmul(left: list[list[float]], right: list[list[float]]) -> list[list[float]]:
9 for matrix in (left, right):
10 if (not matrix or not matrix[0] or any(len(row) != len(matrix[0]) for row in matrix)
11 or any(not math.isfinite(x) for row in matrix for x in row)):
12 raise ValueError("matrices must be nonempty, rectangular, and finite")
13 if len(left[0]) != len(right):
14 raise ValueError("inner dimensions must match")
15 inner = len(left[0])
16 cols = len(right[0])
17 out = [[0.0] * cols for _ in left]
18 for i, row in enumerate(left):
19 for j in range(cols):
20 out[i][j] = sum(row[k] * right[k][j] for k in range(inner))
21 return out
22
23def apply_gelu(matrix: list[list[float]]) -> list[list[float]]:
24 return [[gelu(value) for value in row] for row in matrix]
25
26up = [[((row + 1) * (col + 3) % 7) / 10.0 for col in range(I)] for row in range(H)]
27down = [[((row + 2) * (col + 1) % 5) / 10.0 for col in range(H)] for row in range(I)]
28hidden = [[((col % 5) - 2) / 4.0 for col in range(H)]]
29
30serial = matmul(apply_gelu(matmul(hidden, up)), down)
31
32if I % TP:
33 raise ValueError("intermediate width must divide by TP")
34width = I // TP
35partials = []
36for shard in range(TP):
37 up_shard = [row[shard * width : (shard + 1) * width] for row in up]
38 down_shard = down[shard * width : (shard + 1) * width]
39 local = apply_gelu(matmul(hidden, up_shard))
40 partials.append(matmul(local, down_shard))
41
42reduced = [sum(partials[shard][0][col] for shard in range(TP)) for col in range(H)]
43assert all(abs(left - right) < 1e-12 for left, right in zip(serial[0], reduced))
44print("Z", [round(value, 4) for value in reduced])1Z [-0.1782, -0.2062, -0.1582, -0.1862, 0.0, -0.1782, -0.2062, -0.1582]In backward, each column shard contributes part of the gradient with respect to the replicated input X; those contributions must be summed. A bias on the second projection is added after summing partial outputs, otherwise TP copies would count it repeatedly. Sequence-parallel paths replace some replicated-activation reductions with reduce-scatter and all-gather operations, so the one-forward-all-reduce diagram isn't a complete training communication count.[5]
Attention uses the same idea with heads. Query, key, and value projections are partitioned, local attention runs on head shards, and the output projection combines results. Check the chosen attention implementation's hidden-width, query-head, and query-group divisibility rules. Grouped-query attention can make key/value partitioning the limiting constraint; specialized replication paths shouldn't be assumed from ordinary TP arithmetic.
That local arithmetic has a price. Each rank stores roughly 1/TP of a sharded weight, but every layer may issue all-reduce, all-gather, or reduce-scatter collectives. Predict what changes if TP grows from 2 to 8 while the batch and model stay fixed: shards get smaller, the group gets larger, and the collective may cross a node boundary. Measure local GEMM time and TP collective time on the same batch before increasing TP again. The exact count changes with sequence parallelism, overlap settings, and kernel path, so keep TP inside high-bandwidth links when possible.[7]
Pipeline parallelism: split depth
PP assigns contiguous transformer layers to stages. In PP=2, stage 0 owns early layers and sends its hidden activation to stage 1, which owns later layers. During training, microbatches flow through a schedule such as 1F1B (one forward pass paired with one backward pass) so stages stay busy. Virtual pipeline stages interleave smaller layer chunks on each rank and can reduce bubbles when layer counts and batch sizes permit.
Run the prediction with one microbatch: stage 1 must sit idle until stage 0 sends its first activation. Add enough microbatches and the stages can overlap, but the first and last waves still leave a bubble. PP lowers per-rank weight memory without all-reducing every layer across nodes, yet it adds activation send/receive latency. A topology with four TP ranks inside each node and two PP stages across nodes can beat a single eight-way TP group over a slower fabric. In the 64-GPU lattice, that's why rank 3's pipeline peer is 35 rather than 4.
For an idealized non-interleaved schedule with equally timed stages, p stages and m microbatches have bubble fraction approximately (p-1)/(m+p-1). At p=2, one microbatch wastes half the slots; eight reduce that fraction to 1/9. This ignores communication, stage imbalance, and kernel overlap. Count microbatches from independent samples, not TP or CP copies of a sample.[8]
For the 64-rank example, global batch 64 with microbatch size 2 and DP=4 gives eight microbatches per replica. Fixed sequence length 2,048 means 131,072 token positions per update before padding or loss masks. The function rejects a batch that would require a fractional microbatch; it doesn't model ramp-up or variable-length packing.
1def batch_budget(global_batch: int, microbatch: int, dp: int, sequence: int, pp: int) -> tuple[int, int, float]:
2 if any(type(x) is not int or x <= 0 for x in (global_batch, microbatch, dp, sequence, pp)):
3 raise ValueError("counts must be positive integers")
4 if global_batch % (microbatch * dp):
5 raise ValueError("global batch must divide by microbatch times DP")
6 microbatches = global_batch // (microbatch * dp)
7 bubble = (pp - 1) / (microbatches + pp - 1)
8 return microbatches, global_batch * sequence, bubble
9
10m, tokens, bubble = batch_budget(64, 2, 4, 2048, 2)
11print("microbatches per replica:", m)
12print("token positions per update:", tokens)
13print("idealized bubble:", round(bubble, 3))1microbatches per replica: 8
2token positions per update: 131072
3idealized bubble: 0.111Data parallelism: split samples, then maybe shard state
DP repeats the model-parallel mesh for independent samples, then combines gradients across replicas. Standard DDP keeps a copy of each rank's assigned model shard across its DP peers; TP and PP may already have partitioned the full model. Megatron's distributed optimizer changes optimizer-state ownership while keeping those copies synchronized. Increasing DP at fixed global batch also reduces microbatches per replica and can worsen the PP bubble.
Megatron-FSDP goes further on that replication domain. With --use-megatron-fsdp, --data-parallel-sharding-strategy selects optim, optim_grads, or optim_grads_params. These shard optimizer state, then gradients, then parameters: Zero Redundancy Optimizer (ZeRO)-style ownership, not extra ranks. These are configuration fragments, not complete launch commands. Hierarchical data parallelism (--num-distributed-optimizer-instances) can partition the synchronization domain into smaller sharding domains plus outer replication.[3]
The next chapter, DeepSpeed, treats that ownership ladder as its primary control surface. Here it's one option on Megatron's DP group after the model-parallel mesh is already chosen.
Context and sequence parallelism: split token work
Sequence parallelism is a companion to TP, not another way to split a full attention calculation. It partitions sequence-dimension activation work in operations such as LayerNorm and dropout, often using reduce-scatter and all-gather variants to reduce activation memory. It doesn't mean each rank can run full attention with only its local tokens. Current docs recommend enabling it whenever TP is on.[3]
Now predict with an 8K-token sequence and CP=2: each rank can hold part of the sequence, but attention still needs remote keys and values. Context parallelism exchanges that context while preserving causal masks and global positions; implementations may distribute balanced noncontiguous chunks rather than one contiguous half. Communication can use point-to-point (--cp-comm-type p2p), all-gather, all-to-all, or hierarchical paths. The guide's 8K+ recommendation isn't a minimum correctness requirement. CP duplicates weights, so their gradients must include contributions from all context peers.[3][2] Dynamic CP changes the degree for variable-length batches, but its admissible shapes and performance need separate testing.[1]
Expert parallelism: split sparse capacity
An MoE layer has a router that chooses top-k experts for each token. Predict from a skewed route histogram: if most tokens choose one expert, the total token count hasn't changed, but one rank now owns a queue while others wait. Expert parallelism (EP) places different experts on different GPUs. The dispatcher permutes tokens, sends each token to its expert owners, runs grouped expert GEMMs, then sends outputs back and unpermutes them. Current dispatcher code exposes all-gather and all-to-all paths, plus Flex backends that can specialize communication, including a DeepEP path on supported NVIDIA GPUs.[9]
That prediction names EP's failure surface. An all-to-all can move every token to a different rank, while router skew produces hot experts, uneven work, padding, and queueing. Load-balancing loss, capacity policy, dropless dispatch, grouped GEMM, and overlapping communication with expert compute address different parts of that problem. They don't remove the need to inspect per-expert token counts and all-to-all timing.[9]
Why is a large TP value often a poor choice for a small MoE expert?
Answer
Each expert's matrix can become too narrow for efficient GEMMs, while every expert token still pays TP collectives. Keep attention and expert grids independent when the model and hardware support expert tensor parallelism or parallel folding.
MoE parallel folding
EP exposes a design mismatch: attention wants wide, well-connected tensor shards, while a small expert wants a large local matrix and little tensor communication. Traditional layouts constrain EP to the data-parallel domain, so both layers inherit the same grid. Predict the result of applying TP=4 to a narrow expert: its local GEMM shrinks while each token still pays the extra collective.
Megatron Core's MoE parallel folding decouples those grids. Attention can use TP × CP × DP × PP, while MoE can use ETP × EP × EDP × PP. Folding can break the old EP ≤ DP ceiling, reduce minimum GPU counts when CP and EP share ranks, and keep expert communication inside a high-bandwidth island.[9]
The official example is a 256-GPU attention layout, not the 64-GPU dense lattice above:
| Layout | Attention mesh | Expert mesh | What changes |
|---|---|---|---|
| Traditional | TP=4, CP=2, DP=8, PP=4 | EP ≤ DP = 8 | Expert grid is trapped inside DP |
| Folding | same attention mesh | ETP=1, EP=64, EDP=1 | EP can grow 8× without growing world size |
Both rows describe the same world: attention uses 4 × 2 × 8 × 4 = 256, while folded experts use 1 × 64 × 1 × 4 = 256. EDP means copies of a given expert shard, not the attention layer's DP degree. PP remains four in both grids. EP=64 requires a compatible expert count and rank placement; it doesn't mean 64 experts fit inside an ordinary eight-GPU NVLink domain. The logical factorization doesn't prove the physical communication path is fast.
The table's freedom has a verification cost. Every group needs matching tensor shapes, routing metadata, optimizer ownership, and checkpoint keys. A topology that looks good on paper can fail if the launch order assigns experts across a congested link or if a checkpoint was saved with a different expert grid. Validate expert counts, one routed batch, and a tiny synthetic MoE run before scheduling a long pretraining job.
Distributed optimizer: trade replicated state for collectives
The mesh decides where computation runs, but it doesn't decide who keeps every training byte. For conventional Adam with BF16 model parameters and FP32 main gradients, the guide's ledger is 18 bytes per parameter before activations and temporary buffers. The distributed optimizer changes it to bytes, where is the optimizer sharding group's size.[10]
The six replicated bytes are two for a BF16 parameter and four for its main gradient buffer. The twelve shardable bytes are an FP32 master parameter and two FP32 Adam moments. These are bytes per parameter already assigned to this TP/PP shard, not necessarily per parameter of the entire model.
At , that gives nine bytes. In the running dense mesh, however, DP=4 and CP=2; without hierarchical partitioning, the pinned implementation's DP-with-CP domain has eight ranks, so this ledger gives 7.5 bytes, not nine. Expert parameters can use a different domain. Inspect the actual optimizer group before substituting a number for .[2]
Trace one parameter's ownership across a step:
- Backward fills a gradient buffer on each DP rank.
reduce-scattercombines gradients and gives each rank its shard, with scaling matched to the loss normalization.- Each rank updates its local FP32 master parameters and optimizer moments.
all-gatherrebuilds the BF16 parameter buffer for the next forward.
The parameter and gradient buffers are contiguous so communication can start as buckets become ready. Megatron can overlap gradient reduce-scatter with backward and parameter all-gather with forward when the configuration and backend support it. Overlap hides latency only when compute lasts long enough and the network isn't already saturated. If a larger DP value saves memory but lengthens the reduce-scatter tail, the memory win isn't a throughput win.
Sharding group size d | Distributed bytes per BF16 parameter, FP32 grad | What still costs memory |
|---|---|---|
| 1 | 18 | Full optimizer and gradient state |
| 2 | 12 | Half of optimizer-owned state per rank |
| 4 | 9 | One quarter of sharded state plus model buffers |
| 8 | 7.5 | Diminishing state savings, same activations |
Those numbers are theoretical accounting, not a capacity guarantee. Activation buffers, temporary all-gather storage, communication buckets, embeddings, MoE routing buffers, and framework overhead still need headroom. Use the ledger to choose a candidate DP degree, then check the per-rank high-water mark and step-time tail. A larger DP value can lower throughput even while it saves memory.
This CPU example isolates the ownership change with SGD, not Adam or NCCL. Two replicas each contribute an equally weighted gradient; the mean gradient is partitioned, updated locally, and gathered. Averaging is explicit here. CP contributions and variable valid-token counts need normalization appropriate to the actual loss, not an extra blind division by every group dimension.
1import math
2
3def sharded_sgd(parameters: list[float], gradients: list[list[float]], lr: float) -> list[float]:
4 d, n = len(gradients), len(parameters)
5 if not d or not n or n % d:
6 raise ValueError("this equal-shard example requires a nonempty divisible buffer")
7 if any(len(g) != n for g in gradients):
8 raise ValueError("every replica must contribute the same gradient shape")
9 if not math.isfinite(lr) or lr < 0 or any(not math.isfinite(x) for row in [parameters, *gradients] for x in row):
10 raise ValueError("parameters, gradients, and learning rate must be finite")
11 reduced = [sum(g[i] for g in gradients) / d for i in range(n)]
12 width = n // d
13 owned = [[parameters[i] - lr * reduced[i] for i in range(rank * width, (rank + 1) * width)]
14 for rank in range(d)]
15 return [value for shard in owned for value in shard]
16
17parameters = [1.0, 2.0, 3.0, 4.0]
18gradients = [[1.0, 3.0, 5.0, 7.0], [3.0, 5.0, 7.0, 9.0]]
19print("updated:", sharded_sgd(parameters, gradients, 0.1))
20print("bytes per parameter at d=8:", 6 + 12 / 8)1updated: [0.8, 1.6, 2.4, 3.2]
2bytes per parameter at d=8: 7.5What does a distributed optimizer reduce-scatter, and what does it all-gather?
Answer
It reduce-scatters gradients so each DP rank owns a fully reduced gradient slice, updates its local optimizer and FP32 master slice, then all-gathers updated model parameters before the next forward.
Checkpoints, resharding, and recovery
At step 200,000, a node can fail after the model has reached a useful loss. Restarting from weights alone would lose optimizer history, random-number-generator state, and data position. Large runs checkpoint optimizer moments, data-loader position, iteration counters, and parallel-layout metadata alongside model weights. A checkpoint that can resume after a node failure is part of the training system, not a final export step.
Megatron's distributed checkpointing stores global tensor metadata and rank-local shards. Model tensors, optimizer state, and data-loader state don't necessarily have the same resharding support. A successful weight load alone doesn't prove a full resume. Test the requested topology change with the exact model, checkpoint format, and software versions; compare an optimizer update and data position as well as a forward pass. Keep launch arguments and the Git commit beside the checkpoint.
The Core 0.18.0 optimizer API distinguishes dp_reshardable, whose buffer representation supports DP changes, from fully_reshardable, whose parameter-shaped representation permits model-parallel resharding. That distinction concerns optimizer serialization, not a guarantee that arbitrary architectures or data-loader states can migrate.[11] Bridge's Energon notes provide a concrete limitation: its saved stream state is keyed by pure DP rank, and a missing per-rank state file fails restoration after some DP changes.[4]
Bridge makes the ecosystem boundary explicit. Hugging Face uses dense tensors and its own key names. Megatron uses tensor and pipeline shards, tied embeddings, optimizer partitions, and model-specific state. Bridge maps between those representations and can validate conversion. Exporting a model for inference may require gathering or converting shards, while resuming training should preserve distributed state instead of flattening everything onto one host. The Bridge docs also separate weight-only initialization from a full resume that restores optimizer, RNG, data-loader, and scheduler state.[4]
If a restart begins at the wrong iteration, loads only weights, or fails on one expert shard, inspect the state inventory before rerunning. An asynchronous save request isn't a completed checkpoint: only publish it for recovery after all required shards and metadata finish. Checksums can detect corruption but don't establish compatibility. Load only trusted checkpoints because some formats deserialize executable Python objects.[12] Then run the fixed-batch and next-update smoke test before accepting a long resume.
Precision: FP16, BF16, FP8, and FP4
Megatron Core integrates Transformer Engine for mixed precision. Common baselines use FP16 or BF16. FP8 recipes can increase matrix throughput and reduce activation bandwidth on supported NVIDIA GPUs, but scale tracking, amax history, accumulation precision, and kernel support affect convergence. Newer hardware also supports FP4 and NVFP4 paths for quantized weights or activations; these paths aren't drop-in BF16 replacements on every GPU.[13]
Use low precision where it matches the hardware and model recipe. Keep optimizer masters and sensitive reductions in higher precision when the implementation expects it. Compare loss curves, gradient norms, validation slices, and checkpoint reload behavior against a BF16 baseline. A faster kernel that changes overflow behavior can cost more than it saves, so precision changes need the same fixed-batch correctness gate as topology changes.
The repository README still reports up to 47% model FLOPs utilization (MFU) for specific H100 benchmarks: vocabulary 131,072, sequence length 4,096, overlap flags enabled, and a no-convergence measurement. Weak scaling on that study moves from about 41% MFU on smaller models to 47-48% on the largest; strong scaling a GPT-3-class model from 96 to 4,608 H100 GPUs drops MFU from 47% to 42% as communication becomes more exposed.[1] Treat those numbers as source-scoped measurements, not a promise for every model. Current recipe notes also live in the Megatron Bridge performance guide.[14]
Project identity
Performance numbers are easier to interpret when you know who owns each layer and what can be redistributed. Megatron-LM and Megatron Core are NVIDIA-led. Origin, review ownership, and license boundaries are more specific than a generic "community framework" label.
| Field | Current project fact |
|---|---|
| Origin and steward | NVIDIA created and maintains Megatron-LM and Megatron Core.[1] |
| Founding contributors | The original paper names Mohammad Shoeybi, Mostofa Patwary, Raul Puri, Patrick LeGresley, Jared Casper, and Bryan Catanzaro.[5] |
| Contributor model | Public contributions use tests, pull requests, NVIDIA code ownership, and Developer Certificate of Origin sign-off. Large architectural changes should begin with an issue.[15] |
| Source license | The top-level license applies BSD-3-Clause-style terms to files unless noted, then preserves Apache-2.0, MIT, BSD, and other notices for bundled code. Treat the repository as composite rather than assigning one SPDX identifier to every file.[16] |
| Commercial boundary | Megatron Core is source code, not a hosted endpoint. NVIDIA containers, NeMo products, support, and cloud offerings can carry separate terms.[1] |
| Asset boundary | Model checkpoints, tokenizers, and training data used with Megatron keep their own licenses and access rules. |
Research lineage
The mesh and optimizer choices came from several research threads, not one magic kernel. Megatron follows a research lineage of complementary systems ideas.
| Source | Contribution | How it appears in current code |
|---|---|---|
| Shoeybi et al. (2019) | Intra-layer tensor model parallelism for multi-billion-parameter transformers | Column and row parallel layers; paired forward/backward communication[5] |
| Narayanan et al. (2021) | Scaling model-parallel training across GPU clusters with overlap and scheduling | Combined DP, TP, PP launch planning and communication overlap[8] |
| Korthikanti et al. (2022) | Selective activation recomputation to save memory with bounded extra compute | Fine-grained checkpoint and recomputation controls[17] |
| Megatron Core guide | Current combinations of TP, PP, CP, EP, DP, and Megatron-FSDP | Rank generators, group collections, and launch constraints[3] |
| Current repository | GPU kernels, models, precision recipes, checkpointing, and reference scripts | megatron/core/, examples/, and tests/[1] |
The papers explain why the decomposition works. The repository explains what it costs to keep the decomposition correct as architectures, accelerators, and model families change. Read both in that order: use the paper to predict the needed collective, then use the code and a trace to check that the launch actually performs it.
Strengths and weaknesses
| Dimension | Strength | Weakness or boundary |
|---|---|---|
| Scale | Mature TP, PP, DP, CP, EP, overlap, and checkpoint paths | Group combinations multiply configuration and debugging states |
| Performance | Fused kernels, grouped GEMM, overlap, and topology-aware groups | Small batches or slow links expose collective latency |
| Model coverage | Dense transformers, MoE, multimodal and hybrid architectures | New architectures need model-specific builders and conversion work |
| Memory | Distributed optimizer, Megatron-FSDP, selective recompute, low precision, sharded checkpoints | Activation, routing, and temporary buffers still need headroom |
| Ecosystem | Bridge connects Hugging Face and Megatron representations | Version skew between Core, Bridge, Transformer Engine, and CUDA can break loads |
| Reproducibility | Functional tests and explicit launch arguments | Large runs need disciplined manifests, seeds, and checkpoint audits |
The trade is now visible. This stack pays off when you control a GPU cluster, need pretraining or serious adaptation, and can invest in topology and launch hygiene. It isn't the shortest path for a one-GPU fine-tune, a hosted inference endpoint, or a quick experiment where model fit and iteration speed matter more than maximum cluster utilization.
Production topology and debugging loop
The following is a GPU lab plan, not an executed result. In a compatible environment, begin with a tiny two-GPU model at TP=2, PP=1, CP=1, DP=1 and synthetic data. Check finite loss and gradients, then verify that saving and reloading preserves fixed-batch logits and the next update within declared tolerances. Random synthetic data needn't show meaningful convergence. Use real held-out data separately for learning quality.
Add one axis at a time. TP=4, PP=2 specifies eight model-parallel ranks when CP=1, not 64 total ranks and not eight necessarily-local ranks. A 64-rank launch also needs a consistent DP/CP factorization and launcher placement. Use a recipe shipped with the pinned software and its supported container; the individual axis flags above aren't a complete training command.
A scale-out result is an experiment, not a flag count. Write a run receipt before changing the mesh so a faster step can be compared with an identical workload:
- GPU model, HBM capacity, NVLink or NVSwitch links, and inter-node fabric.
- Core, Transformer Engine, CUDA, NCCL, PyTorch, and Bridge versions.
- TP, PP, CP, EP, DP, microbatch, global batch, sequence length, token count, and precision.
- Rank order, expected process groups, and placement of each group on the fabric.
- Communication overlap flags, NCCL topology logs, and checkpointing or recomputation settings.
- Baseline command, warm-up steps, timing window, synchronization method, and repeated-step summary.
- Metrics for MFU, tokens per second, step time, communication time, memory high-water mark, expert balance, and loss.
- Correctness evidence: fixed-batch outputs, gradient or loss tolerance, router assignments where relevant, and a checkpoint reload smoke test.
NCCL tests separate operation time from algorithm and bus bandwidth, so a bandwidth number is meaningful only with the collective, payload, and rank count attached.[18] NCCL calls are asynchronous on their CUDA stream; use CUDA events or an explicit stream synchronization around the timing boundary rather than timing only the enqueue call.[19]
Read a distributed trace in two passes. Nsight Systems shows host gaps, CUDA streams, collective launches, memory transfers, and whether pipeline stages overlap. Nsight Compute then explains a selected kernel's launch, memory, compute, and occupancy behavior.[20][21] A counter is a clue, not a verdict: high occupancy doesn't prove useful work, and a high collective share doesn't tell you whether the link, rank placement, or message size caused it.
Before accepting a faster topology or precision recipe, send one fixed microbatch through both candidates. Compare loss and logits within a declared tolerance, then compare gradients or one optimizer update. For CP and PP, check that gathered outputs line up; for EP, compare route counts and expert order. Run a short loss and validation slice afterward. Compute Sanitizer can check memory accesses, shared-memory hazards, uninitialized values, and synchronization when a custom kernel or communication path is suspect.[22]
When every step slows after adding GPUs, inspect group placement and collective timing before changing the model. TP across Ethernet can turn each layer into a network barrier. A PP stage that goes idle points to a bubble or too few microbatches; its stage timeline should show that wait. A CP slowdown points to context exchange dominating attention work. An EP tail points to hot experts, padding, or all-to-all congestion; compare per-expert counts with per-rank dispatch time. A DP slowdown after memory improves points to the gradient reduce-scatter or parameter all-gather tail.
When loss changes after a topology or precision change, compare one fixed batch across runs. Check parameter shard reconstruction, seed and data order, optimizer state reload, gradient scaling, and router decisions. Bitwise equality isn't always expected across kernels, but a reproducible run should explain any tolerance and show that validation behavior remains inside its release budget.
What is the first topology check after a TP scale-out slows every step?
Answer
Verify that TP ranks share the intended NVLink or NVSwitch domain and that NCCL didn't fall back to a slower socket path. Then inspect collective timing before changing batch or model settings.
Evaluation rubric
Use the rank lattice and MLP walkthrough to check yourself:
- Factor dense and expert grids separately and write each rank's coordinate under
tp-cp-ep-dp-ppwithout counting EP twice. - Explain why a column-parallel projection can feed a row-parallel projection without an intermediate all-gather.
- Distinguish sequence parallelism's activation-memory role from context parallelism's sequence partition.
- Treat Megatron-FSDP as DP-state sharding, not extra GPUs in the world-size product.
- Trace MoE token dispatch, expert compute, combine, and the communication tax of skew.
- Calculate distributed-optimizer memory from the DP degree and name buffers that aren't in the formula.
- Plan a checkpoint save, reload, and reshard smoke test.
- Treat MFU and speedup claims as configuration-specific measurements with hardware and convergence caveats.
Follow-up questions
In the 64-rank example, why can DP=4 still produce an optimizer sharding domain of size eight?
Answer
CP=2 repeats dense weights across context peers. Without hierarchical partitioning, the DP-with-CP domain contains 4 × 2 = 8 ranks. The pure DP count still determines independent samples, while the optimizer domain determines the sharded-state denominator.
What changes if DP doubles while global batch and microbatch size stay fixed?
Answer
Microbatches per replica halve. That can expose a larger pipeline bubble even if replicated optimizer memory falls. Compare equal global token counts and the same loss normalization before interpreting step-time changes.