| |
| """Decode the packed Mach-2 Additive Medium checkpoint into a bf16 Hugging Face checkpoint. |
| |
| python decode.py materialize --packed-dir . --out /path/ckpt [--workers 8] [--no-base] |
| """ |
| import argparse |
| import hashlib |
| import json |
| import math |
| import os |
|
|
| import numpy as np |
|
|
| HIDDEN, INTER, NEXP, NLAYERS = 2560, 640, 512, 48 |
| BASE_REPO, BASE_REVISION = "Qwen/Qwen3.8-Flash-Next", "de4b8e4d43b917e7706784d8bb445c9af86a3540" |
| FORMAT = "fn_packed_v2" |
|
|
|
|
| def _hmat(q_elems, chi, sub): |
| n = len(q_elems) |
| Q = np.array([[chi(sub(q_elems[i], q_elems[j])) for j in range(n)] for i in range(n)], np.float64) |
| C = np.zeros((n + 1, n + 1)) |
| C[0, 1:] = 1.0 |
| C[1:, 0] = 1.0 |
| C[1:, 1:] = Q |
| A = np.array([[1.0, 1.0], [1.0, -1.0]]) |
| B = np.array([[1.0, -1.0], [-1.0, -1.0]]) |
| H = np.kron(C, A) + np.kron(np.eye(n + 1), B) |
| r = 2 * (n + 1) |
| assert np.array_equal(H, H.T) and np.array_equal(H @ H, r * np.eye(r)) |
| return H.astype(np.float32) |
|
|
|
|
| def hadamard_12(): |
| chi5 = {0: 0, 1: 1, 2: -1, 3: -1, 4: 1} |
| return _hmat(list(range(5)), lambda e: chi5[e], lambda u, v: (u - v) % 5) |
|
|
|
|
| def hadamard_20(): |
| els = [(a, b) for a in range(3) for b in range(3)] |
|
|
| def mul(u, v): |
| (a, b), (c, d) = u, v |
| return ((a * c - b * d) % 3, (a * d + b * c) % 3) |
| squares = {mul(e, e) for e in els if e != (0, 0)} |
| return _hmat(els, lambda e: 0 if e == (0, 0) else (1 if e in squares else -1), |
| lambda u, v: ((u[0] - v[0]) % 3, (u[1] - v[1]) % 3)) |
|
|
|
|
| _HS = { |
| 12: ["+-----------", "++-+---+++-+", "+++-+---+++-", "+-++-+---+++", "++-++-+---++", "+++-++-+---+", |
| "++++-++-+---", "+-+++-++-+--", "+--+++-++-+-", "+---+++-++-+", "++---+++-++-", "+-+---+++-++"], |
| 20: ["+-------------------", "++-++----+-+-++++--+", "+++-++----+-+-++++--", "+-++-++----+-+-++++-", |
| "+--++-++----+-+-++++", "++--++-++----+-+-+++", "+++--++-++----+-+-++", "++++--++-++----+-+-+", |
| "+++++--++-++----+-+-", "+-++++--++-++----+-+", "++-++++--++-++----+-", "+-+-++++--++-++----+", |
| "++-+-++++--++-++----", "+-+-+-++++--++-++---", "+--+-+-++++--++-++--", "+---+-+-++++--++-++-", |
| "+----+-+-++++--++-++", "++----+-+-++++--++-+", "+++----+-+-++++--++-", "+-++----+-+-++++--++"], |
| } |
| _HR = {} |
|
|
|
|
| def _hr(kind, radix): |
| if (kind, radix) not in _HR: |
| if kind == "spine": |
| H = np.array([[1.0 if c == "+" else -1.0 for c in r] for r in _HS[radix]], np.float32) |
| assert np.array_equal(H @ H.T, radix * np.eye(radix, dtype=np.float32)) |
| else: |
| H = hadamard_12() if radix == 12 else hadamard_20() |
| _HR[(kind, radix)] = H |
| return _HR[(kind, radix)] |
|
|
|
|
| def hadamard(x, kind="expert"): |
| x = np.asarray(x, np.float32) |
| shape, N = x.shape, x.shape[-1] |
| if N & (N - 1): |
| radix = 12 if (N % 12 == 0 and (N // 12) & (N // 12 - 1) == 0) else 20 |
| M = N // radix |
| assert N % radix == 0 and M & (M - 1) == 0, f"Hadamard needs 2^k, 12*2^k or 20*2^k, got {N}" |
| if kind == "spine": |
| xb = _butterflies(x.reshape(-1, radix, M)) if M > 1 else x.reshape(-1, radix, 1) |
| return (np.matmul(_hr(kind, radix), xb) / np.float32(math.sqrt(N))).reshape(shape) |
| xb = hadamard(x.reshape(-1, radix, M)) if M > 1 else x.reshape(-1, radix, 1) |
| return (np.matmul(_hr(kind, radix), xb) / np.float32(math.sqrt(radix))).reshape(shape) |
| return (_butterflies(x) / np.float32(math.sqrt(N))).reshape(shape) |
|
|
|
|
| def _butterflies(x): |
| N = x.shape[-1] |
| cur = np.ascontiguousarray(x, np.float32).reshape(-1, N) |
| span = 1 |
| while span < N: |
| blk = cur.reshape(-1, N // (2 * span), 2, span) |
| cur = np.stack([blk[:, :, 0] + blk[:, :, 1], blk[:, :, 0] - blk[:, :, 1]], 2).reshape(-1, N) |
| span *= 2 |
| return cur.reshape(x.shape) |
|
|
|
|
| def bf16_round(x): |
| u = np.ascontiguousarray(x, np.float32).view(np.uint32) |
| return ((u + np.uint32(0x7FFF) + ((u >> np.uint32(16)) & np.uint32(1))) & np.uint32(0xFFFF0000)).view(np.float32) |
|
|
|
|
| def _read(path): |
| from safetensors import safe_open |
| with safe_open(path, framework="np") as fh: |
| return {k: fh.get_tensor(k) for k in fh.keys()}, dict(fh.metadata() or {}) |
|
|
|
|
| EX_V, EX_L, EX_TD = 4, 16, 16 |
| EX_SHAPES = {"gate": (INTER, HIDDEN), "up": (INTER, HIDDEN), "down": (HIDDEN, INTER)} |
|
|
|
|
| def _ex_states(words, K4): |
| words = np.asarray(words).view(np.uint16).astype(np.int64) |
| rows, nstep, step = words.shape[0], EX_TD * EX_TD // EX_V, int(K4) |
| bits = ((words[:, :, None] >> np.arange(15, -1, -1)) & 1).reshape(rows, -1)[:, :step * nstep] |
| bits = np.concatenate([bits, bits[:, :EX_L - step]], axis=1) |
| w = 1 << np.arange(EX_L - 1, -1, -1, dtype=np.int64) |
| idx = np.arange(nstep)[:, None] * step + np.arange(EX_L)[None, :] |
| return bits[:, idx] @ w |
|
|
|
|
| def wave_map(Mb, Nb): |
| starts = [(Mb - i - 1, Nb - 1) for i in range(Mb)] + [(0, Nb - i - 1) for i in range(Nb)] |
| idx = np.zeros((Mb, Nb), dtype=np.int64) |
| for w, (jm, jn) in enumerate(starts): |
| while 0 <= jm < Mb and 0 <= jn < Nb: |
| idx[jm, jn] = w |
| jm, jn = jm + 1, jn - 1 |
| return idx |
|
|
|
|
| def sign_seed(layer, expert, proj, side, ns="canon"): |
| return int.from_bytes(hashlib.sha256(f"{ns}|L{layer}|e{expert}|{proj}|{side}".encode()).digest()[:7], "big") |
|
|
|
|
| _PM0, _PM1 = np.uint64(0xD2511F53), np.uint64(0xCD9E8D57) |
| _PW0, _PW1, _PMASK = np.uint64(0x9E3779B9), np.uint64(0xBB67AE85), np.uint64(0xFFFFFFFF) |
|
|
|
|
| def expert_signs(seeds, dim): |
| E = len(seeds) |
| c2 = np.tile(np.arange(dim, dtype=np.uint64), E) |
| k0 = np.repeat(np.array([s & 0xFFFFFFFF for s in seeds], np.uint64), dim) |
| k1 = np.repeat(np.array([s >> 32 for s in seeds], np.uint64), dim) |
| c0 = c1 = c3 = np.zeros(E * dim, np.uint64) |
| for r in range(10): |
| if r: |
| k0, k1 = (k0 + _PW0) & _PMASK, (k1 + _PW1) & _PMASK |
| p0, p1 = _PM0 * c0, _PM1 * c2 |
| c0, c1, c2, c3 = (p1 >> np.uint64(32)) ^ c1 ^ k0, p1 & _PMASK, (p0 >> np.uint64(32)) ^ c3 ^ k1, p0 & _PMASK |
| inv2pi = np.float32(2 * np.pi / 2 ** 32) |
| v = (c1.astype(np.float32) * inv2pi + inv2pi / np.float32(2)).astype(np.float32) |
| u = (c0.astype(np.float32) * np.float32(2.0 ** -32) + np.float32(2.0 ** -33)).astype(np.float32) |
| return np.where((v <= np.float32(np.pi)) | (u == np.float32(1.0)), 1.0, -1.0).astype(np.float32).reshape(E, dim) |
|
|
|
|
| def _lut(cb, K4): |
| return np.asarray(cb.get(f"lut.k{K4}", cb["lut"]), np.float32) |
|
|
|
|
| def decode_expert(t, cb, proj, e): |
| m, n = EX_SHAPES[proj] |
| K4 = int(t[f"{proj}.rate_k4"][e]) |
| ex = np.asarray(t[f"k{K4}.{proj}.experts"]) |
| row = int(np.searchsorted(ex, e)) |
| assert row < ex.size and ex[row] == e, (proj, e, K4) |
| states = _ex_states(t[f"k{K4}.{proj}.trellis"][row], K4) |
| Mb, Nb = m // EX_TD, n // EX_TD |
| unit = _lut(cb, K4)[states].reshape(Mb, Nb, EX_TD, EX_TD).transpose(0, 2, 1, 3).reshape(m, n) |
| g = np.asarray(t[f"{proj}.wave_gamma"][e], np.float32)[wave_map(Mb, Nb)] |
| unit = (unit.reshape(Mb, EX_TD, Nb, EX_TD) * g[:, None, :, None]).reshape(m, n) |
| unit = unit * np.float32(t[f"{proj}.wscale"][e]) |
| rows = hadamard(unit) * t[f"{proj}.SU"][e] |
| cols = hadamard(np.ascontiguousarray(rows.T)) * t[f"{proj}.SV"][e] |
| return np.ascontiguousarray(cols.T) |
|
|
|
|
| def load_expert_layer(packed_dir, layer): |
| t, meta = _read(os.path.join(packed_dir, "experts", f"L{layer:02d}.safetensors")) |
| assert meta.get("format") == FORMAT, f"L{layer}: not an {FORMAT} expert shard ({meta.get('format')})" |
| cb = _read(os.path.join(packed_dir, "experts", "codebook.safetensors"))[0] |
| return with_signs(t, layer), cb |
|
|
|
|
| def with_signs(t, layer): |
| for p, (m, n) in EX_SHAPES.items(): |
| if f"{p}.SU" not in t: |
| t[f"{p}.SU"] = expert_signs([sign_seed(layer, e, p, "SU") for e in range(NEXP)], n) |
| t[f"{p}.SV"] = expert_signs([sign_seed(layer, e, p, "SV") for e in range(NEXP)], m) |
| return t |
|
|
|
|
| def decode_expert_layer(packed_dir, layer, experts=None): |
| t, cb = load_expert_layer(packed_dir, layer) |
| experts = list(range(NEXP)) if experts is None else list(experts) |
| gu = np.empty((len(experts), 2 * INTER, HIDDEN), np.float32) |
| dn = np.empty((len(experts), HIDDEN, INTER), np.float32) |
| for i, e in enumerate(experts): |
| gu[i, :INTER] = decode_expert(t, cb, "gate", e) |
| gu[i, INTER:] = decode_expert(t, cb, "up", e) |
| dn[i] = decode_expert(t, cb, "down", e) |
| return {"gate_up_proj": gu, "down_proj": dn} |
|
|
|
|
| NE_K, NE_L, NE_V, NE_TLUT_BITS, NE_TD = 4, 16, 2, 9, 16 |
| _FULL_LUT = {} |
|
|
|
|
| def _ne_full_lut(tlut): |
| key = np.asarray(tlut).tobytes() |
| if key not in _FULL_LUT: |
| small = np.asarray(tlut, np.float32) |
| s = np.arange(1 << NE_L, dtype=np.int64) |
| p = s * (s + 1) |
| row = (p >> (16 - NE_TLUT_BITS - 1)) & ((1 << NE_TLUT_BITS) - 1) |
| table = small[row].copy() |
| table[:, 0] *= (1 - ((p >> 15) & 1) * 2).astype(np.float32) |
| _FULL_LUT[key] = table |
| return _FULL_LUT[key] |
|
|
|
|
| def _ne_states(words): |
| words = np.asarray(words).view(np.uint16).astype(np.int64) |
| rows, T = words.shape[0], NE_TD * NE_TD |
| step, nstep = NE_K * NE_V, T // NE_V |
| bits = ((words[:, :, None] >> np.arange(15, -1, -1)) & 1).reshape(rows, -1)[:, :T * NE_K] |
| bits = np.concatenate([bits, bits[:, :NE_L - step]], axis=1) |
| w = 1 << np.arange(NE_L - 1, -1, -1, dtype=np.int64) |
| idx = np.arange(nstep)[:, None] * step + np.arange(NE_L)[None, :] |
| return bits[:, idx] @ w |
|
|
|
|
| def decode_ne_tensor(t, name, m, n, tlut): |
| states = _ne_states(t[f"{name}|trellis"]) |
| vals = _ne_full_lut(tlut)[states] |
| unit = np.ascontiguousarray(vals.reshape(m // NE_TD, n // NE_TD, NE_TD, NE_TD).transpose(0, 2, 1, 3)).reshape(m, n) |
| unit = unit * np.float32(np.asarray(t[f"{name}|Wscale"]).reshape(-1)[0]) |
| su, sv = np.asarray(t[f"{name}|SU"]), np.asarray(t[f"{name}|SV"]) |
| rows = hadamard(unit, "spine") * np.sign(su).astype(np.float32) |
| cols = hadamard(np.ascontiguousarray(rows.T), "spine") * np.sign(sv).astype(np.float32) |
| w = np.ascontiguousarray(cols.T) |
| mx = np.asarray(t[f"{name}|rc_max"], np.float32) |
| r, c = rc_grid(mx[0], np.abs(sv.astype(np.int32))), rc_grid(mx[1], np.abs(su.astype(np.int32))) |
| return (r[:, None] * bf16_round(w)) * c[None, :] |
|
|
|
|
| def rc_grid(mx, k): |
| return (np.float32(mx) * k.astype(np.float32)).astype(np.float32) * np.float32(1.0 / 127.0) |
|
|
|
|
| def decode_ne_shard(packed_dir, layer): |
| t, meta = _read(os.path.join(packed_dir, "ne", f"L{layer:02d}.safetensors")) |
| tlut = _read(os.path.join(packed_dir, "ne", "tlut.safetensors"))[0]["tlut"] |
| dims = json.loads(meta["dims"]) |
| return {name: decode_ne_tensor(t, name, d[2], d[3], tlut)[:d[0], :d[1]] for name, d in dims.items()} |
|
|
|
|
| def unpack_int5(qp, n): |
| m = qp.shape[0] |
| by = np.asarray(qp).reshape(m, n // 8, 5) |
| full = np.zeros((m, n // 8, 8), dtype=np.uint8) |
| full[:, :, :5] = by |
| word = full.reshape(m, n).view("<u8").reshape(m, n // 8) |
| out = np.zeros((m, n // 8, 8), dtype=np.int8) |
| for i in range(8): |
| out[:, :, i] = ((word >> np.uint64(5 * i)) & np.uint64(31)).astype(np.int8) - 16 |
| return out.reshape(m, n) |
|
|
|
|
| def decode_head(packed_dir): |
| d = os.path.join(packed_dir, "head") |
| parts = [] |
| for f in sorted(x for x in os.listdir(d) if x.startswith("head_c") and x.endswith(".safetensors")): |
| t, meta = _read(os.path.join(d, f)) |
| parts += [(int(name.split(":")[1]), m0, n0, int(meta.get("group", 64)), t, name) |
| for name, (m0, n0) in json.loads(meta["dims"]).items()] |
| parts.sort(key=lambda x: x[0]) |
| out = np.empty((sum(p[1] for p in parts), parts[0][2]), np.float32) |
| for r0, m0, n0, g, t, name in parts: |
| q = unpack_int5(t[f"{name}|qp"], n0).astype(np.float32) |
| out[r0:r0 + m0] = q * np.repeat(np.asarray(t[f"{name}|gscale"], np.float32), g, axis=1)[:, :n0] |
| return out |
|
|
|
|
| def decode_embed(packed_dir, bits=4, rows_per=16384): |
| t, _ = _read(os.path.join(packed_dir, "ne", f"embed_int{bits}.safetensors")) |
| rows, ng = t["mn"].shape |
| hid = t["q_packed"].shape[1] * 8 // bits |
| out = np.empty((rows, hid), np.float32) |
| for r0 in range(0, rows, rows_per): |
| sl = slice(r0, min(r0 + rows_per, rows)) |
| b = np.unpackbits(t["q_packed"][sl], axis=1, count=hid * bits).reshape(-1, hid, bits) |
| q = np.zeros(b.shape[:2], np.uint8) |
| for j in range(bits): |
| q = (q << 1) | b[..., j] |
| mn = t["mn"][sl].astype(np.float32)[..., None] |
| mx = t["mx"][sl].astype(np.float32)[..., None] |
| step = np.maximum(mx - mn, np.float32(1e-8)) * np.float32(1.0 / (2 ** bits - 1)) |
| out[sl] = (mn + q.reshape(-1, ng, hid // ng).astype(np.float32) * step).reshape(-1, hid) |
| return out |
|
|
|
|
| def decode_int8_rows(packed_dir): |
| from safetensors import safe_open |
| out = {} |
| with safe_open(os.path.join(packed_dir, "ne", "int8_rows.safetensors"), framework="pt") as fh: |
| names = sorted({k.split("|")[0] for k in fh.keys()}) |
| for name in names: |
| q = fh.get_tensor(f"{name}|q").numpy().astype(np.float32) |
| s = fh.get_tensor(f"{name}|scale").float().numpy() |
| out[name] = q * s |
| return out |
|
|
|
|
| def _bf16(a): |
| import torch |
| return torch.from_numpy(np.ascontiguousarray(a, np.float32)).to(torch.bfloat16) |
|
|
|
|
| def _expert_layer_job(args): |
| packed_dir, out, L = args |
| import torch |
| from safetensors.torch import save_file |
| t, cb = load_expert_layer(packed_dir, L) |
| gu = torch.empty((NEXP, 2 * INTER, HIDDEN), dtype=torch.bfloat16) |
| dn = torch.empty((NEXP, HIDDEN, INTER), dtype=torch.bfloat16) |
| for e in range(NEXP): |
| gu[e, :INTER] = _bf16(decode_expert(t, cb, "gate", e)) |
| gu[e, INTER:] = _bf16(decode_expert(t, cb, "up", e)) |
| dn[e] = _bf16(decode_expert(t, cb, "down", e)) |
| p = f"model.language_model.layers.{L}.mlp.experts." |
| fn = f"experts-L{L:02d}.safetensors" |
| save_file({p + "gate_up_proj": gu, p + "down_proj": dn}, os.path.join(out, fn), metadata={"format": "pt"}) |
| return {p + "gate_up_proj": fn, p + "down_proj": fn} |
|
|
|
|
| def _ne_layer_job(args): |
| packed_dir, out, L = args |
| from safetensors.torch import save_file |
| dec = decode_ne_shard(packed_dir, L) |
| fn = f"spine-L{L:02d}.safetensors" |
| save_file({k: _bf16(v) for k, v in dec.items()}, os.path.join(out, fn), metadata={"format": "pt"}) |
| return {k: fn for k in dec} |
|
|
|
|
| def materialize(packed_dir, out, workers=4, base=True): |
| import shutil |
| from concurrent.futures import ProcessPoolExecutor |
| from safetensors import safe_open |
| from safetensors.torch import save_file |
| os.makedirs(out, exist_ok=True) |
| root = os.path.dirname(os.path.abspath(packed_dir.rstrip("/"))) if os.path.basename(packed_dir.rstrip("/")) == "packed" \ |
| else packed_dir |
| pk = os.path.join(root, "packed") |
| wm = {} |
| with ProcessPoolExecutor(workers) as ex: |
| for r in ex.map(_ne_layer_job, [(pk, out, L) for L in range(NLAYERS)]): |
| wm.update(r) |
| for r in ex.map(_expert_layer_job, [(pk, out, L) for L in range(NLAYERS)]): |
| wm.update(r) |
| save_file({"lm_head.weight": _bf16(decode_head(pk))}, os.path.join(out, "head.safetensors"), metadata={"format": "pt"}) |
| wm["lm_head.weight"] = "head.safetensors" |
| emb = "model.language_model.embed_tokens.weight" |
| save_file({emb: _bf16(decode_embed(pk))}, os.path.join(out, "embed.safetensors"), metadata={"format": "pt"}) |
| wm[emb] = "embed.safetensors" |
| i8 = decode_int8_rows(pk) |
| save_file({k: _bf16(v) for k, v in i8.items()}, os.path.join(out, "int8rows.safetensors"), metadata={"format": "pt"}) |
| wm.update({k: "int8rows.safetensors" for k in i8}) |
| for f in ["extras.safetensors"] + sorted(os.path.join("packed", "table", x) for x in os.listdir(os.path.join(pk, "table")) |
| if x.endswith(".safetensors")): |
| dst = os.path.basename(f) |
| shutil.copyfile(os.path.join(root, f), os.path.join(out, dst)) |
| with safe_open(os.path.join(out, dst), framework="pt") as fh: |
| wm.update({k: dst for k in fh.keys()}) |
| if base: |
| from huggingface_hub import hf_hub_download |
| idx = json.load(open(hf_hub_download(BASE_REPO, "model.safetensors.index.json", revision=BASE_REVISION)))["weight_map"] |
| want = {k: f for k, f in idx.items() if k.startswith("mtp.") or k.startswith("model.visual.")} |
| for f in sorted(set(want.values())): |
| src = hf_hub_download(BASE_REPO, f, revision=BASE_REVISION) |
| with safe_open(src, framework="pt") as fh: |
| ks = [k for k in fh.keys() if k in want] |
| save_file({k: fh.get_tensor(k) for k in ks}, os.path.join(out, f"base-{f}"), metadata={"format": "pt"}) |
| wm.update({k: f"base-{f}" for k in ks}) |
| for f in os.listdir(root): |
| if f.endswith((".json", ".jinja", ".txt")) and f not in ("MANIFEST.json", "model.safetensors.index.json"): |
| shutil.copyfile(os.path.join(root, f), os.path.join(out, f)) |
| json.dump({"metadata": {}, "weight_map": dict(sorted(wm.items()))}, open(os.path.join(out, "model.safetensors.index.json"), "w"), |
| indent=1) |
| print(f"MATERIALIZED {len(wm)} tensors -> {out}", flush=True) |
|
|
|
|
| if __name__ == "__main__": |
| ap = argparse.ArgumentParser(description=__doc__.split("\n\n")[0]) |
| sub = ap.add_subparsers(dest="cmd", required=True) |
| m = sub.add_parser("materialize") |
| m.add_argument("--packed-dir", default=".") |
| m.add_argument("--out", required=True) |
| m.add_argument("--workers", type=int, default=4) |
| m.add_argument("--no-base", action="store_true", help="text-only: skip MTP and vision weights") |
| a = ap.parse_args() |
| materialize(a.packed_dir, a.out, a.workers, not a.no_base) |
|
|
