Upload stream_int8.py with huggingface_hub

stream_int8.py

@@ -0,0 +1,269 @@
+#!/usr/bin/env python3
+"""
+Streaming weight-only INT8 quantizer for large ONNX models.
+
+Implements the same transformation as:
+  quantize_dynamic(..., MatMulConstBOnly=True, per_channel=False, weight_type=QInt8)
+
+Fully streaming: reads and writes one tensor at a time.
+Peak RAM: ~1.5 GB (for the largest single tensor, the embedding table ~1.2 GB).
+
+Usage: python stream_int8.py
+"""
+
+import gc
+from pathlib import Path
+import numpy as np
+import onnx
+from onnx import TensorProto, numpy_helper, helper
+
+FP32_ONNX = Path("/Volumes/backups/ai/zerank_fp32_tmp/model_fp32.onnx")
+FP32_DATA = Path("/Volumes/backups/ai/zerank_fp32_tmp/model_fp32.onnx_data")
+INT8_OUT  = Path("/Volumes/backups/ai/zerank_onnx_int8/model_int8.onnx")
+INT8_DATA = Path("/Volumes/backups/ai/zerank_onnx_int8/model_int8.onnx_data")
+MODEL_ID  = "zeroentropy/zerank-1-small"
+
+INT8_OUT.parent.mkdir(parents=True, exist_ok=True)
+
+
+def quantize_tensor_per_tensor(arr: np.ndarray):
+    """Symmetric per-tensor INT8 quantization (zero_point = 0)."""
+    arr = arr.astype(np.float32)
+    abs_max = np.max(np.abs(arr))
+    if abs_max == 0:
+        scale = np.float32(1.0)
+        quantized = np.zeros_like(arr, dtype=np.int8)
+    else:
+        scale = np.float32(abs_max / 127.0)
+        quantized = np.clip(np.round(arr / scale), -127, 127).astype(np.int8)
+    return quantized, scale
+
+
+def add_external_data(init: onnx.TensorProto, offset: int, length: int, data_file_name: str):
+    """Update an initializer proto to point to external data."""
+    init.data_location = TensorProto.EXTERNAL
+    init.ClearField("external_data")
+    for k, v in [("location", data_file_name), ("offset", str(offset)), ("length", str(length))]:
+        e = init.external_data.add()
+        e.key, e.value = k, v
+
+
+def quantize_model():
+    print(f"Loading proto skeleton (no external data)...")
+    m = onnx.load(str(FP32_ONNX), load_external_data=False)
+    print(f"  Nodes: {len(m.graph.node)}, Initializers: {len(m.graph.initializer)}")
+
+    # Build index of external initializers
+    ext_index = {}  # name → (offset, length, dtype, dims)
+    inline_index = {}  # name → data bytes (for inline tensors)
+    for init in m.graph.initializer:
+        if init.data_location == TensorProto.EXTERNAL:
+            info = {e.key: e.value for e in init.external_data}
+            ext_index[init.name] = {
+                "offset": int(info.get("offset", 0)),
+                "length": int(info.get("length", 0)),
+                "dtype": init.data_type,
+                "dims": list(init.dims),
+            }
+        else:
+            inline_index[init.name] = init
+
+    # Find all MatMul nodes with constant B (initializer)
+    matmul_b_names = set()
+    for node in m.graph.node:
+        if node.op_type == "MatMul" and len(node.input) >= 2:
+            b_name = node.input[1]
+            if b_name in ext_index or b_name in inline_index:
+                matmul_b_names.add(b_name)
+
+    print(f"  MatMul B weights to quantize: {len(matmul_b_names)}")
+    non_matmul = [name for name, meta in ext_index.items() if name not in matmul_b_names]
+    print(f"  Non-MatMul external tensors (kept as FP32): {len(non_matmul)}")
+
+    # ── Phase 1: Stream all tensors to INT8 data file ─────────────────────────
+    print(f"\nPhase 1: Writing tensor data to {INT8_DATA.name}")
+    data_file_name = INT8_DATA.name  # just the filename, not full path
+
+    # Track where each tensor ends up in the output data file
+    # key → (offset, length) for the output
+    out_positions = {}  # name → (offset, length)
+    # For quantized weights: also store scale values (tiny, inline later)
+    scale_values = {}   # weight_name → float32 scale
+
+    try:
+        from tqdm import tqdm
+    except ImportError:
+        tqdm = None
+
+    offset = 0
+    with open(str(FP32_DATA), "rb") as fp32_f, open(str(INT8_DATA), "wb") as int8_f:
+        # 1a. Write quantized MatMul weights (INT8)
+        matmul_list = sorted(matmul_b_names)
+        if tqdm:
+            it = tqdm(matmul_list, desc="  Quantizing MatMul weights")
+        else:
+            it = matmul_list
+
+        for w_name in it:
+            if w_name in ext_index:
+                meta = ext_index[w_name]
+                fp32_f.seek(meta["offset"])
+                raw = fp32_f.read(meta["length"])
+                arr = np.frombuffer(raw, dtype=np.float32).reshape(meta["dims"])
+            else:
+                arr = numpy_helper.to_array(inline_index[w_name]).astype(np.float32)
+
+            q_arr, scale_val = quantize_tensor_per_tensor(arr)
+            del arr
+            scale_values[w_name] = scale_val
+
+            raw_int8 = q_arr.tobytes()
+            int8_f.write(raw_int8)
+            out_positions[w_name + "_quantized"] = (offset, len(raw_int8))
+            offset += len(raw_int8)
+            del q_arr
+
+        # 1b. Copy non-MatMul external tensors verbatim (already FP32/int64/etc.)
+        print(f"  Copying {len(non_matmul)} non-MatMul tensors...")
+        for name in non_matmul:
+            meta = ext_index[name]
+            fp32_f.seek(meta["offset"])
+            raw = fp32_f.read(meta["length"])
+            int8_f.write(raw)
+            out_positions[name] = (offset, len(raw))
+            offset += len(raw)
+
+    print(f"  Data file written: {INT8_DATA.stat().st_size / 1e9:.2f} GB")
+
+    # ── Phase 2: Rebuild the ONNX proto ───────────────────────────────────────
+    print("\nPhase 2: Rebuilding ONNX proto...")
+
+    # Rebuild graph: replace MatMul nodes with DQL → MatMul
+    new_nodes = []
+    dql_inserted = set()
+    for node in m.graph.node:
+        if node.op_type == "MatMul" and node.input[1] in matmul_b_names:
+            b_name = node.input[1]
+            dql_out_name = b_name + "_dequant"
+            if b_name not in dql_inserted:
+                dql_node = helper.make_node(
+                    "DequantizeLinear",
+                    inputs=[b_name + "_quantized", b_name + "_scale", b_name + "_zero_point"],
+                    outputs=[dql_out_name],
+                )
+                new_nodes.append(dql_node)
+                dql_inserted.add(b_name)
+            new_node = helper.make_node(
+                "MatMul",
+                inputs=[node.input[0], dql_out_name],
+                outputs=list(node.output),
+                name=node.name,
+            )
+            new_nodes.append(new_node)
+        else:
+            new_nodes.append(node)
+
+    del m.graph.node[:]
+    m.graph.node.extend(new_nodes)
+
+    # Rebuild initializers
+    new_initializers = []
+
+    # a. Quantized MatMul weights (external data)
+    for w_name in matmul_b_names:
+        meta = ext_index.get(w_name) or {
+            "dims": list(numpy_helper.to_array(inline_index[w_name]).shape)
+        }
+        dims = meta["dims"]
+
+        q_init = TensorProto()
+        q_init.name = w_name + "_quantized"
+        q_init.data_type = TensorProto.INT8
+        q_init.dims.extend(dims)
+        off, length = out_positions[w_name + "_quantized"]
+        add_external_data(q_init, off, length, data_file_name)
+
+        scale_init = numpy_helper.from_array(
+            np.array([scale_values[w_name]], dtype=np.float32), name=w_name + "_scale"
+        )
+        zp_init = numpy_helper.from_array(
+            np.array([0], dtype=np.int8), name=w_name + "_zero_point"
+        )
+
+        new_initializers.extend([q_init, scale_init, zp_init])
+
+    # b. Non-MatMul external tensors (external data, already written)
+    for name in non_matmul:
+        meta = ext_index[name]
+        init = TensorProto()
+        init.name = name
+        init.data_type = meta["dtype"]
+        init.dims.extend(meta["dims"])
+        off, length = out_positions[name]
+        add_external_data(init, off, length, data_file_name)
+        new_initializers.append(init)
+
+    # c. Inline initializers from FP32 model (already inline in proto — not external data)
+    for init in m.graph.initializer:
+        if init.name not in ext_index:  # it's inline
+            new_initializers.append(init)
+
+    del m.graph.initializer[:]
+    m.graph.initializer.extend(new_initializers)
+    del m.graph.value_info[:]  # clear stale type annotations
+
+    print(f"  Saving proto → {INT8_OUT}")
+    onnx.save(m, str(INT8_OUT))
+    print(f"  Proto size: {INT8_OUT.stat().st_size / 1e6:.1f} MB")
+    total_gb = (INT8_OUT.stat().st_size + INT8_DATA.stat().st_size) / 1e9
+    print(f"  Total INT8 size: {total_gb:.2f} GB")
+
+
+def verify():
+    import onnxruntime as ort
+    from transformers import AutoTokenizer
+
+    print(f"\nVerifying {INT8_OUT.name}...")
+    sess_opts = ort.SessionOptions()
+    sess = ort.InferenceSession(
+        str(INT8_OUT), sess_opts, providers=["CPUExecutionProvider"]
+    )
+    for inp in sess.get_inputs():
+        print(f"  in:  {inp.name} {inp.shape}")
+    for out in sess.get_outputs():
+        print(f"  out: {out.name} {out.shape}")
+
+    tok = AutoTokenizer.from_pretrained(MODEL_ID, trust_remote_code=True)
+    pairs = [
+        ("what is a panda?", "A panda is a large black-and-white bear native to China."),
+        ("what is a panda?", "The sky is blue and the grass is green."),
+    ]
+    scores = []
+    for q, d in pairs:
+        enc = tok(q, d, return_tensors="np", truncation=True, max_length=256)
+        logit = sess.run(["logits"], {
+            "input_ids":      enc["input_ids"].astype(np.int64),
+            "attention_mask": enc["attention_mask"].astype(np.int64),
+        })[0]
+        scores.append(float(logit[0][0]))
+
+    print(f"  logits: {[f'{s:.3f}' for s in scores]}")
+    assert scores[0] > scores[1], \
+        f"Relevant doc should score higher: {scores[0]:.3f} vs {scores[1]:.3f}"
+    print("  OK — relevant doc ranked higher")
+
+
+if __name__ == "__main__":
+    for p in [INT8_OUT, INT8_DATA]:
+        if p.exists():
+            p.unlink()
+            print(f"Deleted {p.name}")
+
+    quantize_model()
+    gc.collect()
+    verify()
+
+    print("\nAll done. Upload commands:")
+    print("  huggingface-cli upload cstr/zerank-1-small-ONNX /private/tmp/zerank_export/zerank_onnx . --repo-type model")
+    print(f"  huggingface-cli upload cstr/zerank-1-small-ONNX {INT8_OUT.parent}/ . --commit-message 'add INT8' --repo-type model --include '*.onnx*'")
+    print(f"  huggingface-cli upload cstr/zerank-1-small-ONNX /Volumes/backups/ai/zerank_onnx_int4/model_int4_full.onnx model_int4_full.onnx --repo-type model")