| # jev-judge |
|
|
| Jev-API compatible typed-decision judge distilled from `SargeDev/jev-distill-corpus-v3` onto the |
| `qwen3_5` code path (Qwen3.5-9B for iteration, Qwen3.8-27B as the final target). Design: `DESIGN.md`. |
|
|
| `(state, question, kind, options)` → bare-text template → one forward pass → hidden state at the last |
| token → 24-slot fp32 linear head (initialised from `lm_head` rows, so step 0 ≡ zero-shot restricted |
| decoding) → masked softmax → calibrated distribution aligned with `options`. |
|
|
| ## Layout |
|
|
| ``` |
| configs/ train.yaml (9B S2) · train_27b.yaml · train_s1.yaml · train_smoke.yaml |
| scripts/ prepare_data.py · m0_spike.py · train.py · fit_temperature.py · evaluate.py · export.py |
| review_checkpoint.sh (pause → calibrate → evaluate → resume) · run_scan.sh · bg.sh |
| src/jev_judge/ template.py · head_init.py · model.py · data.py · losses.py · metrics.py · calibration.py |
| train_loop.py · infer.py · checkpointing.py · server.py |
| tests/ test_head_equivalence.py (gate) · test_server_contract.py |
| data/ *.parquet (6 splits, +n_tokens/is_uniform/n_options) · raw/ (jsonl) |
| reports/ data_audit.md · m0_*.md · b0_*.md · review/step*.md · eval_*.md |
| ``` |
|
|
| ## Quick start (1× B200, torch 2.13+cu130, transformers 5.16, peft 0.21, flash-linear-attention 0.5.2) |
|
|
| ```bash |
| uv pip install --system -e ".[dev]" |
| python3 scripts/prepare_data.py --raw data/raw --out data --tokenizer /root/models/Qwen3.5-9B # M1 audit |
| pytest tests/test_head_equivalence.py -v -s -m gpu --model-path /root/models/Qwen3.5-9B # gate (<1e-5) |
| python3 scripts/m0_spike.py --model /root/models/Qwen3.5-9B --out reports/m0_qwen35_9b.md # throughput |
| python3 scripts/evaluate.py --base /root/models/Qwen3.5-9B --out reports/b0_qwen35_9b.md --perm-rows 1000 # B0 |
| |
| export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True |
| scripts/bg.sh logs/s2.log python3 scripts/train.py --config configs/train.yaml --stage s2 --seed 42 \ |
| --out checkpoints/s2_9b_seed42 --set gradient_checkpointing=true max_padded_tokens=24000 |
| scripts/review_checkpoint.sh checkpoints/s2_9b_seed42/best step0500 <trainer_pid> # mid-run gate |
| |
| python3 scripts/fit_temperature.py --checkpoint checkpoints/s2_9b_seed42/best --out checkpoints/s2_9b_seed42/best/calibration.json |
| python3 scripts/evaluate.py --checkpoint checkpoints/s2_9b_seed42/best --temperature checkpoints/s2_9b_seed42/best/calibration.json \ |
| --out reports/eval_s2_9b.md --baseline-json reports/b0_qwen35_9b.json --perm-rows 1000 |
| python3 scripts/export.py --checkpoint checkpoints/s2_9b_seed42/best --calibration checkpoints/s2_9b_seed42/best/calibration.json \ |
| --out exports/jev-judge-qwen35-9b --name jev-judge-qwen35-9b |
| python3 -m jev_judge.server --export exports/jev-judge-qwen35-9b --port 18080 |
| JEV_EXPORT_DIR=exports/jev-judge-qwen35-9b pytest tests/test_server_contract.py -m gpu |
| ``` |
|
|
| ## API (DESIGN §4) |
|
|
| `POST /v1/decisions` · `POST /v1/decisions:batch` (≤256, order preserved) · `GET /healthz` |
|
|
| ```json |
| {"kind": "choice", "state": "...", "question": "...", "options": ["approve", "deny"], "truncate": false} |
| → {"id": "req_…", "kind": "choice", "options": [...], "distribution": [0.94, 0.06], |
| "decision": {"noul": null, "choice": "approve", "score": null, "expected_score": null}, |
| "confidence": 0.94, "model": {"name": "…", "version": "…", "calibrated": true}, "latency_ms": 42.1} |
| ``` |
|
|
| 422 for invalid kind/options · 413 for over-length input with `truncate=false` (or batch > 256) · |
| header `X-Jev-Judge-Version` · serving refuses to start without `calibration.json`. |
|
|
| ## B200 notes (v0.7) |
|
|
| * fla `chunk_gated_delta_rule` Triton kernels work on sm_100; `causal_conv1d` cannot be built against |
| torch cu130 with the system nvcc 12.8 → transformers falls back to `F.conv1d` (~9% of fwd time). |
| * Activation memory of this architecture without checkpointing is ≈10 MB / padded token (9B); use |
| `gradient_checkpointing: true` + `max_padded_tokens: 24000` (peak ≈34 GB, ~9k tok/s) for S2. |
| * Training steps carry a ~300 ms CPU floor (≈5k small kernel launches incl. 248 LoRA modules); |
| keep micro-batches large. Running two trainers concurrently on one GPU is *slower* in aggregate. |
| * LoRA adapters are kept in fp32 (master weights) and the backbone runs under bf16 autocast; the |
| head is always fp32 outside autocast. AdamW is fused. |
|
|
