Upload folder using huggingface_hub

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://<hostname>: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://<hostname>: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://<hostname>: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/
+```
+

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 `<your-hf-username>/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:<port> AND [::]:<port>
+
+# 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 `<HF_USER>/<RUN_NAME>-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 `<HF_USER>/<RUN_NAME>` 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/<HF_USER>/<RUN_NAME>`
+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: <unreachable>` 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 <file>` (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 <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 <your-repo-url> 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:<port> AND [::]:<port>
+
+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/<HF_USER>/<RUN_NAME>
+```
+
+### 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/<HF_USER>/<RUN_NAME>
+```
+
+Everything else (per-shard logs, mongo cache count, GPU util, staging repo,
+upload markers) is plumbing that supports that one line.

gpu-sft/.gitignore

@@ -0,0 +1,6 @@
+__pycache__/
+*.pyc
+runs/
+hf_cache/
+evalout/
+*.npy

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

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/<exp>-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
+```

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
+------
+<root>/
+  Qwen_Qwen3-8B/                                     # baseline: model id, '/' -> '_'
+    MATH500_seed42/**/results_<ISO>.json             # single-seed benchmark
+    AIME24_seed42/**/results_<ISO>.json              # per-seed dirs, seeds 42..51
+    ...
+    AIME24_seed51/**/results_<ISO>.json
+    compile_AIME24_avg10seeds/compiled_results/averaged_results.json   # optional
+  <experiment_name>-step<N>/                         # one dir per checkpoint
+    <same benchmark dirs as above>
+
+A trailing "-<hash>" 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": {"<evalchemy task name>": {"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 <checkpoint_root>/<exp>[-<hash>]/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 ``<exp>-step<N>`` - 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 <root> 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_<ISO>.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 <compile_dir>/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 <dir>/<exp>-<hash>/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: <output dir>/.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()

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 <SIZE>` 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 `- <name>` under an experiment heading must be the EXACT experiment name,
+  which is also the eval output directory prefix: `<name>-step<N>/`. Do not
+  truncate, do not add a `-step<N>` suffix here.
+- Row order inside a suite table follows this file. The `#### Best Checkpoints`
+  table re-sorts by sampling compute (`nofilter` first, then `_n<N>` asc,
+  `_vr<N>` 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

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: '<experiment_name>-step<N>'.
+    """
+    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: <root>/<model_dir>/<task>_seed<seed>/<model_sanitized>/results_<ISO>.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: <root>/<model_dir>/compile_<task>_avg<N>seeds/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")

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 <tools></tools> XML tags:\n<tools>" }}
+    {%- for tool in tools %}
+        {{- "\n" }}
+        {{- tool | tojson }}
+    {%- endfor %}
+    {{- "\n</tools>\n\nFor each function call, return a json object with function name and arguments within <tool_call></tool_call> XML tags:\n<tool_call>\n{\"name\": <function-name>, \"arguments\": <args-json-object>}\n</tool_call><|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('<tool_response>') and message.content.endswith('</tool_response>')) %}
+        {%- 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 '</think>' in content %}
+                {%- set reasoning_content = content.split('</think>')[0].rstrip('\n').split('<think>')[-1].lstrip('\n') %}
+                {%- set content = content.split('</think>')[-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 %}{{- '<think>\n' + reasoning_content.strip('\n') + '\n</think>\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 %}{{- '<tool_call>\n{"name": "' }}
+                {{- tool_call.name }}
+                {{- '", "arguments": ' }}
+                {%- if tool_call.arguments is string %}
+                    {{- tool_call.arguments }}
+                {%- else %}
+                    {{- tool_call.arguments | tojson }}
+                {%- endif %}
+                {{- '}\n</tool_call>' }}{% 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<tool_response>\n' }}
+        {{- content }}
+        {{- '\n</tool_response>' }}
+        {%- 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 %}{{- '<think>\n\n</think>\n\n' }}{% endgeneration %}
+    {%- endif %}
+{%- endif %}"""

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

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" <<ACTEOF
+# source this before running gpu_eval_driver.py / gpu_eval_monitor.py by hand
+export MARIN_GPU_EVAL_ROOT="$ROOT"
+export EVALCHEMY_DIR="$EVALCHEMY_DIR"
+export VIRTUAL_ENV="$VENV"
+export PATH="$VENV/bin:\$PATH"
+
+# HF / dataset behaviour required by LiveCodeBench and the math suite
+export HF_DATASETS_TRUST_REMOTE_CODE=1
+export HF_HOME="$CACHE/huggingface"
+export HF_HUB_CACHE="$CACHE/huggingface/hub"
+export HF_DATASETS_CACHE="$CACHE/datasets"
+export HF_ALLOW_CODE_EVAL=1
+
+# vLLM / torch caches kept off \$HOME so concurrent runs do not fight
+export VLLM_CACHE_ROOT="$CACHE/vllm"
+export VLLM_CONFIG_ROOT="$CACHE/vllm"
+export TORCHINDUCTOR_CACHE_DIR="$CACHE/torchinductor"
+export TRITON_CACHE_DIR="$CACHE/triton"
+export TOKENIZERS_PARALLELISM=false
+
+# Supply your own:
+#   export HF_TOKEN=hf_xxx     # needed for cais/hle and Idavidrein/gpqa
+ACTEOF
+
+log "done"
+log ""
+log "  source $ROOT/activate.sh"
+log "  export HF_TOKEN=...            # only needed for the science suite"
+log "  python $HERE/gpu_eval_driver.py --help"
+log ""
+log "evalchemy: $EVALCHEMY_DIR @ $(git -C "$EVALCHEMY_DIR" rev-parse --short HEAD)"
+log "python:    $PY"

gpu-sft/scripts/gpu_eval/gpu_eval_driver.py

@@ -0,0 +1,996 @@
+#!/usr/bin/env python3
+# Copyright The Marin Authors
+# SPDX-License-Identifier: Apache-2.0
+"""Run Marin's evalchemy benchmark suites against a LOCAL HF checkpoint on one GPU box.
+
+This is the local-disk, CUDA-vLLM replacement for the TPU path
+(``experiments/evals/exp_evalchemy_eval.py`` + ``EvalchemyEvaluator``). It keeps
+the benchmark set, the sampling parameters and -- critically -- the on-disk
+result schema that ``claude/compile_results.py`` reads to build results.md.
+
+    python gpu_eval_driver.py \\
+        --checkpoint /data/runs/exp_sft_qwen3_8b_x/hf/step-300 \\
+        --experiment exp_sft_qwen3_8b_x \\
+        --suite math \\
+        --results-root /data/marin/evaluation/evalchemy
+
+WHAT IS FAITHFUL TO THE TPU PIPELINE
+    * Task set and per-task seed counts (math AIME*/AMC23/HMMT x10, MATH500 and
+      OlympiadBench x1; science x3; code x6), seeds 42..51.
+    * num_fewshot=0, chat template applied, temperature 0.7, top_p 1.0
+      (both are the benchmark/vLLM defaults; Marin's --gen_kwargs never reached
+      evalchemy chat benchmarks on TPU either), max_gen_toks 32768,
+      max_model_len = 32768 + 4096 = 36864.
+    * Real per-benchmark graders (is_equiv, symbolic, execution) -- the naive
+      string re-grade bug in ``evalchemy_results_compiler.py`` is NOT reproduced.
+    * Output layout, so the existing results.md compiler needs zero changes:
+        <root>/<experiment>-step<N>/<TASK>_seed<S>-<h>/<TASK>_0shot/<model>/results_<iso>.json
+        <root>/<experiment>-step<N>/compile_<TASK>_avg<K>seeds-<h>/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 ``<TASK>_seed<S>-`` and
+    ``compile_<TASK>_avg<K>seeds-``, 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 <eval_dir>/_raw/<TASK> 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())

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

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 = <int>' 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())

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 <checkpoint_dir> <experiment> <suite>
+  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"

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

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

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 <hf-dataset-id> <7-char-revision> <run-name>
+#
+# 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 <dataset-id> <revision> <run-name>}"
+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
+# ----------------------------------------------------------------------------------

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 = {"<think>": "<|start_think|>", "</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
+   `'</think>' in content` reasoning-split branch never fires, `reasoning_content` stays
+   empty, and every FINAL assistant turn is rendered as
+
+       <|im_start|>assistant\n<think>\n\n</think>\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 <tools></tools> XML tags:\n<tools>" }}
+    {%- for tool in tools %}
+        {{- "\n" }}
+        {{- tool | tojson }}
+    {%- endfor %}
+    {{- "\n</tools>\n\nFor each function call, return a json object with function name and arguments within <tool_call></tool_call> XML tags:\n<tool_call>\n{\"name\": <function-name>, \"arguments\": <args-json-object>}\n</tool_call><|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('<tool_response>') and message.content.endswith('</tool_response>')) %}
+        {%- 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 '</think>' in content %}
+                {%- set reasoning_content = content.split('</think>')[0].rstrip('\n').split('<think>')[-1].lstrip('\n') %}
+                {%- set content = content.split('</think>')[-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 %}{{- '<think>\n' + reasoning_content.strip('\n') + '\n</think>\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 %}{{- '<tool_call>\n{"name": "' }}
+                {{- tool_call.name }}
+                {{- '", "arguments": ' }}
+                {%- if tool_call.arguments is string %}
+                    {{- tool_call.arguments }}
+                {%- else %}
+                    {{- tool_call.arguments | tojson }}
+                {%- endif %}
+                {{- '}\n</tool_call>' }}{% 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<tool_response>\n' }}
+        {{- content }}
+        {{- '\n</tool_response>' }}
+        {%- 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 %}{{- '<think>\n\n</think>\n\n' }}{% endgeneration %}
+    {%- endif %}
+{%- endif %}"""
+
+DEFAULT_TEXT_REPLACEMENTS = {"<think>": "<|start_think|>", "</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()

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 <repo> 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 <rundir> --repo fzzhang/<exp>_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/<exp>_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()

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 <exp_name> [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/<exp_name> by prepare_sft_data.py.
+# For a FRESH run, delete $RUNS_ROOT/<exp_name> 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 <exp_name> [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

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 `<out>/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 | <path>")
+    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()

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)

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

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

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 `<HF_USER>/<RUN_NAME>-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)

sdg/__init__.py

@@ -0,0 +1 @@
+__version__ = "0.1.0"

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

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)

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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