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 路径否则一律使用devlts是仅当用户明确要求时才使用的选择它是较旧的长期支持底座不是常规第二通道——未经明确要求不要主动附加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两个 stagejetstage 需要内部密钥本地/公开构建应传入--target main。dev 与 lts两种镜像变体的差异与选型仓库中存在两个镜像变体各自拥有独立的 Dockerfile由 PR 标签container::lts决定使用哪一个。两者的本质差异在于基础容器dev跟随最新的 NGC PyTorch 发布而lts长期支持固定在前一个仍受支持的 NGC PyTorch/CUDA 发布上。container::lts的存在意义是验证变更在较旧底座上依然可用——下表列出的依赖差异是由此派生而来的并非目的本身。变体基础镜像固定Dockerfile依赖存放位置使用场景devdocker/.ngc_version.dev最新 NGC 发布docker/Dockerfile.ci.devpyproject.toml的devextra由 uv 解析默认——CI、本地开发、绝大多数 PRltsdocker/.ngc_version.lts较旧的长期支持发布docker/Dockerfile.ci.ltsdocker/lts/requirements.txt固定版本源自 main 分支的uv.lock向后兼容通道——验证变更在较旧 NGC 底座上仍可运行未携带的 extrasModelOpt、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有两个 stagemain和jet。jetstage 需要内部构建密钥没有它构建会失败。务必传入--target main停在公开 stage。# dev 镜像默认 docker build \ --target main \ --build-arg FROM_IMAGE_NAME$(cat docker/.ngc_version.dev) \ --build-arg IMAGE_TYPEdev \ -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_VERSION0.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等 extrasdev 变体额外包含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 打补丁后从源码安装 DeepEPlts变体跳过。镜像默认以非 root 用户nemo-runtimeUID/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 \ --nodes1 --gpus-per-node8 \ --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 \ --nodes1 --gpus-per-node8 \ --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运行时训练 extrasdev完整开发环境TransformerEngine、ModelOpt 等testpytest、coverage、nemo-runlintingruff、black、isort、pylintbuildCython、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 lintingLTS 环境的复现方式是端到端构建docker/Dockerfile.ci.lts不存在等价的uv sync命令因为 LTS 依赖已不在pyproject.toml中。LTS 的顶层固定版本集位于docker/lts/requirements.txt例如tqdm4.67.3、megatron-energon[av_decode]7.3.2、flashinfer-python0.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更新它。新增依赖遵循三步工作流获取容器镜像—— 见上文 Step 1交互式启动容器—— 见上文 Step 2在容器内更新锁文件然后提交# 容器内 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 后出现ModuleNotFoundErrorpip 安装到了 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 adduv 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创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考