From 2c2b7ccc246ed1a798e686f4ba29901740afd6a1 Mon Sep 17 00:00:00 2001 From: yi_lu Date: Wed, 1 Jul 2026 23:10:37 +0800 Subject: [PATCH] Add Megatron-Bridge pretrain launcher --- .../export_pretrain_parquet_text_jsonl.py | 75 +++++++++ .../laoyao_2b_moe_nemo_megatron.yaml | 4 +- .../preprocess_megatron_bridge_pretrain.sh | 38 +++++ scripts/train_megatron_bridge_2b_moe.sh | 48 ++++++ training/megatron_bridge/README.md | 49 ++++++ .../megatron_bridge/laoyao_2b_moe_pretrain.py | 144 ++++++++++++++++++ .../nemo_megatron/pretrain_2b_moe_200b.yaml | 2 +- 7 files changed, 357 insertions(+), 3 deletions(-) create mode 100755 dataset/pretrain/scripts/export_pretrain_parquet_text_jsonl.py create mode 100755 scripts/preprocess_megatron_bridge_pretrain.sh create mode 100755 scripts/train_megatron_bridge_2b_moe.sh create mode 100644 training/megatron_bridge/README.md create mode 100755 training/megatron_bridge/laoyao_2b_moe_pretrain.py diff --git a/dataset/pretrain/scripts/export_pretrain_parquet_text_jsonl.py b/dataset/pretrain/scripts/export_pretrain_parquet_text_jsonl.py new file mode 100755 index 0000000..5ec093a --- /dev/null +++ b/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() diff --git a/model/nemo_megatron/laoyao_2b_moe_nemo_megatron.yaml b/model/nemo_megatron/laoyao_2b_moe_nemo_megatron.yaml index 7fc6498..8b3ebff 100644 --- a/model/nemo_megatron/laoyao_2b_moe_nemo_megatron.yaml +++ b/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 diff --git a/scripts/preprocess_megatron_bridge_pretrain.sh b/scripts/preprocess_megatron_bridge_pretrain.sh new file mode 100755 index 0000000..7cfb144 --- /dev/null +++ b/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" diff --git a/scripts/train_megatron_bridge_2b_moe.sh b/scripts/train_megatron_bridge_2b_moe.sh new file mode 100755 index 0000000..056caed --- /dev/null +++ b/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 + " diff --git a/training/megatron_bridge/README.md b/training/megatron_bridge/README.md new file mode 100644 index 0000000..61a1fe9 --- /dev/null +++ b/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 前缀。 diff --git a/training/megatron_bridge/laoyao_2b_moe_pretrain.py b/training/megatron_bridge/laoyao_2b_moe_pretrain.py new file mode 100755 index 0000000..7c25889 --- /dev/null +++ b/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() diff --git a/training/nemo_megatron/pretrain_2b_moe_200b.yaml b/training/nemo_megatron/pretrain_2b_moe_200b.yaml index 405f79d..148cfaf 100644 --- a/training/nemo_megatron/pretrain_2b_moe_200b.yaml +++ b/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