Something went wrong. Try again.
a tiny pet with a GPT-2 brain, living in your browser
Something went wrong. Try again.
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189"""Quantize an optimum ONNX export of GPT-2 for the browser (WebAssembly).
Usage: quantize.py EXPORT_DIRWrites 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 everyposition, used only by quant_eval.py).
Layout: the 48 body matmuls are stored 4-bit (asymmetric, block 32) anddequantized to fp32 when the session loads (see dq4; the runtime must setsession.disable_quant_qdq=1). The tied token embedding / LM head is 8-bit per row: at 4 bitsGPT-2's head falls apart (loss 4.23 -> 6.98 on the base model), while int8per row costs about 1.5%. The published graph also computes logits for thelast position only, which is all generation reads."""import osimport sys
import numpy as npimport onnxfrom 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])