Add Megatron data manifest and g0050 setup

2026-07-02 20:50:24 +08:00
parent 816eccb5b5
commit 5609b1f8e4
7 changed files with 608 additions and 19 deletions
--- a/dataset/pretrain/scripts/convert_pretrain_parquet_to_megatron.py
+++ b/dataset/pretrain/scripts/convert_pretrain_parquet_to_megatron.py
@@ -0,0 +1,340 @@
+#!/usr/bin/env python3
+from __future__ import annotations
+
+import argparse
+import json
+import os
+import re
+import sys
+import time
+from concurrent.futures import ProcessPoolExecutor, as_completed
+from multiprocessing import Pool
+from pathlib import Path
+from types import SimpleNamespace
+from typing import Iterable
+
+import pyarrow.parquet as pq
+
+_TOKENIZER = None
+_APPEND_EOD = True
+_MAX_SEQ_LEN = 65536
+
+
+def add_megatron_to_path(megatron_dir: str) -> None:
+    path = str(Path(megatron_dir).resolve())
+    if path not in sys.path:
+        sys.path.insert(0, path)
+
+
+def make_tokenizer_args(args: argparse.Namespace) -> SimpleNamespace:
+    return SimpleNamespace(
+        rank=0,
+        make_vocab_size_divisible_by=128,
+        tensor_model_parallel_size=1,
+        pad_vocab_size=True,
+        padded_vocab_size=None,
+        vocab_size=args.vocab_size,
+        vocab_file=args.vocab_file,
+        merge_file=args.merge_file,
+        vocab_extra_ids=0,
+        tokenizer_type=args.tokenizer_type,
+        tokenizer_model=args.tokenizer_model,
+        metadata_path=args.tokenizer_metadata,
+        special_tokens=args.tokenizer_special_tokens,
+        tokenizer_sentencepiece_legacy=args.tokenizer_sentencepiece_legacy,
+        tokenizer_hf_no_use_fast=args.tokenizer_hf_no_use_fast,
+        tokenizer_hf_no_include_special_tokens=args.tokenizer_hf_no_include_special_tokens,
+        trust_remote_code=args.trust_remote_code,
+        tiktoken_pattern=args.tiktoken_pattern,
+        tiktoken_num_special_tokens=args.tiktoken_num_special_tokens,
+        null_tokenizer_eod_id=args.null_tokenizer_eod_id,
+        null_tokenizer_pad_id=args.null_tokenizer_pad_id,
+        tokenizer_prompt_format=None,
+        chat_template=None,
+        image_tag_type=None,
+        force_system_message=False,
+        sft_tokenizer_prompt_format=None,
+    )
+
+
+def build_tokenizer(args: argparse.Namespace):
+    add_megatron_to_path(args.megatron_dir)
+    from megatron.core.tokenizers.utils.build_tokenizer import build_tokenizer
+
+    return build_tokenizer(make_tokenizer_args(args))
+
+
+def init_worker(args: argparse.Namespace) -> None:
+    global _TOKENIZER, _APPEND_EOD, _MAX_SEQ_LEN
+    _APPEND_EOD = args.append_eod
+    _MAX_SEQ_LEN = args.max_seq_len
+    _TOKENIZER = build_tokenizer(args)
+    if _APPEND_EOD and _TOKENIZER.eod is None:
+        raise ValueError("Tokenizer has no EOD/EOS token, but --append-eod is enabled.")
+
+
+def encode_text(text: object):
+    if not isinstance(text, str):
+        return None
+    text = text.strip()
+    if not text:
+        return None
+    token_ids = _TOKENIZER.tokenize(text)
+    if not token_ids:
+        return None
+
+    if _MAX_SEQ_LEN and _MAX_SEQ_LEN > 0:
+        content_chunk_size = _MAX_SEQ_LEN - 1 if _APPEND_EOD else _MAX_SEQ_LEN
+        chunks = []
+        for start in range(0, len(token_ids), content_chunk_size):
+            chunk = token_ids[start : start + content_chunk_size]
+            if _APPEND_EOD:
+                chunk = [*chunk, _TOKENIZER.eod]
+            chunks.append((chunk, [len(chunk)]))
+        return chunks
+
+    if _APPEND_EOD:
+        token_ids.append(_TOKENIZER.eod)
+    return [(token_ids, [len(token_ids)])]
+
+
+def output_paths(output_prefix: Path, text_key: str) -> tuple[Path, Path]:
+    return (
+        Path(str(output_prefix) + f"_{text_key}_document.bin"),
+        Path(str(output_prefix) + f"_{text_key}_document.idx"),
+    )
+
+
+def resolve_inputs(inputs: Iterable[str], max_files: int | None) -> list[Path]:
+    paths: list[Path] = []
+    for item in inputs:
+        path = Path(item)
+        if path.is_dir():
+            paths.extend(path.glob("*.parquet"))
+            paths.extend(path.glob("*.zstd.parquet"))
+            continue
+        matches = sorted(Path().glob(item)) if not path.is_absolute() else sorted(path.parent.glob(path.name))
+        if matches:
+            paths.extend(matches)
+            continue
+        if path.exists():
+            paths.append(path)
+            continue
+        raise FileNotFoundError(item)
+    unique = sorted({p.resolve() for p in paths}, key=str)
+    if max_files is not None:
+        unique = unique[:max_files]
+    return unique
+
+
+def safe_prefix_component(value: str) -> str:
+    value = re.sub(r"[^A-Za-z0-9_.-]+", "_", value.strip())
+    value = value.strip("._-")
+    return value or "source"
+
+
+def convert_one_parquet(task: tuple[str, argparse.Namespace]) -> dict[str, object]:
+    parquet_path = Path(task[0])
+    args = task[1]
+    add_megatron_to_path(args.megatron_dir)
+    from megatron.core.datasets import indexed_dataset
+
+    stem = parquet_path.name.replace(".zstd.parquet", "").replace(".parquet", "")
+    source_name = safe_prefix_component(parquet_path.parent.name)
+    output_prefix = Path(args.output_dir) / f"{args.output_prefix_prefix}_{source_name}_{stem}"
+    bin_file, idx_file = output_paths(output_prefix, args.text_key)
+    if not args.overwrite and bin_file.exists() and idx_file.exists():
+        return {
+            "status": "skip",
+            "source": str(parquet_path),
+            "output_prefix": str(output_prefix),
+            "reason": "existing bin/idx",
+        }
+
+    bin_file.unlink(missing_ok=True)
+    idx_file.unlink(missing_ok=True)
+    output_prefix.parent.mkdir(parents=True, exist_ok=True)
+
+    tokenizer = build_tokenizer(args)
+    dtype = indexed_dataset.DType.optimal_dtype(tokenizer.vocab_size)
+    builder = indexed_dataset.IndexedDatasetBuilder(str(bin_file), dtype=dtype)
+
+    start = time.time()
+    rows = source_docs = output_docs = tokens = split_source_docs = split_extra_docs = 0
+    stop = False
+
+    def consume(encoded) -> None:
+        nonlocal source_docs, output_docs, tokens, split_source_docs, split_extra_docs, stop
+        if encoded is None or stop:
+            return
+        source_docs += 1
+        if len(encoded) > 1:
+            split_source_docs += 1
+            split_extra_docs += len(encoded) - 1
+        for token_ids, lengths in encoded:
+            if args.max_seq_len and len(token_ids) > args.max_seq_len:
+                raise ValueError(
+                    f"internal split error: chunk has {len(token_ids)} tokens, "
+                    f"max_seq_len={args.max_seq_len}"
+                )
+            builder.add_document(token_ids, lengths)
+            output_docs += 1
+            tokens += len(token_ids)
+        if args.log_interval and source_docs % args.log_interval == 0:
+            elapsed = max(time.time() - start, 1e-6)
+            print(
+                f"[{parquet_path.name}] source_docs={source_docs} output_docs={output_docs} "
+                f"tokens={tokens} source_docs/s={source_docs / elapsed:.2f}",
+                flush=True,
+            )
+        if args.max_docs and source_docs >= args.max_docs:
+            stop = True
+
+    pf = pq.ParquetFile(parquet_path)
+    if args.text_key not in pf.schema_arrow.names:
+        raise ValueError(f"{parquet_path} has no column {args.text_key!r}")
+
+    if args.workers_per_file == 1:
+        init_worker(args)
+        for batch in pf.iter_batches(columns=[args.text_key], batch_size=args.batch_size, use_threads=True):
+            texts = batch.column(0).to_pylist()
+            rows += len(texts)
+            for text in texts:
+                consume(encode_text(text))
+                if stop:
+                    break
+            if stop:
+                break
+    else:
+        with Pool(processes=args.workers_per_file, initializer=init_worker, initargs=(args,)) as pool:
+            for batch in pf.iter_batches(columns=[args.text_key], batch_size=args.batch_size, use_threads=True):
+                texts = batch.column(0).to_pylist()
+                rows += len(texts)
+                for encoded in pool.imap(encode_text, texts, chunksize=args.chunksize):
+                    consume(encoded)
+                    if stop:
+                        break
+                if stop:
+                    break
+
+    builder.finalize(str(idx_file))
+    elapsed = max(time.time() - start, 1e-6)
+    return {
+        "status": "ok",
+        "source": str(parquet_path),
+        "output_prefix": str(output_prefix),
+        "bin_file": str(bin_file),
+        "idx_file": str(idx_file),
+        "dtype": dtype.__name__,
+        "rows": rows,
+        "source_docs": source_docs,
+        "docs": output_docs,
+        "output_docs": output_docs,
+        "split_source_docs": split_source_docs,
+        "split_extra_docs": split_extra_docs,
+        "max_seq_len": args.max_seq_len,
+        "tokens": tokens,
+        "elapsed_sec": elapsed,
+        "source_docs_per_sec": source_docs / elapsed,
+        "docs_per_sec": output_docs / elapsed,
+    }
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(description="Convert pretrain parquet directly to Megatron .bin/.idx.")
+    parser.add_argument("--input", action="append", required=True, help="Parquet dir/file/glob. Repeatable.")
+    parser.add_argument("--output-dir", required=True)
+    parser.add_argument("--manifest", default=None)
+    parser.add_argument("--megatron-dir", default="/opt/Megatron-Bridge/3rdparty/Megatron-LM")
+    parser.add_argument("--tokenizer-type", default="HuggingFaceTokenizer")
+    parser.add_argument("--tokenizer-model", required=True)
+    parser.add_argument("--tokenizer-metadata", default=None)
+    parser.add_argument("--tokenizer-special-tokens", nargs="*", default=None)
+    parser.add_argument("--tokenizer-sentencepiece-legacy", action="store_true")
+    parser.add_argument("--tokenizer-hf-no-use-fast", action="store_true")
+    parser.add_argument("--tokenizer-hf-no-include-special-tokens", action="store_true")
+    parser.add_argument("--trust-remote-code", action="store_true")
+    parser.add_argument("--vocab-file", default=None)
+    parser.add_argument("--merge-file", default=None)
+    parser.add_argument("--vocab-size", type=int, default=None)
+    parser.add_argument("--tiktoken-pattern", default=None)
+    parser.add_argument("--tiktoken-num-special-tokens", type=int, default=1000)
+    parser.add_argument("--null-tokenizer-eod-id", type=int, default=None)
+    parser.add_argument("--null-tokenizer-pad-id", type=int, default=-1)
+    parser.add_argument("--text-key", default="text")
+    parser.add_argument("--output-prefix-prefix", default="laoyao_2b_moe_8192")
+    parser.add_argument("--parallel-files", type=int, default=1)
+    parser.add_argument("--workers-per-file", type=int, default=max((os.cpu_count() or 8) // 4, 1))
+    parser.add_argument("--batch-size", type=int, default=8192)
+    parser.add_argument("--chunksize", type=int, default=128)
+    parser.add_argument("--max-files", type=int, default=None)
+    parser.add_argument("--max-docs", type=int, default=0, help="Per-file doc cap; 0 means no cap.")
+    parser.add_argument("--max-seq-len", type=int, default=65536, help="Split documents into chunks no longer than this; 0 disables splitting.")
+    parser.add_argument("--log-interval", type=int, default=10000)
+    parser.add_argument("--append-eod", action=argparse.BooleanOptionalAction, default=True)
+    parser.add_argument("--overwrite", action="store_true")
+    return parser.parse_args()
+
+
+def main() -> None:
+    args = parse_args()
+    if args.parallel_files < 1 or args.workers_per_file < 1:
+        raise ValueError("parallel and worker counts must be positive")
+    if args.max_seq_len < 0:
+        raise ValueError("--max-seq-len must be >= 0")
+    if args.append_eod and args.max_seq_len == 1:
+        raise ValueError("--max-seq-len must be > 1 when --append-eod is enabled")
+    files = resolve_inputs(args.input, args.max_files)
+    if not files:
+        raise FileNotFoundError("no parquet files resolved")
+
+    Path(args.output_dir).mkdir(parents=True, exist_ok=True)
+    manifest_path = Path(args.manifest) if args.manifest else Path(args.output_dir) / "manifest.json"
+    print(
+        json.dumps(
+            {
+                "files": len(files),
+                "parallel_files": args.parallel_files,
+                "workers_per_file": args.workers_per_file,
+                "total_tokenizer_workers": args.parallel_files * args.workers_per_file,
+                "output_dir": args.output_dir,
+                "max_seq_len": args.max_seq_len,
+            },
+            ensure_ascii=False,
+        ),
+        flush=True,
+    )
+
+    tasks = [(str(path), args) for path in files]
+    results: list[dict[str, object]] = []
+    if args.parallel_files == 1:
+        for task in tasks:
+            result = convert_one_parquet(task)
+            results.append(result)
+            print(json.dumps(result, ensure_ascii=False), flush=True)
+    else:
+        with ProcessPoolExecutor(max_workers=args.parallel_files) as executor:
+            futures = [executor.submit(convert_one_parquet, task) for task in tasks]
+            for future in as_completed(futures):
+                result = future.result()
+                results.append(result)
+                print(json.dumps(result, ensure_ascii=False), flush=True)
+
+    manifest = {
+        "args": vars(args),
+        "files": [str(path) for path in files],
+        "results": sorted(results, key=lambda item: str(item.get("source", ""))),
+        "total_source_docs": sum(int(r.get("source_docs", 0)) for r in results),
+        "total_docs": sum(int(r.get("docs", 0)) for r in results),
+        "total_output_docs": sum(int(r.get("output_docs", 0)) for r in results),
+        "total_split_source_docs": sum(int(r.get("split_source_docs", 0)) for r in results),
+        "total_split_extra_docs": sum(int(r.get("split_extra_docs", 0)) for r in results),
+        "total_tokens": sum(int(r.get("tokens", 0)) for r in results),
+        "ok_prefixes": [str(r["output_prefix"]) for r in results if r.get("status") in {"ok", "skip"}],
+    }
+    manifest_path.write_text(json.dumps(manifest, ensure_ascii=False, indent=2), encoding="utf-8")
+    print(f"manifest={manifest_path}", flush=True)
+
+
+if __name__ == "__main__":
+    main()
--- a/scripts/README.md
+++ b/scripts/README.md
@@ -2,4 +2,33 @@

 - `sync_pretrain_data_into_repo.sh`: 200B 数据构建完成后，把数据目录同步到 `dataset/pretrain/data/`，默认优先 hardlink。
 - `wait_and_sync_pretrain_data.sh`: 后台等待当前 200B 构建进程结束，然后自动同步数据。
+- `preprocess_megatron_bridge_pretrain.sh`: 旧的 Megatron indexed dataset 预处理入口，保留用于对照。
+- `preprocess_megatron_bridge_pretrain_direct.sh`: 直接从 parquet 生成 Megatron indexed dataset，不落中间 JSONL。
+- `train_megatron_bridge_2b_moe.sh`: 当前主训练入口，使用 NeMo 26.06 镜像中的 Megatron-Bridge。
 - `train_nemo_megatron_2b_moe.sh`: NeMo/Megatron 训练入口占位，包含 image、mount、路径检查。
+- `g0050_download_and_setup_from_modelscope.sh`: 在 g0050 上一键准备训练环境并从 ModelScope 下载未 tokenize parquet 数据。
+
+## g0050 下载与部署
+
+在 Mac 侧通过 B300 跳转到 g0050：
+
+```bash
+ssh B300 'ssh ubuntu@g0050 "cd /ssd/workspace/yi/laoyao_2b_moe && MODELSCOPE_API_TOKEN=ms-... bash scripts/g0050_download_and_setup_from_modelscope.sh"'
+```
+
+默认行为：
+
+- repo 路径：`/ssd/workspace/yi/laoyao_2b_moe`
+- 数据路径：`/ssd/workspace/yi/laoyao_2b_moe_pretraining_dataset`
+- ModelScope dataset：`eigentom/laoyao_2b_moe_pretrain_parquet_20260702`
+- 训练镜像：`nvcr.io/nvidia/nemo:26.06`
+- 下载镜像：ModelScope CUDA 13.0 / Swift 4.3.1 官方镜像
+- 代理：默认使用 B300/g0050 侧的 `http://100.72.0.101:8888`
+
+私有 Gitea clone 时不要把 token 写进脚本，运行时通过环境变量传入：
+
+```bash
+GIT_REPO_URL=https://yi_lu:<token>@git.deeepseek.net/yi_lu/laoyao_2b_moe.git \
+MODELSCOPE_API_TOKEN=ms-... \
+bash scripts/g0050_download_and_setup_from_modelscope.sh
+```
--- a/scripts/g0050_download_and_setup_from_modelscope.sh
+++ b/scripts/g0050_download_and_setup_from_modelscope.sh
@@ -0,0 +1,92 @@
+#!/usr/bin/env bash
+set -euo pipefail
+
+# Prepare g0050 for Laoyao 2B MoE training and download the pre-tokenization
+# parquet dataset from ModelScope. Keep credentials out of this file:
+#
+#   MODELSCOPE_API_TOKEN=ms-... bash scripts/g0050_download_and_setup_from_modelscope.sh
+#
+# Optional private git URL:
+#
+#   GIT_REPO_URL=https://user:token@git.deeepseek.net/yi_lu/laoyao_2b_moe.git ...
+
+REPO_ROOT="${REPO_ROOT:-/ssd/workspace/yi/laoyao_2b_moe}"
+DATA_ROOT="${DATA_ROOT:-/ssd/workspace/yi/laoyao_2b_moe_pretraining_dataset}"
+DATASET_ID="${DATASET_ID:-eigentom/laoyao_2b_moe_pretrain_parquet_20260702}"
+GIT_REPO_URL="${GIT_REPO_URL:-https://git.deeepseek.net/yi_lu/laoyao_2b_moe.git}"
+GIT_BRANCH="${GIT_BRANCH:-main}"
+
+NEMO_IMAGE="${NEMO_IMAGE:-nvcr.io/nvidia/nemo:26.06}"
+MODELSCOPE_IMAGE="${MODELSCOPE_IMAGE:-modelscope-registry.cn-hangzhou.cr.aliyuncs.com/modelscope-repo/modelscope:ubuntu22.04-cuda13.0.3-py312-torch2.11.0-vllm0.23.0-modelscope1.37.1-swift4.3.1}"
+MODELSCOPE_ENDPOINT="${MODELSCOPE_ENDPOINT:-https://www.modelscope.cn}"
+MAX_WORKERS="${MAX_WORKERS:-8}"
+
+HTTP_PROXY_DEFAULT="${HTTP_PROXY_DEFAULT:-http://100.72.0.101:8888}"
+HTTPS_PROXY_DEFAULT="${HTTPS_PROXY_DEFAULT:-http://100.72.0.101:8888}"
+export http_proxy="${http_proxy:-$HTTP_PROXY_DEFAULT}"
+export https_proxy="${https_proxy:-$HTTPS_PROXY_DEFAULT}"
+export HTTP_PROXY="${HTTP_PROXY:-$http_proxy}"
+export HTTPS_PROXY="${HTTPS_PROXY:-$https_proxy}"
+
+if [[ -z "${MODELSCOPE_API_TOKEN:-}" ]]; then
+  echo "ERROR: set MODELSCOPE_API_TOKEN before running this script." >&2
+  exit 2
+fi
+
+mkdir -p "$(dirname "$REPO_ROOT")" "$DATA_ROOT"
+
+echo "[setup] repo_root=$REPO_ROOT"
+if [[ -d "$REPO_ROOT/.git" ]]; then
+  git -C "$REPO_ROOT" fetch origin "$GIT_BRANCH"
+  git -C "$REPO_ROOT" checkout "$GIT_BRANCH"
+  git -C "$REPO_ROOT" pull --ff-only origin "$GIT_BRANCH"
+elif [[ -d "$REPO_ROOT" ]] && [[ -z "$(find "$REPO_ROOT" -mindepth 1 -maxdepth 1 -print -quit)" ]]; then
+  rmdir "$REPO_ROOT"
+  git clone --branch "$GIT_BRANCH" "$GIT_REPO_URL" "$REPO_ROOT"
+else
+  git clone --branch "$GIT_BRANCH" "$GIT_REPO_URL" "$REPO_ROOT"
+fi
+
+echo "[setup] checking docker images"
+if ! docker image inspect "$NEMO_IMAGE" >/dev/null 2>&1; then
+  docker pull "$NEMO_IMAGE"
+fi
+if ! docker image inspect "$MODELSCOPE_IMAGE" >/dev/null 2>&1; then
+  docker pull "$MODELSCOPE_IMAGE"
+fi
+
+echo "[download] dataset=$DATASET_ID"
+echo "[download] data_root=$DATA_ROOT"
+docker run --rm \
+  --network=host \
+  -e http_proxy="$http_proxy" \
+  -e https_proxy="$https_proxy" \
+  -e HTTP_PROXY="$HTTP_PROXY" \
+  -e HTTPS_PROXY="$HTTPS_PROXY" \
+  -e MODELSCOPE_API_TOKEN="$MODELSCOPE_API_TOKEN" \
+  -v "$DATA_ROOT:$DATA_ROOT" \
+  "$MODELSCOPE_IMAGE" \
+  modelscope download \
+    --dataset "$DATASET_ID" \
+    --token "$MODELSCOPE_API_TOKEN" \
+    --endpoint "$MODELSCOPE_ENDPOINT" \
+    --local_dir "$DATA_ROOT" \
+    --max-workers "$MAX_WORKERS"
+
+echo "[verify] downloaded tree"
+find "$DATA_ROOT" -maxdepth 3 -type f | sed -n '1,20p'
+du -sh "$DATA_ROOT" || true
+
+cat <<EOF
+
+Done.
+
+Expected source parquet roots:
+  $DATA_ROOT/train/pretrain_rebalanced_web40_edu20_chinese10_science10_logic10_math5_code5_200b_v1_20260701
+  $DATA_ROOT/train/logic_topup_proof_pile_17b_v1_20260701
+
+Next steps:
+  cd $REPO_ROOT
+  bash scripts/preprocess_megatron_bridge_pretrain_direct.sh
+  bash scripts/train_megatron_bridge_2b_moe.sh
+EOF
--- a/scripts/preprocess_megatron_bridge_pretrain.sh
+++ b/scripts/preprocess_megatron_bridge_pretrain.sh
@@ -3,36 +3,59 @@ 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}"
+SOURCE_DATA_DIRS="${SOURCE_DATA_DIRS:-/mnt/beegfs/yi/laoyao_2b_moe_pretraining_dataset/train/pretrain_rebalanced_web40_edu20_chinese10_science10_logic10_math5_code5_200b_v1_20260701:/mnt/beegfs/yi/laoyao_2b_moe_pretraining_dataset/train/logic_topup_proof_pile_17b_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}"
+MIN_FREE_GB="${MIN_FREE_GB:-1000}"
+KEEP_JSONL="${KEEP_JSONL:-0}"

 mkdir -p "$WORK_DIR"

+available_gb="$(df -BG "$WORK_DIR" | awk 'NR==2 {gsub("G", "", $4); print $4}')"
+if [[ "$MAX_DOCS" == "0" && "$available_gb" -lt "$MIN_FREE_GB" ]]; then
+  cat >&2 <<EOF
+Refusing full preprocessing: only ${available_gb}GB free at $WORK_DIR.
+Full 200B-token Megatron indexed output is expected to require hundreds of GB,
+and this script also stages a JSONL export. Set MIN_FREE_GB lower only if you
+have confirmed an output filesystem with enough capacity, or run a bounded
+probe with MAX_DOCS.
+EOF
+  exit 2
+fi
+
+input_args=()
+IFS=':' read -r -a source_dirs <<< "$SOURCE_DATA_DIRS"
+for source_dir in "${source_dirs[@]}"; do
+  input_args+=(--input "${source_dir}/*.parquet")
+done
+
 docker run --rm --ipc=host --network=host \
  --ulimit memlock=-1 --ulimit stack=67108864 \
  -v /mnt/beegfs:/mnt/beegfs \
  -w "$REPO_ROOT" \
  "$IMAGE" \
-  bash -lc "
+  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'
+      "$@" \
+      --output '"$JSONL"' \
+      --max-docs '"$MAX_DOCS"'
    python3 /opt/Megatron-Bridge/3rdparty/Megatron-LM/tools/preprocess_data.py \
-      --input '$JSONL' \
+      --input '"$JSONL"' \
      --json-keys text \
      --tokenizer-type HuggingFaceTokenizer \
-      --tokenizer-model '$TOKENIZER_MODEL' \
+      --tokenizer-model '"$TOKENIZER_MODEL"' \
      --append-eod \
-      --output-prefix '$OUTPUT_PREFIX' \
-      --workers '$WORKERS'
-    ls -lh '${OUTPUT_PREFIX}'*
-  "
+      --output-prefix '"$OUTPUT_PREFIX"' \
+      --workers '"$WORKERS"'
+    if [[ "'"$KEEP_JSONL"'" == "0" ]]; then
+      rm -f '"$JSONL"'
+    fi
+    ls -lh '"${OUTPUT_PREFIX}"'*
+  ' bash "${input_args[@]}"

 echo "Megatron indexed dataset prefix: $OUTPUT_PREFIX"
--- a/scripts/preprocess_megatron_bridge_pretrain_direct.sh
+++ b/scripts/preprocess_megatron_bridge_pretrain_direct.sh
@@ -0,0 +1,72 @@
+#!/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_DIRS="${SOURCE_DATA_DIRS:-/mnt/beegfs/yi/laoyao_2b_moe_pretraining_dataset/train/pretrain_rebalanced_web40_edu20_chinese10_science10_logic10_math5_code5_200b_v1_20260701:/mnt/beegfs/yi/laoyao_2b_moe_pretraining_dataset/train/logic_topup_proof_pile_17b_v1_20260701}"
+WORK_DIR="${WORK_DIR:-/mnt/beegfs/yi/laoyao_2b_moe_pretraining_dataset/megatron_bridge/pretrain_8192_direct_smoke_v1}"
+TOKENIZER_MODEL="${TOKENIZER_MODEL:-$REPO_ROOT/tokenizer/glm5.2}"
+OUTPUT_PREFIX_PREFIX="${OUTPUT_PREFIX_PREFIX:-laoyao_2b_moe_8192_direct}"
+TEXT_KEY="${TEXT_KEY:-text}"
+PARALLEL_FILES="${PARALLEL_FILES:-1}"
+WORKERS_PER_FILE="${WORKERS_PER_FILE:-8}"
+BATCH_SIZE="${BATCH_SIZE:-8192}"
+CHUNKSIZE="${CHUNKSIZE:-128}"
+MAX_FILES="${MAX_FILES:-1}"
+MAX_DOCS="${MAX_DOCS:-100000}"
+MAX_SEQ_LEN="${MAX_SEQ_LEN:-65536}"
+MIN_FREE_GB="${MIN_FREE_GB:-20}"
+OVERWRITE="${OVERWRITE:-1}"
+
+mkdir -p "$WORK_DIR"
+
+available_gb="$(df -BG "$WORK_DIR" | awk 'NR==2 {gsub("G", "", $4); print $4}')"
+if [[ "$MAX_FILES" == "0" && "$MAX_DOCS" == "0" && "$available_gb" -lt "$MIN_FREE_GB" ]]; then
+  cat >&2 <<EOF
+Refusing full direct preprocessing: only ${available_gb}GB free at $WORK_DIR.
+Direct Megatron indexed output for 200B GLM5.2 tokens is expected to require
+roughly 800GB plus index/cache overhead. Increase free space or lower MAX_FILES/MAX_DOCS.
+EOF
+  exit 2
+fi
+
+input_args=()
+IFS=':' read -r -a source_dirs <<< "$SOURCE_DATA_DIRS"
+for source_dir in "${source_dirs[@]}"; do
+  input_args+=(--input "${source_dir}/*.parquet")
+done
+
+max_file_args=()
+if [[ "$MAX_FILES" != "0" ]]; then
+  max_file_args+=(--max-files "$MAX_FILES")
+fi
+
+overwrite_args=()
+if [[ "$OVERWRITE" == "1" ]]; then
+  overwrite_args+=(--overwrite)
+fi
+
+docker run --rm --ipc=host --network=host \
+  --ulimit memlock=-1 --ulimit stack=67108864 \
+  -v /mnt/beegfs:/mnt/beegfs \
+  -w "$REPO_ROOT" \
+  "$IMAGE" \
+  python3 dataset/pretrain/scripts/convert_pretrain_parquet_to_megatron.py \
+    "${input_args[@]}" \
+    --output-dir "$WORK_DIR" \
+    --manifest "$WORK_DIR/manifest.json" \
+    --megatron-dir /opt/Megatron-Bridge/3rdparty/Megatron-LM \
+    --tokenizer-type HuggingFaceTokenizer \
+    --tokenizer-model "$TOKENIZER_MODEL" \
+    --text-key "$TEXT_KEY" \
+    --output-prefix-prefix "$OUTPUT_PREFIX_PREFIX" \
+    --parallel-files "$PARALLEL_FILES" \
+    --workers-per-file "$WORKERS_PER_FILE" \
+    --batch-size "$BATCH_SIZE" \
+    --chunksize "$CHUNKSIZE" \
+    --max-docs "$MAX_DOCS" \
+    --max-seq-len "$MAX_SEQ_LEN" \
+    "${max_file_args[@]}" \
+    "${overwrite_args[@]}"
+
+echo "Direct Megatron indexed dataset manifest: $WORK_DIR/manifest.json"
--- a/scripts/train_megatron_bridge_2b_moe.sh
+++ b/scripts/train_megatron_bridge_2b_moe.sh
@@ -4,7 +4,8 @@ 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}"
+DATA_PREFIX="${DATA_PREFIX:-}"
+DATA_MANIFEST="${DATA_MANIFEST:-/mnt/beegfs/yi/laoyao_2b_moe_pretraining_dataset/megatron_bridge/pretrain_65536_direct_v1/manifest.json}"
 TRAIN_ITERS="${TRAIN_ITERS:-10}"
 SEQ_LENGTH="${SEQ_LENGTH:-8192}"
 MICRO_BATCH_SIZE="${MICRO_BATCH_SIZE:-1}"
