From 28f3ab62f8dbfb0869cc4c98a27c9eae9f25c6e2 Mon Sep 17 00:00:00 2001 From: Arik Horodniceanu Date: Tue, 12 May 2026 13:47:39 -0700 Subject: [PATCH] Qualcomm AI Engine Direct - [GenAI Pipeline] PR5: Model preparation & quantization strategy implementations --- .../model_preparation_output_config.py | 9 +- .../configs/quantization_input_config.py | 23 +- .../qualcomm/genai_pipeline/genai_pipeline.py | 2 + .../default_device_runner_adapter.py | 7 +- .../strategies/model_preparation/__init__.py | 4 + .../default_model_loader_adapter.py | 103 ++++- .../executorch_model_preparation_strategy.py | 209 +++++++++- .../model_preparation/model_loader_adapter.py | 27 +- .../quantization/default_quantizer_adapter.py | 43 +- .../executorch_quantization_strategy.py | 203 +++++++++- .../quantization/quantizer_adapter.py | 16 +- .../configs/test_quantization_input_config.py | 1 + .../test_default_model_loader_adapter.py | 85 ++++ ...t_executorch_model_preparation_strategy.py | 374 ++++++++++++++++++ .../test_executorch_quantization_strategy.py | 313 ++++++++++++++- 15 files changed, 1367 insertions(+), 52 deletions(-) create mode 100644 backends/qualcomm/genai_pipeline/tests/strategies/model_preparation/test_executorch_model_preparation_strategy.py diff --git a/backends/qualcomm/genai_pipeline/configs/model_preparation_output_config.py b/backends/qualcomm/genai_pipeline/configs/model_preparation_output_config.py index 25c079ffabd..194616d3563 100644 --- a/backends/qualcomm/genai_pipeline/configs/model_preparation_output_config.py +++ b/backends/qualcomm/genai_pipeline/configs/model_preparation_output_config.py @@ -8,7 +8,7 @@ from dataclasses import dataclass from pathlib import Path -from typing import Any, Iterable, Optional, TYPE_CHECKING +from typing import Any, Iterable, Optional, Tuple, TYPE_CHECKING if TYPE_CHECKING: from torch import nn @@ -27,6 +27,12 @@ class ModelPreparationOutputConfig: Attributes: model_module: The prepared nn.Module ready for quantization. tokenizer: The tokenizer instance for encoding/decoding text. + example_inputs: Positional example inputs for ``torch.export``, derived + from the **model** (never from ``calibration_data``): they carry the + exported graph's signature, its zero-initialized KV caches, and the + AR length baked in because HTP has no dynamic shapes. The dependency + runs model -> dataset, not the reverse -- the calibration dataset's + attention-mask schema is itself derived from this tuple. calibration_data: Calibration samples. Any Iterable[Tuple[Tensor, ...]], including a DataLoader with a custom collate_fn. runtime_tokenizer_path: Path to the runtime tokenizer **file** (not the @@ -36,6 +42,7 @@ class ModelPreparationOutputConfig: model_module: Optional["nn.Module"] = None tokenizer: Any = None + example_inputs: Optional[Tuple[Any, ...]] = None calibration_data: Optional[Iterable[Any]] = None runtime_tokenizer_path: Optional[Path] = None chat_template: Optional[str] = None diff --git a/backends/qualcomm/genai_pipeline/configs/quantization_input_config.py b/backends/qualcomm/genai_pipeline/configs/quantization_input_config.py index d64a9b4950c..30fdedeb904 100644 --- a/backends/qualcomm/genai_pipeline/configs/quantization_input_config.py +++ b/backends/qualcomm/genai_pipeline/configs/quantization_input_config.py @@ -7,7 +7,7 @@ from __future__ import annotations from dataclasses import dataclass, field -from typing import Any, Dict, Iterable, Optional, TYPE_CHECKING +from typing import Any, Dict, Iterable, Optional, Tuple, TYPE_CHECKING if TYPE_CHECKING: from executorch.backends.qualcomm.serialization.qc_schema import ( @@ -21,10 +21,11 @@ class QuantizationInputConfig: """Input configuration for the quantization stage. - ``model_module`` and ``calibration_data`` are ``Optional`` only because the - orchestrator builds this from the previous stage's output, which is empty - when model preparation is skipped. **Both are required once the quantization - stage executes**, and strategies should validate their presence. + ``model_module``, ``example_inputs`` and ``calibration_data`` are + ``Optional`` only because the orchestrator builds this from the previous + stage's output, which is empty when model preparation is skipped. **All + three are required once the quantization stage executes**, and strategies + should validate their presence. Flows needing no quantization (FP16, GPU backends) skip the stage entirely via ``GenAIPipeline.from_proxy(proxy, skip_stages={STAGE_QUANTIZATION})`` @@ -36,8 +37,17 @@ class QuantizationInputConfig: soc_model: The target SoC (e.g., QcomChipset.SM8750). Required. backend_type: QNN backend type (HTP, GPU, LPAI, etc.). Required. model_module: The nn.Module to quantize. Required when the stage runs. + example_inputs: Positional example inputs for ``torch.export``. Required + when the stage runs. Sourced from the **model** via + ``ModelLoaderAdapter.get_example_inputs``, never from + ``calibration_data``: this tuple defines the exported graph's + positional signature, supplies the zero-initialized KV caches a + dataset sample does not carry, and fixes the AR length because HTP + has no dynamic shapes. calibration_data: Calibration samples. Required when the stage runs. Any - Iterable[Tuple[Tensor, ...]], including a DataLoader. + Iterable[Tuple[Tensor, ...]], including a DataLoader. Consumed only + by ``calibrate()`` -- it is never indexed or peeked at, so a + single-use generator stays intact. training_data: Training dataset for quantization-aware training (QAT), typically (features, labels) pairs. Mirrors ``qat_training_data`` in ``build_executorch_binary``. ``None`` selects PTQ. @@ -48,6 +58,7 @@ class QuantizationInputConfig: soc_model: "QcomChipset" backend_type: "QnnExecuTorchBackendType" model_module: Optional["nn.Module"] = None + example_inputs: Optional[Tuple[Any, ...]] = None calibration_data: Optional[Iterable[Any]] = None training_data: Optional[Iterable[Any]] = None quant_recipe: Any = None diff --git a/backends/qualcomm/genai_pipeline/genai_pipeline.py b/backends/qualcomm/genai_pipeline/genai_pipeline.py index 8da56e4ee9f..e419d1ac4ec 100644 --- a/backends/qualcomm/genai_pipeline/genai_pipeline.py +++ b/backends/qualcomm/genai_pipeline/genai_pipeline.py @@ -213,6 +213,8 @@ def _run_quantization( soc_model=context.soc_model, backend_type=self._engine_proxy.backend_type, model_module=model_prep_output.model_module, + # Export inputs come from the model, not from calibration_data. + example_inputs=model_prep_output.example_inputs, calibration_data=model_prep_output.calibration_data, ) output = self._quantization_stage.invoke(context, input_config) diff --git a/backends/qualcomm/genai_pipeline/strategies/inference/default_device_runner_adapter.py b/backends/qualcomm/genai_pipeline/strategies/inference/default_device_runner_adapter.py index 49adab3ed4b..ac41e684bc3 100644 --- a/backends/qualcomm/genai_pipeline/strategies/inference/default_device_runner_adapter.py +++ b/backends/qualcomm/genai_pipeline/strategies/inference/default_device_runner_adapter.py @@ -77,13 +77,18 @@ def execute( ) -> InferenceResult: """Execute the model on device via ADB. + Note: This is a two-step protocol. ``execute()`` runs the model and + returns performance metrics. Call ``pull_results()`` afterward to + retrieve the actual output data files from the device. + Args: inference_options: Engine-specific options. Supported keys: - ``method_index``: Index of the method to execute (default 0). - ``iteration``: Number of inference iterations (default 1). Returns: - InferenceResult with output data and performance metrics. + InferenceResult with performance metrics. ``output_data`` is None + until ``pull_results()`` is called separately. """ if self._adb is None: raise RuntimeError( diff --git a/backends/qualcomm/genai_pipeline/strategies/model_preparation/__init__.py b/backends/qualcomm/genai_pipeline/strategies/model_preparation/__init__.py index 45b2bc3f3c3..5b7aef9cf69 100644 --- a/backends/qualcomm/genai_pipeline/strategies/model_preparation/__init__.py +++ b/backends/qualcomm/genai_pipeline/strategies/model_preparation/__init__.py @@ -4,6 +4,9 @@ # This source code is licensed under the BSD-style license found in the # LICENSE file in the root directory of this source tree. +from executorch.backends.qualcomm.genai_pipeline.strategies.model_preparation.executorch_model_preparation_strategy import ( + ExecuTorchModelPreparationStrategy, +) from executorch.backends.qualcomm.genai_pipeline.strategies.model_preparation.model_loader_adapter import ( ModelLoaderAdapter, ) @@ -12,6 +15,7 @@ ) __all__ = [ + "ExecuTorchModelPreparationStrategy", "ModelLoaderAdapter", "ModelPreparationStrategy", ] diff --git a/backends/qualcomm/genai_pipeline/strategies/model_preparation/default_model_loader_adapter.py b/backends/qualcomm/genai_pipeline/strategies/model_preparation/default_model_loader_adapter.py index bdc6de78182..ef96a17b98f 100644 --- a/backends/qualcomm/genai_pipeline/strategies/model_preparation/default_model_loader_adapter.py +++ b/backends/qualcomm/genai_pipeline/strategies/model_preparation/default_model_loader_adapter.py @@ -8,7 +8,7 @@ import logging from pathlib import Path -from typing import Any, Dict, Optional +from typing import Any, Dict, Optional, Tuple logger = logging.getLogger(__name__) @@ -30,6 +30,17 @@ class DefaultModelLoaderAdapter: ``CalibrationDataAdapter``. """ + #: Batch size and sequence length of the generated example inputs. HTP has no + #: dynamic shapes, so these dimensions are baked into the exported graph. + DEFAULT_BATCH_SIZE = 1 + DEFAULT_AR_LEN = 1 + + #: Preferred runtime tokenizer file names, in priority order. + #: ``pytorch_tokenizers.get_tokenizer`` dispatches on the file extension + #: (``.json`` -> ``HuggingFaceTokenizer``, otherwise Llama2c/Tiktoken), so the + #: file we hand back selects the runtime tokenizer implementation. + RUNTIME_TOKENIZER_NAMES = ("tokenizer.json", "tokenizer.model") + def load_model( self, model_name: str, @@ -94,6 +105,51 @@ def load_tokenizer( logger.info("Tokenizer loaded successfully") return tokenizer + def get_example_inputs( + self, + model: Any, + extra_options: Optional[Dict[str, Any]] = None, + ) -> Tuple[Any, ...]: + """Build example inputs for ``torch.export`` from the model itself. + + Prefers the model's own ``get_example_inputs()`` when it exposes one, so + models that already describe their export signature (the LLM wrappers + build a flat ``(tokens, attn_mask, pos_ids, *k_caches, *v_caches)`` + tuple) stay authoritative. Otherwise a minimal ``(input_ids,)`` is + synthesized, which is the correct signature for a plain HuggingFace + causal LM without an external KV cache. + + Args: + model: The module returned by :meth:`load_model`. + extra_options: Additional options. Supported keys: + - ``batch_size``: Batch dimension (default: + ``DEFAULT_BATCH_SIZE``). + - ``ar_len``: Sequence length / autoregressive window + (default: ``DEFAULT_AR_LEN``). + + Returns: + A flat tuple positionally matching ``model.forward``. + """ + import torch + + extra_options = extra_options or {} + + model_provided = getattr(model, "get_example_inputs", None) + if callable(model_provided): + logger.info("Using example inputs provided by the model") + return tuple(model_provided()) + + batch_size = extra_options.get("batch_size", self.DEFAULT_BATCH_SIZE) + ar_len = extra_options.get("ar_len", self.DEFAULT_AR_LEN) + + logger.info( + "Synthesizing example inputs with batch_size=%d, ar_len=%d", + batch_size, + ar_len, + ) + # int64 token ids: the embedding lookup indexes with them. + return (torch.zeros((batch_size, ar_len), dtype=torch.int64),) + def export_tokenizer( self, tokenizer: Any, @@ -102,11 +158,19 @@ def export_tokenizer( ) -> Path: """Export tokenizer to disk and return the runtime tokenizer file. - ``save_pretrained`` writes several files and returns the tuple of paths - it wrote, with the tokenizer file last. Both ``llm::load_tokenizer`` and - ``pytorch_tokenizers.get_tokenizer`` expect that **single file**, not the - containing directory, so we return it -- mirroring the existing - ``TokenizerWrapper._from_hf`` flow. + ``save_pretrained`` writes several files and returns the tuple of paths it + wrote. Both ``llm::load_tokenizer`` and ``pytorch_tokenizers.get_tokenizer`` + expect a **single file**, not the containing directory, so one artifact has + to be singled out. + + The file is chosen **by name** -- ``tokenizer.json`` first, then + ``tokenizer.model`` -- rather than by position in the returned tuple. + ``get_tokenizer`` dispatches on the extension, so picking the wrong + artifact silently constructs the wrong tokenizer class instead of raising, + and ``save_pretrained``'s ordering is an implementation detail that varies + with the tokenizer (fast vs slow, whether ``added_tokens.json`` is + written). ``artifacts[-1]`` remains a last-resort fallback for tokenizers + that emit neither name, mirroring ``TokenizerWrapper._from_hf``. Args: tokenizer: The tokenizer instance to export. @@ -128,6 +192,31 @@ def export_tokenizer( f"save_pretrained() reported no tokenizer artifacts in {output_dir}." ) - runtime_tokenizer_path = Path(artifacts[-1]) + runtime_tokenizer_path = self._select_runtime_tokenizer(artifacts) logger.info("Tokenizer exported to %s", runtime_tokenizer_path) return runtime_tokenizer_path + + @classmethod + def _select_runtime_tokenizer(cls, artifacts: Any) -> Path: + """Pick the runtime tokenizer file out of ``save_pretrained``'s artifacts. + + Args: + artifacts: The paths reported by ``save_pretrained``. + + Returns: + The first artifact matching :attr:`RUNTIME_TOKENIZER_NAMES`, falling + back to the last artifact when none matches. + """ + paths = [Path(artifact) for artifact in artifacts] + + for name in cls.RUNTIME_TOKENIZER_NAMES: + for path in paths: + if path.name == name: + return path + + logger.warning( + "None of %s found among tokenizer artifacts; falling back to %s.", + ", ".join(cls.RUNTIME_TOKENIZER_NAMES), + paths[-1], + ) + return paths[-1] diff --git a/backends/qualcomm/genai_pipeline/strategies/model_preparation/executorch_model_preparation_strategy.py b/backends/qualcomm/genai_pipeline/strategies/model_preparation/executorch_model_preparation_strategy.py index 8322f7a68b6..1d9ea05a975 100644 --- a/backends/qualcomm/genai_pipeline/strategies/model_preparation/executorch_model_preparation_strategy.py +++ b/backends/qualcomm/genai_pipeline/strategies/model_preparation/executorch_model_preparation_strategy.py @@ -4,25 +4,100 @@ # This source code is licensed under the BSD-style license found in the # LICENSE file in the root directory of this source tree. +from __future__ import annotations + +import logging +from pathlib import Path +from typing import Optional + from executorch.backends.qualcomm.genai_pipeline.configs.model_preparation_input_config import ( ModelPreparationInputConfig, ) from executorch.backends.qualcomm.genai_pipeline.configs.model_preparation_output_config import ( ModelPreparationOutputConfig, ) +from executorch.backends.qualcomm.genai_pipeline.datasets.calibration_data_adapter import ( + CalibrationDataAdapter, +) +from executorch.backends.qualcomm.genai_pipeline.datasets.default_calibration_data_adapter import ( + DEFAULT_NUM_SAMPLES, + DEFAULT_SEQ_LENGTH, +) +from executorch.backends.qualcomm.genai_pipeline.exceptions import StageError from executorch.backends.qualcomm.genai_pipeline.pipeline_context import PipelineContext +from executorch.backends.qualcomm.genai_pipeline.strategies.model_preparation.model_loader_adapter import ( + ModelLoaderAdapter, +) from executorch.backends.qualcomm.genai_pipeline.strategies.model_preparation.model_preparation_strategy import ( ModelPreparationStrategy, ) +logger = logging.getLogger(__name__) + +_STAGE_NAME = "model_preparation" + class ExecuTorchModelPreparationStrategy(ModelPreparationStrategy): """ExecuTorch-based model preparation using HuggingFace transformers. - Loads model weights, tokenizer, and generates calibration data - for downstream quantization and compilation stages. + Delegates to injectable adapters for all external API calls, enabling + dependency injection for testability: + + * ``ModelLoaderAdapter`` acquires the model, its tokenizer, and the example + inputs describing the model's ``torch.export`` signature. + * ``CalibrationDataAdapter`` produces the calibration corpus. Datasets are a + cross-stage concern -- the same corpus feeds PTQ calibration and on-device + evaluation -- so they live in ``genai_pipeline.datasets`` rather than being + a model-loading responsibility. + + The model preparation flow: + 1. Validate the input configuration + 2. Load the model + 3. Load the tokenizer + 4. Build the export example inputs from the model + 5. Generate calibration data + 6. Optionally export the tokenizer for runtime use + 7. Extract the chat template + + Args: + model_loader_adapter: Injectable adapter for model and tokenizer + loading. Defaults to ``DefaultModelLoaderAdapter`` if not provided. + calibration_data_adapter: Injectable adapter for calibration data + generation. Defaults to ``DefaultCalibrationDataAdapter`` if not + provided. """ + def __init__( + self, + model_loader_adapter: Optional[ModelLoaderAdapter] = None, + calibration_data_adapter: Optional[CalibrationDataAdapter] = None, + ) -> None: + if model_loader_adapter is None: + from executorch.backends.qualcomm.genai_pipeline.strategies.model_preparation.default_model_loader_adapter import ( + DefaultModelLoaderAdapter, + ) + + model_loader_adapter = DefaultModelLoaderAdapter() + self._adapter = model_loader_adapter + + if calibration_data_adapter is None: + from executorch.backends.qualcomm.genai_pipeline.datasets.default_calibration_data_adapter import ( + DefaultCalibrationDataAdapter, + ) + + calibration_data_adapter = DefaultCalibrationDataAdapter() + self._calibration_adapter = calibration_data_adapter + + @property + def adapter(self) -> ModelLoaderAdapter: + """The model loader adapter used by this strategy.""" + return self._adapter + + @property + def calibration_data_adapter(self) -> CalibrationDataAdapter: + """The calibration data adapter used by this strategy.""" + return self._calibration_adapter + def invoke( self, context: PipelineContext, @@ -31,13 +106,135 @@ def invoke( """Prepare the model, tokenizer, and calibration data. Args: - context: The pipeline context. + context: The pipeline context with global settings. input_config: The model preparation input configuration. + Supported keys in ``input_config.extra_options``: + - ``model_options``: Dict passed to ``load_model(extra_options=...)``. + - ``tokenizer_options``: Dict passed to ``load_tokenizer(extra_options=...)``. + - ``example_input_options``: Dict passed to + ``get_example_inputs(extra_options=...)``. + - ``num_calibration_samples``: Number of calibration samples + (default: ``DEFAULT_NUM_SAMPLES``). + - ``calibration_seq_length``: Sequence length per sample + (default: ``DEFAULT_SEQ_LENGTH``). + - ``calibration_options``: Dict passed to ``generate_calibration_data(extra_options=...)``. + - ``export_tokenizer``: If True, export tokenizer for runtime (default: False). + - ``tokenizer_export_options``: Dict passed to ``export_tokenizer(extra_options=...)``. + - ``chat_template``: Explicit chat template, used only when the + tokenizer does not carry one. Returns: ModelPreparationOutputConfig with model, tokenizer, and calibration data. + + Raises: + StageError: If model_name is missing or any loading step fails. """ - raise NotImplementedError( - "ExecuTorchModelPreparationStrategy.invoke() is a stub. " - "Implementation will be added in a subsequent PR." + logger.info( + "Starting model preparation for '%s' on SoC=%s", + input_config.model_name, + input_config.soc_model, ) + + self._validate_input(input_config) + + try: + extra = dict(input_config.extra_options) + + # Step 1: Load model + logger.debug("Loading model") + model_module = self._adapter.load_model( + model_name=input_config.model_name, + extra_options=extra.get("model_options"), + ) + + # Step 2: Load tokenizer + logger.debug("Loading tokenizer") + tokenizer = self._adapter.load_tokenizer( + model_name=input_config.model_name, + extra_options=extra.get("tokenizer_options"), + ) + + # Step 3: Build the export example inputs from the model itself. + # These are deliberately *not* taken from the calibration dataset: + # they define the exported graph's positional signature (including + # zero-initialized KV caches, which no dataset sample carries) and + # the dataset's own attention-mask schema is derived from them. + logger.debug("Building example inputs for export") + example_inputs = self._adapter.get_example_inputs( + model=model_module, + extra_options=extra.get("example_input_options"), + ) + + # Step 4: Generate calibration data via the cross-stage dataset adapter + logger.debug("Generating calibration data") + num_samples = extra.get("num_calibration_samples", DEFAULT_NUM_SAMPLES) + seq_length = extra.get("calibration_seq_length", DEFAULT_SEQ_LENGTH) + calibration_data = self._calibration_adapter.generate_calibration_data( + tokenizer=tokenizer, + num_samples=num_samples, + seq_length=seq_length, + extra_options=extra.get("calibration_options"), + ) + + # Step 5: Optionally export tokenizer for runtime + runtime_tokenizer_path = None + if extra.get("export_tokenizer", False): + logger.debug("Exporting tokenizer for runtime use") + output_dir = Path(context.artifact_dir) / "tokenizer" + runtime_tokenizer_path = self._adapter.export_tokenizer( + tokenizer=tokenizer, + output_dir=output_dir, + extra_options=extra.get("tokenizer_export_options"), + ) + + # Step 6: Extract chat_template from tokenizer (for instruct models). + # The tokenizer wins over extra_options: a template shipped with the + # model is authoritative, and extra_options is only a fallback for + # models that carry none. + chat_template = None + if getattr(tokenizer, "chat_template", None): + chat_template = tokenizer.chat_template + logger.debug("Chat template extracted from tokenizer") + elif extra.get("chat_template"): + chat_template = extra["chat_template"] + logger.debug("Chat template provided via extra_options") + + logger.info("Model preparation completed successfully") + + return ModelPreparationOutputConfig( + model_module=model_module, + tokenizer=tokenizer, + example_inputs=example_inputs, + calibration_data=calibration_data, + runtime_tokenizer_path=runtime_tokenizer_path, + chat_template=chat_template, + ) + + except StageError: + raise + except Exception as e: + raise StageError( + stage_name=_STAGE_NAME, + message="Model preparation failed", + original_exception=e, + ) from e + + def _validate_input(self, input_config: ModelPreparationInputConfig) -> None: + """Validate required fields in the input configuration. + + Args: + input_config: The model preparation input configuration. + + Raises: + StageError: If required fields are missing. + """ + if not input_config.model_name: + raise StageError( + stage_name=_STAGE_NAME, + message="model_name is required for model preparation", + ) + if not input_config.soc_model: + raise StageError( + stage_name=_STAGE_NAME, + message="soc_model is required for model preparation", + ) diff --git a/backends/qualcomm/genai_pipeline/strategies/model_preparation/model_loader_adapter.py b/backends/qualcomm/genai_pipeline/strategies/model_preparation/model_loader_adapter.py index 939757a4f8d..fedefd63071 100644 --- a/backends/qualcomm/genai_pipeline/strategies/model_preparation/model_loader_adapter.py +++ b/backends/qualcomm/genai_pipeline/strategies/model_preparation/model_loader_adapter.py @@ -7,7 +7,7 @@ from __future__ import annotations from pathlib import Path -from typing import Any, Dict, Optional, Protocol, runtime_checkable +from typing import Any, Dict, Optional, Protocol, runtime_checkable, Tuple @runtime_checkable @@ -82,6 +82,31 @@ def load_tokenizer( """ ... + def get_example_inputs( + self, + model: Any, + extra_options: Optional[Dict[str, Any]] = None, + ) -> Tuple[Any, ...]: + """Build the positional example inputs for ``torch.export``. + + These come from the **model**, never from the calibration dataset: the + tuple defines the exported graph's positional signature, including the + zero-initialized KV cache entries a dataset sample does not carry, and + it bakes in the AR length (HTP has no dynamic shapes). The dependency + runs model -> dataset, not the reverse: the reference flow derives the + dataset's attention-mask schema *from* the example input + (``LLMWrapper.attn_mask`` returns ``example_input[1]``). + + Args: + model: The module previously returned by :meth:`load_model`. + extra_options: Additional options controlling the example shapes. + + Returns: + A flat tuple positionally matching ``model.forward``, ready to pass + straight to ``torch.export.export(model, example_inputs)``. + """ + ... + def export_tokenizer( self, tokenizer: Any, diff --git a/backends/qualcomm/genai_pipeline/strategies/quantization/default_quantizer_adapter.py b/backends/qualcomm/genai_pipeline/strategies/quantization/default_quantizer_adapter.py index 27a28138c1a..5ce3a5f2200 100644 --- a/backends/qualcomm/genai_pipeline/strategies/quantization/default_quantizer_adapter.py +++ b/backends/qualcomm/genai_pipeline/strategies/quantization/default_quantizer_adapter.py @@ -23,17 +23,30 @@ class DefaultQuantizerAdapter: def make_quantizer( self, - quant_dtype: Any, - backend: Any, - soc_model: Any, + quant_dtype: Any = None, + backend: Any = None, + soc_model: Any = None, + quant_recipe: Any = None, **kwargs: Any, ) -> Any: """Create a QNN quantizer via ``export_utils.make_quantizer``. + ``quant_dtype`` defaults to ``None`` and is only forwarded when set, so + ``export_utils.make_quantizer`` remains the single owner of the default + (``QuantDtype.use_8a8w``) rather than this wrapper duplicating it. + + ``quant_recipe`` is **not** an argument of + ``export_utils.make_quantizer``; a recipe is applied to the constructed + quantizer via ``QnnQuantizer.set_recipe``, so it is consumed here and + never forwarded. + Args: - quant_dtype: Quantization data type. + quant_dtype: Quantization data type. ``None`` leaves the default to + ``export_utils.make_quantizer``. backend: QNN backend type enum. soc_model: Target SoC (string name like "SM8750" or QcomChipset enum). + quant_recipe: Optional recipe applied via ``set_recipe`` after the + quantizer is constructed. **kwargs: Forwarded to ``make_quantizer``. Returns: @@ -48,17 +61,27 @@ def make_quantizer( soc_model_str = soc_model.name if hasattr(soc_model, "name") else str(soc_model) logger.debug( - "Creating quantizer: dtype=%s, backend=%s, soc=%s", + "Creating quantizer: dtype=%s, backend=%s, soc=%s, recipe=%s", quant_dtype, backend, soc_model_str, + quant_recipe, ) - return _make_quantizer( - quant_dtype=quant_dtype, - backend=backend, - soc_model=soc_model_str, + make_quantizer_kwargs = { + "backend": backend, + "soc_model": soc_model_str, **kwargs, - ) + } + if quant_dtype is not None: + make_quantizer_kwargs["quant_dtype"] = quant_dtype + + quantizer = _make_quantizer(**make_quantizer_kwargs) + + if quant_recipe is not None: + logger.debug("Applying quantization recipe via set_recipe") + quantizer.set_recipe(quant_recipe) + + return quantizer def export_model( self, diff --git a/backends/qualcomm/genai_pipeline/strategies/quantization/executorch_quantization_strategy.py b/backends/qualcomm/genai_pipeline/strategies/quantization/executorch_quantization_strategy.py index 998a1d8d0ba..33102f89d3d 100644 --- a/backends/qualcomm/genai_pipeline/strategies/quantization/executorch_quantization_strategy.py +++ b/backends/qualcomm/genai_pipeline/strategies/quantization/executorch_quantization_strategy.py @@ -4,24 +4,89 @@ # This source code is licensed under the BSD-style license found in the # LICENSE file in the root directory of this source tree. +from __future__ import annotations + +import logging +from typing import Optional + from executorch.backends.qualcomm.genai_pipeline.configs.quantization_input_config import ( QuantizationInputConfig, ) from executorch.backends.qualcomm.genai_pipeline.configs.quantization_output_config import ( QuantizationOutputConfig, ) +from executorch.backends.qualcomm.genai_pipeline.exceptions import StageError from executorch.backends.qualcomm.genai_pipeline.pipeline_context import PipelineContext from executorch.backends.qualcomm.genai_pipeline.strategies.quantization.quantization_strategy import ( QuantizationStrategy, ) +from executorch.backends.qualcomm.genai_pipeline.strategies.quantization.quantizer_adapter import ( + QuantizerAdapter, +) + +logger = logging.getLogger(__name__) + +_STAGE_NAME = "quantization" class ExecuTorchQuantizationStrategy(QuantizationStrategy): """ExecuTorch-based quantization using QNN quantizer annotator rules. - Reads backend_type from the input config to select the appropriate annotator (HTP, GPU, or LPAI rules). + Delegates to a ``QuantizerAdapter`` for all external API calls, + enabling dependency injection for testability. + + The quantization flow follows the PT2E pattern: + 1. Export the model via ``torch.export`` using ``input_config.example_inputs`` + 2. Create a QNN quantizer with appropriate backend rules + 3. Prepare the model (insert observers) + 4. Calibrate with provided dataset + 5. Convert to quantized model + + .. note:: + **Single-graph only; multi-graph is a tracked follow-up.** This sequence + runs one ``prepare_pt2e`` / calibrate / ``convert_pt2e`` pass over one + module. Models exported as several graphs from the same weights (the + hybrid AR-N prefill / AR-1 decode pair) need a different, *asymmetric* + orchestration, which is why it is deliberately not attempted here: + + * every graph runs the full ``prepare_pt2e`` -> run -> ``convert_pt2e`` + sequence, but the data fed in between differs per graph: the + calibration-only graph (full AR sequence with KV cache, never + deployed) receives the real dataset, while the deployed prefill and + decode graphs receive their own ``example_inputs`` -- one pass is + still required there, or ``convert_pt2e`` fails on uninitialized observers; + * the scales/zero-points collected on the calibration graph are then + propagated to prefill and decode by an encoding-reconciliation step. + + So the eventual interface is a per-graph map of module to data source, plus a + reconciliation hook, not a bare Dict[str, nn.Module]. + Landing that requires graph-map fields on the quantization configs, which are + additive to this dataclass, so deferring costs nothing structurally. + Per the layering used throughout this package, the fan-out belongs in + this *strategy* -- adapters stay thin 1:1 wrappers over one graph. + + Args: + quantizer_adapter: Injectable adapter for quantization operations. + Defaults to ``DefaultQuantizerAdapter`` if not provided. """ + def __init__( + self, + quantizer_adapter: Optional[QuantizerAdapter] = None, + ) -> None: + if quantizer_adapter is None: + from executorch.backends.qualcomm.genai_pipeline.strategies.quantization.default_quantizer_adapter import ( + DefaultQuantizerAdapter, + ) + + quantizer_adapter = DefaultQuantizerAdapter() + self._adapter = quantizer_adapter + + @property + def adapter(self) -> QuantizerAdapter: + """The quantizer adapter used by this strategy.""" + return self._adapter + def invoke( self, context: PipelineContext, @@ -29,14 +94,140 @@ def invoke( ) -> QuantizationOutputConfig: """Quantize the model using ExecuTorch/QNN quantization. + Executes the full PT2E quantization pipeline: + export → make_quantizer → prepare_pt2e → calibrate → convert_pt2e. + Args: - context: The pipeline context. + context: The pipeline context with global settings. input_config: The quantization input configuration. Returns: - QuantizationOutputConfig with quantized model. + QuantizationOutputConfig with the quantized model. + + Raises: + StageError: If the model, example inputs or calibration data is + missing, or if any quantization step fails. """ - raise NotImplementedError( - "ExecuTorchQuantizationStrategy.invoke() is a stub. " - "Implementation will be added in a subsequent PR." + logger.info( + "Starting quantization for model '%s' on SoC=%s, backend=%s", + context.model_name, + input_config.soc_model, + input_config.backend_type, ) + + self._validate_input(input_config) + + if input_config.training_data is not None: + # QAT is not yet wired up: ``training_data`` is carried on the config + # (mirroring ``qat_training_data`` in ``build_executorch_binary``) so + # the contract is stable, but this strategy only implements PTQ. + logger.warning( + "training_data was provided but quantization-aware training is " + "not implemented by this strategy; proceeding with PTQ." + ) + + try: + # Step 1: Export model. The example inputs describe the model's own + # export signature (see ``QuantizationInputConfig.example_inputs``); + # they are never drawn from ``calibration_data``, which is left + # untouched so that even a single-use generator reaches ``calibrate`` + # with every sample intact. + logger.debug("Exporting model") + exported_model = self._adapter.export_model( + input_config.model_module, + input_config.example_inputs, + ) + + # Step 2: Create quantizer + logger.debug("Creating quantizer") + quant_kwargs = dict(input_config.extra_options) + quant_dtype = quant_kwargs.pop("quant_dtype", None) + quant_recipe = quant_kwargs.pop("quant_recipe", None) or getattr( + input_config, "quant_recipe", None + ) + + # Build make_quantizer arguments — only pass quant_dtype if + # explicitly provided, so the default owned by + # ``export_utils.make_quantizer`` (use_8a8w) applies otherwise + # instead of being shadowed by a value chosen here. + # ``quant_recipe`` is not an argument of that function: the adapter + # consumes it and applies it to the constructed quantizer via + # ``QnnQuantizer.set_recipe``. + make_quantizer_kwargs = { + "backend": input_config.backend_type, + "soc_model": input_config.soc_model, + **quant_kwargs, + } + if quant_dtype is not None: + make_quantizer_kwargs["quant_dtype"] = quant_dtype + if quant_recipe is not None: + make_quantizer_kwargs["quant_recipe"] = quant_recipe + + quantizer = self._adapter.make_quantizer(**make_quantizer_kwargs) + + # Step 3: Prepare (insert observers) + logger.debug("Preparing model for quantization") + annotated_model = self._adapter.prepare_pt2e(exported_model, quantizer) + + # Step 4: Calibrate + logger.debug("Running calibration") + calibrated_model = self._adapter.calibrate( + annotated_model, input_config.calibration_data + ) + + # Step 5: Convert to quantized model + logger.debug("Converting to quantized model") + quantized_model = self._adapter.convert_pt2e(calibrated_model) + + logger.info("Quantization completed successfully") + return QuantizationOutputConfig(quantized_model=quantized_model) + + except StageError: + raise + except Exception as e: + raise StageError( + stage_name=_STAGE_NAME, + message="Quantization failed", + original_exception=e, + ) from e + + def _validate_input(self, input_config: QuantizationInputConfig) -> None: + """Validate required fields in the input configuration. + + Args: + input_config: The quantization input configuration. + + Raises: + StageError: If required fields are missing. + """ + if input_config.model_module is None: + raise StageError( + stage_name=_STAGE_NAME, + message="model_module is required for quantization", + ) + if input_config.example_inputs is None: + raise StageError( + stage_name=_STAGE_NAME, + message=( + "example_inputs is required for quantization; it is produced " + "from the model by ModelLoaderAdapter.get_example_inputs" + ), + ) + # ``is None`` rather than a truthiness test: ``calibration_data`` may be a + # generator or ``DataLoader``, and ``not `` would consume the + # first sample without reliably detecting emptiness. + if input_config.calibration_data is None: + raise StageError( + stage_name=_STAGE_NAME, + message="calibration_data is required for quantization", + ) + if input_config.soc_model is None: + raise StageError( + stage_name=_STAGE_NAME, + message="soc_model is required for quantization", + ) + if input_config.backend_type is None: + raise StageError( + stage_name=_STAGE_NAME, + message="backend_type is required for quantization", + ) diff --git a/backends/qualcomm/genai_pipeline/strategies/quantization/quantizer_adapter.py b/backends/qualcomm/genai_pipeline/strategies/quantization/quantizer_adapter.py index 970eb67e3a4..c03b79554bf 100644 --- a/backends/qualcomm/genai_pipeline/strategies/quantization/quantizer_adapter.py +++ b/backends/qualcomm/genai_pipeline/strategies/quantization/quantizer_adapter.py @@ -33,17 +33,27 @@ class QuantizerAdapter(Protocol): def make_quantizer( self, - quant_dtype: Any, - backend: Any, - soc_model: Any, + quant_dtype: Any = None, + backend: Any = None, + soc_model: Any = None, + quant_recipe: Any = None, **kwargs: Any, ) -> Any: """Create a QNN quantizer with the given configuration. + Every argument defaults to ``None`` so that callers can omit any of them + and let the implementation -- or the API it wraps -- supply the default. + In particular an omitted ``quant_dtype`` must not be forwarded, so the + underlying ``make_quantizer`` default applies rather than being shadowed. + Args: quant_dtype: Quantization data type (e.g., QuantDtype.use_8a8w). + ``None`` selects the implementation's default. backend: QNN backend type (HTP, GPU, LPAI). soc_model: Target SoC chipset. + quant_recipe: Optional quantization recipe. Applied to the + constructed quantizer (``QnnQuantizer.set_recipe``) rather than + passed to ``make_quantizer``, which takes no such argument. **kwargs: Additional quantizer options (per_channel, observers, etc.). Returns: diff --git a/backends/qualcomm/genai_pipeline/tests/configs/test_quantization_input_config.py b/backends/qualcomm/genai_pipeline/tests/configs/test_quantization_input_config.py index 279039cbe22..86693c77ebe 100644 --- a/backends/qualcomm/genai_pipeline/tests/configs/test_quantization_input_config.py +++ b/backends/qualcomm/genai_pipeline/tests/configs/test_quantization_input_config.py @@ -29,6 +29,7 @@ def test_optional_fields_default_to_none(self): soc_model=MagicMock(), backend_type=MagicMock() ) self.assertIsNone(config.model_module) + self.assertIsNone(config.example_inputs) self.assertIsNone(config.calibration_data) self.assertIsNone(config.training_data) self.assertIsNone(config.quant_recipe) diff --git a/backends/qualcomm/genai_pipeline/tests/strategies/model_preparation/test_default_model_loader_adapter.py b/backends/qualcomm/genai_pipeline/tests/strategies/model_preparation/test_default_model_loader_adapter.py index c9b07a96e06..bc9820ee5af 100644 --- a/backends/qualcomm/genai_pipeline/tests/strategies/model_preparation/test_default_model_loader_adapter.py +++ b/backends/qualcomm/genai_pipeline/tests/strategies/model_preparation/test_default_model_loader_adapter.py @@ -16,6 +16,8 @@ TEST_TOKENIZER_CONFIG = "tokenizer_config.json" TEST_SPECIAL_TOKENS_MAP = "special_tokens_map.json" TEST_TOKENIZER_JSON = "tokenizer.json" +TEST_TOKENIZER_MODEL = "tokenizer.model" +TEST_ADDED_TOKENS = "added_tokens.json" class TestExportTokenizer(unittest.TestCase): @@ -60,6 +62,89 @@ def test_raises_when_no_artifacts_written(self): self._make_tokenizer(artifacts), self.output_dir ) + def test_selects_by_name_not_by_position(self): + # save_pretrained's ordering is an implementation detail: a fast + # tokenizer may append added_tokens.json after tokenizer.json. Selecting + # positionally would hand the runtime the wrong file, and because + # get_tokenizer dispatches on the extension it would construct the wrong + # tokenizer class rather than fail. + artifacts = ( + str(self.output_dir / TEST_TOKENIZER_CONFIG), + str(self.output_dir / TEST_TOKENIZER_JSON), + str(self.output_dir / TEST_ADDED_TOKENS), + ) + result = self.adapter.export_tokenizer( + self._make_tokenizer(artifacts), self.output_dir + ) + self.assertEqual(result, self.output_dir / TEST_TOKENIZER_JSON) + + def test_prefers_tokenizer_json_over_tokenizer_model(self): + artifacts = ( + str(self.output_dir / TEST_TOKENIZER_MODEL), + str(self.output_dir / TEST_TOKENIZER_JSON), + ) + result = self.adapter.export_tokenizer( + self._make_tokenizer(artifacts), self.output_dir + ) + self.assertEqual(result, self.output_dir / TEST_TOKENIZER_JSON) + + def test_falls_back_to_tokenizer_model_when_no_json(self): + artifacts = ( + str(self.output_dir / TEST_TOKENIZER_CONFIG), + str(self.output_dir / TEST_TOKENIZER_MODEL), + str(self.output_dir / TEST_SPECIAL_TOKENS_MAP), + ) + result = self.adapter.export_tokenizer( + self._make_tokenizer(artifacts), self.output_dir + ) + self.assertEqual(result, self.output_dir / TEST_TOKENIZER_MODEL) + + def test_falls_back_to_last_artifact_when_no_known_name(self): + # e.g. stories110m ships tokenizer.bin; keep save_pretrained's ordering + # as a last resort rather than failing outright. + artifacts = ( + str(self.output_dir / TEST_TOKENIZER_CONFIG), + str(self.output_dir / "tokenizer.bin"), + ) + result = self.adapter.export_tokenizer( + self._make_tokenizer(artifacts), self.output_dir + ) + self.assertEqual(result, self.output_dir / "tokenizer.bin") + + +class TestGetExampleInputs(unittest.TestCase): + """Export inputs describe the *model's* signature, so a model that already + knows its own signature must win over anything synthesized here.""" + + def setUp(self): + self.adapter = DefaultModelLoaderAdapter() + + def test_prefers_example_inputs_provided_by_the_model(self): + expected = (MagicMock(name="tokens"), MagicMock(name="attn_mask")) + model = MagicMock() + model.get_example_inputs.return_value = expected + + result = self.adapter.get_example_inputs(model) + + self.assertEqual(result, expected) + model.get_example_inputs.assert_called_once_with() + + def test_synthesizes_int64_token_ids_when_model_provides_none(self): + import torch + + # `spec=[]` gives an object with no attributes, so the adapter cannot + # find a `get_example_inputs` to defer to. + model = MagicMock(spec=[]) + + result = self.adapter.get_example_inputs( + model, extra_options={"batch_size": 2, "ar_len": 8} + ) + + self.assertEqual(len(result), 1) + self.assertEqual(tuple(result[0].shape), (2, 8)) + # Token ids index an embedding table, so they must be integral. + self.assertEqual(result[0].dtype, torch.int64) + if __name__ == "__main__": unittest.main() diff --git a/backends/qualcomm/genai_pipeline/tests/strategies/model_preparation/test_executorch_model_preparation_strategy.py b/backends/qualcomm/genai_pipeline/tests/strategies/model_preparation/test_executorch_model_preparation_strategy.py new file mode 100644 index 00000000000..9e02a6f5a7c --- /dev/null +++ b/backends/qualcomm/genai_pipeline/tests/strategies/model_preparation/test_executorch_model_preparation_strategy.py @@ -0,0 +1,374 @@ +# Copyright (c) Qualcomm Innovation Center, Inc. +# All rights reserved +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +import unittest +from pathlib import Path +from unittest.mock import MagicMock, patch + +from executorch.backends.qualcomm.genai_pipeline.configs.model_preparation_input_config import ( + ModelPreparationInputConfig, +) +from executorch.backends.qualcomm.genai_pipeline.configs.model_preparation_output_config import ( + ModelPreparationOutputConfig, +) +from executorch.backends.qualcomm.genai_pipeline.datasets.default_calibration_data_adapter import ( + DEFAULT_NUM_SAMPLES, + DEFAULT_SEQ_LENGTH, +) +from executorch.backends.qualcomm.genai_pipeline.exceptions import StageError +from executorch.backends.qualcomm.genai_pipeline.strategies.model_preparation.executorch_model_preparation_strategy import ( + ExecuTorchModelPreparationStrategy, +) +from executorch.backends.qualcomm.genai_pipeline.strategies.model_preparation.model_preparation_strategy import ( + ModelPreparationStrategy, +) +from executorch.backends.qualcomm.genai_pipeline.tests.test_utils import ( + make_test_context, +) + + +def _make_mock_adapter(): + """Create a mock model loader adapter with sensible defaults.""" + adapter = MagicMock() + adapter.load_model.return_value = MagicMock(name="model_module") + tokenizer = MagicMock(name="tokenizer") + # Default to "no chat template" so tests opt in explicitly. + tokenizer.chat_template = None + adapter.load_tokenizer.return_value = tokenizer + adapter.get_example_inputs.return_value = (MagicMock(name="example_input"),) + adapter.export_tokenizer.return_value = Path("/tmp/tokenizer/tokenizer.json") + return adapter + + +def _make_mock_calibration_adapter(): + """Create a mock calibration data adapter with sensible defaults.""" + adapter = MagicMock() + adapter.generate_calibration_data.return_value = [(MagicMock(),)] + return adapter + + +def _make_strategy(loader=None, calibration=None): + """Build the strategy with both adapters mocked by default.""" + return ExecuTorchModelPreparationStrategy( + model_loader_adapter=loader if loader is not None else _make_mock_adapter(), + calibration_data_adapter=( + calibration if calibration is not None else _make_mock_calibration_adapter() + ), + ) + + +def _make_valid_input_config(**overrides): + """Create a valid ModelPreparationInputConfig with defaults.""" + defaults = { + "model_name": "test_model", + "soc_model": "SM8750", + } + defaults.update(overrides) + return ModelPreparationInputConfig(**defaults) + + +class TestExecuTorchModelPreparationStrategy(unittest.TestCase): + + def test_is_model_preparation_strategy(self): + """Strategy inherits from ModelPreparationStrategy ABC.""" + self.assertIsInstance(_make_strategy(), ModelPreparationStrategy) + + def test_default_adapters_created_when_none_provided(self): + """Both adapters fall back to their default implementations.""" + with patch( + "executorch.backends.qualcomm.genai_pipeline.strategies.model_preparation." + "default_model_loader_adapter.DefaultModelLoaderAdapter" + ) as mock_loader_cls, patch( + "executorch.backends.qualcomm.genai_pipeline.datasets." + "default_calibration_data_adapter.DefaultCalibrationDataAdapter" + ) as mock_calib_cls: + strategy = ExecuTorchModelPreparationStrategy() + + mock_loader_cls.assert_called_once() + mock_calib_cls.assert_called_once() + self.assertIs(strategy.adapter, mock_loader_cls.return_value) + self.assertIs( + strategy.calibration_data_adapter, mock_calib_cls.return_value + ) + + def test_custom_adapters_injected(self): + """Both adapters are used when provided via the constructor.""" + loader = _make_mock_adapter() + calibration = _make_mock_calibration_adapter() + strategy = _make_strategy(loader, calibration) + self.assertIs(strategy.adapter, loader) + self.assertIs(strategy.calibration_data_adapter, calibration) + + def test_invoke_happy_path(self): + """Full model preparation pipeline runs successfully end-to-end.""" + loader = _make_mock_adapter() + calibration = _make_mock_calibration_adapter() + strategy = _make_strategy(loader, calibration) + + result = strategy.invoke(make_test_context(), _make_valid_input_config()) + + self.assertIsInstance(result, ModelPreparationOutputConfig) + self.assertIs(result.model_module, loader.load_model.return_value) + self.assertIs(result.tokenizer, loader.load_tokenizer.return_value) + self.assertEqual( + result.calibration_data, + calibration.generate_calibration_data.return_value, + ) + + def test_invoke_calls_loader_in_correct_order(self): + """The loader is driven in order: load_model → load_tokenizer → example inputs.""" + loader = _make_mock_adapter() + strategy = _make_strategy(loader) + + strategy.invoke(make_test_context(), _make_valid_input_config()) + + self.assertEqual( + [c[0] for c in loader.method_calls], + ["load_model", "load_tokenizer", "get_example_inputs"], + ) + + def test_invoke_example_inputs_derived_from_the_loaded_model(self): + """``example_inputs`` come from the model, not from the calibration data. + + The exported graph's positional signature (zero-initialized KV caches, + fixed AR length) is a property of the model; the calibration dataset is + in fact derived *from* it, so the dependency must not be inverted. + """ + loader = _make_mock_adapter() + calibration = _make_mock_calibration_adapter() + strategy = _make_strategy(loader, calibration) + + result = strategy.invoke(make_test_context(), _make_valid_input_config()) + + loader.get_example_inputs.assert_called_once_with( + model=loader.load_model.return_value, + extra_options=None, + ) + self.assertIs(result.example_inputs, loader.get_example_inputs.return_value) + + def test_invoke_example_input_options_forwarded(self): + """example_input_options from extra_options reach get_example_inputs.""" + loader = _make_mock_adapter() + strategy = _make_strategy(loader) + example_opts = {"ar_len": 128} + input_config = _make_valid_input_config( + extra_options={"example_input_options": example_opts} + ) + + strategy.invoke(make_test_context(), input_config) + + loader.get_example_inputs.assert_called_once_with( + model=loader.load_model.return_value, + extra_options=example_opts, + ) + + def test_invoke_generates_calibration_data_from_dataset_adapter(self): + """Calibration data comes from the dataset adapter, using the loaded tokenizer.""" + loader = _make_mock_adapter() + calibration = _make_mock_calibration_adapter() + strategy = _make_strategy(loader, calibration) + + strategy.invoke(make_test_context(), _make_valid_input_config()) + + calibration.generate_calibration_data.assert_called_once_with( + tokenizer=loader.load_tokenizer.return_value, + num_samples=DEFAULT_NUM_SAMPLES, + seq_length=DEFAULT_SEQ_LENGTH, + extra_options=None, + ) + + def test_invoke_passes_model_name_to_load_model(self): + """load_model receives model_name from input config.""" + loader = _make_mock_adapter() + strategy = _make_strategy(loader) + + strategy.invoke( + make_test_context(), _make_valid_input_config(model_name="llama3_2-1b") + ) + + loader.load_model.assert_called_once_with( + model_name="llama3_2-1b", + extra_options=None, + ) + + def test_invoke_passes_model_name_to_load_tokenizer(self): + """load_tokenizer receives model_name from input config.""" + loader = _make_mock_adapter() + strategy = _make_strategy(loader) + + strategy.invoke( + make_test_context(), _make_valid_input_config(model_name="llama3_2-1b") + ) + + loader.load_tokenizer.assert_called_once_with( + model_name="llama3_2-1b", + extra_options=None, + ) + + def test_invoke_custom_calibration_params_from_extra_options(self): + """Calibration params from extra_options override the defaults.""" + calibration = _make_mock_calibration_adapter() + strategy = _make_strategy(calibration=calibration) + calibration_options = {"dataset": "wikitext"} + input_config = _make_valid_input_config( + extra_options={ + "num_calibration_samples": 64, + "calibration_seq_length": 256, + "calibration_options": calibration_options, + } + ) + + strategy.invoke(make_test_context(), input_config) + + _, kwargs = calibration.generate_calibration_data.call_args + self.assertEqual(kwargs["num_samples"], 64) + self.assertEqual(kwargs["seq_length"], 256) + self.assertEqual(kwargs["extra_options"], calibration_options) + + def test_invoke_no_tokenizer_export_by_default(self): + """export_tokenizer is NOT called when export_tokenizer option is absent.""" + loader = _make_mock_adapter() + strategy = _make_strategy(loader) + + result = strategy.invoke(make_test_context(), _make_valid_input_config()) + + loader.export_tokenizer.assert_not_called() + self.assertIsNone(result.runtime_tokenizer_path) + + def test_invoke_exports_tokenizer_when_requested(self): + """export_tokenizer is called when export_tokenizer=True in extra_options.""" + loader = _make_mock_adapter() + strategy = _make_strategy(loader) + input_config = _make_valid_input_config( + extra_options={"export_tokenizer": True} + ) + context = make_test_context(artifact_dir="/my/artifacts") + + result = strategy.invoke(context, input_config) + + loader.export_tokenizer.assert_called_once_with( + tokenizer=loader.load_tokenizer.return_value, + output_dir=Path("/my/artifacts") / "tokenizer", + extra_options=None, + ) + # The adapter returns the tokenizer file itself, not its directory. + self.assertEqual( + result.runtime_tokenizer_path, Path("/tmp/tokenizer/tokenizer.json") + ) + + def test_invoke_extracts_chat_template_from_tokenizer(self): + """A chat template on the tokenizer is carried into the output config.""" + loader = _make_mock_adapter() + loader.load_tokenizer.return_value.chat_template = "{{ messages }}" + strategy = _make_strategy(loader) + + result = strategy.invoke(make_test_context(), _make_valid_input_config()) + + self.assertEqual(result.chat_template, "{{ messages }}") + + def test_invoke_chat_template_from_extra_options_when_tokenizer_has_none(self): + """extra_options supplies the chat template only as a fallback.""" + strategy = _make_strategy() + input_config = _make_valid_input_config( + extra_options={"chat_template": "fallback"} + ) + + result = strategy.invoke(make_test_context(), input_config) + + self.assertEqual(result.chat_template, "fallback") + + def test_invoke_chat_template_prefers_tokenizer_over_extra_options(self): + """With both present the tokenizer wins: extra_options is only a fallback. + + Pins the precedence itself -- the single-source tests above would still + pass if the two branches were swapped. + """ + loader = _make_mock_adapter() + loader.load_tokenizer.return_value.chat_template = "tokenizer_template" + strategy = _make_strategy(loader) + input_config = _make_valid_input_config( + extra_options={"chat_template": "fallback"} + ) + + result = strategy.invoke(make_test_context(), input_config) + + self.assertEqual(result.chat_template, "tokenizer_template") + + def test_invoke_missing_model_name_raises_stage_error(self): + """StageError raised when model_name is empty.""" + strategy = _make_strategy() + + with self.assertRaises(StageError) as cm: + strategy.invoke( + make_test_context(), _make_valid_input_config(model_name="") + ) + self.assertIn("model_name", str(cm.exception)) + self.assertEqual(cm.exception.stage_name, "model_preparation") + + def test_invoke_missing_soc_model_raises_stage_error(self): + """StageError raised when soc_model is empty.""" + strategy = _make_strategy() + + with self.assertRaises(StageError) as cm: + strategy.invoke(make_test_context(), _make_valid_input_config(soc_model="")) + self.assertIn("soc_model", str(cm.exception)) + self.assertEqual(cm.exception.stage_name, "model_preparation") + + def test_invoke_adapter_exception_wrapped_in_stage_error(self): + """Exceptions from the adapter are wrapped in StageError.""" + loader = _make_mock_adapter() + loader.load_model.side_effect = RuntimeError("model load failed") + strategy = _make_strategy(loader) + + with self.assertRaises(StageError) as cm: + strategy.invoke(make_test_context(), _make_valid_input_config()) + self.assertEqual(cm.exception.stage_name, "model_preparation") + self.assertIsInstance(cm.exception.original_exception, RuntimeError) + self.assertIn("model load failed", str(cm.exception)) + + def test_invoke_calibration_adapter_exception_wrapped_in_stage_error(self): + """Failures in the dataset adapter surface as a model_preparation StageError.""" + calibration = _make_mock_calibration_adapter() + calibration.generate_calibration_data.side_effect = ValueError("no dataset") + strategy = _make_strategy(calibration=calibration) + + with self.assertRaises(StageError) as cm: + strategy.invoke(make_test_context(), _make_valid_input_config()) + self.assertEqual(cm.exception.stage_name, "model_preparation") + self.assertIsInstance(cm.exception.original_exception, ValueError) + + def test_invoke_stage_error_not_double_wrapped(self): + """StageError from adapter is re-raised directly.""" + loader = _make_mock_adapter() + original_error = StageError( + stage_name="model_preparation", message="inner error" + ) + loader.load_model.side_effect = original_error + strategy = _make_strategy(loader) + + with self.assertRaises(StageError) as cm: + strategy.invoke(make_test_context(), _make_valid_input_config()) + self.assertIs(cm.exception, original_error) + + def test_invoke_model_options_forwarded(self): + """model_options from extra_options are forwarded to load_model.""" + loader = _make_mock_adapter() + strategy = _make_strategy(loader) + model_opts = {"torch_dtype": "float16"} + input_config = _make_valid_input_config( + extra_options={"model_options": model_opts} + ) + + strategy.invoke(make_test_context(), input_config) + + loader.load_model.assert_called_once_with( + model_name="test_model", + extra_options=model_opts, + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/backends/qualcomm/genai_pipeline/tests/strategies/quantization/test_executorch_quantization_strategy.py b/backends/qualcomm/genai_pipeline/tests/strategies/quantization/test_executorch_quantization_strategy.py index 5740e46771e..f6fb4281242 100644 --- a/backends/qualcomm/genai_pipeline/tests/strategies/quantization/test_executorch_quantization_strategy.py +++ b/backends/qualcomm/genai_pipeline/tests/strategies/quantization/test_executorch_quantization_strategy.py @@ -5,37 +5,328 @@ # LICENSE file in the root directory of this source tree. import unittest -from unittest.mock import MagicMock +from unittest.mock import create_autospec, MagicMock, patch from executorch.backends.qualcomm.genai_pipeline.configs.quantization_input_config import ( QuantizationInputConfig, ) +from executorch.backends.qualcomm.genai_pipeline.configs.quantization_output_config import ( + QuantizationOutputConfig, +) +from executorch.backends.qualcomm.genai_pipeline.exceptions import StageError from executorch.backends.qualcomm.genai_pipeline.strategies.quantization.executorch_quantization_strategy import ( ExecuTorchQuantizationStrategy, ) from executorch.backends.qualcomm.genai_pipeline.strategies.quantization.quantization_strategy import ( QuantizationStrategy, ) +from executorch.backends.qualcomm.genai_pipeline.strategies.quantization.quantizer_adapter import ( + QuantizerAdapter, +) from executorch.backends.qualcomm.genai_pipeline.tests.test_utils import ( make_test_context, ) +def _make_mock_adapter(): + """Create a mock adapter with all methods returning sensible defaults. + + Autospec'd against ``QuantizerAdapter`` rather than a bare ``MagicMock``: a + bare mock accepts any call whatsoever, so a strategy call that no real + adapter could satisfy -- omitting an argument the Protocol declares as + required -- would pass here and only fail in production. The autospec binds + every call to the Protocol signature, keeping Protocol, adapter and strategy + in step. + + Its limit is worth stating: ``make_quantizer`` takes ``**kwargs`` by design, + to forward ``extra_options`` verbatim, so no signature check can reject a + keyword the *underlying* API does not accept. That the forwarded keywords + are ones ``export_utils.make_quantizer`` actually takes is only observable + against the real adapter, and belongs in the integration tests. + """ + adapter = create_autospec(QuantizerAdapter, instance=True) + adapter.export_model.return_value = MagicMock(name="exported_model") + adapter.make_quantizer.return_value = MagicMock(name="quantizer") + adapter.prepare_pt2e.return_value = MagicMock(name="annotated_model") + adapter.calibrate.return_value = MagicMock(name="calibrated_model") + adapter.convert_pt2e.return_value = MagicMock(name="quantized_model") + return adapter + + +def _make_valid_input_config(**overrides): + """Create a valid QuantizationInputConfig with defaults.""" + defaults = { + "soc_model": MagicMock(name="SM8750"), + "backend_type": MagicMock(name="kHtpBackend"), + "model_module": MagicMock(name="test_model"), + "example_inputs": (MagicMock(name="example_input"),), + "calibration_data": [(MagicMock(),)], + } + defaults.update(overrides) + return QuantizationInputConfig(**defaults) + + class TestExecuTorchQuantizationStrategy(unittest.TestCase): def test_is_quantization_strategy(self): - strategy = ExecuTorchQuantizationStrategy() + """Strategy inherits from QuantizationStrategy ABC.""" + adapter = _make_mock_adapter() + strategy = ExecuTorchQuantizationStrategy(quantizer_adapter=adapter) self.assertIsInstance(strategy, QuantizationStrategy) - def test_invoke_raises_not_implemented(self): - strategy = ExecuTorchQuantizationStrategy() - with self.assertRaises(NotImplementedError): - strategy.invoke( - make_test_context(), - QuantizationInputConfig( - soc_model=MagicMock(), backend_type=MagicMock() - ), - ) + def test_default_adapter_created_when_none_provided(self): + """When no adapter is provided, DefaultQuantizerAdapter is created.""" + with patch( + "executorch.backends.qualcomm.genai_pipeline.strategies.quantization." + "default_quantizer_adapter.DefaultQuantizerAdapter" + ) as mock_cls: + strategy = ExecuTorchQuantizationStrategy(quantizer_adapter=None) + mock_cls.assert_called_once() + self.assertIs(strategy.adapter, mock_cls.return_value) + + def test_custom_adapter_injected(self): + """Custom adapter is used when provided via constructor.""" + adapter = _make_mock_adapter() + strategy = ExecuTorchQuantizationStrategy(quantizer_adapter=adapter) + self.assertIs(strategy.adapter, adapter) + + def test_invoke_happy_path(self): + """Full quantization pipeline runs successfully end-to-end.""" + adapter = _make_mock_adapter() + strategy = ExecuTorchQuantizationStrategy(quantizer_adapter=adapter) + + result = strategy.invoke(make_test_context(), _make_valid_input_config()) + + self.assertIsInstance(result, QuantizationOutputConfig) + self.assertIs(result.quantized_model, adapter.convert_pt2e.return_value) + + def test_invoke_calls_adapter_in_correct_order(self): + """Adapter methods are called in the correct PT2E sequence.""" + adapter = _make_mock_adapter() + strategy = ExecuTorchQuantizationStrategy(quantizer_adapter=adapter) + + strategy.invoke(make_test_context(), _make_valid_input_config()) + + # Verify call order: export → make_quantizer → prepare → calibrate → convert + self.assertEqual( + [c[0] for c in adapter.method_calls], + [ + "export_model", + "make_quantizer", + "prepare_pt2e", + "calibrate", + "convert_pt2e", + ], + ) + + def test_invoke_exports_with_example_inputs_not_calibration_data(self): + """export_model receives ``example_inputs``, never a calibration sample. + + The export signature comes from the model (zero-initialized KV caches, + fixed AR length); a calibration sample has neither, so sourcing it from + the dataset would export the wrong graph. + """ + adapter = _make_mock_adapter() + strategy = ExecuTorchQuantizationStrategy(quantizer_adapter=adapter) + model = MagicMock(name="model") + example_inputs = (MagicMock(name="tokens"), MagicMock(name="attn_mask")) + calibration_sample = (MagicMock(name="calibration_sample"),) + input_config = _make_valid_input_config( + model_module=model, + example_inputs=example_inputs, + calibration_data=[calibration_sample], + ) + + strategy.invoke(make_test_context(), input_config) + + adapter.export_model.assert_called_once_with(model, example_inputs) + + def test_invoke_passes_calibration_data_through_untouched(self): + """``calibration_data`` reaches ``calibrate`` as the very same object. + + Nothing peeks at, indexes or copies it, so a single-use generator keeps + every sample and a DataLoader keeps streaming instead of being pulled + into memory. + """ + adapter = _make_mock_adapter() + strategy = ExecuTorchQuantizationStrategy(quantizer_adapter=adapter) + samples = [(MagicMock(name=f"sample{i}"),) for i in range(3)] + dataset = (sample for sample in samples) + input_config = _make_valid_input_config(calibration_data=dataset) + + strategy.invoke(make_test_context(), input_config) + + self.assertIs(adapter.calibrate.call_args[0][1], dataset) + # Untouched by the strategy, so all three samples are still available. + self.assertEqual(list(adapter.calibrate.call_args[0][1]), samples) + + def test_invoke_missing_example_inputs_raises_stage_error(self): + """StageError raised when example_inputs is None.""" + adapter = _make_mock_adapter() + strategy = ExecuTorchQuantizationStrategy(quantizer_adapter=adapter) + input_config = _make_valid_input_config(example_inputs=None) + + with self.assertRaises(StageError) as cm: + strategy.invoke(make_test_context(), input_config) + self.assertIn("example_inputs", str(cm.exception)) + self.assertEqual(cm.exception.stage_name, "quantization") + adapter.export_model.assert_not_called() + + def test_invoke_passes_correct_args_to_make_quantizer(self): + """make_quantizer receives backend_type and soc_model (no quant_dtype when not set).""" + adapter = _make_mock_adapter() + strategy = ExecuTorchQuantizationStrategy(quantizer_adapter=adapter) + soc = MagicMock(name="soc") + backend = MagicMock(name="backend") + input_config = _make_valid_input_config(soc_model=soc, backend_type=backend) + + strategy.invoke(make_test_context(), input_config) + + # quant_dtype is NOT passed when not explicitly set in extra_options, so + # the default owned by export_utils.make_quantizer applies. + adapter.make_quantizer.assert_called_once_with( + backend=backend, + soc_model=soc, + ) + + def test_invoke_passes_quant_dtype_from_extra_options(self): + """quant_dtype extracted from extra_options and forwarded.""" + adapter = _make_mock_adapter() + strategy = ExecuTorchQuantizationStrategy(quantizer_adapter=adapter) + quant_dtype = MagicMock(name="quant_dtype") + input_config = _make_valid_input_config( + extra_options={"quant_dtype": quant_dtype} + ) + + strategy.invoke(make_test_context(), input_config) + + adapter.make_quantizer.assert_called_once_with( + quant_dtype=quant_dtype, + backend=input_config.backend_type, + soc_model=input_config.soc_model, + ) + + def test_invoke_passes_quant_recipe_from_config(self): + """quant_recipe on the config reaches the adapter as a declared kwarg. + + The adapter consumes it (``QnnQuantizer.set_recipe``) rather than + forwarding it to ``export_utils.make_quantizer``, which takes no such + argument. + """ + adapter = _make_mock_adapter() + strategy = ExecuTorchQuantizationStrategy(quantizer_adapter=adapter) + recipe = MagicMock(name="quant_recipe") + input_config = _make_valid_input_config(quant_recipe=recipe) + + strategy.invoke(make_test_context(), input_config) + + _, kwargs = adapter.make_quantizer.call_args + self.assertIs(kwargs["quant_recipe"], recipe) + + def test_invoke_training_data_does_not_switch_to_qat(self): + """``training_data`` is accepted but this strategy still performs PTQ.""" + adapter = _make_mock_adapter() + strategy = ExecuTorchQuantizationStrategy(quantizer_adapter=adapter) + input_config = _make_valid_input_config(training_data=[(MagicMock(),)]) + + with self.assertLogs( + "executorch.backends.qualcomm.genai_pipeline.strategies.quantization." + "executorch_quantization_strategy", + level="WARNING", + ): + result = strategy.invoke(make_test_context(), input_config) + + # Still the plain PTQ sequence. + adapter.calibrate.assert_called_once() + self.assertIs(result.quantized_model, adapter.convert_pt2e.return_value) + + def test_invoke_missing_model_raises_stage_error(self): + """StageError raised when model_module is None.""" + adapter = _make_mock_adapter() + strategy = ExecuTorchQuantizationStrategy(quantizer_adapter=adapter) + input_config = _make_valid_input_config(model_module=None) + + with self.assertRaises(StageError) as cm: + strategy.invoke(make_test_context(), input_config) + self.assertIn("model_module", str(cm.exception)) + self.assertEqual(cm.exception.stage_name, "quantization") + + def test_invoke_none_calibration_data_raises_stage_error(self): + """StageError raised when calibration_data is None.""" + adapter = _make_mock_adapter() + strategy = ExecuTorchQuantizationStrategy(quantizer_adapter=adapter) + input_config = _make_valid_input_config(calibration_data=None) + + with self.assertRaises(StageError) as cm: + strategy.invoke(make_test_context(), input_config) + self.assertIn("calibration_data", str(cm.exception)) + + def test_invoke_adapter_exception_wrapped_in_stage_error(self): + """Exceptions from the adapter are wrapped in StageError.""" + adapter = _make_mock_adapter() + adapter.export_model.side_effect = RuntimeError("export failed") + strategy = ExecuTorchQuantizationStrategy(quantizer_adapter=adapter) + + with self.assertRaises(StageError) as cm: + strategy.invoke(make_test_context(), _make_valid_input_config()) + self.assertEqual(cm.exception.stage_name, "quantization") + self.assertIsInstance(cm.exception.original_exception, RuntimeError) + self.assertIn("export failed", str(cm.exception)) + + def test_invoke_stage_error_not_double_wrapped(self): + """StageError from adapter is re-raised directly, not wrapped again.""" + adapter = _make_mock_adapter() + original_error = StageError(stage_name="quantization", message="inner error") + adapter.export_model.side_effect = original_error + strategy = ExecuTorchQuantizationStrategy(quantizer_adapter=adapter) + + with self.assertRaises(StageError) as cm: + strategy.invoke(make_test_context(), _make_valid_input_config()) + self.assertIs(cm.exception, original_error) + + def test_invoke_missing_soc_model_raises_stage_error(self): + """StageError raised when soc_model is None.""" + adapter = _make_mock_adapter() + strategy = ExecuTorchQuantizationStrategy(quantizer_adapter=adapter) + input_config = _make_valid_input_config(soc_model=None) + + with self.assertRaises(StageError) as cm: + strategy.invoke(make_test_context(), input_config) + self.assertIn("soc_model", str(cm.exception)) + self.assertEqual(cm.exception.stage_name, "quantization") + + def test_invoke_missing_backend_type_raises_stage_error(self): + """StageError raised when backend_type is None.""" + adapter = _make_mock_adapter() + strategy = ExecuTorchQuantizationStrategy(quantizer_adapter=adapter) + input_config = _make_valid_input_config(backend_type=None) + + with self.assertRaises(StageError) as cm: + strategy.invoke(make_test_context(), input_config) + self.assertIn("backend_type", str(cm.exception)) + self.assertEqual(cm.exception.stage_name, "quantization") + + def test_invoke_extra_options_forwarded_to_make_quantizer(self): + """Extra options (minus quant_dtype) forwarded as kwargs.""" + adapter = _make_mock_adapter() + strategy = ExecuTorchQuantizationStrategy(quantizer_adapter=adapter) + input_config = _make_valid_input_config( + extra_options={ + "quant_dtype": "test_dtype", + "per_channel_conv": True, + "act_symmetric": True, + } + ) + + strategy.invoke(make_test_context(), input_config) + + adapter.make_quantizer.assert_called_once_with( + quant_dtype="test_dtype", + backend=input_config.backend_type, + soc_model=input_config.soc_model, + per_channel_conv=True, + act_symmetric=True, + ) if __name__ == "__main__":