Megatron-LM Documentation

Ongoing research training transformer models at scale

17,380 stars Python Python 1 file · ~4,865 tokens #large-language-models#model-para#transformers
RAW Doc

Docker & Self-Hosted Deployment Verified OCI

Deploy Megatron-LM on any server, homelab, or cloud container platform.

1-Click Terminal Command
bash (docker run) 
docker run -d --name megatron-lm --restart always nvidia/megatron-lm:latest

File: README.md

Megatron-LM and Megatron Core

GPU-optimized library for training transformer models at scale



About

This repository contains two components: Megatron-LM and Megatron Core.

Megatron-LM is a reference example that includes Megatron Core plus pre-configured training scripts, ideal for research teams, learning distributed training, and quick experimentation.

Megatron Core is a composable library with GPU-optimized building blocks for custom training frameworks. It provides transformer building blocks, advanced parallelism strategies (TP, PP, DP, EP, and CP), mixed precision support (FP16, BF16, FP8, and FP4), and model architectures, ideal for framework developers and ML engineers building custom training pipelines.

Megatron Bridge provides bidirectional Hugging Face ↔ Megatron checkpoint conversion with production-ready recipes.

Getting Started

Install from PyPI:

bash
uv pip install megatron-core

Or clone and install from source:

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

Note: Building from source can use a lot of memory. If the build runs out of memory, limit parallel compilation jobs by setting MAX_JOBS (for example, MAX_JOBS=4 uv pip install -e .).

For NVIDIA GPU Cloud (NGC) container setup and all installation options, review the Installation Guide.

Latest News

Previous News

Project Structure

text
Megatron-LM/
├── megatron/
│   ├── core/                    # Megatron Core (kernels, parallelism, building blocks)
│   │   ├── models/              # Transformer models
│   │   ├── transformer/         # Transformer building blocks
│   │   ├── tensor_parallel/     # Tensor parallelism
│   │   ├── pipeline_parallel/   # Pipeline parallelism
│   │   ├── distributed/         # Distributed training (FSDP, DDP)
│   │   ├── optimizer/           # Optimizers
│   │   ├── datasets/            # Dataset loaders
│   │   ├── inference/           # Inference engines and server
│   │   └── export/              # Model export (example: TensorRT-LLM)
│   ├── training/                # Training scripts
│   ├── legacy/                  # Legacy components
│   ├── post_training/           # Post-training (quantization, distillation, pruning, etc.)
│   └── rl/                      # Reinforcement learning (including RLHF)
├── examples/                    # Ready-to-use training examples
├── tools/                       # Utility tools
├── tests/                       # Comprehensive test suite
└── docs/                        # Documentation

Performance Benchmarking

For the latest performance benchmarking results, refer to NVIDIA Megatron Bridge Performance Summary.

The codebase efficiently trains models from 2B to 462B parameters across thousands of GPUs, achieving up to 47% Model FLOP Utilization (MFU) on H100 clusters.

Benchmark Configuration:

  • Vocabulary size: 131,072 tokens
  • Sequence length: 4,096 tokens
  • Model scaling: Varied hidden size, attention heads, and layers to achieve target parameter counts
  • Communication optimizations: Fine-grained overlapping with DP (--overlap-grad-reduce, --overlap-param-gather), TP (--tp-comm-overlap), and PP (enabled by default)

Key Results:

  • 6,144 H100 GPUs: Successfully benchmarked 462B parameter model training.
  • Superlinear scaling: MFU increases from 41% to 47–48% with model size.
  • End-to-end measurement: Throughputs include all operations (data loading, optimizer steps, communication, and logging).
  • Production ready: Full training pipeline with checkpointing and fault tolerance.
  • Note: Performance results measured without training to convergence

Weak Scaling Results

The weak scaled results show superlinear scaling (MFU increases from 41% for the smallest model considered to 47–48% for the largest models); this is because larger GEMMs have higher arithmetic intensity and are consequently more efficient to execute.

Strong Scaling Results

This test strong scales the standard GPT-3 model (slightly more than 175 billion parameters due to larger vocabulary size) from 96 H100 GPUs to 4,608 GPUs, using the same batch size of 1,152 sequences throughout. Communication becomes more exposed at larger scale, leading to a reduction in MFU from 47% to 42%.

Roadmaps

Resources

Getting Help

  • 📖 Documentation - Official guides and API reference
  • 🐛 Issues - Bug reports and feature requests

Contributing

Contributions are welcome. Ways to contribute:

  • 🐛 Report bugs - Help improve reliability
  • 💡 Suggest features - Shape the future of Megatron Core
  • 📝 Improve docs - Make Megatron Core more accessible
  • 🔧 Submit PRs - Contribute code improvements

