diff --git a/README.md b/README.md new file mode 100644 index 0000000000000000000000000000000000000000..9f69bf9a213ee4038e8d3e5460697b1a8c335997 --- /dev/null +++ b/README.md @@ -0,0 +1,274 @@ +# SDG — Distributed Sharded Pipeline + +Generate synthetic data with vLLM + uncertainty-quantification validation, distributed across multiple GPUs via SLURM sharding with a shared MongoDB inference cache. + +## Prerequisites + +```bash +pip install -r sdg/requirements.txt +``` + +Install MongoDB (download binary): + +```bash +curl -O https://fastdl.mongodb.org/linux/mongodb-linux-x86_64-ubuntu2204-8.0.4.tgz +tar xzf mongodb-linux-x86_64-ubuntu2204-8.0.4.tgz +cp mongodb-linux-x86_64-ubuntu2204-8.0.4/bin/mongod /u/nlp/anaconda/main/anaconda3/envs/tonyreasoningtraces/bin/ +mongod --version +``` + +## Quick Start (Single GPU, No Sharding) + +```bash +python -m sdg.generate --config sdg/configs/qwen3_4b_openthoughts4_math30K_instill_n4_valredundancy3_round1.yaml --limit 100 +``` + +This requires `mongo_uri` set in the config YAML (see below). + +## Distributed Sharded Run + +### 1. Start MongoDB + +Launch a MongoDB server on a cluster node (no GPU needed) with max TTL of 21 days: + +```bash +nlprun -a tonyreasoningtraces --time 21-0 -m john3 --exclude john17 -c 4 -g 0 --memory 64g \ + -w /nlp/scr4/nlp/crfm/text2image/text2image-rlhf/reasoning/virtual-world-data \ + --job-name sdg-mongodb "python -m sdg.launch_mongodb 2>&1 | tee mongo.log" + +nlprun -a tonyreasoningtraces --time 21-0 -q jag -p high -m jagupard39 --exclude john17 -c 4 -g 0 --memory 64g \ + -w /nlp/scr4/nlp/crfm/text2image/text2image-rlhf/reasoning/virtual-world-data \ + --job-name sdg-serve "python -m sdg.launch_mongodb 2>&1 | tee mongo.log" +``` + +Or locally: + +```bash +python -m sdg.launch_mongodb --port 27017 --dbpath /tmp/sdg_mongo +``` + +It will print `mongo-uri: mongodb://:27017` to stdout (and `mongo.log`). + +### 2. Set `mongo_uri` and `num_shards` in your config + +Grab the hostname from `mongo.log`, then edit your YAML config (e.g. `sdg/configs/qwen3_4b_openthoughts4_math30K_instill_n4_valredundancy3_round1.yaml`): + +```yaml +mongo_uri: "mongodb://:27017" +num_shards: 32 +``` + +### 3. Launch shards + +```commandline +conda activate tonyreasoningtraces && cd /nlp/scr4/nlp/crfm/text2image/text2image-rlhf/reasoning/virtual-world-data +``` + +Dry run first to verify commands: + +```bash +python -m sdg.launcher --config sdg/configs/qwen3_4b_openthoughts4_math30K_instill_n4_valredundancy3_round1.yaml launch --dry-run +``` + +Then launch for real: + +```bash +python -m sdg.launcher --config sdg/configs/qwen3_4b_openthoughts4_math30K_instill_n4_valredundancy3_round1.yaml launch +``` + +This reads `num_shards` from the config and submits that many nlprun jobs (1 GPU each). Each job processes `ceil(total_seeds / num_shards)` examples and writes to `output/{experiment_name}/shards/shard_NNN/`. You can override with `--num-shards N` on the CLI. + +### 4. Check status + +```bash +python -m sdg.launcher --config sdg/configs/qwen3_4b_openthoughts4_math30K_instill_n4_valredundancy3_round1.yaml check +``` + +Prints a table like: + +``` +Shard Status Seeds Passed Failed +000 completed 2930 2100 830 +001 running - - - +002 failed - - - +... +Overall: 20/32 completed, 8/32 running, 4/32 failed +``` + +When all shards are completed, `check` automatically: +- Combines all shard `output.jsonl` files into `output/{experiment_name}/output.jsonl` +- Aggregates all `stats.json` into a combined `output/{experiment_name}/stats.json` +- Uploads to HuggingFace at `teetone/{experiment_name}` + +### 5. Rerun failed shards + +```bash +python -m sdg.launcher --config sdg/configs/qwen3_4b_openthoughts4_math30K_instill_n4_valredundancy3_round1.yaml rerun +``` + +This re-launches only `failed` shards (skips `completed`, `running`, and `not_started`). + +## Output Directory Structure + +``` +output/{experiment_name}/ + launch_meta.json # Shard metadata + config.yaml # Config copy + output.jsonl # Combined output (after all shards done) + stats.json # Aggregated stats (after all shards done) + logs/ + shard_000.log + shard_001.log + ... + shards/ + shard_000/ + output.jsonl + stats.json + shard_001/ + output.jsonl + stats.json + ... +``` + +## Example: OpenThoughts4 with 32 shards + +```bash +# 1. Start MongoDB on john (no GPU, 16 GB RAM) +nlprun -a tonyreasoningtraces -q john --exclude john17 -c 4 -g 0 --memory 64g \ + -w /nlp/scr4/nlp/crfm/text2image/text2image-rlhf/reasoning/virtual-world-data \ + --job-name sdg-mongodb "python -m sdg.launch_mongodb 2>&1 | tee mongo.log" + +# 2. Grab the hostname from the job output, then set mongo_uri in your config +# mongo_uri: "mongodb://:27017" + +# 3. Launch +python -m sdg.launcher --config sdg/configs/qwen3_4b_openthoughts4_math30K_instill_n4_valredundancy3_round1.yaml launch + +# 4. Monitor +python -m sdg.launcher --config sdg/configs/qwen3_4b_openthoughts4_math30K_instill_n4_valredundancy3_round1.yaml check + +# 5. Rerun any failures +python -m sdg.launcher --config sdg/configs/qwen3_4b_openthoughts4_math30K_instill_n4_valredundancy3_round1.yaml rerun +``` + +## CLI Reference + +| Command | Description | +|---------|-------------| +| `python -m sdg.generate --config CONFIG` | Run pipeline on a single GPU | +| `python -m sdg.generate --config CONFIG --shard-id I --num-shards N` | Run a single shard | +| `python -m sdg.launch_mongodb` | Start MongoDB server | +| `python -m sdg.launcher --config CONFIG launch` | Launch all shards via nlprun | +| `python -m sdg.launcher --config CONFIG launch --dry-run` | Print commands without executing | +| `python -m sdg.launcher --config CONFIG check` | Check shard status, combine if done | +| `python -m sdg.launcher --config CONFIG rerun` | Rerun failed shards only | +| `python -m sdg.launcher --config CONFIG rerun --dry-run` | Show which failed shards would be relaunched | + +## Launch Jobs + +### Qwen3 4B + +#### Launch OpenThoughts3 Math53K + +```bash +# n=1, val_redundancy=3 - DONE +python -m sdg.launcher --config sdg/configs/qwen3_4b_openthoughts3_math53K_instill_n1_valredundancy3_round1.yaml launch + +# n=1, val_redundancy=1 - DONE +python -m sdg.launcher --config sdg/configs/qwen3_4b_openthoughts3_math53K_instill_n1_valredundancy1_round1.yaml launch + +# n=8, val_redundancy=1 - DONE +python -m sdg.launcher --config sdg/configs/qwen3_4b_openthoughts3_math53K_instill_n8_valredundancy1_round1.yaml launch + +# n=8, val_redundancy=3 - DONE +python -m sdg.launcher --config sdg/configs/qwen3_4b_openthoughts3_math53K_instill_n8_valredundancy3_round1.yaml launch + +# n=1, val_redundancy=5 +python -m sdg.launcher --config sdg/configs/qwen3_4b_openthoughts3_math53K_instill_n1_valredundancy5_round1.yaml launch + +# n=4, val_redundancy=5 +python -m sdg.launcher --config sdg/configs/qwen3_4b_openthoughts3_math53K_instill_n4_valredundancy5_round1.yaml launch + +# n=8, val_redundancy=5 - RUNNING +python -m sdg.launcher --config sdg/configs/qwen3_4b_openthoughts3_math53K_instill_n8_valredundancy5_round1.yaml launch + +# n=8, no_filter (no UQ validation) +python -m sdg.launcher --config sdg/configs/qwen3_4b_openthoughts3_math53K_instill_n8_no_filter_round1.yaml launch + +# n=8, val_redundancy=5, all_valid (keep all samples that pass UQ) +python -m sdg.launcher --config sdg/configs/qwen3_4b_openthoughts3_math53K_instill_n8_valredundancy5_allvalid_round1.yaml launch +``` + +#### Launch OpenThoughts4 Science26K + +```bash +# n=4, val_redundancy=3 - DONE +python -m sdg.launcher --config sdg/configs/qwen3_4b_openthoughts4_science26K_instill_n4_valredundancy3_round1.yaml launch + +# n=1, val_redundancy=1 - DONE +python -m sdg.launcher --config sdg/configs/qwen3_4b_openthoughts4_science26K_instill_n1_valredundancy1_round1.yaml launch + +# n=8, val_redundancy=3 - DONE +python -m sdg.launcher --config sdg/configs/qwen3_4b_openthoughts4_science26K_instill_n8_valredundancy3_round1.yaml launch + +# n=8, val_redundancy=5 - DONE +python -m sdg.launcher --config sdg/configs/qwen3_4b_openthoughts4_science26K_instill_n8_valredundancy5_round1.yaml launch + +# n=8, no_filter (no UQ validation) - RUNNING +python -m sdg.launcher --config sdg/configs/qwen3_4b_openthoughts4_science26K_instill_n8_no_filter_round1.yaml launch +``` + +#### Launch OpenThoughts4 Code9K + +```bash +# n=4, val_redundancy=3 - DONE +python -m sdg.launcher --config sdg/configs/qwen3_4b_openthoughts4_code9K_instill_n4_valredundancy3_round1.yaml launch + +# n=1, val_redundancy=1 - DONE +python -m sdg.launcher --config sdg/configs/qwen3_4b_openthoughts4_code9K_instill_n1_valredundancy1_round1.yaml launch + +# n=8, val_redundancy=3 - DONE +python -m sdg.launcher --config sdg/configs/qwen3_4b_openthoughts4_code9K_instill_n8_valredundancy3_round1.yaml launch + +# n=8, val_redundancy=5 - DONE +python -m sdg.launcher --config sdg/configs/qwen3_4b_openthoughts4_code9K_instill_n8_valredundancy5_round1.yaml launch + +# n=8, no_filter (no UQ validation) - RUNNING +python -m sdg.launcher --config sdg/configs/qwen3_4b_openthoughts4_code9K_instill_n8_no_filter_round1.yaml launch +``` + +### Qwen3 8B + + +#### Launch OpenThoughts4 Science26K (Qwen3-8B) + +```bash +# n=8, val_redundancy=5 - DONE +python -m sdg.launcher --config sdg/configs/qwen3_8b_openthoughts4_science26K_instill_n8_valredundancy5_round1.yaml launch + +# n=4, val_redundancy=5 - RUNNING +python -m sdg.launcher --config sdg/configs/qwen3_8b_openthoughts4_science26K_instill_n4_valredundancy5_round1.yaml launch +``` + +#### Launch OpenThoughts4 Code9K (Qwen3-8B) + +```bash +# n=8, val_redundancy=5 - DONE +python -m sdg.launcher --config sdg/configs/qwen3_8b_openthoughts4_code9K_instill_n8_valredundancy5_round1.yaml launch +``` + +#### Launch OpenThoughts4 Math219K (Qwen3-8B) + +```bash +# n=8, val_redundancy=5 - RUNNING +python -m sdg.launcher --config sdg/configs/qwen3_8b_openthoughts4_math219K_instill_n8_valredundancy5_round1.yaml launch +``` + +# Other + +## Check logs + +```bash +scp -r tonyhlee@scdt.stanford.edu:/nlp/scr4/nlp/crfm/text2image/text2image-rlhf/reasoning/virtual-world-data/output/qwen3_0.6b_openthoughts3_math53K_instill_n8_valredundancy5_round1 /Users/tonyhlee/Dev/virtual-world-data/ +``` + diff --git a/RUNBOOK.md b/RUNBOOK.md new file mode 100644 index 0000000000000000000000000000000000000000..fbb82376ea2dfed82a461b8e7373a0c2e2e76e6b --- /dev/null +++ b/RUNBOOK.md @@ -0,0 +1,694 @@ +# Multi-Node SDG Pipeline Runbook + +End-to-end recipe for running the SDG pipeline distributed across **4 worker +nodes × 8 H100 GPUs = 32 shards**, with cross-node MongoDB caching and HuggingFace +output gathering. Tested on ByteDance Arnold (Kubernetes-backed Azure ND96isr_H100_v5). + +- Target config: `sdg/configs/qwen3_4b_openthoughts3_math53K_instill_n8_valredundancy5_round1.yaml` +- Final dataset: published to `/qwen3_4b_openthoughts3_math53K_instill_n8_valredundancy5_round1` +- Wall time: **~40–60 hours** end-to-end +- HF username placeholder used below: `fzzhang` — replace with yours + +If the cluster dies, follow this runbook top-to-bottom on a fresh allocation. + +--- + +## 0. Cluster prerequisites + +- 4 worker nodes, each with 8× H100 (80 GB), shared no-filesystem (per-pod overlay). +- NVIDIA driver supporting at least CUDA 12.4 (driver 535+ is fine). +- HF account with write access to the target dataset namespace. +- Each worker has these Arnold-provided env vars: `ARNOLD_WORKER_0_HOST`, + `ARNOLD_WORKER_0_PORT`, `MY_POD_NAME` (ends with `worker-N`). +- Cross-worker IPv6 traffic on the `ARNOLD_WORKER_*_PORT` ports is open by default — + verify in §3 below. + +--- + +## 1. Per-worker bootstrap (run on EVERY worker, in parallel) + +### 1.1 Clone the repo + +```bash +cd ~/lmc_muon # or your preferred dir +# git clone or rsync this repo into ./self-verified-distillation +cd self-verified-distillation +``` + +### 1.2 Create the conda env and install pinned dependencies + +The driver caps at CUDA 12.4, so vLLM 0.8.5 + torch 2.6.0+cu124 is the highest +stack that runs. Newer vLLM (≥0.9) ships cu126/cu130 wheels which **will not +run on driver 535**. + +```bash +conda create -n sdg python=3.12 -y +conda activate sdg + +pip install -U uv +uv pip install -r sdg/requirements.txt + +# Downgrade to cu12.4-compatible stack +uv pip uninstall vllm torch torchvision torchaudio flashinfer-python triton +uv pip install torch==2.6.0 torchvision==0.21.0 \ + --index-url https://download.pytorch.org/whl/cu124 +uv pip install "vllm==0.8.5" "transformers>=4.51.3,<5.0" + +# Sanity check — MUST print "OK 1.0" +python -c "import torch; x=torch.zeros(1).cuda(); print('OK', (x+1).item())" +``` + +### 1.3 HuggingFace login + +```bash +huggingface-cli login # paste token; or set HF_TOKEN env var +python -c "from huggingface_hub import HfApi; print(HfApi().whoami()['name'])" +# should print your HF username (e.g. fzzhang) +``` + +--- + +## 2. MongoDB on worker_0 (ONLY worker_0) + +### 2.1 Install mongod + +```bash +cd ~/lmc_muon/self-verified-distillation +curl -O https://fastdl.mongodb.org/linux/mongodb-linux-x86_64-ubuntu2204-8.0.4.tgz +tar xzf mongodb-linux-x86_64-ubuntu2204-8.0.4.tgz +export PATH="$PWD/mongodb-linux-x86_64-ubuntu2204-8.0.4/bin:$PATH" +which mongod # must print a path + +# Persist for new shells +echo "export PATH=\"$PWD/mongodb-linux-x86_64-ubuntu2204-8.0.4/bin:\$PATH\"" >> ~/.bashrc +``` + +### 2.2 Start mongod bound to all interfaces with IPv6 + +The `--ipv6` flag is mandatory; without it `--bind_ip_all` is IPv4-only and +cross-worker traffic (IPv6 fabric) cannot reach mongod. + +```bash +mkdir -p /tmp/sdg_mongo +nohup mongod --port $ARNOLD_WORKER_0_PORT --dbpath /tmp/sdg_mongo \ + --bind_ip_all --ipv6 > mongo.log 2>&1 & + +sleep 5 +grep -i "Listening on" mongo.log +# Expected: BOTH 0.0.0.0: AND [::]: + +# Local sanity +python3 -c "from pymongo import MongoClient; \ + MongoClient(f'mongodb://[::1]:$ARNOLD_WORKER_0_PORT', serverSelectionTimeoutMS=5000).admin.command('ping'); \ + print('mongo OK locally')" +``` + +--- + +## 3. Verify cross-worker MongoDB reachability + +On workers 1, 2, 3: + +```bash +python3 -c " +import os +from pymongo import MongoClient +host = os.environ['ARNOLD_WORKER_0_HOST'] +port = os.environ['ARNOLD_WORKER_0_PORT'] +uri = f'mongodb://[{host}]:{port}' +MongoClient(uri, serverSelectionTimeoutMS=5000).admin.command('ping') +print(f'mongo OK from {os.uname().nodename} → {uri}') +" +``` + +You should see `mongo OK from ...` on each. If any worker times out, the +platform is blocking inter-worker traffic on that port — do not proceed. + +--- + +## 4. Patch the config on each worker + +Point `mongo_uri` at the IPv6 MongoDB and set `num_shards=32` to use all GPUs. + +```bash +# Run on EACH worker +python3 <<'EOF' +import os, yaml +host = os.environ['ARNOLD_WORKER_0_HOST'] +port = os.environ['ARNOLD_WORKER_0_PORT'] +path = 'sdg/configs/qwen3_4b_openthoughts3_math53K_instill_n8_valredundancy5_round1.yaml' +with open(path) as f: cfg = yaml.safe_load(f) +cfg['mongo_uri'] = f'mongodb://[{host}]:{port}' +cfg['num_shards'] = 32 +with open(path, 'w') as f: yaml.dump(cfg, f, sort_keys=False) +print('updated:', cfg['mongo_uri']) +EOF +``` + +--- + +## 5. Helper scripts + +All orchestration scripts live in `scripts/` (committed to the repo, so each +worker has them after `git clone`). They are designed to be **runnable as-is**; +override behavior with env vars rather than editing the files: + +| Script | Purpose | Run on | +|---|---|---| +| `scripts/launch_shards.sh` | Launch 8 shards on this worker, one per GPU. Honors `CONFIG`, `LIMIT`, `NUM_SHARDS`. | every worker | +| `scripts/upload_my_shards.py` | Poll for this worker's 8 stats.json, then upload to staging HF repo. Honors `HF_USER`, `RUN_NAME`. | every worker | +| `scripts/combine_and_publish.py` | Poll for 4 worker-done markers, gather, combine, publish to HF (full launcher.py stats schema). Honors `HF_USER`, `RUN_NAME`, `NUM_WORKERS`, `NUM_SHARDS`. | worker_0 only | +| `scripts/monitor.sh` | Per-worker live status (GPUs, current shard progress, completion counts, mongo cache size via cross-worker IPv6). | any worker | + +Defaults target the math53K config + `HF_USER=fzzhang` + 4 workers × 32 shards. +**Why scripts/, not inline heredocs:** during the original run, terminal +autocomplete corrupted multi-line heredoc pastes (Python files came out +unparseable). Committed files bypass paste entirely — you just `git pull` and +run them. + +**Verify they're present and parse on each worker:** + +```bash +ls scripts/ +python -c "import ast; ast.parse(open('scripts/upload_my_shards.py').read()); ast.parse(open('scripts/combine_and_publish.py').read()); print('OK')" +bash -n scripts/launch_shards.sh && bash -n scripts/monitor.sh && echo "shell OK" +``` + +If you need a different HF account or a different config, export the env vars +before running (no file edits needed): + +```bash +export HF_USER=youraccount +export RUN_NAME=your_run_name_matching_your_config +export CONFIG=sdg/configs/your_config.yaml +``` + +--- + +## 6. Per-worker uploader + +`scripts/upload_my_shards.py` runs on **every worker** — sleep-polls every 2 +minutes for that worker's 8 `stats.json` to appear, then uploads each +`shards/shard_NNN/` dir to `/-staging` and writes a +`workers/worker_N.done` marker so the combiner knows that worker is finished. + +Consumes near-zero resources while waiting — safe to launch any time (before, +during, or after the shards run). No file edits needed; override `HF_USER` / +`RUN_NAME` via env if you're not using defaults (see §5). + +--- + +## 7. Combiner (worker_0 only) + +`scripts/combine_and_publish.py` runs on **worker_0 only** — sleep-polls for all +4 worker-done markers, downloads the staging repo, concatenates the 32 +`output.jsonl` files, aggregates stats with the **full launcher.py schema** +(`summary`, `checks`, `failure_breakdown`, `failure_breakdown_examples`, +`num_shards`), enforces a row-count vs `total_passed` integrity check, and +publishes to `/` with a HuggingFace dataset-card README. + +Set `HF_USER` / `RUN_NAME` / `NUM_WORKERS` / `NUM_SHARDS` via env if defaults +don't match. + +--- + +## 8. Optional: smoke test before committing 40+ hours + +Tiny end-to-end validation (~15 min on all 32 GPUs). Confirms multi-worker +plumbing without burning real time. Run **on each worker**: + +```bash +LIMIT=128 ./scripts/launch_shards.sh +``` + +Wait ~15 min, then check completion (on each worker): + +```bash +grep -l "PIPELINE COMPLETE" logs/worker_*/shard_*.log | wc -l # should be 8 +``` + +If all four workers report 8, plumbing is verified. Clear the smoke outputs +before the real run: + +```bash +# On each worker +rm -rf output/qwen3_4b_openthoughts3_math53K_instill_n8_valredundancy5_round1 +rm -rf logs/worker_* +# Optionally wipe the cache too (on worker_0) +pkill mongod && sleep 2 +rm -rf /tmp/sdg_mongo/* +# Restart mongod (see §2.2) +``` + +--- + +## 9. Production launch + +### 9.1 Launch shards on each worker + +```bash +./scripts/launch_shards.sh +``` + +### 9.2 Start the per-worker uploader on each worker + +It will sleep-poll until its 8 shards complete, then upload. + +```bash +nohup python scripts/upload_my_shards.py > upload.log 2>&1 & +disown +``` + +### 9.3 Start the combiner on worker_0 ONLY + +It will sleep-poll until all 4 worker-done markers exist on HF, then combine +and publish. + +```bash +nohup python scripts/combine_and_publish.py > combine.log 2>&1 & +disown +``` + +After this, you can disconnect all terminals. Work continues under `nohup`. + +### 9.4 (Recommended) Verify integrity once shards complete + +Before the uploader kicks off uploads, you can sanity-check each shard locally. +For `selection_policy: first_valid` (the default), `output.jsonl` line count +must equal `total_passed` in `stats.json`. Run **on each worker**: + +```bash +RUN=output/qwen3_4b_openthoughts3_math53K_instill_n8_valredundancy5_round1/shards +for d in $RUN/shard_*; do + lines=$(wc -l < $d/output.jsonl 2>/dev/null) + passed=$(python3 -c "import json; print(json.load(open('$d/stats.json'))['summary']['total_passed'])" 2>/dev/null) + echo "$(basename $d): $lines rows | total_passed=$passed" +done +``` + +Lines must match `total_passed` for every shard. The combiner re-enforces this +check globally and will `raise RuntimeError` if anything diverges — catching it +locally first is faster. + +--- + +## 10. Monitoring + +### 10.1 Per-worker live view: `scripts/monitor.sh` + +```bash +watch -n 30 ./scripts/monitor.sh +# Or if `watch` isn't installed in the container: +while true; do clear; ./scripts/monitor.sh; sleep 30; done +``` + +Each refresh shows: per-GPU util, the most-recent tqdm state for each shard +(decoded from `\r`-overwritten progress bars), per-worker completion counts, +and the cluster-wide MongoDB cache size (queried over IPv6 so it works from any +worker, not just worker_0). + +⚠ **Do not use `tail -1 | cut -cN`** to read shard progress — tqdm writes one +physical line with `\r` overwrites, so `cut -cN` shows the line's *beginning* +(which never changes) while progress advances at the *end*. `monitor.sh` +correctly does `tail -c 500 | tr '\r' '\n' | tail -1`. + +Expected sequence in the per-shard line over the run: +1. `Loading model:` / `Capturing CUDA graphs` (~2 min vLLM warmup) +2. `Processed prompts: N/512 [...]` (Stage 3 generation, batch size 64 × n=8 = 512) +3. `[Generation] N/1656 examples (...%)` (between-batch progress lines) +4. `Round 0: M pending rows` (Stage 4 UQ starts; tqdm batches now sized by pending rows) +5. `Cycle: ... Factual: ... Correctness: ...` (inside a UQ round) +6. `PIPELINE COMPLETE` (shard finished) + +To definitively confirm a shard is in UQ rather than generation, grep the log: + +```bash +grep -E "Round [0-9]+:|Cycle:|Factual:|Correctness:|\[Generation\]|PIPELINE" \ + logs/worker_0/shard_0.log | tail -20 +``` + +You'll see `[Generation] 1656/1656 examples (100.0%)` followed by `Round 0:` if +generation is done. + +### 10.2 Aggregate signal: MongoDB cache size + +`monitor.sh` already prints this. Standalone: + +```bash +python3 -c " +from pymongo import MongoClient +import os +host = os.environ['ARNOLD_WORKER_0_HOST'] +port = os.environ['ARNOLD_WORKER_0_PORT'] +print(MongoClient(f'mongodb://[{host}]:{port}').sdg_cache.inference_cache.count_documents({})) +" +``` + +Grows continuously. Expect hundreds of thousands by mid-run, well over a +million during UQ. **Use the IPv6 host, not `127.0.0.1`** — only worker_0 can +reach loopback mongod. + +### 10.3 Uploader / combiner status + +```bash +tail -f upload.log # on any worker (file lives in working dir) +tail -f combine.log # on worker_0 +``` + +The combiner prints `DONE: https://huggingface.co/datasets//` +when finished — that's the success signal for the entire pipeline. + +--- + +## 11. Troubleshooting + +### Shards exit immediately +- `Exit 127`: `mongod` (or `python`) not on PATH. `which mongod`, re-source `~/.bashrc`. +- `Qwen2Tokenizer has no attribute all_special_tokens_extended`: transformers + too new for vLLM 0.8.5. `uv pip install "transformers>=4.51.3,<5.0"`. +- `CUDA driver version is insufficient`: torch+CUDA wheel doesn't match driver. + Reinstall via §1.2 exactly. + +### MongoDB connection refused from other workers +- mongod log shows only `0.0.0.0:PORT` (no `[::]:PORT`): missing `--ipv6` flag. + `pkill mongod`, restart with `--bind_ip_all --ipv6`. + +### `monitor.sh` shows shards stuck at `0%` for hours +- Almost always a display lie, not real stuckness. `tail -1 | cut -cN` would + show the beginning of a long `\r`-overwritten line forever; the actual + progress is at the *end*. Use `scripts/monitor.sh` (uses `tail -c 500 | tr '\r' '\n' | tail -1`). +- Confirm work is happening: GPU util ≥ 80%, log file `mtime` updates, + MongoDB cache count growing. + +### `monitor.sh` shows `mongo cache: ` from a non-zero worker +- Workers other than worker_0 cannot reach `127.0.0.1:PORT`. The script uses + `ARNOLD_WORKER_0_HOST` over IPv6, but if that env var isn't set in the + current shell, the query fails. Re-source whatever sets Arnold env vars or + run from a fresh login shell. + +### Inline Python heredoc files won't parse +- Terminal autocomplete corrupts `cat > file <<'EOF' ... EOF`. Symptoms: + `unterminated string literal`, `import os, re, time Path` (extra word + appended), or scripts cut off mid-line. +- **Fix**: don't paste — use the committed files in `scripts/`. They survive + `git pull`. If you must hand-write, use `nano ` (editor paste bypasses + shell autocomplete) or base64-encode locally and `echo '...' | base64 -d > file`. + +### Trial gets reclaimed mid-run +- Shard outputs on local disk are lost (no shared FS) and the MongoDB cache at + `/tmp/sdg_mongo` dies with worker_0. +- **Mitigation**: start the uploaders in §9.2 *at launch time*, not at the + end. They sleep-poll for completion and push completed work to HF + incrementally. If a worker dies after some shards finish but before its 8 + are all done, you lose only the unfinished ones. +- Once a worker has written its `worker_N.done` marker on HF, that worker's + data is durable. The combiner can run from any machine afterward, including + a fresh trial. + +### Some shards fail mid-run +- Each shard is independent; just relaunch the failed ones. +- Find failures: `grep -L "PIPELINE COMPLETE" logs/worker_*/shard_*.log` +- Relaunch one: `CUDA_VISIBLE_DEVICES=I nohup python -m sdg.generate --config --shard-id N --num-shards 32 > logs/worker_W/shard_N.log 2>&1 &` +- MongoDB cache means most work is replayed for free on the rerun. + +### Published `stats.json` is missing `failure_breakdown` fields +- Old `combine_and_publish.py` versions wrote a simplified stats schema. + Current `scripts/combine_and_publish.py` includes the full schema. +- If you have a published dataset with the old stats, you can re-aggregate + from `gathered/shards/*/stats.json` (still present on worker_0 after the + combine) and re-upload `stats.json` — same logic as the combiner's stats + section, just runs standalone. + +--- + +## 12. After publish: verify and clean up + +### 12.1 Verify the published dataset loads + +```bash +python3 -c " +from datasets import load_dataset +ds = load_dataset('fzzhang/qwen3_4b_openthoughts3_math53K_instill_n8_valredundancy5_round1', split='train') +print('rows:', len(ds)) +print('keys:', list(ds[0].keys())) +print('roles:', [m['role'] for m in ds[0]['messages']]) +" +``` + +Should print the row count (~21k), `keys: ['messages']`, and the roles +(`['user', 'assistant']`). Confirms the HF dataset-card README correctly +auto-loads `output.jsonl`. + +### 12.2 Clean up local + intermediate state + +```bash +# On each worker — free disk if you plan to reuse the trial +rm -rf output/ logs/ gathered/ combined_output.jsonl combined_stats.json upload.log + +# On worker_0 +pkill mongod +rm -rf /tmp/sdg_mongo combine.log + +# Optionally delete the intermediate staging HF repo +HF_USER=fzzhang RUN_NAME=qwen3_4b_openthoughts3_math53K_instill_n8_valredundancy5_round1 python -c " +import os +from huggingface_hub import HfApi +repo = f\"{os.environ['HF_USER']}/{os.environ['RUN_NAME']}-staging\" +HfApi().delete_repo(repo_id=repo, repo_type='dataset') +print('deleted', repo) +" +``` + +The final HF dataset is independent of the staging repo and the MongoDB +cache — deleting them after the run is purely cosmetic. + +--- + +## Step-by-step playbook + +Every command in execution order. The "Where" column tells you which worker to +run the command on. Commands marked "every worker" should be run on workers 0, +1, 2, AND 3 (in parallel terminals to save wall time). Steps numbered with a +letter (e.g. 5a) are part of the same logical step but run on different +workers. + +### Step 1 — get the code (every worker) + +```bash +cd ~/lmc_muon # or wherever you keep the repo +git clone self-verified-distillation +cd self-verified-distillation +``` + +### Step 2 — Python env + pinned deps (every worker) + +```bash +conda create -n sdg python=3.12 -y +conda activate sdg +pip install -U uv +uv pip install -r sdg/requirements.txt +uv pip uninstall vllm torch torchvision torchaudio flashinfer-python triton +uv pip install torch==2.6.0 torchvision==0.21.0 \ + --index-url https://download.pytorch.org/whl/cu124 +uv pip install "vllm==0.8.5" "transformers>=4.51.3,<5.0" +python -c "import torch; x=torch.zeros(1).cuda(); print('OK', (x+1).item())" +# Must print: OK 1.0 +``` + +### Step 3 — HuggingFace login (every worker) + +```bash +huggingface-cli login # paste your token +python -c "from huggingface_hub import HfApi; print(HfApi().whoami()['name'])" +# Must print your HF username +``` + +### Step 4 — install mongod (worker_0 ONLY) + +```bash +cd ~/lmc_muon/self-verified-distillation +curl -O https://fastdl.mongodb.org/linux/mongodb-linux-x86_64-ubuntu2204-8.0.4.tgz +tar xzf mongodb-linux-x86_64-ubuntu2204-8.0.4.tgz +export PATH="$PWD/mongodb-linux-x86_64-ubuntu2204-8.0.4/bin:$PATH" +echo "export PATH=\"$PWD/mongodb-linux-x86_64-ubuntu2204-8.0.4/bin:\$PATH\"" >> ~/.bashrc +which mongod # must print a path +``` + +### Step 5 — start mongod (worker_0 ONLY) + +```bash +mkdir -p /tmp/sdg_mongo +nohup mongod --port $ARNOLD_WORKER_0_PORT --dbpath /tmp/sdg_mongo \ + --bind_ip_all --ipv6 > mongo.log 2>&1 & +sleep 5 +grep -i "Listening on" mongo.log +# Must show BOTH 0.0.0.0: AND [::]: + +python3 -c "from pymongo import MongoClient; \ + MongoClient(f'mongodb://[::1]:$ARNOLD_WORKER_0_PORT', serverSelectionTimeoutMS=5000).admin.command('ping'); \ + print('mongo OK locally')" +``` + +### Step 6 — verify cross-worker mongo reachability (workers 1, 2, 3) + +```bash +python3 -c " +import os +from pymongo import MongoClient +host = os.environ['ARNOLD_WORKER_0_HOST'] +port = os.environ['ARNOLD_WORKER_0_PORT'] +uri = f'mongodb://[{host}]:{port}' +MongoClient(uri, serverSelectionTimeoutMS=5000).admin.command('ping') +print(f'mongo OK from {os.uname().nodename} → {uri}') +" +``` + +All three must print `mongo OK from ...`. If any times out, the platform is +blocking inter-worker traffic — stop here. + +### Step 7 — patch the config (every worker) + +```bash +python3 <<'EOF' +import os, yaml +host = os.environ['ARNOLD_WORKER_0_HOST'] +port = os.environ['ARNOLD_WORKER_0_PORT'] +path = 'sdg/configs/qwen3_4b_openthoughts3_math53K_instill_n8_valredundancy5_round1.yaml' +with open(path) as f: cfg = yaml.safe_load(f) +cfg['mongo_uri'] = f'mongodb://[{host}]:{port}' +cfg['num_shards'] = 32 +with open(path, 'w') as f: yaml.dump(cfg, f, sort_keys=False) +print('updated:', cfg['mongo_uri']) +EOF +``` + +### Step 8 — verify scripts/ are intact (every worker) + +```bash +python -c "import ast; ast.parse(open('scripts/upload_my_shards.py').read()); ast.parse(open('scripts/combine_and_publish.py').read()); print('OK')" +bash -n scripts/launch_shards.sh && bash -n scripts/monitor.sh && echo "shell OK" +``` + +### Step 9 — (optional) smoke test (every worker) + +```bash +LIMIT=128 ./scripts/launch_shards.sh +# Wait ~15-20 min, then check: +grep -l "PIPELINE COMPLETE" logs/worker_*/shard_*.log | wc -l # must be 8 +``` + +If smoke passes, clear outputs before the real run: + +```bash +# Every worker: +rm -rf output/qwen3_4b_openthoughts3_math53K_instill_n8_valredundancy5_round1 +rm -rf logs/worker_* +# Worker_0 only — wipe cache and restart mongod: +pkill mongod && sleep 2 && rm -rf /tmp/sdg_mongo/* +# Then redo Step 5 on worker_0 +``` + +### Step 10 — launch the real run (every worker) + +```bash +./scripts/launch_shards.sh +# Expected: 8 PIDs printed. ~5 min later, all 8 GPUs should be at ~95% util. +``` + +### Step 11 — start the per-worker uploader (every worker) + +Safe to launch immediately — it sleep-polls until the worker's shards finish. + +```bash +nohup python scripts/upload_my_shards.py > upload.log 2>&1 & +disown +``` + +### Step 12 — start the combiner (worker_0 ONLY) + +```bash +nohup python scripts/combine_and_publish.py > combine.log 2>&1 & +disown +``` + +### Step 13 — monitor (any worker, anytime) + +```bash +watch -n 30 ./scripts/monitor.sh +# Or, if `watch` isn't installed: +while true; do clear; ./scripts/monitor.sh; sleep 30; done +``` + +After this you can disconnect all terminals. The shards, uploaders, and +combiner all run under `nohup`. + +### Step 14 — (recommended) per-shard integrity check after shards finish (every worker) + +When `monitor.sh` shows `stats.json: 8/8` and `PIPELINE COMPLETE: 8/8` on a +worker, sanity-check its outputs: + +```bash +RUN=output/qwen3_4b_openthoughts3_math53K_instill_n8_valredundancy5_round1/shards +for d in $RUN/shard_*; do + lines=$(wc -l < $d/output.jsonl 2>/dev/null) + passed=$(python3 -c "import json; print(json.load(open('$d/stats.json'))['summary']['total_passed'])" 2>/dev/null) + echo "$(basename $d): $lines rows | total_passed=$passed" +done +``` + +`lines` must equal `total_passed` for every shard. If they diverge, the +combiner will refuse to publish — fix the broken shard before its uploader +starts pushing. + +### Step 15 — wait for `DONE` (worker_0) + +```bash +tail -f combine.log +# Watch for: DONE: https://huggingface.co/datasets// +``` + +### Step 16 — verify the published dataset loads (any machine with HF access) + +```bash +python3 -c " +from datasets import load_dataset +ds = load_dataset('fzzhang/qwen3_4b_openthoughts3_math53K_instill_n8_valredundancy5_round1', split='train') +print('rows:', len(ds)) +print('keys:', list(ds[0].keys())) +print('roles:', [m['role'] for m in ds[0]['messages']]) +" +``` + +### Step 17 — (optional) cleanup + +```bash +# Every worker +rm -rf output/ logs/ gathered/ combined_output.jsonl combined_stats.json upload.log + +# Worker_0 only +pkill mongod +rm -rf /tmp/sdg_mongo combine.log + +# Delete the intermediate staging HF repo (worker_0) +HF_USER=fzzhang \ +RUN_NAME=qwen3_4b_openthoughts3_math53K_instill_n8_valredundancy5_round1 \ +python -c " +import os +from huggingface_hub import HfApi +repo = f\"{os.environ['HF_USER']}/{os.environ['RUN_NAME']}-staging\" +HfApi().delete_repo(repo_id=repo, repo_type='dataset') +print('deleted', repo) +" +``` + +--- + +## Success signal + +The whole pipeline succeeds when `combine.log` on worker_0 prints: +``` +DONE: https://huggingface.co/datasets// +``` + +Everything else (per-shard logs, mongo cache count, GPU util, staging repo, +upload markers) is plumbing that supports that one line. diff --git a/gpu-sft/.gitignore b/gpu-sft/.gitignore new file mode 100644 index 0000000000000000000000000000000000000000..ac313b2671fca4ef0d7d7d7d073db2bfae2fc11e --- /dev/null +++ b/gpu-sft/.gitignore @@ -0,0 +1,6 @@ +__pycache__/ +*.pyc +runs/ +hf_cache/ +evalout/ +*.npy diff --git a/gpu-sft/EVAL_RESULTS.md b/gpu-sft/EVAL_RESULTS.md new file mode 100644 index 0000000000000000000000000000000000000000..fa9e71ad09bb8869f6087188e2aee2626e538821 --- /dev/null +++ b/gpu-sft/EVAL_RESULTS.md @@ -0,0 +1,50 @@ +# GPU eval results (PARTIAL) — captured 2026-09-13 before a machine reorg + +Raw per-task JSONs live on HDFS at `/mnt/hdfs/fangzhao_writable/marin-eval/results` +(durable across the reorg). This file is the accessible-anywhere (GitHub) copy of the +headline numbers so they're not gated behind HDFS access. + +**Eval stack:** driver 535, cu124 **vLLM 0.8.5 + evalchemy@7f24168**, `gpu_eval_driver.py`, +tp=8, temp 0.7 / top_p 1.0 / max_gen 32768, seeds 42..51 (AIME*/AMC23/HMMT n=10; MATH500 & +OlympiadBench n=1; LiveCodeBench* n=6). Per the port's design, GPU numbers are **not +bit-comparable** to the TPU tables (different kernels/seeding) — deltas are computed vs the +`qwen3_8b_base` row below. + +## qwen3_8b_base — baseline (the (+Δ) denominator) + +### math +| task | mean | std | seeds | +|---|---|---|---| +| MATH500 | 0.9080 | – | 1 | +| OlympiadBench | 0.6795 | – | 1 | +| AIME24 | 0.7500 | 0.0478 | 10 | +| AIME25 | 0.6467 | 0.0358 | 10 | +| AIME26 | 0.6833 | 0.0572 | 10 | +| AMC23 | 0.9425 | 0.0290 | 10 | +| HMMT | **FAILED** (exited 1 @ 11s — bug, see `_raw/HMMT.evalchemy.log`) | – | – | + +### code +| task | mean | std | seeds | +|---|---|---|---| +| LiveCodeBench | 0.6181 | 0.0230 | 6 | +| LiveCodeBenchv5_official | not finished (killed mid-run) | – | – | +| LiveCodeBenchv6_official | pending | – | – | + +## Gate 2 (eval parity) — PASS +GPU AIME24 **0.750 [±0.048]** vs TPU reference **0.740 [±0.049]** → within ±1σ. GPU baseline +runs slightly higher than the TPU reference across the board (MATH500 90.8 vs 89.2, +OlympiadBench 68.0 vs 62.5, AIME25 64.7 vs 60.0, AIME26 68.3 vs 62.0) — a consistent, small +upward shift consistent with the known kernel/seeding differences (not bit-comparable, as +designed). The eval stack is validated. + +## Still pending (resume on the next machine) +- **baseline:** HMMT (fix the quick failure), LiveCodeBench v5/v6. +- **experiments (needed for the deltas):** `hs_competition_nofilter` (math), + `coding_depth4v3_nofilter` (code). A last-hour salvage of hs AIME24 may have banked — + append here if it did. +- then `claude/compile_results_local.py` over the HDFS results → `results.md` with `(+Δ)`. + +## Trained models (durable on HDFS, final step-1999) +- `/mnt/hdfs/fangzhao_writable/marin-sft/coding_depth4v3_nofilter/hf/step-1999` — DONE (2000 steps) +- `/mnt/hdfs/fangzhao_writable/marin-sft/hs_competition_nofilter/hf/step-1999` — DONE (2000 steps) +- `science_depth4v3_nofilter` — re-running fresh (`_v3`) on a separate machine diff --git a/gpu-sft/README.md b/gpu-sft/README.md new file mode 100644 index 0000000000000000000000000000000000000000..d52cf5ed66eb9bbe5cfb823071360f37e9449fed --- /dev/null +++ b/gpu-sft/README.md @@ -0,0 +1,47 @@ +# gpu-sft — Qwen3 self-instill SFT + eval (TPU→GPU port) + +Downstream of the SDG pipeline: **train Qwen3 on the self-instill datasets** (the +`fzzhang/qwen3_8b_*_instill_...` sets our SDG runs produced) → **eval** on +math/science/code → regenerate **`results.md`**. + +## Provenance +All scripts here were **extracted verbatim** from Tony's handoff doc +`fangzhao.md` (Appendix A), via a fenced-code parser — not hand-typed. +**They have NOT been run on a GPU** (written on macOS, no CUDA). Treat as a strong +first draft to smoke-test on one node; the acceptance tests are the two parity +gates in §10 of the handoff. Do not trust any experiment number until Gate 1 +(training-loss parity) and Gate 2 (eval-baseline parity) pass. + +## Pipeline +``` +HF dataset@rev → scripts/gpu_sft/prepare_sft_data.py (CPU) → packed .npy + assistant mask + → scripts/gpu_sft/train_sft_qwen3.py (GPU FSDP) → hf/step-N/ (bf16) + → scripts/gpu_eval/gpu_eval_driver.py (vLLM+evalchemy) → eval JSON + → claude/compile_results_local.py → results.md +``` + +## Environment (two venvs — mandatory; see handoff §5.1) +- `sft-train`: torch 2.7 cu124 + transformers 4.57 + flash-attn (this dir's `requirements-gpu.txt`). +- `sft-eval`: vLLM (its own torch) + `datasets>=3.1,<4` + pinned evalchemy (via `scripts/gpu_eval/bootstrap_gpu_eval.sh`). +- `HF_HOME` on the big local disk, NOT `$HOME`. Pin `datasets<4.0` in both (v4 breaks LiveCodeBench). +- Gated licences: the `fzzhang` HF account must accept `cais/hle` AND `Idavidrein/gpqa` before any eval. + +## Storage plan (this cluster: Azure ND96isr H100 v5, 4 nodes × 8×H100-80G, not preemptible) +- **No HDFS, no cross-node FS.** 8B@32K fits on ONE node → run **4 independent single-node experiments**, one per node. +- Working dir + `HF_HOME` + active checkpoints → **local NVMe** (`/`, ~3.3 TB free; one 8B experiment ~2 TB fits). +- Durable artifact = **bf16 `hf/step-N` exports → private HF repo** `fzzhang/-stepN` (existing workflow; durable + eval-anywhere). +- Tradeoff: HF export drops the fp32 optimizer state, so no bit-exact training resume. Given not-preemptible + ~1.5-day runs, accept re-run risk. Revisit HDFS write-perm only if bit-exact auto-resume is needed. + +## Status +- **Phase 0 (zero-GPU): DONE** — all 8 Python files compile; the `results.md` compiler is validated end-to-end via `claude/make_fixture_eval_tree.py` → `claude/compile_results_local.py` (correct §9 format). +- **Phase 1 (1 node): pending** — two venvs + gated-licence preflight → `prepare_sft_data.py` on one dataset (with the §10 sanity asserts) → tiny smoke train → parity gate (`coding_depth4v3_n8_vr5`, 2000 steps) vs Tony's W&B (`tonyhlee/instill_ot4`). +- **Phase 2:** GPU baselines (Qwen3-8B) into a fresh results root → experiments. + +## Zero-GPU compiler self-test (reproduce anywhere) +``` +python claude/make_fixture_eval_tree.py /tmp/fixture +python claude/compile_results_local.py --root /tmp/fixture/evalout \ + --models /tmp/fixture/evals.md --output /tmp/fixture/results.md \ + --timestamp "2026-09-01 12:00" --no-cache +cat /tmp/fixture/results.md +``` diff --git a/gpu-sft/claude/compile_results_local.py b/gpu-sft/claude/compile_results_local.py new file mode 100644 index 0000000000000000000000000000000000000000..5466c719ecc8a7082948949ea6a7a485096a9518 --- /dev/null +++ b/gpu-sft/claude/compile_results_local.py @@ -0,0 +1,1220 @@ +#!/usr/bin/env python3 +# Copyright The Marin Authors +# SPDX-License-Identifier: Apache-2.0 + +"""Compile LOCAL evalchemy outputs into results.md. + +Local-disk port of claude/compile_results.py (the GCS/Iris version that produced +the 2897-line results.md). Same output format, byte-for-byte: same header line, +same column sets, same delta/std cell format, same Val Sigma Delta / Test Sigma +Delta formulas, same per-column star rule, same Best Checkpoints and Missing +Evals blocks. + +Only the INPUT layer changed: gcsfs -> os.scandir. See LAYOUT below. + +Usage +----- + # one eval root + python compile_results_local.py --root /data/marin/evaluation/evalchemy + + # several roots (e.g. two boxes rsynced into one tree); FIRST root wins on + # conflicting complete values, exactly like the multi-region merge did + python compile_results_local.py --root /data/a/evaluation/evalchemy \ + --root /data/b/evaluation/evalchemy + + # deterministic regeneration (frozen timestamp) for byte-diffing + python compile_results_local.py --root ... --timestamp "2026-08-08 09:19" + + # only report what is not yet evaluated + python compile_results_local.py --root ... --missing-only + +LAYOUT +------ +/ + Qwen_Qwen3-8B/ # baseline: model id, '/' -> '_' + MATH500_seed42/**/results_.json # single-seed benchmark + AIME24_seed42/**/results_.json # per-seed dirs, seeds 42..51 + ... + AIME24_seed51/**/results_.json + compile_AIME24_avg10seeds/compiled_results/averaged_results.json # optional + -step/ # one dir per checkpoint + + +A trailing "-" on any of those directory names is tolerated +(`AIME24_seed42-1a2b3c`, `compile_AIME24_avg10seeds-009af9`) so a tree copied +down from GCS works unchanged. + +Per-seed file schema (what the reader needs): + {"results": {"": {"accuracy_avg": 0.6333}}} # or "accuracy" +Compile file schema: + [{"correct_mean": 0.6567, "correct_std": 0.0686, "num_seeds": 10}] + +Scores are fractions in [0, 1]; every displayed number is value * 100. +""" + +from __future__ import annotations + +import argparse +import datetime +import json +import os +import re +import sys +import tempfile +from concurrent.futures import ThreadPoolExecutor, as_completed +from dataclasses import dataclass, field +from decimal import ROUND_HALF_UP, Decimal +from pathlib import Path + +# Some eval result JSONs (HMMT) contain integers like (26!)^3 that exceed +# CPython's default int->str digit cap. +sys.set_int_max_str_digits(0) + + +def _round_half_up(x: float, ndigits: int = 1) -> float: + """Round half away from zero (i.e. 62.35 -> 62.4, not 62.3 like banker's).""" + quant = Decimal("1") / (Decimal("10") ** ndigits) + return float(Decimal(repr(x)).quantize(quant, rounding=ROUND_HALF_UP)) + + +# --------------------------------------------------------------------------- +# Configuration - parsed from evals.md (single source of truth) +# --------------------------------------------------------------------------- + +DEFAULT_MODELS_MD = Path(__file__).parent / "evals.md" + + +@dataclass +class ModelSizeGroup: + """A group of experiments sharing the same model size and baseline.""" + + label: str # e.g. "Qwen3 4B" + baseline: str # e.g. "Qwen/Qwen3-4B" + experiments: list[str] = field(default_factory=list) + + +def _detect_size(text: str) -> str | None: + """Detect model size like '4B', '0.6B', or '1.7B' from a string.""" + m = re.search(r"\b(\d+(?:\.\d+)?[Bb])\b", text) + return m.group(1).upper() if m else None + + +def parse_models_md(path: Path) -> tuple[list[str], list[str], list[ModelSizeGroup]]: + """Parse baselines, experiments, and model-size groups from evals.md. + + Expects a ``## Models to Evaluate`` section containing ``### Baselines`` + and one or more ``### ...`` subsections with ``- model_name`` entries. + Returns (all_baselines, all_experiments, model_size_groups). + + NOTE: group order is `int(re.sub(r"\\D", "", size))`, so 0.6B -> 6 and + 1.7B -> 17. That is why results.md orders sections 4B, 0.6B, 8B, 1.7B, 32B. + Preserved deliberately: changing it reorders the whole document. + """ + text = path.read_text() + baselines: list[str] = [] + experiments: list[str] = [] + + size_baselines: dict[str, str] = {} # "4B" -> "Qwen/Qwen3-4B" + size_experiments: dict[str, list[str]] = {} # "4B" -> [...] + + in_models_section = False + current_is_baseline = False + current_size: str | None = None + + for line in text.splitlines(): + stripped = line.strip() + + if stripped.startswith("## "): + in_models_section = stripped == "## Models to Evaluate" + continue + + if not in_models_section: + continue + + if stripped.startswith("### "): + current_is_baseline = "baseline" in stripped.lower() + current_size = None if current_is_baseline else _detect_size(stripped) + continue + + if stripped.startswith("- "): + name = stripped[2:].strip() + if not name: + continue + if current_is_baseline: + baselines.append(name) + size = _detect_size(name) + if size: + size_baselines[size] = name + else: + experiments.append(name) + size = current_size or _detect_size(name) + if size: + size_experiments.setdefault(size, []).append(name) + + groups: list[ModelSizeGroup] = [] + for size in sorted(size_baselines.keys(), key=lambda s: int(re.sub(r"\D", "", s))): + groups.append( + ModelSizeGroup( + label=f"Qwen3 {size}", + baseline=size_baselines[size], + experiments=size_experiments.get(size, []), + ) + ) + + return baselines, experiments, groups + + +# Populated in main(); module-level so the builder functions read them like the +# original script did. +BASELINES: list[str] = [] +EXPERIMENTS: list[str] = [] +MODEL_SIZE_GROUPS: list[ModelSizeGroup] = [] + + +# --------------------------------------------------------------------------- +# Benchmark suites - val benchmarks drive Val Sigma Delta, test drive Test Sigma Delta +# --------------------------------------------------------------------------- + +MATH_VAL = ["MATH500", "OlympiadBench", "AIME24", "AIME25"] +MATH_TEST = ["AIME26", "HMMT"] +MATH_COLUMNS = MATH_VAL + MATH_TEST + +BENCHMARKS_MULTISEED = ["AIME24", "AIME25", "AIME26", "HMMT"] +BENCHMARKS_SINGLE = ["MATH500", "OlympiadBench"] +ALL_BENCHMARKS = BENCHMARKS_SINGLE + BENCHMARKS_MULTISEED + +SCIENCE_VAL = ["JEEBENCH", "OLYMPIADBENCH_PHYSICS"] +SCIENCE_TEST = ["GPQA_DIAMOND", "HUMANITYS_LAST_EXAM"] +SCIENCE_COLUMNS = SCIENCE_VAL + SCIENCE_TEST + +SCIENCE_MULTISEED = ["GPQA_DIAMOND", "JEEBENCH", "HUMANITYS_LAST_EXAM", "OLYMPIADBENCH_PHYSICS"] +SCIENCE_SINGLE: list[str] = [] +ALL_SCIENCE_BENCHMARKS = SCIENCE_SINGLE + SCIENCE_MULTISEED + +CODE_VAL = ["LIVECODEBENCH"] +CODE_TEST = ["LIVECODEBENCH_V5_OFFICIAL", "LIVECODEBENCH_V6_OFFICIAL"] +CODE_COLUMNS = CODE_VAL + CODE_TEST + +CODE_MULTISEED = ["LIVECODEBENCH", "LIVECODEBENCH_V5_OFFICIAL", "LIVECODEBENCH_V6_OFFICIAL"] +CODE_SINGLE: list[str] = [] +ALL_CODE_BENCHMARKS = CODE_SINGLE + CODE_MULTISEED + +# Display column name -> evalchemy task name (the name used in directory names +# AND as the key inside results_*.json["results"]). +BENCH_TO_TASK = { + "GPQA_DIAMOND": "GPQADiamond", + "JEEBENCH": "JEEBench", + "HUMANITYS_LAST_EXAM": "HLE", + "OLYMPIADBENCH_PHYSICS": "OlympiadBench_Physics", + "LIVECODEBENCH": "LiveCodeBench", + "LIVECODEBENCH_V5_OFFICIAL": "LiveCodeBenchv5_official", + "LIVECODEBENCH_V6_OFFICIAL": "LiveCodeBenchv6_official", +} + +# Benchmarks whose evalchemy compile step re-grades with naive string equality +# and is therefore WRONG: symbolic (HMMT), tolerance-based +# (OlympiadBench_Physics), execution-based (LiveCodeBench*, JEEBench). For these +# we never read a compile file; we average per-seed accuracy_avg ourselves. +SKIP_COMPILE_BENCHMARKS = { + "HMMT", + "LIVECODEBENCH", + "LIVECODEBENCH_V5_OFFICIAL", + "LIVECODEBENCH_V6_OFFICIAL", + "JEEBENCH", + "OLYMPIADBENCH_PHYSICS", +} + +# A multiseed cell is emitted only when this many seeds are present. +BENCH_EXPECTED_SEEDS = { + "AIME24": 10, + "AIME25": 10, + "AIME26": 10, + "AMC23": 10, + "HMMT": 10, + "GPQA_DIAMOND": 3, + "JEEBENCH": 3, + "HUMANITYS_LAST_EXAM": 3, + "OLYMPIADBENCH_PHYSICS": 3, + "LIVECODEBENCH": 6, + "LIVECODEBENCH_V5_OFFICIAL": 6, + "LIVECODEBENCH_V6_OFFICIAL": 6, +} + +# Single-seed benchmarks live in a seed42 dir. +SINGLE_SEED = 42 + + +def suite_for_model(name: str) -> str: + """Which eval suite an experiment belongs to. Substring match on the name.""" + lower = name.lower() + if "science" in lower: + return "science" + if "code" in lower or "coding" in lower or "rstarcoder" in lower: + return "code" + if "nemotron_cascade2" in lower: + return "science" + return "math" + + +def benchmarks_for_model(name: str, is_baseline: bool) -> list[str]: + """All benchmarks that should be checked for a model.""" + if is_baseline: + return ALL_BENCHMARKS + ALL_SCIENCE_BENCHMARKS + ALL_CODE_BENCHMARKS + suite = suite_for_model(name) + if suite == "science": + return list(ALL_SCIENCE_BENCHMARKS) + if suite == "code": + return list(ALL_CODE_BENCHMARKS) + return list(ALL_BENCHMARKS) + + +MAX_WORKERS = 8 +DEFAULT_OUTPUT = "results.md" +RESULTS_FILE_RE = re.compile(r"results_.*\.json$") +HASH_SUFFIX = r"(?:-[0-9A-Za-z]+)?" + + +# --------------------------------------------------------------------------- +# Local filesystem helpers +# --------------------------------------------------------------------------- + + +def list_subdirs(path: str) -> list[str]: + """Basenames of immediate subdirectories of ``path``. [] if absent.""" + try: + with os.scandir(path) as it: + return sorted(e.name for e in it if e.is_dir()) + except (FileNotFoundError, NotADirectoryError, PermissionError): + return [] + + +def find_results_files(path: str) -> list[str]: + """Recursively find results_*.json under ``path``, lexicographically sorted.""" + out: list[str] = [] + for dirpath, _dirnames, filenames in os.walk(path): + for fn in filenames: + if RESULTS_FILE_RE.search(fn): + out.append(os.path.join(dirpath, fn)) + return sorted(out) + + +def read_json(path: str): + """Parse a JSON file. Returns None on missing/unreadable/unparseable.""" + try: + with open(path, encoding="utf-8") as fh: + return json.load(fh) + except (OSError, json.JSONDecodeError, ValueError): + return None + + +# --------------------------------------------------------------------------- +# Extracted-scalar cache +# --------------------------------------------------------------------------- +# +# A finished results_*.json is immutable, and LiveCodeBench files run to tens of +# MB while we only ever pull one scalar out. Cache the extracted scalar keyed by +# (path, mtime_ns, size) so a re-run after new evals land re-parses only the new +# files. Directory listings are NOT cached, so newly finished evals are always +# discovered. + + +@dataclass +class ScalarCache: + root: Path | None + hits: int = 0 + misses: int = 0 + + def _entry_path(self, path: str) -> Path: + assert self.root is not None + return self.root / (path.lstrip("/") + ".json") + + @staticmethod + def _stat_key(path: str) -> tuple[int, int] | None: + try: + st = os.stat(path) + except OSError: + return None + return (st.st_mtime_ns, st.st_size) + + @staticmethod + def _write_atomic(path: Path, content: str) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + fd, tmp = tempfile.mkstemp(dir=path.parent, prefix=".tmp-") + try: + with os.fdopen(fd, "w", encoding="utf-8") as fh: + fh.write(content) + os.replace(tmp, path) + except BaseException: + try: + os.unlink(tmp) + except OSError: + pass + raise + + def get(self, path: str, compute): + """Return ``compute()``'s small dict record for ``path``, via cache.""" + key = self._stat_key(path) + if self.root is None or key is None: + self.misses += 1 + return compute() + + entry = self._entry_path(path) + if entry.exists(): + rec = read_json(str(entry)) + if isinstance(rec, dict) and rec.get("_key") == list(key): + self.hits += 1 + return rec.get("rec") + try: + entry.unlink() + except OSError: + pass + + self.misses += 1 + rec = compute() + if rec is not None: + self._write_atomic(entry, json.dumps({"_key": list(key), "rec": rec})) + return rec + + +_CACHE = ScalarCache(root=None) + + +# --------------------------------------------------------------------------- +# Discovery +# --------------------------------------------------------------------------- + + +@dataclass +class ModelInfo: + name: str + is_baseline: bool = False + checkpoint_dir: str | None = None # only set by the checkpoint fallback + checkpoints: list[str] = field(default_factory=list) # ["step-100", ...] + + +def _step_sort_key(step: str) -> int: + m = re.search(r"\d+", step) + return int(m.group()) if m else 0 + + +def _steps_from_checkpoint_root(exp: str, checkpoint_root: str | None) -> tuple[str | None, list[str]]: + """Fallback: list /[-]/hf/step-* . + + Mirrors the GCS script's behaviour of surfacing not-yet-evaluated checkpoints + so they show up under Missing Evals. Picks the folder with the most steps + when a config hash produced more than one output dir. + """ + if not checkpoint_root: + return None, [] + best_dir, best_steps = None, [] + for name in list_subdirs(checkpoint_root): + m = re.match(r"^(.+?)(?:-[0-9a-f]{4,})?$", name) + prefix = m.group(1) if m else name + if not (prefix == exp or exp.startswith(prefix)): + continue + folder = os.path.join(checkpoint_root, name) + steps = sorted( + (d for d in list_subdirs(os.path.join(folder, "hf")) if re.search(r"\d+", d)), + key=_step_sort_key, + ) + if len(steps) > len(best_steps): + best_dir, best_steps = folder, steps + return best_dir, best_steps + + +def discover_models(root: str, checkpoint_root: str | None) -> list[ModelInfo]: + """Discover steps from eval-result directory names under one root. + + Step discovery is driven by eval dirs named ``-step`` - NOT by + checkpoint listings - so an experiment shows up as soon as it has any eval + output. ``checkpoint_root`` is only consulted for experiments with zero eval + dirs. + """ + models: list[ModelInfo] = [ModelInfo(name=b, is_baseline=True) for b in BASELINES] + + all_eval_dirs = list_subdirs(root) + + for exp in EXPERIMENTS: + pattern = re.compile(re.escape(exp) + r"-step(\d+)$") + steps = sorted( + (f"step-{m.group(1)}" for m in (pattern.match(d) for d in all_eval_dirs) if m), + key=_step_sort_key, + ) + if steps: + models.append(ModelInfo(name=exp, checkpoints=steps)) + continue + folder, ckpt_steps = _steps_from_checkpoint_root(exp, checkpoint_root) + models.append(ModelInfo(name=exp, checkpoint_dir=folder, checkpoints=ckpt_steps)) + + return models + + +def merge_models(region_models: list[list[ModelInfo]]) -> list[ModelInfo]: + """Merge per-root discovery, preserving evals.md ordering. Steps are unioned.""" + if not region_models: + return [] + by_root = [{m.name: m for m in models} for models in region_models] + merged: list[ModelInfo] = [] + for model in region_models[0]: + variants = [models[model.name] for models in by_root if model.name in models] + checkpoints = sorted({s for v in variants for s in v.checkpoints}, key=_step_sort_key) + ckpt_dir = next((v.checkpoint_dir for v in variants if v.checkpoint_dir), None) + merged.append( + ModelInfo( + name=model.name, + is_baseline=model.is_baseline, + checkpoint_dir=ckpt_dir, + checkpoints=checkpoints, + ) + ) + return merged + + +def merge_metric_maps( + root_maps: list[dict[tuple[str, str, str], float | None]], + roots: list[str], +) -> dict[tuple[str, str, str], float | None]: + """Merge per-root metrics. First root with a non-None value wins.""" + merged: dict[tuple[str, str, str], float | None] = {} + source: dict[tuple[str, str, str], str] = {} + for root, metrics in zip(roots, root_maps, strict=True): + for key, value in metrics.items(): + existing = merged.get(key) + if existing is None and value is not None: + merged[key] = value + source[key] = root + elif key not in merged: + merged[key] = None + elif existing is not None and value is not None and abs(existing - value) > 1e-12: + print( + f"WARNING: Conflicting complete results for {key}: " + f"{source[key]}={existing}, {root}={value}; keeping {source[key]}", + file=sys.stderr, + ) + return merged + + +# --------------------------------------------------------------------------- +# Reading +# --------------------------------------------------------------------------- + + +def eval_dir_name(model: ModelInfo, step: str) -> str: + """Directory name under holding one model/checkpoint's eval output.""" + if model.is_baseline: + return model.name.replace("/", "_") + m = re.search(r"\d+", step) + return f"{model.name}-step{m.group()}" if m else model.name + + +def _pick_results_file(files: list[str], newest: bool) -> str | None: + """Choose among several results_.json in one seed dir. + + Names embed an ISO timestamp, so lexicographic order is chronological. + ``newest=True`` (default) takes the last re-run; the GCS script took the + first. Only matters when a seed dir was evaluated more than once. + """ + if not files: + return None + return files[-1] if newest else files[0] + + +def read_seed_result(seed_dir: str, task: str, newest: bool) -> float | None: + """Read one seed's accuracy from results_*.json under ``seed_dir``.""" + results_file = _pick_results_file(find_results_files(seed_dir), newest) + if results_file is None: + return None + + def _compute(): + d = read_json(results_file) + if d is None: + return None + try: + r = d["results"][task] + except (KeyError, TypeError): + return None + val = r.get("accuracy_avg", r.get("accuracy")) + return None if val is None else {"value": val} + + rec = _CACHE.get(results_file, _compute) + return rec["value"] if rec is not None else None + + +def read_compile_result(compile_dir: str, expected_seeds: int | None) -> tuple[float, float] | None: + """Read (mean, std) from /compiled_results/averaged_results.json.""" + path = os.path.join(compile_dir, "compiled_results", "averaged_results.json") + + def _compute(): + d = read_json(path) + if d is None: + return None + try: + return { + "mean": d[0]["correct_mean"], + "std": d[0].get("correct_std", 0.0), + "num_seeds": d[0].get("num_seeds", 0), + } + except (KeyError, IndexError, TypeError): + return None + + rec = _CACHE.get(path, _compute) + if rec is None: + return None + if expected_seeds is not None and rec.get("num_seeds", 0) < expected_seeds: + print( + f"WARNING: Incomplete seeds in {compile_dir}: {rec.get('num_seeds', 0)}/{expected_seeds}", + file=sys.stderr, + ) + return None + return (rec["mean"], rec["std"]) + + +def _mean_and_sample_std(values: list[float]) -> tuple[float, float | None]: + """Unweighted mean and SAMPLE (ddof=1) standard deviation.""" + mean = sum(values) / len(values) + if len(values) < 2: + return mean, None + variance = sum((v - mean) ** 2 for v in values) / (len(values) - 1) + return mean, variance**0.5 + + +@dataclass +class CellSources: + """Where one (model, step, benchmark) cell can be read from.""" + + compile_dir: str | None = None + seed_dirs: dict[int, str] = field(default_factory=dict) # seed -> dir + + +def _scan_eval_dir(eval_path: str, benchmarks: list[str]) -> dict[str, CellSources]: + """Index one eval dir's subdirectories by benchmark.""" + entries = list_subdirs(eval_path) + out: dict[str, CellSources] = {} + for bench in benchmarks: + task = BENCH_TO_TASK.get(bench, bench) + src = CellSources() + seed_re = re.compile(rf"^{re.escape(task)}_seed(\d+){HASH_SUFFIX}$") + compile_re = re.compile(rf"^compile_{re.escape(task)}_avg(\d+)seeds{HASH_SUFFIX}$") + expected = BENCH_EXPECTED_SEEDS.get(bench) + best_compile_n = -1 + for name in entries: + m = seed_re.match(name) + if m: + src.seed_dirs[int(m.group(1))] = os.path.join(eval_path, name) + continue + m = compile_re.match(name) + if m: + n = int(m.group(1)) + # Prefer a compile dir whose seed count matches expectations; + # otherwise keep the largest. + score = 10**6 if (expected is not None and n == expected) else n + if score > best_compile_n: + best_compile_n = score + src.compile_dir = os.path.join(eval_path, name) + if src.compile_dir or src.seed_dirs: + out[bench] = src + return out + + +def gather_results( + models: list[ModelInfo], + root: str, + pool: ThreadPoolExecutor, + per_seed_only: bool, + newest_results_file: bool, +) -> tuple[dict[tuple[str, str, str], float | None], dict[tuple[str, str, str], float | None]]: + """Read every (model, step, benchmark) cell under one root. + + Returns (results, std_devs), each {(model_name, step_label, benchmark): value|None}. + Baseline rows use the empty string as the step label. + """ + # Pass 1: index directories. Baselines merge the underscore and hyphen + # variants of the model id (Qwen_Qwen3-8B and Qwen-Qwen3-8B). + sources: dict[tuple[str, str, str], CellSources] = {} + results: dict[tuple[str, str, str], float | None] = {} + std_devs: dict[tuple[str, str, str], float | None] = {} + + for model in models: + benchmarks = benchmarks_for_model(model.name, model.is_baseline) + steps = [""] if model.is_baseline else model.checkpoints + for step in steps: + dir_names = [eval_dir_name(model, step)] + if model.is_baseline: + alt = model.name.replace("/", "-") + if alt != dir_names[0]: + dir_names.append(alt) + per_bench: dict[str, CellSources] = {} + for dn in dir_names: + for bench, src in _scan_eval_dir(os.path.join(root, dn), benchmarks).items(): + merged = per_bench.setdefault(bench, CellSources()) + merged.compile_dir = merged.compile_dir or src.compile_dir + merged.seed_dirs.update(src.seed_dirs) + for bench in benchmarks: + key = (model.name, step, bench) + results[key] = None + std_devs[key] = None + if bench in per_bench: + sources[key] = per_bench[bench] + + # Pass 2: single-seed benchmarks (MATH500, OlympiadBench) - one file, no std. + single_futs = {} + for key, src in sources.items(): + _, _, bench = key + if bench not in BENCHMARKS_SINGLE: + continue + seed_dir = src.seed_dirs.get(SINGLE_SEED) or ( + sorted(src.seed_dirs.items())[0][1] if src.seed_dirs else None + ) + if seed_dir is None: + continue + task = BENCH_TO_TASK.get(bench, bench) + single_futs[pool.submit(read_seed_result, seed_dir, task, newest_results_file)] = key + for fut in as_completed(single_futs): + key = single_futs[fut] + results[key] = fut.result() + std_devs[key] = None + + # Pass 3: multiseed benchmarks read from a compile dir when one is usable. + compile_futs = {} + for key, src in sources.items(): + _, _, bench = key + if bench in BENCHMARKS_SINGLE or bench in SKIP_COMPILE_BENCHMARKS: + continue + if per_seed_only or not src.compile_dir: + continue + expected = BENCH_EXPECTED_SEEDS.get(bench) + compile_futs[pool.submit(read_compile_result, src.compile_dir, expected)] = key + for fut in as_completed(compile_futs): + key = compile_futs[fut] + got = fut.result() + if got is not None: + results[key], std_devs[key] = got + + # Pass 4: everything still unread is averaged from per-seed accuracy_avg with + # a SAMPLE (n-1) std. Requires exactly the expected number of seeds; a + # partial sweep stays missing rather than reporting a biased mean. + seed_futs: dict[object, tuple[tuple[str, str, str], int]] = {} + seed_counts: dict[tuple[str, str, str], int] = {} + for key, src in sources.items(): + _, _, bench = key + if bench in BENCHMARKS_SINGLE or results.get(key) is not None: + continue + if not src.seed_dirs: + continue + task = BENCH_TO_TASK.get(bench, bench) + seed_counts[key] = len(src.seed_dirs) + for seed, seed_dir in sorted(src.seed_dirs.items()): + seed_futs[pool.submit(read_seed_result, seed_dir, task, newest_results_file)] = (key, seed) + + seed_values: dict[tuple[str, str, str], list[float]] = {} + for fut in as_completed(seed_futs): + key, _seed = seed_futs[fut] + val = fut.result() + if val is not None: + seed_values.setdefault(key, []).append(val) + + for key, n_dirs in seed_counts.items(): + _, _, bench = key + values = seed_values.get(key, []) + expected = BENCH_EXPECTED_SEEDS.get(bench, n_dirs) + if len(values) == expected: + mean, std = _mean_and_sample_std(values) + results[key] = mean + std_devs[key] = std + else: + results[key] = None + std_devs[key] = None + if values: + print( + f"WARNING: {key[0]} {key[1] or '-'} {bench}: {len(values)}/{expected} seeds; " + "reporting as missing", + file=sys.stderr, + ) + + return results, std_devs + + +# --------------------------------------------------------------------------- +# Computed columns +# --------------------------------------------------------------------------- + + +def compute_sigma( + model_name: str, + step: str, + results: dict[tuple[str, str, str], float | None], + baseline: dict[str, float | None], + benchmarks: list[str], +) -> float | None: + """Sum of (value - baseline) over ``benchmarks``, in percentage points. + + UNROUNDED - used for star selection and Best-Checkpoints argmax so nearly + tied steps are not collapsed by rounding. None if any benchmark is missing. + """ + if not benchmarks: + return None + total = 0.0 + for b in benchmarks: + v = results.get((model_name, step, b)) + bv = baseline.get(b) + if v is None or bv is None: + return None + total += (v - bv) * 100 + return total + + +def compute_display_sigma( + model_name: str, + step: str, + results: dict[tuple[str, str, str], float | None], + baseline: dict[str, float | None], + benchmarks: list[str], +) -> float | None: + """Display-only Sigma: sum of the per-cell ROUNDED diffs, then rounded again. + + This is what gets printed, so a reader adding up the visible cell deltas + reproduces the Sigma exactly. + """ + if not benchmarks: + return None + total = 0.0 + for b in benchmarks: + v = results.get((model_name, step, b)) + bv = baseline.get(b) + if v is None or bv is None: + return None + total += _round_half_up(v * 100, 1) - _round_half_up(bv * 100, 1) + return _round_half_up(total, 1) + + +def compute_stars( + model: ModelInfo, + results: dict[tuple[str, str, str], float | None], + val_sigmas: dict[str, float | None], + test_sigmas: dict[str, float | None], + columns: list[str], +) -> set[tuple[str, str]]: + """(step, column) pairs that get a star. + + Per-column, per-experiment best value. Never across experiments, never on a + baseline row. Needs >= 2 steps with a value in that column, compares the + RAW float, and stars every step tied at the max. + """ + stars: set[tuple[str, str]] = set() + + for col in columns: + non_none = [(s, results.get((model.name, s, col))) for s in model.checkpoints] + non_none = [(s, v) for s, v in non_none if v is not None] + if len(non_none) < 2: + continue + best = max(v for _, v in non_none) + stars.update((s, col) for s, v in non_none if v == best) + + for label, sigmas in (("Val Σ Δ", val_sigmas), ("Test Σ Δ", test_sigmas)): + non_none = [(s, sigmas.get(s)) for s in model.checkpoints] + non_none = [(s, v) for s, v in non_none if v is not None] + if len(non_none) < 2: + continue + best = max(v for _, v in non_none) + stars.update((s, label) for s, v in non_none if v == best) + + return stars + + +# --------------------------------------------------------------------------- +# Formatting +# --------------------------------------------------------------------------- + + +def _fmt_cell(val: float | None, baseline_val: float | None, is_baseline: bool, std: float | None = None) -> str: + """One table cell: '88.4', '65.7 [6.9]', '61.3 (-4.4)', '66.7 (+1.0) [4.4]', '❌'.""" + if val is None: + return "❌" + pct_rounded = _round_half_up(val * 100, 1) + pct = f"{pct_rounded:.1f}" + if is_baseline or baseline_val is None: + if std is not None: + return f"{pct} [{_round_half_up(std * 100, 1):.1f}]" + return pct + # Delta is the difference of the two ALREADY-ROUNDED percentages. + diff = _round_half_up(pct_rounded - _round_half_up(baseline_val * 100, 1), 1) + sign = "+" if diff >= 0 else "" + if std is not None: + return f"{pct} ({sign}{diff:.1f}) [{_round_half_up(std * 100, 1):.1f}]" + return f"{pct} ({sign}{diff:.1f})" + + +def _fmt_step(step: str) -> str: + m = re.search(r"\d+", step) + return m.group() if m else step + + +def _fmt_sigma(ts: float | None) -> str: + if ts is None: + return "N/A" + ts_rounded = _round_half_up(ts, 1) + sign = "+" if ts_rounded >= 0 else "" + return f"{sign}{ts_rounded:.1f}" + + +def _experiment_compute_order_key(name: str) -> tuple[int, int, int, int, str]: + """Order experiments from least to most sampling/filtering compute.""" + if "nofilter" in name: + return (0, 0, 0, 0, name) + n_match = re.search(r"_n(\d+)(?:_|$)", name) + vr_match = re.search(r"_vr(\d+)(?:_|$)", name) + n_value = int(n_match.group(1)) if n_match else 10**9 + vr_value = int(vr_match.group(1)) if vr_match else 10**9 + variant_rank = 1 if "allvalid" in name else 0 + return (1, n_value, vr_value, variant_rank, name) + + +SUITE_CONFIG = { + "math": { + "title": "Math", + "columns": MATH_COLUMNS, + "val": MATH_VAL, + "test": MATH_TEST, + "all_benchmarks": ALL_BENCHMARKS, + }, + "science": { + "title": "Science", + "columns": SCIENCE_COLUMNS, + "val": SCIENCE_VAL, + "test": SCIENCE_TEST, + "all_benchmarks": ALL_SCIENCE_BENCHMARKS, + }, + "code": { + "title": "Code", + "columns": CODE_COLUMNS, + "val": CODE_VAL, + "test": CODE_TEST, + "all_benchmarks": ALL_CODE_BENCHMARKS, + }, +} + + +def _build_suite_section( + suite: str, + baseline_model: ModelInfo | None, + baseline_vals: dict[str, float | None], + experiment_models: list[ModelInfo], + results: dict[tuple[str, str, str], float | None], + std_devs: dict[tuple[str, str, str], float | None] | None = None, + heading_level: str = "###", +) -> list[str]: + """Build one suite section: main table, Best Checkpoints, Missing Evals.""" + cfg = SUITE_CONFIG[suite] + columns: list[str] = cfg["columns"] + val_benchmarks: list[str] = cfg["val"] + test_benchmarks: list[str] = cfg["test"] + all_benchmarks: list[str] = cfg["all_benchmarks"] + + lines: list[str] = [] + lines.append(f"{heading_level} {cfg['title']}") + lines.append("") + + header_cols = ["Model", "Step"] + columns + ["Val Σ Δ", "Test Σ Δ"] + lines.append("| " + " | ".join(header_cols) + " |") + lines.append("|" + "|".join(["---"] * len(header_cols)) + "|") + + if baseline_model: + cells = [ + _fmt_cell( + baseline_vals.get(b), + None, + True, + std=std_devs.get((baseline_model.name, "", b)) if std_devs else None, + ) + for b in columns + ] + cells += ["-", "-"] + lines.append(f"| {baseline_model.name} | - | " + " | ".join(cells) + " |") + + suite_models = [m for m in experiment_models if suite_for_model(m.name) == suite] + + for model in suite_models: + # Blank separator row before every experiment block (including the first). + lines.append("| | | " + " | ".join([""] * (len(columns) + 2)) + " |") + + if not model.checkpoints: + cells = ["❌"] * len(columns) + ["N/A", "N/A"] + lines.append(f"| {model.name} | - | " + " | ".join(cells) + " |") + continue + + val_sigmas: dict[str, float | None] = {} + test_sigmas: dict[str, float | None] = {} + val_display: dict[str, float | None] = {} + test_display: dict[str, float | None] = {} + for step in model.checkpoints: + val_sigmas[step] = compute_sigma(model.name, step, results, baseline_vals, val_benchmarks) + test_sigmas[step] = compute_sigma(model.name, step, results, baseline_vals, test_benchmarks) + val_display[step] = compute_display_sigma(model.name, step, results, baseline_vals, val_benchmarks) + test_display[step] = compute_display_sigma(model.name, step, results, baseline_vals, test_benchmarks) + + # Stars use the unrounded sigmas so close scores aren't tied by rounding. + stars = compute_stars(model, results, val_sigmas, test_sigmas, columns) + + for step in model.checkpoints: + cells = [] + for bench in columns: + v = results.get((model.name, step, bench)) + sd = std_devs.get((model.name, step, bench)) if std_devs else None + s = _fmt_cell(v, baseline_vals.get(bench), False, std=sd) + if (step, bench) in stars: + s = f"⭐ {s}" + cells.append(s) + + vs_s = _fmt_sigma(val_display.get(step)) + if (step, "Val Σ Δ") in stars: + vs_s = f"⭐ {vs_s}" + cells.append(vs_s) + + ts_s = _fmt_sigma(test_display.get(step)) + if (step, "Test Σ Δ") in stars: + ts_s = f"⭐ {ts_s}" + cells.append(ts_s) + + lines.append(f"| {model.name} | {_fmt_step(step)} | " + " | ".join(cells) + " |") + + # Best Checkpoints - argmax over the UNROUNDED Val Sigma Delta, all ties emitted, + # rows in compute order, no blank separator rows. + lines.append("") + lines.append(f"{heading_level}# Best Checkpoints (by Val Σ Δ)") + lines.append("") + lines.append("| " + " | ".join(header_cols) + " |") + lines.append("|" + "|".join(["---"] * len(header_cols)) + "|") + + for model in sorted(suite_models, key=lambda m: _experiment_compute_order_key(m.name)): + if not model.checkpoints: + continue + best_val: float = -1e9 + best_steps: list[str] = [] + for step in model.checkpoints: + vs = compute_sigma(model.name, step, results, baseline_vals, val_benchmarks) + if vs is None: + continue + if vs > best_val: + best_val, best_steps = vs, [step] + elif vs == best_val: + best_steps.append(step) + for step in best_steps: + cells = [ + _fmt_cell( + results.get((model.name, step, b)), + baseline_vals.get(b), + False, + std=std_devs.get((model.name, step, b)) if std_devs else None, + ) + for b in columns + ] + cells.append(_fmt_sigma(compute_display_sigma(model.name, step, results, baseline_vals, val_benchmarks))) + cells.append(_fmt_sigma(compute_display_sigma(model.name, step, results, baseline_vals, test_benchmarks))) + lines.append(f"| {model.name} | {_fmt_step(step)} | " + " | ".join(cells) + " |") + + # Missing Evals - in evals.md order, not compute order. + lines.append("") + lines.append(f"{heading_level}# Missing Evals") + lines.append("") + any_missing = False + for model in suite_models: + for step in model.checkpoints: + missing = [b for b in all_benchmarks if results.get((model.name, step, b)) is None] + if missing: + lines.append(f"- {model.name} step-{_fmt_step(step)} (missing: {', '.join(missing)})") + any_missing = True + if not any_missing: + lines.append("All evaluations complete!") + + return lines + + +def build_results_md( + models: list[ModelInfo], + results: dict[tuple[str, str, str], float | None], + std_devs: dict[tuple[str, str, str], float | None] | None = None, + timestamp: str | None = None, +) -> str: + """Render the whole document. Exactly one trailing newline.""" + lines: list[str] = [] + + # Header counts EXCLUDE baselines and count experiment x step x suite-benchmark. + total_cells = 0 + complete_cells = 0 + for m in models: + if m.is_baseline: + continue + benchmarks = benchmarks_for_model(m.name, m.is_baseline) + for s in m.checkpoints: + for b in benchmarks: + total_cells += 1 + if results.get((m.name, s, b)) is not None: + complete_cells += 1 + + stamp = timestamp or datetime.datetime.now().strftime("%Y-%m-%d %H:%M") + lines.append("# Evalchemy Evaluation Results") + lines.append("") + lines.append("") + lines.append(f"*{complete_cells} of {total_cells} benchmark results gathered. Generated at {stamp}.*") + + model_by_name: dict[str, ModelInfo] = {m.name: m for m in models} + + for group in MODEL_SIZE_GROUPS: + lines.append("") + lines.append(f"## {group.label}") + + baseline_model = model_by_name.get(group.baseline) + baseline_vals: dict[str, float | None] = {} + if baseline_model: + for bench in MATH_COLUMNS + SCIENCE_COLUMNS + CODE_COLUMNS: + baseline_vals[bench] = results.get((baseline_model.name, "", bench)) + + experiment_models = [model_by_name[e] for e in group.experiments if e in model_by_name] + + suites_present = {suite_for_model(m.name) for m in experiment_models} + if baseline_model: + for suite_key in ["math", "science", "code"]: + if any(baseline_vals.get(b) is not None for b in SUITE_CONFIG[suite_key]["columns"]): + suites_present.add(suite_key) + + for suite_key in ["math", "science", "code"]: + if suite_key not in suites_present: + continue + lines.append("") + lines.extend( + _build_suite_section( + suite_key, + baseline_model, + baseline_vals, + experiment_models, + results, + std_devs=std_devs, + heading_level="###", + ) + ) + + return "\n".join(lines) + "\n" + + +# --------------------------------------------------------------------------- +# Missing-eval report +# --------------------------------------------------------------------------- + + +def print_missing(models: list[ModelInfo], results: dict[tuple[str, str, str], float | None]) -> None: + for model in models: + if model.is_baseline or not model.checkpoints: + continue + benchmarks = benchmarks_for_model(model.name, model.is_baseline) + for step in model.checkpoints: + missing = [b for b in benchmarks if results.get((model.name, step, b)) is None] + if missing: + print(f"{model.name} {step} (missing: {', '.join(missing)}) [{suite_for_model(model.name)}]") + + +# --------------------------------------------------------------------------- +# Entry point +# --------------------------------------------------------------------------- + + +def main() -> None: + global BASELINES, EXPERIMENTS, MODEL_SIZE_GROUPS, _CACHE + + parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + parser.add_argument( + "--root", + dest="roots", + action="append", + required=True, + help="Local eval-output root (repeatable; first root wins on conflicts)", + ) + parser.add_argument( + "--checkpoint-root", + default=None, + help="Optional /-/hf/step-N tree, used only for experiments with no eval output yet", + ) + parser.add_argument("--models", default=str(DEFAULT_MODELS_MD), help="Markdown file with '## Models to Evaluate'") + parser.add_argument("--output", default=DEFAULT_OUTPUT, help="Output path for results.md") + parser.add_argument("--workers", type=int, default=MAX_WORKERS, help="Parallel file readers") + parser.add_argument("--timestamp", default=None, help="Freeze the header timestamp (format: 'YYYY-MM-DD HH:MM')") + parser.add_argument("--missing-only", action="store_true", help="Report missing evals; do not write results.md") + parser.add_argument( + "--per-seed-only", + action="store_true", + help=( + "Ignore compile dirs entirely and average per-seed accuracy for every multiseed benchmark " + "(requires per-seed dirs for all of them; benchmarks that only have a compile dir become missing)" + ), + ) + parser.add_argument( + "--first-results-file", + action="store_true", + help="On re-evaluated seed dirs, read the OLDEST results_*.json (GCS script behaviour) instead of the newest", + ) + parser.add_argument("--cache-dir", default=None, help="Extracted-scalar cache dir (default: /.eval_cache)") + parser.add_argument("--no-cache", action="store_true", help="Disable the scalar cache") + args = parser.parse_args() + + models_md = Path(args.models) + if not models_md.exists(): + raise SystemExit(f"Models file not found: {models_md}") + BASELINES, EXPERIMENTS, MODEL_SIZE_GROUPS = parse_models_md(models_md) + if not MODEL_SIZE_GROUPS: + raise SystemExit(f"No model-size groups parsed from {models_md} (need '## Models to Evaluate' + '### Baselines')") + + if args.no_cache: + _CACHE = ScalarCache(root=None) + else: + cache_dir = args.cache_dir or os.path.join(os.path.dirname(os.path.abspath(args.output)) or ".", ".eval_cache") + _CACHE = ScalarCache(root=Path(cache_dir)) + + roots = [os.path.abspath(r.rstrip("/")) for r in args.roots] + print(f"Roots: {', '.join(roots)}", file=sys.stderr) + print(f"Models: {len(BASELINES)} baselines, {len(EXPERIMENTS)} experiments", file=sys.stderr) + + with ThreadPoolExecutor(max_workers=args.workers) as pool: + per_root_models = [discover_models(root, args.checkpoint_root) for root in roots] + models = merge_models(per_root_models) + + per_root_results = [] + per_root_stds = [] + for root, root_models in zip(roots, per_root_models, strict=True): + # Read each root against the MERGED step list so a step discovered in + # one root is also looked for in the others. + aligned = [ + ModelInfo( + name=m.name, + is_baseline=m.is_baseline, + checkpoint_dir=m.checkpoint_dir, + checkpoints=next((x.checkpoints for x in models if x.name == m.name), m.checkpoints), + ) + for m in root_models + ] + r, s = gather_results( + aligned, + root, + pool, + per_seed_only=args.per_seed_only, + newest_results_file=not args.first_results_file, + ) + per_root_results.append(r) + per_root_stds.append(s) + + results = merge_metric_maps(per_root_results, roots) + std_devs = merge_metric_maps(per_root_stds, roots) + + if _CACHE.root is not None: + print(f"Cache: {_CACHE.hits} hits, {_CACHE.misses} reads", file=sys.stderr) + + if args.missing_only: + print_missing(models, results) + return + + md = build_results_md(models, results, std_devs, timestamp=args.timestamp) + with open(args.output, "w", encoding="utf-8") as fh: + fh.write(md) + print(f"\nWrote {args.output}", file=sys.stderr) + + +if __name__ == "__main__": + main() diff --git a/gpu-sft/claude/evals.md b/gpu-sft/claude/evals.md new file mode 100644 index 0000000000000000000000000000000000000000..8ac1f84dab58227ef7a7fa537aaa7fce0fbd5748 --- /dev/null +++ b/gpu-sft/claude/evals.md @@ -0,0 +1,38 @@ +# Local Evals + +The compiler (`claude/compile_results_local.py`) parses ONLY the +`## Models to Evaluate` section below. Everything else in this file is prose for +humans. + +Rules: + +- `### Baselines` entries are the delta reference. One baseline per model size; + the size token (`0.6B`, `1.7B`, `4B`, `8B`, `32B`) is read out of the name. + A `## Qwen3 ` section is emitted only if a baseline for that size exists. +- Every other `### ...` heading is an experiment group. The size comes from the + heading text, falling back to the experiment name. +- Each `- ` under an experiment heading must be the EXACT experiment name, + which is also the eval output directory prefix: `-step/`. Do not + truncate, do not add a `-step` suffix here. +- Row order inside a suite table follows this file. The `#### Best Checkpoints` + table re-sorts by sampling compute (`nofilter` first, then `_n` asc, + `_vr` asc, `allvalid` last). +- Suite routing is a substring match on the name: `science` -> Science; + `code` / `coding` / `rstarcoder` -> Code; otherwise Math. The suite decides + which benchmarks are read AND which table the row lands in. +- Section order across sizes is `int(re.sub(r"\D", "", size))`, so 0.6B sorts as + 6 and 1.7B as 17: 4B, 0.6B, 8B, 1.7B, 32B. Known quirk, kept for continuity + with the existing results.md. + +## Models to Evaluate + +### Baselines +- Qwen/Qwen3-4B +- Qwen/Qwen3-8B + +### SFT 4B Experiments +- exp_sft_qwen3_4b_selfinstill_ot3_math53k_n8_vr5_round1 + +### SFT 8B Experiments +- exp_sft_qwen3_8b_selfinstill_science_depth4v3_n8_vr5_2k_lr5e6_wd01 +- exp_sft_qwen3_8b_selfinstill_coding_depth4v3_n8_vr5_2k_lr5e6_wd01 diff --git a/gpu-sft/claude/make_fixture_eval_tree.py b/gpu-sft/claude/make_fixture_eval_tree.py new file mode 100644 index 0000000000000000000000000000000000000000..bed033c4e56f60dbee8897b80f9e985e1754d01b --- /dev/null +++ b/gpu-sft/claude/make_fixture_eval_tree.py @@ -0,0 +1,183 @@ +#!/usr/bin/env python3 +# Copyright The Marin Authors +# SPDX-License-Identifier: Apache-2.0 + +"""Reference writer + self-test fixture for the local results pipeline. + +Two jobs: + +1. ``write_seed_result`` / ``write_compile_result`` are the CANONICAL writers. + The local eval runner should call these (or emit byte-identical files) so + compile_results_local.py can read its output. Nothing else defines the + contract. + +2. Running this file as a script materializes a small fake eval tree plus a + matching models markdown, so the whole pipeline can be smoke-tested with no + GPU and no real eval: + + python make_fixture_eval_tree.py /tmp/fixture + python compile_results_local.py --root /tmp/fixture/evalout \\ + --models /tmp/fixture/evals.md --output /tmp/fixture/results.md \\ + --timestamp "2026-09-01 12:00" --no-cache + cat /tmp/fixture/results.md +""" + +from __future__ import annotations + +import json +import os +import sys + +# Display column name -> evalchemy task name. Must match compile_results_local.py. +BENCH_TO_TASK = { + "GPQA_DIAMOND": "GPQADiamond", + "JEEBENCH": "JEEBench", + "HUMANITYS_LAST_EXAM": "HLE", + "OLYMPIADBENCH_PHYSICS": "OlympiadBench_Physics", + "LIVECODEBENCH": "LiveCodeBench", + "LIVECODEBENCH_V5_OFFICIAL": "LiveCodeBenchv5_official", + "LIVECODEBENCH_V6_OFFICIAL": "LiveCodeBenchv6_official", +} + + +def model_dir_name(model: str, step: int | None) -> str: + """Eval output dir for a model/checkpoint. + + Baseline (HF id, step None): 'Qwen/Qwen3-8B' -> 'Qwen_Qwen3-8B'. + Experiment: '-step'. + """ + if step is None: + return model.replace("/", "_") + return f"{model}-step{step}" + + +def write_seed_result( + root: str, + model: str, + step: int | None, + task: str, + seed: int, + accuracy: float, + timestamp: str = "2026-09-01T00-00-00.000000", + extra: dict | None = None, +) -> str: + """Write one (checkpoint, task, seed) result file. ``accuracy`` is a fraction. + + Layout: //_seed//results_.json + The reader takes ``results[task]["accuracy_avg"]``, falling back to + ``["accuracy"]``, so single-pass benchmarks may emit either key. + """ + d = os.path.join(root, model_dir_name(model, step), f"{task}_seed{seed}", model.replace("/", "__")) + os.makedirs(d, exist_ok=True) + payload = {"results": {task: {"accuracy_avg": accuracy, **(extra or {})}}} + path = os.path.join(d, f"results_{timestamp}.json") + with open(path, "w", encoding="utf-8") as fh: + json.dump(payload, fh) + return path + + +def write_compile_result( + root: str, + model: str, + step: int | None, + task: str, + seeds: list[int], + mean: float, + std: float, +) -> str: + """Write the averaged-across-seeds file for one (checkpoint, task). + + Layout: //compile__avgseeds/compiled_results/averaged_results.json + + OPTIONAL. Omit it and compile_results_local.py averages the per-seed files + itself (mean + sample n-1 std), which is the recommended local mode: + evalchemy's own compile step re-grades with naive string equality and is + wrong for HMMT / JEEBench / OlympiadBench_Physics / LiveCodeBench*. + """ + n = len(seeds) + d = os.path.join(root, model_dir_name(model, step), f"compile_{task}_avg{n}seeds", "compiled_results") + os.makedirs(d, exist_ok=True) + payload = [ + { + "base_model_name": model, + "dataset_name": task, + "num_seeds": n, + "seeds": seeds, + "correct_mean": mean, + "correct_std": std, + } + ] + path = os.path.join(d, "averaged_results.json") + with open(path, "w", encoding="utf-8") as fh: + json.dump(payload, fh) + return path + + +# --------------------------------------------------------------------------- +# Fixture +# --------------------------------------------------------------------------- + +EVALS_MD = """# Local Evals + +## Models to Evaluate + +### Baselines +- Qwen/Qwen3-8B + +### SFT 8B Experiments +- exp_sft_qwen3_8b_selfinstill_mynewdata_n8_vr5_2k_lr5e6_wd01 +""" + +BASELINE = "Qwen/Qwen3-8B" +EXPERIMENT = "exp_sft_qwen3_8b_selfinstill_mynewdata_n8_vr5_2k_lr5e6_wd01" +MATH_SEEDS = list(range(42, 52)) # 10 seeds + + +def build(out_dir: str) -> None: + evalout = os.path.join(out_dir, "evalout") + os.makedirs(evalout, exist_ok=True) + with open(os.path.join(out_dir, "evals.md"), "w", encoding="utf-8") as fh: + fh.write(EVALS_MD) + + # Baseline: MATH500/OlympiadBench single-seed, AIME* via compile files, + # HMMT via per-seed files (compile is never trusted for HMMT). + write_seed_result(evalout, BASELINE, None, "MATH500", 42, 0.884) + write_seed_result(evalout, BASELINE, None, "OlympiadBench", 42, 0.617) + write_compile_result(evalout, BASELINE, None, "AIME24", MATH_SEEDS, 0.6566666666666666, 0.06858355272004873) + write_compile_result(evalout, BASELINE, None, "AIME25", MATH_SEEDS, 0.5866666666666667, 0.045460606091052) + write_compile_result(evalout, BASELINE, None, "AIME26", MATH_SEEDS, 0.5933333333333334, 0.05426322800056239) + for seed, acc in zip(MATH_SEEDS, [0.38, 0.40, 0.42, 0.36, 0.44, 0.38, 0.40, 0.42, 0.34, 0.39], strict=True): + write_seed_result(evalout, BASELINE, None, "HMMT", seed, acc) + + steps = { + 100: dict( + math500=0.872, + olympiad=0.596, + aime24=0.650, + aime25=0.567, + aime26=0.523, + hmmt=[0.36, 0.38, 0.40, 0.34, 0.42, 0.36, 0.38, 0.40, 0.32, 0.37], + ), + 200: dict( + math500=0.906, + olympiad=0.642, + aime24=0.737, + aime25=0.620, + aime26=0.637, + hmmt=[0.41, 0.43, 0.45, 0.39, 0.47, 0.41, 0.43, 0.45, 0.37, 0.42], + ), + } + for step, d in steps.items(): + write_seed_result(evalout, EXPERIMENT, step, "MATH500", 42, d["math500"]) + write_seed_result(evalout, EXPERIMENT, step, "OlympiadBench", 42, d["olympiad"]) + write_compile_result(evalout, EXPERIMENT, step, "AIME24", MATH_SEEDS, d["aime24"], 0.0623) + write_compile_result(evalout, EXPERIMENT, step, "AIME25", MATH_SEEDS, d["aime25"], 0.0531) + write_compile_result(evalout, EXPERIMENT, step, "AIME26", MATH_SEEDS, d["aime26"], 0.0688) + for seed, acc in zip(MATH_SEEDS, d["hmmt"], strict=True): + write_seed_result(evalout, EXPERIMENT, step, "HMMT", seed, acc) + + print(f"fixture: {evalout}\nmodels: {os.path.join(out_dir, 'evals.md')}") + + +if __name__ == "__main__": + build(sys.argv[1] if len(sys.argv) > 1 else "fixture") diff --git a/gpu-sft/experiments/qwen3_chat_template.py b/gpu-sft/experiments/qwen3_chat_template.py new file mode 100644 index 0000000000000000000000000000000000000000..f16ac133979e79da3288a8957f24f8014c3c2e8e --- /dev/null +++ b/gpu-sft/experiments/qwen3_chat_template.py @@ -0,0 +1,108 @@ +# Copyright 2025 The Marin Authors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# flake8: noqa + +""" +Updated chat template for Qwen3 that has {% generation %} tags inserted (required by Levanter). +""" + +QWEN_3_CHAT_TEMPLATE = r"""{%- if tools is defined and tools %} + {{- '<|im_start|>system\n' }} + {%- if messages[0].role == 'system' %} + {{- messages[0].content + '\n\n' }} + {%- endif %} + {{- "# Tools\n\nYou may call one or more functions to assist with the user query.\n\nYou are provided with function signatures within XML tags:\n" }} + {%- for tool in tools %} + {{- "\n" }} + {{- tool | tojson }} + {%- endfor %} + {{- "\n\n\nFor each function call, return a json object with function name and arguments within XML tags:\n\n{\"name\": , \"arguments\": }\n<|im_end|>\n" }} +{%- else %} + {%- if messages[0].role == 'system' %} + {{- '<|im_start|>system\n' + messages[0].content + '<|im_end|>\n' }} + {%- endif %} +{%- endif %} +{%- set ns = namespace(multi_step_tool=true, last_query_index=messages|length - 1) %} +{%- for message in messages[::-1] %} + {%- set index = (messages|length - 1) - loop.index0 %} + {%- if ns.multi_step_tool and message.role == "user" and message.content is string and not(message.content.startswith('') and message.content.endswith('')) %} + {%- set ns.multi_step_tool = false %} + {%- set ns.last_query_index = index %} + {%- endif %} +{%- endfor %} +{%- for message in messages %} + {%- if message.content is string %} + {%- set content = message.content %} + {%- else %} + {%- set content = '' %} + {%- endif %} + {%- if (message.role == "user") or (message.role == "system" and not loop.first) %} + {{- '<|im_start|>' + message.role + '\n' + content + '<|im_end|>' + '\n' }} + {%- elif message.role == "assistant" %} + {%- set reasoning_content = '' %} + {%- if message.reasoning_content is string %} + {%- set reasoning_content = message.reasoning_content %} + {%- else %} + {%- if '' in content %} + {%- set reasoning_content = content.split('')[0].rstrip('\n').split('')[-1].lstrip('\n') %} + {%- set content = content.split('')[-1].lstrip('\n') %} + {%- endif %} + {%- endif %} + {%- if loop.index0 > ns.last_query_index %} + {%- if loop.last or (not loop.last and reasoning_content) %} + {{- '<|im_start|>' + message.role + '\n' }}{% generation %}{{- '\n' + reasoning_content.strip('\n') + '\n\n\n' + content.lstrip('\n') }}{% endgeneration %} + {%- else %} + {{- '<|im_start|>' + message.role + '\n' }}{% generation %}{{- content }}{% endgeneration %} + {%- endif %} + {%- else %} + {{- '<|im_start|>' + message.role + '\n' }}{% generation %}{{- content }}{% endgeneration %} + {%- endif %} + {%- if message.tool_calls %} + {%- for tool_call in message.tool_calls %} + {%- if (loop.first and content) or (not loop.first) %} + {% generation %}{{- '\n' }}{% endgeneration %} + {%- endif %} + {%- if tool_call.function %} + {%- set tool_call = tool_call.function %} + {%- endif %} + {% generation %}{{- '\n{"name": "' }} + {{- tool_call.name }} + {{- '", "arguments": ' }} + {%- if tool_call.arguments is string %} + {{- tool_call.arguments }} + {%- else %} + {{- tool_call.arguments | tojson }} + {%- endif %} + {{- '}\n' }}{% endgeneration %} + {%- endfor %} + {%- endif %} + {% generation %}{{- '<|im_end|>\n' }}{% endgeneration %} + {%- elif message.role == "tool" %} + {%- if loop.first or (messages[loop.index0 - 1].role != "tool") %} + {{- '<|im_start|>user' }} + {%- endif %} + {{- '\n\n' }} + {{- content }} + {{- '\n' }} + {%- if loop.last or (messages[loop.index0 + 1].role != "tool") %} + {{- '<|im_end|>\n' }} + {%- endif %} + {%- endif %} +{%- endfor %} +{%- if add_generation_prompt %} + {{- '<|im_start|>assistant\n' }} + {%- if enable_thinking is defined and enable_thinking is false %} + {% generation %}{{- '\n\n\n\n' }}{% endgeneration %} + {%- endif %} +{%- endif %}""" diff --git a/gpu-sft/requirements-gpu.txt b/gpu-sft/requirements-gpu.txt new file mode 100644 index 0000000000000000000000000000000000000000..a79f905d76a8aefeb6d3ca9a906534915698ae78 --- /dev/null +++ b/gpu-sft/requirements-gpu.txt @@ -0,0 +1,19 @@ +# Tested shape for the GPU port. Pin these; transformers in particular is load-bearing +# because the script relies on cu_seq_lens_q/k reaching flash_attn_varlen_func +# (transformers.modeling_flash_attention_utils, "Case 2" padding-free path). +--extra-index-url https://download.pytorch.org/whl/cu128 +torch==2.8.0 +transformers==4.57.5 +tokenizers>=0.22 +datasets>=3.1.0,<4.0.0 # v4 dropped dataset-script support; also what the marin eval branch pins +accelerate>=1.0.0 +safetensors>=0.4.5 +numpy>=1.26 +jinja2>=3.1.0 # only needed by prepare_sft_data.py --verify +sentencepiece +huggingface_hub>=0.26 +# FlashAttention-2 is REQUIRED for the 32K varlen path. Build takes ~20 min, or use the +# prebuilt wheel matching your torch/cuda/python/abi from the flash-attention releases page. +flash-attn==2.8.3 --no-build-isolation +# optional +wandb>=0.18 diff --git a/gpu-sft/scripts/gpu_eval/bootstrap_gpu_eval.sh b/gpu-sft/scripts/gpu_eval/bootstrap_gpu_eval.sh new file mode 100644 index 0000000000000000000000000000000000000000..9dc7b0e77218a17aeeeb7a93b59bdbb8e55aa9d6 --- /dev/null +++ b/gpu-sft/scripts/gpu_eval/bootstrap_gpu_eval.sh @@ -0,0 +1,194 @@ +#!/usr/bin/env bash +# Copyright The Marin Authors +# SPDX-License-Identifier: Apache-2.0 +# +# Bootstrap a self-contained CUDA evalchemy + vLLM environment for evaluating +# LOCAL HF-format checkpoints on a single GPU box. No GCS, no Iris, no TPU. +# +# ./bootstrap_gpu_eval.sh /opt/marin-gpu-eval +# +# Idempotent: re-running skips completed phases. Delete the phase markers under +# $ROOT/.markers to force a phase to re-run. +# +# --------------------------------------------------------------------------- +# WHY THESE PINS +# --------------------------------------------------------------------------- +# evalchemy teetone/evalchemy @ 7f24168 (2026-03-31) +# This is the pin the TPU pipeline used to produce the current results.md. +# Use it if you need numbers comparable to the existing tables. +# 460b022 (2026-05-07) is the newer pin carried on the marin-agents +# `evalfixhfauth` branch; it skips questions whose prompt >= max_model_len +# instead of crashing, and adds the TTC task variants. Set +# EVALCHEMY_COMMIT=460b022 if you hit prompt-length crashes. +# +# lm-eval stanford-crfm/lm-evaluation-harness @ d5e3391f22cde186c827674d5c3ec7c5f4fe0cab +# evalchemy is NOT pip-installable; it is a source checkout that imports +# lm_eval. This exact fork+commit is what Marin pins. Upstream lm-eval 0.5.x +# moved eval_logger and broke evalchemy's eval_tracker import. +# +# datasets >=3.6,<4.0.0 HARD upper bound. +# datasets v4 removed dataset-script support. LiveCodeBench loads +# `livecodebench/code_generation_lite` with trust_remote_code=True, which is +# a dataset script. On datasets>=4 every LCB task fails at load time. +# +# transformers >=4.51,<5 Qwen3 architecture support landed in 4.51.0. +# +# vllm >=0.10 CUDA wheel. Do NOT install tpu-inference / the +# `vllm` dependency group from lib/marin/pyproject.toml: that group pins +# jax==0.10.0 + libtpu==0.0.40 + tpu-inference==0.22.1 and is TPU-only. +# +# --------------------------------------------------------------------------- +# HUGGING FACE AUTH -- what the `evalfixhfauth` branch actually does +# --------------------------------------------------------------------------- +# The marin-agents branch name is stale/aspirational: no commit on it touches HF +# authentication. The real HF-related fixes on that branch are three, and a +# local user must reproduce all three: +# +# 1. HF_DATASETS_TRUST_REMOTE_CODE=1 in the eval process env (LiveCodeBench +# needs it). This script exports it in the generated activate snippet, and +# gpu_eval_driver.py sets it on every subprocess. +# 2. datasets<4.0.0 pin (above). +# 3. HF_TOKEN supplied as a process env var at submit time -- it is never in +# source. On TPU it was threaded through `iris job run -e HF_TOKEN ...`. +# LOCALLY YOU MUST DO THIS INSTEAD: +# huggingface-cli login # writes ~/.cache/huggingface/token +# or export HF_TOKEN=hf_xxx # per shell / systemd unit +# Required for gated datasets: cais/hle (HLE) and Idavidrein/gpqa +# (GPQADiamond) -- you must ALSO click "Agree" on those dataset pages with +# the same account. The math suite (MATH500, OlympiadBench, AIME24/25/26, +# HMMT) and the code suite (LiveCodeBench*) are public and need no token, +# though an anonymous box will hit HF rate limits during a 10-seed sweep. +# +set -euo pipefail + +ROOT="${1:-/opt/marin-gpu-eval}" +EVALCHEMY_REPO="${EVALCHEMY_REPO:-https://github.com/teetone/evalchemy.git}" +EVALCHEMY_COMMIT="${EVALCHEMY_COMMIT:-7f24168}" +PYTHON_VERSION="${PYTHON_VERSION:-3.11}" +LM_EVAL_REQ="lm-eval[math]@git+https://github.com/stanford-crfm/lm-evaluation-harness@d5e3391f22cde186c827674d5c3ec7c5f4fe0cab" + +VENV="$ROOT/venv" +EVALCHEMY_DIR="$ROOT/evalchemy" +CACHE="$ROOT/cache" +MARKERS="$ROOT/.markers" +HERE="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" + +mkdir -p "$ROOT" "$CACHE" "$MARKERS" + +log() { printf '[bootstrap] %s\n' "$*" >&2; } + +if ! command -v uv >/dev/null 2>&1; then + log "uv not found on PATH; install it with: curl -LsSf https://astral.sh/uv/install.sh | sh" + exit 1 +fi + +# -------------------------------------------------------------------------- +# 1. venv +# -------------------------------------------------------------------------- +if [[ ! -x "$VENV/bin/python" ]]; then + log "creating venv at $VENV (python $PYTHON_VERSION)" + uv venv --python "$PYTHON_VERSION" --seed "$VENV" +fi +PY="$VENV/bin/python" + +# -------------------------------------------------------------------------- +# 2. dependencies +# -------------------------------------------------------------------------- +if [[ ! -f "$MARKERS/deps" ]]; then + log "installing CUDA vLLM + lm-eval + evalchemy runtime deps" + # Driver-aware: on driver 535 (no CUDA 12.6+) invoke with + # TORCH_PRE="torch==2.6.0 torchvision==0.21.0 --index-url https://download.pytorch.org/whl/cu124" + # VLLM_SPEC="vllm==0.8.5" + # vLLM >=0.9 ships cu126/cu130 wheels that will NOT run on driver 535. + if [[ -n "${TORCH_PRE:-}" ]]; then + log "pre-installing pinned torch: $TORCH_PRE" + UV_CACHE_DIR="$CACHE/uv" uv pip install --python "$PY" $TORCH_PRE + fi + UV_CACHE_DIR="$CACHE/uv" uv pip install --python "$PY" \ + "${VLLM_SPEC:-vllm>=0.10}" \ + "transformers>=4.51,<5" \ + "datasets>=3.6,<4.0.0" \ + "fsspec[http]>=2023.1.0" \ + "accelerate>=1.8" \ + "safetensors" \ + "sentencepiece" \ + "$LM_EVAL_REQ" \ + "sympy>=1.12.1,<1.14" \ + "antlr4-python3-runtime==4.11" \ + "bespokelabs-curator==0.1.16" \ + "sqlalchemy" \ + "fire" \ + "pandas" \ + "numpy" + touch "$MARKERS/deps" +else + log "deps marker present; skipping install" +fi + +# Guard: datasets v4 silently breaks every LiveCodeBench task. +"$PY" - <<'PYCHECK' +import sys +from importlib.metadata import version +dv = version("datasets") +major = int(dv.split(".")[0]) +if major >= 4: + sys.exit( + f"datasets {dv} installed but LiveCodeBench requires <4.0.0 " + "(v4 dropped dataset-script support). Reinstall with 'datasets>=3.6,<4.0.0'." + ) +print(f"[bootstrap] datasets {dv} OK") +PYCHECK + +# -------------------------------------------------------------------------- +# 3. evalchemy checkout, pinned + patched +# -------------------------------------------------------------------------- +if [[ ! -d "$EVALCHEMY_DIR/.git" ]]; then + log "cloning evalchemy -> $EVALCHEMY_DIR" + git clone "$EVALCHEMY_REPO" "$EVALCHEMY_DIR" +fi + +log "pinning evalchemy to $EVALCHEMY_COMMIT and reverting previous patches" +git -C "$EVALCHEMY_DIR" fetch --all --quiet || true +git -C "$EVALCHEMY_DIR" checkout --quiet "$EVALCHEMY_COMMIT" +# Hard-reset the files the patcher rewrites so patching is deterministic. +git -C "$EVALCHEMY_DIR" checkout --quiet "$EVALCHEMY_COMMIT" -- eval/ + +log "applying GPU patches" +"$PY" "$HERE/patch_evalchemy_gpu.py" "$EVALCHEMY_DIR" + +# -------------------------------------------------------------------------- +# 4. activate snippet +# -------------------------------------------------------------------------- +cat > "$ROOT/activate.sh" </-step/_seed-/_0shot//results_.json + /-step/compile__avgseeds-/compiled_results/averaged_results.json + +WHAT DELIBERATELY DIFFERS + * ONE vLLM engine per task instead of one Iris job per (task, seed). + TPU/JAX cannot do per-request seeds, so Marin faked multi-seed by launching + K whole jobs with K engine seeds. CUDA vLLM *can* do per-request seeds, so + we set EVALCHEMY_N_REPEAT=K and pass ``--seed 42,42,42,42``; evalchemy's + benchmarks then use per-request seed ``42 + repetition_index``, which maps + exactly onto Marin seeds 42..42+K-1. One model load instead of K. + Consequence: numbers are NOT bit-comparable to the TPU tables (different + hardware, different kernels, per-request vs engine seeding). Rankings are. + * tensor_parallel_size defaults to the number of visible GPUs, not 32. + * max_num_seqs defaults to 32-64 per the GPU profile, not 256. 256 on a + single 80GB card guarantees KV-cache thrash and preemption at 36864 ctx. +""" + +from __future__ import annotations + +import argparse +import csv +import dataclasses +import hashlib +import json +import logging +import os +import re +import shutil +import statistics +import subprocess +import sys +import time +from datetime import datetime +from pathlib import Path +from typing import Any + +logger = logging.getLogger("gpu_eval_driver") + +# -------------------------------------------------------------------------- +# Task tables. Mirrors experiments/evals/exp_evalchemy_eval.py and +# experiments/evals/evalchemy_task_configs.py. +# -------------------------------------------------------------------------- + +SEED_BASE = 42 + +# task name (as evalchemy knows it) -> number of seeds +SUITES: dict[str, dict[str, int]] = { + "math": { + "MATH500": 1, + "OlympiadBench": 1, + "AIME24": 10, + "AIME25": 10, + "AIME26": 10, + "AMC23": 10, + "HMMT": 10, + }, + "science": { + "GPQADiamond": 3, + "JEEBench": 3, + "HLE": 3, + "OlympiadBench_Physics": 3, + }, + "code": { + "LiveCodeBench": 6, + "LiveCodeBenchv5_official": 6, + "LiveCodeBenchv6_official": 6, + }, +} + +ALL_TASK_SEEDS: dict[str, int] = {task: n for suite in SUITES.values() for task, n in suite.items()} + +# Benchmarks that are single-pass by construction: they carry no ``n_repeat`` +# attribute, so EVALCHEMY_N_REPEAT is inert and they emit ``accuracy`` rather +# than ``accuracy_avg``/``run_stats``. +SINGLE_PASS_TASKS = frozenset({"MATH500", "OlympiadBench", "OlympiadBench_Physics"}) + +# Generation / engine constants, matching the TPU pipeline. +MAX_GEN_TOKS = 32768 +CONTEXT_BUFFER = 4096 +MAX_MODEL_LEN = MAX_GEN_TOKS + CONTEXT_BUFFER # 36864 +TEMPERATURE = 0.7 +TOP_P = 1.0 +NUM_FEWSHOT = 0 + +RESULTS_FILE_RE = re.compile(r"^results_.*\.json$") +STEP_RE = re.compile(r"step-(\d+)") + +# -------------------------------------------------------------------------- +# GPU engine profiles +# -------------------------------------------------------------------------- + + +@dataclasses.dataclass(frozen=True) +class GpuProfile: + """vLLM engine settings for an 8B bf16 model at 36864 context.""" + + name: str + gpu_memory_utilization: float + max_num_seqs: int + max_num_batched_tokens: int + # Minimum bytes vLLM must be allowed to claim on each GPU before the run is + # worth starting: weights + activations + a few full-length sequences of KV. + min_budget_gib_per_gpu: float + min_gpus: int + notes: str + + +GPU_PROFILES: dict[str, GpuProfile] = { + # 1x 80GB is the reference config: 16.4 GiB bf16 weights, ~4-6 GiB + # activations/CUDA graphs, ~48 GiB KV = ~340k cached tokens = ~9 concurrent + # full-length (36864) sequences, far more at realistic generation lengths. + "a100-80g": GpuProfile( + name="a100-80g", + gpu_memory_utilization=0.90, + max_num_seqs=64, + max_num_batched_tokens=8192, + min_budget_gib_per_gpu=40.0, + min_gpus=1, + notes="A100-SXM/PCIe 80GB. 0.90 leaves ~8 GiB for the driver+NCCL. Do not exceed 0.92.", + ), + "h100-80g": GpuProfile( + name="h100-80g", + gpu_memory_utilization=0.90, + max_num_seqs=64, + max_num_batched_tokens=16384, + min_budget_gib_per_gpu=40.0, + min_gpus=1, + notes="H100 80GB HBM3. Higher prefill budget than A100; FlashAttention-3 kernels.", + ), + "h200-141g": GpuProfile( + name="h200-141g", + gpu_memory_utilization=0.90, + max_num_seqs=128, + max_num_batched_tokens=16384, + min_budget_gib_per_gpu=40.0, + min_gpus=1, + notes="H200 141GB. KV cache stops being the bottleneck; raise max_num_seqs.", + ), + # 40GB: 0.92*39.6 = 36.4 GiB budget, minus 16.4 weights minus ~4 overhead + # leaves ~16 GiB KV = ~113k tokens = 3 full-length sequences. Runs, but slow + # and preemption-heavy. Use two cards. + "a100-40g": GpuProfile( + name="a100-40g", + gpu_memory_utilization=0.92, + max_num_seqs=16, + max_num_batched_tokens=4096, + min_budget_gib_per_gpu=34.0, + min_gpus=1, + notes="A100 40GB. Tight for 8B at 36864 ctx; strongly prefer --tensor-parallel-size 2.", + ), + "l40s-48g": GpuProfile( + name="l40s-48g", + gpu_memory_utilization=0.92, + max_num_seqs=24, + max_num_batched_tokens=4096, + min_budget_gib_per_gpu=34.0, + min_gpus=1, + notes="L40S 48GB, no NVLink. Keep tensor_parallel_size=1; PCIe TP is slower than 1 card.", + ), +} + +# nvidia-smi product name substring -> profile +_NAME_TO_PROFILE = ( + ("H200", "h200-141g"), + ("H100", "h100-80g"), + ("A100-SXM4-40GB", "a100-40g"), + ("A100 40GB", "a100-40g"), + ("A100", "a100-80g"), + ("L40S", "l40s-48g"), +) + + +@dataclasses.dataclass(frozen=True) +class GpuInfo: + index: int + name: str + total_mib: int + used_mib: int + free_mib: int + + +def query_gpus() -> list[GpuInfo]: + """Read nvidia-smi. Honours CUDA_VISIBLE_DEVICES ordering when set.""" + smi = shutil.which("nvidia-smi") + if smi is None: + raise RuntimeError("nvidia-smi not found; this driver requires an NVIDIA GPU box") + out = subprocess.run( + [smi, "--query-gpu=index,name,memory.total,memory.used,memory.free", "--format=csv,noheader,nounits"], + check=True, + capture_output=True, + text=True, + ).stdout + gpus = [] + for line in out.strip().splitlines(): + index, name, total, used, free = (part.strip() for part in line.split(",")) + gpus.append(GpuInfo(int(index), name, int(total), int(used), int(free))) + + visible = os.environ.get("CUDA_VISIBLE_DEVICES") + if visible: + wanted = [int(x) for x in visible.split(",") if x.strip() != ""] + by_index = {g.index: g for g in gpus} + missing = [i for i in wanted if i not in by_index] + if missing: + raise RuntimeError(f"CUDA_VISIBLE_DEVICES names GPUs {missing} that nvidia-smi does not report") + gpus = [by_index[i] for i in wanted] + return gpus + + +def detect_profile(gpus: list[GpuInfo]) -> GpuProfile: + name = gpus[0].name + for needle, key in _NAME_TO_PROFILE: + if needle.lower() in name.lower(): + profile = GPU_PROFILES[key] + # An "A100" with <60GB is the 40GB SKU regardless of the name string. + if key == "a100-80g" and gpus[0].total_mib < 60_000: + profile = GPU_PROFILES["a100-40g"] + logger.info("detected GPU %r -> profile %s", name, profile.name) + return profile + raise RuntimeError( + f"no built-in profile for GPU {name!r}; pass --gpu-profile explicitly " + f"(one of: {', '.join(sorted(GPU_PROFILES))})" + ) + + +def resolve_gpu_memory_utilization(profile: GpuProfile, gpus: list[GpuInfo], headroom_mib: int) -> float: + """Clamp gpu_memory_utilization so co-tenants on the same card do not OOM us. + + THE PITFALL THIS EXISTS FOR: vLLM interprets ``gpu_memory_utilization`` as a + fraction of TOTAL device memory, not of FREE device memory. If an embedding + server (or a JAX training process, which preallocates 75% by default) already + holds 10 GiB of an 80 GiB card and you ask for 0.90, vLLM will size its KV + pool as if it owned 72 GiB. Total demand becomes 82 GiB and one of the two + processes dies with CUDA OOM -- and it usually dies *mid-run*, after the KV + pool is grown, not at init, so you lose hours of generation. + + So: measure what is already resident, and cap the fraction at what is + actually claimable. If that leaves less than the profile's minimum budget, + refuse to start rather than crash later. + """ + worst = min(gpus, key=lambda g: g.free_mib) + claimable_mib = worst.free_mib - headroom_mib + if claimable_mib <= 0: + raise RuntimeError( + f"GPU {worst.index} ({worst.name}) has only {worst.free_mib} MiB free " + f"(headroom {headroom_mib} MiB). Something else owns this card." + ) + allowed_fraction = claimable_mib / worst.total_mib + gmu = min(profile.gpu_memory_utilization, allowed_fraction) + + budget_gib = gmu * worst.total_mib / 1024.0 + if budget_gib < profile.min_budget_gib_per_gpu: + raise RuntimeError( + f"only {budget_gib:.1f} GiB claimable on GPU {worst.index} " + f"({worst.used_mib} MiB already in use by another process) but profile " + f"{profile.name} needs {profile.min_budget_gib_per_gpu:.1f} GiB. " + "Free the card, pin the other process to a different GPU with " + "CUDA_VISIBLE_DEVICES, or raise --tensor-parallel-size." + ) + if gmu < profile.gpu_memory_utilization: + logger.warning( + "clamping gpu_memory_utilization %.2f -> %.3f: GPU %d already has %d MiB resident", + profile.gpu_memory_utilization, + gmu, + worst.index, + worst.used_mib, + ) + return round(gmu, 3) + + +def kv_bytes_per_token(checkpoint: Path) -> int | None: + """Estimate bf16 KV cache cost per token from the checkpoint's config.json.""" + config_path = checkpoint / "config.json" + if not config_path.exists(): + return None + cfg = json.loads(config_path.read_text()) + layers = cfg.get("num_hidden_layers") + kv_heads = cfg.get("num_key_value_heads") or cfg.get("num_attention_heads") + head_dim = cfg.get("head_dim") + if head_dim is None and cfg.get("hidden_size") and cfg.get("num_attention_heads"): + head_dim = cfg["hidden_size"] // cfg["num_attention_heads"] + if not (layers and kv_heads and head_dim): + return None + return layers * kv_heads * head_dim * 2 * 2 # K and V, 2 bytes each + + +# -------------------------------------------------------------------------- +# Output layout -- must stay byte-compatible with claude/compile_results.py +# -------------------------------------------------------------------------- + + +def sanitize_model_name(name: str) -> str: + """Reproduce lm-eval's GeneralConfigTracker.model_name_sanitized.""" + return re.sub(r"[\"<>:/\|\\?\*\[\]]+", "__", name) + + +def _hash6(*parts: str) -> str: + """Stable 6-hex suffix, standing in for the Marin executor's config hash. + + The results.md compiler only substring-matches ``_seed-`` and + ``compile__avgseeds-``, so the value is opaque -- but it must be + STABLE so re-runs land in the same directory instead of piling up. + """ + return hashlib.md5("\x00".join(parts).encode()).hexdigest()[:6] + + +def step_of(checkpoint: Path) -> str | None: + match = STEP_RE.search(str(checkpoint)) + return match.group(1) if match else None + + +@dataclasses.dataclass(frozen=True) +class Layout: + """Every path a single (checkpoint, task) evaluation reads or writes.""" + + results_root: Path + experiment: str + step: str | None + task: str + seeds: tuple[int, ...] + model_dir_name: str + + @property + def eval_dir(self) -> Path: + name = f"{self.experiment}-step{self.step}" if self.step else self.experiment + return self.results_root / name + + def seed_dir(self, seed: int) -> Path: + suffix = _hash6(self.experiment, self.step or "", self.task, str(seed)) + return self.eval_dir / f"{self.task}_seed{seed}-{suffix}" + + def seed_results_dir(self, seed: int) -> Path: + return self.seed_dir(seed) / f"{self.task}_{NUM_FEWSHOT}shot" / self.model_dir_name + + @property + def compile_dir(self) -> Path | None: + if len(self.seeds) < 2: + return None + suffix = _hash6(self.experiment, self.step or "", self.task, "compile") + return self.eval_dir / f"compile_{self.task}_avg{len(self.seeds)}seeds-{suffix}" + + @property + def averaged_results(self) -> Path | None: + compile_dir = self.compile_dir + return None if compile_dir is None else compile_dir / "compiled_results" / "averaged_results.json" + + @property + def raw_dir(self) -> Path: + return self.eval_dir / "_raw" / self.task + + @property + def marker(self) -> Path: + return self.eval_dir / "_markers" / f"{self.task}.complete.json" + + +def find_results_file(root: Path) -> Path | None: + candidates = [p for p in root.rglob("results_*.json") if RESULTS_FILE_RE.match(p.name)] + if not candidates: + return None + return max(candidates, key=lambda p: p.stat().st_mtime) + + +def read_task_accuracy(results_file: Path, task: str) -> float | None: + """Read the metric the results.md compiler reads: accuracy_avg, else accuracy.""" + try: + payload = json.loads(results_file.read_text()) + except (OSError, json.JSONDecodeError): + return None + entry = payload.get("results", {}).get(task) + if not isinstance(entry, dict): + return None + for key in ("accuracy_avg", "accuracy"): + value = entry.get(key) + if isinstance(value, (int, float)): + return float(value) + return None + + +def is_task_complete(layout: Layout) -> bool: + """Same three-part definition the eval monitor and results.md compiler use.""" + for seed in layout.seeds: + results_file = find_results_file(layout.seed_dir(seed)) + if results_file is None or read_task_accuracy(results_file, layout.task) is None: + return False + averaged = layout.averaged_results + if averaged is not None: + if not averaged.exists(): + return False + try: + payload = json.loads(averaged.read_text()) + if payload[0]["num_seeds"] != len(layout.seeds): + return False + except (OSError, json.JSONDecodeError, KeyError, IndexError): + return False + return True + + +# -------------------------------------------------------------------------- +# Running one task +# -------------------------------------------------------------------------- + + +def build_model_args(checkpoint: Path, *, profile: GpuProfile, gmu: float, tensor_parallel_size: int) -> str: + parts = [ + f"pretrained={checkpoint}", + "dtype=bfloat16", + f"tensor_parallel_size={tensor_parallel_size}", + f"gpu_memory_utilization={gmu}", + f"max_model_len={MAX_MODEL_LEN}", + f"max_gen_toks={MAX_GEN_TOKS}", + f"max_num_seqs={profile.max_num_seqs}", + f"max_num_batched_tokens={profile.max_num_batched_tokens}", + "trust_remote_code=True", + # Deterministic-ish prefix reuse across the K repetitions of a task: + # every repetition sends the identical prompt, so the prefix cache turns + # K prefills into 1. Big win at 10 seeds. + "enable_prefix_caching=True", + ] + return ",".join(parts) + + +def subprocess_env( + *, + evalchemy_dir: Path, + n_repeat: int, + num_proc: int, + hf_cache: Path | None, + extra: dict[str, str] | None = None, +) -> dict[str, str]: + env = os.environ.copy() + env.update( + { + "EVALCHEMY_N_REPEAT": str(n_repeat), + "EVALCHEMY_NUM_PROC": str(num_proc), + # LiveCodeBench loads a dataset SCRIPT; without this it refuses. + "HF_DATASETS_TRUST_REMOTE_CODE": "1", + # lm-eval refuses to execute generated code without this. + "HF_ALLOW_CODE_EVAL": "1", + "TOKENIZERS_PARALLELISM": "false", + "PYTHONUNBUFFERED": "1", + # The LCB grader spawns children; keep them off the GPU entirely. + "VLLM_WORKER_MULTIPROC_METHOD": "spawn", + "PYTHONPATH": os.pathsep.join( + [str(evalchemy_dir), *([os.environ["PYTHONPATH"]] if os.environ.get("PYTHONPATH") else [])] + ), + } + ) + if hf_cache is not None: + env.update( + { + "HF_HOME": str(hf_cache), + "HF_HUB_CACHE": str(hf_cache / "hub"), + "HF_DATASETS_CACHE": str(hf_cache / "datasets"), + } + ) + if extra: + env.update(extra) + return env + + +def run_evalchemy( + *, + evalchemy_dir: Path, + checkpoint: Path, + task: str, + n_repeat: int, + out_dir: Path, + profile: GpuProfile, + gmu: float, + tensor_parallel_size: int, + batch_size: int, + hf_cache: Path | None, + num_proc: int, + debug: bool, + timeout: int | None, +) -> Path: + """Invoke evalchemy once for one task, producing one results_*.json.""" + entry = evalchemy_dir / "_marin_gpu_entry.py" + if not entry.exists(): + raise RuntimeError(f"{entry} missing -- run patch_evalchemy_gpu.py against {evalchemy_dir}") + + if out_dir.exists(): + # Exactly one results file per run, so downstream globs are unambiguous. + shutil.rmtree(out_dir) + out_dir.mkdir(parents=True, exist_ok=True) + + command = [ + sys.executable, + str(entry), + "--model", + "vllm", + "--tasks", + task, + "--model_args", + build_model_args(checkpoint, profile=profile, gmu=gmu, tensor_parallel_size=tensor_parallel_size), + "--batch_size", + str(batch_size), + "--output_path", + str(out_dir), + "--verbosity", + "INFO", + "--apply_chat_template", + "--confirm_run_unsafe_code", + "--max_tokens", + str(MAX_GEN_TOKS), + # evalchemy hands seeds[0] to vLLM as the per-request seed and adds the + # repetition index, so repetition i runs at seed SEED_BASE + i. + "--seed", + ",".join([str(SEED_BASE)] * 4), + "--gen_kwargs", + f"temperature={TEMPERATURE},top_p={TOP_P},max_gen_toks={MAX_GEN_TOKS}", + ] + if debug: + command.append("--debug") + + env = subprocess_env(evalchemy_dir=evalchemy_dir, n_repeat=n_repeat, num_proc=num_proc, hf_cache=hf_cache) + + log_path = out_dir.parent / f"{task}.evalchemy.log" + log_path.parent.mkdir(parents=True, exist_ok=True) + logger.info("launching %s (n_repeat=%d) -> %s", task, n_repeat, log_path) + logger.debug("command: %s", " ".join(command)) + + started = time.time() + with log_path.open("w") as log_file: + log_file.write(f"# {' '.join(command)}\n\n") + log_file.flush() + process = subprocess.Popen( + command, + cwd=str(evalchemy_dir), + env=env, + stdout=log_file, + stderr=subprocess.STDOUT, + ) + try: + returncode = process.wait(timeout=timeout) + except subprocess.TimeoutExpired: + process.kill() + process.wait() + raise RuntimeError(f"{task} exceeded timeout of {timeout}s; see {log_path}") from None + + elapsed = time.time() - started + results_file = find_results_file(out_dir) + + # Exit 0 with no results file means scoring hung or crashed silently -- the + # same guard the TPU evaluator carries. Treat it as a hard failure. + if returncode != 0: + raise RuntimeError(f"{task} exited {returncode} after {elapsed:.0f}s; see {log_path}") + if results_file is None: + raise RuntimeError( + f"{task} exited 0 after {elapsed:.0f}s but wrote no results_*.json under {out_dir}; " + f"scoring likely hung or crashed silently. See {log_path}" + ) + logger.info("%s finished in %.0fs -> %s", task, elapsed, results_file) + return results_file + + +# -------------------------------------------------------------------------- +# Fan-out: one combined run -> per-seed result files + a compile file +# -------------------------------------------------------------------------- + + +def per_repetition_accuracies(task_entry: dict[str, Any], n_repeat: int, task: str) -> list[float]: + run_stats = task_entry.get("run_stats") + if isinstance(run_stats, list) and len(run_stats) == n_repeat: + return [float(stat["accuracy"]) for stat in run_stats] + if n_repeat == 1: + for key in ("accuracy_avg", "accuracy"): + value = task_entry.get(key) + if isinstance(value, (int, float)): + return [float(value)] + raise RuntimeError( + f"{task}: expected {n_repeat} repetitions but the result carries " + f"run_stats={type(run_stats).__name__} of length " + f"{len(run_stats) if isinstance(run_stats, list) else 'n/a'}. " + "The n_repeat patch probably did not apply -- re-run patch_evalchemy_gpu.py." + ) + + +def slice_examples(examples: list[dict[str, Any]], index: int) -> list[dict[str, Any]] | None: + """Project a multi-repetition example list down to one repetition. + + Returns None when the benchmark does not keep per-repetition outputs + (LiveCodeBench* keep only the last run's examples), in which case the + per-seed file simply omits ``examples``. + """ + sliced: list[dict[str, Any]] = [] + for example in examples: + answers = example.get("model_answers") + outputs = example.get("model_outputs") + if not isinstance(answers, list) or index >= len(answers): + return None + projected = dict(example) + projected["model_answers"] = [answers[index]] + if isinstance(outputs, list) and index < len(outputs): + projected["model_outputs"] = [outputs[index]] + sliced.append(projected) + return sliced + + +def write_per_seed_results( + *, + raw_results_file: Path, + layout: Layout, + task: str, + checkpoint: Path, +) -> dict[int, float]: + """Split one combined run into K schema-compatible per-seed result files.""" + raw = json.loads(raw_results_file.read_text()) + entry = raw.get("results", {}).get(task) + if not isinstance(entry, dict): + available = sorted(raw.get("results", {})) + raise RuntimeError(f"{raw_results_file} has no results[{task!r}]; found {available}") + + seeds = layout.seeds + accuracies = per_repetition_accuracies(entry, len(seeds), task) + run_stats = entry.get("run_stats") if isinstance(entry.get("run_stats"), list) else None + examples = entry.get("examples") if isinstance(entry.get("examples"), list) else None + + per_seed: dict[int, float] = {} + timestamp = datetime.now().isoformat().replace(":", "-") + + for index, seed in enumerate(seeds): + accuracy = accuracies[index] + stat = run_stats[index] if run_stats and index < len(run_stats) else {} + num_total = stat.get("num_total", entry.get("num_total")) + num_solved = stat.get("num_solved", entry.get("num_solved")) + + seed_entry: dict[str, Any] = { + "num_total": num_total, + "num_solved": num_solved, + # Both keys, because the results.md compiler prefers accuracy_avg and + # falls back to accuracy, while some downstream readers expect the + # single-pass shape. + "accuracy": accuracy, + "accuracy_avg": accuracy, + "accuracy_std_err": 0.0, + "num_repeat": 1, + "solved_avg": num_solved, + "run_stats": [ + { + "repetition": 1, + "num_total": num_total, + "num_solved": num_solved, + "accuracy": accuracy, + } + ], + } + if examples is not None: + projected = slice_examples(examples, index) + if projected is not None: + seed_entry["examples"] = projected + # Carry through any extra per-difficulty metrics LCB emits. + for key, value in entry.items(): + if key.startswith("accuracy_") and key.endswith("_avg") and key != "accuracy_avg": + seed_entry[key] = value + + payload = {k: v for k, v in raw.items() if k != "results"} + payload["results"] = {task: seed_entry} + payload["_marin_gpu_port"] = { + "source_results_file": str(raw_results_file), + "repetition_index": index, + "seed": seed, + "note": ( + "Synthesised from a single K-repetition evalchemy run. On GPU the " + "seed is per-request (SEED_BASE + repetition_index), not an engine " + "seed as on TPU." + ), + "checkpoint": str(checkpoint), + } + + out_dir = layout.seed_results_dir(seed) + out_dir.mkdir(parents=True, exist_ok=True) + for stale in out_dir.glob("results_*.json"): + stale.unlink() + (out_dir / f"results_{timestamp}.json").write_text(json.dumps(payload, indent=2, default=str)) + per_seed[seed] = accuracy + + logger.info("%s per-seed accuracies: %s", task, {s: round(a, 4) for s, a in per_seed.items()}) + return per_seed + + +def write_compile_results(*, layout: Layout, per_seed: dict[int, float], checkpoint: Path) -> None: + """Write the averaged_results.json that claude/compile_results.py reads. + + Schema matches ``experiments/evals/evalchemy_results_compiler.py`` exactly: + a one-element JSON array with base_model_name, dataset_name, num_seeds, + seeds[], correct_mean, correct_std, correct_per_seed{}. ``correct_std`` is a + SAMPLE standard deviation (ddof=1), matching pandas Series.std(), and it is + what renders as the bracketed [x.x] in results.md. + + Unlike the TPU compile step, the numbers here come from each benchmark's own + grader rather than a naive string re-match, so HMMT / JEEBench / + OlympiadBench_Physics / LiveCodeBench* are correct in this file too. The + results.md compiler still bypasses the compile dir for those six and averages + the per-seed files instead; both paths now agree. + """ + compile_dir = layout.compile_dir + if compile_dir is None: + return + + seeds = sorted(per_seed) + values = [per_seed[s] for s in seeds] + mean = statistics.fmean(values) + std = statistics.stdev(values) if len(values) > 1 else 0.0 + + record = { + "base_model_name": layout.model_dir_name.lower(), + "dataset_name": layout.task.lower(), + "num_seeds": len(seeds), + "seeds": seeds, + "correct_mean": mean, + "correct_std": std, + "correct_per_seed": {str(s): per_seed[s] for s in seeds}, + } + + out_dir = compile_dir / "compiled_results" + out_dir.mkdir(parents=True, exist_ok=True) + (out_dir / "averaged_results.json").write_text(json.dumps([record], indent=2)) + + with (out_dir / "averaged_results.csv").open("w", newline="") as handle: + writer = csv.writer(handle) + writer.writerow(["base_model_name", "dataset_name", "num_seeds", "seeds", "correct_mean", "correct_std"]) + writer.writerow( + [record["base_model_name"], record["dataset_name"], record["num_seeds"], seeds, mean, std] + ) + + # compiled_results.{json,csv} exist for parity with the TPU compile step. + # NOTE: the TPU version stores one row per graded EXAMPLE; we store one row + # per seed, because we take accuracy from the benchmark's own grader rather + # than re-grading examples with string equality. Nothing downstream reads it. + rows = [ + { + "dataset_name": layout.task.lower(), + "model_name": layout.model_dir_name.lower(), + "seed": seed, + "accuracy": per_seed[seed], + "checkpoint": str(checkpoint), + } + for seed in seeds + ] + (out_dir / "compiled_results.json").write_text(json.dumps(rows, indent=2)) + with (out_dir / "compiled_results.csv").open("w", newline="") as handle: + writer = csv.DictWriter(handle, fieldnames=list(rows[0])) + writer.writeheader() + writer.writerows(rows) + + logger.info( + "%s compiled: mean=%.4f std=%.4f over %d seeds -> %s", + layout.task, + mean, + std, + len(seeds), + out_dir / "averaged_results.json", + ) + + +# -------------------------------------------------------------------------- +# Checkpoint validation +# -------------------------------------------------------------------------- + + +def validate_checkpoint(checkpoint: Path) -> None: + """Reject half-written HF exports before burning an hour of GPU on them. + + Levanter writes hf/step-N incrementally while training continues; a monitor + that races the export will otherwise load a truncated shard set. + """ + if not checkpoint.is_dir(): + raise RuntimeError(f"{checkpoint} is not a directory") + if not (checkpoint / "config.json").exists(): + raise RuntimeError(f"{checkpoint}/config.json missing") + if not any((checkpoint / name).exists() for name in ("tokenizer_config.json", "tokenizer.json")): + raise RuntimeError(f"{checkpoint} has no tokenizer files") + + index_path = checkpoint / "model.safetensors.index.json" + if index_path.exists(): + index = json.loads(index_path.read_text()) + shards = sorted(set(index.get("weight_map", {}).values())) + missing = [s for s in shards if not (checkpoint / s).exists()] + if missing: + raise RuntimeError(f"{checkpoint} is incomplete: {len(missing)} of {len(shards)} shards missing: {missing[:3]}") + return + if not (checkpoint / "model.safetensors").exists(): + raise RuntimeError(f"{checkpoint} has neither model.safetensors nor a shard index") + + +# -------------------------------------------------------------------------- +# Entry point +# -------------------------------------------------------------------------- + + +def resolve_tasks(args: argparse.Namespace) -> dict[str, int]: + if args.tasks: + selected = {} + for name in args.tasks: + if name not in ALL_TASK_SEEDS: + raise SystemExit(f"unknown task {name!r}; known: {', '.join(sorted(ALL_TASK_SEEDS))}") + selected[name] = ALL_TASK_SEEDS[name] + elif args.suite: + selected = dict(SUITES[args.suite]) + else: + raise SystemExit("pass --suite or --tasks") + if args.seeds is not None: + selected = {task: args.seeds for task in selected} + return selected + + +def evaluate_task(args: argparse.Namespace, task: str, n_seeds: int, profile: GpuProfile, gmu: float) -> dict[str, Any]: + checkpoint = args.checkpoint.resolve() + effective_seeds = 1 if task in SINGLE_PASS_TASKS else n_seeds + seeds = tuple(SEED_BASE + i for i in range(effective_seeds)) + + layout = Layout( + results_root=args.results_root.resolve(), + experiment=args.experiment, + step=step_of(checkpoint), + task=task, + seeds=seeds, + model_dir_name=sanitize_model_name(str(checkpoint)), + ) + + if not args.force and is_task_complete(layout): + logger.info("%s already complete for %s; skipping", task, layout.eval_dir.name) + return {"task": task, "status": "cached", "seeds": list(seeds)} + + started = time.time() + raw_results_file = run_evalchemy( + evalchemy_dir=args.evalchemy_dir.resolve(), + checkpoint=checkpoint, + task=task, + n_repeat=effective_seeds, + out_dir=layout.raw_dir, + profile=profile, + gmu=gmu, + tensor_parallel_size=args.tensor_parallel_size, + batch_size=args.batch_size, + hf_cache=args.hf_cache, + num_proc=args.num_proc, + debug=args.debug, + timeout=args.task_timeout, + ) + per_seed = write_per_seed_results( + raw_results_file=raw_results_file, layout=layout, task=task, checkpoint=checkpoint + ) + write_compile_results(layout=layout, per_seed=per_seed, checkpoint=checkpoint) + + summary = { + "task": task, + "status": "ok", + "seeds": list(seeds), + "mean": statistics.fmean(per_seed.values()), + "std": statistics.stdev(per_seed.values()) if len(per_seed) > 1 else 0.0, + "per_seed": {str(k): v for k, v in per_seed.items()}, + "elapsed_seconds": round(time.time() - started, 1), + } + layout.marker.parent.mkdir(parents=True, exist_ok=True) + layout.marker.write_text(json.dumps({**summary, "checkpoint": str(checkpoint)}, indent=2)) + if args.prune_raw and layout.raw_dir.exists(): + shutil.rmtree(layout.raw_dir) + logger.info("pruned %s", layout.raw_dir) + return summary + + +def build_parser() -> argparse.ArgumentParser: + parser = argparse.ArgumentParser( + description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter + ) + parser.add_argument("--checkpoint", type=Path, required=True, help="Local HF checkpoint dir (.../hf/step-N)") + parser.add_argument( + "--experiment", + required=True, + help="Experiment name WITHOUT a -stepN suffix; the step is taken from --checkpoint", + ) + parser.add_argument("--suite", choices=sorted(SUITES), help="math | science | code") + parser.add_argument("--tasks", nargs="+", help="Explicit evalchemy task names (overrides --suite)") + parser.add_argument("--seeds", type=int, help="Override the per-task seed count (default: TPU parity)") + parser.add_argument("--results-root", type=Path, required=True, help="Local stand-in for gs://.../evaluation/evalchemy") + parser.add_argument( + "--evalchemy-dir", + type=Path, + default=Path(os.environ.get("EVALCHEMY_DIR", "/opt/marin-gpu-eval/evalchemy")), + help="Patched evalchemy checkout (default: $EVALCHEMY_DIR)", + ) + parser.add_argument("--gpu-profile", choices=sorted(GPU_PROFILES), help="Default: auto-detect from nvidia-smi") + parser.add_argument("--tensor-parallel-size", type=int, help="Default: number of visible GPUs") + parser.add_argument("--batch-size", type=int, default=64, help="lm-eval batch size (default 64)") + parser.add_argument( + "--gpu-memory-utilization", + type=float, + help="Override the profile value. Still clamped by free VRAM unless --no-vram-clamp.", + ) + parser.add_argument( + "--vram-headroom-mib", + type=int, + default=2048, + help="MiB left unclaimed on the busiest GPU when clamping (default 2048)", + ) + parser.add_argument("--no-vram-clamp", action="store_true", help="Trust the profile value verbatim (unsafe)") + parser.add_argument("--hf-cache", type=Path, help="HF_HOME for the eval subprocess") + parser.add_argument("--num-proc", type=int, default=8, help="Dataset map/filter workers (default 8)") + parser.add_argument("--task-timeout", type=int, default=None, help="Seconds per task before SIGKILL") + parser.add_argument("--debug", action="store_true", help="Smoke mode: evalchemy limits each task to 10 examples") + parser.add_argument( + "--prune-raw", + action="store_true", + help="Delete /_raw/ after a successful fan-out. The raw file holds all K " + "repetitions at full 32768-token outputs and is the largest artifact on disk.", + ) + parser.add_argument("--force", action="store_true", help="Re-run tasks that already have complete results") + parser.add_argument("--dry-run", action="store_true", help="Print the plan and exit") + parser.add_argument("--log-level", default="INFO") + return parser + + +def main(argv: list[str] | None = None) -> int: + args = build_parser().parse_args(argv) + logging.basicConfig( + level=getattr(logging, args.log_level.upper()), + format="%(asctime)s %(levelname)-7s %(name)s | %(message)s", + ) + + tasks = resolve_tasks(args) + validate_checkpoint(args.checkpoint.resolve()) + + gpus = query_gpus() + profile = GPU_PROFILES[args.gpu_profile] if args.gpu_profile else detect_profile(gpus) + if args.gpu_memory_utilization is not None: + profile = dataclasses.replace(profile, gpu_memory_utilization=args.gpu_memory_utilization) + if args.tensor_parallel_size is None: + args.tensor_parallel_size = len(gpus) + + gmu = ( + profile.gpu_memory_utilization + if args.no_vram_clamp + else resolve_gpu_memory_utilization(profile, gpus, args.vram_headroom_mib) + ) + + logger.info("checkpoint %s (step %s)", args.checkpoint, step_of(args.checkpoint.resolve())) + logger.info("gpus %s", [f"{g.index}:{g.name}:{g.free_mib}MiB free" for g in gpus]) + logger.info("profile %s (%s)", profile.name, profile.notes) + logger.info( + "engine tp=%d gmu=%.3f max_num_seqs=%d max_model_len=%d", + args.tensor_parallel_size, + gmu, + profile.max_num_seqs, + MAX_MODEL_LEN, + ) + per_token = kv_bytes_per_token(args.checkpoint.resolve()) + if per_token: + seq_gib = per_token * MAX_MODEL_LEN / 1024**3 + logger.info("kv cache %d B/token -> %.2f GiB per full-length sequence", per_token, seq_gib) + logger.info("tasks %s", {t: (1 if t in SINGLE_PASS_TASKS else n) for t, n in tasks.items()}) + + if args.dry_run: + return 0 + + summaries: list[dict[str, Any]] = [] + failures: list[tuple[str, str]] = [] + for task, n_seeds in tasks.items(): + try: + summaries.append(evaluate_task(args, task, n_seeds, profile, gmu)) + except Exception as error: # one bad task must not abandon the rest + logger.exception("task %s failed", task) + failures.append((task, str(error))) + + print(json.dumps({"summaries": summaries, "failures": failures}, indent=2)) + return 1 if failures else 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/gpu-sft/scripts/gpu_eval/gpu_eval_monitor.py b/gpu-sft/scripts/gpu_eval/gpu_eval_monitor.py new file mode 100644 index 0000000000000000000000000000000000000000..c6084f6f933ecc2947d9e1e37e1461496e853aa2 --- /dev/null +++ b/gpu-sft/scripts/gpu_eval/gpu_eval_monitor.py @@ -0,0 +1,486 @@ +#!/usr/bin/env python3 +# Copyright The Marin Authors +# SPDX-License-Identifier: Apache-2.0 +"""Eval-during-training monitor for a single local GPU box. + +Local replacement for ``claude/eval_monitor.py``: GCS listing becomes a +filesystem walk, ``iris job run`` becomes a ``gpu_eval_driver.py`` subprocess, +and the "is another job already running this?" question becomes an OS file lock. + + python gpu_eval_monitor.py \\ + --run-dir /data/runs/exp_sft_qwen3_8b_x \\ + --experiment exp_sft_qwen3_8b_x \\ + --suite math \\ + --results-root /data/marin/evaluation/evalchemy \\ + --gpu-lock /var/tmp/marin-gpu.lock + +HOW GPU ACCESS IS SERIALISED AGAINST TRAINING + One advisory ``flock(2)`` on --gpu-lock. The lock is taken and released + around EACH TASK, not around a whole suite, so a training restart never + waits hours behind a 10-seed AIME sweep. + + The training launcher MUST take the same lock. Wrap it: + + flock /var/tmp/marin-gpu.lock -c 'python -m levanter.main.train_lm ...' + + (util-linux ``flock`` and Python's ``fcntl.flock`` both call flock(2) on the + file, so they interlock correctly. ``flock`` blocks by default; add ``-w N`` + for a timeout.) + + The lock alone is not sufficient, because JAX preallocates ~75% of every + visible GPU at import and holds it until the process exits, and because + nothing forces a stray notebook or embedding server to cooperate. So the + monitor ALSO refuses to launch until the busiest visible GPU reports at + least --min-free-vram-mib free. Both gates must pass. + + If you have spare cards, the better answer is physical separation: give + training ``CUDA_VISIBLE_DEVICES=0,1,2,3`` and run this monitor with + ``--cuda-visible-devices 4,5,6,7 --no-gpu-lock``. Then nothing contends and + evals never stall the run. + +DEDUPE + Three layers, matching the TPU monitor: + 1. Artifact probe -- gpu_eval_driver.is_task_complete() checks that every + expected per-seed results file parses and the compile file has the full + seed count. This is the source of truth; markers are only a fast path. + 2. In-process set of (step, task) pairs launched this session. + 3. The lock itself, which prevents two monitors from racing one GPU. + +CHECKPOINT READINESS + A step is only eligible once its HF export is structurally complete: every + shard named in model.safetensors.index.json exists, config.json exists, and + the tokenizer is present. Levanter writes hf/step-N while training + continues, so a naive directory-mtime check will load a truncated export. + A --settle-seconds quiet period on the newest file guards the tail. +""" + +from __future__ import annotations + +import argparse +import contextlib +import errno +import fcntl +import json +import logging +import os +import re +import signal +import subprocess +import sys +import time +from collections.abc import Iterator +from datetime import UTC, datetime +from pathlib import Path +from typing import Any + +sys.path.insert(0, str(Path(__file__).resolve().parent)) + +import gpu_eval_driver as driver # noqa: E402 + +logger = logging.getLogger("gpu_eval_monitor") + +STEP_DIR_RE = re.compile(r"^(?:step-|checkpoint-)(\d+)$") +DEFAULT_POLL_SECONDS = 300 +DEFAULT_MAX_RUNTIME_DAYS = 14 + +_SHUTDOWN = False + + +def _handle_signal(signum: int, _frame: Any) -> None: + global _SHUTDOWN + logger.warning("received signal %d; finishing the current task then exiting", signum) + _SHUTDOWN = True + + +# -------------------------------------------------------------------------- +# GPU lock +# -------------------------------------------------------------------------- + + +@contextlib.contextmanager +def gpu_lock(path: Path | None, *, poll_seconds: int, wait_seconds: int | None) -> Iterator[bool]: + """Exclusive flock(2) on ``path``. Yields True when held, False when disabled.""" + if path is None: + yield False + return + + path.parent.mkdir(parents=True, exist_ok=True) + deadline = None if wait_seconds is None else time.time() + wait_seconds + handle = path.open("a+") + try: + announced = False + while True: + try: + fcntl.flock(handle.fileno(), fcntl.LOCK_EX | fcntl.LOCK_NB) + break + except OSError as error: + if error.errno not in (errno.EACCES, errno.EAGAIN): + raise + if deadline is not None and time.time() >= deadline: + raise TimeoutError(f"could not acquire {path} within {wait_seconds}s") from None + if not announced: + logger.info("waiting for GPU lock %s (held by the training job?)", path) + announced = True + if _SHUTDOWN: + raise TimeoutError("shutdown requested while waiting for the GPU lock") from None + time.sleep(min(poll_seconds, 15)) + handle.seek(0) + handle.truncate() + handle.write(f"{os.getpid()} gpu_eval_monitor {datetime.now(UTC).isoformat()}\n") + handle.flush() + logger.info("acquired GPU lock %s", path) + yield True + finally: + with contextlib.suppress(OSError): + fcntl.flock(handle.fileno(), fcntl.LOCK_UN) + handle.close() + + +def free_vram_mib() -> int: + """Free MiB on the busiest visible GPU.""" + return min(gpu.free_mib for gpu in driver.query_gpus()) + + +# -------------------------------------------------------------------------- +# Checkpoint discovery +# -------------------------------------------------------------------------- + + +def discover_steps(run_dir: Path, subdir: str) -> list[tuple[int, Path]]: + """Return sorted (step, path) for every checkpoint directory under run_dir/subdir.""" + root = run_dir / subdir if subdir else run_dir + if not root.is_dir(): + return [] + found: list[tuple[int, Path]] = [] + for child in root.iterdir(): + if not child.is_dir(): + continue + match = STEP_DIR_RE.match(child.name) + if match: + found.append((int(match.group(1)), child)) + return sorted(found) + + +def checkpoint_is_settled(checkpoint: Path, settle_seconds: int) -> bool: + """True when nothing under the checkpoint has been written for settle_seconds.""" + if settle_seconds <= 0: + return True + newest = 0.0 + for path in checkpoint.rglob("*"): + if path.is_file(): + newest = max(newest, path.stat().st_mtime) + return newest > 0 and (time.time() - newest) >= settle_seconds + + +def checkpoint_is_ready(checkpoint: Path, settle_seconds: int) -> tuple[bool, str]: + try: + driver.validate_checkpoint(checkpoint) + except RuntimeError as error: + return False, str(error) + if not checkpoint_is_settled(checkpoint, settle_seconds): + return False, f"still being written (quiet period {settle_seconds}s not met)" + return True, "ready" + + +# -------------------------------------------------------------------------- +# Work items +# -------------------------------------------------------------------------- + + +def pending_tasks( + *, + experiment: str, + checkpoint: Path, + tasks: dict[str, int], + results_root: Path, +) -> list[str]: + """Tasks for this checkpoint whose artifacts are not already complete.""" + step = driver.step_of(checkpoint) + model_dir_name = driver.sanitize_model_name(str(checkpoint.resolve())) + pending = [] + for task, n_seeds in tasks.items(): + effective = 1 if task in driver.SINGLE_PASS_TASKS else n_seeds + layout = driver.Layout( + results_root=results_root.resolve(), + experiment=experiment, + step=step, + task=task, + seeds=tuple(driver.SEED_BASE + i for i in range(effective)), + model_dir_name=model_dir_name, + ) + if not driver.is_task_complete(layout): + pending.append(task) + return pending + + +def run_driver(args: argparse.Namespace, checkpoint: Path, task: str) -> tuple[int, float]: + """Run gpu_eval_driver.py for exactly one task in a fresh process. + + A fresh process per task is deliberate: vLLM does not reliably return all + device memory to the allocator on engine teardown, so a long-lived driver + would slowly starve itself (and the training job) across a 5-task suite. + """ + command = [ + sys.executable, + str(Path(__file__).resolve().parent / "gpu_eval_driver.py"), + "--checkpoint", + str(checkpoint), + "--experiment", + args.experiment, + "--tasks", + task, + "--results-root", + str(args.results_root), + "--evalchemy-dir", + str(args.evalchemy_dir), + "--num-proc", + str(args.num_proc), + ] + if args.gpu_profile: + command += ["--gpu-profile", args.gpu_profile] + if args.tensor_parallel_size: + command += ["--tensor-parallel-size", str(args.tensor_parallel_size)] + if args.hf_cache: + command += ["--hf-cache", str(args.hf_cache)] + if args.task_timeout: + command += ["--task-timeout", str(args.task_timeout)] + if args.prune_raw: + command.append("--prune-raw") + if args.debug: + command.append("--debug") + + env = os.environ.copy() + if args.cuda_visible_devices is not None: + env["CUDA_VISIBLE_DEVICES"] = args.cuda_visible_devices + + logger.info("launching driver: %s %s", task, checkpoint) + started = time.time() + process = subprocess.run(command, env=env, check=False) + return process.returncode, time.time() - started + + +# -------------------------------------------------------------------------- +# Event log +# -------------------------------------------------------------------------- + + +class EventLog: + def __init__(self, path: Path | None) -> None: + self.path = path + if path is not None: + path.parent.mkdir(parents=True, exist_ok=True) + + def emit(self, event: str, **fields: Any) -> None: + record = {"ts": datetime.now(UTC).isoformat(), "event": event, **fields} + logger.info("%s %s", event, {k: v for k, v in fields.items() if k != "ts"}) + if self.path is not None: + with self.path.open("a") as handle: + handle.write(json.dumps(record, default=str) + "\n") + + +# -------------------------------------------------------------------------- +# Main loop +# -------------------------------------------------------------------------- + + +def select_steps(steps: list[tuple[int, Path]], args: argparse.Namespace) -> list[tuple[int, Path]]: + selected = [ + (step, path) + for step, path in steps + if step >= args.min_step + and (args.max_step is None or step <= args.max_step) + and (args.step_stride <= 1 or step % args.step_stride == 0) + ] + if args.newest_first: + selected.reverse() + if args.latest_only and selected: + selected = selected[:1] if args.newest_first else selected[-1:] + return selected + + +def monitor(args: argparse.Namespace) -> int: + events = EventLog(args.event_log) + tasks = driver.SUITES[args.suite] if args.suite else {t: driver.ALL_TASK_SEEDS[t] for t in args.tasks} + if args.seeds is not None: + tasks = {task: args.seeds for task in tasks} + + deadline = time.time() + args.max_runtime_days * 24 * 3600 + launched: set[tuple[int, str]] = set() + + events.emit( + "start", + run_dir=str(args.run_dir), + experiment=args.experiment, + tasks=list(tasks), + poll_seconds=args.poll, + gpu_lock=str(args.gpu_lock) if args.gpu_lock else None, + min_free_vram_mib=args.min_free_vram_mib, + ) + + while not _SHUTDOWN and time.time() < deadline: + steps = select_steps(discover_steps(args.run_dir, args.checkpoint_subdir), args) + if not steps: + events.emit("no_checkpoints", root=str(args.run_dir / args.checkpoint_subdir)) + + for step, checkpoint in steps: + if _SHUTDOWN: + break + + ready, reason = checkpoint_is_ready(checkpoint, args.settle_seconds) + if not ready: + events.emit("checkpoint_not_ready", step=step, reason=reason) + continue + + todo = pending_tasks( + experiment=args.experiment, + checkpoint=checkpoint, + tasks=tasks, + results_root=args.results_root, + ) + todo = [t for t in todo if (step, t) not in launched or args.retry_failed] + if not todo: + continue + + events.emit("step_pending", step=step, tasks=todo) + + for task in todo: + if _SHUTDOWN: + break + + free = free_vram_mib() + if free < args.min_free_vram_mib: + events.emit( + "vram_blocked", + step=step, + task=task, + free_mib=free, + required_mib=args.min_free_vram_mib, + hint="training or another process still holds the card", + ) + break # nothing on this box will run; go back to sleep + + try: + with gpu_lock(args.gpu_lock, poll_seconds=args.poll, wait_seconds=args.lock_wait_seconds): + # Re-check under the lock: free VRAM can change while we + # queued behind the training job, and another monitor may + # have finished this task in the meantime. + free = free_vram_mib() + if free < args.min_free_vram_mib: + events.emit("vram_blocked_after_lock", step=step, task=task, free_mib=free) + break + if task not in pending_tasks( + experiment=args.experiment, + checkpoint=checkpoint, + tasks={task: tasks[task]}, + results_root=args.results_root, + ): + events.emit("raced", step=step, task=task) + continue + + launched.add((step, task)) + returncode, elapsed = run_driver(args, checkpoint, task) + except TimeoutError as error: + events.emit("lock_timeout", step=step, task=task, detail=str(error)) + break + + if returncode == 0: + events.emit("task_done", step=step, task=task, elapsed_seconds=round(elapsed, 1)) + else: + events.emit( + "task_failed", + step=step, + task=task, + returncode=returncode, + elapsed_seconds=round(elapsed, 1), + hint=f"see {args.results_root}/{args.experiment}-step{step}/_raw/{task}.evalchemy.log", + ) + if args.stop_on_failure: + events.emit("stop_on_failure", step=step, task=task) + return 1 + + if args.once: + break + events.emit("sleep", seconds=args.poll) + slept = 0 + while slept < args.poll and not _SHUTDOWN: + time.sleep(min(5, args.poll - slept)) + slept += 5 + + events.emit("exit", shutdown=_SHUTDOWN, expired=time.time() >= deadline) + return 0 + + +def build_parser() -> argparse.ArgumentParser: + parser = argparse.ArgumentParser( + description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter + ) + parser.add_argument("--run-dir", type=Path, required=True, help="Training output dir containing hf/step-N") + parser.add_argument("--checkpoint-subdir", default="hf", help="Subdir holding checkpoints (default: hf)") + parser.add_argument("--experiment", required=True, help="Experiment name WITHOUT a -stepN suffix") + parser.add_argument("--results-root", type=Path, required=True) + parser.add_argument("--suite", choices=sorted(driver.SUITES)) + parser.add_argument("--tasks", nargs="+", help="Explicit task names (overrides --suite)") + parser.add_argument("--seeds", type=int, help="Override per-task seed count") + parser.add_argument( + "--evalchemy-dir", + type=Path, + default=Path(os.environ.get("EVALCHEMY_DIR", "/opt/marin-gpu-eval/evalchemy")), + ) + + parser.add_argument("--gpu-lock", type=Path, default=Path("/var/tmp/marin-gpu.lock")) + parser.add_argument("--no-gpu-lock", action="store_true", help="Skip locking (use with disjoint --cuda-visible-devices)") + parser.add_argument("--lock-wait-seconds", type=int, default=None, help="Give up waiting for the lock (default: wait forever)") + parser.add_argument( + "--min-free-vram-mib", + type=int, + default=70_000, + help="Refuse to launch below this much free VRAM on the busiest visible GPU. " + "70000 suits an 8B bf16 model at 36864 ctx on an 80GB card.", + ) + parser.add_argument("--cuda-visible-devices", help="CUDA_VISIBLE_DEVICES for eval subprocesses") + parser.add_argument("--gpu-profile", choices=sorted(driver.GPU_PROFILES)) + parser.add_argument("--tensor-parallel-size", type=int) + + parser.add_argument("--poll", type=int, default=DEFAULT_POLL_SECONDS, help="Seconds between passes") + parser.add_argument("--settle-seconds", type=int, default=120, help="Quiet period a checkpoint must show") + parser.add_argument("--step-stride", type=int, default=1, help="Only eval steps divisible by this") + parser.add_argument("--min-step", type=int, default=0) + parser.add_argument("--max-step", type=int, default=None) + parser.add_argument("--latest-only", action="store_true", help="Only ever evaluate the newest eligible step") + parser.add_argument("--newest-first", action="store_true", help="Walk steps newest-first") + parser.add_argument("--retry-failed", action="store_true", help="Re-attempt tasks that failed earlier this session") + parser.add_argument("--stop-on-failure", action="store_true") + parser.add_argument("--once", action="store_true", help="Single pass, then exit") + parser.add_argument("--max-runtime-days", type=float, default=DEFAULT_MAX_RUNTIME_DAYS) + parser.add_argument("--task-timeout", type=int, default=None, help="Seconds per task before the driver is killed") + parser.add_argument("--num-proc", type=int, default=8) + parser.add_argument("--hf-cache", type=Path) + parser.add_argument("--prune-raw", action="store_true") + parser.add_argument("--debug", action="store_true", help="Smoke mode (10 examples per task)") + parser.add_argument("--event-log", type=Path, help="Append JSONL events here") + parser.add_argument("--log-level", default="INFO") + return parser + + +def main(argv: list[str] | None = None) -> int: + args = build_parser().parse_args(argv) + logging.basicConfig( + level=getattr(logging, args.log_level.upper()), + format="%(asctime)s %(levelname)-7s %(name)s | %(message)s", + ) + if not args.suite and not args.tasks: + raise SystemExit("pass --suite or --tasks") + if args.tasks: + unknown = [t for t in args.tasks if t not in driver.ALL_TASK_SEEDS] + if unknown: + raise SystemExit(f"unknown tasks: {unknown}") + if args.no_gpu_lock: + args.gpu_lock = None + + signal.signal(signal.SIGINT, _handle_signal) + signal.signal(signal.SIGTERM, _handle_signal) + return monitor(args) + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/gpu-sft/scripts/gpu_eval/patch_evalchemy_gpu.py b/gpu-sft/scripts/gpu_eval/patch_evalchemy_gpu.py new file mode 100644 index 0000000000000000000000000000000000000000..4af693af28bc8f42ef29076ccf59ac477c86b898 --- /dev/null +++ b/gpu-sft/scripts/gpu_eval/patch_evalchemy_gpu.py @@ -0,0 +1,274 @@ +#!/usr/bin/env python3 +# Copyright The Marin Authors +# SPDX-License-Identifier: Apache-2.0 +"""Patch a cloned evalchemy checkout for single-GPU offline vLLM evaluation. + +This is the GPU analogue of the runtime patching that +``lib/marin/src/marin/evaluation/evaluators/evalchemy_evaluator.py`` does on TPU. +The TPU-only patches (per-request seed stripping, fake ``vllm-tpu`` version +metadata, ``MODEL_IMPL_TYPE=vllm``, GCS ``AutoConfig`` redirection) are dropped; +two GPU-specific ones are added. + +Patches applied (all idempotent, all fail loudly if their marker is missing so a +pin bump cannot silently no-op): + +1. ``eval/eval_tracker.py`` -- lm-eval ``eval_logger`` import shim. +2. ``eval/eval.py`` -- ``utils.eval_logger`` attribute shim. +3. ``*/eval_instruct.py`` -- ``self.n_repeat`` becomes env-driven + (``EVALCHEMY_N_REPEAT``) so one vLLM engine can produce all repetitions. +4. ``*/livecodebench_utils.py`` -- fork -> spawn for the sandboxed test runner. + REQUIRED on GPU: CUDA context and vLLM/filelock state are not fork-safe and a + forked child either deadlocks or corrupts the parent's CUDA context. +5. ``LiveCodeBench*/eval_instruct.py`` -- dataset ``num_proc`` becomes env-driven + (``EVALCHEMY_NUM_PROC``) so LCB's dataset filter does not oversubscribe the + box while a training job is running. +6. Writes ``_marin_gpu_entry.py`` -- the launcher the driver invokes. It applies + the ``vllm.utils.get_open_port`` compatibility shim and lifts Python's + int->str digit cap before delegating to ``eval.eval``. + +Usage: + python patch_evalchemy_gpu.py /opt/evalchemy +""" + +from __future__ import annotations + +import argparse +import logging +import re +import sys +from pathlib import Path + +logger = logging.getLogger("patch_evalchemy_gpu") + +# Every benchmark that carries a ``self.n_repeat`` assignment at the pinned +# commits. Mirrors N_REPEAT_BENCHMARK_PATHS in Marin's TPU evaluator. +N_REPEAT_BENCHMARKS = ( + "AIME24", + "AIME25", + "AIME26", + "AMC23", + "HMMT", + "GPQADiamond", + "JEEBench", + "HLE", + "LiveCodeBench", + "LiveCodeBenchv5_official", + "LiveCodeBenchv6_official", + "CodeForces", + "CodeElo", +) + +# Benchmarks with no ``n_repeat`` at all -- single pass by construction. +SINGLE_PASS_BENCHMARKS = ("MATH500", "OlympiadBench", "OlympiadBench_Physics") + +LCB_BENCHMARKS = ("LiveCodeBench", "LiveCodeBenchv5_official", "LiveCodeBenchv6_official") + +LCB_UTILITY_FILES = ( + "eval/chat_benchmarks/LiveCodeBench/livecodebench_utils.py", + "eval/chat_benchmarks/LiveCodeBenchv5/livecodebench_utils.py", + "eval/chat_benchmarks/LiveCodeBenchv5_official/livecodebench_utils.py", + "eval/chat_benchmarks/LiveCodeBenchv6_official/livecodebench_utils.py", +) + +OLD_TRACKER_IMPORT = ( + "from lm_eval.utils import eval_logger, handle_non_serializable, hash_string, simple_parse_args_string" +) + +NEW_TRACKER_IMPORT = '''try: + from lm_eval.utils import eval_logger, handle_non_serializable, hash_string, simple_parse_args_string +except ImportError: + try: + from lm_eval.logging_utils import eval_logger + from lm_eval.utils import handle_non_serializable, hash_string, simple_parse_args_string + except ImportError: + import logging + eval_logger = logging.getLogger("lm-eval") + from lm_eval.utils import handle_non_serializable, hash_string, simple_parse_args_string''' + +EVAL_PY_MARKER = "from lm_eval.tasks import TaskManager as PretrainTaskManager" + +EVAL_PY_LOGGER_PATCH = ''' +# Patch utils.eval_logger for lm-eval compatibility (marin gpu port) +if not hasattr(utils, "eval_logger"): + try: + from lm_eval.logging_utils import eval_logger as _eval_logger + utils.eval_logger = _eval_logger + except ImportError: + import logging as _logging + utils.eval_logger = _logging.getLogger("lm-eval") +''' + +OLD_LCB_FORK = """ manager = multiprocessing.Manager() + result = manager.list() + p = multiprocessing.Process(target=run_tests_for_one_example, args=(test_cases, completion, result, is_extracted))""" + +NEW_LCB_SPAWN = """ # marin gpu port: CUDA/vLLM/filelock state is not safe to inherit through + # fork. A forked grader child either deadlocks on the CUDA driver mutex or + # corrupts the parent engine's context. Spawn is the only safe context here. + multiprocessing_context = multiprocessing.get_context("spawn") + manager = multiprocessing_context.Manager() + result = manager.list() + p = multiprocessing_context.Process( + target=run_tests_for_one_example, + args=(test_cases, completion, result, is_extracted), + )""" + +ENTRY_MODULE = '''#!/usr/bin/env python3 +"""Launcher for evalchemy under a local CUDA vLLM install (marin gpu port). + +Generated by patch_evalchemy_gpu.py -- do not edit by hand. + +Applies two process-wide fixes that must happen before ``eval.eval`` imports +lm-eval, then delegates. Invoked as:: + + python _marin_gpu_entry.py --model vllm --tasks AIME24 ... +""" + +import runpy +import socket +import sys + +# HMMT and some OlympiadBench answers are enormous integers (e.g. (26!)^3). +# json.dumps calls str() on them; without this the results file is never written +# and the eval exits 0 with no output. +sys.set_int_max_str_digits(0) + +# lm-eval's vllm_causallms imports get_open_port from vllm.utils, which vLLM +# moved/removed in newer releases. Re-provide it before lm-eval is imported. +try: + import vllm.utils as _vllm_utils +except ImportError: + _vllm_utils = None + +if _vllm_utils is not None and not hasattr(_vllm_utils, "get_open_port"): + def _get_open_port() -> int: + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock: + sock.bind(("", 0)) + return sock.getsockname()[1] + + _vllm_utils.get_open_port = _get_open_port + +sys.argv = ["eval.eval", *sys.argv[1:]] +runpy.run_module("eval.eval", run_name="__main__", alter_sys=True) +''' + + +def _patch_eval_tracker(root: Path) -> None: + path = root / "eval" / "eval_tracker.py" + source = path.read_text() + if "from lm_eval.logging_utils import eval_logger" in source: + logger.info("eval_tracker.py already patched") + return + if OLD_TRACKER_IMPORT not in source: + raise RuntimeError(f"marker missing in {path}: lm-eval import shim cannot be applied") + path.write_text(source.replace(OLD_TRACKER_IMPORT, NEW_TRACKER_IMPORT, 1)) + logger.info("patched %s", path) + + +def _patch_eval_py(root: Path) -> None: + path = root / "eval" / "eval.py" + source = path.read_text() + if "# Patch utils.eval_logger" in source: + logger.info("eval.py already patched") + return + if EVAL_PY_MARKER not in source: + raise RuntimeError(f"marker missing in {path}: eval_logger shim cannot be applied") + path.write_text(source.replace(EVAL_PY_MARKER, EVAL_PY_MARKER + EVAL_PY_LOGGER_PATCH, 1)) + logger.info("patched %s", path) + + +def _patch_n_repeat(root: Path) -> list[str]: + patched: list[str] = [] + for name in N_REPEAT_BENCHMARKS: + path = root / "eval" / "chat_benchmarks" / name / "eval_instruct.py" + if not path.exists(): + logger.warning("benchmark %s not present at this pin; skipping n_repeat patch", name) + continue + source = path.read_text() + if "EVALCHEMY_N_REPEAT" in source: + patched.append(name) + continue + # __import__ rather than a bare ``os`` reference: several benchmark + # modules do not import os, and inserting a top-level import is fragile + # across pin bumps (it would shadow a module docstring). + new_source, count = re.subn( + r"self\.n_repeat\s*=\s*\d+", + 'self.n_repeat = int(__import__("os").environ.get("EVALCHEMY_N_REPEAT", "1"))', + source, + count=1, + ) + if count != 1: + raise RuntimeError(f"no 'self.n_repeat = ' assignment found in {path}") + path.write_text(new_source) + patched.append(name) + logger.info("patched n_repeat in %s", path) + return patched + + +def _patch_lcb_num_proc(root: Path) -> None: + for name in LCB_BENCHMARKS: + path = root / "eval" / "chat_benchmarks" / name / "eval_instruct.py" + if not path.exists(): + continue + source = path.read_text() + if "EVALCHEMY_NUM_PROC" in source: + continue + if "num_proc=cpu_count" not in source: + raise RuntimeError(f"'num_proc=cpu_count' not found in {path}") + source = source.replace( + "num_proc=cpu_count", + 'num_proc=int(__import__("os").environ.get("EVALCHEMY_NUM_PROC", str(cpu_count)))', + 1, + ) + path.write_text(source) + logger.info("patched num_proc in %s", path) + + +def _patch_lcb_spawn(root: Path) -> None: + for relative in LCB_UTILITY_FILES: + path = root / relative + if not path.exists(): + logger.warning("%s not present at this pin; skipping spawn patch", relative) + continue + source = path.read_text() + if "multiprocessing_context = multiprocessing.get_context" in source: + continue + if OLD_LCB_FORK not in source: + raise RuntimeError(f"LiveCodeBench fork block not found in {path}") + path.write_text(source.replace(OLD_LCB_FORK, NEW_LCB_SPAWN, 1)) + logger.info("patched fork->spawn in %s", path) + + +def _write_entry(root: Path) -> Path: + path = root / "_marin_gpu_entry.py" + path.write_text(ENTRY_MODULE) + path.chmod(0o755) + logger.info("wrote %s", path) + return path + + +def patch(root: Path) -> None: + if not (root / "eval" / "eval.py").exists(): + raise RuntimeError(f"{root} does not look like an evalchemy checkout (no eval/eval.py)") + _patch_eval_tracker(root) + _patch_eval_py(root) + patched = _patch_n_repeat(root) + _patch_lcb_num_proc(root) + _patch_lcb_spawn(root) + _write_entry(root) + logger.info("n_repeat is env-driven for: %s", ", ".join(patched)) + logger.info("single-pass benchmarks (no n_repeat): %s", ", ".join(SINGLE_PASS_BENCHMARKS)) + logger.info("evalchemy at %s is ready for single-GPU vLLM evaluation", root) + + +def main() -> int: + logging.basicConfig(level=logging.INFO, format="%(levelname)s %(message)s") + parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + parser.add_argument("evalchemy_dir", type=Path, help="Path to the cloned evalchemy checkout") + args = parser.parse_args() + patch(args.evalchemy_dir.resolve()) + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/gpu-sft/scripts/gpu_eval/run_all_evals.sh b/gpu-sft/scripts/gpu_eval/run_all_evals.sh new file mode 100644 index 0000000000000000000000000000000000000000..474d644b6c44a60a0f0c3117a5495330f55c8a5b --- /dev/null +++ b/gpu-sft/scripts/gpu_eval/run_all_evals.sh @@ -0,0 +1,60 @@ +#!/usr/bin/env bash +# Run the Qwen3-8B baseline + the completed SFT checkpoints through evalchemy, sequentially, +# writing results to HDFS so a machine kill doesn't lose completed evals. +# +# source /opt/tiger/lmc_muon/gpu-eval/activate.sh # (this script also sources it) +# nohup bash gpu-sft/scripts/gpu_eval/run_all_evals.sh \ +# > /mnt/hdfs/fangzhao_writable/marin-eval/eval_all.log 2>&1 & +# +# Each eval uses the full node (tp defaults to #GPUs). Jobs run one at a time. Results go to +# $RESULTS on HDFS; a completed task's JSON is durable immediately, so a kill only costs the +# in-flight task (re-run the script — the driver skips finished tasks unless --force). +source /opt/tiger/lmc_muon/gpu-eval/activate.sh 2>/dev/null || true +set -uo pipefail + +DRIVER=/opt/tiger/lmc_muon/self-verified-distillation/gpu-sft/scripts/gpu_eval/gpu_eval_driver.py +RESULTS="${RESULTS:-/mnt/hdfs/fangzhao_writable/marin-eval/results}" +HDFS_SFT=/mnt/hdfs/fangzhao_writable/marin-sft +MODELS=/opt/tiger/lmc_muon/models +PROFILE="${PROFILE:-h100-80g}" +mkdir -p "$RESULTS" "$MODELS" + +run() { # run + echo "=========== eval $2 / $3 $(date -u) ===========" + python "$DRIVER" --checkpoint "$1" --experiment "$2" --suite "$3" \ + --gpu-profile "$PROFILE" --results-root "$RESULTS" \ + || echo "!! $2/$3 FAILED (continuing to next)" +} + +# Stage the completed checkpoints from HDFS -> local, keeping a 'step-1999' leaf so the driver +# parses the step number (16 GB each, ~30 s at 574 MB/s). +for e in coding_depth4v3_nofilter hs_competition_nofilter; do + DST="$MODELS/$e/step-1999" + if [ ! -f "$DST/config.json" ]; then + echo "staging $e/hf/step-1999 from HDFS -> $DST" + mkdir -p "$MODELS/$e" && cp -r "$HDFS_SFT/$e/hf/step-1999" "$DST" + fi +done + +if [ "${PARALLEL:-0}" = "1" ]; then + LOGDIR="$RESULTS/_logs"; mkdir -p "$LOGDIR" + echo "PARALLEL mode: 4 jobs x tp=2 across GPU pairs (per-job logs in $LOGDIR)" + CUDA_VISIBLE_DEVICES=0,1 python "$DRIVER" --checkpoint "$MODELS/Qwen3-8B" --experiment qwen3_8b_base --suite math --tensor-parallel-size 2 --gpu-profile "$PROFILE" --results-root "$RESULTS" > "$LOGDIR/base_math.log" 2>&1 & + CUDA_VISIBLE_DEVICES=2,3 python "$DRIVER" --checkpoint "$MODELS/Qwen3-8B" --experiment qwen3_8b_base --suite code --tensor-parallel-size 2 --gpu-profile "$PROFILE" --results-root "$RESULTS" > "$LOGDIR/base_code.log" 2>&1 & + CUDA_VISIBLE_DEVICES=4,5 python "$DRIVER" --checkpoint "$MODELS/hs_competition_nofilter/step-1999" --experiment hs_competition_nofilter --suite math --tensor-parallel-size 2 --gpu-profile "$PROFILE" --results-root "$RESULTS" > "$LOGDIR/hs_math.log" 2>&1 & + CUDA_VISIBLE_DEVICES=6,7 python "$DRIVER" --checkpoint "$MODELS/coding_depth4v3_nofilter/step-1999" --experiment coding_depth4v3_nofilter --suite code --tensor-parallel-size 2 --gpu-profile "$PROFILE" --results-root "$RESULTS" > "$LOGDIR/coding_code.log" 2>&1 & + wait + echo "=========== ALL PARALLEL EVALS DONE $(date -u) ===========" + echo "results -> $RESULTS ; per-job logs -> $LOGDIR" + exit 0 +fi + +# 1) baseline (the denominator for every (+delta)) on both suites +run "$MODELS/Qwen3-8B" qwen3_8b_base math +run "$MODELS/Qwen3-8B" qwen3_8b_base code +# 2) the two completed checkpoints on their matching suites +run "$MODELS/hs_competition_nofilter/step-1999" hs_competition_nofilter math +run "$MODELS/coding_depth4v3_nofilter/step-1999" coding_depth4v3_nofilter code + +echo "=========== ALL EVALS DONE $(date -u) ===========" +echo "results -> $RESULTS" diff --git a/gpu-sft/scripts/gpu_sft/accelerate_fsdp_qwen3.yaml b/gpu-sft/scripts/gpu_sft/accelerate_fsdp_qwen3.yaml new file mode 100644 index 0000000000000000000000000000000000000000..d41a256a4c5593ff63cb761d6d1f5b7ef1fb8ffb --- /dev/null +++ b/gpu-sft/scripts/gpu_sft/accelerate_fsdp_qwen3.yaml @@ -0,0 +1,41 @@ +# Optional: run train_sft_qwen3.py under `accelerate launch` instead of `torchrun`. +# The script builds its own FSDP wrapper, so accelerate is used ONLY as a launcher here +# (it sets RANK/LOCAL_RANK/WORLD_SIZE). These fsdp_config values document the exact +# equivalent settings for anyone who prefers accelerate's Accelerator.prepare(). +# +# accelerate launch --config_file accelerate_fsdp_qwen3.yaml train_sft_qwen3.py ... +# +# If you DO switch to Accelerator.prepare(), you must keep: +# * mixed_precision bf16 with fp32 master weights and fp32 gradient reduction +# * FULL_SHARD + transformer-layer wrapping of Qwen3DecoderLayer +# * use_orig_params: true (the weight-decay parameter groups depend on it) +# * accelerator.clip_grad_norm_ AFTER all accumulation micro-steps +# * your own loss normalisation; accelerate's default 1/gradient_accumulation_steps +# is NOT the recipe's per-microbatch token-weighted mean. +compute_environment: LOCAL_MACHINE +distributed_type: FSDP +downcast_bf16: 'no' +machine_rank: 0 +main_training_function: main +mixed_precision: bf16 +num_machines: 1 +num_processes: 8 +rdzv_backend: static +same_network: true +use_cpu: false +enable_cpu_affinity: true +fsdp_config: + fsdp_version: 1 + fsdp_sharding_strategy: FULL_SHARD + fsdp_auto_wrap_policy: TRANSFORMER_BASED_WRAP + fsdp_transformer_layer_cls_to_wrap: Qwen3DecoderLayer + fsdp_backward_prefetch: BACKWARD_PRE + fsdp_forward_prefetch: false + fsdp_use_orig_params: true + fsdp_cpu_ram_efficient_loading: true + fsdp_sync_module_states: true + fsdp_offload_params: false + fsdp_state_dict_type: SHARDED_STATE_DICT + fsdp_activation_checkpointing: false # the script calls gradient_checkpointing_enable itself + fsdp_reshard_after_forward: true + fsdp_limit_all_gathers: true diff --git a/gpu-sft/scripts/gpu_sft/deepspeed_zero3_qwen3.json b/gpu-sft/scripts/gpu_sft/deepspeed_zero3_qwen3.json new file mode 100644 index 0000000000000000000000000000000000000000..7510812fa0f755e6a5b0a062d58c7f7cd16846d8 --- /dev/null +++ b/gpu-sft/scripts/gpu_sft/deepspeed_zero3_qwen3.json @@ -0,0 +1,40 @@ +{ + "_comment": [ + "Fallback for 40 GB-class GPUs or shops standardised on DeepSpeed. Only the sharding", + "backend changes; the recipe (LR schedule, clipping, loss normalisation, packing) still", + "comes from train_sft_qwen3.py, so run it with --skip-checks and replace the FSDP block", + "with deepspeed.initialize, or drive it through accelerate's DeepSpeed plugin.", + "bf16.enabled with fp32 optimizer states reproduces jmp p=f32,c=bfloat16.", + "gradient_clipping is DISABLED here because the script clips explicitly before step().", + "Optimizer/scheduler are declared 'none' so the script's AdamW + LambdaLR stay in charge." + ], + "train_micro_batch_size_per_gpu": 1, + "gradient_accumulation_steps": 8, + "gradient_clipping": 0.0, + "steps_per_print": 1, + "wall_clock_breakdown": false, + "bf16": { "enabled": true }, + "fp16": { "enabled": false }, + "zero_optimization": { + "stage": 3, + "overlap_comm": true, + "contiguous_gradients": true, + "reduce_bucket_size": 67108864, + "stage3_prefetch_bucket_size": 67108864, + "stage3_param_persistence_threshold": 65536, + "stage3_max_live_parameters": 1000000000, + "stage3_max_reuse_distance": 1000000000, + "stage3_gather_16bit_weights_on_model_save": true, + "offload_optimizer": { "device": "none" }, + "offload_param": { "device": "none" } + }, + "activation_checkpointing": { + "partition_activations": false, + "cpu_checkpointing": false, + "contiguous_memory_optimization": false, + "number_checkpoints": 36, + "synchronize_checkpoint_boundary": false, + "profile": false + }, + "_offload_note": "On 8x40 GB set offload_optimizer.device='cpu' with pin_memory=true: that frees ~12 GB/GPU at a 20-40% throughput cost and does not change the math." +} diff --git a/gpu-sft/scripts/gpu_sft/launch_sft_8b.sh b/gpu-sft/scripts/gpu_sft/launch_sft_8b.sh new file mode 100644 index 0000000000000000000000000000000000000000..9dd7bcd7ce99549cb2e22878979a3914615a497c --- /dev/null +++ b/gpu-sft/scripts/gpu_sft/launch_sft_8b.sh @@ -0,0 +1,108 @@ +#!/usr/bin/env bash +# Qwen3-8B self-instill SFT on a single NVIDIA node. Local disk only: no GCS, no Iris. +# +# ./launch_sft_8b.sh <7-char-revision> +# +# Stage 1 (CPU, ~10-30 min) is cached; stage 2 is the training run. +set -euo pipefail + +DATASET_ID="${1:?usage: launch_sft_8b.sh }" +REVISION="${2:?}" +RUN_NAME="${3:?}" + +ROOT="${SFT_ROOT:-/data/sft}" +DATA_DIR="${ROOT}/data/$(echo "${DATASET_ID}" | tr '/' '_')-${REVISION}" +OUT_DIR="${ROOT}/runs/${RUN_NAME}" +NPROC="${NPROC:-$(nvidia-smi --list-gpus | wc -l)}" + +export HF_HOME="${HF_HOME:-${ROOT}/hf}" +export TOKENIZERS_PARALLELISM=false +export OMP_NUM_THREADS="${OMP_NUM_THREADS:-8}" +export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True +export NCCL_DEBUG="${NCCL_DEBUG:-WARN}" +export PYTHONUNBUFFERED=1 +# W&B is optional; the TPU runs used WANDB_MODE=offline on the workers. +export WANDB_MODE="${WANDB_MODE:-offline}" + +mkdir -p "${OUT_DIR}" + +# ---------------------------------------------------------------------------------- +# Stage 1: HF dataset -> packed local .npy (idempotent, skipped if manifest exists) +# ---------------------------------------------------------------------------------- +if [[ ! -f "${DATA_DIR}/manifest.json" ]]; then + python prepare_sft_data.py \ + --dataset-id "${DATASET_ID}" \ + --revision "${REVISION}" \ + --tokenizer Qwen/Qwen3-8B \ + --max-seq-len 32768 \ + --text-replacements marin \ + --num-proc "$(nproc)" \ + --out "${DATA_DIR}" 2>&1 | tee "${OUT_DIR}/prepare.log" + + # one-time fidelity gate: our encoder must agree with levanter's segment-split encoder + python prepare_sft_data.py \ + --dataset-id "${DATASET_ID}" --revision "${REVISION}" \ + --tokenizer Qwen/Qwen3-8B --verify 200 2>&1 | tee -a "${OUT_DIR}/prepare.log" +fi + +# ---------------------------------------------------------------------------------- +# Stage 1.5: print the plan (LR schedule, epochs, export cadence) before burning GPUs +# ---------------------------------------------------------------------------------- +python train_sft_qwen3.py --data "${DATA_DIR}" --out "${OUT_DIR}" --recipe qwen3-8b --dry-run + +# ---------------------------------------------------------------------------------- +# Stage 2: training +# +# Batch geometry (must reproduce global batch 64 x 32768 tokens = 2,097,152 tok/step): +# NPROC=8 : --per-device-batch 1, loss-groups 2 -> 4 micro-steps/group, 8 accum/step +# NPROC=4 : --per-device-batch 1, loss-groups 2 -> 8 micro-steps/group, 16 accum/step +# NPROC=16: --per-device-batch 1, loss-groups 2 -> 2 micro-steps/group, 4 accum/step +# +# Memory per GPU, Qwen3-8B (8.19e9 params), seq 32768, FSDP FULL_SHARD over 8 GPUs, +# fp32 master + bf16 compute + full activation checkpointing + chunked CE(2048): +# fp32 master params 4 * 8.19e9 / 8 = 4.1 GB +# Adam m + v (fp32) 8 * 8.19e9 / 8 = 8.2 GB +# fp32 grad shard 4 * 8.19e9 / 8 = 4.1 GB +# bf16 all-gather buffers root(1.24e9)+layers = 3.7 GB +# checkpointed layer inputs 36*32768*4096*2 = 9.7 GB +# one layer recompute working set = 3.2 GB +# chunked CE (2048 x 151936 fp32, fwd+grad) = 2.5 GB +# last hidden + grad = 0.5 GB +# ---------------------------------------------------------- +# ~36 GB + NCCL/fragmentation -> ~40 GB (fits 80 GB comfortably) +# +# 4 GPUs : optimiser/param/grad term triples to 32.8 GB -> ~50 GB, still fits 80 GB. +# 8x40 GB A100: ~30-34 GB, tight. Use --loss-chunk-size 1024 --wrap-embeddings, +# or 16 GPUs, or the 8K fallback below. +# +# Wall clock: ~3.9e20 FLOPs total (1.75e20 dense + 1.2e20 attention at 32K, x4/3 for +# recompute). 8xH100 at 30-35% MFU -> 40-50 h. 8xB200/H200 -> ~20-25 h. +# ---------------------------------------------------------------------------------- +torchrun --standalone --nproc_per_node="${NPROC}" train_sft_qwen3.py \ + --data "${DATA_DIR}" \ + --out "${OUT_DIR}" \ + --recipe qwen3-8b \ + --train-batch-size 64 \ + --per-device-batch 1 \ + --loss-groups 2 \ + --num-train-steps 2000 \ + --steps-per-hf-export 100 \ + --steps-per-checkpoint 20 \ + --loss-chunk-size 2048 \ + --hf-export-dtype bfloat16 \ + --resume auto \ + 2>&1 | tee -a "${OUT_DIR}/train.log" + +# ---------------------------------------------------------------------------------- +# 8K-context fallback (OOM at 32K, or no FlashAttention build). +# Keeps tokens/step, step count and LR schedule identical: 256 x 8192 = 2,097,152. +# Documents longer than 8192 are LEFT-SLICED (kept from the beginning) exactly as +# levanter's slice_strategy="left" does; check `docs_over_max_seq_len` in the +# prepare-stage stats before accepting this. +# +# python prepare_sft_data.py ... --max-seq-len 8192 --out ${DATA_DIR}-8k +# torchrun --standalone --nproc_per_node=8 train_sft_qwen3.py \ +# --data ${DATA_DIR}-8k --out ${OUT_DIR}-8k --recipe qwen3-8b \ +# --train-batch-size 256 --per-device-batch 4 --loss-groups 2 \ +# --num-train-steps 2000 --loss-chunk-size 2048 +# ---------------------------------------------------------------------------------- diff --git a/gpu-sft/scripts/gpu_sft/prepare_sft_data.py b/gpu-sft/scripts/gpu_sft/prepare_sft_data.py new file mode 100644 index 0000000000000000000000000000000000000000..2fdbd4dfab03a65c8770dd8b1a1a5ca5a14910fa --- /dev/null +++ b/gpu-sft/scripts/gpu_sft/prepare_sft_data.py @@ -0,0 +1,525 @@ +#!/usr/bin/env python3 +# ruff: noqa: E501 +"""Stage 1 of the GPU port of the Marin/Levanter Qwen3 self-instill SFT recipe. + +HF dataset (pinned revision) + -> Marin conversation transform (DEFAULT_TEXT_REPLACEMENTS) + -> Qwen3 chat template with {% generation %} assistant masking + -> Levanter greedy *contiguous* packing (max 64 segments / 32768 tokens) + -> local .npy memmaps + manifest.json + +Every step here is a deliberate port of the TPU pipeline. Where the TPU pipeline has a +bug or a surprising behaviour, this script reproduces it by default and exposes a flag to +turn it off. Read the "FIDELITY NOTES" block before changing anything. + +FIDELITY NOTES +-------------- +1. TEXT REPLACEMENTS (reproduced by default; `--text-replacements none` to disable). + marin/transform/conversation/transform_conversation.py applies + DEFAULT_TEXT_REPLACEMENTS = {"": "<|start_think|>", "": "<|end_think|>"} + to every message because multi_turn_adapter() leaves `replacements=None`. + Those two strings are Marin/Llama3 tokens, NOT Qwen3 tokens, so under the Qwen3 + tokenizer they encode as ordinary multi-token text. Consequence: the chat template's + `'' in content` reasoning-split branch never fires, `reasoning_content` stays + empty, and every FINAL assistant turn is rendered as + + <|im_start|>assistant\n\n\n\n\n<|start_think|>...<|end_think|>...<|im_end|>\n + + i.e. an EMPTY Qwen3 think block followed by the real reasoning in non-special text. + This is what the TPU checkpoints were trained on. Disabling it changes the token + layout and the assistant mask, and the resulting model will not match the TPU runs. + +2. DOCUMENT ORDER. Packing is greedy over *adjacent* documents in cache order. We use HF + dataset row order. The TPU cache order came from the zephyr shard writer; if that + interleaved shards, pack boundaries differ. Impact is limited: cross-document + attention is blocked and per-token loss weights are unchanged, so only the grouping + of documents into 32768-token windows (and hence the microbatch denominators) moves. + +3. VOCAB. We do NOT resize embeddings. The HF checkpoint already has vocab_size=151936 + while len(tokenizer) is 151669; Levanter padded the *tokenizer* up to the model. We + emit the 267 <|padding_i|> tokens only at HF-export time (see train_sft_qwen3.py) so + the exported tokenizer matches the TPU exports. Token ids are unaffected. + +4. NO SHUFFLE. The TPU experiment passes `shuffle=DATASET_SIZE` (an int) where levanter + expects bool | BlockShuffleConfig, which makes shuffling a silent no-op. The trainer + therefore walks packed examples in order and wraps around. Nothing to do here. + +Usage +----- + python prepare_sft_data.py \\ + --dataset-id fzzhang/qwen3_8b_hs_competition_depth2_probval_instill_n8_valredundancy5_round1 \\ + --revision <7-char-sha> \\ + --tokenizer Qwen/Qwen3-8B \\ + --max-seq-len 32768 \\ + --out /data/sft/hs_competition_depth2 \\ + --num-proc 16 + + # prove the encoder matches Levanter's segment-split encoder on 200 sampled docs + python prepare_sft_data.py ... --verify 200 + + # packing statistics only (answers "how many epochs is 2000 steps?") + python prepare_sft_data.py ... --stats-only +""" + +from __future__ import annotations + +import argparse +import hashlib +import json +import os +import re +import sys +import time +from itertools import chain +from typing import Any + +import numpy as np + +# ---------------------------------------------------------------------------------- +# Verbatim copy of marin experiments/qwen3_chat_template.py :: QWEN_3_CHAT_TEMPLATE +# (upstream Qwen3 template + {% generation %} tags around assistant content, tool-call +# bodies and the trailing '<|im_end|>\n'). +# ---------------------------------------------------------------------------------- +QWEN3_CHAT_TEMPLATE = r"""{%- if tools is defined and tools %} + {{- '<|im_start|>system\n' }} + {%- if messages[0].role == 'system' %} + {{- messages[0].content + '\n\n' }} + {%- endif %} + {{- "# Tools\n\nYou may call one or more functions to assist with the user query.\n\nYou are provided with function signatures within XML tags:\n" }} + {%- for tool in tools %} + {{- "\n" }} + {{- tool | tojson }} + {%- endfor %} + {{- "\n\n\nFor each function call, return a json object with function name and arguments within XML tags:\n\n{\"name\": , \"arguments\": }\n<|im_end|>\n" }} +{%- else %} + {%- if messages[0].role == 'system' %} + {{- '<|im_start|>system\n' + messages[0].content + '<|im_end|>\n' }} + {%- endif %} +{%- endif %} +{%- set ns = namespace(multi_step_tool=true, last_query_index=messages|length - 1) %} +{%- for message in messages[::-1] %} + {%- set index = (messages|length - 1) - loop.index0 %} + {%- if ns.multi_step_tool and message.role == "user" and message.content is string and not(message.content.startswith('') and message.content.endswith('')) %} + {%- set ns.multi_step_tool = false %} + {%- set ns.last_query_index = index %} + {%- endif %} +{%- endfor %} +{%- for message in messages %} + {%- if message.content is string %} + {%- set content = message.content %} + {%- else %} + {%- set content = '' %} + {%- endif %} + {%- if (message.role == "user") or (message.role == "system" and not loop.first) %} + {{- '<|im_start|>' + message.role + '\n' + content + '<|im_end|>' + '\n' }} + {%- elif message.role == "assistant" %} + {%- set reasoning_content = '' %} + {%- if message.reasoning_content is string %} + {%- set reasoning_content = message.reasoning_content %} + {%- else %} + {%- if '' in content %} + {%- set reasoning_content = content.split('')[0].rstrip('\n').split('')[-1].lstrip('\n') %} + {%- set content = content.split('')[-1].lstrip('\n') %} + {%- endif %} + {%- endif %} + {%- if loop.index0 > ns.last_query_index %} + {%- if loop.last or (not loop.last and reasoning_content) %} + {{- '<|im_start|>' + message.role + '\n' }}{% generation %}{{- '\n' + reasoning_content.strip('\n') + '\n\n\n' + content.lstrip('\n') }}{% endgeneration %} + {%- else %} + {{- '<|im_start|>' + message.role + '\n' }}{% generation %}{{- content }}{% endgeneration %} + {%- endif %} + {%- else %} + {{- '<|im_start|>' + message.role + '\n' }}{% generation %}{{- content }}{% endgeneration %} + {%- endif %} + {%- if message.tool_calls %} + {%- for tool_call in message.tool_calls %} + {%- if (loop.first and content) or (not loop.first) %} + {% generation %}{{- '\n' }}{% endgeneration %} + {%- endif %} + {%- if tool_call.function %} + {%- set tool_call = tool_call.function %} + {%- endif %} + {% generation %}{{- '\n{"name": "' }} + {{- tool_call.name }} + {{- '", "arguments": ' }} + {%- if tool_call.arguments is string %} + {{- tool_call.arguments }} + {%- else %} + {{- tool_call.arguments | tojson }} + {%- endif %} + {{- '}\n' }}{% endgeneration %} + {%- endfor %} + {%- endif %} + {% generation %}{{- '<|im_end|>\n' }}{% endgeneration %} + {%- elif message.role == "tool" %} + {%- if loop.first or (messages[loop.index0 - 1].role != "tool") %} + {{- '<|im_start|>user' }} + {%- endif %} + {{- '\n\n' }} + {{- content }} + {{- '\n' }} + {%- if loop.last or (messages[loop.index0 + 1].role != "tool") %} + {{- '<|im_end|>\n' }} + {%- endif %} + {%- endif %} +{%- endfor %} +{%- if add_generation_prompt %} + {{- '<|im_start|>assistant\n' }} + {%- if enable_thinking is defined and enable_thinking is false %} + {% generation %}{{- '\n\n\n\n' }}{% endgeneration %} + {%- endif %} +{%- endif %}""" + +DEFAULT_TEXT_REPLACEMENTS = {"": "<|start_think|>", "": "<|end_think|>"} + +# levanter/tokenizers.py sentinels, used only by --verify. +_GEN_START = "__MARIN_GEN_START_7f3a9c__" +_GEN_END = "__MARIN_GEN_END_7f3a9c__" + +MAX_SEGMENTS_PER_EXAMPLE = 64 # levanter datasets.py: 64 when pack is True +FORMAT_VERSION = 2 + + +# ---------------------------------------------------------------------------------- +# message normalisation (port of multi_turn_adapter + transform_row) +# ---------------------------------------------------------------------------------- +def normalize_messages( + row: dict[str, Any], + conversation_column: str, + role_key: str, + content_key: str, + replacements: dict[str, str] | None, +) -> list[dict[str, Any]]: + raw = row[conversation_column] + out: list[dict[str, Any]] = [] + for m in raw: + role = m[role_key] + content = m.get(content_key) + if isinstance(content, str) and replacements: + for old, new in replacements.items(): + content = content.replace(old, new) + msg: dict[str, Any] = {"role": role, "content": content} + # marin's OpenAIChatMessage carries these; the template only reads tool_calls. + if m.get("tool_calls"): + msg["tool_calls"] = m["tool_calls"] + out.append(msg) + return out + + +# ---------------------------------------------------------------------------------- +# encoding +# ---------------------------------------------------------------------------------- +def encode_conversation(tokenizer, messages: list[dict[str, Any]]) -> tuple[list[int], list[int]]: + """HF path: render once, tokenize once, map {% generation %} char spans to tokens.""" + enc = tokenizer.apply_chat_template( + messages, + chat_template=QWEN3_CHAT_TEMPLATE, + tokenize=True, + add_generation_prompt=False, + return_dict=True, + return_assistant_tokens_mask=True, + ) + return list(enc["input_ids"]), list(enc["assistant_masks"]) + + +def _levanter_env(): + """Recreate levanter._make_jinja_env([_GenerationSentinelExtension]).""" + import time as _time + + import jinja2 + import jinja2.ext + import jinja2.sandbox + + class _GenerationSentinelExtension(jinja2.ext.Extension): + tags = {"generation"} + + def parse(self, parser): + lineno = next(parser.stream).lineno + body = parser.parse_statements(("name:endgeneration",), drop_needle=True) + call = self.call_method("_wrap_generation", []) + return jinja2.nodes.CallBlock(call, [], [], body).set_lineno(lineno) + + @staticmethod + def _wrap_generation(caller): + return _GEN_START + caller() + _GEN_END + + def _raise(message): + raise jinja2.exceptions.TemplateError(message) + + # Match transformers.apply_chat_template, which renders with the DEFAULT lenient + # Undefined: the Qwen3 template references optional fields (message.tool_calls, + # message.reasoning_content) that plain {role, content} messages lack, and under + # StrictUndefined `{% if message.tool_calls %}` raises UndefinedError. Lenient + # Undefined evaluates missing attrs as falsy, exactly as the real encoding path does. + env = jinja2.sandbox.SandboxedEnvironment( + undefined=jinja2.Undefined, + trim_blocks=True, + lstrip_blocks=True, + extensions=[_GenerationSentinelExtension], + ) + env.globals.update({"raise_exception": _raise, "strftime_now": lambda fmt: _time.strftime(fmt)}) + return env + + +def encode_conversation_levanter(tokenizer, messages: list[dict[str, Any]]) -> tuple[list[int], list[int]]: + """Reference path: levanter's segment-split encoder, used only by --verify.""" + env = _levanter_env() + rendered = env.from_string(QWEN3_CHAT_TEMPLATE).render( + messages=messages, + add_generation_prompt=False, + bos_token=tokenizer.bos_token or "", + eos_token=tokenizer.eos_token or "", + ) + parts = re.split(f"({re.escape(_GEN_START)}|{re.escape(_GEN_END)})", rendered) + ids: list[int] = [] + mask: list[int] = [] + is_assistant = False + for part in parts: + if part == _GEN_START: + is_assistant = True + continue + if part == _GEN_END: + is_assistant = False + continue + if not part: + continue + seg = tokenizer.encode(part, add_special_tokens=False) + ids.extend(seg) + mask.extend([1 if is_assistant else 0] * len(seg)) + return ids, mask + + +# ---------------------------------------------------------------------------------- +# packing (line-by-line port of levanter/data/packing.py::pack_documents, +# slice_strategy="left", pad_with_zeros=True, max_segments_per_example=64) +# ---------------------------------------------------------------------------------- +def pack_document_ranges(lengths: np.ndarray, max_length: int, max_segments: int) -> list[tuple[int, int]]: + n = int(len(lengths)) + ranges: list[tuple[int, int]] = [] + i = 0 + while i < n: + start = i + total_segments = 0 + running = 0 + while i < n: + if total_segments + 1 > max_segments: + break + candidate = running + int(lengths[i]) + end_pack_after_this = False + if candidate > max_length: + if i == start: + # single oversized document: keep it and slice from the left + end_pack_after_this = True + else: + break + running = candidate + total_segments += 1 + i += 1 + if end_pack_after_this: + break + if i == start: # defensive; unreachable with max_segments >= 1 + i = start + 1 + ranges.append((start, i)) + return ranges + + +# ---------------------------------------------------------------------------------- +def build(args) -> None: + from datasets import load_dataset + from transformers import AutoTokenizer + + t0 = time.time() + replacements = DEFAULT_TEXT_REPLACEMENTS if args.text_replacements == "marin" else None + print(f"[prepare] text replacements: {replacements}", flush=True) + + tokenizer = AutoTokenizer.from_pretrained(args.tokenizer) + if not tokenizer.is_fast: + raise RuntimeError("A fast tokenizer is required for return_assistant_tokens_mask.") + + ds = load_dataset(args.dataset_id, revision=args.revision, split=args.split) + if args.limit: + ds = ds.select(range(min(args.limit, len(ds)))) + print(f"[prepare] loaded {len(ds)} rows from {args.dataset_id}@{args.revision}", flush=True) + + conversation_column = args.conversation_column + role_key, content_key = args.role_key, args.content_key + + def _map(batch, indices): + tok = tokenizer + ids_out, mask_out, len_out = [], [], [] + for k in range(len(indices)): + row = {conversation_column: batch[conversation_column][k]} + messages = normalize_messages(row, conversation_column, role_key, content_key, replacements) + ids, mask = encode_conversation(tok, messages) + if not any(mask): + raise ValueError( + f"Row {indices[k]} produced no assistant tokens. levanter's ChatProcessor raises here too." + ) + ids_out.append(ids) + mask_out.append(mask) + len_out.append(len(ids)) + return {"input_ids": ids_out, "assistant_mask": mask_out, "length": len_out} + + enc = ds.map( + _map, + batched=True, + batch_size=16, + with_indices=True, + num_proc=args.num_proc, + remove_columns=ds.column_names, + desc="tokenize", + ) + + lengths = np.asarray(enc["length"], dtype=np.int64) + L = args.max_seq_len + ranges = pack_document_ranges(lengths, L, MAX_SEGMENTS_PER_EXAMPLE) + n_packs = len(ranges) + + stats = { + "n_documents": int(len(lengths)), + "doc_tokens_total": int(lengths.sum()), + "doc_len_mean": float(lengths.mean()), + "doc_len_p50": float(np.percentile(lengths, 50)), + "doc_len_p90": float(np.percentile(lengths, 90)), + "doc_len_p99": float(np.percentile(lengths, 99)), + "doc_len_max": int(lengths.max()), + "docs_over_max_seq_len": int((lengths > L).sum()), + "tokens_truncated": int(np.clip(lengths - L, 0, None).sum()), + "n_packed_examples": n_packs, + "docs_per_pack_mean": float(len(lengths) / max(n_packs, 1)), + "pack_fill_ratio": float(min(lengths.sum(), n_packs * L) / max(n_packs * L, 1)), + } + for gb in (64, 128, 256): + stats[f"epochs_at_2000_steps_bs{gb}"] = round(2000 * gb / max(n_packs, 1), 3) + print("[prepare] stats: " + json.dumps(stats, indent=2), flush=True) + if args.stats_only: + return + + os.makedirs(args.out, exist_ok=True) + tokens = np.lib.format.open_memmap( + os.path.join(args.out, "tokens.npy"), mode="w+", dtype=np.uint32, shape=(n_packs, L) + ) + amask = np.lib.format.open_memmap( + os.path.join(args.out, "assistant_mask.npy"), mode="w+", dtype=np.uint8, shape=(n_packs, L) + ) + doc_lens_flat: list[int] = [] + pack_offsets = np.zeros(n_packs + 1, dtype=np.int64) + assistant_counts = np.zeros(n_packs, dtype=np.int64) + real_tokens = np.zeros(n_packs, dtype=np.int64) + + for p, (s, e) in enumerate(ranges): + rows = enc[s:e] + ids = list(chain.from_iterable(rows["input_ids"])) + msk = list(chain.from_iterable(rows["assistant_mask"])) + dlens = [len(x) for x in rows["input_ids"]] + if len(ids) > L: + assert e - s == 1, "levanter never packs two docs when one alone overflows" + ids, msk, dlens = ids[:L], msk[:L], [L] + n = len(ids) + tokens[p, :n] = np.asarray(ids, dtype=np.uint32) + tokens[p, n:] = 0 # pad_with_zeros=True + amask[p, :n] = np.asarray(msk, dtype=np.uint8) + amask[p, n:] = 0 + doc_lens_flat.extend(dlens) + pack_offsets[p + 1] = len(doc_lens_flat) + # levanter loss weight = roll(mask, -1) * not_last_mask => sum(mask[1:]) + assistant_counts[p] = int(sum(msk)) - int(msk[0]) + real_tokens[p] = n + if p % 2000 == 0: + print(f"[prepare] packed {p}/{n_packs}", flush=True) + + tokens.flush() + amask.flush() + np.save(os.path.join(args.out, "doc_lens.npy"), np.asarray(doc_lens_flat, dtype=np.int32)) + np.save(os.path.join(args.out, "pack_offsets.npy"), pack_offsets) + np.save(os.path.join(args.out, "assistant_counts.npy"), assistant_counts) + np.save(os.path.join(args.out, "real_tokens.npy"), real_tokens) + + stats["assistant_tokens_total"] = int(assistant_counts.sum()) + stats["assistant_token_fraction"] = float(assistant_counts.sum() / max(real_tokens.sum(), 1)) + cfg = { + "format_version": FORMAT_VERSION, + "dataset_id": args.dataset_id, + "revision": args.revision, + "split": args.split, + "tokenizer": args.tokenizer, + "max_seq_len": L, + "max_segments_per_example": MAX_SEGMENTS_PER_EXAMPLE, + "text_replacements": replacements, + "chat_template_sha256": hashlib.sha256(QWEN3_CHAT_TEMPLATE.encode()).hexdigest(), + "tokenizer_len_unpadded": len(tokenizer), + "n_packs": n_packs, + "stats": stats, + "built_at": time.strftime("%Y-%m-%dT%H:%M:%S"), + } + with open(os.path.join(args.out, "manifest.json"), "w") as f: + json.dump(cfg, f, indent=2) + print(f"[prepare] wrote {n_packs} packed examples to {args.out} in {time.time() - t0:.0f}s", flush=True) + + +# ---------------------------------------------------------------------------------- +def verify(args) -> None: + """Assert the HF encoder == levanter's segment-split encoder on sampled documents.""" + from datasets import load_dataset + from transformers import AutoTokenizer + + replacements = DEFAULT_TEXT_REPLACEMENTS if args.text_replacements == "marin" else None + tokenizer = AutoTokenizer.from_pretrained(args.tokenizer) + ds = load_dataset(args.dataset_id, revision=args.revision, split=args.split) + rng = np.random.default_rng(0) + idx = rng.choice(len(ds), size=min(args.verify, len(ds)), replace=False) + + bad = 0 + for j, i in enumerate(idx): + messages = normalize_messages( + ds[int(i)], args.conversation_column, args.role_key, args.content_key, replacements + ) + a_ids, a_mask = encode_conversation(tokenizer, messages) + b_ids, b_mask = encode_conversation_levanter(tokenizer, messages) + if a_ids != b_ids or a_mask != b_mask: + bad += 1 + n = min(len(a_ids), len(b_ids)) + first = next((k for k in range(n) if a_ids[k] != b_ids[k] or a_mask[k] != b_mask[k]), n) + print( + f"[verify] MISMATCH row={int(i)} len_hf={len(a_ids)} len_lev={len(b_ids)} first_diff={first}", + file=sys.stderr, + ) + print(f" hf : {a_ids[max(0, first - 5) : first + 5]}", file=sys.stderr) + print(f" lev: {b_ids[max(0, first - 5) : first + 5]}", file=sys.stderr) + if j == 0: + txt = tokenizer.decode(a_ids[:120]) + print(f"[verify] sample head: {txt!r}", flush=True) + print(f"[verify] assistant tokens: {sum(a_mask)}/{len(a_mask)}", flush=True) + if bad: + raise SystemExit(f"[verify] {bad}/{len(idx)} documents disagree with the levanter encoder") + print(f"[verify] OK: {len(idx)} documents byte-identical to the levanter encoder", flush=True) + + +def main() -> None: + p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + p.add_argument("--dataset-id", required=True) + p.add_argument("--revision", required=True, help="7-char HF sha, pinned exactly like instruction_datasets.py") + p.add_argument("--split", default="train") + p.add_argument("--tokenizer", default="Qwen/Qwen3-8B") + p.add_argument("--max-seq-len", type=int, default=32768) + p.add_argument("--out", default=None) + p.add_argument("--num-proc", type=int, default=max(1, (os.cpu_count() or 8) // 2)) + p.add_argument("--limit", type=int, default=0) + p.add_argument("--conversation-column", default="messages") + p.add_argument("--role-key", default="role") + p.add_argument("--content-key", default="content") + p.add_argument("--text-replacements", choices=("marin", "none"), default="marin") + p.add_argument("--verify", type=int, default=0, help="verify N sampled docs against the levanter encoder and exit") + p.add_argument("--stats-only", action="store_true") + args = p.parse_args() + + if args.verify: + verify(args) + return + if not args.out and not args.stats_only: + p.error("--out is required unless --stats-only/--verify") + build(args) + + +if __name__ == "__main__": + main() diff --git a/gpu-sft/scripts/gpu_sft/push_exports.py b/gpu-sft/scripts/gpu_sft/push_exports.py new file mode 100644 index 0000000000000000000000000000000000000000..06acca04672461c7405a122493c1464a257d17dc --- /dev/null +++ b/gpu-sft/scripts/gpu_sft/push_exports.py @@ -0,0 +1,98 @@ +#!/usr/bin/env python +"""Watch a training run's hf/ dir and upload each new step-N export to a private HF +repo, so ephemeral-disk checkpoints survive a machine kill and are eval-able anywhere. + +train_sft_qwen3.py writes a bf16 hf/step-N/ every steps_per_hf_export. This uploads each +one (once it has finished writing) to under path_in_repo=step-N/. One repo per +experiment; steps are subfolders. Idempotent via a local .pushed_steps.json marker, so a +restart of this watcher (or of the whole run) re-uploads nothing. + +Run in the background alongside training (run_sft.sh does this for you): + python push_exports.py --out --repo fzzhang/_sft --until-step 2000 + +HF auth: uses your `huggingface-cli login` token (~/.cache/huggingface/token). +""" +from __future__ import annotations + +import argparse +import json +import time +from pathlib import Path + + +def step_dirs(hf_dir: Path) -> dict[int, Path]: + out: dict[int, Path] = {} + if not hf_dir.is_dir(): + return out + for d in hf_dir.glob("step-*"): + if d.is_dir(): + try: + out[int(d.name.split("-", 1)[1])] = d + except (IndexError, ValueError): + pass + return out + + +def is_stable(d: Path, stable_sec: float) -> bool: + """True if no file in the dir was modified within the last stable_sec (write finished).""" + files = [p for p in d.rglob("*") if p.is_file()] + if not files: + return False + return (time.time() - max(p.stat().st_mtime for p in files)) > stable_sec + + +def main() -> None: + ap = argparse.ArgumentParser(description="Upload hf/step-N exports to HF as they land") + ap.add_argument("--out", required=True, help="training run dir (contains hf/)") + ap.add_argument("--repo", required=True, help="target HF repo id, e.g. fzzhang/_sft") + ap.add_argument("--poll", type=int, default=120, help="seconds between scans") + ap.add_argument("--stable-sec", type=float, default=60.0, + help="a step dir must be idle this long before upload (guards partial writes)") + ap.add_argument("--every", type=int, default=1, + help="upload every Nth export (exports are every 100 steps; 1=all). " + "The --until-step export is always uploaded regardless.") + ap.add_argument("--until-step", type=int, default=0, + help="exit after this step is pushed (0 = run forever / until killed)") + ap.add_argument("--once", action="store_true", + help="scan once, push all unpushed steps (ignore --stable-sec), then exit") + args = ap.parse_args() + + from huggingface_hub import HfApi + + api = HfApi() + api.create_repo(args.repo, repo_type="model", private=True, exist_ok=True) + + hf_dir = Path(args.out) / "hf" + marker = Path(args.out) / ".pushed_steps.json" + pushed = set(json.loads(marker.read_text())) if marker.exists() else set() + print(f"[push] {hf_dir} -> {args.repo} (private); already pushed: {sorted(pushed)}", flush=True) + + while True: + for step in sorted(step_dirs(hf_dir)): + if step in pushed: + continue + if args.every > 1 and (step // 100) % args.every != 0 and step != args.until_step: + continue + d = step_dirs(hf_dir)[step] + if not args.once and not is_stable(d, args.stable_sec): + continue # still being written; try next poll + print(f"[push] uploading step-{step} ...", flush=True) + try: + api.upload_folder(folder_path=str(d), repo_id=args.repo, repo_type="model", + path_in_repo=f"step-{step}", commit_message=f"add step-{step}") + pushed.add(step) + marker.write_text(json.dumps(sorted(pushed))) + print(f"[push] DONE step-{step} -> {args.repo}/step-{step}", flush=True) + except Exception as e: # network / rate-limit / partial — retry next poll + print(f"[push] FAILED step-{step}: {type(e).__name__}: {e} (retry next poll)", flush=True) + if args.once: + print("[push] --once sweep complete; exiting.", flush=True) + return + if args.until_step and args.until_step in pushed: + print(f"[push] final step-{args.until_step} pushed; exiting.", flush=True) + return + time.sleep(args.poll) + + +if __name__ == "__main__": + main() diff --git a/gpu-sft/scripts/gpu_sft/run_sft.sh b/gpu-sft/scripts/gpu_sft/run_sft.sh new file mode 100644 index 0000000000000000000000000000000000000000..c651578518d93412c47a7f22a9c4773da958b3d4 --- /dev/null +++ b/gpu-sft/scripts/gpu_sft/run_sft.sh @@ -0,0 +1,131 @@ +#!/usr/bin/env bash +# Copyright The Marin Authors +# SPDX-License-Identifier: Apache-2.0 +# +# Self-healing launcher for train_sft_qwen3.py on ONE 8xH100 node. +# +# bash run_sft.sh [num_train_steps] +# +# Loops torchrun: on a transient crash (e.g. an NCCL watchdog collective timeout) +# it waits 30s and relaunches with --resume auto (continues from checkpoints-temp) +# and a FIXED --wandb-id (continues the SAME W&B run instead of forking a new one). +# Exits 0 when training reaches the step target. Bails out if it fast-fails +# repeatedly -- that means a real error (bad config / OOM / disk), not a blip. +# +# Data must already be prepared at $DATA_ROOT/ by prepare_sft_data.py. +# For a FRESH run, delete $RUNS_ROOT/ first (else --resume auto continues +# the old checkpoints). Launch under nohup so it survives your shell: +# +# nohup bash gpu-sft/scripts/gpu_sft/run_sft.sh science_depth4v3_nofilter \ +# > /opt/tiger/lmc_muon/gpu-sft-runs/science_depth4v3_nofilter/wrapper.log 2>&1 & +# +# NOTE: no `set -e` on purpose -- the whole point is to survive torchrun failures. +set -uo pipefail + +EXP="${1:?usage: run_sft.sh [num_train_steps]}" +NUM_STEPS="${2:-2000}" +NPROC="${NPROC:-8}" + +DATA_ROOT="${DATA_ROOT:-/opt/tiger/lmc_muon/gpu-sft-data}" +RUNS_ROOT="${RUNS_ROOT:-/opt/tiger/lmc_muon/gpu-sft-runs}" +REPO="${REPO:-/opt/tiger/lmc_muon/self-verified-distillation}" +WANDB_PROJECT="${WANDB_PROJECT:-marin-instill-gpu}" +WANDB_ID="${WANDB_ID:-$EXP}" # override to FORCE a fresh W&B run id (avoids reattaching to a prior partial run of the same name) + +# Durability: async HDFS mirror of hf/ exports + native checkpoints/. Local NVMe stays the +# training --out (fast, no FUSE I/O stall that could trip the NCCL watchdog); this only +# copies in the background so a machine kill can't erase progress. Set MIRROR=0 to disable. +MIRROR="${MIRROR:-1}" +HDFS_ROOT="${HDFS_ROOT:-/mnt/hdfs/fangzhao_writable/marin-sft}" +MIRROR_EVERY="${MIRROR_EVERY:-300}" + +DATA="$DATA_ROOT/$EXP" +OUT="$RUNS_ROOT/$EXP" +LOG="$OUT/train.log" + +MAX_RESTARTS="${MAX_RESTARTS:-30}" # absolute cap on relaunches +FAST_FAIL_SEC="${FAST_FAIL_SEC:-180}" # a run dying sooner than this = a hard fail +MAX_FAST_FAILS="${MAX_FAST_FAILS:-3}" # bail after this many CONSECUTIVE fast fails + +export HF_HOME="${HF_HOME:-/opt/tiger/lmc_muon/hf_cache}" +export TOKENIZERS_PARALLELISM=false +export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True +export WANDB_MODE="${WANDB_MODE:-online}" +export NCCL_DEBUG="${NCCL_DEBUG:-WARN}" + +if [[ ! -d "$DATA" ]]; then + echo "[run_sft] ERROR: prepared data not found at $DATA -- run prepare_sft_data.py first" >&2 + exit 2 +fi +mkdir -p "$OUT" +cd "$REPO" || { echo "[run_sft] ERROR: cannot cd $REPO" >&2; exit 2; } + +echo "[run_sft] exp=$EXP steps=$NUM_STEPS nproc=$NPROC wandb-id=$WANDB_ID out=$OUT" | tee -a "$LOG" + +# --- durability: background HDFS mirror (hf/ exports + native checkpoints/; skip the +# rolling checkpoints-temp/ which is overwritten and not rsync-snapshot-safe) ---------- +DEST="$HDFS_ROOT/$EXP" +mirror_once() { rsync -a --exclude 'checkpoints-temp' "$OUT/" "$DEST/" 2>>"$OUT/mirror.log"; } +mirror_pid=""; push_pid="" +if [ "$MIRROR" = "1" ] && [ -d "$(dirname "$HDFS_ROOT")" ]; then + mkdir -p "$DEST" + echo "[run_sft] HDFS mirror ON: $OUT -> $DEST every ${MIRROR_EVERY}s" | tee -a "$LOG" + ( while true; do sleep "$MIRROR_EVERY"; mirror_once; done ) & + mirror_pid=$! +elif [ -n "${PUSH_HF_REPO:-}" ]; then + # No HDFS (e.g. A800 boxes): push each hf/step-N to a private HF repo instead. + echo "[run_sft] no HDFS -> HF push: hf/step-N -> $PUSH_HF_REPO" | tee -a "$LOG" + nohup python gpu-sft/scripts/gpu_sft/push_exports.py --out "$OUT" --repo "$PUSH_HF_REPO" >> "$OUT/push.log" 2>&1 & + push_pid=$! +else + echo "[run_sft] WARNING: NO durability -- no HDFS mount and no PUSH_HF_REPO; checkpoints on ephemeral disk only!" | tee -a "$LOG" +fi +trap '[ -n "${mirror_pid:-}" ] && kill "$mirror_pid" 2>/dev/null; [ -n "${push_pid:-}" ] && kill "$push_pid" 2>/dev/null' EXIT + +fast_fails=0 +for attempt in $(seq 1 "$MAX_RESTARTS"); do + echo "[run_sft] === attempt $attempt/$MAX_RESTARTS $(date -u) ===" | tee -a "$LOG" + start=$(date +%s) + + torchrun --standalone --nproc_per_node="$NPROC" \ + gpu-sft/scripts/gpu_sft/train_sft_qwen3.py \ + --data "$DATA" --out "$OUT" --recipe qwen3-8b \ + --train-batch-size 64 --per-device-batch 1 --loss-groups 2 \ + --num-train-steps "$NUM_STEPS" \ + --steps-per-hf-export 100 --steps-per-checkpoint 500 \ + --loss-chunk-size 2048 --hf-export-dtype bfloat16 \ + --init-mode all_ranks --resume auto \ + --wandb-project "$WANDB_PROJECT" --wandb-name "$EXP" --wandb-id "$WANDB_ID" \ + --log-every 1 >> "$LOG" 2>&1 + rc=$? + dur=$(( $(date +%s) - start )) + + if [[ $rc -eq 0 ]]; then + echo "[run_sft] SUCCESS: $EXP finished (rc=0) after $attempt attempt(s)" | tee -a "$LOG" + if [ -n "$mirror_pid" ]; then echo "[run_sft] final HDFS mirror..." | tee -a "$LOG"; mirror_once; fi + if [ -n "$push_pid" ]; then + echo "[run_sft] final HF push..." | tee -a "$LOG" + kill "$push_pid" 2>/dev/null + python gpu-sft/scripts/gpu_sft/push_exports.py --out "$OUT" --repo "$PUSH_HF_REPO" --once >> "$OUT/push.log" 2>&1 + fi + exit 0 + fi + + echo "[run_sft] exit rc=$rc after ${dur}s" | tee -a "$LOG" + if [[ $dur -lt $FAST_FAIL_SEC ]]; then + fast_fails=$((fast_fails + 1)) + echo "[run_sft] fast-fail $fast_fails/$MAX_FAST_FAILS (died <${FAST_FAIL_SEC}s -> likely a real error, not transient)" | tee -a "$LOG" + if [[ $fast_fails -ge $MAX_FAST_FAILS ]]; then + echo "[run_sft] ABORT: $MAX_FAST_FAILS consecutive fast fails -- inspect $LOG (config/OOM/disk, not a blip)" | tee -a "$LOG" + exit 1 + fi + else + fast_fails=0 # ran a while then died -> treat as transient, reset the counter + fi + + echo "[run_sft] relaunching in 30s (will --resume auto from checkpoints-temp)..." | tee -a "$LOG" + sleep 30 +done + +echo "[run_sft] ABORT: hit MAX_RESTARTS=$MAX_RESTARTS without finishing" | tee -a "$LOG" +exit 1 diff --git a/gpu-sft/scripts/gpu_sft/train_sft_qwen3.py b/gpu-sft/scripts/gpu_sft/train_sft_qwen3.py new file mode 100644 index 0000000000000000000000000000000000000000..082f871f0b4bdb75b107a43e27f8f146e19f318e --- /dev/null +++ b/gpu-sft/scripts/gpu_sft/train_sft_qwen3.py @@ -0,0 +1,820 @@ +#!/usr/bin/env python3 +# ruff: noqa: E501 +"""Stage 2: single-node/multi-node NVIDIA GPU port of the Marin/Levanter Qwen3 SFT recipe. + +PyTorch + HF transformers + torch FSDP. No JAX, no GCS, no Iris, no marin credentials. +Reads the packed dataset produced by prepare_sft_data.py, writes native (resumable) +checkpoints and HF-format exports to local disk on the same cadence as the TPU run, so +the existing evalchemy tooling can consume `/hf/step-N` unchanged. + +Launch (8 GPUs, one node): + torchrun --standalone --nproc_per_node=8 train_sft_qwen3.py --data /data/sft/foo --out /data/runs/foo + +Plan only, no GPUs needed: + python train_sft_qwen3.py --data /data/sft/foo --out /tmp/x --dry-run + +WHAT IS REPRODUCED EXACTLY +-------------------------- +* optimiser chain order: global-norm clip -> Adam(b1,b2,eps) -> decoupled wd -> scale(-lr) +* weight-decay mask: no decay on *norm*, embeddings and biases; lm_head IS decayed +* LR: linear warmup 0->lr over int(warmup*T) steps, constant, then cosine to lr*min_lr_ratio + over int(decay*T) steps (clamped). LR at optimiser step 0 is exactly 0.0. +* mixed precision: fp32 master params / bf16 compute / fp32 gradient reduction +* loss: per-microbatch token-weighted mean over assistant tokens, then arithmetic mean + across the `--loss-groups` microbatches of an optimiser step (levanter grad_accum uses + ReductionType.MEAN over 2 microbatches of 32 on v5p-64). No z-loss. +* RoPE positions are CONTIGUOUS 0..L-1 across a packed example (levanter passes + `pos_ids = arange(Pos)`; it does NOT reset per document), while attention is still + blocked across documents with varlen cu_seq_lens. Most HF/TRL packing implementations + reset position_ids per document -- that is a real divergence and we do not do it. +* data order: sequential over packed examples with wraparound (shuffle is a silent no-op + in the TPU config), so `--data-seed` only matters if you pass --shuffle. +* checkpoint/export naming: levanter's StepInfo.step == completed_steps - 1, and hooks + fire when `step > 1 and step % every == 0`, plus a forced final hook. A 2000-step run + therefore writes hf/step-100 ... hf/step-1900 and a final hf/step-1999. + +WHAT DIFFERS (and why) -- see also the handoff doc: +* per-device microbatch is 1 packed example (not 32-across-32-chips); grouping into + `--loss-groups` groups restores the exact loss normalisation, but the *gradient* + accumulation order differs, so results are not bitwise identical. +* attention kernel is FlashAttention-2 (varlen), not the TPU splash/Pallas kernel. +* cross-entropy is a chunked torch implementation, not levanter's fused Pallas CE. +* HF export defaults to bf16 here (the TPU export was fp32; vLLM casts to bf16 anyway). +""" + +from __future__ import annotations + +import argparse +import contextlib +import json +import math +import os +import shutil +import time +from dataclasses import dataclass, replace + +import numpy as np + +MODEL_VOCAB_SIZE = 151936 # Qwen3 config.vocab_size for 1.7B / 4B / 8B / 32B + + +# ---------------------------------------------------------------------------------- +# recipes (frozen per model size, copied from the marin experiment files) +# ---------------------------------------------------------------------------------- +@dataclass +class Recipe: + model_id: str + learning_rate: float + weight_decay: float + warmup: float + decay: float + min_lr_ratio: float + max_grad_norm: float + beta1: float + beta2: float + eps: float + train_batch_size: int + num_train_steps: int + max_seq_len: int + steps_per_checkpoint: int + steps_per_hf_export: int + lr_schedule: str = "cosine" + + +RECIPES: dict[str, Recipe] = { + # experiments/exp_sft_qwen3_8b_selfinstill_*_2k_lr5e6_wd01.py + "qwen3-8b": Recipe( + model_id="Qwen/Qwen3-8B", + learning_rate=5e-6, + weight_decay=0.01, + warmup=0.05, + decay=0.9, + min_lr_ratio=0.1, + max_grad_norm=1.0, + beta1=0.9, + beta2=0.999, + eps=1e-8, + train_batch_size=64, + num_train_steps=2000, + max_seq_len=32768, + steps_per_checkpoint=20, + steps_per_hf_export=100, + ), + # experiments/exp_*_sft_qwen3_4b_* (OT4 recipe: higher LR, no wd, tight clip) + "qwen3-4b": Recipe( + model_id="Qwen/Qwen3-4B", + learning_rate=2e-5, + weight_decay=0.0, + warmup=0.03, + decay=0.9, + min_lr_ratio=0.1, + max_grad_norm=0.2, + beta1=0.9, + beta2=0.999, + eps=1e-8, + train_batch_size=128, + num_train_steps=0, # 0 => ceil(8 * n_packs / batch); override with --num-train-steps + max_seq_len=32768, + steps_per_checkpoint=20, + steps_per_hf_export=100, + ), + "qwen3-1.7b": Recipe( + model_id="Qwen/Qwen3-1.7B", + learning_rate=2e-5, + weight_decay=0.0, + warmup=0.03, + decay=0.9, + min_lr_ratio=0.1, + max_grad_norm=0.2, + beta1=0.9, + beta2=0.999, + eps=1e-8, + train_batch_size=128, + num_train_steps=0, + max_seq_len=32768, + steps_per_checkpoint=20, + steps_per_hf_export=100, + ), +} + + +# ---------------------------------------------------------------------------------- +# LR schedule (port of levanter OptimizerConfig.lr_scheduler for a single cycle) +# ---------------------------------------------------------------------------------- +def _frac_or_steps(v: float, total: int) -> int: + if v < 0.0 or (v > 1.0 and v % 1 != 0): + raise ValueError(f"Invalid fraction {v}") + return int(v * total) if v <= 1.0 else int(v) + + +def make_lr_multiplier(recipe: Recipe, num_train_steps: int): + warmup_steps = min(_frac_or_steps(recipe.warmup, num_train_steps), num_train_steps) + max_decay = max(num_train_steps - warmup_steps, 0) + decay_steps = min(max(_frac_or_steps(recipe.decay, num_train_steps), 0), max_decay) + stable_steps = num_train_steps - warmup_steps - decay_steps + alpha = recipe.min_lr_ratio + sched = recipe.lr_schedule + + def fn(step: int) -> float: + if warmup_steps > 0 and step < warmup_steps: + return step / warmup_steps # optax.linear_schedule(0.0, lr, warmup_steps) + if step < warmup_steps + stable_steps: + return 1.0 + if decay_steps == 0: + return 1.0 + t = min(step - warmup_steps - stable_steps, decay_steps) + if sched == "cosine": + cos = 0.5 * (1.0 + math.cos(math.pi * t / decay_steps)) + return (1.0 - alpha) * cos + alpha + if sched == "linear": + return 1.0 + (alpha - 1.0) * (t / decay_steps) + if sched == "constant": + return 1.0 + raise ValueError(f"unsupported lr_schedule {sched}") + + return fn, (warmup_steps, stable_steps, decay_steps) + + +# ---------------------------------------------------------------------------------- +# weight decay mask (port of AdamConfig.build_weight_decay_mask reasonable_default) +# ---------------------------------------------------------------------------------- +def is_no_decay(name: str) -> bool: + """levanter excludes LayerNorm/RMSNorm/RmsNorm/Embedding modules and *.bias. + + lm_head is an ordinary Linear in levanter, so it IS decayed. With tied embeddings + (1.7B/4B) the single shared tensor is named model.embed_tokens.weight and is excluded, + which matches levanter (lm_head is None when tie_word_embeddings=True). + """ + if name.endswith(".bias"): + return True + if "norm" in name.lower(): # input_layernorm, post_attention_layernorm, model.norm, q_norm, k_norm + return True + if name.endswith("embed_tokens.weight"): + return True + return False + + +# ---------------------------------------------------------------------------------- +# packed dataset +# ---------------------------------------------------------------------------------- +class PackedSFTData: + def __init__(self, path: str): + with open(os.path.join(path, "manifest.json")) as f: + self.manifest = json.load(f) + self.tokens = np.load(os.path.join(path, "tokens.npy"), mmap_mode="r") + self.mask = np.load(os.path.join(path, "assistant_mask.npy"), mmap_mode="r") + self.doc_lens = np.load(os.path.join(path, "doc_lens.npy")) + self.pack_offsets = np.load(os.path.join(path, "pack_offsets.npy")) + self.assistant_counts = np.load(os.path.join(path, "assistant_counts.npy")) + self.max_seq_len = int(self.manifest["max_seq_len"]) + self.n_packs = int(self.tokens.shape[0]) + assert self.tokens.shape[1] == self.max_seq_len + + def doc_lengths(self, i: int) -> list[int]: + s, e = int(self.pack_offsets[i]), int(self.pack_offsets[i + 1]) + return [int(x) for x in self.doc_lens[s:e]] + + +def collate(data: PackedSFTData, indices, device): + """Flatten `indices` packed examples into one padding-free row of B*L tokens. + + Returns tensors shaped [1, B*L] plus varlen metadata. batch dim is always 1 so the + HF flash-attention path uses flash_attn_varlen_func with our explicit cu_seq_lens. + """ + import torch + + L = data.max_seq_len + ids = np.concatenate([np.asarray(data.tokens[i], dtype=np.int64) for i in indices]) + msk = np.concatenate([np.asarray(data.mask[i], dtype=np.int64) for i in indices]) + labels = np.where(msk == 1, ids, -100) + for b in range(len(indices)): + labels[b * L] = -100 # levanter's not_last_mask: position L-1 of the PREVIOUS example + targets = np.empty_like(labels) + targets[:-1] = labels[1:] + targets[-1] = -100 + + seg_lens: list[int] = [] + for i in indices: + dl = data.doc_lengths(i) + seg_lens.extend(dl) + pad = L - int(sum(dl)) + if pad > 0: + seg_lens.append(pad) # padding is its own attention segment + cu = np.concatenate([[0], np.cumsum(np.asarray(seg_lens, dtype=np.int64))]).astype(np.int32) + + pos = np.tile(np.arange(L, dtype=np.int64), len(indices)) # contiguous per example, like levanter + n_targets = int((targets != -100).sum()) + return { + "input_ids": torch.from_numpy(ids).unsqueeze(0).to(device, non_blocking=True), + "position_ids": torch.from_numpy(pos).unsqueeze(0).to(device, non_blocking=True), + "targets": torch.from_numpy(targets).unsqueeze(0).to(device, non_blocking=True), + "cu_seq_lens": torch.from_numpy(cu).to(device, non_blocking=True), + "max_length": int(max(seg_lens)), + }, n_targets + + +# ---------------------------------------------------------------------------------- +# model wrapper with chunked (memory-bounded) cross entropy +# ---------------------------------------------------------------------------------- +def _build_sft_module(): + import torch + import torch.nn as nn + import torch.nn.functional as F + from torch.utils.checkpoint import checkpoint + + def _ce_chunk(hidden_chunk, target_chunk, lm_head): + logits = lm_head(hidden_chunk).float() + return F.cross_entropy(logits, target_chunk, ignore_index=-100, reduction="sum") + + class SFTModule(nn.Module): + """hf_model + a chunked lm_head/CE that never materialises [L, 151936] logits.""" + + def __init__(self, hf_model, loss_chunk_size: int, checkpoint_loss: bool = True): + super().__init__() + self.hf_model = hf_model + self.loss_chunk_size = loss_chunk_size + self.checkpoint_loss = checkpoint_loss + + def forward(self, input_ids, position_ids, targets, cu_seq_lens, max_length, return_hidden: bool = False): + out = self.hf_model.model( + input_ids=input_ids, + position_ids=position_ids, + use_cache=False, + cu_seq_lens_q=cu_seq_lens, + cu_seq_lens_k=cu_seq_lens, + max_length_q=max_length, + max_length_k=max_length, + ) + hidden = out.last_hidden_state[0] + if return_hidden: + return hidden + tgt = targets[0] + total = torch.zeros((), device=hidden.device, dtype=torch.float32) + n = hidden.shape[0] + cs = self.loss_chunk_size or n + for s in range(0, n, cs): + e = min(s + cs, n) + if self.checkpoint_loss and self.training: + part = checkpoint(_ce_chunk, hidden[s:e], tgt[s:e], self.hf_model.lm_head, use_reentrant=False) + else: + part = _ce_chunk(hidden[s:e], tgt[s:e], self.hf_model.lm_head) + total = total + part + return total + + return SFTModule + + +# ---------------------------------------------------------------------------------- +# distributed helpers +# ---------------------------------------------------------------------------------- +def dist_info(): + return ( + int(os.environ.get("RANK", 0)), + int(os.environ.get("LOCAL_RANK", 0)), + int(os.environ.get("WORLD_SIZE", 1)), + ) + + +def log0(rank: int, msg: str) -> None: + if rank == 0: + print(msg, flush=True) + + +def fmt_hms(seconds: float) -> str: + seconds = max(int(seconds), 0) + h, rem = divmod(seconds, 3600) + m, s = divmod(rem, 60) + return f"{h}:{m:02d}:{s:02d}" + + +def detect_grad_reduce_factor(device, world_size: int, rank: int) -> float: + """FSDP averages gradients over the DP group; loss scales must undo that. + + Empirically determined with a 1-layer probe so a torch behaviour change cannot + silently rescale the effective learning rate by world_size. + """ + import torch + import torch.nn as nn + from torch.distributed.fsdp import FullyShardedDataParallel as FSDP + from torch.distributed.fsdp import ShardingStrategy + + if world_size == 1: + return 1.0 + lin = nn.Linear(64, 64, bias=False) + with torch.no_grad(): + lin.weight.zero_() + probe = FSDP( + lin, + sharding_strategy=ShardingStrategy.FULL_SHARD, + device_id=device, + use_orig_params=True, + ) + x = torch.ones(1, 64, device=device) + probe(x).sum().backward() + with FSDP.summon_full_params(probe, with_grads=True, writeback=False): + observed = float(lin.weight.grad.detach().float().mean().item()) + del probe, lin + if abs(observed - 1.0) < 1e-3: + log0(rank, "[init] FSDP averages gradients across ranks -> loss scale multiplied by world_size") + return float(world_size) + if abs(observed - world_size) < 1e-3: + log0(rank, "[init] FSDP sums gradients across ranks -> no loss rescale") + return 1.0 + raise RuntimeError(f"Unexpected FSDP gradient reduction: mean grad {observed}, world_size {world_size}") + + +def verify_document_isolation(model, device, rank: int) -> None: + """Prove cross-document attention is actually blocked before burning GPU-days. + + If transformers stops forwarding cu_seq_lens_* to the attention kernel this test + fails loudly instead of silently training with cross-document leakage. + """ + import torch + + L, d0 = 64, 32 + ids = torch.randint(0, 1000, (1, L), device=device) + pos = torch.arange(L, device=device).unsqueeze(0) + cu = torch.tensor([0, d0, L], dtype=torch.int32, device=device) + tgt = torch.full((1, L), -100, dtype=torch.long, device=device) + was_training = model.training + model.eval() + with torch.no_grad(): + h1 = model(ids, pos, tgt, cu, d0, return_hidden=True)[:d0].clone() + ids2 = ids.clone() + ids2[0, d0:] = torch.randint(0, 1000, (L - d0,), device=device) + h2 = model(ids2, pos, tgt, cu, d0, return_hidden=True)[:d0].clone() + if was_training: + model.train() + diff = (h1 - h2).abs().max().item() + if diff != 0.0: + raise RuntimeError( + "Cross-document attention is NOT blocked (max hidden-state delta " + f"{diff:.3e} after perturbing the second document). The varlen cu_seq_lens " + "kwargs are not reaching flash-attention. Check the transformers version " + "and that attn_implementation='flash_attention_2'." + ) + log0(rank, "[init] cross-document attention isolation verified (delta == 0)") + + +# ---------------------------------------------------------------------------------- +def build_model(args, recipe: Recipe, rank: int, local_rank: int, world_size: int): + import torch + from torch.distributed.fsdp import BackwardPrefetch, MixedPrecision, ShardingStrategy + from torch.distributed.fsdp import FullyShardedDataParallel as FSDP + from torch.distributed.fsdp.wrap import ModuleWrapPolicy + from transformers import AutoConfig, AutoModelForCausalLM + + config = AutoConfig.from_pretrained(recipe.model_id) + config.use_cache = False + if config.vocab_size != MODEL_VOCAB_SIZE: + log0(rank, f"[init] WARNING model vocab_size={config.vocab_size}, expected {MODEL_VOCAB_SIZE}") + + if args.init_mode == "rank0_meta" and world_size > 1: + if rank == 0: + hf_model = AutoModelForCausalLM.from_pretrained( + recipe.model_id, dtype=torch.float32, attn_implementation=args.attn_impl, low_cpu_mem_usage=True + ) + else: + with torch.device("meta"): + hf_model = AutoModelForCausalLM.from_config(config, attn_implementation=args.attn_impl) + else: + hf_model = AutoModelForCausalLM.from_pretrained( + recipe.model_id, dtype=torch.float32, attn_implementation=args.attn_impl, low_cpu_mem_usage=True + ) + hf_model.config.use_cache = False + hf_model.gradient_checkpointing_enable(gradient_checkpointing_kwargs={"use_reentrant": False}) + + sft = _build_sft_module()(hf_model, args.loss_chunk_size, checkpoint_loss=not args.no_loss_checkpoint) + + layer_cls = type(hf_model.model.layers[0]) + wrap_classes = {layer_cls} + if args.wrap_embeddings: + wrap_classes.add(torch.nn.Embedding) + + mp = MixedPrecision( + param_dtype=torch.bfloat16, + reduce_dtype=torch.float32 if args.grad_reduce_dtype == "float32" else torch.bfloat16, + buffer_dtype=torch.float32, + ) + + def param_init_fn(module): + module.to_empty(device=torch.device("cuda", local_rank), recurse=False) + + model = FSDP( + sft, + auto_wrap_policy=ModuleWrapPolicy(wrap_classes), + mixed_precision=mp, + sharding_strategy=ShardingStrategy.FULL_SHARD, + backward_prefetch=BackwardPrefetch.BACKWARD_PRE, + device_id=torch.device("cuda", local_rank), + use_orig_params=True, + limit_all_gathers=True, + sync_module_states=(args.init_mode == "rank0_meta" and world_size > 1), + param_init_fn=param_init_fn if (args.init_mode == "rank0_meta" and world_size > 1 and rank != 0) else None, + ) + return model, hf_model + + +def build_optimizer(model, recipe: Recipe, args): + import torch + + decay, no_decay = [], [] + dn, nn_ = [], [] + for name, p in model.named_parameters(): + if not p.requires_grad: + continue + clean = name.replace("_fsdp_wrapped_module.", "").replace("hf_model.", "") + (no_decay if is_no_decay(clean) else decay).append(p) + (nn_ if is_no_decay(clean) else dn).append(clean) + groups = [ + {"params": decay, "weight_decay": recipe.weight_decay}, + {"params": no_decay, "weight_decay": 0.0}, + ] + opt = torch.optim.AdamW( + groups, + lr=recipe.learning_rate, + betas=(recipe.beta1, recipe.beta2), + eps=recipe.eps, + fused=args.fused_adam, + foreach=not args.fused_adam, + ) + return opt, dn, nn_ + + +# ---------------------------------------------------------------------------------- +# checkpoint / export +# ---------------------------------------------------------------------------------- +def save_native(model, optimizer, step: int, path: str, extra: dict) -> None: + import torch.distributed as dist + import torch.distributed.checkpoint as dcp + from torch.distributed.checkpoint.state_dict import get_state_dict + + msd, osd = get_state_dict(model, optimizer) + tmp = path + ".partial" + dcp.save({"model": msd, "optim": osd}, checkpoint_id=tmp) + dist.barrier() + if dist.get_rank() == 0: + with open(os.path.join(tmp, "meta.json"), "w") as f: + json.dump({"step": step, **extra}, f) + if os.path.exists(path): + shutil.rmtree(path) + os.replace(tmp, path) + dist.barrier() + + +def load_native(model, optimizer, path: str) -> int: + import torch.distributed.checkpoint as dcp + from torch.distributed.checkpoint.state_dict import get_state_dict, set_state_dict + + msd, osd = get_state_dict(model, optimizer) + state = {"model": msd, "optim": osd} + dcp.load(state, checkpoint_id=path) + set_state_dict(model, optimizer, model_state_dict=state["model"], optim_state_dict=state["optim"]) + with open(os.path.join(path, "meta.json")) as f: + return int(json.load(f)["step"]) + + +def export_hf(model, hf_model, tokenizer, out_dir: str, label: int, dtype_name: str, keep_gen_cfg: bool) -> None: + import torch + import torch.distributed as dist + from torch.distributed.fsdp import FullStateDictConfig, StateDictType + from torch.distributed.fsdp import FullyShardedDataParallel as FSDP + + cfg = FullStateDictConfig(offload_to_cpu=True, rank0_only=True) + with FSDP.state_dict_type(model, StateDictType.FULL_STATE_DICT, cfg): + sd = model.state_dict() + if dist.get_rank() == 0: + dtype = {"bfloat16": torch.bfloat16, "float32": torch.float32, "float16": torch.float16}[dtype_name] + clean = {} + for k, v in sd.items(): + k = k[len("hf_model.") :] if k.startswith("hf_model.") else k + clean[k] = v.to(dtype) + target = os.path.join(out_dir, "hf", f"step-{label}") + tmp = target + ".partial" + if os.path.exists(tmp): + shutil.rmtree(tmp) + os.makedirs(tmp, exist_ok=True) + hf_model.save_pretrained(tmp, state_dict=clean, safe_serialization=True, max_shard_size="5GB") + tokenizer.save_pretrained(tmp) + cfg_path = os.path.join(tmp, "config.json") + with open(cfg_path) as f: + cfg_json = json.load(f) + cfg_json["torch_dtype"] = dtype_name # vLLM reads this; keep it honest + cfg_json["dtype"] = dtype_name + with open(cfg_path, "w") as f: + json.dump(cfg_json, f, indent=2) + gen_cfg = os.path.join(tmp, "generation_config.json") + if not keep_gen_cfg and os.path.exists(gen_cfg): + os.remove(gen_cfg) # the levanter export writes no generation_config.json + if os.path.exists(target): + shutil.rmtree(target) + os.replace(tmp, target) + print(f"[export] wrote {target}", flush=True) + del sd + dist.barrier() + + +def export_tokenizer(recipe: Recipe): + """267 <|padding_i|> tokens so len(tokenizer) == model vocab, exactly like + HFCheckpointConverter.with_tokenizer_padded_to_match_model().""" + from transformers import AutoTokenizer + + tok = AutoTokenizer.from_pretrained(recipe.model_id) + missing = MODEL_VOCAB_SIZE - len(tok) + if missing > 0: + tok.add_tokens([f"<|padding_{i}|>" for i in range(missing)]) + return tok, missing + + +# ---------------------------------------------------------------------------------- +def dry_run(args, recipe: Recipe, data: PackedSFTData, num_train_steps: int) -> None: + lr_fn, (w, s, d) = make_lr_multiplier(recipe, num_train_steps) + L = recipe.max_seq_len + tokens_per_step = recipe.train_batch_size * L + print("=" * 78) + print(f"model {recipe.model_id}") + print(f"packed examples {data.n_packs} (seq len {L})") + print(f"global batch {recipe.train_batch_size} examples = {tokens_per_step:,} tokens/step") + print(f"steps {num_train_steps} => {tokens_per_step * num_train_steps:,} tokens presented") + print(f"epochs over the data {num_train_steps * recipe.train_batch_size / data.n_packs:.2f}") + print(f"lr schedule warmup {w} / stable {s} / cosine {d} to {recipe.learning_rate * recipe.min_lr_ratio:.3e}") + for st in [0, 1, w - 1, w, w + s - 1, w + s, num_train_steps // 2, num_train_steps - 1]: + if 0 <= st < num_train_steps: + print(f" step {st:>5} lr = {recipe.learning_rate * lr_fn(st):.6e}") + print(f"weight decay {recipe.weight_decay} (excluded: *norm*, embed_tokens, *.bias; lm_head decayed)") + print(f"grad clip / betas {recipe.max_grad_norm} / ({recipe.beta1}, {recipe.beta2}) eps {recipe.eps}") + print(f"loss groups per step {args.loss_groups} (TPU used 2 microbatches of 32)") + print(f"assistant tokens {int(data.assistant_counts.sum()):,} over the dataset") + print(f"exports hf/step-N every {recipe.steps_per_hf_export}, native every {recipe.steps_per_checkpoint}") + print("=" * 78) + + +# ---------------------------------------------------------------------------------- +def main() -> None: + p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + p.add_argument("--data", required=True, help="output dir of prepare_sft_data.py") + p.add_argument("--out", required=True, help="run dir; checkpoints/, checkpoints-temp/, hf/ live here") + p.add_argument("--recipe", default="qwen3-8b", choices=sorted(RECIPES)) + p.add_argument("--model-id", default=None, help="override the recipe's HF model id") + p.add_argument("--num-train-steps", type=int, default=0) + p.add_argument("--train-batch-size", type=int, default=0) + p.add_argument("--per-device-batch", type=int, default=1, help="packed examples per GPU per micro-step") + p.add_argument("--loss-groups", type=int, default=2, help="microbatches per optimiser step; TPU used 2") + p.add_argument("--steps-per-checkpoint", type=int, default=0) + p.add_argument("--steps-per-hf-export", type=int, default=0) + p.add_argument("--temp-checkpoint-minutes", type=float, default=10.0) + p.add_argument("--loss-chunk-size", type=int, default=2048) + p.add_argument("--no-loss-checkpoint", action="store_true") + p.add_argument("--attn-impl", default="flash_attention_2") + p.add_argument("--grad-reduce-dtype", default="float32", choices=("float32", "bfloat16")) + p.add_argument("--init-mode", default="rank0_meta", choices=("rank0_meta", "all_ranks")) + p.add_argument("--wrap-embeddings", action="store_true") + p.add_argument("--fused-adam", action="store_true") + p.add_argument("--hf-export-dtype", default="bfloat16", choices=("bfloat16", "float32", "float16")) + p.add_argument("--keep-generation-config", action="store_true") + p.add_argument("--trainer-seed", type=int, default=0, help="levanter TrainerConfig.seed (stays 0 on TPU)") + p.add_argument("--data-seed", type=int, default=42, help="only used when --shuffle is set") + p.add_argument("--shuffle", action="store_true", help="NOT what the TPU run did; shuffle is a no-op there") + p.add_argument("--resume", default="auto", help="auto | none | ") + p.add_argument("--log-every", type=int, default=1) + p.add_argument("--skip-checks", action="store_true") + p.add_argument("--wandb-project", default=None) + p.add_argument("--wandb-name", default=None) + p.add_argument("--wandb-id", default=None, + help="fixed W&B run id; reused across restarts (with resume=allow) so an " + "auto-resumed run continues ONE W&B run instead of forking a new one") + p.add_argument("--dry-run", action="store_true") + args = p.parse_args() + + recipe = RECIPES[args.recipe] + if args.model_id: + recipe = replace(recipe, model_id=args.model_id) + if args.train_batch_size: + recipe = replace(recipe, train_batch_size=args.train_batch_size) + if args.steps_per_checkpoint: + recipe = replace(recipe, steps_per_checkpoint=args.steps_per_checkpoint) + if args.steps_per_hf_export: + recipe = replace(recipe, steps_per_hf_export=args.steps_per_hf_export) + + data = PackedSFTData(args.data) + recipe = replace(recipe, max_seq_len=data.max_seq_len) + num_train_steps = args.num_train_steps or recipe.num_train_steps + if not num_train_steps: + num_train_steps = math.ceil(8 * data.n_packs / recipe.train_batch_size) + + if args.dry_run: + dry_run(args, recipe, data, num_train_steps) + return + + import torch + import torch.distributed as dist + + rank, local_rank, world_size = dist_info() + torch.cuda.set_device(local_rank) + dist.init_process_group("nccl") + device = torch.device("cuda", local_rank) + torch.manual_seed(args.trainer_seed) + np.random.seed(args.trainer_seed) + + B = recipe.train_batch_size + G = args.loss_groups + per_group = B // G + micro_per_group = per_group // (world_size * args.per_device_batch) + if B % G or per_group % (world_size * args.per_device_batch): + raise ValueError( + f"train_batch_size={B} must divide by loss_groups={G} and then by " + f"world_size*per_device_batch={world_size * args.per_device_batch}" + ) + + if recipe.steps_per_hf_export <= 0 or recipe.steps_per_checkpoint <= 0: + raise ValueError("steps_per_hf_export and steps_per_checkpoint must be > 0") + log0(rank, f"[init] world_size={world_size} batch={B} groups={G} micro/group={micro_per_group} seq={recipe.max_seq_len}") + grad_factor = 1.0 if args.skip_checks else detect_grad_reduce_factor(device, world_size, rank) + + model, hf_model = build_model(args, recipe, rank, local_rank, world_size) + if not args.skip_checks: + verify_document_isolation(model, device, rank) + optimizer, decayed, undecayed = build_optimizer(model, recipe, args) + log0(rank, f"[init] weight decay applied to {len(decayed)} tensors, excluded {len(undecayed)}") + log0(rank, f"[init] excluded sample: {undecayed[:4]} ... lm_head decayed: {any(n.endswith('lm_head.weight') for n in decayed)}") + + lr_fn, (w_steps, s_steps, d_steps) = make_lr_multiplier(recipe, num_train_steps) + scheduler = torch.optim.lr_scheduler.LambdaLR(optimizer, lr_lambda=lr_fn) + tokenizer, n_pad = export_tokenizer(recipe) + log0(rank, f"[init] export tokenizer padded with {n_pad} tokens -> len={len(tokenizer)}") + + ckpt_dir = os.path.join(args.out, "checkpoints") + temp_dir = os.path.join(args.out, "checkpoints-temp") + if rank == 0: + os.makedirs(ckpt_dir, exist_ok=True) + os.makedirs(os.path.join(args.out, "hf"), exist_ok=True) + dist.barrier() + + start_step = 0 + resume_path = None + if args.resume == "auto": + cands = [] + if os.path.exists(os.path.join(temp_dir, "meta.json")): + cands.append(temp_dir) + if os.path.isdir(ckpt_dir): + cands += [os.path.join(ckpt_dir, d) for d in os.listdir(ckpt_dir) if d.startswith("step-")] + cands = [c for c in cands if os.path.exists(os.path.join(c, "meta.json"))] + if cands: + resume_path = max(cands, key=lambda c: json.load(open(os.path.join(c, "meta.json")))["step"]) + elif args.resume != "none": + resume_path = args.resume + if resume_path: + start_step = load_native(model, optimizer, resume_path) + for _ in range(start_step): + scheduler.step() + log0(rank, f"[init] resumed from {resume_path} at step {start_step}") + + order = np.arange(data.n_packs) + if args.shuffle: + np.random.default_rng(args.data_seed).shuffle(order) + + run = None + if args.wandb_project and rank == 0: + import wandb + + run = wandb.init( + project=args.wandb_project, + name=args.wandb_name or os.path.basename(args.out), + id=args.wandb_id, + resume="allow" if args.wandb_id else None, + ) + run.config.update({**vars(args), **recipe.__dict__, "num_train_steps": num_train_steps}, + allow_val_change=True) + + model.train() + t_start = time.time() + last_temp = time.time() + log0(rank, "First batch loaded, starting first train step (includes CUDA graph/kernel warmup)...") + + for step in range(start_step, num_train_steps): + t_step = time.time() + base = step * B + batch_idx = [int(order[(base + i) % data.n_packs]) for i in range(B)] + optimizer.zero_grad(set_to_none=True) + + group_loss_sums = torch.zeros(G, device=device, dtype=torch.float32) + group_tokens = torch.zeros(G, device=device, dtype=torch.float32) + for g in range(G): + gidx = batch_idx[g * per_group : (g + 1) * per_group] + denom = float(sum(int(data.assistant_counts[i]) for i in gidx)) + if denom <= 0: + raise RuntimeError(f"group {g} of step {step} has zero assistant tokens") + scale = grad_factor / (G * denom) + for m in range(micro_per_group): + off = (m * world_size + rank) * args.per_device_batch + mine = gidx[off : off + args.per_device_batch] + batch, n_tok = collate(data, mine, device) + loss_sum = model( + batch["input_ids"], + batch["position_ids"], + batch["targets"], + batch["cu_seq_lens"], + batch["max_length"], + ) + (loss_sum * scale).backward() + group_loss_sums[g] += loss_sum.detach() + group_tokens[g] += n_tok + + dist.all_reduce(group_loss_sums) + dist.all_reduce(group_tokens) + if step == start_step: + expect = [float(sum(int(data.assistant_counts[i]) for i in batch_idx[g * per_group : (g + 1) * per_group])) for g in range(G)] + got = [float(x) for x in group_tokens.tolist()] + if any(abs(a - b) > 0.5 for a, b in zip(expect, got)): + raise RuntimeError(f"loss denominator mismatch: manifest={expect} observed={got}") + log0(rank, f"[init] loss denominators verified: {got}") + loss_value = float((group_loss_sums / group_tokens.clamp(min=1)).mean().item()) + + gnorm = model.clip_grad_norm_(recipe.max_grad_norm) + optimizer.step() + scheduler.step() + + completed = step + 1 + label = completed - 1 # levanter StepInfo.step + dt = time.time() - t_step + if rank == 0 and (completed % args.log_every == 0 or completed == num_train_steps): + done = completed - start_step + total = num_train_steps - start_step + rate = (time.time() - t_start) / max(done, 1) + print( + f"Progress on:train {completed}it/{num_train_steps / 1000:.2f}kit " + f"rate:{rate:.1f}s/it remaining:{fmt_hms(rate * (total - done))} " + f"elapsed:{fmt_hms(time.time() - t_start)} postfix:loss={loss_value:.3f}", + flush=True, + ) + if run is not None: + run.log( + { + "train/loss": loss_value, + "train/lr": scheduler.get_last_lr()[0], + "train/grad_norm": float(gnorm), + "train/tokens": int(group_tokens.sum().item()), + "train/step_time": dt, + }, + step=completed, + ) + + is_final = completed == num_train_steps + if (label > 1 and label % int(recipe.steps_per_hf_export) == 0) or is_final: + export_hf(model, hf_model, tokenizer, args.out, label, args.hf_export_dtype, args.keep_generation_config) + if (label > 1 and label % int(recipe.steps_per_checkpoint) == 0) or is_final: + save_native(model, optimizer, completed, os.path.join(ckpt_dir, f"step-{label}"), {"label": label}) + elif (time.time() - last_temp) > args.temp_checkpoint_minutes * 60: + save_native(model, optimizer, completed, temp_dir, {"label": label}) + last_temp = time.time() + + log0(rank, f"[done] {num_train_steps} steps in {fmt_hms(time.time() - t_start)}") + if run is not None: + run.finish() + dist.barrier() + dist.destroy_process_group() + + +if __name__ == "__main__": + with contextlib.suppress(KeyboardInterrupt): + main() diff --git a/scripts/combine_and_publish.py b/scripts/combine_and_publish.py new file mode 100644 index 0000000000000000000000000000000000000000..ca21b21c643889c592435a098472297addb4cd07 --- /dev/null +++ b/scripts/combine_and_publish.py @@ -0,0 +1,171 @@ +"""Wait for all workers to upload, gather, combine, and publish the final dataset. + +Aggregates stats with the FULL launcher.py schema (`summary`, `checks`, +`failure_breakdown`, `failure_breakdown_examples`, `num_shards`) so the published +`stats.json` matches the original `sdg/launcher.py` output structure. + +Env vars: + HF_USER — HuggingFace username (default: fzzhang) + RUN_NAME — experiment name (default: math53K run name) + NUM_WORKERS — total worker count (default: 4) + NUM_SHARDS — total shards (default: 32) +""" +import json +import os +import tempfile +import time +from pathlib import Path + +from huggingface_hub import HfApi, snapshot_download + +HF_USER = os.environ.get("HF_USER", "fzzhang") +RUN_NAME = os.environ.get( + "RUN_NAME", "qwen3_4b_openthoughts3_math53K_instill_n8_valredundancy5_round1" +) +NUM_WORKERS = int(os.environ.get("NUM_WORKERS", "4")) +NUM_SHARDS = int(os.environ.get("NUM_SHARDS", "32")) +STAGING_REPO = f"{HF_USER}/{RUN_NAME}-staging" +FINAL_REPO = f"{HF_USER}/{RUN_NAME}" + +api = HfApi() + +print(f"waiting for {NUM_WORKERS} workers to upload...", flush=True) +while True: + try: + files = api.list_repo_files(STAGING_REPO, repo_type="dataset") + except Exception as e: + print(f" list error (staging repo may not exist yet): {e}", flush=True) + files = [] + done = [ + f for f in files + if f.startswith("workers/worker_") and f.endswith(".done") + ] + print( + f" [{time.strftime('%H:%M:%S')}] {len(done)}/{NUM_WORKERS} workers done", + flush=True, + ) + if len(done) >= NUM_WORKERS: + break + time.sleep(120) + +print("downloading staging repo...", flush=True) +local = snapshot_download( + repo_id=STAGING_REPO, repo_type="dataset", local_dir="./gathered" +) +shards_dir = Path(local) / "shards" +present = sorted(shards_dir.glob("shard_*")) +print(f" {len(present)} shard dirs present", flush=True) + +# Concatenate output.jsonl files in shard order +combined = Path("combined_output.jsonl") +total = 0 +with open(combined, "w") as out: + for sid in range(NUM_SHARDS): + f = shards_dir / f"shard_{sid:03d}" / "output.jsonl" + if not f.exists(): + print(f" WARN: missing {f}") + continue + with open(f) as src: + for line in src: + out.write(line) + total += 1 +print(f"combined {total} rows", flush=True) + +# Aggregate stats with the FULL launcher.py schema +agg = { + "summary": { + "total_examples": 0, + "total_passed": 0, + "total_passed_samples": 0, + "total_failed": 0, + "pass_rate_pct": 0.0, + }, + "checks": { + "static_check": {"passed": 0, "failed": 0}, + "cycle_consistency": {"passed": 0, "failed": 0}, + "factual_accuracy": {"passed": 0, "failed": 0}, + "correctness": {"passed": 0, "failed": 0}, + }, + "failure_breakdown": { + "failed_static_check": 0, + "failed_cycle_consistency": 0, + "failed_factual_accuracy": 0, + "failed_correctness": 0, + "exhausted_samples": 0, + "total_failed": 0, + }, + "failure_breakdown_examples": { + "failed_static_check": 0, + "failed_cycle_consistency": 0, + "failed_factual_accuracy": 0, + "failed_correctness": 0, + "total_failed": 0, + }, + "num_shards": NUM_SHARDS, +} +for sid in range(NUM_SHARDS): + sf = shards_dir / f"shard_{sid:03d}" / "stats.json" + if not sf.exists(): + continue + s = json.load(open(sf)) + for k in ["total_examples", "total_passed", "total_passed_samples", "total_failed"]: + agg["summary"][k] += s["summary"].get(k, 0) + for cn in agg["checks"]: + if cn in s.get("checks", {}): + for k in ["passed", "failed"]: + agg["checks"][cn][k] += s["checks"][cn][k] + for k in agg["failure_breakdown"]: + if k in s.get("failure_breakdown", {}): + agg["failure_breakdown"][k] += s["failure_breakdown"][k] + for k in agg["failure_breakdown_examples"]: + if k in s.get("failure_breakdown_examples", {}): + agg["failure_breakdown_examples"][k] += s["failure_breakdown_examples"][k] +te = agg["summary"]["total_examples"] +tp = agg["summary"]["total_passed"] +agg["summary"]["pass_rate_pct"] = round(tp / te * 100, 2) if te else 0.0 + +stats_path = Path("combined_stats.json") +json.dump(agg, open(stats_path, "w"), indent=2) +print(f"stats: {agg['summary']}", flush=True) + +# Sanity: row count should match total_passed (first_valid policy) +expected = agg["summary"].get("total_passed_samples") or agg["summary"]["total_passed"] +if total != expected: + raise RuntimeError( + f"INTEGRITY: combined output.jsonl has {total} rows but stats expect " + f"{expected}. Shard outputs may be corrupted." + ) + +print(f"publishing to {FINAL_REPO}", flush=True) +api.create_repo(FINAL_REPO, repo_type="dataset", exist_ok=True) +readme = ( + "---\n" + "configs:\n" + "- config_name: default\n" + " data_files:\n" + " - split: train\n" + " path: output.jsonl\n" + "---\n" +) +with tempfile.NamedTemporaryFile(mode="w", suffix=".md", delete=False) as f: + f.write(readme) + rp = f.name +api.upload_file( + path_or_fileobj=rp, + path_in_repo="README.md", + repo_id=FINAL_REPO, + repo_type="dataset", +) +api.upload_file( + path_or_fileobj=str(combined), + path_in_repo="output.jsonl", + repo_id=FINAL_REPO, + repo_type="dataset", +) +api.upload_file( + path_or_fileobj=str(stats_path), + path_in_repo="stats.json", + repo_id=FINAL_REPO, + repo_type="dataset", +) +print(f"DONE: https://huggingface.co/datasets/{FINAL_REPO}", flush=True) diff --git a/scripts/launch_shards.sh b/scripts/launch_shards.sh new file mode 100644 index 0000000000000000000000000000000000000000..41a4e1df6c1f36e12e99246f4432ae966c41320b --- /dev/null +++ b/scripts/launch_shards.sh @@ -0,0 +1,36 @@ +#!/bin/bash +# Launch 8 shards on this worker, one per GPU. +# Worker index is derived from $MY_POD_NAME (Arnold pod name ends with worker-N). +# +# Env vars: +# CONFIG — path to the SDG YAML config (default: math53K config) +# LIMIT — optional --limit value for smoke tests (e.g. LIMIT=128) +# NUM_SHARDS — total shards across all workers (default: 32 = 4 nodes × 8 GPU) +# +# Outputs: +# logs/worker_${WORKER_IDX}/shard_NNN.log +set -e + +WORKER_IDX=$(echo "$MY_POD_NAME" | grep -oE 'worker-[0-9]+$' | grep -oE '[0-9]+') +if [ -z "$WORKER_IDX" ]; then + echo "ERROR: could not infer worker index from MY_POD_NAME=$MY_POD_NAME" >&2 + exit 1 +fi + +CONFIG="${CONFIG:-sdg/configs/qwen3_4b_openthoughts3_math53K_instill_n8_valredundancy5_round1.yaml}" +NUM_SHARDS="${NUM_SHARDS:-32}" +GPUS_PER_WORKER=8 +START=$((WORKER_IDX * GPUS_PER_WORKER)) + +mkdir -p "logs/worker_${WORKER_IDX}" +LIMIT_FLAG="" +if [ -n "$LIMIT" ]; then LIMIT_FLAG="--limit $LIMIT"; fi + +echo "worker $WORKER_IDX launching shards $START..$((START+GPUS_PER_WORKER-1)) (config=$CONFIG, limit=${LIMIT:-none})" +for i in $(seq 0 $((GPUS_PER_WORKER - 1))); do + SHARD_ID=$((START + i)) + CUDA_VISIBLE_DEVICES=$i nohup python -m sdg.generate \ + --config "$CONFIG" --shard-id "$SHARD_ID" --num-shards "$NUM_SHARDS" $LIMIT_FLAG \ + > "logs/worker_${WORKER_IDX}/shard_${SHARD_ID}.log" 2>&1 & +done +echo "launched 8 shards; PIDs: $(jobs -p)" diff --git a/scripts/monitor.sh b/scripts/monitor.sh new file mode 100644 index 0000000000000000000000000000000000000000..e8e5209b7aff7b076830d42391dc3a19f9679e69 --- /dev/null +++ b/scripts/monitor.sh @@ -0,0 +1,44 @@ +#!/bin/bash +# Per-worker monitor: GPU util, current tqdm state for each shard, completion counts, +# and the cluster-wide MongoDB cache size. +# +# Works on any worker (uses $ARNOLD_WORKER_0_HOST for cross-worker mongo reachability, +# so it doesn't matter whether you're on worker_0 or not). +# +# Usage: +# ./scripts/monitor.sh # one shot +# watch -n 30 ./scripts/monitor.sh # auto-refresh +# while true; do clear; ./scripts/monitor.sh; sleep 30; done # if `watch` unavailable + +RUN_NAME="${RUN_NAME:-qwen3_4b_openthoughts3_math53K_instill_n8_valredundancy5_round1}" + +echo "=== $(hostname) $(date +%H:%M:%S) ===" +nvidia-smi --query-gpu=index,utilization.gpu --format=csv,noheader + +echo +echo "=== Shards on this worker ===" +# Show the most recent tqdm update on each shard's log. +# tqdm uses \r to overwrite a single line — naive `tail -1` returns the entire +# line and `cut -cN` truncates from the wrong end. Convert \r → \n and take +# the actual most-recent state. +for f in logs/worker_*/shard_*.log; do + [ -e "$f" ] || continue + recent=$(tail -c 500 "$f" 2>/dev/null | tr '\r' '\n' | tail -1) + echo "$(basename "$f"): $recent" +done + +echo +STATS=$(ls "output/${RUN_NAME}/shards"/*/stats.json 2>/dev/null | wc -l) +COMPLETE=$(grep -l "PIPELINE COMPLETE" logs/worker_*/shard_*.log 2>/dev/null | wc -l) + +# Mongo cache count — use ARNOLD_WORKER_0_HOST (IPv6) so this works from any worker, +# not just worker_0. 127.0.0.1 would only work on worker_0. +MONGO=$(python3 -c " +from pymongo import MongoClient +import os +host = os.environ['ARNOLD_WORKER_0_HOST'] +port = os.environ['ARNOLD_WORKER_0_PORT'] +print(MongoClient(f'mongodb://[{host}]:{port}', serverSelectionTimeoutMS=3000).sdg_cache.inference_cache.count_documents({})) +" 2>/dev/null) + +echo "stats.json: $STATS/8 | PIPELINE COMPLETE: $COMPLETE/8 | mongo cache: ${MONGO:-}" diff --git a/scripts/upload_my_shards.py b/scripts/upload_my_shards.py new file mode 100644 index 0000000000000000000000000000000000000000..fb0bc32c230c614ac2494ee3a8960f890494aa50 --- /dev/null +++ b/scripts/upload_my_shards.py @@ -0,0 +1,66 @@ +"""Wait for this worker's 8 shards to complete, then upload to the staging HF repo. + +Env vars: + HF_USER — HuggingFace username (default: fzzhang) + RUN_NAME — experiment name (default: math53K run name) + +Polls every 2 minutes until all 8 stats.json files appear, then uploads each +shard folder to `/-staging` and writes a `workers/worker_N.done` +marker that the combiner watches for. +""" +import os +import re +import time +from pathlib import Path + +from huggingface_hub import HfApi + +HF_USER = os.environ.get("HF_USER", "fzzhang") +RUN_NAME = os.environ.get( + "RUN_NAME", "qwen3_4b_openthoughts3_math53K_instill_n8_valredundancy5_round1" +) +GPUS_PER_WORKER = 8 +SHARDS_DIR = Path(f"output/{RUN_NAME}/shards") +STAGING_REPO = f"{HF_USER}/{RUN_NAME}-staging" + +m = re.search(r"worker-(\d+)$", os.environ.get("MY_POD_NAME", "")) +if not m: + raise SystemExit("cannot infer worker idx from MY_POD_NAME") +widx = int(m.group(1)) +my_shards = list(range(widx * GPUS_PER_WORKER, (widx + 1) * GPUS_PER_WORKER)) +print(f"worker {widx} owns shards {my_shards}", flush=True) + + +def done(sid: int) -> bool: + return (SHARDS_DIR / f"shard_{sid:03d}" / "stats.json").exists() + + +while True: + d = [s for s in my_shards if done(s)] + print(f"[wait] {len(d)}/{len(my_shards)} shards complete", flush=True) + if len(d) == len(my_shards): + break + time.sleep(120) + +api = HfApi() +api.create_repo(STAGING_REPO, repo_type="dataset", exist_ok=True, private=True) +print(f"uploading to {STAGING_REPO}", flush=True) + +for sid in my_shards: + sp = SHARDS_DIR / f"shard_{sid:03d}" + api.upload_folder( + folder_path=str(sp), + path_in_repo=f"shards/shard_{sid:03d}", + repo_id=STAGING_REPO, + repo_type="dataset", + commit_message=f"upload shard {sid:03d}", + ) + print(f" shard_{sid:03d} done", flush=True) + +api.upload_file( + path_or_fileobj=f"worker {widx} done\n".encode(), + path_in_repo=f"workers/worker_{widx}.done", + repo_id=STAGING_REPO, + repo_type="dataset", +) +print(f"worker {widx} fully uploaded", flush=True) diff --git a/sdg/__init__.py b/sdg/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..3dc1f76bc69e3f559bee6253b24fc93acee9e1f9 --- /dev/null +++ b/sdg/__init__.py @@ -0,0 +1 @@ +__version__ = "0.1.0" diff --git a/sdg/__main__.py b/sdg/__main__.py new file mode 100644 index 0000000000000000000000000000000000000000..9ef019bdaa06a07cbee2b888dac6d2140bd0ac34 --- /dev/null +++ b/sdg/__main__.py @@ -0,0 +1,4 @@ +"""Allow ``python -m sdg`` as an alias for ``python -m sdg.generate``.""" +from sdg.generate import main + +main() diff --git a/sdg/config.py b/sdg/config.py new file mode 100644 index 0000000000000000000000000000000000000000..479f3cafefcd2a71627f577c1083bcb610b51b89 --- /dev/null +++ b/sdg/config.py @@ -0,0 +1,94 @@ +""" +SDGConfig dataclass and YAML serialization. +""" + +from dataclasses import dataclass, field, fields, asdict +from pathlib import Path + +import yaml + + +@dataclass +class SDGConfig: + # Model + model_name: str = "Qwen/Qwen3-4B" + is_instruction_tuned: bool = True + tensor_parallel_size: int = 1 + max_model_len: int = 32768 + gpu_memory_utilization: float = 0.95 + + # Dataset + dataset_hf_id: str = "open-thoughts/OpenThoughts-114k-math" + dataset_hf_config: str | None = None # HF dataset config/subset name + dataset_hf_revision: str | None = None # HF dataset revision (commit hash) + dataset_split: str = "train" + dataset_format: str = "openthoughts4" # "openthoughts4" | "messages" + check_boxed: bool = True + limit: int | None = None + + # Generation + num_generations: int = 4 + gen_temperature: float = 0.8 + gen_top_p: float = 0.95 + gen_max_tokens: int = 32768 + gen_batch_size: int = 64 + + # Validation + validation_type: str | None = "uq" # "uq" | "simple" | null (None = skip validation) + selection_policy: str = "first_valid" # "first_valid" | "all_valid" + num_validation_votes: int = 3 + val_temperature: float = 0.6 + val_top_p: float = 0.95 + val_max_tokens: int = 32768 + val_batch_size: int = 64 + + # Output + experiment_name: str = "" + output_dir: str = "output" + + # MongoDB cache + mongo_uri: str = "" + + # Sharding + num_shards: int = 1 + + # Machines to exclude when launching Slurm jobs + machines_to_exclude: list[str] = field(default_factory=list) + + # ── derived properties ────────────────────────────────────────────── + + @property + def run_name(self) -> str: + if not self.experiment_name: + raise ValueError("experiment_name must be set in the config YAML") + return self.experiment_name + + @property + def mongo_db_name(self) -> str: + return "sdg_cache" + + @property + def output_path(self) -> Path: + return Path(self.output_dir) / self.run_name + + # ── YAML I/O ──────────────────────────────────────────────────────── + + @classmethod + def from_yaml(cls, path: str | Path) -> "SDGConfig": + with open(path) as f: + data = yaml.safe_load(f) or {} + valid_keys = {f.name for f in fields(cls)} + filtered = {k: v for k, v in data.items() if k in valid_keys} + config = cls(**filtered) + if not config.mongo_uri: + raise ValueError( + f"mongo_uri is not set in {path}. " + f"Add a line like `mongo_uri: \"mongodb://HOST:PORT\"` to the config. " + f"(Common mistake: the key must be `mongo_uri` with an underscore, not `mongo-uri`.)" + ) + return config + + def to_yaml(self, path: str | Path) -> None: + Path(path).parent.mkdir(parents=True, exist_ok=True) + with open(path, "w") as f: + yaml.dump(asdict(self), f, default_flow_style=False, sort_keys=False) diff --git a/sdg/configs/qwen3_0.6b_openthoughts3_math53K_instill_n8_valredundancy5_round1.yaml b/sdg/configs/qwen3_0.6b_openthoughts3_math53K_instill_n8_valredundancy5_round1.yaml new file mode 100644 index 0000000000000000000000000000000000000000..8b8673d884e5696f30c491881af48abcca5e1bf4 --- /dev/null +++ b/sdg/configs/qwen3_0.6b_openthoughts3_math53K_instill_n8_valredundancy5_round1.yaml @@ -0,0 +1,41 @@ +experiment_name: "qwen3_0.6b_openthoughts3_math53K_instill_n8_valredundancy5_round1" + +model_name: "Qwen/Qwen3-0.6B" +is_instruction_tuned: true +tensor_parallel_size: 1 +max_model_len: 32768 +gpu_memory_utilization: 0.95 + +dataset_hf_id: "teetone/openthoughts3-math-deduped-53K" +dataset_split: "train" +dataset_format: "openthoughts4" +check_boxed: true +limit: null + +num_generations: 8 +gen_temperature: 0.8 +gen_top_p: 0.95 +gen_max_tokens: 32768 +gen_batch_size: 64 + +validation_type: uq +selection_policy: first_valid +num_validation_votes: 5 +val_temperature: 0.8 +val_top_p: 0.95 +val_max_tokens: 32768 +val_batch_size: 64 + +output_dir: "output" +mongo_uri: "mongodb://john3.stanford.edu:27017" +num_shards: 30 + +machines_to_exclude: + - sphinx1 + - sphinx2 + - sphinx3 + - sphinx4 + - sphinx5 + - sphinx6 + - sphinx7 + - sphinx8 diff --git a/sdg/configs/qwen3_0.6b_openthoughts4_code9K_instill_n8_valredundancy5_round1.yaml b/sdg/configs/qwen3_0.6b_openthoughts4_code9K_instill_n8_valredundancy5_round1.yaml new file mode 100644 index 0000000000000000000000000000000000000000..39a118b8461c27ab16781538d6ef7f2965fa24ee --- /dev/null +++ b/sdg/configs/qwen3_0.6b_openthoughts4_code9K_instill_n8_valredundancy5_round1.yaml @@ -0,0 +1,42 @@ +experiment_name: "qwen3_0.6b_openthoughts4_code9K_instill_n8_valredundancy5_round1" + +model_name: "Qwen/Qwen3-0.6B" +is_instruction_tuned: true +tensor_parallel_size: 1 +max_model_len: 32768 +gpu_memory_utilization: 0.95 + +dataset_hf_id: "teetone/ot4-code-deduped" +dataset_hf_revision: "840633b" +dataset_split: "train" +dataset_format: "openthoughts4" +check_boxed: false +limit: null + +num_generations: 8 +gen_temperature: 0.8 +gen_top_p: 0.95 +gen_max_tokens: 32768 +gen_batch_size: 64 + +validation_type: uq +selection_policy: first_valid +num_validation_votes: 5 +val_temperature: 0.8 +val_top_p: 0.95 +val_max_tokens: 32768 +val_batch_size: 64 + +output_dir: "output" +mongo_uri: "mongodb://john3.stanford.edu:27017" +num_shards: 30 + +machines_to_exclude: + - sphinx1 + - sphinx2 + - sphinx3 + - sphinx4 + - sphinx5 + - sphinx6 + - sphinx7 + - sphinx8 diff --git a/sdg/configs/qwen3_0.6b_openthoughts4_science26K_instill_n8_valredundancy5_round1.yaml b/sdg/configs/qwen3_0.6b_openthoughts4_science26K_instill_n8_valredundancy5_round1.yaml new file mode 100644 index 0000000000000000000000000000000000000000..8356d597f332aa54814d3ee2e68829058974a5c4 --- /dev/null +++ b/sdg/configs/qwen3_0.6b_openthoughts4_science26K_instill_n8_valredundancy5_round1.yaml @@ -0,0 +1,42 @@ +experiment_name: "qwen3_0.6b_openthoughts4_science26K_instill_n8_valredundancy5_round1" + +model_name: "Qwen/Qwen3-0.6B" +is_instruction_tuned: true +tensor_parallel_size: 1 +max_model_len: 32768 +gpu_memory_utilization: 0.95 + +dataset_hf_id: "teetone/ot4-science-deduped" +dataset_hf_revision: "ddae07c" +dataset_split: "train" +dataset_format: "openthoughts4" +check_boxed: false +limit: null + +num_generations: 8 +gen_temperature: 0.8 +gen_top_p: 0.95 +gen_max_tokens: 32768 +gen_batch_size: 64 + +validation_type: uq +selection_policy: first_valid +num_validation_votes: 5 +val_temperature: 0.8 +val_top_p: 0.95 +val_max_tokens: 32768 +val_batch_size: 64 + +output_dir: "output" +mongo_uri: "mongodb://john3.stanford.edu:27017" +num_shards: 30 + +machines_to_exclude: + - sphinx1 + - sphinx2 + - sphinx3 + - sphinx4 + - sphinx5 + - sphinx6 + - sphinx7 + - sphinx8 diff --git a/sdg/configs/qwen3_0.6b_rstarcoder_seedsft_instill_n8_valredundancy5_round1.yaml b/sdg/configs/qwen3_0.6b_rstarcoder_seedsft_instill_n8_valredundancy5_round1.yaml new file mode 100644 index 0000000000000000000000000000000000000000..35a370052e5dc0c1f098aa3ac7b693d828d48467 --- /dev/null +++ b/sdg/configs/qwen3_0.6b_rstarcoder_seedsft_instill_n8_valredundancy5_round1.yaml @@ -0,0 +1,42 @@ +experiment_name: "qwen3_0.6b_rstarcoder_instill_n8_valredundancy5_round1" + +model_name: "Qwen/Qwen3-0.6B" +is_instruction_tuned: true +tensor_parallel_size: 1 +max_model_len: 32768 +gpu_memory_utilization: 0.95 + +dataset_hf_id: "teetone/rstarcoder-seedsft-deduped" +dataset_hf_revision: "1242238" +dataset_split: "train" +dataset_format: "openthoughts4" +check_boxed: false +limit: null + +num_generations: 8 +gen_temperature: 0.8 +gen_top_p: 0.95 +gen_max_tokens: 32768 +gen_batch_size: 64 + +validation_type: uq +selection_policy: first_valid +num_validation_votes: 5 +val_temperature: 0.8 +val_top_p: 0.95 +val_max_tokens: 32768 +val_batch_size: 64 + +output_dir: "output" +mongo_uri: "mongodb://jagupard39.stanford.edu:27017" +num_shards: 30 + +machines_to_exclude: + - sphinx1 + - sphinx2 + - sphinx3 + - sphinx4 + - sphinx5 + - sphinx6 + - sphinx7 + - sphinx8 diff --git a/sdg/configs/qwen3_1.7b_openthoughts3_math53K_instill_n8_valredundancy5_round1.yaml b/sdg/configs/qwen3_1.7b_openthoughts3_math53K_instill_n8_valredundancy5_round1.yaml new file mode 100644 index 0000000000000000000000000000000000000000..3acdf1ef58c99e1956f136d78254705ae9686a99 --- /dev/null +++ b/sdg/configs/qwen3_1.7b_openthoughts3_math53K_instill_n8_valredundancy5_round1.yaml @@ -0,0 +1,41 @@ +experiment_name: "qwen3_1.7b_openthoughts3_math53K_instill_n8_valredundancy5_round1" + +model_name: "Qwen/Qwen3-1.7B" +is_instruction_tuned: true +tensor_parallel_size: 1 +max_model_len: 32768 +gpu_memory_utilization: 0.95 + +dataset_hf_id: "teetone/openthoughts3-math-deduped-53K" +dataset_split: "train" +dataset_format: "openthoughts4" +check_boxed: true +limit: null + +num_generations: 8 +gen_temperature: 0.8 +gen_top_p: 0.95 +gen_max_tokens: 32768 +gen_batch_size: 64 + +validation_type: uq +selection_policy: first_valid +num_validation_votes: 5 +val_temperature: 0.8 +val_top_p: 0.95 +val_max_tokens: 32768 +val_batch_size: 64 + +output_dir: "output" +mongo_uri: "mongodb://john3.stanford.edu:27017" +num_shards: 30 + +machines_to_exclude: + - sphinx1 + - sphinx2 + - sphinx3 + - sphinx4 + - sphinx5 + - sphinx6 + - sphinx7 + - sphinx8 diff --git a/sdg/configs/qwen3_1.7b_openthoughts4_code9K_instill_n8_valredundancy5_round1.yaml b/sdg/configs/qwen3_1.7b_openthoughts4_code9K_instill_n8_valredundancy5_round1.yaml new file mode 100644 index 0000000000000000000000000000000000000000..215a9b45f9b5624c9119878e7cd6bb4f93caa236 --- /dev/null +++ b/sdg/configs/qwen3_1.7b_openthoughts4_code9K_instill_n8_valredundancy5_round1.yaml @@ -0,0 +1,42 @@ +experiment_name: "qwen3_1.7b_openthoughts4_code9K_instill_n8_valredundancy5_round1" + +model_name: "Qwen/Qwen3-1.7B" +is_instruction_tuned: true +tensor_parallel_size: 1 +max_model_len: 32768 +gpu_memory_utilization: 0.95 + +dataset_hf_id: "teetone/ot4-code-deduped" +dataset_hf_revision: "840633b" +dataset_split: "train" +dataset_format: "openthoughts4" +check_boxed: false +limit: null + +num_generations: 8 +gen_temperature: 0.8 +gen_top_p: 0.95 +gen_max_tokens: 32768 +gen_batch_size: 64 + +validation_type: uq +selection_policy: first_valid +num_validation_votes: 5 +val_temperature: 0.8 +val_top_p: 0.95 +val_max_tokens: 32768 +val_batch_size: 64 + +output_dir: "output" +mongo_uri: "mongodb://john3.stanford.edu:27017" +num_shards: 30 + +machines_to_exclude: + - sphinx1 + - sphinx2 + - sphinx3 + - sphinx4 + - sphinx5 + - sphinx6 + - sphinx7 + - sphinx8 diff --git a/sdg/configs/qwen3_1.7b_openthoughts4_science26K_instill_n8_valredundancy5_round1.yaml b/sdg/configs/qwen3_1.7b_openthoughts4_science26K_instill_n8_valredundancy5_round1.yaml new file mode 100644 index 0000000000000000000000000000000000000000..6dca0ba72f67584e6dc3376860e00bcdacd0d5be --- /dev/null +++ b/sdg/configs/qwen3_1.7b_openthoughts4_science26K_instill_n8_valredundancy5_round1.yaml @@ -0,0 +1,42 @@ +experiment_name: "qwen3_1.7b_openthoughts4_science26K_instill_n8_valredundancy5_round1" + +model_name: "Qwen/Qwen3-1.7B" +is_instruction_tuned: true +tensor_parallel_size: 1 +max_model_len: 32768 +gpu_memory_utilization: 0.95 + +dataset_hf_id: "teetone/ot4-science-deduped" +dataset_hf_revision: "ddae07c" +dataset_split: "train" +dataset_format: "openthoughts4" +check_boxed: false +limit: null + +num_generations: 8 +gen_temperature: 0.8 +gen_top_p: 0.95 +gen_max_tokens: 32768 +gen_batch_size: 64 + +validation_type: uq +selection_policy: first_valid +num_validation_votes: 5 +val_temperature: 0.8 +val_top_p: 0.95 +val_max_tokens: 32768 +val_batch_size: 64 + +output_dir: "output" +mongo_uri: "mongodb://john3.stanford.edu:27017" +num_shards: 30 + +machines_to_exclude: + - sphinx1 + - sphinx2 + - sphinx3 + - sphinx4 + - sphinx5 + - sphinx6 + - sphinx7 + - sphinx8 diff --git a/sdg/configs/qwen3_32b_openthoughts3_math53K_instill_n8_valredundancy5_round1.yaml b/sdg/configs/qwen3_32b_openthoughts3_math53K_instill_n8_valredundancy5_round1.yaml new file mode 100644 index 0000000000000000000000000000000000000000..c45932737840f7ec0f240ed347c0e25d9d756155 --- /dev/null +++ b/sdg/configs/qwen3_32b_openthoughts3_math53K_instill_n8_valredundancy5_round1.yaml @@ -0,0 +1,41 @@ +experiment_name: "qwen3_32b_openthoughts3_math53K_instill_n8_valredundancy5_round1" + +model_name: "Qwen/Qwen3-32B" +is_instruction_tuned: true +tensor_parallel_size: 1 +max_model_len: 32768 +gpu_memory_utilization: 0.95 + +dataset_hf_id: "teetone/openthoughts3-math-deduped-53K" +dataset_split: "train" +dataset_format: "openthoughts4" +check_boxed: true +limit: null + +num_generations: 8 +gen_temperature: 0.8 +gen_top_p: 0.95 +gen_max_tokens: 32768 +gen_batch_size: 64 + +validation_type: uq +selection_policy: first_valid +num_validation_votes: 5 +val_temperature: 0.8 +val_top_p: 0.95 +val_max_tokens: 32768 +val_batch_size: 64 + +output_dir: "output" +mongo_uri: "mongodb://jagupard39.stanford.edu:27017" +num_shards: 30 + +machines_to_exclude: + - sphinx1 + - sphinx2 + - sphinx3 + - sphinx4 + - sphinx5 + - sphinx6 + - sphinx7 + - sphinx8 diff --git a/sdg/configs/qwen3_32b_openthoughts4_code9K_instill_n8_valredundancy5_round1.yaml b/sdg/configs/qwen3_32b_openthoughts4_code9K_instill_n8_valredundancy5_round1.yaml new file mode 100644 index 0000000000000000000000000000000000000000..cf60c849838829351268a9467d8ff968e4990cc4 --- /dev/null +++ b/sdg/configs/qwen3_32b_openthoughts4_code9K_instill_n8_valredundancy5_round1.yaml @@ -0,0 +1,42 @@ +experiment_name: "qwen3_32b_openthoughts4_code9K_instill_n8_valredundancy5_round1" + +model_name: "Qwen/Qwen3-32B" +is_instruction_tuned: true +tensor_parallel_size: 1 +max_model_len: 32768 +gpu_memory_utilization: 0.95 + +dataset_hf_id: "teetone/ot4-code-deduped" +dataset_hf_revision: "840633b" +dataset_split: "train" +dataset_format: "openthoughts4" +check_boxed: false +limit: null + +num_generations: 8 +gen_temperature: 0.8 +gen_top_p: 0.95 +gen_max_tokens: 32768 +gen_batch_size: 64 + +validation_type: uq +selection_policy: first_valid +num_validation_votes: 5 +val_temperature: 0.8 +val_top_p: 0.95 +val_max_tokens: 32768 +val_batch_size: 64 + +output_dir: "output" +mongo_uri: "mongodb://john3.stanford.edu:27017" +num_shards: 30 + +machines_to_exclude: + - sphinx1 + - sphinx2 + - sphinx3 + - sphinx4 + - sphinx5 + - sphinx6 + - sphinx7 + - sphinx8 diff --git a/sdg/configs/qwen3_32b_openthoughts4_science26K_instill_n8_valredundancy5_round1.yaml b/sdg/configs/qwen3_32b_openthoughts4_science26K_instill_n8_valredundancy5_round1.yaml new file mode 100644 index 0000000000000000000000000000000000000000..205665c41398c59a2fa8c4d291a836d29b3e8a34 --- /dev/null +++ b/sdg/configs/qwen3_32b_openthoughts4_science26K_instill_n8_valredundancy5_round1.yaml @@ -0,0 +1,42 @@ +experiment_name: "qwen3_32b_openthoughts4_science26K_instill_n8_valredundancy5_round1" + +model_name: "Qwen/Qwen3-32B" +is_instruction_tuned: true +tensor_parallel_size: 1 +max_model_len: 32768 +gpu_memory_utilization: 0.95 + +dataset_hf_id: "teetone/ot4-science-deduped" +dataset_hf_revision: "ddae07c" +dataset_split: "train" +dataset_format: "openthoughts4" +check_boxed: false +limit: null + +num_generations: 8 +gen_temperature: 0.8 +gen_top_p: 0.95 +gen_max_tokens: 32768 +gen_batch_size: 64 + +validation_type: uq +selection_policy: first_valid +num_validation_votes: 5 +val_temperature: 0.8 +val_top_p: 0.95 +val_max_tokens: 32768 +val_batch_size: 64 + +output_dir: "output" +mongo_uri: "mongodb://john3.stanford.edu:27017" +num_shards: 30 + +machines_to_exclude: + - sphinx1 + - sphinx2 + - sphinx3 + - sphinx4 + - sphinx5 + - sphinx6 + - sphinx7 + - sphinx8 diff --git a/sdg/configs/qwen3_32b_rstarcoder_seedsft_instill_n8_valredundancy5_round1.yaml b/sdg/configs/qwen3_32b_rstarcoder_seedsft_instill_n8_valredundancy5_round1.yaml new file mode 100644 index 0000000000000000000000000000000000000000..81532f3350084ddb3c803dbd759338d7631dc1c3 --- /dev/null +++ b/sdg/configs/qwen3_32b_rstarcoder_seedsft_instill_n8_valredundancy5_round1.yaml @@ -0,0 +1,42 @@ +experiment_name: "qwen3_32b_rstarcoder_instill_n8_valredundancy5_round1" + +model_name: "Qwen/Qwen3-32B" +is_instruction_tuned: true +tensor_parallel_size: 1 +max_model_len: 32768 +gpu_memory_utilization: 0.95 + +dataset_hf_id: "teetone/rstarcoder-seedsft-deduped" +dataset_hf_revision: "1242238" +dataset_split: "train" +dataset_format: "openthoughts4" +check_boxed: false +limit: null + +num_generations: 8 +gen_temperature: 0.8 +gen_top_p: 0.95 +gen_max_tokens: 32768 +gen_batch_size: 64 + +validation_type: uq +selection_policy: first_valid +num_validation_votes: 5 +val_temperature: 0.8 +val_top_p: 0.95 +val_max_tokens: 32768 +val_batch_size: 64 + +output_dir: "output" +mongo_uri: "mongodb://jagupard39.stanford.edu:27017" +num_shards: 30 + +machines_to_exclude: + - sphinx1 + - sphinx2 + - sphinx3 + - sphinx4 + - sphinx5 + - sphinx6 + - sphinx7 + - sphinx8 diff --git a/sdg/configs/qwen3_4b_mot_math_instill_n4_valredundancy3_round1.yaml b/sdg/configs/qwen3_4b_mot_math_instill_n4_valredundancy3_round1.yaml new file mode 100644 index 0000000000000000000000000000000000000000..015946db4c0e7d99ef792bc7d3e4d7a85e44056c --- /dev/null +++ b/sdg/configs/qwen3_4b_mot_math_instill_n4_valredundancy3_round1.yaml @@ -0,0 +1,42 @@ +model_name: "Qwen/Qwen3-4B" +is_instruction_tuned: true +tensor_parallel_size: 1 +max_model_len: 32768 +gpu_memory_utilization: 0.95 + +dataset_hf_id: "open-r1/Mixture-of-Thoughts" +dataset_hf_config: "math" +dataset_split: "train" +dataset_format: "messages" +check_boxed: true +limit: null + +num_generations: 4 +gen_temperature: 0.8 +gen_top_p: 0.95 +gen_max_tokens: 32768 +gen_batch_size: 64 + +validation_type: uq +selection_policy: first_valid +num_validation_votes: 3 +val_temperature: 0.8 +val_top_p: 0.95 +val_max_tokens: 32768 +val_batch_size: 64 + +experiment_name: "qwen3_4b_mot_math_instill_n4_valredundancy3_round1" +output_dir: "output" + +mongo_uri: "mongodb://john3.stanford.edu:27017" +num_shards: 30 + +machines_to_exclude: + - sphinx1 + - sphinx2 + - sphinx3 + - sphinx4 + - sphinx5 + - sphinx6 + - sphinx7 + - sphinx8 diff --git a/sdg/configs/qwen3_4b_nemotron_cascade2_science100k_instill_n8_valredundancy5_round1.yaml b/sdg/configs/qwen3_4b_nemotron_cascade2_science100k_instill_n8_valredundancy5_round1.yaml new file mode 100644 index 0000000000000000000000000000000000000000..bd08945487d61311e9dc8c42f3b3eb2549116a9f --- /dev/null +++ b/sdg/configs/qwen3_4b_nemotron_cascade2_science100k_instill_n8_valredundancy5_round1.yaml @@ -0,0 +1,42 @@ +experiment_name: "qwen3_4b_nemotron_cascade2_science100k_instill_n8_valredundancy5_round1" + +model_name: "Qwen/Qwen3-4B" +is_instruction_tuned: true +tensor_parallel_size: 1 +max_model_len: 32768 +gpu_memory_utilization: 0.95 + +dataset_hf_id: "teetone/nemotron-cascade-2-science-deduped-100k" +dataset_hf_revision: "adb996f" +dataset_split: "train" +dataset_format: "openthoughts4" +check_boxed: false +limit: null + +num_generations: 8 +gen_temperature: 0.8 +gen_top_p: 0.95 +gen_max_tokens: 32768 +gen_batch_size: 64 + +validation_type: uq +selection_policy: first_valid +num_validation_votes: 5 +val_temperature: 0.8 +val_top_p: 0.95 +val_max_tokens: 32768 +val_batch_size: 64 + +output_dir: "output" +mongo_uri: "mongodb://jagupard39.stanford.edu:27017" +num_shards: 30 + +machines_to_exclude: + - sphinx1 + - sphinx2 + - sphinx3 + - sphinx4 + - sphinx5 + - sphinx6 + - sphinx7 + - sphinx8 diff --git a/sdg/configs/qwen3_4b_openthoughts3_math53K_instill_n1_valredundancy1_round1.yaml b/sdg/configs/qwen3_4b_openthoughts3_math53K_instill_n1_valredundancy1_round1.yaml new file mode 100644 index 0000000000000000000000000000000000000000..cfa8377513b060abb2cfb1af49f853564c99536e --- /dev/null +++ b/sdg/configs/qwen3_4b_openthoughts3_math53K_instill_n1_valredundancy1_round1.yaml @@ -0,0 +1,41 @@ +experiment_name: "qwen3_4b_openthoughts3_math53K_instill_n1_valredundancy1_round1" + +model_name: "Qwen/Qwen3-4B" +is_instruction_tuned: true +tensor_parallel_size: 1 +max_model_len: 32768 +gpu_memory_utilization: 0.95 + +dataset_hf_id: "teetone/openthoughts3-math-deduped-53K" +dataset_split: "train" +dataset_format: "openthoughts4" +check_boxed: true +limit: null + +num_generations: 1 +gen_temperature: 0.8 +gen_top_p: 0.95 +gen_max_tokens: 32768 +gen_batch_size: 64 + +validation_type: uq +selection_policy: first_valid +num_validation_votes: 1 +val_temperature: 0.8 +val_top_p: 0.95 +val_max_tokens: 32768 +val_batch_size: 64 + +output_dir: "output" +mongo_uri: "mongodb://john3.stanford.edu:27017" +num_shards: 30 + +machines_to_exclude: + - sphinx1 + - sphinx2 + - sphinx3 + - sphinx4 + - sphinx5 + - sphinx6 + - sphinx7 + - sphinx8 diff --git a/sdg/configs/qwen3_4b_openthoughts3_math53K_instill_n1_valredundancy3_round1.yaml b/sdg/configs/qwen3_4b_openthoughts3_math53K_instill_n1_valredundancy3_round1.yaml new file mode 100644 index 0000000000000000000000000000000000000000..a80821b011a59b3d3300939f1cd8ce4bd22baf90 --- /dev/null +++ b/sdg/configs/qwen3_4b_openthoughts3_math53K_instill_n1_valredundancy3_round1.yaml @@ -0,0 +1,41 @@ +experiment_name: "qwen3_4b_openthoughts3_math53K_instill_n1_valredundancy3_round1" + +model_name: "Qwen/Qwen3-4B" +is_instruction_tuned: true +tensor_parallel_size: 1 +max_model_len: 32768 +gpu_memory_utilization: 0.95 + +dataset_hf_id: "teetone/openthoughts3-math-deduped-53K" +dataset_split: "train" +dataset_format: "openthoughts4" +check_boxed: true +limit: null + +num_generations: 1 +gen_temperature: 0.8 +gen_top_p: 0.95 +gen_max_tokens: 32768 +gen_batch_size: 64 + +validation_type: uq +selection_policy: first_valid +num_validation_votes: 3 +val_temperature: 0.8 +val_top_p: 0.95 +val_max_tokens: 32768 +val_batch_size: 64 + +output_dir: "output" +mongo_uri: "mongodb://john3.stanford.edu:27017" +num_shards: 30 + +machines_to_exclude: + - sphinx1 + - sphinx2 + - sphinx3 + - sphinx4 + - sphinx5 + - sphinx6 + - sphinx7 + - sphinx8 diff --git a/sdg/configs/qwen3_4b_openthoughts3_math53K_instill_n1_valredundancy5_round1.yaml b/sdg/configs/qwen3_4b_openthoughts3_math53K_instill_n1_valredundancy5_round1.yaml new file mode 100644 index 0000000000000000000000000000000000000000..b4bb4277121d9d64a597d008346b2d1347aaa505 --- /dev/null +++ b/sdg/configs/qwen3_4b_openthoughts3_math53K_instill_n1_valredundancy5_round1.yaml @@ -0,0 +1,41 @@ +experiment_name: "qwen3_4b_openthoughts3_math53K_instill_n1_valredundancy5_round1" + +model_name: "Qwen/Qwen3-4B" +is_instruction_tuned: true +tensor_parallel_size: 1 +max_model_len: 32768 +gpu_memory_utilization: 0.95 + +dataset_hf_id: "teetone/openthoughts3-math-deduped-53K" +dataset_split: "train" +dataset_format: "openthoughts4" +check_boxed: true +limit: null + +num_generations: 1 +gen_temperature: 0.8 +gen_top_p: 0.95 +gen_max_tokens: 32768 +gen_batch_size: 64 + +validation_type: uq +selection_policy: first_valid +num_validation_votes: 5 +val_temperature: 0.8 +val_top_p: 0.95 +val_max_tokens: 32768 +val_batch_size: 64 + +output_dir: "output" +mongo_uri: "mongodb://john3.stanford.edu:27017" +num_shards: 30 + +machines_to_exclude: + - sphinx1 + - sphinx2 + - sphinx3 + - sphinx4 + - sphinx5 + - sphinx6 + - sphinx7 + - sphinx8 diff --git a/sdg/configs/qwen3_4b_openthoughts3_math53K_instill_n4_valredundancy1_round1.yaml b/sdg/configs/qwen3_4b_openthoughts3_math53K_instill_n4_valredundancy1_round1.yaml new file mode 100644 index 0000000000000000000000000000000000000000..afad9e6a4620f9ea8a0a0823caaa5850f878c9b4 --- /dev/null +++ b/sdg/configs/qwen3_4b_openthoughts3_math53K_instill_n4_valredundancy1_round1.yaml @@ -0,0 +1,41 @@ +experiment_name: "qwen3_4b_openthoughts3_math53K_instill_n4_valredundancy1_round1" + +model_name: "Qwen/Qwen3-4B" +is_instruction_tuned: true +tensor_parallel_size: 1 +max_model_len: 32768 +gpu_memory_utilization: 0.95 + +dataset_hf_id: "teetone/openthoughts3-math-deduped-53K" +dataset_split: "train" +dataset_format: "openthoughts4" +check_boxed: true +limit: null + +num_generations: 4 +gen_temperature: 0.8 +gen_top_p: 0.95 +gen_max_tokens: 32768 +gen_batch_size: 64 + +validation_type: uq +selection_policy: first_valid +num_validation_votes: 1 +val_temperature: 0.8 +val_top_p: 0.95 +val_max_tokens: 32768 +val_batch_size: 64 + +output_dir: "output" +mongo_uri: "mongodb://john3.stanford.edu:27017" +num_shards: 30 + +machines_to_exclude: + - sphinx1 + - sphinx2 + - sphinx3 + - sphinx4 + - sphinx5 + - sphinx6 + - sphinx7 + - sphinx8 diff --git a/sdg/configs/qwen3_4b_openthoughts3_math53K_instill_n4_valredundancy3_round1.yaml b/sdg/configs/qwen3_4b_openthoughts3_math53K_instill_n4_valredundancy3_round1.yaml new file mode 100644 index 0000000000000000000000000000000000000000..8a62a742f8d8d87d2af480e987729a0e909d6f7e --- /dev/null +++ b/sdg/configs/qwen3_4b_openthoughts3_math53K_instill_n4_valredundancy3_round1.yaml @@ -0,0 +1,41 @@ +experiment_name: "qwen3_4b_openthoughts3_math53K_instill_n4_valredundancy3_round1" + +model_name: "Qwen/Qwen3-4B" +is_instruction_tuned: true +tensor_parallel_size: 1 +max_model_len: 32768 +gpu_memory_utilization: 0.95 + +dataset_hf_id: "teetone/openthoughts3-math-deduped-53K" +dataset_split: "train" +dataset_format: "openthoughts4" +check_boxed: true +limit: null + +num_generations: 4 +gen_temperature: 0.8 +gen_top_p: 0.95 +gen_max_tokens: 32768 +gen_batch_size: 64 + +validation_type: uq +selection_policy: first_valid +num_validation_votes: 3 +val_temperature: 0.8 +val_top_p: 0.95 +val_max_tokens: 32768 +val_batch_size: 64 + +output_dir: "output" +mongo_uri: "mongodb://john3.stanford.edu:27017" +num_shards: 30 + +machines_to_exclude: + - sphinx1 + - sphinx2 + - sphinx3 + - sphinx4 + - sphinx5 + - sphinx6 + - sphinx7 + - sphinx8 diff --git a/sdg/configs/qwen3_4b_openthoughts3_math53K_instill_n4_valredundancy5_round1.yaml b/sdg/configs/qwen3_4b_openthoughts3_math53K_instill_n4_valredundancy5_round1.yaml new file mode 100644 index 0000000000000000000000000000000000000000..1e5157baa38a9dbf54741b324dc0ddd8c84db7ca --- /dev/null +++ b/sdg/configs/qwen3_4b_openthoughts3_math53K_instill_n4_valredundancy5_round1.yaml @@ -0,0 +1,41 @@ +experiment_name: "qwen3_4b_openthoughts3_math53K_instill_n4_valredundancy5_round1" + +model_name: "Qwen/Qwen3-4B" +is_instruction_tuned: true +tensor_parallel_size: 1 +max_model_len: 32768 +gpu_memory_utilization: 0.95 + +dataset_hf_id: "teetone/openthoughts3-math-deduped-53K" +dataset_split: "train" +dataset_format: "openthoughts4" +check_boxed: true +limit: null + +num_generations: 4 +gen_temperature: 0.8 +gen_top_p: 0.95 +gen_max_tokens: 32768 +gen_batch_size: 64 + +validation_type: uq +selection_policy: first_valid +num_validation_votes: 5 +val_temperature: 0.8 +val_top_p: 0.95 +val_max_tokens: 32768 +val_batch_size: 64 + +output_dir: "output" +mongo_uri: "mongodb://john3.stanford.edu:27017" +num_shards: 30 + +machines_to_exclude: + - sphinx1 + - sphinx2 + - sphinx3 + - sphinx4 + - sphinx5 + - sphinx6 + - sphinx7 + - sphinx8 diff --git a/sdg/configs/qwen3_4b_openthoughts3_math53K_instill_n8_no_filter_round1.yaml b/sdg/configs/qwen3_4b_openthoughts3_math53K_instill_n8_no_filter_round1.yaml new file mode 100644 index 0000000000000000000000000000000000000000..e68e8a30e7c6c4ae1dc8d653fbf3a802054526c6 --- /dev/null +++ b/sdg/configs/qwen3_4b_openthoughts3_math53K_instill_n8_no_filter_round1.yaml @@ -0,0 +1,41 @@ +experiment_name: "qwen3_4b_openthoughts3_math53K_instill_n8_no_filter_round1" + +model_name: "Qwen/Qwen3-4B" +is_instruction_tuned: true +tensor_parallel_size: 1 +max_model_len: 32768 +gpu_memory_utilization: 0.95 + +dataset_hf_id: "teetone/openthoughts3-math-deduped-53K" +dataset_split: "train" +dataset_format: "openthoughts4" +check_boxed: true +limit: null + +num_generations: 8 +gen_temperature: 0.8 +gen_top_p: 0.95 +gen_max_tokens: 32768 +gen_batch_size: 64 + +validation_type: null +selection_policy: first_valid +num_validation_votes: 0 +val_temperature: 0.8 +val_top_p: 0.95 +val_max_tokens: 32768 +val_batch_size: 64 + +output_dir: "output" +mongo_uri: "mongodb://john3.stanford.edu:27017" +num_shards: 30 + +machines_to_exclude: + - sphinx1 + - sphinx2 + - sphinx3 + - sphinx4 + - sphinx5 + - sphinx6 + - sphinx7 + - sphinx8 diff --git a/sdg/configs/qwen3_4b_openthoughts3_math53K_instill_n8_simple_valredundancy5_round1.yaml b/sdg/configs/qwen3_4b_openthoughts3_math53K_instill_n8_simple_valredundancy5_round1.yaml new file mode 100644 index 0000000000000000000000000000000000000000..98bc37399b19599f7c7f8970375fb437b4c2262b --- /dev/null +++ b/sdg/configs/qwen3_4b_openthoughts3_math53K_instill_n8_simple_valredundancy5_round1.yaml @@ -0,0 +1,41 @@ +experiment_name: "qwen3_4b_openthoughts3_math53K_instill_n8_simple_valredundancy5_round1" + +model_name: "Qwen/Qwen3-4B" +is_instruction_tuned: true +tensor_parallel_size: 1 +max_model_len: 32768 +gpu_memory_utilization: 0.95 + +dataset_hf_id: "teetone/openthoughts3-math-deduped-53K" +dataset_split: "train" +dataset_format: "openthoughts4" +check_boxed: true +limit: null + +num_generations: 8 +gen_temperature: 0.8 +gen_top_p: 0.95 +gen_max_tokens: 32768 +gen_batch_size: 64 + +validation_type: simple +selection_policy: first_valid +num_validation_votes: 5 +val_temperature: 0.8 +val_top_p: 0.95 +val_max_tokens: 32768 +val_batch_size: 64 + +output_dir: "output" +mongo_uri: "mongodb://john3.stanford.edu:27017" +num_shards: 30 + +machines_to_exclude: + - sphinx1 + - sphinx2 + - sphinx3 + - sphinx4 + - sphinx5 + - sphinx6 + - sphinx7 + - sphinx8 diff --git a/sdg/configs/qwen3_4b_openthoughts3_math53K_instill_n8_valredundancy1_round1.yaml b/sdg/configs/qwen3_4b_openthoughts3_math53K_instill_n8_valredundancy1_round1.yaml new file mode 100644 index 0000000000000000000000000000000000000000..220e760187a5ec41d7a698ec1567aebf017deaa6 --- /dev/null +++ b/sdg/configs/qwen3_4b_openthoughts3_math53K_instill_n8_valredundancy1_round1.yaml @@ -0,0 +1,41 @@ +experiment_name: "qwen3_4b_openthoughts3_math53K_instill_n8_valredundancy1_round1" + +model_name: "Qwen/Qwen3-4B" +is_instruction_tuned: true +tensor_parallel_size: 1 +max_model_len: 32768 +gpu_memory_utilization: 0.95 + +dataset_hf_id: "teetone/openthoughts3-math-deduped-53K" +dataset_split: "train" +dataset_format: "openthoughts4" +check_boxed: true +limit: null + +num_generations: 8 +gen_temperature: 0.8 +gen_top_p: 0.95 +gen_max_tokens: 32768 +gen_batch_size: 64 + +validation_type: uq +selection_policy: first_valid +num_validation_votes: 1 +val_temperature: 0.8 +val_top_p: 0.95 +val_max_tokens: 32768 +val_batch_size: 64 + +output_dir: "output" +mongo_uri: "mongodb://john3.stanford.edu:27017" +num_shards: 30 + +machines_to_exclude: + - sphinx1 + - sphinx2 + - sphinx3 + - sphinx4 + - sphinx5 + - sphinx6 + - sphinx7 + - sphinx8 diff --git a/sdg/configs/qwen3_4b_openthoughts3_math53K_instill_n8_valredundancy3_round1.yaml b/sdg/configs/qwen3_4b_openthoughts3_math53K_instill_n8_valredundancy3_round1.yaml new file mode 100644 index 0000000000000000000000000000000000000000..f41e43348607ae5476065e53a3633887565b7942 --- /dev/null +++ b/sdg/configs/qwen3_4b_openthoughts3_math53K_instill_n8_valredundancy3_round1.yaml @@ -0,0 +1,41 @@ +experiment_name: "qwen3_4b_openthoughts3_math53K_instill_n8_valredundancy3_round1" + +model_name: "Qwen/Qwen3-4B" +is_instruction_tuned: true +tensor_parallel_size: 1 +max_model_len: 32768 +gpu_memory_utilization: 0.95 + +dataset_hf_id: "teetone/openthoughts3-math-deduped-53K" +dataset_split: "train" +dataset_format: "openthoughts4" +check_boxed: true +limit: null + +num_generations: 8 +gen_temperature: 0.8 +gen_top_p: 0.95 +gen_max_tokens: 32768 +gen_batch_size: 64 + +validation_type: uq +selection_policy: first_valid +num_validation_votes: 3 +val_temperature: 0.8 +val_top_p: 0.95 +val_max_tokens: 32768 +val_batch_size: 64 + +output_dir: "output" +mongo_uri: "mongodb://john3.stanford.edu:27017" +num_shards: 30 + +machines_to_exclude: + - sphinx1 + - sphinx2 + - sphinx3 + - sphinx4 + - sphinx5 + - sphinx6 + - sphinx7 + - sphinx8 diff --git a/sdg/configs/qwen3_4b_openthoughts3_math53K_instill_n8_valredundancy5_allvalid_round1.yaml b/sdg/configs/qwen3_4b_openthoughts3_math53K_instill_n8_valredundancy5_allvalid_round1.yaml new file mode 100644 index 0000000000000000000000000000000000000000..d3393fff79e612c9ab7966f8ededb902dbeee45c --- /dev/null +++ b/sdg/configs/qwen3_4b_openthoughts3_math53K_instill_n8_valredundancy5_allvalid_round1.yaml @@ -0,0 +1,41 @@ +experiment_name: "qwen3_4b_openthoughts3_math53K_instill_n8_valredundancy5_allvalid_round1" + +model_name: "Qwen/Qwen3-4B" +is_instruction_tuned: true +tensor_parallel_size: 1 +max_model_len: 32768 +gpu_memory_utilization: 0.95 + +dataset_hf_id: "teetone/openthoughts3-math-deduped-53K" +dataset_split: "train" +dataset_format: "openthoughts4" +check_boxed: true +limit: null + +num_generations: 8 +gen_temperature: 0.8 +gen_top_p: 0.95 +gen_max_tokens: 32768 +gen_batch_size: 64 + +validation_type: uq +selection_policy: all_valid +num_validation_votes: 5 +val_temperature: 0.8 +val_top_p: 0.95 +val_max_tokens: 32768 +val_batch_size: 64 + +output_dir: "output" +mongo_uri: "mongodb://john3.stanford.edu:27017" +num_shards: 30 + +machines_to_exclude: + - sphinx1 + - sphinx2 + - sphinx3 + - sphinx4 + - sphinx5 + - sphinx6 + - sphinx7 + - sphinx8 diff --git a/sdg/configs/qwen3_4b_openthoughts3_math53K_instill_n8_valredundancy5_round1.yaml b/sdg/configs/qwen3_4b_openthoughts3_math53K_instill_n8_valredundancy5_round1.yaml new file mode 100644 index 0000000000000000000000000000000000000000..981e22e3388108e7dd81bfce1a2c1d6817179442 --- /dev/null +++ b/sdg/configs/qwen3_4b_openthoughts3_math53K_instill_n8_valredundancy5_round1.yaml @@ -0,0 +1,41 @@ +experiment_name: "qwen3_4b_openthoughts3_math53K_instill_n8_valredundancy5_round1" + +model_name: "Qwen/Qwen3-4B" +is_instruction_tuned: true +tensor_parallel_size: 1 +max_model_len: 32768 +gpu_memory_utilization: 0.95 + +dataset_hf_id: "teetone/openthoughts3-math-deduped-53K" +dataset_split: "train" +dataset_format: "openthoughts4" +check_boxed: true +limit: null + +num_generations: 8 +gen_temperature: 0.8 +gen_top_p: 0.95 +gen_max_tokens: 32768 +gen_batch_size: 64 + +validation_type: uq +selection_policy: first_valid +num_validation_votes: 5 +val_temperature: 0.8 +val_top_p: 0.95 +val_max_tokens: 32768 +val_batch_size: 64 + +output_dir: "output" +mongo_uri: "mongodb://john3.stanford.edu:27017" +num_shards: 30 + +machines_to_exclude: + - sphinx1 + - sphinx2 + - sphinx3 + - sphinx4 + - sphinx5 + - sphinx6 + - sphinx7 + - sphinx8 diff --git a/sdg/configs/qwen3_4b_openthoughts4_code9K_instill_n1_valredundancy1_round1.yaml b/sdg/configs/qwen3_4b_openthoughts4_code9K_instill_n1_valredundancy1_round1.yaml new file mode 100644 index 0000000000000000000000000000000000000000..85213d7e719eb17a09b2430d7fa5b43470916640 --- /dev/null +++ b/sdg/configs/qwen3_4b_openthoughts4_code9K_instill_n1_valredundancy1_round1.yaml @@ -0,0 +1,42 @@ +experiment_name: "qwen3_4b_openthoughts4_code9K_instill_n1_valredundancy1_round1" + +model_name: "Qwen/Qwen3-4B" +is_instruction_tuned: true +tensor_parallel_size: 1 +max_model_len: 32768 +gpu_memory_utilization: 0.95 + +dataset_hf_id: "teetone/ot4-code-deduped" +dataset_hf_revision: "840633b" +dataset_split: "train" +dataset_format: "openthoughts4" +check_boxed: false +limit: null + +num_generations: 1 +gen_temperature: 0.8 +gen_top_p: 0.95 +gen_max_tokens: 32768 +gen_batch_size: 64 + +validation_type: uq +selection_policy: first_valid +num_validation_votes: 1 +val_temperature: 0.8 +val_top_p: 0.95 +val_max_tokens: 32768 +val_batch_size: 64 + +output_dir: "output" +mongo_uri: "mongodb://john3.stanford.edu:27017" +num_shards: 30 + +machines_to_exclude: + - sphinx1 + - sphinx2 + - sphinx3 + - sphinx4 + - sphinx5 + - sphinx6 + - sphinx7 + - sphinx8 diff --git a/sdg/configs/qwen3_4b_openthoughts4_code9K_instill_n1_valredundancy3_round1.yaml b/sdg/configs/qwen3_4b_openthoughts4_code9K_instill_n1_valredundancy3_round1.yaml new file mode 100644 index 0000000000000000000000000000000000000000..41b676ea8c58d163fd62fc9960bda75ecbc62def --- /dev/null +++ b/sdg/configs/qwen3_4b_openthoughts4_code9K_instill_n1_valredundancy3_round1.yaml @@ -0,0 +1,42 @@ +experiment_name: "qwen3_4b_openthoughts4_code9K_instill_n1_valredundancy3_round1" + +model_name: "Qwen/Qwen3-4B" +is_instruction_tuned: true +tensor_parallel_size: 1 +max_model_len: 32768 +gpu_memory_utilization: 0.95 + +dataset_hf_id: "teetone/ot4-code-deduped" +dataset_hf_revision: "840633b" +dataset_split: "train" +dataset_format: "openthoughts4" +check_boxed: false +limit: null + +num_generations: 1 +gen_temperature: 0.8 +gen_top_p: 0.95 +gen_max_tokens: 32768 +gen_batch_size: 64 + +validation_type: uq +selection_policy: first_valid +num_validation_votes: 3 +val_temperature: 0.8 +val_top_p: 0.95 +val_max_tokens: 32768 +val_batch_size: 64 + +output_dir: "output" +mongo_uri: "mongodb://john3.stanford.edu:27017" +num_shards: 30 + +machines_to_exclude: + - sphinx1 + - sphinx2 + - sphinx3 + - sphinx4 + - sphinx5 + - sphinx6 + - sphinx7 + - sphinx8 diff --git a/sdg/configs/qwen3_4b_openthoughts4_code9K_instill_n1_valredundancy5_round1.yaml b/sdg/configs/qwen3_4b_openthoughts4_code9K_instill_n1_valredundancy5_round1.yaml new file mode 100644 index 0000000000000000000000000000000000000000..6e9f13613a602cb4f326bffacfcdbef37821aaa5 --- /dev/null +++ b/sdg/configs/qwen3_4b_openthoughts4_code9K_instill_n1_valredundancy5_round1.yaml @@ -0,0 +1,42 @@ +experiment_name: "qwen3_4b_openthoughts4_code9K_instill_n1_valredundancy5_round1" + +model_name: "Qwen/Qwen3-4B" +is_instruction_tuned: true +tensor_parallel_size: 1 +max_model_len: 32768 +gpu_memory_utilization: 0.95 + +dataset_hf_id: "teetone/ot4-code-deduped" +dataset_hf_revision: "840633b" +dataset_split: "train" +dataset_format: "openthoughts4" +check_boxed: false +limit: null + +num_generations: 1 +gen_temperature: 0.8 +gen_top_p: 0.95 +gen_max_tokens: 32768 +gen_batch_size: 64 + +validation_type: uq +selection_policy: first_valid +num_validation_votes: 5 +val_temperature: 0.8 +val_top_p: 0.95 +val_max_tokens: 32768 +val_batch_size: 64 + +output_dir: "output" +mongo_uri: "mongodb://john3.stanford.edu:27017" +num_shards: 30 + +machines_to_exclude: + - sphinx1 + - sphinx2 + - sphinx3 + - sphinx4 + - sphinx5 + - sphinx6 + - sphinx7 + - sphinx8 diff --git a/sdg/configs/qwen3_4b_openthoughts4_code9K_instill_n4_valredundancy1_round1.yaml b/sdg/configs/qwen3_4b_openthoughts4_code9K_instill_n4_valredundancy1_round1.yaml new file mode 100644 index 0000000000000000000000000000000000000000..4cb2c102799c1947fa9d98363405f7f2445aa249 --- /dev/null +++ b/sdg/configs/qwen3_4b_openthoughts4_code9K_instill_n4_valredundancy1_round1.yaml @@ -0,0 +1,42 @@ +experiment_name: "qwen3_4b_openthoughts4_code9K_instill_n4_valredundancy1_round1" + +model_name: "Qwen/Qwen3-4B" +is_instruction_tuned: true +tensor_parallel_size: 1 +max_model_len: 32768 +gpu_memory_utilization: 0.95 + +dataset_hf_id: "teetone/ot4-code-deduped" +dataset_hf_revision: "840633b" +dataset_split: "train" +dataset_format: "openthoughts4" +check_boxed: false +limit: null + +num_generations: 4 +gen_temperature: 0.8 +gen_top_p: 0.95 +gen_max_tokens: 32768 +gen_batch_size: 64 + +validation_type: uq +selection_policy: first_valid +num_validation_votes: 1 +val_temperature: 0.8 +val_top_p: 0.95 +val_max_tokens: 32768 +val_batch_size: 64 + +output_dir: "output" +mongo_uri: "mongodb://john3.stanford.edu:27017" +num_shards: 30 + +machines_to_exclude: + - sphinx1 + - sphinx2 + - sphinx3 + - sphinx4 + - sphinx5 + - sphinx6 + - sphinx7 + - sphinx8 diff --git a/sdg/configs/qwen3_4b_openthoughts4_code9K_instill_n4_valredundancy3_round1.yaml b/sdg/configs/qwen3_4b_openthoughts4_code9K_instill_n4_valredundancy3_round1.yaml new file mode 100644 index 0000000000000000000000000000000000000000..f0cf8477a85963ed8eb394ec344d28eba843e5dc --- /dev/null +++ b/sdg/configs/qwen3_4b_openthoughts4_code9K_instill_n4_valredundancy3_round1.yaml @@ -0,0 +1,42 @@ +experiment_name: "qwen3_4b_openthoughts4_code9K_instill_n4_valredundancy3_round1" + +model_name: "Qwen/Qwen3-4B" +is_instruction_tuned: true +tensor_parallel_size: 1 +max_model_len: 32768 +gpu_memory_utilization: 0.95 + +dataset_hf_id: "teetone/ot4-code-deduped" +dataset_hf_revision: "840633b" +dataset_split: "train" +dataset_format: "openthoughts4" +check_boxed: false +limit: null + +num_generations: 4 +gen_temperature: 0.8 +gen_top_p: 0.95 +gen_max_tokens: 32768 +gen_batch_size: 64 + +validation_type: uq +selection_policy: first_valid +num_validation_votes: 3 +val_temperature: 0.8 +val_top_p: 0.95 +val_max_tokens: 32768 +val_batch_size: 64 + +output_dir: "output" +mongo_uri: "mongodb://john3.stanford.edu:27017" +num_shards: 30 + +machines_to_exclude: + - sphinx1 + - sphinx2 + - sphinx3 + - sphinx4 + - sphinx5 + - sphinx6 + - sphinx7 + - sphinx8 diff --git a/sdg/configs/qwen3_4b_openthoughts4_code9K_instill_n4_valredundancy5_round1.yaml b/sdg/configs/qwen3_4b_openthoughts4_code9K_instill_n4_valredundancy5_round1.yaml new file mode 100644 index 0000000000000000000000000000000000000000..90489e77b3ab93a8cfedb02892993ae53ae23c89 --- /dev/null +++ b/sdg/configs/qwen3_4b_openthoughts4_code9K_instill_n4_valredundancy5_round1.yaml @@ -0,0 +1,42 @@ +experiment_name: "qwen3_4b_openthoughts4_code9K_instill_n4_valredundancy5_round1" + +model_name: "Qwen/Qwen3-4B" +is_instruction_tuned: true +tensor_parallel_size: 1 +max_model_len: 32768 +gpu_memory_utilization: 0.95 + +dataset_hf_id: "teetone/ot4-code-deduped" +dataset_hf_revision: "840633b" +dataset_split: "train" +dataset_format: "openthoughts4" +check_boxed: false +limit: null + +num_generations: 4 +gen_temperature: 0.8 +gen_top_p: 0.95 +gen_max_tokens: 32768 +gen_batch_size: 64 + +validation_type: uq +selection_policy: first_valid +num_validation_votes: 5 +val_temperature: 0.8 +val_top_p: 0.95 +val_max_tokens: 32768 +val_batch_size: 64 + +output_dir: "output" +mongo_uri: "mongodb://john3.stanford.edu:27017" +num_shards: 30 + +machines_to_exclude: + - sphinx1 + - sphinx2 + - sphinx3 + - sphinx4 + - sphinx5 + - sphinx6 + - sphinx7 + - sphinx8 diff --git a/sdg/configs/qwen3_4b_openthoughts4_code9K_instill_n8_no_filter_round1.yaml b/sdg/configs/qwen3_4b_openthoughts4_code9K_instill_n8_no_filter_round1.yaml new file mode 100644 index 0000000000000000000000000000000000000000..e0d0a2a31c17a67ce407819ccf858f6ba38dfed0 --- /dev/null +++ b/sdg/configs/qwen3_4b_openthoughts4_code9K_instill_n8_no_filter_round1.yaml @@ -0,0 +1,42 @@ +experiment_name: "qwen3_4b_openthoughts4_code9K_instill_n8_no_filter_round1" + +model_name: "Qwen/Qwen3-4B" +is_instruction_tuned: true +tensor_parallel_size: 1 +max_model_len: 32768 +gpu_memory_utilization: 0.95 + +dataset_hf_id: "teetone/ot4-code-deduped" +dataset_hf_revision: "840633b" +dataset_split: "train" +dataset_format: "openthoughts4" +check_boxed: false +limit: null + +num_generations: 8 +gen_temperature: 0.8 +gen_top_p: 0.95 +gen_max_tokens: 32768 +gen_batch_size: 64 + +validation_type: null +selection_policy: first_valid +num_validation_votes: 0 +val_temperature: 0.8 +val_top_p: 0.95 +val_max_tokens: 32768 +val_batch_size: 64 + +output_dir: "output" +mongo_uri: "mongodb://john3.stanford.edu:27017" +num_shards: 30 + +machines_to_exclude: + - sphinx1 + - sphinx2 + - sphinx3 + - sphinx4 + - sphinx5 + - sphinx6 + - sphinx7 + - sphinx8 diff --git a/sdg/configs/qwen3_4b_openthoughts4_code9K_instill_n8_valredundancy1_round1.yaml b/sdg/configs/qwen3_4b_openthoughts4_code9K_instill_n8_valredundancy1_round1.yaml new file mode 100644 index 0000000000000000000000000000000000000000..6938b6d77f249557960b0450b63904fc8edd11f7 --- /dev/null +++ b/sdg/configs/qwen3_4b_openthoughts4_code9K_instill_n8_valredundancy1_round1.yaml @@ -0,0 +1,42 @@ +experiment_name: "qwen3_4b_openthoughts4_code9K_instill_n8_valredundancy1_round1" + +model_name: "Qwen/Qwen3-4B" +is_instruction_tuned: true +tensor_parallel_size: 1 +max_model_len: 32768 +gpu_memory_utilization: 0.95 + +dataset_hf_id: "teetone/ot4-code-deduped" +dataset_hf_revision: "840633b" +dataset_split: "train" +dataset_format: "openthoughts4" +check_boxed: false +limit: null + +num_generations: 8 +gen_temperature: 0.8 +gen_top_p: 0.95 +gen_max_tokens: 32768 +gen_batch_size: 64 + +validation_type: uq +selection_policy: first_valid +num_validation_votes: 1 +val_temperature: 0.8 +val_top_p: 0.95 +val_max_tokens: 32768 +val_batch_size: 64 + +output_dir: "output" +mongo_uri: "mongodb://john3.stanford.edu:27017" +num_shards: 30 + +machines_to_exclude: + - sphinx1 + - sphinx2 + - sphinx3 + - sphinx4 + - sphinx5 + - sphinx6 + - sphinx7 + - sphinx8 diff --git a/sdg/configs/qwen3_4b_openthoughts4_code9K_instill_n8_valredundancy3_round1.yaml b/sdg/configs/qwen3_4b_openthoughts4_code9K_instill_n8_valredundancy3_round1.yaml new file mode 100644 index 0000000000000000000000000000000000000000..abd94d0309c879bf9b86fda4ea175515e7c1a5d9 --- /dev/null +++ b/sdg/configs/qwen3_4b_openthoughts4_code9K_instill_n8_valredundancy3_round1.yaml @@ -0,0 +1,42 @@ +experiment_name: "qwen3_4b_openthoughts4_code9K_instill_n8_valredundancy3_round1" + +model_name: "Qwen/Qwen3-4B" +is_instruction_tuned: true +tensor_parallel_size: 1 +max_model_len: 32768 +gpu_memory_utilization: 0.95 + +dataset_hf_id: "teetone/ot4-code-deduped" +dataset_hf_revision: "840633b" +dataset_split: "train" +dataset_format: "openthoughts4" +check_boxed: false +limit: null + +num_generations: 8 +gen_temperature: 0.8 +gen_top_p: 0.95 +gen_max_tokens: 32768 +gen_batch_size: 64 + +validation_type: uq +selection_policy: first_valid +num_validation_votes: 3 +val_temperature: 0.8 +val_top_p: 0.95 +val_max_tokens: 32768 +val_batch_size: 64 + +output_dir: "output" +mongo_uri: "mongodb://john3.stanford.edu:27017" +num_shards: 30 + +machines_to_exclude: + - sphinx1 + - sphinx2 + - sphinx3 + - sphinx4 + - sphinx5 + - sphinx6 + - sphinx7 + - sphinx8 diff --git a/sdg/configs/qwen3_4b_openthoughts4_code9K_instill_n8_valredundancy5_round1.yaml b/sdg/configs/qwen3_4b_openthoughts4_code9K_instill_n8_valredundancy5_round1.yaml new file mode 100644 index 0000000000000000000000000000000000000000..9c6e69fd693dbce2993cd10b98f9f3b9bcee3bd1 --- /dev/null +++ b/sdg/configs/qwen3_4b_openthoughts4_code9K_instill_n8_valredundancy5_round1.yaml @@ -0,0 +1,42 @@ +experiment_name: "qwen3_4b_openthoughts4_code9K_instill_n8_valredundancy5_round1" + +model_name: "Qwen/Qwen3-4B" +is_instruction_tuned: true +tensor_parallel_size: 1 +max_model_len: 32768 +gpu_memory_utilization: 0.95 + +dataset_hf_id: "teetone/ot4-code-deduped" +dataset_hf_revision: "840633b" +dataset_split: "train" +dataset_format: "openthoughts4" +check_boxed: false +limit: null + +num_generations: 8 +gen_temperature: 0.8 +gen_top_p: 0.95 +gen_max_tokens: 32768 +gen_batch_size: 64 + +validation_type: uq +selection_policy: first_valid +num_validation_votes: 5 +val_temperature: 0.8 +val_top_p: 0.95 +val_max_tokens: 32768 +val_batch_size: 64 + +output_dir: "output" +mongo_uri: "mongodb://john3.stanford.edu:27017" +num_shards: 30 + +machines_to_exclude: + - sphinx1 + - sphinx2 + - sphinx3 + - sphinx4 + - sphinx5 + - sphinx6 + - sphinx7 + - sphinx8 diff --git a/sdg/configs/qwen3_4b_openthoughts4_math219K_instill_n4_valredundancy3_round1.yaml b/sdg/configs/qwen3_4b_openthoughts4_math219K_instill_n4_valredundancy3_round1.yaml new file mode 100644 index 0000000000000000000000000000000000000000..7b96a75d7f1719e8d2bf6c61faf1d563b7537ecd --- /dev/null +++ b/sdg/configs/qwen3_4b_openthoughts4_math219K_instill_n4_valredundancy3_round1.yaml @@ -0,0 +1,42 @@ +experiment_name: "qwen3_4b_openthoughts4_math219K_instill_n4_valredundancy3_round1" + +model_name: "Qwen/Qwen3-4B" +is_instruction_tuned: true +tensor_parallel_size: 1 +max_model_len: 32768 +gpu_memory_utilization: 0.95 + +dataset_hf_id: "teetone/ot4-math-deduped" +dataset_hf_revision: "716443d" +dataset_split: "train" +dataset_format: "openthoughts4" +check_boxed: true +limit: null + +num_generations: 4 +gen_temperature: 0.8 +gen_top_p: 0.95 +gen_max_tokens: 32768 +gen_batch_size: 64 + +validation_type: uq +selection_policy: first_valid +num_validation_votes: 3 +val_temperature: 0.8 +val_top_p: 0.95 +val_max_tokens: 32768 +val_batch_size: 64 + +output_dir: "output" +mongo_uri: "mongodb://john3.stanford.edu:27017" +num_shards: 30 + +machines_to_exclude: + - sphinx1 + - sphinx2 + - sphinx3 + - sphinx4 + - sphinx5 + - sphinx6 + - sphinx7 + - sphinx8 diff --git a/sdg/configs/qwen3_4b_openthoughts4_math219K_instill_n8_valredundancy5_round1.yaml b/sdg/configs/qwen3_4b_openthoughts4_math219K_instill_n8_valredundancy5_round1.yaml new file mode 100644 index 0000000000000000000000000000000000000000..b3d18083432434410f01703fae7044c4ba4e481c --- /dev/null +++ b/sdg/configs/qwen3_4b_openthoughts4_math219K_instill_n8_valredundancy5_round1.yaml @@ -0,0 +1,42 @@ +experiment_name: "qwen3_4b_openthoughts4_math219K_instill_n8_valredundancy5_round1" + +model_name: "Qwen/Qwen3-4B" +is_instruction_tuned: true +tensor_parallel_size: 1 +max_model_len: 32768 +gpu_memory_utilization: 0.95 + +dataset_hf_id: "teetone/ot4-math-deduped" +dataset_hf_revision: "716443d" +dataset_split: "train" +dataset_format: "openthoughts4" +check_boxed: true +limit: null + +num_generations: 8 +gen_temperature: 0.8 +gen_top_p: 0.95 +gen_max_tokens: 32768 +gen_batch_size: 64 + +validation_type: uq +selection_policy: first_valid +num_validation_votes: 5 +val_temperature: 0.8 +val_top_p: 0.95 +val_max_tokens: 32768 +val_batch_size: 64 + +output_dir: "output" +mongo_uri: "mongodb://john3.stanford.edu:27017" +num_shards: 30 + +machines_to_exclude: + - sphinx1 + - sphinx2 + - sphinx3 + - sphinx4 + - sphinx5 + - sphinx6 + - sphinx7 + - sphinx8 \ No newline at end of file diff --git a/sdg/configs/qwen3_4b_openthoughts4_math30K_instill_n4_valredundancy3_round1.yaml b/sdg/configs/qwen3_4b_openthoughts4_math30K_instill_n4_valredundancy3_round1.yaml new file mode 100644 index 0000000000000000000000000000000000000000..f4e753482ce2d2b3af43ceb983546704633abff2 --- /dev/null +++ b/sdg/configs/qwen3_4b_openthoughts4_math30K_instill_n4_valredundancy3_round1.yaml @@ -0,0 +1,42 @@ +experiment_name: "qwen3_4b_openthoughts4_math30K_instill_n4_valredundancy3_round1" + +model_name: "Qwen/Qwen3-4B" +is_instruction_tuned: true +tensor_parallel_size: 1 +max_model_len: 32768 +gpu_memory_utilization: 0.95 + +dataset_hf_id: "marin-community/open-thoughts-4-30k-math-qwen3-32b-annotated-32768-tokens" +dataset_hf_revision: "6a05237" +dataset_split: "train" +dataset_format: "openthoughts4" +check_boxed: true +limit: null + +num_generations: 4 +gen_temperature: 0.8 +gen_top_p: 0.95 +gen_max_tokens: 32768 +gen_batch_size: 64 + +validation_type: uq +selection_policy: first_valid +num_validation_votes: 3 +val_temperature: 0.8 +val_top_p: 0.95 +val_max_tokens: 32768 +val_batch_size: 64 + +output_dir: "output" +mongo_uri: "mongodb://john3.stanford.edu:27017" +num_shards: 30 + +machines_to_exclude: + - sphinx1 + - sphinx2 + - sphinx3 + - sphinx4 + - sphinx5 + - sphinx6 + - sphinx7 + - sphinx8 diff --git a/sdg/configs/qwen3_4b_openthoughts4_science26K_instill_n1_valredundancy1_round1.yaml b/sdg/configs/qwen3_4b_openthoughts4_science26K_instill_n1_valredundancy1_round1.yaml new file mode 100644 index 0000000000000000000000000000000000000000..8fdf1587da4f48b1fab89e3dbd945f4ee951c03c --- /dev/null +++ b/sdg/configs/qwen3_4b_openthoughts4_science26K_instill_n1_valredundancy1_round1.yaml @@ -0,0 +1,42 @@ +experiment_name: "qwen3_4b_openthoughts4_science26K_instill_n1_valredundancy1_round1" + +model_name: "Qwen/Qwen3-4B" +is_instruction_tuned: true +tensor_parallel_size: 1 +max_model_len: 32768 +gpu_memory_utilization: 0.95 + +dataset_hf_id: "teetone/ot4-science-deduped" +dataset_hf_revision: "ddae07c" +dataset_split: "train" +dataset_format: "openthoughts4" +check_boxed: false +limit: null + +num_generations: 1 +gen_temperature: 0.8 +gen_top_p: 0.95 +gen_max_tokens: 32768 +gen_batch_size: 64 + +validation_type: uq +selection_policy: first_valid +num_validation_votes: 1 +val_temperature: 0.8 +val_top_p: 0.95 +val_max_tokens: 32768 +val_batch_size: 64 + +output_dir: "output" +mongo_uri: "mongodb://john3.stanford.edu:27017" +num_shards: 30 + +machines_to_exclude: + - sphinx1 + - sphinx2 + - sphinx3 + - sphinx4 + - sphinx5 + - sphinx6 + - sphinx7 + - sphinx8 diff --git a/sdg/configs/qwen3_4b_openthoughts4_science26K_instill_n1_valredundancy3_round1.yaml b/sdg/configs/qwen3_4b_openthoughts4_science26K_instill_n1_valredundancy3_round1.yaml new file mode 100644 index 0000000000000000000000000000000000000000..dd328115a49e482cf52fbf1bd03a6cb8cf9f9732 --- /dev/null +++ b/sdg/configs/qwen3_4b_openthoughts4_science26K_instill_n1_valredundancy3_round1.yaml @@ -0,0 +1,42 @@ +experiment_name: "qwen3_4b_openthoughts4_science26K_instill_n1_valredundancy3_round1" + +model_name: "Qwen/Qwen3-4B" +is_instruction_tuned: true +tensor_parallel_size: 1 +max_model_len: 32768 +gpu_memory_utilization: 0.95 + +dataset_hf_id: "teetone/ot4-science-deduped" +dataset_hf_revision: "ddae07c" +dataset_split: "train" +dataset_format: "openthoughts4" +check_boxed: false +limit: null + +num_generations: 1 +gen_temperature: 0.8 +gen_top_p: 0.95 +gen_max_tokens: 32768 +gen_batch_size: 64 + +validation_type: uq +selection_policy: first_valid +num_validation_votes: 3 +val_temperature: 0.8 +val_top_p: 0.95 +val_max_tokens: 32768 +val_batch_size: 64 + +output_dir: "output" +mongo_uri: "mongodb://john3.stanford.edu:27017" +num_shards: 30 + +machines_to_exclude: + - sphinx1 + - sphinx2 + - sphinx3 + - sphinx4 + - sphinx5 + - sphinx6 + - sphinx7 + - sphinx8 diff --git a/sdg/configs/qwen3_4b_openthoughts4_science26K_instill_n1_valredundancy5_round1.yaml b/sdg/configs/qwen3_4b_openthoughts4_science26K_instill_n1_valredundancy5_round1.yaml new file mode 100644 index 0000000000000000000000000000000000000000..71e779142addb1b3b5ad72b1b9b97382043a3a2e --- /dev/null +++ b/sdg/configs/qwen3_4b_openthoughts4_science26K_instill_n1_valredundancy5_round1.yaml @@ -0,0 +1,42 @@ +experiment_name: "qwen3_4b_openthoughts4_science26K_instill_n1_valredundancy5_round1" + +model_name: "Qwen/Qwen3-4B" +is_instruction_tuned: true +tensor_parallel_size: 1 +max_model_len: 32768 +gpu_memory_utilization: 0.95 + +dataset_hf_id: "teetone/ot4-science-deduped" +dataset_hf_revision: "ddae07c" +dataset_split: "train" +dataset_format: "openthoughts4" +check_boxed: false +limit: null + +num_generations: 1 +gen_temperature: 0.8 +gen_top_p: 0.95 +gen_max_tokens: 32768 +gen_batch_size: 64 + +validation_type: uq +selection_policy: first_valid +num_validation_votes: 5 +val_temperature: 0.8 +val_top_p: 0.95 +val_max_tokens: 32768 +val_batch_size: 64 + +output_dir: "output" +mongo_uri: "mongodb://john3.stanford.edu:27017" +num_shards: 30 + +machines_to_exclude: + - sphinx1 + - sphinx2 + - sphinx3 + - sphinx4 + - sphinx5 + - sphinx6 + - sphinx7 + - sphinx8 diff --git a/sdg/configs/qwen3_4b_openthoughts4_science26K_instill_n4_valredundancy1_round1.yaml b/sdg/configs/qwen3_4b_openthoughts4_science26K_instill_n4_valredundancy1_round1.yaml new file mode 100644 index 0000000000000000000000000000000000000000..688afe8785bd2640892c917836f474f4b6708271 --- /dev/null +++ b/sdg/configs/qwen3_4b_openthoughts4_science26K_instill_n4_valredundancy1_round1.yaml @@ -0,0 +1,42 @@ +experiment_name: "qwen3_4b_openthoughts4_science26K_instill_n4_valredundancy1_round1" + +model_name: "Qwen/Qwen3-4B" +is_instruction_tuned: true +tensor_parallel_size: 1 +max_model_len: 32768 +gpu_memory_utilization: 0.95 + +dataset_hf_id: "teetone/ot4-science-deduped" +dataset_hf_revision: "ddae07c" +dataset_split: "train" +dataset_format: "openthoughts4" +check_boxed: false +limit: null + +num_generations: 4 +gen_temperature: 0.8 +gen_top_p: 0.95 +gen_max_tokens: 32768 +gen_batch_size: 64 + +validation_type: uq +selection_policy: first_valid +num_validation_votes: 1 +val_temperature: 0.8 +val_top_p: 0.95 +val_max_tokens: 32768 +val_batch_size: 64 + +output_dir: "output" +mongo_uri: "mongodb://john3.stanford.edu:27017" +num_shards: 30 + +machines_to_exclude: + - sphinx1 + - sphinx2 + - sphinx3 + - sphinx4 + - sphinx5 + - sphinx6 + - sphinx7 + - sphinx8 diff --git a/sdg/configs/qwen3_4b_openthoughts4_science26K_instill_n4_valredundancy3_round1.yaml b/sdg/configs/qwen3_4b_openthoughts4_science26K_instill_n4_valredundancy3_round1.yaml new file mode 100644 index 0000000000000000000000000000000000000000..414cf2416febf389fd2a29cf3dcb15620322fe4e --- /dev/null +++ b/sdg/configs/qwen3_4b_openthoughts4_science26K_instill_n4_valredundancy3_round1.yaml @@ -0,0 +1,42 @@ +experiment_name: "qwen3_4b_openthoughts4_science26K_instill_n4_valredundancy3_round1" + +model_name: "Qwen/Qwen3-4B" +is_instruction_tuned: true +tensor_parallel_size: 1 +max_model_len: 32768 +gpu_memory_utilization: 0.95 + +dataset_hf_id: "teetone/ot4-science-deduped" +dataset_hf_revision: "ddae07c" +dataset_split: "train" +dataset_format: "openthoughts4" +check_boxed: false +limit: null + +num_generations: 4 +gen_temperature: 0.8 +gen_top_p: 0.95 +gen_max_tokens: 32768 +gen_batch_size: 64 + +validation_type: uq +selection_policy: first_valid +num_validation_votes: 3 +val_temperature: 0.8 +val_top_p: 0.95 +val_max_tokens: 32768 +val_batch_size: 64 + +output_dir: "output" +mongo_uri: "mongodb://john3.stanford.edu:27017" +num_shards: 30 + +machines_to_exclude: + - sphinx1 + - sphinx2 + - sphinx3 + - sphinx4 + - sphinx5 + - sphinx6 + - sphinx7 + - sphinx8 diff --git a/sdg/configs/qwen3_4b_openthoughts4_science26K_instill_n4_valredundancy5_round1.yaml b/sdg/configs/qwen3_4b_openthoughts4_science26K_instill_n4_valredundancy5_round1.yaml new file mode 100644 index 0000000000000000000000000000000000000000..4e60bfadc4b8ddc1c70b3059f926c323a77f4c72 --- /dev/null +++ b/sdg/configs/qwen3_4b_openthoughts4_science26K_instill_n4_valredundancy5_round1.yaml @@ -0,0 +1,42 @@ +experiment_name: "qwen3_4b_openthoughts4_science26K_instill_n4_valredundancy5_round1" + +model_name: "Qwen/Qwen3-4B" +is_instruction_tuned: true +tensor_parallel_size: 1 +max_model_len: 32768 +gpu_memory_utilization: 0.95 + +dataset_hf_id: "teetone/ot4-science-deduped" +dataset_hf_revision: "ddae07c" +dataset_split: "train" +dataset_format: "openthoughts4" +check_boxed: false +limit: null + +num_generations: 4 +gen_temperature: 0.8 +gen_top_p: 0.95 +gen_max_tokens: 32768 +gen_batch_size: 64 + +validation_type: uq +selection_policy: first_valid +num_validation_votes: 5 +val_temperature: 0.8 +val_top_p: 0.95 +val_max_tokens: 32768 +val_batch_size: 64 + +output_dir: "output" +mongo_uri: "mongodb://john3.stanford.edu:27017" +num_shards: 30 + +machines_to_exclude: + - sphinx1 + - sphinx2 + - sphinx3 + - sphinx4 + - sphinx5 + - sphinx6 + - sphinx7 + - sphinx8 diff --git a/sdg/configs/qwen3_4b_openthoughts4_science26K_instill_n8_no_filter_round1.yaml b/sdg/configs/qwen3_4b_openthoughts4_science26K_instill_n8_no_filter_round1.yaml new file mode 100644 index 0000000000000000000000000000000000000000..5ad660232012ecb6404b8603c7c166247579b733 --- /dev/null +++ b/sdg/configs/qwen3_4b_openthoughts4_science26K_instill_n8_no_filter_round1.yaml @@ -0,0 +1,42 @@ +experiment_name: "qwen3_4b_openthoughts4_science26K_instill_n8_no_filter_round1" + +model_name: "Qwen/Qwen3-4B" +is_instruction_tuned: true +tensor_parallel_size: 1 +max_model_len: 32768 +gpu_memory_utilization: 0.95 + +dataset_hf_id: "teetone/ot4-science-deduped" +dataset_hf_revision: "ddae07c" +dataset_split: "train" +dataset_format: "openthoughts4" +check_boxed: false +limit: null + +num_generations: 8 +gen_temperature: 0.8 +gen_top_p: 0.95 +gen_max_tokens: 32768 +gen_batch_size: 64 + +validation_type: null +selection_policy: first_valid +num_validation_votes: 0 +val_temperature: 0.8 +val_top_p: 0.95 +val_max_tokens: 32768 +val_batch_size: 64 + +output_dir: "output" +mongo_uri: "mongodb://john3.stanford.edu:27017" +num_shards: 30 + +machines_to_exclude: + - sphinx1 + - sphinx2 + - sphinx3 + - sphinx4 + - sphinx5 + - sphinx6 + - sphinx7 + - sphinx8 diff --git a/sdg/configs/qwen3_4b_openthoughts4_science26K_instill_n8_valredundancy1_round1.yaml b/sdg/configs/qwen3_4b_openthoughts4_science26K_instill_n8_valredundancy1_round1.yaml new file mode 100644 index 0000000000000000000000000000000000000000..d8fccb555a26f03fef77ac0c9547fed8066da3fc --- /dev/null +++ b/sdg/configs/qwen3_4b_openthoughts4_science26K_instill_n8_valredundancy1_round1.yaml @@ -0,0 +1,42 @@ +experiment_name: "qwen3_4b_openthoughts4_science26K_instill_n8_valredundancy1_round1" + +model_name: "Qwen/Qwen3-4B" +is_instruction_tuned: true +tensor_parallel_size: 1 +max_model_len: 32768 +gpu_memory_utilization: 0.95 + +dataset_hf_id: "teetone/ot4-science-deduped" +dataset_hf_revision: "ddae07c" +dataset_split: "train" +dataset_format: "openthoughts4" +check_boxed: false +limit: null + +num_generations: 8 +gen_temperature: 0.8 +gen_top_p: 0.95 +gen_max_tokens: 32768 +gen_batch_size: 64 + +validation_type: uq +selection_policy: first_valid +num_validation_votes: 1 +val_temperature: 0.8 +val_top_p: 0.95 +val_max_tokens: 32768 +val_batch_size: 64 + +output_dir: "output" +mongo_uri: "mongodb://john3.stanford.edu:27017" +num_shards: 30 + +machines_to_exclude: + - sphinx1 + - sphinx2 + - sphinx3 + - sphinx4 + - sphinx5 + - sphinx6 + - sphinx7 + - sphinx8 diff --git a/sdg/configs/qwen3_4b_openthoughts4_science26K_instill_n8_valredundancy3_round1.yaml b/sdg/configs/qwen3_4b_openthoughts4_science26K_instill_n8_valredundancy3_round1.yaml new file mode 100644 index 0000000000000000000000000000000000000000..31225add9d870bd317a2a4dc176d225d0fc221b4 --- /dev/null +++ b/sdg/configs/qwen3_4b_openthoughts4_science26K_instill_n8_valredundancy3_round1.yaml @@ -0,0 +1,41 @@ +experiment_name: "qwen3_4b_openthoughts4_science26K_instill_n8_valredundancy3_round1" + +model_name: "Qwen/Qwen3-4B" +is_instruction_tuned: true +tensor_parallel_size: 1 +max_model_len: 32768 +gpu_memory_utilization: 0.95 + +dataset_hf_id: "teetone/ot4-science-deduped" +dataset_hf_revision: "ddae07c" +dataset_split: "train" +dataset_format: "openthoughts4" +check_boxed: false +limit: null + +num_generations: 8 +gen_temperature: 0.8 +gen_top_p: 0.95 +gen_max_tokens: 32768 +gen_batch_size: 64 + +validation_type: uq +selection_policy: first_valid +num_validation_votes: 3 +val_temperature: 0.8 +val_top_p: 0.95 +val_max_tokens: 32768 +val_batch_size: 64 + +output_dir: "output" +mongo_uri: "mongodb://john3.stanford.edu:27017" +num_shards: 30 +machines_to_exclude: + - sphinx1 + - sphinx2 + - sphinx3 + - sphinx4 + - sphinx5 + - sphinx6 + - sphinx7 + - sphinx8 diff --git a/sdg/configs/qwen3_4b_openthoughts4_science26K_instill_n8_valredundancy5_round1.yaml b/sdg/configs/qwen3_4b_openthoughts4_science26K_instill_n8_valredundancy5_round1.yaml new file mode 100644 index 0000000000000000000000000000000000000000..45c246538a2ead46ca2b3568c43a403a180ce1b7 --- /dev/null +++ b/sdg/configs/qwen3_4b_openthoughts4_science26K_instill_n8_valredundancy5_round1.yaml @@ -0,0 +1,42 @@ +experiment_name: "qwen3_4b_openthoughts4_science26K_instill_n8_valredundancy5_round1" + +model_name: "Qwen/Qwen3-4B" +is_instruction_tuned: true +tensor_parallel_size: 1 +max_model_len: 32768 +gpu_memory_utilization: 0.95 + +dataset_hf_id: "teetone/ot4-science-deduped" +dataset_hf_revision: "ddae07c" +dataset_split: "train" +dataset_format: "openthoughts4" +check_boxed: false +limit: null + +num_generations: 8 +gen_temperature: 0.8 +gen_top_p: 0.95 +gen_max_tokens: 32768 +gen_batch_size: 64 + +validation_type: uq +selection_policy: first_valid +num_validation_votes: 5 +val_temperature: 0.8 +val_top_p: 0.95 +val_max_tokens: 32768 +val_batch_size: 64 + +output_dir: "output" +mongo_uri: "mongodb://john3.stanford.edu:27017" +num_shards: 30 + +machines_to_exclude: + - sphinx1 + - sphinx2 + - sphinx3 + - sphinx4 + - sphinx5 + - sphinx6 + - sphinx7 + - sphinx8 diff --git a/sdg/configs/qwen3_4b_rstarcoder_seedsft_instill_n8_valredundancy5_round1.yaml b/sdg/configs/qwen3_4b_rstarcoder_seedsft_instill_n8_valredundancy5_round1.yaml new file mode 100644 index 0000000000000000000000000000000000000000..3dd947540f91d26fa1524045a90f1d6dc910c041 --- /dev/null +++ b/sdg/configs/qwen3_4b_rstarcoder_seedsft_instill_n8_valredundancy5_round1.yaml @@ -0,0 +1,42 @@ +experiment_name: "qwen3_4b_rstarcoder_instill_n8_valredundancy5_round1" + +model_name: "Qwen/Qwen3-4B" +is_instruction_tuned: true +tensor_parallel_size: 1 +max_model_len: 32768 +gpu_memory_utilization: 0.95 + +dataset_hf_id: "teetone/rstarcoder-seedsft-deduped" +dataset_hf_revision: "1242238" +dataset_split: "train" +dataset_format: "openthoughts4" +check_boxed: false +limit: null + +num_generations: 8 +gen_temperature: 0.8 +gen_top_p: 0.95 +gen_max_tokens: 32768 +gen_batch_size: 64 + +validation_type: uq +selection_policy: first_valid +num_validation_votes: 5 +val_temperature: 0.8 +val_top_p: 0.95 +val_max_tokens: 32768 +val_batch_size: 64 + +output_dir: "output" +mongo_uri: "mongodb://john3.stanford.edu:27017" +num_shards: 30 + +machines_to_exclude: + - sphinx1 + - sphinx2 + - sphinx3 + - sphinx4 + - sphinx5 + - sphinx6 + - sphinx7 + - sphinx8 diff --git a/sdg/configs/qwen3_8b_openthoughts3_math53K_instill_n8_valredundancy5_round1.yaml b/sdg/configs/qwen3_8b_openthoughts3_math53K_instill_n8_valredundancy5_round1.yaml new file mode 100644 index 0000000000000000000000000000000000000000..b3c356ac177fc79d438fbf79fe808163aa60139b --- /dev/null +++ b/sdg/configs/qwen3_8b_openthoughts3_math53K_instill_n8_valredundancy5_round1.yaml @@ -0,0 +1,41 @@ +experiment_name: "qwen3_8b_openthoughts3_math53K_instill_n8_valredundancy5_round1" + +model_name: "Qwen/Qwen3-8B" +is_instruction_tuned: true +tensor_parallel_size: 1 +max_model_len: 32768 +gpu_memory_utilization: 0.95 + +dataset_hf_id: "teetone/openthoughts3-math-deduped-53K" +dataset_split: "train" +dataset_format: "openthoughts4" +check_boxed: true +limit: null + +num_generations: 8 +gen_temperature: 0.8 +gen_top_p: 0.95 +gen_max_tokens: 32768 +gen_batch_size: 64 + +validation_type: uq +selection_policy: first_valid +num_validation_votes: 5 +val_temperature: 0.8 +val_top_p: 0.95 +val_max_tokens: 32768 +val_batch_size: 64 + +output_dir: "output" +mongo_uri: "mongodb://john3.stanford.edu:27017" +num_shards: 30 + +machines_to_exclude: + - sphinx1 + - sphinx2 + - sphinx3 + - sphinx4 + - sphinx5 + - sphinx6 + - sphinx7 + - sphinx8 diff --git a/sdg/configs/qwen3_8b_openthoughts4_code9K_instill_n8_valredundancy5_round1.yaml b/sdg/configs/qwen3_8b_openthoughts4_code9K_instill_n8_valredundancy5_round1.yaml new file mode 100644 index 0000000000000000000000000000000000000000..567310eaa5c93972f8a2eed2c9ee88043579eedf --- /dev/null +++ b/sdg/configs/qwen3_8b_openthoughts4_code9K_instill_n8_valredundancy5_round1.yaml @@ -0,0 +1,42 @@ +experiment_name: "qwen3_8b_openthoughts4_code9K_instill_n8_valredundancy5_round1" + +model_name: "Qwen/Qwen3-8B" +is_instruction_tuned: true +tensor_parallel_size: 1 +max_model_len: 32768 +gpu_memory_utilization: 0.95 + +dataset_hf_id: "teetone/ot4-code-deduped" +dataset_hf_revision: "840633b" +dataset_split: "train" +dataset_format: "openthoughts4" +check_boxed: false +limit: null + +num_generations: 8 +gen_temperature: 0.8 +gen_top_p: 0.95 +gen_max_tokens: 32768 +gen_batch_size: 64 + +validation_type: uq +selection_policy: first_valid +num_validation_votes: 5 +val_temperature: 0.8 +val_top_p: 0.95 +val_max_tokens: 32768 +val_batch_size: 64 + +output_dir: "output" +mongo_uri: "mongodb://john3.stanford.edu:27017" +num_shards: 30 + +machines_to_exclude: + - sphinx1 + - sphinx2 + - sphinx3 + - sphinx4 + - sphinx5 + - sphinx6 + - sphinx7 + - sphinx8 diff --git a/sdg/configs/qwen3_8b_openthoughts4_math219K_instill_n4_valredundancy3_round1.yaml b/sdg/configs/qwen3_8b_openthoughts4_math219K_instill_n4_valredundancy3_round1.yaml new file mode 100644 index 0000000000000000000000000000000000000000..a32300acc8112a831959f8dcab43d49ec5c23790 --- /dev/null +++ b/sdg/configs/qwen3_8b_openthoughts4_math219K_instill_n4_valredundancy3_round1.yaml @@ -0,0 +1,42 @@ +experiment_name: "qwen3_8b_openthoughts4_math219K_instill_n4_valredundancy3_round1" + +model_name: "Qwen/Qwen3-8B" +is_instruction_tuned: true +tensor_parallel_size: 1 +max_model_len: 32768 +gpu_memory_utilization: 0.95 + +dataset_hf_id: "teetone/ot4-math-deduped" +dataset_hf_revision: "716443d" +dataset_split: "train" +dataset_format: "openthoughts4" +check_boxed: true +limit: null + +num_generations: 4 +gen_temperature: 0.8 +gen_top_p: 0.95 +gen_max_tokens: 32768 +gen_batch_size: 64 + +validation_type: uq +selection_policy: first_valid +num_validation_votes: 3 +val_temperature: 0.8 +val_top_p: 0.95 +val_max_tokens: 32768 +val_batch_size: 64 + +output_dir: "output" +mongo_uri: "mongodb://john3.stanford.edu:27017" +num_shards: 30 + +machines_to_exclude: + - sphinx1 + - sphinx2 + - sphinx3 + - sphinx4 + - sphinx5 + - sphinx6 + - sphinx7 + - sphinx8 diff --git a/sdg/configs/qwen3_8b_openthoughts4_math219K_instill_n8_valredundancy3_round1.yaml b/sdg/configs/qwen3_8b_openthoughts4_math219K_instill_n8_valredundancy3_round1.yaml new file mode 100644 index 0000000000000000000000000000000000000000..5ea4dc1d003a69a1ce38595461a0da422ceed2e3 --- /dev/null +++ b/sdg/configs/qwen3_8b_openthoughts4_math219K_instill_n8_valredundancy3_round1.yaml @@ -0,0 +1,42 @@ +gexperiment_name: "qwen3_8b_openthoughts4_math219K_instill_n8_valredundancy3_round1" + +model_name: "Qwen/Qwen3-8B" +is_instruction_tuned: true +tensor_parallel_size: 1 +max_model_len: 32768 +gpu_memory_utilization: 0.95 + +dataset_hf_id: "teetone/ot4-math-deduped" +dataset_hf_revision: "716443d" +dataset_split: "train" +dataset_format: "openthoughts4" +check_boxed: true +limit: null + +num_generations: 8 +gen_temperature: 0.8 +gen_top_p: 0.95 +gen_max_tokens: 32768 +gen_batch_size: 64 + +validation_type: uq +selection_policy: first_valid +num_validation_votes: 3 +val_temperature: 0.8 +val_top_p: 0.95 +val_max_tokens: 32768 +val_batch_size: 64 + +output_dir: "output" +mongo_uri: "mongodb://john3.stanford.edu:27017" +num_shards: 30 + +machines_to_exclude: + - sphinx1 + - sphinx2 + - sphinx3 + - sphinx4 + - sphinx5 + - sphinx6 + - sphinx7 + - sphinx8 diff --git a/sdg/configs/qwen3_8b_openthoughts4_math219K_instill_n8_valredundancy5_round1.yaml b/sdg/configs/qwen3_8b_openthoughts4_math219K_instill_n8_valredundancy5_round1.yaml new file mode 100644 index 0000000000000000000000000000000000000000..e74c675073efab585df3d8f2b5098362ef61e1eb --- /dev/null +++ b/sdg/configs/qwen3_8b_openthoughts4_math219K_instill_n8_valredundancy5_round1.yaml @@ -0,0 +1,42 @@ +experiment_name: "qwen3_8b_openthoughts4_math219K_instill_n8_valredundancy5_round1" + +model_name: "Qwen/Qwen3-8B" +is_instruction_tuned: true +tensor_parallel_size: 1 +max_model_len: 32768 +gpu_memory_utilization: 0.95 + +dataset_hf_id: "teetone/ot4-math-deduped" +dataset_hf_revision: "716443d" +dataset_split: "train" +dataset_format: "openthoughts4" +check_boxed: true +limit: null + +num_generations: 8 +gen_temperature: 0.8 +gen_top_p: 0.95 +gen_max_tokens: 32768 +gen_batch_size: 64 + +validation_type: uq +selection_policy: first_valid +num_validation_votes: 5 +val_temperature: 0.8 +val_top_p: 0.95 +val_max_tokens: 32768 +val_batch_size: 64 + +output_dir: "output" +mongo_uri: "mongodb://john3.stanford.edu:27017" +num_shards: 30 + +machines_to_exclude: + - sphinx1 + - sphinx2 + - sphinx3 + - sphinx4 + - sphinx5 + - sphinx6 + - sphinx7 + - sphinx8 diff --git a/sdg/configs/qwen3_8b_openthoughts4_science26K_instill_n4_valredundancy5_round1.yaml b/sdg/configs/qwen3_8b_openthoughts4_science26K_instill_n4_valredundancy5_round1.yaml new file mode 100644 index 0000000000000000000000000000000000000000..74f7f2350aea55a9f12c5b2cf0d031c44f5f8fbc --- /dev/null +++ b/sdg/configs/qwen3_8b_openthoughts4_science26K_instill_n4_valredundancy5_round1.yaml @@ -0,0 +1,42 @@ +experiment_name: "qwen3_8b_openthoughts4_science26K_instill_n4_valredundancy5_round1" + +model_name: "Qwen/Qwen3-8B" +is_instruction_tuned: true +tensor_parallel_size: 1 +max_model_len: 32768 +gpu_memory_utilization: 0.95 + +dataset_hf_id: "teetone/ot4-science-deduped" +dataset_hf_revision: "ddae07c" +dataset_split: "train" +dataset_format: "openthoughts4" +check_boxed: false +limit: null + +num_generations: 4 +gen_temperature: 0.8 +gen_top_p: 0.95 +gen_max_tokens: 32768 +gen_batch_size: 64 + +validation_type: uq +selection_policy: first_valid +num_validation_votes: 5 +val_temperature: 0.8 +val_top_p: 0.95 +val_max_tokens: 32768 +val_batch_size: 64 + +output_dir: "output" +mongo_uri: "mongodb://john3.stanford.edu:27017" +num_shards: 30 + +machines_to_exclude: + - sphinx1 + - sphinx2 + - sphinx3 + - sphinx4 + - sphinx5 + - sphinx6 + - sphinx7 + - sphinx8 diff --git a/sdg/configs/qwen3_8b_openthoughts4_science26K_instill_n8_valredundancy3_round1.yaml b/sdg/configs/qwen3_8b_openthoughts4_science26K_instill_n8_valredundancy3_round1.yaml new file mode 100644 index 0000000000000000000000000000000000000000..5a987a945833b258cae03c49aa3fadd42d081938 --- /dev/null +++ b/sdg/configs/qwen3_8b_openthoughts4_science26K_instill_n8_valredundancy3_round1.yaml @@ -0,0 +1,42 @@ +experiment_name: "qwen3_8b_openthoughts4_science26K_instill_n8_valredundancy3_round1" + +model_name: "Qwen/Qwen3-8B" +is_instruction_tuned: true +tensor_parallel_size: 1 +max_model_len: 32768 +gpu_memory_utilization: 0.95 + +dataset_hf_id: "teetone/ot4-science-deduped" +dataset_hf_revision: "ddae07c" +dataset_split: "train" +dataset_format: "openthoughts4" +check_boxed: false +limit: null + +num_generations: 8 +gen_temperature: 0.8 +gen_top_p: 0.95 +gen_max_tokens: 32768 +gen_batch_size: 64 + +validation_type: uq +selection_policy: first_valid +num_validation_votes: 3 +val_temperature: 0.8 +val_top_p: 0.95 +val_max_tokens: 32768 +val_batch_size: 64 + +output_dir: "output" +mongo_uri: "mongodb://john3.stanford.edu:27017" +num_shards: 30 + +machines_to_exclude: + - sphinx1 + - sphinx2 + - sphinx3 + - sphinx4 + - sphinx5 + - sphinx6 + - sphinx7 + - sphinx8 diff --git a/sdg/configs/qwen3_8b_openthoughts4_science26K_instill_n8_valredundancy5_round1.yaml b/sdg/configs/qwen3_8b_openthoughts4_science26K_instill_n8_valredundancy5_round1.yaml new file mode 100644 index 0000000000000000000000000000000000000000..72fed2e0aafdff4703baaae150af54774f9d00b7 --- /dev/null +++ b/sdg/configs/qwen3_8b_openthoughts4_science26K_instill_n8_valredundancy5_round1.yaml @@ -0,0 +1,42 @@ +experiment_name: "qwen3_8b_openthoughts4_science26K_instill_n8_valredundancy5_round1" + +model_name: "Qwen/Qwen3-8B" +is_instruction_tuned: true +tensor_parallel_size: 1 +max_model_len: 32768 +gpu_memory_utilization: 0.95 + +dataset_hf_id: "teetone/ot4-science-deduped" +dataset_hf_revision: "ddae07c" +dataset_split: "train" +dataset_format: "openthoughts4" +check_boxed: false +limit: null + +num_generations: 8 +gen_temperature: 0.8 +gen_top_p: 0.95 +gen_max_tokens: 32768 +gen_batch_size: 64 + +validation_type: uq +selection_policy: first_valid +num_validation_votes: 5 +val_temperature: 0.8 +val_top_p: 0.95 +val_max_tokens: 32768 +val_batch_size: 64 + +output_dir: "output" +mongo_uri: "mongodb://john3.stanford.edu:27017" +num_shards: 30 + +machines_to_exclude: + - sphinx1 + - sphinx2 + - sphinx3 + - sphinx4 + - sphinx5 + - sphinx6 + - sphinx7 + - sphinx8 diff --git a/sdg/configs/qwen3_8b_rstarcoder_seedsft_instill_n8_valredundancy5_round1.yaml b/sdg/configs/qwen3_8b_rstarcoder_seedsft_instill_n8_valredundancy5_round1.yaml new file mode 100644 index 0000000000000000000000000000000000000000..2d0e712ea740d08eda51a78f46a9d995c9c9e9b6 --- /dev/null +++ b/sdg/configs/qwen3_8b_rstarcoder_seedsft_instill_n8_valredundancy5_round1.yaml @@ -0,0 +1,42 @@ +experiment_name: "qwen3_8b_rstarcoder_instill_n8_valredundancy5_round1" + +model_name: "Qwen/Qwen3-8B" +is_instruction_tuned: true +tensor_parallel_size: 1 +max_model_len: 32768 +gpu_memory_utilization: 0.95 + +dataset_hf_id: "teetone/rstarcoder-seedsft-deduped" +dataset_hf_revision: "1242238" +dataset_split: "train" +dataset_format: "openthoughts4" +check_boxed: false +limit: null + +num_generations: 8 +gen_temperature: 0.8 +gen_top_p: 0.95 +gen_max_tokens: 32768 +gen_batch_size: 64 + +validation_type: uq +selection_policy: first_valid +num_validation_votes: 5 +val_temperature: 0.8 +val_top_p: 0.95 +val_max_tokens: 32768 +val_batch_size: 64 + +output_dir: "output" +mongo_uri: "mongodb://jagupard39.stanford.edu:27017" +num_shards: 30 + +machines_to_exclude: + - sphinx1 + - sphinx2 + - sphinx3 + - sphinx4 + - sphinx5 + - sphinx6 + - sphinx7 + - sphinx8 diff --git a/sdg/generate.py b/sdg/generate.py new file mode 100644 index 0000000000000000000000000000000000000000..3a399e6cfc2f5f176937a9a31a411a3e6f0b9fa0 --- /dev/null +++ b/sdg/generate.py @@ -0,0 +1,66 @@ +""" +CLI entry point for the SDG pipeline. + +Usage: + python -m sdg.generate --config sdg/configs/qwen3_4b_mot.yaml + python -m sdg.generate --config sdg/configs/qwen3_4b_openthoughts4.yaml --limit 100 +""" + +import argparse + +from sdg.config import SDGConfig +from sdg.pipeline import run_pipeline + + +def main() -> None: + parser = argparse.ArgumentParser( + description="Standalone SDG pipeline (generate + UQ validation)" + ) + parser.add_argument( + "--config", required=True, help="Path to YAML config file" + ) + + # CLI overrides + parser.add_argument("--model", dest="model_name") + parser.add_argument("--dataset", dest="dataset_hf_id") + parser.add_argument("--dataset-config", dest="dataset_hf_config") + parser.add_argument("--limit", type=int) + parser.add_argument("--num-generations", type=int) + parser.add_argument("--num-validation-votes", type=int) + parser.add_argument("--max-rounds", type=int) + parser.add_argument("--output-dir") + parser.add_argument("--tensor-parallel-size", type=int) + + # Sharding + parser.add_argument("--shard-id", type=int, default=None, help="Shard index (0-based)") + parser.add_argument("--num-shards", type=int, default=None, help="Total number of shards") + + args = parser.parse_args() + + # Load base config from YAML + config = SDGConfig.from_yaml(args.config) + + # Apply CLI overrides + overrides = { + "model_name": args.model_name, + "dataset_hf_id": args.dataset_hf_id, + "dataset_hf_config": args.dataset_hf_config, + "limit": args.limit, + "num_generations": args.num_generations, + "num_validation_votes": args.num_validation_votes, + "max_rounds": args.max_rounds, + "output_dir": args.output_dir, + "tensor_parallel_size": args.tensor_parallel_size, + } + for key, value in overrides.items(): + if value is not None: + setattr(config, key, value) + + if (args.shard_id is None) != (args.num_shards is None): + parser.error("--shard-id and --num-shards must be provided together") + + run_pipeline(config, shard_id=args.shard_id, num_shards=args.num_shards) + + +if __name__ == "__main__": + main() diff --git a/sdg/inference.py b/sdg/inference.py new file mode 100644 index 0000000000000000000000000000000000000000..625260b54c2f75946e0d6e247e316687c3fef816 --- /dev/null +++ b/sdg/inference.py @@ -0,0 +1,161 @@ +""" +vLLM wrapper for generation and validation inference. + +Pattern follows filter_with_uq.py — single LLM instance shared for both +generation (n=num_generations) and validation (n=num_validation_votes). + +Inference results are cached per-prompt in MongoDB so that interrupted runs +can resume without re-doing expensive vLLM calls, and multiple shards can +share a single cache. +""" + +from __future__ import annotations + +import hashlib +import json + +from pymongo import MongoClient +from vllm import LLM, SamplingParams + +from sdg.config import SDGConfig + + +class MongoCache: + """Dict-like interface over a MongoDB collection for inference caching.""" + + def __init__(self, uri: str, db_name: str): + self.client = MongoClient(uri) + self.db = self.client[db_name] + self.collection = self.db["inference_cache"] + self.collection.create_index("key", unique=True) + # Fail fast if MongoDB is unreachable + self.client.admin.command("ping") + + def get(self, key: str, default=None): + doc = self.collection.find_one({"key": key}) + if doc is not None: + return doc["value"] + return default + + def __setitem__(self, key: str, value): + self.collection.update_one( + {"key": key}, + {"$set": {"key": key, "value": value}}, + upsert=True, + ) + + +class VLLMEngine: + """Thin wrapper around vllm.LLM with two pre-built SamplingParams.""" + + def __init__(self, config: SDGConfig): + print(f"Loading model: {config.model_name}") + self.model_name = config.model_name + self.llm = LLM( + model=config.model_name, + tensor_parallel_size=config.tensor_parallel_size, + max_model_len=config.max_model_len, + gpu_memory_utilization=config.gpu_memory_utilization, + trust_remote_code=True, + ) + + self.gen_params = SamplingParams( + temperature=config.gen_temperature, + top_p=config.gen_top_p, + max_tokens=config.gen_max_tokens, + n=config.num_generations, + ) + + if config.num_validation_votes > 0: + self.val_params = SamplingParams( + temperature=config.val_temperature, + top_p=config.val_top_p, + max_tokens=config.val_max_tokens, + n=config.num_validation_votes, + ) + else: + self.val_params = None + + # ── inference cache ────────────────────────────────────────────── + db_name = config.mongo_db_name + print(f"Connecting to MongoDB: {config.mongo_uri} db={db_name}") + self.cache = MongoCache(config.mongo_uri, db_name) + print("MongoDB cache connected successfully") + print("Model loaded successfully") + + # ── helpers ────────────────────────────────────────────────────────── + + def _cache_key(self, prompt: str, params: SamplingParams) -> str: + raw = json.dumps( + [self.model_name, params.temperature, params.top_p, + params.max_tokens, params.n, prompt], + ensure_ascii=False, + ) + return hashlib.sha256(raw.encode()).hexdigest() + + def apply_chat_template(self, prompt: str) -> str: + tokenizer = self.llm.get_tokenizer() + messages = [{"role": "user", "content": prompt}] + return tokenizer.apply_chat_template( + messages, tokenize=False, add_generation_prompt=True + ) + + def _generate_cached( + self, prompts: list[str], params: SamplingParams, + ) -> list[list[str]]: + """Generate with per-prompt caching. Returns list[list[str]].""" + templated = [self.apply_chat_template(p) for p in prompts] + + # Split into cached hits vs misses + results: list[list[str] | None] = [None] * len(prompts) + miss_indices: list[int] = [] + miss_templated: list[str] = [] + cache_hits = 0 + + for i, (prompt, tmpl) in enumerate(zip(prompts, templated)): + key = self._cache_key(tmpl, params) + cached = self.cache.get(key) + if cached is not None: + results[i] = cached + cache_hits += 1 + else: + miss_indices.append(i) + miss_templated.append(tmpl) + + # Run vLLM only for misses + if miss_templated: + outputs = self.llm.generate(miss_templated, params) + assert len(outputs) == len(miss_indices), ( + f"vLLM returned {len(outputs)} outputs for {len(miss_indices)} prompts" + ) + for idx, tmpl, out in zip(miss_indices, miss_templated, outputs): + texts = [o.text for o in out.outputs] + key = self._cache_key(tmpl, params) + self.cache[key] = texts + results[idx] = texts + + if cache_hits: + print(f" Cache: {cache_hits} hits, {len(miss_indices)} misses") + + return results # type: ignore[return-value] + + # ── public API ─────────────────────────────────────────────────────── + + def generate_multi(self, prompts: list[str]) -> list[list[str]]: + """Generate *num_generations* samples per prompt (candidate generation).""" + return self._generate_cached(prompts, self.gen_params) + + def generate_with_votes(self, prompts: list[str]) -> list[list[str]]: + """Generate *num_validation_votes* samples per prompt (UQ voting).""" + return self._generate_cached(prompts, self.val_params) + + def generate_single(self, prompts: list[str]) -> list[str]: + """Generate one sample per prompt (e.g. inferred-question generation).""" + single_params = SamplingParams( + temperature=self.val_params.temperature, + top_p=self.val_params.top_p, + max_tokens=self.val_params.max_tokens, + n=1, + ) + results = self._generate_cached(prompts, single_params) + return [r[0] for r in results] diff --git a/sdg/launch_mongodb.py b/sdg/launch_mongodb.py new file mode 100644 index 0000000000000000000000000000000000000000..6ea9daffb0ca97e9fb3a0b8f7848d226eaabee7c --- /dev/null +++ b/sdg/launch_mongodb.py @@ -0,0 +1,52 @@ +""" +Standalone script that starts a MongoDB server locally and prints the mongo-uri. + +Uses ``mongod`` with a configurable port and dbpath. Runs indefinitely +(foreground process) so that worker shards can connect to it. + +Usage: + python -m sdg.launch_mongodb + python -m sdg.launch_mongodb --port 27018 --dbpath /tmp/sdg_mongo + +nlprun command (john queue, no GPU, 16 GB RAM): + nlprun -a tonyreasoningtraces -q john --exclude john17 -c 4 -g 0 --memory 64g \ + -w /nlp/scr4/nlp/crfm/text2image/text2image-rlhf/reasoning/virtual-world-data \ + --job-name sdg-mongodb "python -m sdg.launch_mongodb 2>&1 | tee mongo.log" +""" + +import argparse +import shutil +import socket +import subprocess +import sys +from pathlib import Path + + +def main() -> None: + parser = argparse.ArgumentParser(description="Start a MongoDB server for SDG cache") + parser.add_argument("--port", type=int, default=27017, help="MongoDB port (default: 27017)") + parser.add_argument("--dbpath", default="/nlp/scr4/nlp/crfm/text2image/text2image-rlhf/reasoning/virtual-world-data/cache", help="Data directory") + args = parser.parse_args() + + # Ensure mongod is available + mongod = shutil.which("mongod") + if mongod is None: + print("ERROR: mongod not found on PATH. See README.md for install instructions.", file=sys.stderr) + sys.exit(1) + + # Create dbpath if needed + Path(args.dbpath).mkdir(parents=True, exist_ok=True) + + hostname = socket.gethostname() + uri = f"mongodb://{hostname}:{args.port}" + print(f"mongo-uri: {uri}", flush=True) + + # Launch mongod in foreground (replaces this process on Unix) + cmd = [mongod, "--port", str(args.port), "--dbpath", args.dbpath, "--bind_ip_all"] + print(f"Running: {' '.join(cmd)}", flush=True) + proc = subprocess.run(cmd) + sys.exit(proc.returncode) + + +if __name__ == "__main__": + main() diff --git a/sdg/launcher.py b/sdg/launcher.py new file mode 100644 index 0000000000000000000000000000000000000000..52946010d69c8cad7c70a5465a690c9ca636efd5 --- /dev/null +++ b/sdg/launcher.py @@ -0,0 +1,575 @@ +""" +Orchestrator for distributed sharded SDG pipeline. + +Usage: + # Launch all shards + python -m sdg.launcher --config sdg/configs/my_experiment.yaml launch + + # Dry run (print commands without executing) + python -m sdg.launcher --config sdg/configs/my_experiment.yaml launch --dry-run + + # Check status of all shards + python -m sdg.launcher --config sdg/configs/my_experiment.yaml check + + # Rerun failed/not-started shards + python -m sdg.launcher --config sdg/configs/my_experiment.yaml rerun +""" + +from __future__ import annotations + +import argparse +import json +import os +import subprocess +import sys +from datetime import datetime, timezone +from pathlib import Path + +from pymongo import MongoClient + +from sdg.config import SDGConfig + +WORKDIR = "/nlp/scr4/nlp/crfm/text2image/text2image-rlhf/reasoning/virtual-world-data" + + +# ============================================================================= +# Helpers +# ============================================================================= + +def _verify_mongo(uri: str) -> None: + """Fail fast if MongoDB is unreachable.""" + try: + client = MongoClient(uri, serverSelectionTimeoutMS=5000) + client.admin.command("ping") + print(f"MongoDB reachable at {uri}") + except Exception as e: + print(f"ERROR: Cannot reach MongoDB at {uri}: {e}", file=sys.stderr) + sys.exit(1) + + +def _shard_status(run_dir: Path, shard_id: int, run_name: str, num_shards: int) -> str: + """Determine status of a single shard: completed / running / failed / not_started.""" + shard_dir = run_dir / "shards" / f"shard_{shard_id:03d}" + stats_file = shard_dir / "stats.json" + log_file = run_dir / "logs" / f"shard_{shard_id:03d}.log" + + if stats_file.exists(): + return "completed" + + # Check if job is in the queue via squeue + job_name = f"{run_name}-shard-{shard_id:03d}" + try: + result = subprocess.run( + ["squeue", "-u", os.environ.get("USER", ""), "--name", job_name, + "--noheader", "-o", "%.18i %.100j %.8T"], + capture_output=True, text=True, timeout=10, + ) + line = result.stdout.strip() + if line: + state = line.split()[-1] # RUNNING, PENDING, etc. + return state.lower() + except (FileNotFoundError, subprocess.TimeoutExpired): + pass # squeue not available or timed out + + if log_file.exists(): + return "failed" + + return "not_started" + + +def _resolve_run_dir(args: argparse.Namespace) -> Path: + """Resolve run directory from --config (preferred) or --run-dir.""" + if args.config: + config = SDGConfig.from_yaml(args.config) + return config.output_path + return Path(args.run_dir) + + +def _parse_log_progress(log_file: Path) -> str: + """Parse a shard log file and return a one-line progress summary.""" + if not log_file.exists(): + return "" + lines = log_file.read_text().splitlines() + + stage = "starting" + detail = "" + passed_total = 0 + + for line in lines: + stripped = line.strip() + if "Loading model:" in stripped: + stage = "loading model" + elif "Capturing CUDA graphs" in stripped: + stage = "loading model (CUDA graphs)" + elif "Model loaded successfully" in stripped: + stage = "model loaded" + elif "Loading dataset:" in stripped: + stage = "loading dataset" + elif "Processed prompts:" in stripped: + stage = "generating (vLLM)" + detail = stripped + elif "[Generation]" in stripped: + stage = "generation" + detail = stripped + elif stripped.startswith("Round ") and "pending rows" in stripped: + stage = "validation" + detail = stripped + elif "newly passed" in stripped: + try: + parts = stripped.split() + idx = parts.index("newly") + passed_total += int(parts[idx - 1]) + except (ValueError, IndexError): + pass + elif "PIPELINE COMPLETE" in stripped: + stage = "done" + + if stage == "generation": + return detail + elif stage == "generating (vLLM)": + return detail + elif stage == "validation": + return f"{detail} | {passed_total} passed so far" + elif stage == "done": + return "pipeline complete" + return stage + + +def _load_shard_stats(run_dir: Path, shard_id: int) -> dict | None: + """Load stats.json for a completed shard.""" + stats_file = run_dir / "shards" / f"shard_{shard_id:03d}" / "stats.json" + if stats_file.exists(): + with open(stats_file) as f: + return json.load(f) + return None + + +# ============================================================================= +# --launch +# ============================================================================= + +def cmd_launch(args: argparse.Namespace) -> None: + config = SDGConfig.from_yaml(args.config) + num_shards = args.num_shards if args.num_shards else config.num_shards + run_name = config.run_name + + if not config.mongo_uri: + print("ERROR: mongo_uri not set in config. Set it in your YAML.", file=sys.stderr) + sys.exit(1) + + if not args.dry_run: + _verify_mongo(config.mongo_uri) + + run_dir = config.output_path + run_dir.mkdir(parents=True, exist_ok=True) + log_dir = run_dir / "logs" + log_dir.mkdir(parents=True, exist_ok=True) + + shard_commands = [] + nlprun_commands = [] + + exclude_flag = "" + if config.machines_to_exclude: + exclude_flag = f" --exclude {','.join(config.machines_to_exclude)}" + + for i in range(num_shards): + cmd = f"python -m sdg.generate --config {args.config} --shard-id {i} --num-shards {num_shards}" + job_name = f"{run_name}-shard-{i:03d}" + nlprun_cmd = ( + f"nlprun -a tonyreasoningtraces -q sphinx{exclude_flag} -c 4 -g 1 --memory 100g " + f"-w {WORKDIR} --job-name {job_name} \"{cmd}\"" + ) + shard_commands.append(cmd) + nlprun_commands.append(nlprun_cmd) + + # Save launch metadata + meta = { + "run_name": run_name, + "config_path": str(args.config), + "num_shards": num_shards, + "mongo_uri": config.mongo_uri, + "timestamp": datetime.now(timezone.utc).isoformat(), + "shard_commands": shard_commands, + "nlprun_commands": nlprun_commands, + } + meta_path = run_dir / "launch_meta.json" + with open(meta_path, "w") as f: + json.dump(meta, f, indent=2) + print(f"Wrote {meta_path}") + + # Copy config + config.to_yaml(run_dir / "config.yaml") + + # Execute or print + for i, nlprun_cmd in enumerate(nlprun_commands): + if args.dry_run: + print(f"[DRY RUN] {nlprun_cmd}") + else: + print(f"Launching shard {i:03d}: {nlprun_cmd}") + subprocess.run(nlprun_cmd, shell=True, check=True) + + print(f"\nLaunched {num_shards} shards for run: {run_name}") + + +# ============================================================================= +# --check +# ============================================================================= + +def cmd_check(args: argparse.Namespace) -> None: + run_dir = _resolve_run_dir(args) + meta_path = run_dir / "launch_meta.json" + if not meta_path.exists(): + print(f"ERROR: {meta_path} not found", file=sys.stderr) + sys.exit(1) + + with open(meta_path) as f: + meta = json.load(f) + + run_name = meta["run_name"] + num_shards = meta["num_shards"] + + # Collect all statuses first + statuses = {} + for i in range(num_shards): + statuses[i] = _shard_status(run_dir, i, run_name, num_shards) + + # Show tail of logs for running shards first + running_shards = [i for i, s in statuses.items() if s == "running"] + if running_shards: + print(f"{'='*60}") + print(f"Logs for {len(running_shards)} running shard(s)") + print(f"{'='*60}") + for i in running_shards: + log_file = run_dir / "logs" / f"shard_{i:03d}.log" + print(f"\n{'─'*60}") + print(f" shard {i:03d} ({log_file})") + print(f"{'─'*60}") + if log_file.exists(): + lines = log_file.read_text().splitlines() + tail = lines[-20:] if len(lines) > 20 else lines + for line in tail: + print(f" {line}") + else: + print(" (no log file yet)") + print() + + # Status table + print(f"Run: {run_name} ({num_shards} shards)") + print(f"{'Shard':<6} {'Status':<12} {'Seeds':>7} {'Passed':>8} {'Failed':>8} {'Progress'}") + print("-" * 90) + + for i in range(num_shards): + status = statuses[i] + stats = _load_shard_stats(run_dir, i) if status == "completed" else None + seeds = str(stats["summary"]["total_examples"]) if stats else "-" + passed = str(stats["summary"]["total_passed"]) if stats else "-" + failed = str(stats["summary"]["total_failed"]) if stats else "-" + progress = "" + if status == "running": + log_file = run_dir / "logs" / f"shard_{i:03d}.log" + progress = _parse_log_progress(log_file) + print(f"{i:03d} {status:<12} {seeds:>7} {passed:>8} {failed:>8} {progress}") + + # Summary + counts = {} + for s in statuses.values(): + counts[s] = counts.get(s, 0) + 1 + + parts = [] + for label in ["completed", "running", "pending", "failed", "not_started"]: + if label in counts: + parts.append(f"{counts[label]}/{num_shards} {label}") + print(f"\nOverall: {', '.join(parts)}") + + # If all completed, combine outputs and prompt for upload + force_upload = getattr(args, "force_upload", False) + completed_count = counts.get("completed", 0) + failed_count = counts.get("failed", 0) + terminal_count = completed_count + failed_count + + if completed_count == num_shards: + print(f"\n{'='*60}") + print("All shards completed!") + print(f"{'='*60}") + _combine_outputs(run_dir, num_shards, meta) + elif force_upload and terminal_count == num_shards and completed_count > 0: + print(f"\n{'='*60}") + print(f"--force-upload: {completed_count}/{num_shards} completed, " + f"{failed_count}/{num_shards} failed. Uploading completed shards only.") + print(f"{'='*60}") + _combine_outputs(run_dir, num_shards, meta) + elif force_upload: + print(f"\n--force-upload requested but {num_shards - terminal_count} shard(s) " + f"are still running/pending. Skipping upload.") + + +def _aggregate_stats(run_dir: Path, num_shards: int) -> dict: + """Aggregate stats from all shards into a combined dict.""" + combined_stats = { + "summary": {"total_examples": 0, "total_passed": 0, "total_passed_samples": 0, "total_failed": 0, "pass_rate_pct": 0.0}, + "checks": { + "static_check": {"passed": 0, "failed": 0}, + "cycle_consistency": {"passed": 0, "failed": 0}, + "factual_accuracy": {"passed": 0, "failed": 0}, + "correctness": {"passed": 0, "failed": 0}, + }, + "failure_breakdown": { + "failed_static_check": 0, + "failed_cycle_consistency": 0, + "failed_factual_accuracy": 0, + "failed_correctness": 0, + "exhausted_samples": 0, + "total_failed": 0, + }, + "failure_breakdown_examples": { + "failed_static_check": 0, + "failed_cycle_consistency": 0, + "failed_factual_accuracy": 0, + "failed_correctness": 0, + "total_failed": 0, + }, + "num_shards": num_shards, + } + + for i in range(num_shards): + stats = _load_shard_stats(run_dir, i) + if stats is None: + continue + for key in ["total_examples", "total_passed", "total_passed_samples", "total_failed"]: + combined_stats["summary"][key] += stats["summary"].get(key, 0) + for check_name in combined_stats["checks"]: + if check_name in stats.get("checks", {}): + for k in ["passed", "failed"]: + combined_stats["checks"][check_name][k] += stats["checks"][check_name][k] + for key in combined_stats["failure_breakdown"]: + if key in stats.get("failure_breakdown", {}): + combined_stats["failure_breakdown"][key] += stats["failure_breakdown"][key] + for key in combined_stats["failure_breakdown_examples"]: + if key in stats.get("failure_breakdown_examples", {}): + combined_stats["failure_breakdown_examples"][key] += stats["failure_breakdown_examples"][key] + + total_ex = combined_stats["summary"]["total_examples"] + total_pass = combined_stats["summary"]["total_passed"] + combined_stats["summary"]["pass_rate_pct"] = round(total_pass / total_ex * 100, 2) if total_ex else 0.0 + + return combined_stats + + +def _combine_outputs(run_dir: Path, num_shards: int, meta: dict) -> None: + """Combine shard outputs into a single output.jsonl and stats.json, then prompt for HF upload.""" + # Combine JSONL + combined_jsonl = run_dir / "output.jsonl" + total_lines = 0 + with open(combined_jsonl, "w") as out_f: + for i in range(num_shards): + shard_jsonl = run_dir / "shards" / f"shard_{i:03d}" / "output.jsonl" + if shard_jsonl.exists(): + with open(shard_jsonl) as in_f: + for line in in_f: + out_f.write(line) + total_lines += 1 + print(f"Combined {total_lines} rows into {combined_jsonl}") + + # Aggregate and write stats + combined_stats = _aggregate_stats(run_dir, num_shards) + + # Sanity check: JSONL rows must match expected output count from stats. + # With all_valid selection policy, total_passed_samples (individual passing + # samples) is the correct count; total_passed counts seeds, not rows. + expected_rows = combined_stats["summary"].get("total_passed_samples") or combined_stats["summary"]["total_passed"] + if total_lines != expected_rows: + raise RuntimeError( + f"INTEGRITY ERROR: output.jsonl has {total_lines} rows but " + f"stats report {expected_rows} expected output rows. " + f"Shard outputs may be corrupted or out of sync with stats." + ) + + stats_path = run_dir / "stats.json" + with open(stats_path, "w") as f: + json.dump(combined_stats, f, indent=2) + print(f"Aggregated stats to {stats_path}") + + # Print summary for review + s = combined_stats["summary"] + checks = combined_stats["checks"] + fb = combined_stats["failure_breakdown"] + fbe = combined_stats["failure_breakdown_examples"] + + print(f"\n{'='*60}") + print("COMBINED STATISTICS") + print(f"{'='*60}") + print(f" Total examples: {s['total_examples']:,}") + print(f" Total passed: {s['total_passed']:,} (seeds)") + if s.get("total_passed_samples"): + print(f" Passed samples: {s['total_passed_samples']:,} (individual)") + print(f" Total failed: {s['total_failed']:,}") + print(f" Pass rate: {s['pass_rate_pct']}%") + print(f"\n --- Checks (sample-level, summed across all rounds) ---") + print(f" Static check: {checks['static_check']['passed']:,} passed, {checks['static_check']['failed']:,} failed") + print(f" Cycle consistency: {checks['cycle_consistency']['passed']:,} passed, {checks['cycle_consistency']['failed']:,} failed") + print(f" Factual accuracy: {checks['factual_accuracy']['passed']:,} passed, {checks['factual_accuracy']['failed']:,} failed") + print(f" Correctness: {checks['correctness']['passed']:,} passed, {checks['correctness']['failed']:,} failed") + print(f"\n --- Failure Breakdown (sample-level) ---") + print(f" Failed static: {fb['failed_static_check']:,}") + print(f" Failed cycle: {fb['failed_cycle_consistency']:,}") + print(f" Failed factual: {fb['failed_factual_accuracy']:,}") + print(f" Failed correctness:{fb['failed_correctness']:,}") + print(f" Exhausted samples: {fb['exhausted_samples']:,}") + print(f"\n --- Failure Breakdown (example-level, by deepest check reached) ---") + print(f" Failed static: {fbe['failed_static_check']:,}") + print(f" Failed cycle: {fbe['failed_cycle_consistency']:,}") + print(f" Failed factual: {fbe['failed_factual_accuracy']:,}") + print(f" Failed correctness:{fbe['failed_correctness']:,}") + print(f" Total failed: {fbe['total_failed']:,}") + print(f"\n --- Output ---") + print(f" output.jsonl: {total_lines:,} rows") + print(f" Path: {combined_jsonl}") + print(f"{'='*60}") + + # Prompt for HF upload + repo_id = f"teetone/{meta['run_name']}" + print(f"\nUpload to HuggingFace as: https://huggingface.co/datasets/{repo_id} ?") + answer = input("[y/N] ").strip().lower() + if answer != "y": + print("Skipping upload.") + return + + try: + import tempfile + from huggingface_hub import HfApi + api = HfApi() + api.create_repo(repo_id, repo_type="dataset", exist_ok=True) + + # Upload a README with dataset card so HF only loads output.jsonl as data + readme_content = ( + "---\n" + "configs:\n" + "- config_name: default\n" + " data_files:\n" + " - split: train\n" + " path: output.jsonl\n" + "---\n" + ) + with tempfile.NamedTemporaryFile(mode="w", suffix=".md", delete=False) as tmp: + tmp.write(readme_content) + tmp_path = tmp.name + api.upload_file( + path_or_fileobj=tmp_path, + path_in_repo="README.md", + repo_id=repo_id, + repo_type="dataset", + ) + os.unlink(tmp_path) + + api.upload_file( + path_or_fileobj=str(combined_jsonl), + path_in_repo="output.jsonl", + repo_id=repo_id, + repo_type="dataset", + ) + api.upload_file( + path_or_fileobj=str(stats_path), + path_in_repo="stats.json", + repo_id=repo_id, + repo_type="dataset", + ) + print(f"Uploaded to HuggingFace: https://huggingface.co/datasets/{repo_id}") + except Exception as e: + print(f"WARNING: HuggingFace upload failed: {e}") + + +# ============================================================================= +# --rerun +# ============================================================================= + +def cmd_rerun(args: argparse.Namespace) -> None: + config = SDGConfig.from_yaml(args.config) + run_dir = config.output_path + meta_path = run_dir / "launch_meta.json" + if not meta_path.exists(): + print(f"ERROR: {meta_path} not found", file=sys.stderr) + sys.exit(1) + + with open(meta_path) as f: + meta = json.load(f) + + run_name = meta["run_name"] + num_shards = meta["num_shards"] + + exclude_flag = "" + if config.machines_to_exclude: + exclude_flag = f" --exclude {','.join(config.machines_to_exclude)}" + + relaunch = [] + for i in range(num_shards): + status = _shard_status(run_dir, i, run_name, num_shards) + if status in ("failed", "not_started"): + relaunch.append(i) + + if not relaunch: + print("No failed or not_started shards to rerun.") + return + + if args.dry_run: + print(f"[DRY RUN] Would relaunch {len(relaunch)} failed shards: {relaunch}") + for i in relaunch: + cmd = f"python -m sdg.generate --config {args.config} --shard-id {i} --num-shards {num_shards}" + job_name = f"{run_name}-shard-{i:03d}" + nlprun_cmd = ( + f"nlprun -a tonyreasoningtraces -q sphinx{exclude_flag} -c 4 -g 1 --memory 100g " + f"-w {WORKDIR} --job-name {job_name} \"{cmd}\"" + ) + print(f" shard {i:03d}: {nlprun_cmd}") + return + + print(f"Relaunching {len(relaunch)} failed shards: {relaunch}") + for i in relaunch: + cmd = f"python -m sdg.generate --config {args.config} --shard-id {i} --num-shards {num_shards}" + job_name = f"{run_name}-shard-{i:03d}" + nlprun_cmd = ( + f"nlprun -a tonyreasoningtraces -q sphinx{exclude_flag} -c 4 -g 1 --memory 100g " + f"-w {WORKDIR} --job-name {job_name} \"{cmd}\"" + ) + print(f"Relaunching shard {i:03d}: {nlprun_cmd}") + subprocess.run(nlprun_cmd, shell=True, check=True) + + print(f"\nRelaunched {len(relaunch)} shards") + + +# ============================================================================= +# CLI +# ============================================================================= + +def main() -> None: + parser = argparse.ArgumentParser(description="SDG distributed pipeline orchestrator") + parser.add_argument("--config", required=True, help="Path to YAML config") + parser.add_argument("--run-dir", help="Run directory (alternative to --config for check/rerun)") + + subparsers = parser.add_subparsers(dest="command", required=True) + + launch_parser = subparsers.add_parser("launch", help="Launch all shards") + launch_parser.add_argument("--num-shards", type=int, help="Override num_shards from config") + launch_parser.add_argument("--dry-run", action="store_true", help="Print commands without executing") + + check_parser = subparsers.add_parser("check", help="Check shard status") + check_parser.add_argument( + "--force-upload", + action="store_true", + help="Upload to HuggingFace as long as every shard is either completed or failed " + "(no running/pending). Only completed shards' data is uploaded.", + ) + rerun_parser = subparsers.add_parser("rerun", help="Rerun failed shards") + rerun_parser.add_argument("--dry-run", action="store_true", help="Print commands without executing") + + args = parser.parse_args() + + if args.command == "launch": + cmd_launch(args) + elif args.command == "check": + cmd_check(args) + elif args.command == "rerun": + cmd_rerun(args) + + +if __name__ == "__main__": + main() diff --git a/sdg/mongo_example.py b/sdg/mongo_example.py new file mode 100644 index 0000000000000000000000000000000000000000..5c8863c4c9ca9510b511f3536d3aefd82e2b8b0c --- /dev/null +++ b/sdg/mongo_example.py @@ -0,0 +1,49 @@ +""" +Quick smoke-test: write a document to MongoDB and read it back. + +Usage: + python -m sdg.test_mongo # localhost:27017 + python -m sdg.test_mongo --mongo-uri mongodb://HOST:PORT +""" + +import argparse +import sys +from pymongo import MongoClient + + +def main() -> None: + parser = argparse.ArgumentParser(description="Test MongoDB read/write") + parser.add_argument("--mongo-uri", default="mongodb://john5.stanford.edu:27017", + help="MongoDB connection URI") + args = parser.parse_args() + + print(f"Connecting to {args.mongo_uri} ...") + client = MongoClient(args.mongo_uri, serverSelectionTimeoutMS=5000) + + # Verify connectivity + client.admin.command("ping") + print("Ping OK") + + db = client["sdg_test"] + col = db["smoke_test"] + + # Write + doc = {"key": "hello", "value": "world"} + col.replace_one({"key": "hello"}, doc, upsert=True) + print(f"Wrote: {doc}") + + # Read + result = col.find_one({"key": "hello"}) + assert result is not None, "Read back None — write may have failed" + print(f"Read: key={result['key']} value={result['value']}") + + # Cleanup + col.delete_one({"key": "hello"}) + print("Cleaned up test document") + + client.close() + print("All good!") + + +if __name__ == "__main__": + main() diff --git a/sdg/pipeline.py b/sdg/pipeline.py new file mode 100644 index 0000000000000000000000000000000000000000..d2caf623956c1dd2e9a04507c3793d615cfe771c --- /dev/null +++ b/sdg/pipeline.py @@ -0,0 +1,771 @@ +""" +Main SDG pipeline orchestration. + +Stages: + 1. Load HF dataset → extract instruction_seed + 2. Init VLLMEngine + 3. Generate N candidates per prompt, static-check, collect valid samples + 4. Iterative UQ (max_rounds): cycle → factual → correctness, early-exit + 5. Write output JSONL + stats.json + config.yaml +""" + +from __future__ import annotations + +import json +import math +import sys +from dataclasses import dataclass, field, asdict +from pathlib import Path + +from datasets import load_dataset +from tqdm import tqdm + +from sdg.config import SDGConfig +from sdg.inference import VLLMEngine +from sdg.prompts import REASONING_INSTRUCTION +from sdg.validation import ( + collect_valid_samples, + validate_batch_correctness, + validate_batch_cycle, + validate_batch_factual, + validate_batch_simple, +) + +# Valid values for SDGConfig.selection_policy +VALID_SELECTION_POLICIES = {"first_valid", "all_valid"} + + +# ============================================================================= +# TeeStream — duplicate stdout to a log file +# ============================================================================= + +class TeeStream: + """Write to both the original stream and a log file.""" + + def __init__(self, log_path: Path, original_stream, share_file=None): + self.log_file = share_file if share_file is not None else open(log_path, "w") + self._owns_file = share_file is None + self.original = original_stream + + def write(self, data): + self.original.write(data) + self.log_file.write(data) + self.log_file.flush() + + def flush(self): + self.original.flush() + self.log_file.flush() + + def fileno(self): + return self.original.fileno() + + def isatty(self): + return self.original.isatty() + + def close(self): + if self._owns_file: + self.log_file.close() + + +# ============================================================================= +# Row state tracker +# ============================================================================= + +@dataclass +class RowState: + row_id: int + instruction_seed: str + valid_samples: list[str] = field(default_factory=list) + valid_indices: list[int] = field(default_factory=list) + passed_sample: str | None = None # first_valid: single sample that passed all checks + passed_samples: list[str] = field(default_factory=list) # all_valid: all samples that passed + passed_round: int | None = None + best_stage: int = 0 # deepest check reached: 0=none, 1=cycle, 2=factual, 3=correctness + + +# ============================================================================= +# Dataset loading +# ============================================================================= + +def load_seeds(config: SDGConfig) -> list[str]: + """Load instruction seeds from HuggingFace dataset.""" + print(f"Loading dataset: {config.dataset_hf_id} (config={config.dataset_hf_config}, split={config.dataset_split}, revision={config.dataset_hf_revision})") + ds = load_dataset(config.dataset_hf_id, name=config.dataset_hf_config, split=config.dataset_split, revision=config.dataset_hf_revision) + + if config.limit: + ds = ds.select(range(min(config.limit, len(ds)))) + print(f" Limited to {len(ds)} rows") + + seeds: list[str] = [] + for row in ds: + if config.dataset_format == "openthoughts4": + convos = row.get("conversations", []) + if convos: + seeds.append(convos[0]["value"]) + elif config.dataset_format == "messages": + msgs = row.get("messages", []) + if msgs: + seeds.append(msgs[0]["content"]) + else: + raise ValueError(f"Unknown dataset_format: {config.dataset_format}") + + print(f" Extracted {len(seeds)} instruction seeds") + return seeds + + +# ============================================================================= +# Stage 3: Generation + static check +# ============================================================================= + +def generate_candidates( + engine: VLLMEngine, + seeds: list[str], + config: SDGConfig, +) -> list[RowState]: + """Generate N candidates per seed, run static check, build RowState list.""" + rows: list[RowState] = [] + total = len(seeds) + processed = 0 + + for batch_start in range(0, total, config.gen_batch_size): + batch_seeds = seeds[batch_start : batch_start + config.gen_batch_size] + if config.check_boxed: + prompts = [f"{seed}\n{REASONING_INSTRUCTION}\n" for seed in batch_seeds] + else: + prompts = [f"{seed}\n" for seed in batch_seeds] + + # generate N samples per prompt + all_samples = engine.generate_multi(prompts) + + for i, (seed, samples) in enumerate(zip(batch_seeds, all_samples)): + valid, indices = collect_valid_samples( + samples, config.is_instruction_tuned, config.check_boxed + ) + if valid: + rows.append( + RowState( + row_id=batch_start + i, + instruction_seed=seed, + valid_samples=valid, + valid_indices=indices, + ) + ) + + processed += len(batch_seeds) + print(f"[Generation] {processed}/{total} examples ({processed/total*100:.1f}%) | {len(rows)} valid") + + print( + f" {len(rows)} / {len(seeds)} rows have at least one valid sample" + ) + return rows + + +# ============================================================================= +# Stage 4: Iterative UQ +# ============================================================================= + +@dataclass +class RoundStats: + round_idx: int = 0 + pending: int = 0 + exhausted: int = 0 + cycle_passed: int = 0 + cycle_failed: int = 0 + factual_passed: int = 0 + factual_failed: int = 0 + correctness_passed: int = 0 + correctness_failed: int = 0 + newly_passed: int = 0 + + +def _validate_round( + engine: VLLMEngine, + round_items: list[tuple[RowState, str]], + config: SDGConfig, + stats: RoundStats, +) -> tuple[list[tuple[RowState, str]], list[RowState], list[RowState], list[RowState]]: + """Run cycle → factual → correctness on a batch of (row, sample) pairs. + + Returns: + (correctness_passed_items, cycle_failed_rows, factual_failed_rows, correctness_failed_rows) + """ + # ── cycle consistency ──────────────────────────────────────── + cycle_items = [ + (row.row_id, row.instruction_seed, sample) + for row, sample in round_items + ] + cycle_results = validate_batch_cycle( + engine, cycle_items, config.val_batch_size + ) + + cycle_passed_items: list[tuple[RowState, str]] = [] + cycle_failed_rows: list[RowState] = [] + for row, sample in round_items: + if cycle_results.get(row.row_id, False): + cycle_passed_items.append((row, sample)) + stats.cycle_passed += 1 + row.best_stage = max(row.best_stage, 1) + else: + cycle_failed_rows.append(row) + stats.cycle_failed += 1 + + print( + f" Cycle: {stats.cycle_passed} passed, {stats.cycle_failed} failed" + ) + + # ── factual check (cycle-passed only) ──────────────────────── + factual_passed_items: list[tuple[RowState, str]] = [] + factual_failed_rows: list[RowState] = [] + if cycle_passed_items: + factual_items = [ + (row.row_id, row.instruction_seed, sample) + for row, sample in cycle_passed_items + ] + factual_results = validate_batch_factual( + engine, factual_items, config.val_batch_size + ) + for row, sample in cycle_passed_items: + if factual_results.get(row.row_id, False): + factual_passed_items.append((row, sample)) + stats.factual_passed += 1 + row.best_stage = max(row.best_stage, 2) + else: + factual_failed_rows.append(row) + stats.factual_failed += 1 + print( + f" Factual: {stats.factual_passed} passed, {stats.factual_failed} failed" + ) + + # ── correctness check (factual-passed only) ───────────────── + correctness_passed_items: list[tuple[RowState, str]] = [] + correctness_failed_rows: list[RowState] = [] + if factual_passed_items: + correctness_items = [ + (row.row_id, row.instruction_seed, sample) + for row, sample in factual_passed_items + ] + correctness_results = validate_batch_correctness( + engine, correctness_items, config.val_batch_size + ) + for row, sample in factual_passed_items: + if correctness_results.get(row.row_id, False): + correctness_passed_items.append((row, sample)) + stats.correctness_passed += 1 + row.best_stage = max(row.best_stage, 3) + else: + correctness_failed_rows.append(row) + stats.correctness_failed += 1 + print( + f" Correctness: {stats.correctness_passed} passed, " + f"{stats.correctness_failed} failed" + ) + + return correctness_passed_items, cycle_failed_rows, factual_failed_rows, correctness_failed_rows + + +def run_uq_rounds( + engine: VLLMEngine, + rows: list[RowState], + config: SDGConfig, +) -> tuple[list[RowState], list[RowState], list[RoundStats]]: + """ + Iterative UQ loop (first_valid policy). + + Tests one sample per row per round. Once a row passes all three checks + it is moved to passed_rows and not tested again. + + Returns: + (passed_rows, remaining_rows, per_round_stats) + """ + passed_rows: list[RowState] = [] + pending_rows: list[RowState] = rows[:] + exhausted_rows: list[RowState] = [] + all_stats: list[RoundStats] = [] + + for round_idx in range(config.num_generations): + if not pending_rows: + print(f"Round {round_idx}: no pending rows — done early") + break + + stats = RoundStats(round_idx=round_idx, pending=len(pending_rows)) + print(f"\n{'='*60}") + print(f"Round {round_idx}: {len(pending_rows)} pending rows") + print(f"{'='*60}") + + # ── extract sample at round_idx ────────────────────────────── + round_items: list[tuple[RowState, str]] = [] # (row, current_sample) + skipped: list[RowState] = [] + + for row in pending_rows: + if round_idx < len(row.valid_samples): + round_items.append((row, row.valid_samples[round_idx])) + else: + skipped.append(row) # exhausted samples + + if not round_items: + print(f" No rows have a sample at index {round_idx} — stopping") + break + + stats.exhausted = len(skipped) + exhausted_rows.extend(skipped) + print(f" {len(round_items)} rows have sample at index {round_idx}") + if skipped: + print(f" {len(skipped)} rows exhausted all samples (dropped)") + + # ── validate ───────────────────────────────────────────────── + correctness_passed, cycle_failed, factual_failed, correctness_failed = ( + _validate_round(engine, round_items, config, stats) + ) + + newly_passed: list[RowState] = [] + for row, sample in correctness_passed: + row.passed_sample = sample + row.passed_round = round_idx + newly_passed.append(row) + + stats.newly_passed = len(newly_passed) + passed_rows.extend(newly_passed) + + # ── remaining pending = all failures from this round ───────── + pending_rows = cycle_failed + factual_failed + correctness_failed + + print( + f" Round {round_idx} summary: {stats.newly_passed} newly passed, " + f"{len(pending_rows)} still pending" + ) + all_stats.append(stats) + + # remaining = still pending after last round + exhausted across all rounds + all_failed = pending_rows + exhausted_rows + return passed_rows, all_failed, all_stats + + +def _validate_round_simple( + engine: VLLMEngine, + round_items: list[tuple[RowState, str]], + config: SDGConfig, + stats: RoundStats, +) -> tuple[list[tuple[RowState, str]], list[RowState]]: + """Single-prompt yes/no validation. Reuses correctness_* stat fields.""" + items = [ + (row.row_id, row.instruction_seed, sample) + for row, sample in round_items + ] + results = validate_batch_simple(engine, items, config.val_batch_size) + + passed_items: list[tuple[RowState, str]] = [] + failed_rows: list[RowState] = [] + for row, sample in round_items: + if results.get(row.row_id, False): + passed_items.append((row, sample)) + stats.correctness_passed += 1 + row.best_stage = max(row.best_stage, 3) + else: + failed_rows.append(row) + stats.correctness_failed += 1 + # treat simple-fail as a correctness-stage failure for reporting + row.best_stage = max(row.best_stage, 2) + + print( + f" Simple: {stats.correctness_passed} passed, " + f"{stats.correctness_failed} failed" + ) + return passed_items, failed_rows + + +def run_simple_rounds( + engine: VLLMEngine, + rows: list[RowState], + config: SDGConfig, +) -> tuple[list[RowState], list[RowState], list[RoundStats]]: + """Iterative simple validation loop (first_valid policy).""" + passed_rows: list[RowState] = [] + pending_rows: list[RowState] = rows[:] + exhausted_rows: list[RowState] = [] + all_stats: list[RoundStats] = [] + + for round_idx in range(config.num_generations): + if not pending_rows: + print(f"Round {round_idx}: no pending rows — done early") + break + + stats = RoundStats(round_idx=round_idx, pending=len(pending_rows)) + print(f"\n{'='*60}") + print(f"Round {round_idx}: {len(pending_rows)} pending rows") + print(f"{'='*60}") + + round_items: list[tuple[RowState, str]] = [] + skipped: list[RowState] = [] + for row in pending_rows: + if round_idx < len(row.valid_samples): + round_items.append((row, row.valid_samples[round_idx])) + else: + skipped.append(row) + + if not round_items: + print(f" No rows have a sample at index {round_idx} — stopping") + break + + stats.exhausted = len(skipped) + exhausted_rows.extend(skipped) + print(f" {len(round_items)} rows have sample at index {round_idx}") + if skipped: + print(f" {len(skipped)} rows exhausted all samples (dropped)") + + passed_items, failed_rows = _validate_round_simple( + engine, round_items, config, stats + ) + + newly_passed: list[RowState] = [] + for row, sample in passed_items: + row.passed_sample = sample + row.passed_round = round_idx + newly_passed.append(row) + + stats.newly_passed = len(newly_passed) + passed_rows.extend(newly_passed) + pending_rows = failed_rows + + print( + f" Round {round_idx} summary: {stats.newly_passed} newly passed, " + f"{len(pending_rows)} still pending" + ) + all_stats.append(stats) + + all_failed = pending_rows + exhausted_rows + return passed_rows, all_failed, all_stats + + +def run_simple_rounds_all_valid( + engine: VLLMEngine, + rows: list[RowState], + config: SDGConfig, +) -> tuple[list[RowState], list[RowState], list[RoundStats]]: + """Iterative simple validation loop (all_valid policy).""" + all_stats: list[RoundStats] = [] + + for round_idx in range(config.num_generations): + round_items: list[tuple[RowState, str]] = [] + for row in rows: + if round_idx < len(row.valid_samples): + round_items.append((row, row.valid_samples[round_idx])) + + if not round_items: + print(f"Round {round_idx}: no rows have a sample at this index — stopping") + break + + stats = RoundStats(round_idx=round_idx, pending=len(round_items)) + print(f"\n{'='*60}") + print(f"Round {round_idx}: {len(round_items)} rows") + print(f"{'='*60}") + + passed_items, _ = _validate_round_simple( + engine, round_items, config, stats + ) + + for row, sample in passed_items: + row.passed_samples.append(sample) + if row.passed_round is None: + row.passed_round = round_idx + + stats.newly_passed = len(passed_items) + print( + f" Round {round_idx} summary: {stats.newly_passed} samples passed" + ) + all_stats.append(stats) + + passed_rows = [row for row in rows if row.passed_samples] + failed_rows = [row for row in rows if not row.passed_samples] + return passed_rows, failed_rows, all_stats + + +def run_uq_rounds_all_valid( + engine: VLLMEngine, + rows: list[RowState], + config: SDGConfig, +) -> tuple[list[RowState], list[RowState], list[RoundStats]]: + """ + Iterative UQ loop (all_valid policy). + + Tests every valid sample for every row. A row is never removed early; + all samples that pass all three checks are collected in row.passed_samples. + + Returns: + (passed_rows, failed_rows, per_round_stats) + """ + all_stats: list[RoundStats] = [] + + for round_idx in range(config.num_generations): + # ── collect rows that have a sample at this index ──────────── + round_items: list[tuple[RowState, str]] = [] + for row in rows: + if round_idx < len(row.valid_samples): + round_items.append((row, row.valid_samples[round_idx])) + + if not round_items: + print(f"Round {round_idx}: no rows have a sample at this index — stopping") + break + + stats = RoundStats(round_idx=round_idx, pending=len(round_items)) + print(f"\n{'='*60}") + print(f"Round {round_idx}: {len(round_items)} rows") + print(f"{'='*60}") + + # ── validate ───────────────────────────────────────────────── + correctness_passed, _, _, _ = _validate_round( + engine, round_items, config, stats, + ) + + for row, sample in correctness_passed: + row.passed_samples.append(sample) + if row.passed_round is None: + row.passed_round = round_idx + + stats.newly_passed = len(correctness_passed) + print( + f" Round {round_idx} summary: {stats.newly_passed} samples passed" + ) + all_stats.append(stats) + + passed_rows = [row for row in rows if row.passed_samples] + failed_rows = [row for row in rows if not row.passed_samples] + return passed_rows, failed_rows, all_stats + + +# ============================================================================= +# Stage 5: Write outputs +# ============================================================================= + +def _build_stats( + passed_rows: list[RowState], + remaining_rows: list[RowState], + round_stats: list[RoundStats], + total_seeds: int, + rows_after_static: int, + config: SDGConfig | None = None, +) -> dict: + """Build a well-organized stats dict for stats.json.""" + static_failed = total_seeds - rows_after_static + + # Aggregate check totals across all rounds + total_cycle_passed = sum(s.cycle_passed for s in round_stats) + total_cycle_failed = sum(s.cycle_failed for s in round_stats) + total_factual_passed = sum(s.factual_passed for s in round_stats) + total_factual_failed = sum(s.factual_failed for s in round_stats) + total_correctness_passed = sum(s.correctness_passed for s in round_stats) + total_correctness_failed = sum(s.correctness_failed for s in round_stats) + total_exhausted = sum(s.exhausted for s in round_stats) + + pass_rate = (len(passed_rows) / total_seeds * 100) if total_seeds else 0.0 + + total_passed_samples = sum(len(r.passed_samples) for r in passed_rows) + + # Example-level failure breakdown: categorize each failed example by the + # deepest validation stage it ever reached across all rounds/samples. + # best_stage 0 → never passed cycle (stuck at cycle) + # best_stage 1 → passed cycle but never factual + # best_stage 2 → passed factual but never correctness + ex_failed_cycle = sum(1 for r in remaining_rows if r.best_stage == 0) + ex_failed_factual = sum(1 for r in remaining_rows if r.best_stage == 1) + ex_failed_correctness = sum(1 for r in remaining_rows if r.best_stage == 2) + + return { + "summary": { + "total_examples": total_seeds, + "total_passed": len(passed_rows), + "total_passed_samples": total_passed_samples, + "total_failed": total_seeds - len(passed_rows), + "pass_rate_pct": round(pass_rate, 2), + }, + "checks": { + "static_check": { + "passed": rows_after_static, + "failed": static_failed, + }, + "cycle_consistency": { + "passed": total_cycle_passed, + "failed": total_cycle_failed, + }, + "factual_accuracy": { + "passed": total_factual_passed, + "failed": total_factual_failed, + }, + "correctness": { + "passed": total_correctness_passed, + "failed": total_correctness_failed, + }, + }, + "failure_breakdown": { + "failed_static_check": static_failed, + "failed_cycle_consistency": total_cycle_failed, + "failed_factual_accuracy": total_factual_failed, + "failed_correctness": total_correctness_failed, + "exhausted_samples": total_exhausted, + "total_failed": total_seeds - len(passed_rows), + }, + "failure_breakdown_examples": { + "failed_static_check": static_failed, + "failed_cycle_consistency": ex_failed_cycle, + "failed_factual_accuracy": ex_failed_factual, + "failed_correctness": ex_failed_correctness, + "total_failed": total_seeds - len(passed_rows), + }, + "rounds": [asdict(s) for s in round_stats], + } + + +def write_outputs( + passed_rows: list[RowState], + remaining_rows: list[RowState], + round_stats: list[RoundStats], + total_seeds: int, + rows_after_static: int, + config: SDGConfig, + output_path: Path | None = None, +) -> None: + out = output_path if output_path is not None else config.output_path + out.mkdir(parents=True, exist_ok=True) + + # ── output.jsonl ───────────────────────────────────────────────── + jsonl_path = out / "output.jsonl" + total_records = 0 + with open(jsonl_path, "w") as f: + for row in passed_rows: + if config.check_boxed: + user_content = row.instruction_seed.strip() + "\n" + REASONING_INSTRUCTION + "\n" + else: + user_content = row.instruction_seed.strip() + "\n" + if config.selection_policy == "all_valid" and row.passed_samples: + samples = row.passed_samples + else: + samples = [row.passed_sample] + for sample in samples: + record = { + "messages": [ + {"role": "user", "content": user_content}, + {"role": "assistant", "content": sample}, + ] + } + f.write(json.dumps(record, ensure_ascii=False) + "\n") + total_records += 1 + print(f"Wrote {total_records} records ({len(passed_rows)} examples) to {jsonl_path}") + + # ── stats.json ─────────────────────────────────────────────────── + stats_path = out / "stats.json" + stats_data = _build_stats( + passed_rows, remaining_rows, round_stats, + total_seeds, rows_after_static, config, + ) + with open(stats_path, "w") as f: + json.dump(stats_data, f, indent=2) + print(f"Wrote stats to {stats_path}") + + # ── config.yaml copy ───────────────────────────────────────────── + config.to_yaml(out / "config.yaml") + print(f"Wrote config to {out / 'config.yaml'}") + + +# ============================================================================= +# Public entry point +# ============================================================================= + +def run_pipeline( + config: SDGConfig, + shard_id: int | None = None, + num_shards: int | None = None, +) -> None: + """Run the full SDG pipeline end-to-end.""" + if config.selection_policy not in VALID_SELECTION_POLICIES: + raise ValueError( + f"Invalid selection_policy={config.selection_policy!r}, " + f"must be one of {VALID_SELECTION_POLICIES}" + ) + + is_sharded = shard_id is not None and num_shards is not None + + print(f"Run name: {config.run_name}") + print(f"Output: {config.output_path}") + if is_sharded: + print(f"Shard: {shard_id}/{num_shards}") + + # ── Per-shard logging via TeeStream ────────────────────────────── + tee_out = None + tee_err = None + if is_sharded: + log_dir = config.output_path / "logs" + log_dir.mkdir(parents=True, exist_ok=True) + log_path = log_dir / f"shard_{shard_id:03d}.log" + tee_out = TeeStream(log_path, sys.stdout) + tee_err = TeeStream(log_path, sys.stderr, share_file=tee_out.log_file) + sys.stdout = tee_out + sys.stderr = tee_err + print(f"Logging to {log_path}") + + try: + # Stage 1 — load data + seeds = load_seeds(config) + + # ── Shard slicing ──────────────────────────────────────────── + if is_sharded: + total = len(seeds) + chunk_size = math.ceil(total / num_shards) + start = shard_id * chunk_size + end = min(start + chunk_size, total) + seeds = seeds[start:end] + print(f"Shard {shard_id}: seeds [{start}:{end}] ({len(seeds)} examples)") + + # Stage 2 — init engine + engine = VLLMEngine(config) + + # Stage 3 — generate candidates + static check + rows = generate_candidates(engine, seeds, config) + + # Stage 4 — iterative UQ (skip if validation_type is None) + rows_after_static = len(rows) + if config.validation_type is None: + print(f"\nSkipping validation (validation_type=null)") + for row in rows: + if config.selection_policy == "all_valid": + row.passed_samples = list(row.valid_samples) + else: + row.passed_sample = row.valid_samples[0] + row.passed_round = 0 + row.best_stage = 3 + passed = rows + remaining = [] + stats = [] + elif config.validation_type == "simple": + if config.selection_policy == "all_valid": + passed, remaining, stats = run_simple_rounds_all_valid(engine, rows, config) + else: + passed, remaining, stats = run_simple_rounds(engine, rows, config) + elif config.selection_policy == "all_valid": + passed, remaining, stats = run_uq_rounds_all_valid(engine, rows, config) + else: + passed, remaining, stats = run_uq_rounds(engine, rows, config) + + # Stage 5 — write output + total_seeds = len(seeds) + output_path = None + if is_sharded: + output_path = config.output_path / "shards" / f"shard_{shard_id:03d}" + write_outputs(passed, remaining, stats, total_seeds, rows_after_static, config, output_path=output_path) + + # Final summary + print(f"\n{'='*60}") + print("PIPELINE COMPLETE") + print(f"{'='*60}") + print(f"Total seeds: {total_seeds}") + print(f"Total passed: {len(passed)}") + print(f"Total dropped: {total_seeds - len(passed)}") + if total_seeds: + print(f"Pass rate: {len(passed) / total_seeds * 100:.1f}%") + if output_path: + print(f"Output dir: {output_path}") + else: + print(f"Output dir: {config.output_path}") + finally: + if tee_out is not None: + sys.stdout = tee_out.original + sys.stderr = tee_err.original + tee_out.close() diff --git a/sdg/preprocessing/__init__.py b/sdg/preprocessing/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/sdg/preprocessing/dedupe/__init__.py b/sdg/preprocessing/dedupe/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..cb33ccfe8e608e38416d823c83d72436b42590c9 --- /dev/null +++ b/sdg/preprocessing/dedupe/__init__.py @@ -0,0 +1,29 @@ +"""Deduplication package: MinHash near-dup + semantic paraphrase dedup. + +Standard SFT dedup pipeline order: + + exact (caller's responsibility) + | + v + MinHashDeduplicator - lexical near-dups (formatting variants, edits) + | + v + SemanticDeduplicator - paraphrases (different wording, same meaning) + +Both classes operate on a list[str] and return list[int] of indices to KEEP, +so callers re-index back into their own record list. An optional +`key_fn(idx) -> sortable` controls which member of a duplicate cluster wins +(smaller key wins). Default: smallest original index ("first seen"). +""" + +from .minhash import MinHashDeduplicator +from .semantic import SemanticDeduplicator +from .templates import find_common_prefixes, strip_template, template_hit_stats + +__all__ = [ + "MinHashDeduplicator", + "SemanticDeduplicator", + "find_common_prefixes", + "strip_template", + "template_hit_stats", +] diff --git a/sdg/preprocessing/dedupe/base.py b/sdg/preprocessing/dedupe/base.py new file mode 100644 index 0000000000000000000000000000000000000000..2f089e9b7d0de9a9b65abd88d121fae14091d957 --- /dev/null +++ b/sdg/preprocessing/dedupe/base.py @@ -0,0 +1,87 @@ +"""Shared primitives: union-find, cluster -> representative selection, normalize.""" + +from __future__ import annotations + +import re +from collections import defaultdict +from typing import Any, Callable, Iterable + + +class UnionFind: + """Union-find with path compression and union-by-rank.""" + + def __init__(self, n: int): + self.parent = list(range(n)) + self.rank = [0] * n + + def find(self, x: int) -> int: + root = x + while self.parent[root] != root: + root = self.parent[root] + while self.parent[x] != root: + self.parent[x], x = root, self.parent[x] + return root + + def union(self, x: int, y: int) -> None: + rx, ry = self.find(x), self.find(y) + if rx == ry: + return + if self.rank[rx] < self.rank[ry]: + rx, ry = ry, rx + self.parent[ry] = rx + if self.rank[rx] == self.rank[ry]: + self.rank[rx] += 1 + + +def cluster_and_pick( + n: int, + pairs: Iterable[tuple[int, int]], + key_fn: Callable[[int], Any] | None = None, +) -> tuple[list[int], dict[int, list[int]]]: + """Cluster items in [0, n) via `pairs`, pick one rep per cluster. + + Returns (keep, clusters): + keep: sorted list of representative indices (one per cluster). + clusters: dict {rep_idx -> sorted list of ALL member indices in that + cluster, including the rep}. Singletons appear as {i: [i]}. + """ + uf = UnionFind(n) + for i, j in pairs: + uf.union(i, j) + + members_by_root: dict[int, list[int]] = defaultdict(list) + for i in range(n): + members_by_root[uf.find(i)].append(i) + + clusters: dict[int, list[int]] = {} + for members in members_by_root.values(): + rep = min(members) if key_fn is None else min(members, key=key_fn) + clusters[rep] = sorted(members) + + keep = sorted(clusters.keys()) + return keep, clusters + + +def pick_representatives( + n: int, + pairs: Iterable[tuple[int, int]], + key_fn: Callable[[int], Any] | None = None, +) -> list[int]: + """Cluster items in [0, n) via `pairs`, then keep one rep per cluster. + + Representative is the cluster member with the smallest `key_fn(i)`. With + no key_fn, falls back to smallest index ("first seen"). Returned indices + are sorted ascending so caller can re-index in original order. + + Thin wrapper over `cluster_and_pick` that drops the cluster-membership map. + """ + keep, _ = cluster_and_pick(n, pairs, key_fn) + return keep + + +_WHITESPACE_RE = re.compile(r"\s+") + + +def normalize_text(text: str) -> str: + """Lowercase + collapse internal whitespace + strip.""" + return _WHITESPACE_RE.sub(" ", text.lower()).strip() diff --git a/sdg/preprocessing/dedupe/fixtures.py b/sdg/preprocessing/dedupe/fixtures.py new file mode 100644 index 0000000000000000000000000000000000000000..43bbea54514efa0931a04a328c4349dd8b2c29c1 --- /dev/null +++ b/sdg/preprocessing/dedupe/fixtures.py @@ -0,0 +1,97 @@ +"""Realistic science-prompt fixtures used by test_invariants.py. + +The list `REALISTIC_SCIENCE_PROMPTS` is hand-crafted to contain four classes +of items so tests can pin exact post-stage counts: + + Class A: byte-identical duplicates (3 copies) + -> caught by Stage 1 (caller's exact-prompt reservoir dedup) + + Class B: formatting variants of one prompt (3 prompts: extra whitespace, + upper-case differences, trailing newline) + -> NOT caught by exact dedup; caught by Stage 2 (MinHash with + normalize=True) + + Class C: near-duplicate paraphrases of one prompt (3 prompts: each is the + same long base text with one single-word edit) + -> NOT caught by exact or formatting normalization; caught by Stage 2 + (MinHash with threshold=0.8) + + Class D: 7 distinct singleton prompts (math / biology / chemistry / physics) + -> none of the dedup stages should touch these + +Total = 3 + 3 + 3 + 7 = 16 prompts. +After exact dedup: 1 + 3 + 3 + 7 = 14 +After MinHash: 1 + 1 + 1 + 7 = 10 + +Class C uses a ~140-word base text so that single-word edits keep Jaccard +robustly above MinHash's threshold of 0.8. With n=140 tokens (≈136 5-grams) +and a single-word edit affecting 5 5-grams, Jaccard = 131/141 ≈ 0.929 — well +into LSH's high-detection regime, so the test is deterministic. +""" + +from __future__ import annotations + + +# ── Class A: byte-identical ───────────────────────────────────────────────── +_EXACT_DUP = ( + "Solve the equation x squared minus four equals zero. Show all steps " + "clearly." +) + +# ── Class B: formatting variants of one prompt ────────────────────────────── +_FMT_BASE = "Calculate the derivative of f(x) = x^3 + 2x with respect to x." +_FORMATTING_VARIANTS = [ + _FMT_BASE, + " Calculate the derivative of f(x) = x^3 + 2x with respect to x. ", + "Calculate the DERIVATIVE of f(x) = x^3 + 2x with respect to x.\n", +] + +# ── Class C: near-duplicate paraphrases (single-word edits over ~80 words) ── +_NEAR_DUP_BASE = ( + "Explain in detail the process of photosynthesis in green plants. " + "Describe step by step how light energy is captured by chlorophyll " + "molecules in the thylakoid membrane, how electrons are transported " + "through photosystems II and I, how ATP and NADPH are generated through " + "photophosphorylation, and how the Calvin cycle uses these energy " + "carriers to fix carbon dioxide into glucose. Include the role of stomata " + "in gas exchange and discuss how environmental factors like light " + "intensity, water availability, soil nutrients, atmospheric carbon " + "dioxide concentration, and temperature affect the overall rate of " + "photosynthesis. Discuss the differences between C3, C4, and CAM " + "photosynthetic pathways and their evolutionary adaptations to different " + "environmental conditions including drought stress and high light " + "intensity. Finally explain how mitochondrial respiration interacts with " + "chloroplast metabolism in the cells of plants and how this energy " + "balance is maintained throughout the day and night cycle of plant " + "tissues across different seasons in temperate and tropical biomes." +) +_NEAR_DUP_VARIANTS = [ + _NEAR_DUP_BASE, + _NEAR_DUP_BASE.replace("process", "mechanism", 1), # 1-word edit near start + _NEAR_DUP_BASE.replace("glucose", "sugar", 1), # 1-word edit in middle +] + +# ── Class D: distinct singletons ──────────────────────────────────────────── +_SINGLETONS = [ + "Compute the integral of sin(x) cos(x) from 0 to pi/2 using substitution.", + "Describe the structure of a eukaryotic cell and its main organelles in detail.", + "What is the wavelength of light with frequency 5 x 10^14 Hz in vacuum?", + "Find the eigenvalues of the 2x2 matrix [[2, 1], [1, 2]] step by step.", + "Explain how DNA replication occurs during cell division and list the enzymes involved.", + "Balance the chemical equation H2 + O2 produces H2O and explain conservation of atoms.", + "Determine the molarity of a solution containing 5 grams of NaCl in 250 mL of water.", +] + + +REALISTIC_SCIENCE_PROMPTS: list[str] = ( + [_EXACT_DUP, _EXACT_DUP, _EXACT_DUP] + + _FORMATTING_VARIANTS + + _NEAR_DUP_VARIANTS + + _SINGLETONS +) + +# Index ranges for assertions in tests. +EXACT_DUP_RANGE = range(0, 3) # indices 0, 1, 2 +FORMATTING_RANGE = range(3, 6) # indices 3, 4, 5 +NEAR_DUP_RANGE = range(6, 9) # indices 6, 7, 8 +SINGLETON_RANGE = range(9, 16) # indices 9..15 diff --git a/sdg/preprocessing/dedupe/minhash.py b/sdg/preprocessing/dedupe/minhash.py new file mode 100644 index 0000000000000000000000000000000000000000..d55399f0f7115ce548e779a03be1e42cd7c5f501 --- /dev/null +++ b/sdg/preprocessing/dedupe/minhash.py @@ -0,0 +1,139 @@ +"""Lexical near-duplicate dedup via MinHash + LSH (datasketch). + +Catches near-duplicates that share most of their token n-grams: copy-paste +with edits, formatting changes, slight rewordings. Misses true paraphrases +(use SemanticDeduplicator for those). + +Defaults follow the recipe used in FineWeb / Dolma / DataComp-LM: +threshold=0.8 (Jaccard), num_perm=128, 5-word shingles. +""" + +from __future__ import annotations + +from typing import Any, Callable + +from tqdm import tqdm + +from .base import cluster_and_pick, normalize_text + + +class MinHashDeduplicator: + def __init__( + self, + threshold: float = 0.8, + num_perm: int = 128, + shingle_size: int = 5, + shingle_kind: str = "word", # "word" or "char" + normalize: bool = True, + seed: int = 42, + ): + if shingle_kind not in {"word", "char"}: + raise ValueError(f"shingle_kind must be 'word' or 'char', got {shingle_kind!r}") + if not 0.0 < threshold < 1.0: + raise ValueError( + f"threshold must be in (0, 1) (datasketch's LSH cannot pick " + f"valid bands at threshold=1.0); got {threshold}" + ) + if shingle_size < 1: + raise ValueError(f"shingle_size must be >= 1, got {shingle_size}") + self.threshold = threshold + self.num_perm = num_perm + self.shingle_size = shingle_size + self.shingle_kind = shingle_kind + self.normalize = normalize + self.seed = seed + + def shingles(self, text: str) -> set[bytes]: + if self.shingle_kind == "word": + tokens = text.split() + if len(tokens) < self.shingle_size: + return {" ".join(tokens).encode("utf-8")} + return { + " ".join(tokens[i : i + self.shingle_size]).encode("utf-8") + for i in range(len(tokens) - self.shingle_size + 1) + } + if len(text) < self.shingle_size: + return {text.encode("utf-8")} + return { + text[i : i + self.shingle_size].encode("utf-8") + for i in range(len(text) - self.shingle_size + 1) + } + + def dedup( + self, + texts: list[str], + key_fn: Callable[[int], Any] | None = None, + show_progress: bool = True, + return_clusters: bool = False, + ) -> "list[int] | tuple[list[int], dict[int, list[int]]]": + from datasketch import MinHash, MinHashLSH + + n = len(texts) + if n == 0: + return ([], {}) if return_clusters else [] + + iter_sigs = ( + tqdm(texts, desc="MinHash signatures", unit=" rec", smoothing=0.05) + if show_progress + else texts + ) + minhashes: list[MinHash] = [] + for text in iter_sigs: + norm = normalize_text(text) if self.normalize else text + shingles = self.shingles(norm) + m = MinHash(num_perm=self.num_perm, seed=self.seed) + for sh in shingles: + m.update(sh) + minhashes.append(m) + + lsh = MinHashLSH(threshold=self.threshold, num_perm=self.num_perm) + iter_insert = ( + tqdm(enumerate(minhashes), total=n, desc="MinHash LSH insert", unit=" rec", smoothing=0.05) + if show_progress + else enumerate(minhashes) + ) + for i, m in iter_insert: + lsh.insert(str(i), m) + + iter_query = ( + tqdm(enumerate(minhashes), total=n, desc="MinHash query", unit=" rec", smoothing=0.05) + if show_progress + else enumerate(minhashes) + ) + # LSH banding returns *candidate* pairs from its probabilistic S-curve, + # which includes sub-threshold false positives. Without verification, a + # dense near-threshold blob (distinct items sharing a long boilerplate + # block, each pair ~0.65 < threshold) gets transitively collapsed by + # union-find. We verify each candidate with EXACT Jaccard before it + # becomes a merge edge. (The signature-estimated Jaccard is too noisy + # near the boundary — ~±0.1 at num_perm=128 — to use here.) + # + # Shingles are recomputed lazily only for candidate-involved indices and + # cached. Pairs are always (i, j) with j > i and queries run in + # increasing i, so index i is never referenced again after its own + # query — we evict it immediately to keep the cache bounded. + norm_fn = normalize_text if self.normalize else (lambda t: t) + shingle_cache: dict[int, set[bytes]] = {} + + def shingles_for(idx: int) -> set[bytes]: + cached = shingle_cache.get(idx) + if cached is None: + cached = self.shingles(norm_fn(texts[idx])) + shingle_cache[idx] = cached + return cached + + pairs: list[tuple[int, int]] = [] + for i, m in iter_query: + candidates = [j for j in (int(c) for c in lsh.query(m)) if j > i] + if candidates: + si = shingles_for(i) + for j in candidates: + sj = shingles_for(j) + if len(si & sj) / len(si | sj) >= self.threshold: + pairs.append((i, j)) + shingle_cache.pop(i, None) + + keep, clusters = cluster_and_pick(n, pairs, key_fn) + if return_clusters: + return keep, clusters + return keep diff --git a/sdg/preprocessing/dedupe/semantic.py b/sdg/preprocessing/dedupe/semantic.py new file mode 100644 index 0000000000000000000000000000000000000000..e4e122d0d9bc442884860b34e22c9e3a19b4fc86 --- /dev/null +++ b/sdg/preprocessing/dedupe/semantic.py @@ -0,0 +1,212 @@ +"""Paraphrase dedup via sentence embeddings + FAISS ANN (SemDeDup-style). + +Catches semantic near-duplicates that share few n-grams ("Solve x^2-4=0" vs +"Find roots of x squared minus four"). Embeds with sentence-transformers, +indexes with faiss, queries top-k, and clusters pairs with cosine >= threshold. + +Defaults: BAAI/bge-small-en-v1.5 (384-dim, MPS-friendly), cosine >= 0.92. +Useful threshold range: 0.90-0.95. Below 0.88 starts dropping legitimately +distinct prompts; above 0.95 only catches near-exact paraphrases. + +On Apple Silicon (M-series), `device="auto"` selects MPS for encoding. +FAISS always runs on CPU here (no faiss-gpu on macOS). +""" + +from __future__ import annotations + +import time +from typing import Any, Callable + +from tqdm import tqdm + +from .base import cluster_and_pick + + +class SemanticDeduplicator: + def __init__( + self, + model_name: str = "BAAI/bge-small-en-v1.5", + threshold: float = 0.92, + batch_size: int = 128, + device: str = "auto", # "auto" | "mps" | "cuda" | "cpu" + topk: int = 10, + hnsw_threshold: int = 50_000, # use IndexFlatIP below this size + ): + if not 0.0 < threshold <= 1.0: + raise ValueError(f"threshold must be in (0, 1], got {threshold}") + if topk < 1: + raise ValueError(f"topk must be >= 1, got {topk}") + self.model_name = model_name + self.threshold = threshold + self.batch_size = batch_size + self.device = device + self.topk = topk + self.hnsw_threshold = hnsw_threshold + + def resolve_device(self) -> str: + if self.device != "auto": + return self.device + try: + import torch + + if torch.backends.mps.is_available(): + return "mps" + if torch.cuda.is_available(): + return "cuda" + except ImportError: + pass + return "cpu" + + def encode(self, texts: list[str], show_progress: bool = True): + """Encode texts to L2-normalized float32 embeddings (np.ndarray).""" + import gc + import numpy as np + from sentence_transformers import SentenceTransformer + + device = self.resolve_device() + if show_progress: + print( + f" Encoding {len(texts):,} texts with {self.model_name} on " + f"device={device} (batch_size={self.batch_size})" + ) + t0 = time.perf_counter() + model = SentenceTransformer(self.model_name, device=device) + embeddings = model.encode( + texts, + batch_size=self.batch_size, + convert_to_numpy=True, + normalize_embeddings=True, + show_progress_bar=show_progress, + ) + # Take ownership of the buffer so FAISS doesn't segfault on a + # torch-backed memoryview, then drop the model + flush MPS cache + GC + # to release allocations before FAISS starts grabbing memory. + embeddings = np.array(embeddings, dtype=np.float32, copy=True, order="C") + del model + try: + import torch + + if torch.backends.mps.is_available(): + torch.mps.empty_cache() + elif torch.cuda.is_available(): + torch.cuda.empty_cache() + except ImportError: + pass + gc.collect() + if show_progress: + print(f" encoded in {time.perf_counter() - t0:.1f}s") + return embeddings + + def dedup_from_embeddings( + self, + embeddings, + key_fn: Callable[[int], Any] | None = None, + show_progress: bool = True, + return_clusters: bool = False, + ) -> "list[int] | tuple[list[int], dict[int, list[int]]]": + """Cluster pre-computed (L2-normalized) embeddings via FAISS top-k.""" + import faiss + import gc + import numpy as np + + # macOS arm64: PyTorch's libomp and FAISS's libomp can clash and + # segfault under concurrent thread pools. Single-threading FAISS is + # the standard workaround. (Tiny perf hit at our scale; eliminates + # the segfault.) + faiss.omp_set_num_threads(1) + gc.collect() + + n = len(embeddings) + if n == 0: + return ([], {}) if return_clusters else [] + + emb = np.ascontiguousarray(embeddings, dtype=np.float32) + + # Defensive: NaN/Inf in input embeddings can segfault FAISS internals. + # Replace with zero vectors (which have inner-product 0 with everything, + # well below any reasonable threshold, so they end up as singletons). + nonfinite_rows = ~np.isfinite(emb).all(axis=1) + n_nonfinite = int(nonfinite_rows.sum()) + if n_nonfinite > 0: + if show_progress: + print( + f" WARNING: {n_nonfinite:,} of {n:,} embeddings contained " + f"NaN/Inf — replacing with zero vectors." + ) + emb = np.nan_to_num(emb, nan=0.0, posinf=0.0, neginf=0.0) + + dim = emb.shape[1] + + if n < self.hnsw_threshold: + index = faiss.IndexFlatIP(dim) + index_kind = "Flat" + else: + # METRIC_INNER_PRODUCT so search() returns cosine sim (vectors are + # L2-normalized). Default is L2, which would silently invert the + # threshold semantics. + index = faiss.IndexHNSWFlat(dim, 32, faiss.METRIC_INNER_PRODUCT) + index.hnsw.efConstruction = 200 + index.hnsw.efSearch = 64 + index_kind = "HNSW" + + if show_progress: + print(f" Building {index_kind} FAISS index over {n:,} vectors (dim={dim}) ...") + t0 = time.perf_counter() + index.add(emb) + if show_progress: + print(f" done in {time.perf_counter() - t0:.2f}s") + + topk = min(self.topk, n) + if show_progress: + print(f" Searching top-{topk} neighbors for {n:,} queries ...") + t0 = time.perf_counter() + distances, indices = index.search(emb, topk) + if show_progress: + print(f" done in {time.perf_counter() - t0:.2f}s") + + pairs: list[tuple[int, int]] = [] + t0 = time.perf_counter() + iter_pairs = ( + tqdm(range(n), desc=" Extracting pairs", unit=" rec", smoothing=0.05, leave=False) + if show_progress else range(n) + ) + for i in iter_pairs: + for d, j in zip(distances[i], indices[i]): + if j == -1 or j <= i: + continue + if d >= self.threshold: + pairs.append((i, int(j))) + if show_progress: + print( + f" Found {len(pairs):,} candidate pairs (cos>={self.threshold}) " + f"in {time.perf_counter() - t0:.2f}s" + ) + + t0 = time.perf_counter() + keep, clusters = cluster_and_pick(n, pairs, key_fn) + if show_progress: + print( + f" Built {len(keep):,} clusters via union-find " + f"in {time.perf_counter() - t0:.2f}s" + ) + if return_clusters: + return keep, clusters + return keep + + def dedup( + self, + texts: list[str], + key_fn: Callable[[int], Any] | None = None, + show_progress: bool = True, + return_clusters: bool = False, + ) -> "list[int] | tuple[list[int], dict[int, list[int]]]": + n = len(texts) + if n == 0: + return ([], {}) if return_clusters else [] + embeddings = self.encode(texts, show_progress=show_progress) + return self.dedup_from_embeddings( + embeddings, + key_fn=key_fn, + show_progress=show_progress, + return_clusters=return_clusters, + ) diff --git a/sdg/preprocessing/dedupe/templates.py b/sdg/preprocessing/dedupe/templates.py new file mode 100644 index 0000000000000000000000000000000000000000..8a2e015da261480116e5156595ba911564b5bd60 --- /dev/null +++ b/sdg/preprocessing/dedupe/templates.py @@ -0,0 +1,79 @@ +"""Auto-detect and strip common instruction-template prefixes for dedup similarity. + +Many SFT datasets (especially MCQ-heavy ones) use a small number of fixed +instruction templates as prompt prefixes, e.g. + + "Answer the following multiple choice question. The last line of your + response should be in the following format: 'Answer: A/B/C/D' (e.g. + 'Answer: A')." + +These shared prefixes inflate Jaccard similarity (MinHash) and embedding +cosine similarity (semantic), causing distinct questions that share a +template to be falsely merged as duplicates. + +This module provides: + - find_common_prefixes(): scan a corpus and auto-detect frequent prefixes + - strip_template(): remove the longest matching template from a text + - template_hit_stats(): audit helper showing per-template hit counts + +Intended usage: pass the *stripped* prompts to MinHash/Semantic for the +similarity computation, but keep the *original* (unstripped) prompts in +the records that get persisted/uploaded. +""" + +from __future__ import annotations + +from collections import Counter + + +def find_common_prefixes( + texts: list[str], + min_count: int = 10, + prefix_lengths: tuple[int, ...] = (50, 100, 150, 200, 300, 400, 500, 600), +) -> list[str]: + """Find character-prefix substrings that begin >= `min_count` texts. + + Tries each length in `prefix_lengths` independently — this lets the + detector find both short and long template variants. The result is sorted + longest-first so that during stripping the most-specific (longest) match + wins (see `strip_template`). + """ + counter: Counter = Counter() + for t in texts: + for length in prefix_lengths: + if len(t) >= length: + counter[t[:length]] += 1 + templates = [pfx for pfx, c in counter.items() if c >= min_count] + templates.sort(key=len, reverse=True) + return templates + + +def strip_template(text: str, templates: list[str]) -> str: + """Return `text` with the longest matching template prefix removed. + + `templates` should be sorted longest-first (as returned by + `find_common_prefixes`). Leading whitespace is stripped from the result. + Returns `text` unchanged when no template matches or `templates` is empty. + """ + for tmpl in templates: + if text.startswith(tmpl): + return text[len(tmpl):].lstrip() + return text + + +def template_hit_stats(texts: list[str], templates: list[str]) -> dict[str, int]: + """For each template, count how many texts have it as the longest matching prefix. + + Texts that don't match any template are counted under ''. + Sum of counts equals len(texts). + """ + counts: dict[str, int] = {tmpl: 0 for tmpl in templates} + counts[""] = 0 + for t in texts: + for tmpl in templates: + if t.startswith(tmpl): + counts[tmpl] += 1 + break + else: + counts[""] += 1 + return counts diff --git a/sdg/preprocessing/dedupe/test_base.py b/sdg/preprocessing/dedupe/test_base.py new file mode 100644 index 0000000000000000000000000000000000000000..1e04a542e614991a2c5aabaaaec58f217e9349c6 --- /dev/null +++ b/sdg/preprocessing/dedupe/test_base.py @@ -0,0 +1,88 @@ +"""Tests for the union-find / representative-picker / normalization helpers.""" + +from __future__ import annotations + +from sdg.preprocessing.dedupe.base import UnionFind, normalize_text, pick_representatives + + +# ── UnionFind ──────────────────────────────────────────────────────────────── + +def test_unionfind_each_singleton(): + uf = UnionFind(5) + assert {uf.find(i) for i in range(5)} == {0, 1, 2, 3, 4} + + +def test_unionfind_simple_merge(): + uf = UnionFind(5) + uf.union(0, 1) + assert uf.find(0) == uf.find(1) + assert uf.find(2) != uf.find(0) + + +def test_unionfind_transitive_chain(): + uf = UnionFind(5) + uf.union(0, 1) + uf.union(1, 2) + uf.union(2, 3) + assert uf.find(0) == uf.find(3) + assert uf.find(0) != uf.find(4) + + +def test_unionfind_idempotent_self_union(): + uf = UnionFind(3) + uf.union(1, 1) + uf.union(0, 1) + uf.union(0, 1) # repeat + assert uf.find(0) == uf.find(1) + + +# ── pick_representatives ───────────────────────────────────────────────────── + +def test_pick_reps_no_pairs_keeps_all(): + keep = pick_representatives(4, [], key_fn=None) + assert keep == [0, 1, 2, 3] + + +def test_pick_reps_collapses_pair(): + keep = pick_representatives(3, [(0, 1)], key_fn=None) + assert keep == [0, 2] # (0,1) collapsed to smallest index 0; 2 alone + + +def test_pick_reps_collapses_chain(): + keep = pick_representatives(5, [(0, 1), (1, 2), (3, 4)], key_fn=None) + assert keep == [0, 3] + + +def test_pick_reps_key_fn_selects_max_by_negative_length(): + lengths = [10, 50, 5, 100] + key_fn = lambda i: -lengths[i] # longest wins + keep = pick_representatives(4, [(0, 1), (2, 3)], key_fn=key_fn) + assert keep == [1, 3] # idx 1 (len 50) > idx 0 (10); idx 3 (100) > idx 2 (5) + + +def test_pick_reps_returns_sorted(): + keep = pick_representatives(6, [(2, 5), (0, 3)], key_fn=None) + assert keep == sorted(keep) + + +def test_pick_reps_empty(): + assert pick_representatives(0, [], key_fn=None) == [] + + +# ── normalize_text ─────────────────────────────────────────────────────────── + +def test_normalize_lowercases(): + assert normalize_text("Hello WORLD") == "hello world" + + +def test_normalize_collapses_whitespace(): + assert normalize_text("a b\t\tc\n\nd") == "a b c d" + + +def test_normalize_strips_edges(): + assert normalize_text(" hi ") == "hi" + + +def test_normalize_idempotent(): + text = "Hello\t\tWorld\n" + assert normalize_text(normalize_text(text)) == normalize_text(text) diff --git a/sdg/preprocessing/dedupe/test_benchmarks.py b/sdg/preprocessing/dedupe/test_benchmarks.py new file mode 100644 index 0000000000000000000000000000000000000000..7437c953beec63c18efeb2ee095e14ea81a85567 --- /dev/null +++ b/sdg/preprocessing/dedupe/test_benchmarks.py @@ -0,0 +1,401 @@ +"""Speed benchmarks for MinHash + Semantic deduplicators. + +Prints timing tables to stdout (uses `capsys.disabled()` so output appears +even without `pytest -s`). All benchmarks are bounded so the suite stays +under ~60s on an M-series machine. + +To include the (slower) sentence-transformers encoding benchmark, set +SDG_BENCH_ENCODE=1 in the environment. That one downloads the model on +first run. + +Run: + pytest sdg/preprocessing/dedupe/test_benchmarks.py -v + SDG_BENCH_ENCODE=1 pytest sdg/preprocessing/dedupe/test_benchmarks.py -v +""" + +from __future__ import annotations + +import gc +import os +import platform +import random +import resource +import time + +import pytest + +ds = pytest.importorskip("datasketch") +faiss = pytest.importorskip("faiss") +import numpy as np + +from sdg.preprocessing.dedupe.minhash import MinHashDeduplicator +from sdg.preprocessing.dedupe.semantic import SemanticDeduplicator + + +# ──────────────────────────────────────────────────────────────────────────── +# Helpers +# ──────────────────────────────────────────────────────────────────────────── + +_VOCAB = ( + "alpha beta gamma delta epsilon zeta eta theta iota kappa lambda mu nu xi " + "omicron pi rho sigma tau upsilon phi chi psi omega the quick brown fox " + "jumps over lazy dog cat house tree mountain river ocean sun moon star " + "planet galaxy universe science math physics chemistry biology equation " + "matrix vector function derivative integral hypothesis theorem proof " + "experiment observation analysis synthesis conclusion problem solution" +).split() + + +def _distinct_prompts(n: int, words: int = 25, seed: int = 42) -> list[str]: + rng = random.Random(seed) + return [" ".join(rng.sample(_VOCAB, words)) + f" id{i}" for i in range(n)] + + +def _normalize(v: np.ndarray) -> np.ndarray: + norms = np.linalg.norm(v, axis=1, keepdims=True) + return (v / np.clip(norms, 1e-12, None)).astype(np.float32) + + +def _print_machine_header(out, label: str) -> None: + out(f"\n{'=' * 78}") + out(f" {label}") + out(f" host={platform.machine()} python={platform.python_version()} " + f"darwin={platform.release()}") + out(f"{'=' * 78}") + + +def _peak_rss_mb() -> float: + """Peak resident memory in MB (macOS reports bytes; Linux reports KB).""" + raw = resource.getrusage(resource.RUSAGE_SELF).ru_maxrss + if platform.system() == "Darwin": + return raw / 1e6 + return raw / 1e3 # Linux: KB -> MB + + +# ──────────────────────────────────────────────────────────────────────────── +# MinHash benchmarks +# ──────────────────────────────────────────────────────────────────────────── + +def test_bench_minhash_full_dedup_increasing_n(capsys): + """End-to-end MinHash dedup throughput at increasing N (distinct prompts).""" + sizes = [100, 500, 2_000, 10_000] + rows = [] + for n in sizes: + texts = _distinct_prompts(n) + t0 = time.perf_counter() + keep = MinHashDeduplicator(threshold=0.8).dedup(texts, show_progress=False) + elapsed = time.perf_counter() - t0 + rps = n / elapsed if elapsed > 0 else float("inf") + rows.append((n, elapsed, rps, len(keep))) + + with capsys.disabled(): + out = print + _print_machine_header(out, "MinHash full dedup (signatures + LSH + cluster)") + out(f" {'N':>10s} {'time':>10s} {'rec/sec':>14s} {'kept':>10s}") + for n, elapsed, rps, kept in rows: + out(f" {n:>10,} {elapsed:>9.3f}s {rps:>12,.0f} {kept:>10,}") + out("=" * 78) + + +def test_bench_minhash_signature_only(capsys): + """Signature creation alone (no LSH/query) — isolates shingling+hashing cost.""" + from datasketch import MinHash + + sizes = [100, 1_000, 10_000] + rows = [] + deduper = MinHashDeduplicator() + for n in sizes: + texts = _distinct_prompts(n) + t0 = time.perf_counter() + for text in texts: + shingles = deduper.shingles(text.lower()) + m = MinHash(num_perm=128, seed=42) + for sh in shingles: + m.update(sh) + elapsed = time.perf_counter() - t0 + rps = n / elapsed if elapsed > 0 else float("inf") + rows.append((n, elapsed, rps)) + + with capsys.disabled(): + out = print + _print_machine_header(out, "MinHash signature creation only (no LSH)") + out(f" {'N':>10s} {'time':>10s} {'rec/sec':>14s}") + for n, elapsed, rps in rows: + out(f" {n:>10,} {elapsed:>9.3f}s {rps:>12,.0f}") + out("=" * 78) + + +def test_bench_minhash_threshold_sweep(capsys): + """Same N, vary threshold — high threshold should be faster (fewer LSH bands).""" + n = 5_000 + texts = _distinct_prompts(n) + rows = [] + for threshold in [0.5, 0.7, 0.8, 0.9, 0.95]: + t0 = time.perf_counter() + keep = MinHashDeduplicator(threshold=threshold).dedup(texts, show_progress=False) + elapsed = time.perf_counter() - t0 + rows.append((threshold, elapsed, len(keep))) + + with capsys.disabled(): + out = print + _print_machine_header(out, f"MinHash threshold sweep (N={n:,} distinct prompts)") + out(f" {'threshold':>10s} {'time':>10s} {'kept':>10s}") + for threshold, elapsed, kept in rows: + out(f" {threshold:>10.2f} {elapsed:>9.3f}s {kept:>10,}") + out("=" * 78) + + +def test_bench_minhash_num_perm_sweep(capsys): + """Higher num_perm = better accuracy but slower.""" + n = 2_000 + texts = _distinct_prompts(n) + rows = [] + for num_perm in [32, 64, 128, 256, 512]: + t0 = time.perf_counter() + MinHashDeduplicator(num_perm=num_perm).dedup(texts, show_progress=False) + elapsed = time.perf_counter() - t0 + rows.append((num_perm, elapsed)) + + with capsys.disabled(): + out = print + _print_machine_header(out, f"MinHash num_perm sweep (N={n:,})") + out(f" {'num_perm':>10s} {'time':>10s}") + for num_perm, elapsed in rows: + out(f" {num_perm:>10d} {elapsed:>9.3f}s") + out("=" * 78) + + +# ──────────────────────────────────────────────────────────────────────────── +# Semantic benchmarks (skip the encoder; bench dedup_from_embeddings) +# ──────────────────────────────────────────────────────────────────────────── + +def test_bench_semantic_dedup_from_embeddings_increasing_n(capsys): + """Pure FAISS+cluster throughput at increasing N (random 384-dim vectors). + + Capped at 25K to stay clear of a faiss-cpu 1.8.0 macOS-arm64 memory-pressure + regime that can segfault when the suite has accumulated state from prior + tests. (50K runs fine in isolation; not when chained.) + """ + sizes = [500, 2_000, 10_000, 25_000] + rows = [] + for n in sizes: + rng = np.random.default_rng(42) + emb = _normalize(rng.standard_normal((n, 384)).astype(np.float32)) + dedup = SemanticDeduplicator(threshold=0.9, hnsw_threshold=5_000) + kind = "Flat" if n < dedup.hnsw_threshold else "HNSW" + t0 = time.perf_counter() + keep = dedup.dedup_from_embeddings(emb, show_progress=False) + elapsed = time.perf_counter() - t0 + rps = n / elapsed if elapsed > 0 else float("inf") + rows.append((n, kind, elapsed, rps, len(keep))) + + with capsys.disabled(): + out = print + _print_machine_header(out, "Semantic dedup_from_embeddings (FAISS index + cluster)") + out(f" {'N':>10s} {'index':>8s} {'time':>10s} {'rec/sec':>14s} {'kept':>10s}") + for n, kind, elapsed, rps, kept in rows: + out(f" {n:>10,} {kind:>8s} {elapsed:>9.3f}s {rps:>12,.0f} {kept:>10,}") + out("=" * 78) + + +def test_bench_semantic_threshold_sweep(capsys): + """Lower threshold = more clusters to materialize but FAISS time is the same.""" + n = 5_000 + rng = np.random.default_rng(42) + emb = _normalize(rng.standard_normal((n, 384)).astype(np.float32)) + rows = [] + for threshold in [0.7, 0.8, 0.9, 0.95, 0.99]: + dedup = SemanticDeduplicator(threshold=threshold, hnsw_threshold=1_000) + t0 = time.perf_counter() + keep = dedup.dedup_from_embeddings(emb, show_progress=False) + elapsed = time.perf_counter() - t0 + rows.append((threshold, elapsed, len(keep))) + + with capsys.disabled(): + out = print + _print_machine_header(out, f"Semantic threshold sweep (N={n:,}, 384-dim)") + out(f" {'threshold':>10s} {'time':>10s} {'kept':>10s}") + for threshold, elapsed, kept in rows: + out(f" {threshold:>10.2f} {elapsed:>9.3f}s {kept:>10,}") + out("=" * 78) + + +def test_bench_semantic_dim_sweep(capsys): + """Embedding dimensionality vs throughput at fixed N.""" + n = 5_000 + rows = [] + for dim in [128, 256, 384, 768, 1024]: + rng = np.random.default_rng(42) + emb = _normalize(rng.standard_normal((n, dim)).astype(np.float32)) + dedup = SemanticDeduplicator(threshold=0.95, hnsw_threshold=1_000) + t0 = time.perf_counter() + dedup.dedup_from_embeddings(emb, show_progress=False) + elapsed = time.perf_counter() - t0 + rows.append((dim, elapsed)) + + with capsys.disabled(): + out = print + _print_machine_header(out, f"Semantic embedding-dim sweep (N={n:,}, HNSW)") + out(f" {'dim':>10s} {'time':>10s}") + for dim, elapsed in rows: + out(f" {dim:>10d} {elapsed:>9.3f}s") + out("=" * 78) + + +# ──────────────────────────────────────────────────────────────────────────── +# Optional: real-encoder benchmark (gated by env var) +# ──────────────────────────────────────────────────────────────────────────── + +@pytest.mark.skipif( + os.getenv("SDG_BENCH_ENCODE") != "1", + reason="set SDG_BENCH_ENCODE=1 to run the real-encoder benchmark " + "(downloads BAAI/bge-small-en-v1.5 on first run)", +) +def test_bench_semantic_real_encoder_small(capsys): + pytest.importorskip("sentence_transformers") + from sentence_transformers import SentenceTransformer + + sizes = [128, 512, 1024] + dedup = SemanticDeduplicator() + device = dedup.resolve_device() + model = SentenceTransformer(dedup.model_name, device=device) + + rows = [] + for n in sizes: + texts = _distinct_prompts(n) + t0 = time.perf_counter() + emb = model.encode( + texts, + batch_size=128, + convert_to_numpy=True, + normalize_embeddings=True, + show_progress_bar=False, + ) + elapsed = time.perf_counter() - t0 + rps = n / elapsed if elapsed > 0 else float("inf") + rows.append((n, elapsed, rps)) + assert emb.shape == (n, 384) + + with capsys.disabled(): + out = print + _print_machine_header( + out, f"Semantic encoding (BAAI/bge-small-en-v1.5, device={device}, batch=128)" + ) + out(f" {'N':>10s} {'time':>10s} {'rec/sec':>14s}") + for n, elapsed, rps in rows: + out(f" {n:>10,} {elapsed:>9.3f}s {rps:>12,.0f}") + out("=" * 78) + + +# ──────────────────────────────────────────────────────────────────────────── +# Large-scale projection: how long would 3M prompts take? +# ──────────────────────────────────────────────────────────────────────────── +# +# The Nemotron-Cascade-2 science split has ~3M unique prompts after exact +# dedup. We benchmark MinHash and semantic-FAISS at 50K/100K/250K and +# extrapolate. Encoder cost (~30 min for 3M on MPS) is a separate stage — +# enable with SDG_BENCH_ENCODE=1 to measure. +# +# Memory: 250K signatures + LSH ~500-700 MB; 250K x 384 float32 = 384 MB. +# Both fit easily on a laptop. We `gc.collect()` between sizes defensively +# because faiss-cpu 1.8.0 has shown memory-pressure segfaults under pytest. + +def test_bench_3m_projection_minhash(capsys): + """Project MinHash dedup time to 3M using N=50K/100K/250K data points.""" + sizes = [50_000, 100_000, 250_000] + rows = [] + for n in sizes: + gc.collect() + texts = _distinct_prompts(n) + t0 = time.perf_counter() + keep = MinHashDeduplicator(threshold=0.8).dedup(texts, show_progress=False) + elapsed = time.perf_counter() - t0 + rps = n / elapsed if elapsed > 0 else float("inf") + rss = _peak_rss_mb() + rows.append((n, elapsed, rps, len(keep), rss)) + del texts, keep + gc.collect() + + n_ref, _, rps_ref, _, _ = rows[-1] + proj_sec = 3_000_000 / rps_ref + + with capsys.disabled(): + out = print + _print_machine_header(out, "MinHash large-scale projection -> 3M target") + out(f" {'N':>10s} {'time':>10s} {'rec/sec':>12s} {'kept':>10s} {'peak RSS MB':>14s}") + for n, elapsed, rps, kept, rss in rows: + out(f" {n:>10,} {elapsed:>9.2f}s {rps:>10,.0f} {kept:>10,} {rss:>13,.0f}") + out(f" {'PROJ 3M':>10s} {proj_sec:>9.0f}s ({proj_sec/60:.1f} min, " + f"linear extrap from N={n_ref:,})") + out("=" * 78) + + +def test_bench_3m_projection_semantic_faiss(capsys): + """Project semantic dedup_from_embeddings time to 3M (FAISS only, no encoding).""" + sizes = [50_000, 100_000, 250_000] + rows = [] + for n in sizes: + gc.collect() + rng = np.random.default_rng(42) + emb = _normalize(rng.standard_normal((n, 384)).astype(np.float32)) + dedup = SemanticDeduplicator(threshold=0.92, hnsw_threshold=10_000) + t0 = time.perf_counter() + keep = dedup.dedup_from_embeddings(emb, show_progress=False) + elapsed = time.perf_counter() - t0 + rps = n / elapsed if elapsed > 0 else float("inf") + rss = _peak_rss_mb() + rows.append((n, elapsed, rps, len(keep), rss)) + del emb, keep, dedup + gc.collect() + + # HNSW per-record cost grows with N (~log N). Use ratio from largest point + # but warn that this is optimistic. + n_ref, t_ref, rps_ref, _, _ = rows[-1] + proj_sec_linear = 3_000_000 / rps_ref + # NlogN scaling estimate: t_3M ≈ t_ref * (3M/N_ref) * log(3M)/log(N_ref) + import math + proj_sec_nlogn = t_ref * (3_000_000 / n_ref) * (math.log(3_000_000) / math.log(n_ref)) + + with capsys.disabled(): + out = print + _print_machine_header(out, "Semantic FAISS large-scale projection -> 3M target") + out(f" {'N':>10s} {'time':>10s} {'rec/sec':>12s} {'peak RSS MB':>14s}") + for n, elapsed, rps, kept, rss in rows: + out(f" {n:>10,} {elapsed:>9.2f}s {rps:>10,.0f} {rss:>13,.0f}") + out(f" {'PROJ 3M (linear)':>20s} {proj_sec_linear:>9.0f}s ({proj_sec_linear/60:.1f} min)") + out(f" {'PROJ 3M (N log N)':>20s} {proj_sec_nlogn:>9.0f}s ({proj_sec_nlogn/60:.1f} min)") + out("=" * 78) + + +# ──────────────────────────────────────────────────────────────────────────── +# End-to-end pipeline benchmark +# ──────────────────────────────────────────────────────────────────────────── + +def test_bench_full_pipeline_minhash_then_semantic(capsys): + """End-to-end MinHash + semantic dedup on synthetic 5K-prompt dataset. + + Uses random embeddings for the semantic stage to avoid model download. + """ + n = 5_000 + texts = _distinct_prompts(n) + + t0 = time.perf_counter() + keep_mh = MinHashDeduplicator(threshold=0.8).dedup(texts, show_progress=False) + t_mh = time.perf_counter() - t0 + + rng = np.random.default_rng(0) + emb = _normalize(rng.standard_normal((len(keep_mh), 384)).astype(np.float32)) + + dedup = SemanticDeduplicator(threshold=0.92, hnsw_threshold=1_000) + t0 = time.perf_counter() + keep_sem = dedup.dedup_from_embeddings(emb, show_progress=False) + t_sem = time.perf_counter() - t0 + + with capsys.disabled(): + out = print + _print_machine_header(out, f"End-to-end pipeline (N={n:,} prompts)") + out(f" {'stage':<32s} {'time':>10s} {'kept':>10s}") + out(f" {'MinHash (Jaccard 0.8)':<32s} {t_mh:>9.3f}s {len(keep_mh):>10,}") + out(f" {'Semantic (cosine 0.92, HNSW)':<32s} {t_sem:>9.3f}s {len(keep_sem):>10,}") + out(f" {'TOTAL':<32s} {t_mh + t_sem:>9.3f}s {len(keep_sem):>10,}") + out("=" * 78) diff --git a/sdg/preprocessing/dedupe/test_edge_cases.py b/sdg/preprocessing/dedupe/test_edge_cases.py new file mode 100644 index 0000000000000000000000000000000000000000..d5f27cb51efe8fca43fced2092c8b72c8e379dac --- /dev/null +++ b/sdg/preprocessing/dedupe/test_edge_cases.py @@ -0,0 +1,394 @@ +"""Edge cases and limits for MinHash + Semantic deduplicators. + +Covers: + - Degenerate inputs (empty strings, whitespace, single char, unicode, emoji) + - Threshold boundaries (1.0, near-zero, just-above-cutoff) + - Shingle/permutation extremes + - Large clusters, many disjoint clusters, mixed + - HNSW vs Flat boundary behavior + - Numerical edge cases for embeddings (zero vector, antipodal, very high dim) +""" + +from __future__ import annotations + +import pytest + +ds = pytest.importorskip("datasketch") +faiss = pytest.importorskip("faiss") +import numpy as np + +from sdg.preprocessing.dedupe.minhash import MinHashDeduplicator +from sdg.preprocessing.dedupe.semantic import SemanticDeduplicator + + +# ──────────────────────────────────────────────────────────────────────────── +# MinHash: degenerate inputs +# ──────────────────────────────────────────────────────────────────────────── + +def test_minhash_all_empty_strings_collapse(): + keep = MinHashDeduplicator().dedup(["", "", ""], show_progress=False) + assert len(keep) == 1 + + +def test_minhash_one_empty_with_distinct_others(): + keep = MinHashDeduplicator().dedup( + ["", "alpha beta gamma delta epsilon zeta", "completely different content here"], + show_progress=False, + ) + # Empty doesn't share shingles with the others -> kept as singleton + assert len(keep) == 3 + + +def test_minhash_whitespace_only_strings_collapse_with_normalize(): + keep = MinHashDeduplicator(normalize=True).dedup( + [" ", "\t\t", "\n\n", " \t \n "], + show_progress=False, + ) + assert len(keep) == 1 + + +def test_minhash_single_character_strings(): + """Each char becomes its own short shingle; identical chars collapse.""" + keep = MinHashDeduplicator().dedup(["a", "a", "b"], show_progress=False) + assert len(keep) == 2 + + +def test_minhash_string_shorter_than_shingle_size(): + """When text is shorter than shingle size, the whole text is a single shingle.""" + keep = MinHashDeduplicator(shingle_size=10).dedup( + ["short text", "short text", "very different brief"], + show_progress=False, + ) + assert len(keep) == 2 + + +def test_minhash_unicode_cjk(): + text_a = "解决方程 x 平方 减去 四 等于 零 的 详细 步骤 解释" + text_b = text_a # identical + text_c = "光合作用 是 植物 利用 光能 制造 有机物 的 过程 详解" + keep = MinHashDeduplicator().dedup([text_a, text_b, text_c], show_progress=False) + assert len(keep) == 2 + + +def test_minhash_emoji_inputs(): + a = "Solve this 🎯 equation x squared minus four equals zero step by step" + b = "Solve this 🎯 equation x squared minus four equals zero step by step" + c = "Cook the 🍝 pasta in salted boiling water for nine to ten minutes" + keep = MinHashDeduplicator().dedup([a, b, c], show_progress=False) + assert len(keep) == 2 + + +def test_minhash_repeated_token_pattern(): + """Many copies of the same token -> all shingles are identical.""" + a = ("yes " * 50).strip() + b = ("yes " * 100).strip() + c = ("no " * 50).strip() + keep = MinHashDeduplicator(threshold=0.5).dedup([a, b, c], show_progress=False) + # a and b share their only shingle "yes yes yes yes yes" -> Jaccard 1.0 + # c is disjoint from them + assert len(keep) == 2 + + +# ──────────────────────────────────────────────────────────────────────────── +# MinHash: threshold limits +# ──────────────────────────────────────────────────────────────────────────── + +def test_minhash_threshold_1_rejected_by_validator(): + """datasketch's LSH cannot find valid bands at threshold=1.0; we reject early.""" + with pytest.raises(ValueError, match="threshold"): + MinHashDeduplicator(threshold=1.0) + + +def test_minhash_high_threshold_keeps_near_dups_apart(): + """At a strict threshold (0.95), a clearly-below pair (Jaccard ~0.4) does not collapse.""" + base = " ".join(f"word{i}" for i in range(20)) # 20 tokens, 16 5-grams + a = base + # Append 20 unrelated tokens -> b has 36 5-grams, 16 shared with a, Jaccard 16/36 ≈ 0.44 + b = base + " " + " ".join(f"extra{i}" for i in range(20)) + keep = MinHashDeduplicator(threshold=0.95).dedup([a, b], show_progress=False) + assert len(keep) == 2 + + +def test_minhash_high_threshold_collapses_exact_duplicates(): + """At a strict threshold (0.95), byte-identical strings still collapse.""" + a = " ".join(f"word{i}" for i in range(20)) + keep = MinHashDeduplicator(threshold=0.95).dedup([a, a, a], show_progress=False) + assert len(keep) == 1 + + +def test_minhash_low_threshold_aggressive_collapse(): + """At very low threshold, weak similarity is enough to merge.""" + texts = [ + "alpha beta gamma delta epsilon zeta eta theta iota kappa lambda mu", + "alpha xxxxx yyyyy zzzzz wwwww vvvvv uuuuu ttttt sssss rrrrr qqqqq pppp", # shares only 'alpha' as 1-gram + ] + # threshold 0.05 + 1-word shingles -> very low bar; "alpha" alone may suffice + keep = MinHashDeduplicator(threshold=0.05, shingle_size=1).dedup(texts, show_progress=False) + # Tolerant: result is 1 (collapsed) or 2 (LSH missed). Either is acceptable + # at extreme threshold; just don't crash. + assert len(keep) in {1, 2} + + +def test_minhash_invalid_threshold_zero_raises(): + with pytest.raises(ValueError): + MinHashDeduplicator(threshold=0.0) + + +def test_minhash_invalid_threshold_negative_raises(): + with pytest.raises(ValueError): + MinHashDeduplicator(threshold=-0.5) + + +# ──────────────────────────────────────────────────────────────────────────── +# MinHash: shingle and permutation extremes +# ──────────────────────────────────────────────────────────────────────────── + +def test_minhash_shingle_size_one(): + """1-word shingles == bag-of-words Jaccard.""" + a = "the quick brown fox" + b = "fox the brown quick" # same words, different order + keep = MinHashDeduplicator(shingle_size=1, threshold=0.9).dedup([a, b], show_progress=False) + # Same shingle set -> Jaccard 1.0 -> collapse + assert len(keep) == 1 + + +def test_minhash_shingle_size_one_word_order_matters_with_5gram(): + a = "the quick brown fox jumps over the lazy dog quickly today" + b = "today quickly dog lazy the over jumps fox brown quick the" + keep = MinHashDeduplicator(shingle_size=5, threshold=0.7).dedup([a, b], show_progress=False) + # Word order changes -> different 5-grams -> Jaccard near 0 -> kept apart + assert len(keep) == 2 + + +def test_minhash_invalid_shingle_size_zero_raises(): + with pytest.raises(ValueError): + MinHashDeduplicator(shingle_size=0) + + +def test_minhash_low_num_perm_still_works(): + """Low num_perm == noisier similarity estimates but still functional.""" + a = "alpha beta gamma delta epsilon zeta eta theta iota kappa lambda" + keep = MinHashDeduplicator(num_perm=16, threshold=0.7).dedup([a, a], show_progress=False) + assert len(keep) == 1 + + +def test_minhash_high_num_perm_still_works(): + a = "alpha beta gamma delta epsilon zeta eta theta iota kappa lambda mu" + keep = MinHashDeduplicator(num_perm=512, threshold=0.7).dedup([a, a], show_progress=False) + assert len(keep) == 1 + + +# ──────────────────────────────────────────────────────────────────────────── +# MinHash: cluster topology +# ──────────────────────────────────────────────────────────────────────────── + +def test_minhash_large_identical_cluster_collapses_to_one(): + a = "Solve x squared minus four equals zero showing all your work step by step in detail" + keep = MinHashDeduplicator().dedup([a] * 200, show_progress=False) + assert len(keep) == 1 + + +def test_minhash_many_disjoint_pairs_all_collapse(): + """50 pairs of identical strings -> 50 kept.""" + pairs = [] + for i in range(50): + # each pair shares its own unique vocabulary, no cross-pair overlap + prompt = f"unique prompt vocabulary words alpha{i} beta{i} gamma{i} delta{i} epsilon{i} zeta{i}" + pairs.append(prompt) + pairs.append(prompt) + keep = MinHashDeduplicator().dedup(pairs, show_progress=False) + assert len(keep) == 50 + + +def test_minhash_chain_of_similar_strings_clusters(): + """Three near-identical strings (each one-word edit from the base) form one cluster. + + Pure transitivity (A-B above, B-C above, A-C below) is hard to construct + reliably under MinHash's probabilistic LSH at moderate num_perm. Union-find + behavior on transitivity is independently tested in test_base.py; here we + just verify cluster formation works for chained near-dups. + """ + base = " ".join(f"word{i}" for i in range(40)) # 40 tokens + a = base + b = a.replace("word0 ", "REPLACED_B ", 1) + c = a.replace("word39", "REPLACED_C", 1) + # Pairwise Jaccard ~0.95+, well above threshold. + keep = MinHashDeduplicator(threshold=0.85).dedup([a, b, c], show_progress=False) + assert len(keep) == 1 + + +# ──────────────────────────────────────────────────────────────────────────── +# MinHash: normalization corner cases +# ──────────────────────────────────────────────────────────────────────────── + +def test_minhash_normalize_handles_extreme_whitespace(): + a = "alpha beta gamma delta epsilon zeta eta theta iota kappa" + b = "ALPHA\tBETA GAMMA\n\nDELTA\rEPSILON ZETA ETA\tTHETA IOTA KAPPA" + keep = MinHashDeduplicator(normalize=True, threshold=0.9).dedup([a, b], show_progress=False) + assert len(keep) == 1 + + +def test_minhash_normalize_off_keeps_case_variants_apart(): + a = "ALPHA BETA GAMMA DELTA EPSILON ZETA ETA THETA IOTA KAPPA" + b = "alpha beta gamma delta epsilon zeta eta theta iota kappa" + keep = MinHashDeduplicator(normalize=False, threshold=0.9).dedup([a, b], show_progress=False) + # Different shingles entirely (case-sensitive) -> Jaccard 0 -> both kept + assert len(keep) == 2 + + +# ──────────────────────────────────────────────────────────────────────────── +# Semantic: degenerate vectors +# ──────────────────────────────────────────────────────────────────────────── + +def _normalize(v: np.ndarray) -> np.ndarray: + norms = np.linalg.norm(v, axis=1, keepdims=True) + return (v / np.clip(norms, 1e-12, None)).astype(np.float32) + + +def test_semantic_zero_vector_does_not_crash(): + """A zero vector has undefined direction; clipping keeps it numerically safe.""" + emb = np.zeros((3, 8), dtype=np.float32) + # All-zero inner product is 0 -> below any reasonable threshold -> all kept + keep = SemanticDeduplicator(threshold=0.5).dedup_from_embeddings(emb) + assert len(keep) == 3 + + +def test_semantic_antipodal_vectors_not_clustered(): + """Cosine = -1 must not be treated as similar.""" + a = np.array([[1.0, 0.0, 0.0]], dtype=np.float32) + b = np.array([[-1.0, 0.0, 0.0]], dtype=np.float32) + emb = np.vstack([a, b]) + keep = SemanticDeduplicator(threshold=0.5).dedup_from_embeddings(emb) + assert keep == [0, 1] + + +def test_semantic_one_dim_embeddings(): + emb = np.array([[1.0], [1.0], [-1.0]], dtype=np.float32) + keep = SemanticDeduplicator(threshold=0.99).dedup_from_embeddings(emb) + # Two +1 collapse, -1 stays separate + assert len(keep) == 2 + + +def test_semantic_high_dim_1024(): + rng = np.random.default_rng(0) + emb = _normalize(rng.standard_normal((100, 1024)).astype(np.float32)) + keep = SemanticDeduplicator(threshold=0.95).dedup_from_embeddings(emb) + assert len(keep) == 100 # nearly orthogonal in 1024D + + +# ──────────────────────────────────────────────────────────────────────────── +# Semantic: topk extremes +# ──────────────────────────────────────────────────────────────────────────── + +def test_semantic_topk_one_misses_cross_matches(): + """topk=1 returns each vector's nearest neighbor = self. j>i is never + satisfied, so even highly similar vectors are not merged. + + Uses slightly distinct (cos≈0.996) vectors to avoid the tied-IP nondeterminism + you'd get from two byte-identical embeddings. + """ + emb = _normalize( + np.array( + [[1.0, 0.0], [np.cos(np.radians(5)), np.sin(np.radians(5))]], + dtype=np.float32, + ) + ) + keep = SemanticDeduplicator(threshold=0.95, topk=1).dedup_from_embeddings(emb) + assert len(keep) == 2 + + +def test_semantic_topk_equal_to_n_full_pairwise(): + rng = np.random.default_rng(1) + base = rng.standard_normal((1, 16)).astype(np.float32) + emb = _normalize(np.vstack([base, base, base])) # all identical + keep = SemanticDeduplicator(threshold=0.99, topk=3).dedup_from_embeddings(emb) + assert len(keep) == 1 + + +def test_semantic_topk_larger_than_n_clamped_safely(): + emb = _normalize(np.array([[1.0, 0.0], [1.0, 0.0]], dtype=np.float32)) + # topk=100 > n=2 -> internal min() prevents overrun + keep = SemanticDeduplicator(threshold=0.99, topk=100).dedup_from_embeddings(emb) + assert len(keep) == 1 + + +# ──────────────────────────────────────────────────────────────────────────── +# Semantic: threshold limits +# ──────────────────────────────────────────────────────────────────────────── + +def test_semantic_threshold_1_only_exactly_identical(): + a = _normalize(np.array([[1.0, 0.0]], dtype=np.float32)) + b = _normalize(np.array([[np.cos(np.radians(2)), np.sin(np.radians(2))]], dtype=np.float32)) + emb = np.vstack([a, a, b]) + keep = SemanticDeduplicator(threshold=1.0).dedup_from_embeddings(emb) + # a == a collapses (cos == 1), b stays apart (cos ≈ 0.999 < 1.0) + assert len(keep) == 2 + + +def test_semantic_low_threshold_collapses_many(): + """At cos>=0.5, mildly related vectors collapse together.""" + angles = [0, 10, 20, 30, 40, 50] # all within 60 degrees + emb = _normalize( + np.array( + [[np.cos(np.radians(a)), np.sin(np.radians(a))] for a in angles], + dtype=np.float32, + ) + ) + keep = SemanticDeduplicator(threshold=0.5, topk=6).dedup_from_embeddings(emb) + # All pairwise cos > 0.5 -> all merge transitively + assert len(keep) == 1 + + +# ──────────────────────────────────────────────────────────────────────────── +# Semantic: HNSW vs Flat boundary +# ──────────────────────────────────────────────────────────────────────────── + +def test_semantic_just_below_hnsw_threshold_uses_flat(): + """Boundary check: n < hnsw_threshold -> Flat path; just verify no crash + correctness.""" + rng = np.random.default_rng(7) + emb = _normalize(rng.standard_normal((49, 64)).astype(np.float32)) + dedup = SemanticDeduplicator(threshold=0.99, hnsw_threshold=50) + keep = dedup.dedup_from_embeddings(emb) + assert len(keep) == 49 + + +def test_semantic_just_above_hnsw_threshold_uses_hnsw(): + rng = np.random.default_rng(7) + emb = _normalize(rng.standard_normal((51, 64)).astype(np.float32)) + dedup = SemanticDeduplicator(threshold=0.99, hnsw_threshold=50) + keep = dedup.dedup_from_embeddings(emb) + assert len(keep) == 51 + + +def test_semantic_hnsw_path_finds_planted_cluster_among_distractors(): + """HNSW must still detect a planted duplicate cluster despite approximation.""" + rng = np.random.default_rng(99) + distractors = rng.standard_normal((300, 128)).astype(np.float32) + seed_vec = rng.standard_normal((1, 128)).astype(np.float32) + planted = np.repeat(seed_vec, 30, axis=0) + raw = np.concatenate([distractors, planted], axis=0) + emb = _normalize(raw) + dedup = SemanticDeduplicator(threshold=0.95, hnsw_threshold=50, topk=40) + keep = dedup.dedup_from_embeddings(emb) + # 300 distractors + 1 cluster representative + assert len(keep) == 301 + + +# ──────────────────────────────────────────────────────────────────────────── +# Semantic: representative selection across many clusters +# ──────────────────────────────────────────────────────────────────────────── + +def test_semantic_key_fn_picks_max_across_many_clusters(): + """5 clusters of 3 identical vectors each; key_fn picks the longest in each.""" + rng = np.random.default_rng(11) + seeds = rng.standard_normal((5, 32)).astype(np.float32) + raw = np.repeat(seeds, 3, axis=0) # 15 vectors, 5 distinct directions x 3 + emb = _normalize(raw) + + # response_lengths designed so that within each cluster of 3 (idx triples + # 0-2, 3-5, ...), the LAST index has the largest length. + response_lengths = [10, 20, 30, 40, 50, 60, 70, 80, 90, 100, 110, 120, 130, 140, 150] + key_fn = lambda i: -response_lengths[i] + dedup = SemanticDeduplicator(threshold=0.99, topk=5) + keep = dedup.dedup_from_embeddings(emb, key_fn=key_fn) + assert keep == [2, 5, 8, 11, 14] diff --git a/sdg/preprocessing/dedupe/test_invariants.py b/sdg/preprocessing/dedupe/test_invariants.py new file mode 100644 index 0000000000000000000000000000000000000000..c2f0f60e7e878b06bf8cd4db3e0e08dded21cba2 --- /dev/null +++ b/sdg/preprocessing/dedupe/test_invariants.py @@ -0,0 +1,347 @@ +"""Invariants, end-to-end pipeline simulation, realistic data, and snapshots. + +Sections: + 1. MinHash invariants — properties dedup must always satisfy + 2. Semantic invariants — same, for SemanticDeduplicator + 3. End-to-end preprocess-script pipeline simulation + 4. Realistic-data tests using fixtures.REALISTIC_SCIENCE_PROMPTS + 5. Backward-compat snapshots — pinned outputs to detect default changes +""" + +from __future__ import annotations + +import json +import random + +import pytest + +ds = pytest.importorskip("datasketch") +faiss = pytest.importorskip("faiss") +import numpy as np + +from sdg.preprocessing.dedupe.fixtures import ( + EXACT_DUP_RANGE, + FORMATTING_RANGE, + NEAR_DUP_RANGE, + REALISTIC_SCIENCE_PROMPTS, + SINGLETON_RANGE, +) +from sdg.preprocessing.dedupe.minhash import MinHashDeduplicator +from sdg.preprocessing.dedupe.semantic import SemanticDeduplicator + + +def _normalize(v: np.ndarray) -> np.ndarray: + norms = np.linalg.norm(v, axis=1, keepdims=True) + return (v / np.clip(norms, 1e-12, None)).astype(np.float32) + + +def _make_records(prompts: list[str]) -> list[dict]: + """Build records that mirror the preprocess script's schema.""" + return [ + { + "conversations": [ + {"from": "human", "value": p}, + {"from": "gpt", "value": "x" * (50 + i * 10)}, + ], + "source": "test", + } + for i, p in enumerate(prompts) + ] + + +# ════════════════════════════════════════════════════════════════════════════ +# Section 1: MinHash invariants +# ════════════════════════════════════════════════════════════════════════════ + +def _sample_minhash_inputs() -> list[str]: + rng = random.Random(7) + base_vocab = "alpha beta gamma delta epsilon zeta eta theta iota kappa lambda mu nu xi pi rho sigma tau".split() + distinct = [" ".join(rng.sample(base_vocab, 12)) + f" id{i}" for i in range(40)] + duplicates = distinct[:5] * 3 # 15 dup copies of the first 5 + return distinct + duplicates + + +def test_minhash_keep_indices_are_sorted_unique_in_range(): + inputs = _sample_minhash_inputs() + n = len(inputs) + keep = MinHashDeduplicator(threshold=0.8).dedup(inputs, show_progress=False) + assert keep == sorted(keep) + assert len(set(keep)) == len(keep) + assert all(0 <= i < n for i in keep) + assert len(keep) <= n + + +def test_minhash_keep_indices_are_python_ints_json_serializable(): + inputs = _sample_minhash_inputs() + keep = MinHashDeduplicator(threshold=0.8).dedup(inputs, show_progress=False) + assert all(type(i) is int for i in keep) + json.dumps(keep) # must not raise + + +def test_minhash_idempotent_on_survivors(): + """Running dedup on the survivors of a previous dedup is a no-op.""" + inputs = _sample_minhash_inputs() + dedup = MinHashDeduplicator(threshold=0.8) + keep1 = dedup.dedup(inputs, show_progress=False) + survivors = [inputs[i] for i in keep1] + keep2 = dedup.dedup(survivors, show_progress=False) + assert keep2 == list(range(len(survivors))) + + +def test_minhash_permutation_invariance_keep_count(): + """Shuffling input doesn't change the number of survivors.""" + inputs = _sample_minhash_inputs() + dedup = MinHashDeduplicator(threshold=0.8) + keep_a = dedup.dedup(inputs, show_progress=False) + + rng = random.Random(99) + perm = list(range(len(inputs))) + rng.shuffle(perm) + shuffled = [inputs[p] for p in perm] + keep_b = dedup.dedup(shuffled, show_progress=False) + + assert len(keep_a) == len(keep_b) + + +def test_minhash_appending_duplicates_does_not_increase_kept_count(): + """Adding more byte-identical copies of existing items can only keep len(keep) the same.""" + inputs = _sample_minhash_inputs() + dedup = MinHashDeduplicator(threshold=0.8) + keep_before = dedup.dedup(inputs, show_progress=False) + kept_texts_before = {inputs[i] for i in keep_before} + + # Append more dups of items that already exist + augmented = inputs + inputs[:10] + inputs[:10] + keep_after = dedup.dedup(augmented, show_progress=False) + assert len(keep_after) == len(keep_before) + # The set of kept *texts* should be identical (representatives may have shifted by index) + kept_texts_after = {augmented[i] for i in keep_after} + assert kept_texts_after == kept_texts_before + + +# ════════════════════════════════════════════════════════════════════════════ +# Section 2: Semantic invariants +# ════════════════════════════════════════════════════════════════════════════ + +def _sample_embeddings(n: int = 50, dim: int = 64, seed: int = 7) -> np.ndarray: + rng = np.random.default_rng(seed) + distinct = rng.standard_normal((n, dim)).astype(np.float32) + # Plant 5 dup-clusters of 4 each + seeds = rng.standard_normal((5, dim)).astype(np.float32) + planted = np.repeat(seeds, 4, axis=0) + return _normalize(np.concatenate([distinct, planted], axis=0)) + + +def test_semantic_keep_indices_are_sorted_unique_in_range(): + emb = _sample_embeddings() + n = len(emb) + keep = SemanticDeduplicator(threshold=0.95).dedup_from_embeddings(emb) + assert keep == sorted(keep) + assert len(set(keep)) == len(keep) + assert all(0 <= i < n for i in keep) + assert len(keep) <= n + + +def test_semantic_keep_indices_are_python_ints_json_serializable(): + emb = _sample_embeddings() + keep = SemanticDeduplicator(threshold=0.95).dedup_from_embeddings(emb) + assert all(type(i) is int for i in keep) + json.dumps(keep) + + +def test_semantic_idempotent_on_survivors(): + emb = _sample_embeddings() + dedup = SemanticDeduplicator(threshold=0.95, topk=20) + keep1 = dedup.dedup_from_embeddings(emb) + survivors = emb[keep1] + keep2 = dedup.dedup_from_embeddings(survivors) + assert keep2 == list(range(len(survivors))) + + +def test_semantic_permutation_invariance_keep_count(): + emb = _sample_embeddings() + dedup = SemanticDeduplicator(threshold=0.95, topk=20) + keep_a = dedup.dedup_from_embeddings(emb) + + rng = np.random.default_rng(42) + perm = rng.permutation(len(emb)) + shuffled = emb[perm] + keep_b = dedup.dedup_from_embeddings(shuffled) + + assert len(keep_a) == len(keep_b) + + +def test_semantic_doubling_planted_cluster_does_not_increase_kept_count(): + """Adding more copies of an existing planted cluster can only keep len(keep) the same.""" + emb = _sample_embeddings() + dedup = SemanticDeduplicator(threshold=0.95, topk=30) + keep_before = dedup.dedup_from_embeddings(emb) + + # Duplicate the last 4 rows (one of the planted clusters) three more times + extra = np.tile(emb[-4:], (3, 1)) + augmented = np.concatenate([emb, extra], axis=0) + keep_after = dedup.dedup_from_embeddings(augmented) + assert len(keep_after) == len(keep_before) + + +# ════════════════════════════════════════════════════════════════════════════ +# Section 3: End-to-end preprocess-script pipeline simulation +# ════════════════════════════════════════════════════════════════════════════ + +def test_pipeline_exact_then_minhash_matches_preprocess_script_shape(): + """Exact reservoir dedup + MinHash with longest-response key_fn — same shape + as preprocess_nemotron_cascade_science.py.""" + records_input = _make_records(REALISTIC_SCIENCE_PROMPTS) + + # Stage 1: Exact dedup (mimic the reservoir loop's "first seen wins" net effect) + seen: dict[str, dict] = {} + for rec in records_input: + prompt = rec["conversations"][0]["value"] + if prompt not in seen: + seen[prompt] = rec + records = list(seen.values()) + assert len(records) == 14 # 3 byte-identical -> 1; everything else distinct strings + + # Stage 2: MinHash with longest-response key_fn + prompts = [r["conversations"][0]["value"] for r in records] + rep_key = lambda i: -len(records[i]["conversations"][1]["value"]) + keep = MinHashDeduplicator(threshold=0.8).dedup(prompts, key_fn=rep_key, show_progress=False) + records = [records[i] for i in keep] + assert len(records) == 10 # formatting cluster -> 1; near-dup cluster -> 1 + + +def test_pipeline_exact_then_minhash_then_semantic_with_orthogonal_embeddings(): + """Three-stage pipeline. Use orthogonal (random high-dim) embeddings so the + semantic stage shouldn't merge anything — verifies the wiring, not the encoder.""" + records_input = _make_records(REALISTIC_SCIENCE_PROMPTS) + + seen: dict[str, dict] = {} + for rec in records_input: + seen.setdefault(rec["conversations"][0]["value"], rec) + records = list(seen.values()) + + prompts = [r["conversations"][0]["value"] for r in records] + rep_key = lambda i: -len(records[i]["conversations"][1]["value"]) + keep = MinHashDeduplicator(threshold=0.8).dedup(prompts, key_fn=rep_key, show_progress=False) + records = [records[i] for i in keep] + n_after_minhash = len(records) + + # Stage 3: Semantic with random high-dim orthogonal embeddings + rng = np.random.default_rng(0) + mock_emb = _normalize(rng.standard_normal((len(records), 384)).astype(np.float32)) + rep_key = lambda i: -len(records[i]["conversations"][1]["value"]) + keep = SemanticDeduplicator(threshold=0.92).dedup_from_embeddings(mock_emb, key_fn=rep_key) + records = [records[i] for i in keep] + + assert len(records) == n_after_minhash # random embeddings don't merge + + +def test_pipeline_longest_response_representative_actually_wins(): + """Verify the longest-response key_fn picks the right rep within a cluster.""" + # Build 3 byte-identical prompts but with response lengths [10, 1000, 100] + prompt = "Same prompt text used three times to form a cluster of duplicates." + records = [ + {"conversations": [{"from": "human", "value": prompt}, {"from": "gpt", "value": "x" * 10}], "source": "t"}, + {"conversations": [{"from": "human", "value": prompt}, {"from": "gpt", "value": "x" * 1000}], "source": "t"}, + {"conversations": [{"from": "human", "value": prompt}, {"from": "gpt", "value": "x" * 100}], "source": "t"}, + ] + prompts = [r["conversations"][0]["value"] for r in records] + rep_key = lambda i: -len(records[i]["conversations"][1]["value"]) + keep = MinHashDeduplicator(threshold=0.8).dedup(prompts, key_fn=rep_key, show_progress=False) + assert keep == [1] # index 1 has the longest response + + +# ════════════════════════════════════════════════════════════════════════════ +# Section 4: Realistic-data tests +# ════════════════════════════════════════════════════════════════════════════ + +def test_minhash_on_realistic_prompts_collapses_classes_correctly(): + """When given all 16 realistic prompts directly (no exact-dedup pre-step), + MinHash should collapse: byte-identical (3->1), formatting (3->1), + near-dup (3->1), and leave 7 singletons. Total kept = 10.""" + keep = MinHashDeduplicator(threshold=0.8).dedup( + REALISTIC_SCIENCE_PROMPTS, show_progress=False + ) + assert len(keep) == 10 + + +def test_minhash_keeps_one_representative_from_each_realistic_cluster(): + """Each cluster (A/B/C) contributes exactly one representative to `keep`.""" + keep = MinHashDeduplicator(threshold=0.8).dedup( + REALISTIC_SCIENCE_PROMPTS, show_progress=False + ) + keep_set = set(keep) + # Class A (byte-identical, indices 0-2): exactly 1 kept + assert sum(1 for i in EXACT_DUP_RANGE if i in keep_set) == 1 + # Class B (formatting variants, indices 3-5): exactly 1 kept + assert sum(1 for i in FORMATTING_RANGE if i in keep_set) == 1 + # Class C (near-dup paraphrases, indices 6-8): exactly 1 kept + assert sum(1 for i in NEAR_DUP_RANGE if i in keep_set) == 1 + # Class D (singletons): all 7 kept + assert sum(1 for i in SINGLETON_RANGE if i in keep_set) == 7 + + +# ════════════════════════════════════════════════════════════════════════════ +# Section 5: Backward-compat snapshots +# ════════════════════════════════════════════════════════════════════════════ +# +# These tests pin the EXACT output for known input + seed. If anything in the +# default behavior changes (shingling, normalization, LSH band selection, FAISS +# search defaults, representative-tie-breaking), one of these will break and +# you'll know immediately. +# +# To intentionally update the snapshot, run the test, copy the new output into +# the assertion. Treat this as a deliberate behavioral change requiring review. + +def test_minhash_snapshot_default_behavior_seed42(): + """Pinned: defaults (threshold=0.8, num_perm=128, 5-word shingles, normalize=True, seed=42). + + The "fox" text is ~120 tokens so a single-word edit (dawn->dusk) gives + Jaccard ≈ 0.92, robustly above threshold for LSH detection. + """ + fox = ( + "the quick brown fox jumps over the lazy dog at sunrise on a cool " + "morning before dawn arrives in the misty meadow where the rabbits " + "nest among the tall grasses by the gentle stream that flows down " + "from the distant hills covered in spring wildflowers and where " + "songbirds sing as the sun rises over the eastern ridge bringing " + "warm light to the valley below where farmers begin their daily work " + "tending fields of corn and wheat that stretch toward the horizon " + "under a clear blue sky filled with drifting white clouds today" + ) + different = ( + "completely unrelated content here today tomorrow next week and the " + "month after that with absolutely no shared vocabulary at all and " + "no overlapping themes or ideas to be found in this paragraph that " + "talks about something entirely separate from foxes and meadows or " + "anything natural touching instead on abstract concepts in pure " + "mathematics like topology category theory and algebraic geometry " + "as well as functional programming languages such as Haskell OCaml " + "and Standard ML which influence modern type systems profoundly today" + ) + texts = [ + fox, + fox, # exact dup + different, + fox.replace("dawn", "dusk", 1), # 1-word edit (dawn->dusk) + ] + keep = MinHashDeduplicator(seed=42).dedup(texts, show_progress=False) + assert keep == [0, 2] + + +def test_semantic_snapshot_default_behavior(): + """Pinned: defaults (threshold=0.92, topk=10) on handcrafted normalized embeddings.""" + emb = _normalize( + np.array( + [ + [1.0, 0.0, 0.0, 0.0], + [1.0, 0.0, 0.0, 0.0], # identical to 0 + [0.0, 1.0, 0.0, 0.0], # orthogonal to 0 + [np.cos(np.radians(15)), np.sin(np.radians(15)), 0.0, 0.0], # cos=0.966 with 0 + ], + dtype=np.float32, + ) + ) + keep = SemanticDeduplicator().dedup_from_embeddings(emb) + # Cluster {0, 1, 3} via cos>=0.92; singleton {2} + assert keep == [0, 2] diff --git a/sdg/preprocessing/dedupe/test_large.py b/sdg/preprocessing/dedupe/test_large.py new file mode 100644 index 0000000000000000000000000000000000000000..347bd7411a35d251fbc0c00c122e295b6be01911 --- /dev/null +++ b/sdg/preprocessing/dedupe/test_large.py @@ -0,0 +1,179 @@ +"""Large-text and large-N correctness tests. + +Verifies the deduplicators handle realistic-scale inputs without false +positives (over-merging distinct items) or false negatives (missing planted +duplicates). + +Sizes are kept modest enough to run inside pytest in <30s on M-series. +For pure throughput numbers see test_benchmarks.py. +""" + +from __future__ import annotations + +import random + +import pytest + +ds = pytest.importorskip("datasketch") +faiss = pytest.importorskip("faiss") +import numpy as np + +from sdg.preprocessing.dedupe.minhash import MinHashDeduplicator +from sdg.preprocessing.dedupe.semantic import SemanticDeduplicator + + +# ──────────────────────────────────────────────────────────────────────────── +# Synthetic prompt generators +# ──────────────────────────────────────────────────────────────────────────── + +_VOCAB = ( + "alpha beta gamma delta epsilon zeta eta theta iota kappa lambda mu nu xi " + "omicron pi rho sigma tau upsilon phi chi psi omega the quick brown fox " + "jumps over lazy dog cat house tree mountain river ocean sun moon star " + "planet galaxy universe science math physics chemistry biology equation " + "matrix vector function derivative integral hypothesis theorem proof " + "experiment observation analysis synthesis conclusion problem solution " + "method approach strategy algorithm structure pattern model framework" +).split() + + +def _distinct_prompts(n: int, words_per_prompt: int = 25, seed: int = 42) -> list[str]: + """Generate n distinct random-word prompts (very low cross-similarity).""" + rng = random.Random(seed) + return [ + " ".join(rng.sample(_VOCAB, words_per_prompt)) + f" uniqueid{i}" + for i in range(n) + ] + + +def _planted_dup_prompts( + n_distinct: int, dup_factor: int, words_per_prompt: int = 25, seed: int = 42 +) -> tuple[list[str], int]: + """Generate n_distinct base prompts, each duplicated dup_factor times. + + Returns (texts, expected_unique). Total len(texts) == n_distinct * dup_factor. + """ + base = _distinct_prompts(n_distinct, words_per_prompt, seed) + texts: list[str] = [] + for prompt in base: + for _ in range(dup_factor): + texts.append(prompt) + rng = random.Random(seed + 1) + rng.shuffle(texts) + return texts, n_distinct + + +# ──────────────────────────────────────────────────────────────────────────── +# MinHash: large-N correctness +# ──────────────────────────────────────────────────────────────────────────── + +def test_minhash_5k_distinct_no_false_positives(): + texts = _distinct_prompts(5_000) + keep = MinHashDeduplicator(threshold=0.8).dedup(texts, show_progress=False) + # Up to a tiny LSH false-positive rate; expect >= 99% kept. + assert len(keep) >= 4_950, f"Too many false-positive merges: kept {len(keep)}" + + +def test_minhash_5k_with_planted_clusters_collapses_correctly(): + texts, n_unique = _planted_dup_prompts(n_distinct=1_000, dup_factor=5) + keep = MinHashDeduplicator(threshold=0.8).dedup(texts, show_progress=False) + # Allow tiny LSH miss tolerance (rare) -> within +/- 1% of n_unique + assert n_unique - 10 <= len(keep) <= n_unique + 50 + + +def test_minhash_long_string_100k_chars_does_not_crash(): + """A single >100K-char prompt should be processed without error.""" + rng = random.Random(42) + # 20K random tokens => ~120K chars, with many distinct 5-grams + tokens = [rng.choice(_VOCAB) for _ in range(20_000)] + long_a = " ".join(tokens) + long_b = long_a + " trailing differentiator words here" + keep = MinHashDeduplicator(threshold=0.85).dedup([long_a, long_b], show_progress=False) + # 4 extra tokens out of 20K -> Jaccard ~1.0 -> should collapse + assert len(keep) == 1 + + +def test_minhash_mixed_short_and_long_prompts(): + short = ["short prompt one alpha", "short prompt two beta", "short prompt three gamma"] + longs = _distinct_prompts(10, words_per_prompt=40, seed=99) + texts = short + longs + short # short prompts duplicated, longs distinct + keep = MinHashDeduplicator(threshold=0.8).dedup(texts, show_progress=False) + # 3 unique short + 10 unique long = 13 + assert 12 <= len(keep) <= 14 + + +# ──────────────────────────────────────────────────────────────────────────── +# Semantic: large-N correctness (using random embeddings to avoid model load) +# ──────────────────────────────────────────────────────────────────────────── + +def _normalize(v: np.ndarray) -> np.ndarray: + norms = np.linalg.norm(v, axis=1, keepdims=True) + return (v / np.clip(norms, 1e-12, None)).astype(np.float32) + + +def test_semantic_5k_random_embeddings_no_false_positives(): + rng = np.random.default_rng(42) + emb = _normalize(rng.standard_normal((5_000, 384)).astype(np.float32)) + dedup = SemanticDeduplicator(threshold=0.95, hnsw_threshold=1_000) + keep = dedup.dedup_from_embeddings(emb) + # 384-dim random gaussians at 0.95 threshold should produce nearly no merges. + assert len(keep) >= 4_995 + + +def test_semantic_10k_with_planted_clusters_finds_them(): + rng = np.random.default_rng(123) + distractors = rng.standard_normal((9_000, 384)).astype(np.float32) + # 100 planted clusters of 10 each + seeds = rng.standard_normal((100, 384)).astype(np.float32) + planted = np.repeat(seeds, 10, axis=0) + raw = np.concatenate([distractors, planted], axis=0) + emb = _normalize(raw) + dedup = SemanticDeduplicator(threshold=0.95, hnsw_threshold=1_000, topk=15) + keep = dedup.dedup_from_embeddings(emb) + # 9_000 distractors + 100 cluster reps = 9_100; allow small approximation tolerance + assert 9_080 <= len(keep) <= 9_120 + + +def test_semantic_one_giant_cluster_plus_singletons(): + rng = np.random.default_rng(7) + distractors = rng.standard_normal((500, 128)).astype(np.float32) + seed_vec = rng.standard_normal((1, 128)).astype(np.float32) + giant = np.repeat(seed_vec, 500, axis=0) # 500 identical + raw = np.concatenate([distractors, giant], axis=0) + emb = _normalize(raw) + dedup = SemanticDeduplicator(threshold=0.95, hnsw_threshold=100, topk=600) + keep = dedup.dedup_from_embeddings(emb) + # 500 distractors + 1 representative + assert len(keep) == 501 + + +# ──────────────────────────────────────────────────────────────────────────── +# End-to-end pipeline: MinHash -> Semantic stack on synthetic data +# ──────────────────────────────────────────────────────────────────────────── + +def test_pipeline_minhash_then_semantic_on_synthetic_records(): + """Two-stage pipeline behaves like SDG preprocess script. + + Build records, run MinHash on prompts, then semantic dedup on the survivors + (using random embeddings as a stand-in for real encoding). Verify the + record list shrinks monotonically and final count is plausible. + """ + texts, _ = _planted_dup_prompts(n_distinct=500, dup_factor=3) + records = [{"prompt": t, "response": f"resp for {i}"} for i, t in enumerate(texts)] + + n_initial = len(records) + prompts = [r["prompt"] for r in records] + + keep_mh = MinHashDeduplicator(threshold=0.8).dedup(prompts, show_progress=False) + records = [records[i] for i in keep_mh] + assert len(records) < n_initial + assert 490 <= len(records) <= 530 # ~500 unique + + # Use random embeddings sized to surviving records + rng = np.random.default_rng(0) + emb = _normalize(rng.standard_normal((len(records), 384)).astype(np.float32)) + dedup = SemanticDeduplicator(threshold=0.95) + keep_sem = dedup.dedup_from_embeddings(emb) + records = [records[i] for i in keep_sem] + # Random embeddings are nearly orthogonal -> semantic stage shouldn't merge much. + assert len(records) >= 480 diff --git a/sdg/preprocessing/dedupe/test_minhash.py b/sdg/preprocessing/dedupe/test_minhash.py new file mode 100644 index 0000000000000000000000000000000000000000..58e8f785edb5d55171307faff4cb657b5587c3a3 --- /dev/null +++ b/sdg/preprocessing/dedupe/test_minhash.py @@ -0,0 +1,180 @@ +"""Tests for MinHashDeduplicator. Skipped if datasketch is not installed.""" + +from __future__ import annotations + +import pytest + +pytest.importorskip("datasketch") + +from sdg.preprocessing.dedupe.minhash import MinHashDeduplicator + + +# ── Boundary cases ─────────────────────────────────────────────────────────── + +def test_empty_input_returns_empty(): + assert MinHashDeduplicator().dedup([], show_progress=False) == [] + + +def test_single_item_returns_single(): + assert MinHashDeduplicator().dedup(["only one"], show_progress=False) == [0] + + +# ── Identity / disjoint ────────────────────────────────────────────────────── + +def test_two_identical_items_collapse_to_one(): + text = "Solve the equation x squared minus four equals zero showing every step" + keep = MinHashDeduplicator().dedup([text, text], show_progress=False) + assert len(keep) == 1 + + +def test_two_completely_different_items_both_kept(): + keep = MinHashDeduplicator().dedup( + [ + "describe the lifecycle of a butterfly in detail", + "compute the integral of sine x from zero to pi", + ], + show_progress=False, + ) + assert len(keep) == 2 + + +# ── Threshold behavior ─────────────────────────────────────────────────────── + +def test_near_duplicate_above_threshold_collapses(): + # Tail-word edit only (one shingle differs) -> exact Jaccard ~0.93, safely + # above the 0.7 threshold even after exact-Jaccard candidate verification. + base = ( + "Solve the equation x squared minus four equals zero. Show all of your work " + "step by step and explain your reasoning so that a beginner can follow it easily today." + ) + near = ( + "Solve the equation x squared minus four equals zero. Show all of your work " + "step by step and explain your reasoning so that a beginner can follow it easily now." + ) + keep = MinHashDeduplicator(threshold=0.7).dedup([base, near], show_progress=False) + assert len(keep) == 1 + + +def test_distinct_items_below_threshold_both_kept(): + keep = MinHashDeduplicator(threshold=0.95).dedup( + [ + "explain how photosynthesis works in plants", + "describe the process of cellular respiration in animals", + ], + show_progress=False, + ) + assert len(keep) == 2 + + +# ── Cluster of three ───────────────────────────────────────────────────────── + +def test_cluster_of_three_collapses_to_one(): + base = "the quick brown fox jumps over the lazy dog under a bright moon at night" + near1 = base + " quietly" + near2 = base + " softly" + keep = MinHashDeduplicator(threshold=0.7).dedup([base, near1, near2], show_progress=False) + assert len(keep) == 1 + + +# ── LSH candidate verification (sub-threshold false positives) ─────────────── + +def test_distinct_prompts_sharing_long_boilerplate_block_not_merged(): + """Distinct questions that share a long fixed boilerplate block sit just + below the threshold (~0.65) pairwise. LSH may surface them as candidates; + without verifying the candidate's actual Jaccard, union-find transitively + collapses them. Each distinct question must survive as its own record. + """ + suffix = ( + "\n\nA) If both assertion and reason are true and reason is the correct " + "explanation of assertion\n" + "B) If both assertion and reason are true but reason is not the correct " + "explanation of assertion\n" + "C) If assertion is true but reason is false\n" + "D) If both assertion and reason are false" + ) + bodies = [ + "F is more electronegative than Cl F has high electron affinity than Cl.", + "Leucoplasts perform photosynthesis Chloroplasts store fats, starch and proteins", + "CN ion is an ambient nucleophile. Nucleophiles are electron rich species.", + "Fructose is a reducing sugar. It has a ketonic group.", + "The lactic acid shows the geometrical isomerism. Lactic acid has a double bond.", + ] + prompts = [b + suffix for b in bodies] + keep = MinHashDeduplicator(threshold=0.8).dedup(prompts, show_progress=False) + assert len(keep) == len(prompts) + + +def test_true_duplicates_still_merge_after_verification(): + """Verification must not suppress genuine near-duplicates: same question, + trivial rewording, sharing the same boilerplate block -> still one record. + """ + suffix = ( + "\n\nAnswer Choices:\n(A) 0.05 uC\n(B) 0.5 uC\n(C) 5 uC\n(D) 50 uC\n" + "(E) 500 uC\n(F) 5000 uC\n(G) 0.005 uC\n(H) 0.0005 uC\n(I) 50 uC\n(J) 5 uC" + ) + # Same question, one-word rewording (the->this) -> exact Jaccard ~0.83, + # above 0.8 even after verification; the shared boilerplate doesn't prevent it. + base = "A capacitor with a capacitance of 5 microfarads is charged to a voltage of 10 volts. What is the charge stored on the capacitor?" + near = "A capacitor with a capacitance of 5 microfarads is charged to a voltage of 10 volts. What is the charge stored on this capacitor?" + keep = MinHashDeduplicator(threshold=0.8).dedup([base + suffix, near + suffix], show_progress=False) + assert len(keep) == 1 + + +# ── Representative selection via key_fn ────────────────────────────────────── + +def test_key_fn_selects_longest_response_proxy(): + text = "Solve the equation x squared minus four equals zero showing every step" + response_lengths = [10, 500, 50] + key_fn = lambda i: -response_lengths[i] # longest wins + keep = MinHashDeduplicator().dedup([text, text, text], key_fn=key_fn, show_progress=False) + assert keep == [1] + + +# ── Normalization ──────────────────────────────────────────────────────────── + +def test_normalization_collapses_case_and_whitespace_variants(): + a = "Solve the equation x squared minus four equals zero showing every step" + b = " solve THE equation x squared minus four equals zero showing every step " + keep = MinHashDeduplicator(normalize=True, threshold=0.7).dedup([a, b], show_progress=False) + assert len(keep) == 1 + + +def test_normalization_off_keeps_case_variants(): + """Without normalization, formatting variants may not cross threshold.""" + a = "ALPHA BETA GAMMA DELTA EPSILON" + b = "alpha beta gamma delta epsilon" + # With small inputs and high threshold, expect both kept when normalize=False + # (different shingles -> Jaccard 0) + dedup = MinHashDeduplicator(normalize=False, threshold=0.9, shingle_size=2) + keep = dedup.dedup([a, b], show_progress=False) + assert len(keep) == 2 + + +# ── Shingle kinds ──────────────────────────────────────────────────────────── + +def test_char_shingles_work(): + a = "abcdefghijklmnopqrstuvwxyz" * 3 + b = "abcdefghijklmnopqrstuvwxyz" * 3 + "ZZZ" + keep = MinHashDeduplicator(shingle_kind="char", shingle_size=8, threshold=0.7).dedup( + [a, b], show_progress=False + ) + assert len(keep) == 1 + + +def test_invalid_shingle_kind_raises(): + with pytest.raises(ValueError): + MinHashDeduplicator(shingle_kind="bogus") + + +def test_invalid_threshold_raises(): + with pytest.raises(ValueError): + MinHashDeduplicator(threshold=1.5) + + +# ── Determinism ────────────────────────────────────────────────────────────── + +def test_determinism_same_seed_same_result(): + texts = ["alpha beta gamma delta epsilon"] * 5 + ["unrelated content here please"] + a = MinHashDeduplicator(seed=7).dedup(texts, show_progress=False) + b = MinHashDeduplicator(seed=7).dedup(texts, show_progress=False) + assert a == b diff --git a/sdg/preprocessing/dedupe/test_semantic.py b/sdg/preprocessing/dedupe/test_semantic.py new file mode 100644 index 0000000000000000000000000000000000000000..6796262dcac9edad91ee55fb392247d08b46ec30 --- /dev/null +++ b/sdg/preprocessing/dedupe/test_semantic.py @@ -0,0 +1,156 @@ +"""Tests for SemanticDeduplicator. + +Avoids downloading any embedding model by exercising `dedup_from_embeddings` +directly with handcrafted L2-normalized vectors. Skipped if faiss is not +installed. +""" + +from __future__ import annotations + +import pytest + +pytest.importorskip("faiss") +import numpy as np + +from sdg.preprocessing.dedupe.semantic import SemanticDeduplicator + + +def _normalize(v: np.ndarray) -> np.ndarray: + norms = np.linalg.norm(v, axis=1, keepdims=True) + return (v / np.clip(norms, 1e-12, None)).astype(np.float32) + + +def _vec(*components: float) -> np.ndarray: + return np.array(components, dtype=np.float32) + + +# ── Boundary cases ─────────────────────────────────────────────────────────── + +def test_empty_returns_empty(): + dedup = SemanticDeduplicator() + assert dedup.dedup_from_embeddings(np.zeros((0, 4), dtype=np.float32)) == [] + + +def test_single_returns_single(): + dedup = SemanticDeduplicator() + emb = _normalize(np.array([[1.0, 0.0, 0.0, 0.0]])) + assert dedup.dedup_from_embeddings(emb) == [0] + + +# ── Identity / orthogonality ───────────────────────────────────────────────── + +def test_identical_embeddings_collapse(): + dedup = SemanticDeduplicator(threshold=0.9) + emb = _normalize(np.stack([_vec(1, 0, 0), _vec(1, 0, 0)])) + keep = dedup.dedup_from_embeddings(emb) + assert len(keep) == 1 + + +def test_orthogonal_embeddings_both_kept(): + dedup = SemanticDeduplicator(threshold=0.5) + emb = _normalize(np.stack([_vec(1, 0, 0), _vec(0, 1, 0)])) + keep = dedup.dedup_from_embeddings(emb) + assert keep == [0, 1] + + +# ── Threshold behavior ─────────────────────────────────────────────────────── + +def test_above_threshold_collapses(): + # cos(a, b) ≈ 0.95 (small angle) + a = _vec(1.0, 0.0) + b = _vec(np.cos(np.radians(18)), np.sin(np.radians(18))) # ~0.951 + emb = _normalize(np.stack([a, b])) + dedup = SemanticDeduplicator(threshold=0.92) + keep = dedup.dedup_from_embeddings(emb) + assert len(keep) == 1 + + +def test_below_threshold_both_kept(): + # cos(a, b) ≈ 0.866 (30 degrees) + a = _vec(1.0, 0.0) + b = _vec(np.cos(np.radians(30)), np.sin(np.radians(30))) + emb = _normalize(np.stack([a, b])) + dedup = SemanticDeduplicator(threshold=0.92) + keep = dedup.dedup_from_embeddings(emb) + assert keep == [0, 1] + + +# ── Cluster of three ───────────────────────────────────────────────────────── + +def test_cluster_of_three_via_chain_above_threshold(): + a = _vec(1.0, 0.0) + b = _vec(np.cos(np.radians(10)), np.sin(np.radians(10))) + c = _vec(np.cos(np.radians(20)), np.sin(np.radians(20))) + emb = _normalize(np.stack([a, b, c])) + dedup = SemanticDeduplicator(threshold=0.95) + keep = dedup.dedup_from_embeddings(emb) + # All three angles within ~20 degrees -> all pairwise cos > 0.93, + # so they form one cluster -> one kept. + assert len(keep) == 1 + + +# ── Representative selection ───────────────────────────────────────────────── + +def test_key_fn_selects_representative(): + emb = _normalize(np.stack([_vec(1, 0), _vec(1, 0), _vec(1, 0)])) + response_lengths = [10, 500, 50] + key_fn = lambda i: -response_lengths[i] + dedup = SemanticDeduplicator(threshold=0.99) + keep = dedup.dedup_from_embeddings(emb, key_fn=key_fn) + assert keep == [1] + + +# ── HNSW path ──────────────────────────────────────────────────────────────── + +def test_hnsw_path_runs_above_threshold_size(): + """Force the HNSW code path with a tiny hnsw_threshold. + + Uses 384-dim random gaussians (matching our default embed model). At that + dimensionality, random pairs have expected cosine ~0 with std ~1/sqrt(384), + so a threshold of 0.95 is deep in the tail and nothing should collapse. + """ + rng = np.random.default_rng(42) + raw = rng.standard_normal((200, 384)).astype(np.float32) + emb = _normalize(raw) + dedup = SemanticDeduplicator(threshold=0.95, hnsw_threshold=10) + keep = dedup.dedup_from_embeddings(emb) + assert len(keep) == 200 + + +def test_hnsw_path_collapses_planted_duplicates(): + """HNSW should still find planted near-duplicates among many distractors.""" + rng = np.random.default_rng(123) + distractors = rng.standard_normal((180, 384)).astype(np.float32) + # Plant 20 copies of one vector + seed_vec = rng.standard_normal((1, 384)).astype(np.float32) + planted = np.repeat(seed_vec, 20, axis=0) + raw = np.concatenate([distractors, planted], axis=0) + emb = _normalize(raw) + dedup = SemanticDeduplicator(threshold=0.95, hnsw_threshold=10, topk=25) + keep = dedup.dedup_from_embeddings(emb) + # 180 distractors + 1 representative of the planted cluster = 181 + assert len(keep) == 181 + + +# ── Validation ─────────────────────────────────────────────────────────────── + +def test_invalid_threshold_raises(): + with pytest.raises(ValueError): + SemanticDeduplicator(threshold=0.0) + + +def test_invalid_topk_raises(): + with pytest.raises(ValueError): + SemanticDeduplicator(topk=0) + + +# ── Device resolution ──────────────────────────────────────────────────────── + +def test_device_explicit_passthrough(): + assert SemanticDeduplicator(device="cpu").resolve_device() == "cpu" + + +def test_device_auto_returns_known_string(): + """auto must resolve to one of mps/cuda/cpu (depending on host).""" + resolved = SemanticDeduplicator(device="auto").resolve_device() + assert resolved in {"mps", "cuda", "cpu"} diff --git a/sdg/preprocessing/dedupe/test_templates.py b/sdg/preprocessing/dedupe/test_templates.py new file mode 100644 index 0000000000000000000000000000000000000000..bf57334442deb5f8e6825728b8079591b119a6b4 --- /dev/null +++ b/sdg/preprocessing/dedupe/test_templates.py @@ -0,0 +1,149 @@ +"""Tests for template auto-detection and stripping (no external deps).""" + +from __future__ import annotations + +from sdg.preprocessing.dedupe.templates import ( + find_common_prefixes, + strip_template, + template_hit_stats, +) + + +# ── find_common_prefixes ───────────────────────────────────────────────────── + +def test_find_common_prefixes_empty_input(): + assert find_common_prefixes([], min_count=2) == [] + + +def test_find_common_prefixes_returns_empty_when_no_common_prefix(): + texts = ["alpha question one", "beta question two", "gamma question three"] + assert find_common_prefixes(texts, min_count=2, prefix_lengths=(5, 10)) == [] + + +def test_find_common_prefixes_detects_simple_repeated_prefix(): + template = "Answer the question: " + texts = [template + f"unique body number {i}" for i in range(20)] + found = find_common_prefixes(texts, min_count=10, prefix_lengths=(len(template),)) + assert template in found + + +def test_find_common_prefixes_respects_min_count(): + rare = "Rare prefix: " + common = "Common prefix: " + texts = ( + [rare + f"a{i}" for i in range(5)] + + [common + f"b{i}" for i in range(20)] + ) + found = find_common_prefixes(texts, min_count=10, prefix_lengths=(len(common),)) + assert common in found + assert rare not in found + + +def test_find_common_prefixes_returns_sorted_longest_first(): + long_tmpl = "A" * 100 + " specific" + short_tmpl = "A" * 50 + texts = [long_tmpl + " body " + str(i) for i in range(20)] + found = find_common_prefixes(texts, min_count=10, prefix_lengths=(50, 100, 150)) + assert found == sorted(found, key=len, reverse=True) + + +def test_find_common_prefixes_handles_short_texts(): + """Texts shorter than a probed prefix length must not crash.""" + texts = ["xy"] * 20 # all shorter than every probed length + found = find_common_prefixes(texts, min_count=2, prefix_lengths=(50, 100)) + assert found == [] + + +def test_find_common_prefixes_finds_multiple_distinct_templates(): + t1 = "Template A: " + t2 = "Different template B: " + texts = ( + [t1 + f"x{i}" for i in range(15)] + + [t2 + f"y{i}" for i in range(15)] + + [f"singleton {i}" for i in range(5)] + ) + found = find_common_prefixes(texts, min_count=10, prefix_lengths=(len(t1), len(t2))) + assert t1 in found + assert t2 in found + + +# ── strip_template ─────────────────────────────────────────────────────────── + +def test_strip_template_returns_text_unchanged_with_no_match(): + assert strip_template("hello world", ["FOO: "]) == "hello world" + + +def test_strip_template_returns_text_unchanged_with_empty_templates(): + assert strip_template("hello world", []) == "hello world" + + +def test_strip_template_simple_match(): + assert strip_template("PREFIX: actual content", ["PREFIX: "]) == "actual content" + + +def test_strip_template_strips_leading_whitespace_after_removal(): + assert strip_template("PREFIX: \n\ncontent", ["PREFIX:"]) == "content" + + +def test_strip_template_longest_match_wins_when_sorted_correctly(): + """Templates list sorted longest-first as find_common_prefixes returns.""" + templates = ["PREFIX EXTRA: ", "PREFIX: "] + assert strip_template("PREFIX EXTRA: content", templates) == "content" + assert strip_template("PREFIX: content", templates) == "content" + + +def test_strip_template_does_not_match_substring_in_middle(): + """Template must occur as a *prefix*, not anywhere in the text.""" + assert strip_template("hello PREFIX: world", ["PREFIX: "]) == "hello PREFIX: world" + + +# ── template_hit_stats ─────────────────────────────────────────────────────── + +def test_template_hit_stats_partitions_texts_correctly(): + texts = [ + "MCQ: question 1", + "MCQ: question 2", + "MATH: question 3", + "totally distinct prompt", + ] + templates = ["MCQ: ", "MATH: "] + stats = template_hit_stats(texts, templates) + assert stats["MCQ: "] == 2 + assert stats["MATH: "] == 1 + assert stats[""] == 1 + assert sum(stats.values()) == len(texts) + + +def test_template_hit_stats_uses_longest_match_when_templates_nest(): + """If templates list has nested prefixes, longest-first prevents double counting.""" + texts = ["PREFIX EXTRA: x"] * 10 + ["PREFIX: y"] * 5 + templates = ["PREFIX EXTRA: ", "PREFIX: "] + stats = template_hit_stats(texts, templates) + assert stats["PREFIX EXTRA: "] == 10 + assert stats["PREFIX: "] == 5 + assert stats[""] == 0 + + +def test_template_hit_stats_empty_templates_all_unmatched(): + texts = ["a", "b", "c"] + stats = template_hit_stats(texts, []) + assert stats == {"": 3} + + +# ── End-to-end: detect then strip ──────────────────────────────────────────── + +def test_detect_then_strip_round_trip(): + """Realistic sequence: detect templates, then strip them from the corpus.""" + template = "Answer the following multiple choice question: " + questions = [ + "What is 2+2?", + "What color is the sky?", + "Who wrote Hamlet?", + "What is the capital of France?", + "What is the speed of light?", + ] * 4 # 20 questions + texts = [template + q for q in questions] + templates = find_common_prefixes(texts, min_count=10, prefix_lengths=(len(template),)) + stripped = [strip_template(t, templates) for t in texts] + # Every text should have been stripped down to just its question. + assert all(s in [q for q in questions] for s in stripped) diff --git a/sdg/preprocessing/preprocess_nemotron_cascade_science.py b/sdg/preprocessing/preprocess_nemotron_cascade_science.py new file mode 100644 index 0000000000000000000000000000000000000000..e13f4731daddedab69c1438c41f9831e3ecb77f8 --- /dev/null +++ b/sdg/preprocessing/preprocess_nemotron_cascade_science.py @@ -0,0 +1,726 @@ +""" +Preprocess Nemotron-Cascade-2-SFT-Data (science subset). + +Source: https://huggingface.co/datasets/nvidia/Nemotron-Cascade-2-SFT-Data + +Each row has top-level columns `domain`, `source`, `messages` (list of +{role, content} dicts: system, user, assistant), and `generator`. This script: + - Extracts the user-role message as the prompt. + - Deduplicates by user prompt via reservoir sampling (one random rollout + per unique prompt). + - Keeps the top-level `source` column (e.g. "Nemotron-Cascade-1", + "Nemotron-Science-v1"). + - Emits records in the `openthoughts4` conversations format so SDG configs + can consume the dataset directly. +""" + +from __future__ import annotations + +import argparse +import json +import os +import random +import sys +import time +from pathlib import Path + +# ── HF offline mode ────────────────────────────────────────────────────── +# Must run BEFORE `from datasets import ...` because datasets/huggingface_hub +# read these env vars into module-level constants at import time. Forcing them +# here makes a cached run do ZERO network calls (no revision check that could +# hang/fail when offline). +# --offline : force offline. +# --no-offline : force online (overrides everything). +# Default mode is download-to-cache (NOT streaming), so offline auto-enables +# IFF the dataset is already cached — the first download still needs network, +# and explicit --stream always stays online. A full (non --dry-run) run uploads +# to HF and therefore MUST stay online, so offline is only auto-enabled for dry +# runs (explicit --offline is still honored everywhere). +def _resolve_offline_flags() -> None: + argv = sys.argv[1:] + if "--no-offline" in argv: + return + want_offline = "--offline" in argv + streaming = "--stream" in argv + is_dry_run = "--dry-run" in argv + if not want_offline and is_dry_run and not streaming: + hf_home = os.environ.get("HF_HOME") or os.path.expanduser("~/.cache/huggingface") + # Cache folder name HF derives from "nvidia/Nemotron-Cascade-2-SFT-Data". + cache_dir = Path(hf_home) / "hub" / "datasets--nvidia--Nemotron-Cascade-2-SFT-Data" + want_offline = cache_dir.exists() + if want_offline: + os.environ.setdefault("HF_DATASETS_OFFLINE", "1") + os.environ.setdefault("HF_HUB_OFFLINE", "1") + + +_resolve_offline_flags() + +from datasets import load_dataset +from huggingface_hub import HfApi, create_repo +from tqdm import tqdm + +_REPO_ROOT = Path(__file__).resolve().parents[2] +if str(_REPO_ROOT) not in sys.path: + sys.path.insert(0, str(_REPO_ROOT)) +from sdg.preprocessing.dedupe import ( + MinHashDeduplicator, + SemanticDeduplicator, + find_common_prefixes, + strip_template, + template_hit_stats, +) + + +DATASET_ID = "nvidia/Nemotron-Cascade-2-SFT-Data" +DATASET_SUBSET = "science" +HF_REPO_ID = "teetone/nemotron-cascade-2-science-deduped" +EXPECTED_UNIQUE_PROMPTS: int | None = None + + +def extract_role(messages: list[dict], role: str) -> str | None: + """Return the content of the first message matching `role`, or None.""" + for m in messages: + if m.get("role") == role: + content = m.get("content") + if isinstance(content, str): + return content + return None + + +def _write_dedup_audit_section( + f, + stage_name: str, + threshold: float, + records_before: list[dict], + keep_indices: list[int], + clusters: dict, + max_text_chars: int = 600, +) -> None: + """Append one stage's per-cluster audit info to file handle f. + + `clusters` is the {rep_idx -> [member_idx, ...]} dict returned by + dedup(..., return_clusters=True). We show each non-singleton cluster with + its KEPT representative followed by the DROP'd members so a human can + eyeball *why* the algorithm merged them. + """ + n_in = len(records_before) + n_out = len(keep_indices) + non_singleton = {rep: m for rep, m in clusters.items() if len(m) > 1} + + f.write(f"\n{'=' * 78}\n") + f.write(f"STAGE: {stage_name} (threshold={threshold})\n") + f.write(f" Input: {n_in:,} records\n") + f.write(f" Output: {n_out:,} records\n") + f.write(f" Removed: {n_in - n_out:,} records\n") + f.write(f" Clusters of size >= 2: {len(non_singleton):,}\n") + f.write(f"{'=' * 78}\n\n") + + if not non_singleton: + f.write("(No clusters of size >= 2 — nothing was deduped at this stage.)\n") + return + + # Sort by cluster size descending so the most aggressive merges show first. + sorted_clusters = sorted(non_singleton.items(), key=lambda kv: -len(kv[1])) + for cluster_idx, (rep, members) in enumerate(sorted_clusters, 1): + f.write(f"[CLUSTER {cluster_idx}] size={len(members)}\n") + for idx in members: + rec = records_before[idx] + prompt = rec["conversations"][0]["value"] + response = rec["conversations"][1]["value"] + marker = "KEPT" if idx == rep else "DROP" + if len(prompt) > max_text_chars: + shown = prompt[:max_text_chars] + f"...[+{len(prompt) - max_text_chars} chars]" + else: + shown = prompt + f.write( + f" {marker} [idx={idx:>6} resp_len={len(response):>6,}]\n" + f" {shown!r}\n" + ) + f.write("\n") + + +def preprocess_and_upload( + seed: int = 42, + dry_run: bool = False, + stream: bool = False, + limit: int | None = None, + audit_path: str | None = None, + dump_embeddings: str | None = None, + auto_strip_templates: bool = True, + template_min_count: int = 10, + skip_minhash: bool = False, + minhash_threshold: float = 0.8, + minhash_num_perm: int = 128, + minhash_shingle_size: int = 5, + skip_semantic: bool = False, + semantic_threshold: float = 0.95, + embed_model: str = "BAAI/bge-small-en-v1.5", + embed_batch_size: int = 128, + semantic_topk: int = 10, + device: str = "auto", +) -> None: + random.seed(seed) + + mode = "streaming" if stream else "download-to-cache" + offline = os.environ.get("HF_HUB_OFFLINE") == "1" + print(f"Loading {DATASET_ID} (subset={DATASET_SUBSET}, {mode}, " + f"offline={'on' if offline else 'off'}) ...") + ds = load_dataset(DATASET_ID, DATASET_SUBSET, split="train", streaming=stream) + + prompt_reservoir: dict[str, tuple[dict, int]] = {} + total = 0 + skipped_no_user = 0 + skipped_no_assistant = 0 + + pbar = tqdm(ds, desc="Streaming Nemotron-Cascade-2 (science)", unit=" rows", smoothing=0.05) + for row in pbar: + total += 1 + if limit is not None and total > limit: + total -= 1 # we counted this one but won't process it + break + + messages = row.get("messages") or [] + user_prompt = extract_role(messages, "user") + if not user_prompt: + skipped_no_user += 1 + pbar.set_postfix(unique=len(prompt_reservoir), no_user=skipped_no_user, no_asst=skipped_no_assistant) + continue + + assistant_response = extract_role(messages, "assistant") + if assistant_response is None: + skipped_no_assistant += 1 + pbar.set_postfix(unique=len(prompt_reservoir), no_user=skipped_no_user, no_asst=skipped_no_assistant) + continue + + source = row.get("source") + + rec = { + "conversations": [ + {"from": "human", "value": user_prompt}, + {"from": "gpt", "value": assistant_response}, + ], + "source": source, + } + + if user_prompt not in prompt_reservoir: + prompt_reservoir[user_prompt] = (rec, 1) + else: + _, count = prompt_reservoir[user_prompt] + count += 1 + if random.randint(1, count) == 1: + prompt_reservoir[user_prompt] = (rec, count) + else: + prompt_reservoir[user_prompt] = (prompt_reservoir[user_prompt][0], count) + + pbar.set_postfix(unique=len(prompt_reservoir), no_user=skipped_no_user, no_asst=skipped_no_assistant) + pbar.close() + + print(f" Total rows streamed: {total:,}") + print(f" Skipped (no user message): {skipped_no_user:,}") + print(f" Skipped (no assistant message): {skipped_no_assistant:,}") + print(f" Unique user prompts: {len(prompt_reservoir):,}") + + records = [row for row, _ in prompt_reservoir.values()] + n_after_exact = len(records) + print(f" Exact-deduped rows: {n_after_exact:,}") + + # Resolve audit file path: --limit auto-enables auditing unless --skip-audit-on-limit. + write_audit = audit_path is not None or limit is not None + audit_handle = None + if write_audit: + if audit_path is None: + ts = time.strftime("%Y%m%d_%H%M%S") + audit_path = f"output/logs/dedupe_audit_{ts}.txt" + Path(audit_path).parent.mkdir(parents=True, exist_ok=True) + audit_handle = open(audit_path, "w") + audit_handle.write( + f"Dedup audit for nemotron-cascade-2-science (limit={limit}, dry_run={dry_run})\n" + f"Pipeline: exact reservoir -> MinHash -> Semantic\n" + f"After exact dedup: {n_after_exact:,} unique prompts\n" + ) + audit_handle.write(f"\n{'=' * 78}\n") + audit_handle.write("RUN CONFIG\n") + audit_handle.write(f"{'=' * 78}\n") + audit_handle.write(f" seed = {seed}\n") + audit_handle.write(f" stream = {stream}\n") + audit_handle.write(f" offline = {offline}\n") + audit_handle.write(f" limit = {limit}\n") + audit_handle.write(f" dry_run = {dry_run}\n") + audit_handle.write(f" auto_strip_templates = {auto_strip_templates}\n") + audit_handle.write(f" template_min_count = {template_min_count}\n") + audit_handle.write(f" skip_minhash = {skip_minhash}\n") + audit_handle.write(f" minhash_threshold = {minhash_threshold}\n") + audit_handle.write(f" minhash_num_perm = {minhash_num_perm}\n") + audit_handle.write(f" minhash_shingle_size = {minhash_shingle_size}\n") + audit_handle.write(f" skip_semantic = {skip_semantic}\n") + audit_handle.write(f" semantic_threshold = {semantic_threshold}\n") + audit_handle.write(f" embed_model = {embed_model}\n") + audit_handle.write(f" embed_batch_size = {embed_batch_size}\n") + audit_handle.write(f" semantic_topk = {semantic_topk}\n") + audit_handle.write(f" device = {device}\n") + audit_handle.flush() + print(f"\n[audit] Writing per-cluster audit to: {audit_path}") + + # ── Optional: auto-detect templates and strip before similarity comp ── + # Original prompts in `records` are NEVER mutated; stripping only affects + # the prompts passed to MinHash / Semantic for similarity computation. + templates: list[str] = [] + if auto_strip_templates: + print( + f"\n[template-strip] Auto-detecting common prefix templates " + f"(min_count={template_min_count}) ..." + ) + raw_prompts = [r["conversations"][0]["value"] for r in records] + templates = find_common_prefixes(raw_prompts, min_count=template_min_count) + if templates: + hit_counts = template_hit_stats(raw_prompts, templates) + print(f" Detected {len(templates)} template(s):") + for tmpl in templates: + preview = tmpl[:80].replace("\n", "\\n") + ellipsis = "..." if len(tmpl) > 80 else "" + print( + f" [{hit_counts[tmpl]:>6,} hits, {len(tmpl):>4} chars] " + f"{preview!r}{ellipsis}" + ) + print(f" [{hit_counts['']:>6,} prompts unchanged]") + if audit_handle is not None: + audit_handle.write(f"\n{'=' * 78}\n") + audit_handle.write( + f"TEMPLATE STRIPPING (auto-detected, min_count={template_min_count})\n" + ) + audit_handle.write(f" Found {len(templates)} template(s):\n") + audit_handle.write(f"{'=' * 78}\n\n") + for tmpl in templates: + audit_handle.write( + f"[{hit_counts[tmpl]:>6,} hits, {len(tmpl):>4} chars]\n" + f" {tmpl!r}\n\n" + ) + audit_handle.write( + f"[{hit_counts['']:>6,} prompts unchanged by stripping]\n" + ) + audit_handle.flush() + else: + print(f" No common templates found at min_count={template_min_count}.") + + # ── Stage 2: MinHash near-duplicate dedup ───────────────────────── + if skip_minhash: + print("\nSkipping MinHash near-dup stage (--skip-minhash)") + n_after_minhash = n_after_exact + else: + print( + f"\n── MinHash near-dup " + f"(threshold={minhash_threshold}, num_perm={minhash_num_perm}, " + f"shingle_size={minhash_shingle_size}) ──" + ) + records_before_minhash = records + prompts = [ + strip_template(r["conversations"][0]["value"], templates) for r in records + ] + rep_key = lambda i: -len(records[i]["conversations"][1]["value"]) + result = MinHashDeduplicator( + threshold=minhash_threshold, + num_perm=minhash_num_perm, + shingle_size=minhash_shingle_size, + seed=seed, + ).dedup(prompts, key_fn=rep_key, return_clusters=write_audit) + if write_audit: + keep, clusters_mh = result + else: + keep, clusters_mh = result, None + records = [records[i] for i in keep] + n_after_minhash = len(records) + print( + f" MinHash kept {n_after_minhash:,} / {n_after_exact:,} " + f"({n_after_exact - n_after_minhash:,} removed)" + ) + if write_audit and audit_handle is not None: + _write_dedup_audit_section( + audit_handle, + stage_name=f"MinHash near-dup (num_perm={minhash_num_perm}, " + f"shingle_size={minhash_shingle_size})", + threshold=minhash_threshold, + records_before=records_before_minhash, + keep_indices=keep, + clusters=clusters_mh, + ) + audit_handle.flush() + + # ── Stage 3: Semantic paraphrase dedup ──────────────────────────── + if skip_semantic: + print("\nSkipping semantic dedup stage (--skip-semantic)") + n_after_semantic = n_after_minhash + else: + print( + f"\n── Semantic dedup " + f"(model={embed_model}, threshold={semantic_threshold}) ──" + ) + records_before_semantic = records + prompts = [ + strip_template(r["conversations"][0]["value"], templates) for r in records + ] + rep_key = lambda i: -len(records[i]["conversations"][1]["value"]) + deduper = SemanticDeduplicator( + model_name=embed_model, + threshold=semantic_threshold, + batch_size=embed_batch_size, + device=device, + topk=semantic_topk, + ) + if dump_embeddings is not None: + # Persist the (expensive) embeddings + per-record metadata so the + # post-MinHash set can be re-clustered at any threshold and analyzed + # (topic coverage, threshold sweep) offline without re-embedding. + import numpy as np + Path(dump_embeddings).mkdir(parents=True, exist_ok=True) + embeddings = deduper.encode(prompts) + np.save(str(Path(dump_embeddings) / "embeddings.npy"), embeddings) + with open(Path(dump_embeddings) / "records.jsonl", "w") as _df: + for r in records: + _df.write(json.dumps({ + "prompt": r["conversations"][0]["value"], + "stripped": strip_template(r["conversations"][0]["value"], templates), + "resp_len": len(r["conversations"][1]["value"]), + "source": r.get("source"), + }, ensure_ascii=False) + "\n") + print(f" [dump-embeddings] saved {len(records):,} embeddings + records to {dump_embeddings}") + result = deduper.dedup_from_embeddings( + embeddings, key_fn=rep_key, return_clusters=write_audit + ) + else: + result = deduper.dedup(prompts, key_fn=rep_key, return_clusters=write_audit) + if write_audit: + keep, clusters_sem = result + else: + keep, clusters_sem = result, None + records = [records[i] for i in keep] + n_after_semantic = len(records) + print( + f" Semantic kept {n_after_semantic:,} / {n_after_minhash:,} " + f"({n_after_minhash - n_after_semantic:,} removed)" + ) + if write_audit and audit_handle is not None: + _write_dedup_audit_section( + audit_handle, + stage_name=f"Semantic paraphrase (model={embed_model}, topk={semantic_topk})", + threshold=semantic_threshold, + records_before=records_before_semantic, + keep_indices=keep, + clusters=clusters_sem, + ) + audit_handle.flush() + + if audit_handle is not None: + audit_handle.close() + print(f"\n[audit] Wrote per-cluster dedup audit to: {audit_path}") + + # ── Assertions ──────────────────────────────────────────────────── + if EXPECTED_UNIQUE_PROMPTS is not None: + assert len(records) == EXPECTED_UNIQUE_PROMPTS, ( + f"Expected {EXPECTED_UNIQUE_PROMPTS} rows, got {len(records)}" + ) + + seen_prompts: set[str] = set() + for rec in records: + prompt = rec["conversations"][0]["value"] + assert prompt not in seen_prompts, f"Duplicate prompt found: {prompt[:80]}..." + seen_prompts.add(prompt) + print(" All assertions passed.") + + # ── Source distribution ────────────────────────────────────────── + source_counts: dict[str, int] = {} + for rec in records: + src = rec.get("source") or "" + source_counts[src] = source_counts.get(src, 0) + 1 + print() + print("──────── Source breakdown (final dedup'd records) ────────") + for src, n in sorted(source_counts.items(), key=lambda x: -x[1]): + print(f" {src}: {n:,}") + print("──────────────────────────────────────────────────────────") + + # ── Multiple-choice detection ──────────────────────────────────── + mc_count = 0 + non_mc_count = 0 + for rec in records: + prompt = rec["conversations"][0]["value"] + if "multiple choice" in prompt.lower(): + mc_count += 1 + else: + non_mc_count += 1 + pct_mc = (mc_count / len(records) * 100) if records else 0.0 + print() + print("──────── Multiple-choice breakdown ────────") + print(f" Contains 'multiple choice': {mc_count:,} ({pct_mc:.1f}%)") + print(f" Does NOT contain 'multiple choice': {non_mc_count:,} ({100 - pct_mc:.1f}%)") + print("───────────────────────────────────────────") + + # ── Sample preview ──────────────────────────────────────────────── + if dry_run: + n_preview = min(5, len(records)) + print() + print(f"──────── Sample preview ({n_preview} records that WOULD upload) ────────") + for i, rec in enumerate(random.sample(records, n_preview)): + print(f"\n[Sample {i + 1}/{n_preview}]") + blob = json.dumps(rec, ensure_ascii=False, indent=2) + print(blob[:2000]) + if len(blob) > 2000: + print(" ... (truncated to 2000 chars)") + print("─────────────────────────────────────────────────────────────────") + + # ── Write final jsonl + upload to HuggingFace ───────────────────── + # Write to a stable local path FIRST so a failed/blocked upload never throws + # away the (expensive) pipeline output — it can be re-uploaded from this file + # without re-running the whole pipeline. + out_path = Path("output/final") / (HF_REPO_ID.split("/")[-1] + ".jsonl") + out_path.parent.mkdir(parents=True, exist_ok=True) + print(f"\n Writing {len(records):,} final records to {out_path} ...") + t0 = time.perf_counter() + with open(out_path, "w") as f: + for rec in tqdm(records, desc=" jsonl write", unit=" rec", smoothing=0.05): + f.write(json.dumps(rec, ensure_ascii=False) + "\n") + file_size_mb = out_path.stat().st_size / 1e6 + print(f" wrote {file_size_mb:,.1f} MB in {time.perf_counter() - t0:.1f}s") + + if dry_run: + print("[dry-run] Skipping upload to HuggingFace.") + else: + print(f"\nUploading {file_size_mb:,.1f} MB to {HF_REPO_ID} ...") + api = HfApi() + create_repo(HF_REPO_ID, repo_type="dataset", exist_ok=True) + t0 = time.perf_counter() + api.upload_file( + path_or_fileobj=str(out_path), + path_in_repo="data/train-00000-of-00001.jsonl", + repo_id=HF_REPO_ID, + repo_type="dataset", + ) + print(f" uploaded in {time.perf_counter() - t0:.1f}s") + + print() + print("──────── Pipeline attrition ────────") + print(f" {'Stage':<38s} {'Kept':>12s} {'Removed':>12s} {'% kept vs prev':>17s}") + + def _attrition_row(label: str, kept: int, prev: int | None) -> None: + if prev is None: + print(f" {label:<38s} {kept:>12,} {'-':>12s} {'-':>17s}") + return + removed = prev - kept + pct = (kept / prev * 100) if prev else 0.0 + print(f" {label:<38s} {kept:>12,} {removed:>12,} {pct:>16.2f}%") + + _attrition_row("Streamed from HF", total, None) + _attrition_row("After exact dedup (reservoir)", n_after_exact, total) + _attrition_row("After MinHash near-dup", n_after_minhash, n_after_exact) + _attrition_row("After semantic dedup", n_after_semantic, n_after_minhash) + print("──────────────────────────────────────────────────────────────────────────────────") + + print() + print("──────── Final stats ────────") + print(f" Total rows streamed: {total:,}") + print(f" Skipped (no user message): {skipped_no_user:,}") + print(f" Skipped (no assistant message): {skipped_no_assistant:,}") + print(f" Unique user prompts (final): {len(records):,}") + print(f" Contains 'multiple choice': {mc_count:,} ({pct_mc:.1f}%)") + print(f" Does NOT contain 'multiple choice': {non_mc_count:,} ({100 - pct_mc:.1f}%)") + if dry_run: + print(f" Rows that WOULD upload: {len(records):,} (dry-run, not uploaded)") + print(f" HF dataset (target): https://huggingface.co/datasets/{HF_REPO_ID}") + else: + print(f" Rows uploaded to HF: {len(records):,}") + print(f" HF dataset: https://huggingface.co/datasets/{HF_REPO_ID}") + print("─────────────────────────────") + + +def sample_and_upload( + n_sample: int, + seed: int = 42, + dry_run: bool = False, + source_path: str | None = None, + target_repo: str | None = None, +) -> None: + """Shuffle + randomly sample N records from the already-deduped dataset and + upload to a separate HF repo. Independent of the dedup pipeline. + + Source is the local final jsonl produced by a full run + (output/final/.jsonl) if present, else the HF dataset HF_REPO_ID. + Records are copied verbatim, so the output format is identical to the source. + """ + random.seed(seed) + + if target_repo is None: + suffix = f"{n_sample // 1000}k" if n_sample % 1000 == 0 else str(n_sample) + target_repo = f"{HF_REPO_ID}-{suffix}" + + local_final = ( + Path(source_path) if source_path + else Path("output/final") / (HF_REPO_ID.split("/")[-1] + ".jsonl") + ) + + # Reservoir sampling: a uniform random sample of N in a single pass, holding + # only N lines in memory (no need to count the source first). + if local_final.exists(): + print(f"Sampling {n_sample:,} from local file: {local_final}") + reservoir: list[str] = [] + total = 0 + with open(local_final) as f: + for i, line in enumerate(f): + total += 1 + if i < n_sample: + reservoir.append(line) + else: + j = random.randint(0, i) + if j < n_sample: + reservoir[j] = line + else: + print(f"Local file not found; loading from HF: {HF_REPO_ID}") + ds = load_dataset(HF_REPO_ID, split="train") + total = len(ds) + if n_sample > total: + raise ValueError(f"--sample {n_sample} exceeds source size {total:,}") + idx = random.sample(range(total), n_sample) + reservoir = [json.dumps(ds[i], ensure_ascii=False) + "\n" for i in idx] + + print(f" source records: {total:,}") + if n_sample > total: + raise ValueError(f"--sample {n_sample} exceeds source size {total:,}") + random.shuffle(reservoir) + + out_path = Path("output/final") / (target_repo.split("/")[-1] + ".jsonl") + out_path.parent.mkdir(parents=True, exist_ok=True) + with open(out_path, "w") as f: + f.writelines(reservoir) + mb = out_path.stat().st_size / 1e6 + print(f" wrote {len(reservoir):,} sampled records ({mb:,.1f} MB) to {out_path}") + + src_counts: dict[str, int] = {} + for line in reservoir: + s = json.loads(line).get("source") or "" + src_counts[s] = src_counts.get(s, 0) + 1 + print(" sample source breakdown:", + {k: f"{v:,}" for k, v in sorted(src_counts.items(), key=lambda x: -x[1])}) + + if dry_run: + print(f"[dry-run] Skipping upload to {target_repo}.") + return + + print(f"\nUploading {mb:,.1f} MB to {target_repo} ...") + create_repo(target_repo, repo_type="dataset", exist_ok=True) + t0 = time.perf_counter() + HfApi().upload_file( + path_or_fileobj=str(out_path), + path_in_repo="data/train-00000-of-00001.jsonl", + repo_id=target_repo, + repo_type="dataset", + ) + print(f" uploaded in {time.perf_counter() - t0:.1f}s " + f"-> https://huggingface.co/datasets/{target_repo}") + + +if __name__ == "__main__": + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument( + "--dry-run", + action="store_true", + help="Run preprocessing and assertions but skip the HuggingFace upload.", + ) + parser.add_argument("--seed", type=int, default=42, help="Random seed for reservoir sampling.") + parser.add_argument("--stream", action=argparse.BooleanOptionalAction, default=False, + help="Stream the dataset over the network instead of downloading it to " + "the local HF cache first. Default: OFF (download-to-cache, then " + "iterate from disk — faster and re-runnable offline). Use --stream " + "to opt into streaming, or --no-stream to force download mode.") + parser.add_argument("--offline", action="store_true", + help="Force HF offline mode (HF_DATASETS_OFFLINE=1, HF_HUB_OFFLINE=1): " + "read dataset + embed model purely from cache, zero network calls. " + "Handled before importing datasets; listed here for --help/validation.") + parser.add_argument("--no-offline", action="store_true", + help="Force online mode (e.g. to refresh the cache). In download mode, " + "offline auto-enables once cached; this overrides that behavior.") + + # ── Sampling mode (separate code path; does NOT run the dedup pipeline) ── + parser.add_argument("--sample", type=int, default=None, metavar="N", + help="SAMPLING MODE: shuffle + randomly sample N records from the already-" + f"deduped dataset ({HF_REPO_ID}) and upload to a separate repo. " + "Reads output/final/.jsonl if present, else loads from HF. " + "Does not run the dedup pipeline. Honors --seed and --dry-run.") + parser.add_argument("--sample-repo", type=str, default=None, + help="Target HF repo for --sample (default: '-k').") + parser.add_argument("--sample-source", type=str, default=None, + help="Override the source jsonl for --sample (default: " + "output/final/.jsonl).") + + parser.add_argument("--limit", type=int, default=None, + help="Process at most N rows from the stream (debug aid). " + "When set, also writes a per-cluster audit file by default.") + parser.add_argument("--audit-path", type=str, default=None, + help="Path for the per-cluster dedup audit file. Auto-generated " + "under output/logs/ when --limit is set; specify here to override " + "or to enable audit on a full run.") + parser.add_argument("--dump-embeddings", type=str, default=None, + help="Directory to save the semantic-stage embeddings.npy + records.jsonl " + "(post-MinHash set). Lets you re-cluster at any threshold and run " + "topic/coverage analysis offline without re-embedding.") + + parser.add_argument("--auto-strip-templates", action=argparse.BooleanOptionalAction, + default=True, + help="Auto-detect and strip frequent prompt-template prefixes before " + "computing similarity. Original prompts in records are unchanged; " + "only the dedup signal sees stripped versions. Default: ON (prevents " + "boilerplate-driven over-merging). Use --no-auto-strip-templates " + "to disable.") + parser.add_argument("--template-min-count", type=int, default=10, + help="Min number of prompts a prefix must appear in to be treated " + "as a template (default 10).") + + parser.add_argument("--skip-minhash", action="store_true", help="Skip MinHash near-dup stage.") + parser.add_argument("--minhash-threshold", type=float, default=0.8, + help="Jaccard threshold for MinHash (default 0.8).") + parser.add_argument("--minhash-num-perm", type=int, default=128, + help="Number of MinHash permutations (default 128).") + parser.add_argument("--minhash-shingle-size", type=int, default=5, + help="Word-shingle size for MinHash (default 5).") + + parser.add_argument("--skip-semantic", action="store_true", help="Skip semantic dedup stage.") + parser.add_argument("--semantic-threshold", type=float, default=0.95, + help="Cosine threshold for semantic dedup (default 0.95). 0.92 was found " + "to over-merge distinct STEM questions into topic blobs; 0.95 keeps " + "~109K more records and removes ~70%% fewer from large clusters.") + parser.add_argument("--embed-model", type=str, default="BAAI/bge-small-en-v1.5", + help="Sentence-transformers model for embedding.") + parser.add_argument("--embed-batch-size", type=int, default=128, help="Embedding batch size.") + parser.add_argument("--semantic-topk", type=int, default=10, + help="FAISS top-k neighbors per query.") + parser.add_argument("--device", type=str, default="auto", + choices=["auto", "mps", "cuda", "cpu"], help="Embedding device.") + + args = parser.parse_args() + + if args.sample is not None: + sample_and_upload( + n_sample=args.sample, + seed=args.seed, + dry_run=args.dry_run, + source_path=args.sample_source, + target_repo=args.sample_repo, + ) + sys.exit(0) + + preprocess_and_upload( + seed=args.seed, + dry_run=args.dry_run, + stream=args.stream, + limit=args.limit, + audit_path=args.audit_path, + dump_embeddings=args.dump_embeddings, + auto_strip_templates=args.auto_strip_templates, + template_min_count=args.template_min_count, + skip_minhash=args.skip_minhash, + minhash_threshold=args.minhash_threshold, + minhash_num_perm=args.minhash_num_perm, + minhash_shingle_size=args.minhash_shingle_size, + skip_semantic=args.skip_semantic, + semantic_threshold=args.semantic_threshold, + embed_model=args.embed_model, + embed_batch_size=args.embed_batch_size, + semantic_topk=args.semantic_topk, + device=args.device, + ) diff --git a/sdg/preprocessing/preprocess_ot3_1_2m_code.py b/sdg/preprocessing/preprocess_ot3_1_2m_code.py new file mode 100644 index 0000000000000000000000000000000000000000..f957b29e3c81ace1e93412b4554e17fc293f5415 --- /dev/null +++ b/sdg/preprocessing/preprocess_ot3_1_2m_code.py @@ -0,0 +1,125 @@ +""" +Preprocess OpenThoughts3-1.2M (code domain only). + +Source: https://huggingface.co/datasets/open-thoughts/OpenThoughts3-1.2M + +This script filters the code domain, deduplicates by user prompt +(picking one random rollout per prompt via reservoir sampling), +asserts correctness, and uploads to HuggingFace. +""" + +from __future__ import annotations + +import json +import random +import tempfile +from pathlib import Path + +from datasets import load_dataset +from huggingface_hub import HfApi, create_repo + + +DATASET_ID = "open-thoughts/OpenThoughts3-1.2M" +EXPECTED_UNIQUE_PROMPTS = 5_714 +HF_REPO_ID = "teetone/openthoughts3-code-deduped" + + +def extract_prompt(conversations: list[dict]) -> str | None: + """Return the first human turn from a conversation.""" + for turn in conversations: + if turn["from"] == "human": + return turn["value"] + return None + + +def preprocess_and_upload(seed: int = 42) -> None: + random.seed(seed) + + print(f"Loading {DATASET_ID} (streaming) ...") + ds = load_dataset(DATASET_ID, split="train", streaming=True) + + # Single streaming pass: reservoir-sample one row per unique prompt. + # For each prompt we track (chosen_row, count_seen) so we can do + # uniform random selection without storing all rollouts. + prompt_reservoir: dict[str, tuple[dict, int]] = {} + total = 0 + code_count = 0 + + for row in ds: + total += 1 + if row["domain"] != "code": + continue + code_count += 1 + + prompt = extract_prompt(row["conversations"]) + if prompt is None: + continue + + rec = { + "difficulty": row["difficulty"], + "source": row["source"], + "domain": row["domain"], + "conversations": row["conversations"], + } + + if prompt not in prompt_reservoir: + prompt_reservoir[prompt] = (rec, 1) + else: + _, count = prompt_reservoir[prompt] + count += 1 + # Reservoir sampling: replace with probability 1/count + if random.randint(1, count) == 1: + prompt_reservoir[prompt] = (rec, count) + else: + prompt_reservoir[prompt] = (prompt_reservoir[prompt][0], count) + + if code_count % 100_000 == 0: + print(f" Processed {code_count:,} code rows, {len(prompt_reservoir):,} unique prompts ...") + + print(f" Total rows streamed: {total:,}") + print(f" Code rows: {code_count:,}") + print(f" Unique prompts: {len(prompt_reservoir):,}") + + records = [row for row, _ in prompt_reservoir.values()] + print(f" Deduplicated rows: {len(records):,}") + + # ── Assertions ──────────────────────────────────────────────────── + assert len(records) == EXPECTED_UNIQUE_PROMPTS, ( + f"Expected {EXPECTED_UNIQUE_PROMPTS} rows, got {len(records)}" + ) + + seen_prompts: set[str] = set() + for rec in records: + assert rec["domain"] == "code", f"Non-code row found: domain={rec['domain']}" + prompt = extract_prompt(rec["conversations"]) + assert prompt not in seen_prompts, f"Duplicate prompt found: {prompt[:80]}..." + seen_prompts.add(prompt) + + assert len(seen_prompts) == EXPECTED_UNIQUE_PROMPTS, ( + f"Expected {EXPECTED_UNIQUE_PROMPTS} unique prompts, got {len(seen_prompts)}" + ) + print(" All assertions passed.") + + # ── Upload to HuggingFace ───────────────────────────────────────── + print(f"Uploading to {HF_REPO_ID} ...") + api = HfApi() + create_repo(HF_REPO_ID, repo_type="dataset", exist_ok=True) + + with tempfile.NamedTemporaryFile(mode="w", suffix=".jsonl", delete=False) as f: + temp_path = f.name + for rec in records: + f.write(json.dumps(rec, ensure_ascii=False) + "\n") + + api.upload_file( + path_or_fileobj=temp_path, + path_in_repo="data/train-00000-of-00001.jsonl", + repo_id=HF_REPO_ID, + repo_type="dataset", + ) + Path(temp_path).unlink() + + print(f"Done! Dataset available at: https://huggingface.co/datasets/{HF_REPO_ID}") + + +if __name__ == "__main__": + preprocess_and_upload() diff --git a/sdg/preprocessing/preprocess_ot3_1_2m_math.py b/sdg/preprocessing/preprocess_ot3_1_2m_math.py new file mode 100644 index 0000000000000000000000000000000000000000..8f2f3f64fa8a0e8480e1bf4dd6305055ffe89f64 --- /dev/null +++ b/sdg/preprocessing/preprocess_ot3_1_2m_math.py @@ -0,0 +1,130 @@ +""" +Preprocess OpenThoughts3-1.2M (math domain only). + +Source: https://huggingface.co/datasets/open-thoughts/OpenThoughts3-1.2M + +Dataset stats (math domain): + - 850,000 rows from source 'ai2-adapt-dev/openmath-2-math' + - 53,125 unique user prompts (~16 rollouts per prompt) + - Conversations: 2 turns (human prompt + gpt response with reasoning) + +This script filters the math domain, deduplicates by user prompt +(picking one random rollout per prompt via reservoir sampling), +asserts correctness, and uploads to HuggingFace. +""" + +from __future__ import annotations + +import json +import random +import tempfile +from pathlib import Path + +from datasets import load_dataset +from huggingface_hub import HfApi, create_repo + + +DATASET_ID = "open-thoughts/OpenThoughts3-1.2M" +EXPECTED_UNIQUE_PROMPTS = 53_125 +HF_REPO_ID = "teetone/openthoughts3-math-deduped-53K" + + +def extract_prompt(conversations: list[dict]) -> str | None: + """Return the first human turn from a conversation.""" + for turn in conversations: + if turn["from"] == "human": + return turn["value"] + return None + + +def preprocess_and_upload(seed: int = 42) -> None: + random.seed(seed) + + print(f"Loading {DATASET_ID} (streaming) ...") + ds = load_dataset(DATASET_ID, split="train", streaming=True) + + # Single streaming pass: reservoir-sample one row per unique prompt. + # For each prompt we track (chosen_row, count_seen) so we can do + # uniform random selection without storing all rollouts. + prompt_reservoir: dict[str, tuple[dict, int]] = {} + total = 0 + math_count = 0 + + for row in ds: + total += 1 + if row["domain"] != "math": + continue + math_count += 1 + + prompt = extract_prompt(row["conversations"]) + if prompt is None: + continue + + rec = { + "difficulty": row["difficulty"], + "source": row["source"], + "domain": row["domain"], + "conversations": row["conversations"], + } + + if prompt not in prompt_reservoir: + prompt_reservoir[prompt] = (rec, 1) + else: + _, count = prompt_reservoir[prompt] + count += 1 + # Reservoir sampling: replace with probability 1/count + if random.randint(1, count) == 1: + prompt_reservoir[prompt] = (rec, count) + else: + prompt_reservoir[prompt] = (prompt_reservoir[prompt][0], count) + + if math_count % 100_000 == 0: + print(f" Processed {math_count:,} math rows, {len(prompt_reservoir):,} unique prompts ...") + + print(f" Total rows streamed: {total:,}") + print(f" Math rows: {math_count:,}") + print(f" Unique prompts: {len(prompt_reservoir):,}") + + records = [row for row, _ in prompt_reservoir.values()] + print(f" Deduplicated rows: {len(records):,}") + + # ── Assertions ──────────────────────────────────────────────────── + assert len(records) == EXPECTED_UNIQUE_PROMPTS, ( + f"Expected {EXPECTED_UNIQUE_PROMPTS} rows, got {len(records)}" + ) + + seen_prompts: set[str] = set() + for rec in records: + assert rec["domain"] == "math", f"Non-math row found: domain={rec['domain']}" + prompt = extract_prompt(rec["conversations"]) + assert prompt not in seen_prompts, f"Duplicate prompt found: {prompt[:80]}..." + seen_prompts.add(prompt) + + assert len(seen_prompts) == EXPECTED_UNIQUE_PROMPTS, ( + f"Expected {EXPECTED_UNIQUE_PROMPTS} unique prompts, got {len(seen_prompts)}" + ) + print(" All assertions passed.") + + # ── Upload to HuggingFace ───────────────────────────────────────── + print(f"Uploading to {HF_REPO_ID} ...") + api = HfApi() + create_repo(HF_REPO_ID, repo_type="dataset", exist_ok=True) + + with tempfile.NamedTemporaryFile(mode="w", suffix=".jsonl", delete=False) as f: + temp_path = f.name + for rec in records: + f.write(json.dumps(rec, ensure_ascii=False) + "\n") + + api.upload_file( + path_or_fileobj=temp_path, + path_in_repo="data/train-00000-of-00001.jsonl", + repo_id=HF_REPO_ID, + repo_type="dataset", + ) + Path(temp_path).unlink() + + print(f"Done! Dataset available at: https://huggingface.co/datasets/{HF_REPO_ID}") + + +if __name__ == "__main__": + preprocess_and_upload() diff --git a/sdg/preprocessing/preprocess_ot3_1_2m_science.py b/sdg/preprocessing/preprocess_ot3_1_2m_science.py new file mode 100644 index 0000000000000000000000000000000000000000..86fb4e901623653cc85f2d8bd85dd6bcdfd8a819 --- /dev/null +++ b/sdg/preprocessing/preprocess_ot3_1_2m_science.py @@ -0,0 +1,125 @@ +""" +Preprocess OpenThoughts3-1.2M (science domain only). + +Source: https://huggingface.co/datasets/open-thoughts/OpenThoughts3-1.2M + +This script filters the science domain, deduplicates by user prompt +(picking one random rollout per prompt via reservoir sampling), +asserts correctness, and uploads to HuggingFace. +""" + +from __future__ import annotations + +import json +import random +import tempfile +from pathlib import Path + +from datasets import load_dataset +from huggingface_hub import HfApi, create_repo + + +DATASET_ID = "open-thoughts/OpenThoughts3-1.2M" +EXPECTED_UNIQUE_PROMPTS = 6_250 +HF_REPO_ID = "teetone/openthoughts3-science-deduped" + + +def extract_prompt(conversations: list[dict]) -> str | None: + """Return the first human turn from a conversation.""" + for turn in conversations: + if turn["from"] == "human": + return turn["value"] + return None + + +def preprocess_and_upload(seed: int = 42) -> None: + random.seed(seed) + + print(f"Loading {DATASET_ID} (streaming) ...") + ds = load_dataset(DATASET_ID, split="train", streaming=True) + + # Single streaming pass: reservoir-sample one row per unique prompt. + # For each prompt we track (chosen_row, count_seen) so we can do + # uniform random selection without storing all rollouts. + prompt_reservoir: dict[str, tuple[dict, int]] = {} + total = 0 + science_count = 0 + + for row in ds: + total += 1 + if row["domain"] != "science": + continue + science_count += 1 + + prompt = extract_prompt(row["conversations"]) + if prompt is None: + continue + + rec = { + "difficulty": row["difficulty"], + "source": row["source"], + "domain": row["domain"], + "conversations": row["conversations"], + } + + if prompt not in prompt_reservoir: + prompt_reservoir[prompt] = (rec, 1) + else: + _, count = prompt_reservoir[prompt] + count += 1 + # Reservoir sampling: replace with probability 1/count + if random.randint(1, count) == 1: + prompt_reservoir[prompt] = (rec, count) + else: + prompt_reservoir[prompt] = (prompt_reservoir[prompt][0], count) + + if science_count % 100_000 == 0: + print(f" Processed {science_count:,} science rows, {len(prompt_reservoir):,} unique prompts ...") + + print(f" Total rows streamed: {total:,}") + print(f" Science rows: {science_count:,}") + print(f" Unique prompts: {len(prompt_reservoir):,}") + + records = [row for row, _ in prompt_reservoir.values()] + print(f" Deduplicated rows: {len(records):,}") + + # ── Assertions ──────────────────────────────────────────────────── + assert len(records) == EXPECTED_UNIQUE_PROMPTS, ( + f"Expected {EXPECTED_UNIQUE_PROMPTS} rows, got {len(records)}" + ) + + seen_prompts: set[str] = set() + for rec in records: + assert rec["domain"] == "science", f"Non-science row found: domain={rec['domain']}" + prompt = extract_prompt(rec["conversations"]) + assert prompt not in seen_prompts, f"Duplicate prompt found: {prompt[:80]}..." + seen_prompts.add(prompt) + + assert len(seen_prompts) == EXPECTED_UNIQUE_PROMPTS, ( + f"Expected {EXPECTED_UNIQUE_PROMPTS} unique prompts, got {len(seen_prompts)}" + ) + print(" All assertions passed.") + + # ── Upload to HuggingFace ───────────────────────────────────────── + print(f"Uploading to {HF_REPO_ID} ...") + api = HfApi() + create_repo(HF_REPO_ID, repo_type="dataset", exist_ok=True) + + with tempfile.NamedTemporaryFile(mode="w", suffix=".jsonl", delete=False) as f: + temp_path = f.name + for rec in records: + f.write(json.dumps(rec, ensure_ascii=False) + "\n") + + api.upload_file( + path_or_fileobj=temp_path, + path_in_repo="data/train-00000-of-00001.jsonl", + repo_id=HF_REPO_ID, + repo_type="dataset", + ) + Path(temp_path).unlink() + + print(f"Done! Dataset available at: https://huggingface.co/datasets/{HF_REPO_ID}") + + +if __name__ == "__main__": + preprocess_and_upload() diff --git a/sdg/preprocessing/preprocess_ot4_3_5m_math.py b/sdg/preprocessing/preprocess_ot4_3_5m_math.py new file mode 100644 index 0000000000000000000000000000000000000000..a0d32f8c602957378405f3baf116ea66efce0826 --- /dev/null +++ b/sdg/preprocessing/preprocess_ot4_3_5m_math.py @@ -0,0 +1,115 @@ +""" +Preprocess hero_run_4_math (math domain). + +Source: https://huggingface.co/datasets/mlfoundations-dev/hero_run_4_math + +Dataset stats: + - ~3,500,000 rows + - Deduplicated by `instruction_seed` + +This script deduplicates by instruction_seed (picking one random rollout +per prompt via reservoir sampling), asserts correctness, and uploads +to HuggingFace. +""" + +from __future__ import annotations + +import json +import random +import tempfile +from pathlib import Path + +from datasets import load_dataset +from huggingface_hub import HfApi, create_repo + + +DATASET_ID = "mlfoundations-dev/hero_run_4_math" +EXPECTED_UNIQUE_PROMPTS = 218_750 +HF_REPO_ID = "teetone/ot4-math-deduped" + + +def preprocess_and_upload(seed: int = 42) -> None: + random.seed(seed) + + print(f"Loading {DATASET_ID} (streaming) ...") + ds = load_dataset(DATASET_ID, split="train", streaming=True) + + # Single streaming pass: reservoir-sample one row per unique instruction_seed. + # For each prompt we track (chosen_row, count_seen) so we can do + # uniform random selection without storing all rollouts. + prompt_reservoir: dict[str, tuple[dict, int]] = {} + total = 0 + + for row in ds: + total += 1 + + prompt = row.get("instruction_seed") + if prompt is None: + continue + + rec = { + "conversations": [ + {"from": "human", "value": prompt}, + {"from": "gpt", "value": row["response_seed"]}, + ], + } + + if prompt not in prompt_reservoir: + prompt_reservoir[prompt] = (rec, 1) + else: + _, count = prompt_reservoir[prompt] + count += 1 + # Reservoir sampling: replace with probability 1/count + if random.randint(1, count) == 1: + prompt_reservoir[prompt] = (rec, count) + else: + prompt_reservoir[prompt] = (prompt_reservoir[prompt][0], count) + + if total % 100_000 == 0: + print(f" Processed {total:,} rows, {len(prompt_reservoir):,} unique prompts ...") + + print(f" Total rows streamed: {total:,}") + print(f" Unique prompts: {len(prompt_reservoir):,}") + + records = [row for row, _ in prompt_reservoir.values()] + print(f" Deduplicated rows: {len(records):,}") + + # ── Assertions ──────────────────────────────────────────────────── + assert len(records) == EXPECTED_UNIQUE_PROMPTS, ( + f"Expected {EXPECTED_UNIQUE_PROMPTS} rows, got {len(records)}" + ) + + seen_prompts: set[str] = set() + for rec in records: + prompt = rec["conversations"][0]["value"] + assert prompt not in seen_prompts, f"Duplicate prompt found: {prompt[:80]}..." + seen_prompts.add(prompt) + + assert len(seen_prompts) == EXPECTED_UNIQUE_PROMPTS, ( + f"Expected {EXPECTED_UNIQUE_PROMPTS} unique prompts, got {len(seen_prompts)}" + ) + print(" All assertions passed.") + + # ── Upload to HuggingFace ───────────────────────────────────────── + print(f"Uploading to {HF_REPO_ID} ...") + api = HfApi() + create_repo(HF_REPO_ID, repo_type="dataset", exist_ok=True) + + with tempfile.NamedTemporaryFile(mode="w", suffix=".jsonl", delete=False) as f: + temp_path = f.name + for rec in records: + f.write(json.dumps(rec, ensure_ascii=False) + "\n") + + api.upload_file( + path_or_fileobj=temp_path, + path_in_repo="data/train-00000-of-00001.jsonl", + repo_id=HF_REPO_ID, + repo_type="dataset", + ) + Path(temp_path).unlink() + + print(f"Done! Dataset available at: https://huggingface.co/datasets/{HF_REPO_ID}") + + +if __name__ == "__main__": + preprocess_and_upload() \ No newline at end of file diff --git a/sdg/preprocessing/preprocess_ot4_417k_science.py b/sdg/preprocessing/preprocess_ot4_417k_science.py new file mode 100644 index 0000000000000000000000000000000000000000..bbcf964ac6379edb7352eec3b8c60b18af345777 --- /dev/null +++ b/sdg/preprocessing/preprocess_ot4_417k_science.py @@ -0,0 +1,115 @@ +""" +Preprocess hero_run_4_science (science domain). + +Source: https://huggingface.co/datasets/mlfoundations-dev/hero_run_4_science + +Dataset stats: + - 417,000 rows + - Deduplicated by `instruction_seed` + +This script deduplicates by instruction_seed (picking one random rollout +per prompt via reservoir sampling), asserts correctness, and uploads +to HuggingFace. +""" + +from __future__ import annotations + +import json +import random +import tempfile +from pathlib import Path + +from datasets import load_dataset +from huggingface_hub import HfApi, create_repo + + +DATASET_ID = "mlfoundations-dev/hero_run_4_science" +EXPECTED_UNIQUE_PROMPTS = 26_041 +HF_REPO_ID = "teetone/ot4-science-deduped" + + +def preprocess_and_upload(seed: int = 42) -> None: + random.seed(seed) + + print(f"Loading {DATASET_ID} (streaming) ...") + ds = load_dataset(DATASET_ID, split="train", streaming=True) + + # Single streaming pass: reservoir-sample one row per unique instruction_seed. + # For each prompt we track (chosen_row, count_seen) so we can do + # uniform random selection without storing all rollouts. + prompt_reservoir: dict[str, tuple[dict, int]] = {} + total = 0 + + for row in ds: + total += 1 + + prompt = row.get("instruction_seed") + if prompt is None: + continue + + rec = { + "conversations": [ + {"from": "human", "value": prompt}, + {"from": "gpt", "value": row["gpt41_mini_response"]}, + ], + } + + if prompt not in prompt_reservoir: + prompt_reservoir[prompt] = (rec, 1) + else: + _, count = prompt_reservoir[prompt] + count += 1 + # Reservoir sampling: replace with probability 1/count + if random.randint(1, count) == 1: + prompt_reservoir[prompt] = (rec, count) + else: + prompt_reservoir[prompt] = (prompt_reservoir[prompt][0], count) + + if total % 100_000 == 0: + print(f" Processed {total:,} rows, {len(prompt_reservoir):,} unique prompts ...") + + print(f" Total rows streamed: {total:,}") + print(f" Unique prompts: {len(prompt_reservoir):,}") + + records = [row for row, _ in prompt_reservoir.values()] + print(f" Deduplicated rows: {len(records):,}") + + # ── Assertions ──────────────────────────────────────────────────── + assert len(records) == EXPECTED_UNIQUE_PROMPTS, ( + f"Expected {EXPECTED_UNIQUE_PROMPTS} rows, got {len(records)}" + ) + + seen_prompts: set[str] = set() + for rec in records: + prompt = rec["conversations"][0]["value"] + assert prompt not in seen_prompts, f"Duplicate prompt found: {prompt[:80]}..." + seen_prompts.add(prompt) + + assert len(seen_prompts) == EXPECTED_UNIQUE_PROMPTS, ( + f"Expected {EXPECTED_UNIQUE_PROMPTS} unique prompts, got {len(seen_prompts)}" + ) + print(" All assertions passed.") + + # ── Upload to HuggingFace ───────────────────────────────────────── + print(f"Uploading to {HF_REPO_ID} ...") + api = HfApi() + create_repo(HF_REPO_ID, repo_type="dataset", exist_ok=True) + + with tempfile.NamedTemporaryFile(mode="w", suffix=".jsonl", delete=False) as f: + temp_path = f.name + for rec in records: + f.write(json.dumps(rec, ensure_ascii=False) + "\n") + + api.upload_file( + path_or_fileobj=temp_path, + path_in_repo="data/train-00000-of-00001.jsonl", + repo_id=HF_REPO_ID, + repo_type="dataset", + ) + Path(temp_path).unlink() + + print(f"Done! Dataset available at: https://huggingface.co/datasets/{HF_REPO_ID}") + + +if __name__ == "__main__": + preprocess_and_upload() diff --git a/sdg/preprocessing/preprocess_ot4_959k_code.py b/sdg/preprocessing/preprocess_ot4_959k_code.py new file mode 100644 index 0000000000000000000000000000000000000000..11a54b1a082db49fd997b3d81abce430a3380a91 --- /dev/null +++ b/sdg/preprocessing/preprocess_ot4_959k_code.py @@ -0,0 +1,115 @@ +""" +Preprocess hero_run_4_code (code domain). + +Source: https://huggingface.co/datasets/mlfoundations-dev/hero_run_4_code + +Dataset stats: + - 959,024 rows + - Deduplicated by `instruction_seed` + +This script deduplicates by instruction_seed (picking one random rollout +per prompt via reservoir sampling), asserts correctness, and uploads +to HuggingFace. +""" + +from __future__ import annotations + +import json +import random +import tempfile +from pathlib import Path + +from datasets import load_dataset +from huggingface_hub import HfApi, create_repo + + +DATASET_ID = "mlfoundations-dev/hero_run_4_code" +EXPECTED_UNIQUE_PROMPTS = 9_168 +HF_REPO_ID = "teetone/ot4-code-deduped" + + +def preprocess_and_upload(seed: int = 42) -> None: + random.seed(seed) + + print(f"Loading {DATASET_ID} (streaming) ...") + ds = load_dataset(DATASET_ID, split="train", streaming=True) + + # Single streaming pass: reservoir-sample one row per unique instruction_seed. + # For each prompt we track (chosen_row, count_seen) so we can do + # uniform random selection without storing all rollouts. + prompt_reservoir: dict[str, tuple[dict, int]] = {} + total = 0 + + for row in ds: + total += 1 + + prompt = row.get("instruction_seed") + if prompt is None: + continue + + rec = { + "conversations": [ + {"from": "human", "value": prompt}, + {"from": "gpt", "value": row["output"]}, + ], + } + + if prompt not in prompt_reservoir: + prompt_reservoir[prompt] = (rec, 1) + else: + _, count = prompt_reservoir[prompt] + count += 1 + # Reservoir sampling: replace with probability 1/count + if random.randint(1, count) == 1: + prompt_reservoir[prompt] = (rec, count) + else: + prompt_reservoir[prompt] = (prompt_reservoir[prompt][0], count) + + if total % 100_000 == 0: + print(f" Processed {total:,} rows, {len(prompt_reservoir):,} unique prompts ...") + + print(f" Total rows streamed: {total:,}") + print(f" Unique prompts: {len(prompt_reservoir):,}") + + records = [row for row, _ in prompt_reservoir.values()] + print(f" Deduplicated rows: {len(records):,}") + + # ── Assertions ──────────────────────────────────────────────────── + assert len(records) == EXPECTED_UNIQUE_PROMPTS, ( + f"Expected {EXPECTED_UNIQUE_PROMPTS} rows, got {len(records)}" + ) + + seen_prompts: set[str] = set() + for rec in records: + prompt = rec["conversations"][0]["value"] + assert prompt not in seen_prompts, f"Duplicate prompt found: {prompt[:80]}..." + seen_prompts.add(prompt) + + assert len(seen_prompts) == EXPECTED_UNIQUE_PROMPTS, ( + f"Expected {EXPECTED_UNIQUE_PROMPTS} unique prompts, got {len(seen_prompts)}" + ) + print(" All assertions passed.") + + # ── Upload to HuggingFace ───────────────────────────────────────── + print(f"Uploading to {HF_REPO_ID} ...") + api = HfApi() + create_repo(HF_REPO_ID, repo_type="dataset", exist_ok=True) + + with tempfile.NamedTemporaryFile(mode="w", suffix=".jsonl", delete=False) as f: + temp_path = f.name + for rec in records: + f.write(json.dumps(rec, ensure_ascii=False) + "\n") + + api.upload_file( + path_or_fileobj=temp_path, + path_in_repo="data/train-00000-of-00001.jsonl", + repo_id=HF_REPO_ID, + repo_type="dataset", + ) + Path(temp_path).unlink() + + print(f"Done! Dataset available at: https://huggingface.co/datasets/{HF_REPO_ID}") + + +if __name__ == "__main__": + preprocess_and_upload() diff --git a/sdg/preprocessing/preprocess_rstarcoder.py b/sdg/preprocessing/preprocess_rstarcoder.py new file mode 100644 index 0000000000000000000000000000000000000000..07fc7e03a39e8c90443b69e8c89a03b0c0126fb2 --- /dev/null +++ b/sdg/preprocessing/preprocess_rstarcoder.py @@ -0,0 +1,189 @@ +""" +Preprocess rStar-Coder (seed_sft subset). + +Source: https://huggingface.co/datasets/microsoft/rStar-Coder + +Dataset fields: + - question_id: Unique identifier linking problems with their test cases + - question: The original competitive programming problem description + - starter_code: Initial code template provided with the problem + - response: The long CoT solution + - code: The implementation code for the solution + - verified: Whether the solution has been verified with test cases + - is_passed: Whether the solution has passed the test cases + +This script deduplicates by `question_id` (picking one random rollout per +question via reservoir sampling), asserts correctness, and uploads to +HuggingFace in the `openthoughts4` conversations format so the SDG configs +can consume it. +""" + +from __future__ import annotations + +import argparse +import json +import random +import tempfile +from pathlib import Path + +from datasets import load_dataset +from huggingface_hub import HfApi, create_repo +from tqdm import tqdm + + +DATASET_ID = "microsoft/rStar-Coder" +DATASET_SUBSET = "seed_sft" +# Set to an int to assert the expected number of unique question_ids; leave as +# None on first run, then update after observing the printed count. +EXPECTED_UNIQUE_PROMPTS: int | None = None +HF_REPO_ID = "teetone/rstarcoder-seedsft-deduped" + + +def build_prompt(question: str, starter_code: str | None) -> str: + """Combine the problem description with the starter code template.""" + question = (question or "").strip() + starter = (starter_code or "").strip() + if starter: + return f"{question}\n\n{starter}" + return question + + +def preprocess_and_upload(seed: int = 42, dry_run: bool = False) -> None: + random.seed(seed) + + print(f"Loading {DATASET_ID} (subset={DATASET_SUBSET}, streaming) ...") + ds = load_dataset(DATASET_ID, DATASET_SUBSET, split="train", streaming=True) + + # Single streaming pass: reservoir-sample one row per unique question_id. + # For each question we track (chosen_row, count_seen) so we can do uniform + # random selection without storing all rollouts. + qid_reservoir: dict[str, tuple[dict, int]] = {} + total = 0 + skipped_no_qid = 0 + + pbar = tqdm(ds, desc="Streaming rStar-Coder", unit="row", smoothing=0.05) + for row in pbar: + total += 1 + + qid = row.get("question_id") + if qid is None: + skipped_no_qid += 1 + continue + qid = str(qid) + + prompt = build_prompt(row.get("question"), row.get("starter_code")) + if not prompt: + continue + + rec = { + "conversations": [ + {"from": "human", "value": prompt}, + {"from": "gpt", "value": row.get("response", "")}, + ], + } + + if qid not in qid_reservoir: + qid_reservoir[qid] = (rec, 1) + else: + _, count = qid_reservoir[qid] + count += 1 + # Reservoir sampling: replace with probability 1/count + if random.randint(1, count) == 1: + qid_reservoir[qid] = (rec, count) + else: + qid_reservoir[qid] = (qid_reservoir[qid][0], count) + + pbar.set_postfix(unique=len(qid_reservoir), skipped=skipped_no_qid) + pbar.close() + + print(f" Total rows streamed: {total:,}") + print(f" Skipped (no question_id): {skipped_no_qid:,}") + print(f" Unique question_ids: {len(qid_reservoir):,}") + + records_by_qid = [row for row, _ in qid_reservoir.values()] + print(f" After question_id dedup: {len(records_by_qid):,}") + + # ── Second-pass dedup on question text ──────────────────────────── + # Some rows can share an identical question while having different + # question_ids (e.g. the same problem reused under multiple sources). + # Drop those so only one record per unique question survives. + records: list[dict] = [] + seen_questions: set[str] = set() + duplicate_questions = 0 + for rec in records_by_qid: + prompt = rec["conversations"][0]["value"] + if prompt in seen_questions: + duplicate_questions += 1 + continue + seen_questions.add(prompt) + records.append(rec) + print(f" Removed by question-text dedup: {duplicate_questions:,}") + print(f" Final unique questions: {len(records):,}") + + # ── Assertions ──────────────────────────────────────────────────── + if EXPECTED_UNIQUE_PROMPTS is not None: + assert len(records) == EXPECTED_UNIQUE_PROMPTS, ( + f"Expected {EXPECTED_UNIQUE_PROMPTS} rows, got {len(records)}" + ) + assert len(seen_questions) == len(records), "Question-text dedup invariant failed" + print(" All assertions passed.") + + # ── Sample preview ──────────────────────────────────────────────── + if dry_run: + n_preview = min(5, len(records)) + print() + print(f"──────── Sample preview ({n_preview} records that WOULD upload) ────────") + for i, rec in enumerate(random.sample(records, n_preview)): + print(f"\n[Sample {i + 1}/{n_preview}]") + print(json.dumps(rec, ensure_ascii=False, indent=2)[:2000]) + if len(json.dumps(rec, ensure_ascii=False)) > 2000: + print(" ... (truncated to 2000 chars)") + print("─────────────────────────────────────────────────────────────────") + + # ── Upload to HuggingFace ───────────────────────────────────────── + if dry_run: + print("[dry-run] Skipping upload to HuggingFace.") + else: + print(f"Uploading to {HF_REPO_ID} ...") + api = HfApi() + create_repo(HF_REPO_ID, repo_type="dataset", exist_ok=True) + + with tempfile.NamedTemporaryFile(mode="w", suffix=".jsonl", delete=False) as f: + temp_path = f.name + for rec in records: + f.write(json.dumps(rec, ensure_ascii=False) + "\n") + + api.upload_file( + path_or_fileobj=temp_path, + path_in_repo="data/train-00000-of-00001.jsonl", + repo_id=HF_REPO_ID, + repo_type="dataset", + ) + Path(temp_path).unlink() + + print() + print("──────── Final stats ────────") + print(f" Total rows streamed: {total:,}") + print(f" Skipped (no question_id): {skipped_no_qid:,}") + print(f" Unique question_ids: {len(qid_reservoir):,}") + print(f" Removed by question-text dedup: {duplicate_questions:,}") + print(f" Final unique questions: {len(records):,}") + if dry_run: + print(f" Rows that WOULD upload: {len(records):,} (dry-run, not uploaded)") + print(f" HF dataset (target): https://huggingface.co/datasets/{HF_REPO_ID}") + else: + print(f" Rows uploaded to HF: {len(records):,}") + print(f" HF dataset: https://huggingface.co/datasets/{HF_REPO_ID}") + print("─────────────────────────────") + + +if __name__ == "__main__": + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument( + "--dry-run", + action="store_true", + help="Run preprocessing and assertions but skip the HuggingFace upload.", + ) + parser.add_argument("--seed", type=int, default=42, help="Random seed for reservoir sampling.") + args = parser.parse_args() + preprocess_and_upload(seed=args.seed, dry_run=args.dry_run) diff --git a/sdg/prompts.py b/sdg/prompts.py new file mode 100644 index 0000000000000000000000000000000000000000..546b0b492e6d6a59f80a7f0a14d050a1bbf3761b --- /dev/null +++ b/sdg/prompts.py @@ -0,0 +1,89 @@ +""" +Prompt templates for the SDG pipeline. + +All prompts target instruction-tuned (thinking) models and are copied verbatim +from experiments/self_instill/prompts.py (_INSTRUCT variants). +""" + +# ============================================================================= +# GENERATION +# ============================================================================= + +REASONING_INSTRUCTION = "Please reason step by step, and put your final answer within \\boxed{{}}." + +# ============================================================================= +# VALIDATION - Cycle Consistency +# ============================================================================= + +CYCLE_QUESTION_GENERATION_PROMPT = """Given an answer, please generate the most likely question that would have prompted this answer. Focus on inferring the core question that this answer is addressing. Output only the inferred question, without any additional explanation. + +Answer: +{answer} + +Inferred Question:""" + +CYCLE_COMPARISON_PROMPT = """You are evaluating whether an answer is relevant to the original question and touches the core of the question by comparing the original question with an inferred question derived only from the answer. + +Original Question: {original_question} +Inferred Question: {inferred_question} + +Compare the two questions and determine: +1. If the original question and inferred question are asking about the same core topic +2. If the original question and inferred question share the same key elements and requirements +3. If answering one question would effectively address the other question + +After your analysis, provide your decision: [[Y]] if the questions are semantically equivalent and address the same core problem, or [[N]] if they are asking about different things.""" + +# ============================================================================= +# VALIDATION - Factual Error Check +# ============================================================================= + +FACTUAL_ERROR_PROMPT = """Please act as an impartial judge and carefully analyze the following answer for any factual errors, logical flaws, or misleading information. + +Question: {question} +Answer: {answer} + +Consider the credibility of the claims made in the answer and determine if they align with established knowledge. Evaluate: +1. Are there any incorrect facts, dates, numbers, formulas, or claims? +2. Is there any faulty logic, reasoning, or problem-solving approach? +3. Are there any misleading, incomplete, or ambiguous explanations? +4. Does the answer introduce any misconceptions or propagate common errors? + +Minor typos or grammatical errors are acceptable. But be strict about any factual error, calculation error, or logical flaw. When unsure, lean toward accepting statements unless they contain clear errors. + +After a thorough analysis, provide your decision: [[Y]] if the answer has no factual errors or major flaws, or [[N]] if it contains important factual errors or logical flaws that would mislead the user.""" + +# ============================================================================= +# VALIDATION - Total Correctness +# ============================================================================= + +TOTAL_CORRECTNESS_PROMPT = """Please act as an impartial judge and evaluate whether the response is completely correct in both process and conclusion. + +Question: {question} +Answer: {answer} + +Consider correctness, usefulness, completeness and depth in your assessment. Consider whether this answer completely solves the question. + +You should rely on your own reasoning to form a reference solution and compare the answer to your reasoning. + +Begin your evaluation by giving a brief summary of your thoughts on the response. Focus on whether it is accurate, addresses the question well, and is reasonably detailed. Be precise about any errors or gaps you notice. + +Notes: +1. If the answer is partial, high-level, or just states that this is an open problem, you should not accept it. +2. If the answer lacks details or is not comprehensive, you should not accept it. +3. If the answer contains any errors, you should not accept it. +4. You should only accept the answer if it is at least 95% correct and solves the question. + +After providing your explanation, decide whether this answer is correct. Think twice about whether this answer solves the question. +Format: Accepted: [[Y]] if you accept the answer or Accepted: [[N]] if you do not accept.""" + +# ============================================================================= +# VALIDATION - Simple (single yes/no judgement) +# ============================================================================= + +SIMPLE_CORRECTNESS_PROMPT = """Please act as an impartial judge and evaluate whether the following answer is accurate and correct given the question. + +Question: {question} +Answer: {answer} + +Provide your decision: [[Y]] if the answer is accurate and correct, or [[N]] if it is not.""" diff --git a/sdg/questions/__init__.py b/sdg/questions/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/sdg/questions/chimera_stats.json b/sdg/questions/chimera_stats.json new file mode 100644 index 0000000000000000000000000000000000000000..c01bb674f679cfdf32d4a7e49e34ac4d24e65469 --- /dev/null +++ b/sdg/questions/chimera_stats.json @@ -0,0 +1,58 @@ +{ + "Biology": { + "avg_questions_per_topic": 7.98, + "max_num_questions_per_topic": 8, + "min_num_questions_per_topic": 7, + "num_questions": 383, + "num_unique_topics": 48 + }, + "Chemistry": { + "avg_questions_per_topic": 7.93, + "max_num_questions_per_topic": 8, + "min_num_questions_per_topic": 7, + "num_questions": 1102, + "num_unique_topics": 139 + }, + "Computer Science": { + "avg_questions_per_topic": 7.9, + "max_num_questions_per_topic": 8, + "min_num_questions_per_topic": 6, + "num_questions": 1303, + "num_unique_topics": 165 + }, + "History": { + "avg_questions_per_topic": 7.96, + "max_num_questions_per_topic": 8, + "min_num_questions_per_topic": 7, + "num_questions": 422, + "num_unique_topics": 53 + }, + "Linguistics": { + "avg_questions_per_topic": 7.92, + "max_num_questions_per_topic": 8, + "min_num_questions_per_topic": 7, + "num_questions": 317, + "num_unique_topics": 40 + }, + "Literature": { + "avg_questions_per_topic": 7.88, + "max_num_questions_per_topic": 8, + "min_num_questions_per_topic": 7, + "num_questions": 504, + "num_unique_topics": 64 + }, + "Mathematics": { + "avg_questions_per_topic": 7.76, + "max_num_questions_per_topic": 8, + "min_num_questions_per_topic": 4, + "num_questions": 4452, + "num_unique_topics": 574 + }, + "Physics": { + "avg_questions_per_topic": 7.73, + "max_num_questions_per_topic": 8, + "min_num_questions_per_topic": 6, + "num_questions": 742, + "num_unique_topics": 96 + } +} \ No newline at end of file diff --git a/sdg/questions/chimera_topics.json b/sdg/questions/chimera_topics.json new file mode 100644 index 0000000000000000000000000000000000000000..73c3ea7d7060a24da62489378a8fa32df667bcc0 --- /dev/null +++ b/sdg/questions/chimera_topics.json @@ -0,0 +1,1197 @@ +{ + "Biology": [ + "Aging, senescence, and geroscience", + "Agricultural biology and crop science", + "Astrobiology and origins of life", + "Behavioral ecology and ethology", + "Bioethics and policy in biotechnology", + "Biomaterials and tissue engineering", + "Biomechanics and mechanobiology", + "Biophysics and single-molecule biophysics", + "CRISPR technologies and gene editing", + "Cancer biology and tumor microenvironment", + "Cell biology and organelle dynamics", + "Circadian biology and chronobiology", + "Cognitive neuroscience and connectomics", + "Computational neuroscience and brain-inspired models", + "Conservation biology and biodiversity", + "Developmental biology and morphogenesis", + "Ecology (population, community, ecosystem)", + "Endocrinology and metabolic disease", + "Epidemiology and disease ecology", + "Epigenetics and chromatin biology", + "Evolutionary biology and evo-devo", + "Gene drives and population engineering", + "Genetics and genomics", + "High-throughput screening and functional genomics", + "Immunology and immuno-oncology", + "Interactomics and network biology", + "Marine and freshwater biology", + "Metabolomics and systems metabolism", + "Microbiology (bacteriology, virology, mycology, parasitology)", + "Microbiome and host–microbe interactions", + "Molecular biology and gene regulation", + "Neurobiology and systems neuroscience", + "Paleobiology and evolutionary paleontology", + "Pharmacology and toxicology", + "Photosynthesis and plant metabolism", + "Phylogenetics and comparative genomics", + "Plant biology and plant immunity", + "Population genetics and quantitative genetics", + "Proteomics and protein science", + "Quantitative and single-cell biology", + "Spatial omics and imaging genomics", + "Stem cell biology and regenerative medicine", + "Structural biology (X-ray, NMR, cryo-EM)", + "Synthetic biology and genome engineering", + "Systems biology and dynamical modeling", + "Systems immunology and immunometabolism", + "Transcriptomics and RNA biology", + "Virology and viral evolution" + ], + "Chemistry": [ + "Allostery and enzyme regulation", + "Analytical chemistry and instrumentation", + "Astrochemistry and prebiotic chemistry", + "Asymmetric synthesis and chiral pool strategies", + "Atmospheric chemistry and aerosols", + "Autophagy and apoptosis", + "Battery chemistry and solid electrolytes", + "Biocatalysis and industrial biochemistry", + "Biochemistry and chemical biology", + "Bioenergetics and oxidative phosphorylation", + "Bioinorganic chemistry and metalloenzymes", + "Biosynthetic pathway-inspired synthesis", + "Carbocation and carbanion chemistry", + "Carbohydrate and peptide synthesis", + "Carbohydrate chemistry and glycosylation", + "Cascade and multicomponent reactions", + "Catalysis (heterogeneous, homogeneous, biocatalysis)", + "Catalyst design and ligand development", + "Chaperones and protein quality control", + "Chemical biology and bioorthogonal chemistry", + "Chemical genetics and epigenetic probes", + "Chemical kinetics and reaction dynamics", + "Cheminformatics and QSAR", + "Chemoenzymatic synthesis and biocatalysis", + "Chromatography and separations", + "Click chemistry", + "Click chemistry and bioorthogonal reactions", + "Coenzymes and cofactors", + "Computational biochemistry and molecular dynamics", + "Computational reaction mechanisms and QM/MM", + "Coordination chemistry and ligand design", + "Covalent inhibitors and PROTACs", + "Covalent inhibitors and targeted degradation (PROTACs)", + "Cross-coupling (Suzuki, Heck, Negishi, Buchwald–Hartwig)", + "Crystal engineering and MOFs/COFs", + "Crystallization and polymorphism", + "Cytokine signaling and immunometabolism", + "C–H activation and functionalization", + "C–H activation and functionalization strategies", + "DNA repair mechanisms and genome stability", + "Density functional theory and ab initio methods", + "Directed evolution and enzyme engineering", + "Docking and structure-based drug design", + "Electrochemistry and electrocatalysis", + "Electroorganic synthesis", + "Energy catalysis and CO2 utilization", + "Environmental chemistry and toxicology", + "Enzyme mechanisms and catalysis", + "Enzymology and enzyme kinetics", + "Epigenetics and chromatin remodeling", + "Flow chemistry and microfluidic synthesis", + "Flow chemistry and microreactors", + "Flux analysis and metabolic control analysis", + "Forensic chemistry and analytical forensics", + "Fragment-based and structure-based design", + "Fragment-based drug discovery", + "Geochemistry and cosmochemistry", + "Glycobiology and glycoproteins", + "Glycolipids and glycoconjugates", + "Green chemistry and sustainable processes", + "Green solvents and sustainable synthesis", + "Heterocyclic chemistry", + "Heterocyclic chemistry (N, O, S heterocycles)", + "Hydrogenation and asymmetric catalysis", + "Hypervalent iodine chemistry", + "Immunochemistry and antigen–antibody interactions", + "Industrial catalysis and reactor design", + "Inorganic chemistry", + "Ionic liquids and deep eutectic solvents", + "Lipid nanoparticles and drug delivery", + "Machine learning in chemistry and autonomous labs", + "Macrocyclization and supramolecular synthesis", + "Main group chemistry and hypervalent compounds", + "Mass spectrometry and proteomics", + "Medicinal chemistry and SAR", + "Medicinal chemistry and drug design", + "Membrane structure and lipidomics", + "Membrane transporters and ion channels", + "Metabolic engineering and synthetic pathways", + "Metabolism and metabolic pathways", + "Mitochondrial dynamics and mitophagy", + "NMR, EPR, IR, Raman, UV-Vis spectroscopy", + "Nanochemistry and colloids", + "Nitrogen metabolism and assimilation", + "Nuclear chemistry and radiochemistry", + "Nucleic acids biochemistry (DNA and RNA)", + "Nucleoside and nucleotide chemistry", + "Olefin metathesis and ring-closing metathesis", + "Organic chemistry and synthesis", + "Organic electronics (OLED, OPV, OFET)", + "Organic materials and conjugated polymers", + "Organocatalysis (enamine, iminium, NHC catalysis)", + "Organocatalysis and enamine/iminium catalysis", + "Organofluorine and perfluoroalkylation", + "Organometallic chemistry and cross-coupling", + "Organometallics in organic synthesis", + "Peptide synthesis and solid-phase methods", + "Pericyclic reactions and orbital symmetry", + "Perovskites and photovoltaic materials", + "Pharmaceutical biotechnology and biologics", + "Photochemistry and photophysics", + "Photopharmacology and light-controlled drugs", + "Photoredox and electroorganic synthesis", + "Photoredox catalysis and energy transfer catalysis", + "Photosynthesis biochemistry", + "Physical chemistry (thermodynamics, kinetics, spectroscopy)", + "Physical organic chemistry and reaction mechanisms", + "Polymer chemistry and soft materials", + "Post-translational modifications and proteostasis", + "Prions, amyloids, and protein aggregation", + "Process chemistry and scale-up", + "Protecting group strategies and chemoselectivity", + "Protein structure, dynamics, and folding", + "Proteomics and mass spectrometry", + "QM/MM and free energy calculations", + "Quantum chemistry and electronic structure theory", + "RNA processing, splicing, and RNA-based regulation", + "Radical chemistry and persistent radicals", + "Reaction discovery with ML and autonomous labs", + "Reaction mechanisms and physical organic chemistry", + "Redox biology and ROS signaling", + "Replication, transcription, and translation", + "Retrosynthetic analysis and strategy", + "Safety, regulation, and chemical education", + "Second messengers (cAMP, IP3, Ca2+)", + "Self-assembly and host–guest chemistry", + "Signal transduction (GPCRs, kinases, phosphatases)", + "Solid-state chemistry and materials", + "Statistical mechanics and molecular simulation", + "Stereochemistry and asymmetric catalysis", + "Stereochemistry and conformational analysis", + "Structural methods (X-ray, NMR, cryo-EM)", + "Structure, bonding, and stereoelectronics", + "Supramolecular chemistry and host–guest systems", + "Surface science and catalysis characterization", + "Systems biochemistry and metabolic regulation", + "Theoretical and computational chemistry", + "Total synthesis of natural products", + "X-ray crystallography and diffraction" + ], + "Computer Science": [ + "AI for systems and systems for AI", + "AR, VR, and XR interfaces", + "Accessibility and inclusive design", + "Active learning and optimal data selection", + "Algorithms and data structures", + "Animation and physics-based simulation", + "Approximation algorithms and PTAS", + "AutoML and hyperparameter optimization", + "Bayesian deep learning and uncertainty estimation", + "Bilevel optimization and hyperparameter tuning", + "Binary analysis and software hardening", + "Blockchain and distributed ledger technologies", + "Category theory in CS", + "Causal inference and causal discovery in CS", + "Causal representation learning", + "Causality in ML and causal discovery", + "Cloud computing and virtualization", + "Computability and automata theory", + "Computational biology and bioinformatics", + "Computational chemistry and materials informatics", + "Computational complexity theory", + "Computational economics and market design", + "Computational imaging and computational photography", + "Computational neuroscience", + "Computer architecture and hardware-software co-design", + "Computer graphics and rendering", + "Computer networks and programmable data planes", + "Computer vision and multimodal learning", + "Computer vision and pattern recognition", + "Computer-supported cooperative work", + "Convolutional, recurrent, and attention-based networks", + "Cryptography and cryptographic protocols", + "Cyber-physical systems and real-time computing", + "Data mining and knowledge discovery", + "Data-centric AI and data governance", + "Data-centric AI, data governance, and quality", + "Database systems and data engineering", + "Datacenter architectures and disaggregated systems", + "Deep learning and representation learning", + "Deep learning architectures", + "Differential privacy and federated learning", + "Differential privacy and secure computation", + "Digital currencies and cryptoeconomics", + "Digital twins and cyber-physical security", + "Distributed systems and consensus", + "Domain generalization and distribution shift", + "Edge AI and on-device learning", + "Edge computing and IoT", + "Education technologies and computing education research", + "Embedded systems and robotics platforms", + "Energy-aware and green computing", + "Energy-based models and score matching", + "Fairness, accountability, and ethics", + "Fairness, accountability, and transparency in AI", + "Federated and decentralized learning", + "Formal languages and parsing", + "Formal methods and model checking", + "Foundation models and large language models", + "GPU and accelerator computing", + "Generative chemistry and materials discovery", + "Generative modeling", + "Geometric and topological data analysis", + "Geometry processing and CAD", + "Graph algorithms and network algorithms", + "Graph and relational learning for science", + "Graph neural networks and geometric deep learning", + "Hierarchical and multi-agent RL", + "Human-AI interaction and interactive ML", + "Human-computer interaction and UX", + "Implicit neural representations and neural fields", + "In-memory and approximate computing", + "Information retrieval and search", + "Information-theoretic learning bounds", + "Instruction tuning, RLHF, and RLAIF", + "Interpretability and explainability", + "Inverse RL and imitation learning", + "Knowledge representation and reasoning", + "Lambda calculus and functional programming theory", + "Large language models and foundation models", + "Learning PDE solvers and neural operators", + "Logic in computer science", + "MLOps and ML systems", + "MLOps, deployment, and monitoring", + "Machine learning and data-driven methods", + "Mechanistic interpretability of neural networks", + "Meta-learning and few-shot learning", + "Model compression, pruning, quantization, distillation", + "Multi-task and continual learning", + "Multiagent systems and game-theoretic AI", + "Natural language processing and LLMs", + "Natural language processing and computational linguistics", + "Networking for AI workloads", + "Neural architecture search and AutoML", + "Neurosymbolic AI and program synthesis", + "Nonconvex optimization and landscape analysis", + "Offline and batch RL", + "Online algorithms and competitive analysis", + "Online learning and bandits", + "Operating systems and kernels", + "Optimization for ML", + "PAC learning and VC theory", + "Parallel and high-performance computing", + "Parameterized and fine-grained complexity", + "Perception, SLAM, and control in robotics", + "Physics-informed ML and operator learning", + "Planning, search, and constraint programming", + "Policy gradients and actor-critic methods", + "Post-Moore architectures and neuromorphic computing", + "Post-quantum cryptography", + "Privacy engineering and compliance", + "Probabilistic modeling and graphical models", + "Probabilistic programming and inference engines", + "Program analysis and synthesis", + "Program synthesis and neuro-symbolic systems", + "Programming languages and compilers", + "Proof assistants and program verification", + "Protein structure prediction and generative biology", + "Quantum programming languages and compilers", + "Randomized algorithms and probabilistic methods", + "Recommender systems and personalization", + "Reinforcement learning (model-free and model-based)", + "Reinforcement learning and decision-making", + "Responsible and safe AI", + "Retrieval-augmented generation and tool use", + "Robotics and autonomous systems", + "Robustness and adversarial ML", + "Runtime systems and garbage collection", + "Safe and risk-sensitive RL", + "Scientific computing and numerical methods", + "Scientific machine learning", + "Second-order and adaptive optimization", + "Secure hardware and side-channel defenses", + "Security and privacy", + "Self-supervised and contrastive learning", + "Semi-supervised and weakly supervised learning", + "Serverless computing and microservices", + "Smart contracts and formal verification", + "Social computing and computational social science", + "Software engineering and software architecture", + "Software testing, fuzzing, and reliability", + "Software-defined networking and NFV", + "Speech and audio processing", + "Speech, audio, and music ML", + "Stability and compression-based generalization", + "Stochastic gradient methods and variance reduction", + "Storage systems and filesystems", + "Supervised learning (classification and regression)", + "Surrogate modeling and uncertainty quantification", + "Systems and network security", + "Theoretical ML and generalization", + "Theoretical ML and statistical learning theory", + "Time series forecasting and anomaly detection", + "Transaction processing and concurrency control", + "Transfer learning and domain adaptation", + "Transformers and sequence modeling", + "Trustworthy AI and robustness", + "Type systems and program semantics", + "Ubiquitous and wearable computing", + "Unsupervised learning (clustering, density estimation)", + "Usable security and human factors", + "VAEs, GANs, normalizing flows, and diffusion models", + "Visualization and visual analytics", + "Web and internet protocols", + "Wireless, mobile, and 5G/6G networks", + "World models and model-based planning" + ], + "History": [ + "African history (regional and thematic)", + "Ancient Near Eastern history", + "Archival studies and paleography", + "Big history and deep history", + "Business history and financial history", + "Classical Greek and Roman history", + "Cliometrics and quantitative history", + "Contemporary history", + "Cultural heritage, archives, and public history", + "Cultural history and history of the senses", + "Demographic history and family history", + "Digital history and computational humanities", + "Diplomatic history and international relations", + "Early modern history", + "East Asian history (China, Japan, Korea)", + "Economic history and the history of capitalism", + "Environmental and climate history", + "Epigraphy and numismatics", + "European history (Western, Central, Eastern)", + "Food history and culinary cultures", + "Gender history and women’s history", + "Genocide studies and Holocaust history", + "Geographic information systems in history", + "Historiography and theory of history", + "History of medicine and public health", + "History of science and technology", + "Indigenous and First Nations histories", + "Intellectual history and history of ideas", + "Labor history and working-class studies", + "Late antiquity and early medieval history", + "Latin American and Caribbean history", + "Legal and constitutional history", + "Material culture and consumption", + "Medieval history", + "Memory studies and memorialization", + "Methods in historical research and source criticism", + "Microhistory and local history", + "Middle Eastern and Islamic history", + "Military history and war studies", + "Modern history", + "North American history", + "Oral history and life histories", + "Political history and state formation", + "Prehistory and prehistoric archaeology", + "Queer histories and sexuality", + "Religious history and church history", + "Remote sensing and landscape archaeology", + "Russian and Soviet history", + "Social history and everyday life", + "South Asian and Southeast Asian history", + "Transnational and diaspora studies", + "Urban and rural history", + "World history and global history" + ], + "Linguistics": [ + "Anthropological linguistics and ethnography of speaking", + "Bilingualism and multilingualism", + "Cognitive linguistics and conceptual metaphor", + "Comparative and contrastive linguistics", + "Computational linguistics and NLP", + "Corpus linguistics and big data", + "Dialectology and microvariation", + "Discourse analysis and conversation analysis", + "Distributional semantics and embeddings", + "Educational linguistics and literacy", + "Endangered languages and revitalization", + "Experimental methods (eye-tracking, EEG/MEG, fMRI)", + "Field linguistics and language documentation", + "Forensic linguistics and language policy", + "Gesture and multimodal communication", + "Grammaticalization and contact-induced change", + "Historical linguistics and language change", + "Information structure and focus", + "Language acquisition (first and second)", + "Language typology and universals", + "Lexicography and lexicology", + "Morphology (inflectional and derivational)", + "Neurolinguistics and language in the brain", + "Parsing, tagging, and morphological analysis", + "Philosophy of language and formal pragmatics", + "Phonetics (articulatory, acoustic, auditory)", + "Phonology (theory and phonology-phonetics interface)", + "Pragmatics (speech acts, implicature, relevance)", + "Pragmatics in conversational AI", + "Prosody, stress, and intonation", + "Psycholinguistics and sentence processing", + "Semantic change and diachronic semantics", + "Semantics (formal, lexical, and dynamic)", + "Sign language linguistics", + "Sociolinguistics and language variation", + "Sociophonetics and indexical meaning", + "Speech recognition and text-to-speech", + "Syntax (generative, dependency, construction grammar)", + "Word formation and lexical morphology", + "Writing systems, orthography, and graphemics" + ], + "Literature": [ + "Adaptation studies and intermediality", + "African, Caribbean, and diasporic literatures", + "American and European literatures", + "Ancient Near Eastern literatures", + "Archive studies and critical bibliography", + "Book history and print culture", + "Children’s and young adult literature", + "Classical literatures (Greek, Latin)", + "Cognitive literary studies and neuro-literary approaches", + "Contemporary and global anglophone literature", + "Corpus approaches to literature", + "Critical race and ethnic studies", + "Cultural studies and semiotics", + "Detective and crime fiction", + "Digital humanities and computational literary studies", + "Disability studies and medical humanities", + "East Asian literatures (Chinese, Japanese, Korean)", + "Ecocriticism and environmental humanities", + "Enlightenment literature", + "Evolutionary literary studies", + "Feminist, gender, and queer theory", + "Film, theatre, and performance studies", + "Game narratives and interactive fiction", + "Genre studies (poetry, drama, novel, short story)", + "Graphic narratives and comics studies", + "Hermeneutics and interpretation", + "Holocaust and genocide literatures", + "Horror and Gothic literatures", + "Indigenous and First Nations literatures", + "Latin American and Iberian literatures", + "Law and literature", + "Life writing, autobiography, and memoir", + "Literary theory and criticism", + "Manuscript studies, codicology, and paleography", + "Medieval literature", + "Middle Eastern and Islamic literatures", + "Modernist and postmodernist literature", + "Narratology and discourse analysis", + "New formalism and historical poetics", + "New historicism and cultural materialism", + "Open scholarship and digital editions", + "Orality, folklore, and myth", + "Philology and historical linguistics interface", + "Platform studies and e-literature", + "Poetics and rhetoric", + "Poetry and prosody", + "Postcolonial studies and decolonial theory", + "Reader-response and affect theory", + "Renaissance and early modern literature", + "Romanticism and Victorian literature", + "Science and literature", + "Science fiction, fantasy, and speculative fiction", + "Shakespeare and early modern drama", + "Social media narratives and fan fiction", + "South Asian and Southeast Asian literatures", + "Structuralism, poststructuralism, and deconstruction", + "Stylometry and authorship attribution", + "Text mining and distant reading", + "Textual criticism, stemmatics, and editorial theory", + "Translation studies and comparative poetics", + "Trauma and memory studies", + "Travel writing and nature writing", + "World literature and comparative literature", + "World-systems and global literary circulation" + ], + "Mathematics": [ + "3-manifolds and hyperbolic structures", + "4-manifolds and smooth structures", + "Abelian varieties, Jacobians, and Pryms", + "Absoluteness and reflection principles", + "Actuarial science and risk theory", + "Additive combinatorics", + "Additive combinatorics and sum-product phenomena", + "Additive number theory and sumset problems", + "Adic and perfectoid spaces", + "Algebra", + "Algebraic coding theory", + "Algebraic combinatorics", + "Algebraic cycles and Chow groups", + "Algebraic geometry and moduli", + "Algebraic geometry and tropical combinatorics", + "Algebraic number theory", + "Algebraic spaces and stacks", + "Algebraic topology", + "Algebraic topology and homotopy theory", + "Algorithmic class field theory", + "Algorithms and randomized algorithms", + "Analysis", + "Analytic combinatorics and singularity analysis", + "Analytic number theory", + "Applications to analysis and Banach spaces", + "Applied combinatorics in networks and data", + "Applied probability", + "Applied topology and persistent homology", + "Approximate counting and correlation decay", + "Approximation algorithms for graph problems", + "Approximation theory and splines", + "Arithmetic combinatorics", + "Arithmetic geometry and Arakelov theory", + "Arithmetic geometry and Diophantine geometry", + "Arithmetic statistics", + "Arithmetic statistics and Cohen–Lenstra heuristics", + "Association schemes and distance-regular graphs", + "Automated reasoning and theorem proving", + "Automated theorem proving for mathematics", + "Automorphic forms and the Langlands program", + "Automorphism groups and Cayley graphs", + "Axiomatic frameworks and foundations of mathematics", + "Axiomatic set theory (ZFC and alternatives)", + "BSD conjecture and L-series of elliptic curves", + "Backward SDE and FBSDE", + "Bandits and online learning", + "Bayesian nonparametrics and Dirichlet processes", + "Bayesian statistics and Bayesian nonparametrics", + "Bijective combinatorics", + "Bioinformatics and phylogenetics", + "Biophysics and neuroscience models", + "Biostatistics and clinical trials", + "Bipartite and multipartite graphs", + "Birational geometry and minimal model program", + "Blockchain primitives and number-theoretic protocols", + "Borel equivalence relations and classification theory", + "Branching processes and coalescent theory", + "Brauer groups and central simple algebras", + "Bridgeland stability conditions", + "Brill–Noether theory and syzygies", + "Brownian motion and Gaussian processes", + "CAT(0) and Alexandrov spaces", + "Calabi–Yau manifolds and special holonomy", + "Calabi–Yau varieties and K3 surfaces", + "Calculus of variations", + "Calibrations and special Lagrangians", + "Cardinal characteristics of the continuum", + "Category theory and topos theory", + "Causal inference under uncertainty", + "Chabauty–Coleman methods and p-adic integration", + "Chaos and bifurcation theory", + "Character sheaves and Springer theory", + "Cheeger inequalities and isoperimetry on graphs", + "Chip-firing games and sandpile groups", + "Chordal, interval, and comparability graphs", + "Chromatic invariants and Tutte polynomials", + "Class field theory", + "Classical algebraic geometry (varieties and maps)", + "Classifying spaces and characteristic classes", + "Clique-width and rank-width", + "CoHAs and categorification in topology", + "Coarse and controlled topology", + "Cobordism theory and Thom spectra", + "Coding theory and combinatorial designs", + "Coding theory and compressed sensing", + "Coding theory and number-theoretic constructions", + "Combinatorial commutative algebra", + "Combinatorial optimization", + "Combinatorial representation theory", + "Combinatorial set theory", + "Combinatorics for ML and graph neural networks", + "Combinatorics in statistical physics (Ising, Potts, dimers)", + "Combinatorics of polynomials and real stable polynomials", + "Commutative algebra", + "Comparison geometry and Ricci curvature", + "Complex and Kähler geometry", + "Computability and recursion theory", + "Computational algebra and Gröbner bases", + "Computational algebraic geometry and Gröbner bases", + "Computational and algorithmic number theory", + "Computational and applied mathematics", + "Computational geometry", + "Concentration of measure and isoperimetry", + "Conformal and CR geometry", + "Conformal and topological quantum field theories", + "Conformal prediction and uncertainty quantification in ML", + "Constructible universe L and fine structure", + "Constructive and intuitionistic mathematics", + "Contact geometry and Reeb dynamics", + "Container method and hypergraph containers", + "Continuum theory and Peano continua", + "Control and observability on networks", + "Control theory and systems theory", + "Convex geometry and Brunn–Minkowski theory", + "Core model theory and extender models", + "Coupling methods and contractivity", + "Covering spaces and fundamental group", + "Cryptography and algebraic cryptography", + "Cryptography and post-quantum cryptography", + "Data science and manifold learning", + "Decorated character varieties and higher Teichmüller theory", + "Deep learning theory and generalization", + "Deligne–Mumford and Artin stacks", + "Derived categories of coherent sheaves", + "Descriptive set theory", + "Descriptive set theory (Borel and projective hierarchies)", + "Descriptive set-theoretic dynamics", + "Design theory and block designs", + "Design theory and finite geometries", + "Determinacy axioms (AD, ADL(R))", + "Differential geometry", + "Diophantine approximation and transcendence", + "Diophantine geometry and arithmetic geometry", + "Discrete Morse theory", + "Discrete and computational geometry", + "Discrete conformal geometry and circle packings", + "Discrete geometry and incidence geometry", + "Discrete geometry and polyhedral combinatorics", + "Discrete mathematics", + "Disordered systems and spin glasses", + "Distance-regular graphs and association schemes", + "Distributed and local graph algorithms", + "Distribution of rational points and unlikely intersections", + "Donaldson and Seiberg–Witten theory", + "Donaldson–Thomas and Pandharipande–Thomas invariants", + "Econophysics and financial mathematics", + "Edge coloring and total coloring", + "Edge expansion, vertex expansion, and conductance", + "Effective descriptive set theory", + "Electrical networks and effective resistance", + "Elementary number theory", + "Elliptic curve methods and isogeny computations", + "Elliptic curves and abelian varieties", + "Embeddings and metric distortion", + "Emerging directions", + "Empirical processes and Donsker theorems", + "Empirical processes and concentration bounds", + "Enumerative and extremal combinatorics", + "Enumerative combinatorics and generating functions", + "Enumerative geometry and Hurwitz theory", + "Environmental and climate risk modeling", + "Epidemic models and branching processes", + "Epidemics and contagion on networks", + "Epidemiology and infectious disease modeling", + "Erdos distinct distances and Szemerédi–Trotter", + "Erdos–Ko–Rado and intersecting families", + "Ergodic methods on homogeneous spaces", + "Ergodic theory and dynamical systems", + "Ergodic theory and mixing", + "Euclidean and classical geometry", + "Exchangeability and de Finetti theorems", + "Expander constructions and pseudorandomness", + "Expander constructions via number theory", + "Expanders and expander constructions", + "Exponential sums and the circle method", + "Extremal combinatorics", + "Extreme value theory and heavy tails", + "Extreme value theory and regular variation", + "Face numbers, f-vectors, and h-vectors", + "Facility location and scheduling", + "Fairness and privacy in algorithms", + "Fairness and robustness via probabilistic modeling", + "Faltings’ theorem and Mordell-type problems", + "Fano varieties and rationality problems", + "Federated Bayesian learning and distributed inference", + "Fibre bundles and principal bundles", + "Field theory and Galois theory", + "Filtering and state estimation (Kalman, particle filters)", + "Financial mathematics and derivatives pricing", + "Fine-grained lower bounds for combinatorial problems", + "Finite element and spectral methods", + "Finite geometries and incidence structures", + "Floer homology and pseudo-holomorphic curves", + "Flows, cuts, and circulations", + "Flows, cuts, and sparsifiers", + "Foliations, laminations, and measured laminations", + "Forcing and iterated forcing", + "Formal methods in cyber-physical systems", + "Fractal geometry", + "Fractional coloring and circular coloring", + "Free probability and noncommutative probability", + "Free probability in applications", + "Functional CLTs and invariance principles", + "Functional analysis and operator theory", + "Functional inequalities (Poincaré, log-Sobolev, transport)", + "Galois representations and deformation theory", + "Galois theory of number fields", + "Game theory and mechanism design", + "Gaussian processes and kernel methods", + "General relativity and geometric PDE", + "Generalized metric spaces and compactness", + "Generators, Dirichlet forms, and potential theory", + "Generic absoluteness and forcing axioms (MA, PFA, MM)", + "Geometric Langlands program", + "Geometric analysis", + "Geometric analysis and PDE on manifolds", + "Geometric deep learning and graph neural networks", + "Geometric graphs and unit disk graphs", + "Geometric group theory", + "Geometric group theory and quasi-isometries", + "Geometric measure theory", + "Geometric measure theory and currents", + "Geometric set cover and epsilon-nets", + "Geometric topology", + "Geometric topology interfaces", + "Geometrization and 3-manifold geometry", + "Geometry and topology", + "Geometry of Banach spaces", + "Geometry of numbers and lattice point problems", + "Gibbs measures, spin systems, and phase transitions", + "Graph coloring and list coloring", + "Graph compression and summarization", + "Graph databases and knowledge graphs", + "Graph decomposition and minors structure", + "Graph drawing and visualization", + "Graph embeddings and metric embeddings", + "Graph homomorphisms and constraint satisfaction", + "Graph isomorphism and canonical labeling", + "Graph limits and graphons", + "Graph minors and Robertson–Seymour theory", + "Graph polynomials (Tutte, chromatic, interlace)", + "Graph products and powers", + "Graph representation learning and GNNs", + "Graph theory and spectral graph theory", + "Graph theory and structural graph theory", + "Greedoid theory", + "Gromov–Witten theory and quantum cohomology", + "Group theory", + "Hamiltonian systems and symplectic dynamics", + "Hardness of approximation and PCP connections", + "Harmonic analysis and Fourier analysis", + "Heegaard Floer homology and contact invariants", + "Heights, Arakelov geometry, and Northcott property", + "High-dimensional probability and concentration inequalities", + "High-dimensional probability for data science", + "Hilbert schemes and Quot schemes", + "Hodge theory", + "Hodge theory and period mappings", + "Hodge theory and variations of Hodge structure", + "Holographic and Holant frameworks", + "Homological algebra", + "Homological mirror symmetry", + "Homology, cohomology, and cohomology operations", + "Homotopical algebra and model categories", + "Homotopy theory and Postnikov towers", + "Homotopy type theory in proof assistants", + "Hopf algebras and quantum groups", + "Hydrology and extreme events", + "Hyperbolic geometry and Kleinian groups", + "Hypergraph theory", + "Hypergraph theory and set systems", + "Hypergraphs and uniformity", + "Ideals on omega and forcing with ideals", + "Imaging science and tomography", + "Index theory and Atiyah–Singer theorem", + "Infinity categories and higher topos theory", + "Information geometry and statistical manifolds", + "Information geometry for learning", + "Information theory and entropy methods", + "Information-theoretic inequalities and entropy methods", + "Inner model theory and mice", + "Insurance mathematics and solvency modeling", + "Integrable probability and last passage percolation", + "Integrable systems", + "Integral geometry and geometric probability", + "Integral polyhedra and total unimodularity", + "Interacting particle systems and hydrodynamic limits", + "Interacting particle systems and mean-field models", + "Interactive proof assistants and formalization of mathematics", + "Interdisciplinary and applied areas", + "Intersection theory and Poincaré duality", + "Intersection theory and Riemann–Roch", + "Invariant theory", + "Inverse problems and regularization", + "Iwasawa theory", + "K-theory (algebraic and topological)", + "K-theory (topological and algebraic)", + "KPZ universality and integrable probability", + "Kazhdan–Lusztig theory and Hecke algebras", + "Knot theory and link invariants", + "Knot theory and low-dimensional topology", + "L-functions and functional equations", + "L-functions and zeta functions", + "Langlands program (classical and geometric)", + "Large cardinal axioms", + "Large cardinals and elementary embeddings", + "Large deviations and Gartner–Ellis theorem", + "Latin squares and orthogonal arrays", + "Lattice algorithms and LLL", + "Levy processes and jump diffusions", + "Lie groups and Lie algebras", + "Linear systems and base loci", + "Liouville quantum gravity and random geometry", + "Local fields and ramification", + "Local limit theorems and Berry–Esseen bounds", + "Logarithmic geometry", + "Low-dimensional topology", + "Manifold topology and smoothing theory", + "Manin conjecture and Peyre constants", + "Mapping class groups and surfaces", + "Markov chain Monte Carlo and mixing analysis", + "Markov chains and Markov processes", + "Markov processes and semigroup theory", + "Martingales and optional stopping", + "Matching theory and matroid intersection", + "Matching, matroid parity, and factors", + "Mathematical biology and epidemiology", + "Mathematical physics", + "Matroid theory and oriented matroids", + "Mean curvature flow and Ricci flow", + "Mean field games", + "Mean-field models and McKean–Vlasov equations", + "Measurable cardinals and ultrafilters", + "Measure theory and integration", + "Measure-theoretic probability and integration", + "Measure-theoretic probability and stochastic processes", + "Metric geometry and Gromov hyperbolicity", + "Microlocal analysis and pseudodifferential operators", + "Minimal surfaces and geometric flows", + "Minimal surfaces and mean curvature flow", + "Mirror symmetry and Fukaya categories", + "Mirror symmetry and SYZ", + "Mixed Hodge structures and degenerations", + "Model theory", + "Model theory in number theory and o-minimality", + "Moderate deviations and LIL", + "Modular forms and automorphic forms", + "Moduli of curves and stable maps", + "Moduli problems and geometric invariant theory", + "Monte Carlo and Markov chain mixing", + "Monte Carlo methods and variance reduction", + "Motives and motivic cohomology", + "Multiplier ideals and singular metrics", + "Network design and survivable networks", + "Network design, reliability, and survivability", + "Network inference and community detection", + "Network science and complex systems", + "Network science and multilayer networks", + "Neural differential equations and neural operators", + "Non-Euclidean geometry (hyperbolic and spherical)", + "Noncommutative algebra", + "Noncommutative algebraic geometry", + "Noncommutative topology and C*-algebraic topology", + "Nonlinear analysis", + "Nonstandard analysis", + "Number theory", + "Numerical analysis and scientific computing", + "Numerical linear algebra", + "O-minimality and definable sets in AG", + "Obstruction theory", + "Operations research and mathematical programming", + "Operations research under uncertainty", + "Operator algebras (C*-algebras, von Neumann algebras)", + "Optimal experimental design and active learning", + "Optimal transport and Wasserstein geometry", + "Optimal transport and probability metrics", + "Optimal transport for probabilistic modeling", + "Optimization (convex, nonconvex, integer, semidefinite)", + "Orbifolds and stacks", + "Order dimension and comparability graphs", + "PCF theory and singular cardinals", + "Parameterized algorithms and kernels on graphs", + "Parameterized complexity and kernelization", + "Partial differential equations (elliptic, parabolic, hyperbolic)", + "Partition calculus and infinite combinatorics", + "Percolation and interacting particle systems", + "Percolation and phase transitions", + "Percolation and phase transitions on networks", + "Perfect graphs and the strong perfect graph theorem", + "Perfectoid spaces and p-adic Hodge theory", + "Permutation patterns and sorting networks", + "Perpetuities and fixed-point equations in distribution", + "Perverse sheaves and D-modules", + "Perverse sheaves and microlocal sheaf theory", + "Philosophy and history of mathematics", + "Physics-informed machine learning and scientific ML", + "Planar graphs and graphs on surfaces", + "Point-set topology and separation axioms", + "Poisson geometry", + "Polyhedral combinatorics", + "Polyhedral geometry and polytopes", + "Polytopes, zonotopes, and associahedra", + "Posets, lattices, and Möbius inversion", + "Positivity of line bundles and vanishing theorems", + "Post-quantum cryptography and isogeny-based crypto", + "Primality testing and integer factorization", + "Prime number theory and zero-density estimates", + "Privacy-preserving analytics and differential privacy", + "Probabilistic combinatorics", + "Probabilistic combinatorics and random structures", + "Probabilistic numerics", + "Probabilistic potential theory", + "Probabilistic programming and approximate inference", + "Probability and statistics", + "Projective and affine geometry", + "Proof theory", + "Propagation of chaos and chaos decompositions", + "Property testing and sublinear algorithms", + "Property testing for graph properties", + "Pseudo-Riemannian and Lorentzian geometry", + "Pseudorandomness and extractor constructions", + "Quadratic forms and modular curves", + "Quantum field theory (constructive, algebraic, axiomatic)", + "Quantum gravity and string theory", + "Quantum information theory and quantum error correction", + "Quantum mechanics and spectral analysis", + "Quantum unique ergodicity and arithmetic quantum chaos", + "Queueing theory and network queues", + "Quiver varieties and symplectic resolutions", + "Ramsey theory", + "Ramsey theory and van der Waerden-type problems", + "Random Schrödinger operators", + "Random Schrödinger operators and Anderson localization", + "Random fields and Gaussian free fields", + "Random geometry and SLE in applications", + "Random graphs and complex networks", + "Random graphs and hypergraphs", + "Random graphs and percolation theory", + "Random graphs and stochastic block models", + "Random matrices and free probability", + "Random matrices and spectral statistics", + "Random matrices and universality", + "Random matrix theory and L-function statistics", + "Random multiplicative functions", + "Random multiplicative functions and pretentious methods", + "Random permutations and descents", + "Random tilings and dimer models", + "Random walk in random environment", + "Random walks and diffusion on graphs", + "Random walks on groups and heat kernel bounds", + "Random walks, mixing times, and cover times", + "Random, Cohen, and Sacks forcing", + "Rational connectedness and Campana program", + "Rational homotopy theory", + "Rational points on higher-dimensional varieties", + "Real and complex analysis", + "Reinforcement learning with probabilistic models", + "Reliability engineering and maintenance", + "Reliability theory and survival analysis", + "Representation theory", + "Reverse mathematics", + "Ricci flow and comparison geometry", + "Riemannian and pseudo-Riemannian geometry", + "Riemannian geometry", + "Rigid analytic and Berkovich spaces", + "Ring theory and module theory", + "Risk-sensitive and robust decision-making", + "Rough paths and regularity structures", + "Sato–Tate equidistribution and automorphic Sato–Tate", + "Schemes and morphisms of schemes", + "Schmidt subspace theorem and Roth’s theorem", + "Schramm–Loewner evolution and conformal invariance", + "Schubert calculus and Coxeter combinatorics", + "Semimartingales and stochastic integration", + "Sequential Monte Carlo and particle methods", + "Set theory and foundations", + "Set-theoretic topology and selection principles", + "Shape theory and ANR spaces", + "Sheaf theory and cohomology", + "Shimura varieties and automorphic sheaves", + "Shimura varieties and modularity", + "Sieve methods and additive number theory", + "Sieve methods and distribution in arithmetic progressions", + "Signal processing and compressed sensing", + "Simplicial complexes and collapsibility", + "Singularity theory and resolutions", + "Spatial statistics and geostatistics", + "Spatial-temporal modeling and point processes", + "Spectral geometry and isoperimetric inequalities", + "Spectral graph theory and Laplacians", + "Spectral graph theory and expanders", + "Spectral hypergraph theory", + "Spectral sequences and derived functors", + "Spectral sparsification", + "Spectral theory", + "Sperner theory and lattice of subsets", + "Spin geometry and Dirac operators", + "Square and diamond principles", + "Stable homotopy and spectra", + "Stable homotopy category and spectra", + "Stable laws and domains of attraction", + "Stable pairs and wall crossing", + "Stanley–Reisner theory and face rings", + "Stationary sets and club guessing", + "Statistical decision theory and inference", + "Statistical mechanics and phase transitions", + "Statistical mechanics and probabilistic combinatorics", + "Steiner systems and finite projective planes", + "Stein’s method and normal approximation", + "Stochastic PDE and random fields", + "Stochastic PDE in applications", + "Stochastic analysis on manifolds", + "Stochastic calculus (Itô and Stratonovich)", + "Stochastic calculus and SDE", + "Stochastic control and optimal stopping", + "Stochastic control, HJB, and viscosity solutions", + "Stochastic differential equations and flows", + "Stochastic networks and communication systems", + "Stochastic optimization and stochastic gradient methods", + "Stochastic partial differential equations", + "Stratified spaces and intersection homology", + "Streaming algorithms for combinatorial problems", + "Structural graph theory", + "Sub-Riemannian and Finsler geometry", + "Sublinear and streaming graph algorithms", + "Submodular optimization", + "Sumset inequalities and Freiman-type results", + "Surface classification and higher-dimensional birational geometry", + "Suslin trees and Aronszajn trees", + "Symplectic and contact geometry", + "Symplectic and contact topology", + "Symplectic geometry and Hamiltonian dynamics", + "Systems biology and network biology", + "TQFT and modular tensor categories", + "Teichmüller theory and moduli of curves", + "Teichmüller theory and moduli spaces", + "Temporal and dynamic graphs", + "Theoretical computer science and complexity", + "Theta functions and moduli of vector bundles", + "Threshold phenomena and sharp thresholds", + "Tilings, packings, and sphere packing", + "Time series analysis and stochastic forecasting", + "Time-frequency analysis and wavelets", + "Topological combinatorics", + "Topological data analysis", + "Topological data analysis and persistent homology", + "Topological dynamics and minimal sets", + "Topological graph theory and crossing number", + "Topological groups and Lie groups", + "Topological methods in combinatorics", + "Topological methods in data analysis", + "Topological quantum computing", + "Topological robotics and configuration spaces", + "Topology (point-set, geometric, and algebraic)", + "Toric and tropical geometry", + "Toric varieties and fans", + "Tournament and oriented graphs", + "Transcendence and Diophantine approximation", + "Treewidth, clique-width, and decompositions", + "Treewidth, pathwidth, and decompositions", + "Tropical geometry and tropical intersection", + "Type theory and homotopy type theory", + "Uncertainty quantification", + "VC theory and uniform laws of large numbers", + "Variational inequalities and complementarity problems", + "Waring’s problem and Vinogradov methods", + "Weak convergence and Skorokhod topology", + "Young tableaux and Schur functions", + "Zero-free regions and the Riemann hypothesis", + "Zeta functions of graphs", + "p-adic Hodge theory and Fontaine modules", + "p-adic analysis and Hida families", + "p-adic number theory and local fields", + "q-series, partitions, and symmetric functions" + ], + "Physics": [ + "2D materials and van der Waals heterostructures", + "Ab initio and many-body methods", + "Accelerator physics and beam dynamics", + "Acoustics and aeroacoustics", + "Amorphous solids and glasses", + "Astrophysical and space plasmas", + "Astrophysics and cosmology", + "Atomic, molecular, and optical physics", + "Battery physics and solid-state electrolytes", + "Biophysics and physical biology", + "Chemical physics and physical chemistry", + "Classical mechanics and analytical dynamics", + "Cold atoms, Bose-Einstein condensates, and quantum gases", + "Colloids, emulsions, and granular matter", + "Computational physics and scientific computing", + "Condensed matter physics (theory and experiment)", + "Continuum mechanics and elasticity", + "Cosmology (inflation, CMB, dark matter, dark energy)", + "Data-driven physics and physics-informed ML", + "Digital twins for plasmas and fusion", + "Econophysics and sociophysics", + "Electromagnetism and classical electrodynamics", + "Electronic structure theory and density functional theory", + "Energy materials and photovoltaics", + "Equilibrium and nonequilibrium statistical physics", + "Floquet engineering and time crystals", + "Fluid dynamics and turbulence", + "Galactic dynamics and black holes", + "Gauge theories and the Standard Model", + "Geophysics and seismology", + "Gravitation and gravitational waves", + "Gravitational lensing and weak lensing", + "Hadron physics and effective field theories", + "Heavy-ion collisions and quark-gluon plasma", + "High-energy astrophysics and cosmic rays", + "High-intensity lasers and free-electron lasers", + "Higher-order topological phases", + "Hybrid quantum devices and atomtronics", + "Information thermodynamics and Maxwell’s demon", + "Instrumentation and detector physics", + "Inverse problems in physics", + "Large-scale structure and galaxy formation", + "Lattice gauge theory", + "Machine learning for materials discovery", + "Magnetism and spintronics", + "Majorana modes and non-Abelian anyons", + "Materials physics and nanoscience", + "Medical physics and imaging", + "Metamaterials and metasurfaces", + "Microfluidics and lab-on-chip physics", + "Moiré materials and twisted bilayers", + "Monte Carlo and molecular dynamics", + "Neutrino physics and oscillations", + "Non-Hermitian and PT-symmetric physics", + "Nonequilibrium physics and fluctuation theorems", + "Nonlinear and ultrafast optics (attosecond physics)", + "Nonlinear dynamics and chaos", + "Nuclear physics and nuclear structure", + "Numerical relativity", + "Open quantum systems and decoherence", + "Particle physics (theory and experiment)", + "Photonics and nanophotonics", + "Physical oceanography and climate physics", + "Physics education research and outreach", + "Planetary science and heliophysics", + "Plasma physics and magnetohydrodynamics", + "Polaritonic chemistry and cavity QED materials", + "Polymer physics and biophysics of macromolecules", + "Quantum chromodynamics and QCD phenomenology", + "Quantum computing and quantum algorithms", + "Quantum control and quantum thermodynamics", + "Quantum error correction and fault tolerance", + "Quantum field theory", + "Quantum information science", + "Quantum mechanics and quantum foundations", + "Quantum networks and quantum internet", + "Quantum optics and quantum metrology", + "Quantum sensing and precision metrology", + "Quantum simulators with atoms, ions, and photons", + "Quantum transport and mesoscopic physics", + "Quantum-enhanced imaging and spectroscopy", + "Relativity (special and general)", + "Rydberg atoms and quantum simulation", + "Soft matter physics and complex fluids", + "Spectroscopy across the EM spectrum", + "Stellar evolution and nucleosynthesis", + "Strongly correlated electrons", + "Superconductivity and superfluidity", + "Surface physics and scanning probe microscopy", + "Synchrotron and neutron scattering", + "Thermodynamics and statistical mechanics", + "Thermoelectrics and heat transport", + "Topological insulators and semimetals", + "Topological phases of matter", + "Topological photonics and phononics", + "Topological quantum matter out of equilibrium" + ] +} \ No newline at end of file diff --git a/sdg/questions/extract_chimera_topics.py b/sdg/questions/extract_chimera_topics.py new file mode 100644 index 0000000000000000000000000000000000000000..fc64d9c81e31ddd0ad6e85835633972d2536d2e7 --- /dev/null +++ b/sdg/questions/extract_chimera_topics.py @@ -0,0 +1,156 @@ +""" +Extract the subject -> topics mapping from the CHIMERA dataset. + +Source: https://huggingface.co/datasets/TianHongZXY/CHIMERA + +The dataset has a single config ("Qwen3-235B-2507") and split ("train") with +9,225 rows. Each row is one question carrying a `subject` and a `topic` field. + +This script streams just those two columns and writes two files: + + * chimera_topics.json — maps each unique subject to the sorted list of its + unique topics: + + { + "Physics": ["Classical Mechanics", "Electromagnetism", ...], + "Chemistry": ["Organic Chemistry", ...], + ... + } + + * chimera_stats.json — per-subject statistics: + + { + "Physics": { + "num_questions": 1234, + "num_unique_topics": 96, + "avg_questions_per_topic": 12.85, + "max_num_questions_per_topic": 30, + "min_num_questions_per_topic": 1 + }, + ... + } +""" + +from __future__ import annotations + +import argparse +import json +from collections import Counter, defaultdict +from pathlib import Path + +from datasets import load_dataset +from tqdm import tqdm + + +DATASET_ID = "TianHongZXY/CHIMERA" +DEFAULT_CONFIG = "Qwen3-235B-2507" +DEFAULT_SPLIT = "train" +DEFAULT_TOPICS_OUTPUT = Path(__file__).parent / "chimera_topics.json" +DEFAULT_STATS_OUTPUT = Path(__file__).parent / "chimera_stats.json" + + +def scan_dataset( + dataset_id: str = DATASET_ID, + config: str = DEFAULT_CONFIG, + split: str = DEFAULT_SPLIT, +) -> dict[str, Counter]: + """Stream the dataset and count questions per (subject, topic). + + Returns {subject: Counter({topic: num_questions})}. + """ + print(f"Loading {dataset_id} (config={config}, split={split}, streaming) ...") + ds = load_dataset(dataset_id, name=config, split=split, streaming=True) + # Only decode the two columns we care about (avoids the large solution fields). + ds = ds.select_columns(["subject", "topic"]) + + subject_topic_counts: dict[str, Counter] = defaultdict(Counter) + total = 0 + skipped = 0 + + for row in tqdm(ds, desc="Scanning rows"): + total += 1 + subject = row.get("subject") + topic = row.get("topic") + if not subject or not topic: + skipped += 1 + continue + subject_topic_counts[subject][topic] += 1 + + print(f" Rows scanned: {total:,}") + if skipped: + print(f" Rows skipped (missing subject/topic): {skipped:,}") + print(f" Unique subjects: {len(subject_topic_counts):,}") + print(f" Total unique topics: {sum(len(c) for c in subject_topic_counts.values()):,}") + return subject_topic_counts + + +def build_topics(subject_topic_counts: dict[str, Counter]) -> dict[str, list[str]]: + """{subject: sorted unique topics}.""" + return { + subject: sorted(counts) + for subject, counts in sorted(subject_topic_counts.items()) + } + + +def build_stats(subject_topic_counts: dict[str, Counter]) -> dict[str, dict]: + """Per-subject question counts. + + {subject: {num_questions, num_unique_topics, avg_questions_per_topic, + max_num_questions_per_topic, min_num_questions_per_topic}} + """ + stats: dict[str, dict] = {} + for subject, counts in sorted(subject_topic_counts.items()): + num_questions = sum(counts.values()) + num_topics = len(counts) + per_topic = counts.values() + stats[subject] = { + "num_questions": num_questions, + "num_unique_topics": num_topics, + "avg_questions_per_topic": round(num_questions / num_topics, 2) + if num_topics + else 0.0, + "max_num_questions_per_topic": max(per_topic) if num_topics else 0, + "min_num_questions_per_topic": min(per_topic) if num_topics else 0, + } + return stats + + +def _write_json(data: dict, path: Path) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + with path.open("w", encoding="utf-8") as f: + json.dump(data, f, ensure_ascii=False, indent=2, sort_keys=True) + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--dataset-id", default=DATASET_ID) + parser.add_argument("--config", default=DEFAULT_CONFIG) + parser.add_argument("--split", default=DEFAULT_SPLIT) + parser.add_argument( + "--topics-output", + type=Path, + default=DEFAULT_TOPICS_OUTPUT, + help="Path to write the subject->topics JSON (default: chimera_topics.json).", + ) + parser.add_argument( + "--stats-output", + type=Path, + default=DEFAULT_STATS_OUTPUT, + help="Path to write the per-subject stats JSON (default: chimera_stats.json).", + ) + args = parser.parse_args() + + subject_topic_counts = scan_dataset(args.dataset_id, args.config, args.split) + + topics = build_topics(subject_topic_counts) + stats = build_stats(subject_topic_counts) + + _write_json(topics, args.topics_output) + _write_json(stats, args.stats_output) + + print(f"Wrote {len(topics):,} subjects to {args.topics_output}") + print(f"Wrote stats for {len(stats):,} subjects to {args.stats_output}") + + +if __name__ == "__main__": + main() diff --git a/sdg/questions/prompts.py b/sdg/questions/prompts.py new file mode 100644 index 0000000000000000000000000000000000000000..248e07894391103e80bdfcbb77bc563d7a3151c3 --- /dev/null +++ b/sdg/questions/prompts.py @@ -0,0 +1,53 @@ +""" +Prompt templates for question generation from CHIMERA subjects/topics. + +Prompts are copied verbatim from the CHIMERA paper (subject expansion, problem +generation, and problem validation). +""" + +# ============================================================================= +# SUBJECT EXPANSION +# ============================================================================= + +# Usage: SUBJECT_EXPANSION_PROMPT.format(subject=) +# {subject} — a broad subject to expand into specific topics (e.g. "Physics"). +SUBJECT_EXPANSION_PROMPT = ( + "I have a broad subject. Generate a list of more specific topics that fall under this " + "subject. Each topic should be unique and different from the other topics. The topics " + "should be an important field of study within this subject. Here is the broad subject: " + "{subject}. Be as specific and diverse as possible when listing topics. For example, the " + 'subject "Probability Theory" could have topics like "Probability Distributions", ' + '"Random Variable", "Conditional Probability", "Bayesian Statistics", etc. Output should ' + "be in the format of a list of topics: \\n\\n\\n..." +) + +# ============================================================================= +# PROBLEM GENERATION +# ============================================================================= + +# Usage: PROBLEM_GENERATION_PROMPT.format(topic=) +# {topic} — a specific topic to generate a PhD-level problem for (e.g. "Random Variable"). +PROBLEM_GENERATION_PROMPT = """Given a topic, generate a new problem that falls under this topic. The problem should be at a PhD level or higher. Here is the topic: {topic}. +Your answer must strictly follow the format and criteria below, without any markdown formatting: +Problem: A clearly stated problem that requires advanced knowledge, techniques, or reasoning related to the topic. The problem must have a single, unambiguous, and objectively verifiable answer (e.g., a specific numeric value, closed-form expression, proof, or well-defined result). Clearly define all variables, parameters, and assumptions to avoid ambiguity. +Reasoning: A detailed, logical reasoning process that shows how to solve the problem step by step, ensuring each step is verifiable and reproducible. +Solution: A complete and rigorous solution derived from the reasoning steps above, including all necessary calculations or proofs. +Answer: The final, concise answer to the problem. It must be clearly checkable and easy to verify as correct.""" + +# ============================================================================= +# PROBLEM VALIDATOR +# ============================================================================= + +# Usage: PROBLEM_VALIDATOR_PROMPT.format(topic=, problem=) +# {topic} — the topic the problem should fall under. +# {problem} — the generated problem statement to validate. +PROBLEM_VALIDATOR_PROMPT = """Given a topic and a problem, verify whether the problem is valid and satisfies the requirements. +Topic: {topic} +Problem: {problem} +Check the following: +- The problem matches the topic and is at a PhD level or higher. +- The statement is clear, self-contained, and unambiguous. +- There is exactly one objectively verifiable answer. +Your answer must strictly follow the format and criteria below, without any markdown formatting: +Verdict: VALID or INVALID +Reason: one short sentence explaining the decision""" diff --git a/sdg/requirements.txt b/sdg/requirements.txt new file mode 100644 index 0000000000000000000000000000000000000000..7b9b0249bab6ea761898b224ca87e31bd866eb93 --- /dev/null +++ b/sdg/requirements.txt @@ -0,0 +1,9 @@ +vllm>=0.9.0 +datasets>=4.0.0 +tqdm>=4.60.0 +pyyaml>=6.0 +pymongo>=4.6.0 +huggingface_hub>=0.20.0 +datasketch>=1.6.0 +sentence-transformers>=3.0.0 +faiss-cpu>=1.7.4,<1.9 # 1.9+ requires numpy 2.x; this venv uses numpy 1.x diff --git a/sdg/upload_to_hf.py b/sdg/upload_to_hf.py new file mode 100644 index 0000000000000000000000000000000000000000..04a4b5d7795948e29fbdd5732295b999bea4c87f --- /dev/null +++ b/sdg/upload_to_hf.py @@ -0,0 +1,85 @@ +#!/usr/bin/env python3 +"""Upload SDG output JSONL to Hugging Face as a dataset. + +Usage: + python -m sdg.upload_to_hf output/my_experiment/output.jsonl + python -m sdg.upload_to_hf output/my_experiment/output.jsonl --org teetone --private +""" + +import argparse +import json +import tempfile +from pathlib import Path + +from huggingface_hub import HfApi, create_repo + + +def load_jsonl(filepath: str) -> list[dict]: + """Load JSONL file and return list of records.""" + records = [] + with open(filepath, "r") as f: + for line in f: + if line.strip(): + records.append(json.loads(line)) + return records + + +def main(): + parser = argparse.ArgumentParser(description="Upload SDG output JSONL to Hugging Face") + parser.add_argument("jsonl_file", help="Path to the JSONL file") + parser.add_argument("--repo-name", default=None, help="Dataset repo name (default: parent directory name)") + parser.add_argument("--org", default="teetone", help="HF organization/username") + parser.add_argument("--private", action="store_true", help="Make dataset private") + args = parser.parse_args() + + jsonl_path = Path(args.jsonl_file) + if not jsonl_path.exists(): + raise FileNotFoundError(f"File not found: {jsonl_path}") + + # Derive dataset name from parent directory (experiment_name) or filename + dataset_name = args.repo_name or jsonl_path.parent.name + repo_id = f"{args.org}/{dataset_name}" + + print(f"Loading {jsonl_path}...") + records = load_jsonl(str(jsonl_path)) + print(f"Loaded {len(records):,} records") + + # Write clean JSONL with just messages + print("Preparing data...") + with tempfile.NamedTemporaryFile(mode="w", suffix=".jsonl", delete=False) as f: + temp_path = f.name + for r in records: + clean_record = {"messages": r["messages"]} + f.write(json.dumps(clean_record, ensure_ascii=False) + "\n") + + api = HfApi() + + print(f"Creating repo {repo_id}...") + create_repo(repo_id, repo_type="dataset", private=args.private, exist_ok=True) + + print(f"Uploading to {repo_id}...") + api.upload_file( + path_or_fileobj=temp_path, + path_in_repo="data/train-00000-of-00001.jsonl", + repo_id=repo_id, + repo_type="dataset", + ) + + # Also upload stats.json if it exists alongside the JSONL + stats_path = jsonl_path.parent / "stats.json" + if stats_path.exists(): + print(f"Uploading {stats_path}...") + api.upload_file( + path_or_fileobj=str(stats_path), + path_in_repo="stats.json", + repo_id=repo_id, + repo_type="dataset", + ) + + Path(temp_path).unlink() + + print(f"Done! Dataset available at: https://huggingface.co/datasets/{repo_id}") + + +if __name__ == "__main__": + main() diff --git a/sdg/validation.py b/sdg/validation.py new file mode 100644 index 0000000000000000000000000000000000000000..310329da6d00c5e1d10386b4339502c909725284 --- /dev/null +++ b/sdg/validation.py @@ -0,0 +1,215 @@ +""" +Validation logic: static checks, decision extraction, and batch UQ validators. + +Ported from experiments/self_instill/uq_verification.py and filter_with_uq.py. +""" + +from __future__ import annotations + +import re +from typing import TYPE_CHECKING + +from sdg.prompts import ( + CYCLE_COMPARISON_PROMPT, + CYCLE_QUESTION_GENERATION_PROMPT, + FACTUAL_ERROR_PROMPT, + SIMPLE_CORRECTNESS_PROMPT, + TOTAL_CORRECTNESS_PROMPT, +) + +if TYPE_CHECKING: + from sdg.inference import VLLMEngine + +# ============================================================================= +# Regex patterns +# ============================================================================= +THINK_PATTERN = re.compile(r"\s*\S.*?\s*\S", re.DOTALL) +BOXED_PATTERN = re.compile(r"\\boxed\{[^}]+\}") +DECISION_RE = re.compile(r"\[\[\s*([YN])\s*\]\]", re.IGNORECASE) + + +# ============================================================================= +# Static check +# ============================================================================= + +def passes_static_check(text: str, is_instruction_tuned: bool, check_boxed: bool = True) -> bool: + if not text: + return False + if check_boxed and not BOXED_PATTERN.search(text): + return False + if is_instruction_tuned and not THINK_PATTERN.search(text): + return False + return True + + +# ============================================================================= +# Decision extraction +# ============================================================================= + +def extract_decision_robust(text: str) -> bool: + """Priority: [[Y]]/[[N]] → standalone Y/N → YES/NO → False.""" + if not text: + return False + text_upper = text.upper() + + matches = list(DECISION_RE.finditer(text_upper)) + if matches: + return matches[-1].group(1) == "Y" + + standalone = re.findall(r"(? bool: + if not vote_texts: + return False + return all(extract_decision_robust(t) for t in vote_texts) + + +# ============================================================================= +# Text helpers +# ============================================================================= + +def extract_final_answer(answer: str) -> str: + """Content after (or full text if no tag).""" + if "" in answer: + return answer.split("")[-1].strip() + return answer + + +def clean_inferred_question(text: str) -> str: + """After , first line only.""" + if not text: + return "" + if "" in text: + text = text.split("")[-1].strip().split("\n")[0].strip() + else: + text = text.strip().split("\n")[0].strip() + return text + + +def collect_valid_samples( + samples: list[str], is_instruction_tuned: bool, check_boxed: bool = True +) -> tuple[list[str], list[int]]: + """Return (valid_samples, valid_indices) that pass static check.""" + valid_samples: list[str] = [] + valid_indices: list[int] = [] + for idx, s in enumerate(samples): + if passes_static_check(s, is_instruction_tuned, check_boxed): + valid_samples.append(s) + valid_indices.append(idx) + return valid_samples, valid_indices + + +# ============================================================================= +# Batch validation functions +# ============================================================================= + +def validate_batch_cycle( + engine: "VLLMEngine", + items: list[tuple[int, str, str]], + val_batch_size: int, +) -> dict[int, bool]: + """ + Two-step cycle consistency: generate inferred Q, then compare. + + Args: + items: list of (row_id, question, answer) + + Returns: + {row_id: passed} + """ + results: dict[int, bool] = {} + + for batch_start in range(0, len(items), val_batch_size): + batch = items[batch_start : batch_start + val_batch_size] + + # Step 1 — generate inferred questions + gen_prompts = [ + CYCLE_QUESTION_GENERATION_PROMPT.format(answer=extract_final_answer(a)) + for _, _, a in batch + ] + inferred_raw = engine.generate_single(gen_prompts) + inferred_clean = [clean_inferred_question(q) for q in inferred_raw] + + # Step 2 — compare + compare_prompts = [ + CYCLE_COMPARISON_PROMPT.format( + original_question=q, inferred_question=iq + ) + for (_, q, _), iq in zip(batch, inferred_clean) + ] + vote_lists = engine.generate_with_votes(compare_prompts) + + for (row_id, _, _), votes in zip(batch, vote_lists): + results[row_id] = check_unanimous(votes) + + return results + + +def validate_batch_factual( + engine: "VLLMEngine", + items: list[tuple[int, str, str]], + val_batch_size: int, +) -> dict[int, bool]: + results: dict[int, bool] = {} + + for batch_start in range(0, len(items), val_batch_size): + batch = items[batch_start : batch_start + val_batch_size] + prompts = [ + FACTUAL_ERROR_PROMPT.format(question=q, answer=extract_final_answer(a)) + for _, q, a in batch + ] + vote_lists = engine.generate_with_votes(prompts) + for (row_id, _, _), votes in zip(batch, vote_lists): + results[row_id] = check_unanimous(votes) + + return results + + +def validate_batch_correctness( + engine: "VLLMEngine", + items: list[tuple[int, str, str]], + val_batch_size: int, +) -> dict[int, bool]: + results: dict[int, bool] = {} + + for batch_start in range(0, len(items), val_batch_size): + batch = items[batch_start : batch_start + val_batch_size] + prompts = [ + TOTAL_CORRECTNESS_PROMPT.format(question=q, answer=extract_final_answer(a)) + for _, q, a in batch + ] + vote_lists = engine.generate_with_votes(prompts) + for (row_id, _, _), votes in zip(batch, vote_lists): + results[row_id] = check_unanimous(votes) + + return results + + +def validate_batch_simple( + engine: "VLLMEngine", + items: list[tuple[int, str, str]], + val_batch_size: int, +) -> dict[int, bool]: + """Single yes/no judgement; passes only if all num_validation_votes votes say yes.""" + results: dict[int, bool] = {} + + for batch_start in range(0, len(items), val_batch_size): + batch = items[batch_start : batch_start + val_batch_size] + prompts = [ + SIMPLE_CORRECTNESS_PROMPT.format(question=q, answer=extract_final_answer(a)) + for _, q, a in batch + ] + vote_lists = engine.generate_with_votes(prompts) + for (row_id, _, _), votes in zip(batch, vote_lists): + results[row_id] = check_unanimous(votes) + + return results