Megatron-LM Install Guide: Deploy v0.17 for LLM Training

A common question in ML Discords right now: “I cloned Megatron-LM, ran pip install, and nothing works – what did I miss?” Almost always, the answer is that Megatron isn’t really one package. It’s a training reference plus a fast-moving core library plus an APEX build that’s picky about your CUDA version. Getting large model training running takes about 20 minutes if you pick the right install path, or two frustrated evenings if you don’t.

This guide covers deploying Megatron-LM the way it’s actually shipped in 2026 – with the caveats the README skips over.

Two packages, one repo

Two things live inside the NVIDIA/Megatron-LM repo, and most tutorials conflate them. Megatron-LM: reference training scripts for research teams, quick experimentation, and copy-paste starting points. Megatron Core: a composable library with GPU-optimized building blocks – tensor, pipeline, data, expert, and context parallelism, plus mixed-precision support for FP16, BF16, FP8, and FP4 (per the official README). You install both, but in production you’ll be importing from megatron.core, not the training scripts.

The moving target as of mid-2026 is Megatron Core 0.17. That release drops Python 3.10 support – downstream applications must raise their minimum to Python 3.12. If your cluster CI still pins 3.10, that’s your first blocker, and it will fail silently during dependency resolution rather than throwing an obvious error.

System requirements (honest version)

The docs don’t publish a strict minimum, so here’s what actually works:

The RAM figure matters more than most tutorials admit. When building from source, ninja spawns parallel nvcc compilation jobs with no cap by default. On a 16 GB machine it’ll happily launch a dozen at once and OOM your kernel – no error message, just a killed process. The fix is MAX_JOBS=4, documented in a brief note in the README that’s easy to miss.

Pick your install path

Three options. They’re not equal.

1. NGC PyTorch container – recommended for anyone who doesn’t enjoy debugging APEX at midnight.
2. uv pip install from source – works if you already have a matching CUDA + PyTorch stack.
3. Bare pip on a fresh machine – technically possible, usually painful.

The catch: NVIDIA doesn’t tell you to grab the newest container. Releases are always built against the previous month’s NGC container for stability – grabbing :latest routinely breaks builds against nightly PyTorch changes. The recommended tag as of this writing is nvcr.io/nvidia/pytorch:25.04-py3.

Install via NGC container (the sane path)

Production teams do this. The container ships with PyTorch, CUDA, cuDNN, NCCL, and APEX pre-compiled – no wheel roulette.

# Pull the container that Megatron was built against
docker pull nvcr.io/nvidia/pytorch:25.04-py3

# Run with GPU passthrough and your workspace mounted
docker run --gpus all -it --rm 
 --ipc=host --ulimit memlock=-1 --ulimit stack=67108864 
 -v $(pwd):/workspace/megatron 
 -v /path/to/dataset:/workspace/dataset 
 -v /path/to/checkpoints:/workspace/checkpoints 
 -e PIP_CONSTRAINT= 
 nvcr.io/nvidia/pytorch:25.04-py3

# Inside the container
cd /workspace/megatron
git clone https://github.com/NVIDIA/Megatron-LM.git
cd Megatron-LM
pip install --no-build-isolation .[mlm,dev]

The -e PIP_CONSTRAINT= flag is easy to miss. The NGC PyTorch container pins global package versions through a PIP_CONSTRAINT environment variable, which will fight your Megatron install if left set. Passing it empty in the docker run command unsets it for that session.

Install via pip (source, without container)

Only take this route with a matched CUDA/PyTorch environment already in place. From the official quickstart:

# Fresh Python 3.12 venv
python3.12 -m venv .venv && source .venv/bin/activate

# Install PyTorch matching your CUDA (example: CUDA 12.4)
pip install torch --index-url https://download.pytorch.org/whl/cu124

# Install Megatron Core with training extras
pip install --no-build-isolation megatron-core[mlm,dev]

# Clone repo for examples and scripts
git clone https://github.com/NVIDIA/Megatron-LM.git
cd Megatron-LM
MAX_JOBS=4 uv pip install -e .[mlm,dev]

MAX_JOBS=4 is the difference between a successful build and a killed process on any machine under 64 GB.

On APEX: Skip installing NVIDIA APEX unless a training script explicitly imports it. Modern Megatron Core uses Transformer Engine kernels for fused ops, and APEX is now optional for most recipes. Adding it “just in case” is how you inherit CUDA-version headaches.

First-time verification

