Megatron-LM walkthrough · Stanley Jacob

Part I: The mental model

torchrun --nproc_per_node=8   one process per GPU, SPMD from here on
      |
      v
pretrain_gpt.py               model_provider + forward_step, then pretrain()
      |
      v
megatron/training/training.py initialize_megatron, build model/opt/data, train loop
      |
      v
parallel_state.py             initialize_model_parallel(tp, pp, cp, ...) builds groups
      |                         rank order puts tp innermost (node-local), pp outermost
      v
GPTModel for this stage       only this rank's layers, every linear TP-sharded
      |
      v
get_forward_backward_func()   no-pipeline / 1F1B / interleaved-1F1B schedule
      |
      |  per microbatch: forward -> 2 TP all-reduces per layer -> P2P to next stage
      v
backward                      mirrors forward, P2P activation-gradients upstream
      |
      v
finalize_model_grads          DP all-reduce (or reduce-scatter), shared-embedding
      |                         and sequence-parallel layernorm grad syncs
      v
optimizer.step()              distributed optimizer all-gathers updated params

The one-sentence identity: Megatron-LM is the reference that showed a transformer can be split across GPUs by splitting its two big matrix multiplies in complementary directions, column-parallel then row-parallel, so a single all-reduce per block collects the result, and then stacked that intra-layer tensor parallelism with pipeline, sequence, and data parallelism into a layout you choose to match your interconnect. Where torchtitan keeps the model single-device and applies every parallelism from outside as a transformation, Megatron takes the opposite stance. The parallelism lives inside the model. A Megatron transformer block is built from ColumnParallelLinear and RowParallelLinear layers that already know they are sharded, and the collectives are placed by hand in exactly the spots the math demands. That is more invasive, and it is also why Megatron has historically extracted the highest model-flops utilization on NVIDIA clusters.

Two ideas carry the whole design. The first is that a nonlinearity decides where you may shard for free. In an MLP the first projection is split by columns, so each rank computes a full slice of the hidden activation and can apply GeLU to it with no communication, because GeLU is element-wise and never mixes columns. The second projection is split by rows, which produces a partial sum that one all-reduce turns into the answer. Column then row means one synchronization point per block instead of two. The second idea is that the four parallelism dimensions have wildly different communication profiles, and the job of a good configuration is to map each dimension onto a wire fast enough for it. Tensor parallelism fires every layer and needs NVLink. Pipeline and data parallelism fire once per microbatch or once per step and tolerate Ethernet or InfiniBand between nodes. Most of this chapter is teaching you to see that mapping. Everything here is verified against recent main in July 2026, the repository reorganized the training harness into megatron/training/ and the reusable primitives into megatron/core/, so where a path is likely to differ on an older fork I say so.

Part II: Using it

Megatron-LM is a Linux-and-NVIDIA-GPUs project. The path of least resistance is the NGC PyTorch container, which already has CUDA, cuDNN, NCCL, Apex, and Transformer Engine built and matched:

git clone https://github.com/NVIDIA/Megatron-LM
cd Megatron-LM
# run inside NVIDIA's PyTorch container, which ships the matched CUDA stack
docker run --gpus all -it --rm -v $PWD:/workspace/megatron \
  nvcr.io/nvidia/pytorch:24.05-py3
# the reusable library is also pip-installable on its own
pip install megatron-core

The reference training scripts live at the top level (pretrain_gpt.py, pretrain_bert.py, pretrain_t5.py, plus mixture-of-experts and multimodal variants), and worked examples with real hyperparameters live under examples/. Training reads a preprocessed indexed dataset rather than raw text, so the first real step is tokenizing a corpus into the binary .bin and .idx pair that Megatron's dataset code memory-maps:

python tools/preprocess_data.py \
  --input my_corpus.jsonl \
  --output-prefix my-gpt2 \
  --tokenizer-type GPT2BPETokenizer \
  --vocab-file gpt2-vocab.json --merge-file gpt2-merges.txt \
  --append-eod --workers 32
# produces my-gpt2_text_document.bin and my-gpt2_text_document.idx

A run is launched with torchrun, and the whole configuration is command-line flags. There is no config file format to learn. A small single-node example that exercises tensor and sequence parallelism looks like this:

torchrun --nproc_per_node=8 --nnodes=1 pretrain_gpt.py \
  --tensor-model-parallel-size 8 \
  --pipeline-model-parallel-size 1 \
  --sequence-parallel \
  --num-layers 24 --hidden-size 2048 --num-attention-heads 16 \
  --seq-length 2048 --max-position-embeddings 2048 \
  --micro-batch-size 4 --global-batch-size 256 \
  --train-iters 500000 --lr 1.5e-4 --min-lr 1e-5 \
  --lr-decay-style cosine --weight-decay 0.1 --clip-grad 1.0 \
  --data-path my-gpt2_text_document \
  --vocab-file gpt2-vocab.json --merge-file gpt2-merges.txt \
  --use-distributed-optimizer --bf16

The parallelism dimensions are the first three flags. --tensor-model-parallel-size is the TP degree, --pipeline-model-parallel-size is the PP degree, and --sequence-parallel turns on sequence parallelism, which is only legal when TP is greater than one because it rides on the tensor-parallel group. Data parallelism is the remainder, world size divided by the product of the model-parallel degrees, and Megatron fills it in automatically. To move from this toy to a multi-node run you change --nnodes, add a rendezvous endpoint, and grow the degrees so their product still divides the world size.

The model itself is assembled from parallel layers. You rarely write them directly, the transformer block does, but seeing them is the fastest way to understand what tensor parallelism actually is. A Megatron MLP is essentially this:

