training-llms-megatron

Building a model from scratch at the largest sizes means splitting the work across thousands of graphics cards in several directions at once — this is the machinery that does the splitting, tuned so as little of that hardware as possible sits idle.

WHEN TO USE

SCOPE

USE FOR

Pretrain a 70B-plus parameter LLaMA-style model across a multi-node GPU cluster.
Train a sparse Mixture-of-Experts model like Mixtral with expert parallelism.
Push GPU efficiency toward 40%+ Model FLOP Utilization on H100s.
Choose the right mix of tensor, pipeline, and data parallelism for a given model size.
Diagnose low GPU utilization, out-of-memory errors, or a diverging loss curve.

NOT FOR

Training a model under 70B parameters with a simpler, PyTorch-native API; FSDP fits that better.
Wanting an easier setup for a sub-100B model; DeepSpeed is simpler to get running.
Prototyping quickly with a simpler workflow; HuggingFace Accelerate suits that better.
Learning how large-model training works from a single-file implementation; LitGPT is built for that.

This is a saved copy of the skill's SKILL.md instructions. It may differ from the upstream file or the version you install.

In this document (8)

Megatron-Core - Large-Scale LLM Training

Quick start

Megatron-Core trains LLMs from 2B to 462B parameters with up to 47% Model FLOP Utilization on H100 GPUs through advanced parallelism strategies.

Installation:

# Docker (recommended)
docker run --gpus all -it --rm nvcr.io/nvidia/pytorch:25.04-py3

# Or pip
pip install megatron-core

Simple distributed training:

# Train with 2 GPUs using data parallelism
torchrun --nproc_per_node=2 examples/run_simple_mcore_train_loop.py

# Or LLaMA-3 8B training
./examples/llama/train_llama3_8b_fp8.sh

Common workflows

Workflow 1: Train LLaMA-style model with 3D parallelism

Copy this checklist:

LLaMA Training Setup:
- [ ] Step 1: Choose parallelism configuration
- [ ] Step 2: Configure training hyperparameters
- [ ] Step 3: Launch distributed training
- [ ] Step 4: Monitor performance metrics

Step 1: Choose parallelism configuration

Model size determines parallelism strategy:

Model Size	GPUs	Tensor Parallel	Pipeline Parallel	Data Parallel	Context Parallel
7B	8	1	1	8	1
13B	8	2	1	4	1
70B	64	4	4	4	1
405B	128	8	8	2	2

Step 2: Configure training hyperparameters

#!/bin/bash
# train_llama_70b.sh

GPUS_PER_NODE=8
NNODES=8  # 64 GPUs total
TP=4      # Tensor parallel
PP=4      # Pipeline parallel
CP=1      # Context parallel

# LLaMA 70B configuration
MODEL_SIZE=70  # Billion parameters
HIDDEN_SIZE=8192
NUM_LAYERS=80
NUM_HEADS=64
SEQ_LENGTH=4096

# Training hyperparameters
MICRO_BATCH=1
GLOBAL_BATCH=1024
LR=3e-4

torchrun \
  --nproc_per_node=$GPUS_PER_NODE \
  --nnodes=$NNODES \
  pretrain_gpt.py \
  --tensor-model-parallel-size $TP \
  --pipeline-model-parallel-size $PP \
  --context-parallel-size $CP \
  --sequence-parallel \
  --num-layers $NUM_LAYERS \
  --hidden-size $HIDDEN_SIZE \
  --num-attention-heads $NUM_HEADS \
  --seq-length $SEQ_LENGTH \
  --max-position-embeddings $SEQ_LENGTH \
  --micro-batch-size $MICRO_BATCH \
  --global-batch-size $GLOBAL_BATCH \
  --lr $LR \
  --train-iters 100000 \
  --lr-decay-style cosine \
  --lr-warmup-iters 2000 \
  --weight-decay 0.1 \
  --clip-grad 1.0 \
  --bf16 \
  --use-mcore-models \
  --transformer-impl transformer_engine \
  --data-path /path/to/data \
  --vocab-file /path/to/vocab.json \
  --merge-file /path/to/merges.txt

Step 3: Launch distributed training

# Single node (8 GPUs)
bash train_llama_70b.sh

# Multi-node with SLURM
sbatch --nodes=8 --gpus-per-node=8 train_llama_70b.sh

Step 4: Monitor performance metrics

Key metrics to track:

Model FLOP Utilization (MFU): Target >40% on H100
Throughput: Tokens/sec/GPU
Memory usage: <80GB per GPU for 70B model
Loss: Should decrease steadily

Workflow 2: Configure Mixture of Experts (MoE) training

For sparse MoE models like Mixtral.

MoE Training:
- [ ] Step 1: Configure expert parallelism
- [ ] Step 2: Set MoE hyperparameters
- [ ] Step 3: Launch training with EP

Step 1: Configure expert parallelism

# Mixtral 8x7B example
TENSOR_PARALLEL=2
PIPELINE_PARALLEL=1
EXPERT_PARALLEL=4  # Split 8 experts across 4 GPUs
DATA_PARALLEL=4

TOTAL_GPUS=$((TENSOR_PARALLEL * PIPELINE_PARALLEL * EXPERT_PARALLEL * DATA_PARALLEL))
# = 2 * 1 * 4 * 4 = 32 GPUs

Step 2: Set MoE hyperparameters

torchrun \
  --nproc_per_node=8 \
  pretrain_gpt.py \
  --tensor-model-parallel-size 2 \
  --pipeline-model-parallel-size 1 \
  --expert-model-parallel-size 4 \
  --num-experts 8 \
  --moe-router-topk 2 \
  --moe-router-load-balancing-type aux_loss \
  --moe-aux-loss-coeff 0.01 \
  --hidden-size 4096 \
  --num-layers 32 \
  --num-attention-heads 32 \
  --seq-length 4096 \
  --max-position-embeddings 4096 \
  --bf16 \
  --use-mcore-models \
  --transformer-impl transformer_engine \
  --data-path /path/to/data \
  --vocab-file /path/to/vocab.json \
  --merge-file /path/to/merges.txt

Step 3: Launch training with EP

Expert parallelism distributes different experts across GPUs, reducing memory while maintaining capacity.

Memory without EP: 8 experts × 7B = 56GB per GPU
Memory with EP=4: 2 experts × 7B = 14GB per GPU
Savings: 75% memory reduction

Workflow 3: Optimize for maximum throughput

Achieve 47% MFU on H100.

Performance Optimization:
- [ ] Step 1: Enable Flash Attention
- [ ] Step 2: Use FP8 precision (H100)
- [ ] Step 3: Optimize micro-batch size
- [ ] Step 4: Tune parallelism degrees

Step 1: Enable optimizations

--use-mcore-models  # Use Megatron Core models
--transformer-impl transformer_engine  # Use Transformer Engine
--sequence-parallel  # Reduce activation memory (use with TP)

Step 2: Use FP8 precision (H100 only)

--fp8-hybrid  # FP8 mixed precision training
# Transformer Engine handles FP8 automatically

Result: 1.5-2x speedup on H100 vs BF16.

Step 3: Optimize micro-batch size

Find largest micro-batch that fits in memory:

# Start with 1, increase until OOM
for MBS in 1 2 4 8; do
  echo "Testing micro-batch-size=$MBS"
  torchrun ... --micro-batch-size $MBS
done

Typical values:

7B model: 4-8
70B model: 1-2
405B model: 1

Step 4: Tune parallelism degrees

Rules of thumb:

Tensor Parallel: Use ≤8 (limited by NVLink within node)
Pipeline Parallel: Use for >70B models
Context Parallel: Use for sequences >8K tokens
Data Parallel: Fill remaining GPUs

Example 405B on 128 H100s:

TP=8 (1 node)
PP=8 (across nodes)
CP=2 (long sequences)
DP=1
Total = 8 × 8 × 2 × 1 = 128 GPUs

When to use vs alternatives

Use Megatron-Core when:

Training models >10B parameters
Need maximum efficiency (target >40% MFU)
Using NVIDIA GPUs (A100, H100)
Production training at scale
Want fine-grained parallelism control

Use alternatives instead:

PyTorch FSDP: Models <70B, simpler API, PyTorch native
DeepSpeed: Easier setup, good for <100B models
HuggingFace Accelerate: Prototyping, simpler workflows
LitGPT: Educational, single-file implementations

Common issues

Issue: Low GPU utilization (<30% MFU)

Causes:

Micro-batch too small
Too much parallelism overhead
Not using Flash Attention

Fixes:

# Increase micro-batch
--micro-batch-size 4  # Was 1

# Enable optimizations
--use-flash-attn
--sequence-parallel

# Reduce TP if >8
--tensor-model-parallel-size 4  # Was 16

Issue: Out of memory

Reduce memory with:

