Optimiser

The product verb: bnn optimise. Prefer this over the legacy wrap entry points.

optimise_model

optimise_model(
    model: Module,
    calib_inputs: Tensor | None = None,
    config: OptimiseConfig | None = None,
    *,
    teacher: Module | None = None,
    **kwargs: Any,
) -> OptimiseResult

Calibrate → (optional distill/QAT) → wrap → effectiveness → optional .bnnpack.

Parameters

model: FP (or mixed) nn.Module to optimise for CPU/edge packed inference. calib_inputs: Optional batch used for agreement metrics, distill, and light QAT. Shape must match the model's forward. config: OptimiseConfig; keyword overrides also accepted via kwargs. fuse_bn folds BN before pack; distill_steps runs STE KD when calib_inputs (and a teacher) are available. teacher: Optional FP teacher for agreement / distill; defaults to a deepcopy taken before wrap when calib_inputs is provided.

Returns

OptimiseResult Wrapped model, WrapReport, and a bnn_optimise_report_v1 payload.

Source code in bnn/optimise.py

def optimise_model(
    model: nn.Module,
    calib_inputs: Tensor | None = None,
    config: OptimiseConfig | None = None,
    *,
    teacher: nn.Module | None = None,
    **kwargs: Any,
) -> OptimiseResult:
    """Calibrate → (optional distill/QAT) → wrap → effectiveness → optional ``.bnnpack``.

    Parameters
    ----------
    model:
        FP (or mixed) ``nn.Module`` to optimise for CPU/edge packed inference.
    calib_inputs:
        Optional batch used for agreement metrics, distill, and light QAT. Shape
        must match the model's forward.
    config:
        ``OptimiseConfig``; keyword overrides also accepted via ``kwargs``.
        ``fuse_bn`` folds BN before pack; ``distill_steps`` runs STE KD when
        ``calib_inputs`` (and a teacher) are available.
    teacher:
        Optional FP teacher for agreement / distill; defaults to a deepcopy
        taken *before* wrap when ``calib_inputs`` is provided.

    Returns
    -------
    OptimiseResult
        Wrapped model, ``WrapReport``, and a ``bnn_optimise_report_v1`` payload.
    """
    cfg = config or OptimiseConfig()
    # Allow kwargs overrides for ergonomics
    for k, v in kwargs.items():
        if not hasattr(cfg, k):
            raise TypeError(f"unknown OptimiseConfig field: {k}")
        setattr(cfg, k, v)

    work = model if cfg.inplace else copy.deepcopy(model)
    teacher_mod = teacher
    if calib_inputs is not None and teacher_mod is None:
        teacher_mod = copy.deepcopy(work)
        teacher_mod.eval()

    distill_info: dict[str, Any] | None = None
    mode = cfg.mode
    policy = cfg.policy
    use_binary_distill = (
        cfg.distill_steps > 0
        and policy not in ("ternary_wo",)
        and mode not in ("ternary_weight_only",)
    )
    if use_binary_distill:
        if calib_inputs is None or teacher_mod is None:
            warnings.warn(
                "OptimiseConfig.distill_steps>0 requires calib_inputs "
                "(and a teacher); distillation skipped.",
                stacklevel=2,
            )
        else:
            from .wrap.distill import DistillConfig, distill_binary_student

            d_report = distill_binary_student(
                work,
                teacher_mod,
                calib_inputs,
                cfg=DistillConfig(
                    steps=cfg.distill_steps,
                    lr=cfg.distill_lr,
                    temperature=cfg.distill_temperature,
                    layer_names=cfg.distill_layer_names,
                    drop_in_threshold=cfg.drop_in_threshold,
                    logit_loss="mse",
                    fold_alpha=True,
                ),
            )
            distill_info = d_report.to_dict()

    qat_info: dict[str, Any] | None = None
    use_binary_qat = (
        cfg.qat_steps > 0
        and policy not in ("ternary_wo",)
        and mode not in ("ternary_weight_only",)
    )
    if use_binary_qat and calib_inputs is not None:
        layer_names = cfg.qat_layer_names
        qat_info = light_qat_recover(
            work,
            calib_inputs,
            teacher=teacher_mod,
            steps=cfg.qat_steps,
            lr=1e-3,
            layer_names=layer_names,
            logit_loss=cfg.qat_logit_loss,  # type: ignore[arg-type]
            fold_alpha=cfg.qat_fold_alpha,
            hidden_mse=cfg.qat_hidden_mse,
            binarize_activations=cfg.qat_binarize_activations,
            sign_mode=cfg.qat_sign_mode,  # type: ignore[arg-type]
        )

    sensitivity_payload: dict[str, Any] | None = None
    exclude_exact: list[str] | None = None
    if cfg.sensitivity:
        if calib_inputs is None:
            warnings.warn(
                "OptimiseConfig.sensitivity=True requires calib_inputs; "
                "sensitivity scoring skipped.",
                stacklevel=2,
            )
        else:
            from .wrap.sensitivity import score_layer_sensitivity

            sens_mode = "ternary_weight_only" if cfg.accuracy_first else "binary_xnor"
            if mode in ("ternary_weight_only", "binary_xnor"):
                sens_mode = mode
            sens = score_layer_sensitivity(
                work,
                calib_inputs,
                mode=sens_mode,  # type: ignore[arg-type]
                policy="all_large_linear",
                min_in_features=cfg.min_in_features,
                min_out_features=cfg.min_out_features,
                drop_in_threshold=cfg.drop_in_threshold,
                fragile_drop=cfg.sensitivity_fragile_drop,
                calib=cfg.calib,
            )
            sensitivity_payload = sens.to_dict()
            exclude_exact = list(sens.skip_suggested)

    before = model_param_bytes(work)
    wrapped, report = wrap_model(
        work,
        mode=mode,
        policy=policy,  # type: ignore[arg-type]
        skip_name_substr=cfg.skip_name_substr,
        min_in_features=cfg.min_in_features,
        min_out_features=cfg.min_out_features,
        skip_attn=cfg.skip_attn,
        calib=cfg.calib,
        inplace=True,
        accuracy_first=cfg.accuracy_first,
        exclude_exact=exclude_exact,
        force_narrow=cfg.force,
        fuse_bn=cfg.fuse_bn,
        drop_in_threshold=cfg.drop_in_threshold,
    )
    after = model_param_bytes(wrapped)

    if distill_info is not None:
        report.distill = distill_info
    if qat_info is not None:
        report.qat = qat_info

    status = "OK"
    if calib_inputs is not None and teacher_mod is not None:
        with torch.no_grad():
            t_logits = teacher_mod(calib_inputs)
            s_logits = wrapped(calib_inputs)
        eff = measure_agreement(
            t_logits, s_logits, drop_in_threshold=cfg.drop_in_threshold
        )
        attach_effectiveness(report, eff, force=cfg.force)
        if not report.drop_in_ok and not cfg.force:
            status = "REFUSE_DROP_IN_CLAIM"
        elif report.forced:
            status = "FORCED"
    elif cfg.force:
        status = "FORCED_NO_CALIB"
        report.forced = True

    pack_path: Path | None = None
    if cfg.encode_path is not None:
        from .codec import encode_file

        pack_path = Path(cfg.encode_path)
        encode_file(
            wrapped,
            pack_path,
            meta={
                "source": "bnn.optimise",
                "policy": report.policy,
                "mode": report.mode,
            },
            min_in_features=cfg.encode_min_width,
            include_binary_linear=True,
            include_fp_linear=False,
            include_packed=True,
        )

    payload = envelope(
        policy=str(report.policy),
        mode=str(report.mode),
        replaced=report.replaced,
        skipped=report.skipped,
        compression_replaced_weights=report.compression,
        fp32_weight_bytes_replaced=report.fp32_weight_bytes_replaced,
        packed_weight_bytes=report.packed_weight_bytes,
        native_kernel=bool(report.native_kernel),
        drop_in_ok=report.drop_in_ok,
        forced=bool(report.forced),
        status=status,
        policy_reason=report.policy_reason,
        calib_method=report.calib_method,
        effectiveness=report.effectiveness,
        qat=report.qat,
        distill=report.distill,
        fuse=report.fuse,
        sensitivity=sensitivity_payload,
        param_bytes_before=before,
        param_bytes_after=after,
        pack_path=str(pack_path) if pack_path else None,
    )
    errs = validate_optimise_report(payload)
    if errs:
        warnings.warn(f"optimise report schema issues: {errs}", stacklevel=2)

    return OptimiseResult(
        model=wrapped,
        report=report,
        payload=payload,
        pack_path=pack_path,
        bytes_before=before,
        bytes_after=after,
    )

