Skip to content

Commit 8470e22

Browse files
Ronald Tseronaldtse
authored andcommitted
feat: production int8 head-fp32 rebuild with the full gate stack
rebuild_int8_head32: rebuilds a shipped int8 zip with the tied head in fp32 from the fp32 graphs, then gates it — CER parity written into the zip, margin analysis, and the 1% confident-flip budget. Corrected artifacts land as {mid}-int8-head32.zip; swapping them into the release name is a separate version decision.
1 parent da8e3db commit 8470e22

1 file changed

Lines changed: 97 additions & 0 deletions

File tree

‎src/gpu/modal_export.py‎

Lines changed: 97 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -392,3 +392,100 @@ def margins(model: str, precisions: str = "fp32,fp16,int8", limit: int = 0) -> N
392392
reports = margin_model.remote(model, precisions.split(","), limit)
393393
for precision, status in reports.items():
394394
print(f"{model} [{precision}] {status}")
395+
396+
@app.function(
397+
cpu=8,
398+
memory=32 * 1024,
399+
timeout=5 * 3600,
400+
volumes={**CHECKPOINT_VOLUMES, **DATASET_VOLUMES, "/outputs": MODELS_VOLUME},
401+
)
402+
def rebuild_int8_head32(model_id: str, limit: int = 0) -> dict:
403+
"""Rebuild a shipped int8 zip with the head kept in fp32 (the E1
404+
fix) and run the full release gate stack on it: CER parity (written
405+
into the zip) + margin analysis + the confident-flip budget.
406+
407+
The corrected artifact lands as {mid}-int8-head32.zip next to the
408+
shipped one; swapping it into the release name is a version-number
409+
decision, made separately."""
410+
import sys
411+
import tempfile
412+
import zipfile
413+
414+
sys.path.insert(0, "/root/interscript-ml/src")
415+
416+
spec = MODELS[model_id]
417+
checkpoint = Path(spec["volume"]) / spec["checkpoint"]
418+
test_path = Path(spec["test_volume"]) / spec["test_data"]
419+
420+
from imf.export import head_matmul_names, load_byte_seq2seq, quantize_int8
421+
from imf.parity import (
422+
reference_decode,
423+
run_margin_analysis,
424+
run_parity,
425+
write_margin_report,
426+
write_parity,
427+
)
428+
429+
model = load_byte_seq2seq(checkpoint)
430+
pairs = _load_pairs(test_path)
431+
if limit:
432+
pairs = pairs[:limit]
433+
434+
out_dir = Path("/outputs/imf") / model_id
435+
meta_path = Path("/root/interscript-ml", spec["metadata"])
436+
mid = re.search(r"^id:\s*(\S+)", meta_path.read_text(encoding="utf-8"), re.M).group(1)
437+
fp32_zip = out_dir / f"{mid}-fp32.zip"
438+
int8_zip = out_dir / f"{mid}-int8.zip"
439+
for path in (fp32_zip, int8_zip):
440+
if not path.exists():
441+
raise RuntimeError(f"{path.name} missing on the volume")
442+
443+
with tempfile.TemporaryDirectory() as tmp:
444+
tmp = Path(tmp)
445+
with zipfile.ZipFile(fp32_zip) as zf:
446+
zf.extract("encoder.onnx", tmp)
447+
dec = "decoder-kv.onnx" if "decoder-kv.onnx" in zf.namelist() else "decoder.onnx"
448+
zf.extract(dec, tmp)
449+
enc_q, dec_q = tmp / "encoder-h32.onnx", tmp / dec.replace(".onnx", "-h32.onnx")
450+
quantize_int8(tmp / "encoder.onnx", enc_q)
451+
quantize_int8(tmp / dec, dec_q, nodes_to_exclude=head_matmul_names(tmp / dec))
452+
453+
new_zip = out_dir / f"{mid}-int8-head32.zip"
454+
with zipfile.ZipFile(int8_zip) as src, zipfile.ZipFile(
455+
new_zip, "w", zipfile.ZIP_DEFLATED
456+
) as dst:
457+
for name in src.namelist():
458+
if name == "encoder.onnx":
459+
dst.writestr(name, enc_q.read_bytes())
460+
elif name == dec:
461+
dst.writestr(name, dec_q.read_bytes())
462+
else:
463+
dst.writestr(name, src.read(name))
464+
465+
reference = reference_decode(model, [s for s, _ in pairs], max_len=128)
466+
report = run_parity(model, new_zip, pairs, max_len=128, reference=reference)
467+
if not report.passed:
468+
raise RuntimeError(f"parity gate FAILED for {new_zip.name}: {report}")
469+
write_parity(new_zip, report)
470+
margins = run_margin_analysis(model, new_zip, pairs, max_len=128)
471+
write_margin_report(margins, out_dir / f"{mid}-int8-head32-margins.json")
472+
confident = margins.flip_rate * (1 - margins.flip_low_margin_share)
473+
if confident > 0.01:
474+
raise RuntimeError(
475+
f"margin gate FAILED for {new_zip.name}: {confident:.2%} confident flips"
476+
)
477+
MODELS_VOLUME.commit()
478+
return {
479+
"model": model_id, "zip": new_zip.name,
480+
"parity": {"samples": report.samples, "cer_delta": report.cer_delta},
481+
"margins": {"flip_rate": margins.flip_rate, "kld": margins.kld_mean,
482+
"low_share": margins.flip_low_margin_share,
483+
"confident_flip_rate": round(confident, 6)},
484+
"size_bytes": new_zip.stat().st_size,
485+
}
486+
487+
488+
@app.local_entrypoint()
489+
def rebuild_int8(model: str, limit: int = 0) -> None:
490+
print(rebuild_int8_head32.remote(model, limit))
491+

0 commit comments

Comments
 (0)