初始化项目，由ModelHub XC社区提供模型

Model: AliceYin/l20-edu-135m
Source: Original Platform

training_artifacts/configs/l20_edu_135m_stage2_replay_polish_8k.yaml

@@ -0,0 +1,54 @@
+run_name: l20-edu-135m-stage2-replay-polish-8k
+output_dir: runs/l20-edu-135m-stage2-replay-polish-8k
+seed: 20260610
+tokenizer_name: AliceYin/l20-edu-135m
+init_model_name_or_path: runs/l20-edu-135m-stage2-math-code-textbook-replay-8k/step-001850
+
+dataset:
+  tokenized_path: data/l20_stage2_math_code_textbook_replay_8k
+  split: train
+  text_column: text
+  append_eos: true
+
+model:
+  block_size: 8192
+  hidden_size: 576
+  intermediate_size: 1536
+  num_hidden_layers: 30
+  num_attention_heads: 9
+  num_key_value_heads: 3
+  rope_theta: 10000.0
+  rms_norm_eps: 0.000001
+  attention_dropout: 0.0
+  tie_word_embeddings: true
+  vocab_multiple: 64
+  attn_implementation: sdpa
+  rope_scaling:
+    rope_type: yarn
+    factor: 4.0
+    original_max_position_embeddings: 2048
+
+trainer:
+  micro_batch_size: 3
+  gradient_accumulation_steps: 22
+  max_steps: 300
+  warmup_steps: 20
+  learning_rate: 0.000003
+  min_lr_ratio: 0.3
+  weight_decay: 0.1
+  beta1: 0.9
+  beta2: 0.95
+  grad_clip: 1.0
+  dtype: bfloat16
+  compile: false
+  compile_mode: null
+  compile_fullgraph: null
+  liger_kernel: true
+  gradient_checkpointing: false
+  log_interval: 1
+  eval_interval: 50
+  eval_batches: 16
+  save_interval: 50
+  keep_last_checkpoints: 4
+  num_workers: 0
+  mfu_peak_tflops: 119.5

training_artifacts/configs/l20_edu_135m_stage3_dclm_edu_replay_8k.yaml

@@ -0,0 +1,54 @@
+run_name: l20-edu-135m-stage3-dclm-edu-replay-8k
+output_dir: runs/l20-edu-135m-stage3-dclm-edu-replay-8k
+seed: 20260611
+tokenizer_name: AliceYin/l20-edu-135m
+init_model_name_or_path: runs/l20-edu-135m-stage2-replay-polish-8k/final
+
+dataset:
+  tokenized_path: data/l20_stage3_dclm_edu_replay_8k
+  split: train
+  text_column: text
+  append_eos: true
+
+model:
+  block_size: 8192
+  hidden_size: 576
+  intermediate_size: 1536
+  num_hidden_layers: 30
+  num_attention_heads: 9
+  num_key_value_heads: 3
+  rope_theta: 10000.0
+  rms_norm_eps: 0.000001
+  attention_dropout: 0.0
+  tie_word_embeddings: true
+  vocab_multiple: 64
+  attn_implementation: sdpa
+  rope_scaling:
+    rope_type: yarn
+    factor: 4.0
+    original_max_position_embeddings: 2048
+
+trainer:
+  micro_batch_size: 3
+  gradient_accumulation_steps: 22
+  max_steps: 3700
+  warmup_steps: 100
+  learning_rate: 0.000006
+  min_lr_ratio: 0.15
+  weight_decay: 0.1
+  beta1: 0.9
+  beta2: 0.95
+  grad_clip: 1.0
+  dtype: bfloat16
+  compile: false
+  compile_mode: null
+  compile_fullgraph: null
+  liger_kernel: true
+  gradient_checkpointing: false
+  log_interval: 1
+  eval_interval: 100
+  eval_batches: 16
+  save_interval: 250
+  keep_last_checkpoints: 4
+  num_workers: 0
+  mfu_peak_tflops: 119.5

training_artifacts/configs/mixtures/l20_stage3_dclm_edu_replay.yaml

