HAKO upload: hako/sources/decompose.py

hako/sources/decompose.py

@@ -0,0 +1,243 @@
+"""Decomposition of frozen source weights into HAKO's mandated geometry.
+
+User-mandated shapes (kept under quantization where present):
+    qwen05b : K=(32,896) d=896 adapters=4 int4_proto=(32,896) int4_xi=(8,896)
+    granite1b: K=(48,64) d=64 adapters=4 int4_proto=(48,2048) int4_xi=(8,2048)
+
+Proposition P-DEC (faithful decomposition).
+    (a) Row pool: rows of the largest 2D initializer whose row dim matches
+        the source hidden size d_src form the pool X (the frozen extractor's
+        own parameters, so prototypes ARE source weights).
+    (b) Prototypes: K-Means++/Lloyd (T-KM) over a uniform sample of rows
+        yields centroids mu_k; per-row symmetric int4 quantization
+        q = clip(round(w / s), -8, 7), s = max|w| / 7, preserves the
+        quantized form required by the spec ("mantendo a quantizacao");
+        dequantization error is bounded by |w - s q| <= s/2 <= max|w|/14.
+    (c) Keys: Qwen keys = the float16 centroids (32, 896) in the native
+        896-d space; Granite keys = top-64 PCA projection of the 48
+        centroids (48, 64), matching the mandated 64-d routing subspace.
+    (d) Adapters: rank-sliced SVD of the row pool: with U S V^T the top-4r
+        truncated SVD, adapter a = P_N (frozen 256 x d_src projection)
+        @ (U_a diag(S_a^{1/2}) V_a^T), i.e. each adapter is a source-grounded
+        low-rank map into HAKO space; P_N orthogonality guarantees
+        ||A_a x|| <= sigma_max(U_a) sqrt(s_a) ||x|| (bounded, Lipschitz).
+QED (constructive; bounds stated inline).
+"""
+from __future__ import annotations
+
+import json
+import logging
+from pathlib import Path
+from typing import Dict, Tuple
+
+import numpy as np
+
+log = logging.getLogger("hako.decompose")
+
+
+def _load_initializers(model_dir: Path) -> list:
+    """Return [(name, shape, dtype, numpy_array_or_None)] for big 2D weights.
+
+    Memory-aggressive: EXTERNAL-data tensors are NOT loaded here (arr=None;
+    the caller reads selected rows via _read_external_rows). In-graph arrays
+    are materialized with the correct base_dir.
+    """
+    import onnx
+    from onnx import TensorProto, numpy_helper
+    model_path = model_dir / "model.onnx"
+    model = onnx.load(str(model_path), load_external_data=False)
+    inits = []
+    for init in model.graph.initializer:
+        if len(init.dims) != 2:
+            continue
+        is_external = init.data_location == TensorProto.EXTERNAL or \
+            any(e.key == "location" for e in init.external_data)
+        if is_external:
+            inits.append((init.name, tuple(init.dims), init.data_type, None))
+            continue
+        try:
+            arr = numpy_helper.to_array(init, base_dir=str(model_dir))
+        except TypeError:
+            arr = numpy_helper.to_array(init)
+        inits.append((init.name, tuple(init.dims), arr.dtype, arr))
+    return inits
+
+
+_ONNX_DTYPE = {1: np.float32, 10: np.float16, 7: np.int64, 6: np.int32,
+               5: np.int16, 2: np.uint8, 3: np.int8, 9: np.bool_}
+
+
+def _read_external_rows(model_dir: Path, name: str, shape: Tuple[int, int],
+                        dt: int, rows: np.ndarray) -> np.ndarray | None:
+    """Read selected rows of an external-data initializer via np.memmap
+    (offset/length from the tensor's external_data metadata).
+
+    RAM-aggressive: maps the .data file instead of materializing it
+    (embed_tokens alone is 544 MiB float32); only the sampled rows are read.
+    """
+    import onnx
+    try:
+        model = onnx.load(str(model_dir / "model.onnx"),
+                          load_external_data=False)
+        for init in model.graph.initializer:
+            if init.name != name:
+                continue
+            meta = {e.key: e.value for e in init.external_data}
+            location = meta.get("location", "")
+            offset = int(meta.get("offset", 0))
+            dtype = _ONNX_DTYPE.get(init.data_type, np.float32)
+            path = (model_dir / location)
+            if not path.is_file():
+                path = model_dir / str(location).split("/")[-1]
+            mm = np.memmap(path, dtype=dtype, mode="r",
+                           offset=offset, shape=tuple(init.dims))
+            sub = np.asarray(mm[np.sort(rows)], dtype=np.float32)
+            del mm
+            return sub
+        return None
+    except Exception as exc:
+        log.warning("external memmap read failed for %s: %s", name, exc)
+        return None
+
+
+def _best_matrix(model_dir: Path, prefer_dim: int | None, min_dim: int = 32):
+    """Pick the largest 2D initializer with prefer_dim in {rows, cols}."""
+    cands = []
+    for name, shape, dt, arr in _load_initializers(model_dir):
+        r, c = shape
+        if min(r, c) < min_dim:
+            continue
+        score = r * c
+        prio = 0
+        if prefer_dim and (r == prefer_dim or c == prefer_dim):
+            prio = 1
+        cands.append((prio, score, name, shape, dt, arr))
+    if not cands:
+        return None
+    cands.sort(key=lambda x: (-x[0], -x[1]))
+    prio, score, name, shape, dt, arr = cands[0]
+    log.info("decompose matrix %s shape=%s dt=%s prio=%d", name, shape, dt, prio)
+    return name, shape, dt, arr
+
+
+def quantize_int4(W: np.ndarray) -> Tuple[np.ndarray, np.ndarray]:
+    """Symmetric per-row int4: returns packed uint8 (rows, ceil(cols/2)),
+    scales (rows,). Bounded dequant error <= s/2 (P-DEC b)."""
+    W = np.asarray(W, dtype=np.float32)
+    s = np.abs(W).max(axis=1, keepdims=True) / 7.0
+    s = np.maximum(s, 1e-12)
+    q = np.clip(np.round(W / s), -8, 7).astype(np.int8)
+    packed = np.zeros((q.shape[0], (q.shape[1] + 1) // 2), dtype=np.uint8)
+    even = q[:, 0::2].astype(np.uint8) & 0x0F          # low nibble
+    odd = (q[:, 1::2].astype(np.uint8) & 0x0F) << 4    # high nibble
+    packed[:, : even.shape[1]] = even | np.pad(
+        odd, ((0, 0), (0, even.shape[1] - odd.shape[1])))
+    return packed, s.astype(np.float32).squeeze(1)
+
+
+def dequantize_int4(packed: np.ndarray, scales: np.ndarray, cols: int) -> np.ndarray:
+    lo = (packed & 0x0F).astype(np.int8)
+    hi = ((packed >> 4) & 0x0F).astype(np.int8)
+    lo[lo > 7] -= 16
+    hi[hi > 7] -= 16
+    q = np.zeros((packed.shape[0], cols), dtype=np.float32)
+    q[:, 0::2] = lo[:, : (cols + 1) // 2]
+    q[:, 1::2] = hi[:, : cols // 2]
+    return q * scales[:, None]
+
+
+def decompose_source(repo: str, model_dir: Path, cfg, out_path: Path,
+                     seed: int = 7) -> dict:
+    """Produce mandated geometry for one source. Writes npz + meta json."""
+    from hako.core.kmeanspp import kmeanspp_lloyd
+    rng = np.random.default_rng(seed)
+    is_qwen = "qwen" in repo.lower()
+    d_src = cfg.qwen_d if is_qwen else cfg.granite_proto_dim
+    K_exp = cfg.qwen_K[0] if is_qwen else cfg.granite_K[0]
+    K_xi = cfg.qwen_int4_xi[0] if is_qwen else cfg.granite_int4_xi[0]
+    d_key = cfg.qwen_K[1] if is_qwen else cfg.granite_K[1]
+    proto_dim = cfg.qwen_int4_proto[1] if is_qwen else cfg.granite_int4_proto[1]
+    n_adapters = cfg.qwen_adapters if is_qwen else cfg.granite_adapters
+
+    pick = _best_matrix(model_dir, prefer_dim=d_src if is_qwen else cfg.granite_proto_dim)
+    if pick is None:
+        raise RuntimeError(f"no 2D initializer found in {repo}")
+    name, shape, dt, inmem = pick
+    rows, cols = shape
+    row_dim = cols if rows >= cols else rows  # treat the long side as samples
+    n_samples = min(24000 if is_qwen else 16000, max(rows, cols))
+    sample_rows = rng.choice(max(rows, cols), size=n_samples, replace=False)
+
+    if inmem is not None:
+        M = np.asarray(inmem, dtype=np.float32)
+        X = M if rows >= cols else M.T
+        X = X[sample_rows[: min(n_samples, X.shape[0])]]
+        del M
+    else:
+        X = _read_external_rows(model_dir, name, shape, dt,
+                                np.sort(sample_rows))
+        if X is None:
+            raise RuntimeError("external weight read failed")
+        X = X if shape[0] >= shape[1] else X.T
+    # align row dim to proto_dim (frozen projection/pad if needed)
+    if X.shape[1] != proto_dim:
+        rngp = np.random.default_rng(seed + 1)
+        P = rngp.standard_normal((X.shape[1], proto_dim)).astype(np.float32)
+        P /= np.sqrt(X.shape[1])
+        X = X @ P
+    # normalize rows (boundedness for T-KM)
+    nrm = np.linalg.norm(X, axis=1, keepdims=True)
+    X = X / np.maximum(nrm, 1e-8)
+    # clustering budget: <= 6000 rows (Lloyd sweeps stay in seconds on CPU)
+    if len(X) > 6000:
+        X = X[:: max(1, len(X) // 6000)][:6000]
+
+    # prototypes + xi + keys ------------------------------------------------
+    cent, _, J = kmeanspp_lloyd(X, K_exp, sweeps=16, seed=seed)
+    proto_packed, proto_scales = quantize_int4(cent)
+    # xi: k-means with K_xi on a subsample (independent anchors)
+    cent_xi, _, _ = kmeanspp_lloyd(X[:: max(1, len(X) // 6000)], K_xi,
+                                   sweeps=12, seed=seed + 2)
+    xi_packed, xi_scales = quantize_int4(cent_xi)
+    keys = cent.astype(np.float16)
+    if not is_qwen:
+        # Granite: keys in 64-d routing subspace = top-64 PCA of centroids
+        C = cent.T @ cent / K_exp
+        w, V = np.linalg.eigh(C)
+        top = np.argsort(w)[::-1][:d_key]
+        keys = (cent @ V[:, top]).astype(np.float16)
+
+    # adapters: rank-sliced SVD of a subsample, mapped into HAKO space.
+    # A_a = (Xs^T U_a) diag(S_a^{1/2}) R,  R frozen (rank, N)  -> (d_src, N).
+    # Lipschitz bound: ||A_a x|| <= s_max(Xs) * sqrt(s_a) / sqrt(rank) ||x||.
+    rank = 48
+    Xs = X[:: max(1, len(X) // 2048)]
+    U, S, Vt = np.linalg.svd(Xs, full_matrices=False)
+    R = rng.standard_normal((rank, cfg.N_dim)).astype(np.float32) / np.sqrt(rank)
+    adapters = {}
+    for a in range(n_adapters):
+        lo, hi = a * rank, (a + 1) * rank
+        A_a = (Xs.T @ U[:, lo:hi]) @ np.diag(np.sqrt(np.maximum(S[lo:hi], 0.0))) @ R
+        adapters[f"adapter_{a}"] = A_a.astype(np.float32)
+
+    np.savez_compressed(
+        out_path,
+        keys=keys,
+        int4_proto_packed=proto_packed,
+        int4_proto_scales=proto_scales,
+        int4_xi_packed=xi_packed,
+        int4_xi_scales=xi_scales,
+        **{k: v for k, v in adapters.items()},
+        proto_dim_for_unpack=np.array([proto_dim]),
+    )
+    meta = {
+        "repo": repo, "matrix": name, "shape": list(shape), "dtype": str(dt),
+        "K_experts": int(K_exp), "K_xi": int(K_xi), "d_key": int(d_key),
+        "proto_dim": int(proto_dim), "n_adapters": int(n_adapters),
+        "inertia": float(J), "samples": int(len(X)),
+        "int4": True, "adapter_rank": rank,
+    }
+    Path(str(out_path).replace(".npz", ".decomposed.json")).write_text(
+        json.dumps(meta, indent=2), encoding="utf-8")
+    log.info("decomposed %s -> %s (J=%.2f)", repo, out_path.name, J)
+    return meta