Weight codec

Portable .bnnpack container for packed weights.

encode_linear_state

encode_linear_state(
    weight: Tensor,
    bias: Tensor | None = None,
    *,
    alpha: Tensor | None = None,
    name: str = "linear",
    with_hash: bool = True,
) -> dict[str, Any]

Encode one Linear weight into a portable packed blob (binary XNOR path).

Accepts FP nn.Linear weights or BinaryLinear latents (signed via STE). Compression is exact 32× when in_features % 64 == 0 (no pad words); otherwise slightly lower due to uint64 padding.

Source code in bnn/codec/packfile.py

def encode_linear_state(
    weight: torch.Tensor,
    bias: torch.Tensor | None = None,
    *,
    alpha: torch.Tensor | None = None,
    name: str = "linear",
    with_hash: bool = True,
) -> dict[str, Any]:
    """Encode one Linear weight into a portable packed blob (binary XNOR path).

    Accepts FP ``nn.Linear`` weights or ``BinaryLinear`` latents (signed via STE).
    Compression is **exact 32×** when ``in_features % 64 == 0`` (no pad words);
    otherwise slightly lower due to uint64 padding.
    """
    w = weight.detach().float().cpu()
    out_f, in_f = int(w.shape[0]), int(w.shape[1])
    if alpha is None:
        alpha_t = w.abs().mean(dim=1).clamp(min=1e-4)
    else:
        alpha_t = alpha.detach().float().reshape(-1).cpu()
        if alpha_t.numel() == 1:
            alpha_t = alpha_t.expand(out_f)
        if alpha_t.numel() != out_f:
            raise ValueError(f"alpha length {alpha_t.numel()} != out_features {out_f}")
    w_pm1 = sign_pm1(w).numpy().astype(np.float32)
    packed, n = pack_binary_pm1(w_pm1, axis=1)
    assert n == in_f
    wp_i64 = torch.from_numpy(np.ascontiguousarray(packed).view(np.int64).copy())
    blob: dict[str, Any] = {
        "kind": KIND_BINARY_XNOR,
        "name": name,
        "in_features": in_f,
        "out_features": out_f,
        "n": int(n),
        "weight_packed_i64": wp_i64,
        "alpha": alpha_t.contiguous().clone(),
        "fp32_bytes": int(w.numel() * 4),
        "packed_bytes": int(packed.nbytes),
        "compression": float((w.numel() * 4) / max(packed.nbytes, 1)),
    }
    if bias is not None:
        blob["bias"] = bias.detach().float().cpu().contiguous().clone()
    else:
        blob["bias"] = None
    return _attach_hash(blob, wp_i64) if with_hash else blob

encode_model_linears

encode_model_linears(
    model: Module,
    *,
    skip_name_substr: tuple[str, ...] | None = None,
    min_in_features: int = 1,
    include_packed: bool = True,
    include_binary_linear: bool = True,
    include_fp_linear: bool = False,
    include_ternary: bool = False,
    include_conv: bool = False,
) -> dict[str, Any]

Encode modules into a layers dict.

Defaults favor the thesis wrap story: - Already-packed PackedBinaryXNORLinear (post-wrap FFN) - BinaryLinear STE modules - Not arbitrary FP nn.Linear (avoids silently binary-packing attn/embed/head)

Set include_fp_linear=True only when you intentionally want cold PTQ of FP Linears; then skip_name_substr defaults to HYBRID_FFN_SKIP. Opt in to ternary / Conv2d with include_ternary / include_conv.

Source code in bnn/codec/packfile.py

