Wrapping models

Replace nn.Linear / nn.Conv2d with packed equivalents.

wrap_model

wrap_model(
    model: Module,
    mode: WrapMode | str | None = None,
    *,
    policy: WrapPolicy = "hybrid_ffn",
    skip_name_substr: Iterable[str] | None = None,
    min_in_features: int = 64,
    min_out_features: int = 0,
    skip_attn: bool = True,
    calib: CalibConfig | None = None,
    inplace: bool = True,
    accuracy_first: bool = False,
    exclude_exact: Iterable[str] | None = None,
    force_narrow: bool = False,
    fuse_bn: bool = False,
    drop_in_threshold: float = 0.85,
) -> tuple[nn.Module, WrapReport]

Product wrap API with hybrid / aggressive / ternary_wo / auto policies.

mode=None means unspecified (default binary_xnor, or recommender when policy='auto' / mode='auto').

exclude_exact: full dotted module names to never wrap (sensitivity). force_narrow: allow binary_xnor on shapes guardrails would refuse. fuse_bn: fold Linear+BN1d / BiReal BN before packing (W3.T09).

Source code in bnn/wrap/api.py

def wrap_model(
    model: nn.Module,
    mode: WrapMode | str | None = None,
    *,
    policy: WrapPolicy = "hybrid_ffn",
    skip_name_substr: Iterable[str] | None = None,
    min_in_features: int = 64,
    min_out_features: int = 0,
    skip_attn: bool = True,
    calib: CalibConfig | None = None,
    inplace: bool = True,
    accuracy_first: bool = False,
    exclude_exact: Iterable[str] | None = None,
    force_narrow: bool = False,
    fuse_bn: bool = False,
    drop_in_threshold: float = 0.85,
) -> tuple[nn.Module, WrapReport]:
    """Product wrap API with hybrid / aggressive / ternary_wo / auto policies.

    ``mode=None`` means unspecified (default binary_xnor, or recommender when
    ``policy='auto'`` / ``mode='auto'``).

    ``exclude_exact``: full dotted module names to never wrap (sensitivity).
    ``force_narrow``: allow binary_xnor on shapes guardrails would refuse.
    ``fuse_bn``: fold Linear+BN1d / BiReal BN before packing (W3.T09).
    """
    import copy as _copy

    from .guardrails import check_linear_wrap_guardrails

    hw = detect_hardware()
    resolved_mode: WrapMode
    # W3.T03 — policy_reason is always a non-empty string
    policy_reason: str

    if policy == "auto" or mode == "auto":
        decision = recommend_wrap_policy(None, hw, accuracy_first=accuracy_first)
        if policy == "auto":
            policy = decision.policy
        if mode is None or mode == "auto":
            resolved_mode = decision.mode
        else:
            resolved_mode = mode  # type: ignore[assignment]
        policy_reason = decision.reason or (
            f"auto → policy={policy} mode={resolved_mode}"
        )
        if min_in_features == 64:
            min_in_features = decision.min_in_features
        if min_out_features == 0:
            min_out_features = decision.min_out_features
        skip_attn = decision.skip_attn
    elif policy == "ternary_wo":
        resolved_mode = "ternary_weight_only"
        policy_reason = "policy=ternary_wo (accurate-first weight-only)"
    else:
        resolved_mode = (mode or "binary_xnor")  # type: ignore[assignment]
        policy_reason = f"policy={policy} mode={resolved_mode}"

    if not inplace:
        model = _copy.deepcopy(model)

    fuse_payload: dict | None = None
    if fuse_bn:
        from .fuse import fuse_bn_for_wrap_

        fuse_payload = fuse_bn_for_wrap_(model).to_dict()

    to_replace, skipped = select_linears(
        model,
        policy=policy,
        skip_name_substr=skip_name_substr,
        min_in_features=min_in_features,
        min_out_features=min_out_features,
        skip_attn=skip_attn,
        exclude_exact=exclude_exact,
    )

    report = WrapReport(
        mode=resolved_mode,
        policy=policy,
        skipped=skipped,
        native_kernel=native_kernel_available(),
        calib_method=(calib.method if calib else "absmean"),
        policy_reason=policy_reason,
        fuse=fuse_payload,
        # W3.T02 — effectiveness always present (stub until attach_effectiveness)
        effectiveness=unmeasured_effectiveness(
            drop_in_threshold=drop_in_threshold
        ).to_dict(),
        drop_in_ok=False,
    )

    for name, lin in to_replace:
        verdict = check_linear_wrap_guardrails(
            lin, mode=str(resolved_mode), force=force_narrow
        )
        if not verdict.ok:
            report.skipped.append(f"{name} ({verdict.code}: {verdict.message})")
            continue
        fp_bytes = int(lin.weight.numel() * 4)
        new, packed_b = _build_wrapped(lin, resolved_mode, calib=calib)
        report.replaced.append(name)
        report.fp32_weight_bytes_replaced += fp_bytes
        report.packed_weight_bytes += packed_b
        if resolved_mode == "binary_xnor":
            report.native_kernel = getattr(new, "uses_native", False)
        _set_module(model, name, new)

    return model, report

wrap_linear_modules

wrap_linear_modules(
    model: Module,
    mode: WrapMode = "binary_xnor",
    *,
    skip_name_substr: Iterable[str] = DEFAULT_SKIP,
    min_in_features: int = 64,
    min_out_features: int = 0,
    calib: CalibConfig | None = None,
    inplace: bool = True,
) -> tuple[nn.Module, WrapReport]

