Megatron-LM: Scalable LLM Training

• Megatron-LM is an open-source suite that employs key parallelism strategies to scale Transformer training to multi-billion and trillion-parameter sizes.
• It utilizes tensor, pipeline, and data parallelism along with expert parallelism to optimize communication and memory efficiency across hundreds of GPUs.
• System-level optimizations including fused kernels, mixed precision, and dynamic scheduling contribute to near-linear scaling and high throughput in large-scale training.

The Megatron-LM Project is a collection of open-source libraries, parallelism strategies, and system-level tools designed to enable efficient, scalable training of very large Transformer-based LLMs, including dense and Mixture-of-Expert architectures, at multi-billion and trillion-parameter scale. Originating from NVIDIA with ongoing co-development from the broader academic and open-source community, Megatron-LM provides foundational techniques in model and data parallelism, as well as extensive system optimizations for distributed deep learning.

1. Foundational Design and Motivation

Megatron-LM emerged to address scaling barriers in training Transformer models with parameters far exceeding single-GPU memory, and with computational demands too great for conventional data parallelism alone. The core observation is that transformer layers—especially the large weight matrices in self-attention and MLP blocks—are highly amenable to partitioning across devices with judicious placement of communication operations. The resulting intra-layer ("tensor") model parallelism allows models to scale from 1 billion to beyond 1 trillion parameters, circumventing the limitations of naive data-parallel setups, and avoiding the need for new compilers or custom C++ kernels (Shoeybi et al., 2019).

Critical system-level challenges motivating the development include:

• Memory constraints for weights, optimizer state, and activations.
• Exponential compute scaling with model capacity.
• The need for linear or near-linear scaling efficiency on GPU clusters up to thousands of nodes.
• Diagnosability and reliability issues introduced by multi-axis parallelism and complex failure modes.

Megatron-LM’s approach is recognized for its modularity, allowing parallelism strategies to be composed orthogonally, and for its empirical demonstration of state-of-the-art scaling and accuracy on open language modeling benchmarks (Shoeybi et al., 2019, Narayanan et al., 2021, Smith et al., 2022).

2. Parallelism Methodologies

2.1 Tensor (Model) Parallelism

Within each transformer layer, large dense GEMMs are split across $P$ devices:

• Column-parallel splits partition matrix A in the MLP across devices. Each GPU computes $Y_i = \mathrm{GeLU}(X A_i)$ for its shard, with no communication before or after activation.
• Row-parallel splits handle the output GEMM and projection layers, requiring an all-reduce to aggregate per-device outputs.

The PyTorch implementation uses model-parallel groups, with collective operations (all_reduce, all_gather) orchestrated via torch.distributed and NCCL (Shoeybi et al., 2019). The method also supports efficient MHA projection and aggregation with per-head partitioning, minimizing synchronization.

2.2 Pipeline Parallelism

Model layers are partitioned into $P_p$ contiguous stages, each assigned to one or more GPUs. Microbatches are streamed through the stages with a 1F1B ("one-forward, one-backward") or interleaved schedule to minimize pipeline bubbles:

• The interleaved pipeline schedule splits each stage into $v$ chunks, reducing forward/backward idle ("bubble") time by a factor $v$: bubble fraction = $((P_p - 1)/m)/v$.
• This strategy is critical for sustaining high utilization across hundreds or thousands of nodes, especially when the optimal batch size per GPU is small (Narayanan et al., 2021).

2.3 Data Parallelism

Once the model is partitioned across tensor and pipeline axes, data parallelism replicates the composite shard across $P_d$ groups, each ingesting a distinct subset of the training data. After backward propagation, gradient synchronization is carried out via all-reduce.

Parameters and optimizer states are further sharded using DeepSpeed's ZeRO optimizer, allowing efficient distributed training with advanced memory optimizations (Smith et al., 2022).

2.4 Multi-Dimensional and MoE Parallelism

Megatron-LM supports decoupling parallelism axes for different layer types:

• Parallel Folding separates (TP, CP, DP, PP) for dense layers and (ETP, EP, EDP, PP) for MoE layers, allowing attention and MoE blocks to be optimally mapped across hardware (Yan et al., 8 Mar 2026).
• Expert Parallelism for MoE: expert parameters are sharded using their own parallel groups, and custom all-to-all dispatchers schedule tokens to experts efficiently, often with group-based overlap for communication/computation.

A summary table of main parallelism axes appears below.

Parallelism Axis	Function	Principal Use
TP (Tensor)	Intra-layer weight partitioning	Dense layers
PP (Pipeline)	Inter-layer model partitioning	Layer striping
DP (Data)	Replica batch splitting	All layers
EP (Expert)	Expert weight and compute sharding	MoE/multi-expert
CP (Context)	Long sequence partitioning across heads	Long-context LMs

