总目录 › Megatron-LM › 第 2 章
Chapter 02

环境准备与安装

📌 本章命令均在 NVIDIA/Megatron-LM · commit 432d76b2b8bf 上验证（仓库自报版本号 0.15.0 ，见 README 顶部 badge）。Megatron 迭代很快，HEAD 与本文不一致时请先 cd repos/megatron-lm && git checkout 432d76b2b8bf 切回基线。

2.1 装之前要先承认的事

Megatron-LM 不是"装好就能跑"的工具，它对环境有非常具体的要求：

2.2 三条安装路线

仓库 README 给出的安装方式有三条，按"省心程度"由低到高：

路线 A：NGC 容器（NVIDIA 官方推荐，最省心）

NVIDIA 自己发的 PyTorch 容器里 CUDA、cuDNN、NCCL、apex、TransformerEngine 都已经装好了，进去 pip install megatron-core 就直接能用：

docker run --gpus all -it --rm \
  -v $(pwd):/workspace \
  nvcr.io/nvidia/pytorch:24.10-py3 bash

# 进容器后
pip install megatron-core

NGC 容器版本号要和你卡片的 CUDA 兼容（H100/A100 用 24.xx 系列没问题，老的 V100 可能要选更老 tag）。这是生产部署的标准做法。

路线 B：纯 pip / uv 安装（开发用）

仓库 README 推荐用 uv（比 pip 快很多）：

# 装 uv（如果还没装）
curl -LsSf https://astral.sh/uv/install.sh | sh

# 直接装 release 版
uv pip install megatron-core

# 或克隆 + editable 安装（要改 megatron 内部代码就走这条）
git clone https://github.com/NVIDIA/Megatron-LM.git
cd Megatron-LM
uv pip install -e .
注意 README 原话："Building from source can use a lot of memory. If the build runs out of memory, limit parallel compilation jobs by setting MAX_JOBS (e.g. MAX_JOBS=4 uv pip install -e .)."
编译时容易内存炸，机器内存 < 32 GB 一定要带 MAX_JOBS=4，否则会被 OOM kill。

路线 C：从你已克隆的 repos/megatron-lm 装

本仓库已经在 /Users/zjw/Documents/LLM/train/repos/megatron-lm/ 准备好了源码（commit 432d76b2b8bf）。如果你只是想读源码、不打算训东西，不需要装任何东西——文件已经在那儿，可以直接 grep / 读。

如果要真训练，路线选 A 或 B。

2.3 验证安装

不管走哪条路线，进 Python 跑这三行确认装好：

import megatron.core
print(megatron.core.__version__)            # 应输出版本号
from megatron.core import parallel_state    # 不能抛 ImportError
from megatron.core.transformer.transformer_config import TransformerConfig
print(TransformerConfig)                    # 不能抛 ImportError

任意一行报错都说明环境没准备好，重点查：

  1. PyTorch 版本是否 ≥ 2.4：python -c "import torch; print(torch.__version__)"
  2. CUDA 是否可用：python -c "import torch; print(torch.cuda.is_available(), torch.cuda.device_count())"
  3. NCCL 是否在：python -c "import torch; print(torch.cuda.nccl.version())"

2.4 可选但强烈建议的额外组件

这些不装能跑，装上能显著加速：

组件作用安装
TransformerEngineFP8 训练、fused attention、fused LayerNormNGC 容器自带；外部 pip install transformer-engine[pytorch]
Apex NVIDIA 的 fused optimizer / fused softmax NGC 容器自带；外部要从源码装
FlashAttention 标准 flash-attn 替换 PyTorch SDPA pip install flash-attn --no-build-isolation
fused kernel 仓库自带的 C++ kernel，第一次跑会自动 JIT 编译不用装，但要确保 nvcc 可用
提示 用 NGC 容器 = 上面这些全都不用管。

2.5 第一次跑出错怎么办

最常见的几种失败：

#报错 / 现象解决思路
1ImportError: cannot import name 'xxx' from 'megatron.core'八成是 megatron-core 版本和教程/代码不一致。pip show megatron-core 看版本，对照 README 的 CHANGELOG.md。
2NCCL hang / 卡在 init_process_group检查 MASTER_ADDR/MASTER_PORT/WORLD_SIZE/RANK 四个环境变量；防火墙 / 容器网络通常是真凶。
3CUDA out of memory 一启动就炸不是你模型大，是 fused kernel JIT 编译占了显存。先 export TORCH_CUDA_ARCH_LIST=8.0+PTX（或你卡的 arch）限制编译范围。
4MAX_JOBS 没设、机器卡死看 2.2 节 的警告。

完整调试手册留到第 11 章。