def encode_model_linears(
    model: nn.Module,
    *,
    skip_name_substr: tuple[str, ...] | None = None,
    min_in_features: int = 1,
    include_packed: bool = True,
    include_binary_linear: bool = True,
    include_fp_linear: bool = False,
    include_ternary: bool = False,
    include_conv: bool = False,
) -> dict[str, Any]:
    """Encode modules into a layers dict.

    Defaults favor the thesis wrap story:
    - Already-packed ``PackedBinaryXNORLinear`` (post-wrap FFN)
    - ``BinaryLinear`` STE modules
    - **Not** arbitrary FP ``nn.Linear`` (avoids silently binary-packing attn/embed/head)

    Set ``include_fp_linear=True`` only when you intentionally want cold PTQ of FP
    Linears; then ``skip_name_substr`` defaults to ``HYBRID_FFN_SKIP``.
    Opt in to ternary / Conv2d with ``include_ternary`` / ``include_conv``.
    """
    if skip_name_substr is None:
        skip_name_substr = HYBRID_FFN_SKIP if include_fp_linear else ()
    layers: dict[str, Any] = {}
    for name, mod in model.named_modules():
        lname = name.lower()
        if any(s.lower() in lname for s in skip_name_substr):
            continue
        if include_packed and isinstance(mod, PackedBinaryXNORLinear):
            if mod.in_features < min_in_features:
                continue
            layers[name] = encode_from_packed_module(mod, name=name)
            continue
        if include_ternary and isinstance(mod, TernaryWeightOnlyLinear):
            if mod.in_features < min_in_features:
                continue
            layers[name] = encode_from_ternary_module(mod, name=name)
            continue
        if include_conv and isinstance(mod, PackedBinaryConv2d):
            layers[name] = encode_from_packed_conv(mod, name=name)
            continue
        if include_binary_linear and isinstance(mod, BinaryLinear):
            if mod.weight.shape[1] < min_in_features:
                continue
            a = mod.alpha.detach() if hasattr(mod, "alpha") else None
            layers[name] = encode_linear_state(
                mod.weight,
                _optional_bias_tensor(mod.bias),
                alpha=a,
                name=name,
            )
            continue
        if include_fp_linear and isinstance(mod, nn.Linear):
            if mod.weight.shape[1] < min_in_features:
                continue
            layers[name] = encode_linear_state(
                mod.weight,
                _optional_bias_tensor(mod.bias),
                name=name,
            )
            continue
        if include_conv and isinstance(mod, nn.Conv2d) and mod.groups == 1:
            layers[name] = encode_conv_state(
                mod.weight,
                _optional_bias_tensor(mod.bias),
                stride=_symmetric_hw_int(mod.stride, name="stride"),
                padding=_symmetric_hw_int(mod.padding, name="padding"),
                name=name,
            )
    return layers

encode_file

encode_file(
    model: Module,
    path: Path | str,
    *,
    meta: dict[str, Any] | None = None,
    version: int = BNNPACK_VERSION,
    **kwargs: Any,
) -> Path

Source code in bnn/codec/packfile.py

def encode_file(
    model: nn.Module,
    path: Path | str,
    *,
    meta: dict[str, Any] | None = None,
    version: int = BNNPACK_VERSION,
    **kwargs: Any,
) -> Path:
    layers = encode_model_linears(model, **kwargs)
    return save_bnnpack(layers, path, meta=meta, version=version)

decode_file

decode_file(
    path: Path | str,
) -> tuple[dict[str, nn.Module], dict[str, Any]]

Load .bnnpack → mapping name → packed module + meta.

Source code in bnn/codec/packfile.py

def decode_file(
    path: Path | str,
) -> tuple[dict[str, nn.Module], dict[str, Any]]:
    """Load ``.bnnpack`` → mapping name → packed module + meta."""
    payload = load_bnnpack(path)
    modules = {name: decode_layer(blob) for name, blob in payload["layers"].items()}
    meta = payload.get("meta") or {}
    return modules, meta if isinstance(meta, dict) else {}

load_bnnpack

load_bnnpack(
    path: Path | str, *, verify_hashes: bool = True
) -> dict[str, Any]

Load .bnnpack with weights_only=True only (no unsafe pickle fallback).

Soft-warns when the path sits outside lab results/ / checkpoints/ / data/ (W10.T06). Never falls back to unsafe pickle.

When verify_hashes is True (default) and the file is v2+, recompute per-layer content_sha256 and raise if any mismatch.

Source code in bnn/codec/packfile.py