--tensor-model-parallel-size 2  # Split model across GPUs
--recompute-granularity full  # Gradient checkpointing
--recompute-method block  # Checkpoint transformer blocks
--recompute-num-layers 1  # Checkpoint every layer

Or use CPU/NVMe offloading:

--cpu-optimizer  # Offload optimizer to CPU
--cpu-optimizer-type ADAM  # CPU Adam variant

Issue: Training slower than expected

Check:

Network bottleneck: Ensure InfiniBand/NVLink enabled
Pipeline bubbles: Use interleaved pipeline schedule
```
--num-layers-per-virtual-pipeline-stage 2
```
Data loading: Use fast data loader
```
--dataloader-type cyclic
```

Issue: Diverging loss

Stabilize training:

--lr-warmup-iters 2000  # Longer warmup
--clip-grad 1.0  # Gradient clipping
--init-method-std 0.006  # Smaller init
--attention-dropout 0.0  # No dropout in attention
--hidden-dropout 0.0  # No dropout in FFN

Advanced topics

Parallelism strategies: See references/parallelism-guide.md for detailed comparison of TP/PP/DP/CP/EP with performance analysis and when to use each.

Performance benchmarks: See references/benchmarks.md for MFU numbers across different model sizes and GPU configurations.

Production configurations: See references/production-examples.md for real-world setups from LLaMA 3 405B, Nemotron-4 340B, and DeepSeek-V3 671B.

Training recipes: See references/training-recipes.md for complete hyperparameter configurations for GPT/LLaMA/Mixtral architectures.

Hardware requirements

GPU: NVIDIA Ampere+ (A100, H100, B200)
- Turing works but slower
- FP8 requires Hopper/Ada/Blackwell
Network: InfiniBand or 400Gb+ Ethernet for multi-node
Memory per GPU:
- 7B model: 40GB+
- 70B model: 80GB (with TP=4)
- 405B model: 80GB (with TP=8, PP=8)
Storage: Fast NVMe for checkpoints (1TB+ for 70B+ models)

Resources

Docs: https://docs.nvidia.com/megatron-core/
GitHub: https://github.com/NVIDIA/Megatron-LM
Papers:
- "Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism" (2019)
- "Efficient Large-Scale Language Model Training on GPU Clusters Using Megatron-LM" (2021)
NeMo Framework: https://docs.nvidia.com/nemo-framework/ (built on Megatron-Core)

This skill ships inside the AI Research Skills plugin.

Run this outside your AI tool, using Node.js 22.20.0 or newer. Skills CLI 1.7.2 skips skill selection; prompts cancel installation without an interactive terminal. If choosing Project when asked about installation scope, open your project's folder before running this.

Install checklist

Check selected tools if a picker appears; Claude Code may be unchecked under Additional agents. ↑/↓ moves, Space toggles selections, Enter confirms.
Without a tool picker, check Installing to: for your intended tool. Stop if it names a different tool.
If asked, choose Symlink (Recommended) for shared files or Copy to all agents for separate copies; review the summary before confirming installation.
Install the find-skills skill? defaults to Yes, adding another skill for all projects. Choose No if you don't want that extra installation.

Skills CLI 1.7.2 sends install data to skills.sh by default, including the source, skill names, target tools and file details; GitHub sources are reported only when GitHub confirms the repository is public, while other remote sources can also be reported. Confirmed-public GitHub sources also trigger a security-audit request. Before installing, set the environment variable DO_NOT_TRACK=1 or DISABLE_TELEMETRY=1 to disable both.

npx skills add orchestra-research/ai-research-skills --skill "training-llms-megatron" --full-depth

After installing for Claude Code, open it in the installation project if you chose Project; Global makes the skill available across projects on this computer. Once Claude Code loads the skill, its description helps Claude decide when to use it.

If installation created a top-level skills folder that did not exist when your Claude Code session started, run /reload-skills in that session. Run it again after later changes in that folder.

What's inside

Files included with this skill.

View on GitHub

Megatron-Core - Large-Scale LLM Training

Quick start

Common workflows

Workflow 1: Train LLaMA-style model with 3D parallelism

Workflow 2: Configure Mixture of Experts (MoE) training

Workflow 3: Optimize for maximum throughput

When to use vs alternatives

Common issues

Advanced topics

Hardware requirements

Resources

Local tool required

What's inside

More skills

academic-plotting

ara-compiler

ara-research-manager

ara-rigor-reviewer

audiocraft-audio-generation

autogpt-agents