@@ -0,0 +1,163 @@
+name: l20-stage3-dclm-edu-replay
+output_dir: data/l20_stage3_dclm_edu_replay_8k
+tokenizer: AliceYin/l20-edu-135m
+target_tokens: 2000000000
+val_tokens: 4194304
+block_size: 8192
+report_interval: 1000
+
+# Stage 3 targets the current gap to SmolLM/SmolLM2 on commonsense and
+# language-understanding tasks. DCLM-Edu and FineWeb-Edu are the backbone;
+# replay limits drift, while small math/code/textbook tails preserve gains.
+sources:
+  - name: dclm-edu-3plus
+    family: edu_web
+    kind: hf_text
+    dataset: HuggingFaceTB/dclm-edu
+    split: train
+    text_column: text
+    weight: 0.45
+    min_chars: 700
+    max_chars: 45000
+    min_score: 3.0
+    min_int_score: 3
+    quality:
+      min_alpha_ratio: 0.48
+      max_digit_ratio: 0.28
+      min_unique_word_ratio: 0.18
+      max_repeated_ngram_fraction: 0.16
+
+  - name: smollm-fineweb-edu-dedup-4plus
+    family: edu_web
+    kind: hf_text
+    dataset: HuggingFaceTB/smollm-corpus
+    config_name: fineweb-edu-dedup
+    split: train
+    text_column: text
+    weight: 0.25
+    min_chars: 700
+    max_chars: 45000
+    min_score: 4.0
+    min_int_score: 4
+    quality:
+      min_alpha_ratio: 0.48
+      max_digit_ratio: 0.28
+      min_unique_word_ratio: 0.18
+      max_repeated_ngram_fraction: 0.16
+
+  - name: edu-hq-replay
+    family: edu_replay
+    kind: tokenized_replay
+    tokenized_path: data/l20_edu_hq_8k
+    tokenized_split: train
+    weight: 0.12
+    sample_seed: 20260611
+
+  - name: smollm-cosmopedia-v2
+    family: synthetic_textbook
+    kind: hf_text
+    dataset: HuggingFaceTB/smollm-corpus
+    config_name: cosmopedia-v2
+    split: train
+    text_column: text
+    weight: 0.08
+    min_chars: 700
+    max_chars: 45000
+    quality:
+      min_alpha_ratio: 0.45
+      max_digit_ratio: 0.35
+      min_unique_word_ratio: 0.16
+      max_repeated_ngram_fraction: 0.18
+
+  - name: finemath-4plus
+    family: math
+    kind: hf_text
+    dataset: HuggingFaceTB/finemath
+    config_name: finemath-4plus
+    split: train
+    text_column: text
+    weight: 0.06
+    min_chars: 500
+    max_chars: 45000
+    min_score: 4.0
+    min_int_score: 4
+    quality:
+      min_alpha_ratio: 0.25
+      max_digit_ratio: 0.45
+      min_unique_word_ratio: 0.08
+      max_repeated_ngram_fraction: 0.22
+
+  - name: stack-edu-python
+    family: code
+    kind: stack_edu
+    dataset: HuggingFaceTB/stack-edu
+    config_name: Python
+    split: train
+    text_column: text
+    weight: 0.02
+    min_chars: 250
+    max_chars: 60000
+    min_score: 3.5
+    min_int_score: 4
+    require_license_type: permissive
+    include_metadata_header: true
+
+  - name: stack-edu-javascript
+    family: code
+    kind: stack_edu
+    dataset: HuggingFaceTB/stack-edu
+    config_name: JavaScript
+    split: train
+    text_column: text
+    weight: 0.008
+    min_chars: 250
+    max_chars: 60000
+    min_score: 3.5
+    min_int_score: 4
+    require_license_type: permissive
+    include_metadata_header: true
+
+  - name: stack-edu-typescript
+    family: code
+    kind: stack_edu
+    dataset: HuggingFaceTB/stack-edu
+    config_name: TypeScript
+    split: train
+    text_column: text
+    weight: 0.005
+    min_chars: 250
+    max_chars: 60000
+    min_score: 3.5
+    min_int_score: 4
+    require_license_type: permissive
+    include_metadata_header: true
+
+  - name: stack-edu-cpp
+    family: code
+    kind: stack_edu
+    dataset: HuggingFaceTB/stack-edu
+    config_name: Cpp
+    split: train
+    text_column: text
+    weight: 0.004
+    min_chars: 250
+    max_chars: 60000
+    min_score: 3.5
+    min_int_score: 4
+    require_license_type: permissive
+    include_metadata_header: true
+
+  - name: stack-edu-rust
+    family: code
+    kind: stack_edu
+    dataset: HuggingFaceTB/stack-edu
+    config_name: Rust
+    split: train
+    text_column: text
+    weight: 0.003
+    min_chars: 250
+    max_chars: 60000
+    min_score: 3.5
+    min_int_score: 4
+    require_license_type: permissive
+    include_metadata_header: true

training_artifacts/scripts/eval_smollm_benchmark.sh