from megatron.core.tensor_parallel import ColumnParallelLinear, RowParallelLinear

class ParallelMLP(nn.Module):
    def __init__(self, config):
        super().__init__()
        # first projection: split the ffn (output) dimension across TP ranks
        self.fc1 = ColumnParallelLinear(
            config.hidden_size, config.ffn_hidden_size,
            config=config, gather_output=False,   # keep the output sharded
            bias=True, skip_bias_add=True,
        )
        # second projection: split the ffn (input) dimension across TP ranks
        self.fc2 = RowParallelLinear(
            config.ffn_hidden_size, config.hidden_size,
            config=config, input_is_parallel=True, # input is already sharded
            bias=True, skip_bias_add=True,
        )

    def forward(self, x):
        y, _ = self.fc1(x)   # f operator: identity forward, all-reduce in backward
        y = gelu(y)          # element-wise on a column shard, no communication
        z, _ = self.fc2(y)   # g operator: all-reduce forward, identity in backward
        return z

Now the mistakes people make. First, --sequence-parallel without --tensor-model-parallel-size greater than one is either rejected or a silent no-op depending on the version, because sequence parallelism shards the same regions the tensor-parallel group owns. Second, the degrees have divisibility rules that are easy to trip. The number of attention heads must be divisible by the tensor-parallel size, because heads are what TP splits, and the number of transformer layers must be divisible by the pipeline-parallel size, and when interleaving is on it must be divisible again by the number of virtual stages. Third, batch-size confusion. --micro-batch-size is what one pipeline stage processes at a time, --global-batch-size is the optimizer-step batch summed over all data-parallel ranks and microbatches, and the number of microbatches is the global batch divided by the microbatch size divided by the data-parallel degree. Set a global batch that is not divisible by that product and startup fails with an assertion. Fourth, and most consequential: if your tensor-parallel group spans two nodes, the every-layer all-reduces cross the slow link and step time collapses, with no error to tell you why. Keep TP inside a node.

Part III: When it is the right tool

Megatron-LM is the right tool when you are pretraining large transformers on NVIDIA hardware and want the configuration that squeezes the most model-flops utilization out of the cluster, including the last few percent that come from Transformer Engine kernels and FP8 on Hopper and Blackwell, communication-computation overlap, and a distributed optimizer tuned over years of trillion-token runs. It is also the substrate under a good deal of the ecosystem. NVIDIA NeMo builds its LLM training on Megatron-Core, and many frontier-lab stacks either use it directly or borrow its tensor-parallel layout wholesale. If your goal is peak throughput on a known GPU fleet, this is the honest default.

The honest cases for alternatives. torchtitan when you want a smaller, PyTorch-native codebase you can read in an afternoon, where every parallelism is applied to a plain model from outside via DTensor, and portability across hardware and clean composition matter more than the last percent of MFU. DeepSpeed when you are invested in its ZeRO sharding and its Hugging Face integration, or when CPU and NVMe offload for fitting a model on modest hardware is the point. Ordinary fine-tuning stacks when you are adapting an existing checkpoint rather than pretraining, which is a different problem with different defaults. The trade Megatron asks you to accept is that it owns your model code. You express a model as Megatron parallel layers and specs, not as a plain nanoGPT-style module, and the reward for that coupling is control over exactly where every collective fires.

The architecture-shaped warning is the same one that governs every multi-dimensional trainer, and Megatron is where it was first written down clearly. Match each parallelism to the interconnect it can survive on. Tensor and sequence parallelism communicate inside every layer and only tolerate NVLink-class bandwidth. Pipeline and data parallelism communicate far less often and tolerate the network between nodes. Composing them in the wrong orientation is the classic scaling failure of this domain:

dangerous:  TP group spans nodes
   node0 [g0 g1 g2 g3] --InfiniBand/Ethernet-- node1 [g4 g5 g6 g7]
            <----- TP all-reduces cross the slow link, twice per layer ---->

safe:       TP inside the node, PP/DP across nodes
   node0 [g0 g1 g2 g3]  = one TP group of 4 on NVLink
   node1 [g4 g5 g6 g7]  = another TP group of 4 on NVLink
            <-- only pipeline P2P or the once-per-step DP all-reduce cross -->

Megatron's rank ordering is built to make the safe layout the default. initialize_model_parallel in parallel_state.py lays ranks out with the tensor-parallel dimension innermost, so consecutive ranks on one node form a TP group, and the pipeline dimension outermost, so pipeline stages sit on different nodes. Nothing stops a bad degree choice from spilling a TP group across a node boundary, and the symptom is simply catastrophic step time.

Part IV: The full life of one training step

The specimen is one step of a GPT model launched with the pretrain_gpt.py command above, on a mesh that has tensor, pipeline, and data parallelism all greater than one. Most of the path runs identically with only one dimension active, and where it forks I follow the interesting branch.

Stage 1: torchrun and pretrain_gpt.py

torchrun spawns one process per GPU and sets RANK, LOCAL_RANK, and WORLD_SIZE. From here everything is SPMD, the same program on every rank. pretrain_gpt.py is short. It defines a model_provider that returns a GPTModel, a forward_step that reads a batch and returns the loss, and then hands both to pretrain(...), the driver in megatron/training/training.py (on older forks this lived at megatron/training.py). Nothing about the model is built yet.

Stage 2: initialize_megatron and the process groups

