# NVIDIA NeMo / Megatron-Bridge Backend

本项目的训练后端目标是 NVIDIA 官方 NeMo/Megatron 栈，而不是旧的手写 PyTorch trainer。

## 默认训练镜像

g0050 上当前默认训练镜像：

```text
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 上执行：

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