Legacy API: skip-list based wrap (still used by demos/tests).

Source code in bnn/wrap/api.py

def wrap_linear_modules(
    model: nn.Module,
    mode: WrapMode = "binary_xnor",
    *,
    skip_name_substr: Iterable[str] = DEFAULT_SKIP,
    min_in_features: int = 64,
    min_out_features: int = 0,
    calib: CalibConfig | None = None,
    inplace: bool = True,
) -> tuple[nn.Module, WrapReport]:
    """Legacy API: skip-list based wrap (still used by demos/tests)."""
    return wrap_model(
        model,
        mode,
        policy="default",
        skip_name_substr=skip_name_substr,
        min_in_features=min_in_features,
        min_out_features=min_out_features,
        calib=calib,
        inplace=inplace,
    )

model_param_bytes

model_param_bytes(model: Module) -> dict

Source code in bnn/wrap/api.py

def model_param_bytes(model: nn.Module) -> dict:
    p_bytes = sum(p.numel() * p.element_size() for p in model.parameters())
    b_bytes = sum(b.numel() * b.element_size() for b in model.buffers())
    return {"param_bytes": p_bytes, "buffer_bytes": b_bytes, "total_bytes": p_bytes + b_bytes}

Packed modules

PackedBinaryXNORLinear

Bases: Module

Inference Linear: packed ±1 weights + signed activations → XNOR GEMM.

Weights are packed once at construction and cached on the module.

Source code in bnn/wrap/packed_linear.py

class PackedBinaryXNORLinear(nn.Module):
    """Inference Linear: packed ±1 weights + signed activations → XNOR GEMM.

    Weights are packed **once** at construction and cached on the module.
    """

    # nn.Module.__getattr__ is typed as returning Tensor | Module, so buffers
    # must be declared for a type checker to see them as tensors.
    weight_packed_i64: Tensor
    alpha: Tensor
    bias: Tensor | None
    _wp_np: np.ndarray
    _alpha_np: np.ndarray
    _bias_np: np.ndarray | None

    def __init__(
        self,
        weight: Tensor,
        bias: Tensor | None = None,
        *,
        alpha: Tensor | None = None,
        calib: CalibConfig | None = None,
    ):
        super().__init__()
        out_f, in_f = weight.shape
        self.in_features = in_f
        self.out_features = out_f
        w_pm1 = sign_pm1(weight.detach().float().cpu()).numpy().astype(np.float32)
        packed, n = pack_binary_pm1(w_pm1, axis=1)
        assert n == in_f
        # Persist packed words in state_dict (int64 view of uint64 bits)
        wp_i64 = torch.from_numpy(np.ascontiguousarray(packed).view(np.int64).copy())
        self.register_buffer("weight_packed_i64", wp_i64)
        self._n = in_f
        self._packed_once = True
        if alpha is None:
            alpha_t = calibrate_linear_scales(weight, cfg=calib or CalibConfig(per_channel=True))
            if alpha_t.ndim == 0:
                alpha_t = alpha_t.expand(out_f)
            alpha_t = alpha_t.float().cpu()
        else:
            alpha_t = alpha.detach().float().reshape(-1).cpu()
            if alpha_t.numel() == 1:
                alpha_t = alpha_t.expand(out_f)
        self.register_buffer("alpha", alpha_t.contiguous().clone())
        if bias is not None:
            self.register_buffer("bias", bias.detach().float().cpu().contiguous().clone())
        else:
            self.bias = None
        self.uses_native = native_kernel_available()
        self._sync_numpy_views()

    def _sync_numpy_views(self) -> None:
        """Rebuild fast NumPy views after init / load_state_dict."""
        wp = self.weight_packed_i64.detach().cpu().numpy()
        self._wp_np = np.ascontiguousarray(wp.view(np.uint64))
        self._alpha_np = np.ascontiguousarray(
            self.alpha.detach().cpu().numpy(), dtype=np.float32
        )
        self._bias_np = (
            None
            if self.bias is None
            else np.ascontiguousarray(self.bias.detach().cpu().numpy(), dtype=np.float32)
        )

    def _load_from_state_dict(self, *args, **kwargs) -> None:
        super()._load_from_state_dict(*args, **kwargs)
        self._sync_numpy_views()
        self.uses_native = native_kernel_available()
        self._packed_once = True

    def extra_repr(self) -> str:
        return (
            f"in={self.in_features}, out={self.out_features}, "
            f"native={self.uses_native}, mode=binary_xnor, packed_once={self._packed_once}"
        )

    def forward(self, x: Tensor) -> Tensor:
        orig = x.shape
        x_cpu = x.detach().to(dtype=torch.float32, device="cpu").contiguous()
        x2 = x_cpu.reshape(-1, self.in_features).numpy()
        xp = _pack_activations_fast(x2, self._n)
        y = None
        if self.uses_native:
            # Preferred: alpha/bias folded into the kernel, so the (B, M)
            # output is written once instead of re-read twice by NumPy.
            y = binary_gemm_native_scaled(
                xp, self._wp_np, self._n, self._alpha_np, self._bias_np
            )
        if y is None:
            if self.uses_native:
                y = binary_gemm_native_prepacked(xp, self._wp_np, self._n)
                assert y is not None
            else:
                y = binary_gemm_numpy_or_blas(xp, self._wp_np, self._n)
            # Unfused fallback: scale (+ bias) in-place on numpy
            y *= self._alpha_np
            if self._bias_np is not None:
                y += self._bias_np
        out = torch.from_numpy(np.ascontiguousarray(y))
        if x.device.type != "cpu":
            out = out.to(x.device)
        return out.reshape(*orig[:-1], self.out_features)

    def packed_weight_bytes(self) -> int:
        return int(self._wp_np.nbytes)

    def gemm_only(self, x_pm1: np.ndarray) -> np.ndarray:
        """Microbench: x already ±1 float (B, N); uses cached packed weights."""
        xp, _ = pack_binary_pm1(x_pm1, axis=1)
        if self.uses_native:
            y = binary_gemm_native_prepacked(xp, self._wp_np, self._n)
            assert y is not None
        else:
            y = binary_gemm_numpy_or_blas(xp, self._wp_np, self._n, x_pm1=x_pm1)
        return y * self._alpha_np

