Something went wrong. Try again.
a tiny pet with an even tinier brain, living in your browser
Something went wrong. Try again.
21 kB · 442 lines
Python
at main
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443"""Weight-only candidate exports for standard GPT-2, GPTNeoX and Llama.
Usage: quantize.py EXPORT_DIR [q4|q8|fp32] [--block-size 32]
The browser graph computes only the last logits; its full-logit twin is formasked loss. q4 uses blocked uint4 body weights and int8 token tables/head;q8 uses per-output-channel int8 weights. Both use DequantizeLinear with fp32scales, never dynamic activation quantization. QDQ fusion and load-time fp32folding are separate runtime experiments, not universal speed recommendations.Topology and config identify token tables/head; untied NeoX tables stay separate."""import argparseimport collectionsimport copyimport jsonfrom pathlib import Path
import numpy as npimport onnxfrom onnx import helper, numpy_helper
from onnx_utils import DTYPES, config_dimensions, sha256
class Graph: def __init__(self, model): self.model = model self.refresh() self.names = {name for node in model.graph.node for name in (*node.input, *node.output)} self.names.update(self.initializers) self.names.update(node.name for node in model.graph.node)
def refresh(self): self.initializers = {i.name: i for i in self.model.graph.initializer} self.producers = {output: node for node in self.model.graph.node for output in node.output}
def unique(self, stem): name, suffix = stem, 0 while name in self.names: suffix += 1 name = f"{stem}_{suffix}" self.names.add(name) return name
def add_array(self, array, stem): name = self.unique(stem) self.model.graph.initializer.append(numpy_helper.from_array(np.ascontiguousarray(array), name)) self.refresh() return name
def constant(self, name, visited=None): """Resolve real constant topology, not a model-specific weight-name shim.""" visited = set() if visited is None else visited if name in visited: raise ValueError(f"cyclic constant topology at {name}") visited.add(name) if name in self.initializers: return numpy_helper.to_array(self.initializers[name]), name node = self.producers.get(name) if node is None: return None attrs = attributes(node) if node.op_type == "Constant" and "value" in attrs: return numpy_helper.to_array(attrs["value"]), name if node.op_type not in ("Identity", "Transpose", "Cast") or len(node.input) != 1: return None source = self.constant(node.input[0], visited) if source is None: return None array, root = source if node.op_type == "Transpose": array = array.transpose(attrs.get("perm")) elif node.op_type == "Cast": array = array.astype(helper.tensor_dtype_to_np_dtype(attrs["to"]), copy=False) return array, root
def depends_on(self, name, target, visited=None): if name == target: return True visited = set() if visited is None else visited if name in visited: return False visited.add(name) node = self.producers.get(name) # Follow token *values*, not Shape(input_ids) used to construct positions # or rotary tables. Such shape-only dependencies are not token lookups. if node is None or node.op_type not in ("Identity", "Cast", "Reshape", "Flatten", "Unsqueeze", "Squeeze"): return False return self.depends_on(node.input[0], target, visited)
def rank(self, name): for value in (*self.model.graph.input, *self.model.graph.value_info, *self.model.graph.output): if value.name == name and value.type.tensor_type.HasField("shape"): return len(value.type.tensor_type.shape.dim) if name in self.initializers: return len(self.initializers[name].dims) node = self.producers.get(name) if node is None: return None if node.op_type in ("Identity", "Cast", "Transpose"): return self.rank(node.input[0]) if node.op_type == "Reshape": shape = self.constant(node.input[1]) return int(shape[0].size) if shape is not None else None if node.op_type == "Flatten": return 2 return None
def attributes(node): return {attribute.name: helper.get_attribute_value(attribute) for attribute in node.attribute}
def prune(model): """Remove obsolete constant/transposed tables, including exported duplicates.""" outputs = {output.name for output in model.graph.output} nodes = list(model.graph.node) while True: used = {name for node in nodes for name in node.input} | outputs live = [node for node in nodes if any(output in used for output in node.output)] if len(live) == len(nodes): break nodes = live used = {name for node in nodes for name in node.input} | outputs inits = [init for init in model.graph.initializer if init.name in used] del model.graph.node[:] model.graph.node.extend(nodes) del model.graph.initializer[:] model.graph.initializer.extend(inits) # Intermediate metadata can otherwise retain the pre-slicing sequence shape. del model.graph.value_info[:] return model
def gemm_to_matmul(model): """Lower supported constant-weight Gemms without mutating a shared table.""" graph = Graph(model) nodes = [] materialized = {} for node in model.graph.node: if node.op_type != "Gemm": nodes.append(node) continue attrs = attributes(node) source = graph.constant(node.input[1]) if source is None or attrs.get("transA", 0) or attrs.get("alpha", 1.0) != 1.0 or attrs.get("beta", 1.0) != 1.0: raise ValueError(f"unsupported Gemm semantics at {node.name}") weight, _ = source if weight.ndim != 2: raise ValueError(f"Gemm {node.name} weight must be a matrix") transposed = bool(attrs.get("transB", 0)) key = (node.input[1], transposed) if key not in materialized: if not transposed and node.input[1] in graph.initializers: materialized[key] = node.input[1] else: materialized[key] = graph.add_array(weight.T if transposed else weight, "__linear_weight") output = graph.unique("__gemm_matmul") nodes.append(helper.make_node("MatMul", [node.input[0], materialized[key]], [output], graph.unique("__Gemm_MatMul"))) if len(node.input) > 2 and node.input[2]: nodes.append(helper.make_node("Add", [output, node.input[2]], list(node.output), graph.unique("__Gemm_Add"))) else: nodes.append(helper.make_node("Identity", [output], list(node.output), graph.unique("__Gemm_Identity"))) del model.graph.node[:] model.graph.node.extend(nodes) return model
def find_head(graph, vocab, hidden): """Trace logits back through safe output-only casts/identity/bias operations.""" if not any(output.name == "logits" for output in graph.model.graph.output): raise ValueError("causal export must expose a logits output") value, visited = "logits", set() while value not in visited: visited.add(value) node = graph.producers.get(value) if node is None: break if node.op_type == "MatMul": constant = graph.constant(node.input[1]) if constant is None or constant[0].shape != (hidden, vocab): raise ValueError(f"LM head weight must be [{hidden}, {vocab}], at {node.name}") if constant[0].dtype != np.float32: raise ValueError("candidate export must use fp32 model weights") if graph.rank(node.input[0]) != 3: raise ValueError(f"LM head activation must be [batch, sequence, hidden], at {node.name}") return node, constant if node.op_type in ("Identity", "Cast"): value = node.input[0] elif node.op_type == "Add": dynamic = [name for name in node.input if graph.constant(name) is None] if len(dynamic) != 1: break value = dynamic[0] else: break raise ValueError("unsupported logits topology: expected a matrix head and only output casts/identity/bias")
def find_embedding(graph, vocab, hidden): matches = [] for node in graph.model.graph.node: if node.op_type != "Gather" or len(node.input) < 2: continue if not graph.depends_on(node.input[1], "input_ids"): continue constant = graph.constant(node.input[0]) if constant is None: continue if attributes(node).get("axis", 0) != 0 or constant[0].shape != (vocab, hidden): raise ValueError(f"token embedding {node.name} has the wrong axis/shape: {constant[0].shape}") if constant[0].dtype != np.float32: raise ValueError("candidate export must use fp32 token embeddings") matches.append((node, constant)) if len(matches) != 1: raise ValueError(f"expected exactly one token-embedding Gather, found {len(matches)}") return matches[0]
def nibbles(values): values = values.reshape(-1) if values.size % 2: values = np.append(values, np.uint8(0)) return (values[0::2] | (values[1::2] << 4)).astype(np.uint8).tobytes()
def stored_weight(graph, weight, dtype, axis, role, records, block=32): """Create one fp32 table or one constant weight-only dequantization.""" output = graph.unique("__" + role + "_weight") shape = list(weight.shape) record = {"name": output, "role": role, "shape": shape, "elements": int(weight.size)} if dtype == "fp32": graph.model.graph.initializer.append(numpy_helper.from_array(np.ascontiguousarray(weight), output)) record["storage"] = "fp32" records.append(record) graph.refresh() return output, [] if dtype == "q4" and axis == 0 and weight.shape[0] % block == 0: k, n = weight.shape blocks = weight.reshape(k // block, block, n) low, high = np.minimum(blocks.min(1), 0), np.maximum(blocks.max(1), 0) scale = np.maximum((high - low) / 15, 1e-8).astype(np.float32) zero = np.clip(np.round(-low / scale), 0, 15).astype(np.uint8) quant = np.clip(np.round(blocks / scale[:, None]) + zero[:, None], 0, 15).astype(np.uint8) quant_name, zero_name, scale_name = (graph.unique(output + suffix) for suffix in ("_uint4", "_zero", "_scale")) graph.model.graph.initializer.extend([ helper.make_tensor(quant_name, onnx.TensorProto.UINT4, shape, nibbles(quant), raw=True), helper.make_tensor(zero_name, onnx.TensorProto.UINT4, list(scale.shape), nibbles(zero), raw=True), numpy_helper.from_array(scale, scale_name), ]) node = helper.make_node("DequantizeLinear", [quant_name, scale_name, zero_name], [output], graph.unique(output + "_Dequantize"), axis=0, block_size=block) record.update(storage="uint4", axis=0, block_size=block) else: # Tables/head deliberately remain int8 in a q4 experiment. A body K not # divisible by the requested block has an explicit int8 fallback, never # an unreported fp32 weight or an invalid blocked scale shape. if dtype == "q4": record["fallback"] = f"K={weight.shape[0]} is not divisible by block_size={block}" axis = 1 reduce_axis = 1 - axis scale = (np.maximum(np.abs(weight).max(reduce_axis), 1e-8) / 127).astype(np.float32) broadcast = scale[:, None] if axis == 0 else scale[None, :] quant = np.clip(np.round(weight / broadcast), -127, 127).astype(np.int8) quant_name = graph.add_array(quant, output + "_int8") scale_name = graph.add_array(scale, output + "_scale") node = helper.make_node("DequantizeLinear", [quant_name, scale_name], [output], graph.unique(output + "_Dequantize"), axis=axis) record.update(storage="int8", axis=axis) records.append(record) graph.refresh() return output, [node]
def quantize_model(model, config, dtype, block): dimensions = config_dimensions(config) model = gemm_to_matmul(model) # Shape inference supplies ranks for GPT-2's flattened embedding indices and # for the head activation; it must not substitute assumptions for bad graphs. model = onnx.shape_inference.infer_shapes(model, strict_mode=True) graph = Graph(model) head, (head_weight, head_root) = find_head(graph, dimensions["vocab"], dimensions["hidden"]) embedding, (embedding_weight, embedding_root) = find_embedding(graph, dimensions["vocab"], dimensions["hidden"]) equal = np.array_equal(embedding_weight.T, head_weight) # GPT-2 defaults to tying even when save_pretrained omits that default; # NeoX and standard Llama default to separate tables. configured_tie = config.get("tie_word_embeddings", config["model_type"] == "gpt2") if configured_tie and not equal: raise ValueError("config declares tied embeddings but exported embedding/head differ") tied = equal and (embedding_root == head_root or bool(configured_tie)) records, prefix = [], [] table_dtype = "fp32" if dtype == "fp32" else "q8" head_output = head.output[0] head_name, nodes = stored_weight(graph, head_weight, table_dtype, 1, "tied_table" if tied else "lm_head", records) prefix.extend(nodes) head.input[1] = head_name if tied: rank = graph.rank(embedding.input[1]) if rank not in (1, 2): raise ValueError(f"unsupported token-embedding indices rank: {rank}") columns = graph.unique("__token_columns") replacement = [helper.make_node("Gather", [head_name, embedding.input[1]], [columns], graph.unique("__Token_Gather"), axis=1), helper.make_node("Transpose", [columns], list(embedding.output), graph.unique("__Token_Transpose"), perm=list(range(1, rank + 1)) + [0])] body = [] for node in model.graph.node: body.extend(replacement if node is embedding else [node]) del model.graph.node[:] model.graph.node.extend(body) else: name, nodes = stored_weight(graph, embedding_weight, table_dtype, 0, "token_embedding", records) prefix.extend(nodes) embedding.input[0] = name graph.refresh() materialized = {} for node in model.graph.node: if node.op_type != "MatMul" or node.output[0] == head_output: continue source = graph.constant(node.input[1]) if source is None: # attention's runtime QK/value multiplications continue weight, _ = source if weight.ndim != 2 or weight.dtype != np.float32: raise ValueError(f"linear weight {node.name} must be an fp32 matrix, got {weight.shape} {weight.dtype}") old = node.input[1] if old not in materialized: name, nodes = stored_weight(graph, weight, dtype, 0 if dtype == "q4" else 1, "body", records, block) materialized[old] = name prefix.extend(nodes) node.input[1] = materialized[old] nodes = list(model.graph.node) del model.graph.node[:] model.graph.node.extend(prefix + nodes) for opset in model.opset_import: if opset.domain in ("", "ai.onnx"): opset.version = max(opset.version, 21 if dtype == "q4" else 13) model = prune(model) return model, {"dtype": dtype, "dimensions": dimensions, "tied_embedding_head": tied, "token_table_storage": table_dtype, "body_block_size": block if dtype == "q4" else None, "weights": records, "storage_counts": dict(collections.Counter(record["storage"] for record in records)), "runtime_experiments": {"fp32_fold": {"disable_quant_qdq": "1", "disable_prepacking": "1"}, "qdq_fusion": {"disable_quant_qdq": "0"}}}
def last_token_logits(model, config): model = onnx.shape_inference.infer_shapes(model, strict_mode=True) graph = Graph(model) d = config_dimensions(config) head, _ = find_head_for_slice(graph, d) activation = head.input[0] names = [graph.add_array(np.asarray(value, dtype=np.int64), "__last_" + name) for name, value in (("starts", [-1]), ("ends", [np.iinfo(np.int64).max]), ("axes", [1]))] last = graph.unique("__last_hidden") index = next(i for i, node in enumerate(model.graph.node) if node is head) model.graph.node.insert(index, helper.make_node("Slice", [activation, *names], [last], graph.unique("__Last_Token_Slice"))) head.input[0] = last logits = next(output for output in model.graph.output if output.name == "logits") logits.type.tensor_type.shape.dim[1].ClearField("dim_param") logits.type.tensor_type.shape.dim[1].dim_value = 1 return prune(model)
def find_head_for_slice(graph, dimensions): """The head's RHS may now be DQ rather than a floating initializer.""" value = "logits" while True: node = graph.producers.get(value) if node is None: break if node.op_type == "MatMul": if graph.rank(node.input[0]) != 3: raise ValueError("last-token slicing requires a rank-three head activation") weight = graph.producers.get(node.input[1]) shape = tuple(graph.initializers[node.input[1]].dims) if node.input[1] in graph.initializers else ( tuple(graph.initializers[weight.input[0]].dims) if weight is not None and weight.op_type == "DequantizeLinear" else None) if shape != (dimensions["hidden"], dimensions["vocab"]): raise ValueError(f"last-token head has unexpected weight shape {shape}") return node, shape if node.op_type in ("Identity", "Cast"): value = node.input[0] elif node.op_type == "Add": dynamic = [name for name in node.input if graph.constant(name) is None] if len(dynamic) != 1: break value = dynamic[0] else: break raise ValueError("cannot locate an unambiguous LM head for last-token slicing")
def graph_metadata(model): def describe(values): return [{"name": value.name, "dtype": onnx.TensorProto.DataType.Name(value.type.tensor_type.elem_type), "shape": [dim.dim_value if dim.HasField("dim_value") else dim.dim_param or None for dim in value.type.tensor_type.shape.dim]} for value in values] return {"inputs": describe(model.graph.input), "outputs": describe(model.graph.output)}
def main(): parser = argparse.ArgumentParser(description=__doc__) parser.add_argument("export") parser.add_argument("dtype", nargs="?", choices=DTYPES, default="q4") parser.add_argument("--block-size", type=int, default=32) args = parser.parse_args() if args.block_size < 16 or args.block_size & (args.block_size - 1): parser.error("block-size must be a power of two of at least 16") directory = Path(args.export) source = directory / "model.onnx" if not source.is_file(): source = directory / "onnx" / "model.onnx" if not source.is_file(): parser.error("no model.onnx in export root or onnx/; a single unmerged with-past export is required") browser_name = "model.onnx" if args.dtype == "fp32" else "model_quantized.onnx" if args.dtype == "q8" else "model_q4.onnx" browser = directory / "onnx" / browser_name full = directory / "onnx" / f"model_{args.dtype}_full.onnx" receipt = directory / "quantization.json" if any(path.exists() for path in (browser, full, receipt)): parser.error("candidate output already exists; use a fresh exporter directory") with open(directory / "config.json") as f: config = json.load(f) model = onnx.load(str(source), load_external_data=True) onnx.external_data_helper.convert_model_from_external_data(model) result, report = quantize_model(model, config, args.dtype, args.block_size) onnx.checker.check_model(result) sliced = last_token_logits(copy.deepcopy(result), config) onnx.checker.check_model(sliced) report.update(source={"path": str(source), "sha256": sha256(source)}, model_type=config["model_type"], full_graph=graph_metadata(result), browser_graph=graph_metadata(sliced)) browser.parent.mkdir(parents=True, exist_ok=True) with open(full, "xb") as f: onnx.save_model(result, f) with open(browser, "xb") as f: onnx.save_model(sliced, f) report["artifacts"] = {path.name: {"bytes": path.stat().st_size, "sha256": sha256(path)} for path in (browser, full)} with open(receipt, "x") as f: json.dump(report, f, indent=2, allow_nan=False) f.write("\n") print(json.dumps(report, allow_nan=False))
if __name__ == "__main__": main()