OptimiseConfig dataclass

Knob set for optimise_model (stable defaults).

Source code in bnn/optimise.py

@dataclass
class OptimiseConfig:
    """Knob set for ``optimise_model`` (stable defaults)."""

    policy: WrapPolicy | str = "auto"
    mode: WrapMode | str | None = None
    min_in_features: int = 64
    min_out_features: int = 0
    skip_attn: bool = True
    skip_name_substr: Iterable[str] | None = None
    calib: CalibConfig | None = field(default_factory=CalibConfig)
    qat_steps: int = 0
    qat_layer_names: list[str] | None = None
    qat_logit_loss: str = "mse"
    qat_fold_alpha: bool = True
    qat_hidden_mse: float = 0.0
    qat_binarize_activations: bool = True
    qat_sign_mode: str | None = None
    drop_in_threshold: float = 0.85
    force: bool = False
    accuracy_first: bool = False
    inplace: bool = False
    encode_path: Path | str | None = None
    encode_min_width: int = 64
    # W3.T05 — optional layer-wise sensitivity before wrap
    sensitivity: bool = False
    sensitivity_fragile_drop: float = 0.05
    # W3.T09 — fold Linear+BN1d / BiReal BN before packing
    fuse_bn: bool = False
    # W3.T08 — optional STE KD before wrap (0 = skip; toy/demo scale)
    distill_steps: int = 0
    distill_lr: float = 1e-3
    distill_temperature: float = 2.0
    distill_layer_names: list[str] | None = None

