Qualcomm AI Engine Direct - [GenAI Pipeline] Multimodal LLM & LLM Preparation and Quantization

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, Tuple, TYPE_CHECKING
+from typing import Any, Dict, Optional, Tuple, TYPE_CHECKING
 
 if TYPE_CHECKING:
     from torch import nn
@@ -24,25 +24,37 @@ class ModelPreparationOutputConfig:
     not a valid post-execution state: when the stage runs, ``model_module`` and
     ``tokenizer`` are always populated.
 
+    ``model_module`` is keyed by component. ``example_inputs`` and ``meta`` add
+    the inner graph axis, because their values vary by graph.
+
     Attributes:
-        model_module: The prepared nn.Module ready for quantization.
+        model_module: The prepared modules, keyed by component.
         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.
+        example_inputs: Per-component, per-graph positional example inputs for
+            ``torch.export`` (``{component: {graph: tuple}}``), derived from the
+            **model** (never from the calibration data): they carry each 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 data's
+            attention-mask schema is itself derived from these tuples.
         runtime_tokenizer_path: Path to the runtime tokenizer **file** (not the
             containing directory) for on-device inference.
         chat_template: Optional chat template for instruct models.
+        meta: Per-graph ``get_metadata()`` constants (layer count, head dim,
+            context/AR lengths). Feeds logits / KV-cache shape reconstruction in
+            quantization and is baked into the ``.pte``.
+        inference: The ``ModelInference`` bound to the calibration graph, built
+            by the loader adapter. Drives PTQ calibration in the quantization
+            stage; ``None`` for flows that build no inference (e.g. multimodal,
+            not yet supported).
+        num_shardings: Optional per-component number of shardings for the model.
     """
 
-    model_module: Optional["nn.Module"] = None
+    model_module: Optional[Dict[str, "nn.Module"]] = None
     tokenizer: Any = None
-    example_inputs: Optional[Tuple[Any, ...]] = None
-    calibration_data: Optional[Iterable[Any]] = None
+    example_inputs: Optional[Dict[str, Dict[str, Tuple[Any, ...]]]] = None
     runtime_tokenizer_path: Optional[Path] = None
     chat_template: Optional[str] = None
+    meta: Optional[Dict[str, Dict[str, Any]]] = None
+    inference: Any = None
+    num_shardings: Optional[Dict[str, int]] = None

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, Tuple, TYPE_CHECKING
+from typing import Any, Dict, Optional, Tuple, TYPE_CHECKING
 
 if TYPE_CHECKING:
     from executorch.backends.qualcomm.serialization.qc_schema import (
@@ -21,45 +21,53 @@
 class QuantizationInputConfig:
     """Input configuration for the quantization stage.
 
-    ``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.
+    ``model_module`` and ``example_inputs`` 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.
 
     Flows needing no quantization (FP16, GPU backends) skip the stage entirely
     via ``GenAIPipeline.from_proxy(proxy, skip_stages={STAGE_QUANTIZATION})``
     rather than entering it with a no-op strategy; the orchestrator then returns
     an empty ``QuantizationOutputConfig()`` and compilation receives the
     unquantized module.
 
+    ``model_module`` is keyed by component. ``example_inputs`` and ``meta`` add
+    the inner graph axis, because their values vary by graph.
+
     Attributes:
         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. Consumed only
-            by ``calibrate()`` -- it is never indexed or peeked at, so a
-            single-use generator stays intact.
+        backend_type: One QNN backend type (HTP, GPU, LPAI, etc.) shared by all
+            components. Per-component backend routing is unsupported.
+        model_module: Prepared modules, keyed by component. Required when the
+            stage runs.
+        example_inputs: Positional example inputs for ``torch.export``, per
+            component and graph. Required when the stage runs. Sourced from the **model** via
+            ``ModelLoaderAdapter.get_example_inputs``, never from a dataset
+            sample: 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.
+        tokenizer: The TokenizerWrapper from model preparation, used by the
+            strategy to build calibration data.
         training_data: Training dataset for quantization-aware training (QAT),
             typically (features, labels) pairs. Mirrors ``qat_training_data`` in
             ``build_executorch_binary``. ``None`` selects PTQ.
-        quant_recipe: Quantization recipe (per-layer bit widths, group sizes, etc.).
+        quant_recipe: Quantization recipe, or a per-component map of them.
+        meta: Per-graph constant metadata from model preparation.
+        inference: Optional model-specific inference instance or callable used
+            by PTQ calibration. Created during model preparation; ``None`` when
+            the model does not provide one.
         extra_options: Additional quantization-specific options.
     """
 
     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
+    model_module: Optional[Dict[str, "nn.Module"]] = None
+    example_inputs: Optional[Dict[str, Dict[str, Tuple[Any, ...]]]] = None
+    tokenizer: Any = None
     quant_recipe: Any = None
+    meta: Optional[Dict[str, Any]] = None
+    inference: Optional[Any] = None
     extra_options: Dict[str, Any] = field(default_factory=dict)

backends/qualcomm/genai_pipeline/configs/quantization_output_config.py

@@ -5,15 +5,20 @@
 # LICENSE file in the root directory of this source tree.
 
 from dataclasses import dataclass
-from typing import Any
+from typing import Dict
+
+from executorch.backends.qualcomm.genai_pipeline.graph_bundle import GraphBundle
 
 
 @dataclass
 class QuantizationOutputConfig:
     """Output produced by the quantization stage.
 
     Attributes:
