Add Megatron data manifest and g0050 setup

training/megatron_bridge/laoyao_2b_moe_pretrain.py

@@ -2,6 +2,7 @@
 from __future__ import annotations
 
 import argparse
+import json
 from pathlib import Path
 
 import torch
@@ -18,6 +19,20 @@ REPO_ROOT = Path(__file__).resolve().parents[2]
 TOKENIZER_DIR = REPO_ROOT / "tokenizer/glm5.2"
 
 
+def load_data_blend(args: argparse.Namespace) -> list[tuple[str, float]] | None:
+    if args.data_manifest:
+        manifest_path = Path(args.data_manifest)
+        manifest = json.loads(manifest_path.read_text(encoding="utf-8"))
+        prefixes = manifest.get("ok_prefixes") or []
+        if not prefixes:
+            raise ValueError(f"{manifest_path} has no ok_prefixes")
+        weight = 1.0 / len(prefixes)
+        return [(str(prefix), weight) for prefix in prefixes]
+    if args.data_prefix:
+        return [(args.data_prefix, 1.0)]
+    return None
+
+
 def build_config(args: argparse.Namespace) -> ConfigContainer:
     cfg = _pretrain_common()
 
@@ -62,9 +77,10 @@ def build_config(args: argparse.Namespace) -> ConfigContainer:
     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)]
+    data_blend = load_data_blend(args)
+    if data_blend:
+        cfg.dataset.data_path = data_blend[0][0]
+        cfg.dataset.blend = data_blend
     else:
         cfg.dataset.blend = None
 
@@ -97,6 +113,7 @@ def build_config(args: argparse.Namespace) -> ConfigContainer:
 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("--data-manifest", default=None, help="Manifest produced by convert_pretrain_parquet_to_megatron.py.")
     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)
@@ -136,6 +153,8 @@ def main() -> None:
         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'}")
+        print(f"dataset_manifest={args.data_manifest or 'none'}")
+        print(f"dataset_blend_size={len(cfg.dataset.blend or [])}")
         return
     pretrain(config=cfg, forward_step_func=forward_step)