pretrain first calls initialize_megatron, which parses the hundreds of flags in megatron/training/arguments.py into a single args object, seeds RNGs, and calls initialize_model_parallel in megatron/core/parallel_state.py. This is the step that turns a flat list of ranks into a grid. It creates the NCCL process groups for each dimension, the tensor-model-parallel group, the pipeline-model-parallel group, the data-parallel group, and, when enabled, context-parallel and expert-parallel groups, and it stores them behind accessors like get_tensor_model_parallel_group() and get_pipeline_model_parallel_group() that the rest of the code reaches for. The rank ordering here is what puts TP inside a node and PP across nodes, and it is configurable through an order string for exotic topologies. The parallel layers never take a process group as an argument, they ask parallel_state for the right one, which is how a single ColumnParallelLinear can be dropped into any model and still shard along the correct axis.

Stage 3: building only this rank's slice of the model

model_provider now builds a GPTModel (megatron/core/models/gpt/gpt_model.py), but not the whole thing. Pipeline parallelism means each rank constructs only the transformer layers assigned to its stage, so the first stage holds the input embedding and the early blocks, the last stage holds the final norm and the output projection, and the middle stages hold only blocks. Within a stage, every linear is a parallel layer sharded across the tensor-model-parallel group, the attention heads split across TP, the MLP hidden dimension split across TP, and the vocabulary embedding split by VocabParallelEmbedding. The concrete layer is chosen by a spec (megatron/core/models/gpt/gpt_layer_specs.py), which selects either Transformer Engine modules for fused, FP8-capable kernels or the local pure-PyTorch modules, so the same model definition runs with or without Transformer Engine. Weights are initialized so that each rank holds exactly its shard, and no rank ever allocates the full model.

Stage 4: the data, and the forward-backward schedule

The optimizer, learning-rate schedule, and the indexed dataset and its dataloader are built, then the loop in train() begins. Each iteration first pulls a global batch and splits it into microbatches. Only the first and last pipeline stages actually need the tokens, and within a tensor-parallel group the data is broadcast from one rank so every TP rank sees the same input. The step is then handed to the schedule returned by get_forward_backward_func() in megatron/core/pipeline_parallel/schedules.py, which is one of three functions. With no pipeline parallelism it is a plain loop over microbatches. With pipeline parallelism it is forward_backward_pipelining_without_interleaving, the 1F1B schedule, and when virtual stages are configured it is forward_backward_pipelining_with_interleaving. The schedule owns the choreography of which microbatch runs forward, which runs backward, and when each stage sends or receives.

Stage 5: what fires inside one layer's forward

Follow one microbatch through one transformer block on a middle stage. The activation arrives from the previous stage over a P2P receive. Inside attention, the QKV projection is a ColumnParallelLinear, so each TP rank computes Q, K, and V for its own subset of heads with no communication. Attention runs locally per head. The output projection is a RowParallelLinear, which produces a partial sum that the g operator all-reduces across the tensor-parallel group to give the true attention output. The MLP repeats the pattern, column-parallel up-projection, element-wise GeLU on the shard with no communication, row-parallel down-projection closed by a second all-reduce. So one block forward fires two all-reduces on the TP group. If sequence parallelism is on, each of those all-reduces is instead an all-gather entering the tensor-parallel region and a reduce-scatter leaving it, which moves the same number of bytes but keeps the layernorm and dropout activations sharded along the sequence dimension to save memory. After the block, the output is sent by P2P to the next stage.

Stage 6: the backward pass and the mirror collectives

When the schedule runs a microbatch backward, every collective from the forward pass appears transposed. The f operator, which was identity in the forward, all-reduces the input gradient of each column-parallel layer so that gradient is correct across the TP group. The row-parallel layer's g, which all-reduced in the forward, is identity in the backward. Activation gradients flow upstream between pipeline stages over P2P sends, the reverse of the forward P2P. In the 1F1B schedule these forwards and backwards are interleaved so that a stage does one microbatch forward, then one microbatch backward, keeping at most a pipeline-depth of activations alive at once instead of all of them, which is the memory win over an all-forward-then-all-backward GPipe schedule.

Stage 7: finalize_model_grads, the grads nobody sees coming

After the schedule drains, gradients are local to each model-parallel rank and still need three separate reductions, handled together in megatron/core/distributed/finalize_model_grads.py. The big one is the data-parallel reduction, an all-reduce of every gradient across the data-parallel group so every replica agrees, or a reduce-scatter when the distributed optimizer is on. The subtle two are the reason reference implementations exist. First, GPT ties the input embedding and the output projection to the same weight, but those two uses live on the first and last pipeline stages, so their gradients must be all-reduced between just those two stages (unless --untie-embeddings-and-output-weights is set). Second, with sequence parallelism the layernorm parameters are duplicated across the tensor-parallel group, yet each rank computed its gradient from a different sequence shard, so those gradients must be all-reduced across the TP group. Miss either of these and the model trains, slowly diverging, with no crash to point at, which is exactly the class of bug a battle-tested trainer earns its keep by getting right.

Stage 8: the optimizer step and the distributed optimizer

With gradients finalized, optimizer.step() runs. The gradient norm is computed across every model-parallel dimension and clipped, then AdamW updates the parameters. When --use-distributed-optimizer is set (megatron/core/optimizer/distrib_optimizer.py), the optimizer states and the fp32 master parameters are sharded across the data-parallel group in the style of ZeRO stage one, so each rank updates only its slice, then an all-gather rebuilds the full updated parameters on every rank for the next forward. That trades a plain gradient all-reduce for a reduce-scatter plus an all-gather of equal total volume, in exchange for dividing optimizer memory by the data-parallel degree. Overlap flags like --overlap-grad-reduce and --overlap-param-gather hide that traffic under the backward and the next forward. The learning-rate schedule advances, metrics are logged, and a checkpoint is written on interval into a per-rank sharded format, with a distributed format that supports resharding across a different parallel layout. That closes one step, tokens in, gradients reduced across three groups, updated weights gathered, bytes out.