@@ -0,0 +1,71 @@
+#!/usr/bin/env bash
+set -euo pipefail
+
+TASKS="${TASKS:-arc_challenge,arc_easy,hellaswag,lambada_openai,piqa,winogrande}"
+DEVICE="${DEVICE:-cuda:0}"
+DTYPE="${DTYPE:-bfloat16}"
+BATCH_SIZE="${BATCH_SIZE:-auto}"
+OUTPUT_ROOT="${OUTPUT_ROOT:-eval_results/smollm_target_$(date +%Y%m%d_%H%M%S)}"
+REQUESTED_CANDIDATE="${CANDIDATE:-}"
+
+export HF_ENDPOINT="${HF_ENDPOINT:-https://hf-mirror.com}"
+export HF_HUB_DISABLE_XET="${HF_HUB_DISABLE_XET:-1}"
+
+if ! command -v lm_eval >/dev/null 2>&1; then
+  echo "lm_eval is not on PATH. Activate the eval environment first, for example: source .venv-eval/bin/activate" >&2
+  exit 2
+fi
+
+sanitize_name() {
+  echo "$1" | tr '/: ' '___'
+}
+
+declare -a MODELS
+if [ "$#" -gt 0 ]; then
+  MODELS=("$@")
+  CANDIDATE="${REQUESTED_CANDIDATE:-${1%%=*}}"
+else
+  MODELS=(
+    "ours-stage2=runs/l20-edu-135m-stage2-math-code-textbook-replay-8k/step-001850"
+    "smollm-135m=HuggingFaceTB/SmolLM-135M"
+    "smollm2-135m=HuggingFaceTB/SmolLM2-135M"
+  )
+  if [ -e "runs/l20-edu-135m-stage2-replay-polish-8k/final" ] || [ -e "runs/l20-edu-135m-stage2-replay-polish-8k/step-000300" ]; then
+    MODELS=("ours-polish=runs/l20-edu-135m-stage2-replay-polish-8k/final" "${MODELS[@]}")
+    CANDIDATE="${REQUESTED_CANDIDATE:-ours-polish}"
+  else
+    CANDIDATE="${REQUESTED_CANDIDATE:-ours-stage2}"
+  fi
+fi
+
+mkdir -p "$OUTPUT_ROOT"
+declare -a RESULTS
+
+for entry in "${MODELS[@]}"; do
+  if [[ "$entry" != *=* ]]; then
+    echo "Expected model entry NAME=MODEL_PATH_OR_HF_ID, got: $entry" >&2
+    exit 2
+  fi
+  name="${entry%%=*}"
+  model="${entry#*=}"
+  out_dir="$OUTPUT_ROOT/$(sanitize_name "$name")"
+  echo "==> Evaluating $name: $model"
+  lm_eval \
+    --model hf \
+    --model_args "pretrained=${model},dtype=${DTYPE}" \
+    --tasks "$TASKS" \
+    --device "$DEVICE" \
+    --batch_size "$BATCH_SIZE" \
+    --output_path "$out_dir" \
+    --log_samples
+  RESULTS+=("--result" "${name}=${out_dir}")
+done
+
+python scripts/summarize_smollm_benchmark.py \
+  "${RESULTS[@]}" \
+  --candidate "$CANDIDATE" \
+  --baseline smollm-135m \
+  --baseline smollm2-135m \
+  --out-md "$OUTPUT_ROOT/summary.md" \
+  --out-json "$OUTPUT_ROOT/summary.json" \
+  --out-csv "$OUTPUT_ROOT/summary.csv"

training_artifacts/scripts/prepare_l20_stage3_dclm_edu_replay_8k.sh

