Megatron-LM

Overview#

Megatron-LM is both a paper series (Shoeybi 2019, Narayanan 2021, Korthikanti 2022) and an open-source repository at github.com/NVIDIA/Megatron-LM. The first paper introduced tensor parallelism and showed 8.3B-parameter training inside a single DGX-2 (16x V100) box. The second introduced interleaved 1F1B pipeline parallelism and demonstrated near-linear weak scaling to 3,072 A100s for a 1T-parameter target. The third added sequence parallelism and selective activation recomputation, removing the last big activation-memory hotspot and cutting recompute cost by 5x.

Functionally, Megatron-LM is a CUDA-Python codebase wrapping PyTorch with custom autograd-aware collectives (NCCL AllReduce / AllGather / ReduceScatter / P2P), a Transformer Engine integration for FP8 and FP4 math on Hopper and Blackwell, a Distributed Adam optimiser that shards moment state across the DP group (ZeRO-1 in spirit), and a corpus of launch scripts under `examples/` that show how the pieces compose for GPT, BERT, T5, RETRO, Llama, Mistral, Mixtral and DeepSeek-style architectures.

By 2026 Megatron Core (the refactored library form, `megatron.core`) is the embedded engine inside NVIDIA NeMo Framework, NeMo-Aligner for SFT/DPO/RLHF, the NVIDIA Nemotron training pipeline, Pax and MaxText derivatives on Google Cloud, Colossal-AI's high-performance path, and many internal lab forks. If your training run is north of ~30B parameters and on NVIDIA hardware, you are almost certainly running Megatron primitives — directly or one wrapper away. Yobitel NeoCloud customers training 70B+ models commonly use Megatron-LM (or NeMo, which wraps Megatron Core) as the default pretraining engine on multi-node H100, H200, and B200 training pods in the UK and EU regions.

This entry documents the production surface: the CLI and flag set, the four parallelism axes plus optimiser sharding, the data-pipeline contract, sizing tables at the common scales (70B / 175B / 405B), the recommended Hopper and Blackwell recipes, and the migration paths from FSDP, DeepSpeed and NeMo. This entry helps you choose and operate Megatron-LM for training pods on Yobitel NeoCloud or your own multi-GPU cluster.

Quick start#

The example below pretrains a GPT-3-style 1.3B model on 8x H100 SXM5 using tensor parallelism within the node, BF16 weights, FlashAttention-3, sequence parallelism on the TP group, and the Megatron Distributed Adam optimiser. The first block clones Megatron-LM, installs apex and Transformer Engine, and preprocesses a sample corpus into the indexed binary format Megatron expects. The second block launches the pretraining run with `torchrun`. The third block converts the resulting checkpoint to HuggingFace format for downstream serving.

# 1. Clone Megatron-LM and prepare a sample indexed dataset
git clone https://github.com/NVIDIA/Megatron-LM && cd Megatron-LM
pip install -r requirements.txt
pip install "transformer-engine[pytorch]" apex  # CUDA 12.4+

python tools/preprocess_data.py \
    --input ./data/oscar-sample.jsonl \
    --output-prefix ./data/oscar-sample \
    --vocab-file ./vocab/gpt2-vocab.json \
    --merge-file ./vocab/gpt2-merges.txt \
    --tokenizer-type GPT2BPETokenizer \
    --workers 32 --append-eod
# Produces oscar-sample_text_document.{bin,idx}

# 2. Pretrain a 1.3B GPT on 8x H100 with TP=2, DP=4, BF16, FA3, SP
GPUS_PER_NODE=8
torchrun --nproc_per_node=$GPUS_PER_NODE \
    pretrain_gpt.py \
    --num-layers 24 --hidden-size 2048 --num-attention-heads 16 \
    --seq-length 4096 --max-position-embeddings 4096 \
    --micro-batch-size 4 --global-batch-size 256 \
    --train-iters 50000 --lr 2.0e-4 --min-lr 2.0e-5 \
    --lr-decay-style cosine --lr-warmup-iters 2000 \
    --tensor-model-parallel-size 2 \
    --pipeline-model-parallel-size 1 \
    --sequence-parallel \
    --use-flash-attn \
    --bf16 \
    --use-distributed-optimizer \
    --recompute-granularity selective \
    --data-path ./data/oscar-sample_text_document \
    --vocab-file ./vocab/gpt2-vocab.json \
    --merge-file ./vocab/gpt2-merges.txt \
    --save ./checkpoints/gpt-1b3 --load ./checkpoints/gpt-1b3 \
    --save-interval 5000 --log-interval 10 \
    --tensorboard-dir ./tb

