From 323452b9bd96f1ad334f987597ce56eda4925390 Mon Sep 17 00:00:00 2001 From: yi_lu Date: Thu, 2 Jul 2026 21:10:58 +0800 Subject: [PATCH] Add g0050 NeMo flash-attn4 training image --- docker/nemo/Dockerfile.nemo-megatron | 31 +++++++++++++ docker/nemo/README.md | 58 +++++++++++++++++++++++-- scripts/build_nemo_megatron_image.sh | 37 ++++++++++++++++ scripts/train_megatron_bridge_2b_moe.sh | 16 +++++-- 4 files changed, 134 insertions(+), 8 deletions(-) create mode 100644 docker/nemo/Dockerfile.nemo-megatron create mode 100755 scripts/build_nemo_megatron_image.sh diff --git a/docker/nemo/Dockerfile.nemo-megatron b/docker/nemo/Dockerfile.nemo-megatron new file mode 100644 index 0000000..101271d --- /dev/null +++ b/docker/nemo/Dockerfile.nemo-megatron @@ -0,0 +1,31 @@ +ARG BASE_IMAGE=nvcr.io/nvidia/nemo:26.06 +FROM ${BASE_IMAGE} + +ARG INSTALL_FLASH_ATTN4=0 +ARG PIP_INDEX_URL=https://mirrors.aliyun.com/pypi/simple/ +ARG PIP_EXTRA_INDEX_URL=https://pypi.org/simple + +ENV PIP_INDEX_URL=${PIP_INDEX_URL} +ENV PIP_EXTRA_INDEX_URL=${PIP_EXTRA_INDEX_URL} +ENV PIP_DISABLE_PIP_VERSION_CHECK=1 + +RUN if [ "${INSTALL_FLASH_ATTN4}" = "1" ]; then \ + python3 -m pip uninstall -y flash-attn flash-attn-3 flash-attn-4 || true; \ + MAX_JOBS="${MAX_JOBS:-16}" python3 -m pip install --no-cache-dir --no-build-isolation \ + "flash-attn-4" "nvidia-cutlass-dsl"; \ + fi + +RUN python3 - <<'PY' +import importlib.metadata as metadata +required = ["torch", "transformer-engine", "megatron-core", "megatron-bridge"] +for package in required: + print(package, metadata.version(package)) +try: + print("flash-attn-4", metadata.version("flash-attn-4")) +except metadata.PackageNotFoundError: + print("flash-attn-4 MISSING") +try: + print("nvidia-cutlass-dsl", metadata.version("nvidia-cutlass-dsl")) +except metadata.PackageNotFoundError: + print("nvidia-cutlass-dsl MISSING") +PY diff --git a/docker/nemo/README.md b/docker/nemo/README.md index b6b85a3..74c2a89 100644 --- a/docker/nemo/README.md +++ b/docker/nemo/README.md @@ -1,11 +1,61 @@ -# NVIDIA NeMo Backend +# NVIDIA NeMo / Megatron-Bridge Backend 本项目的训练后端目标是 NVIDIA 官方 NeMo/Megatron 栈，而不是旧的手写 PyTorch trainer。 -默认镜像： +## 默认训练镜像 + +g0050 上当前默认训练镜像： ```text -nvcr.io/nvidia/nemo:26.06 +laoyao/nemo-megatron:26.06-flashattn4 ``` -如果 g0033 拉取 NGC 镜像受限，可以先在可联网机器拉取后导入，或在本机配置 Docker daemon 代理。训练脚本只依赖 `/mnt/beegfs` 挂载，不把模型权重、数据 shard 或日志提交到 git。 +该镜像基于 NeMo 26.06/Megatron-Bridge 0.5.0，当前关键包版本： + +| package | version | +|---|---| +| torch | 2.12.0a0 NVIDIA build | +| transformer-engine | 2.16.0 | +| megatron-core | 0.18.0 | +| megatron-bridge | 0.5.0 | +| flash-attn-4 | 4.0.0b11 | +| nvidia-cutlass-dsl | 4.5.2 | + +`flash-attn` v2 已从该镜像中移除，避免 H200/B300 上旧 attention backend 的兼容问题。 + +## 构建镜像 + +在 g0050 上执行： + +```bash +cd /ssd/workspace/yi/laoyao_2b_moe +bash scripts/build_nemo_megatron_image.sh +``` + +默认会复用 g0050 已存在的 `ti-coding-agent/nemo-bridge-flashattn4:26.06` 作为 base image，并打出： + +```text +laoyao/nemo-megatron:26.06-flashattn4 +``` + +如果换机器没有这个缓存镜像，可以从 NVIDIA 官方 NeMo 26.06 镜像补装 flash-attn4： + +```bash +BASE_IMAGE=nvcr.io/nvidia/nemo:26.06 \ +INSTALL_FLASH_ATTN4=1 \ +bash scripts/build_nemo_megatron_image.sh +``` + +脚本默认使用 B300/g0050 代理 `http://100.72.0.101:8888`，pip 默认走阿里云源并保留 PyPI fallback。 + +## 训练入口 + +```bash +bash scripts/train_megatron_bridge_2b_moe.sh +``` + +训练脚本默认使用 `laoyao/nemo-megatron:26.06-flashattn4`。如需临时切回官方镜像，可覆盖： + +```bash +IMAGE=nvcr.io/nvidia/nemo:26.06 bash scripts/train_megatron_bridge_2b_moe.sh +``` diff --git a/scripts/build_nemo_megatron_image.sh b/scripts/build_nemo_megatron_image.sh new file mode 100755 index 0000000..870abe1 --- /dev/null +++ b/scripts/build_nemo_megatron_image.sh @@ -0,0 +1,37 @@ +#!/usr/bin/env bash +set -euo pipefail + +REPO_ROOT="${REPO_ROOT:-$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)}" +IMAGE_TAG="${IMAGE_TAG:-laoyao/nemo-megatron:26.06-flashattn4}" +BASE_IMAGE="${BASE_IMAGE:-ti-coding-agent/nemo-bridge-flashattn4:26.06}" +INSTALL_FLASH_ATTN4="${INSTALL_FLASH_ATTN4:-0}" + +export http_proxy="${http_proxy:-http://100.72.0.101:8888}" +export https_proxy="${https_proxy:-http://100.72.0.101:8888}" +export HTTP_PROXY="${HTTP_PROXY:-$http_proxy}" +export HTTPS_PROXY="${HTTPS_PROXY:-$https_proxy}" + +docker build \ + --network=host \ + --build-arg BASE_IMAGE="$BASE_IMAGE" \ + --build-arg INSTALL_FLASH_ATTN4="$INSTALL_FLASH_ATTN4" \ + --build-arg PIP_INDEX_URL="${PIP_INDEX_URL:-https://mirrors.aliyun.com/pypi/simple/}" \ + --build-arg PIP_EXTRA_INDEX_URL="${PIP_EXTRA_INDEX_URL:-https://pypi.org/simple}" \ + -f "$REPO_ROOT/docker/nemo/Dockerfile.nemo-megatron" \ + -t "$IMAGE_TAG" \ + "$REPO_ROOT" + +docker run --rm --gpus all --ipc=host --network=host "$IMAGE_TAG" bash -lc ' +python3 - <