gemm_only

gemm_only(x_pm1: ndarray) -> np.ndarray

Microbench: x already ±1 float (B, N); uses cached packed weights.

Source code in bnn/wrap/packed_linear.py

def gemm_only(self, x_pm1: np.ndarray) -> np.ndarray:
    """Microbench: x already ±1 float (B, N); uses cached packed weights."""
    xp, _ = pack_binary_pm1(x_pm1, axis=1)
    if self.uses_native:
        y = binary_gemm_native_prepacked(xp, self._wp_np, self._n)
        assert y is not None
    else:
        y = binary_gemm_numpy_or_blas(xp, self._wp_np, self._n, x_pm1=x_pm1)
    return y * self._alpha_np

TernaryWeightOnlyLinear

Bases: Module

Accurate-first weight-only ternary (FP activations, FP GEMM after dequant).

Source code in bnn/wrap/packed_linear.py

class TernaryWeightOnlyLinear(nn.Module):
    """Accurate-first weight-only ternary (FP activations, FP GEMM after dequant)."""

    # Buffers declared for the type checker: nn.Module.__getattr__ is
    # typed as Tensor | Module.
    weight_q: Tensor
    scale: Tensor
    bias: Tensor | None

    def __init__(
        self,
        weight: Tensor,
        bias: Tensor | None = None,
        *,
        per_channel: bool = True,
        calib: CalibConfig | None = None,
    ):
        super().__init__()
        out_f, in_f = weight.shape
        self.in_features = in_f
        self.out_features = out_f
        cfg = calib or CalibConfig(per_channel=per_channel)
        if cfg.per_channel:
            q, scale = absmean_ternary_per_channel(weight.float())
        else:
            q, scale = absmean_ternary(weight.float())
        self.register_buffer("weight_q", q.cpu())
        if scale.ndim == 0:
            self.register_buffer("scale", scale.cpu().reshape(()))
            self._per_channel = False
        else:
            self.register_buffer("scale", scale.cpu().contiguous())
            self._per_channel = True
        if bias is not None:
            self.register_buffer("bias", bias.detach().float().cpu().clone())
        else:
            self.bias = None

    def extra_repr(self) -> str:
        return (
            f"in={self.in_features}, out={self.out_features}, "
            f"mode=ternary_weight_only, per_channel={self._per_channel}"
        )

    def forward(self, x: Tensor) -> Tensor:
        if self._per_channel:
            w = self.weight_q.float() * self.scale.unsqueeze(1)
        else:
            w = self.weight_q.float() * self.scale
        return F.linear(
            x.float(),
            w.to(x.device),
            None if self.bias is None else self.bias.to(x.device),
        )

    def packed_weight_bytes(self) -> int:
        # Theoretical 2-bit pack size (actual buffer is int8 — report theoretical)
        return max(self.weight_q.numel() * 2 // 8, 1)

    @property
    def compression_kind(self) -> str:
        return "theoretical_2bit"

BinaryWeightOnlyDequantLinear

Bases: Module

Source code in bnn/wrap/packed_linear.py

class BinaryWeightOnlyDequantLinear(nn.Module):
    # Buffers declared for the type checker: nn.Module.__getattr__ is
    # typed as Tensor | Module.
    weight_pm1: Tensor
    alpha: Tensor
    bias: Tensor | None
    _wp_np: np.ndarray

    def __init__(self, weight: Tensor, bias: Tensor | None = None):
        super().__init__()
        out_f, in_f = weight.shape
        self.in_features = in_f
        self.out_features = out_f
        w_pm1 = sign_pm1(weight.detach().float().cpu()).numpy().astype(np.float32)
        packed, _ = pack_binary_pm1(w_pm1, axis=1)
        self._wp_np = packed
        self.register_buffer("weight_pm1", torch.from_numpy(w_pm1))
        alpha = weight.detach().abs().mean().clamp(min=1e-4)
        self.register_buffer("alpha", alpha.cpu().reshape(()))
        if bias is not None:
            self.register_buffer("bias", bias.detach().float().cpu().clone())
        else:
            self.bias = None

    def forward(self, x: Tensor) -> Tensor:
        w = self.weight_pm1.to(x.device) * self.alpha.to(x.device)
        return F.linear(
            x.float(),
            w,
            None if self.bias is None else self.bias.to(x.device),
        )

    def packed_weight_bytes(self) -> int:
        return int(self._wp_np.nbytes)

PackedBinaryConv2d

Bases: Module

Packed ±1 Conv2d weights (size win). Forward = dequant + F.conv2d.

Thesis: this is a size path (uint64 pack of ±1 kernels), not an XNOR popcount Conv claim. Packed words live in weight_packed_i64 for .bnnpack / state_dict round-trips (W5.T09).

Source code in bnn/wrap/packed_linear.py

class PackedBinaryConv2d(nn.Module):
    """Packed ±1 Conv2d weights (size win). Forward = dequant + F.conv2d.

    Thesis: this is a **size** path (uint64 pack of ±1 kernels), not an XNOR
    popcount Conv claim. Packed words live in ``weight_packed_i64`` for
    ``.bnnpack`` / state_dict round-trips (W5.T09).
    """

    # Buffers declared for the type checker: nn.Module.__getattr__ is
    # typed as Tensor | Module.
    weight_pm1: Tensor
    weight_packed_i64: Tensor
    alpha: Tensor
    bias: Tensor | None
    _wp_np: np.ndarray

    def __init__(
        self,
        weight: Tensor,
        bias: Tensor | None = None,
        *,
        stride: int = 1,
        padding: int = 0,
        dilation: int = 1,
        groups: int = 1,
        alpha: Tensor | None = None,
    ):
        super().__init__()
        if groups != 1:
            raise ValueError("PackedBinaryConv2d supports groups=1 only")
        if dilation != 1:
            raise ValueError("PackedBinaryConv2d supports dilation=1 only")
        out_c, in_c, kh, kw = weight.shape
        self.in_channels = in_c
        self.out_channels = out_c
        self.kernel_size = (kh, kw)
        self.stride = int(stride)
        self.padding = int(padding)
        self.dilation = int(dilation)
        self.groups = int(groups)
        w_pm1 = sign_pm1(weight.detach().float().cpu()).numpy().astype(np.float32)
        flat = w_pm1.reshape(out_c, -1)
        packed, n = pack_binary_pm1(flat, axis=1)
        self._n = n
        wp_i64 = torch.from_numpy(np.ascontiguousarray(packed).view(np.int64).copy())
        self.register_buffer("weight_packed_i64", wp_i64)
        self.register_buffer("weight_pm1", torch.from_numpy(w_pm1))
        self._sync_numpy_views()
        if alpha is None:
            a = weight.detach().abs().mean(dim=(1, 2, 3)).clamp(min=1e-4).float().cpu()
        else:
            a = alpha.detach().float().reshape(-1).cpu()
            if a.numel() == 1:
                a = a.expand(out_c)
        self.register_buffer("alpha", a.contiguous().clone())
        if bias is not None:
            self.register_buffer("bias", bias.detach().float().cpu().clone())
        else:
            self.bias = None

    def _sync_numpy_views(self) -> None:
        wp = self.weight_packed_i64.detach().cpu().numpy()
        self._wp_np = np.ascontiguousarray(wp.view(np.uint64))

    def _load_from_state_dict(self, *args, **kwargs) -> None:
        super()._load_from_state_dict(*args, **kwargs)
        self._sync_numpy_views()

    def extra_repr(self) -> str:
        return (
            f"in={self.in_channels}, out={self.out_channels}, "
            f"k={self.kernel_size}, mode=binary_conv_packed_dequant, "
            f"packed_once=True"
        )

    def forward(self, x: Tensor) -> Tensor:
        w = self.weight_pm1.to(x.device) * self.alpha.view(-1, 1, 1, 1).to(x.device)
        # Out-of-place ±1 activations — do not mutate caller tensors.
        x_b = x.gt(0).to(x.dtype).mul(2).sub(1)
        return F.conv2d(
            x_b,
            w,
            None if self.bias is None else self.bias.to(x.device),
            stride=self.stride,
            padding=self.padding,
            dilation=self.dilation,
            groups=self.groups,
        )

    def packed_weight_bytes(self) -> int:
        return int(self._wp_np.nbytes)

Policy, calibration, guardrails

recommend_wrap_policy

recommend_wrap_policy(
    layer: Linear | None = None,
    hw: HardwareInfo | None = None,
    *,
    accuracy_first: bool = False,
) -> PolicyDecision

Recommend wrap mode/policy for a layer (or globally if layer is None).

Source code in bnn/wrap/policy.py

def recommend_wrap_policy(
    layer: nn.Linear | None = None,
    hw: HardwareInfo | None = None,
    *,
    accuracy_first: bool = False,
) -> PolicyDecision:
    """Recommend wrap mode/policy for a layer (or globally if layer is None)."""
    hw = hw or detect_hardware()
    if hw.prefer_gpu_int4:
        return PolicyDecision(
            policy="hybrid_ffn",
            mode="ternary_weight_only",
            reason="CUDA present — prefer torchao/AWQ INT4/FP8 for GPU; ternary_wo only as size demo",
            min_in_features=MIN_WIDTH_DEFAULT,
            min_out_features=MIN_WIDTH_DEFAULT,
            skip_attn=True,
            fallback_note="Use GGUF/AWQ/torchao for production GPU; do not claim binary 32× on GPU",
        )

    in_f = out_f = None
    if layer is not None:
        in_f, out_f = int(layer.in_features), int(layer.out_features)

    wide = True
    if in_f is not None and out_f is not None:
        wide = in_f >= MIN_WIDTH_BINARY_EFFICIENT and out_f >= MIN_WIDTH_BINARY_EFFICIENT

    if accuracy_first or not hw.native_binary_gemm:
        return PolicyDecision(
            policy="ternary_wo" if accuracy_first else "hybrid_ffn",
            mode="ternary_weight_only",
            reason=(
                "accuracy_first"
                if accuracy_first
                else "native binary GEMM missing — ternary weight-only (size; FP GEMM)"
            ),
            min_in_features=MIN_WIDTH_DEFAULT,
            min_out_features=MIN_WIDTH_DEFAULT,
            skip_attn=True,
            fallback_note="For CPU LLM speed without BitNet kernels prefer GGUF Q4_K",
        )

    if layer is not None and not wide:
        return PolicyDecision(
            policy="hybrid_ffn",
            mode="ternary_weight_only",
            reason=f"narrow layer ({in_f}×{out_f}) — binary pack overhead; prefer ternary/INT8",
            min_in_features=MIN_WIDTH_DEFAULT,
            min_out_features=MIN_WIDTH_DEFAULT,
            skip_attn=True,
            fallback_note="INT8 dynamic or keep FP for narrow Linears",
        )

    return PolicyDecision(
        policy="hybrid_ffn",
        mode="binary_xnor",
        reason="CPU + native XNOR DLL + wide FFN → binary_xnor hybrid",
        min_in_features=MIN_WIDTH_BINARY_EFFICIENT,
        min_out_features=MIN_WIDTH_DEFAULT,
        skip_attn=True,
    )

CalibConfig dataclass

Source code in bnn/wrap/calibrate.py

@dataclass
class CalibConfig:
    method: ScaleMethod = "absmean"
    percentile: float = 99.0
    per_channel: bool = True
    max_batches: int = 4

calibrate_linear_scales

calibrate_linear_scales(
    weight: Tensor,
    *,
    cfg: CalibConfig | None = None,
    activation_batches: list[Tensor] | None = None,
) -> Tensor

Return alpha/scale for a Linear weight.

If activation_batches is provided, optionally blend with activation absmean — still weight-primary for PTQ wrap.

Activation nudge (honest): when per_channel scales are used, the act factor is a global scalar sqrt(mean(|act|)) clamped to [0.5, 2.0] and multiplied onto every channel. It is not per-token or per-channel activation calibration — only a mild distribution nudge so absmean weight scales are not wildly off under atypical input ranges.

Source code in bnn/wrap/calibrate.py

def calibrate_linear_scales(
    weight: Tensor,
    *,
    cfg: CalibConfig | None = None,
    activation_batches: list[Tensor] | None = None,
) -> Tensor:
    """Return alpha/scale for a Linear weight.

    If ``activation_batches`` is provided, optionally blend with activation
    absmean — still weight-primary for PTQ wrap.

    **Activation nudge (honest):** when ``per_channel`` scales are used, the
    act factor is a *global* scalar ``sqrt(mean(|act|))`` clamped to
    ``[0.5, 2.0]`` and multiplied onto every channel. It is not per-token or
    per-channel activation calibration — only a mild distribution nudge so
    absmean weight scales are not wildly off under atypical input ranges.
    """
    cfg = cfg or CalibConfig()
    alpha = scale_from_weight(weight, cfg)
    if activation_batches:
        acts = torch.cat(
            [a.detach().float().reshape(-1, a.shape[-1]) for a in activation_batches[: cfg.max_batches]],
            dim=0,
        )
        act_s = acts.abs().mean().clamp(min=1e-8)
        if alpha.ndim == 0:
            alpha = (alpha * act_s).sqrt()
        else:
            # Global act factor on per-channel weight scales (see docstring).
            alpha = alpha * (act_s.sqrt().clamp(0.5, 2.0))
    return alpha

check_linear_wrap_guardrails

check_linear_wrap_guardrails(
    lin: Linear,
    *,
    mode: str = "binary_xnor",
    force: bool = False,
) -> GuardrailVerdict

Return whether wrapping this Linear is advisable.

Hard-refuse binary XNOR on pathologically narrow dims. Widths below MIN_WIDTH_BINARY_EFFICIENT remain allowed when the caller set a low min_in_features (demos / pedagogy) — efficiency tip only.

Source code in bnn/wrap/guardrails.py

def check_linear_wrap_guardrails(
    lin: nn.Linear,
    *,
    mode: str = "binary_xnor",
    force: bool = False,
) -> GuardrailVerdict:
    """Return whether wrapping this Linear is advisable.

    Hard-refuse binary XNOR on pathologically narrow dims. Widths below
    ``MIN_WIDTH_BINARY_EFFICIENT`` remain allowed when the caller set a low
    ``min_in_features`` (demos / pedagogy) — efficiency tip only.
    """
    inn, out = int(lin.in_features), int(lin.out_features)
    if mode != "binary_xnor":
        return GuardrailVerdict(True, "OK", "non-binary mode")

    if inn < HARD_REFUSE_IN or out < HARD_REFUSE_OUT:
        msg = (
            f"Refuse binary_xnor on pathological Linear ({inn}×{out}): "
            f"use ternary_weight_only / skip / INT8. "
            f"(Efficient binary usually wants ≥{MIN_WIDTH_BINARY_EFFICIENT}.) "
            f"Pass force_narrow/force=True to override."
        )
        if force:
            return GuardrailVerdict(True, "FORCED_NARROW", msg)
        return GuardrailVerdict(False, "NARROW_BINARY", msg)

    if inn < MIN_WIDTH_BINARY_EFFICIENT or out < MIN_WIDTH_BINARY_EFFICIENT:
        return GuardrailVerdict(
            True,
            "SUBOPTIMAL_WIDTH",
            (
                f"binary_xnor on {inn}×{out} may lose to FP/INT8 on wall-clock; "
                f"efficient regime is typically ≥{MIN_WIDTH_BINARY_EFFICIENT}."
            ),
        )
    return GuardrailVerdict(True, "OK", "shape acceptable")

light_qat_recover

light_qat_recover(
    model: Module,
    calib_x: Tensor,
    *,
    teacher: Module | None = None,
    steps: int = 50,
    lr: float = 0.001,
    layer_names: list[str] | None = None,
    loss_fn: Callable[[Tensor, Tensor], Tensor]
    | None = None,
    logit_loss: LogitLoss = "kd",
    temperature: float = 2.0,
    fold_alpha: bool = True,
    train_targets_only: bool = False,
    hidden_mse: float = 0.0,
    binarize_activations: bool = True,
    sign_mode: SignMode | None = None,
) -> dict

Short STE fine-tune on named Linears (default: modules named ffn / mlp).

If teacher is given, distill via logit_loss (kd / mse / cosine); else requires loss_fn. Learned STE alpha is folded into restored Linear magnitudes by default so wrap calib matches QAT.

Source code in bnn/wrap/qat.py

def light_qat_recover(
    model: nn.Module,
    calib_x: torch.Tensor,
    *,
    teacher: nn.Module | None = None,
    steps: int = 50,
    lr: float = 1e-3,
    layer_names: list[str] | None = None,
    loss_fn: Callable[[torch.Tensor, torch.Tensor], torch.Tensor] | None = None,
    logit_loss: LogitLoss = "kd",
    temperature: float = 2.0,
    fold_alpha: bool = True,
    train_targets_only: bool = False,
    hidden_mse: float = 0.0,
    binarize_activations: bool = True,
    sign_mode: SignMode | None = None,
) -> dict:
    """Short STE fine-tune on named Linears (default: modules named *ffn* / *mlp*).

    If ``teacher`` is given, distill via ``logit_loss`` (kd / mse / cosine);
    else requires ``loss_fn``. Learned STE ``alpha`` is folded into restored
    Linear magnitudes by default so wrap calib matches QAT.
    """
    if steps <= 0:
        return {"steps": 0, "skipped": True}

    if teacher is None and loss_fn is None:
        raise ValueError(
            "light_qat_recover requires teacher=... or loss_fn=... "
            "(self-argmax CE fallback removed — it was a no-op / harmful)"
        )

    targets = _collect_target_linears(model, layer_names)
    if not targets:
        return {"steps": 0, "skipped": True, "reason": "no target Linears"}

    target_names = [n for n, _ in targets]
    for name, lin in targets:
        _set_module(
            model, name, _swap_linear_to_binary(lin, binarize_activations=binarize_activations)
        )

    model.train()
    if teacher is not None:
        teacher.eval()

    if train_targets_only:
        name_set = set(target_names)
        params = []
        for n, p in model.named_parameters():
            owner = n.rsplit(".", 1)[0] if "." in n else n
            if owner in name_set or n in name_set:
                params.append(p)
            else:
                p.requires_grad_(False)
        opt = torch.optim.Adam(params, lr=lr)
    else:
        opt = torch.optim.Adam(
            [p for p in model.parameters() if p.requires_grad],
            lr=lr,
        )

    s_cache: dict[str, Tensor] = {}
    t_cache: dict[str, Tensor] = {}
    s_hooks: list[torch.utils.hooks.RemovableHandle] = []
    t_hooks: list[torch.utils.hooks.RemovableHandle] = []
    if hidden_mse > 0 and teacher is not None:
        s_cache, s_hooks = _hidden_hooks(model, target_names)
        t_cache, t_hooks = _hidden_hooks(teacher, target_names)

    last_loss = 0.0
    try:
        with temporary_sign_mode(sign_mode):
            for _ in range(steps):
                opt.zero_grad(set_to_none=True)
                student_out = model(calib_x)
                if loss_fn is not None:
                    loss = loss_fn(student_out, calib_x)
                elif teacher is not None:
                    with torch.no_grad():
                        t_out = teacher(calib_x)
                    loss = agreement_loss(
                        student_out, t_out, kind=logit_loss, temperature=temperature
                    )
                    if hidden_mse > 0 and s_cache and t_cache:
                        hid = torch.zeros((), device=student_out.device, dtype=student_out.dtype)
                        n_h = 0
                        for key in target_names:
                            if key in s_cache and key in t_cache:
                                hid = hid + F.mse_loss(s_cache[key], t_cache[key].detach())
                                n_h += 1
                        if n_h:
                            loss = loss + hidden_mse * (hid / n_h)
                else:
                    raise ValueError("unreachable: teacher/loss_fn required")

                loss.backward()
                opt.step()
                clip_weights_(model)
                last_loss = float(loss.detach().item())
    finally:
        for h in s_hooks + t_hooks:
            h.remove()
        if train_targets_only:
            for p in model.parameters():
                p.requires_grad_(True)

    model.eval()
    restored: list[str] = []
    named = dict(model.named_modules())
    for name, _ in targets:
        ste = named.get(name)
        if not isinstance(ste, (BinaryLinear, WeightOnlySTELinear)):
            continue
        _set_module(model, name, _restore_binary_to_linear(ste, fold_alpha=fold_alpha))
        restored.append(name)

    return {
        "steps": steps,
        "skipped": False,
        "last_loss": last_loss,
        "restored_linears": restored,
        "logit_loss": logit_loss if loss_fn is None else "custom",
        "fold_alpha": fold_alpha,
        "binarize_activations": binarize_activations,
        "sign_mode": sign_mode or get_sign_mode(),
        "note": "Light STE only; production needs BitDistill-scale QAT on real data",
    }

score_layer_sensitivity

score_layer_sensitivity(
    model: Module,
    calib_inputs: Tensor,
    *,
    mode: ScoreMode = "binary_xnor",
    policy: str = "all_large_linear",
    min_in_features: int = 32,
    min_out_features: int = 0,
    drop_in_threshold: float = 0.85,
    skip_fragile: bool = True,
    fragile_drop: float = 0.05,
    calib: CalibConfig | None = None,
) -> SensitivityReport

Score each eligible Linear by cosine drop when wrapped alone.

fragile_drop: if baseline_cosine - layer_cosine >= fragile_drop, suggest skip (also if cosine falls below drop_in_threshold).

Source code in bnn/wrap/sensitivity.py

def score_layer_sensitivity(
    model: nn.Module,
    calib_inputs: Tensor,
    *,
    mode: ScoreMode = "binary_xnor",
    policy: str = "all_large_linear",
    min_in_features: int = 32,
    min_out_features: int = 0,
    drop_in_threshold: float = 0.85,
    skip_fragile: bool = True,
    fragile_drop: float = 0.05,
    calib: CalibConfig | None = None,
) -> SensitivityReport:
    """Score each eligible Linear by cosine drop when wrapped alone.

    ``fragile_drop``: if ``baseline_cosine - layer_cosine >= fragile_drop``,
    suggest skip (also if cosine falls below ``drop_in_threshold``).
    """
    teacher = copy.deepcopy(model)
    teacher.eval()
    model = model.eval()

    with torch.no_grad():
        t_logits = teacher(calib_inputs)
        base_logits = model(calib_inputs)
    base = measure_agreement(t_logits, base_logits, drop_in_threshold=drop_in_threshold)
    baseline_cos = float(base.cosine)

    candidates, _skipped = select_linears(
        model,
        policy=policy,  # type: ignore[arg-type]
        min_in_features=min_in_features,
        min_out_features=min_out_features,
        skip_attn=True,
    )

    layers: list[LayerSensitivity] = []
    skip_suggested: list[str] = []
    wrap_suggested: list[str] = []

    for name, lin in candidates:
        probe = copy.deepcopy(model)
        _set_module(probe, name, _wrap_one_linear(lin, mode, calib=calib))
        probe.eval()
        with torch.no_grad():
            s_logits = probe(calib_inputs)
        eff = measure_agreement(t_logits, s_logits, drop_in_threshold=drop_in_threshold)
        drop = max(0.0, baseline_cos - float(eff.cosine))
        fragile = drop >= fragile_drop or float(eff.cosine) < drop_in_threshold
        if fragile and skip_fragile:
            rec: ScoreMode | Literal["skip"] = "skip"
            reason = f"fragile: cosine={eff.cosine:.4f} drop={drop:.4f}"
            skip_suggested.append(name)
        else:
            rec = mode
            reason = f"ok: cosine={eff.cosine:.4f} drop={drop:.4f}"
            wrap_suggested.append(name)
        layers.append(
            LayerSensitivity(
                name=name,
                in_features=int(lin.in_features),
                out_features=int(lin.out_features),
                cosine=float(eff.cosine),
                cosine_drop=float(drop),
                top1_agreement=eff.top1_agreement,
                recommended=rec,
                reason=reason,
            )
        )

    layers.sort(key=lambda L: L.cosine_drop, reverse=True)
    return SensitivityReport(
        baseline_cosine=baseline_cos,
        layers=layers,
        skip_suggested=skip_suggested,
        wrap_suggested=wrap_suggested,
        drop_in_threshold=drop_in_threshold,
        mode_scored=mode,
    )

search_layer_modes

search_layer_modes(
    model: Module,
    calib_inputs: Tensor,
    *,
    quality_floor: float = 0.9,
    policy: str = "all_large_linear",
    min_in_features: int = 32,
    min_out_features: int = 0,
    calib: CalibConfig | None = None,
    max_relaxations: int | None = None,
) -> ModeSearchReport

Pick binary / ternary / skip per layer to maximise theoretical compression while keeping measured output cosine at or above quality_floor (W3.T06).

Strategy: start from the most aggressive assignment (everything binary), then repeatedly relax the single layer that is costing the most quality — binary → ternary → skip — remeasuring the whole model each time. Relaxing greedily by measured damage is what makes this better than a per-layer threshold: layer interactions only show up in the joint measurement.

Cost is O(L) probes in the common case rather than the 3**L of an exhaustive search, which is why it stays usable on real stacks.

Source code in bnn/wrap/sensitivity.py

def search_layer_modes(
    model: nn.Module,
    calib_inputs: Tensor,
    *,
    quality_floor: float = 0.90,
    policy: str = "all_large_linear",
    min_in_features: int = 32,
    min_out_features: int = 0,
    calib: CalibConfig | None = None,
    max_relaxations: int | None = None,
) -> ModeSearchReport:
    """Pick binary / ternary / skip **per layer** to maximise theoretical
    compression while keeping measured output cosine at or above
    ``quality_floor`` (W3.T06).

    Strategy: start from the most aggressive assignment (everything binary),
    then repeatedly relax the single layer that is costing the most quality —
    binary → ternary → skip — remeasuring the *whole* model each time. Relaxing
    greedily by measured damage is what makes this better than a per-layer
    threshold: layer interactions only show up in the joint measurement.

    Cost is ``O(L)`` probes in the common case rather than the ``3**L`` of an
    exhaustive search, which is why it stays usable on real stacks.
    """
    teacher = copy.deepcopy(model).eval()
    base_model = model.eval()

    with torch.no_grad():
        t_logits = teacher(calib_inputs)
        baseline_cos = float(
            measure_agreement(t_logits, base_model(calib_inputs)).cosine
        )

    candidates, _skipped = select_linears(
        base_model,
        policy=policy,  # type: ignore[arg-type]
        min_in_features=min_in_features,
        min_out_features=min_out_features,
        skip_attn=True,
    )
    if not candidates:
        return ModeSearchReport(
            baseline_cosine=baseline_cos,
            final_cosine=baseline_cos,
            quality_floor=quality_floor,
            met_floor=baseline_cos >= quality_floor,
        )

    elems = {name: int(lin.weight.numel()) for name, lin in candidates}
    order: list[str] = [name for name, _ in candidates]
    chosen: dict[str, LayerMode] = dict.fromkeys(order, "binary_xnor")
    reasons: dict[str, str] = {}
    probes = 0

    def build(assignment: dict[str, LayerMode]) -> nn.Module:
        probe = copy.deepcopy(base_model)
        for name, lin in candidates:
            mode = assignment[name]
            if mode == "skip":
                continue
            _set_module(probe, name, _wrap_one_linear(lin, mode, calib=calib))
        return probe.eval()

    def cosine_of(assignment: dict[str, LayerMode]) -> float:
        nonlocal probes
        probes += 1
        with torch.no_grad():
            return float(measure_agreement(t_logits, build(assignment)(calib_inputs)).cosine)

    current = cosine_of(chosen)
    # Each layer can be relaxed at most twice (binary→ternary→skip).
    budget = 2 * len(order) if max_relaxations is None else max_relaxations

    while current < quality_floor and budget > 0:
        # Which single relaxation buys the most quality right now?
        best_gain, best_name, best_mode, best_cos = 0.0, None, None, current
        for name in order:
            nxt = _relax(chosen[name])
            if nxt is None:
                continue
            trial = dict(chosen)
            trial[name] = nxt
            cos = cosine_of(trial)
            if cos - current > best_gain:
                best_gain, best_name, best_mode, best_cos = cos - current, name, nxt, cos
        if best_name is None or best_mode is None or best_gain <= 0.0:
            break  # nothing left that helps — report honestly below
        reasons[best_name] = (
            f"relaxed to {best_mode}: cosine {current:.4f} -> {best_cos:.4f}"
        )
        chosen[best_name] = best_mode
        current = best_cos
        budget -= 1

    assignments = [
        SearchAssignment(
            name=name,
            mode=chosen[name],
            cosine=current,
            weight_elems=elems[name],
            packed_bytes=elems[name] * _BYTES_PER_ELEM[chosen[name]],
            reason=reasons.get(name, "kept most aggressive mode (floor already met)"),
        )
        for name in order
    ]
    return ModeSearchReport(
        baseline_cosine=baseline_cos,
        final_cosine=current,
        quality_floor=quality_floor,
        met_floor=current >= quality_floor,
        assignments=assignments,
        probes=probes,
    )

ModeSearchReport dataclass

Result of the per-layer binary / ternary / skip search (W3.T06).

Source code in bnn/wrap/sensitivity.py

@dataclass
class ModeSearchReport:
    """Result of the per-layer binary / ternary / skip search (W3.T06)."""

    baseline_cosine: float
    final_cosine: float
    quality_floor: float
    met_floor: bool
    assignments: list[SearchAssignment] = field(default_factory=list)
    probes: int = 0

    @property
    def binary(self) -> list[str]:
        return [a.name for a in self.assignments if a.mode == "binary_xnor"]

    @property
    def ternary(self) -> list[str]:
        return [a.name for a in self.assignments if a.mode == "ternary_weight_only"]

    @property
    def skipped(self) -> list[str]:
        return [a.name for a in self.assignments if a.mode == "skip"]

    def compression(self) -> float:
        """Theoretical weight compression over the searched layers only."""
        elems = sum(a.weight_elems for a in self.assignments)
        packed = sum(a.packed_bytes for a in self.assignments)
        if not elems or packed <= 0:
            return 1.0
        return (elems * 4.0) / packed

    def to_dict(self) -> dict[str, Any]:
        return {
            "baseline_cosine": self.baseline_cosine,
            "final_cosine": self.final_cosine,
            "quality_floor": self.quality_floor,
            "met_floor": self.met_floor,
            "probes": self.probes,
            "binary": self.binary,
            "ternary": self.ternary,
            "skipped": self.skipped,
            "theoretical_compression_searched_layers": self.compression(),
            "assignments": [a.to_dict() for a in self.assignments],
            "thesis_note": (
                "Compression is a theoretical pack ratio; cosine is measured. "
                "Never report compression as an end-to-end speedup."
            ),
        }

compression

compression() -> float

Theoretical weight compression over the searched layers only.

Source code in bnn/wrap/sensitivity.py

def compression(self) -> float:
    """Theoretical weight compression over the searched layers only."""
    elems = sum(a.weight_elems for a in self.assignments)
    packed = sum(a.packed_bytes for a in self.assignments)
    if not elems or packed <= 0:
        return 1.0
    return (elems * 4.0) / packed

SearchAssignment dataclass

Chosen mode for one Linear, with the evidence behind the choice.

Source code in bnn/wrap/sensitivity.py

@dataclass
class SearchAssignment:
    """Chosen mode for one Linear, with the evidence behind the choice."""

    name: str
    mode: LayerMode
    cosine: float
    weight_elems: int
    packed_bytes: float
    reason: str

    def to_dict(self) -> dict[str, Any]:
        return asdict(self)