Add g0050 NeMo flash-attn4 training image

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

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
+```

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 - <<PY
+import importlib.metadata as metadata
+for package in ["torch", "transformer-engine", "flash-attn", "flash-attn-4", "nvidia-cutlass-dsl", "megatron-core", "megatron-bridge"]:
+    try:
+        print(package, metadata.version(package))
+    except metadata.PackageNotFoundError:
+        print(package, "MISSING")
+import cutlass
+print("cutlass import OK")
+PY
+'
+
+echo "Built image: $IMAGE_TAG"

scripts/train_megatron_bridge_2b_moe.sh

@@ -2,10 +2,10 @@
 set -euo pipefail
 
 REPO_ROOT="${REPO_ROOT:-/mnt/beegfs/yi/laoyao_2b_moe}"
-IMAGE="${IMAGE:-nvcr.io/nvidia/nemo:26.06}"
+IMAGE="${IMAGE:-laoyao/nemo-megatron:26.06-flashattn4}"
 NPROC_PER_NODE="${NPROC_PER_NODE:-8}"
 DATA_PREFIX="${DATA_PREFIX:-}"
-DATA_MANIFEST="${DATA_MANIFEST:-/mnt/beegfs/yi/laoyao_2b_moe_pretraining_dataset/megatron_bridge/pretrain_65536_direct_v1/manifest.json}"
+DATA_MANIFEST="${DATA_MANIFEST-/mnt/beegfs/yi/laoyao_2b_moe_pretraining_dataset/megatron_bridge/pretrain_65536_direct_v1/manifest.json}"
 TRAIN_ITERS="${TRAIN_ITERS:-10}"
 SEQ_LENGTH="${SEQ_LENGTH:-8192}"
 MICRO_BATCH_SIZE="${MICRO_BATCH_SIZE:-1}"
@@ -36,13 +36,21 @@ fi
 DATA_ARGS=()
 if [[ -n "$DATA_MANIFEST" ]]; then
   DATA_ARGS=(--data-manifest "$DATA_MANIFEST")
-else
+elif [[ -n "$DATA_PREFIX" ]]; then
   DATA_ARGS=(--data-prefix "$DATA_PREFIX")
 fi
 
+DOCKER_MOUNTS=()
+if [[ -d /mnt/beegfs ]]; then
+  DOCKER_MOUNTS+=(-v /mnt/beegfs:/mnt/beegfs)
+fi
+if [[ -d /ssd ]]; then
+  DOCKER_MOUNTS+=(-v /ssd:/ssd)
+fi
+
 docker run --rm --gpus all --ipc=host --network=host \
   --ulimit memlock=-1 --ulimit stack=67108864 \
-  -v /mnt/beegfs:/mnt/beegfs \
+  "${DOCKER_MOUNTS[@]}" \
   -w "$REPO_ROOT" \
   "$IMAGE" \
   bash -lc "