diff --git a/changelog.d/67.added b/changelog.d/67.added new file mode 100644 index 0000000..95e88d0 --- /dev/null +++ b/changelog.d/67.added @@ -0,0 +1 @@ +Add utility to compute analytical bits per weight (BPW) of a prepared eager-mode quantized or palettized model diff --git a/src/coreai_opt/inspection/__init__.py b/src/coreai_opt/inspection/__init__.py index e345d64..2105ee9 100644 --- a/src/coreai_opt/inspection/__init__.py +++ b/src/coreai_opt/inspection/__init__.py @@ -20,6 +20,7 @@ print(inspector.format_summary()) """ +from .bits_per_weight import BitsPerWeightResult, bits_per_weight from .model_inspector import ModelInspector from .types import ( BoundaryEdge, @@ -32,6 +33,7 @@ ) __all__ = [ + "BitsPerWeightResult", "BoundaryEdge", "InputEdge", "ModelInspector", @@ -40,4 +42,5 @@ "ModuleInfo", "OpInfo", "SourceFrame", + "bits_per_weight", ] diff --git a/src/coreai_opt/inspection/_bpw_utils.py b/src/coreai_opt/inspection/_bpw_utils.py new file mode 100644 index 0000000..c67109f --- /dev/null +++ b/src/coreai_opt/inspection/_bpw_utils.py @@ -0,0 +1,280 @@ +# Copyright 2026 Apple Inc. +# +# Use of this source code is governed by a BSD-3-Clause license that can +# be found in the LICENSE file or at https://opensource.org/licenses/BSD-3-Clause + +"""Module-tree walking and storage-cost helpers need for ``bits_per_weight``.""" + +import math +from collections.abc import Iterator + +import torch +from torch.nn.utils.parametrize import ParametrizationList as _ParametrizationList + +from coreai_opt.config.spec import ( + CompressionSimulatorBase as _CompressionSimulatorBase, + CompressionTargetTensor as _CompressionTargetTensor, +) +from coreai_opt.palettization.spec.fake_palettize import _FakePalettizeImplBase +from coreai_opt.pruning.spec import PruneImplBase as _PruneImplBase +from coreai_opt.quantization.spec import QuantizationScheme as _QuantizationScheme +from coreai_opt.quantization.spec.fake_quantize import FakeQuantizeImplBase as _FakeQuantizeImplBase +from coreai_opt.quantization.spec.qformulation import ( + QuantizationFormulation as _QuantizationFormulation, +) + +_DEFAULT_SCALE_BITS = 32 + +WeightCompressor = _FakeQuantizeImplBase | _FakePalettizeImplBase + + +def named_modules_excluding_compression_machinery( + module: torch.nn.Module, name: str = "" +) -> Iterator[tuple[str, torch.nn.Module]]: + """Yield ``(dotted_name, module)`` for every module that owns logical tensors. + + Compression machinery owns no logical weights of its own, so a machinery + module and everything below it is left out by not descending into it. + + Args: + module (torch.nn.Module): The subtree root to walk. + name (str): Dotted path of ``module`` from the model root + + Yields: + tuple[str, torch.nn.Module]: Name and module, parents before children. + """ + if isinstance(module, (_CompressionSimulatorBase, _ParametrizationList)): + return + + yield name, module + + for child_name, child in module.named_children(): + yield from named_modules_excluding_compression_machinery( + child, f"{name}.{child_name}" if name else child_name + ) + + +def get_weight_compressor(param_list: _ParametrizationList) -> WeightCompressor | None: + """Return the weight-targeting compressor in a parametrization list, if any. + + Args: + param_list (ParametrizationList): Parametrizations registered on a parameter. + + Returns: + WeightCompressor | None: The first ``_FakePalettizeImplBase`` or + weight-target ``FakeQuantizeImplBase`` in the list, or ``None`` if the + list contains no recognized weight compressor. + """ + for entry in param_list: + if isinstance(entry, _FakePalettizeImplBase): + return entry + if ( + isinstance(entry, _FakeQuantizeImplBase) + and entry.quantization_target == _CompressionTargetTensor.WEIGHT + ): + return entry + return None + + +def ensure_single_original( + param_list: _ParametrizationList, module_name: str, tensor_name: str +) -> None: + """Raise if a parametrization stores its dense tensor as multiple originals. + + A ``right_inverse`` returning a sequence makes PyTorch register ``original0``, + ``original1``, ... instead of a single ``original`` (as + ``torch.nn.utils.parametrizations.weight_norm`` does), so there is no one + dense tensor whose storage cost we can attribute. + """ + if not param_list.is_tensor: + raise NotImplementedError( + f"bits_per_weight cannot size the parametrization on " + f"'{module_name}.{tensor_name}': it stores multiple original tensors " + f"(e.g. weight_norm / spectral_norm) rather than a single dense one." + ) + + +def ensure_not_pruned(param_list: _ParametrizationList, module_name: str, tensor_name: str) -> None: + """Raise if a weight carries a pruning parametrization.""" + for entry in param_list: + if isinstance(entry, _PruneImplBase): + raise NotImplementedError( + f"bits_per_weight cannot compute the storage cost of a pruned " + f"weight '{module_name}.{tensor_name}'." + ) + + +def tensor_storage_bits(weight: torch.Tensor, compressor: WeightCompressor | None) -> int: + """Return the storage cost in bits of a (possibly compressed) weight tensor.""" + if compressor is None: + return full_precision_bits(weight) + if isinstance(compressor, _FakePalettizeImplBase): + return _palettized_bits(weight, compressor) + return _quantized_bits(weight, compressor) + + +def full_precision_bits(tensor: torch.Tensor) -> int: + """Return the dense storage cost of a tensor in bits.""" + return int(tensor.numel() * tensor.element_size() * 8) + + +def _quantized_bits(weight: torch.Tensor, quantization_fq: _FakeQuantizeImplBase) -> int: + """Return the storage cost of a quantized weight including scale / offset overhead. + + Args: + weight (torch.Tensor): The dense original weight tensor. + fq (FakeQuantizeImplBase): The weight fake-quantize parametrization. + + Returns: + int: ``payload_bits + scale_bits + offset_bits`` in bits, where the payload + is ``numel * n_bits`` and the per-block scale and offset overhead is + amortized across the weight. + + Note: + ``num_blocks`` is read directly from the materialized + ``qparams_calculator.scale`` buffer. This is the canonical + per-granularity block count (per-tensor, per-channel, per-block, and + multi-axis per-block all reduce to ``scale.numel()``). When that buffer + is not yet materialized it is empty,``num_blocks`` is + derived analytically from ``granularity.get_block_size``. + """ + num_elements = weight.numel() + + scale = quantization_fq.qparams_calculator.scale + if scale is not None and scale.numel() > 0: + num_blocks = scale.numel() + else: + block_size = quantization_fq.granularity.get_block_size(weight.shape) + num_blocks = num_elements // math.prod(block_size) + + payload_bits = num_elements * quantization_fq.n_bits + scale_bits = num_blocks * _float_qparam_bits(quantization_fq) + + return int(payload_bits + scale_bits + _offset_bits(quantization_fq, num_blocks)) + + +def _float_qparam_bits(fq: _FakeQuantizeImplBase) -> int: + """Return the per-element bit width of the float qparams this weight exports with.""" + dtype = fq.qparams_calculator._compute_dtype_for_export + return int(dtype.itemsize * 8) + + +def _offset_bits(fq: _FakeQuantizeImplBase, num_blocks: int) -> int: + """Return the per-block dequantization offset cost of a quantized weight, in bits. + + - ``ZP``: the export ships ``zero_point``, packed at ``n_bits``. + - ``MINVAL``: the export ships ``minval`` instead, and drops the zero-point. + ``minval`` is a float, so it costs a full float per block. + """ + if fq.qformulation == _QuantizationFormulation.MINVAL: + return num_blocks * _float_qparam_bits(fq) + if fq.qscheme == _QuantizationScheme.ASYMMETRIC: + return num_blocks * fq.n_bits + return 0 + + +def _palettized_bits(weight: torch.Tensor, palettization_fq: _FakePalettizeImplBase) -> int: + """Return the storage cost of a palettized weight including LUT / per-channel-scale overhead. + + Args: + weight (torch.Tensor): The dense original weight tensor. + pal (_FakePalettizeImplBase): The weight fake-palettize parametrization. + + Returns: + int: Effective storage cost of a palettized weight tensor along with + overhead. + """ + num_elements = weight.numel() + + # Indices: one n_bits index per cluster_dim-sized group along axis 0. + indices_bits = (num_elements // palettization_fq.cluster_dim) * palettization_fq.n_bits + + # LUT (centroids): shape after _reshape_lut_tensor is + # (num_blocks_axis0, num_blocks_axis1, 2**n_bits, cluster_dim). + lut = palettization_fq.lut + if lut is not None and lut.numel() > 0: + lut_elements = lut.numel() + else: + lut_elements = ( + palettization_fq.granularity.num_blocks_to_cluster(weight) + * (2**palettization_fq.n_bits) + * palettization_fq.cluster_dim + ) + + if palettization_fq.lut_qspec is not None: + lut_dtype_bits = palettization_fq.lut_qspec.n_bits + else: + lut_dtype_bits = _buffer_dtype_bits(lut, weight.element_size() * 8) + lut_bits = lut_elements * lut_dtype_bits + + # Per-channel scale: one weight-dtype value per output channel + # (weight.shape[0]); amortized when enabled, regardless of calibration. + num_channels = weight.shape[0] + per_channel_scale_bits = 0 + if palettization_fq.enable_per_channel_scale: + per_channel_scale = palettization_fq.per_channel_scale + if per_channel_scale is not None and per_channel_scale.numel() > 0: + num_channels = per_channel_scale.numel() + per_channel_scale_bits = num_channels * per_channel_scale.element_size() * 8 + else: + per_channel_scale_bits = num_channels * weight.element_size() * 8 + + return int( + indices_bits + + lut_bits + + per_channel_scale_bits + + _lut_quant_bits(weight, palettization_fq, num_channels) + ) + + +def _lut_quant_bits( + weight: torch.Tensor, palettization_fq: _FakePalettizeImplBase, num_channels: int +) -> int: + """Return the qparams cost of a quantized LUT, in bits. + + Dequantizing a quantized LUT needs one scale per palettization block: the LUT + fake-quantizer overrides ``lut_qspec``'s per-tensor granularity to per-channel + over the stacked LUT. Two behaviors of the export shape the cost: + + - A symmetric zero-point is a single repeated value, which the export emits as + one shared constant at no per-element cost, so only asymmetric zero-points + are counted. + - With per-channel scaling enabled, the LUT scale is fused into the per-channel + scale (see ``palettization/kmeans/_prepare_for_export.py``), so a single scale + tensor ships and is already accounted for by the caller. only the zero-point, + expanded the same way, is extra. + """ + if palettization_fq.lut_qspec is None: + return 0 + + lut_scale = palettization_fq.lut_quantization_scale + if lut_scale is not None and lut_scale.numel() > 0: + num_lut_blocks = lut_scale.numel() + else: + num_lut_blocks = palettization_fq.granularity.num_blocks_to_cluster(weight) + + if palettization_fq.enable_per_channel_scale: + scale_elements, zero_point_elements = 0, num_channels + else: + scale_elements = zero_point_elements = num_lut_blocks + + bits = scale_elements * _buffer_dtype_bits(lut_scale, _DEFAULT_SCALE_BITS) + if palettization_fq.lut_qspec.qscheme == _QuantizationScheme.ASYMMETRIC: + bits += zero_point_elements * palettization_fq.lut_qspec.n_bits + return int(bits) + + +def _buffer_dtype_bits(buffer: torch.Tensor | None, default: int) -> int: + """Return the per-element bit width of a buffer, or a default if unmaterialized. + + Args: + buffer (torch.Tensor | None): A scale, zero-point, or LUT buffer. + default (int): Bit width to assume when the buffer is missing or empty. + + Returns: + int: ``buffer.element_size() * 8`` if the buffer holds data, else + ``default``. + """ + if buffer is not None and buffer.numel() > 0: + return int(buffer.element_size() * 8) + return default diff --git a/src/coreai_opt/inspection/bits_per_weight.py b/src/coreai_opt/inspection/bits_per_weight.py new file mode 100644 index 0000000..d479918 --- /dev/null +++ b/src/coreai_opt/inspection/bits_per_weight.py @@ -0,0 +1,206 @@ +# Copyright 2026 Apple Inc. +# +# Use of this source code is governed by a BSD-3-Clause license that can +# be found in the LICENSE file or at https://opensource.org/licenses/BSD-3-Clause + +"""Compute the average bits-per-weight (bpw) of a prepared ``coreai-opt`` model. + +Estimates the average bit width of the tensors a model carries, amortizing +compression overhead (quantization scales / zero-points, palettization LUTs and +per-channel scales). Compressed +tensors count at their effective compressed cost, everything else (biases, +norms, untargeted weights, buffers such as BatchNorm running stats) counts at +its full-precision dtype cost. + +This is an analytical estimate, not a bit-exact proxy of the model asset. +It answers "what does this compression config cost, in principle?" from +a prepared model. + +Only eager-mode prepared models are supported currently. + +Supported model shapes and compression: + +- Full-precision ``torch.nn.Module``. +- Eager-mode integer weight quantization (int8 / int4 / int2 and their unsigned + variants), symmetric or asymmetric, at any granularity. Sub-byte payloads and + zero-points are packed at ``n_bits`` with no padding, matching the export. +- Palettization at any spec-supported ``n_bits`` (1, 2, 3, 4, 6, 8), including a + quantized LUT (``lut_qspec``). + +Unsupported (raises ``NotImplementedError``): + +- Floating-point weight quantization (FP8 / FP4): the deploy-time storage math + is not yet validated for these formats. +- Pruned models +- Weight parametrizations whose dense tensor is not a single ``original`` + tensor (e.g. ``torch.nn.utils.parametrizations.weight_norm``). +- Graph-mode / ``torch.fx.GraphModule`` models. + +**Notes:** This is an estimate, not a measurement. A real export may differ: + +- It may be *smaller*, because a backend ships only what its graph needs while + this counts every parameter and buffer the model owns, and because a tensor may + be representable more compactly than its shape and dtype imply. A tensor that + ``forward`` never reads or a module that is never called both are accounted for + here but may be skipped in the export. +- It may be *larger*, because a serialized artifact carries structural metadata + that is not modelled here. +- This utility is intended to be used with a prepared ``coreai-opt`` model. Passing + a finalized model to it may result in unpredictable behavior or a wrong bpw value. + +Example: + >>> from coreai_opt.inspection import bits_per_weight + >>> result = bits_per_weight(prepared_model) + >>> result.bpw + 8.86 +""" + +from collections import defaultdict +from dataclasses import dataclass + +import torch +from torch.nn.utils import parametrize as _parametrize + +from coreai_opt._utils.torch_utils import is_float_quant_dtype as _is_float_quant_dtype +from coreai_opt.base_model_compressor import _COREAI_OPT_PREPARED_ATTR as _PREPARED_MARKER +from coreai_opt.quantization.spec.fake_quantize import FakeQuantizeImplBase as _FakeQuantizeImplBase + +from ._bpw_utils import ( + ensure_not_pruned as _ensure_not_pruned, + ensure_single_original as _ensure_single_original, + full_precision_bits as _full_precision_bits, + get_weight_compressor as _get_weight_compressor, + named_modules_excluding_compression_machinery as _named_modules_excluding_compression_machinery, + tensor_storage_bits as _tensor_storage_bits, +) + +__all__ = ["BitsPerWeightResult", "bits_per_weight"] + + +@dataclass +class BitsPerWeightResult: + """Result of a bits-per-weight computation. + + Attributes: + bpw (float): Overall average bits per weight across all parameters + (``total_bits / total_weights``); ``0.0`` if the model has no + parameters. + per_module_map (dict[str, float]): Map from module name to that module's own + average bits per weight. Modules with no logical tensors are omitted. + total_bits (int): Total storage cost in bits, including amortized + compression overhead. + total_weights (int): Total number of logical parameter elements. + """ + + bpw: float + per_module_map: dict[str, float] + total_bits: int + total_weights: int + + def __repr__(self) -> str: + return ( + f"BitsPerWeightResult(bpw={self.bpw:.4f}, " + f"total_bits={self.total_bits}, total_weights={self.total_weights})" + ) + + +def bits_per_weight(model: torch.nn.Module) -> BitsPerWeightResult: + """Compute the average bits-per-weight of a prepared ``coreai-opt`` model. + + Walks the module tree once. For each parametrized weight, the dense original + tensor is counted at its effective compressed cost (quantization or + palettization). Every other directly-owned parameter (biases, norms, + untargeted weights) and every buffer (BatchNorm running stats, RoPE caches, + etc.) are counted at their full-precision dtype cost, regardless of + ``persistent=``, i.e., the metric covers every tensor the model carries. + + Args: + model (torch.nn.Module): A full-precision, eager-mode integer-quantized, + or palettized prepared model. + + Returns: + BitsPerWeightResult: Overall bpw, per-module breakdown, and the totals + used to derive them. + + Raises: + NotImplementedError: If ``model`` is a graph-mode prepared model (a + ``torch.fx.GraphModule``) or a ``torch.export.ExportedProgram``, or if it + contains a weight compression whose storage cost this utility cannot + compute: floating-point (FP8 / FP4) quantization, pruning, or a + parametrization storing multiple original tensors. + """ + if isinstance(model, (torch.fx.GraphModule, torch.export.ExportedProgram)): + raise NotImplementedError( + f"Graph mode prepared models are not supported currently, got {type(model)}. " + "Only full-precision, eager-mode integer quantized, and palettized " + "nn.Modules are handled." + ) + + module_bits: dict[str, int] = defaultdict(int) + module_weights: dict[str, int] = defaultdict(int) + + # id() of every Parameter / Buffer already counted, so a tied tensor + # is counted once + seen_ids: set[int] = set() + + for name, module in _named_modules_excluding_compression_machinery(model): + # Parametrized weights: count the dense original at its compressed cost. + if _parametrize.is_parametrized(module): + for tensor_name, param_list in module.parametrizations.items(): + _ensure_single_original(param_list, name, tensor_name) + _ensure_not_pruned(param_list, name, tensor_name) + + original = param_list.original + if id(original) in seen_ids: + continue + seen_ids.add(id(original)) + + compressor = _get_weight_compressor(param_list) + + if isinstance(compressor, _FakeQuantizeImplBase) and _is_float_quant_dtype( + compressor.target_dtype + ): + raise NotImplementedError( + f"bits_per_weight cannot compute the storage cost of floating-point " + f"weight quantization (dtype {compressor.target_dtype}) on module " + f"'{name}'. Only integer quantization (int8 / int4 / int2 and " + f"their unsigned variants) and palettization are supported currently." + ) + + module_bits[name] += _tensor_storage_bits(original, compressor) + module_weights[name] += original.numel() + + # Directly-owned plain parameters: bias, untargeted weights, norms, etc. + # A parametrized weight is no longer in _parameters (PyTorch moves it into + # the ParametrizationList), so it is not re-counted here. + for param in module.parameters(recurse=False): + if id(param) in seen_ids: + continue + seen_ids.add(id(param)) + module_bits[name] += _full_precision_bits(param) + module_weights[name] += param.numel() + + # Buffers (BatchNorm running stats, RoPE caches, ...) + # recurse=False keeps buf_name un-prefixed, so the marker comparison is + # bare-to-bare. + for buf_name, buf in module.named_buffers(recurse=False): + if buf_name == _PREPARED_MARKER or id(buf) in seen_ids: + continue + seen_ids.add(id(buf)) + module_bits[name] += _full_precision_bits(buf) + module_weights[name] += buf.numel() + + total_bits = sum(module_bits.values()) + total_weights = sum(module_weights.values()) + per_module_map = { + name: bits / module_weights[name] + for name, bits in module_bits.items() + if module_weights.get(name, 0) > 0 + } + bpw = total_bits / total_weights if total_weights > 0 else 0.0 + return BitsPerWeightResult( + bpw=bpw, + per_module_map=per_module_map, + total_bits=total_bits, + total_weights=total_weights, + ) diff --git a/tests/export/export_utils.py b/tests/export/export_utils.py index d622324..4d4d3f6 100644 --- a/tests/export/export_utils.py +++ b/tests/export/export_utils.py @@ -607,3 +607,28 @@ def convert_and_verify( ) return converted_model + + +def coreai_export_size_bytes( + finalized_model: torch.nn.Module, + input_data: torch.Tensor, +) -> int: + """Export a finalized model to a Core AI asset and return its weight payload size. + + Args: + finalized_model: A model finalized to the Core AI backend. + input_data: Example input used to trace the model. + + Returns: + int: Size in bytes of the serialized program. + """ + converter = create_converter(ExportBackend.CoreAI) + traced_model = converter.trace(finalized_model, input_data, expected_ops={}) + program = converter.convert(traced_model, input_data) + with tempfile.TemporaryDirectory(suffix=".aimodel") as tmpdir: + program.save_asset(Path(tmpdir)) + payload_files = list(Path(tmpdir).rglob("*.mlirb")) + assert len(payload_files) == 1, ( + f"Expected exactly one serialized program in the exported bundle, found {payload_files}" + ) + return payload_files[0].stat().st_size diff --git a/tests/export/test_bpw_export_size.py b/tests/export/test_bpw_export_size.py new file mode 100644 index 0000000..a23d019 --- /dev/null +++ b/tests/export/test_bpw_export_size.py @@ -0,0 +1,286 @@ +# Copyright 2026 Apple Inc. +# +# Use of this source code is governed by a BSD-3-Clause license that can +# be found in the LICENSE file or at https://opensource.org/licenses/BSD-3-Clause + +"""Validate the bits-per-weight prediction against the actual Core AI exported +asset size.""" + +import pytest +import torch +import torch.nn as nn + +from coreai_opt import ExportBackend +from coreai_opt.inspection import BitsPerWeightResult, bits_per_weight +from coreai_opt.palettization import ( + KMeansPalettizer, + KMeansPalettizerConfig, + ModuleKMeansPalettizerConfig, +) +from coreai_opt.palettization.spec import ( + PalettizationSpec, + PerGroupedChannelGranularity, + default_weight_palettization_spec, +) +from coreai_opt.quantization import ModuleQuantizerConfig, Quantizer, QuantizerConfig +from coreai_opt.quantization.spec import ( + PerBlockGranularity, + PerChannelGranularity, + QuantizationSpec, +) +from tests.models.simple import GatedMLPModel + +from . import export_utils + +# Fractional budget for what a serialized artifact carries beyond the weight payload. +_MAX_COREAI_EXPORT_OVERHEAD = 0.025 + +# GatedMLPModel at this width holds ~3.1M params (12 MiB fp32, 1.5 MiB at int4). +_MLP_DIM = 1024 + + +def _gated_mlp( + bias: bool = False, extra_buffers: bool = False, persistent_buffers: bool = True +) -> nn.Module: + return GatedMLPModel( + dim=_MLP_DIM, + hidden_dim=_MLP_DIM, + bias=bias, + extra_buffers=extra_buffers, + persistent_buffers=persistent_buffers, + ) + + +def _mlp_input() -> torch.Tensor: + return torch.rand(1, 4, _MLP_DIM) + + +def _eager_quant_config( + dtype: torch.dtype, qscheme: str, granularity: object | None = None +) -> QuantizerConfig: + """Build an eager-mode weight-quantization config, per-channel by default.""" + spec = QuantizationSpec( + dtype=dtype, + qscheme=qscheme, + granularity=granularity or PerChannelGranularity(axis=0), + ) + return QuantizerConfig( + global_config=ModuleQuantizerConfig(op_state_spec={"weight": spec}, op_input_spec=None), + execution_mode="eager", + ) + + +def _assert_prediction_matches_export( + finalized_model: nn.Module, + input_data: torch.Tensor, + result: BitsPerWeightResult, +) -> None: + """Assert the bits-per-weight prediction matches the exported asset size. + + The lower bound holds only for fixtures whose tensors are all reachable from + ``forward`` and randomly initialized. The export ships what the graph needs, + and may shrink tensors it can represent more compactly, so a fixture with a + dead buffer or uniform-valued weights would export below the prediction. + """ + predicted_bytes = result.total_bits / 8 + actual_bytes = export_utils.coreai_export_size_bytes(finalized_model, input_data) + actual_bpw = actual_bytes * 8 / result.total_weights + overhead_bytes = actual_bytes - predicted_bytes + context = f"predicted bpw={result.bpw:.4f}, actual bpw={actual_bpw:.4f}" + + assert overhead_bytes >= 0, ( + f"exported asset {actual_bytes:,} bytes fell below predicted " + f"{predicted_bytes:,.0f} by {-overhead_bytes:,.0f} bytes ({context})" + ) + # The only excess should be structural metadata, which stays a small fraction + # of the payload. + assert overhead_bytes <= predicted_bytes * _MAX_COREAI_EXPORT_OVERHEAD, ( + f"exported asset {actual_bytes:,} bytes exceeded predicted " + f"{predicted_bytes:,.0f} by {overhead_bytes:,.0f} bytes " + f"({overhead_bytes / predicted_bytes:.2%}, budget {_MAX_COREAI_EXPORT_OVERHEAD:.1%}) " + f"({context})" + ) + + +def _assert_quantized_matches_export( + model: nn.Module, + input_data: torch.Tensor, + dtype: torch.dtype, + qscheme: str = "symmetric", + granularity: object | None = None, +) -> None: + """Quantize and check the prediction against the export.""" + quantizer = Quantizer(model, _eager_quant_config(dtype, qscheme, granularity)) + prepared_model = quantizer.prepare((input_data,)) + + result = bits_per_weight(prepared_model) + finalized_model = quantizer.finalize(backend=ExportBackend.CoreAI) + _assert_prediction_matches_export(finalized_model, input_data, result) + + +def _assert_palettized_matches_export( + model: nn.Module, input_data: torch.Tensor, spec: PalettizationSpec +) -> None: + """Palettize and check the prediction against the export.""" + config = KMeansPalettizerConfig( + global_config=ModuleKMeansPalettizerConfig(op_state_spec={"weight": spec}) + ) + palettizer = KMeansPalettizer(model, config) + prepared_model = palettizer.prepare((input_data,)) + + result = bits_per_weight(prepared_model) + finalized_model = palettizer.finalize(backend=ExportBackend.CoreAI) + _assert_prediction_matches_export(finalized_model, input_data, result) + + +@pytest.mark.parametrize( + ("dtype", "qscheme"), + [ + (torch.int8, "symmetric"), + (torch.int8, "asymmetric"), + (torch.int4, "symmetric"), + (torch.int4, "asymmetric"), + ], + ids=["int8_symmetric", "int8_asymmetric", "int4_symmetric", "int4_asymmetric"], +) +def test_eager_quant_prediction_matches_export(dtype: torch.dtype, qscheme: str) -> None: + """Per-channel weight quantization: predicted deploy size matches the export.""" + # bias=True so the fp32 biases are amortized into the prediction too. + _assert_quantized_matches_export(_gated_mlp(bias=True), _mlp_input(), dtype, qscheme) + + +@pytest.mark.parametrize( + "spec", + [ + PalettizationSpec(n_bits=2), + PalettizationSpec(n_bits=4), + PalettizationSpec(n_bits=8), + PalettizationSpec(n_bits=4, granularity=PerGroupedChannelGranularity(axis=0, group_size=8)), + PalettizationSpec( + n_bits=4, granularity=PerGroupedChannelGranularity(axis=0, group_size=32) + ), + ], + ids=[ + "per_tensor_n2", + "per_tensor_n4", + "per_tensor_n8", + "per_grouped_channel_group8", + "per_grouped_channel_group32", + ], +) +def test_palettized_prediction_matches_export(spec: PalettizationSpec) -> None: + """Weight palettization across bit widths and granularities. + + Per-grouped-channel multiplies the LUT count by the number of channel groups, so + it is what exercises the ``num_blocks_to_cluster`` term in the LUT cost. + bias=False isolates the palettized weights so the deploy size is dominated by the + indices and LUTs, not amortized fp32 biases. + """ + _assert_palettized_matches_export(_gated_mlp(bias=False), _mlp_input(), spec) + + +@pytest.mark.parametrize( + "weight_dtype", [torch.float32, torch.float16], ids=["fp32_weights", "fp16_weights"] +) +@pytest.mark.parametrize("dtype", [torch.int4, torch.int2], ids=["int4", "int2"]) +def test_perblock_asymmetric_subbyte_prediction_matches_export( + dtype: torch.dtype, weight_dtype: torch.dtype +) -> None: + """Per-block asymmetric sub-byte weights: zero-points are packed at ``n_bits``. + + Per-channel granularity keeps the zero-point term small enough to hide its + width; per-block ``block_size=32`` makes it ~3% of the payload, so this is what + catches a zero-point charged at ``target_dtype``'s byte width (``element_size()`` + is 1 for int4 and int2 alike) instead of at ``n_bits``. + + The ``fp16_weights`` case additionally pins the scale width to the weight dtype. + """ + _assert_quantized_matches_export( + _gated_mlp(bias=False).to(weight_dtype), + _mlp_input().to(weight_dtype), + dtype, + "asymmetric", + granularity=PerBlockGranularity(axis=1, block_size=32), + ) + + +def test_quantized_lut_prediction_matches_export() -> None: + """Palettization with a quantized LUT: the LUT's own qparams are amortized too. + + One group per channel at 1 bit maximizes the qparam-to-payload ratio (one LUT + scale and zero-point per palettization block against a 1-bit index payload), so + omitting the LUT-quantization qparams shows up as a ~4% shortfall here while + staying under 0.1% for the default per-tensor 4-bit config. + """ + _assert_palettized_matches_export( + _gated_mlp(bias=False), + _mlp_input(), + PalettizationSpec( + n_bits=1, + granularity=PerGroupedChannelGranularity(axis=0, group_size=1), + lut_qspec=QuantizationSpec(dtype=torch.uint8, qscheme="asymmetric"), + ), + ) + + +@pytest.mark.parametrize("dtype", [torch.int8, torch.int4], ids=["int8", "int4"]) +def test_resnet18_quant_prediction_matches_export( + resnet18_model: nn.Module, + resnet_example_input: torch.Tensor, + dtype: torch.dtype, +) -> None: + """Pretrained ResNet-18 quantized: a deep real model with ~50 leaf modules.""" + _assert_quantized_matches_export(resnet18_model, resnet_example_input, dtype) + + +def test_resnet18_palettized_prediction_matches_export( + resnet18_model: nn.Module, + resnet_example_input: torch.Tensor, +) -> None: + """Pretrained ResNet-18 palettized at the default 4 bits. + + Only the default config, matching ``test_kmeans_export.test_resnet_export``: the + per-``n_bits`` matrix runs on the faster synthetic model instead. + """ + _assert_palettized_matches_export( + resnet18_model, resnet_example_input, default_weight_palettization_spec() + ) + + +@pytest.mark.parametrize("dtype", [torch.float32, torch.float16], ids=["fp32", "fp16"]) +def test_mnist_dense_prediction_matches_export( + custom_test_mnist_model: nn.Module, + mnist_example_input: torch.Tensor, + dtype: torch.dtype, +) -> None: + """Real conv/BN/linear model, dense fp32 and fp16: the export matches prediction. + + Exercises persistent BatchNorm buffers (running_mean / running_var), which are + counted by bits_per_weight and survive export as constants rather than being + folded away. + """ + model = custom_test_mnist_model.to(dtype) + model.eval() + input_data = mnist_example_input.to(dtype) + + result = bits_per_weight(model) + _assert_prediction_matches_export(model, input_data, result) + + +@pytest.mark.parametrize("dtype", [torch.int8, torch.int4], ids=["int8", "int4"]) +def test_mnist_quant_prediction_matches_export( + custom_test_mnist_model: nn.Module, + mnist_example_input: torch.Tensor, + dtype: torch.dtype, +) -> None: + """Real conv/BN/linear model, int8 and int4 per-channel weights.""" + _assert_quantized_matches_export(custom_test_mnist_model, mnist_example_input, dtype) + + +@pytest.mark.parametrize("persistent", [True, False], ids=["persistent", "non_persistent"]) +def test_buffers_are_accounted_for_in_export(persistent: bool) -> None: + """Buffers ship at full precision even when the weights are compressed and + it reflects in BPW.""" + _assert_quantized_matches_export( + _gated_mlp(extra_buffers=True, persistent_buffers=persistent), _mlp_input(), torch.int4 + ) diff --git a/tests/inspection/test_bits_per_weight.py b/tests/inspection/test_bits_per_weight.py new file mode 100644 index 0000000..3cda834 --- /dev/null +++ b/tests/inspection/test_bits_per_weight.py @@ -0,0 +1,440 @@ +# Copyright 2026 Apple Inc. +# +# Use of this source code is governed by a BSD-3-Clause license that can +# be found in the LICENSE file or at https://opensource.org/licenses/BSD-3-Clause + +"""Tests for the bits_per_weight utility.""" + +import pytest +import torch +import torch.nn as nn + +from coreai_opt.inspection import bits_per_weight +from coreai_opt.palettization import ( + KMeansPalettizer, + KMeansPalettizerConfig, + ModuleKMeansPalettizerConfig, +) +from coreai_opt.palettization.spec import ( + PalettizationSpec, + PerGroupedChannelGranularity, + PerTensorGranularity as PalettPerTensorGranularity, +) +from coreai_opt.pruning import MagnitudePruner +from coreai_opt.quantization import ModuleQuantizerConfig, Quantizer, QuantizerConfig +from coreai_opt.quantization.spec import ( + PerBlockGranularity, + PerChannelGranularity, + PerTensorGranularity, + QuantizationSpec, +) +from tests.models.simple import LinearBatchNormModel, SharedParamsModel, SimpleLinearModel + +# SimpleLinearModel and LinearBatchNormModel are both Linear(64, 128) -> Linear(128, 64), +# so their weights are 128 x 64 and 64 x 128. +# Expected bits and num_weights are derived by +# hand from these shapes and used as the golden values for testing the utility. +_IN_FEATURES = 64 +_HIDDEN_FEATURES = 128 +_OUT_FEATURES = 64 +_WEIGHT_ELEMS = _HIDDEN_FEATURES * _IN_FEATURES + _OUT_FEATURES * _HIDDEN_FEATURES + +_BIAS_ELEMS = _HIDDEN_FEATURES + _OUT_FEATURES +# One per-channel qparam (quantization scale, or palettization per-channel scale) per +# output channel of each layer. +_OUT_CHANNELS = _HIDDEN_FEATURES + _OUT_FEATURES +_FP32_BITS = 32 +_INT64_BITS = 64 + +# Sanity envelope for the qparam / LUT / bias overhead a sane config adds on top of the +# nominal bit width, in bpw. See _assert_bpw_is_plausible for why 2 and not 1. +_MAX_EXPECTED_OVERHEAD_BPW = 2.0 + +_EXAMPLE_INPUT = torch.rand(4, _IN_FEATURES) + + +def _expected_elems(bias: bool) -> int: + """Logical parameter elements: the dense weights, plus biases when present.""" + return _WEIGHT_ELEMS + (_BIAS_ELEMS if bias else 0) + + +def _expected_quant_bits( + n_bits: int, + num_qparam_blocks: int, + qscheme: str, + bias: bool, + qformulation: str = "zp", +) -> int: + """Hand-derived cost of the weight-quantized model, in bits. + + Payload is ``n_bits`` per weight. Overhead is one fp32 scale per qparam block, the + per-block dequantization offset, and the fp32 biases, which quantization does not + target. + """ + scale_bits = num_qparam_blocks * _FP32_BITS + if qformulation == "minval": + offset_bits = num_qparam_blocks * _FP32_BITS + elif qscheme == "asymmetric": + offset_bits = num_qparam_blocks * n_bits + else: + offset_bits = 0 + bias_bits = _BIAS_ELEMS * _FP32_BITS if bias else 0 + return _WEIGHT_ELEMS * n_bits + scale_bits + offset_bits + bias_bits + + +def _expected_palettized_bits( + n_bits: int, num_lut_blocks: int, bias: bool, per_channel_scale: bool +) -> int: + """Hand-derived cost of the palettized model, in bits. + + Payload is one ``n_bits`` index per weight (``cluster_dim=1``). Overhead is one LUT + of ``2**n_bits`` fp32 centroids per LUT block, one fp32 per-channel scale per output + channel when that is enabled, and the fp32 biases, which palettization does not + target. + """ + lut_bits = num_lut_blocks * (2**n_bits) * _FP32_BITS + per_channel_scale_bits = _OUT_CHANNELS * _FP32_BITS if per_channel_scale else 0 + bias_bits = _BIAS_ELEMS * _FP32_BITS if bias else 0 + return _WEIGHT_ELEMS * n_bits + lut_bits + per_channel_scale_bits + bias_bits + + +def _assert_bpw_is_plausible(bpw: float, n_bits: int) -> None: + """Structural sanity bounds on bpw, independent of how the cost is modelled.""" + assert n_bits <= bpw < n_bits + _MAX_EXPECTED_OVERHEAD_BPW + + +def _prepare_eager_quant( + model: nn.Module, + dtype: torch.dtype, + qscheme: str = "symmetric", + granularity: object | None = None, + qformulation: str = "zp", +) -> nn.Module: + """Prepare an eager-mode weight-quantized model.""" + spec = QuantizationSpec( + dtype=dtype, + qscheme=qscheme, + granularity=granularity or PerChannelGranularity(axis=0), + qformulation=qformulation, + ) + config = QuantizerConfig( + global_config=ModuleQuantizerConfig(op_state_spec={"weight": spec}, op_input_spec=None), + execution_mode="eager", + ) + return Quantizer(model, config).prepare(example_inputs=(_EXAMPLE_INPUT,)) + + +def _prepare_palettized(model: nn.Module, spec: PalettizationSpec) -> nn.Module: + """Prepare a palettized model.""" + config = KMeansPalettizerConfig( + global_config=ModuleKMeansPalettizerConfig(op_state_spec={"weight": spec}) + ) + return KMeansPalettizer(model, config).prepare((_EXAMPLE_INPUT,)) + + +def test_full_precision_bits_per_weight(): + assert bits_per_weight(SimpleLinearModel()).bpw == 32.0 + assert bits_per_weight(SimpleLinearModel().half()).bpw == 16.0 + assert bits_per_weight(SimpleLinearModel().bfloat16()).bpw == 16.0 + + +@pytest.mark.parametrize("bias", [False, True], ids=["no_bias", "bias"]) +@pytest.mark.parametrize("qscheme", ["symmetric", "asymmetric"]) +@pytest.mark.parametrize( + ("granularity", "num_qparam_blocks", "qformulation"), + [ + pytest.param(PerTensorGranularity(axis=None), 2, "zp", id="per_tensor_zp"), + pytest.param(PerTensorGranularity(axis=None), 2, "minval", id="per_tensor_minval"), + pytest.param(PerChannelGranularity(axis=0), _OUT_CHANNELS, "zp", id="per_channel_zp"), + pytest.param( + PerChannelGranularity(axis=0), _OUT_CHANNELS, "minval", id="per_channel_minval" + ), + # One qparam block per 32 weights along the input-feature axis of each layer. + pytest.param( + PerBlockGranularity(axis=1, block_size=32), + _HIDDEN_FEATURES * (_IN_FEATURES // 32) + _OUT_FEATURES * (_HIDDEN_FEATURES // 32), + "zp", + id="per_block_zp", + ), + ], +) +@pytest.mark.parametrize( + ("dtype", "n_bits"), + [ + (torch.int8, 8), + (torch.int4, 4), + (torch.int2, 2), + (torch.uint8, 8), + (torch.uint4, 4), + (torch.uint2, 2), + ], + ids=["int8", "int4", "int2", "uint8", "uint4", "uint2"], +) +def test_eager_quant_matches_analytical( + dtype: torch.dtype, + n_bits: int, + granularity: object, + num_qparam_blocks: int, + qformulation: str, + qscheme: str, + bias: bool, +): + prepared = _prepare_eager_quant( + SimpleLinearModel(bias=bias), dtype, qscheme, granularity, qformulation + ) + result = bits_per_weight(prepared) + + expected_bits = _expected_quant_bits(n_bits, num_qparam_blocks, qscheme, bias, qformulation) + expected_elems = _expected_elems(bias) + assert result.total_bits == expected_bits + # The inserted scale / zero-point buffers must not inflate the denominator: only the + # original dense parameters are counted. + assert result.total_weights == expected_elems + assert result.bpw == expected_bits / expected_elems + _assert_bpw_is_plausible(result.bpw, n_bits) + + +@pytest.mark.parametrize("bias", [False, True], ids=["no_bias", "bias"]) +@pytest.mark.parametrize("per_channel_scale", [False, True], ids=["no_pcs", "pcs"]) +@pytest.mark.parametrize( + ("n_bits", "granularity", "num_lut_blocks"), + [ + pytest.param(n_bits, granularity, num_lut_blocks, id=f"n{n_bits}_{granularity_id}") + for n_bits in (1, 2, 4) + for granularity, num_lut_blocks, granularity_id in ( + (PalettPerTensorGranularity(axis=None), 2, "per_tensor"), + ( + PerGroupedChannelGranularity(axis=0, group_size=8), + _HIDDEN_FEATURES // 8 + _OUT_FEATURES // 8, + "group8", + ), + ( + PerGroupedChannelGranularity(axis=0, group_size=32), + _HIDDEN_FEATURES // 32 + _OUT_FEATURES // 32, + "group32", + ), + ) + ] + # n_bits=8 only at per-tensor granularity. Its grouped variants carry more LUT than + # payload, so they break the sanity envelope and live in + # test_overhead_heavy_configs_exceed_bitwidth instead. + + [pytest.param(8, PalettPerTensorGranularity(axis=None), 2, id="n8_per_tensor")], +) +def test_palettized_matches_analytical( + n_bits: int, + granularity: object, + num_lut_blocks: int, + per_channel_scale: bool, + bias: bool, +): + prepared = _prepare_palettized( + SimpleLinearModel(bias=bias), + PalettizationSpec( + n_bits=n_bits, + granularity=granularity, + enable_per_channel_scale=per_channel_scale, + ), + ) + result = bits_per_weight(prepared) + + expected_bits = _expected_palettized_bits(n_bits, num_lut_blocks, bias, per_channel_scale) + expected_elems = _expected_elems(bias) + assert result.total_bits == expected_bits + # The inserted LUT buffers must not inflate the denominator. + assert result.total_weights == expected_elems + assert result.bpw == expected_bits / expected_elems + _assert_bpw_is_plausible(result.bpw, n_bits) + + +@pytest.mark.parametrize( + ("compression", "dtype", "n_bits", "granularity", "num_blocks"), + [ + # One fp32 scale and zero-point per 8 weights: 4.25 bpw of qparams over a 2 bpw + # payload. + pytest.param( + "quantization", + torch.int2, + 2, + PerBlockGranularity(axis=1, block_size=8), + _HIDDEN_FEATURES * (_IN_FEATURES // 8) + _OUT_FEATURES * (_HIDDEN_FEATURES // 8), + id="quant_int2_block8", + ), + pytest.param( + "quantization", + torch.int4, + 4, + PerBlockGranularity(axis=1, block_size=4), + _HIDDEN_FEATURES * (_IN_FEATURES // 4) + _OUT_FEATURES * (_HIDDEN_FEATURES // 4), + id="quant_int4_block4", + ), + # One 2**8-entry fp32 LUT per 8 output channels: 12 bpw of LUT over an 8 bpw + # payload. Palettization takes n_bits directly, so it needs no dtype. + pytest.param( + "palettization", + None, + 8, + PerGroupedChannelGranularity(axis=0, group_size=8), + _HIDDEN_FEATURES // 8 + _OUT_FEATURES // 8, + id="palett_n8_group8", + ), + pytest.param( + "palettization", + None, + 8, + PerGroupedChannelGranularity(axis=0, group_size=32), + _HIDDEN_FEATURES // 32 + _OUT_FEATURES // 32, + id="palett_n8_group32", + ), + ], +) +def test_overhead_heavy_configs_exceed_bitwidth( + compression: str, + dtype: torch.dtype | None, + n_bits: int, + granularity: object, + num_blocks: int, +): + """Configs whose qparam or LUT overhead outweighs the payload it serves. + + So the bpw will exceed n_bits by a non-trivial amount. These are in-efficient + compression configs which the utility should be agnostic to. And their bpw values + will not be close to n_bits, so they will not adhere to the + ``_assert_bpw_is_plausible`` check that the above tests perform. + """ + model = SimpleLinearModel(bias=False) + if compression == "quantization": + prepared = _prepare_eager_quant(model, dtype, "asymmetric", granularity) + expected_bits = _expected_quant_bits(n_bits, num_blocks, "asymmetric", bias=False) + else: + prepared = _prepare_palettized( + model, PalettizationSpec(n_bits=n_bits, granularity=granularity) + ) + expected_bits = _expected_palettized_bits( + n_bits, num_blocks, bias=False, per_channel_scale=False + ) + result = bits_per_weight(prepared) + + assert result.total_bits == expected_bits + assert result.total_weights == _expected_elems(bias=False) + assert n_bits + _MAX_EXPECTED_OVERHEAD_BPW < result.bpw < _FP32_BITS + + +def test_persistent_buffers_counted(): + """BatchNorm running stats ship in ``state_dict()``, so they are amortized in.""" + result = bits_per_weight(LinearBatchNormModel()) + + # Params: the two fp32 weights (the model is bias-free), plus BatchNorm's fp32 + # weight and bias. + param_elems = _WEIGHT_ELEMS + 2 * _HIDDEN_FEATURES + # Buffers: fp32 running_mean and running_var, one of each per hidden feature, plus + # a scalar int64 num_batches_tracked. + buffer_elems = 2 * _HIDDEN_FEATURES + 1 + + assert result.total_weights == param_elems + buffer_elems + assert result.total_bits == ( + param_elems * _FP32_BITS + 2 * _HIDDEN_FEATURES * _FP32_BITS + _INT64_BITS + ) + + +def test_multi_original_parametrization_is_unsupported(): + """``weight_norm`` stores original0 / original1 rather than a single original. + + There is no one dense tensor to attribute a storage cost to, so this must raise + rather than fail with ``AttributeError``. + """ + model = nn.utils.parametrizations.weight_norm(nn.Linear(_IN_FEATURES, _HIDDEN_FEATURES)) + + with pytest.raises(NotImplementedError, match="multiple original tensors"): + bits_per_weight(model) + + +def test_pruned_weight_is_unsupported(): + """A pruned weight raises rather than being priced as a dense tensor.""" + model = SimpleLinearModel(bias=False) + prepared = MagnitudePruner(model).prepare((_EXAMPLE_INPUT,)) + + with pytest.raises(NotImplementedError, match="pruned"): + bits_per_weight(prepared) + + +def test_non_persistent_buffer_counted(): + """Buffers count regardless of ``persistent=``: every tensor the model carries.""" + + class _ModuleWithScratch(nn.Module): + def __init__(self) -> None: + super().__init__() + self.l = nn.Linear(8, 8, bias=False) + self.register_buffer("scratch", torch.zeros(1000), persistent=False) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return self.l(x) + self.scratch[: x.shape[-1]] + + result = bits_per_weight(_ModuleWithScratch()) + assert result.total_weights == 8 * 8 + 1000 + assert result.total_bits == (8 * 8 + 1000) * 32 + + +@pytest.mark.parametrize( + ("dtype", "granularity"), + [ + (torch.float8_e4m3fn, PerChannelGranularity(axis=0)), + (torch.float4_e2m1fn_x2, PerBlockGranularity(axis=1, block_size=32)), + ], + ids=["fp8_per_channel", "fp4_per_block"], +) +def test_float_quant_is_unsupported(dtype, granularity): + spec = QuantizationSpec(dtype=dtype, qscheme="symmetric", granularity=granularity) + config = QuantizerConfig( + global_config=ModuleQuantizerConfig(op_state_spec={"weight": spec}, op_input_spec=None), + execution_mode="eager", + ) + prepared = Quantizer(SimpleLinearModel(), config).prepare(example_inputs=(_EXAMPLE_INPUT,)) + + with pytest.raises(NotImplementedError, match="floating-point weight quantization"): + bits_per_weight(prepared) + + +def test_per_module_attributes_cost_to_the_owning_module(): + """Under mixed precision, each module's bpw lands on that module.""" + + spec = QuantizationSpec( + dtype=torch.int8, qscheme="symmetric", granularity=PerChannelGranularity(axis=0) + ) + config = QuantizerConfig( + global_config=ModuleQuantizerConfig(op_state_spec={"weight": spec}, op_input_spec=None), + # l2 stays fp32 while l1 is int8 + module_name_configs={"l2": None}, + execution_mode="eager", + ) + prepared = Quantizer(SimpleLinearModel(bias=False), config).prepare( + example_inputs=(_EXAMPLE_INPUT,) + ) + prepared(_EXAMPLE_INPUT) + result = bits_per_weight(prepared) + + # l1 carries 8 bits per weight plus one fp32 scale per output channel. + l1_elems = _HIDDEN_FEATURES * _IN_FEATURES + expected_l1_bits = l1_elems * 8 + _HIDDEN_FEATURES * _FP32_BITS + l2_elems = _OUT_FEATURES * _HIDDEN_FEATURES + + assert result.per_module_map == { + "l1": pytest.approx(expected_l1_bits / l1_elems), + "l2": _FP32_BITS, + } + + assert result.total_bits == expected_l1_bits + l2_elems * _FP32_BITS + assert result.bpw == result.total_bits / (l1_elems + l2_elems) + + assert result.per_module_map["l1"] < result.bpw < result.per_module_map["l2"] + + +def test_per_module_counts_a_tied_weight_once(): + """A weight shared by several modules is charged only to the first one.""" + per_module_map = bits_per_weight(SharedParamsModel()).per_module_map + + assert per_module_map["shared_linear"] == _FP32_BITS + assert "layer1" not in per_module_map + assert "layer2" not in per_module_map + # Modules that own tensors of their own are unaffected. + assert per_module_map["input_layer"] == _FP32_BITS + assert per_module_map["output"] == _FP32_BITS diff --git a/tests/models/simple.py b/tests/models/simple.py index 46b7375..bfea1e9 100644 --- a/tests/models/simple.py +++ b/tests/models/simple.py @@ -102,16 +102,31 @@ class GatedMLPModel(nn.Module): Inspired by the MLP block in Qwen3 and similar transformer architectures. """ - def __init__(self, dim: int = 32, hidden_dim: int = 64) -> None: + def __init__( + self, + dim: int = 32, + hidden_dim: int = 64, + bias: bool = False, + extra_buffers: bool = False, + persistent_buffers: bool = True, + ) -> None: super().__init__() - self.gate_proj = nn.Linear(dim, hidden_dim, bias=False) - self.up_proj = nn.Linear(dim, hidden_dim, bias=False) - self.down_proj = nn.Linear(hidden_dim, dim, bias=False) + self.gate_proj = nn.Linear(dim, hidden_dim, bias=bias) + self.up_proj = nn.Linear(dim, hidden_dim, bias=bias) + self.down_proj = nn.Linear(hidden_dim, dim, bias=bias) + + self.extra_buffers = extra_buffers + if extra_buffers: + self.register_buffer("proj", torch.randn(dim, dim), persistent=persistent_buffers) + self.register_buffer("shift", torch.randn(dim), persistent=persistent_buffers) def forward(self, x: torch.Tensor) -> torch.Tensor: up_tensor = self.up_proj(x) gate_tensor = nn.functional.silu(self.gate_proj(x)) - return self.down_proj(up_tensor * gate_tensor) + out = self.down_proj(up_tensor * gate_tensor) + if self.extra_buffers: + out = torch.matmul(out, self.proj) + self.shift + return out @pytest.fixture @@ -127,10 +142,10 @@ def gated_mlp_model_input(): class SimpleLinearModel(torch.nn.Module): - def __init__(self): + def __init__(self, bias: bool = True): super().__init__() - self.l1 = nn.Linear(64, 128) - self.l2 = nn.Linear(128, 64) + self.l1 = nn.Linear(64, 128, bias=bias) + self.l2 = nn.Linear(128, 64, bias=bias) def forward(self, x): x = self.l1(x) @@ -150,6 +165,30 @@ def simple_linear_model_input(): return torch.randn(4, 64) +class LinearBatchNormModel(torch.nn.Module): + """Linear model with a BatchNorm between its two layers. + + Same layer shapes as :class:`SimpleLinearModel`, plus BatchNorm's persistent + buffers (``running_mean``, ``running_var``, ``num_batches_tracked``), which makes + it useful for testing behavior that depends on buffers and not on parameters alone. + """ + + def __init__(self, bias: bool = False): + super().__init__() + self.l1 = nn.Linear(64, 128, bias=bias) + self.bn = nn.BatchNorm1d(128) + self.l2 = nn.Linear(128, 64, bias=bias) + + def forward(self, x): + return self.l2(self.bn(self.l1(x))) + + +@pytest.fixture +def linear_batchnorm_model(): + """Fixture providing a linear model carrying BatchNorm running-stat buffers.""" + return LinearBatchNormModel() + + class SimpleMHAModel(nn.Module): def __init__(self, embed_dim=64, num_heads=4): super().__init__()