Add proof-pile logic inputs to pretrain builder

2026-07-01 22:42:18 +08:00
parent cdd8e2f85b
commit f0f72767f5
11 changed files with 122 additions and 54 deletions
--- 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` 忽略，仅用于小规模临时样本。
--- 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
+```
--- 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."
  ]
 }
--- a/dataset/pretrain/manifests/rebalanced_182b_20260701_summary.json
+++ 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."
+  ]
+}
--- 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"
--- 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}" \
--- 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" \
--- 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"

--- 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}"
--- 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

--- 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