-        quantized_model: The quantized nn.Module or path to saved QDQ model.
+        graphs: Component- and graph-keyed quantized modules with their export
+            inputs, metadata, and optional quantized IO dtypes. The compilation
+            stage consumes the bundles and bakes metadata-derived KV-cache and
+            logits quantization attributes into the ``.pte``.
     """
 
-    quantized_model: Any = None
+    graphs: Dict[str, Dict[str, GraphBundle]] = None

backends/qualcomm/genai_pipeline/datasets/__init__.py

@@ -11,24 +11,26 @@
 evaluation during inference (including pre-built ``.pte`` flows, where no model
 preparation runs at all). They therefore live here rather than under
 ``strategies/model_preparation/``.
+
+Dataset loaders are organized by modality under ``loaders/`` (``loaders/llm/``,
+``loaders/mllm/``) and are purpose-agnostic: each just loads a component-keyed
+dataset dict. The purpose lives in the per-purpose adapter packages
+(``calibration/``, ``training/``, ``evaluation/``), which compose a set of
+loaders with a component-aware collector.
 """
 
-from executorch.backends.qualcomm.genai_pipeline.datasets.calibration_data_adapter import (
-    CalibrationDataAdapter,
-)
-from executorch.backends.qualcomm.genai_pipeline.datasets.default_calibration_data_adapter import (
-    DefaultCalibrationDataAdapter,
-)
-from executorch.backends.qualcomm.genai_pipeline.datasets.default_training_data_adapter import (
-    DefaultTrainingDataAdapter,
-)
-from executorch.backends.qualcomm.genai_pipeline.datasets.training_data_adapter import (
-    TrainingDataAdapter,
+from executorch.backends.qualcomm.genai_pipeline.datasets.dataset_lookup import (
+    get_calibration_dataset_adapter,
+    get_collector,
+    get_dataset_adapter,
+    get_eval_dataset_adapter,
+    get_training_dataset_adapter,
 )
 
 __all__ = [
-    "CalibrationDataAdapter",
-    "DefaultCalibrationDataAdapter",
-    "DefaultTrainingDataAdapter",
-    "TrainingDataAdapter",
+    "get_dataset_adapter",
+    "get_collector",
+    "get_calibration_dataset_adapter",
+    "get_training_dataset_adapter",
+    "get_eval_dataset_adapter",
 ]

backends/qualcomm/genai_pipeline/datasets/calibration/__init__.py

@@ -0,0 +1,27 @@
+# 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.
+
+"""Calibration purpose adapters (protocol + per-modality implementations)."""
+
+from executorch.backends.qualcomm.genai_pipeline.datasets.calibration.calibration_data_adapter import (
+    CalibrationDataAdapter,
+)
+from executorch.backends.qualcomm.genai_pipeline.datasets.calibration.default_calibration_data_adapter import (
+    DefaultCalibrationDataAdapter,
+)
+from executorch.backends.qualcomm.genai_pipeline.datasets.calibration.llm_calibration_data_adapter import (
+    LLMCalibrationDataAdapter,
+)
+from executorch.backends.qualcomm.genai_pipeline.datasets.calibration.mllm_calibration_data_adapter import (
+    MLLMCalibrationDataAdapter,
+)
+
+__all__ = [
+    "CalibrationDataAdapter",
+    "DefaultCalibrationDataAdapter",
+    "LLMCalibrationDataAdapter",
+    "MLLMCalibrationDataAdapter",
+]

backends/qualcomm/genai_pipeline/datasets/calibration/calibration_data_adapter.py

@@ -0,0 +1,46 @@
+# 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.
+
+"""Calibration data adapter protocol.
+
+A purpose adapter assembles calibration data from raw dataset loaders. Its output
+is always component-keyed so quantization strategies can consume one contract for
+LLM and MLLM models.
+"""
+
+from __future__ import annotations
+
+from typing import Any, Dict, Iterable, Optional, Protocol, runtime_checkable
+
+
+@runtime_checkable
+class CalibrationDataAdapter(Protocol):
+    """Protocol for assembling calibration data for quantization."""
+
+    def generate_calibration_data(
+        self,
+        tokenizer: Any,
+        example_inputs: Optional[Dict[str, Any]] = None,
+        extra_options: Optional[Dict[str, Any]] = None,
+    ) -> Dict[str, Iterable[Any]]:
+        """Return ``{component: iterable}`` calibration inputs.
+
+        Args:
+            tokenizer: Tokenizer used by dataset loaders.
+            example_inputs: Model calibration-graph signatures used by adapters
+                that construct collators. Adapters without collators may ignore it.
+            extra_options: Optional data-source settings. ``max_context_len`` is
+                the shared sequence-length source of truth for dataset loaders
+                and collators.
+
+        Returns:
+            A map from artifact component keys to calibration iterables.
+
+        Note:
+            Sizing such as sample count, batch size, and seed is fixed when the
+            adapter is constructed rather than passed to this method.
+        """
+        ...

backends/qualcomm/genai_pipeline/datasets/calibration/default_calibration_data_adapter.py

@@ -0,0 +1,113 @@
+# 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.
+
+from __future__ import annotations
+
+import logging
+from typing import Any, Dict, Iterable, List, Optional
+
+logger = logging.getLogger(__name__)
+
+DEFAULT_NUM_SAMPLES = 128
+DEFAULT_SEQ_LENGTH = 1024
+DEFAULT_BATCH_SIZE = 1
+DEFAULT_SEED = 42
+
+
+class DefaultCalibrationDataAdapter:
+    """Pipeline sanity-check calibration provider.
+
+    Returns caller-supplied data when present; otherwise random token
+    sequences. This adapter deliberately does **not** wrap datasets in
+    DataLoaders or run them through a collator -- it exists to confirm the
+    quantization pipeline is wired up, not to produce corpus-backed calibration.
+    For production quantization accuracy, use the modality-specific
+    ``LLMCalibrationDataAdapter`` / ``MLLMCalibrationDataAdapter`` (selected by
+    ``get_calibration_dataset_adapter`` when dataset sources are configured) or
+    pass ready-made data via ``extra_options["dataset"]``.
+
+    The returned object is always ``{component: iterable}``.
+
+    Args:
+        num_samples: How many random samples to generate.
+        batch_size: Leading batch dimension of each random sample.
+        seed: Manual seed for reproducible generation.
+    """
+
+    def __init__(
+        self,
+        num_samples: int = DEFAULT_NUM_SAMPLES,
+        batch_size: int = DEFAULT_BATCH_SIZE,
+        seed: int = DEFAULT_SEED,
+    ) -> None:
+        self._num_samples = num_samples
+        self._batch_size = batch_size
+        self._seed = seed
+
+    def generate_calibration_data(
+        self,
+        tokenizer: Any,
+        example_inputs: Optional[Dict[str, Any]] = None,
+        extra_options: Optional[Dict[str, Any]] = None,
+    ) -> Dict[str, Iterable[Any]]:
+        """Return caller-supplied data or random fallback data.
+
+        Args:
+            tokenizer: Tokenizer whose vocabulary size bounds random token IDs.
+            example_inputs: Ignored because this adapter does not use collators.
+            extra_options: Optional calibration settings. ``dataset`` supplies
+                ready-made calibration data; ``max_context_len`` and ``seed``
+                configure random fallback data.
+
+        Returns:
+            A map containing caller-supplied data or random decoder inputs.
+
+        Note:
+            Random samples are ``(input_ids, attention_mask)`` tuples keyed by
+            ``ARTIFACT_TEXT_DECODER`` and can be consumed directly without a
+            collator. ``DEFAULT_SEQ_LENGTH`` is used when no sequence length is
+            supplied.
+        """
+        import torch
+
+        from executorch.backends.qualcomm.genai_pipeline.artifact_keys import (
+            ARTIFACT_TEXT_DECODER,
+        )
+
+        extra_options = extra_options or {}
+
+        if "dataset" in extra_options:
+            logger.info("Using caller-supplied dataset for calibration")
+            data = extra_options["dataset"]
+            return data if isinstance(data, dict) else {ARTIFACT_TEXT_DECODER: data}
+
+        seq_length = extra_options.get("max_context_len", DEFAULT_SEQ_LENGTH)
+        seed = extra_options.get("seed", self._seed)
+        torch.manual_seed(seed)
+
+        logger.info(
+            "Generating %d random calibration samples (seq_length=%d, batch_size=%d, seed=%d)",
+            self._num_samples,
+            seq_length,
+            self._batch_size,
+            seed,
+        )
+
+        vocab_size = getattr(tokenizer, "vocab_size", None)
+        if vocab_size is None or vocab_size <= 0:
+            raise ValueError(
+                "Tokenizer does not have a valid vocab_size attribute. "
+                "Cannot generate random calibration data. Supply a dataset "
+                "via extra_options['dataset'] instead."
+            )
+
+        calibration_data: List[Any] = []
+        for _ in range(self._num_samples):
+            input_ids = torch.randint(0, vocab_size, (self._batch_size, seq_length))
+            attention_mask = torch.ones(self._batch_size, seq_length, dtype=torch.long)
+            calibration_data.append((input_ids, attention_mask))
+
+        return {ARTIFACT_TEXT_DECODER: calibration_data}