Part V: Internals deep dives

Deep dive: tensor parallelism and the f and g operators

Tensor parallelism is the original Megatron idea, from the 2019 paper, and it rests on two conjugate operators that the code calls f and g. They live in megatron/core/tensor_parallel/mappings.py as autograd functions, and in essence they are this:

# conceptual, megatron/core/tensor_parallel/mappings.py
class _CopyToModelParallelRegion(torch.autograd.Function):     # f
    @staticmethod
    def forward(ctx, x):
        return x                          # identity in the forward
    @staticmethod
    def backward(ctx, grad):
        return _all_reduce(grad)          # all-reduce the gradient

class _ReduceFromModelParallelRegion(torch.autograd.Function): # g
    @staticmethod
    def forward(ctx, x):
        return _all_reduce(x)             # all-reduce the activation
    @staticmethod
    def backward(ctx, grad):
        return grad                       # identity in the backward

ColumnParallelLinear wraps its input in f and returns a sharded output. RowParallelLinear takes a sharded input and wraps its output in g. Chain a column-parallel layer into a row-parallel layer and you get exactly one all-reduce in the forward, from g, and exactly one in the backward, from f. The ordering is the whole trick. Splitting the first matrix by columns means the nonlinearity between the two matrices acts on independent columns and needs no communication, so the only synchronization is at the block's output. Do it the other way, row-parallel first, and you would have to all-reduce before the nonlinearity, doubling the collectives. The attention block uses the same pairing, the QKV projection is column-parallel and splits the heads, the output projection is row-parallel and closes with the all-reduce.

Layer	Weight split	Forward collective	Backward collective
ColumnParallelLinear	by output columns	none (output stays sharded)	all-reduce input grad (f)
RowParallelLinear	by input rows	all-reduce output (g)	none
Attention block	heads across TP	1 all-reduce (output proj)	1 all-reduce (QKV input)
MLP block	ffn dim across TP	1 all-reduce (fc2)	1 all-reduce (fc1)

So a transformer layer under tensor parallelism performs four all-reduces per step, two in the forward and two in the backward, each on an activation of shape (microbatch, sequence, hidden). Two more details close the loop. The vocabulary embedding is split across TP by VocabParallelEmbedding, each rank owns a slice of the vocabulary and masks tokens outside its range, and the partial embeddings are all-reduced. And the loss uses vocab_parallel_cross_entropy in megatron/core/tensor_parallel/cross_entropy.py, which computes the softmax denominator by all-reducing only the per-row max and sum across the TP group rather than gathering the full logits, so the enormous vocabulary-by-batch logit tensor never has to be materialized on one rank. That is a small idea with a large memory payoff, and it is the kind of thing you only get from a trainer that owns its model.

Deep dive: sequence parallelism, memory for free

Plain tensor parallelism shards the linear layers but not the regions between them. The layernorms, the dropouts, and the residual adds run on the full, unsharded activation, replicated identically on every TP rank. For long sequences that replicated activation is a large fraction of the memory budget, and it is pure waste since every rank holds the same bytes. Sequence parallelism, from the 2022 activation-recomputation paper, shards those regions along the sequence dimension instead. Rank zero holds the first slice of sequence positions, rank one the next, and so on, for the parts of the block that act pointwise along the sequence.

The elegant part is what happens at the boundary between a sequence-parallel region and a tensor-parallel region. A ring all-reduce is exactly a reduce-scatter followed by an all-gather, so Megatron replaces the tensor-parallel all-reduce with those two halves and places one at each boundary, moving the same total bytes while keeping the between-layer activations sharded. Entering the tensor-parallel region you all-gather the sequence-sharded activation into the full activation the linear needs. Leaving it you reduce-scatter back to a sequence shard. The operators for this live alongside f and g in mappings.py as gather_from_sequence_parallel_region and reduce_scatter_to_sequence_parallel_region. The bandwidth bill is identical to tensor parallelism, and the reward is that the layernorm and dropout activations, and the activations you stash for the backward pass, are divided by the tensor-parallel degree. This is why sequence parallelism is almost always on whenever tensor parallelism is, and why the flag is coupled to it.

Its companion is selective activation recomputation. Full recomputation reruns an entire layer in the backward pass to avoid storing its activations, trading compute for memory. Megatron observed that the attention softmax and its dropout are cheap in flops but expensive in stored activations, so selective recomputation stores everything except that region and recomputes only the softmax path. You get most of the memory saving for a small fraction of the recompute cost. The flags are --recompute-activations for the selective policy and --recompute-granularity full --recompute-method uniform for the heavier full policy.

Deep dive: pipeline parallelism and interleaving

Pipeline parallelism splits the stack of transformer layers into contiguous stages, one group of layers per pipeline rank, and streams microbatches through them like an assembly line. The cost is the bubble, the idle time while the pipeline fills at the start of a step and drains at the end. With p stages and m microbatches the fraction of time spent in the bubble is (p - 1) / m, which is why you want many more microbatches than stages. The naive GPipe schedule runs all forwards then all backwards, which needs enough memory to hold m microbatches of activations. Megatron's default is 1F1B, one forward one backward, in forward_backward_pipelining_without_interleaving, which reaches a steady state where each stage alternates a forward and a backward and keeps only about p microbatches of activations alive. Same bubble, far less memory.

