| """Fused backbone ONNX: 28 Qwen3 decoder layers + final RMSNorm + cb0 head, with KV cache in/out. |
| |
| Build (one layer resident at a time, so it fits the 4 GB limit): |
| 1. each layer: transformers Qwen3DecoderLayer loaded from the checkpoint (tools/maya_export/streamed_backbone.py), |
| exported to ONNX (opset 17, dynamo=False), saved with external data; |
| for wo8 quantized with MatMulNBitsQuantizer(bits=8, block_size=32, is_symmetric=True, accuracy_level=4). |
| 2. head: final RMSNorm + cb0_head (fp32). |
| 3. merge: graph structure only (weights stay in the per-part data files), tensors renamed, then the weight bytes |
| are copied into one data file next to the merged model. |
| |
| Output: work_dir()/assets/v1/backbone_{precision}.onnx + backbone_{precision}.data |
| Per-part files: work_dir()/backbone_build/{precision}/ (kept for the layer-by-layer reference in the tests). |
| |
| .venv/bin/python tools/memguard.py --rss-limit-mb 4000 --min-free-mb 5000 -- \ |
| .venv-tts/bin/python -m tools.maya_export.export_backbone --precision wo8 |
| """ |
| from __future__ import annotations |
|
|
| import argparse |
| import shutil |
| from pathlib import Path |
|
|
| import numpy as np |
| import onnx |
| import torch |
| import torch.nn as nn |
| from onnx import TensorProto, helper |
|
|
| from tools.maya_export.common import ckpt_paths, work_dir |
| from tools.maya_export.streamed_backbone import StreamedBackbone, shard_files |
|
|
| N_LAYERS = 28 |
| N_KV_HEADS = 8 |
| HEAD_DIM = 128 |
| D_MODEL = 2048 |
| N_CB0 = 2049 |
| OPSET = 17 |
| PREFILL = 8 |
| COPY_CHUNK = 1 << 24 |
|
|
|
|
| class LayerKV(nn.Module): |
| """One decoder layer, KV cache as flat tensors, additive attention mask passed in.""" |
|
|
| def __init__(self, layer): |
| super().__init__() |
| self.layer = layer |
|
|
| def forward(self, h, cos, sin, mask, past_k, past_v): |
| from transformers import DynamicCache |
|
|
| cache = DynamicCache() |
| cache.update(past_k, past_v, 0) |
| out = self.layer(h, attention_mask=mask, position_ids=None, past_key_values=cache, |
| use_cache=True, cache_position=None, position_embeddings=(cos, sin)) |
| if isinstance(out, tuple): |
| out = out[0] |
| return out, cache.layers[0].keys, cache.layers[0].values |
|
|
|
|
| class Head(nn.Module): |
| def __init__(self, norm: nn.Module, cb0_head: nn.Linear) -> None: |
| super().__init__() |
| self.norm = norm |
| self.cb0_head = cb0_head |
|
|
| def forward(self, hid): |
| x = self.norm(hid) |
| return x, self.cb0_head(x) |
|
|
|
|
| def masks(T: int, past: int) -> torch.Tensor: |
| """Additive causal mask (1, 1, T, past+T): query i sees keys j <= past + i.""" |
| i = torch.arange(T)[:, None] |
| j = torch.arange(past + T)[None, :] |
| return torch.zeros(T, past + T).masked_fill(j > past + i, float("-inf"))[None, None] |
|
|
|
|
| def rope(rot, T: int, past: int) -> tuple[torch.Tensor, torch.Tensor]: |
| pos = torch.arange(past, past + T)[None] |
| return rot(torch.zeros(1, T, 1), pos) |
|
|
|
|
| def _export(module: nn.Module, args: tuple, path: Path, **kw) -> None: |
| with torch.inference_mode(): |
| torch.onnx.export(module, args, str(path), opset_version=OPSET, dynamo=False, |
| do_constant_folding=True, **kw) |
|
|
|
|
| def _save_external(model: onnx.ModelProto, path: Path) -> None: |
| |
| path.with_name(path.name + ".data").unlink(missing_ok=True) |
| onnx.save_model(model, str(path), save_as_external_data=True, all_tensors_to_one_file=True, |
| location=path.name + ".data", size_threshold=1024) |
|
|
|
|
| def _quantize_wo8(model: onnx.ModelProto, exclude: list[str] | None = None) -> onnx.ModelProto: |
| from onnxruntime.quantization.matmul_nbits_quantizer import MatMulNBitsQuantizer |
|
|
| q = MatMulNBitsQuantizer(model, bits=8, block_size=32, is_symmetric=True, accuracy_level=4, |
| nodes_to_exclude=exclude or []) |
| q.process() |
| return q.model.model |
|
|
|
|
| def export_layer(sb: StreamedBackbone, i: int, precision: str, build: Path) -> Path: |
| sb._load_layer(i) |
| sb.layer.self_attn.layer_idx = 0 |
| wrap = LayerKV(sb.layer).eval() |
| h0 = torch.randn(1, PREFILL, D_MODEL) * 0.02 |
| cos, sin = rope(sb.rotary, PREFILL, 0) |
| pk = torch.zeros(1, N_KV_HEADS, 0, HEAD_DIM) |
| pv = torch.zeros(1, N_KV_HEADS, 0, HEAD_DIM) |
| tmp = build / f"layer_{i:02d}_raw.onnx" |
| _export(wrap, (h0, cos, sin, masks(PREFILL, 0), pk, pv), tmp, |
| input_names=["h", "cos", "sin", "mask", "past_k", "past_v"], |
| output_names=["out", "new_k", "new_v"], |
| dynamic_axes={"h": {1: "t"}, "cos": {1: "t"}, "sin": {1: "t"}, "mask": {2: "t", 3: "total"}, |
| "past_k": {2: "past"}, "past_v": {2: "past"}, |
| "out": {1: "t"}, "new_k": {2: "total"}, "new_v": {2: "total"}}) |
| model = onnx.load(str(tmp)) |
| if precision == "wo8": |
| model = _quantize_wo8(model) |
| out = build / f"layer_{i:02d}.onnx" |
| _save_external(model, out) |
| tmp.unlink() |
| tmp.with_suffix(".onnx.data").unlink(missing_ok=True) |
| return out |
|
|
|
|
| def export_head(sb: StreamedBackbone, build: Path) -> Path: |
| cb0 = nn.Linear(D_MODEL, N_CB0, bias=False) |
| cb0.weight.data.copy_(sb.weights.get("backbone.cb0_head.weight").float()) |
| head = Head(sb.norm, cb0).eval() |
| tmp = build / "head_raw.onnx" |
| _export(head, (torch.randn(1, PREFILL, D_MODEL),), tmp, input_names=["hid"], |
| output_names=["hidden", "cb0_logits"], |
| dynamic_axes={"hid": {1: "t"}, "hidden": {1: "t"}, "cb0_logits": {1: "t"}}) |
| out = build / "head.onnx" |
| _save_external(onnx.load(str(tmp)), out) |
| tmp.unlink() |
| tmp.with_suffix(".onnx.data").unlink(missing_ok=True) |
| return out |
|
|
|
|
| def _rename(model: onnx.ModelProto, prefix: str, mapping: dict[str, str], keep: set[str]): |
| """Rename every tensor and node of a part. Names in `mapping` are replaced; names in `keep` are shared |
| across parts (cos/sin/mask); everything else gets `prefix`.""" |
| def new(name: str) -> str: |
| if name == "": |
| return name |
| if name in mapping: |
| return mapping[name] |
| if name in keep: |
| return name |
| return prefix + name |
|
|
| g = model.graph |
| nodes = [] |
| for node in g.node: |
| n = onnx.NodeProto() |
| n.CopyFrom(node) |
| del n.input[:] |
| n.input.extend(new(x) for x in node.input) |
| del n.output[:] |
| n.output.extend(new(x) for x in node.output) |
| n.name = prefix + node.name if node.name else "" |
| nodes.append(n) |
| inits = [] |
| for init in g.initializer: |
| t = TensorProto() |
| t.CopyFrom(init) |
| t.name = new(init.name) |
| inits.append(t) |
| return nodes, inits |
|
|
|
|
| def merge(layer_paths: list[Path], head_path: Path, out_path: Path, precision: str) -> None: |
| all_nodes, all_inits, opsets = [], [], {} |
| for i, path in enumerate(layer_paths): |
| m = onnx.load(str(path), load_external_data=False) |
| src_in = "embeds" if i == 0 else f"hid_{i - 1}" |
| mapping = {"h": src_in, "past_k": f"past_k_{i}", "past_v": f"past_v_{i}", |
| "out": f"hid_{i}", "new_k": f"present_k_{i}", "new_v": f"present_v_{i}"} |
| nodes, inits = _rename(m, f"L{i}_", mapping, keep={"cos", "sin", "mask"}) |
| all_nodes += nodes |
| all_inits += inits |
| for op in m.opset_import: |
| opsets[op.domain] = max(opsets.get(op.domain, 0), op.version) |
| hm = onnx.load(str(head_path), load_external_data=False) |
| nodes, inits = _rename(hm, "H_", {"hid": f"hid_{len(layer_paths) - 1}", "hidden": "hidden", |
| "cb0_logits": "cb0_logits"}, keep=set()) |
| all_nodes += nodes |
| all_inits += inits |
| for op in hm.opset_import: |
| opsets[op.domain] = max(opsets.get(op.domain, 0), op.version) |
|
|
| data_name = out_path.with_suffix(".data").name |
| _consolidate(all_inits, layer_paths + [head_path], out_path.parent, data_name) |
|
|
| inputs = [helper.make_tensor_value_info("embeds", TensorProto.FLOAT, [1, "T", D_MODEL]), |
| helper.make_tensor_value_info("cos", TensorProto.FLOAT, [1, "T", HEAD_DIM]), |
| helper.make_tensor_value_info("sin", TensorProto.FLOAT, [1, "T", HEAD_DIM]), |
| helper.make_tensor_value_info("mask", TensorProto.FLOAT, [1, 1, "T", "total"])] |
| outputs = [helper.make_tensor_value_info("hidden", TensorProto.FLOAT, [1, "T", D_MODEL]), |
| helper.make_tensor_value_info("cb0_logits", TensorProto.FLOAT, [1, "T", N_CB0])] |
| n = len(layer_paths) |
| for i in range(n): |
| inputs.append(helper.make_tensor_value_info(f"past_k_{i}", TensorProto.FLOAT, [1, N_KV_HEADS, "P", HEAD_DIM])) |
| inputs.append(helper.make_tensor_value_info(f"past_v_{i}", TensorProto.FLOAT, [1, N_KV_HEADS, "P", HEAD_DIM])) |
| outputs.append(helper.make_tensor_value_info(f"present_k_{i}", TensorProto.FLOAT, [1, N_KV_HEADS, "total", HEAD_DIM])) |
| outputs.append(helper.make_tensor_value_info(f"present_v_{i}", TensorProto.FLOAT, [1, N_KV_HEADS, "total", HEAD_DIM])) |
|
|
| graph = helper.make_graph(all_nodes, f"maya2_backbone_{precision}", inputs, outputs, all_inits) |
| opset_imports = [helper.make_opsetid(d, v) for d, v in sorted(opsets.items())] |
| ir_version = onnx.load(str(layer_paths[0]), load_external_data=False).ir_version |
| model = helper.make_model(graph, opset_imports=opset_imports, ir_version=ir_version) |
| onnx.save_model(model, str(out_path)) |
|
|
|
|
| def _consolidate(inits: list[TensorProto], parts: list[Path], out_dir: Path, data_name: str) -> None: |
| """Copy every external tensor's bytes into one data file, updating offset and location in place.""" |
| offset = 0 |
| with open(out_dir / data_name, "wb") as dst: |
| for t in inits: |
| if t.data_location != TensorProto.EXTERNAL: |
| continue |
| entries = {e.key: e.value for e in t.external_data} |
| loc = entries["location"] |
| src = _source_file(loc, parts) |
| length = int(entries["length"]) |
| start = int(entries.get("offset", 0)) |
| with open(src, "rb") as f: |
| f.seek(start) |
| remaining = length |
| while remaining: |
| chunk = f.read(min(COPY_CHUNK, remaining)) |
| if not chunk: |
| raise IOError(f"short read for {t.name} in {src}") |
| dst.write(chunk) |
| remaining -= len(chunk) |
| del t.external_data[:] |
| for k, v in (("location", data_name), ("offset", str(offset)), ("length", str(length))): |
| e = t.external_data.add() |
| e.key, e.value = k, v |
| offset += length |
|
|
|
|
| def _source_file(location: str, parts: list[Path]) -> Path: |
| for p in parts: |
| if (p.parent / location).exists(): |
| return p.parent / location |
| raise FileNotFoundError(location) |
|
|
|
|
| def build(precision: str) -> Path: |
| if precision not in ("fp32", "wo8"): |
| raise ValueError(precision) |
| build_dir = work_dir() / "backbone_build" / precision |
| build_dir.mkdir(parents=True, exist_ok=True) |
| paths = ckpt_paths() |
| sb = StreamedBackbone(paths["backbone_dir"], shard_files(paths["ckpt"]), n_layers=N_LAYERS) |
| layer_paths = [] |
| for i in range(N_LAYERS): |
| layer_paths.append(export_layer(sb, i, precision, build_dir)) |
| print(f"layer {i:02d} {precision} ok", flush=True) |
| head = export_head(sb, build_dir) |
| out_dir = work_dir() / "assets" / "v1" |
| out_dir.mkdir(parents=True, exist_ok=True) |
| out = out_dir / f"backbone_{precision}.onnx" |
| merge(layer_paths, head, out, precision) |
| print(f"wrote {out} and {out.with_suffix('.data').name}", flush=True) |
| return out |
|
|
|
|
| if __name__ == "__main__": |
| ap = argparse.ArgumentParser() |
| ap.add_argument("--precision", choices=("fp32", "wo8"), required=True) |
| build(ap.parse_args().precision) |
|
|
