File size: 9,153 Bytes
61e8439 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 | ---
license: apache-2.0
tags:
- proof-pilot
- olmo3
- dflash
- speculative-decoding
- on-policy-distillation
- gptq
- quantization
---
# Proof Pilot deploy bundle
Deployment-ready model artifacts for the AIMO Proof Pilot submission. Every subfolder is an
`Olmo3SinkForCausalLM` (olmo3_sink, DeepSeek-V4-Flash transplanted tokenizer, vocab 129,280)
family member and they share one tokenizer.
Related repositories: [code](https://github.com/ycchen-tw/proof-pilot-codes) ·
[intermediate checkpoints](https://huggingface.co/ycchen/proof-pilot-checkpoints) ·
[training-data artifacts](https://huggingface.co/datasets/ycchen/proof-pilot-datasets).
| Subfolder | Model | Notes |
|---|---|---|
| `soft-distill-7b-deploy/` | stage1-v2-7b soft-distill v2 (deploy) | 7B off-policy soft distillation (bf16), with tokenizer + chat template |
| `soft-distill-32b-deploy/` | stage1-v2-32b soft-distill v2 (deploy) | 32B off-policy soft distillation (bf16, rope-legacy deploy config), with tokenizer + chat template |
| `opd-32b-deploy/` | **32B agentic OPD v2 — v33 / job 135076 / step_200** | **On-policy distillation** (teacher = DeepSeek-V4-Flash), student lineage = stage1-v2-32b-softdistill-v2test (GQA-8, YaRN). bf16 deploy (rope-legacy config + hybrid-SWA) + tokenizer + chat template. **The deployed target.** IMO-ProofBench v2 agentic loop (prove→verify→refine→select): **4.48/7 under an independent Claude cross-check grader**. (An earlier DeepSeek-Flash ranking-grader pass scored 3.808/7; its raw outputs were not retained, so the retained number is the Claude cross-check.) See "OPD 32B versions" below. |
| `opd-32b-v33-s150/` | 32B agentic OPD v2 — v33 / step_150 | Earlier checkpoint of the **same healthy run** as `opd-32b-deploy` (step_200), for s150-vs-s200 comparisons. Same deploy format + tokenizer + chat template. |
| `opd-32b-v33-s200-gptq-w4a16/` | **GPTQ-w4a16 quantization of `opd-32b-deploy` (step_200)** | int4 weight-only (compressed-tensors, int4 sym group-128 GPTQ) + **calibrated fp8 KV scales**. **18.74 GB** (bf16 65 GB → int4; four shards + index). Calibration matches inference: sink-on + long-ctx (10,240) + factor-32 YaRN. Serve = sglang `olmo2_sink` + triton (sm120) / fa3 (H200) + `--kv-cache-dtype fp8_e4m3`. Details below. |
| `opd-32b-v33-s150-gptq-w4a16/` | GPTQ-w4a16 quantization of step_150 | Same recipe as the s200 GPTQ (int4 sym g128 GPTQ + sink-on calib + long-ctx 10,240 + factor-32 YaRN + calibrated fp8 KV scales). 18.74 GB. Same serving. |
| `dflash-7b-draft/` | DFlash draft for the 7B target | SGLang-deployable speculative-decoding draft |
| `dflash-32b-draft/` | DFlash draft for **stage1-v2-32b** (s5317) | SGLang-deployable draft aligned to the older 32B deploy target |
| `dflash-32b-draft-v2test/` | DFlash draft for stage1-v2-32b-softdistill-v2test — **phase-1 warm-up (not final)** | Curriculum **phase-1 short-context warm-up** (step-10000 snapshot). SWA512 / block_size 11 / 8L / GQA-8. **Deploy the phase-L draft below instead.** |
| `dflash-32b-draft-v2test-phaseL/` | **DFlash draft — phase-2 final (recommended for deployment)** | Curriculum **phase-2 long-context specialization** (job 140680, warm-started from phase-1): training data = the real long-proof deployment distribution (OPD 32B rollouts with finish_reason=length filtered out + dsflash-v2-test teacher proofs, micro 65,536), GAMMA 20. Clean step_3000 finish: acc 0.605 / greedy mean prefix 4.90. SWA512 / block_size 11 / 8L / GQA-8. Serving accept ~3.1–4.1 (single stream, dev H200). |
| `dflash-32b-draft-v2test-phaseL-int4mlp/` | **int4-MLP quantization of the phase-L draft** | MLP (gate/up/down) quantized to compressed-tensors int4 (RTN W4A16 g128); qkv/o + sink/fc/mask_embed stay bf16 (preserves DFlash fused-KV). **4.82 → 2.30 GB (−55%)**. Measured with sglang DFlash: loads as int4 (weight mem 2.16 GB), `fused KV materialization ENABLED`, accept 3.1–4.1 (== bf16 draft), single-stream **+2–15% tok/s** vs the bf16 draft. Serving needs the patched `dflash_sink.py` (threads quant_config→MLP) + `--speculative-draft-model-quantization compressed-tensors`. |
| `dflash-32b-draft-v2test-phaseL-int4mlp-gptq/` | int4-MLP GPTQ variant | Same int4-MLP as above but **GPTQ** (full-rank target-hidden Hessian, 26,130 rows) instead of RTN — strictly more accurate weights (weighted error −69% vs RTN). Identical deployment; accept ~3.0–3.8 (== RTN == bf16; draft accept is already at the bf16 lossless-verify ceiling, so RTN/GPTQ are statistically equivalent in τ). 2.30 GB. |
DFlash draft folders contain only `config.json` + `model.safetensors` (resharded) and must be
paired with **their matching** target model (the v2test drafts pair with the v2test target, not
with the older `dflash-32b-draft` target).
### DFlash 32B v2test curriculum (phase-1 → phase-2)
- **phase-1 (`dflash-32b-draft-v2test/`)**: short-context warm-up (DATA=l4-g2-ml4096, micro
8,192); cheap warm-up, acc ~0.64. **Not the deployment draft — a curriculum snapshot.**
- **phase-2 (`dflash-32b-draft-v2test-phaseL/`)**: warm-started from phase-1, specialized on the
real long-proof deployment distribution (OPD 32B rollouts + dsflash teacher proofs, micro
65,536). **This is the draft to deploy against the OPD/soft-distill 32B target.**
## OPD 32B versions (which is which)
The agentic semi-on-policy OPD 32B (student = stage1-v2-32b-softdistill-v2test, teacher =
DeepSeek-V4-Flash) was trained twice:
- **V32 = job 134244**: collapsed around step ~148–158 through length self-amplification
(eos 90%→13%, cap-hit→87%) and was stopped; its last checkpoint is step_150. **No V32 weights
are in this bundle.**
- **V33 = job 135076**: added a cap-hit admission filter (+ fast sharded save + topology
rebalance), ran healthily to step 237 before being stopped, saving step_150 and step_200.
**Both `opd-32b-deploy` (= step_200) and `opd-32b-v33-s150` (= step_150) come from this run.**
> ⚠️ V33's cap-hit filter mitigates length self-amplification on the **training** side; the
> distilled student **still carries the OPD reverse-KL loop tendency at inference** (e.g. refine
> reasoning falling into repetition attractors and running to the token cap). This is a general
> property of reverse-KL OPD, independent of which step the run stopped at.
### How the deploy weights were produced (same recipe for s150/s200)
```bash
# step_NNN training checkpoint (DCP + consolidated hf/) → serve-ready deploy dir
python deploy/make_olmo3sink_deploy.py \
--src training/opd_v2/runs/agentic_32b_lc140k_v33/checkpoints/step_000NNN/hf \
--dst outputs/agentic_32b_lc140k_v33-sNNN-deploy
python deploy/sm120/enable_swa_config.py outputs/agentic_32b_lc140k_v33-sNNN-deploy
```
## GPTQ target (`opd-32b-v33-s200-gptq-w4a16`)
The 4-bit deployment build of `opd-32b-deploy` (step_200 bf16) for VRAM-constrained
environments such as the Kaggle RTX 6000 Pro (sm120). **18.74 GB**.
**Recipe (deliberately aligned with the serving distribution to eliminate calib/infer
mismatch)**: llm-compressor GPTQ, scheme W4A16 (int4 / symmetric / group_size 128 / Hessian
error compensation); `lm_head`/embeddings/norms/sinks stay bf16. Calibration:
- **sink-on**: the gpt-oss attention sink participates in the eager calibration forward (the
same math as the sglang serving path; the trained sink logit mean is ~+6.7 — not negligible,
so calibration must include it)
- **long-context seqlen 10,240**: beyond YaRN original_max 8,192 and the 4,096 sliding window,
so calibration sees long-range / high-position post-YaRN key distributions
- **factor-32 YaRN rope** (deploy config verbatim)
- calibration data = the SFT L4 training bins (same 129,280 transplant vocab), n=64
**KV cache**: ships **calibrated fp8 per-tensor static k_scale / v_scale** (written into the
config's `kv_cache_scheme`) instead of uncalibrated scale=1.0. Measured k_scale 0.032–0.169
(mean 0.061), v_scale 0.0069–0.436 (mean 0.161; deep-layer V much larger than shallow).
**Quality** (sink-on, serving regime, teacher-forced @8,192, vs the bf16 reference):
ppl **+0.13%**, top-1 agreement **0.975**, KL(bf16‖q) 0.011 — near-lossless weights.
**Serving notes**:
- sglang with `deploy/target/olmo2_sink.py` bind-mounted (in-kernel sink),
`--attention-backend triton` (the only sink-correct backend on sm120) or `fa3` (H200), and
**`--kv-cache-dtype fp8_e4m3`** so the calibrated k/v scales load. sglang auto-detects
compressed-tensors from `quantization_config`; no `--quantization` flag needed.
- ⚠️ Calibrated (non-unit) KV scales are mutually exclusive with the **DFlash fused-KV ring**
(the DFlash path disables fused-KV when it sees non-unit k/v scales). Use unit-scale fp8 KV
when running DFlash speculative decoding.
- `actorder=static` but the checkpoint has **zero `g_idx`** (the reordering is baked in; no
permutation to apply) → both the marlin W4A16 and humming W4A8 paths are safe.
- ⚠️ KV scales were calibrated at length 10,240; at 256k serving, keys at extreme positions may
slightly exceed the calibrated amax and saturate in fp8 (saturation, not an error).
|