diff --git a/python/src/coreai_models/export/_constants.py b/python/src/coreai_models/_constants.py similarity index 64% rename from python/src/coreai_models/export/_constants.py rename to python/src/coreai_models/_constants.py index 94e1e588..7e790c5e 100644 --- a/python/src/coreai_models/export/_constants.py +++ b/python/src/coreai_models/_constants.py @@ -3,15 +3,21 @@ # 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 -"""Constants for the export pipeline.""" +"""Graph and runner contract constants. + +A leaf module: imported by both ``models/`` and ``export/``, imports nothing from +either. +""" + +# Graph name for a single-graph (macOS) export. iOS uses its own entrypoint names. +MAIN_GRAPH_NAME = "main" # KV cache names used by the Swift runner KEY_CACHE_NAME = "keyCache" VALUE_CACHE_NAME = "valueCache" -# Trace-time KV cache sequence length. Used only for export/quantization tracing -# to bound peak memory; at inference the actual cache size is determined -# dynamically. +# Trace-time KV cache sequence length, to bound peak trace memory. At inference the +# cache size is dynamic. TRACE_KV_CACHE_SEQ_LEN = 2048 # Trace-time `input_ids` length and `position_ids` offset for export/quantization diff --git a/python/src/coreai_models/export/compression.py b/python/src/coreai_models/export/compression.py index 4f5cfd41..bd83f513 100644 --- a/python/src/coreai_models/export/compression.py +++ b/python/src/coreai_models/export/compression.py @@ -11,15 +11,18 @@ """ import logging -from collections.abc import Callable +from collections.abc import Callable, Sequence import torch import torch.nn as nn -from coreai_models.export._constants import ( +from coreai_models._constants import ( + MAIN_GRAPH_NAME, QUANT_TRACE_OFFSET, QUANT_TRACE_QUERY_LEN, + TRACE_KV_CACHE_SEQ_LEN, ) +from coreai_models.models.base import BaseForCausalLM, TraceSpec logger = logging.getLogger(__name__) @@ -100,6 +103,8 @@ def quantize_pytorch_model( inputs: tuple, dynamic_shapes: dict, quantization_config: dict, + cache_seq_len: int, + state_indices: Sequence[int], calibration_data_fn: Callable[[], list] | None = None, export_backend: object | None = None, mmap_dir: str | None = None, @@ -119,6 +124,10 @@ def quantize_pytorch_model( coreai-opt expects under `quantization_config`. Includes a `calibrate_activations` key (popped here before constructing the coreai-opt config). + cache_seq_len: Sequence-dim length the caches in ``inputs`` were traced at, + used to bound the calibration query length. + state_indices: Positions in ``inputs`` that are state and must be reset + between calibration samples. calibration_data_fn: Optional function that returns calibration data samples. Required when calibrate_activations is enabled. export_backend: Backend for the finalized quantized model. @@ -142,7 +151,11 @@ def quantize_pytorch_model( # When doing activation quantization, run real calibration data through the # prepared model so the activation observers see representative ranges. - # `inputs` follows the model forward contract: (input_ids, position_ids, k_cache, v_cache). + # + # `inputs[0]` must be input_ids and `inputs[1]` position_ids -- token-based + # calibration cannot do anything else. Which of the rest are state comes from + # `state_indices`; a non-state input keeps its traced tensor, which is only correct + # if its shape does not depend on query length. if run_calibration: if calibration_data_fn is None: raise ValueError( @@ -151,9 +164,21 @@ def quantize_pytorch_model( calibration_data = calibration_data_fn() device = next(model.parameters()).device - cache_seq_len = inputs[2].shape[-2] - # Match the dynamic-shape upper bound declared by the caller: - # position_ids.shape[1] <= cache_seq_len - 1 (see pipeline.py `seq_pos` Dim) + reset_positions = set(state_indices) + for pos in reset_positions: + if pos < 2: + raise ValueError( + "States cannot occupy the first two input positions. " + "Those must be reserved for input_ids and position_ids" + ) + if pos >= len(inputs): + raise IndexError( + f"State index out of bounds, got {pos}, while the number of inputs is " + f"{len(inputs)}" + ) + + # Match the caller's declared bound: position_ids.shape[1] <= cache_seq_len - 1 + # (the `seq_pos` Dim in `BaseForCausalLM.build_dynamic_shapes`). # position_ids has length QUANT_TRACE_OFFSET + query_len, so: # query_len <= cache_seq_len - QUANT_TRACE_OFFSET - 1 max_calib_query_len = cache_seq_len - QUANT_TRACE_OFFSET - 1 @@ -162,16 +187,20 @@ def quantize_pytorch_model( min_calib_query_len = QUANT_TRACE_QUERY_LEN - QUANT_TRACE_OFFSET def _prep_calib_inputs(sample: torch.Tensor) -> tuple: - sample = sample[:, :max_calib_query_len].to(device) - position_ids = ( - torch.arange(QUANT_TRACE_OFFSET + sample.shape[1], dtype=torch.int32) + prepared = list(inputs) + prepared[0] = sample[:, :max_calib_query_len].to(device) + prepared[1] = ( + torch.arange(QUANT_TRACE_OFFSET + prepared[0].shape[1], dtype=torch.int32) .unsqueeze(0) .to(device) ) - zero_cache = tuple( - torch.zeros(inp.shape, dtype=inp.dtype, device=device) for inp in inputs[2:] - ) - return (sample, position_ids, *zero_cache) + for i in range(2, len(prepared)): + inp = inputs[i] + if i in reset_positions: + prepared[i] = torch.zeros(inp.shape, dtype=inp.dtype, device=device) + elif isinstance(inp, torch.Tensor): + prepared[i] = inp.to(device) + return tuple(prepared) calibration_data = [s for s in calibration_data if s.shape[1] >= min_calib_query_len] if not calibration_data: @@ -203,6 +232,66 @@ def _prep_calib_inputs(sample: torch.Tensor) -> tuple: return finalized_model +def quantize_for_export( + model: BaseForCausalLM, + config, + target_dtype: torch.dtype, + quantization_config: dict, + calibration_data_fn: Callable[[], list] | None = None, + mmap_dir: str | None = None, +) -> nn.Module: + """Apply pre-export torch quantization using the model's own graph contract. + + Builds the calibration trace from the export hooks rather than hardcoding a forward + signature, so a model with extra inputs or states calibrates without the caller + knowing about them, and activation calibration resets exactly the states. + + Args: + model: The loaded model, in eval mode. + config: The config the model was built from. + target_dtype: Dtype for the trace's cache tensors. + quantization_config: Inner coreai-opt ``quantization_config`` dict. + calibration_data_fn: Calibration samples; required when the recipe enables + ``calibrate_activations``. + mmap_dir: Directory for the quantizer's disk checkpointing. + """ + spec = TraceSpec(max_context_length=TRACE_KV_CACHE_SEQ_LEN) + reference_inputs = model.build_reference_inputs(config, target_dtype, spec) + dynamic_shapes = model.build_dynamic_shapes(config, spec) + # Same check the export path runs, so a bad contract fails identically on both. + model.validate_export_contract(reference_inputs, dynamic_shapes) + + graph_inputs = reference_inputs[MAIN_GRAPH_NAME] + keys = list(graph_inputs) + if quantization_config.get("calibrate_activations") and keys[:2] != [ + "input_ids", + "position_ids", + ]: + raise ValueError( + f"{type(model).__name__}: activation calibration feeds tokenized samples as " + f"input_ids and rebuilds position_ids, so those must be the first two " + f"parameters of forward; got {tuple(keys[:2])}." + ) + + # Which *positions* are state cannot be read off the contract: the name lists carry + # only relative order. So this assumes the declared inputs precede the states, which + # holds for the macOS graph -- the only graph calibration runs on. A model that + # interleaved a non-state parameter after a cache would need them passed in. + n_inputs = len(model.export_input_names()[MAIN_GRAPH_NAME]) + state_indices = tuple(range(n_inputs, len(keys))) + + return quantize_pytorch_model( + model, + model.reference_inputs_as_args(graph_inputs), + dynamic_shapes[MAIN_GRAPH_NAME], + quantization_config, + calibration_data_fn=calibration_data_fn, + mmap_dir=mmap_dir, + cache_seq_len=spec.cache_seq_len, + state_indices=state_indices, + ) + + def palettize_pytorch_model( model: nn.Module, example_inputs: tuple, diff --git a/python/src/coreai_models/export/macos.py b/python/src/coreai_models/export/macos.py index c8798301..da6d31b4 100644 --- a/python/src/coreai_models/export/macos.py +++ b/python/src/coreai_models/export/macos.py @@ -11,24 +11,19 @@ """ import logging +from typing import Any import coreai_torch import coreai_torch.composite_ops import torch from coreai.authoring import AIProgram -from coreai_models.export._constants import ( - KEY_CACHE_NAME, - QUANT_TRACE_OFFSET, - QUANT_TRACE_QUERY_LEN, - TRACE_KV_CACHE_SEQ_LEN, - VALUE_CACHE_NAME, -) +from coreai_models._constants import MAIN_GRAPH_NAME, TRACE_KV_CACHE_SEQ_LEN from coreai_models.export.mlir_ops import ( register_custom_torch_lowering, remove_functionalization, ) -from coreai_models.primitives.macos.cache import KVCache +from coreai_models.models.base import BaseForCausalLM, TraceSpec logger = logging.getLogger(__name__) @@ -64,69 +59,31 @@ def _build_reference_inputs( - model: torch.nn.Module, + model: BaseForCausalLM, config, target_dtype: torch.dtype, max_context_length: int, -) -> tuple[dict[str, torch.Tensor], dict]: - """Build reference inputs and dynamic shapes for macOS model export. - - Args: - model: The PyTorch model (used only to read config). - config: HuggingFace model config. - target_dtype: Data type for cache tensors. - max_context_length: Maximum context length for the model. +) -> tuple[dict[str, Any], dict]: + """Reference inputs and dynamic shapes for macOS export. - Returns: - Tuple of (reference_inputs dict, dynamic_shapes dict). + Thin wrapper over the model's export-contract hooks, where the per-model variation + lives. Returns ``(reference_inputs, dynamic_shapes)``. """ - batch_size = 1 - vocab_size = config.vocab_size - - input_ids = torch.randint(1, vocab_size, (batch_size, QUANT_TRACE_QUERY_LEN), dtype=torch.int32) - position_ids = ( - torch.arange(QUANT_TRACE_QUERY_LEN + QUANT_TRACE_OFFSET, dtype=torch.int32) - .unsqueeze(0) - .expand(batch_size, QUANT_TRACE_QUERY_LEN + QUANT_TRACE_OFFSET) + # The trace cache length only bounds peak memory, so cap it at the context it serves. + spec = TraceSpec( + max_context_length=max_context_length, + cache_seq_len=min(TRACE_KV_CACHE_SEQ_LEN, max_context_length), ) - - # Clamp `max_position_embeddings` so KVCache.create_cache_tensors doesn't - # allocate a full-context cache for huge models - saved_max_pos = config.max_position_embeddings - config.max_position_embeddings = TRACE_KV_CACHE_SEQ_LEN - k_cache, v_cache = KVCache.create_cache_tensors(config, dtype=target_dtype) - config.max_position_embeddings = saved_max_pos - - reference_inputs = { - "input_ids": input_ids, - "position_ids": position_ids, - "k_cache": k_cache, - "v_cache": v_cache, - } - - dynamic_shapes = { - "input_ids": {1: torch.export.Dim("seq_ids", max=max_context_length - 2)}, - "position_ids": { - 1: torch.export.Dim("seq_pos", min=QUANT_TRACE_QUERY_LEN, max=max_context_length - 1) - }, - "k_cache": { - KVCache.seq_len_dim(): torch.export.Dim( - "k_seq_len", min=TRACE_KV_CACHE_SEQ_LEN, max=max_context_length - ) - }, - "v_cache": { - KVCache.seq_len_dim(): torch.export.Dim( - "v_seq_len", min=TRACE_KV_CACHE_SEQ_LEN, max=max_context_length - ) - }, - } - - return reference_inputs, dynamic_shapes + reference_inputs = model.build_reference_inputs(config, target_dtype, spec) + dynamic_shapes = model.build_dynamic_shapes(config, spec) + model.validate_export_contract(reference_inputs, dynamic_shapes) + # A macOS model has exactly one graph. + return reference_inputs[MAIN_GRAPH_NAME], dynamic_shapes[MAIN_GRAPH_NAME] def export_to_coreai( model: torch.nn.Module, - reference_inputs: dict[str, torch.Tensor], + reference_inputs: dict[str, Any], dynamic_shapes: dict | None = None, input_names: tuple[str, ...] | None = None, output_names: tuple[str, ...] | None = None, @@ -197,7 +154,7 @@ def export_fn(module: torch.nn.Module) -> torch.export.ExportedProgram: def export_macos_model( - model: torch.nn.Module, + model: BaseForCausalLM, config, export_config, ) -> AIProgram: @@ -209,7 +166,8 @@ def export_macos_model( 3. Optimizes the resulting AIProgram Args: - model: A loaded PyTorch model (already in the correct dtype). + model: A loaded PyTorch model (already in the correct dtype). Its + export-contract hooks supply the graph's inputs, states, and names. config: HuggingFace model config (used for cache dimensions, vocab size, etc.). export_config: An ExportConfig instance (used for max_context_length, etc.). @@ -231,18 +189,14 @@ def export_macos_model( model, config, target_dtype, max_context_length ) - input_names = ("input_ids", "position_ids") - output_names = ("logits",) - state_names = (KEY_CACHE_NAME, VALUE_CACHE_NAME) - logger.info("Exporting model to Core AI dialect...") coreai_program = export_to_coreai( model, reference_inputs, dynamic_shapes=dynamic_shapes, - input_names=input_names, - output_names=output_names, - state_names=state_names, + input_names=model.export_input_names()[MAIN_GRAPH_NAME], + output_names=model.export_output_names()[MAIN_GRAPH_NAME], + state_names=model.export_state_names()[MAIN_GRAPH_NAME], ) logger.info("Optimizing AIProgram...") diff --git a/python/src/coreai_models/export/pipeline.py b/python/src/coreai_models/export/pipeline.py index 7e2fd63d..abe49367 100644 --- a/python/src/coreai_models/export/pipeline.py +++ b/python/src/coreai_models/export/pipeline.py @@ -24,17 +24,15 @@ from coreai_opt.palettization.config.palettization_config import KMeansPalettizerConfig from transformers import AutoConfig, AutoTokenizer -from coreai_models.export._constants import ( +from coreai_models._constants import ( IOS_DEFAULT_MAX_CONTEXT_LENGTH, - QUANT_TRACE_OFFSET, - QUANT_TRACE_QUERY_LEN, TRACE_KV_CACHE_SEQ_LEN, ) from coreai_models.export.bundle import bundle_llm_asset from coreai_models.export.compression import ( get_c4, palettize_pytorch_model, - quantize_pytorch_model, + quantize_for_export, ) from coreai_models.export.ios import export_ios_model from coreai_models.export.macos import export_macos_model @@ -44,7 +42,6 @@ get_preset, ) from coreai_models.models.registry import get_model_entry -from coreai_models.primitives.macos.cache import KVCache logger = logging.getLogger(__name__) @@ -241,32 +238,6 @@ async def _async_export_model(config: ExportConfig) -> str: if torch_quantization_config is not None: logger.info(f"Applying pre-export torch quantization (preset={config.compression})") - input_ids = torch.randint( - 1, vocab_size, (batch_size, QUANT_TRACE_QUERY_LEN), dtype=torch.int32 - ) - position_ids = ( - torch.arange(QUANT_TRACE_QUERY_LEN + QUANT_TRACE_OFFSET, dtype=torch.int32) - .unsqueeze(0) - .expand(batch_size, QUANT_TRACE_QUERY_LEN + QUANT_TRACE_OFFSET) - ) - - saved_max_pos = hf_config.max_position_embeddings - hf_config.max_position_embeddings = TRACE_KV_CACHE_SEQ_LEN - k_cache, v_cache = KVCache.create_cache_tensors(hf_config, dtype=target_dtype) - hf_config.max_position_embeddings = saved_max_pos - - quantization_inputs = (input_ids, position_ids, k_cache, v_cache) - quantization_dynamic_shapes = { - "input_ids": {1: torch.export.Dim("seq_ids", max=TRACE_KV_CACHE_SEQ_LEN - 2)}, - "position_ids": { - 1: torch.export.Dim( - "seq_pos", min=QUANT_TRACE_QUERY_LEN, max=TRACE_KV_CACHE_SEQ_LEN - 1 - ) - }, - "k_cache": None, - "v_cache": None, - } - def get_calibration_data(): # type: ignore[no-untyped-def] tokenizer = AutoTokenizer.from_pretrained(config.hf_model_id) return get_c4(tokenizer) @@ -284,10 +255,10 @@ def get_calibration_data(): # type: ignore[no-untyped-def] if not isinstance(torch_quantization_config, dict) else dict(torch_quantization_config) ) - model = quantize_pytorch_model( + model = quantize_for_export( model, - quantization_inputs, - quantization_dynamic_shapes, + hf_config, + target_dtype, quant_cfg, calibration_data_fn=get_calibration_data, mmap_dir=quantizer_mmap_dir, diff --git a/python/src/coreai_models/model_registry.py b/python/src/coreai_models/model_registry.py index cc6a0bc9..d5b5717a 100644 --- a/python/src/coreai_models/model_registry.py +++ b/python/src/coreai_models/model_registry.py @@ -27,7 +27,7 @@ from dataclasses import asdict, dataclass from pathlib import Path -from coreai_models.export._constants import IOS_DEFAULT_MAX_CONTEXT_LENGTH +from coreai_models._constants import IOS_DEFAULT_MAX_CONTEXT_LENGTH # --------------------------------------------------------------------------- # Data model diff --git a/python/src/coreai_models/models/base.py b/python/src/coreai_models/models/base.py index 647e4043..783ed22f 100644 --- a/python/src/coreai_models/models/base.py +++ b/python/src/coreai_models/models/base.py @@ -7,13 +7,15 @@ import collections.abc import gc +import inspect import json import os import re from abc import abstractmethod from collections.abc import Callable +from dataclasses import dataclass from functools import wraps -from typing import TypeVar, cast +from typing import Any, NoReturn, TypeVar, cast import torch from huggingface_hub import snapshot_download @@ -21,14 +23,69 @@ from safetensors.torch import save_file from transformers import AutoConfig from transformers.modeling_utils import PreTrainedModel -from typing_extensions import Self - +from typing_extensions import Self, override + +from coreai_models._constants import ( + KEY_CACHE_NAME, + MAIN_GRAPH_NAME, + QUANT_TRACE_OFFSET, + QUANT_TRACE_QUERY_LEN, + TRACE_KV_CACHE_SEQ_LEN, + VALUE_CACHE_NAME, +) from coreai_models.primitives.ios.embedding import GatherEmbeddings, LoadEmbeddings from coreai_models.primitives.macos.cache import KVCache T = TypeVar("T", bound="BaseForCausalLM") +@dataclass(frozen=True) +class TraceSpec: + """Shapes for one ``torch.export`` trace of a causal LM forward. + + Passed to both ``build_reference_inputs`` and ``build_dynamic_shapes`` so the + tensors and their declared dims cannot disagree. + + Attributes: + max_context_length: Upper bound for the dynamic seq/cache dims. Must be at + least ``query_len + 2``. + cache_seq_len: Length the caches are *traced* at, to bound peak trace memory. + Unrelated to the inference cache size. Must not exceed + ``max_context_length``. + query_len: Trace-time ``input_ids`` length. + offset: Already-cached positions, so ``position_ids`` is + ``query_len + offset`` long. + """ + + max_context_length: int + cache_seq_len: int = TRACE_KV_CACHE_SEQ_LEN + query_len: int = QUANT_TRACE_QUERY_LEN + offset: int = QUANT_TRACE_OFFSET + + def __post_init__(self) -> None: + # Below this there is no legal `Dim(min=query_len, max=max_context_length - 1)`. + if self.max_context_length < self.query_len + 2: + raise ValueError( + f"max_context_length={self.max_context_length} is too small to trace: " + f"it must be at least query_len + 2 = {self.query_len + 2}." + ) + if self.cache_seq_len > self.max_context_length: + raise ValueError( + "cache_seq_len must not be greater than max_context_length. Received " + f"cache_seq_len = {self.cache_seq_len}, " + f"max_context_length = {self.max_context_length}" + ) + + @property + def caches_are_static(self) -> bool: + """Whether the cache dims must be pinned rather than declared dynamic. + + A cache traced at the full context has nowhere to grow, and ``Dim(min=max=n)`` + is illegal anyway. + """ + return self.cache_seq_len == self.max_context_length + + def _is_layer_key_beyond(key: str, num_layers: int) -> bool: """Return True if `key` refers to a transformer layer with index >= num_layers. @@ -300,6 +357,164 @@ def _mutate_state_dict(self: Self, state_dict: dict[str, torch.Tensor]) -> None: """ ... + # ------------------------------------------------------------------ + # Export contract + # + # Everything the exporters need to trace this model, keyed by graph name. A macOS + # model has one graph; iOS has several. These hooks supply only names and tensors; + # which callable each graph traces stays the exporter's business. + # + # Reference inputs bind to the traced signature, so they must be in its EXACT + # order. Names are looked up by name, so each list carries only the RELATIVE order + # of its own kind: for forward(input_ids, key_cache, position_ids, value_cache), + # input_names is (input_ids, position_ids) and state_names is (key_cache, + # value_cache). + # ------------------------------------------------------------------ + + @classmethod + def export_input_names(cls) -> dict[str, tuple[str, ...]]: + """Graph input names per graph, in relative order among the non-state args.""" + return {MAIN_GRAPH_NAME: ("input_ids", "position_ids")} + + @classmethod + def export_state_names(cls) -> dict[str, tuple[str, ...]]: + """Runner-visible state names per graph, in relative order among the state args. + + State args are mutated in place and surfaced through the runtime ``state=`` + kwarg rather than as ordinary inputs/outputs. + """ + return {MAIN_GRAPH_NAME: (KEY_CACHE_NAME, VALUE_CACHE_NAME)} + + @classmethod + def export_output_names(cls) -> dict[str, tuple[str, ...]]: + """Graph output names per graph, in return order.""" + return {MAIN_GRAPH_NAME: ("logits",)} + + def build_reference_inputs( + self, + config, + target_dtype: torch.dtype, + spec: TraceSpec, + ) -> dict[str, dict[str, Any]]: + """Reference tensors to trace with, per graph, keyed by parameter name. + + The inner dicts bind to the traced callable, so their keys must be its + parameters in *exact* signature order. Pass ``spec`` to + :meth:`build_dynamic_shapes` too, so the tensors and their dims cannot disagree. + """ + input_ids = torch.randint(1, config.vocab_size, (1, spec.query_len), dtype=torch.int32) + position_ids = ( + torch.arange(spec.query_len + spec.offset, dtype=torch.int32) + .unsqueeze(0) + .expand(1, spec.query_len + spec.offset) + ) + k_cache, v_cache = KVCache.create_cache_tensors( + config, dtype=target_dtype, seq_len=spec.cache_seq_len + ) + return { + MAIN_GRAPH_NAME: { + "input_ids": input_ids, + "position_ids": position_ids, + "k_cache": k_cache, + "v_cache": v_cache, + } + } + + def build_dynamic_shapes(self, config, spec: TraceSpec) -> dict[str, Any]: + """``dynamic_shapes`` per graph, matching :meth:`build_reference_inputs`. + + Keyed like the reference inputs; ``None`` pins that input to its traced shape. + """ + max_ctx = spec.max_context_length + shapes: dict[str, Any] = { + "input_ids": {1: torch.export.Dim("seq_ids", max=max_ctx - 2)}, + "position_ids": {1: torch.export.Dim("seq_pos", min=spec.query_len, max=max_ctx - 1)}, + } + seq_dim = KVCache.seq_len_dim() + if spec.caches_are_static: + shapes["k_cache"] = None + shapes["v_cache"] = None + else: + shapes["k_cache"] = { + seq_dim: torch.export.Dim("k_seq_len", min=spec.cache_seq_len, max=max_ctx) + } + shapes["v_cache"] = { + seq_dim: torch.export.Dim("v_seq_len", min=spec.cache_seq_len, max=max_ctx) + } + return {MAIN_GRAPH_NAME: shapes} + + def validate_export_contract( + self, + reference_inputs: dict[str, dict[str, Any]], + dynamic_shapes: dict[str, Any], + ) -> None: + """Check the five hooks agree with each other. + + The converter compares name *counts* only, so it would accept a contract whose + graphs disagree. Called by the exporters before tracing. + + Raises: + ValueError: On any disagreement between the hooks. + """ + cls_name = type(self).__name__ + inputs, states, outputs = ( + self.export_input_names(), + self.export_state_names(), + self.export_output_names(), + ) + named = { + "export_input_names": set(inputs), + "export_state_names": set(states), + "export_output_names": set(outputs), + "build_reference_inputs": set(reference_inputs), + "build_dynamic_shapes": set(dynamic_shapes), + } + graphs = named["export_input_names"] + for hook, keys in named.items(): + if keys != graphs: + raise ValueError( + f"{cls_name}: {hook} covers graphs {sorted(keys)} but " + f"export_input_names covers {sorted(graphs)}. Every hook must " + "describe the same graphs." + ) + + for graph in sorted(graphs): + refs = reference_inputs[graph] + declared = len(inputs[graph]) + len(states[graph]) + if declared != len(refs): + raise ValueError( + f"{cls_name}, graph {graph!r}: {len(inputs[graph])} input names + " + f"{len(states[graph])} state names = {declared}, but " + f"build_reference_inputs supplies {len(refs)} tensors " + f"{tuple(refs)}." + ) + shapes = dynamic_shapes[graph] + if shapes is not None and set(shapes) != set(refs): + raise ValueError( + f"{cls_name}, graph {graph!r}: dynamic_shapes keys " + f"{tuple(shapes)} do not match reference inputs {tuple(refs)}." + ) + + def reference_inputs_as_args(self, reference_inputs: dict[str, Any]) -> tuple[Any, ...]: + """One graph's reference inputs as positional args, validating the order. + + For the coreai-opt quantizer, which takes a tuple rather than kwargs; a dict + ordered differently from the signature would silently bind the wrong tensors. + + Raises: + ValueError: If the keys are not a contiguous in-order prefix of ``forward``. + """ + params = list(inspect.signature(self.forward).parameters) + keys = list(reference_inputs) + if keys != params[: len(keys)]: + raise ValueError( + f"{type(self).__name__}.build_reference_inputs returned keys {keys}, " + f"which are not a contiguous in-order prefix of forward's parameters " + f"{params}. Positional conversion would bind tensors to the wrong " + f"parameters; fix the dict's insertion order or the signature." + ) + return tuple(reference_inputs.values()) + @classmethod def _get_reauthored_config( cls, @@ -605,3 +820,48 @@ def __init__(self: Self, config, model_device: str, disable_embedding_quantizati def set_prefill_mode(self, prefill_mode: bool): self.extend.prefill_mode = prefill_mode + + # ------------------------------------------------------------------ + # Export contract -- not implemented for iOS + # + # The iOS graph matches the macOS defaults in no respect (different inputs, cache + # rank and output name -- see export/ios.py), and export_ios_model builds its own. + # Inheriting the defaults would report a plausible but wrong contract, and they are + # reachable: `--variant iOS --compression 4bit` resolves a macOS quantization preset + # and the quantize branch is not variant-gated. + # ------------------------------------------------------------------ + + @classmethod + def _export_contract_not_implemented(cls, hook: str) -> NoReturn: + raise NotImplementedError( + f"{cls.__name__}.{hook} is not yet implemented for iOS models: the iOS graph " + "has a different signature from the macOS one it would otherwise inherit." + ) + + @classmethod + @override + def export_input_names(cls) -> dict[str, tuple[str, ...]]: + cls._export_contract_not_implemented("export_input_names") + + @classmethod + @override + def export_state_names(cls) -> dict[str, tuple[str, ...]]: + cls._export_contract_not_implemented("export_state_names") + + @classmethod + @override + def export_output_names(cls) -> dict[str, tuple[str, ...]]: + cls._export_contract_not_implemented("export_output_names") + + @override + def build_reference_inputs( + self, + config, + target_dtype: torch.dtype, + spec: TraceSpec, + ) -> dict[str, dict[str, Any]]: + self._export_contract_not_implemented("build_reference_inputs") + + @override + def build_dynamic_shapes(self, config, spec: TraceSpec) -> dict[str, Any]: + self._export_contract_not_implemented("build_dynamic_shapes") diff --git a/python/src/coreai_models/primitives/macos/cache.py b/python/src/coreai_models/primitives/macos/cache.py index ece971fc..a3d51cde 100644 --- a/python/src/coreai_models/primitives/macos/cache.py +++ b/python/src/coreai_models/primitives/macos/cache.py @@ -36,15 +36,22 @@ def create_cache_tensors( cls, config, dtype: torch.dtype = torch.float32, + seq_len: int | None = None, ) -> tuple[torch.Tensor, torch.Tensor]: """Create zero-initialized KV cache tensors from a model config. + Args: + config: Model config supplying the layer/head dimensions. + dtype: Cache dtype. + seq_len: Sequence-dim length; defaults to ``config.max_position_embeddings``. + Pass explicitly to build a trace-sized cache without mutating config. + Returns: - (k_cache, v_cache) tensors of shape (n_layers, 1, n_kv_heads, max_seq_len, head_dim). + (k_cache, v_cache) of shape (n_layers, 1, n_kv_heads, max_seq_len, head_dim). """ n_kv_heads = config.num_key_value_heads n_layers = config.num_hidden_layers - max_seq_len = config.max_position_embeddings + max_seq_len = config.max_position_embeddings if seq_len is None else seq_len if hasattr(config, "head_dim") and config.head_dim is not None: head_dim = config.head_dim else: diff --git a/python/src/coreai_models/primitives/macos/cache_scatter.py b/python/src/coreai_models/primitives/macos/cache_scatter.py index 6bba2214..a55c51c3 100644 --- a/python/src/coreai_models/primitives/macos/cache_scatter.py +++ b/python/src/coreai_models/primitives/macos/cache_scatter.py @@ -36,10 +36,11 @@ def create_cache_tensors( cls, config, dtype: torch.dtype = torch.float32, + seq_len: int | None = None, ) -> tuple[torch.Tensor, torch.Tensor]: n_kv_heads = config.num_key_value_heads n_layers = config.num_hidden_layers - max_seq_len = config.max_position_embeddings + max_seq_len = config.max_position_embeddings if seq_len is None else seq_len if hasattr(config, "head_dim") and config.head_dim is not None: head_dim = config.head_dim else: diff --git a/python/tests/_runner_infra/testing_utils.py b/python/tests/_runner_infra/testing_utils.py index 8a5d7527..57439232 100644 --- a/python/tests/_runner_infra/testing_utils.py +++ b/python/tests/_runner_infra/testing_utils.py @@ -957,20 +957,14 @@ def test_weight_activation_quantization(self, activation_quantization) -> None: if activation_quantization: pytest.skip("Activation quantization temporarily disabled with eager mode quantization") - # We replicate the relevant parts of - # ``coreai_models.export.pipeline._async_export_model`` here: - # load HF -> apply torch quantization -> run macOS export. The output - # asset write is intentionally skipped; we only want to confirm + # Same building blocks as `_async_export_model`: load HF -> + # `quantize_for_export` -> `export_macos_model`. Calling those rather than + # reimplementing the calibration trace keeps this from drifting from + # production, as it had. The asset write is skipped; we only confirm # quantize + export produces a non-None AIProgram. - from coreai_models.export._constants import ( - QUANT_TRACE_OFFSET, - QUANT_TRACE_QUERY_LEN, - TRACE_KV_CACHE_SEQ_LEN, - ) - from coreai_models.export.compression import quantize_pytorch_model + from coreai_models.export.compression import quantize_for_export from coreai_models.export.macos import export_macos_model from coreai_models.export.pipeline import ExportConfig - from coreai_models.primitives.macos.cache import KVCache hf_config = transformers.AutoConfig.from_pretrained(self._toy_model_id) is_gemma = "gemma" in self._model_class.__name__.lower() @@ -1034,39 +1028,12 @@ def test_weight_activation_quantization(self, activation_quantization) -> None: hf_state_dict_prefix=hf_state_dict_prefix, ).eval() - # Build calibration / trace inputs for quantization - vocab_size = getattr(hf_config, "vocab_size", 32000) - input_ids = torch.randint(1, vocab_size, (1, QUANT_TRACE_QUERY_LEN), dtype=torch.int32) - position_ids = ( - torch.arange(QUANT_TRACE_QUERY_LEN + QUANT_TRACE_OFFSET, dtype=torch.int32) - .unsqueeze(0) - .expand(1, QUANT_TRACE_QUERY_LEN + QUANT_TRACE_OFFSET) - ) - saved_max_pos = hf_config.max_position_embeddings - hf_config.max_position_embeddings = TRACE_KV_CACHE_SEQ_LEN - k_cache, v_cache = KVCache.create_cache_tensors(hf_config, dtype=target_dtype) - hf_config.max_position_embeddings = saved_max_pos - - quantization_inputs = (input_ids, position_ids, k_cache, v_cache) - quantization_dynamic_shapes = { - "input_ids": {1: torch.export.Dim("seq_ids", max=max_context_length - 2)}, - "position_ids": { - 1: torch.export.Dim( - "seq_pos", - min=QUANT_TRACE_QUERY_LEN, - max=max_context_length - 1, - ) - }, - "k_cache": None, - "v_cache": None, - } - quantizer_mmap_dir = f"{tmpdir}/quantized" os.makedirs(quantizer_mmap_dir, exist_ok=True) - model = quantize_pytorch_model( + model = quantize_for_export( model, - quantization_inputs, - quantization_dynamic_shapes, + hf_config, + target_dtype, dict(torch_quantization_config), calibration_data_fn=None, mmap_dir=quantizer_mmap_dir, diff --git a/python/tests/test_model_conversion/test_infra.py b/python/tests/test_model_conversion/test_infra.py index cfe755ee..5cd714f0 100644 --- a/python/tests/test_model_conversion/test_infra.py +++ b/python/tests/test_model_conversion/test_infra.py @@ -17,8 +17,8 @@ Qwen3ForCausalLM as HFQwen3ForCausalLM, ) +from coreai_models._constants import IOS_DEFAULT_MAX_CONTEXT_LENGTH from coreai_models.export import pipeline as export_pipeline -from coreai_models.export._constants import IOS_DEFAULT_MAX_CONTEXT_LENGTH from coreai_models.export.ios import KEY_CACHE_INPUT_NAME, VALUE_CACHE_INPUT_NAME from coreai_models.export.pipeline import ExportConfig, _async_export_model diff --git a/python/tests/test_model_units/test_export/test_export_contract.py b/python/tests/test_model_units/test_export/test_export_contract.py new file mode 100644 index 00000000..abebc774 --- /dev/null +++ b/python/tests/test_model_units/test_export/test_export_contract.py @@ -0,0 +1,408 @@ +# 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 export contract hooks on ``BaseForCausalLM``. + +Cover the graph contract independently of any real architecture, plus the +append-extras path. No hardware or HuggingFace weights required. +""" + +from types import SimpleNamespace + +import pytest +import torch +from typing_extensions import override + +from coreai_models._constants import ( + KEY_CACHE_NAME, + QUANT_TRACE_OFFSET, + QUANT_TRACE_QUERY_LEN, + TRACE_KV_CACHE_SEQ_LEN, + VALUE_CACHE_NAME, +) +from coreai_models._constants import ( + MAIN_GRAPH_NAME as MAIN, +) +from coreai_models.models.base import BaseForCausalLM, TraceSpec +from coreai_models.primitives.macos.cache import KVCache + +MAX_CONTEXT_LENGTH = 8192 + + +def _tiny_config() -> SimpleNamespace: + """The smallest config the contract hooks read.""" + return SimpleNamespace( + vocab_size=128, + num_hidden_layers=2, + num_attention_heads=4, + num_key_value_heads=2, + hidden_size=32, + head_dim=8, + max_position_embeddings=MAX_CONTEXT_LENGTH, + ) + + +class _StandardLM(BaseForCausalLM): + """A model with the default contract: (input_ids, position_ids, k_cache, v_cache).""" + + @override + def _init_model(self, config) -> None: + self.lm_head = torch.nn.Linear(config.hidden_size, config.vocab_size, bias=False) + + @override + def _mutate_state_dict(self, state_dict) -> None: + pass + + def forward( + self, + input_ids: torch.Tensor, + position_ids: torch.IntTensor, + k_cache: torch.Tensor, + v_cache: torch.Tensor, + ) -> torch.Tensor: + raise NotImplementedError("shape contract only; never traced in these tests") + + +class _ExtraStateLM(_StandardLM): + """A model that appends state beyond the standard KV pair. + + The base KV pair first, then extra state args. + """ + + @override + def forward( + self, + input_ids: torch.Tensor, + position_ids: torch.IntTensor, + k_cache: torch.Tensor, + v_cache: torch.Tensor, + extra_state: torch.Tensor = None, + ) -> torch.Tensor: + raise NotImplementedError("shape contract only; never traced in these tests") + + @classmethod + @override + def export_state_names(cls) -> dict[str, tuple[str, ...]]: + return {MAIN: (*super().export_state_names()[MAIN], "extraState")} + + @override + def build_reference_inputs(self, config, target_dtype, spec): + graphs = super().build_reference_inputs(config, target_dtype, spec) + graphs[MAIN]["extra_state"] = torch.zeros(2, spec.cache_seq_len, dtype=target_dtype) + return graphs + + @override + def build_dynamic_shapes(self, config, spec): + graphs = super().build_dynamic_shapes(config, spec) + graphs[MAIN]["extra_state"] = None + return graphs + + +@pytest.fixture +def config() -> SimpleNamespace: + return _tiny_config() + + +@pytest.fixture +def model(config) -> _StandardLM: + return _StandardLM(config) + + +class TestMacOSContract: + """The macOS model describes one graph, keyed ``main``.""" + + def _built(self, model, config, spec=None): + spec = spec or TraceSpec(max_context_length=MAX_CONTEXT_LENGTH) + return ( + model.build_reference_inputs(config, torch.float16, spec), + model.build_dynamic_shapes(config, spec), + ) + + def test_every_hook_describes_exactly_the_main_graph(self, model, config) -> None: + refs, shapes = self._built(model, config) + for hook in ( + model.export_input_names(), + model.export_state_names(), + model.export_output_names(), + refs, + shapes, + ): + assert list(hook) == [MAIN] + + def test_names(self, model) -> None: + assert model.export_input_names()[MAIN] == ("input_ids", "position_ids") + assert model.export_state_names()[MAIN] == (KEY_CACHE_NAME, VALUE_CACHE_NAME) + assert model.export_output_names()[MAIN] == ("logits",) + + def test_reference_inputs_are_in_exact_signature_order(self, model, config) -> None: + """They bind to the traced callable, so order is exact, not relative.""" + import inspect + + refs, _ = self._built(model, config) + params = list(inspect.signature(model.forward).parameters) + keys = list(refs[MAIN]) + assert keys == params[: len(keys)] + + def test_reference_input_shapes(self, model, config) -> None: + refs, _ = self._built(model, config) + graph = refs[MAIN] + assert graph["input_ids"].shape == (1, QUANT_TRACE_QUERY_LEN) + assert graph["input_ids"].dtype == torch.int32 + assert graph["position_ids"].shape == (1, QUANT_TRACE_QUERY_LEN + QUANT_TRACE_OFFSET) + expected = (2, 1, 2, TRACE_KV_CACHE_SEQ_LEN, 8) + for name in ("k_cache", "v_cache"): + assert graph[name].shape == expected + assert graph[name].dtype == torch.float16 + + def test_caches_traced_at_cache_seq_len(self, model, config) -> None: + spec = TraceSpec(max_context_length=MAX_CONTEXT_LENGTH, cache_seq_len=512) + refs, _ = self._built(model, config, spec) + assert refs[MAIN]["k_cache"].shape[KVCache.seq_len_dim()] == 512 + + def test_config_is_not_mutated(self, model, config) -> None: + """Regression: sizing the cache used to mutate and restore the config.""" + self._built( + model, config, TraceSpec(max_context_length=MAX_CONTEXT_LENGTH, cache_seq_len=512) + ) + assert config.max_position_embeddings == MAX_CONTEXT_LENGTH + + def test_dynamic_shape_bounds(self, model, config) -> None: + _, shapes = self._built(model, config) + graph = shapes[MAIN] + assert graph["input_ids"][1].max == MAX_CONTEXT_LENGTH - 2 + assert graph["position_ids"][1].min == QUANT_TRACE_QUERY_LEN + assert graph["position_ids"][1].max == MAX_CONTEXT_LENGTH - 1 + seq_dim = KVCache.seq_len_dim() + for name in ("k_cache", "v_cache"): + assert graph[name][seq_dim].min == TRACE_KV_CACHE_SEQ_LEN + assert graph[name][seq_dim].max == MAX_CONTEXT_LENGTH + + +class TestSmallContext: + """Contexts at or below the default trace cache length. + + Regression: the cache dim was built as ``Dim(min=TRACE_KV_CACHE_SEQ_LEN, + max=max_context_length)`` unconditionally, which raises from inside ``torch.export`` + whenever the context is <= the trace length. + """ + + def test_cache_seq_len_may_equal_the_context(self) -> None: + assert TraceSpec(max_context_length=512, cache_seq_len=512).cache_seq_len == 512 + + def test_cache_seq_len_above_the_context_is_rejected(self) -> None: + # A cache longer than the context it serves is meaningless, so the spec + # rejects it rather than quietly shrinking it. + with pytest.raises(ValueError, match="must not be greater than"): + TraceSpec(max_context_length=512, cache_seq_len=513) + + def test_larger_context_leaves_trace_length_alone(self) -> None: + assert TraceSpec(max_context_length=8192).cache_seq_len == TRACE_KV_CACHE_SEQ_LEN + + def test_context_too_small_to_trace_is_rejected(self) -> None: + limit = TraceSpec(max_context_length=8192).query_len + 2 + TraceSpec(max_context_length=limit, cache_seq_len=limit) + for bad in (limit - 1, 2, 0, -5): + with pytest.raises(ValueError, match="too small to trace"): + TraceSpec(max_context_length=bad) + + @pytest.mark.parametrize("max_ctx", [512, TRACE_KV_CACHE_SEQ_LEN]) + def test_cache_dims_pin_instead_of_raising(self, model, max_ctx) -> None: + config = _tiny_config() + config.max_position_embeddings = max_ctx + spec = TraceSpec( + max_context_length=max_ctx, cache_seq_len=min(TRACE_KV_CACHE_SEQ_LEN, max_ctx) + ) + assert spec.caches_are_static + shapes = model.build_dynamic_shapes(config, spec)[MAIN] + assert shapes["k_cache"] is None + assert shapes["v_cache"] is None + + def test_dims_stay_dynamic_when_there_is_room(self, model, config) -> None: + spec = TraceSpec(max_context_length=MAX_CONTEXT_LENGTH) + assert not spec.caches_are_static + assert model.build_dynamic_shapes(config, spec)[MAIN]["k_cache"] is not None + + +class TestValidateExportContract: + """Cross-checks the five hooks against each other.""" + + def _built(self, model, config): + spec = TraceSpec(max_context_length=MAX_CONTEXT_LENGTH) + return ( + model.build_reference_inputs(config, torch.float16, spec), + model.build_dynamic_shapes(config, spec), + ) + + def test_accepts_the_default_contract(self, model, config) -> None: + model.validate_export_contract(*self._built(model, config)) + + def test_accepts_appended_state(self, config) -> None: + m = _ExtraStateLM(config) + m.validate_export_contract(*self._built(m, config)) + + def test_rejects_a_hook_covering_different_graphs(self, model, config) -> None: + refs, shapes = self._built(model, config) + refs["extra_graph"] = {} + with pytest.raises(ValueError, match="must describe the same graphs"): + model.validate_export_contract(refs, shapes) + + def test_rejects_a_name_count_mismatch(self, config) -> None: + class _TooFewNames(_StandardLM): + @classmethod + @override + def export_state_names(cls) -> dict[str, tuple[str, ...]]: + return {MAIN: (KEY_CACHE_NAME,)} + + m = _TooFewNames(config) + with pytest.raises(ValueError, match="build_reference_inputs supplies"): + m.validate_export_contract(*self._built(m, config)) + + def test_rejects_dynamic_shape_key_mismatch(self, model, config) -> None: + refs, shapes = self._built(model, config) + del shapes[MAIN]["position_ids"] + with pytest.raises(ValueError, match="do not match reference inputs"): + model.validate_export_contract(refs, shapes) + + +class TestReferenceInputsAsArgs: + """Positional conversion for the quantizer, which takes a tuple not kwargs.""" + + def _graph(self, model, config): + spec = TraceSpec(max_context_length=MAX_CONTEXT_LENGTH) + return model.build_reference_inputs(config, torch.float16, spec)[MAIN] + + def test_returns_signature_order(self, model, config) -> None: + graph = self._graph(model, config) + args = model.reference_inputs_as_args(graph) + assert len(args) == 4 + for arg, expected in zip(args, graph.values(), strict=True): + assert arg is expected + + def test_rejects_reordered_keys(self, model, config) -> None: + graph = self._graph(model, config) + swapped = {"position_ids": graph["position_ids"], "input_ids": graph["input_ids"]} + swapped.update({k: graph[k] for k in ("k_cache", "v_cache")}) + with pytest.raises(ValueError, match="not a contiguous in-order prefix"): + model.reference_inputs_as_args(swapped) + + def test_sees_through_the_logits_cast_decorator(self, config) -> None: + """Every real subclass decorates forward; introspection needs functools.wraps.""" + + class _Decorated(_StandardLM): + @BaseForCausalLM.cast_logits_bfloat16_to_float16 + @override + def forward(self, input_ids, position_ids, k_cache, v_cache): + raise NotImplementedError + + m = _Decorated(config) + assert len(m.reference_inputs_as_args(self._graph(m, config))) == 4 + + +class TestRegisteredModelsSatisfyTheContract: + """Pins the real registry, so renaming a forward parameter fails here.""" + + @staticmethod + def _macos_entries(): + from coreai_models.models.registry import _get_registry + + return sorted( + ((mt, e) for mt, e in _get_registry().items() if e.macos_class is not None), + key=lambda kv: kv[0], + ) + + def test_registry_is_not_empty(self) -> None: + assert self._macos_entries() + + def test_every_registered_model_validates(self) -> None: + import inspect + + from transformers import AutoConfig + + for model_type, entry in self._macos_entries(): + raw = AutoConfig.for_model(model_type) + cfg = ( + getattr(raw, entry.hf_config_attr) + if entry.hf_config_attr and hasattr(raw, entry.hf_config_attr) + else raw + ) + cfg.num_hidden_layers = 2 + cfg.max_position_embeddings = MAX_CONTEXT_LENGTH + m = entry.macos_class(cfg, model_device="meta") + + spec = TraceSpec(max_context_length=MAX_CONTEXT_LENGTH) + refs = m.build_reference_inputs(cfg, torch.float16, spec) + shapes = m.build_dynamic_shapes(cfg, spec) + m.validate_export_contract(refs, shapes) + + params = list(inspect.signature(m.forward).parameters) + keys = list(refs[MAIN]) + assert keys == params[: len(keys)], model_type + + +class TestIOSHooksRaise: + """iOS models must not inherit the macOS contract. + + ``export/ios.py`` never calls these hooks, but the defaults were reachable: + ``--variant iOS --compression 4bit`` resolves a macOS quantization preset and the + quantize branch is not variant-gated. + """ + + HOOKS = ( + "export_input_names", + "export_state_names", + "export_output_names", + "build_reference_inputs", + "build_dynamic_shapes", + ) + + @pytest.fixture + def ios_model(self): + from transformers import AutoConfig + + from coreai_models.models.ios.qwen3 import Qwen3ForCausalLMForiOS + + cfg = AutoConfig.for_model("qwen3") + cfg.num_hidden_layers = 2 + cfg.max_position_embeddings = MAX_CONTEXT_LENGTH + return Qwen3ForCausalLMForiOS(cfg, model_device="meta") + + def _call(self, model, hook): + spec = TraceSpec(max_context_length=MAX_CONTEXT_LENGTH) + args = { + "build_reference_inputs": (model.config, torch.float16, spec), + "build_dynamic_shapes": (model.config, spec), + }.get(hook, ()) + return getattr(model, hook)(*args) + + @pytest.mark.parametrize("hook", HOOKS) + def test_hook_raises_not_implemented(self, ios_model, hook) -> None: + with pytest.raises(NotImplementedError) as exc: + self._call(ios_model, hook) + message = str(exc.value) + assert "not yet implemented for iOS models" in message + assert hook in message, "the error should name the hook that was called" + assert type(ios_model).__name__ in message + + def test_every_base_hook_is_overridden(self) -> None: + """Guard against a hook being added to the base and silently inherited.""" + from coreai_models.models.base import BaseForCausalLMForiOS + + for hook in self.HOOKS: + assert hook in vars(BaseForCausalLMForiOS), ( + f"{hook} is not overridden on BaseForCausalLMForiOS; iOS models would " + "inherit the macOS default" + ) + + def test_the_reachable_quantization_path_raises_clearly(self, ios_model) -> None: + """``--variant iOS --compression 4bit`` reaches ``quantize_for_export``.""" + from coreai_models.export.compression import quantize_for_export + from coreai_models.export.presets import get_preset + + quantization_config = get_preset("4bit").get("torch_quantization_config") + assert quantization_config is not None + with pytest.raises(NotImplementedError, match="not yet implemented for iOS models"): + quantize_for_export( + ios_model, ios_model.config, torch.float16, dict(quantization_config) + )