OptimiseResult dataclass

Product result: wrapped model + versioned report (+ optional pack path).

Source code in bnn/optimise.py

@dataclass
class OptimiseResult:
    """Product result: wrapped model + versioned report (+ optional pack path)."""

    model: nn.Module
    report: WrapReport
    payload: dict[str, Any]
    pack_path: Path | None = None
    bytes_before: dict[str, int] | None = None
    bytes_after: dict[str, int] | None = None

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

Report schema

validate_optimise_report

validate_optimise_report(
    payload: dict[str, Any], *, strict: bool = False
) -> list[str]

Return a list of validation errors (empty ⇒ OK).

strict=True also requires recommended dual-metric keys when latency was measured.

Source code in bnn/wrap/schema.py

def validate_optimise_report(payload: dict[str, Any], *, strict: bool = False) -> list[str]:
    """Return a list of validation errors (empty ⇒ OK).

    ``strict=True`` also requires recommended dual-metric keys when latency was measured.
    """
    errors: list[str] = []
    if not isinstance(payload, dict):
        return ["report must be a dict"]

    schema = payload.get("schema")
    if schema not in (SCHEMA_ID, "ultra_wrap_report_v1"):
        # Accept legacy ultra_wrap_report_v1 as alias during transition
        errors.append(f"schema must be {SCHEMA_ID!r} (or legacy ultra_wrap_report_v1), got {schema!r}")

    if payload.get("schema") == SCHEMA_ID:
        ver = payload.get("schema_version")
        if ver != SCHEMA_VERSION:
            errors.append(f"schema_version must be {SCHEMA_VERSION}, got {ver!r}")
        for key in REQUIRED_KEYS:
            if key not in payload:
                errors.append(f"missing required key: {key}")

    if strict:
        for key in ("effectiveness", "policy_reason"):
            if key not in payload:
                errors.append(f"strict: missing recommended key: {key}")

    # Honesty: compression alone must not be labeled as e2e speedup
    if "e2e_speedup" in payload and payload.get("e2e_speedup") is not None:
        try:
            if float(payload["e2e_speedup"]) > 20 and float(payload.get("compression_replaced_weights") or 0) > 20:
                # Soft warning only — store as note, not hard fail
                note = str(payload.get("thesis_note") or "")
                if "dual-metric" not in note.lower() and "theory" not in note.lower():
                    errors.append(
                        "strict honesty: high e2e_speedup + high compression requires "
                        "thesis_note mentioning dual-metric / theory vs wall-clock"
                    )
        except (TypeError, ValueError):
            pass

    return errors

is_valid_optimise_report

is_valid_optimise_report(
    payload: dict[str, Any], *, strict: bool = False
) -> bool

Source code in bnn/wrap/schema.py

def is_valid_optimise_report(payload: dict[str, Any], *, strict: bool = False) -> bool:
    return not validate_optimise_report(payload, strict=strict)

envelope

envelope(
    *,
    policy: str,
    mode: str,
    replaced: list[str],
    skipped: list[str],
    compression_replaced_weights: float,
    fp32_weight_bytes_replaced: int,
    packed_weight_bytes: int,
    native_kernel: bool,
    drop_in_ok: bool | None,
    forced: bool,
    status: str,
    **extra: Any,
) -> dict[str, Any]

Build a minimal valid bnn_optimise_report_v1 dict.

Source code in bnn/wrap/schema.py

def envelope(
    *,
    policy: str,
    mode: str,
    replaced: list[str],
    skipped: list[str],
    compression_replaced_weights: float,
    fp32_weight_bytes_replaced: int,
    packed_weight_bytes: int,
    native_kernel: bool,
    drop_in_ok: bool | None,
    forced: bool,
    status: str,
    **extra: Any,
) -> dict[str, Any]:
    """Build a minimal valid ``bnn_optimise_report_v1`` dict."""
    out: dict[str, Any] = {
        "schema": SCHEMA_ID,
        "schema_version": SCHEMA_VERSION,
        "policy": policy,
        "mode": mode,
        "replaced": list(replaced),
        "skipped": list(skipped),
        "compression_replaced_weights": float(compression_replaced_weights),
        "fp32_weight_bytes_replaced": int(fp32_weight_bytes_replaced),
        "packed_weight_bytes": int(packed_weight_bytes),
        "native_kernel": bool(native_kernel),
        "drop_in_ok": drop_in_ok,
        "forced": bool(forced),
        "status": status,
        "thesis_note": (
            "Compression is theoretical pack ratio; latency fields are wall-clock. "
            "Never claim GPU 32× from sign()/STE."
        ),
    }
    out.update(extra)
    return out