→ Contributing Guide

Citation

If you use Megatron in your research or project, use the following citation:

bibtex
@article{megatron-lm,
  title={Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism},
  author={Shoeybi, Mohammad and Patwary, Mostofa and Puri, Raul and LeGresley, Patrick and Casper, Jared and Catanzaro, Bryan},
  journal={arXiv preprint arXiv:1909.08053},
  year={2019}
}

File: docs/advanced/index.md

Discussions

In-depth technical discussions and optimization guides:

File: docs/api-guide/core/dist_checkpointing.md

dist_checkpointing package

A library for saving and loading the distributed checkpoints.
A distributed checkpoint in Megatron Core uses the torch_dist format,
a custom checkpointing mechanism built on top of PyTorch's native
checkpointing capabilities.

A key property of distributed checkpoints is that a checkpoint saved under one
parallel configuration (tensor, pipeline, or data parallelism) can be loaded
under a different parallel configuration. This enables flexible scaling and
resharding of models across heterogeneous training setups.

Using the library requires defining sharded state_dict dictionaries with functions from mapping and optimizer modules.
Those state dicts can be saved or loaded with a serialization module using strategies from strategies module.

Safe Checkpoint Loading

Since PyTorch 2.6, the default behavior of torch.load is weights_only=True.
This ensures that only tensors and allow-listed classes are loaded, reducing the risk of arbitrary code execution.

If you encounter an error such as:

bash
WeightsUnpickler error: Unsupported global: GLOBAL argparse.Namespace was not an allowed global by default.

you can fix it by explicitly allow-listing the missing class in your script:

python
import torch, argparse

torch.serialization.add_safe_globals([argparse.Namespace])

Checkpointing Distributed Optimizer

Checkpoint Compatibility and Optimizer State Formats

Beginning with mcore v0.14, the flattened_range attribute was removed from dist_checkpointing. As a result:

  • Optimizer states saved with mcore versions <= 0.14 can no longer be loaded directly. Loading these legacy optimizer states is not supported because the required sharded metadata is no longer available. If you need to continue training from older checkpoints, refer to the workaround described below.
  • Model weights from older checkpoints remain fully compatible. No extra steps are needed—model weights from checkpoints created by earlier versions load automatically; simply add the --no-load-optim flag.

Workaround: Loading legacy optimizer states with ToT MCore

Step 1: Convert the legacy checkpoint using mcore v0.15.0

Run a dummy training job with mcore v0.15.0 to re-save the checkpoint with new optimizer states format.

bash
MODEL_TRAIN_PARAMS=(
    # Define model architecture and training parameters here
)
OLD_CKPT=/workspace/mcore_ckpt_old
CONVERTED_CKPT=/workspace/mcore_ckpt_0.15.0

torchrun --nproc_per_node=8 /opt/megatron-lm/pretrain_gpt.py \
   --save-interval 1 \
   --eval-interval 1 \
   --exit-interval 1 \
   --eval-iters 1 \
   --use-distributed-optimizer \
   --save ${CONVERTED_CKPT} \
   --load ${OLD_CKPT} \
   --ckpt-format torch_dist \
   "${MODEL_TRAIN_PARAMS[@]}"

Step 2: Load the converted checkpoint with ToT MCore

Use the converted checkpoint as the input for continued training with ToT MCore.

bash
MODEL_TRAIN_PARAMS=(
    # Define model architecture and training parameters here
)
NEW_CKPT=/workspace/mcore_ckpt_new
CONVERTED_CKPT=/workspace/mcore_ckpt_0.15.0

torchrun --nproc_per_node=8 /opt/megatron-lm/pretrain_gpt.py \
   --use-distributed-optimizer \
   --save ${NEW_CKPT} \
   --load ${CONVERTED_CKPT} \
   --ckpt-format torch_dist \
   "${MODEL_TRAIN_PARAMS[@]}"

After this step, training can proceed normally using ToT MCore with fully supported optimizer state loading.

Distributed Optimizer Checkpoint Formats

The refactor of the Distributed Optimizer introduces two checkpoint formats:

  • dp_reshardable (Default)
    • Fast save/load performance.
    • Not reshardable — not possible to change model parallelism when using this format.
    • Recommended for general training when model parallelism changes are not needed.
  • fully_reshardable
    • Fully reshardable — supports arbitrary changes in model parallelism.
    • Slower than dp_reshardable.
    • Enabled via the --dist-ckpt-optim-fully-reshardable flag.

Workflow for Changing Model Parallelism

