Commit graph

1 commit

Author SHA1 Message Date
Tejas Chopra
74392b238e
feat: switch Kompress default to kompress-v2-base with weight-only int8 ONNX (#799)
## Summary

Replaces `chopratejas/kompress-base` with
**`chopratejas/kompress-v2-base`** as the default Kompress
text-compression model (the fallback for content not handled by
structured compressors), using a new **weight-only int8 ONNX** artifact
that is fp32-equivalent at 2.2x less memory.

## Why

v2 is the same dual-head ModernBERT (token classifier + span CNN),
LoRA-trained on the v2 dataset. The HF repo originally shipped PyTorch
weights only — pointing Headroom at it naively would have forced the
heavy `[ml]` (torch) path on every `[proxy]` install. We exported ONNX
artifacts reproducing the v1 loader contract (single `final_scores`
output) and published them to the HF repo.

## Eval (labeled dataset_v2 test split, n=500, threshold 0.5)

| artifact | size | f1 | must_keep_recall | keep_rate | fp32 agreement |
|---|---|---|---|---|---|
| fp32 (reference) | 601 MB | 0.9128 | 0.9770 | 0.8100 | 100% |
| **int8-wo (default)** | **261 MB** | **0.9130** | **0.9765** |
**0.8097** | **99.6%** |
| fp16 | 301 MB | 0.9129 | 0.9771 | 0.8099 | 99.9% |
| int4-wo | 230 MB | 0.9102 | 0.9825 | 0.8354 ⚠️ | 94.5% |
| int8-dynamic | 271 MB | 0.9041 | 0.9868 | 0.8857 ⚠️ | 90.5% |

Weight-only int8 (MatMulNBits) keeps activations fp32, avoiding the
upward score bias that makes dynamic int8 keep ~7% more tokens (≈40%
less compression savings). Quantized candidates were generated and
eval-gated by a Modal job in the kompress repo
(`modal_jobs/export_onnx_v2.py`) against the labeled test split.

## Changes

- Default model id → `chopratejas/kompress-v2-base`
- ONNX artifact resolution tries candidates in order (**int8-wo → fp32 →
v1 int8**), falling through on download miss **or session-load failure**
— onnxruntime builds without the MatMulNBits 8-bit kernel fall back to
fp32 instead of losing Kompress
- `HEADROOM_KOMPRESS_ONNX_FILENAME` env pins an exact artifact
- `scripts/export_kompress_v2_onnx.py`: reproducible fp32 export (loads
the merged v2 checkpoint, traces the `final_scores` contract, verifies
vs PyTorch)
- `.gitignore`: local `onnx/` artifacts dir; allowlist the export script

## Testing

- End-to-end: fresh `KompressCompressor().preload()` resolves int8-wo
from HF, loads on onnxruntime 1.23 (CPU, no torch), compresses with
error/traceback content preserved
- fp32 export verified lossless vs PyTorch (max |Δscore| = 0.0, 100%
keep agreement)
- ruff check + format clean, mypy clean, 63 targeted tests pass
2026-06-09 23:28:40 -07:00