Megatron-LM 依赖升级实战：用 upgrade_dependencies.sh 维护 uv 锁文件与全量依赖

Megatron-LM 依赖升级实战用 upgrade_dependencies.sh 维护 uv 锁文件与全量依赖【免费下载链接】Megatron-LMOngoing research training transformer models at scale项目地址: https://gitcode.com/GitHub_Trending/me/Megatron-LM本文聚焦 Megatron-LM 仓库内的依赖管理工具tools/upgrade_dependencies.sh它基于 Docker 容器封装uv命令用于更新uv.lock锁文件并可选地将所有依赖升级到当前允许的最新版本。读完本文你将掌握该脚本的前置条件、环境变量设置、两种运行模式与全部命令行参数并能够结合仓库内 pyproject.toml、uv.lock 与 CI 镜像定义理解 Megatron-LM 依赖锁定、可复现构建与定期升级的完整工作流。为什么 Megatron-LM 需要专用的依赖升级工具Megatron-LM 是一个大型 Transformer 模型训练框架其依赖集合高度复杂既包含 PyTorch、Transformer Engine 这类重量级计算库也包含大量从 Git 源直接引用的组件如 FlashMLA、DeepGEMM、mamba-ssm、nemo-lens 等见 pyproject.toml 中的[tool.uv.sources]。这种结构给依赖管理带来两个核心挑战可复现性训练大模型时环境中任何一个库的版本漂移都可能改变数值行为甚至破坏算子构建。仓库通过 uv.lock约 6000 行锁定整个依赖解析树的精确版本保证 CI 与开发者环境构建一致。升级复杂度直接在本机执行uv lock --upgrade依赖本机的 CUDA、编译工具链和 Python 环境不同开发者机器上容易产生不一致的解析结果。因此仓库将升级流程封装进一个预置好依赖环境的 CI 基础镜像中保证任何人在任何机器上得到同一份 lockfile。pyproject.toml 中的[tool.uv]段落体现了这种管理思路managed true声明该项目由 uv 全权管理default-groups [linting, build, test]指定默认安装组no-build-isolation-package列出一批必须借用现有环境编译的包如transformer-engine、mamba-ssm、deep_gemmoverride-dependencies则对 torch、torchvision、triton 及存在安全公告的 urllib3 做全局覆盖。uv.lock正是依据这些约束解析生成的最终产物tools/upgrade_dependencies.sh负责让这份产物始终保持最新。前置条件根据 tools/upgrade_dependencies.md 与脚本实现运行前需要满足两项要求Docker脚本通过docker run在容器内执行 uv因此宿主机必须安装并可用 Docker。GITLAB_ENDPOINT环境变量脚本需要从 ADLR/Megatron-LM 项目的内部 GitLab 容器镜像仓库拉取mcore_ci_dev:main镜像该仓库主机名通过此变量提供。环境变量设置GITLAB_ENDPOINT在运行脚本之前需要导出内部 GitLab 镜像仓库主机名。注意格式要求不含 scheme不要写https://末尾不带斜杠export GITLAB_ENDPOINTinternal-gitlab-hostname该值属于内部基础设施信息文档建议向团队成员确认或从本地环境 / CI secrets 中获取。脚本在启动时即检查该变量if [ -z ${GITLAB_ENDPOINT:-} ]; then echo GITLAB_ENDPOINT is not set. Please set the GITLAB_ENDPOINT environment variable to the gitlab endpoint of the Megatron-LM repository. exit 1 fi未设置时脚本会立即退出并给出提示这保证了后续docker pull不会因主机名为空而失败。使用方法与命令行选项脚本提供两种运行模式均从仓库根目录执行# 仅更新锁文件默认行为 ./tools/upgrade_dependencies.sh # 升级全部依赖并更新锁文件 ./tools/upgrade_dependencies.sh --upgrade参数说明如下表完整继承自 tools/upgrade_dependencies.mdFlag说明--upgrade在更新锁文件的同时将依赖升级到允许的最新版本--help显示用法信息命令行解析逻辑位于 tools/upgrade_dependencies.sh脚本以set -eoxu pipefail开头出错即停、打印执行轨迹、未定义变量报错、管道失败传播随后遍历所有参数仅接受--upgrade与--help其余参数一律报Unknown argument并退出。--help输出会额外列出环境变量说明其中明确GITLAB_ENDPOINT的取值示例为不带 scheme 的主机名如gitlab.example.com。两种模式的核心差异uv lock 与 uv lock --upgrade脚本根据UPGRADE标志构造传给 uv 的参数UV_ARGS(lock) if [ $UPGRADE true ]; then UV_ARGS(--upgrade) fi默认模式uv lock仅依据 pyproject.toml 中声明的依赖约束重新解析并更新 uv.lock。已满足约束的依赖保持现有版本不变主要用于修正 lockfile 与实际声明不一致的情况例如新增了依赖但尚未锁定。升级模式uv lock --upgrade忽略 lockfile 中已有的解析结果将所有依赖重新解析到约束允许的最新版本并重写 uv.lock。这是依赖升级主流程仓库维护者定期运行该命令生成新锁文件随后提交 lockfile 变更。源码级解析脚本如何完成容器化升级整个升级过程在预构建的 CI 镜像内完成避免宿主机环境差异。核心命令如下docker run \ --rm \ -v $(pwd):/workdir/ \ -w /workdir/ \ $GITLAB_ENDPOINT/adlr/megatron-lm/mcore_ci_dev:main \ bash -ec export TMS_CUDA_MAJOR$(${CUDA_HOME:-/usr/local/cuda}/bin/nvcc --version | sed -n s/.*release \([0-9][0-9]*\).*/\1/p | head -1) test -n $TMS_CUDA_MAJOR exec uv $ bash ${UV_ARGS[]}逐段拆解--rm容器执行完毕即删除不残留临时容器。-v $(pwd):/workdir/将仓库根目录挂载进容器使uv.lock、pyproject.toml的变更直接写回宿主机工作区。脚本此前通过cd $SCRIPT_DIR/..切换到仓库根目录确保挂载的就是仓库根目录。-w /workdir/容器内工作目录指向挂载点uv 在此识别pyproject.toml与uv.lock。镜像$GITLAB_ENDPOINT/adlr/megatron-lm/mcore_ci_dev:main即 ADLR/Megatron-LM 项目的内部 CI 开发镜像。该镜像由 docker/Dockerfile.ci.dev 构建其中已安装 uv固定UV_VERSION0.7.2含 SHA256 校验并预置了编译所需工具链。bash -ec在容器内执行一段内联脚本。-e保证任一步失败即退出-c接收命令字符串末尾的bash作为$0占位后面的${UV_ARGS[]}作为位置参数传入最终exec uv $以子进程替换方式运行uv lock或uv lock --upgrade。TMS_CUDA_MAJORCUDA 主版本探测容器内内联脚本的第一步是从 nvcc 提取 CUDA 主版本并导出为TMS_CUDA_MAJORexport TMS_CUDA_MAJOR$(${CUDA_HOME:-/usr/local/cuda}/bin/nvcc --version | sed -n s/.*release \([0-9][0-9]*\).*/\1/p | head -1) test -n $TMS_CUDA_MAJOR其含义与用途在仓库多处出现TMS即Torch Memory Saver。在 docker/common/install.sh 中安装脚本注释明确指出 torch-memory-saver builds CUDA-suffixed extensions and requires the CUDA major即该库会构建带 CUDA 后缀的扩展必须知道 CUDA 主版本。docker/Dockerfile.ci.dev在构建阶段也执行了完全相同的探测逻辑。升级锁文件时部分依赖会重新编译因此容器内预先导出该变量保证诸如 torch-memory-saver 这类需要 CUDA 版本信息的包在解析/安装路径中行为一致。test -n用于确保探测结果非空否则整个容器命令直接失败避免在 CUDA 环境缺失时产出错误锁文件。锁文件如何影响 CI 镜像构建升级uv.lock的最终消费者是 CI 镜像构建流程。以 docker/Dockerfile.ci.dev 为例构建时将README.md pyproject.toml uv.lock复制进/workspace/随后执行uv sync --only-group build UV_CONCURRENT_INSTALLS1 uv sync -v \ --extra ${IMAGE_TYPE} --extra inference --extra mlm --extra ssm --extra te ${FLASH_MLA_GROUP} --link-mode copy --locked \ --no-install-package torch \ ...其中--locked要求安装必须严格遵循uv.lock若锁文件与pyproject.toml声明不一致构建会直接失败而非悄悄解析。这一机制把lockfile 必须保持最新且一致从约定变成了 CI 强约束——这正是 tools/upgrade_dependencies.sh 存在的意义在提交新依赖或版本变更前先用容器化脚本刷新锁文件确保 CI 镜像可继续用--locked稳定构建。与 LTS 镜像的区别仓库还维护一条 LTS 发布通道docker/Dockerfile.ci.lts。与 dev 镜像不同LTS 镜像每年仅 bump 一次、刻意滞后于浮动的 dev 标签其 Python 依赖不放在pyproject.toml中而是直接固定在 docker/lts/requirements.txt例如tqdm4.67.3、einops0.8.2、megatron-energon[av_decode]7.3.2以避免与pyproject.toml的模块级 extras 冲突。因此升级脚本所更新的uv.lock主要服务于 dev 通道LTS 固定集的升级则按该文件头注释描述的方式先编辑版本号或基于requirements.in重新uv pip compile再重建 LTS 镜像并运行 LTS CI 流水线。使用建议与注意事项GITLAB_ENDPOINT 格式只填主机名不要带https://或尾部/否则拼出的镜像名无效该主机名对应内部镜像仓库外部网络环境不可用。保持锁文件与声明同步无论是否使用--upgrade脚本生成的锁文件都应连同pyproject.toml变更一并提交CI 的--locked模式会拒绝不一致的锁文件。--upgrade的审查成本该模式会一次性升级全部依赖到最新可能引入算子库、优化器库的兼容性变化。建议升级后在 tests 对应的单元测试与功能测试用例上验证尤其是 pyproject.toml 中no-build-isolation-package列出的源码编译包Transformer Engine、mamba-ssm 等。容器内编译前提镜像基于 NGC PyTorch 基础镜像构建docker/Dockerfile.ci.dev 默认FROM nvcr.io/nvidia/pytorch:26.06-py3自带 nvcc 与 CUDA 工具链因此探测TMS_CUDA_MAJOR才能成功本机直接执行同样命令则需要自行保证 CUDA 环境。只读仓库约束本文仅介绍查看、安装、配置与运行方式锁文件的实际刷新应在你的本地克隆或 CI 环境中进行。总结tools/upgrade_dependencies.sh是 Megatron-LM 依赖生命周期管理的关键工具它以内部 CI 镜像为运行环境把uv lock/uv lock --upgrade封装成跨机器一致的命令输出一份始终与pyproject.toml声明同步的 uv.lock进而支撑 CI 中uv sync --locked的可复现构建。对于希望为 Megatron-LM 贡献新依赖、或者维护衍生训练环境的开发者而言掌握该脚本的用法与底层机制是保障环境一致性和构建稳定性的第一步。【免费下载链接】Megatron-LMOngoing research training transformer models at scale项目地址: https://gitcode.com/GitHub_Trending/me/Megatron-LM创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考