apetersson/v41-quant-worker

1#!/usr/bin/env python32"""Aggregate the per-shard manifests of a mixed-Q2 V4.1 expert checkpoint into one manifest and README.3 4Reads ``plan.json`` plus every ``quant-*.safetensors.manifest.json`` written by5``v41_quant_experts.py convert`` and writes ``manifest.json`` and ``README.md`` into the6artifact directory. Pure aggregation over recorded digests: it does not read the weights.7"""8 9from __future__ import annotations10 11import argparse12import hashlib13import json14from datetime import UTC, datetime15from pathlib import Path16 17LLAMA_COMMIT = "465e49b9cea78a68b9c244ffb48d0ee24a82873d"18 19README_TEMPLATE = """# DeepSeek-V4.1-Flash mixed-Q2 expert checkpoint ({stamp})20 21Quantised derivative of `{repo}` revision `{revision}`, produced on a 32-vCPU Runpod CPU pod22with the pinned llama.cpp ggml reference quantisers (`{llama_commit}`).23 24## What this is25 26The released checkpoint stores routed experts as MXFP4: 16 bytes of packed E2M1 values plus27one E8M0 scale byte per 32 weights (4.25 bits/weight). This artifact decodes exactly those28blocks and stores the 40 main-layer routed expert matrices as ggml blocks:29 30| projection | ggml type | bits/weight | tensors | weights |31|---|---|---:|---:|---:|32| gate/up `w1`, `w3` | IQ2_XXS | 2.0625 | {n_iq2} | {w_iq2} |33| down `w2` | Q2_K | 2.625 | {n_q2k} | {w_q2k} |34 35Block geometry (blocks run along the input/reduction dimension; ggml lists `ne` fastest-first, so36the same bytes are `ne=[in, out]`):37 38| tensor | logical out x in | ggml type | blocks/row | bytes/row | first block offset |39|---|---|---:|---:|---:|---|40| `w1.weight`, `w3.weight` | 2304 x 5120 | IQ2_XXS | 20 | 1320 | 0 |41| `w2.weight` | 5120 x 2304 | Q2_K | 9 | 756 | 0 |42 43Each expert tensor is a contiguous row-major sequence of those rows, so a carrier can read4420 blocks per row for gate/up and 9 per row for down with no extra permutation.45 46Total quantised payload **{quant_bytes} GB** for **{quant_weights}** routed weights47({bits} bits/weight average). With the non-routed weights held at native precision and the48Engram tables resident on the host, this fits one B200 180 GB or two RTX PRO 6000 96 GB cards.49 50Quantisation is **not imatrix-based**. No activation calibration was collected and no51importance matrix was consumed: IQ2_XXS was called with a constant unit importance vector52(the unweighted reference path). That is a deliberate recipe choice, not a calibrated one.53 54## Contents55 56| group | tensors | bytes |57|---|---:|---:|58{preserved_rows}59 60Every tensor other than the requantised experts is copied byte-for-byte from the source with61its original dtype string (`F8_E4M3`, `F8_E8M0`, `BF16`, `F32`, `I8`). Requantised expert62tensors keep their original names and are stored as `U8` block bytes; the per-tensor ggml63type, byte shape and block geometry are recorded in `manifest.json` and in each shard64manifest. MTP/DSpark expert tensors stay at native precision because DSpark stays disabled.65 66The four large Engram tensors (`layers.1.engram.embed.*`, `layers.14.engram.embed.*`,67{external_bytes} GB) are **not** copied: the release plan keeps Engram rows and scales native68on host storage. `manifest.json` records their source shard, byte range and source digest so a69loader can attach them from the immutable source revision. The 46,080 MXFP4 scale tensors of70the requantised experts ({dropped_bytes} GB) are dropped because IQ2_XXS and Q2_K carry their71own scales; their source digests are recorded as well.72 73## Layout74 75| output shard | source shard | tensors | bytes |76|---|---|---:|---:|77{shard_rows}78| **total** | | **{total_tensors}** | **{total_bytes} GB** |79 80`plan.json` holds the full tensor map (source shard, byte range, role, output shard, output81offset, output shape and type). Each `quant-*.safetensors` shard has a matching82`.manifest.json` with per-tensor source and output SHA-256 digests.83 84## Verification85 86```bash87python3 v41_quant_experts.py verify --out <this directory> --sample 0    # all digests88python3 v41_quant_audit.py --source <source dir> --out <this directory> \89    --json audit.json                                                    # fidelity sample90```91 92`audit.json` reports cosine similarity and relative RMS error of the dequantised blocks93against the MXFP4 source decode, computed with the independent gguf-py dequantisers from the94same pinned llama.cpp revision. With `--ggml-lib` it also recomputes the quantisation and95reports whether the stored digests are reproducible.96 97## Reproducibility98 99{environment_rows}100 101The ggml reference quantisers are deterministic for a given build and CPU target (repeated102recomputation reproduces the stored digests), but they are **not** bit-identical across SIMD103architectures: the ARM NEON and x86 AVX-512 paths select different codepoints in the IQ2_XXS104grid and can differ in Q2_K rounding. Re-running on a different architecture reproduces the105recipe, not these exact bytes. Reproduce or audit with the pinned commit above and the same106CPU class, or treat a differing digest as expected and compare fidelity through `audit.json`.107 108## Limits109 110This is a weights artifact. No V4.1 quantised runtime exists in the workspace, so nothing111here has been executed end to end and none of these numbers are behavioural results. The112blocks are the inputs a fused-kernel carrier needs; executing them requires new kernels for113IQ2_XXS/Q2_K experts plus host-resident Engram lookup.114"""115 116 117def human_gb(value: int) -> str:118    return f"{value / 1e9:,.2f}"119 120 121def main() -> int:122    parser = argparse.ArgumentParser(description=__doc__)123    parser.add_argument("--out", type=Path, required=True)124    parser.add_argument("--repo", default="deepseek-ai/DeepSeek-V4.1-Flash")125    parser.add_argument("--stamp", default=None)126    parser.add_argument("--check", action="store_true", help="fail if any planned shard is missing")127    parser.add_argument(128        "--environment",129        type=Path,130        help="JSON file describing the execution environment (CPU, SIMD, host, pod) to embed",131    )132    args = parser.parse_args()133 134    out = args.out135    plan = json.loads((out / "plan.json").read_text())136    role_by_name: dict[str, str] = {}137    type_by_name: dict[str, str] = {}138    for shard in plan["shards"]:139        for tensor in shard["tensors"]:140            role_by_name[tensor["name"]] = tensor["role"]141            if tensor["role"] == "quantised":142                type_by_name[tensor["name"]] = tensor["ggml_type"]143 144    per_type: dict[str, dict[str, int]] = {}145    per_role: dict[str, dict[str, int]] = {}146    shard_rows = []147    digest_index = []148    total_output_bytes = 0149    manifests = sorted(150        (json.loads(p.read_text()) for p in out.glob("quant-*.safetensors.manifest.json")),151        key=lambda m: m["output_shard"],152    )153    for manifest in manifests:154        shard_digest = hashlib.sha256()155        data_bytes = 0156        for tensor in sorted(manifest["tensors"], key=lambda t: t["name"]):157            role = role_by_name.get(tensor["name"], "unknown")158            per_role.setdefault(role, {"tensors": 0, "bytes": 0})159            per_role[role]["tensors"] += 1160            per_role[role]["bytes"] += tensor["output_bytes"]161            if role == "quantised":162                type_name = type_by_name[tensor["name"]]163                entry = per_type.setdefault(type_name, {"tensors": 0, "bytes": 0})164                entry["tensors"] += 1165                entry["bytes"] += tensor["output_bytes"]166            total_output_bytes += tensor["output_bytes"]167            data_bytes += tensor["output_bytes"]168            shard_digest.update(tensor["name"].encode())169            shard_digest.update(tensor["output_sha256"].encode())170        shard_rows.append(171            (172                manifest["output_shard"],173                manifest["source_shard"],174                len(manifest["tensors"]),175                manifest["file_bytes"],176            )177        )178        digest_index.append(179            {180                "output_shard": manifest["output_shard"],181                "source_shard": manifest["source_shard"],182                "tensors": len(manifest["tensors"]),183                "file_bytes": manifest["file_bytes"],184                "data_bytes": data_bytes,185                "digest_of_shard_digests": shard_digest.hexdigest(),186            }187        )188 189    expected = {s["output_shard"] for s in plan["shards"]}190    produced = {m["output_shard"] for m in manifests}191    missing = sorted(expected - produced)192    totals = plan["totals"]193    record = {194        "artifact": "DeepSeek-V4.1-Flash mixed-Q2 routed experts",195        "created_at_utc": args.stamp or datetime.now(UTC).strftime("%Y-%m-%dT%H:%M:%SZ"),196        "source": {197            "repo": args.repo,198            "revision": plan["source_revision"],199            "tensors": totals["source_tensors"],200        },201        "recipe": plan["recipe"],202        "importance": plan["importance"],203        "quantiser": f"llama.cpp ggml reference quantisers, commit {LLAMA_COMMIT}",204        "totals": {205            **totals,206            "shards_planned": len(expected),207            "shards_produced": len(produced),208            "missing_shards": missing,209            "unexpected_shards": sorted(produced - expected),210            "bytes_by_role": per_role,211            "quantised_bytes_by_type": per_type,212            "output_bytes_from_manifests": total_output_bytes,213        },214        "environment": json.loads(args.environment.read_text()) if args.environment else {},215        "external_native": plan["external"],216        "dropped_expert_scales": len(plan["dropped"]),217        "shards": digest_index,218    }219    (out / "manifest.json").write_text(json.dumps(record, indent=1) + "\n")220 221    role_rows = "\n".join(222        f"| {role} | {values['tensors']:,} | {human_gb(values['bytes'])} |"223        for role, values in sorted(per_role.items())224    )225    shard_rows_md = "\n".join(226        f"| `{name}` | `{source}` | {tensors:,} | {human_gb(file_bytes)} |"227        for name, source, tensors, file_bytes in shard_rows228    )229    iq2 = per_type.get("IQ2_XXS", {"tensors": 0, "bytes": 0})230    q2k = per_type.get("Q2_K", {"tensors": 0, "bytes": 0})231    readme = README_TEMPLATE.format(232        stamp=record["created_at_utc"],233        repo=args.repo,234        revision=plan["source_revision"],235        llama_commit=LLAMA_COMMIT,236        n_iq2=f"{iq2['tensors']:,}",237        n_q2k=f"{q2k['tensors']:,}",238        w_iq2=f"{(iq2['bytes'] // 66) * 256:,}",239        w_q2k=f"{(q2k['bytes'] // 84) * 256:,}",240        quant_bytes=human_gb(totals["quantised_bytes"]),241        quant_weights=f"{totals['quantised_logical_weights']:,}",242        bits=round(8 * totals["quantised_bytes"] / totals["quantised_logical_weights"], 4),243        preserved_rows=role_rows,244        external_bytes=human_gb(totals["external_bytes"]),245        dropped_bytes=human_gb(totals["dropped_bytes"]),246        environment_rows="\n".join(247            f"- {key}: `{value}`"248            for key, value in sorted(249                (json.loads(args.environment.read_text()) if args.environment else {}).items()250            )251        )252        or "- not recorded",253        shard_rows=shard_rows_md,254        total_tensors=sum(len(m["tensors"]) for m in manifests),255        total_bytes=human_gb(sum(m["file_bytes"] for m in manifests)),256    )257    (out / "README.md").write_text(readme)258    print(259        json.dumps({k: v for k, v in record["totals"].items() if not isinstance(v, dict)}, indent=1)260    )261    print(json.dumps({"by_role": per_role, "by_type": per_type}, indent=1))262    if missing:263        print(f"WARNING missing shards: {len(missing)}")264    return 1 if (missing and args.check) else 0265 266 267if __name__ == "__main__":268    raise SystemExit(main())269