You can combine formats to optimize both flexibility and performance:

  1. Train using dp_reshardable (default) for faster checkpointing.

  2. When you need to change model parallelism:

    • Stop training.
    • Change model parallelism for train config.
    • Resume training with --dist-ckpt-optim-fully-reshardable.

  3. Save at least one checkpoint under the new model parallel configuration.

  4. (Optional) To continue the training with updated model parallelism and better checkpointing performance, stop training and switch back to dp_reshardable format by removing --dist-ckpt-optim-fully-reshardable.

Async Checkpoint Saving Strategy

The framework supports asynchronous checkpoint saving to improve training performance by offloading I/O operations.

We are transitioning to a new async saving implementation based on the NVRx (NVIDIA Resiliency Extension) package. The legacy async strategy (referred to as mcore) is being deprecated.

Migration to NVRx

  • The NVRx-based async strategy will become the default in mcore v0.17.
  • The existing mcore async strategy is deprecated and will be removed in future versions.
  • A deprecation warning is emitted when using the mcore strategy.

Selecting Async Strategy

--async-strategy flag is introduced to control the async strategy. To use legacy async strategy (mcore), set:

bash
--async-strategy mcore

NVRx Dependency

To use the NVRx-based async strategy, you must install the nvidia-resiliency-ext package.

bash
git clone https://github.com/NVIDIA/nvidia-resiliency-ext
cd nvidia-resiliency-ext
pip install .

NOTE

  • If nvidia-resiliency-ext is not installed, the NVRx async strategy will not be available.
  • The mcore strategy will remain temporarily to ensure a smooth transition but will be removed in future releases.
  • It is strongly recommended to migrate to the NVRx strategy as soon as possible.

Async Saving for `fsdp_dtensor` and `torch_dcp` checkpoints

Starting from mcore v0.17, asynchronous checkpoint saving is supported for fsdp_dtensor and torch_dcp formats.

Note that async saving for these formats requires the nvidia-resiliency-ext package. As a result, the only supported async_strategy in this context is nvrx.

Subpackages

{toctree}
:maxdepth: 4

dist_checkpointing.strategies

File: docs/api-guide/core/dist_checkpointing.strategies.md

dist_checkpointing.strategies package

Package defining different checkpoint formats (backends) and saving/loading algorithms (strategies).

Strategies can be used for implementing new checkpoint formats or implementing new (more optimal for a given use case) ways of saving/loading of existing formats.
Strategies are passed to dist_checkpointing.load and dist_checkpointing.save functions and control the actual saving/loading procedure.

File: docs/api-guide/core/distributed.md

distributed package

This package contains various utilities to finalize model weight gradients
on each rank before the optimizer step. This includes a distributed data
parallelism wrapper to all-reduce or reduce-scatter the gradients across
data-parallel replicas, and a finalize_model_grads method to
synchronize gradients across different parallelism modes (e.g., 'tied'
layers on different pipeline stages, or gradients for experts in a MoE on
different ranks due to expert parallelism).

File: docs/api-guide/core/fusions.md

fusions package

This package provides modules that provide commonly fused
operations. Fusing operations improves compute efficiency by
increasing the amount of work done each time a tensor is read from
memory. To perform the fusion, modules in this either rely on PyTorch
functionality for doing just-in-time compilation
(i.e. torch.jit.script in older PyTorch versions of torch.compile
in recent versions), or call into custom kernels in external libraries
such as Apex or TransformerEngine.

File: docs/api-guide/core/index.md

Core APIs

Low-level API reference for core Megatron components.

{toctree}
:maxdepth: 2

transformer
tensor_parallel
pipeline_parallel
fusions
distributed
datasets
dist_checkpointing
dist_checkpointing.strategies

File: docs/api-guide/core/pipeline_parallel.md

pipeline_parallel package

This package contains implementations for two different pipeline parallelism
schedules (one without interleaving and one with interleaving, see Efficient Large-Scale Language Model Training on GPU Clusters Using Megatron-LM
for details), and a default no-pipelining schedule. It also contains methods
for the point-to-point communication that is needed between pipeline stages.

File: docs/api-guide/core/tensor_parallel.md

tensor_parallel package

This package contains an implementation for tensor parallelism in transformer
models (see Megatron-LM: Training Multi-Billion Parameter Language Models
Using Model Parallelism and Reducing
Activation Recomputation in Large Transformer Models
for details).

Explore Related Technologies & Topics

View All 9,111+ Repositories →

Discover similar high-velocity repositories, agent skills, and OpenAPI specifications across the ecosystem.

#large-language-models Browse large-language-models catalog → #model-para Browse model-para catalog → #transformers Browse transformers catalog → ⚡ MCP Servers Model Context Protocol → 🧠 Agent Skills Prompt & rule files → 🔌 OpenAPI Specs REST & Swagger schemas →