@@ -15,9 +16,15 @@ 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
+if [[ -n "$DATA_MANIFEST" ]]; then
+  if [[ ! -f "$DATA_MANIFEST" ]]; then
+    echo "missing Megatron indexed dataset manifest: $DATA_MANIFEST" >&2
+    echo "run scripts/preprocess_megatron_bridge_pretrain_direct.sh first, or set DATA_MANIFEST" >&2
+    exit 1
+  fi
+elif [[ "$DRY_RUN" != "1" && ( -z "$DATA_PREFIX" || ! -f "${DATA_PREFIX}.idx" ) ]]; then
+  echo "missing Megatron indexed data prefix: ${DATA_PREFIX:-<empty>}" >&2
+  echo "set DATA_MANIFEST or DATA_PREFIX" >&2
  exit 1
 fi

@@ -26,6 +33,13 @@ if [[ "$DRY_RUN" == "1" ]]; then
  DRY_RUN_ARG="--dry-run"
 fi

+DATA_ARGS=()
+if [[ -n "$DATA_MANIFEST" ]]; then
+  DATA_ARGS=(--data-manifest "$DATA_MANIFEST")
+else
+  DATA_ARGS=(--data-prefix "$DATA_PREFIX")
+fi
+
 docker run --rm --gpus all --ipc=host --network=host \
  --ulimit memlock=-1 --ulimit stack=67108864 \
  -v /mnt/beegfs:/mnt/beegfs \
@@ -35,7 +49,7 @@ docker run --rm --gpus all --ipc=host --network=host \
    set -euo pipefail
    torchrun --nproc_per_node='$NPROC_PER_NODE' \
      training/megatron_bridge/laoyao_2b_moe_pretrain.py \
-      --data-prefix '$DATA_PREFIX' \
+      ${DATA_ARGS[*]} \
      --seq-length '$SEQ_LENGTH' \
      --train-iters '$TRAIN_ITERS' \
      --micro-batch-size '$MICRO_BATCH_SIZE' \
--- a/training/megatron_bridge/laoyao_2b_moe_pretrain.py
+++ b/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)