Megatron-LM 构建与依赖管理实战：基于 CI 容器的 uv 工作流、镜像变体与 uv.lock 维护

Megatron-LM 构建与依赖管理实战：基于 CI 容器的 uv 工作流、镜像变体与 uv.lock 维护

【免费下载链接】Megatron-LMOngoing research training transformer models at scale项目地址: https://gitcode.com/GitHub_Trending/me/Megatron-LM

本篇技术指南以 Megatron-LM 仓库中的构建与依赖管理规范（skills/mcore-build-and-dependency/SKILL.md）为核心，系统讲解「为何必须在容器内开发与构建」「dev 与 lts 两种镜像变体的差异与选型」「如何获取并启动 CI 容器」「如何用 uv 管理依赖组并更新 uv.lock」以及「常见构建陷阱排查」。读完本文，你将掌握一套可复制的容器化开发环境搭建流程，能独立完成新增依赖、解决 uv.lock 冲突、本地构建镜像并在 Slurm 集群上运行容器等实战操作。

核心原则：一切构建与依赖操作都在容器内进行

Megatron-LM 的核心构建原则可以概括为一句话：构建和开发都在容器内完成。CI 容器预先打包了正确的 CUDA 工具链、PyTorch 构建以及无法在裸机上复现的预编译原生扩展（如 TransformerEngine、DeepEP 等）。在宿主机上直接安装 CUDA、NCCL、带 GPU 支持的 PyTorch、TransformerEngine 以及 ModelOpt、DeepEP 等可选组件，既脆弱又难以复现；而仓库随附的 Dockerfile 对每一个依赖都做了精确固定。

容器化开发环境带来的三个关键保证：

• 所有开发者和 CI 使用完全一致的 CUDA / NCCL / cuDNN 版本；
• uv.lock在本地与 CI 中解析出相同的结果；
• 依赖 GPU 的操作（训练、测试）开箱即用。

在进入完整工作流之前，以下仓库特定事实可以直接作为快速回答的依据：

• 依赖相关工作必须在 Megatron-LM 的 CI 容器内执行，而不是宿主机；
• 容器内的虚拟环境位于/opt/venv，且已加入PATH；
• 默认的dev变体使用 docker/.ngc_version.dev（最新 NGC 发布）和devuv 组；lts变体使用 docker/.ngc_version.lts（较早的长期支持发布）和ltsuv 组。CI 中由 PR 标签container::lts选择 LTS 路径，否则一律使用dev；
• lts是仅当用户明确要求时才使用的选择，它是较旧的长期支持底座，不是常规第二通道——未经明确要求，不要主动附加container::lts标签、构建 LTS 镜像或运行ltsuv 组，即便是在做容器或依赖变更时也不例外；
• 容器内的安装命令：uv sync --locked --group dev --group test、uv sync --locked --only-group linting、uv sync --locked --group lts --group test；
• 依赖编辑使用uv add <package>加uv lock，两者都在容器内执行；
• docker/Dockerfile.ci.dev 包含main与jet两个 stage，jetstage 需要内部密钥，本地/公开构建应传入--target main。

dev 与 lts：两种镜像变体的差异与选型

仓库中存在两个镜像变体，各自拥有独立的 Dockerfile，由 PR 标签container::lts决定使用哪一个。两者的本质差异在于基础容器：dev跟随最新的 NGC PyTorch 发布，而lts（长期支持）固定在前一个仍受支持的 NGC PyTorch/CUDA 发布上。container::lts的存在意义是验证变更在较旧底座上依然可用——下表列出的依赖差异是由此派生而来的，并非目的本身。

变体	基础镜像固定	Dockerfile	依赖存放位置	使用场景
dev	docker/.ngc_version.dev（最新 NGC 发布）	docker/Dockerfile.ci.dev	pyproject.toml的devextra（由 uv 解析）	默认——CI、本地开发、绝大多数 PR
lts	docker/.ngc_version.lts（较旧的长期支持发布）	docker/Dockerfile.ci.lts	docker/lts/requirements.txt（固定版本，源自 main 分支的uv.lock）	向后兼容通道——验证变更在较旧 NGC 底座上仍可运行；未携带的 extras（ModelOpt、CUDA-13 的 TransformerEngine 构建）会被丢弃

历史背景：LTS 依赖原本位于pyproject.toml的[project.optional-dependencies].lts中，后被迁移到 docker/lts/requirements.txt，使pyproject.toml可以承载有实际意义的模块级 extras 而不与 LTS 固定版本集冲突。如需升级某个 LTS 依赖，直接编辑该 requirements.txt 中的版本并重新构建docker/Dockerfile.ci.lts。当前pyproject.toml中保留的lts = []只是一个空别名，用于避免pip install megatron-core[lts]报错。