# 1. GPU visible?
python -c "import torch; print(torch.cuda.device_count(), torch.cuda.get_device_name(0))"

# 2. Megatron Core imports?
python -c "from megatron.core import parallel_state; print('ok')"

# 3. Smoke test with mock data on 2 GPUs
torchrun --nproc_per_node=2 examples/run_simple_mcore_train_loop.py

That third command is straight from the official quickstart’s distributed training example. If it prints loss values for a few steps and exits cleanly, your install is real. If it hangs silently, check NCCL_DEBUG=INFO output – usually a networking or IPC flag issue.

Common errors, real fixes

1. No module named 'fused_layer_norm_cuda' or No module named 'amp_C'

You installed the wrong APEX. There’s a random PyPI package called apex that isn’t NVIDIA’s – a known confusion point documented in multiple GitHub issues. Uninstall it and build from source:

pip uninstall -y apex
git clone https://github.com/NVIDIA/apex
cd apex
pip install -v --disable-pip-version-check --no-cache-dir --no-build-isolation 
 --config-settings "--build-option=--cpp_ext" 
 --config-settings "--build-option=--cuda_ext" ./

2. APEX RuntimeError: CUDA version mismatch

This one eats afternoons. APEX setup.py runs a strict check between the CUDA version PyTorch was compiled with and the bare-metal CUDA version found via nvcc – any mismatch raises a RuntimeError immediately during installation, before a single training token is processed. Run nvcc --version and python -c "import torch; print(torch.version.cuda)". If they differ, either reinstall PyTorch with matching CUDA wheels, or go back to the NGC container where everything is pre-matched.

3. ImportError: No module named 'transformer_engine'

Transformer Engine requires CUDA 12.x. Users on CUDA 11.x hit this after installing Megatron-LM because TE isn’t included in older CUDA environments. There’s no clean workaround on old CUDA – upgrade CUDA or use the pinned NGC container where TE is already bundled.

Upgrading and cleanup

Upgrades are refreshingly boring if you installed via pip install -e .:

cd Megatron-LM
git fetch --tags
git checkout core_r0.17.0 # or whichever tag you want
MAX_JOBS=4 pip install --no-build-isolation -e .[mlm,dev] --force-reinstall

Coming from a pre-0.17 install on Python 3.10? Upgrade Python first – otherwise the resolver silently pins you to an old wheel with no warning.

To uninstall completely: pip uninstall megatron-core apex transformer-engine, then delete the cloned repo and mounted checkpoint directories. If you used Docker, docker rmi nvcr.io/nvidia/pytorch:25.04-py3 reclaims significant disk space.

Is any of this worth it in 2026?

Fair question. As of 2026, FSDP2 in vanilla PyTorch is meaningfully closer to Megatron’s performance than it was a couple of years ago. For a 7B model on 8 GPUs, plain PyTorch gets you moving faster with far less setup pain. Megatron earns its keep at scale: 47% Model FLOP Utilization on H100 clusters, training models from 2B to 462B parameters across thousands of GPUs (per the README performance section). Under 10B params on a single node? Honestly, the install complexity isn’t repaid.

Your next action: spin up the NGC container above, run the 2-GPU mock-data smoke test, and only then think about which model recipe to try. The original 2019 Megatron-LM paper (Shoeybi et al., arXiv:1909.08053) is still the best background reading while your container downloads.

FAQ

Do I need NVIDIA APEX to run Megatron-LM in 2026?

No. Transformer Engine now covers the fused kernels APEX used to provide. Only add APEX if a specific script fails on an apex import – otherwise you’re adding a fragile build step for no gain.

Can I run Megatron-LM on AMD GPUs?

AMD maintains a ROCm-compatible Megatron-LM fork with Docker images for Instinct GPUs (MI300X and similar). As of early 2026, feature parity lags the NVIDIA branch by a release or two, and FP8 support is NVIDIA-only. If you’re evaluating AMD for a new cluster build, check AMD’s current developer hub for the latest image tags – they update frequently enough that any specific tag listed here could already be stale.

Why does the quickstart use uv instead of regular pip?

Speed matters when your dependency tree includes torch, transformer-engine, flash-attn, and a handful of CUDA extension packages. uv resolves and installs that tree noticeably faster than pip, and it handles --no-build-isolation more gracefully in practice. That said, regular pip works fine – expect a longer wait, and make sure you’re not accidentally running inside a virtualenv that has PIP_CONSTRAINT set from a prior session. That constraint file will override your install flags silently, and debugging it is exactly as annoying as it sounds.