1F1B on 4 stages, time flows right, F=forward B=backward of a microbatch

stage0  F1 F2 F3 F4 B1 F5 B2 F6 B3 ...        fills, then steady 1F1B, then drains
stage1     F1 F2 F3 B1 F4 B2 F5 B3 ...
stage2        F1 F2 B1 F3 B2 F4 B3 ...
stage3           F1 B1 F2 B2 F3 B3 ...         last stage starts backward first
              <-- fill bubble --><-- steady -->

Interleaving attacks the bubble directly. Instead of giving each device one contiguous block of layers, give it several non-contiguous chunks, called virtual pipeline stages or model chunks. With v chunks per device the bubble fraction shrinks to (p - 1) / (v * m), because the pipeline fills and drains in smaller increments. The price is that pipeline P2P communication grows by a factor of v, since a microbatch now visits each device v times per direction instead of once. That is a good trade when the P2P links have headroom and the bubble is your bottleneck. The interleaved schedule is forward_backward_pipelining_with_interleaving, and it is turned on with --num-layers-per-virtual-pipeline-stage, which sets how many layers each chunk holds and therefore how many chunks there are. Megatron also ships newer low-bubble schedules in the same module, but 1F1B and interleaved 1F1B are the two you must understand first. The point of pipeline parallelism, unlike tensor parallelism, is that stages only exchange the boundary activation once per microbatch over point-to-point sends, so it tolerates the slower links between nodes, which is why it sits on the outer mesh axis.

Deep dive: the communication-volume arithmetic

This is the calculation that decides every real configuration. Let b be the microbatch size, s the sequence length, h the hidden size, L the number of layers, P the parameter count, m the number of microbatches, and t, p, d the tensor, pipeline, and data-parallel degrees. A bandwidth-optimal ring all-reduce of a message of M elements moves about 2(t - 1)/t * M elements of traffic per GPU, since it is a reduce-scatter of (t-1)/t * M followed by an all-gather of the same. With that, each dimension's per-GPU communication per step is:

Dimension	What crosses	Volume per GPU	Fires	Wire
Tensor (TP)	activations, b·s·h	~ 8·L·(t−1)/t · b·s·h per microbatch	every layer, both passes	NVLink
Sequence (SP)	activations, b·s·h	same total as TP (all-reduce split in two)	every layer	NVLink
Pipeline (PP)	boundary act/grad, b·s·h	~ 2·m · b·s·h (× v with interleaving)	once per microbatch, P2P	inter-node ok
Data (DP)	gradients, P params	~ 2(d−1)/d · P	once per step, all-reduce	inter-node ok

Read the shapes of these formulas, not the constants. Tensor parallelism scales with the activation size b·s·h and fires four all-reduces in every one of L layers, so it is by far the heaviest and most frequent traffic, and it is the only dimension whose volume grows with sequence length. That is the arithmetic reason TP must live on NVLink and stay inside a node, usually at degree eight or less. Sequence parallelism moves the same bytes as TP, so it changes memory, not bandwidth. Pipeline parallelism moves one boundary activation per microbatch and does not scale with L at all, so its volume is smaller by roughly the number of layers per stage and it is point-to-point, which is why it happily crosses node boundaries, and its real cost is the bubble rather than the bytes. Data parallelism moves the parameter count once per step, independent of batch or sequence, so it is cheap per token when microbatches are large and it also tolerates inter-node links. Stack the dimensions so the heaviest, most frequent traffic sits on the fastest wire, TP innermost on NVLink, then PP and DP across the network, and you have the entire art of configuring Megatron in one sentence.

Deep dive: Megatron-Core as a reusable library

Early Megatron-LM was a monolithic training script. The important structural change over the last few years was carving the reusable primitives out into megatron/core, published as the pip-installable megatron-core package, while the top-level scripts and megatron/training became a reference harness that consumes it. The split matters because Megatron-Core is what NeMo and many other stacks actually depend on, and it is the part worth learning even if you never run pretrain_gpt.py.

Its shape is worth memorizing. megatron/core/parallel_state.py owns every process group and is the single source of truth for who talks to whom. megatron/core/tensor_parallel/ holds the parallel linear layers, the vocab-parallel embedding and cross entropy, and the f and g mappings. megatron/core/pipeline_parallel/ holds the schedules and the P2P communication. megatron/core/transformer/ holds the building blocks, the transformer block and layer, attention, and MLP, along with the TransformerConfig dataclass that every one of them reads. megatron/core/models/ assembles those blocks into concrete models like GPTModel. megatron/core/distributed/ holds the data-parallel wrapper and the gradient finalization, and megatron/core/optimizer/ holds the distributed optimizer.

The pattern that ties it together is the spec, or ModuleSpec. Rather than hard-coding which attention or norm class a layer uses, Megatron-Core describes a layer as a spec, a small record naming the submodule classes and their arguments, and a build_module helper instantiates it. That is how the same GPTModel can be built from Transformer Engine modules for fused FP8 kernels or from local PyTorch modules for portability, by passing a different spec, without forking the model. The spec system is Megatron-Core's answer to the tension it lives with, deep coupling between the model and the parallelism, made tolerable by making the coupling declarative and swappable. If torchtitan's lesson is that layout can be a property of the tensor, Megatron-Core's is that when the model must know about its own parallelism, the least you can do is make that knowledge a small, named, replaceable specification.

Part VI: Reading the repository

The repository is large, but the path through it is short if you follow the layers of parallelism rather than the directory listing. Paths are given as they sit on recent main, and the training harness versus core split is the one reorganization to keep in mind.