结论：日常一律使用dev。CI 默认运行dev，这也是你主动触碰的唯一变体。把container::lts视为一道高门槛而非兜底方案：除非用户明确请求 LTS 验证，否则不要附加该标签、不要构建docker/Dockerfile.ci.lts、也不要运行ltsuv 组——即便是容器或依赖变更也不例外。当用户确实提出要求时，container::lts用于验证变更在 LTS 用户所运行的较旧长期支持 PyTorch/CUDA 底座上依然工作。

与镜像变体配套，测试标记也按环境区分：@pytest.mark.flaky_in_dev标记的测试在dev环境中被跳过，@pytest.mark.flaky标记的测试在lts中被跳过。两个标记在 pyproject.toml 的[tool.pytest.ini_options]中注册，并在 tests/unit_tests 的众多用例中实际使用，例如tests/unit_tests/dist_checkpointing/test_global_metadata_reuse.py中的@pytest.mark.flaky_in_dev # Issue #2856。

Step 1 — 获取镜像

获取 CI 镜像有两条路径：拉取 NVIDIA 内部 CI 构建好的镜像，或从零开始本地构建。

方案 A — NVIDIA 内部：拉取 CI 构建的镜像

⚠️ 需要访问内部 GitLab 实例，配置方法见 tools/trigger_internal_ci.md（添加 git remote、获取 token）。

内部 GitLab CI 会向容器仓库发布镜像。镜像仓库主机名可以从你配置的gitlabremote 推导——与trigger_internal_ci.py使用的同一个主机：

# 从 'gitlab' remote 推导主机名： GITLAB_HOST=$(git remote get-url gitlab | sed 's/.*@\(.*\):.*/\1/') docker pull ${GITLAB_HOST}/adlr/megatron-lm/mcore_ci_dev:main

方案 B — 从零构建（对所有人可用）

⚠️Dockerfile.ci.dev有两个 stage：main和jet。jetstage 需要内部构建密钥，没有它构建会失败。务必传入--target main停在公开 stage。

# dev 镜像（默认） docker build \ --target main \ --build-arg FROM_IMAGE_NAME=$(cat docker/.ngc_version.dev) \ --build-arg IMAGE_TYPE=dev \ -f docker/Dockerfile.ci.dev \ -t megatron-lm:local . # lts 镜像（使用专用 Dockerfile，无 IMAGE_TYPE 参数） docker build \ --target main \ --build-arg FROM_IMAGE_NAME=$(cat docker/.ngc_version.lts) \ -f docker/Dockerfile.ci.lts \ -t megatron-lm:local-lts .

从源码看，docker/Dockerfile.ci.dev 的mainstage 会完成以下关键步骤：安装yq与uv（默认UV_VERSION=0.7.2，并校验 SHA256）；设置UV_PROJECT_ENVIRONMENT=/opt/venv并将/opt/venv/bin加入PATH；通过uv venv ${UV_PROJECT_ENVIRONMENT} --system-site-packages在 NGC PyTorch 基础镜像之上叠加项目虚拟环境；随后uv sync安装build组以及dev、inference、mlm、ssm、te等 extras（dev 变体额外包含no_pypi_wheels组以从源码构建 flash-mla），同时用--no-install-package跳过 torch/torchvision/triton/transformer-engine 及所有nvidia-*-cu12包，直接复用基础镜像中已预装的 GPU 栈。构建过程中还把 TransformerEngine 的 NCCL EP JIT 头文件持久化到/opt/nccl-ep/include（经NCCL_EP_JIT_SOURCE_DIR等环境变量暴露），并基于 docker/patches/deepep.patch 与 docker/patches/deepep-nvshmem.patch 打补丁后从源码安装 DeepEP（lts变体跳过）。镜像默认以非 root 用户nemo-runtime（UID/GID 65532）运行，并内置了对 cache 目录可写性的非 root 契约验证。jetstage 则额外安装 jet-api 与 one-logger 等 NVIDIA 内部工具，需要JET_INDEX_URLS、LOGGER_INDEX_URL等 secret 挂载。

使用哪个镜像变体由 PR 标签container::lts控制；没有该标签时使用dev。

Step 2 — 启动容器

方案 A — 本地 Docker 运行时

docker run --rm --gpus all \ -v $(pwd):/workspace \ -w /workspace \ megatron-lm:local \ bash -c "<your command>"

方案 B — Slurm 集群（适用于没有本地 Docker 运行时的场景）

NVIDIA 集群通常使用 Pyxis + enroot 组合。申请一个交互式会话：

srun \ --nodes=1 --gpus-per-node=8 \ --container-image megatron-lm:local \ --container-mounts $(pwd):/workspace \ --container-workdir /workspace \ --pty bash

对于要求先提供.sqsh归档的集群：

enroot import -o megatron-lm.sqsh dockerd://megatron-lm:local srun \ --nodes=1 --gpus-per-node=8 \ --container-image $(pwd)/megatron-lm.sqsh \ --container-mounts $(pwd):/workspace \ --container-workdir /workspace \ --pty bash

