Megatron-LM 与 Megatron Core 安装全指南：PyPI、源码构建与 NGC 容器三种方式详解

Megatron-LM 与 Megatron Core 安装全指南PyPI、源码构建与 NGC 容器三种方式详解【免费下载链接】Megatron-LMOngoing research training transformer models at scale项目地址: https://gitcode.com/GitHub_Trending/me/Megatron-LM本指南以 docs/get-started/install.md 为骨架系统讲解 Megatron-LM / Megatron Core 的三种安装路径——PyPI 安装、源码构建、NGC 容器——并深入到当前仓库的 pyproject.toml、setup.py、Dockerfile 等源码级细节帮助你在自己的 GPU 集群上一次性搭建出可运行大模型训练的完整环境。读完本文你将掌握版本约束核对、uv使用、extras 依赖选择、MAX_JOBS编译调优、NGC 容器PIP_CONSTRAINT坑点规避等全套实战技能并知道如何用最小示例验证安装结果。安装前须知Megatron-LM 与 Megatron Core 的关系在动手安装前先明确安装对象。根据 docs/get-started/overview.mdMegatron-LM 是一个开源的参考实现内置了轻量级的大规模 LLM 训练框架与可定制的原生 PyTorch 训练循环而 Megatron Core 则是其底层可组合库提供了张量并行、流水线并行、数据并行、专家并行、上下文并行等高级并行策略以及 FP16/BF16/FP8 混合精度、GPU 优化 kernel、高性能数据加载器等模块化构件。本文安装的核心 Python 包是megatron-corePyPI 包名它对应仓库根目录下的 pyproject.tomlname megatron-core而整个 Megatron-LM 仓库则同时包含 Megatron Core 源码、pretrain_gpt.py 等训练入口和 examples 下的完整示例脚本。安装 Megatron Core 之后即可配合这些入口文件与示例启动真实训练。系统要求硬件要求推荐 GPUNVIDIA Turing 架构或更新Turing、Ampere、Hopper、Ada、Blackwell 等支持所有 NVIDIA Tensor Core GPU。FP8 支持FP8 混合精度训练要求 NVIDIA Hopper、Ada 或 Blackwell 架构 GPU。仓库内 examples/llama/train_llama3_8b_h100_fp8.sh 即针对 H100 设计通过--fp8-format hybrid等参数启用 FP8可见 FP8 能力与硬件强相关。软件要求Python官方安装文档声明 3.10推荐 3.12。需要说明的是当前仓库 pyproject.toml 中已把requires-python收紧为3.12且 isort 配置py_version 312同样以 Python 3.12 为对齐目标因此实际操作中建议直接使用 Python 3.12兼容性最稳。PyTorch 2.6.0这也是 pyproject.toml 中dependencies [torch2.6.0, ...]的强制声明。CUDA Toolkit建议使用最新稳定版。仓库构建流程docker/common/install.sh在 Ubuntu 基础镜像上安装的是 CUDA 12.8 工具链cuda-toolkit-12-8、cudnn-cuda-12可作为参考基线。前置条件安装 uvuv是一个极快的 Python 包安装器/依赖管理器Megatron Core 官方安装流程以它为默认工具。其核心优势在于解析依赖快、缓存命中率高并且天然支持--no-build-isolation、--group等与 Megatron Core 构建流程配合的选项。curl -LsSf https://astral.sh/uv/install.sh | sh安装完成后uv可执行文件会被放到 shell 的 PATH 中脚本默认写入~/.local/bin必要时手动加入 PATH。仓库自身也使用 uv 托管依赖根目录存在 uv.lock 锁文件pyproject.toml 中[tool.uv] managed true并在no-build-isolation-package中列出了mamba-ssm、transformer-engine、deep_gemm、flash_mla等需要免隔离编译的 CUDA 扩展包。此外 docker/Dockerfile.ci.dev 将 uv 固定到0.7.2版本ARG UV_VERSION0.7.2日常使用uv最新稳定版即可。提示如果没有安装 uv也可以退化为pip install megatron-core但--no-build-isolation、--group等参数在 uv 下行为最一致官方文档与仓库 CI见 docker/Dockerfile.ci.dev 中的uv sync --only-group build均以 uv 为准。Option APyPI 安装推荐PyPI 安装是最快的起步路径直接拉取最新稳定版 wheel无需本地编译适合大多数训练用户。基础安装uv pip install megatron-core该命令只安装核心包与硬性依赖torch、numpy、packaging不含任何可选组件适合仅需最小运行时验证的场景。带训练依赖安装要使用训练相关的常用组件Weights Biases 实验追踪、SentencePiece 分词、Hugging Face Transformers 等使用trainingextrauv pip install megatron-core[training]根据 pyproject.toml 的[project.optional-dependencies]trainingextra 实际包含flask-restful、sentencepiece、tiktoken、wandb、transformers、accelerate、omegaconf、gigatoken等包。注意mlmextra 是已废弃的旧别名内容与training完全一致新项目请统一使用training。含全部 extras 安装含 Transformer Engine需要完整开发/训练能力尤其是 Transformer Engine 提供的优化 kernel 与 FP8 支持时uv pip install --group build uv pip install --no-build-isolation megatron-core[training,dev]这里有两步关键操作务必理解其含义uv pip install --group build先安装build依赖组。从 pyproject.toml 的[dependency-groups]看build组包含setuptools80、packaging、hatchling、pybind11、Cython3.0.0、torch、nvidia-cudnn-frontend供 TE 使用、nvidia-mathdx供 TE 使用等它们是后续免隔离编译所必需的基础工具链。--no-build-isolation表示在编译时直接复用当前环境中的构建依赖而不是临时创建隔离的构建环境。官方文档明确指出torch必须预先安装因为dev组中的多个包mamba-ssm、nv-grouped-gemm、transformer-engine在构建期会 import torch 来编译 CUDA kernel。这一步骤取决于硬件性能预计耗时 20 分钟以上。如果你更倾向于预编译二进制直接选用下文 Option C 的 NGC 容器即可容器内已预编译好这些组件。devextra 的内容非常丰富从 pyproject.toml 可以完整看到包括nvidia-modelopt[torch]模型优化、nvidia-resiliency-ext容错训练、megatron-energon多模态数据加载、tensorstore分布式 checkpoint 存储、flashinfer-python注意力 kernel、nvidia-cudnn-frontend、fastapi、datasets、emerging_optimizers、openai、orjson、zstandard等覆盖推理服务、数据管线、可观测性与实验管理。编译内存警示从源码编译会占用大量内存。默认情况下每个 CPU 核会启动一个编译任务在多核机器上极易触发 OOM。请在安装前通过环境变量限制并行编译任务数例如 export MAX_JOBS4MAX_JOBS会传递给底层 ninja/setuptools 构建系统控制同时进行的编译作业数量。该参数在仓库的 docker/Dockerfile.ci.dev 中也有体现——它在 aarch64 平台上显式设置MAX_JOBS8以限制每包编译并发印证了这一点。更轻量的开发依赖组合如果不想引入 Transformer Engine 与 ModelOpt可以改用 lts extra 获得更精简的开发依赖集 uv pip install --no-build-isolation megatron-core[training,lts]需要特别留意pyproject.toml 中lts []目前是空别名——LTS 的 Python 依赖已经整体迁移到 docker/Dockerfile.ci.lts 与 docker/lts/requirements.txt 中如einops0.8.2、tensorstore0.1.84、megatron-energon[av_decode]7.3.2等固定版本。也就是说[lts]与[dev]是互斥的依赖集合选择lts时这些包要么由你的环境自行提供、要么参照 docker/lts/requirements.txt 的锁定版本单独安装。克隆仓库获取示例安装完 Python 包后训练还需要仓库中的入口脚本与示例克隆仓库git clone https://gitcode.com/GitHub_Trending/me/Megatron-LM.gitOption B从源码安装源码安装适用于三种场景向项目贡献代码、抢先使用未发布的新特性、以及调试时单步跟踪包内部实现。它把仓库克隆到本地并以可编辑editable模式安装包代码改动即时生效无需重装。git clone https://gitcode.com/GitHub_Trending/me/Megatron-LM.git cd Megatron-LM uv pip install -e .可编辑安装会触发一次本地的 C 扩展编译。从 setup.py 可以看到包通过pybind11.setup_helpers.Pybind11Extension编译megatron.core.datasets.helpers_cpp源码为 megatron/core/datasets/helpers.cpp编译参数-O3 -Wall -stdc17该扩展是数据集采样等高性能路径的一部分。因此源码安装环境需要具备可用的 C 编译器g/clang与 pybind11。如需连同全部开发依赖一起安装包含 Transformer Engine需要预先装好构建依赖uv pip install --group build uv pip install --no-build-isolation -e .[training,dev]同样的内存控制技巧依然适用如果构建过程中内存不足用 MAX_JOBS 限制并行编译 MAX_JOBS4 uv pip install --no-build-isolation -e .[training,dev]Option CNVIDIA GPU CloudNGC容器NGC 容器是一条零手工依赖管理的路径PyTorch、CUDA、cuDNN、NCCL、Transformer Engine 等全部预装且版本经过 NVIDIA 测试矩阵校验适合追求开箱即用、不愿自己编译 CUDA 扩展的用户。版本选择建议优先使用上一个月份的 NGC PyTorch 容器而不是最新版以确保与当前 Megatron Core 发布版本及测试矩阵兼容。启动命令docker run --gpus all -it --rm \ -v /path/to/dataset:/workspace/dataset \ -v /path/to/checkpoints:/workspace/checkpoints \ -e PIP_CONSTRAINT \ nvcr.io/nvidia/pytorch:26.01-py3命令要点说明-v将宿主机数据集与 checkpoint 目录挂载进容器实现数据与权重的持久化-e PIP_CONSTRAINT必须保留。NGC PyTorch 容器会通过PIP_CONSTRAINT对全局 Python 环境施加约束导致 Megatron Core 及其依赖被强制降级或拒绝安装该参数显式清空此约束是容器内安装成功的关键前提。仓库的 docker/Dockerfile.ci.dev 同样在镜像层设置ENV PIP_CONSTRAINT来解除这一全局限制。进入容器后安装 Megatron CoreNGC 镜像已内置 torch因此可直接使用--no-build-isolationpip install uv uv pip install --no-build-isolation megatron-core[training,dev]容器内torch已可用--no-build-isolation的构建前提天然满足若仍遇到编译内存问题同样可用MAX_JOBS环境变量控制。至于 NGC 容器标签的具体月份号以你安装时的官方发布为准仓库 docker/Dockerfile.ci.dev 第 4 行即以nvcr.io/nvidia/pytorch:26.06-py3作为基础镜像的示例可作为命名规范的参考。依赖体系源码解读安装到底装了什么在三种安装方式之外理解megatron-core的依赖设计能帮你精准裁剪环境。以下是当前仓库 pyproject.toml 中可验证的 extras 全景Extra / 组主要内容适用场景trainingwandb、sentencepiece、tiktoken、transformers、accelerate、omegaconf、flask-restful、gigatoken常规训练推荐devnvidia-modelopt、nvidia-resiliency-ext、megatron-energon、tensorstore、flashinfer-python、fastapi、datasets、openai、orjson、zstandard 等完整开发/实验环境lts空别名真实锁定版本见 docker/lts/requirements.txt轻量、避开 TE/ModelOpttetransformer-engine[pytorch,core_cu13]单独启用 Transformer Enginessmmamba-ssm、causal-conv1d、flash-linear-attentionMamba 等状态空间模型inferencetorch-memory-saver推理内存优化otelnemo-lens[sdk]OpenTelemetry 可观测性build依赖组setuptools、pybind11、Cython、torch、nvidia-cudnn-frontend、nvidia-mathdx免隔离编译工具链几个值得注意的实现细节构建系统pyproject.toml 的[build-system]要求setuptools80、pybind11、packaging24.2配合 setup.py 的 Pybind11 扩展定义决定了源码安装一定会编译helpers_cpp模块。版本号机制版本号来自 megatron/core/package_info.py当前开发版本为0.20.0且当仓库是 git 检出时会自动把 commit 短哈希追加到版本串如0.20.0abc1234可通过环境变量NO_VCS_VERSION1关闭该行为。uv 锁文件uv.lock 与[tool.uv]配置把transformer-engine、mamba-ssm、flash_mla、deep_gemm、fast-hadamard-transform等固定到精确的 git 提交保证可复现构建——这也是官方 CI 镜像使用uv sync --locked的原因。容器内安装的权威参考docker/Dockerfile.ci.dev 展示了 NVIDIA 官方在 NGC 容器内构建 dev 镜像的完整流程创建带--system-site-packages的 venv、uv sync --only-group build预装构建依赖、再以uv sync --extra dev --extra inference --extra mlm --extra ssm --extra te安装各类 extras并通过--no-install-package torch等参数跳过 torch 与全部nvidia-*-cu12运行时包复用 NGC 镜像自带组件。如果你要在自己的容器里复刻同样环境这份 Dockerfile 是最完整的模板docker/common/install.sh 则提供了在裸机Ubuntu 或已装 PyTorch 的环境上构建 dev/lts 环境的脚本化路径。安装验证与下一步跑起第一个训练安装完成后建议先跑最小训练验证环境再进入真实模型。详见 docs/get-started/quickstart.md最小分布式训练2 张 GPUmock 数据torchrun --nproc_per_node2 examples/run_simple_mcore_train_loop.py该示例不依赖真实数据集用于确认并行通信、模型构建与训练循环在本机配置下全部正常。生产级示例LLaMA-3 8B FP88 张 GPU./examples/llama/train_llama3_8b_h100_fp8.sh该脚本在 8 张 H100 上用 mock 数据训练 LLaMA-3 8B--fp8-format hybrid、--bf16、GQA、RoPE、RMSNorm、sequence parallel、分布式优化器同时演示了 TP/CP/PP 参数、数据缓存路径benchmark_cache_llama3_8b_fp8与 checkpoint/tensorboard 配置是安装后检验完整能力的理想基线也可作为编写自有训练脚本的蓝本。准备真实数据Megatron 训练需要预处理后的二进制数据文件.bin与.idx先准备每行包含text字段的 JSONL 文件{text: Your training text here...} {text: Another training sample...}运行 tools/preprocess_data.py 完成分词与二进制转换python tools/preprocess_data.py \ --input data.jsonl \ --output-prefix processed_data \ --tokenizer-type HuggingFaceTokenizer \ --tokenizer-model /path/to/tokenizer.model \ --workers 8 \ --append-eod关键参数速查--input输入 JSON/JSONL 路径--output-prefix输出.bin/.idx文件前缀--tokenizer-type分词器类型HuggingFaceTokenizer、GPT2BPETokenizer等--tokenizer-model分词器模型文件--workers并行处理进程数--append-eod在样本末尾追加 end-of-document 标记。常见问题排查现象原因解决办法编译时 OOM / 机器卡死默认每 CPU 核一个编译任务内存被打满设置MAX_JOBS4或更小后重装--no-build-isolation安装报错环境中缺少 torch 等构建期依赖先执行uv pip install --group build并确认 torch 已预装容器内安装的依赖被降级或异常NGC 容器的全局PIP_CONSTRAINT生效启动时加-e PIP_CONSTRAINT或在 Dockerfile 中ENV PIP_CONSTRAINTpip install megatron-core[lts]后缺少包ltsextra 当前为空别名按 docker/lts/requirements.txt 的锁定版本手动补齐依赖装完包但跑示例找不到脚本只装了 PyPI 包未克隆仓库git clone当前仓库并按 docs/get-started/quickstart.md 执行Python 版本过低无法安装当前仓库requires-python 3.12使用 Python 3.12与 pyproject.toml 声明一致结语三条安装路径对应三种工作流PyPI 安装适合快速上手与常规训练源码安装适合二次开发与调试NGC 容器适合追求零依赖管理、依赖官方测试矩阵的环境。无论选择哪条路径本文提到的MAX_JOBS、--no-build-isolation、PIP_CONSTRAINT三个关键点都是绕不开的实战细节。安装完成后即可进入 docs/get-started/quickstart.md 开始第一次训练并在需要扩展规模时参考 并行策略指南、数据准备最佳实践 与 高级特性索引FP8 训练、上下文并行等继续深入。【免费下载链接】Megatron-LMOngoing research training transformer models at scale项目地址: https://gitcode.com/GitHub_Trending/me/Megatron-LM创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考