Stage 0, orientation. Read the top-level README.md, then pretrain_gpt.py, then one script under examples/ with real numbers. Questions to hold. What are the three parts a pretraining script must supply, and which flags set the parallelism degrees, and where does the world size get divided among them.

Stage 1, the process groups. Read megatron/core/parallel_state.py, specifically initialize_model_parallel and the accessor functions. Questions. How does the rank ordering keep tensor parallelism inside a node, what groups exist, and how does a layer far away in the code get the right group without being handed one.

Stage 2, tensor parallelism. Read megatron/core/tensor_parallel/layers.py for ColumnParallelLinear, RowParallelLinear, and VocabParallelEmbedding, then mappings.py for f and g, then cross_entropy.py. Questions. Why column then row, where exactly does the all-reduce sit in each layer, and how does the vocab-parallel cross entropy avoid gathering the full logits.

Stage 3, the transformer and the specs. Read megatron/core/transformer/, the transformer block and layer, attention.py, mlp.py, and transformer_config.py, then the GPT model and its layer specs under megatron/core/models/gpt/. Questions. How is a block assembled from parallel layers, what does a spec name, and how does swapping the spec swap Transformer Engine for local modules.

Stage 4, pipeline parallelism. Read megatron/core/pipeline_parallel/schedules.py, the three forward-backward functions, and p2p_communication.py. Questions. What does 1F1B keep alive that GPipe does not, how does interleaving change the bubble and the communication, and which ranks compute the loss.

Stage 5, gradients, the optimizer, and the harness. Read megatron/core/distributed/finalize_model_grads.py, megatron/core/optimizer/distrib_optimizer.py, and then megatron/training/training.py to see train_step tie it together. Questions. What are the three gradient reductions and why does each exist, what does the distributed optimizer shard, and where in the step do the overlap flags hide their traffic.

Where not to start. The mixture-of-experts path adds expert parallelism and its own token-routing collectives, a second grid on top of the dense story, and it will only confuse a first read. The Transformer Engine and FP8 internals, the userbuffers-based TP communication overlap, and the multimodal and retrieval models are all worth knowing later and noise at the start. Read the dense GPT path end to end first.

Part VII: Hands-on labs

Lab 1 needs no GPU. Labs 2 through 5 want two or more GPUs on one node, ideally with NVLink. Exact log strings drift with the version, so match on shape, not text.

Lab 1: derive a legal configuration on paper. Concept: the divisibility and volume arithmetic of Part V.

Given: 16 GPUs, 8 per node, a 32-layer model, 32 attention heads.
Pick tp, pp, dp with tp * pp * dp = 16 such that:
  - tp divides 32 (heads)          -> tp in {1,2,4,8,16}
  - pp divides 32 (layers)         -> pp in {1,2,4,8,16,32}
  - tp stays within one node (<= 8 GPUs) so its all-reduces hit NVLink
A good answer: tp=8 (one node), pp=2 (across the two nodes), dp=1.
A bad answer: tp=16, which spans both nodes and puts the per-layer
all-reduce on the inter-node link.

Write out, using the Part V table, why tp=16 is slow even though it is legal. Then redo the exercise for 64 GPUs and a 48-layer, 64-head model, and justify each degree by the wire its traffic lands on.

Lab 2: watch tensor parallelism move memory and time. Concept: the every-layer all-reduce.

# same model, two layouts on 2 GPUs
torchrun --nproc_per_node=2 pretrain_gpt.py --tensor-model-parallel-size 2 \
  --num-layers 12 --hidden-size 1024 --num-attention-heads 16 \
  --seq-length 1024 --max-position-embeddings 1024 \
  --micro-batch-size 8 --global-batch-size 16 --train-iters 50 \
  --data-path my-gpt2_text_document \
  --vocab-file gpt2-vocab.json --merge-file gpt2-merges.txt --bf16
# compare to --tensor-model-parallel-size 1 (pure data parallel on 2 GPUs)

Compare per-GPU peak memory and step time between the two runs. The tensor-parallel run uses less memory per GPU because the weights and activations are split, and pays for it in the per-layer all-reduces. Now add --sequence-parallel to the TP run and confirm memory drops again while step time barely moves, which is the communication-neutral memory win of Part V made visible.

Lab 3: see the pipeline bubble. Concept: (p−1)/m and why microbatch count matters.

# 4 GPUs, pipeline degree 4, sweep the number of microbatches
torchrun --nproc_per_node=4 pretrain_gpt.py --pipeline-model-parallel-size 4 \
  --num-layers 8 --hidden-size 1024 --num-attention-heads 16 \
  --seq-length 1024 --max-position-embeddings 1024 \
  --micro-batch-size 1 --global-batch-size 4  --train-iters 50 ...   # m = 4
# then rerun with --global-batch-size 32                              # m = 32

With four microbatches and four stages the bubble is (4-1)/4, about three quarters idle, and throughput is poor. With thirty-two microbatches it falls to (4-1)/32, under ten percent, and throughput jumps. Predict the ratio before you measure it.

Lab 4: interleaving trades communication for bubble. Concept: (p−1)/(v·m).

# take the m=4 case above and add virtual pipeline stages
torchrun --nproc_per_node=4 pretrain_gpt.py --pipeline-model-parallel-size 4 \
  --num-layers 8 --num-layers-per-virtual-pipeline-stage 1 \
  --micro-batch-size 1 --global-batch-size 4 ...   # v = 2 chunks per device

With eight layers over four stages and one layer per virtual stage you get two chunks per device, so the bubble target drops from 3/4 to 3/8. Watch throughput improve on the small-microbatch case, and watch the P2P volume roughly double, which is the trade the schedule is making.