启动后，仓库源码被挂载到容器的/workspace，/opt/venv中的 Python 环境已就绪，可以直接执行训练、测试或依赖管理命令。

依赖管理：pyproject.toml + uv 锁文件

依赖在 pyproject.toml 中声明。容器内的虚拟环境位于/opt/venv（已加入PATH）。

所有 uv 操作都必须在容器内执行。永远不要在宿主机上运行uv sync/uv pip install。

uv 依赖组

组	用途
training	运行时训练 extras
dev	完整开发环境（TransformerEngine、ModelOpt 等）
test	pytest、coverage、nemo-run
linting	ruff、black、isort、pylint
build	Cython、pybind11、nvidia-mathdx

原有的ltsextra 已被清空。LTS 依赖固定于 docker/lts/requirements.txt 而非pyproject.toml。不要在[project.optional-dependencies].lts下新增包。

安装命令（在容器内执行）：

# 完整 dev + test 环境 uv sync --locked --group dev --group test # 仅 linting uv sync --locked --only-group linting

LTS 环境的复现方式是端到端构建docker/Dockerfile.ci.lts，不存在等价的uv sync命令，因为 LTS 依赖已不在pyproject.toml中。LTS 的顶层固定版本集位于docker/lts/requirements.txt（例如tqdm==4.67.3、megatron-energon[av_decode]==7.3.2、flashinfer-python==0.6.11.post3等）；升级版本要在该文件中修改并重新构建镜像。

从 pyproject.toml 的源码细节看，除[dependency-groups]外，uv 相关的核心配置包括：[tool.uv]中managed = true、默认组default-groups = ["linting", "build", "test"]、link-mode = "copy"，以及no-build-isolation-package列表（causal-conv1d、flash_mla、mamba-ssm、transformer-engine、deep_gemm、fast-hadamard-transform等需无隔离构建的包）；override-dependencies用sys_platform == 'never'技巧禁止在基础镜像中重复安装 torch/torchvision/triton。多份依赖直接以 git 源引入（[tool.uv.sources]）：TransformerEngine、nemo-run、FlashMLA、DeepGEMM、Emerging-Optimizers、mamba-ssm、nemo-lens、torch-memory-saver 等均固定到精确的 commit/rev。锁定的 uv.lock 文件固定了这些 git 依赖的确切修订版本，修改pyproject.toml后必须用uv lock更新它。

新增依赖

遵循三步工作流：

1. 获取容器镜像—— 见上文 Step 1；
2. 交互式启动容器—— 见上文 Step 2；
3. 在容器内更新锁文件，然后提交：

# 容器内： uv add <package> # 写入 pyproject.toml 并解析 uv lock # 重新生成 uv.lock # 退出容器后，在宿主机： git add pyproject.toml uv.lock git commit -S -s -m "build: add <package> dependency"

解决 uv.lock 的合并冲突

uv.lock是机器生成的文件，绝不要手动解决冲突。正确做法是：

git checkout origin/main -- uv.lock # 以 main 的版本为基准 # 然后在容器内： uv lock # 基于你的 pyproject.toml 变更重新解析

常见陷阱速查

问题	原因	解决办法
uv sync --locked失败	依赖冲突或uv.lock过期	在容器内重新运行uv lock并提交更新后的锁文件
pip install 后出现ModuleNotFoundError	pip 安装到了 uv 管理的 venv 之外	使用uv add和uv sync，绝不使用裸pip install
容器内uv: command not found	用错了容器镜像	使用由Dockerfile.ci.dev构建的megatron-lm镜像
uv 操作时No space left on device	缓存填满容器的/root/.cache/	通过-v $HOME/.cache/uv:/root/.cache/uv挂载宿主机缓存目录
docker build报 secret 相关错误	Dockerfile.ci.dev的jetstage 需要内部密钥	加--target main，在jetstage 之前停止
pull 时报access forbidden	镜像仓库 URL 带显式端口（如:5005）	使用不含端口的${GITLAB_HOST}/adlr/...—— sed 只提取主机名

小结：一套完整的容器化依赖工作流

将本文要点串起来，Megatron-LM 的标准开发闭环是：在dev变体的 CI 容器（由 docker/Dockerfile.ci.dev 构建，基础镜像取自 docker/.ngc_version.dev）中工作，依赖声明于 pyproject.toml，由uv与 uv.lock 精确锁定，任何依赖变更都以「uv add+uv lock+ 提交锁文件」的容器内流程完成，而 LTS 验证则严格遵循用户明确请求的前提，通过构建 docker/Dockerfile.ci.lts 并在container::lts标签下运行 CI 完成。这套流程同时保证了本地与 CI 环境的可复现性，以及 GPU 相关操作的开箱即用。

【免费下载链接】Megatron-LMOngoing research training transformer models at scale项目地址: https://gitcode.com/GitHub_Trending/me/Megatron-LM