62 lines
1.5 KiB
Markdown
62 lines
1.5 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 已从该镜像中移除，避免 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
|
||
```
|