Lab 5: compose all three dimensions. Concept: the full mesh and the safe layout.

# 8 GPUs: tp=2, pp=2, dp=2, with the memory-savers on
torchrun --nproc_per_node=8 pretrain_gpt.py \
  --tensor-model-parallel-size 2 --pipeline-model-parallel-size 2 \
  --sequence-parallel --use-distributed-optimizer \
  --num-layers 8 --hidden-size 2048 --num-attention-heads 16 \
  --seq-length 2048 --max-position-embeddings 2048 \
  --micro-batch-size 2 --global-batch-size 32 --train-iters 50 --bf16 ...

Confirm the startup log reports tp=2 pp=2 dp=2 and that their product is the world size. Map each of the three collectives from Part IV onto a dimension, the every-layer all-reduce onto TP, the microbatch P2P onto PP, and the once-per-step gradient reduction onto DP, and note that on a single node all three ride NVLink here, which is exactly why single-node debugging hides interconnect problems that only appear at multi-node scale.

Part VIII: Questions and model answers

Understanding checks. Answer aloud before reading.

1. What is Megatron-LM in one sentence?

NVIDIA's reference trainer for large transformers, which shards a model across GPUs by splitting each block's two matrix multiplies column-parallel then row-parallel with one all-reduce to collect the result, and stacks that tensor parallelism with pipeline, sequence, and data parallelism, all provided by the reusable Megatron-Core library.

2. Why column-parallel then row-parallel, and not the reverse?

Splitting the first matrix by columns lets each rank compute a full slice of the hidden activation and apply the nonlinearity to it with no communication, because the nonlinearity is element-wise and never mixes columns. The row-parallel second matrix then needs a single all-reduce at the block's output. Reversing the order would force an all-reduce before the nonlinearity, doubling the collectives per block.

3. What are the f and g operators?

Conjugate autograd functions. f is identity in the forward and an all-reduce in the backward, and it wraps the input of a column-parallel layer. g is an all-reduce in the forward and identity in the backward, and it wraps the output of a row-parallel layer. Together they place exactly one all-reduce in each direction per block.

4. How many all-reduces does tensor parallelism perform per layer per step?

Four. Two in the forward, one closing attention and one closing the MLP, and two in the backward, from the f operators of the two column-parallel layers. Each is on an activation of shape (microbatch, sequence, hidden), which is why the volume scales with the activation and with sequence length.

5. What does sequence parallelism change, and what does it not change?

It shards the layernorm, dropout, and residual regions along the sequence dimension instead of replicating them, which cuts activation memory by the tensor-parallel degree. It does not change communication volume, because it splits each tensor-parallel all-reduce into a reduce-scatter and an all-gather that together move the same bytes.

6. Why must tensor parallelism stay inside a node while pipeline parallelism can cross nodes?

Tensor parallelism fires four all-reduces per layer on activation-sized messages that grow with sequence length, so it needs NVLink-class bandwidth and low latency on the critical path. Pipeline parallelism sends one boundary activation per microbatch as a point-to-point message and does not scale with the number of layers, so it tolerates the slower links between nodes.

7. What is the pipeline bubble, and how do microbatches and interleaving affect it?

The bubble is the idle time while the pipeline fills and drains, a fraction (p−1)/m of the step for p stages and m microbatches. More microbatches shrink it directly. Interleaving gives each device several non-contiguous chunks, cutting the bubble to (p−1)/(v·m) for v chunks, at the cost of multiplying pipeline communication by v.

8. What memory advantage does 1F1B have over GPipe?

GPipe runs all forwards then all backwards, so it must store roughly m microbatches of activations. 1F1B reaches a steady state of one forward then one backward per stage and keeps only about p microbatches alive at once, giving the same bubble with far less activation memory.

9. What three gradient reductions happen after the backward pass, and why?

A data-parallel reduction so every replica agrees on the gradient, an all-reduce of the tied input-embedding and output-projection gradient between the first and last pipeline stages because those two uses share one weight, and, under sequence parallelism, an all-reduce of the layernorm gradients across the tensor-parallel group because each rank computed them from a different sequence shard.

10. What does the distributed optimizer shard, and what does it trade?

It shards the optimizer states and the fp32 master parameters across the data-parallel group, ZeRO-stage-one style, so each rank updates only its slice. It trades a plain gradient all-reduce for a reduce-scatter of gradients plus an all-gather of updated parameters of equal total volume, in exchange for dividing optimizer memory by the data-parallel degree.

11. What is Megatron-Core, and who uses it?

The reusable library carved out of Megatron-LM, pip-installable as megatron-core, containing the parallel layers, the pipeline schedules, the transformer building blocks, the process-group management, and the distributed optimizer. NVIDIA NeMo and many other training stacks depend on it, while the top-level scripts are a reference harness that consumes it.

12. When would you choose torchtitan or DeepSpeed over Megatron-LM?

torchtitan when you want a small PyTorch-native codebase that applies every parallelism to a plain model from outside, with portability and readability over peak MFU. DeepSpeed when you are invested in its ZeRO sharding and Hugging Face integration or need CPU and NVMe offload. Megatron when maximum throughput on NVIDIA clusters is worth a framework that owns your model code.

13. Why does the vocab-parallel cross entropy matter?

The logits are a batch-by-vocabulary tensor that is enormous for a large vocabulary, and gathering it onto one rank would blow the memory budget. The vocab-parallel cross entropy keeps the logits sharded across the tensor-parallel group and computes the softmax denominator by all-reducing only the per-row max and sum, so the full logit tensor is never materialized.

