If you have ever tried to train a large language model on a single GPU and watched it crash with an out-of-memory error, you already know the problem. Models that matter today — the ones with tens or hundreds of billions of parameters — simply do not fit on one device. Megatron-LM is NVIDIA's answer to that problem, and it has been quietly powering some of the most serious LLM research and production training runs in the world.
This article walks you through what Megatron-LM actually is, how its internals work, and how to go from zero to a real training run.
What Is Megatron-LM, and Why Should You Care
The repository at github.com/NVIDIA/Megatron-LM contains two things that are worth keeping distinct in your head.
Megatron-LM is the research-oriented layer — it bundles ready-to-run training scripts for GPT, BERT, T5, LLaMA, and multimodal models. If you want to get something running quickly without writing glue code from scratch, this is your starting point.
Megatron Core is the composable library underneath. It exposes the GPU-optimized building blocks (attention layers, parallelism strategies, optimizers, dataset loaders) as importable Python modules so you can assemble your own training framework on top of them. If you are a framework engineer or ML infrastructure developer, this is the layer you actually care about.
The benchmark numbers give you a sense of the scale this is designed for. The codebase has been used to train models ranging from 2B to 462B parameters across thousands of H100 GPUs, achieving up to 47% Model FLOP Utilization. That number matters because MFU is how you measure how efficiently you are using the hardware you paid for.
Understanding the Parallelism — This Is the Core Idea
Before you touch a line of code, you need to understand how Megatron distributes a model that does not fit in memory. There are four strategies, and Megatron uses them in combination.
Tensor Parallelism (TP) splits individual weight matrices across GPUs. A single attention layer's weight matrix gets sliced column-wise or row-wise across multiple devices. During the forward pass, each GPU handles a slice of the computation, and a small all-reduce communication syncs the results. The math is elegant: for a linear layer Y = XA, you can split A column-wise so GPU 0 computes XA_0 and GPU 1 computes XA_1, then concatenate. This is why Megatron can fit a single transformer layer that is too large for one device.
Pipeline Parallelism (PP) assigns different layers to different GPUs. GPU 0 handles layers 0–7, GPU 1 handles layers 8–15, and so on. Data flows through the pipeline in micro-batches, and Megatron's interleaved scheduling keeps the GPUs from sitting idle while waiting for the previous stage to finish.
Data Parallelism (DP) is the familiar one: you run copies of the model on multiple GPUs with different data shards, then average the gradients. Megatron supports both standard DDP and a distributed optimizer that shards optimizer states to reduce memory further.
Context Parallelism (CP) is newer and specifically useful for long sequences. It distributes the sequence dimension across GPUs so you can train on sequences that would otherwise blow your memory budget. The recent Dynamic Context Parallelism feature pushes this further by adapting the parallel size per batch based on actual sequence lengths, yielding up to 1.48x speedup for variable-length training.
Expert Parallelism (EP) is relevant if you are training Mixture-of-Experts models like DeepSeek-V3 or Qwen3. Different experts in the MoE layer live on different GPUs.
In practice you combine these. A typical 70B model training run might use TP=4, PP=4, DP=8, giving you 128 GPUs working together. Megatron handles the communication scheduling so you do not have to wire it up yourself.
Project Structure at a Glance
The repository is well organized once you know where to look:
Megatron-LM/
├── megatron/
│ ├── core/ # The library — import this in your own code
│ │ ├── models/ # GPT, BERT, T5, multimodal
│ │ ├── transformer/ # Attention, MLP, layer building blocks
│ │ ├── tensor_parallel/
│ │ ├── pipeline_parallel/
│ │ ├── distributed/ # FSDP, DDP
│ │ ├── optimizer/
│ │ ├── datasets/
│ │ └── inference/
│ ├── training/ # High-level training scripts
│ └── post_training/ # Quantization, distillation, pruning
├── examples/ # Shell scripts for GPT, LLaMA, Mixtral, etc.
├── tools/ # Data preprocessing utilities
└── tests/
The split between megatron/core/ and megatron/training/ mirrors the Megatron Core vs Megatron-LM distinction described above. If you are building something custom, spend most of your time in core/. If you are running experiments, the examples/ directory has working shell scripts you can adapt.
Getting It Running
Prerequisites
You need:
- One or more NVIDIA GPUs (Ampere or later recommended; H100 for FP8 training)
- CUDA 12.x
- Python 3.10+
- PyTorch 2.x
- NCCL (usually installed with CUDA)
The easiest path is NVIDIA's NGC container, which ships with all dependencies pre-installed:
docker pull nvcr.io/nvidia/pytorch:24.01-py3
If you are not using a container, you will also need Transformer Engine for FP8 support and Flash Attention for efficient attention computation.
Installation
# Install Megatron Core with its language model dependencies
pip install --no-build-isolation megatron-core[mlm,dev]
# Clone the repo for examples and training scripts
git clone https://github.com/NVIDIA/Megatron-LM.git
cd Megatron-LM
pip install --no-build-isolation .[mlm,dev]
The --no-build-isolation flag matters here. Megatron builds some CUDA extensions during install and needs access to your environment's CUDA headers.
Your First Training Run
Once installed, the simplest way to verify everything works is:
# Run a distributed training loop on 2 GPUs with mock data
torchrun --nproc_per_node=2 examples/run_simple_mcore_train_loop.py
If you want to try something closer to a real use case, the LLaMA-3 example uses 8 GPUs with FP8 precision:
./examples/llama/train_llama3_8b_fp8.sh
These scripts handle the argument wiring for you. Once you understand what the flags mean, you can start modifying them.
Data Preparation
This is where most people get tripped up the first time. Megatron does not consume raw text files. It needs data preprocessed into a binary indexed format (.bin + .idx files) for memory-mapped, high-throughput loading during training.
Step 1: Format Your Data as JSONL
Each line is a JSON object with a text field:
{"text": "Your first training document goes here."}
{"text": "Each document is a separate JSON line."}
{"text": "The tokenizer will handle the rest."}
Step 2: Run the Preprocessor
python tools/preprocess_data.py \
--input data.jsonl \
--output-prefix /path/to/processed_data \
--tokenizer-type HuggingFaceTokenizer \
--tokenizer-model /path/to/your/tokenizer \
--workers 8 \
--append-eod
Key flags to know:
-
--output-prefix: The path prefix for the.binand.idxoutput files. Pass this same prefix to--data-pathin your training script. -
--tokenizer-type: UseHuggingFaceTokenizerfor any tokenizer that follows the Hugging Face interface.GPT2BPETokenizeris available for GPT-2's original BPE tokenizer. -
--workers: Parallelizes tokenization. Set this to the number of CPU cores you can spare. -
--append-eod: Adds an end-of-document token between documents. Almost always what you want.
This step can take a while for large datasets. The output is worth the wait — the indexed binary format lets training data loaders access tokens at random offsets in O(1) time, which is critical when you have trillions of tokens.
Writing a Custom Training Loop with Megatron Core
If you want to plug Megatron's parallelism into your own training code rather than using the provided scripts, here is the minimal structure. This is what examples/run_simple_mcore_train_loop.py does under the hood.
import torch
from megatron.core import parallel_state
from megatron.core.models.gpt import GPTModel
from megatron.core.transformer.transformer_config import TransformerConfig
# Initialize distributed and parallelism groups
parallel_state.initialize_model_parallel(
tensor_model_parallel_size=2,
pipeline_model_parallel_size=1,
)
# Define your model configuration
config = TransformerConfig(
num_layers=12,
hidden_size=768,
num_attention_heads=12,
use_cpu_initialization=True,
pipeline_dtype=torch.float16,
)
# Build the model — Megatron handles sharding based on parallel_state
from megatron.core.models.gpt.gpt_layer_specs import get_gpt_layer_local_spec
model = GPTModel(
config=config,
transformer_layer_spec=get_gpt_layer_local_spec(),
vocab_size=50257,
max_sequence_length=1024,
)
# From here, training looks mostly like standard PyTorch
optimizer = torch.optim.Adam(model.parameters(), lr=1e-4)
The key thing to notice is that after initialize_model_parallel, the model construction is aware of which GPU it is running on. When you call GPTModel(...), Megatron automatically places the right layers on the right devices based on your TP and PP settings. You do not have to manually slice weight matrices.
Parallelism Configuration in Practice
When you run the provided shell scripts, you will see flags like these:
--tensor-model-parallel-size 4 \
--pipeline-model-parallel-size 4 \
--num-layers 32 \
--hidden-size 4096 \
--num-attention-heads 32 \
--micro-batch-size 2 \
--global-batch-size 1024 \
--seq-length 4096 \
A few rules of thumb:
Tensor parallel size should evenly divide num-attention-heads and the hidden size's intermediate MLP size. For a model with 32 attention heads, TP sizes of 1, 2, 4, 8, or 16 all work.
Pipeline parallel size should evenly divide num-layers. 32 layers with PP=4 means 8 layers per pipeline stage.
Global batch size equals micro-batch-size * gradient-accumulation-steps * data-parallel-size. Megatron enforces this math and will error if your settings are inconsistent.
Communication overlap flags are worth enabling once things are working. These flags let Megatron overlap gradient reduction with backward computation:
--overlap-grad-reduce \
--overlap-param-gather \
--tp-comm-overlap
They typically improve throughput by 5–15% at large scale.
Checkpoint Conversion: Getting In and Out of the Megatron Format
One practical concern is that Megatron's checkpoint format is not the same as Hugging Face's. If you want to start from a pretrained HuggingFace model or publish your trained weights back, you need to convert.
The Megatron Bridge project handles this bidirectionally. It supports popular models and is the recommended path for production checkpoint management.
For LLaMA specifically, the Megatron-LM repository includes conversion scripts under tools/ that have been used for Llama, Mistral, and other Llama-derived architectures.
FP8 Training
If you are running on H100 or later GPUs, FP8 training is worth trying. It reduces memory usage and can increase throughput significantly because Tensor Cores on Hopper are much faster in FP8 than BF16.
Enabling it is mostly a matter of setting flags:
--fp8-format hybrid \
--fp8-amax-history-len 1024 \
--fp8-amax-compute-algo max
You also need Transformer Engine installed, which handles the actual FP8 kernels. The train_llama3_8b_fp8.sh example shows a working configuration if you want to see all the pieces together.
MoE Models
The MoE (Mixture of Experts) support in Megatron Core has been one of the most active development areas recently. DeepSeek-V3 and Qwen3-style MoE architectures are explicitly supported, and there is an active roadmap for further optimizations through 2025.
The key additional parallelism dimension for MoE is Expert Parallelism:
--expert-model-parallel-size 8 \
--num-experts 64 \
--moe-router-topk 2
With 64 experts and EP=8, each GPU holds 8 experts. During the forward pass, tokens are routed to their top-2 experts, and the MoE communication layer handles moving tokens across GPU boundaries as needed.
Summary
The thing that makes Megatron-LM worth learning is not any single feature — it is that the entire system is designed around the constraint that you are always working at the edge of what the hardware can do.
Top comments (1)
Megatron-LM is the research-oriented layer — it bundles ready-to-run training scripts for GPT, BERT, T5, LLaMA, and multimodal models. If you want to get something running quickly without writing glue code from scratch, this is your starting point.