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)

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

{"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.