# 3. Convert the Megatron checkpoint to HuggingFace format for serving
python tools/checkpoint/convert.py \
    --model-type GPT --loader megatron --saver llama_mistral \
    --load-dir ./checkpoints/gpt-1b3 \
    --save-dir ./hf/gpt-1b3-hf \
    --tokenizer-model ./vocab/gpt2-vocab.json

How it works#

Megatron organises a training run around five orthogonal parallelism axes that compose multiplicatively. World size = DP x TP x PP x CP (sequence parallel is a sub-mode of TP, not a separate axis). Each axis trades a different cost: TP buys per-layer compute scaling at the cost of high-bandwidth intra-block AllReduce; PP buys layer-count scaling at the cost of pipeline-bubble idle time; DP buys batch-size scaling at the cost of one gradient AllReduce per step; CP buys sequence-length scaling at the cost of attention-time P2P traffic.

Inside the forward pass, the standard transformer block becomes: (1) column-parallel QKV projection — every TP rank computes a slice; (2) FA3 attention with the head dimension sharded across the TP group; (3) row-parallel output projection followed by an AllReduce (replaced by ReduceScatter + AllGather when sequence parallel is on); (4) column-parallel up-projection; (5) GELU/SwiGLU; (6) row-parallel down-projection with the same collective. Sequence parallelism sharded LayerNorm and dropout activations along the sequence dimension and replaces the per-block AllReduce with AllGather+ReduceScatter — same total bytes, but the activation that gets replicated drops by N=TP.

Across nodes, pipeline parallelism partitions the layer stack into K stages. The interleaved 1F1B schedule chops each stage into virtual sub-stages, reducing the bubble fraction to (K-1)/(M+K-1) where M is the number of micro-batches per pipeline fill. A 405B model on 1,024 H100s with TP=8, PP=16 and M=128 typically holds the bubble below 5 percent. P2P sends/receives between adjacent stages are NCCL-backed and overlap with the next forward / backward, which is why PP tolerates 400Gb InfiniBand bandwidth where TP demands NVLink.

The Distributed Adam optimiser (ZeRO-1 in DeepSpeed nomenclature) shards the BF16 master copy, FP32 first-moment and second-moment buffers across the DP group. Each rank computes the update for its 1/DP slice and AllGathers the resulting parameter delta. This drops optimiser-state memory from ~12 bytes/param per rank to ~12/DP — for a 70B BF16 model with DP=32 that is 26 GB per rank instead of 840 GB.

• Tensor parallel group: shards weights of one transformer block, AllReduce per block. TP=2/4/8 inside one NVLink island.
• Pipeline parallel group: layers split into K stages, P2P sends across stages. Interleaved 1F1B for minimal bubble.
• Sequence parallel: TP-group sub-mode; shards LayerNorm/dropout activations along sequence dim. Free at long context.
• Context parallel: shards the sequence itself across CP ranks; required at L > ~64K.
• Data parallel: replicates everything else; AllReduce gradients each step. Distributed Adam shards the optimiser inside the DP group.
• Transformer Engine: BF16/FP16 default, FP8 (E4M3/E5M2) on Hopper and Blackwell, FP4 (MXFP4) on Blackwell.
• Selective recomputation: only the largest activations are dropped and recomputed; ~30 percent FLOPs overhead for 5-10x activation-memory savings.

Reference and specifications#

Megatron-LM exposes its surface as a flag set on the `pretrain_gpt.py` / `pretrain_bert.py` / `pretrain_t5.py` entry-point scripts. The table below is the canonical reference for the flags that govern parallelism, precision, memory, and the optimiser as of Megatron-LM 0.8 / Megatron Core 0.11 (June 2026). Flags marked with an asterisk are also available as `TransformerConfig` fields in the Megatron Core library API.

Workload patterns#

Three workload shapes cover the bulk of Megatron-LM production usage: dense LLM pretraining from random weights at the 7B / 70B / 405B scales, continued pretraining of an open model on domain data, and large-scale SFT on instruction corpora. Each has its own parallelism and precision profile, and each maps cleanly to a Yobitel NeoCloud training-pod size — Pattern A on a 32-node (256-GPU) H100 pod with InfiniBand NDR, Pattern B on a 16-node (128-GPU) pod, Pattern C on a single 8x H100 node.