@@ -0,0 +1,36 @@
+#!/usr/bin/env bash
+set -euo pipefail
+
+export TOKENIZERS_PARALLELISM=false
+export HF_ENDPOINT="${HF_ENDPOINT:-https://hf-mirror.com}"
+export HF_HUB_DISABLE_XET="${HF_HUB_DISABLE_XET:-1}"
+export HF_HUB_DOWNLOAD_TIMEOUT="${HF_HUB_DOWNLOAD_TIMEOUT:-120}"
+export HF_HUB_ETAG_TIMEOUT="${HF_HUB_ETAG_TIMEOUT:-60}"
+export PYTHONUNBUFFERED=1
+
+PYTHON="${PYTHON:-python}"
+RECIPE="${RECIPE:-configs/mixtures/l20_stage3_dclm_edu_replay.yaml}"
+TARGET_TOKENS="${TARGET_TOKENS:-2000000000}"
+VAL_TOKENS="${VAL_TOKENS:-4194304}"
+OUTPUT_DIR="${OUTPUT_DIR:-data/l20_stage3_dclm_edu_replay_8k}"
+
+if [ ! -f "data/l20_edu_hq_8k/train.bin" ]; then
+  echo "Stage-1 tokenized replay shard is missing: data/l20_edu_hq_8k/train.bin" >&2
+  exit 2
+fi
+
+set +e
+"$PYTHON" -m l20_pretrain.prepare_mixture_shards \
+  --recipe "$RECIPE" \
+  --output-dir "$OUTPUT_DIR" \
+  --target-tokens "$TARGET_TOKENS" \
+  --val-tokens "$VAL_TOKENS"
+STATUS=$?
+set -e
+
+if [ "$STATUS" -ne 0 ] && [ -f "$OUTPUT_DIR/metadata.json" ]; then
+  echo "prepare_mixture_python_exit_code=$STATUS after metadata was written; treating shards as complete"
+  exit 0
+fi
+
+exit "$STATUS"

training_artifacts/scripts/run_stage3_dclm_edu_replay_after_polish.sh

@@ -0,0 +1,67 @@
+#!/usr/bin/env bash
+set -euo pipefail
+
+export TOKENIZERS_PARALLELISM=false
+export HF_ENDPOINT="${HF_ENDPOINT:-https://hf-mirror.com}"
+export HF_HUB_DISABLE_XET="${HF_HUB_DISABLE_XET:-1}"
+export HF_HUB_DOWNLOAD_TIMEOUT="${HF_HUB_DOWNLOAD_TIMEOUT:-120}"
+export HF_HUB_ETAG_TIMEOUT="${HF_HUB_ETAG_TIMEOUT:-60}"
+export PYTORCH_CUDA_ALLOC_CONF="${PYTORCH_CUDA_ALLOC_CONF:-expandable_segments:True}"
+export PYTHONUNBUFFERED=1
+
+PYTHON="${PYTHON:-python}"
+POLL_SECONDS="${POLL_SECONDS:-300}"
+LOG_DIR="${LOG_DIR:-logs}"
+DATA_DIR="${DATA_DIR:-data/l20_stage3_dclm_edu_replay_8k}"
+POLISH_FINAL="${POLISH_FINAL:-runs/l20-edu-135m-stage2-replay-polish-8k/final}"
+TRAIN_CONFIG="${TRAIN_CONFIG:-configs/l20_edu_135m_stage3_dclm_edu_replay_8k.yaml}"
+PREPARE_LOG="${PREPARE_LOG:-$LOG_DIR/l20_stage3_dclm_edu_replay_prepare_latest.log}"
+TRAIN_LOG="${TRAIN_LOG:-$LOG_DIR/l20_stage3_dclm_edu_replay_train_latest.log}"
+
+mkdir -p "$LOG_DIR"
+
+active_train_pids() {
+  pgrep -af "python -m l20_pretrain.train" || true
+}
+
+data_complete() {
+  "$PYTHON" - "$DATA_DIR" <<'PY'
+from pathlib import Path
+import json
+import sys
+
+root = Path(sys.argv[1])
+metadata_path = root / "metadata.json"
+if not metadata_path.is_file() or not (root / "train.bin").is_file() or not (root / "val.bin").is_file():
+    raise SystemExit(1)
+metadata = json.loads(metadata_path.read_text())
+target = int(metadata.get("target_tokens") or 0)
+train_tokens = int(metadata.get("train_tokens") or 0)
+val_tokens = int(metadata.get("val_tokens") or 0)
+if target <= 0 or train_tokens < target or val_tokens <= 0:
+    raise SystemExit(1)
+PY
+}
+
+echo "stage3_wait_for_polish $(date -Is)"
+while [ ! -e "$POLISH_FINAL" ]; do
+  echo "waiting for polish final: $POLISH_FINAL"
+  sleep "$POLL_SECONDS"
+done
+
+while active_train_pids | grep -q .; do
+  echo "waiting for active training to finish"
+  active_train_pids
+  sleep "$POLL_SECONDS"
+done
+
+if data_complete; then
+  echo "stage3_data_complete $(date -Is)"
+else
+  echo "stage3_prepare_start $(date -Is)"
+  PYTHON="$PYTHON" bash scripts/prepare_l20_stage3_dclm_edu_replay_8k.sh 2>&1 | tee "$PREPARE_LOG"
+fi
+
+echo "stage3_train_start $(date -Is)"
+PYTHON="$PYTHON" bash scripts/train_l20_stage2_math_code_textbook_replay_8k.sh "$TRAIN_CONFIG" 2>&1 | tee "$TRAIN_LOG"
+echo "stage3_train_done $(date -Is)"

