84 lines
2.8 KiB
Markdown
84 lines
2.8 KiB
Markdown
# 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 必须从该镜像中移除。NeMo 26.06 将 FA2 安装在系统
|
||
`/usr/local/lib/python3.12/dist-packages`，而 FA4 安装在 `/opt/venv`；如果两者并存，
|
||
系统 FA2 会抢占顶层 `flash_attn` 包，Transformer Engine 导入
|
||
`flash_attn.cute.interface` 时会失败。
|
||
|
||
构建脚本固定并验证以下组合：
|
||
|
||
- `INSTALL_FLASH_ATTN4=1`，默认启用 FA4 安装。
|
||
- `pip --pre`，允许安装 FA4 的 beta 发行版。
|
||
- `flash-attn-4==4.0.0b11`。
|
||
- `nvidia-cutlass-dsl`、`nvidia-cutlass-dsl-libs-base` 和
|
||
`nvidia-cutlass-dsl-libs-cu13` 均为 `4.5.2`。
|
||
- 同时从系统 Python 和 `/opt/venv` 删除旧 FA2/FA3，再安装 FA4。
|
||
- 构建后必须通过 `flash_attn.cute.interface`、`cutlass` 和
|
||
`triton_kernels.matmul_ogs` 的真实导入检查；仅看到 distribution metadata 不算成功。
|
||
|
||
## 构建镜像
|
||
|
||
在 g0050 上执行：
|
||
|
||
```bash
|
||
cd /ssd/workspace/yi/laoyao_2b_moe
|
||
bash scripts/build_nemo_megatron_image.sh
|
||
```
|
||
|
||
如果机器上已有 `ti-coding-agent/nemo-bridge-flashattn4:26.06`，可以将其作为 base image；
|
||
无论 base 是否已包含 FA4，构建脚本默认都会重新校验并固化上述版本。最终镜像为：
|
||
|
||
```text
|
||
laoyao/nemo-megatron:26.06-flashattn4
|
||
```
|
||
|
||
如果换机器没有这个缓存镜像，从 NVIDIA 官方 NeMo 26.06 镜像构建：
|
||
|
||
```bash
|
||
BASE_IMAGE=nvcr.io/nvidia/nemo:26.06 \
|
||
bash scripts/build_nemo_megatron_image.sh
|
||
```
|
||
|
||
`INSTALL_FLASH_ATTN4` 默认已设为 `1`，上面的命令无需额外传入该参数。只有明确进行
|
||
官方 NeMo 原始环境对照实验时，才应设置 `INSTALL_FLASH_ATTN4=0`。
|
||
|
||
FA4 和 CUTLASS DSL 是预发布/快速迭代包。不要将版本改成未固定的最新版，否则可能造成
|
||
CUTLASS base 与 `libs-cu13` 版本不一致。阿里云 simple index 未完整列出这两个 CUTLASS
|
||
wheel，因此 Dockerfile 使用固定的阿里云 wheel URL，避免回退到低速的
|
||
`files.pythonhosted.org`。
|
||
|
||
脚本默认使用 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
|
||
```
|