Add Megatron-Bridge pretrain launcher

training/megatron_bridge/README.md

@@ -0,0 +1,49 @@
+# 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 前缀。

training/megatron_bridge/laoyao_2b_moe_pretrain.py

@@ -0,0 +1,144 @@
+#!/usr/bin/env python3
+from __future__ import annotations
+
+import argparse
+from pathlib import Path
+
+import torch
+import torch.nn.functional as F
+
+from megatron.bridge.models.gpt_provider import GPTModelProvider
+from megatron.bridge.recipes.common import _pretrain_common
+from megatron.bridge.training.config import ConfigContainer
+from megatron.bridge.training.gpt_step import forward_step
+from megatron.bridge.training.pretrain import pretrain
+
+
+REPO_ROOT = Path(__file__).resolve().parents[2]
+TOKENIZER_DIR = REPO_ROOT / "tokenizer/glm5.2"
+
+
+def build_config(args: argparse.Namespace) -> ConfigContainer:
+    cfg = _pretrain_common()
+
+    cfg.model = GPTModelProvider(
+        num_layers=12,
+        hidden_size=1536,
+        num_attention_heads=24,
+        num_query_groups=4,
+        ffn_hidden_size=4608,
+        seq_length=args.seq_length,
+        vocab_size=154820,
+        should_pad_vocab=True,
+        share_embeddings_and_output_weights=False,
+        position_embedding_type="rope",
+        normalization="RMSNorm",
+        gated_linear_unit=True,
+        activation_func=F.silu,
+        num_moe_experts=12,
+        moe_layer_freq=[1 if idx in {2, 4, 6, 8, 10} else 0 for idx in range(12)],
+        moe_ffn_hidden_size=6144,
+        moe_router_topk=4,
+        moe_router_load_balancing_type="aux_loss",
+        moe_aux_loss_coeff=0.02,
+        moe_z_loss_coeff=0.001,
+        moe_token_dispatcher_type="alltoall",
+        moe_expert_capacity_factor=1.25,
+        moe_router_enable_expert_bias=True,
+        moe_router_bias_update_rate=0.02,
+        moe_grouped_gemm=True,
+        tensor_model_parallel_size=args.tensor_parallel,
+        pipeline_model_parallel_size=args.pipeline_parallel,
+        expert_model_parallel_size=args.expert_parallel,
+        context_parallel_size=args.context_parallel,
+        sequence_parallel=args.tensor_parallel > 1,
+        transformer_impl="transformer_engine",
+        attention_backend="flash",
+        init_method_std=0.02,
+    )
+
+    cfg.tokenizer.tokenizer_type = "HuggingFaceTokenizer"
+    cfg.tokenizer.tokenizer_model = str(TOKENIZER_DIR)
+    cfg.dataset.seq_length = args.seq_length
+    cfg.dataset.split = args.split
+    cfg.dataset.num_workers = args.dataset_workers
+    if args.data_prefix:
+        cfg.dataset.data_path = args.data_prefix
+        cfg.dataset.blend = [(args.data_prefix, 1.0)]
+    else:
+        cfg.dataset.blend = None
+
+    cfg.train.train_iters = args.train_iters
+    cfg.train.micro_batch_size = args.micro_batch_size
+    cfg.train.global_batch_size = args.global_batch_size
+    cfg.optimizer.lr = args.lr
+    cfg.optimizer.min_lr = args.min_lr
+    cfg.optimizer.weight_decay = args.weight_decay
+    cfg.scheduler.lr_warmup_fraction = args.warmup_fraction
+
+    cfg.checkpoint.save = args.save_dir
+    cfg.checkpoint.load = args.load_dir
+    cfg.checkpoint.save_interval = args.save_interval
+
+    cfg.logger.tensorboard_dir = args.tensorboard_dir
+    cfg.logger.log_interval = args.log_interval
+    cfg.validation.eval_interval = args.eval_interval
+    cfg.validation.eval_iters = args.eval_iters
+
+    cfg.ddp.use_megatron_fsdp = False
+    cfg.ddp.overlap_grad_reduce = True
+    cfg.ddp.overlap_param_gather = True
+    cfg.ddp.check_for_nan_in_grad = True
+    cfg.ddp.use_distributed_optimizer = True
+
+    return cfg
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(description="Laoyao 2B MoE Megatron-Bridge pretrain launcher.")
+    parser.add_argument("--data-prefix", default=None, help="Megatron indexed dataset prefix without .bin/.idx.")
+    parser.add_argument("--seq-length", type=int, default=8192)
+    parser.add_argument("--train-iters", type=int, default=10)
+    parser.add_argument("--micro-batch-size", type=int, default=1)
+    parser.add_argument("--global-batch-size", type=int, default=1024)
+    parser.add_argument("--tensor-parallel", type=int, default=1)
+    parser.add_argument("--pipeline-parallel", type=int, default=1)
+    parser.add_argument("--expert-parallel", type=int, default=1)
+    parser.add_argument("--context-parallel", type=int, default=1)
+    parser.add_argument("--dataset-workers", type=int, default=8)
+    parser.add_argument("--split", default="9999,8,2")
+    parser.add_argument("--lr", type=float, default=3.0e-4)
+    parser.add_argument("--min-lr", type=float, default=5.0e-6)
+    parser.add_argument("--weight-decay", type=float, default=0.1)
+    parser.add_argument("--warmup-fraction", type=float, default=0.01)
+    parser.add_argument("--save-dir", default="/mnt/beegfs/yi/laoyao_2b_moe/runs/megatron_bridge/checkpoints")
+    parser.add_argument("--load-dir", default=None)
+    parser.add_argument("--tensorboard-dir", default="/mnt/beegfs/yi/laoyao_2b_moe/runs/megatron_bridge/tensorboard")
+    parser.add_argument("--save-interval", type=int, default=1000)
+    parser.add_argument("--log-interval", type=int, default=10)
+    parser.add_argument("--eval-interval", type=int, default=1000)
+    parser.add_argument("--eval-iters", type=int, default=10)
+    parser.add_argument("--dry-run", action="store_true")
+    return parser.parse_args()
+
+
+def main() -> None:
+    args = parse_args()
+    cfg = build_config(args)
+    if args.dry_run:
+        print("laoyao_megatron_bridge_config_ok")
+        print(f"seq_length={cfg.model.seq_length}")
+        print(f"hidden_size={cfg.model.hidden_size}")
+        print(f"num_moe_experts={cfg.model.num_moe_experts}")
+        print(f"moe_router_topk={cfg.model.moe_router_topk}")
+        print(f"moe_layer_freq={cfg.model.moe_layer_freq}")
+        print(f"moe_router_load_balancing_type={cfg.model.moe_router_load_balancing_type}")
+        print(f"moe_aux_loss_coeff={cfg.model.moe_aux_loss_coeff}")
+        print(f"tokenizer_model={cfg.tokenizer.tokenizer_model}")
+        print(f"dataset_prefix={args.data_prefix or 'mock'}")
+        return
+    pretrain(config=cfg, forward_step_func=forward_step)
+
+
+if __name__ == "__main__":
+    main()

training/nemo_megatron/pretrain_2b_moe_200b.yaml

@@ -14,7 +14,7 @@ training:
   target_tokens: 200000000000
   global_batch_size: 1024
   micro_batch_size: 1
-  seq_length: 16384
+  seq_length: 8192
   optimizer: adamw
   learning_rate: 3.0e-4
   min_learning_rate: 5.0e-6