用途:为 Megatron Bridge 搭建可复现的开发/构建环境。核心原则是两件事——在容器内构建与开发,并始终使用 uv。
为什么用容器:项目依赖 CUDA、NCCL、带 GPU 支持的 PyTorch、Transformer Engine,以及 TRT-LLM、vLLM、DeepEP 等可选组件,直接在宿主机安装脆弱且难复现,因此项目提供固定依赖的生产级 Dockerfile。容器可保证 CUDA/NCCL/cuDNN 版本一致、uv.lock 本地与 CI 解析一致、GPU 相关操作开箱可用。
容器选项:一是使用 NeMo Framework 容器(最快捷),标签可在 NGC 目录页查询,并用 docker run --rm -it --gpus all --shm-size=24g nvcr.io/nvidia/nemo:<tag> bash 启动;二是自行构建 Megatron Bridge 容器,构建命令、构建参数与完整镜像栈见 @docker/README.md。运行容器常挂载源码目录与 $HOME/.cache/uv 以避免重复下载 wheel。在带 Enroot/Pyxis 的 Slurm 集群上,可通过 srun 的 --container-image、--container-mounts 等参数直接传入容器;自定义源码树绑定时仅 rank 0 执行 uv sync,其余等待。
依赖管理:禁止 pip install、conda 或裸 python,所有依赖操作走 uv 且必须在容器内进行。常用命令包括 uv sync --locked、uv sync --locked --all-extras --all-groups、uv run python、uv run python -m torch.distributed.run、uv add <package>、uv add --optional --extra <group> <package>、uv lock(仅 Linux/容器)。安装 pre-commit 钩子用 uv run --group dev pre-commit install。
新增依赖:优先作为可选依赖提交,并先于功能 PR 单独提交依赖变更 PR;同时提交 pyproject.toml 与 uv.lock。uv.lock 仅限 Linux 生成(需解析 CUDA wheel),应在 Docker 内运行 uv lock。切换 MCore 分支后,dev 分支用不带 --locked 的 uv sync,切回 main 再用 uv sync --locked。
快速开始:克隆仓库并初始化 3rdparty/Megatron-LM 子模块,用 docker/Dockerfile.ci 构建镜像,启动开发 shell,安装 pre-commit 钩子,再用 run_recipe.py 的 vanilla_gpt_pretrain_config 做小规模 sanity check。
限制与常见问题:macOS 上 uv sync --locked 会失败(锁文件解析 CUDA wheel);pip 安装后出现 ModuleNotFoundError 需改用 uv add 与 uv sync;MCore 分支切换后 --locked 失败改用普通 uv sync;镜像内无 uv 需使用由 Dockerfile.ci 构建的镜像;磁盘不足需设置 UV_CACHE_DIR;pre-commit 的 ruff 报错可用 uv run ruff check --fix . 与 uv run ruff format . 修复。价格、收费、平台等信息暂未提供。