From 5609b1f8e4d3276df344d4de78a51d82f58f40c8 Mon Sep 17 00:00:00 2001 From: yi_lu Date: Thu, 2 Jul 2026 20:50:24 +0800 Subject: [PATCH] Add Megatron data manifest and g0050 setup --- .../convert_pretrain_parquet_to_megatron.py | 340 ++++++++++++++++++ scripts/README.md | 29 ++ ...0050_download_and_setup_from_modelscope.sh | 92 +++++ .../preprocess_megatron_bridge_pretrain.sh | 45 ++- ...process_megatron_bridge_pretrain_direct.sh | 72 ++++ scripts/train_megatron_bridge_2b_moe.sh | 24 +- .../megatron_bridge/laoyao_2b_moe_pretrain.py | 25 +- 7 files changed, 608 insertions(+), 19 deletions(-) create mode 100755 dataset/pretrain/scripts/convert_pretrain_parquet_to_megatron.py create mode 100755 scripts/g0050_download_and_setup_from_modelscope.sh create mode 100755 scripts/preprocess_megatron_bridge_pretrain_direct.sh diff --git a/dataset/pretrain/scripts/convert_pretrain_parquet_to_megatron.py b/dataset/pretrain/scripts/convert_pretrain_parquet_to_megatron.py new file mode 100755 index 0000000..136a5a4 --- /dev/null +++ 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() diff --git a/scripts/README.md b/scripts/README.md index d761d03..44518ab 100644 --- 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:@git.deeepseek.net/yi_lu/laoyao_2b_moe.git \ +MODELSCOPE_API_TOKEN=ms-... \ +bash scripts/g0050_download_and_setup_from_modelscope.sh +``` diff --git a/scripts/g0050_download_and_setup_from_modelscope.sh b/scripts/g0050_download_and_setup_from_modelscope.sh new file mode 100755 index 0000000..acae625 --- /dev/null +++ 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 <&2 <&2 <&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:-}" >&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' \ diff --git a/training/megatron_bridge/laoyao_2b_moe_pretrain.py b/training/megatron_bridge/laoyao_2b_moe_pretrain.py index 3c3c9ab..2cfc374 100755 --- 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)