Merge pull request #217 from interscript/feat/static-export-path

feat(export): static-int8 artifact path (calibration + builder)

Author: ronaldtse

src/gpu/modal_export.py

@@ -584,6 +584,84 @@ def rebuild_int8_head32(model_id: str, limit: int = 0) -> dict:
     }
 
 
+@app.function(
+    cpu=8,
+    memory=32 * 1024,
+    timeout=5 * 3600,
+    volumes={**CHECKPOINT_VOLUMES, **DATASET_VOLUMES, "/outputs": MODELS_VOLUME},
+)
+def export_int8_static(model_id: str, limit: int = 0) -> dict:
+    """Build the static-activation int8 artifact (TODO.impl/11's
+    positive branch): fp32 encoder + static-int8 decoder, calibrated on
+    the model's own eval pairs across BOTH decode framings.
+
+    The MEASURED composition (scored 4.6241 full-set vs dynamic 4.5701,
+    +8% CPU decode): the fp32 encoder keeps the artifact server-sized;
+    a browser-sized int8-encoder + static-decoder composition needs its
+    own gate run before it ships. Lands as {mid}-int8static.zip; the
+    release swap is a version decision."""
+    import re
+    import sys
+    import tempfile
+    import zipfile
+    from pathlib import Path
+
+    sys.path.insert(0, "/root/interscript-ml/src")
+    from imf.export import (
+        collect_decode_calibration,
+        head_matmul_names,
+        quantize_int8_static,
+        refresh_member_shas,
+    )
+
+    spec = MODELS[model_id]
+    test_path = Path(spec["test_volume"]) / spec["test_data"]
+    pairs = _load_pairs(test_path)[: limit or None]
+
+    out_dir = Path("/outputs/imf") / model_id
+    meta_path = Path("/root/interscript-ml", spec["metadata"])
+    mid = re.search(r"^id:\s*(\S+)", meta_path.read_text(encoding="utf-8"), re.M).group(1)
+    fp32_zip = out_dir / f"{mid}-fp32.zip"
+    if not fp32_zip.exists():
+        raise RuntimeError(f"{fp32_zip.name} missing on the volume")
+
+    import onnxruntime as ort
+
+    with tempfile.TemporaryDirectory() as tmp:
+        tmp = Path(tmp)
+        with zipfile.ZipFile(fp32_zip) as zf:
+            zf.extract("encoder.onnx", tmp)
+            dec = "decoder-kv.onnx" if "decoder-kv.onnx" in zf.namelist() else "decoder.onnx"
+            zf.extract(dec, tmp)
+        opts = ort.SessionOptions()
+        opts.graph_optimization_level = ort.GraphOptimizationLevel.ORT_DISABLE_ALL
+        enc_sess = ort.InferenceSession(str(tmp / "encoder.onnx"), opts)
+        dec_sess = ort.InferenceSession(str(tmp / dec), opts)
+        calibration = collect_decode_calibration(enc_sess, dec_sess, [s for s, _ in pairs])
+        print(f"[{model_id}] calibration: {len(calibration)} feeds", flush=True)
+        dec_static = tmp / dec.replace(".onnx", "-static.onnx")
+        quantize_int8_static(
+            tmp / dec, dec_static, calibration,
+            nodes_to_exclude=head_matmul_names(tmp / dec),
+        )
+        new_zip = out_dir / f"{mid}-int8static.zip"
+        with zipfile.ZipFile(fp32_zip) as src, zipfile.ZipFile(
+            new_zip, "w", zipfile.ZIP_DEFLATED
+        ) as dst:
+            for member in src.namelist():
+                if member == "metadata.yaml":
+                    meta = src.read(member).decode("utf-8")
+                    meta = re.sub(r"^precision:\s*\S+", "precision: int8-static", meta, flags=re.M)
+                    meta = re.sub(r"^id:\s*\S+", f"id: {mid}-int8static", meta, flags=re.M)
+                    dst.writestr(member, meta)
+                elif member == dec:
+                    dst.writestr(member, dec_static.read_bytes())
+                else:
+                    dst.writestr(member, src.read(member))
+        refresh_member_shas(new_zip)
+    return {"artifact": str(new_zip), "calibration": len(calibration)}
+
+
 @app.local_entrypoint()
 def rebuild_int8(model: str, limit: int = 0) -> None:
     print(rebuild_int8_head32.remote(model, limit))
@@ -618,3 +696,8 @@ def zip_meta(model_id: str, precision: str) -> dict:
 def zmeta(model: str, precisions: str = "fp32,fp16,int8") -> None:
     for precision in precisions.split(","):
         print(precision, zip_meta.remote(model, precision))