def load_bnnpack(
    path: Path | str,
    *,
    verify_hashes: bool = True,
) -> dict[str, Any]:
    """Load ``.bnnpack`` with ``weights_only=True`` only (no unsafe pickle fallback).

    Soft-warns when the path sits outside lab ``results/`` / ``checkpoints/`` /
    ``data/`` (W10.T06). Never falls back to unsafe pickle.

    When ``verify_hashes`` is True (default) and the file is v2+, recompute
    per-layer ``content_sha256`` and raise if any mismatch.
    """
    from ..paths import warn_untrusted_pack

    path = Path(path)
    if not path.is_file():
        raise FileNotFoundError(path)
    warn_untrusted_pack(path, kind=".bnnpack")
    try:
        payload = torch.load(path, map_location="cpu", weights_only=True)
    except Exception as exc:
        raise ValueError(
            f"{path}: failed weights_only load ({exc}). "
            "Refusing unsafe pickle fallback for .bnnpack — regenerate the pack "
            "with a current bnn encode, or use only trusted sources."
        ) from exc
    if not isinstance(payload, dict):
        raise ValueError(f"{path} is not a bnnpack dict")
    if payload.get("magic") != BNNPACK_MAGIC:
        raise ValueError(
            f"{path}: bad magic {payload.get('magic')!r}; expected {BNNPACK_MAGIC}"
        )
    ver = int(payload.get("version", 0))
    if ver not in SUPPORTED_VERSIONS:
        raise ValueError(
            f"{path}: unsupported version {payload.get('version')}; "
            f"supported={sorted(SUPPORTED_VERSIONS)}"
        )
    if "layers" not in payload or not isinstance(payload["layers"], dict):
        raise ValueError(f"{path}: missing layers")
    if verify_hashes and ver >= BNNPACK_VERSION_V2:
        bad = verify_layer_hashes(payload)
        if bad:
            raise ValueError(
                f"{path}: content_sha256 mismatch for layers {bad}; "
                "file may be corrupted or tampered"
            )
        # v2 writers must publish a container hashes map covering every layer.
        top = payload.get("hashes")
        layers = payload["layers"]
        if not isinstance(top, dict) or set(top) != set(layers):
            raise ValueError(
                f"{path}: v2 hashes map missing or incomplete "
                f"(hashes={None if not isinstance(top, dict) else sorted(top)}; "
                f"layers={sorted(layers)})"
            )
    return payload

save_bnnpack

save_bnnpack(
    layers: dict[str, Any],
    path: Path | str,
    *,
    meta: dict[str, Any] | None = None,
    version: int = BNNPACK_VERSION,
) -> Path

Source code in bnn/codec/packfile.py

def save_bnnpack(
    layers: dict[str, Any],
    path: Path | str,
    *,
    meta: dict[str, Any] | None = None,
    version: int = BNNPACK_VERSION,
) -> Path:
    if version not in SUPPORTED_VERSIONS:
        raise ValueError(
            f"unsupported bnnpack version {version}; want {sorted(SUPPORTED_VERSIONS)}"
        )
    path = Path(path)
    path.parent.mkdir(parents=True, exist_ok=True)
    if version >= BNNPACK_VERSION_V2:
        for blob in layers.values():
            if isinstance(blob, dict) and "content_sha256" not in blob:
                _attach_hash(blob)
    payload: dict[str, Any] = {
        "magic": BNNPACK_MAGIC,
        "version": int(version),
        "layers": layers,
        "meta": meta or {},
    }
    if version >= BNNPACK_VERSION_V2:
        payload["hashes"] = _layer_hashes(layers)
    torch.save(payload, path)
    return path

roundtrip_gemm_err

roundtrip_gemm_err(
    weight: Tensor, *, batch: int = 4, seed: int = 0
) -> dict[str, float]

Encode → decode → compare packed GEMM vs ±1 FP reference; expect err=0.

Source code in bnn/codec/packfile.py

def roundtrip_gemm_err(
    weight: torch.Tensor,
    *,
    batch: int = 4,
    seed: int = 0,
) -> dict[str, float]:
    """Encode → decode → compare packed GEMM vs ±1 FP reference; expect err=0."""
    from ..kernels.packed import native_kernel_available

    blob = encode_linear_state(weight)
    mod = decode_to_packed_linear(blob)
    err = packed_module_fp_err(mod, batch=batch, seed=seed)
    return {
        "max_abs_err": err,
        "compression": float(blob["compression"]),
        "native": float(native_kernel_available()),
        "packed_bytes": float(blob["packed_bytes"]),
        "fp32_bytes": float(blob["fp32_bytes"]),
    }