3. System-Level Optimizations and Extensions

Megatron-LM and its ecosystem incorporate targeted optimizations to balance memory, computation, and communication:

• Memory: reduced-precision activations (FP8, FP4), fine-grained activation offloading (layer outputs to DRAM), mixed-precision optimizer states (BF16/FP16), and selective recomputation (strategic activation checkpointing).
• Computation: grouped GEMM fusion, custom CUDA kernels for routers and MLP, CUDA graphs for forward/backward pass capture, and kernel fusion for operator efficiency.
• Communication: high-performance all-to-all for expert routing (DeepEP, HybridEP), scatter/gather optimization for cross-node point-to-point, token-based dispatchers to merge communication phases, and overlap of communication/computation using multiple CUDA streams.

These optimizations enable per-GPU performance achieving $1,233$ TFLOPs on GB300 hardware (DeepSeek-V3-685B, MoE) and up to $1,048$ TFLOPs on GB200 for similar models (Yan et al., 8 Mar 2026).

4. Training at Trillion-Parameter Scale

Megatron-LM has been used to train models at unprecedented scale:

• MT-NLG 530B (Megatron-Turing NLG): trained with 3D parallelism ($P_m \times P_p \times P_d$), achieving 126–113 TFLOPS/GPU and 90%+ parallel efficiency from 2,240 to 3,360 A100 GPUs. Memory was reduced using ZeRO stage 2, fused kernels, and activation offload, with a custom 339B-token corpus and rigorous data deduplication pipeline (Smith et al., 2022).
• 1T-parameter GPT: Trained on 3,072 A100 GPUs, achieving 502 PFLOPS aggregate (52% of peak), with end-to-end throughput and scaling empirically validated up to 1T parameters (Narayanan et al., 2021).

Models such as DeepSeek V3-685B and Qwen3-235B demonstrate this scale with TFLOPS bandwidths validated on early Blackwell GPUs, and various large-scale MoE models have also been realized using the Megatron-Core MoE extension (Yan et al., 8 Mar 2026).

5. System Management and Diagnostics: MegatronApp

Distributed training at scale exposes new challenges in diagnosability, efficiency, and reliability. The MegatronApp toolchain introduces four composable modules for comprehensive management (Zhao et al., 26 Jul 2025):

• MegaScan: Fine-grained, Chrome-tracing-compatible operator tracing with straggler detection via heuristics on DP/TP/PP groups; enables root-cause isolation of slowdowns and cross-rank performance anomalies.
• MegaFBD: Decouples forward and backward execution, enabling heterogeneity-aware scheduling (e.g., assigning forward passes to slower devices where compute is light but communication is heavy, and vice versa). Fosters high device utilization on mixed hardware.
• MegaDPP: Adaptive scheduling and ordering for 3D parallelism, specializing placement to maximize pipeline throughput and minimize bubble, including "standby" scheduling when idle resources are available.
• MegaScope: Real-time activation and attention visualization, with live capture designed to minimize I/O overhead.

These modules integrate with Megatron-LM via a runtime flag and provide low-overhead diagnostic capability. For instance, MegaScan's operator-level tracing can be visualized in Chrome, and its ranking script outputs candidate bottleneck ranks.

6. Runtime Adaptation and Hybrid Orchestration: Galvatron

Galvatron extends Megatron-LM and DeepSpeed with dynamic, profiler-driven selection of parallelism configurations. This system continuously monitors cluster hardware and model characteristics, recomputes optimal DP/TP/PP splits, and reconfigures process groups on the fly to adapt to utilization or communication bottlenecks (Gumaan, 13 Mar 2025). Empirical results indicate Galvatron achieves 1.2–1.3× gains in throughput over static best practices and near-linear scaling to hundreds of GPUs.

7. Applications, Empirical Results, and Best Practices

Megatron-LM-based systems have set state-of-the-art results in both zero-shot and fine-tuned settings including WikiText-103, LAMBADA, RACE, and large-scale RLHF/LLM training. Empirical tuning guidelines derived from Megatron’s scaling studies suggest:

• Use intra-server tensor parallelism before pipeline parallelism to avoid slow inter-node all-reduce.
• Set pipeline stages to maximize per-GPU memory fit; use data parallelism to fill larger clusters.
• Tune microbatch size for balance between pipeline bubble and GEMM efficiency.
• Enable activation recomputation to unlock higher global batch sizes.
• For long-context or MoE LMs, decouple attention and expert layer parallelism for optimal shard utilization.

A plausible implication is that the composable architecture and tooling of the Megatron-LM ecosystem are accelerating the development, training, and deployment of ever-larger LLMs, with empirical scalability validated up to tens of thousands of GPUs for both dense and sparse (MoE) models (Narayanan et al., 2021, Yan et al., 8 Mar 2026).