d1-2B for LiteRT: decision graphs at seven lengths, the eight-sequence and eight-signature files, seven picture towers, five SM8850 NPU files, the PyTorch-free host, the fixtures and the conversion scripts
df44977 verified | """Export the d1-2B decision graph as ONE .tflite holding all seven buckets as signatures (decide_L128, decide_L256, | |
| decide_L512, decide_L1024, decide_L2048, decide_L4096_qb1024, decide_L6144_qb1024), then cast its FULLY_CONNECTED | |
| weights to float16 with the single-bucket recipe (export_decide.fp16fc), then count how the weights are stored. | |
| The seven D1Decide modules are built over ONE weight dict: d1_graph.linear keeps the tensors themselves, so every | |
| signature's FC weight is the same torch storage, and litert-torch's constant cache (backend/inline_consts.py, keyed by | |
| storage pointer, shape, stride and offset) can lower it once. Whether the converter and the quantizer keep one copy is | |
| what this script measures (`buffers` in the record): the FC weight buffers each subgraph reads, how many are distinct | |
| across subgraphs, and their bytes. | |
| PYTHONHASHSEED=0 $EXPORT conversion/export_multi.py [--keep-fp32] | |
| Writes out/graphs/d1-2b_decide_multi7_fp16fc.tflite (+ the float32 file, deleted after the fp16 one unless | |
| --keep-fp32) and results/export_decide_multi7.json. Existing outputs are never overwritten. | |
| """ | |
| from __future__ import annotations | |
| import argparse | |
| import collections | |
| import importlib.metadata | |
| import json | |
| import mmap | |
| import os | |
| import resource | |
| import sys | |
| import time | |
| HERE = os.path.dirname(os.path.abspath(__file__)) | |
| WORK = os.path.normpath(os.path.join(HERE, "..")) | |
| sys.path.insert(0, HERE) | |
| BUCKETS = ((128, 0), (256, 0), (512, 0), (1024, 0), (2048, 0), (4096, 1024), (6144, 1024)) | |
| def sig_name(L, qb): | |
| return f"decide_L{L}" + (f"_qb{qb}" if qb else "") | |
| def buffers(path): | |
| """FC weight storage across subgraphs: per subgraph the buffer each FULLY_CONNECTED reads (through a DEQUANTIZE | |
| when the weight is float16), the distinct ones over the file, their bytes; and every buffer's bytes.""" | |
| from ai_edge_litert import schema_py_generated as schema | |
| f = open(path, "rb") | |
| mm = mmap.mmap(f.fileno(), 0, access=mmap.ACCESS_READ) | |
| model = schema.Model.GetRootAs(mm, 0) | |
| op_names = {v: k for k, v in vars(schema.BuiltinOperator).items() if isinstance(v, int)} | |
| codes = [op_names.get(max(model.OperatorCodes(i).BuiltinCode(), model.OperatorCodes(i).DeprecatedBuiltinCode())) | |
| for i in range(model.OperatorCodesLength())] | |
| def nbytes(bi): | |
| b = model.Buffers(bi) | |
| if b is None: | |
| return 0 | |
| return b.DataLength() if b.DataLength() > 0 else (b.Size() if b.Offset() > 1 else 0) | |
| per_sg, users = [], collections.defaultdict(set) | |
| for si in range(model.SubgraphsLength()): | |
| sg = model.Subgraphs(si) | |
| producer = {} | |
| for oi in range(sg.OperatorsLength()): | |
| op = sg.Operators(oi) | |
| for k in range(op.OutputsLength()): | |
| producer[op.Outputs(k)] = op | |
| fc_bufs = [] | |
| for oi in range(sg.OperatorsLength()): | |
| op = sg.Operators(oi) | |
| if codes[op.OpcodeIndex()] != "FULLY_CONNECTED" or op.InputsLength() < 2: | |
| continue | |
| ti = op.Inputs(1) | |
| t = sg.Tensors(ti) | |
| bi = t.Buffer() | |
| if nbytes(bi) == 0 and ti in producer and codes[producer[ti].OpcodeIndex()] == "DEQUANTIZE": | |
| bi = sg.Tensors(producer[ti].Inputs(0)).Buffer() | |
| if nbytes(bi) == 0: | |
| continue # a runtime operand (the mask's sees_q . sees_k) | |
| fc_bufs.append(bi) | |
| users[bi].add(si) | |
| name = sg.Name().decode() if sg.Name() else str(si) | |
| per_sg.append({"subgraph": si, "name": name, "fc_weight_buffers": len(fc_bufs), | |
| "fc_weight_bytes": sum(nbytes(b) for b in fc_bufs)}) | |
| distinct = sorted(users) | |
| all_bytes = sum(nbytes(i) for i in range(model.BuffersLength())) | |
| return {"file_bytes": os.path.getsize(path), "buffers": model.BuffersLength(), "all_buffer_bytes": all_bytes, | |
| "subgraphs": model.SubgraphsLength(), | |
| "signatures": [model.SignatureDefs(i).SignatureKey().decode() for i in range(model.SignatureDefsLength())], | |
| "per_subgraph": per_sg, "fc_weight_buffers_distinct": len(distinct), | |
| "fc_weight_bytes_distinct": sum(nbytes(b) for b in distinct), | |
| "fc_weight_buffers_read_by_every_subgraph": sum(1 for b in distinct if len(users[b]) == model.SubgraphsLength()), | |
| "fc_weight_buffers_read_by_one_subgraph": sum(1 for b in distinct if len(users[b]) == 1)} | |
| def main(): | |
| ap = argparse.ArgumentParser() | |
| ap.add_argument("--keep-fp32", action="store_true") | |
| a = ap.parse_args() | |
| import torch | |
| import d1_graph as G | |
| from export_decide import fp16fc, sha256_file | |
| gdir = os.path.join(WORK, "out/graphs") | |
| p32 = os.path.join(gdir, "d1-2b_decide_multi7_fp32.tflite") | |
| p16 = os.path.join(gdir, "d1-2b_decide_multi7_fp16fc.tflite") | |
| rec_path = os.path.join(WORK, "results/export_decide_multi7.json") | |
| for p in (p16, rec_path): | |
| assert not os.path.exists(p), f"never overwrite {p}" | |
| torch.set_num_threads(8) | |
| t0 = time.time() | |
| cfg = G.Cfg() | |
| W = G.load_text_weights(dtype=torch.float32) | |
| W.pop("embed_tokens.weight", None) # the graph takes embeddings; the table stays on the host | |
| import litert_torch | |
| conv, sigs = None, [] | |
| for L, qb in BUCKETS: | |
| model = G.D1Decide(W, cfg, L, qb or None, torch.float32).eval().requires_grad_(False) | |
| T = L // 2 + 1 # the single-bucket export's sample (export_decide.py) | |
| tags = [-1] * (T // 2) + [1] * (T - T // 2 - 3) + [2] * 3 | |
| q, k = G.make_sees(T // 2, tags, L) | |
| g = torch.Generator().manual_seed(0) | |
| emb = torch.zeros(1, L, cfg.hidden) | |
| emb[0, :T] = torch.randn(T, cfg.hidden, generator=g) * 0.02 | |
| sample = {"embeds": emb, "sees_q": torch.from_numpy(q), "sees_k": torch.from_numpy(k)} | |
| name = sig_name(L, qb) | |
| conv = (litert_torch.signature(name, model, sample_kwargs=sample) if conv is None | |
| else conv.signature(name, model, sample_kwargs=sample)) | |
| sigs.append(name) | |
| rec = {"signatures": sigs, "graph": "d1_graph.D1Decide x 7 over one weight dict", | |
| "pythonhashseed": os.environ.get("PYTHONHASHSEED"), | |
| "versions": {p: importlib.metadata.version(p) for p in | |
| ("torch", "litert-torch", "litert-converter", "ai-edge-litert", "ai-edge-quantizer", "numpy")}, | |
| "load_build_s": round(time.time() - t0, 1)} | |
| t1 = time.time() | |
| if not os.path.exists(p32): | |
| lrt = conv.convert() | |
| rec["convert_s"] = round(time.time() - t1, 1) | |
| t2 = time.time() | |
| lrt.export(p32) | |
| rec["write_s"] = round(time.time() - t2, 1) | |
| del lrt | |
| else: | |
| rec["fp32_reused"] = True | |
| del conv, W, model | |
| rec["fp32"] = {"file": os.path.relpath(p32, WORK), "bytes": os.path.getsize(p32), "sha256": sha256_file(p32)} | |
| rec["buffers_fp32"] = buffers(p32) | |
| print(json.dumps({"fp32_bytes": rec["fp32"]["bytes"], **{k: rec["buffers_fp32"][k] for k in ( | |
| "subgraphs", "fc_weight_buffers_distinct", "fc_weight_bytes_distinct", | |
| "fc_weight_buffers_read_by_every_subgraph")}}), flush=True) | |
| t4 = time.time() | |
| rec["recipe"] = fp16fc(p32, p16) | |
| rec["quant_s"] = round(time.time() - t4, 1) | |
| rec["fp16fc"] = {"file": os.path.relpath(p16, WORK), "bytes": os.path.getsize(p16), "sha256": sha256_file(p16)} | |
| rec["buffers_fp16fc"] = buffers(p16) | |
| if not a.keep_fp32: | |
| os.remove(p32) | |
| rec["fp32"]["deleted"] = True | |
| rec["peak_rss_bytes"] = resource.getrusage(resource.RUSAGE_SELF).ru_maxrss | |
| rec["wall_s"] = round(time.time() - t0, 1) | |
| json.dump(rec, open(rec_path, "w"), indent=1) | |
| b = rec["buffers_fp16fc"] | |
| print(json.dumps({"fp16fc_bytes": rec["fp16fc"]["bytes"], "signatures": b["signatures"], | |
| "fc_weight_buffers_distinct": b["fc_weight_buffers_distinct"], | |
| "fc_weight_bytes_distinct": b["fc_weight_bytes_distinct"], | |
| "read_by_every_subgraph": b["fc_weight_buffers_read_by_every_subgraph"], | |
| "per_subgraph": [(s["name"], s["fc_weight_buffers"]) for s in b["per_subgraph"]], | |
| "peak_rss_gb": round(rec["peak_rss_bytes"] / 2**30, 1), "wall_s": rec["wall_s"]})) | |
| if __name__ == "__main__": | |
| main() | |