training_artifacts/scripts/summarize_smollm_benchmark.py

@@ -0,0 +1,176 @@
+#!/usr/bin/env python3
+from __future__ import annotations
+
+import argparse
+import csv
+import json
+from pathlib import Path
+from typing import Any
+
+
+TASK_METRICS = {
+    "arc_challenge": ("ARC-Challenge", "acc_norm,none"),
+    "arc_easy": ("ARC-Easy", "acc_norm,none"),
+    "hellaswag": ("HellaSwag", "acc_norm,none"),
+    "lambada_openai": ("LAMBADA OpenAI", "acc,none"),
+    "piqa": ("PIQA", "acc_norm,none"),
+    "winogrande": ("WinoGrande", "acc,none"),
+}
+
+
+def parse_named_path(value: str) -> tuple[str, Path]:
+    if "=" not in value:
+        raise argparse.ArgumentTypeError("expected NAME=PATH")
+    name, path = value.split("=", 1)
+    if not name or not path:
+        raise argparse.ArgumentTypeError("expected NAME=PATH")
+    return name, Path(path)
+
+
+def find_result_json(path: Path) -> Path:
+    if path.is_file():
+        return path
+    candidates = sorted(path.rglob("results_*.json"))
+    if not candidates:
+        candidates = sorted(path.rglob("*.json"))
+    for candidate in candidates:
+        try:
+            payload = json.loads(candidate.read_text(encoding="utf-8"))
+        except Exception:
+            continue
+        if isinstance(payload, dict) and isinstance(payload.get("results"), dict):
+            return candidate
+    raise FileNotFoundError(f"No lm-eval result JSON found under {path}")
+
+
+def load_scores(path: Path) -> dict[str, float | None]:
+    result_json = find_result_json(path)
+    payload = json.loads(result_json.read_text(encoding="utf-8"))
+    results = payload["results"]
+    scores: dict[str, float | None] = {}
+    for task, (_, metric) in TASK_METRICS.items():
+        value = results.get(task, {}).get(metric)
+        scores[task] = float(value) if isinstance(value, int | float) else None
+    return scores
+
+
+def mean_score(scores: dict[str, float | None]) -> float | None:
+    values = [value for value in scores.values() if value is not None]
+    return sum(values) / len(values) if values else None
+
+
+def fmt(value: float | None) -> str:
+    return "" if value is None else f"{value:.4f}"
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(
+        description="Summarize the six SmolLM target lm-eval tasks and show gaps."
+    )
+    parser.add_argument("--result", action="append", default=[], type=parse_named_path)
+    parser.add_argument("--candidate", default=None, help="Candidate model name for gap columns.")
+    parser.add_argument("--baseline", action="append", default=[], help="Baseline names to compare against.")
+    parser.add_argument("--out-md", default="eval_results/smollm_benchmark.md")
+    parser.add_argument("--out-json", default="eval_results/smollm_benchmark.json")
+    parser.add_argument("--out-csv", default="eval_results/smollm_benchmark.csv")
+    args = parser.parse_args()
+
+    if not args.result:
+        raise SystemExit("At least one --result NAME=PATH is required.")
+
+    score_by_name = {name: load_scores(path) for name, path in args.result}
+    candidate_name = args.candidate or next(iter(score_by_name))
+    if candidate_name not in score_by_name:
+        raise SystemExit(f"Candidate {candidate_name!r} was not provided in --result.")
+    baseline_names = [name for name in args.baseline if name in score_by_name]
+
+    rows: list[dict[str, Any]] = []
+    for task, (display, metric) in TASK_METRICS.items():
+        row: dict[str, Any] = {
+            "task": task,
+            "display": display,
+            "metric": metric,
+        }
+        for name, scores in score_by_name.items():
+            row[name] = scores.get(task)
+        for baseline_name in baseline_names:
+            candidate_value = row.get(candidate_name)
+            baseline_value = row.get(baseline_name)
+            row[f"gap_vs_{baseline_name}"] = (
+                None
+                if candidate_value is None or baseline_value is None
+                else float(candidate_value) - float(baseline_value)
+            )
+        rows.append(row)
+
+    means = {name: mean_score(scores) for name, scores in score_by_name.items()}
+    payload = {
+        "candidate": candidate_name,
+        "baselines": baseline_names,
+        "tasks": rows,
+        "means": means,
+    }
+
+    out_json = Path(args.out_json)
+    out_json.parent.mkdir(parents=True, exist_ok=True)
+    out_json.write_text(json.dumps(payload, indent=2, sort_keys=True), encoding="utf-8")
+
+    fieldnames = ["Task", "Metric", *score_by_name.keys(), *[f"Gap vs {name}" for name in baseline_names]]
+    out_csv = Path(args.out_csv)
+    out_csv.parent.mkdir(parents=True, exist_ok=True)
+    with out_csv.open("w", encoding="utf-8", newline="") as handle:
+        writer = csv.DictWriter(handle, fieldnames=fieldnames)
+        writer.writeheader()
+        for row in rows:
+            writer.writerow(
+                {
+                    "Task": row["display"],
+                    "Metric": row["metric"],
+                    **{name: fmt(row.get(name)) for name in score_by_name},
+                    **{f"Gap vs {name}": fmt(row.get(f"gap_vs_{name}")) for name in baseline_names},
+                }
+            )
+
+    lines = ["# SmolLM Target Benchmark", ""]
+    lines.append("| Task | Metric | " + " | ".join(score_by_name.keys()) + " |")
+    lines.append("| --- | --- | " + " | ".join(["---"] * len(score_by_name)) + " |")
+    for row in rows:
+        lines.append(
+            "| "
+            + " | ".join(
+                [row["display"], row["metric"], *[fmt(row.get(name)) for name in score_by_name]]
+            )
+            + " |"
+        )
+    lines.extend(["", "## Means", ""])
+    for name, value in means.items():
+        lines.append(f"- `{name}`: {fmt(value)}")
+    if baseline_names:
+        lines.extend(["", "## Candidate Gaps", ""])
+        for baseline_name in baseline_names:
+            baseline_mean = means.get(baseline_name)
+            candidate_mean = means.get(candidate_name)
+            gap = (
+                None
+                if baseline_mean is None or candidate_mean is None
+                else candidate_mean - baseline_mean
+            )
+            lines.append(f"- `{candidate_name}` vs `{baseline_name}` mean gap: {fmt(gap)}")
+            ranked = sorted(
+                (
+                    (row["display"], row.get(f"gap_vs_{baseline_name}"))
+                    for row in rows
+                    if row.get(f"gap_vs_{baseline_name}") is not None
+                ),
+                key=lambda item: item[1],
+            )
+            for task_name, task_gap in ranked:
+                lines.append(f"  - {task_name}: {fmt(task_gap)}")
+    out_md = Path(args.out_md)
+    out_md.parent.mkdir(parents=True, exist_ok=True)
+    out_md.write_text("\n".join(lines) + "\n", encoding="utf-8")
+    print("\n".join(lines))
+
+
+if __name__ == "__main__":
+    main()

