Skip to content

feat(imf): WO02 ONNX export pipeline — KV decoder, fp16/int8, Modal runner - #3

Merged
ronaldtse merged 9 commits into
mainfrom
feat/imf-export
Aug 16, 2026
Merged

feat(imf): WO02 ONNX export pipeline — KV decoder, fp16/int8, Modal runner#3
ronaldtse merged 9 commits into
mainfrom
feat/imf-export

Conversation

@ronaldtse

Copy link
Copy Markdown

Summary

Work order TODO.runtime-arch/02: the ONNX export pipeline that turns any
HF byte-level seq2seq checkpoint into IMF v1 zips — fp32 + fp16 + int8,
with a KV-cache decoder as the default artifact and the plain decoder as
fallback. Generalizes secryst PR #44's scripts/export_onnx_byt5.py.

Stacked on #2 (IMF v1 spec + validator). Base is feat/imf-v1; retarget
to main when #2 merges.

Hard-won export findings (all encoded in code + comments)

  • Tracing mutates the traced module. torch.onnx.export leaves the
    model's runtime behavior subtly altered; the WO03 parity harness needs
    the reference pristine, so export_graphs traces a deepcopy.
  • KV cache design. decoder-kv.onnx caches self-attention only.
    Cross-attention K/V are a deterministic projection of
    encoder_hidden_states — caching them breaks the attention-mask length
    invariant for any fed past length, and EncoderDecoderCache.is_updated
    is a Python bool baked at trace time. Recomputing cross K/V per step is
    correct for all past lengths and removes all cross-cache bookkeeping
    from the three runtimes. Step 0 feeds zero-length self pasts.
  • fp16 must keep LayerNorm/softmax math in fp32. ORT's session-time
    SimplifiedLayerNormFusion crashes on half-converted LN subgraphs
    (InsertPrecisionFreeCast name mismatch). Blocking the LN decomposition
    ops keeps graphs loadable with default session options (old Ruby ORT
    included). Weights — the bulk — still halve.
  • int8 must be MatMul-only dynamic quantization. Quantizing more ops
    inserts casts that hit the same fusion bug; quant_pre_process pins
    concrete example shapes into DynamicQuantizeLinear buffers.

Verification

  • Fixture (tiny random T5, seed 42) exported end-to-end locally:
    • ONNX plain == ONNX KV == untraced torch reference, token-exact,
      4 probe strings
    • fp16 and int8 zips decode token-identical to fp32
    • all three zips pass imf validate (kv contract, opset 14, sha256)
  • tests/test_imf_export.py (4 specs) — full suite 61 passed, ruff clean
  • New CI job export-fixture installs train+export extras and runs the
    torch/ORT parity specs

Modal runner

src/gpu/modal_export.py — CPU-only (exports never compete with A100
training), watchdog invocation documented:

until modal run --detach src/gpu/modal_export.py --model khm-latn; do sleep 60; done

Registry: khm-latn, urd-g2p, urd-diac → secryst-models:/imf/<model>/.
The khm-latn run re-exports fp32 from scratch, replacing the CRC-corrupt
zip found on the checkpoint volume.

Test plan

  • pytest tests/ — 61 passed locally
  • ruff clean
  • CI green (incl. new export-fixture job)
  • Modal exports for khm-latn / urd-g2p / urd-diac complete + validated
    (running next; results in this PR)

@ronaldtse

Copy link
Copy Markdown
Author

