add utility to compute analytical bits per weight

[pkmandke]

Add utility to compute the average bits-per-weight (bpw) of a prepared coreai-opt model in eager mode quantization and palettization.

Public API

Imported from coreai_opt.inspection:

def bits_per_weight(model: torch.nn.Module) -> BitsPerWeightResult

@dataclass
class BitsPerWeightResult:
    bpw: float                      # total_bits / total_weights, 0.0 if no params
    per_module: dict[str, float]    # module name -> module's own average bpw
    total_bits: int                 # total storage cost in bits, incl. amortized overhead
    total_weights: int              # total logical parameter elements

Usage:

from coreai_opt.inspection import BitsPerWeightResult, bits_per_weight

result = bits_per_weight(prepared_model)
result.bpw          # 4.13
result.per_module   # {"l1": 4.06, "l2": 4.25, ...}

Supported: eager-mode integer weight quantization (int8/int4/int2 and unsigned variants, symmetric or asymmetric, any granularity, sub-byte payloads packed at n_bits), and palettization at any spec-supported n_bits including a quantized LUT.

Raises NotImplementedError for: graph-mode / torch.fx.GraphModule models, floating-point (FP8/FP4) weight quantization, and parametrizations that store multiple original tensors (weight_norm, spectral_norm).

Note: This is an analytical estimate, not a measurement. It is meant for prepared models, not finalized ones.

Testing

• tests/inspection/test_bits_per_weight.py: unit tests against hand-derived golden bit counts across quantization dtypes, qschemes, and granularities, plus palettization n_bits x granularity combinations. Also covers persistent and non-persistent buffers, tied weights, per-module attribution, and the unsupported-config errors.
• tests/export/test_bpw_export_size.py: cross-checks total_bits / 8 against the measured payload of an actual Core AI export. The prediction must be a lower bound and land within a 2.5 percent structural-metadata budget.
• tests/export/export_utils.py: adds coreai_export_size_bytes helper.
• tests/models/simple.py: adds LinearBatchNormModel fixture (buffer coverage) and optional bias to two existing fixtures.

---

TODO:

• Add export tests to validate analytical bpw against exported asset size
• Add tests for per-module bpw

[u-simha]

Have left some comments; my main suggestion would be to have a documentation page (even if it is brief) rather than the doc string at the top of the file

[u-simha]

My suggestion would be to have a doc page explaining the usage of this tool; doc strings in the headings of the file isn't often read, compared to the documentation.

[pkmandke]

@u-simha any suggestions on where to add this in the doc? I'm thinking a new page under OTHER WORKFLOWS AND UTILITIES but if there's other similar utils mentioned elsewhere, it can belong there instead of a new page.

[u-simha]

This is quite a large upper bound; didn't find more information in _assert_bpw_is_plausible which explains this?

[pkmandke]

Ah so this is to basically accommodate for cases where the scale/LUTs would add non-trivial overhead. Even per-channel scales and minval add quite a bit of overhead. This gets highlighted here due to the models being relatively small. In real models, these should get amortized and the overhead (ideally) won't be this much. But it does depend on the config chosen. Like for a really small block/group size, the scale/LUT overhead will be non-trivial.

For eg, in this test file, the max overhead is 1.637 for the palettized[n8_per_tensor-pcs-bias] case. So I chose 2 as the bound..

[pkmandke]

Lmk if/how we should address this. We could only use configs for this smallish model that don't add such a large overhead, but that won't ensure config coverage for the actual BPW logic..

[crowbat]

Nit: Can we replace all usages of "8" with defining a _BITS_PER_BYTE = 8 and using that instead?

[crowbat]

Nit, argument in the docstring is misnamed