From f0f72767f5b46c9a9cfb1ed1c1294124e995f1b8 Mon Sep 17 00:00:00 2001 From: yi_lu Date: Wed, 1 Jul 2026 22:42:18 +0800 Subject: [PATCH] Add proof-pile logic inputs to pretrain builder --- README.md | 10 +---- dataset/pretrain/README.md | 23 ++++++++--- .../rebalanced_pretrain_recipe_20260630.json | 2 +- .../rebalanced_182b_20260701_summary.json | 26 ++++++++++++ .../build_rebalanced_pretrain_dataset.py | 41 ++++++++++++++++--- .../start_rebalanced_pretrain_dataset.sh | 4 +- .../start_rebalanced_source_download.sh | 4 +- ...wait_and_build_rebalanced_pretrain_200b.sh | 29 +++++++++++-- scripts/sync_pretrain_data_into_repo.sh | 31 ++++---------- scripts/train_nemo_megatron_2b_moe.sh | 4 +- .../nemo_megatron/pretrain_2b_moe_200b.yaml | 2 +- 11 files changed, 122 insertions(+), 54 deletions(-) create mode 100644 dataset/pretrain/manifests/rebalanced_182b_20260701_summary.json diff --git a/README.md b/README.md index f1e82c4..8bf9c67 100644 --- a/README.md +++ b/README.md @@ -44,16 +44,10 @@ repo 路径： /mnt/beegfs/yi/laoyao_2b_moe ``` -200B 数据当前构建源路径： +预训练数据实际读取路径： ```bash /mnt/beegfs/yi/laoyao_2b_moe_pretraining_dataset/train/pretrain_rebalanced_web40_edu20_chinese10_science10_logic10_math5_code5_200b_v1_20260701 ``` -构建完成后同步到 repo 内： - -```bash -bash scripts/sync_pretrain_data_into_repo.sh -``` - -同步后的数据目录 `dataset/pretrain/data/` 被 `.gitignore` 忽略，不提交到 GitHub。 +数据不再复制到 repo 内；训练直接读取上面的源输出目录。`dataset/pretrain/data/` 仍被 `.gitignore` 忽略，仅用于小规模临时样本。 diff --git a/dataset/pretrain/README.md b/dataset/pretrain/README.md index 839ad9f..a0cc512 100644 --- a/dataset/pretrain/README.md +++ b/dataset/pretrain/README.md @@ -10,22 +10,33 @@ | english_edu | 20% | FineWeb-Edu / Ultra-FineWeb 高质量教育文本 | | chinese_clean | 10% | Ultra-FineWeb zh，后续可接 CCI3-HQ/SkyPile | | science | 10% | medmcqa、proofwriter、scienceqa、sciq、qasc、openbookqa；不足由 english_edu 补齐 | -| logic | 10% | Jiayi scored education/text 数据中的逻辑类样本 | +| logic | 10% | Jiayi scored education/text 数据中的逻辑类样本 + Proof-Pile-2/OpenWebMath/arXiv/AlgebraicStack 的高推理密度文本 | | math | 5% | Jiayi scored math/knowledge 数据 | | code | 5% | Jiayi cleaned multilingual code 数据 | ## 数据本体 -真实数据应放在： +真实数据实际保留在： ```bash -dataset/pretrain/data/pretrain_rebalanced_web40_edu20_chinese10_science10_logic10_math5_code5_200b_v1_20260701 +/mnt/beegfs/yi/laoyao_2b_moe_pretraining_dataset/train/pretrain_rebalanced_web40_edu20_chinese10_science10_logic10_math5_code5_200b_v1_20260701 ``` -该目录被 `.gitignore` 忽略。当前 g0033 构建完成后执行： +不要再把 405G 数据复制到 repo 内。BeegFS hardlink 不被允许，rsync 会造成空间翻倍并可能失败。训练脚本应直接读取该源目录。 + +## Logic 补充源 + +`wait_and_build_rebalanced_pretrain_200b.sh` 默认将以下本地 Proof-Pile-2 子集作为显式 `--logic-input`： ```bash -bash scripts/sync_pretrain_data_into_repo.sh +/mnt/beegfs/cjy/cjy-training/laoyao_model/data/raw/hf_datasets/EleutherAI__proof-pile-2/open-web-math/train/*.jsonl.zst +/mnt/beegfs/cjy/cjy-training/laoyao_model/data/raw/hf_datasets/EleutherAI__proof-pile-2/arxiv/train/*.jsonl.zst +/mnt/beegfs/cjy/cjy-training/laoyao_model/data/raw/hf_datasets/EleutherAI__proof-pile-2/algebraic-stack/train/*.jsonl.zst ``` -同步脚本优先使用 hardlink，避免在同一个 BeegFS 上重复占用大规模空间。 +这些 `.jsonl.zst` 文件由 `build_rebalanced_pretrain_dataset.py` 通过系统 `zstd -dc` 流式读取，不需要先解压到 BeegFS。后续如果下载了 StackExchange 或 PhilPapers，可用冒号分隔的 `LOGIC_INPUT_EXTRA` 追加： + +```bash +LOGIC_INPUT_EXTRA='/path/to/stackexchange/**/*.jsonl.zst:/path/to/philpapers/**/*.jsonl.zst' \ +bash dataset/pretrain/scripts/wait_and_build_rebalanced_pretrain_200b.sh +``` diff --git a/dataset/pretrain/configs/rebalanced_pretrain_recipe_20260630.json b/dataset/pretrain/configs/rebalanced_pretrain_recipe_20260630.json index 63890f0..aac8136 100644 --- a/dataset/pretrain/configs/rebalanced_pretrain_recipe_20260630.json +++ b/dataset/pretrain/configs/rebalanced_pretrain_recipe_20260630.json @@ -98,6 +98,6 @@ "This recipe intentionally reduces code to 5% for stage-1 pretraining.", "The previous code40/text60 mixture is useful as cleaned input inventory, not as the target mixture.", "Science shortfall is explicitly reassigned to english_edu and recorded in report.json. Other missing quotas should remain visible.", - "Run scripts/data/start_rebalanced_pretrain_dataset.sh to materialize parquet shards and report.json." + "Run dataset/pretrain/scripts/start_rebalanced_pretrain_dataset.sh to materialize parquet shards and report.json." ] } diff --git a/dataset/pretrain/manifests/rebalanced_182b_20260701_summary.json b/dataset/pretrain/manifests/rebalanced_182b_20260701_summary.json new file mode 100644 index 0000000..bc955ec --- /dev/null +++ b/dataset/pretrain/manifests/rebalanced_182b_20260701_summary.json @@ -0,0 +1,26 @@ +{ + "name": "pretrain_rebalanced_web40_edu20_chinese10_science10_logic10_math5_code5_200b_v1_20260701", + "status": "materialized_shortfall", + "target_tokens": 200000000000, + "actual_tokens": 182667460903, + "actual_docs": 158109491, + "shards": 1582, + "source_output_path": "/mnt/beegfs/yi/laoyao_2b_moe_pretraining_dataset/train/pretrain_rebalanced_web40_edu20_chinese10_science10_logic10_math5_code5_200b_v1_20260701", + "repo_data_path_policy": "do_not_copy_large_dataset_into_git_repo; train directly from source_output_path", + "source_output_size_observed": "405G", + "category_tokens": { + "english_web": 80000000277, + "english_edu": 59954167427, + "chinese_clean": 20000007279, + "math": 10000000436, + "code": 10000003541, + "logic": 2667446796, + "science": 45835147 + }, + "notes": [ + "The 200B builder exhausted available sources at 182.667B tokens.", + "Hardlink sync into repo failed with Operation not permitted on BeegFS.", + "Rsync fallback failed with No space left on device after an incomplete 172G copy; that partial copy was removed.", + "Training configs should point directly to source_output_path." + ] +} diff --git a/dataset/pretrain/scripts/build_rebalanced_pretrain_dataset.py b/dataset/pretrain/scripts/build_rebalanced_pretrain_dataset.py index 89e5739..1418593 100644 --- a/dataset/pretrain/scripts/build_rebalanced_pretrain_dataset.py +++ b/dataset/pretrain/scripts/build_rebalanced_pretrain_dataset.py @@ -11,6 +11,7 @@ import os from pathlib import Path import random import re +import subprocess import sys from typing import Iterable, Mapping, Sequence @@ -233,7 +234,12 @@ def parse_args() -> argparse.Namespace: default=[], help="Science QA source, e.g. medmcqa/proofwriter/scienceqa/sciq/qasc/openbookqa.", ) - parser.add_argument("--logic-input", action="append", default=[], help="Optional explicit logic source path/dir/glob.") + parser.add_argument( + "--logic-input", + action="append", + default=[], + help="Optional explicit logic source path/dir/glob. Supports parquet, jsonl, and jsonl.zst.", + ) parser.add_argument("--math-input", action="append", default=[], help="Optional explicit math source path/dir/glob.") parser.add_argument("--code-input", action="append", default=[], help="Code source path/dir/glob.") parser.add_argument( @@ -393,10 +399,10 @@ def build_source_specs(args: argparse.Namespace) -> tuple[SourceSpec, ...]: ("science", args.science_input), ("english_edu", args.english_edu_input), ("chinese_clean", args.chinese_clean_input), - ("auto_text", args.jiayi_scored_input or [DEFAULT_JIAYI_SCORED_INPUT]), + ("auto_text", args.jiayi_scored_input), ("logic", args.logic_input), ("math", args.math_input), - ("code", args.code_input or [DEFAULT_CODE_INPUT]), + ("code", args.code_input), ) specs: list[SourceSpec] = [] for family, patterns in raw_specs: @@ -415,8 +421,10 @@ def resolve_paths(patterns: Sequence[str]) -> tuple[Path, ...]: path = Path(pattern) if path.is_dir(): paths.extend(path.glob("*.jsonl")) + paths.extend(path.glob("*.jsonl.zst")) paths.extend(path.glob("*.parquet")) paths.extend(path.glob("rank_*/*.jsonl")) + paths.extend(path.glob("rank_*/*.jsonl.zst")) paths.extend(path.glob("rank_*/*.parquet")) continue matches = glob.glob(pattern, recursive=True) @@ -427,10 +435,10 @@ def resolve_paths(patterns: Sequence[str]) -> tuple[Path, ...]: paths.append(path) continue raise FileNotFoundError(pattern) - supported = [path.resolve() for path in paths if file_format(path) in {"jsonl", "parquet"}] + supported = [path.resolve() for path in paths if file_format(path) in {"jsonl", "jsonl_zst", "parquet"}] unique = tuple(sorted(set(supported), key=str)) if not unique: - raise FileNotFoundError(f"no JSONL/Parquet files matched: {patterns!r}") + raise FileNotFoundError(f"no JSONL/JSONL.ZST/Parquet files matched: {patterns!r}") return unique @@ -446,6 +454,27 @@ def iter_rows(path: Path) -> Iterable[Mapping[str, object]]: if isinstance(row, Mapping): yield row return + if fmt == "jsonl_zst": + with subprocess.Popen( + ["zstd", "-dc", str(path)], + stdout=subprocess.PIPE, + stderr=subprocess.DEVNULL, + text=True, + encoding="utf-8", + errors="replace", + ) as proc: + assert proc.stdout is not None + for line in proc.stdout: + try: + row = json.loads(line) + except json.JSONDecodeError: + continue + if isinstance(row, Mapping): + yield row + returncode = proc.wait() + if returncode: + raise RuntimeError(f"zstd failed with exit code {returncode}: {path}") + return parquet_file = pq.ParquetFile(path) columns = [ name @@ -662,6 +691,8 @@ def prepare_output_dir(path: Path, *, overwrite: bool) -> None: def file_format(path: Path) -> str: + if path.name.lower().endswith(".jsonl.zst"): + return "jsonl_zst" suffix = path.suffix.lower() if suffix == ".jsonl": return "jsonl" diff --git a/dataset/pretrain/scripts/start_rebalanced_pretrain_dataset.sh b/dataset/pretrain/scripts/start_rebalanced_pretrain_dataset.sh index c75b339..3627eae 100755 --- a/dataset/pretrain/scripts/start_rebalanced_pretrain_dataset.sh +++ b/dataset/pretrain/scripts/start_rebalanced_pretrain_dataset.sh @@ -1,7 +1,7 @@ #!/usr/bin/env bash set -euo pipefail -ROOT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")/../.." && pwd)" +ROOT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")/../../.." && pwd)" cd "$ROOT_DIR" export http_proxy="${http_proxy:-http://10.20.34.2:3128}" @@ -18,7 +18,7 @@ echo "[rebalanced-data] root=$ROOT_DIR" echo "[rebalanced-data] default_target_tokens=$TARGET_TOKENS (CLI --target-tokens overrides this)" echo "[rebalanced-data] note=$WORKERS_NOTE" -python scripts/data/build_rebalanced_pretrain_dataset.py \ +python dataset/pretrain/scripts/build_rebalanced_pretrain_dataset.py \ --target-tokens "$TARGET_TOKENS" \ --output-dir "${OUTPUT_DIR:-data/train/pretrain_rebalanced_knowledge80_logic10_math5_code5_v1_20260630}" \ --seed "${SEED:-20260630}" \ diff --git a/dataset/pretrain/scripts/start_rebalanced_source_download.sh b/dataset/pretrain/scripts/start_rebalanced_source_download.sh index 75b6173..3df3b34 100755 --- a/dataset/pretrain/scripts/start_rebalanced_source_download.sh +++ b/dataset/pretrain/scripts/start_rebalanced_source_download.sh @@ -1,7 +1,7 @@ #!/usr/bin/env bash set -euo pipefail -ROOT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")/../.." && pwd)" +ROOT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")/../../.." && pwd)" cd "$ROOT_DIR" export http_proxy="${http_proxy:-http://10.20.34.2:3128}" @@ -21,7 +21,7 @@ SEED="${SEED:-42}" OUTPUT_ROOT="${OUTPUT_ROOT:-/mnt/beegfs/yi/laoyao_2b_moe_pretraining_dataset}" MAX_MBPS="${MAX_MBPS:-0}" -python scripts/data/download_rebalanced_sources.py \ +python dataset/pretrain/scripts/download_rebalanced_sources.py \ --seed "$SEED" \ --output-root "$OUTPUT_ROOT" \ --max-mbps "$MAX_MBPS" \ diff --git a/dataset/pretrain/scripts/wait_and_build_rebalanced_pretrain_200b.sh b/dataset/pretrain/scripts/wait_and_build_rebalanced_pretrain_200b.sh index f247f06..cd9a28b 100755 --- a/dataset/pretrain/scripts/wait_and_build_rebalanced_pretrain_200b.sh +++ b/dataset/pretrain/scripts/wait_and_build_rebalanced_pretrain_200b.sh @@ -1,7 +1,7 @@ #!/usr/bin/env bash set -euo pipefail -ROOT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")/../.." && pwd)" +ROOT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")/../../.." && pwd)" cd "$ROOT_DIR" export http_proxy="${http_proxy:-http://10.20.34.2:3128}" @@ -23,11 +23,27 @@ CHINESE_CLEAN_INPUT="${CHINESE_CLEAN_INPUT:-$DATA_ROOT/ms_ultra_fineweb_zh/**/*. SCIENCE_INPUT="${SCIENCE_INPUT:-$DATA_ROOT/science_*/*/*.parquet}" JIAYI_SCORED_INPUT="${JIAYI_SCORED_INPUT:-/mnt/beegfs/cjy/cjy-training/laoyao_model/data/scored/edu_fineweb_math_knowledge_cross_near_h1_v1_20260620/rank_*/*.jsonl}" CODE_INPUT="${CODE_INPUT:-/mnt/beegfs/cjy/cjy-training/laoyao_model/data/dedup/pretrain_python_clean_code_multilang_cross_near_h1_min256_v1_20260623/*.parquet}" +PROOF_PILE_ROOT="${PROOF_PILE_ROOT:-/mnt/beegfs/cjy/cjy-training/laoyao_model/data/raw/hf_datasets/EleutherAI__proof-pile-2}" + +LOGIC_INPUTS=() +LOGIC_INPUTS+=("${PROOF_PILE_ROOT}/open-web-math/train/*.jsonl.zst") +LOGIC_INPUTS+=("${PROOF_PILE_ROOT}/arxiv/train/*.jsonl.zst") +LOGIC_INPUTS+=("${PROOF_PILE_ROOT}/algebraic-stack/train/*.jsonl.zst") +if [[ -n "${LOGIC_INPUT_EXTRA:-}" ]]; then + IFS=':' read -r -a EXTRA_LOGIC_INPUTS <<< "$LOGIC_INPUT_EXTRA" + LOGIC_INPUTS+=("${EXTRA_LOGIC_INPUTS[@]}") +fi +LOGIC_ARGS=() +for logic_input in "${LOGIC_INPUTS[@]}"; do + LOGIC_ARGS+=(--logic-input "$logic_input") +done echo "[wait-build] root=$ROOT_DIR" echo "[wait-build] data_root=$DATA_ROOT" echo "[wait-build] output_dir=$OUTPUT_DIR" echo "[wait-build] target_tokens=$TARGET_TOKENS" +echo "[wait-build] proof_pile_root=$PROOF_PILE_ROOT" +printf '[wait-build] logic_input=%s\n' "${LOGIC_INPUTS[@]}" echo "[wait-build] waiting for zh shards: $EXPECTED_ZH_SHARDS" while true; do @@ -40,23 +56,30 @@ while true; do done echo "[wait-build] running full dry-run" -bash scripts/data/start_rebalanced_pretrain_dataset.sh \ +bash dataset/pretrain/scripts/start_rebalanced_pretrain_dataset.sh \ --dry-run \ --english-web-input "$ENGLISH_WEB_INPUT" \ --english-edu-input "$ENGLISH_EDU_INPUT" \ --chinese-clean-input "$CHINESE_CLEAN_INPUT" \ --science-input "$SCIENCE_INPUT" \ + "${LOGIC_ARGS[@]}" \ --jiayi-scored-input "$JIAYI_SCORED_INPUT" \ --code-input "$CODE_INPUT" +if [[ "${DRY_RUN_ONLY:-0}" == "1" ]]; then + echo "[wait-build] dry-run only; skip materialization" + exit 0 +fi + echo "[wait-build] starting full materialization" OUTPUT_DIR="$OUTPUT_DIR" TARGET_TOKENS="$TARGET_TOKENS" \ -bash scripts/data/start_rebalanced_pretrain_dataset.sh \ +bash dataset/pretrain/scripts/start_rebalanced_pretrain_dataset.sh \ --overwrite \ --english-web-input "$ENGLISH_WEB_INPUT" \ --english-edu-input "$ENGLISH_EDU_INPUT" \ --chinese-clean-input "$CHINESE_CLEAN_INPUT" \ --science-input "$SCIENCE_INPUT" \ + "${LOGIC_ARGS[@]}" \ --jiayi-scored-input "$JIAYI_SCORED_INPUT" \ --code-input "$CODE_INPUT" diff --git a/scripts/sync_pretrain_data_into_repo.sh b/scripts/sync_pretrain_data_into_repo.sh index fffc03c..16bb0e8 100755 --- a/scripts/sync_pretrain_data_into_repo.sh +++ b/scripts/sync_pretrain_data_into_repo.sh @@ -1,32 +1,15 @@ #!/usr/bin/env bash set -euo pipefail -SRC="${SRC:-/mnt/beegfs/yi/laoyao_2b_moe_pretraining_dataset/train/pretrain_rebalanced_web40_edu20_chinese10_science10_logic10_math5_code5_200b_v1_20260701}" -DST="${DST:-/mnt/beegfs/yi/laoyao_2b_moe/dataset/pretrain/data/pretrain_rebalanced_web40_edu20_chinese10_science10_logic10_math5_code5_200b_v1_20260701}" -LOG="${LOG:-/mnt/beegfs/yi/laoyao_2b_moe_pretraining_dataset/logs/wait_build_rebalanced_200b_20260701_130058.log}" +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}" -mkdir -p "$(dirname "$DST")" +echo "Large pretraining data is intentionally not copied into this repo." +echo "Use source data path directly: $SOURCE_DATA" -if pgrep -f "build_rebalanced_pretrain_dataset.py --target-tokens 200000000000" >/dev/null; then - echo "pretrain builder is still running; not syncing yet" >&2 - exit 2 -fi - -if [ ! -d "$SRC" ]; then - echo "source dataset directory does not exist: $SRC" >&2 +if [ ! -d "$SOURCE_DATA" ]; then + echo "missing source data: $SOURCE_DATA" >&2 exit 1 fi -if ! grep -q "rebalanced_pretrain_done\|completed\|wrote manifest\|finished" "$LOG" 2>/dev/null; then - echo "warning: completion marker not found in log; syncing existing directory anyway" >&2 -fi - -rm -rf "$DST" -mkdir -p "$DST" -if cp -al "$SRC"/. "$DST"/ 2>/tmp/laoyao_sync_cp_al.err; then - echo "synced with hardlinks: $DST" -else - echo "hardlink copy failed, falling back to rsync copy" >&2 - cat /tmp/laoyao_sync_cp_al.err >&2 || true - rsync -a --info=progress2 "$SRC"/ "$DST"/ -fi +du -sh "$SOURCE_DATA" || true +find "$SOURCE_DATA" -maxdepth 1 -type f | wc -l | awk "{print \"file_count=\" \$1}" diff --git a/scripts/train_nemo_megatron_2b_moe.sh b/scripts/train_nemo_megatron_2b_moe.sh index 5296eb2..2557972 100755 --- a/scripts/train_nemo_megatron_2b_moe.sh +++ b/scripts/train_nemo_megatron_2b_moe.sh @@ -4,10 +4,10 @@ set -euo pipefail REPO_ROOT="${REPO_ROOT:-/mnt/beegfs/yi/laoyao_2b_moe}" IMAGE="${IMAGE:-nvcr.io/nvidia/nemo:26.06}" CONFIG="${CONFIG:-$REPO_ROOT/training/nemo_megatron/pretrain_2b_moe_200b.yaml}" -DATA_DIR="$REPO_ROOT/dataset/pretrain/data/pretrain_rebalanced_web40_edu20_chinese10_science10_logic10_math5_code5_200b_v1_20260701" +DATA_DIR="/mnt/beegfs/yi/laoyao_2b_moe_pretraining_dataset/train/pretrain_rebalanced_web40_edu20_chinese10_science10_logic10_math5_code5_200b_v1_20260701" if [ ! -d "$DATA_DIR" ]; then - echo "pretrain data is missing under repo; run scripts/sync_pretrain_data_into_repo.sh after the 200B build finishes" >&2 + echo "pretrain source data is missing; check dataset/pretrain/manifests/rebalanced_182b_20260701_summary.json" >&2 exit 1 fi diff --git a/training/nemo_megatron/pretrain_2b_moe_200b.yaml b/training/nemo_megatron/pretrain_2b_moe_200b.yaml index d5147b8..405f79d 100644 --- a/training/nemo_megatron/pretrain_2b_moe_200b.yaml +++ b/training/nemo_megatron/pretrain_2b_moe_200b.yaml @@ -5,7 +5,7 @@ experiment: paths: model_config: /mnt/beegfs/yi/laoyao_2b_moe/model/nemo_megatron/laoyao_2b_moe_nemo_megatron.yaml - train_data: /mnt/beegfs/yi/laoyao_2b_moe/dataset/pretrain/data/pretrain_rebalanced_web40_edu20_chinese10_science10_logic10_math5_code5_200b_v1_20260701 + train_data: /mnt/beegfs/yi/laoyao_2b_moe_pretraining_dataset/train/pretrain_rebalanced_web40_edu20_chinese10_science10_logic10_math5_code5_200b_v1_20260701 validation_messages: /mnt/beegfs/yi/laoyao_2b_moe/dataset/val/data/heldout_2p8k.jsonl validation_sft: /mnt/beegfs/yi/laoyao_2b_moe/dataset/val/data/heldout_2p8k_sft_prompt_completion.jsonl validation_packed_text: /mnt/beegfs/yi/laoyao_2b_moe/dataset/val/data/heldout_2p8k_packed_text.jsonl