Pattern A — Dense pretraining of a Llama-style 70B from random weights on a 256-GPU H100 cluster (the canonical Yobitel NeoCloud 32-node training pod). Pattern B — Continued pretraining of Llama-3.1 70B on 100B tokens of domain corpus, starting from a HuggingFace checkpoint converted to Megatron format, on a 16-node NeoCloud pod. Pattern C — Large-scale supervised fine-tuning of Llama-3.1 8B on a 5M-example instruction set using packing, on a single 8x H100 NeoCloud node.

# A — Llama-style 70B pretrain on 32 nodes x 8 H100 (256 GPUs)
#     TP=8 (intra-node NVLink) x PP=4 (cross-node IB) x DP=8
torchrun --nproc_per_node=8 --nnodes=32 \
    --rdzv_backend=c10d --rdzv_endpoint=${MASTER_ADDR}:29500 \
    pretrain_gpt.py \
    --num-layers 80 --hidden-size 8192 --num-attention-heads 64 \
    --num-query-groups 8 --ffn-hidden-size 28672 \
    --seq-length 8192 --max-position-embeddings 8192 \
    --position-embedding-type rope --swiglu --normalization RMSNorm \
    --micro-batch-size 1 --global-batch-size 1024 \
    --train-iters 480000 \
    --tensor-model-parallel-size 8 \
    --pipeline-model-parallel-size 4 \
    --virtual-pipeline-model-parallel-size 5 \
    --sequence-parallel --use-flash-attn \
    --bf16 --fp8 hybrid \
    --use-distributed-optimizer \
    --overlap-grad-reduce --overlap-param-gather \
    --recompute-granularity selective \
    --lr 1.5e-4 --min-lr 1.5e-5 --lr-decay-style cosine --lr-warmup-iters 4000 \
    --data-path ${DATA_PREFIX} \
    --tokenizer-type Llama3Tokenizer \
    --save ${CKPT_DIR} --load ${CKPT_DIR} --save-interval 2000 \
    --ckpt-format torch_dist

# B — Continued pretrain of Llama-3.1 70B on domain corpus, 16 nodes
#     Start from a HF -> Megatron-converted checkpoint
python tools/checkpoint/convert.py \
    --model-type GPT --loader llama_mistral --saver megatron \
    --load-dir ${HF_DIR} --save-dir ${MEGATRON_INIT} \
    --target-tensor-parallel-size 8 --target-pipeline-parallel-size 2

torchrun --nproc_per_node=8 --nnodes=16 ... \
    pretrain_gpt.py \
    --load ${MEGATRON_INIT} --no-load-optim --no-load-rng \
    --finetune \
    --global-batch-size 512 --micro-batch-size 1 \
    --train-iters 100000 --lr 5e-5 --min-lr 5e-6 \
    --tensor-model-parallel-size 8 --pipeline-model-parallel-size 2 \
    --sequence-parallel --use-flash-attn --bf16 \
    --use-distributed-optimizer

# C — SFT of Llama-3.1 8B on instruction corpus, single 8x H100 node
#     Use NeMo or NeMo-Aligner for the loss/packing surface
torchrun --nproc_per_node=8 \
    examples/nlp/language_modeling/megatron_gpt_finetune.py \
    --config-path conf --config-name megatron_gpt_sft \
    model.restore_from_path=${LLAMA_8B_CKPT} \
    model.tensor_model_parallel_size=2 \
    model.pipeline_model_parallel_size=1 \
    model.data.train_ds.packed_sequence=True \
    model.data.train_ds.max_seq_length=8192 \
    model.optim.lr=2e-6 \
    trainer.max_steps=20000

Sizing and capacity planning#

The two questions that drive Megatron sizing are: which parallelism shape fits the model, and how many GPUs do I need to finish in a fixed wall-clock? The table below gives reference parallelism configs and observed throughput on H100 SXM5 clusters with 400Gb InfiniBand NDR, BF16+FP8 mixed precision, FA3 attention, sequence parallel on, distributed optimiser on, and selective recomputation. Throughput is per-GPU sustained tokens-per-second from internal training-run telemetry and the published Nemotron and NVIDIA H100 MLPerf records; treat as planning anchors.

