Add g0050 NeMo flash-attn4 training image

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 "