training_artifacts/scripts/train_l20_stage2_math_code_textbook_replay_8k.sh

@@ -0,0 +1,65 @@
+#!/usr/bin/env bash
+set -euo pipefail
+
+export TOKENIZERS_PARALLELISM=false
+export HF_ENDPOINT="${HF_ENDPOINT:-https://hf-mirror.com}"
+export HF_HUB_DISABLE_XET="${HF_HUB_DISABLE_XET:-1}"
+export HF_HUB_DOWNLOAD_TIMEOUT="${HF_HUB_DOWNLOAD_TIMEOUT:-120}"
+export HF_HUB_ETAG_TIMEOUT="${HF_HUB_ETAG_TIMEOUT:-60}"
+export PYTORCH_CUDA_ALLOC_CONF="${PYTORCH_CUDA_ALLOC_CONF:-expandable_segments:True}"
+export PYTHONUNBUFFERED=1
+
+PYTHON="${PYTHON:-python}"
+CONFIG="${1:-configs/l20_edu_135m_stage2_math_code_textbook_replay_8k.yaml}"
+OUTPUT_DIR="$("$PYTHON" - <<'PY' "$CONFIG"
+from pathlib import Path
+import sys
+import yaml
+
+with Path(sys.argv[1]).open("r", encoding="utf-8") as handle:
+    print(yaml.safe_load(handle)["output_dir"])
+PY
+)"
+
+INIT_DIR="$("$PYTHON" - <<'PY' "$CONFIG"
+from pathlib import Path
+import sys
+import yaml
+
+with Path(sys.argv[1]).open("r", encoding="utf-8") as handle:
+    print(yaml.safe_load(handle)["init_model_name_or_path"])
+PY
+)"
+
+DATA_DIR="$("$PYTHON" - <<'PY' "$CONFIG"
+from pathlib import Path
+import sys
+import yaml
+
+with Path(sys.argv[1]).open("r", encoding="utf-8") as handle:
+    print(yaml.safe_load(handle)["dataset"]["tokenized_path"])
+PY
+)"
+
+if [ ! -d "$INIT_DIR" ]; then
+  echo "Stage-1 final checkpoint is missing: $INIT_DIR" >&2
+  exit 2
+fi
+
+if [ ! -f "$DATA_DIR/train.bin" ]; then
+  echo "Stage-2 replay tokenized train shard is missing: $DATA_DIR/train.bin" >&2
+  exit 2
+fi
+
+RESUME_DIR=""
+if [ -d "$OUTPUT_DIR" ]; then
+  RESUME_DIR="$(find "$OUTPUT_DIR" -maxdepth 1 -type d -name 'step-*' | sort | tail -n 1)"
+fi
+
+if [ -n "$RESUME_DIR" ] && [ -f "$RESUME_DIR/trainer_state.pt" ]; then
+  echo "Resuming stage2 replay from $RESUME_DIR"
+  "$PYTHON" -m l20_pretrain.train "$CONFIG" --resume "$RESUME_DIR"
+else
+  echo "Starting stage2 replay from $INIT_DIR"
+  "$PYTHON" -m l20_pretrain.train "$CONFIG"
+fi