+@app.local_entrypoint()
+def static(model_id: str, limit: int = 0) -> None:
+    print(export_int8_static.remote(model_id, limit))
+
+

src/imf/export.py

@@ -523,3 +523,115 @@ def export_zips(
             )
     return zips
 
+
+
+def collect_decode_calibration(
+    enc_sess, dec_sess, texts: list[str], steps: int = 64,
+) -> list[dict]:
+    """Decoder feeds across framings for static activation calibration:
+    prefills, incremental single-token steps, and 8-token windows —
+    the shapes whose scales must hold (the framing axis). Mirrors
+    scripts/static_int8_experiment.py, which measured the recipe."""
+    import numpy as np
+
+    EOS_ID, PAD_ID = 1, 0
+    encode = encode_bytes
+
+    past_names = [i.name for i in dec_sess.get_inputs() if i.name.startswith("past_")]
+    meta = {i.name: i for i in dec_sess.get_inputs()}
+    present_of = {
+        o.name.replace("present_", "past_"): o.name
+        for o in dec_sess.get_outputs()
+        if o.name.startswith("present_")
+    }
+
+    def run(tokens, hidden, pasts):
+        feed = {
+            "input_ids": np.array([tokens], dtype=np.int64),
+            "encoder_hidden_states": hidden,
+        }
+        for n in past_names:
+            feed[n] = pasts.get(n) if pasts else np.zeros(
+                (1, meta[n].shape[1], 0, meta[n].shape[3]), dtype=np.float32
+            )
+        out = dec_sess.run(None, feed)
+        names = [o.name for o in dec_sess.get_outputs()]
+        return dict(zip(names, out, strict=True))
+
+    def split(out):
+        return int(np.argmax(out["logits"][0, -1])), {
+            k: out[v] for k, v in present_of.items()
+        }
+
+    def trim(pasts, n):
+        return {k: v[:, :, :n, :].copy() for k, v in pasts.items()}
+
+    samples: list[dict] = []
+
+    def record(tokens, hidden, pasts):
+        feed = {"input_ids": np.array([tokens], dtype=np.int64),
+                "encoder_hidden_states": hidden}
+        for n in past_names:
+            feed[n] = pasts.get(n) if pasts else np.zeros(
+                (1, meta[n].shape[1], 0, meta[n].shape[3]), dtype=np.float32
+            )
+        samples.append(feed)
+
+    for text in texts:
+        ids = encode(text)
+        hidden = enc_sess.run(None, {"input_ids": np.array([ids], dtype=np.int64)})[0]
+        tok, pasts = split(run([PAD_ID], hidden, None))
+        record([PAD_ID], hidden, None)
+        window = []
+        for _ in range(steps):
+            window.append(tok)
+            record([tok], hidden, pasts)
+            tok, pasts = split(run([tok], hidden, pasts))
+            if tok == EOS_ID:
+                break
+            if len(window) == 8:
+                seq_len = next(iter(pasts.values())).shape[2]
+                record(window, hidden, trim(pasts, max(seq_len - len(window), 0)))
+                window = []
+    return samples
+
+
+def quantize_int8_static(
+    src: Path | str, dst: Path | str, calibration: list[dict],
+    nodes_to_exclude: list[str] | None = None,
+) -> Path:
+    """fp32 -> static-int8 (calibrated QUInt8 activations, MatMul-only,
+    head fp32). Activation scales live in the graph — the export-side
+    half of TODO.impl/11: quality-clean at full set (4.6241 vs dynamic
+    4.5701) and +8% CPU decode speed. Note the measured composition:
+    fp32 encoder + static decoder; an int8-encoder + static-decoder
+    (browser-size) composition needs its own gate run before shipping.
+    """
+    from onnxruntime.quantization import (
+        CalibrationDataReader,
+        QuantFormat,
+        QuantType,
+        quantize_static,
+    )
+
+    class _Reader(CalibrationDataReader):
+        def __init__(self, data):
+            self.data = list(data)
+
+        def get_next(self):
+            return self.data.pop(0) if self.data else None
+
+        def rewind(self):
+            pass  # one pass; the corpus is the calibration set
+
+    quantize_static(
+        str(src),
+        str(dst),
+        calibration_data_reader=_Reader(calibration),
+        quant_format=QuantFormat.QOperator,
+        activation_type=QuantType.QUInt8,
+        weight_type=QuantType.QInt8,
+        op_types_to_quantize=["MatMul"],
+        nodes_to_exclude=nodes_to_exclude or [],
+    )
+    return Path(dst)