• Weak-scaling efficiency on H100 with the above recipe holds at 85-92 percent from 256 to 4,096 GPUs; above 4,096 it drops to 75-85 percent as DP AllReduce starts to dominate.
• Move to Blackwell (B200) and FP4 MXFP for a roughly 2.7x throughput uplift on dense 70B-class training versus the same shape on H100 FP8.
• MoE models add expert parallelism (EP) inside the TP group; EP=8 is the typical choice for 8-expert top-2 routing. Expert AllToAll dominates at EP>8 unless on NVLink Switch.
• Context parallelism becomes necessary above ~64K sequence length; below that, sequence parallel inside TP is the simpler choice.

Limits and quotas#

Megatron itself has few hard limits; what bounds a run are GPU memory, NCCL group counts, and the data-loader contract. The table below summarises the constraints worth knowing before designing a parallelism shape.

Observability#

Megatron-LM emits TensorBoard scalars and optional Weights & Biases logs covering per-step loss, learning rate, gradient norm, samples-per-second, TFLOPs-per-GPU, FP8 amax, optimiser-state norms, and (under --log-memory-to-tensorboard) GPU memory peaks per phase. The metrics that matter operationally at training-cluster scale are throughput stability, gradient-norm sanity, and FP8 scaling-factor health.

On Hopper / Blackwell, pair Megatron logs with NVIDIA DCGM exporter (`DCGM_FI_DEV_GPU_UTIL`, `DCGM_FI_DEV_MEM_COPY_UTIL`, `DCGM_FI_PROF_SM_ACTIVE`, `DCGM_FI_PROF_PIPE_TENSOR_ACTIVE`, `DCGM_FI_PROF_NVLINK_TX_BYTES`) and NCCL profiling (`NCCL_DEBUG=INFO`, `NCCL_DEBUG_SUBSYS=COLL`). The three signals that detect 90 percent of training-cluster problems are: iteration time variance, NCCL AllReduce p99 latency, and per-rank step-loss divergence.

• iteration-time / samples-per-second: holds within +-3 percent in steady state; >10 percent dip means stragglers, IB packet loss, or thermal throttling.
• TFLOPs-per-GPU: should hit 40-60 percent of the device peak; drops correlate with attention shape or recompute config.
• grad-norm: spikes above 10x baseline indicate divergence; correlate with LR schedule.
• loss curve: per-rank loss should match within numerical noise; divergence means a DP rank dropped or got bad data.
• FP8 amax_history_max: persistent saturation indicates clip-and-rescale should fire; never-saturating means FP8 is wasted.
• DCGM PIPE_TENSOR_ACTIVE: the most honest tensor-core utilisation signal; pair with SM_ACTIVE.
• NCCL_DEBUG=WARN in steady state; flip to INFO when investigating slow steps.

# Prometheus rules for a Megatron-LM pretraining job
groups:
  - name: megatron-training
    interval: 60s
    rules:
      - alert: MegatronIterationTimeRegression
        expr: |
          (avg_over_time(megatron:iteration_time_seconds[5m])
           / avg_over_time(megatron:iteration_time_seconds[1h] offset 30m)) > 1.10
        for: 10m
        labels: { severity: warning, team: training }
        annotations:
          summary: "Iteration time +10% vs 1h baseline on {{ $labels.job_name }}"

      - alert: MegatronGradNormSpike
        expr: megatron:grad_norm > 10 * avg_over_time(megatron:grad_norm[1h] offset 30m)
        for: 5m
        labels: { severity: warning }
        annotations:
          summary: "Grad-norm spike — investigate LR or data corruption"

      - alert: MegatronTFLOPsCollapse
        expr: avg_over_time(megatron:tflops_per_gpu[10m]) < 0.5 *
              avg_over_time(megatron:tflops_per_gpu[1h] offset 30m)
        for: 15m
        labels: { severity: critical }
        annotations:
          summary: "Per-GPU TFLOPs halved — check NCCL, FA3 kernel, recompute config"

      - alert: MegatronNCCLAllReduceP99
        expr: histogram_quantile(0.99,
                rate(nccl:allreduce_seconds_bucket[5m])) > 0.2
        for: 10m
        annotations:
          summary: "NCCL AllReduce p99 > 200ms — fabric or topology regression"

      - alert: MegatronFP8AmaxSaturation
        expr: megatron:fp8_amax_history_max / megatron:fp8_amax_history_threshold > 0.95
        for: 30m
        annotations:
          summary: "FP8 amax persistently saturated — recompute scale or fall back to BF16"

Cost and FinOps#