training_artifacts/src/l20_pretrain/prepare_shards.py

@@ -0,0 +1,219 @@
+from __future__ import annotations
+
+import argparse
+from collections import Counter
+from collections.abc import Iterable
+import json
+import os
+from pathlib import Path
+import time
+from typing import Any
+
+import numpy as np
+
+from .env import set_default_hf_home
+
+set_default_hf_home()
+
+from transformers import AutoTokenizer
+
+from .data import create_source, tokenize_without_specials
+from .config import DatasetConfig
+from .quality import normalize_text, quality_filter, stable_hash
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(description="Clean, deduplicate, tokenize, and pack pretraining shards.")
+    parser.add_argument("--output-dir", required=True, help="Directory containing train.bin, val.bin, metadata.json.")
+    parser.add_argument("--tokenizer", default="AliceYin/l20-edu-135m")
+    parser.add_argument("--dataset", default="HuggingFaceFW/fineweb-edu")
+    parser.add_argument("--config-name", default="sample-10BT")
+    parser.add_argument("--split", default="train")
+    parser.add_argument("--text-column", default="text")
+    parser.add_argument("--local-text-path", default=None)
+    parser.add_argument("--target-tokens", type=int, default=100_000_000)
+    parser.add_argument("--val-tokens", type=int, default=2_000_000)
+    parser.add_argument("--block-size", type=int, default=8192)
+    parser.add_argument("--min-chars", type=int, default=500)
+    parser.add_argument("--max-chars", type=int, default=40_000)
+    parser.add_argument("--min-score", type=float, default=3.0)
+    parser.add_argument("--min-int-score", type=int, default=3)
+    parser.add_argument("--report-interval", type=int, default=1000)
+    return parser.parse_args()
+
+
+def get_text(example: Any, text_column: str) -> str | None:
+    if isinstance(example, str):
+        return example
+    if isinstance(example, dict):
+        value = example.get(text_column)
+        return value if isinstance(value, str) else None
+    return None
+
+
+def passes_dataset_score(example: Any, *, min_score: float | None, min_int_score: int | None) -> bool:
+    if not isinstance(example, dict):
+        return True
+    metadata = example.get("metadata")
+    metadata = metadata if isinstance(metadata, dict) else {}
+    if min_score is not None:
+        score = example.get("score")
+        if score is None:
+            score = example.get("edu_score")
+        if score is None:
+            score = metadata.get("score")
+        if score is None:
+            score = metadata.get("edu_score")
+        if score is not None and float(score) < min_score:
+            return False
+    if min_int_score is not None:
+        int_score = example.get("int_score")
+        if int_score is None:
+            int_score = example.get("edu_int_score")
+        if int_score is None:
+            int_score = metadata.get("int_score")
+        if int_score is None:
+            int_score = metadata.get("edu_int_score")
+        if int_score is not None and int(int_score) < min_int_score:
+            return False
+    return True
+
+
+def write_tokens(handle: Any, ids: list[int]) -> int:
+    if not ids:
+        return 0
+    array = np.asarray(ids, dtype=np.uint32)
+    array.tofile(handle)
+    return int(array.size)
+
+
+def iter_examples(args: argparse.Namespace) -> Iterable[Any]:
+    config = DatasetConfig(
+        name=args.dataset,
+        config_name=args.config_name,
+        split=args.split,
+        streaming=True,
+        text_column=args.text_column,
+        min_chars=args.min_chars,
+        max_chars=args.max_chars,
+        min_score=args.min_score,
+        min_int_score=args.min_int_score,
+        append_eos=True,
+        shuffle_buffer=0,
+        local_text_path=args.local_text_path,
+    )
+    return create_source(config)
+
+
+def main() -> None:
+    args = parse_args()
+    output_dir = Path(args.output_dir)
+    output_dir.mkdir(parents=True, exist_ok=True)
+
+    tokenizer = AutoTokenizer.from_pretrained(args.tokenizer, use_fast=True)
+    eos_token_id = tokenizer.eos_token_id
+    if eos_token_id is None:
+        raise ValueError("Tokenizer must provide eos_token_id")
+
+    train_path = output_dir / "train.bin"
+    val_path = output_dir / "val.bin"
+    metadata_path = output_dir / "metadata.json"
+    seen: set[str] = set()
+    counters: Counter[str] = Counter()
+    train_tokens = 0
+    val_tokens = 0
+    started_at = time.time()
+
+    with train_path.open("wb") as train_handle, val_path.open("wb") as val_handle:
+        for example in iter_examples(args):
+            counters["seen"] += 1
+            if not passes_dataset_score(
+                example,
+                min_score=args.min_score,
+                min_int_score=args.min_int_score,
+            ):
+                counters["score_reject"] += 1
+                continue
+
+            raw_text = get_text(example, args.text_column)
+            if not raw_text:
+                counters["empty"] += 1
+                continue
+            text = normalize_text(raw_text, max_chars=args.max_chars)
+            decision = quality_filter(text, min_chars=args.min_chars)
+            if not decision.keep:
+                counters[f"quality_{decision.reason}"] += 1
+                continue
+
+            digest = stable_hash(text)
+            if digest in seen:
+                counters["duplicate"] += 1
+                continue
+            seen.add(digest)
+
+            ids = tokenize_without_specials(tokenizer, text)
+            if len(ids) < 64:
+                counters["too_few_tokens"] += 1
+                continue
+            ids.append(int(eos_token_id))
+
+            if val_tokens < args.val_tokens and int(digest[:8], 16) % 97 == 0:
+                val_tokens += write_tokens(val_handle, ids)
+            else:
+                train_tokens += write_tokens(train_handle, ids)
+            counters["kept"] += 1
+
+            total_tokens = train_tokens + val_tokens
+            if counters["seen"] % args.report_interval == 0:
+                elapsed = max(time.time() - started_at, 1e-9)
+                print(
+                    json.dumps(
+                        {
+                            "event": "prepare",
+                            "seen_docs": counters["seen"],
+                            "kept_docs": counters["kept"],
+                            "train_tokens": train_tokens,
+                            "val_tokens": val_tokens,
+                            "tokens_per_sec": total_tokens / elapsed,
+                            "rejects": {
+                                key: value
+                                for key, value in counters.items()
+                                if key not in {"seen", "kept"}
+                            },
+                        },
+                        ensure_ascii=True,
+                    ),
+                    flush=True,
+                )
+            if train_tokens >= args.target_tokens and val_tokens >= args.val_tokens:
+                break
+
+    metadata = {
+        "dtype": "uint32",
+        "tokenizer": args.tokenizer,
+        "dataset": args.dataset,
+        "config_name": args.config_name,
+        "split": args.split,
+        "block_size": args.block_size,
+        "target_tokens": args.target_tokens,
+        "train_tokens": train_tokens,
+        "val_tokens": val_tokens,
+        "train_blocks": train_tokens // args.block_size,
+        "val_blocks": val_tokens // args.block_size,
+        "filters": {
+            "min_chars": args.min_chars,
+            "max_chars": args.max_chars,
+            "min_score": args.min_score,
+            "min_int_score": args.min_int_score,
+        },
+        "counters": dict(counters),
+        "elapsed_sec": time.time() - started_at,
+        "hf_endpoint": os.environ.get("HF_ENDPOINT"),
+    }
+    with metadata_path.open("w", encoding="utf-8") as handle:
+        json.dump(metadata, handle, indent=2, sort_keys=True)
+    print(json.dumps({"event": "done", **metadata}, ensure_ascii=True), flush=True)
+
+
+if __name__ == "__main__":
+    main()