NVIDIA NeMo / Megatron-Bridge Backend
本项目的训练后端目标是 NVIDIA 官方 NeMo/Megatron 栈，而不是旧的手写 PyTorch trainer。
默认训练镜像
g0050 上当前默认训练镜像：
laoyao/nemo-megatron:26.06-flashattn4
该镜像基于 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 上执行：
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，并打出：
laoyao/nemo-megatron:26.06-flashattn4
如果换机器没有这个缓存镜像，可以从 NVIDIA 官方 NeMo 26.06 镜像补装 flash-attn4：
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 scripts/train_megatron_bridge_2b_moe.sh
训练脚本默认使用 laoyao/nemo-megatron:26.06-flashattn4。如需临时切回官方镜像，可覆盖：
IMAGE=nvcr.io/nvidia/nemo:26.06 bash scripts/train_megatron_bridge_2b_moe.sh