# 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。
- 训练数据 blend 优先读取 `prefix_category_stats.json`，按目标 category mix 计算 shard 权重；如果没有 category stats，则退回按 manifest token 数加权。
- validation 使用单独 heldout 2.8k Megatron indexed prefix，不再从训练 split 里切 valid。

## 文件

- `laoyao_2b_moe_pretrain.py`：自定义 Megatron-Bridge recipe/launcher。包含 GLM-5.2 vocab、MoE grouped GEMM、heldout validation、checkpoint 保留、recompute、distributed optimizer/overlap 参数。
- `../../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。

当前 g0050 主训练使用直接 parquet -> Megatron indexed dataset 的路径：

```bash
SOURCE_DATA_DIRS=/ssd/workspace/yi/laoyao_2b_moe_pretraining_dataset/train/pretrain_rebalanced_web40_edu20_chinese10_science10_logic10_math5_code5_200b_v1_20260701:/ssd/workspace/yi/laoyao_2b_moe_pretraining_dataset/train/logic_topup_proof_pile_17b_v1_20260701 \
WORK_DIR=/ssd/workspace/yi/laoyao_2b_moe_pretraining_dataset/megatron_bridge/pretrain_8192_glm52_direct_v1 \
MAX_SEQ_LEN=8192 \
bash scripts/preprocess_megatron_bridge_pretrain_direct.sh
```

该转换脚本会在 `manifest.json` 中写入可直接传给 Megatron 的 `_text_document` prefix；不要手动拼接错误的 prefix。

## 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 前缀。
