export the compressed browser format from train.py instead of int8 activations

v2/export.py

@@ -6,11 +6,13 @@
 #   "tokenizers>=0.20",
 #   "onnx>=1.17",
 #   "onnxruntime>=1.20",
+#   "onnx-ir",
 #   "safetensors>=0.4",
 #   "huggingface_hub>=0.30",
 # ]
 # ///
-"""Re-exports a trained tinynext checkpoint to browser ONNX (float and int8) with train.py's export, without retraining.
+"""Re-exports a trained tinynext checkpoint with train.py's export (float ONNX and the compressed browser format in
+web/), without retraining.
 
   uv run export.py --weights v2-release/v2 --out v2-release/browser
 """
@@ -49,7 +51,9 @@ def main():
     model = train.TinyNext(config["vocab"], config["dim"], config["layers"], config["heads"], config["ffn"])
     model.load_state_dict(load_file(weights / "model.safetensors"))
     tokenizer = Tokenizer.from_file(str(weights / "tokenizer.json"))
-    report = train.export(model, out, torch.device("cpu"), train.parity_prompts(ev, tokenizer))
+    (out / "tokenizer.json").write_bytes((weights / "tokenizer.json").read_bytes())
+    report = train.export(model, out, torch.device("cpu"), train.parity_prompts(ev, tokenizer), load("compress"))
+    (out / "web" / "config.json").write_bytes((weights / "config.json").read_bytes())
     (out / "export.json").write_text(json.dumps(report, indent=1))
     print(json.dumps(report, indent=1))
 

v2/train.py

@@ -61,15 +61,28 @@ def cpus() -> int:
 # ---- Shared eval code ------------------------------------------------------------------------------------------
 
 
+def fetch_shared(repo: str, folder: str, *names: str) -> None:
+    """Shared files (eval.py, eval_set.json, compress.py) come from next to this script or from the repo."""
+    from huggingface_hub import hf_hub_download
+
+    for name in names:
+        if not (HERE / name).exists():
+            shutil.copy(hf_hub_download(repo, f"{folder}/{name}"), HERE / name)
+
+
+def load_compress_module(repo: str, folder: str):
+    fetch_shared(repo, folder, "compress.py")
+    spec = importlib.util.spec_from_file_location("tinynext_compress", HERE / "compress.py")
+    module = importlib.util.module_from_spec(spec)
+    spec.loader.exec_module(module)
+    return module
+
+
 def load_eval_module(repo: str, folder: str):
     """Uses eval.py and eval_set.json next to this script, or fetches them from the repo, so every model is scored
     by one implementation on the same text. Registered as a module so worker processes can use it too."""
+    fetch_shared(repo, folder, "eval.py", "eval_set.json")
     path = HERE / "eval.py"
-    if not path.exists():
-        from huggingface_hub import hf_hub_download
-
-        for name in ("eval.py", "eval_set.json"):
-            shutil.copy(hf_hub_download(repo, f"{folder}/{name}"), HERE / name)
     spec = importlib.util.spec_from_file_location("tinynext_eval", path)
     module = importlib.util.module_from_spec(spec)
     sys.modules["tinynext_eval"] = module
@@ -437,20 +450,22 @@ def greedy_tokens(step, ids: list[int], empty, steps: int) -> list[int]:
     return tops
 
 
-def export(model: TinyNext, output: Path, device, prompts: list[list[int]]) -> dict:
-    """Float and int8 ONNX for the browser, checked against PyTorch on real prompts over several cached steps."""
+def export(model: TinyNext, output: Path, device, prompts: list[list[int]], compressor) -> dict:
+    """Float ONNX, and the browser format: the same graph with int8 weights (one scale per 64 values) that the page
+    expands back to float32. Both are checked against PyTorch on real prompts over several cached steps.
+
+    onnxruntime-web's int8 kernels were 8-11x slower than its float ones for this model and int8 activations changed
+    the top word too often, so only the download is compressed."""
     import onnxruntime as ort
-    from onnxruntime.quantization import QuantType, quantize_dynamic
 
     model = model.float().eval().to("cpu")
-    model.head = nn.Parameter(model.embed.weight.detach().T.contiguous().clone())
     cfg = model.config
     head_dim = cfg["dim"] // cfg["heads"]
     empty = np.zeros((cfg["layers"], 1, cfg["heads"], 0, head_dim), np.float32)
     # Trace with a non-empty cache so no shape gets fixed at zero; the checks below start from an empty one.
     with torch.no_grad():
         _, trace_k, trace_v = model(torch.tensor([[1, 2]]))
-    float_path, int8_path = output / "model.onnx", output / "model_int8.onnx"
+    float_path, web = output / "model.onnx", output / "web"
     torch.onnx.export(
         BrowserStep(model), (torch.tensor([[5, 6, 7]]), trace_k, trace_v), float_path, opset_version=17,
         input_names=["input_ids", "past_k", "past_v"], output_names=["logits", "present_k", "present_v"],
@@ -458,7 +473,9 @@ def export(model: TinyNext, output: Path, device, prompts: list[list[int]]) -> d
                       "present_k": {3: "total"}, "present_v": {3: "total"}},
         dynamo=False,
     )
-    quantize_dynamic(float_path, int8_path, weight_type=QuantType.QInt8, op_types_to_quantize=["MatMul", "Gather"], per_channel=True)
+    sizes = compressor.compress(float_path, web)
+    expanded = compressor.expand(web)
+    shutil.copy(output / "tokenizer.json", web / "tokenizer.json")
 
     @torch.no_grad()
     def torch_step(ids, k, v):
@@ -467,27 +484,18 @@ def export(model: TinyNext, output: Path, device, prompts: list[list[int]]) -> d
         return logits[0, -1].numpy(), nk.numpy(), nv.numpy()
 
     expected = [greedy_tokens(torch_step, ids, empty, 4) for ids in prompts]
-    report = {"parity_prompts": len(prompts), "parity_steps_per_prompt": 5}
-    for name, path in (("float", float_path), ("int8", int8_path)):
+    report = {"parity_prompts": len(prompts), "parity_steps_per_prompt": 5, "web_sizes": sizes}
+    for name, path in (("float", float_path), ("web_q8", web / "graph.onnx")):
         session = ort.InferenceSession(str(path), providers=["CPUExecutionProvider"])
         onnx_step = lambda ids, k, v: (lambda out: (out[0][0], out[1], out[2]))(
             session.run(None, {"input_ids": ids.astype(np.int64), "past_k": k, "past_v": v}))
         got = [greedy_tokens(onnx_step, ids, empty, 4) for ids in prompts]
-        first = np.mean([g[0] == e[0] for g, e in zip(got, expected)])
-        all_steps = np.mean([g == e for g, e in zip(got, expected)])
-        _, pk, pv = onnx_step(np.array([prompts[0]]), empty, empty)
-        timings = []
-        for _ in range(200):
-            began = time.perf_counter()
-            onnx_step(np.array([[8]]), pk, pv)
-            timings.append((time.perf_counter() - began) * 1000)
         report[name] = {
-            "megabytes": round(path.stat().st_size / 1e6, 2),
-            "same_first_token": round(float(first), 4),
-            "same_five_greedy_tokens": round(float(all_steps), 4),
-            "cached_step_ms_p50_cpu": round(float(np.median(timings)), 3),
+            "megabytes": round(sizes["weights_q8_mb"] if name == "web_q8" else path.stat().st_size / 1e6, 2),
+            "same_first_token": round(float(np.mean([g[0] == e[0] for g, e in zip(got, expected)])), 4),
+            "same_five_greedy_tokens": round(float(np.mean([g == e for g, e in zip(got, expected)])), 4),
         }
-    model.head = None
+    expanded.unlink()  # the browser rebuilds this; only weights.q8 is published
     model.to(device)
     return report
 
@@ -655,13 +663,15 @@ def main():
     (output / "metrics").mkdir(exist_ok=True)
     (output / "metrics" / f"{args.name}.json").write_text(json.dumps(scored, indent=1))
     tracker.log(headline(scored))
-    exported = export(model, output, device, parity_prompts(ev, tokenizer))
+    exported = export(model, output, device, parity_prompts(ev, tokenizer), load_compress_module(args.repo, args.eval_folder))
+
+    config = {  # also written next to the browser files, which need it
 
-    config = {
         **model.config, "parameters": parameters, "seq": args.seq, "special_tokens": dict(zip(SPECIAL, range(len(SPECIAL)))),
         "prompt": "No text after the cursor: the text before it. Otherwise: <|suf|> after <|pre|> before <|mid|>.",
     }
     (output / "config.json").write_text(json.dumps(config, indent=1))
+    shutil.copy(output / "config.json", output / "web" / "config.json")
     report = {"args": vars(args), "parameters": parameters, "data": counts, "data_minutes": round(data_minutes, 1),
               "train_minutes": round(train_minutes, 1), "history": history, "exports": exported}
     (output / "report.json").write_text(json.dumps(report, indent=1))
@@ -669,7 +679,7 @@ def main():
     print(ev.table([output / "metrics" / f"{args.name}.json"]))
     tracker.finish()
 
-    for name in ("train.py", "eval.py", "eval_set.json"):
+    for name in ("train.py", "eval.py", "eval_set.json", "compress.py"):
         shutil.copy(HERE / name, output / name)
     if (HERE / "metrics").exists():
         for path in (HERE / "metrics").glob("*.json"):