Deep Dive - Megatron

Megatron-LM is a reference training stack for transformer models that don't fit on one GPU. Megatron Core is its reusable library of parallel layers, process groups, schedules, optimizers, checkpoint utilities, and model definitions. NVIDIA's current repository calls the split explicit: Megatron-LM gives research teams runnable recipes, while Megatron Core gives framework authors composable building blocks.^{[1]Reference 1Megatron-LMhttps://github.com/NVIDIA/Megatron-LM}

Treat a process rank as a coordinate rather than only rank=27. Its axes include data replica, context lane, pipeline stage, and tensor shard. Megatron creates communication groups from those coordinates, then every forward and backward operation chooses 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.

Why this project exists

Modern language models hit several limits at once. Parameters, gradients, optimizer state, and activations can exceed one GPU's memory. Long sequences can make attention intermediates too large even when weights fit, while a sparse mixture-of-experts (MoE) model stores many experts but activates only a few per token. Compute, memory bandwidth, network latency, and fault recovery all matter.

Megatron's answer is to split different dimensions for different reasons:

These axes compose, but they aren't interchangeable. TP lowers per-layer memory and keeps a single token on several GPUs. DP doesn't make one sequence faster; it admits more independent batches. PP can make a model fit across nodes, but it inserts stage-to-stage latency. CP reduces sequence-local memory and attention work, but each context group must exchange information. A good launch starts with a memory and topology budget, then picks axes that match it.

Two projects in one repository

Open a Megatron-LM clone and you'll see two layers.

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 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. The API surface is broad because each production model needs a slightly different mesh.

Megatron-LM reference scripts

Top-level training scripts and examples/ turn Core into repeatable experiments. pretrain_gpt.py and 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. The scripts are useful for reproducing papers and for learning what a full training run needs, but they aren't a drop-in serving API.

Megatron Bridge

Megatron Bridge lives in the companion NVIDIA-NeMo repository. It converts Hugging Face checkpoints to Megatron layouts and back, with model recipes and parallelism-aware conversion.^{[1]Reference 1Megatron-LMhttps://github.com/NVIDIA/Megatron-LM} A training run can therefore start from an ecosystem checkpoint, use Megatron's sharded layout for throughput, then export a portable checkpoint for evaluation or serving. Conversion isn't a rename operation: tensor and pipeline shards must be mapped to global tensors, and MoE expert placement must be preserved.

The separation is a strength and a cost. Core can evolve without forcing every framework into one CLI, but users must 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

Suppose a run has 64 GPUs and uses TP=4, PP=2, CP=2, and DP=4. The dense model grid is:

$N_{dense} = TP \times PP \times CP \times DP = 4 \times 2 \times 2 \times 4 = 64.$

Each rank gets a tuple (dp, cp, pp, tp). Rank numbering is an implementation detail; group membership is the contract. Megatron's initialize_model_parallel builds rank generators from a configurable order such as tp-cp-ep-dp-pp. It checks that world size is divisible by the model-parallel product before creating groups.^{[2]Reference 2Parallelism Strategies Guide.https://docs.nvidia.com/megatron-core/developer-guide/latest/user-guide/parallelism-guide.html}

The rank lattice visual makes this arithmetic concrete. Start at the 64-rank world, then factor into data replicas, sequence lanes, depth stages, and width shards. The rightmost TP group is usually placed inside one NVLink or NVSwitch domain because every transformer layer can use it. PP groups can cross nodes because they exchange boundary activations less often than TP exchanges partial layer results. DP groups can use the inter-node fabric when gradient traffic is overlapped with backward compute.

EP is a second grid for MoE layers. In a dense run, DP commonly fills the remaining world after TP, PP, and CP. In an MoE run, the expert grid can use expert tensor parallelism (ETP), EP, expert data parallelism (EDP), and PP. Current Megatron Core also exposes generalized tensor-parallel rematerialization axes, so the complete world-size equation can gain more factors. Don't silently call every configuration "3D"; count the actual axes in the launch.

What each parallel axis does

Tensor parallelism: split width

TP cuts a matrix along a dimension that keeps local GEMMs large enough to use the GPU. For a column-parallel first MLP projection, A with shape [H, 4H] becomes four local matrices A_i with shape [H, H] when TP=4. 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.^{[3]Reference 3Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism.https://arxiv.org/abs/1909.08053}

For a tiny worked example, let H=8, intermediate size be 16, and TP=4. Each rank stores [8,4] for the up projection and [4,8] for the down projection. Input X has shape [B,S,8]. Rank i computes Y_i = GeLU(X A_i) with shape [B,S,4], then Z_i = Y_i B_i with shape [B,S,8]. The output is Z = Σ_i Z_i. The GeLU never needs an all-gather because it acts independently on each channel.

Attention uses the same idea with heads. Query, key, and value projections are split across TP ranks, each rank computes attention for its head shard, and the output projection combines them. Head counts and hidden dimensions must divide cleanly by TP. Grouped-query attention can make key/value head divisibility the tighter constraint.

TP trades memory for communication. Each rank stores roughly 1/TP of a sharded weight, but every layer may issue all-reduce, all-gather, or reduce-scatter collectives. The exact count changes with sequence parallelism, overlap settings, and kernel path. Keep TP inside high-bandwidth links when possible.

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, one backward) so stages stay busy. Virtual pipeline stages interleave smaller layer chunks on each rank and reduce bubbles when layer counts and batch sizes permit.

PP lowers per-rank weight memory without all-reducing every layer across nodes. It adds activation send/receive latency and creates a bubble when not enough microbatches keep stages occupied. A topology with four TP ranks inside each node and two PP stages across nodes often beats a single eight-way TP group over a slower fabric.

Data parallelism: split samples

DP replicates model computation and gives each replica different samples. After backward, gradients from replicas are reduced. Standard DDP keeps full model state per rank. Megatron's distributed optimizer changes optimizer-state ownership while keeping DP replicas synchronized, which is useful when optimizer state, not weights, causes the memory wall.

Context and sequence parallelism: split token work

Sequence parallelism works alongside TP. 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.

Context parallelism partitions the input sequence itself. Current Core docs describe each CP rank exchanging other sequence chunks because attention needs context outside its local slice.^{[2]Reference 2Parallelism Strategies Guide.https://docs.nvidia.com/megatron-core/developer-guide/latest/user-guide/parallelism-guide.html} Communication can use point-to-point, all-gather, all-to-all, or hierarchical paths depending on model and hardware. CP is useful when long prompts or long training sequences dominate memory. CP duplicates weights across context peers, so the weight-gradient reduction must account for those replicas.

Expert parallelism: split sparse capacity

An MoE layer has a router that chooses top-k experts for each token. 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. The current dispatcher code exposes all-gather and all-to-all paths, plus Flex backends that can fuse or specialize communication.^{[4]Reference 4Mixture of Expertshttps://docs.nvidia.com/megatron-core/developer-guide/0.15.0/api-guide/moe.html}

EP's risk is visible in its collective: an all-to-all can move every token to a different rank. 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.

MoE parallel folding

Traditional layouts often constrain EP to the data-parallel domain. That couples an attention-friendly grid to an MoE-friendly grid. High TP can help attention's wide matrices but hurt each small expert. High CP can help long-context attention while offering little value to per-token expert MLPs.

Megatron Core's MoE parallel folding decouples the two 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.^{[4]Reference 4Mixture of Expertshttps://docs.nvidia.com/megatron-core/developer-guide/0.15.0/api-guide/moe.html}

This flexibility costs mental and operational complexity. 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 a tiny synthetic MoE run before scheduling a long pretraining job.

Distributed optimizer: trade replicated state for collectives

Adam-style training keeps model parameters, gradients, and first and second moments. With BF16 model weights and FP32 optimizer state, replicated optimizer memory can dominate. Megatron's distributed optimizer shards optimizer-owned state across DP ranks. The official guide lists a theoretical 18 bytes per parameter for a non-distributed BF16 setup versus 6 + 12/d bytes with distributed optimizer at DP size d.^{[5]Reference 5Distributed Optimizerhttps://docs.nvidia.com/megatron-core/developer-guide/latest/user-guide/features/dist_optimizer.html}

The step is easier to reason about as a four-part loop:

1. Backward fills a gradient buffer on each DP rank.
2. reduce-scatter sums gradients and gives each rank its shard.
3. Each rank updates its local FP32 master parameters and optimizer moments.
4. all-gather rebuilds 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.

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. If an optimizer shard is too small to hide communication, a larger DP value can lower throughput even while it saves memory.

Checkpoints, resharding, and recovery

Large runs checkpoint more than model weights. They need optimizer moments, random-number-generator state, data-loader position, iteration counters, and parallel-layout metadata. 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. Fully reshardable formats support selected TP, PP, EP, ETP, and DP layout changes, while the default DP-reshardable format can't change model parallelism. CP usually changes process groups rather than weight-shard shapes. Compatibility still depends on checkpoint format, model code, expert layout, and version. Test a save at one topology and load at a smaller topology before relying on it for recovery. Keep the exact launch arguments and Git commit beside the checkpoint.

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.

Failure modes include an interrupted asynchronous save, an optimizer state saved with a different DP size, a missing expert shard, or a model class that changed its key layout. Use checksums, manifest files, and a small load-and-forward smoke test for every checkpoint artifact. A training job that can't prove which state it loaded isn't reproducible.

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.

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.

The repository's current README advertises FP16, BF16, FP8, and FP4 support and reports up to 47% MFU for specific H100 benchmarks.^{[1]Reference 1Megatron-LMhttps://github.com/NVIDIA/Megatron-LM} Treat that number as a source-scoped measurement, not a promise for every model. The benchmark uses particular sequence lengths, vocabulary, model sizes, overlap flags, and a no-convergence run. Hardware, software versions, batch shape, and model architecture all change MFU.

Project identity

Megatron-LM and Megatron Core are NVIDIA-led open-source projects. Their origin, review ownership, and license boundaries are more specific than a generic "community framework" label.

Research lineage

Megatron follows a research lineage of complementary systems ideas rather than one kernel.

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: a paper's clean grid can hide launch constraints, while a code path can hide the reason a collective is necessary.

Strengths and weaknesses

Megatron is a strong choice 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

Start with a single-node fragment. Run a tiny model with TP=2, PP=1, CP=1, DP=1, and mock data. Confirm loss decreases, checkpoints save, and a reload produces the same next-token logits within your precision tolerance. Then add one axis at a time. A command that names TP=4 and PP=2 isn't a 64-GPU proof unless eight local ranks and the required network are present.

For a real cluster, document:

1. GPU model, HBM capacity, NVLink or NVSwitch links, and inter-node fabric.
2. Core, Transformer Engine, CUDA, NCCL, PyTorch, and Bridge versions.
3. TP, PP, CP, EP, DP, microbatch, global batch, sequence length, and precision.
4. Rank order and expected process groups.
5. Communication overlap flags and NCCL topology logs.
6. Checkpoint format, save interval, resharding smoke test, and restart owner.
7. Metrics for MFU, tokens per second, step time, communication time, memory high-water mark, expert balance, and loss.

When throughput drops after adding GPUs, check topology before changing the model. TP across Ethernet can turn every layer into a network barrier. PP can show a pipeline bubble if the microbatch count is too low. CP can spend more time exchanging context than computing attention. EP can stall on hot experts or all-to-all congestion. DP can hide communication until the network becomes the shared bottleneck.

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.

Mastery check

By this point, you can:

• Factor a world size into TP, PP, CP, EP, and DP axes and write each rank's coordinate.
• 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.
• 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.