diff --git a/docs/src/palettization/config.md b/docs/src/palettization/config.md index 90a88b2a..b494a936 100644 --- a/docs/src/palettization/config.md +++ b/docs/src/palettization/config.md @@ -4,6 +4,8 @@ Palettization Configs follow the same philosophy as the [Quantization Config](.. They are simpler as palettization applies only to the weights in the model. (Hence there are no `op_input_spec` and `op_output_spec` fields in the {class}`~coreai_opt.palettization.config.ModuleKMeansPalettizerConfig` and {class}`~coreai_opt.palettization.config.OpKMeansPalettizerConfig`.) +To estimate what a config costs in terms of bits per weight (BPW), use the {func}`~coreai_opt.inspection.bits_per_weight` utility to get an analytical BPW estimate for a *prepared* `coreai-opt` model. See [Utility for computing analytical Bits Per Weight (BPW)](../utils/mixed_precision.md#utility-for-computing-analytical-bits-per-weight-bpw) for details. + ## PalettizationSpec {class}`~coreai_opt.palettization.spec.PalettizationSpec` defines the following key properties, among others (for full list see API reference): diff --git a/docs/src/quantization/config.md b/docs/src/quantization/config.md index 4f0f2d17..a884b6e7 100644 --- a/docs/src/quantization/config.md +++ b/docs/src/quantization/config.md @@ -15,6 +15,8 @@ Regarding the terminology of ops and modules, as used in these APIs, here is wha {class}`~coreai_opt.quantization.config.QuantizerConfig` lets you customize quantization settings for different parts of the model, via either modules or ops, or a combination of both. +To estimate what a config costs in terms of bits per weight (BPW), use the {func}`~coreai_opt.inspection.bits_per_weight` utility to get an analytical BPW estimate for a model *prepared* by `coreai-opt`. See [Utility for computing analytical Bits Per Weight (BPW)](../utils/mixed_precision.md#utility-for-computing-analytical-bits-per-weight-bpw) for details. + We will first take a look at the {class}`~coreai_opt.quantization.spec.QuantizationSpec` class and then through examples, will walk over how to define the config classes (and thereby covering how they are structured). ## QuantizationSpec diff --git a/docs/src/utils/mixed_precision.md b/docs/src/utils/mixed_precision.md index 58b76c49..ba73e900 100644 --- a/docs/src/utils/mixed_precision.md +++ b/docs/src/utils/mixed_precision.md @@ -38,6 +38,10 @@ Once we have the sensitivity scores of all the layers and candidate configs, we Several strategies can be applied for determining the mixed-precision recipe given a constraint — for example, a target average bits-per-weight (BPW) — and aim to balance model size reduction with minimal accuracy loss. +:::{note} +To obtain the analytical average BPW of a model, see [Utility for computing analytical Bits Per Weight (BPW)](#utility-for-computing-analytical-bits-per-weight-bpw) below. +::: + A simple greedy approach often works well: 1. Sort all `(layer, config)` tuples by sensitivity in descending order (least quality loss first). @@ -69,3 +73,46 @@ Likewise, the per-layer setting being varied across candidate configs does not h - [Mixed-precision palettization with ResNet50](../examples/model_examples/mixed_precision_palettization.md) — applies palettization with 2/4/6-bit per-tensor candidate configs and the greedy approach targeting a BPW of 4. - [coreai-models](https://github.com/apple/coreai-models) — the same workflow is applied to a few LLMs in this repository to produce mixed precision compression recipes. Users can find the mixed precision configs in the repo and apply them with `coreai-opt`. + +## Utility for computing analytical Bits Per Weight (BPW) + +{func}`~coreai_opt.inspection.bits_per_weight` is a utility that computes an analytical BPW estimate from a *prepared* `coreai-opt` model. + +It estimates the average bitwidth of a model, accounting for compression overhead such as quantization scales and zero-points as well as palettization look-up tables and per-channel scales. Compressed tensors are counted at their effective compressed cost and uncompressed tensors are counted at their full-precision dtype cost. + +**Usage:** + +```python +from coreai_opt.quantization import Quantizer, QuantizerConfig +from coreai_opt.quantization.config import ExecutionMode +from coreai_opt.inspection import BitsPerWeightResult, bits_per_weight + +# Prepare an eager-mode weight-quantized model +quantizer = Quantizer( + model, QuantizerConfig.presets.w8(execution_mode=ExecutionMode.EAGER) +) +prepared_model = quantizer.prepare(example_inputs) + +result: BitsPerWeightResult = bits_per_weight(prepared_model) +print(result.bpw) # e.g. 8.86 +``` + +The returned {class}`~coreai_opt.inspection.BitsPerWeightResult` also exposes `per_module_map`, a mapping from module name to the module's own average BPW, which is useful for inspecting how bits are distributed across different parts of the model. + +The utility supports: + +- **Eager-mode integer weight quantization**: any spec-supported integer dtype, any quantization scheme (symmetric or asymmetric), at any granularity. +- **Eager-mode floating-point weight quantization**: FP8 (`e4m3`, `e5m2`) and FP4 (`e2m1`). +- **Palettization** at any spec-supported `n_bits`, including a quantized LUT. + +Sub-byte payloads and zero-points are packed at `n_bits` with no padding. + +:::{note} +Sparsity (pruning) and models quantized using graph mode (`ExecutionMode.GRAPH`) are currently not supported. +::: + +:::{note} +This is an analytical estimate computed on a *prepared* `coreai-opt` model and as such does not reflect the exported asset size of any [`finalize()` backend](../introduction/integration_coreai.md). Passing a finalized model to the utility is not supported. +::: + +For the full API, see the {func}`~coreai_opt.inspection.bits_per_weight` reference.