Root cause found for the decode convention: the stock byt5 tokenizer maps byte b -> id b+3 with a trailing EOS — the helpers (and PR #44's Ruby engine) fed raw bytes. Fixed in 83c163d (canonical table in the spec + helpers + tests); khm golden now decodes exactly ('រោក' -> 'rok', 'សង់ផ្ទះ' -> 'sangphteah'). Graphs were always correct — no re-export; parity gates re-running with the fix.

@ronaldtse
ronaldtse changed the base branch from feat/imf-v1 to main August 16, 2026 09:05
Ronald Tse added 9 commits August 16, 2026 17:06
Generalizes secryst PR #44's export_onnx_byt5.py into the IMF pipeline:

- decoder-kv.onnx: self-attention KV cache (past_*/present_* contract).
  Cross-attention K/V are recomputed each step — a deterministic
  projection of encoder states, so runtimes never track cross caches.
- opset pinned 14; plain decoder kept as fallback in every zip
- fp16: LayerNorm/softmax math kept fp32 (ORT's session-time
  SimplifiedLayerNormFusion crashes on half-converted LN subgraphs);
  weights still halve. int8: MatMul-only dynamic quantization
  (quantizing more inserts casts that break that same fusion, and
  quant_pre_process pins concrete shapes into activation buffers)
- export traces a deepcopy: tracing mutates the traced module's runtime
  behavior, and the WO03 parity harness needs the reference pristine
- fixture mode: tiny random T5; CI job runs torch-vs-ORT parity on all
  three precisions, plain and KV paths, token-exact
- Modal app (CPU, no A100 contention): khm-latn / urd-g2p / urd-diac
  from their checkpoint volumes into secryst-models:/imf/, with the
  until-retry watchdog invocation documented in the module docstring
- metadata + tri-API READMEs for urd-g2p-1.0 and urd-diac-1.0; khm
  metadata now declares decoder: kv

khm-latn re-export also replaces the CRC-corrupt fp32 zip on the volume.
- imf.parity: ONNX KV greedy vs the transformers decoder loop (the exact
  math the export wraps; generate() is config-dependent and no runtime
  implements it) over the test split; gate = cer_delta <= 0.2pp on >=
  500 samples; write_parity rewrites the zip's parity block and enforces
  strict validation — a zip cannot leave the gate non-strict
- imf golden: cross-runtime golden JSONL (fixed inputs + reference
  outputs from ONNX decode; Python is the reference implementation)
- CLI: 'imf parity' and 'imf golden'; accepts src/tgt and input/target
  pair keys
- Modal: 'parity' entrypoint on the export app (torch reference vs the
  exported zips, in place, on the models volume)
- CI: export-fixture job now runs the parity gate end-to-end on the
  fixture (export -> parity -> strict-validate), headless
- 4 new specs (65 total): gate exactness on fp32, cer_delta rejection,
  small-sample rejection, golden roundtrip
Metrics record BOTH decode paths: greedy DER 29.0 (what v1 runtimes
produce — beam search is not in the runtimes) and beam=4 DER 17.46
(reference quality), both sourced to rababa RESULTS.md. Nakdimon test
data mounted for the parity gate.
Nakdimon test.txt is raw diacritized text; the gate needs (stripped,
diacritized) pairs. 1,864 sentence pairs on rababa-datasets:/nakdimon/
test-imf.jsonl (nikud/cantillation U+0591-U+05C7 stripped for input).
The stock google/byt5 tokenizer (which every byte-level checkpoint in
this campaign was trained with) maps UTF-8 byte b to token id b+3 and
appends EOS(1) to inputs; pad=0, unk=2. Feeding text.bytes directly —
as the decode helpers and PR #44's Ruby engine did — silently produces
garbage on real models while looking perfectly healthy on synthetic
fixtures (both sides of a comparison share the wrong convention).

- encode_bytes(): [b+3 for b in bytes] + [1]; decode: id-3, stop at 1
- export/parity/golden helpers + tests now use the canonical table;
  the fixture parity test shares imf.parity's reference decoder
- khm-latn golden regenerated: 'រោក'->'rok', 'សង់ផ្ទះ'->'sangphteah'
  (exact gold matches); spec documents the table and the trap
- graphs themselves were always correct — no re-export needed; the
  vacuous parity runs from before the fix were killed and are re-run
os.replace across filesystems raises EXDEV — on Modal the zips live on
a volume mount while tempfile defaults to /tmp. Same-dir temp keeps the
atomic rename.
The reference is precision-independent; multi-zip gates were paying
3x the torch decode (hours on the 12k-pair Urdu splits).
… real ByT5

The onnxruntime float16 converter produced all-zero encoder hiddens on
the real khm-latn checkpoint (CER 1939pp, every sample mismatched) while
looking fine on the tiny fixture. Exporting the model under .half() is
exact on gold pairs; graph IO becomes float16 (int64 ids unchanged) and
_zero_pasts follows session dtypes.

Measured on 300 khm test samples: fp16 delta 0.43pp, int8 0.84pp —
quantization noise (argmax flips cascading under greedy decode), not
breakage. Whether lossy precisions may exceed the 0.2pp export-fidelity
bar is a policy call pending; fp32 measures 0.0pp.
@ronaldtse
ronaldtse merged commit dce303d into main Aug 16, 2026
8 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant