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