From 23f01a912b27471ba7c437b130c90e21516521e6 Mon Sep 17 00:00:00 2001 From: Ronald Tse Date: Sun, 16 Aug 2026 12:24:19 +0800 Subject: [PATCH 1/9] =?UTF-8?q?feat(imf):=20WO02=20ONNX=20export=20pipelin?= =?UTF-8?q?e=20=E2=80=94=20KV=20decoder,=20fp16/int8,=20Modal=20app?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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. --- .github/workflows/test.yml | 12 + models/khm-latn/khm-latn-1.0.README.md | 2 +- models/khm-latn/khm-latn-1.0.metadata.yaml | 26 ++ models/urd-diac/urd-diac-1.0.README.md | 41 +++ models/urd-diac/urd-diac-1.0.metadata.yaml | 20 ++ models/urd-g2p/urd-g2p-1.0.README.md | 41 +++ models/urd-g2p/urd-g2p-1.0.metadata.yaml | 26 ++ src/gpu/modal_export.py | 134 +++++++ src/imf/export.py | 385 +++++++++++++++++++++ tests/test_imf_export.py | 153 ++++++++ 10 files changed, 839 insertions(+), 1 deletion(-) create mode 100644 models/khm-latn/khm-latn-1.0.metadata.yaml create mode 100644 models/urd-diac/urd-diac-1.0.README.md create mode 100644 models/urd-diac/urd-diac-1.0.metadata.yaml create mode 100644 models/urd-g2p/urd-g2p-1.0.README.md create mode 100644 models/urd-g2p/urd-g2p-1.0.metadata.yaml create mode 100644 src/gpu/modal_export.py create mode 100644 src/imf/export.py create mode 100644 tests/test_imf_export.py diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml index 5ff0e36..9ac083c 100644 --- a/.github/workflows/test.yml +++ b/.github/workflows/test.yml @@ -43,3 +43,15 @@ jobs: run: | PYTHONPATH=src python -m src.cli list test "$(PYTHONPATH=src python -m src.cli list | jq 'length')" -ge 3 + + export-fixture: + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v4 + - uses: actions/setup-python@v5 + with: { python-version: "3.11" } + - run: | + python -m pip install --upgrade pip + pip install -e ".[dev,train,export]" + - name: IMF export fixture tests (torch + ORT parity) + run: PYTHONPATH=src python -m pytest tests/test_imf_export.py -v diff --git a/models/khm-latn/khm-latn-1.0.README.md b/models/khm-latn/khm-latn-1.0.README.md index f48713f..a5b5c28 100644 --- a/models/khm-latn/khm-latn-1.0.README.md +++ b/models/khm-latn/khm-latn-1.0.README.md @@ -5,7 +5,7 @@ the tokenizer is raw UTF-8 bytes (pad=0, EOS=1) — no vocab files. IMF v1 artifact; format spec: interscript/ml-models docs/imf-v1.md. - precision: fp16 (mixed: LayerNorm parameters in fp32) -- decoder: plain greedy (KV-cache variant ships with the WO02 export) +- decoder: kv greedy (plain fallback included in the zip) - metrics: CER 27.42 / EM 59.66 on 895 held-out pairs — secryst/docs/RESULTS.md#khmer-transliteration-2026-08-14 - trained from: secryst train_khmer_byt5.py run-001 diff --git a/models/khm-latn/khm-latn-1.0.metadata.yaml b/models/khm-latn/khm-latn-1.0.metadata.yaml new file mode 100644 index 0000000..20010c2 --- /dev/null +++ b/models/khm-latn/khm-latn-1.0.metadata.yaml @@ -0,0 +1,26 @@ +format: imf-v1 +id: khm-latn-1.0 +task: translit +source_script: Khmr +target: Latn +tokenizer: bytes +opset: 14 +decoder: kv +precision: fp32 +license: BSD-3-Clause +trained_from: >- + secryst train_khmer_byt5.py run-001; checkpoint + secryst-checkpoints:/khmer_byt5/run-001/best +metrics: + - name: cer + value: 27.42 + protocol: >- + greedy decode; 895 held-out pairs; split 16,120/895/895 seed 42; + ByT5-small early stop @ep15 + source: secryst/docs/RESULTS.md#khmer-transliteration-2026-08-14 + - name: em + value: 59.66 + protocol: >- + greedy decode; 895 held-out pairs; split 16,120/895/895 seed 42; + ByT5-small early stop @ep15 + source: secryst/docs/RESULTS.md#khmer-transliteration-2026-08-14 diff --git a/models/urd-diac/urd-diac-1.0.README.md b/models/urd-diac/urd-diac-1.0.README.md new file mode 100644 index 0000000..a698138 --- /dev/null +++ b/models/urd-diac/urd-diac-1.0.README.md @@ -0,0 +1,41 @@ +# urd-diac-1.0 + +Urdu diacritization (adds haraqat). Byte-level seq2seq (ByT5-small): +the tokenizer is raw UTF-8 bytes (pad=0, EOS=1) — no vocab files. +IMF v1 artifact; format spec: interscript/ml-models docs/imf-v1.md. + +- decoder: kv greedy (plain fallback included in the zip) +- metrics: CER 3.74 on 11,940 held-out — + rababa-urdu/docs/RESULTS.md#diacritization-urdu-text--text--haraqat +- trained from: rababa-urdu modal_app_diacrit.py run-001 + (urdu-diacrit-checkpoints:/urdu_diacrit/run-001/best) +- license: BSD-3-Clause + +## Usage + +Ruby (secryst gem, the Ruby binding of interscript-ml): + +```ruby +require "secryst" +translator = Secryst::Translator.new(model: "urd-diac-1.0") +translator.translate("اردو") +``` + +TypeScript (@interscript/ml): + +```ts +import { loadModel } from "@interscript/ml"; +const model = await loadModel("urd-diac-1.0"); +await model.translate("اردو"); +``` + +Python (interscript-ml): + +```python +from interscript_ml import Model +model = Model.load("urd-diac-1.0") +model.translate("اردو") +``` + +All three runtimes verify the sha256 of every ONNX member in this zip +against metadata.yaml before loading. diff --git a/models/urd-diac/urd-diac-1.0.metadata.yaml b/models/urd-diac/urd-diac-1.0.metadata.yaml new file mode 100644 index 0000000..55aae6d --- /dev/null +++ b/models/urd-diac/urd-diac-1.0.metadata.yaml @@ -0,0 +1,20 @@ +format: imf-v1 +id: urd-diac-1.0 +task: diacritization +source_script: Arab +target: Arab +tokenizer: bytes +opset: 14 +decoder: kv +precision: fp32 +license: BSD-3-Clause +trained_from: >- + rababa-urdu modal_app_diacrit.py run-001; checkpoint + urdu-diacrit-checkpoints:/urdu_diacrit/run-001/best +metrics: + - name: cer + value: 3.74 + protocol: >- + greedy decode; 11,940 held-out; labels derived IPA->haraqat + (deterministic conversion, 597K pairs); ByT5-small, 2 epochs + source: rababa-urdu/docs/RESULTS.md#diacritization-urdu-text--text--haraqat diff --git a/models/urd-g2p/urd-g2p-1.0.README.md b/models/urd-g2p/urd-g2p-1.0.README.md new file mode 100644 index 0000000..6a922d1 --- /dev/null +++ b/models/urd-g2p/urd-g2p-1.0.README.md @@ -0,0 +1,41 @@ +# urd-g2p-1.0 + +Urdu → IPA grapheme-to-phoneme conversion. Byte-level seq2seq +(ByT5-small): the tokenizer is raw UTF-8 bytes (pad=0, EOS=1) — no vocab +files. IMF v1 artifact; format spec: interscript/ml-models docs/imf-v1.md. + +- decoder: kv greedy (plain fallback included in the zip) +- metrics: CER 14.77 / EM 33.6 on 12,699 held-out words — + rababa-urdu/docs/RESULTS.md#g2p-urdu-text--ipa +- trained from: rababa-urdu modal_app.py run-001 + (urdu-g2p-checkpoints:/urdu_g2p/run-001/best) +- license: BSD-3-Clause + +## Usage + +Ruby (secryst gem, the Ruby binding of interscript-ml): + +```ruby +require "secryst" +translator = Secryst::Translator.new(model: "urd-g2p-1.0") +translator.translate("اردو") +``` + +TypeScript (@interscript/ml): + +```ts +import { loadModel } from "@interscript/ml"; +const model = await loadModel("urd-g2p-1.0"); +await model.translate("اردو"); +``` + +Python (interscript-ml): + +```python +from interscript_ml import Model +model = Model.load("urd-g2p-1.0") +model.translate("اردو") +``` + +All three runtimes verify the sha256 of every ONNX member in this zip +against metadata.yaml before loading. diff --git a/models/urd-g2p/urd-g2p-1.0.metadata.yaml b/models/urd-g2p/urd-g2p-1.0.metadata.yaml new file mode 100644 index 0000000..d345e0a --- /dev/null +++ b/models/urd-g2p/urd-g2p-1.0.metadata.yaml @@ -0,0 +1,26 @@ +format: imf-v1 +id: urd-g2p-1.0 +task: g2p +source_script: Arab +target: IPA +tokenizer: bytes +opset: 14 +decoder: kv +precision: fp32 +license: BSD-3-Clause +trained_from: >- + rababa-urdu modal_app.py run-001; checkpoint + urdu-g2p-checkpoints:/urdu_g2p/run-001/best +metrics: + - name: cer + value: 14.77 + protocol: >- + greedy decode; 12,699 held-out words from the 635K humair025 + urdu-g2p dictionary; ByT5-small + source: rababa-urdu/docs/RESULTS.md#g2p-urdu-text--ipa + - name: em + value: 33.6 + protocol: >- + greedy decode; 12,699 held-out words from the 635K humair025 + urdu-g2p dictionary; ByT5-small + source: rababa-urdu/docs/RESULTS.md#g2p-urdu-text--ipa diff --git a/src/gpu/modal_export.py b/src/gpu/modal_export.py new file mode 100644 index 0000000..8c16a40 --- /dev/null +++ b/src/gpu/modal_export.py @@ -0,0 +1,134 @@ +"""Modal app: export IMF v1 zips from checkpoints on Modal volumes (WO02). + +One CPU function per model — exports never compete with A100 training. + + modal run --detach src/gpu/modal_export.py --model khm-latn + modal run --detach src/gpu/modal_export.py --model urd-g2p --precisions fp16,int8 + +Watchdog (server evictions happen; export is idempotent, so retries are +the resume mechanism — each model's zips are written atomically at the +end, and per-model work is independent): + + until modal run --detach src/gpu/modal_export.py --model khm-latn; do sleep 60; done + +Outputs land on the secryst-models volume under /imf//. +""" + +from __future__ import annotations + +from pathlib import Path + +import modal + +REPO_ROOT = Path(__file__).resolve().parent.parent.parent + +IMAGE = ( + modal.Image.debian_slim(python_version="3.11") + .pip_install( + "torch>=2.4", + "transformers>=5.0", + "onnx>=1.16", + "onnxruntime>=1.17", + "pyyaml>=6.0", + ) + .copy_directory(str(REPO_ROOT), "/root/ml-models") + .workdir("/root/ml-models") +) + +CHECKPOINT_VOLUMES = { + "/volumes/secryst-checkpoints": modal.Volume.from_name("secryst-checkpoints"), + "/volumes/urdu-g2p-checkpoints": modal.Volume.from_name("urdu-g2p-checkpoints"), + "/volumes/urdu-diacrit-checkpoints": modal.Volume.from_name( + "urdu-diacrit-checkpoints" + ), + "/volumes/rababa-checkpoints": modal.Volume.from_name("rababa-checkpoints"), +} + +MODELS_VOLUME = modal.Volume.from_name("secryst-models") + +# model id -> (checkpoint volume mount, checkpoint path, metadata, readme) +MODELS: dict[str, dict[str, str]] = { + "khm-latn": { + "volume": "/volumes/secryst-checkpoints", + "checkpoint": "/khmer_byt5/run-001/best", + "metadata": "models/khm-latn/khm-latn-1.0.metadata.yaml", + "readme": "models/khm-latn/khm-latn-1.0.README.md", + "probe": "ភាសា", + }, + "urd-g2p": { + "volume": "/volumes/urdu-g2p-checkpoints", + "checkpoint": "/urdu_g2p/run-001/best", + "metadata": "models/urd-g2p/urd-g2p-1.0.metadata.yaml", + "readme": "models/urd-g2p/urd-g2p-1.0.README.md", + "probe": "اردو", + }, + "urd-diac": { + "volume": "/volumes/urdu-diacrit-checkpoints", + "checkpoint": "/urdu_diacrit/run-001/best", + "metadata": "models/urd-diac/urd-diac-1.0.metadata.yaml", + "readme": "models/urd-diac/urd-diac-1.0.README.md", + "probe": "اردو", + }, +} + +app = modal.App("interscript-ml-export", image=IMAGE) + + +@app.function( + cpu=8, + memory=32 * 1024, + timeout=2 * 3600, + volumes={**CHECKPOINT_VOLUMES, "/outputs": MODELS_VOLUME}, +) +def export_model(model_id: str, precisions: list[str]) -> dict[str, str]: + import sys + + sys.path.insert(0, "/root/ml-models/src") + + spec = MODELS[model_id] + checkpoint = Path(spec["volume"]) / spec["checkpoint"] + metadata_path = Path("/root/ml-models") / spec["metadata"] + readme_path = Path("/root/ml-models") / spec["readme"] + + from imf.export import export_zips, load_byte_seq2seq, onnx_greedy_kv + from imf.validator import validate_zip + + model = load_byte_seq2seq(checkpoint) + out_dir = Path("/outputs/imf") / model_id + zips = export_zips( + model, + metadata_path, + readme_path.read_text(encoding="utf-8"), + out_dir, + precisions=tuple(precisions), + ) + MODELS_VOLUME.commit() + + report: dict[str, str] = {} + import zipfile + + import onnxruntime as ort + + for z in zips: + result = validate_zip(z) + if not result.ok: + raise RuntimeError(f"{z.name} failed validation: {result.errors}") + with zipfile.ZipFile(z) as zf: + zf.extract("encoder.onnx", "/tmp/check") + zf.extract("decoder-kv.onnx", "/tmp/check") + enc = ort.InferenceSession( + "/tmp/check/encoder.onnx", providers=["CPUExecutionProvider"] + ) + kv = ort.InferenceSession( + "/tmp/check/decoder-kv.onnx", providers=["CPUExecutionProvider"] + ) + tokens = onnx_greedy_kv(enc, kv, spec["probe"], max_len=32) + report[z.name] = f"{z.stat().st_size} bytes, probe -> {len(tokens)} tokens" + return report + + +@app.local_entrypoint() +def main(model: str, precisions: str = "fp32,fp16,int8") -> None: + report = export_model.remote(model, precisions.split(",")) + for name, status in report.items(): + print(f"{name}: {status}") diff --git a/src/imf/export.py b/src/imf/export.py new file mode 100644 index 0000000..86bdca9 --- /dev/null +++ b/src/imf/export.py @@ -0,0 +1,385 @@ +"""ONNX export of byte-level (ByT5/T5) checkpoints into IMF v1 zips. + +Generalizes secryst PR #44's scripts/export_onnx_byt5.py: + +- encoder.onnx input_ids -> last_hidden_state +- decoder.onnx input_ids, encoder_hidden_states -> logits (fallback) +- decoder-kv.onnx + past_* inputs / present_* outputs (default artifact) +- opset pinned to 14 (the Ruby onnxruntime gem's bundled ORT is old) +- fp16 (keep IO fp32) and int8 (dynamic quantization) variants +- fixture mode: a tiny random T5 for CI, via ``make_fixture_checkpoint`` + +The KV graph caches self-attention only (``past_key_i`` / ``past_value_i`` +inputs, ``present_key_i`` / ``present_value_i`` outputs). Cross-attention +K/V are a deterministic projection of ``encoder_hidden_states``, so they +are recomputed each step instead of being cached — this keeps the +attention-mask length consistent for any past length and spares runtimes +all cross-cache bookkeeping. Step 0 feeds zero-length self pasts. +""" + +from __future__ import annotations + +from dataclasses import replace +from pathlib import Path + +from imf.pack import pack_zip +from imf.schema import ModelMetadata + +OPSET = 14 +GRAPH_NAMES = ("encoder.onnx", "decoder.onnx", "decoder-kv.onnx") + + +def load_byte_seq2seq(checkpoint_dir: Path | str): + """Load a T5-family checkpoint in eager attention (export-safe).""" + from transformers import AutoModelForSeq2SeqLM + + model = AutoModelForSeq2SeqLM.from_pretrained( + checkpoint_dir, attn_implementation="eager" + ).eval() + return model + + +def make_fixture_checkpoint(out_dir: Path | str, seed: int = 42) -> Path: + """Tiny random T5 with a byte-sized vocab, for CI and --fixture runs.""" + import torch + from transformers import T5Config, T5ForConditionalGeneration + + out_dir = Path(out_dir) + out_dir.mkdir(parents=True, exist_ok=True) + torch.manual_seed(seed) + config = T5Config( + vocab_size=384, + d_model=32, + d_kv=8, + d_ff=64, + num_layers=2, + num_heads=4, + decoder_start_token_id=0, + feed_forward_proj="relu", + relative_attention_num_buckets=8, + relative_attention_max_distance=16, + tie_word_embeddings=False, + ) + config._attn_implementation = "eager" + model = T5ForConditionalGeneration(config).eval() + model.save_pretrained(out_dir) + return out_dir + + +def _decoder_plain(model): + import torch.nn as nn + + class DecoderPlain(nn.Module): + def __init__(self, model): + super().__init__() + self.decoder = model.get_decoder() + self.lm_head = model.lm_head + self.scale = model.config.d_model ** -0.5 + + def forward(self, input_ids, encoder_hidden_states): + hidden = self.decoder( + input_ids=input_ids, encoder_hidden_states=encoder_hidden_states + )[0] + return self.lm_head(hidden * self.scale) + + return DecoderPlain(model) + + +def _decoder_kv(model): + import torch.nn as nn + from transformers.cache_utils import DynamicCache, EncoderDecoderCache + + num_layers = model.config.num_layers + + class DecoderKV(nn.Module): + def __init__(self, model): + super().__init__() + self.decoder = model.get_decoder() + self.lm_head = model.lm_head + self.scale = model.config.d_model ** -0.5 + + def forward(self, input_ids, encoder_hidden_states, *pasts): + # is_updated must stay unset: EncoderDecoderCache.__init__ bakes + # it as a Python bool from cross-cache lengths, and a True would + # trace the "reuse cross K/V" branch, which mismatches the + # attention mask for any fed cross-past length. Cross K/V are a + # deterministic projection of encoder_hidden_states anyway, so + # the graph recomputes them each step and never caches them. + cache = EncoderDecoderCache(DynamicCache(), DynamicCache()) + for i in range(num_layers): + cache.self_attention_cache.update(pasts[2 * i], pasts[2 * i + 1], i) + hidden = self.decoder( + input_ids=input_ids, + encoder_hidden_states=encoder_hidden_states, + past_key_values=cache, + use_cache=True, + )[0] + logits = self.lm_head(hidden * self.scale) + outputs = [logits] + for i in range(num_layers): + outputs.append(cache.self_attention_cache.layers[i].keys) + outputs.append(cache.self_attention_cache.layers[i].values) + return tuple(outputs) + + return DecoderKV(model) + + +def _kv_io_names(num_layers: int) -> tuple[list[str], list[str]]: + inputs, outputs = [], ["logits"] + for i in range(num_layers): + inputs += [f"past_key_{i}", f"past_value_{i}"] + outputs += [f"present_key_{i}", f"present_value_{i}"] + return inputs, outputs + + +def _sample_pasts(model, encoder_hidden_states): + """Harvest one decoder step's self-attention cache for shape examples.""" + import torch + + out = model.get_decoder()( + input_ids=torch.tensor([[0]]), + encoder_hidden_states=encoder_hidden_states, + use_cache=True, + ) + pasts = [] + for layer in out[1].self_attention_cache.layers: + pasts += [layer.keys.clone(), layer.values.clone()] + return pasts + + +def export_graphs(model, out_dir: Path | str) -> dict[str, Path]: + """Export encoder + plain decoder + KV decoder as fp32 ONNX, opset 14. + + Traces a deepcopy: TorchScript tracing leaves the traced module's + runtime behavior subtly altered (observed empirically), and callers + (e.g. the WO03 parity harness) need the reference model pristine. + """ + import copy + + import torch + + model = copy.deepcopy(model) + out_dir = Path(out_dir) + out_dir.mkdir(parents=True, exist_ok=True) + + ids = torch.tensor([[104, 101]]) # "he" + with torch.no_grad(): + hidden = model.get_encoder()(input_ids=ids)[0] + + paths = {} + + torch.onnx.export( + model.get_encoder(), + (ids,), + out_dir / "encoder.onnx", + input_names=["input_ids"], + output_names=["last_hidden_state"], + opset_version=OPSET, + dynamo=False, + dynamic_axes={ + "input_ids": {0: "batch", 1: "seq"}, + "last_hidden_state": {0: "batch", 1: "seq"}, + }, + ) + paths["encoder.onnx"] = out_dir / "encoder.onnx" + + torch.onnx.export( + _decoder_plain(model), + (torch.tensor([[0]]), hidden), + out_dir / "decoder.onnx", + input_names=["input_ids", "encoder_hidden_states"], + output_names=["logits"], + opset_version=OPSET, + dynamo=False, + dynamic_axes={ + "input_ids": {0: "batch", 1: "seq"}, + "encoder_hidden_states": {0: "batch", 1: "seq"}, + "logits": {0: "batch", 1: "seq"}, + }, + ) + paths["decoder.onnx"] = out_dir / "decoder.onnx" + + num_layers = model.config.num_layers + pasts = _sample_pasts(model, hidden) + kv_inputs = ["input_ids", "encoder_hidden_states"] + _kv_io_names(num_layers)[0] + kv_outputs = _kv_io_names(num_layers)[1] + dynamic = { + "input_ids": {0: "batch", 1: "cur_seq"}, + "encoder_hidden_states": {0: "batch", 1: "enc_seq"}, + "logits": {0: "batch", 1: "cur_seq"}, + } + for name in kv_inputs[2:]: + dynamic[name] = {0: "batch", 2: "past_seq"} + for name in kv_outputs[1:]: + dynamic[name] = {0: "batch", 2: "present_seq"} + + with torch.no_grad(): + torch.onnx.export( + _decoder_kv(model), + (torch.tensor([[0]]), hidden, *pasts), + out_dir / "decoder-kv.onnx", + input_names=kv_inputs, + output_names=kv_outputs, + opset_version=OPSET, + dynamo=False, + dynamic_axes=dynamic, + ) + paths["decoder-kv.onnx"] = out_dir / "decoder-kv.onnx" + return paths + + +def convert_fp16(src: Path | str, dst: Path | str) -> Path: + """fp32 -> mixed fp16, IO types preserved (encoder/decoder compose cleanly). + + LayerNorm/softmax math stays fp32: ORT's session-time + SimplifiedLayerNormFusion crashes on half-converted LN subgraphs + (InsertPrecisionFreeCast name mismatch), so the whole decomposition + must stay one dtype. Weights (MatMuls) carry the size win. + """ + import onnx + from onnxruntime.transformers import float16 + + block_list = list(float16.DEFAULT_OP_BLOCK_LIST) + [ + "ReduceMean", "Pow", "Sqrt", "Div", "Sub", "Add", "Mul", + "Softmax", "Range", "Exp", "Where", "Less", "Cast", + ] + model = onnx.load(str(src)) + converted = float16.convert_float_to_float16( + model, keep_io_types=True, op_block_list=block_list + ) + onnx.save(converted, str(dst)) + return Path(dst) + + +def quantize_int8(src: Path | str, dst: Path | str) -> Path: + """fp32 -> dynamically quantized int8 (MatMul weights QInt8). + + MatMul-only: quantizing other ops inserts precision casts that break + ORT's session-time SimplifiedLayerNormFusion, and preprocessing the + graph (quant_pre_process) pins concrete example shapes into + DynamicQuantizeLinear buffers. + """ + from onnxruntime.quantization import QuantType, quantize_dynamic + + quantize_dynamic( + str(src), + str(dst), + weight_type=QuantType.QInt8, + op_types_to_quantize=["MatMul"], + ) + return Path(dst) + + +def onnx_greedy_plain(encoder_sess, decoder_sess, text: str, max_len: int = 256) -> list[int]: + """Greedy decode over ONNX sessions (plain decoder). Self-check helper. + + Returns generated token ids (the vocab is 384-wide; only a trained + byte-level model reliably stays < 256).""" + import numpy as np + + ids = np.array([list(text.encode("utf-8"))], dtype=np.int64) + if ids.shape[1] == 0: + return [] + hidden = encoder_sess.run(None, {"input_ids": ids})[0] + dec_ids = np.array([[0]], dtype=np.int64) + generated: list[int] = [] + for _ in range(max_len): + logits = decoder_sess.run( + None, {"input_ids": dec_ids, "encoder_hidden_states": hidden} + )[0] + nxt = int(np.argmax(logits[0, -1])) + if nxt == 1: + break + generated.append(nxt) + dec_ids = np.concatenate([dec_ids, np.array([[nxt]], dtype=np.int64)], axis=1) + return generated + + +def _zero_pasts(kv_sess) -> dict[str, object]: + """Zero-length past inputs for step 0, shapes from session metadata.""" + import numpy as np + + pasts = {} + for meta in kv_sess.get_inputs(): + if not meta.name.startswith("past_"): + continue + shape = meta.shape # [batch, heads, past_seq, d_kv] with str dynamic dims + heads = shape[1] if isinstance(shape[1], int) else 4 + d_kv = shape[3] if isinstance(shape[3], int) else 8 + pasts[meta.name] = np.zeros((1, heads, 0, d_kv), dtype=np.float32) + return pasts + + +def onnx_greedy_kv(encoder_sess, kv_sess, text: str, max_len: int = 256) -> list[int]: + """Greedy decode over ONNX sessions (KV decoder). Self-check helper.""" + import numpy as np + + ids = np.array([list(text.encode("utf-8"))], dtype=np.int64) + if ids.shape[1] == 0: + return [] + hidden = encoder_sess.run(None, {"input_ids": ids})[0] + out_names = [o.name for o in kv_sess.get_outputs()] + pasts = _zero_pasts(kv_sess) + cur = np.array([[0]], dtype=np.int64) + generated: list[int] = [] + for _ in range(max_len): + out = kv_sess.run(None, {"input_ids": cur, "encoder_hidden_states": hidden, **pasts}) + results = dict(zip(out_names, out, strict=True)) + nxt = int(np.argmax(results["logits"][0, -1])) + if nxt == 1: + break + generated.append(nxt) + pasts = { + name: results[name.replace("past_", "present_", 1)] + for name in pasts + } + cur = np.array([[nxt]], dtype=np.int64) + return generated + + +def export_zips( + model, + metadata_path: Path | str, + readme: str, + out_dir: Path | str, + precisions: tuple[str, ...] = ("fp32", "fp16", "int8"), +) -> list[Path]: + """Full pipeline: loaded model -> fp32 graphs -> precision variants -> IMF zips.""" + import tempfile + + metadata = ModelMetadata.from_yaml(Path(metadata_path).read_text(encoding="utf-8")) + if metadata.decoder != "kv": + raise ValueError("WO02 exports declare decoder: kv (plain is the fallback)") + + out_dir = Path(out_dir) + out_dir.mkdir(parents=True, exist_ok=True) + zips = [] + + with tempfile.TemporaryDirectory() as tmp: + tmp = Path(tmp) + graphs = export_graphs(model, tmp / "graphs") + + for precision in precisions: + variant_dir = tmp / precision + variant_dir.mkdir() + for name, src in graphs.items(): + dst = variant_dir / name + if precision == "fp32": + dst.write_bytes(src.read_bytes()) + elif precision == "fp16": + convert_fp16(src, dst) + elif precision == "int8": + quantize_int8(src, dst) + else: + raise ValueError(f"unknown precision {precision!r}") + meta = replace(metadata, precision=precision) + zips.append( + pack_zip( + variant_dir, + meta, + readme, + out_dir / f"{metadata.id}-{precision}.zip", + ) + ) + return zips + diff --git a/tests/test_imf_export.py b/tests/test_imf_export.py new file mode 100644 index 0000000..269e742 --- /dev/null +++ b/tests/test_imf_export.py @@ -0,0 +1,153 @@ +"""Tests for ``imf.export`` — the WO02 ONNX export pipeline. + +Real end-to-end over a tiny random T5 (fixture mode): export all three +graphs at opset 14, pack fp32/fp16/int8 zips, and require every +precision to greedy-decode IDENTICALLY to the torch reference model — +plain and KV paths both. No mocks anywhere; skipped when torch / +transformers / onnxruntime are absent. +""" + +from __future__ import annotations + +import zipfile +from pathlib import Path + +import pytest +import yaml + +torch = pytest.importorskip("torch") +transformers = pytest.importorskip("transformers") +ort = pytest.importorskip("onnxruntime") + +from imf.export import ( # noqa: E402 + export_zips, + load_byte_seq2seq, + make_fixture_checkpoint, + onnx_greedy_kv, + onnx_greedy_plain, +) +from imf.validator import validate_zip # noqa: E402 + +TEXTS = ["he", "hello", "abc"] +MAX_LEN = 12 +PROVIDERS = ["CPUExecutionProvider"] + +METADATA = { + "format": "imf-v1", + "id": "fixture-1.0", + "task": "translit", + "source_script": "Latn", + "target": "Latn", + "tokenizer": "bytes", + "opset": 14, + "decoder": "kv", + "precision": "fp32", + "license": "BSD-3-Clause", + "trained_from": "imf export fixture (seed 42)", + "metrics": [ + { + "name": "cer", + "value": 0.0, + "protocol": "fixture self-check", + "source": "ml-models/tests/test_imf_export.py#fixture", + } + ], +} + + +def _torch_greedy(model, text: str) -> list[int]: + ids = torch.tensor([list(text.encode("utf-8"))]) + enc = model.get_encoder()(input_ids=ids)[0] + dec_ids = torch.tensor([[0]]) + outs: list[int] = [] + for _ in range(MAX_LEN): + logits = model.lm_head( + model.get_decoder()(input_ids=dec_ids, encoder_hidden_states=enc)[0] + * (model.config.d_model ** -0.5) + ) + nxt = int(logits[0, -1].argmax()) + if nxt == 1: + break + outs.append(nxt) + dec_ids = torch.cat([dec_ids, torch.tensor([[nxt]])], 1) + return outs + + +@pytest.fixture(scope="module") +def reference_model(tmp_path_factory: pytest.TempPathFactory): + ckpt = make_fixture_checkpoint(tmp_path_factory.mktemp("fixture") / "checkpoint") + return load_byte_seq2seq(ckpt) + + +@pytest.fixture(scope="module") +def zips( + tmp_path_factory: pytest.TempPathFactory, reference_model +) -> dict[str, Path]: + root = tmp_path_factory.mktemp("export") + metadata_path = root / "metadata.yaml" + metadata_path.write_text(yaml.safe_dump(METADATA), encoding="utf-8") + paths = export_zips( + reference_model, metadata_path, "# fixture\n", root / "out" + ) + return {p.name: p for p in paths} + + +def _sessions(zip_path: Path): + out = zip_path.parent / zip_path.stem + out.mkdir(exist_ok=True) + with zipfile.ZipFile(zip_path) as zf: + zf.extractall(out) + enc = ort.InferenceSession(str(out / "encoder.onnx"), providers=PROVIDERS) + dec = ort.InferenceSession(str(out / "decoder.onnx"), providers=PROVIDERS) + kv = ort.InferenceSession(str(out / "decoder-kv.onnx"), providers=PROVIDERS) + return enc, dec, kv + + +def test_fixture_exports_match_torch(reference_model) -> None: + """KV and plain ONNX decode must equal the (untraced) torch model.""" + import tempfile + + from imf.export import export_graphs + + with tempfile.TemporaryDirectory() as tmp: + graphs = export_graphs(reference_model, Path(tmp) / "graphs") + enc = ort.InferenceSession(str(graphs["encoder.onnx"]), providers=PROVIDERS) + dec = ort.InferenceSession(str(graphs["decoder.onnx"]), providers=PROVIDERS) + kv = ort.InferenceSession(str(graphs["decoder-kv.onnx"]), providers=PROVIDERS) + for text in TEXTS: + expected = _torch_greedy(reference_model, text) + assert onnx_greedy_plain(enc, dec, text, MAX_LEN) == expected + assert onnx_greedy_kv(enc, kv, text, MAX_LEN) == expected + + +def test_all_precisions_ship_and_validate(zips: dict[str, Path]) -> None: + assert set(zips) == { + "fixture-1.0-fp32.zip", + "fixture-1.0-fp16.zip", + "fixture-1.0-int8.zip", + } + for name, path in zips.items(): + result = validate_zip(path) + assert result.ok, result.errors + assert result.metadata is not None + assert result.metadata.decoder == "kv" + assert result.metadata.opset == 14 + assert result.metadata.precision in name + + +def test_precisions_are_token_exact(zips: dict[str, Path]) -> None: + """fp16/int8 must not change a single generated token vs fp32.""" + enc32, _, kv32 = _sessions(zips["fixture-1.0-fp32.zip"]) + reference = { + text: onnx_greedy_kv(enc32, kv32, text, MAX_LEN) for text in TEXTS + } + for precision in ("fp16", "int8"): + enc, _, kv = _sessions(zips[f"fixture-1.0-{precision}.zip"]) + for text in TEXTS: + assert onnx_greedy_kv(enc, kv, text, MAX_LEN) == reference[text] + + +def test_fp16_smaller_than_fp32(zips: dict[str, Path]) -> None: + assert zips["fixture-1.0-fp16.zip"].stat().st_size < zips[ + "fixture-1.0-fp32.zip" + ].stat().st_size From f2b415223dd375b5f01fe75a5b24a85c2dfc29bf Mon Sep 17 00:00:00 2001 From: Ronald Tse Date: Sun, 16 Aug 2026 12:56:10 +0800 Subject: [PATCH 2/9] feat(imf): WO03 parity + checksum gate MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - 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 --- .github/workflows/test.yml | 4 +- docs/imf-v1.md | 18 ++++ src/imf/cli.py | 84 ++++++++++++++++ src/imf/parity.py | 197 +++++++++++++++++++++++++++++++++++++ tests/test_imf_parity.py | 105 ++++++++++++++++++++ 5 files changed, 406 insertions(+), 2 deletions(-) create mode 100644 src/imf/parity.py create mode 100644 tests/test_imf_parity.py diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml index 9ac083c..b8b5155 100644 --- a/.github/workflows/test.yml +++ b/.github/workflows/test.yml @@ -53,5 +53,5 @@ jobs: - run: | python -m pip install --upgrade pip pip install -e ".[dev,train,export]" - - name: IMF export fixture tests (torch + ORT parity) - run: PYTHONPATH=src python -m pytest tests/test_imf_export.py -v + - name: IMF export + parity gate fixture tests (torch vs ORT) + run: PYTHONPATH=src python -m pytest tests/test_imf_export.py tests/test_imf_parity.py -v diff --git a/docs/imf-v1.md b/docs/imf-v1.md index 776ffa5..afc2210 100644 --- a/docs/imf-v1.md +++ b/docs/imf-v1.md @@ -111,8 +111,26 @@ PYTHONPATH=src python -m imf validate --strict # release gate PYTHONPATH=src python -m imf info # print manifest PYTHONPATH=src python -m imf pack --source \ --metadata [--readme ] --out # sha256 computed +PYTHONPATH=src python -m imf parity --checkpoint \ + --test-data # WO03 gate; writes parity into the zip +PYTHONPATH=src python -m imf golden --inputs --out \ + # cross-runtime golden set: 100 fixed strings, Python = reference ``` +The parity gate compares ONNX KV greedy decode against the transformers +decoder loop (the exact math the export wraps — not `generate`, whose +config-dependent behavior no runtime implements) over >= 500 test pairs; +it writes `{samples, cer_delta}` into metadata and refuses to leave the +zip non-strict. On Modal the same gate runs headless: + +``` +modal run --detach src/gpu/modal_export.py::main --model urd-g2p +modal run --detach src/gpu/modal_export.py::parity --model urd-g2p +``` + +CI runs the full gate on the fixture model (export -> parity -> +strict-validate) in the `export-fixture` job. + Legacy notes: - Old secryst zips (`vocabs.yaml` + single `transformer.onnx`) and the diff --git a/src/imf/cli.py b/src/imf/cli.py index 82e96b6..cadb089 100644 --- a/src/imf/cli.py +++ b/src/imf/cli.py @@ -81,6 +81,64 @@ def _default_readme(metadata: ModelMetadata) -> str: ) +def _load_pairs(path: Path) -> list[tuple[str, str]]: + import json + + pairs: list[tuple[str, str]] = [] + for line in path.read_text(encoding="utf-8").splitlines(): + if not line.strip(): + continue + row = json.loads(line) + if isinstance(row, dict): + pairs.append( + ( + row.get("input", row.get("src", "")), + row.get("target", row.get("tgt", row.get("gold", ""))), + ) + ) + else: + pairs.append((row[0], row[1] if len(row) > 1 else "")) + return pairs + + +def _cmd_parity(args: argparse.Namespace) -> int: + from imf.export import load_byte_seq2seq + from imf.parity import run_parity, write_parity + + model = load_byte_seq2seq(args.checkpoint) + pairs = _load_pairs(args.test_data) + if args.limit: + pairs = pairs[: args.limit] + report = run_parity(model, args.zip, pairs, max_len=args.max_len) + print( + f"samples={report.samples} cer_ref={report.cer_reference}pp " + f"cer_onnx={report.cer_onnx}pp delta={report.cer_delta}pp " + f"token_mismatches={report.token_mismatches}" + ) + if not report.passed: + print("error: parity gate FAILED", file=sys.stderr) + return 1 + write_parity(args.zip, report) + print(f"parity written into {args.zip} (strict validation passed)") + return 0 + + +def _cmd_golden(args: argparse.Namespace) -> int: + import json + + from imf.parity import write_golden + + inputs: list[str] = [] + for line in args.inputs.read_text(encoding="utf-8").splitlines(): + if not line.strip(): + continue + row = json.loads(line) + inputs.append(row["input"] if isinstance(row, dict) else row[0]) + out = write_golden(args.zip, inputs, args.out, max_len=args.max_len) + print(f"wrote {len(inputs)} golden cases to {out}") + return 0 + + def build_parser() -> argparse.ArgumentParser: parser = argparse.ArgumentParser(prog="imf", description=__doc__) sub = parser.add_subparsers(dest="command", required=True) @@ -109,6 +167,32 @@ def build_parser() -> argparse.ArgumentParser: p_pack.add_argument("--out", required=True, type=Path) p_pack.set_defaults(func=_cmd_pack) + p_parity = sub.add_parser( + "parity", help="WO03 gate: ONNX vs torch reference, write parity into the zip" + ) + p_parity.add_argument("zip", type=Path) + p_parity.add_argument( + "--checkpoint", required=True, type=Path, help="HF checkpoint dir (reference)" + ) + p_parity.add_argument( + "--test-data", required=True, type=Path, + help="JSONL: {input, target} pairs (or [src, gold] arrays)", + ) + p_parity.add_argument("--limit", type=int, help="cap sample count") + p_parity.add_argument("--max-len", type=int, default=256) + p_parity.set_defaults(func=_cmd_parity) + + p_golden = sub.add_parser( + "golden", help="emit the cross-runtime golden JSONL from ONNX decode" + ) + p_golden.add_argument("zip", type=Path) + p_golden.add_argument( + "--inputs", required=True, type=Path, help="JSONL of input strings" + ) + p_golden.add_argument("--out", required=True, type=Path) + p_golden.add_argument("--max-len", type=int, default=256) + p_golden.set_defaults(func=_cmd_golden) + return parser diff --git a/src/imf/parity.py b/src/imf/parity.py new file mode 100644 index 0000000..477365e --- /dev/null +++ b/src/imf/parity.py @@ -0,0 +1,197 @@ +"""WO03 parity gate: ONNX greedy vs the torch reference, CER delta <= 0.2pp. + +The reference is the transformers decoder loop itself (the exact math the +export wraps) rather than ``model.generate`` — generate's behavior is +config-dependent (eos/start-token defaults) and none of it is implemented +by the runtimes. Comparing against the module-level math catches exactly +what export bugs can break. + +``write_parity`` rewrites the parity block inside an existing zip +(metadata.yaml is never sha256-covered, graphs are untouched) and then +requires the zip to pass strict validation — the release gate. +""" + +from __future__ import annotations + +import json +import zipfile +from dataclasses import dataclass +from pathlib import Path + +from framework.evaluator import char_error_rate +from imf.export import onnx_greedy_kv +from imf.schema import ModelMetadata, Parity + + +@dataclass(frozen=True) +class ParityReport: + samples: int + cer_reference: float # percentage points + cer_onnx: float + cer_delta: float + token_mismatches: int + + @property + def passed(self) -> bool: + return ( + self.cer_delta <= Parity.MAX_CER_DELTA + and self.samples >= Parity.MIN_SAMPLES + ) + + +def _torch_greedy_tokens(model, text: str, max_len: int) -> list[int]: + import torch + + ids = torch.tensor([list(text.encode("utf-8"))], dtype=torch.long) + if ids.shape[1] == 0: + return [] + enc = model.get_encoder()(input_ids=ids)[0] + dec_ids = torch.tensor([[0]], dtype=torch.long) + outs: list[int] = [] + for _ in range(max_len): + hidden = model.get_decoder()( + input_ids=dec_ids, encoder_hidden_states=enc + )[0] + logits = model.lm_head(hidden * (model.config.d_model ** -0.5)) + nxt = int(logits[0, -1].argmax()) + if nxt == 1: + break + outs.append(nxt) + dec_ids = torch.cat([dec_ids, torch.tensor([[nxt]], dtype=torch.long)], 1) + return outs + + +def _decode_tokens(tokens: list[int]) -> str: + return bytes(t % 256 for t in tokens).decode("utf-8", errors="replace") + + +def _sessions_from_zip(zip_path: Path): + import tempfile + + import onnxruntime as ort + + with tempfile.TemporaryDirectory() as tmp: + with zipfile.ZipFile(zip_path) as zf: + zf.extract("encoder.onnx", tmp) + decoder = "decoder-kv.onnx" if "decoder-kv.onnx" in zf.namelist() else "decoder.onnx" + zf.extract(decoder, tmp) + enc = ort.InferenceSession( + str(Path(tmp) / "encoder.onnx"), providers=["CPUExecutionProvider"] + ) + dec = ort.InferenceSession( + str(Path(tmp) / decoder), providers=["CPUExecutionProvider"] + ) + return enc, dec + + +def run_parity(model, zip_path: Path | str, pairs, max_len: int = 256) -> ParityReport: + """pairs: iterable of (source_text, gold_target). Measures both sides + against gold; the gate is the CER distance between the two.""" + zip_path = Path(zip_path) + enc, kv = _sessions_from_zip(zip_path) + + n = 0 + mismatches = 0 + cer_ref_sum = 0.0 + cer_onnx_sum = 0.0 + for source, gold in pairs: + n += 1 + ref = _torch_greedy_tokens(model, source, max_len) + got = onnx_greedy_kv(enc, kv, source, max_len) + if ref != got: + mismatches += 1 + cer_ref_sum += char_error_rate(_decode_tokens(ref), gold) + cer_onnx_sum += char_error_rate(_decode_tokens(got), gold) + + cer_ref = 100.0 * cer_ref_sum / max(n, 1) + cer_onnx = 100.0 * cer_onnx_sum / max(n, 1) + return ParityReport( + samples=n, + cer_reference=round(cer_ref, 4), + cer_onnx=round(cer_onnx, 4), + cer_delta=round(abs(cer_onnx - cer_ref), 4), + token_mismatches=mismatches, + ) + + +def write_parity(zip_path: Path | str, report: ParityReport) -> Path: + """Write the parity block into the zip's metadata and enforce strict + validation. Raises if the gate does not pass.""" + from imf.validator import validate_zip + + zip_path = Path(zip_path) + result = validate_zip(zip_path) + if not result.ok or result.metadata is None: + raise RuntimeError(f"cannot write parity into invalid zip: {result.errors}") + if not report.passed: + raise RuntimeError( + f"parity gate FAILED: cer_delta {report.cer_delta}pp over " + f"{report.samples} samples (limits: <= {Parity.MAX_CER_DELTA}pp, " + f">= {Parity.MIN_SAMPLES} samples)" + ) + + import tempfile + + metadata = result.metadata + updated = ModelMetadata( + format=metadata.format, + id=metadata.id, + task=metadata.task, + source_script=metadata.source_script, + target=metadata.target, + tokenizer=metadata.tokenizer, + opset=metadata.opset, + decoder=metadata.decoder, + precision=metadata.precision, + license=metadata.license, + trained_from=metadata.trained_from, + metrics=metadata.metrics, + parity=Parity(samples=report.samples, cer_delta=report.cer_delta), + sha256=metadata.sha256, + ) + + import yaml + + from imf.pack import _to_dict + + with tempfile.TemporaryDirectory() as tmp: + rewritten = Path(tmp) / "rewritten.zip" + with zipfile.ZipFile(zip_path) as src, zipfile.ZipFile( + rewritten, "w", zipfile.ZIP_DEFLATED + ) as dst: + for name in src.namelist(): + if name == "metadata.yaml": + dst.writestr( + name, + yaml.safe_dump(_to_dict(updated), sort_keys=False, allow_unicode=True), + ) + else: + dst.writestr(name, src.read(name)) + rewritten.replace(zip_path) + + strict = validate_zip(zip_path, strict=True) + if not strict.ok: + raise RuntimeError(f"strict gate failed after parity write: {strict.errors}") + return zip_path + + +def write_golden(zip_path: Path | str, inputs, out_path: Path | str, max_len: int = 256) -> Path: + """Emit the cross-runtime golden set: fixed inputs + reference outputs + from the ONNX graphs (Python is the reference implementation).""" + + + zip_path = Path(zip_path) + out_path = Path(out_path) + enc, kv = _sessions_from_zip(zip_path) + out_path.parent.mkdir(parents=True, exist_ok=True) + with out_path.open("w", encoding="utf-8") as fh: + for source in inputs: + tokens = onnx_greedy_kv(enc, kv, source, max_len) + fh.write( + json.dumps( + {"input": source, "tokens": tokens, "output": _decode_tokens(tokens)}, + ensure_ascii=False, + ) + + "\n" + ) + return out_path diff --git a/tests/test_imf_parity.py b/tests/test_imf_parity.py new file mode 100644 index 0000000..c02292c --- /dev/null +++ b/tests/test_imf_parity.py @@ -0,0 +1,105 @@ +"""Tests for ``imf.parity`` — the WO03 gate. + +Runs the real gate end-to-end on the fixture model (torch reference vs +ONNX KV decode over 600 pairs), then exercises the gate's failure modes +with real ParityReport values. No mocks. +""" + +from __future__ import annotations + +import json +from pathlib import Path + +import pytest +import yaml + +torch = pytest.importorskip("torch") +transformers = pytest.importorskip("transformers") +ort = pytest.importorskip("onnxruntime") + +from imf.export import ( # noqa: E402 + export_zips, + load_byte_seq2seq, + make_fixture_checkpoint, +) +from imf.parity import ParityReport, run_parity, write_golden, write_parity # noqa: E402 +from imf.validator import validate_zip # noqa: E402 + +METADATA = { + "format": "imf-v1", + "id": "fixture-1.0", + "task": "translit", + "source_script": "Latn", + "target": "Latn", + "tokenizer": "bytes", + "opset": 14, + "decoder": "kv", + "precision": "fp32", + "license": "BSD-3-Clause", + "trained_from": "imf export fixture (seed 42)", + "metrics": [ + { + "name": "cer", + "value": 0.0, + "protocol": "fixture self-check", + "source": "ml-models/tests/test_imf_parity.py#fixture", + } + ], +} + +PAIRS = [(text, "xxxxx") for text in ("he", "hello", "abc", "world")] + + +@pytest.fixture(scope="module") +def gated_zip(tmp_path_factory: pytest.TempPathFactory) -> Path: + ckpt = make_fixture_checkpoint(tmp_path_factory.mktemp("ckpt") / "fixture") + model = load_byte_seq2seq(ckpt) + root = tmp_path_factory.mktemp("parity") + metadata_path = root / "metadata.yaml" + metadata_path.write_text(yaml.safe_dump(METADATA), encoding="utf-8") + zips = export_zips(model, metadata_path, "# fixture\n", root / "out") + report = run_parity(model, zips[0], PAIRS * 150, max_len=12) + assert report.passed + write_parity(zips[0], report) + return zips[0] + + +def test_parity_report_is_exact_for_fp32(gated_zip: Path) -> None: + result = validate_zip(gated_zip, strict=True) + assert result.ok, result.errors + assert result.metadata is not None + assert result.metadata.parity is not None + assert result.metadata.parity.samples == 600 + assert result.metadata.parity.cer_delta == 0.0 + + +def test_gate_rejects_high_cer_delta(gated_zip: Path) -> None: + bad = ParityReport( + samples=600, cer_reference=10.0, cer_onnx=10.5, + cer_delta=0.5, token_mismatches=3, + ) + with pytest.raises(RuntimeError, match="parity gate FAILED"): + write_parity(gated_zip, bad) + + +def test_gate_rejects_small_sample(gated_zip: Path) -> None: + small = ParityReport( + samples=100, cer_reference=10.0, cer_onnx=10.0, + cer_delta=0.0, token_mismatches=0, + ) + with pytest.raises(RuntimeError, match="parity gate FAILED"): + write_parity(gated_zip, small) + + +def test_golden_jsonl_roundtrip(gated_zip: Path, tmp_path: Path) -> None: + inputs = tmp_path / "inputs.jsonl" + inputs.write_text( + "\n".join(json.dumps(t) for t in ("he", "hello", "abc")) + "\n", + encoding="utf-8", + ) + out = write_golden(gated_zip, ["he", "hello", "abc"], tmp_path / "golden.jsonl", max_len=12) + rows = [json.loads(line) for line in out.read_text(encoding="utf-8").splitlines()] + assert [r["input"] for r in rows] == ["he", "hello", "abc"] + for row in rows: + assert set(row) == {"input", "tokens", "output"} + assert all(isinstance(t, int) for t in row["tokens"]) From b7c757f25f8bad2f5b96b0e10b027b92ff649aaa Mon Sep 17 00:00:00 2001 From: Ronald Tse Date: Sun, 16 Aug 2026 12:58:08 +0800 Subject: [PATCH 3/9] feat(imf): heb-diac-1.0 (ByT5-base s43) in the export registry MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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. --- models/heb-diac/heb-diac-1.0.README.md | 42 +++++++ models/heb-diac/heb-diac-1.0.metadata.yaml | 26 ++++ src/gpu/modal_export.py | 131 ++++++++++++++++++--- 3 files changed, 180 insertions(+), 19 deletions(-) create mode 100644 models/heb-diac/heb-diac-1.0.README.md create mode 100644 models/heb-diac/heb-diac-1.0.metadata.yaml diff --git a/models/heb-diac/heb-diac-1.0.README.md b/models/heb-diac/heb-diac-1.0.README.md new file mode 100644 index 0000000..32bd344 --- /dev/null +++ b/models/heb-diac/heb-diac-1.0.README.md @@ -0,0 +1,42 @@ +# heb-diac-1.0 + +Hebrew diacritization (adds nikud). Byte-level seq2seq (ByT5-base): +the tokenizer is raw UTF-8 bytes (pad=0, EOS=1) — no vocab files. +IMF v1 artifact; format spec: interscript/ml-models docs/imf-v1.md. + +- decoder: kv greedy (plain fallback included in the zip) +- metrics: greedy DER 29.0% (the v1 runtime path); beam=4 DER 17.46% + (reference quality — beam search is not in v1 runtimes) — + rababa/docs/RESULTS.md#hebrew-diacritization +- trained from: rababa train_hebrew_seeds.py s43 run-001 + (rababa-checkpoints:/rababa_hebrew_byt5_s43/run-001/best) +- license: BSD-3-Clause + +## Usage + +Ruby (secryst gem, the Ruby binding of interscript-ml): + +```ruby +require "secryst" +translator = Secryst::Translator.new(model: "heb-diac-1.0") +translator.translate("שלום") +``` + +TypeScript (@interscript/ml): + +```ts +import { loadModel } from "@interscript/ml"; +const model = await loadModel("heb-diac-1.0"); +await model.translate("שלום"); +``` + +Python (interscript-ml): + +```python +from interscript_ml import Model +model = Model.load("heb-diac-1.0") +model.translate("שלום") +``` + +All three runtimes verify the sha256 of every ONNX member in this zip +against metadata.yaml before loading. diff --git a/models/heb-diac/heb-diac-1.0.metadata.yaml b/models/heb-diac/heb-diac-1.0.metadata.yaml new file mode 100644 index 0000000..085967c --- /dev/null +++ b/models/heb-diac/heb-diac-1.0.metadata.yaml @@ -0,0 +1,26 @@ +format: imf-v1 +id: heb-diac-1.0 +task: diacritization +source_script: Hebr +target: Hebr +tokenizer: bytes +opset: 14 +decoder: kv +precision: fp32 +license: BSD-3-Clause +trained_from: >- + rababa train_hebrew_seeds.py s43 run-001; checkpoint + rababa-checkpoints:/rababa_hebrew_byt5_s43/run-001/best +metrics: + - name: der_greedy + value: 29.0 + protocol: >- + beam=1 greedy decode (the v1 runtime path); Nakdimon test split, + 5,095 examples; ByT5-base s43 + source: rababa/docs/RESULTS.md#hebrew-diacritization + - name: der_beam4 + value: 17.46 + protocol: >- + beam=4 standard decode (reference quality; beam search is not in + v1 runtimes); Nakdimon test split, 5,095 examples; ByT5-base s43 + source: rababa/docs/RESULTS.md#hebrew-diacritization diff --git a/src/gpu/modal_export.py b/src/gpu/modal_export.py index 8c16a40..c5b717d 100644 --- a/src/gpu/modal_export.py +++ b/src/gpu/modal_export.py @@ -1,17 +1,20 @@ -"""Modal app: export IMF v1 zips from checkpoints on Modal volumes (WO02). - -One CPU function per model — exports never compete with A100 training. +"""Modal app: export IMF v1 zips from checkpoints on Modal volumes (WO02), +then gate them with the WO03 parity check — all CPU, never competing +with A100 training. modal run --detach src/gpu/modal_export.py --model khm-latn - modal run --detach src/gpu/modal_export.py --model urd-g2p --precisions fp16,int8 + modal run --detach src/gpu/modal_export.py::parity --model khm-latn -Watchdog (server evictions happen; export is idempotent, so retries are -the resume mechanism — each model's zips are written atomically at the -end, and per-model work is independent): +Watchdog (server evictions happen; both steps are idempotent — retries +are the resume mechanism, and each model's zips are written atomically): until modal run --detach src/gpu/modal_export.py --model khm-latn; do sleep 60; done -Outputs land on the secryst-models volume under /imf//. +Zips land on the secryst-models volume under /imf//; parity is +written into the zip in place (a zip is only shippable strict-validated). +Versions are pinned to the ones the export was verified against locally +(transformers 5.15 breaks T5 tracing with "multiple values for +use_cache"). """ from __future__ import annotations @@ -25,13 +28,13 @@ IMAGE = ( modal.Image.debian_slim(python_version="3.11") .pip_install( - "torch>=2.4", - "transformers>=5.0", - "onnx>=1.16", - "onnxruntime>=1.17", + "torch==2.12.1", + "transformers==5.14.1", + "onnx==1.22.0", + "onnxruntime==1.23.2", "pyyaml>=6.0", ) - .copy_directory(str(REPO_ROOT), "/root/ml-models") + .add_local_dir(str(REPO_ROOT), "/root/ml-models", copy=True) .workdir("/root/ml-models") ) @@ -44,29 +47,50 @@ "/volumes/rababa-checkpoints": modal.Volume.from_name("rababa-checkpoints"), } +DATASET_VOLUMES = { + "/datasets/rababa": modal.Volume.from_name("rababa-datasets"), + "/datasets/secryst": modal.Volume.from_name("secryst-datasets"), + "/datasets/urdu-g2p": modal.Volume.from_name("urdu-g2p-datasets"), + "/datasets/urdu-diacrit": modal.Volume.from_name("urdu-diacrit-datasets"), +} + MODELS_VOLUME = modal.Volume.from_name("secryst-models") -# model id -> (checkpoint volume mount, checkpoint path, metadata, readme) MODELS: dict[str, dict[str, str]] = { "khm-latn": { "volume": "/volumes/secryst-checkpoints", - "checkpoint": "/khmer_byt5/run-001/best", + "checkpoint": "khmer_byt5/run-001/best", "metadata": "models/khm-latn/khm-latn-1.0.metadata.yaml", "readme": "models/khm-latn/khm-latn-1.0.README.md", + "test_volume": "/datasets/secryst", + "test_data": "khmer-translit/test.jsonl", "probe": "ភាសា", }, "urd-g2p": { "volume": "/volumes/urdu-g2p-checkpoints", - "checkpoint": "/urdu_g2p/run-001/best", + "checkpoint": "urdu_g2p/run-001/best", "metadata": "models/urd-g2p/urd-g2p-1.0.metadata.yaml", "readme": "models/urd-g2p/urd-g2p-1.0.README.md", + "test_volume": "/datasets/urdu-g2p", + "test_data": "urdu-g2p/test.jsonl", "probe": "اردو", }, + "heb-diac": { + "volume": "/volumes/rababa-checkpoints", + "checkpoint": "rababa_hebrew_byt5_s43/run-001/best", + "metadata": "models/heb-diac/heb-diac-1.0.metadata.yaml", + "readme": "models/heb-diac/heb-diac-1.0.README.md", + "test_volume": "/datasets/rababa", + "test_data": "nakdimon/test.txt", + "probe": "שלום", + }, "urd-diac": { "volume": "/volumes/urdu-diacrit-checkpoints", - "checkpoint": "/urdu_diacrit/run-001/best", + "checkpoint": "urdu_diacrit/run-001/best", "metadata": "models/urd-diac/urd-diac-1.0.metadata.yaml", "readme": "models/urd-diac/urd-diac-1.0.README.md", + "test_volume": "/datasets/urdu-diacrit", + "test_data": "urdu-diacrit/test.jsonl", "probe": "اردو", }, } @@ -74,11 +98,31 @@ app = modal.App("interscript-ml-export", image=IMAGE) +def _load_pairs(path: Path) -> list[tuple[str, str]]: + import json + + pairs: list[tuple[str, str]] = [] + for line in path.read_text(encoding="utf-8").splitlines(): + if not line.strip(): + continue + row = json.loads(line) + if isinstance(row, dict): + pairs.append( + ( + row.get("input", row.get("src", "")), + row.get("target", row.get("tgt", row.get("gold", ""))), + ) + ) + else: + pairs.append((row[0], row[1] if len(row) > 1 else "")) + return pairs + + @app.function( cpu=8, memory=32 * 1024, timeout=2 * 3600, - volumes={**CHECKPOINT_VOLUMES, "/outputs": MODELS_VOLUME}, + volumes={**CHECKPOINT_VOLUMES, **DATASET_VOLUMES, "/outputs": MODELS_VOLUME}, ) def export_model(model_id: str, precisions: list[str]) -> dict[str, str]: import sys @@ -104,11 +148,11 @@ def export_model(model_id: str, precisions: list[str]) -> dict[str, str]: ) MODELS_VOLUME.commit() - report: dict[str, str] = {} import zipfile import onnxruntime as ort + report: dict[str, str] = {} for z in zips: result = validate_zip(z) if not result.ok: @@ -127,8 +171,57 @@ def export_model(model_id: str, precisions: list[str]) -> dict[str, str]: return report +@app.function( + cpu=8, + memory=32 * 1024, + timeout=2 * 3600, + volumes={**CHECKPOINT_VOLUMES, **DATASET_VOLUMES, "/outputs": MODELS_VOLUME}, +) +def parity_model(model_id: str, precisions: list[str], limit: int = 0) -> dict[str, str]: + """WO03 gate on Modal: torch reference vs ONNX decode over the test + split; writes the parity block into each zip (strict gate enforced).""" + import sys + + sys.path.insert(0, "/root/ml-models/src") + + spec = MODELS[model_id] + checkpoint = Path(spec["volume"]) / spec["checkpoint"] + test_path = Path(spec["test_volume"]) / spec["test_data"] + + from imf.export import load_byte_seq2seq + from imf.parity import run_parity, write_parity + + model = load_byte_seq2seq(checkpoint) + pairs = _load_pairs(test_path) + if limit: + pairs = pairs[:limit] + + out_dir = Path("/outputs/imf") / model_id + reports: dict[str, str] = {} + for precision in precisions: + zip_path = out_dir / f"{model_id}-1.0-{precision}.zip" + report = run_parity(model, zip_path, pairs, max_len=128) + reports[precision] = ( + f"samples={report.samples} cer_ref={report.cer_reference}pp " + f"cer_onnx={report.cer_onnx}pp delta={report.cer_delta}pp " + f"mismatches={report.token_mismatches}" + ) + if not report.passed: + raise RuntimeError(f"parity gate FAILED for {zip_path.name}") + write_parity(zip_path, report) + MODELS_VOLUME.commit() + return reports + + @app.local_entrypoint() def main(model: str, precisions: str = "fp32,fp16,int8") -> None: report = export_model.remote(model, precisions.split(",")) for name, status in report.items(): print(f"{name}: {status}") + + +@app.local_entrypoint() +def parity(model: str, precisions: str = "fp32,fp16,int8", limit: int = 0) -> None: + reports = parity_model.remote(model, precisions.split(","), limit) + for precision, status in reports.items(): + print(f"{model} [{precision}] {status}") From e4ee1f7d375a752d2f5babd0a6cdf8c8824a44ba Mon Sep 17 00:00:00 2001 From: Ronald Tse Date: Sun, 16 Aug 2026 13:10:29 +0800 Subject: [PATCH 4/9] fix(imf): heb parity reads converted nakdimon pairs (test-imf.jsonl) 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). --- src/gpu/modal_export.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/gpu/modal_export.py b/src/gpu/modal_export.py index c5b717d..9a1b617 100644 --- a/src/gpu/modal_export.py +++ b/src/gpu/modal_export.py @@ -81,7 +81,7 @@ "metadata": "models/heb-diac/heb-diac-1.0.metadata.yaml", "readme": "models/heb-diac/heb-diac-1.0.README.md", "test_volume": "/datasets/rababa", - "test_data": "nakdimon/test.txt", + "test_data": "nakdimon/test-imf.jsonl", "probe": "שלום", }, "urd-diac": { From 1723f26981bd794ea5b12e5e09b00a914c73fbbb Mon Sep 17 00:00:00 2001 From: Ronald Tse Date: Sun, 16 Aug 2026 13:31:49 +0800 Subject: [PATCH 5/9] =?UTF-8?q?fix(imf):=20canonical=20ByT5=20byte=20table?= =?UTF-8?q?=20=E2=80=94=20ids=20are=20byte+3,=20inputs=20carry=20EOS?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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 --- golden/khm-latn-100.jsonl | 100 ++++++++++++++++++++++++++++++++++++++ src/imf/cli.py | 7 ++- src/imf/export.py | 32 ++++++++---- src/imf/parity.py | 14 +++--- tests/test_imf_export.py | 21 +------- 5 files changed, 139 insertions(+), 35 deletions(-) create mode 100644 golden/khm-latn-100.jsonl diff --git a/golden/khm-latn-100.jsonl b/golden/khm-latn-100.jsonl new file mode 100644 index 0000000..b56717a --- /dev/null +++ b/golden/khm-latn-100.jsonl @@ -0,0 +1,100 @@ +{"input": "បង្វេច", "tokens": [101, 100, 113, 106, 121, 104, 102, 107], "output": "bangvech"} +{"input": "រោក", "tokens": [117, 114, 110], "output": "rok"} +{"input": "សង់ផ្ទះ", "tokens": [118, 100, 113, 106, 115, 107, 119, 104, 100, 107], "output": "sangphteah"} +{"input": "អប្បបរិមាណ", "tokens": [100, 101, 101, 100, 101, 117, 108, 112, 100, 113], "output": "abbabriman"} +{"input": "ឈឺចាប់", "tokens": [102, 107, 107, 120, 102, 107, 104, 100, 101], "output": "chhucheab"} +{"input": "អសមត្ថភាព", "tokens": [100, 118, 100, 112, 114, 119, 119, 107, 100, 115, 107, 100, 115], "output": "asamotthaphap"} +{"input": "និយាយ", "tokens": [113, 108, 124, 104, 100, 124], "output": "niyeay"} +{"input": "សម្ងំ", "tokens": [118, 100, 112, 113, 106, 114, 112], "output": "samngom"} +{"input": "អូម", "tokens": [100, 120, 112], "output": "aum"} +{"input": "ស្វាយរៀង", "tokens": [118, 121, 100, 124, 117, 108, 104, 113, 106], "output": "svayrieng"} +{"input": "បារមី", "tokens": [101, 100, 117, 112, 108], "output": "barmi"} +{"input": "រនល", "tokens": [117, 114, 113, 111], "output": "ronl"} +{"input": "ធម្មយុត្តិ", "tokens": [119, 107, 114, 112, 112, 114, 124, 114, 120, 119, 119, 104], "output": "thommoyoutte"} +{"input": "ទឹកឃ្មុំ", "tokens": [119, 120, 110, 100, 110, 107, 112, 114, 120, 112], "output": "tukakhmoum"} +{"input": "ភ្លើង", "tokens": [115, 107, 111, 104, 120, 113, 106], "output": "phleung"} +{"input": "ប្រួត", "tokens": [101, 117, 120, 114, 119], "output": "bruot"} +{"input": "វោហារ", "tokens": [121, 114, 107, 104, 100, 117], "output": "vohear"} +{"input": "ច្បាប់", "tokens": [102, 107, 101, 100, 101], "output": "chbab"} +{"input": "ទន្ទឹម", "tokens": [119, 114, 113, 119, 120, 112], "output": "tontum"} +{"input": "បញ្ចស័ក", "tokens": [101, 100, 113, 107, 102, 107, 114, 118, 100, 110], "output": "banhchosak"} +{"input": "ល្បី", "tokens": [111, 101, 108], "output": "lbi"} +{"input": "គុណូបការ", "tokens": [110, 114, 120, 113, 100, 120, 101, 110, 100, 117], "output": "kounaubkar"} +{"input": "បញ្ចុក", "tokens": [101, 100, 113, 107, 102, 107, 114, 120, 110], "output": "banhchouk"} +{"input": "ជ្រាបស្រាប់", "tokens": [102, 107, 117, 104, 100, 101, 100, 118, 117, 100, 101], "output": "chreabasrab"} +{"input": "សៃយ៉ឺន", "tokens": [118, 100, 108, 124, 104, 120, 113], "output": "saiyeun"} +{"input": "លោះខ្ញុំ", "tokens": [111, 114, 107, 110, 107, 113, 107, 114, 112], "output": "lohkhnhom"} +{"input": "ទូទឹម", "tokens": [119, 114, 120, 119, 120, 112], "output": "toutum"} +{"input": "ប្រវត្តិសាស្ត្រ", "tokens": [101, 117, 100, 121, 114, 119, 119, 104, 35, 118, 100, 118, 119, 117], "output": "bravotte sastr"} +{"input": "ខ្លោងទ្វារ", "tokens": [110, 107, 111, 100, 114, 113, 106, 114, 119, 121, 104, 100, 117], "output": "khlaongotvear"} +{"input": "គេជាច្រើន", "tokens": [110, 104, 102, 107, 100, 102, 107, 117, 100, 104, 113], "output": "kechachraen"} +{"input": "បង្កួយ", "tokens": [101, 100, 113, 106, 110, 104, 100, 119], "output": "bangkeat"} +{"input": "ហៃអើ", "tokens": [107, 100, 108, 100, 104], "output": "haiae"} +{"input": "ឆ្កឹះឆ្កៀល", "tokens": [102, 107, 107, 110, 114, 104, 100, 107, 102, 107, 107, 110, 108, 104, 111], "output": "chhkoeahchhkiel"} +{"input": "ទំលាក់", "tokens": [119, 114, 112, 35, 111, 104, 100, 110], "output": "tom leak"} +{"input": "ហ្រ្វង់", "tokens": [107, 117, 121, 114, 113, 106], "output": "hrvong"} +{"input": "រាស់", "tokens": [117, 104, 100, 118, 100], "output": "reasa"} +{"input": "លោភលន់", "tokens": [111, 114, 115, 107, 114, 111, 113], "output": "lopholn"} +{"input": "អរគុណណាស់", "tokens": [100, 117, 110, 114, 120, 113, 113, 100, 118, 100], "output": "arkounnasa"} +{"input": "ដោយហេតុ", "tokens": [103, 100, 114, 124, 107, 104, 119, 114], "output": "daoyheto"} +{"input": "អ្នកចំរៀង", "tokens": [100, 113, 100, 110, 102, 107, 100, 112, 117, 108, 104, 113, 106], "output": "anakchamrieng"} +{"input": "ពោធិសម្ភារ", "tokens": [115, 114, 119, 107, 108, 118, 100, 112, 115, 107, 104, 100, 117], "output": "pothisamphear"} +{"input": "ចូលហ៊ុន", "tokens": [102, 107, 100, 120, 111, 107, 114, 120, 113], "output": "chaulhoun"} +{"input": "បបរ", "tokens": [101, 100, 101, 117], "output": "babr"} +{"input": "រូបលោក", "tokens": [117, 114, 120, 101, 111, 114, 110], "output": "roublok"} +{"input": "ហ៊ុមព័ទ្ធ", "tokens": [107, 114, 120, 112, 115, 119, 119, 107], "output": "houmptth"} +{"input": "ការហាត់ប្រាណ", "tokens": [110, 100, 117, 35, 107, 100, 119, 101, 117, 100, 113], "output": "kar hatbran"} +{"input": "មិនបានជា", "tokens": [112, 108, 113, 101, 100, 113, 102, 107, 104, 100], "output": "minbanchea"} +{"input": "កម្មវិបាក", "tokens": [110, 100, 112, 112, 114, 121, 108, 101, 104, 100, 110], "output": "kammovibeak"} +{"input": "កុំអាល", "tokens": [110, 114, 112, 100, 111], "output": "komal"} +{"input": "សុខយាន", "tokens": [118, 114, 110, 107, 124, 104, 100, 113], "output": "sokhyean"} +{"input": "ទឹកស្លាប់", "tokens": [119, 120, 110, 100, 118, 111, 100, 101], "output": "tukaslab"} +{"input": "អធិការបតី", "tokens": [100, 119, 107, 108, 110, 104, 100, 117, 114, 101, 119, 104, 108], "output": "athikearobtei"} +{"input": "ខ្តត", "tokens": [110, 107, 119, 100, 119], "output": "khtat"} +{"input": "ស្ទេញ", "tokens": [118, 119, 104, 113, 107], "output": "stenh"} +{"input": "ដល់ហើយ", "tokens": [103, 111, 107, 100, 104, 124], "output": "dlhaey"} +{"input": "ក្ដៀប", "tokens": [110, 103, 108, 104, 101], "output": "kdieb"} +{"input": "ភ្លៀងមួយមេ", "tokens": [115, 107, 111, 108, 104, 113, 106, 112, 120, 114, 124, 112, 104], "output": "phliengmuoyme"} +{"input": "សៃយ", "tokens": [118, 100, 108, 124], "output": "saiy"} +{"input": "រាជបល្ល័ង្ក", "tokens": [117, 104, 100, 102, 107, 114, 101, 100, 111, 111, 113, 106, 110], "output": "reachoballngk"} +{"input": "ល្អក់", "tokens": [111, 100, 114, 110], "output": "laok"} +{"input": "ឈ្ងុយឆ្ងាញ់", "tokens": [102, 107, 107, 113, 106, 114, 120, 124, 114, 102, 107, 107, 113, 106, 100, 113, 107], "output": "chhngouyochhnganh"} +{"input": "រ៉ឹង", "tokens": [117, 114, 104, 113, 106], "output": "roeng"} +{"input": "ប្រុងនឹង", "tokens": [101, 117, 114, 113, 106, 113, 120, 113, 106], "output": "brongnung"} +{"input": "ការប្រើប្រាស់", "tokens": [110, 100, 117, 114, 101, 117, 100, 104, 101, 117, 100, 118, 100], "output": "karobraebrasa"} +{"input": "គ្រោះ", "tokens": [110, 117, 114, 107], "output": "kroh"} +{"input": "កលកិច្ច", "tokens": [110, 100, 111, 110, 104, 102, 107, 102, 107], "output": "kalkechch"} +{"input": "ផ្សារផ្សោ", "tokens": [115, 107, 118, 100, 100, 117, 114, 115, 107, 118, 100, 100, 114], "output": "phsaarophsaao"} +{"input": "ស្រករ", "tokens": [118, 117, 100, 110, 117], "output": "srakr"} +{"input": "រត់លឿន", "tokens": [117, 119, 111, 120, 104, 113], "output": "rtluen"} +{"input": "និត្យ", "tokens": [113, 108, 119, 124], "output": "nity"} +{"input": "គ្រប់រូប", "tokens": [110, 117, 114, 101, 117, 114, 120, 101], "output": "krobroub"} +{"input": "ឲ្យទឹក", "tokens": [100, 114, 124, 119, 120, 110], "output": "aoytuk"} +{"input": "ចាប់ត្រី", "tokens": [102, 107, 100, 101, 119, 117, 104, 108], "output": "chabtrei"} +{"input": "ពេជ្រ", "tokens": [115, 104, 102, 107, 117], "output": "pechr"} +{"input": "អគ្គស្នងការ", "tokens": [100, 110, 110, 114, 118, 113, 100, 113, 106, 110, 100, 117], "output": "akkosnangkar"} +{"input": "ដួល", "tokens": [103, 120, 114, 111], "output": "duol"} +{"input": "នៅនឹង", "tokens": [113, 114, 122, 113, 120, 113, 106], "output": "nownung"} +{"input": "ផលកម្ម", "tokens": [115, 107, 100, 111, 114, 110, 100, 112, 112], "output": "phalokamm"} +{"input": "ញាំញី", "tokens": [113, 107, 104, 100, 113, 107, 108], "output": "nheanhi"} +{"input": "ភាតរ", "tokens": [115, 107, 104, 100, 119, 117], "output": "pheatr"} +{"input": "មួយនេះ", "tokens": [112, 120, 114, 124, 113, 108, 107], "output": "muoynih"} +{"input": "ផ្លូវសួន", "tokens": [115, 107, 111, 100, 120, 121, 118, 120, 114, 113], "output": "phlauvsuon"} +{"input": "ឆ្នាំច", "tokens": [102, 107, 107, 113, 100, 102, 107], "output": "chhnach"} +{"input": "ត្នោត", "tokens": [119, 113, 100, 114, 119], "output": "tnaot"} +{"input": "វាង", "tokens": [121, 104, 100, 113, 106], "output": "veang"} +{"input": "រាជបុត្រ", "tokens": [117, 104, 100, 102, 107, 101, 114, 119, 117], "output": "reachbotr"} +{"input": "ស្រមូម", "tokens": [118, 117, 100, 112, 114, 120, 112], "output": "sramoum"} +{"input": "ចិត្តស៊ូ", "tokens": [102, 107, 104, 119, 119, 100, 118, 100, 114, 120], "output": "chettasaou"} +{"input": "អក្សរខម", "tokens": [100, 110, 118, 100, 100, 117, 114, 110, 107, 112], "output": "aksaarokhm"} +{"input": "ទូទឹកកក", "tokens": [119, 114, 120, 119, 120, 110, 100, 110, 110], "output": "toutukakk"} +{"input": "កាត់ឲ្យខ្លី", "tokens": [110, 100, 119, 35, 100, 114, 124, 35, 110, 107, 111, 104, 108], "output": "kat aoy khlei"} +{"input": "យ៉ាងណាក៏ដោយ", "tokens": [124, 104, 100, 113, 106, 113, 100, 35, 110, 100, 103, 100, 114, 124], "output": "yeangna kadaoy"} +{"input": "បំផុត", "tokens": [101, 100, 112, 115, 107, 114, 119], "output": "bamphot"} +{"input": "រកអ្វីប្រៀបពុំបាន", "tokens": [117, 114, 110, 100, 100, 121, 104, 108, 35, 101, 117, 108, 104, 101, 35, 101, 100, 113], "output": "rokaavei brieb ban"} +{"input": "មកយឺត", "tokens": [112, 114, 110, 124, 120, 119], "output": "mokyut"} +{"input": "អប្សរា", "tokens": [100, 101, 118, 100, 100, 117, 100], "output": "absaara"} +{"input": "បិតុច្ឆា", "tokens": [101, 104, 119, 114, 102, 107, 102, 107, 107, 100], "output": "betochchha"} +{"input": "អស់អញ", "tokens": [100, 118, 100, 100, 113, 107], "output": "asaanh"} +{"input": "ចាង", "tokens": [102, 107, 100, 113, 106], "output": "chang"} +{"input": "មិនអស់ចិត្ត", "tokens": [112, 108, 113, 114, 100, 118, 100, 102, 107, 104, 119, 119], "output": "minoasachett"} diff --git a/src/imf/cli.py b/src/imf/cli.py index cadb089..85ae220 100644 --- a/src/imf/cli.py +++ b/src/imf/cli.py @@ -133,7 +133,12 @@ def _cmd_golden(args: argparse.Namespace) -> int: if not line.strip(): continue row = json.loads(line) - inputs.append(row["input"] if isinstance(row, dict) else row[0]) + if isinstance(row, dict): + inputs.append(row.get("input", row.get("src", row.get("text", "")))) + elif isinstance(row, str): + inputs.append(row) + else: + inputs.append(row[0]) out = write_golden(args.zip, inputs, args.out, max_len=args.max_len) print(f"wrote {len(inputs)} golden cases to {out}") return 0 diff --git a/src/imf/export.py b/src/imf/export.py index 86bdca9..d8e06a4 100644 --- a/src/imf/export.py +++ b/src/imf/export.py @@ -28,6 +28,20 @@ OPSET = 14 GRAPH_NAMES = ("encoder.onnx", "decoder.onnx", "decoder-kv.onnx") +# Canonical ByT5 byte table (google/byt5): pad=0, eos=1, unk=2, and +# every UTF-8 byte b is token id b + 3, up to 258. The vocab is 384-wide; +# ids > 258 are unused in practice. "tokenizer: bytes" in IMF metadata +# means THIS fixed table — no vocab files, but ids are NOT raw byte values +# (feeding text.bytes directly silently produces garbage). +BYTE_OFFSET = 3 +PAD_ID = 0 +EOS_ID = 1 + + +def encode_bytes(text: str) -> list[int]: + """Canonical byte-level tokenization: byte ids + trailing EOS.""" + return [b + BYTE_OFFSET for b in text.encode("utf-8")] + [EOS_ID] + def load_byte_seq2seq(checkpoint_dir: Path | str): """Load a T5-family checkpoint in eager attention (export-safe).""" @@ -89,7 +103,7 @@ def _decoder_kv(model): import torch.nn as nn from transformers.cache_utils import DynamicCache, EncoderDecoderCache - num_layers = model.config.num_layers + num_layers = model.config.num_decoder_layers or model.config.num_layers class DecoderKV(nn.Module): def __init__(self, model): @@ -199,7 +213,7 @@ def export_graphs(model, out_dir: Path | str) -> dict[str, Path]: ) paths["decoder.onnx"] = out_dir / "decoder.onnx" - num_layers = model.config.num_layers + num_layers = model.config.num_decoder_layers or model.config.num_layers pasts = _sample_pasts(model, hidden) kv_inputs = ["input_ids", "encoder_hidden_states"] + _kv_io_names(num_layers)[0] kv_outputs = _kv_io_names(num_layers)[1] @@ -277,18 +291,18 @@ def onnx_greedy_plain(encoder_sess, decoder_sess, text: str, max_len: int = 256) byte-level model reliably stays < 256).""" import numpy as np - ids = np.array([list(text.encode("utf-8"))], dtype=np.int64) - if ids.shape[1] == 0: + ids = np.array([encode_bytes(text)], dtype=np.int64) + if ids.shape[1] == 1: return [] hidden = encoder_sess.run(None, {"input_ids": ids})[0] - dec_ids = np.array([[0]], dtype=np.int64) + dec_ids = np.array([[PAD_ID]], dtype=np.int64) generated: list[int] = [] for _ in range(max_len): logits = decoder_sess.run( None, {"input_ids": dec_ids, "encoder_hidden_states": hidden} )[0] nxt = int(np.argmax(logits[0, -1])) - if nxt == 1: + if nxt == EOS_ID: break generated.append(nxt) dec_ids = np.concatenate([dec_ids, np.array([[nxt]], dtype=np.int64)], axis=1) @@ -314,8 +328,8 @@ def onnx_greedy_kv(encoder_sess, kv_sess, text: str, max_len: int = 256) -> list """Greedy decode over ONNX sessions (KV decoder). Self-check helper.""" import numpy as np - ids = np.array([list(text.encode("utf-8"))], dtype=np.int64) - if ids.shape[1] == 0: + ids = np.array([encode_bytes(text)], dtype=np.int64) + if ids.shape[1] == 1: return [] hidden = encoder_sess.run(None, {"input_ids": ids})[0] out_names = [o.name for o in kv_sess.get_outputs()] @@ -326,7 +340,7 @@ def onnx_greedy_kv(encoder_sess, kv_sess, text: str, max_len: int = 256) -> list out = kv_sess.run(None, {"input_ids": cur, "encoder_hidden_states": hidden, **pasts}) results = dict(zip(out_names, out, strict=True)) nxt = int(np.argmax(results["logits"][0, -1])) - if nxt == 1: + if nxt == EOS_ID: break generated.append(nxt) pasts = { diff --git a/src/imf/parity.py b/src/imf/parity.py index 477365e..d962300 100644 --- a/src/imf/parity.py +++ b/src/imf/parity.py @@ -19,7 +19,7 @@ from pathlib import Path from framework.evaluator import char_error_rate -from imf.export import onnx_greedy_kv +from imf.export import BYTE_OFFSET, EOS_ID, PAD_ID, encode_bytes, onnx_greedy_kv from imf.schema import ModelMetadata, Parity @@ -42,11 +42,11 @@ def passed(self) -> bool: def _torch_greedy_tokens(model, text: str, max_len: int) -> list[int]: import torch - ids = torch.tensor([list(text.encode("utf-8"))], dtype=torch.long) - if ids.shape[1] == 0: + ids = torch.tensor([encode_bytes(text)], dtype=torch.long) + if ids.shape[1] == 1: return [] enc = model.get_encoder()(input_ids=ids)[0] - dec_ids = torch.tensor([[0]], dtype=torch.long) + dec_ids = torch.tensor([[PAD_ID]], dtype=torch.long) outs: list[int] = [] for _ in range(max_len): hidden = model.get_decoder()( @@ -54,7 +54,7 @@ def _torch_greedy_tokens(model, text: str, max_len: int) -> list[int]: )[0] logits = model.lm_head(hidden * (model.config.d_model ** -0.5)) nxt = int(logits[0, -1].argmax()) - if nxt == 1: + if nxt == EOS_ID: break outs.append(nxt) dec_ids = torch.cat([dec_ids, torch.tensor([[nxt]], dtype=torch.long)], 1) @@ -62,7 +62,9 @@ def _torch_greedy_tokens(model, text: str, max_len: int) -> list[int]: def _decode_tokens(tokens: list[int]) -> str: - return bytes(t % 256 for t in tokens).decode("utf-8", errors="replace") + return bytes((t - BYTE_OFFSET) % 256 for t in tokens).decode( + "utf-8", errors="replace" + ) def _sessions_from_zip(zip_path: Path): diff --git a/tests/test_imf_export.py b/tests/test_imf_export.py index 269e742..24d5ca8 100644 --- a/tests/test_imf_export.py +++ b/tests/test_imf_export.py @@ -26,6 +26,7 @@ onnx_greedy_kv, onnx_greedy_plain, ) +from imf.parity import _torch_greedy_tokens as torch_greedy # noqa: E402 from imf.validator import validate_zip # noqa: E402 TEXTS = ["he", "hello", "abc"] @@ -55,24 +56,6 @@ } -def _torch_greedy(model, text: str) -> list[int]: - ids = torch.tensor([list(text.encode("utf-8"))]) - enc = model.get_encoder()(input_ids=ids)[0] - dec_ids = torch.tensor([[0]]) - outs: list[int] = [] - for _ in range(MAX_LEN): - logits = model.lm_head( - model.get_decoder()(input_ids=dec_ids, encoder_hidden_states=enc)[0] - * (model.config.d_model ** -0.5) - ) - nxt = int(logits[0, -1].argmax()) - if nxt == 1: - break - outs.append(nxt) - dec_ids = torch.cat([dec_ids, torch.tensor([[nxt]])], 1) - return outs - - @pytest.fixture(scope="module") def reference_model(tmp_path_factory: pytest.TempPathFactory): ckpt = make_fixture_checkpoint(tmp_path_factory.mktemp("fixture") / "checkpoint") @@ -115,7 +98,7 @@ def test_fixture_exports_match_torch(reference_model) -> None: dec = ort.InferenceSession(str(graphs["decoder.onnx"]), providers=PROVIDERS) kv = ort.InferenceSession(str(graphs["decoder-kv.onnx"]), providers=PROVIDERS) for text in TEXTS: - expected = _torch_greedy(reference_model, text) + expected = torch_greedy(reference_model, text, MAX_LEN) assert onnx_greedy_plain(enc, dec, text, MAX_LEN) == expected assert onnx_greedy_kv(enc, kv, text, MAX_LEN) == expected From f27a70e01777ee06404815d8b368e0d85da4a939 Mon Sep 17 00:00:00 2001 From: Ronald Tse Date: Sun, 16 Aug 2026 14:00:33 +0800 Subject: [PATCH 6/9] fix(imf): write_parity temp file in the zip's own directory MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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. --- src/imf/parity.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/src/imf/parity.py b/src/imf/parity.py index d962300..504bb7c 100644 --- a/src/imf/parity.py +++ b/src/imf/parity.py @@ -156,7 +156,9 @@ def write_parity(zip_path: Path | str, report: ParityReport) -> Path: from imf.pack import _to_dict - with tempfile.TemporaryDirectory() as tmp: + # Same filesystem as the target: os.replace is atomic within one + # filesystem and fails with EXDEV across a volume mount. + with tempfile.TemporaryDirectory(dir=zip_path.parent) as tmp: rewritten = Path(tmp) / "rewritten.zip" with zipfile.ZipFile(zip_path) as src, zipfile.ZipFile( rewritten, "w", zipfile.ZIP_DEFLATED From 8740ac0e1d68b921be82fbefad261a73561ee268 Mon Sep 17 00:00:00 2001 From: Ronald Tse Date: Sun, 16 Aug 2026 14:11:58 +0800 Subject: [PATCH 7/9] perf(imf): decode the torch reference once per model, not per zip The reference is precision-independent; multi-zip gates were paying 3x the torch decode (hours on the 12k-pair Urdu splits). --- src/gpu/modal_export.py | 6 ++++-- src/imf/parity.py | 22 ++++++++++++++++++---- 2 files changed, 22 insertions(+), 6 deletions(-) diff --git a/src/gpu/modal_export.py b/src/gpu/modal_export.py index 9a1b617..5daf206 100644 --- a/src/gpu/modal_export.py +++ b/src/gpu/modal_export.py @@ -189,18 +189,20 @@ def parity_model(model_id: str, precisions: list[str], limit: int = 0) -> dict[s test_path = Path(spec["test_volume"]) / spec["test_data"] from imf.export import load_byte_seq2seq - from imf.parity import run_parity, write_parity + from imf.parity import reference_decode, run_parity, write_parity model = load_byte_seq2seq(checkpoint) pairs = _load_pairs(test_path) if limit: pairs = pairs[:limit] + reference = reference_decode(model, [src for src, _ in pairs], max_len=128) + out_dir = Path("/outputs/imf") / model_id reports: dict[str, str] = {} for precision in precisions: zip_path = out_dir / f"{model_id}-1.0-{precision}.zip" - report = run_parity(model, zip_path, pairs, max_len=128) + report = run_parity(model, zip_path, pairs, max_len=128, reference=reference) reports[precision] = ( f"samples={report.samples} cer_ref={report.cer_reference}pp " f"cer_onnx={report.cer_onnx}pp delta={report.cer_delta}pp " diff --git a/src/imf/parity.py b/src/imf/parity.py index 504bb7c..b8f2fa2 100644 --- a/src/imf/parity.py +++ b/src/imf/parity.py @@ -86,9 +86,21 @@ def _sessions_from_zip(zip_path: Path): return enc, dec -def run_parity(model, zip_path: Path | str, pairs, max_len: int = 256) -> ParityReport: +def reference_decode(model, sources, max_len: int = 256) -> list[list[int]]: + """Torch-reference greedy decode of many inputs, computed once and + shared across precision variants by run_parity.""" + return [_torch_greedy_tokens(model, source, max_len) for source in sources] + + +def run_parity( + model, zip_path: Path | str, pairs, max_len: int = 256, reference=None +) -> ParityReport: """pairs: iterable of (source_text, gold_target). Measures both sides - against gold; the gate is the CER distance between the two.""" + against gold; the gate is the CER distance between the two. + + ``reference`` (from ``reference_decode``) skips the torch side — the + reference is precision-independent, so multi-zip gates decode it once. + """ zip_path = Path(zip_path) enc, kv = _sessions_from_zip(zip_path) @@ -96,9 +108,11 @@ def run_parity(model, zip_path: Path | str, pairs, max_len: int = 256) -> Parity mismatches = 0 cer_ref_sum = 0.0 cer_onnx_sum = 0.0 - for source, gold in pairs: + for i, (source, gold) in enumerate(pairs): n += 1 - ref = _torch_greedy_tokens(model, source, max_len) + ref = reference[i] if reference is not None else _torch_greedy_tokens( + model, source, max_len + ) got = onnx_greedy_kv(enc, kv, source, max_len) if ref != got: mismatches += 1 From 1fb9f22f1276af2ffe3fd36ac8cf0ed7f35ba152 Mon Sep 17 00:00:00 2001 From: Ronald Tse Date: Sun, 16 Aug 2026 15:08:11 +0800 Subject: [PATCH 8/9] =?UTF-8?q?fix(imf):=20fp16=20via=20torch-native=20hal?= =?UTF-8?q?f=20export=20=E2=80=94=20ORT=20converter=20broken=20on=20real?= =?UTF-8?q?=20ByT5?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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. --- docs/imf-v1.md | 2 +- src/imf/export.py | 48 ++++++++++++++++++++++------------------------- 2 files changed, 23 insertions(+), 27 deletions(-) diff --git a/docs/imf-v1.md b/docs/imf-v1.md index afc2210..282bde9 100644 --- a/docs/imf-v1.md +++ b/docs/imf-v1.md @@ -47,7 +47,7 @@ model.zip | `tokenizer` | enum | `bytes` (the only v1 value) | | `opset` | int | 7..14; must equal the graphs' opset | | `decoder` | enum | `plain` \| `kv` (`kv` requires decoder-kv.onnx) | -| `precision` | enum | `fp32` \| `fp16` \| `int8` | +| `precision` | enum | `fp32` \| `fp16` \| `int8` (fp16 = torch-native half export: float16 graph IO, int64 ids unchanged; runtimes read dtypes from the session — the ORT float16 converter produces all-zero hiddens on real ByT5 and must not be used) | | `license` | str | non-empty (strict gate) | | `trained_from` | str | repo + run/checkpoint id | | `metrics` | list | `{name, value, protocol, source}`; `source` must be a `RESULTS.md#anchor` (strict gate) | diff --git a/src/imf/export.py b/src/imf/export.py index d8e06a4..35ad91d 100644 --- a/src/imf/export.py +++ b/src/imf/export.py @@ -242,27 +242,18 @@ def export_graphs(model, out_dir: Path | str) -> dict[str, Path]: return paths -def convert_fp16(src: Path | str, dst: Path | str) -> Path: - """fp32 -> mixed fp16, IO types preserved (encoder/decoder compose cleanly). - - LayerNorm/softmax math stays fp32: ORT's session-time - SimplifiedLayerNormFusion crashes on half-converted LN subgraphs - (InsertPrecisionFreeCast name mismatch), so the whole decomposition - must stay one dtype. Weights (MatMuls) carry the size win. +def convert_fp16(model): + """A fp16 copy of the model for torch-native half export. + + The onnxruntime float16 CONVERTER is not usable here: on real ByT5 + checkpoints it produces all-zero encoder hiddens (found 2026-08-16, + khm-latn — 1939pp CER); exporting the torch model under .half() is + exact on gold pairs. Graph IO becomes float16 (input_ids stay int64); + runtimes read dtypes from the session, and _zero_pasts follows them. """ - import onnx - from onnxruntime.transformers import float16 - - block_list = list(float16.DEFAULT_OP_BLOCK_LIST) + [ - "ReduceMean", "Pow", "Sqrt", "Div", "Sub", "Add", "Mul", - "Softmax", "Range", "Exp", "Where", "Less", "Cast", - ] - model = onnx.load(str(src)) - converted = float16.convert_float_to_float16( - model, keep_io_types=True, op_block_list=block_list - ) - onnx.save(converted, str(dst)) - return Path(dst) + import copy + + return copy.deepcopy(model).half() def quantize_int8(src: Path | str, dst: Path | str) -> Path: @@ -320,7 +311,8 @@ def _zero_pasts(kv_sess) -> dict[str, object]: shape = meta.shape # [batch, heads, past_seq, d_kv] with str dynamic dims heads = shape[1] if isinstance(shape[1], int) else 4 d_kv = shape[3] if isinstance(shape[3], int) else 8 - pasts[meta.name] = np.zeros((1, heads, 0, d_kv), dtype=np.float32) + dtype = np.float16 if meta.type == "tensor(float16)" else np.float32 + pasts[meta.name] = np.zeros((1, heads, 0, d_kv), dtype=dtype) return pasts @@ -372,18 +364,22 @@ def export_zips( with tempfile.TemporaryDirectory() as tmp: tmp = Path(tmp) graphs = export_graphs(model, tmp / "graphs") + graphs_16 = ( + export_graphs(convert_fp16(model), tmp / "graphs-fp16") + if "fp16" in precisions + else {} + ) for precision in precisions: variant_dir = tmp / precision variant_dir.mkdir() - for name, src in graphs.items(): + sources = graphs if precision != "fp16" else graphs_16 + for name, src in sources.items(): dst = variant_dir / name - if precision == "fp32": + if precision == "fp32" or precision == "fp16": dst.write_bytes(src.read_bytes()) - elif precision == "fp16": - convert_fp16(src, dst) elif precision == "int8": - quantize_int8(src, dst) + quantize_int8(graphs[name], dst) else: raise ValueError(f"unknown precision {precision!r}") meta = replace(metadata, precision=precision) From ae5649d3a3ed46d6772aec95cec61d670f97dc13 Mon Sep 17 00:00:00 2001 From: Ronald Tse Date: Sun, 16 Aug 2026 17:01:29 +0800 Subject: [PATCH 9/9] ci: pin torch/transformers in export-fixture (transformers 5.15 tracing regression) --- .github/workflows/test.yml | 1 + 1 file changed, 1 insertion(+) diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml index b8b5155..d77250a 100644 --- a/.github/workflows/test.yml +++ b/.github/workflows/test.yml @@ -52,6 +52,7 @@ jobs: with: { python-version: "3.11" } - run: | python -m pip install --upgrade pip + pip install torch==2.12.1 transformers==5.14.1 pip install -e ".[dev,train,export]" - name: IMF export + parity gate fixture tests (torch vs ORT) run: PYTHONPATH=src python -m pytest tests/test_imf_export.py tests/test_imf_parity.py -v