Add Megatron-Bridge pretrain launcher

dataset/pretrain/scripts/export_pretrain_parquet_text_jsonl.py

@@ -0,0 +1,75 @@
+#!/usr/bin/env python3
+from __future__ import annotations
+
+import argparse
+import glob
+import json
+from pathlib import Path
+
+import pyarrow.parquet as pq
+
+
+def iter_paths(patterns: list[str]) -> list[Path]:
+    paths: list[Path] = []
+    for pattern in patterns:
+        path = Path(pattern)
+        if path.is_dir():
+            paths.extend(path.glob("*.parquet"))
+            continue
+        matches = glob.glob(pattern)
+        if matches:
+            paths.extend(Path(match) for match in matches)
+            continue
+        if path.exists():
+            paths.append(path)
+            continue
+        raise FileNotFoundError(pattern)
+    return sorted(set(paths), key=str)
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(description="Export pretrain parquet text rows to Megatron JSONL.")
+    parser.add_argument("--input", action="append", required=True, help="Parquet directory/file/glob.")
+    parser.add_argument("--output", required=True, help="Output JSONL path.")
+    parser.add_argument("--text-field", default="text")
+    parser.add_argument("--batch-size", type=int, default=8192)
+    parser.add_argument("--max-docs", type=int, default=0, help="0 means no limit.")
+    parser.add_argument("--progress-every", type=int, default=25)
+    return parser.parse_args()
+
+
+def main() -> None:
+    args = parse_args()
+    paths = iter_paths(args.input)
+    if not paths:
+        raise SystemExit("no parquet inputs resolved")
+
+    output = Path(args.output)
+    output.parent.mkdir(parents=True, exist_ok=True)
+    docs = 0
+    with output.open("w", encoding="utf-8") as handle:
+        for file_idx, path in enumerate(paths, start=1):
+            parquet_file = pq.ParquetFile(path)
+            if args.text_field not in parquet_file.schema_arrow.names:
+                continue
+            for batch in parquet_file.iter_batches(
+                batch_size=args.batch_size,
+                columns=[args.text_field],
+                use_threads=True,
+            ):
+                for row in batch.to_pylist():
+                    text = row.get(args.text_field)
+                    if not isinstance(text, str) or not text.strip():
+                        continue
+                    handle.write(json.dumps({"text": text}, ensure_ascii=False) + "\n")
+                    docs += 1
+                    if args.max_docs and docs >= args.max_docs:
+                        print(f"export_done docs={docs} files_seen={file_idx}/{len(paths)} output={output}")
+                        return
+            if file_idx % args.progress_every == 0:
+                print(f"export_progress files={file_idx}/{len(paths)} docs={docs}")
+    print(f"export_done docs={docs} files_seen={len(paths)} output={output}")
+
+
+if __name__ == "__main__":
+    main()

model/nemo_megatron/laoyao_2b_moe_nemo_megatron.yaml

@@ -15,8 +15,8 @@ model:
   num_attention_heads: 24
   num_query_groups: 4
   ffn_hidden_size: 4608
-  max_position_embeddings: 16384
-  seq_length: 16384
+  max_position_embeddings: 8192
+  seq_length: 8192
   normalization: rmsnorm
   activation: swiglu
   position_embedding_type: rope

scripts/preprocess_megatron_bridge_pretrain.sh

@@ -0,0 +1,38 @@
+#!/usr/bin/env bash
+set -euo pipefail
+
+REPO_ROOT="${REPO_ROOT:-/mnt/beegfs/yi/laoyao_2b_moe}"
+IMAGE="${IMAGE:-nvcr.io/nvidia/nemo:26.06}"
+SOURCE_DATA="${SOURCE_DATA:-/mnt/beegfs/yi/laoyao_2b_moe_pretraining_dataset/train/pretrain_rebalanced_web40_edu20_chinese10_science10_logic10_math5_code5_200b_v1_20260701}"
+WORK_DIR="${WORK_DIR:-/mnt/beegfs/yi/laoyao_2b_moe_pretraining_dataset/megatron_bridge/pretrain_8192_v1}"
+JSONL="${JSONL:-$WORK_DIR/text.jsonl}"
+OUTPUT_PREFIX="${OUTPUT_PREFIX:-$WORK_DIR/laoyao_2b_moe_8192_text_document}"
+TOKENIZER_MODEL="${TOKENIZER_MODEL:-$REPO_ROOT/tokenizer/glm5.2}"
+WORKERS="${WORKERS:-16}"
+MAX_DOCS="${MAX_DOCS:-0}"
+
+mkdir -p "$WORK_DIR"
+
+docker run --rm --ipc=host --network=host \
+  --ulimit memlock=-1 --ulimit stack=67108864 \
+  -v /mnt/beegfs:/mnt/beegfs \
+  -w "$REPO_ROOT" \
+  "$IMAGE" \
+  bash -lc "
+    set -euo pipefail
+    python3 dataset/pretrain/scripts/export_pretrain_parquet_text_jsonl.py \
+      --input '$SOURCE_DATA/*.parquet' \
+      --output '$JSONL' \
+      --max-docs '$MAX_DOCS'
+    python3 /opt/Megatron-Bridge/3rdparty/Megatron-LM/tools/preprocess_data.py \
+      --input '$JSONL' \
+      --json-keys text \
+      --tokenizer-type HuggingFaceTokenizer \
+      --tokenizer-model '$TOKENIZER_MODEL' \
+      --append-eod \
+      --output-prefix '$OUTPUT_PREFIX' \
+      --workers '$WORKERS'
+    ls -lh '${OUTPUT_PREFIX}'*
+  "
+
+echo "Megatron indexed dataset prefix: $OUTPUT_PREFIX"

scripts/train_megatron_bridge_2b_moe.sh

@@ -0,0 +1,48 @@
+#!/usr/bin/env bash
+set -euo pipefail
+
+REPO_ROOT="${REPO_ROOT:-/mnt/beegfs/yi/laoyao_2b_moe}"
+IMAGE="${IMAGE:-nvcr.io/nvidia/nemo:26.06}"
+NPROC_PER_NODE="${NPROC_PER_NODE:-8}"
+DATA_PREFIX="${DATA_PREFIX:-/mnt/beegfs/yi/laoyao_2b_moe_pretraining_dataset/megatron_bridge/pretrain_8192_v1/laoyao_2b_moe_8192_text_document}"
+TRAIN_ITERS="${TRAIN_ITERS:-10}"
+SEQ_LENGTH="${SEQ_LENGTH:-8192}"
+MICRO_BATCH_SIZE="${MICRO_BATCH_SIZE:-1}"
+GLOBAL_BATCH_SIZE="${GLOBAL_BATCH_SIZE:-1024}"
+TP="${TP:-1}"
+PP="${PP:-1}"
+EP="${EP:-1}"
+CP="${CP:-1}"
+DRY_RUN="${DRY_RUN:-0}"
+
+if [[ "$DRY_RUN" != "1" && ! -f "${DATA_PREFIX}.idx" ]]; then
+  echo "missing Megatron indexed data prefix: $DATA_PREFIX" >&2
+  echo "run scripts/preprocess_megatron_bridge_pretrain.sh first, or set DATA_PREFIX" >&2
+  exit 1
+fi
+
+DRY_RUN_ARG=""
+if [[ "$DRY_RUN" == "1" ]]; then
+  DRY_RUN_ARG="--dry-run"
+fi
+
+docker run --rm --gpus all --ipc=host --network=host \
+  --ulimit memlock=-1 --ulimit stack=67108864 \
+  -v /mnt/beegfs:/mnt/beegfs \
+  -w "$REPO_ROOT" \
+  "$IMAGE" \
+  bash -lc "
+    set -euo pipefail
+    torchrun --nproc_per_node='$NPROC_PER_NODE' \
+      training/megatron_bridge/laoyao_2b_moe_pretrain.py \
+      --data-prefix '$DATA_PREFIX' \
+      --seq-length '$SEQ_LENGTH' \
+      --train-iters '$TRAIN_ITERS' \
+      --micro-batch-size '$MICRO_BATCH_SIZE' \
+      --global-batch-size '$GLOBAL_BATCH_SIZE' \
+      --tensor-parallel '$TP' \
+      --pipeline-parallel '$PP' \
+      --expert-parallel '$EP' \
+      --context-parallel '$CP' \
+      $DRY_RUN_ARG
+  "

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