# Megatron-Bridge Training

本目录是 Laoyao 2B MoE 在 NVIDIA Megatron-Bridge 上的训练适配层。

当前策略：

- 不改变原模型参数规模：`hidden_size=1536`、`12 experts`、`topk=4`、`5` 个 MoE layer。
- 训练上下文先用 `seq_length=8192`，不要一开始上 16K。
- tokenizer 使用 repo 内已验证的 GLM-5.2 tokenizer：`tokenizer/glm5.2`。
- 从零预训练使用 Megatron indexed dataset，不能直接把 parquet 喂给 Bridge pretrain。

## 文件

- `laoyao_2b_moe_pretrain.py`：自定义 Megatron-Bridge recipe/launcher。支持 `--dry-run`，用于在不启动训练 loop 的情况下检查配置。
- `../../scripts/preprocess_megatron_bridge_pretrain.sh`：从 parquet 导出 JSONL，并调用 Megatron-LM `preprocess_data.py` 生成 `.bin/.idx`。
- `../../scripts/train_megatron_bridge_2b_moe.sh`：Docker + torchrun 启动入口。

## 数据准备

Bridge 的 LLM pretrain 数据路径必须是 Megatron indexed dataset prefix：

```bash
bash scripts/preprocess_megatron_bridge_pretrain.sh
```

默认输出：

```text
/mnt/beegfs/yi/laoyao_2b_moe_pretraining_dataset/megatron_bridge/pretrain_8192_v1/laoyao_2b_moe_8192_text_document.bin
/mnt/beegfs/yi/laoyao_2b_moe_pretraining_dataset/megatron_bridge/pretrain_8192_v1/laoyao_2b_moe_8192_text_document.idx
```

注意：`preprocess_data.py` 是按文档 tokenization，不在预处理阶段固定切成 8192 行；训练时由 `GPTDatasetConfig.seq_length=8192` 生成固定长度训练 sample。

## Dry Run

H200 被占用时可以先跑单进程 dry-run：

```bash
DRY_RUN=1 NPROC_PER_NODE=1 bash scripts/train_megatron_bridge_2b_moe.sh
```

GPU 空出来后再跑小步数：

```bash
TRAIN_ITERS=5 NPROC_PER_NODE=8 bash scripts/train_megatron_bridge_2b_moe.sh
```

如果还没有构建真实 `.bin/.idx`，可以先把 `DATA_PREFIX` 指向一个小规模 smoke 前缀。
