File size: 4,426 Bytes
b2f3bf4 | 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 | # 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.
|