Megatron-LM

Megatron-LM is NVIDIA's library for training transformer models at scale on GPU clusters. It bundles two parts: Megatron-LM, a reference example with pre-configured training scripts for research teams and quick experimentation, and Megatron Core, a composable library of GPU-optimized building blocks (transformer layers, parallelism strategies, mixed precision kernels) for engineers building custom training pipelines. It's built for ML infrastructure engineers, researchers training large language models, and teams working on mixture-of-experts or hybrid architectures who need control over distributed training internals.

Key features

• Parallelism strategies: tensor parallelism (TP), pipeline parallelism (PP), data parallelism (DP), expert parallelism (EP), and context parallelism (CP), combinable for multi-billion to multi-hundred-billion parameter models.
• Mixed precision support: FP16, BF16, FP8, and FP4 training to reduce memory and increase throughput on supported hardware.
• Transformer building blocks: modular components for assembling custom model architectures, including hybrid transformer-Mamba designs and MoE layers.
• Dynamic context parallelism: adaptive CP sizing for variable-length sequence training.
• Checkpoint interoperability: Megatron Bridge provides bidirectional Hugging Face ↔ Megatron checkpoint conversion with production recipes.
• Distributed training backends: FSDP and DDP support under megatron/core/distributed/.
• Post-training tools: quantization, distillation, and pruning utilities bundled in megatron/post_training/.
• RL and RLHF support: reinforcement learning components in megatron/rl/ for post-training alignment workflows.
• Inference and export: inference engines, a serving component, and export paths including TensorRT-LLM.

Ideal use cases

Megatron-LM fits teams training large transformer models (billions of parameters and up) across multi-GPU or multi-node GPU clusters, where you need fine control over parallelism and memory to hit target throughput. It's a strong fit for research groups experimenting with new architectures (MoE, hybrid transformer-Mamba) or optimizers, and for infrastructure teams building their own training stack on top of Megatron Core's building blocks.

It's not the right tool if you're fine-tuning a small model on a single GPU or a handful of GPUs, where simpler frameworks (Hugging Face Transformers, PyTorch Lightning) involve far less setup. It also assumes familiarity with distributed training concepts; teams new to multi-GPU training will face a learning curve around parallelism configuration, NCCL communication, and cluster scheduling.

Installation

Install Megatron Core from PyPI:

uv pip install megatron-core

Or clone the repository and install from source:

git clone https://github.com/NVIDIA/Megatron-LM.git
cd Megatron-LM
uv pip install -e .

Building from source can consume significant memory. If the build runs out of memory, limit parallel compilation jobs:

MAX_JOBS=4 uv pip install -e .

For NVIDIA GPU Cloud (NGC) container setup and the full list of installation options, check the official Installation Guide linked from the repository's documentation site. After installing, the documentation's "Your First Training Run" guide walks through end-to-end training with data preparation, and the "Parallelism Strategies" guide covers configuring TP, PP, DP, EP, and CP for your cluster size.