14. A multi-node run has good single-node throughput but collapses at scale, with no error. Name the first suspect.

A tensor-parallel group that spilled across a node boundary, so the every-layer all-reduce now crosses the inter-node link. It is legal, it trains correctly, and it is catastrophically slow. The fix is to keep the tensor-parallel degree within one node and push pipeline and data parallelism across nodes.

Part IX: Design lessons

Let the math choose the seams. Tensor parallelism works because the split is placed exactly where a nonlinearity makes it free, column-parallel before the element-wise GeLU, row-parallel after. The best parallel decompositions are not imposed on an algorithm, they are read off its structure. The same instinct shows up in tiling a loop along the axis with no carried dependence and in sharding a database on the key that queries already filter by.

Place communication, do not hide it. Megatron writes the f and g collectives into the model by hand, at the two points the math demands and nowhere else. Owning that placement is what lets it hit peak utilization, and it is the deliberate opposite of the torchtitan philosophy of inferring collectives from tensor layout. Both are defensible, and the trade is explicitness and control against composability and portability.

Know the volume and frequency of every message. The whole discipline of configuring a large run reduces to one table, which dimension moves how many bytes how often, mapped onto which dimension moves data over which wire. Systems that make their communication costs legible get configured well. Systems that hide them get configured by trial and error at great expense.

Trade memory and compute and bubble on purpose. Sequence parallelism spends nothing extra to save activation memory, selective recomputation spends a little compute to save more, and interleaving spends communication to shrink the bubble. Each is a knob with a known price, and scaling is choosing the knobs that fit your bottleneck rather than turning all of them up.

When coupling is unavoidable, make it declarative. Megatron-Core cannot fully decouple the model from its parallelism, so it makes the coupling a swappable spec, a named record of which modules to build. If you must depend on a detail, depend on a small replaceable description of it rather than on a hard-coded class, which is how the same model runs with Transformer Engine or with plain PyTorch.

Extract the library once the patterns are proven. The move from a monolithic script to Megatron-Core plus a thin harness came after the techniques were battle-tested, not before. Reusable abstractions earned from working code outlast reusable abstractions designed up front, which is why NeMo and others could safely build on the extracted core.

Part X: Memorization framework

The one-sentence summary: Megatron-LM shards each transformer block column-parallel then row-parallel so one all-reduce per block collects the result, splits the between-layer regions along the sequence for memory, streams microbatches through pipeline stages with a 1F1B or interleaved schedule, reduces gradients across data-parallel ranks, and provides all of it as the Megatron-Core library.

torchrun -> pretrain_gpt.py -> pretrain() in megatron/training/training.py
  -> initialize_model_parallel (parallel_state.py): tp innermost, pp outermost
  -> GPTModel for this stage: ColumnParallelLinear + RowParallelLinear, f/g
  -> get_forward_backward_func(): no-pipeline / 1F1B / interleaved-1F1B
  -> per microbatch: 2 TP all-reduces per layer, P2P between stages
  -> finalize_model_grads: DP all-reduce + shared-embedding + SP layernorm
  -> optimizer.step(): distributed optimizer all-gathers updated params

The chain mapped to source:

launch          pretrain_gpt.py, examples/
harness         megatron/training/ (training.py, arguments.py, initialize.py)
groups          megatron/core/parallel_state.py
tensor parallel megatron/core/tensor_parallel/ (layers.py, mappings.py, cross_entropy.py)
transformer     megatron/core/transformer/ + megatron/core/models/gpt/
pipeline        megatron/core/pipeline_parallel/ (schedules.py, p2p_communication.py)
grads + optim   megatron/core/distributed/ + megatron/core/optimizer/

Memorize these blocks:

The pairing: column-parallel (split by output columns, free nonlinearity) then row-parallel (split by input rows, one all-reduce). Four all-reduces per layer per step.
The f/g operators: f is identity forward and all-reduce backward, g is all-reduce forward and identity backward.
Sequence parallelism: shards the layernorm and dropout regions along the sequence, splits each all-reduce into reduce-scatter plus all-gather, same bytes, less memory.
Pipeline bubble: (p−1)/m without interleaving, (p−1)/(v·m) with v virtual stages, and interleaving multiplies P2P volume by v.
Communication volumes: TP scales with b·s·h every layer (NVLink), PP with one b·s·h per microbatch (inter-node ok), DP with P once per step (inter-node ok).
Megatron-Core: the reusable library under NeMo, split into parallel_state, tensor_parallel, pipeline_parallel, transformer, models, distributed, optimizer.

Part XI: Final takeaway

If the single-GPU transformer that Megatron shards is the piece you want built from scratch, the ML implementations section derives it, and the attention-memory story behind sequence parallelism and selective recomputation is worked out in the FlashAttention chapter and on the online softmax page. The parallelism arithmetic here is the same arithmetic that governs the parallel computing class, and the inference-time cousin of these ideas lives in the vLLM and torchtitan walkthroughs. Read Megatron once with the communication table in hand, and the code stops looking like a wall of collectives and starts looking like a small set of deliberate choices about which wire carries which message.

Key takeaway: Megatron-LM is the reference for large-scale transformer training because it places every collective by hand where the math allows, column-parallel then row-parallel for one all-reduce per block, sequence-sharded regions for memory, pipelined microbatches for depth, and data-parallel gradient reduction for width. Its enduring lesson is the arithmetic, know the volume and frequency of every message and map it onto the interconnect that can carry it, and the same handful of parallel primitives, packaged as Megatron-Core, will scale a model from one node to a supercomputer.