Megatron-LM cost is almost entirely a function of two numbers: GPU-hour rate and per-GPU sustained throughput. The training-cost-per-trillion-tokens for a given model size is `(cluster_GPUs x hours x hourly_rate)` where hours = (1T / (per_GPU_tok_s x cluster_GPUs x 3600)). The table below uses Yobitel UK list pricing (June 2026) for on-demand H100 SXM5 ($3.10/GPU-hour), reserved H100 ($1.95), and B200 ($5.50 on-demand) and the throughput anchors from the Sizing table.

• Reserved capacity (1-3 year terms) drops the rate ~35-40 percent for predictable multi-month pretraining runs.
• Blackwell B200 doubles to triples per-GPU tok/s on dense models; the rate premium pays back within the first month of training.
• FP8 weights + FP8 activations cut wall-clock 1.5-1.8x on H100 versus pure BF16; FP4 on Blackwell another 1.5x.
• Checkpoint storage is a non-trivial line item — a 405B checkpoint with Adam state is ~5 TB; sharded async writes (`--ckpt-format torch_dist`) and pruning intermediate checkpoints matter.
• Failed/restarted iterations are the silent cost. A 10-percent restart rate doubles the schedule slip. Spend on storage, network, and DCGM monitoring before spending on more GPUs.

Security and compliance#

Megatron-LM has no built-in auth surface — it is a training script, not a serving system. Security controls apply at the cluster boundary: Slurm or Kubernetes job submission auth, private container registry for the NVIDIA / NeMo base images, encrypted weights-at-rest on the checkpoint filesystem (Lustre with kernel crypto or NetApp with NVE), and network isolation of the training fabric from public ingress. Training-time data does not leave the cluster; the only egress is checkpoint writes and metrics shipping.

For UK and EU sovereign training programmes — government foundation models, regulated-industry domain pretrains — Megatron-LM runs inside the same sovereign tenancies that satisfy NCSC Cloud Security Principles and G-Cloud 14 lot definitions. The framework itself is Apache 2.0 (with NVIDIA-specific clauses around brand and trademark) and has no telemetry call-home; the only external dependency at runtime is the W&B / TensorBoard sink if you configure one. For air-gapped environments, mirror the NGC and PyPI dependencies into the internal registry and pin Transformer Engine and Apex versions.

Migration and alternatives#

Three migration paths dominate. From FSDP/HSDP: the trigger is usually scale — FSDP at FULL_SHARD is competitive through ~30B, manageable to ~70B with HSDP, but loses ground above that as the per-layer AllGather amortises poorly. The move to Megatron buys you true tensor parallelism (no per-layer parameter materialisation) and interleaved 1F1B pipeline. From DeepSpeed: similar capability surface but DeepSpeed leans on ZeRO-3 for memory while Megatron leans on TP+PP for compute scaling. From NeMo: NeMo wraps Megatron Core, so 'migrating to Megatron' from NeMo is really 'dropping a layer of recipe ergonomics' — useful when you need to change something NeMo's Hydra configs do not expose.

Troubleshooting#

The error patterns below account for most production Megatron incidents observed on Yobitel-operated training fleets and the public NVIDIA NeMo issue tracker. Each row maps the observable symptom to the underlying mechanism and the minimum-viable fix.

Where this fits in the Yobitel stack#

Megatron-LM (via NeMo Framework) is the standard training engine on Yobitel sovereign GPU tenancies for any pretraining or continued-pretraining workload above 32 GPUs. Yobitel ships NGC-derived NeMo + Megatron Core containers, pre-validated NCCL and InfiniBand configurations for the H100, H200, and B200 fleets, and reference Slurm + Pyxis launch scripts that encode the parallelism-shape recommendations from this entry's Sizing table.

For customers training UK or EU sovereign foundation models — government, defence, regulated-industry — Megatron runs inside London-1 and Frankfurt-1 tenancies satisfying NCSC Cloud Security Principles, G-Cloud 14 lot definitions and the OFFICIAL handling caveat. Weights never leave the customer tenancy; the Yobitel control plane handles only metrics, scheduling, and health telemetry. Where SFT / DPO / RLHF is needed downstream, the same checkpoints flow into NeMo-Aligner on the same cluster without re-sharding.

Yobibyte, Yobitel's managed AI-native platform, exposes Megatron-trained models for inference via its vLLM and TensorRT-LLM endpoints; InferenceBench v3 then scores throughput and latency on the deployed checkpoints across the H100, H200, B200 and MI300X SKUs the customer rents. The training-to-serving handoff is one HF-format conversion away.