"""Quantize an optimum ONNX export of GPT-2 for the browser (WebAssembly). Usage: quantize.py EXPORT_DIR Writes EXPORT_DIR/onnx/model_q4.onnx (published; transformers.js dtype "q4") and EXPORT_DIR/onnx/model_q4_full.onnx (the same weights with logits for every position, used only by quant_eval.py). Layout: the 48 body matmuls are stored 4-bit (asymmetric, block 32) and dequantized to fp32 when the session loads (see dq4; the runtime must set session.disable_quant_qdq=1). The tied token embedding / LM head is 8-bit per row: at 4 bits GPT-2's head falls apart (loss 4.23 -> 6.98 on the base model), while int8 per row costs about 1.5%. The published graph also computes logits for the last position only, which is all generation reads. """ import os import sys import numpy as np import onnx from onnx import helper, numpy_helper def gemm_to_matmul(m, untie_head=True): """Split GPT-2's Conv1D Gemms into MatMul+Add, so every body weight is a plain [K, N] MatMul constant, and (optionally) give lm_head its own weight.""" inits = {i.name: i for i in m.graph.initializer} nodes = [] for n in m.graph.node: attrs = {a.name: onnx.helper.get_attribute_value(a) for a in n.attribute} if n.op_type == "Gemm" and n.input[1] in inits and not attrs.get("transA", 0) \ and attrs.get("alpha", 1.0) == 1.0 and attrs.get("beta", 1.0) == 1.0: w = n.input[1] if attrs.get("transB", 0): arr = numpy_helper.to_array(inits[w]).T.copy() inits[w].CopyFrom(numpy_helper.from_array(arr, w)) mm = n.output[0] + "_mm" nodes.append(onnx.helper.make_node("MatMul", [n.input[0], w], [mm], n.name + "_MatMul")) nodes.append(onnx.helper.make_node("Add", [mm, n.input[2]], [n.output[0]], n.name + "_Add") if len(n.input) > 2 and n.input[2] else onnx.helper.make_node("Identity", [mm], [n.output[0]], n.name + "_Id")) else: nodes.append(n) # The exporter casts the position-embedding Gather to float, a no-op on an # fp32 model; drop it. producers = {o: n for n in nodes for o in n.output} rename = {} for n in list(nodes): attrs = {a.name: onnx.helper.get_attribute_value(a) for a in n.attribute} src_node = producers.get(n.input[0]) if n.op_type == "Cast" else None if src_node is not None and attrs.get("to") == onnx.TensorProto.FLOAT \ and src_node.op_type == "Gather" and src_node.input[0] in inits: rename[n.output[0]] = n.input[0] nodes.remove(n) for n in nodes: for i, x in enumerate(n.input): if x in rename: n.input[i] = rename[x] producers = {o: n for n in nodes for o in n.output} for n in nodes if untie_head else (): if n.op_type == "MatMul" and n.input[1] not in inits: t = producers.get(n.input[1]) if t is not None and t.op_type == "Transpose" and t.input[0] in inits: arr = numpy_helper.to_array(inits[t.input[0]]).T.copy() name = n.name.strip("/").replace("/", "_") + "_weight" m.graph.initializer.append(numpy_helper.from_array(arr, name)) n.input[1] = name # Drop nodes whose outputs nobody reads (the head's old Transpose of the # embedding), then initializers nothing references, so untying does not # keep a full-precision copy of the embedding alive. outputs = {o.name for o in m.graph.output} while True: used = {i for n in nodes for i in n.input} | outputs dead = [n for n in nodes if not any(o in used for o in n.output)] if not dead: break nodes = [n for n in nodes if n not in dead] used = {i for n in nodes for i in n.input} keep = [i for i in m.graph.initializer if i.name in used] del m.graph.initializer[:] m.graph.initializer.extend(keep) del m.graph.node[:] m.graph.node.extend(nodes) return m def nibbles(q): """Pack uint4 values (0..15) two per byte, first value in the low nibble.""" q = q.reshape(-1) if q.size % 2: q = np.append(q, 0) return (q[0::2] | (q[1::2] << 4)).astype(np.uint8).tobytes() def dq4(m, block=32): """4-bit storage for every MatMul weight [K, N]: asymmetric, blocks of `block` rows along K, an fp32 scale and a uint4 zero point per block. One blocked DequantizeLinear (opset 21) restores the fp32 weight under its original name. (fp16 scales save 5 MB of download but need a Cast whose fp16 intermediate raised peak browser memory by about 200 MB.) With the session option session.disable_quant_qdq=1, ONNX Runtime constant-folds these at load instead of fusing them into MatMulNBits, whose WebAssembly kernel is about 5x slower than fp32 (see docs/export-and-runtime.md).""" inits = {i.name: i for i in m.graph.initializer} weights = dict.fromkeys(n.input[1] for n in m.graph.node if n.op_type == "MatMul" and n.input[1] in inits) nodes = [] for name in weights: w = numpy_helper.to_array(inits[name]).astype(np.float32) k, n = w.shape if k % block: continue b = w.reshape(k // block, block, n) lo, hi = np.minimum(b.min(1), 0), np.maximum(b.max(1), 0) scale = np.maximum((hi - lo) / 15, 1e-8).astype(np.float32) zp = np.clip(np.round(-lo / scale), 0, 15).astype(np.uint8) q = np.clip(np.round(b / scale[:, None]) + zp[:, None], 0, 15).astype(np.uint8) m.graph.initializer.remove(inits[name]) m.graph.initializer.extend([ helper.make_tensor(name + "_q4", onnx.TensorProto.UINT4, [k, n], nibbles(q), raw=True), helper.make_tensor(name + "_zp", onnx.TensorProto.UINT4, [k // block, n], nibbles(zp), raw=True), numpy_helper.from_array(scale, name + "_scale")]) nodes.append(helper.make_node("DequantizeLinear", [name + "_q4", name + "_scale", name + "_zp"], [name], name + "_Dequantize", axis=0, block_size=block)) for i, node in enumerate(nodes): m.graph.node.insert(i, node) for o in m.opset_import: if o.domain in ("", "ai.onnx"): o.version = max(o.version, 21) return m def embedding_int8(m, name="transformer.wte.weight"): """Store the tied embedding / LM head once, as int8 with one scale per token, in the head's [hidden, vocab] layout. The LM head reads it directly and the token lookup gathers columns from it (Gather axis 1, then Transpose), so after the load-time fold the browser holds a single fp32 copy instead of the table plus a transposed copy for the head (-154 MB).""" inits = {i.name: i for i in m.graph.initializer} w = numpy_helper.to_array(inits[name]).astype(np.float32) # [vocab, hidden] scale = (np.maximum(np.abs(w).max(1), 1e-8) / 127).astype(np.float32) q = np.clip(np.round(w / scale[:, None]), -127, 127).astype(np.int8) m.graph.initializer.remove(inits[name]) head = name + "_hv" # [hidden, vocab] m.graph.initializer.extend([numpy_helper.from_array(np.ascontiguousarray(q.T), head + "_int8"), numpy_helper.from_array(scale, head + "_scale")]) nodes = [helper.make_node("DequantizeLinear", [head + "_int8", head + "_scale"], [head], head + "_Dequantize", axis=1)] for n in m.graph.node: if n.op_type == "Transpose" and n.input[0] == name: for c in m.graph.node: # the head's MatMul reads the transposed table directly c.input[:] = [head if x == n.output[0] else x for x in c.input] continue if n.op_type == "Gather" and n.input[0] == name: cols = n.output[0] + "_cols" # [hidden, batch, seq] nodes.append(helper.make_node("Gather", [head, n.input[1]], [cols], n.name, axis=1)) nodes.append(helper.make_node("Transpose", [cols], [n.output[0]], n.name + "_T", perm=[1, 2, 0])) continue nodes.append(n) del m.graph.node[:] m.graph.node.extend(nodes) return m def last_token_logits(m): """Slice the LM head's input to the last position: generation only reads the last logits, and full-sequence logits cost a [seq x 50257] matmul.""" head = next(n for n in m.graph.node if "logits" in n.output) x = head.input[0] consts = {"lt_starts": [-1], "lt_ends": [2 ** 62], "lt_axes": [1]} m.graph.initializer.extend(numpy_helper.from_array(np.array(v, np.int64), k) for k, v in consts.items()) m.graph.node.insert(list(m.graph.node).index(head), helper.make_node("Slice", [x, *consts], [x + "_last"], "last_token_Slice")) head.input[0] = x + "_last" return m def main(d): out = os.path.join(d, "onnx") os.makedirs(out, exist_ok=True) m = embedding_int8(dq4(gemm_to_matmul(onnx.load(os.path.join(d, "model.onnx")), untie_head=False))) full = os.path.join(out, "model_q4_full.onnx") onnx.save_model(m, full) q4 = os.path.join(out, "model_q4.onnx") onnx.save_model(last_token_logits(m), q4) print("q4", os.path.getsize(q4) >> 20, "MiB") if __name__ == "__main__": main(sys.argv[1])