From 7fdb147190528eff7667dc0d9349617178439fba Mon Sep 17 00:00:00 2001 From: yuyazhua Date: Mon, 17 Aug 2026 17:16:57 +0800 Subject: [PATCH 1/3] Qualcomm AI Engine Direct - [GenAI Pipeline] LLM/MLLM Model Prep and Quantization Add the LLM/MLLM GenAI pipeline integration for model preparation, dataset-driven calibration, and ExecuTorch quantization. Summary: - Add the GenAI pipeline CLI and stage context wiring. - Add model registry lookup helpers for configs, graph builders, source transforms, checkpoint loaders, quantization settings, and adapters. - Add component-aware LLM/MLLM model preparation adapters for decoder, embedding, vision, and audio modules. - Add dataset adapters, collators, and dataset option for calibration, training, and evaluation. - Add ExecuTorch quantization support for export/prepare, encoding initialization, calibration, encoding override, conversion, and QDQ EP save. - Add source transforms for checkpoint remapping, dtype override, embedding scaling, RoPE layout, RMSNorm offset, and linear-to-conv2d conversion. - Expand unit coverage for model preparation, quantization, source transforms, datasets, and pipeline stage behavior. --- backends/qualcomm/genai_pipeline/cli.py | 754 ++++++++++++++++ .../model_preparation_output_config.py | 38 +- .../configs/quantization_input_config.py | 56 +- .../configs/quantization_output_config.py | 11 +- .../genai_pipeline/datasets/__init__.py | 32 +- .../datasets/calibration/__init__.py | 27 + .../calibration/calibration_data_adapter.py | 46 + .../default_calibration_data_adapter.py | 113 +++ .../llm_calibration_data_adapter.py | 122 +++ .../mllm_calibration_data_adapter.py | 128 +++ .../datasets/calibration_data_adapter.py | 54 -- .../datasets/collators/__init__.py | 24 + .../datasets/collators/llm_collator.py | 64 ++ .../datasets/collators/mllm_collator.py | 71 ++ .../genai_pipeline/datasets/dataset_lookup.py | 239 +++++ .../datasets/dataset_options.py | 115 +++ .../default_calibration_data_adapter.py | 89 -- .../datasets/default_training_data_adapter.py | 65 -- .../datasets/evaluation/__init__.py | 19 + .../default_evaluation_data_adapter.py | 112 +++ .../evaluation/evaluation_data_adapter.py | 31 + .../datasets/loaders/__init__.py | 13 + .../datasets/loaders/llm/__init__.py | 23 + .../loaders/llm/hf_dataset_adapter.py | 67 ++ .../datasets/loaders/llm/lm_eval_adapter.py | 73 ++ .../loaders/llm/message_sample_adapter.py | 66 ++ .../datasets/loaders/mllm/__init__.py | 15 + .../loaders/mllm/message_sample_adapter.py | 104 +++ .../datasets/training/__init__.py | 19 + .../training/default_training_data_adapter.py | 114 +++ .../training/training_data_adapter.py | 31 + .../datasets/training_data_adapter.py | 44 - .../qualcomm/genai_pipeline/genai_pipeline.py | 25 +- .../model_components/__init__.py | 32 + .../model_components/decoder/__init__.py | 86 ++ .../model_components/embedding.py | 25 + .../model_components/encoders/__init__.py | 38 + .../qualcomm/genai_pipeline/model_lookup.py | 675 ++++++++++++++ .../genai_pipeline/models/__init__.py | 34 + .../genai_pipeline/quant_utilities.py | 355 ++++++++ .../genai_pipeline/recipes/__init__.py | 65 ++ .../source_transform/__init__.py | 81 ++ .../source_transform/checkpoint_key_remap.py | 41 + .../source_transform/dtype_override.py | 21 + .../source_transform/embedding_scale.py | 27 + .../source_transform/linear_to_conv2d.py | 43 + .../source_transform/rms_norm_offset.py | 27 + .../source_transform/rope_layout.py | 48 + .../default_model_loader_adapter.py | 70 +- .../executorch_model_preparation_strategy.py | 318 +++++-- .../model_preparation/llm_loader_adapter.py | 379 ++++++++ .../model_preparation/mllm_loader_adapter.py | 520 +++++++++++ .../model_preparation/model_loader_adapter.py | 114 ++- .../quantization/default_quantizer_adapter.py | 120 +-- .../executorch_quantization_strategy.py | 728 +++++++++++++--- .../quantization/llm_quantizer_adapter.py | 148 ++++ .../quantization/mllm_quantizer_adapter.py | 156 ++++ .../quantization/quantizer_adapter.py | 124 +-- .../configs/test_quantization_input_config.py | 14 +- .../test_default_calibration_data_adapter.py | 37 +- .../test_default_training_data_adapter.py | 9 +- .../test_checkpoint_key_remap.py | 56 ++ .../source_transform/test_dtype_override.py | 38 + .../source_transform/test_embedding_scale.py | 51 ++ .../source_transform/test_linear_to_conv2d.py | 134 +++ .../source_transform/test_rms_norm_offset.py | 47 + .../source_transform/test_rope_layout.py | 110 +++ .../tests/stages/test_quantization_stage.py | 4 +- .../test_default_model_loader_adapter.py | 27 +- ...t_executorch_model_preparation_strategy.py | 390 +++++---- .../test_llm_loader_adapter.py | 278 ++++++ .../test_mllm_loader_adapter.py | 372 ++++++++ .../test_default_quantizer_adapter.py | 51 +- .../test_executorch_quantization_strategy.py | 823 ++++++++++++++---- .../test_llm_quantizer_adapter.py | 121 +++ .../test_mllm_quantizer_adapter.py | 187 ++++ .../tests/test_genai_pipeline.py | 67 +- 77 files changed, 8862 insertions(+), 1033 deletions(-) create mode 100644 backends/qualcomm/genai_pipeline/cli.py create mode 100644 backends/qualcomm/genai_pipeline/datasets/calibration/__init__.py create mode 100644 backends/qualcomm/genai_pipeline/datasets/calibration/calibration_data_adapter.py create mode 100644 backends/qualcomm/genai_pipeline/datasets/calibration/default_calibration_data_adapter.py create mode 100644 backends/qualcomm/genai_pipeline/datasets/calibration/llm_calibration_data_adapter.py create mode 100644 backends/qualcomm/genai_pipeline/datasets/calibration/mllm_calibration_data_adapter.py delete mode 100644 backends/qualcomm/genai_pipeline/datasets/calibration_data_adapter.py create mode 100644 backends/qualcomm/genai_pipeline/datasets/collators/__init__.py create mode 100644 backends/qualcomm/genai_pipeline/datasets/collators/llm_collator.py create mode 100644 backends/qualcomm/genai_pipeline/datasets/collators/mllm_collator.py create mode 100644 backends/qualcomm/genai_pipeline/datasets/dataset_lookup.py create mode 100644 backends/qualcomm/genai_pipeline/datasets/dataset_options.py delete mode 100644 backends/qualcomm/genai_pipeline/datasets/default_calibration_data_adapter.py delete mode 100644 backends/qualcomm/genai_pipeline/datasets/default_training_data_adapter.py create mode 100644 backends/qualcomm/genai_pipeline/datasets/evaluation/__init__.py create mode 100644 backends/qualcomm/genai_pipeline/datasets/evaluation/default_evaluation_data_adapter.py create mode 100644 backends/qualcomm/genai_pipeline/datasets/evaluation/evaluation_data_adapter.py create mode 100644 backends/qualcomm/genai_pipeline/datasets/loaders/__init__.py create mode 100644 backends/qualcomm/genai_pipeline/datasets/loaders/llm/__init__.py create mode 100644 backends/qualcomm/genai_pipeline/datasets/loaders/llm/hf_dataset_adapter.py create mode 100644 backends/qualcomm/genai_pipeline/datasets/loaders/llm/lm_eval_adapter.py create mode 100644 backends/qualcomm/genai_pipeline/datasets/loaders/llm/message_sample_adapter.py create mode 100644 backends/qualcomm/genai_pipeline/datasets/loaders/mllm/__init__.py create mode 100644 backends/qualcomm/genai_pipeline/datasets/loaders/mllm/message_sample_adapter.py create mode 100644 backends/qualcomm/genai_pipeline/datasets/training/__init__.py create mode 100644 backends/qualcomm/genai_pipeline/datasets/training/default_training_data_adapter.py create mode 100644 backends/qualcomm/genai_pipeline/datasets/training/training_data_adapter.py delete mode 100644 backends/qualcomm/genai_pipeline/datasets/training_data_adapter.py create mode 100644 backends/qualcomm/genai_pipeline/model_components/__init__.py create mode 100644 backends/qualcomm/genai_pipeline/model_components/decoder/__init__.py create mode 100644 backends/qualcomm/genai_pipeline/model_components/embedding.py create mode 100644 backends/qualcomm/genai_pipeline/model_components/encoders/__init__.py create mode 100644 backends/qualcomm/genai_pipeline/model_lookup.py create mode 100644 backends/qualcomm/genai_pipeline/models/__init__.py create mode 100644 backends/qualcomm/genai_pipeline/quant_utilities.py create mode 100644 backends/qualcomm/genai_pipeline/recipes/__init__.py create mode 100644 backends/qualcomm/genai_pipeline/source_transform/__init__.py create mode 100644 backends/qualcomm/genai_pipeline/source_transform/checkpoint_key_remap.py create mode 100644 backends/qualcomm/genai_pipeline/source_transform/dtype_override.py create mode 100644 backends/qualcomm/genai_pipeline/source_transform/embedding_scale.py create mode 100644 backends/qualcomm/genai_pipeline/source_transform/linear_to_conv2d.py create mode 100644 backends/qualcomm/genai_pipeline/source_transform/rms_norm_offset.py create mode 100644 backends/qualcomm/genai_pipeline/source_transform/rope_layout.py create mode 100644 backends/qualcomm/genai_pipeline/strategies/model_preparation/llm_loader_adapter.py create mode 100644 backends/qualcomm/genai_pipeline/strategies/model_preparation/mllm_loader_adapter.py create mode 100644 backends/qualcomm/genai_pipeline/strategies/quantization/llm_quantizer_adapter.py create mode 100644 backends/qualcomm/genai_pipeline/strategies/quantization/mllm_quantizer_adapter.py create mode 100644 backends/qualcomm/genai_pipeline/tests/source_transform/test_checkpoint_key_remap.py create mode 100644 backends/qualcomm/genai_pipeline/tests/source_transform/test_dtype_override.py create mode 100644 backends/qualcomm/genai_pipeline/tests/source_transform/test_embedding_scale.py create mode 100644 backends/qualcomm/genai_pipeline/tests/source_transform/test_linear_to_conv2d.py create mode 100644 backends/qualcomm/genai_pipeline/tests/source_transform/test_rms_norm_offset.py create mode 100644 backends/qualcomm/genai_pipeline/tests/source_transform/test_rope_layout.py create mode 100644 backends/qualcomm/genai_pipeline/tests/strategies/model_preparation/test_llm_loader_adapter.py create mode 100644 backends/qualcomm/genai_pipeline/tests/strategies/model_preparation/test_mllm_loader_adapter.py create mode 100644 backends/qualcomm/genai_pipeline/tests/strategies/quantization/test_llm_quantizer_adapter.py create mode 100644 backends/qualcomm/genai_pipeline/tests/strategies/quantization/test_mllm_quantizer_adapter.py diff --git a/backends/qualcomm/genai_pipeline/cli.py b/backends/qualcomm/genai_pipeline/cli.py new file mode 100644 index 00000000000..5a050df53d8 --- /dev/null +++ b/backends/qualcomm/genai_pipeline/cli.py @@ -0,0 +1,754 @@ +# 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. + +"""GenAI Pipeline CLI entry point. + +Provides a command-line interface for running the GenAI Pipeline +with LLM models. Replaces the need to understand llama.py's internal +structure — users specify model, SoC, and prompt, and the pipeline +handles the rest. + +Usage: + python -m backends.qualcomm.genai_pipeline.cli \\ + --model llama3_2-1b_instruct \\ + --soc SM8750 \\ + --prompt "Hello, world!" \\ + --artifact-dir ./output + + # Compile only (no device inference): + python -m backends.qualcomm.genai_pipeline.cli \\ + --model llama3_2-1b_instruct \\ + --soc SM8750 \\ + --compile-only + + # List supported models: + python -m backends.qualcomm.genai_pipeline.cli --list-models +""" + +from __future__ import annotations + +import argparse +import logging +import sys +from dataclasses import replace +from typing import Any, Dict, List, Optional + +from executorch.backends.qualcomm.genai_pipeline.control_args import ( + DEFAULT_ARTIFACT_DIR, + DEFAULT_BACKEND, + DEFAULT_BATCH_SIZE, + DEFAULT_CALIB_HF_LIMIT, + DEFAULT_CALIB_LIMIT, + DEFAULT_DTYPE_OVERRIDE, + DEFAULT_EVAL_LIMIT, + DEFAULT_EVAL_METHOD, + DEFAULT_GCAP, + DEFAULT_MAX_SEQ_LEN, + DEFAULT_MODEL_MODE, + DEFAULT_NGRAM, + DEFAULT_PREFILL_AR_LEN, + DEFAULT_SOC_MODEL, + DEFAULT_TEMPERATURE, + DEFAULT_TRAIN_HF_LIMIT, + DEFAULT_TRAIN_LIMIT, + DEFAULT_TRAIN_VAL_RATIO, + DEFAULT_WINDOW, +) +from executorch.backends.qualcomm.genai_pipeline.genai_pipeline import QuantizationStage +from executorch.backends.qualcomm.genai_pipeline.stages.model_preparation_stage import ( + ModelPreparationStage, +) + +logger = logging.getLogger(__name__) + +PROGRAM_NAME = "genai_pipeline" +PROGRAM_DESCRIPTION = ( + "GenAI Pipeline CLI — standardized LLM deployment on Qualcomm platforms" +) + + +def build_parser() -> argparse.ArgumentParser: + """Build the CLI argument parser. + + Returns: + Configured ArgumentParser instance. + """ + parser = argparse.ArgumentParser( + prog=PROGRAM_NAME, + description=PROGRAM_DESCRIPTION, + ) + + # --- Model selection --- + parser.add_argument( + "--model", + type=str, + help="Model name (e.g., 'llama3_2-1b_instruct', 'qwen2_5-0_5b'). " + "Use --list-models to see available options.", + ) + parser.add_argument( + "--list-models", + action="store_true", + help="List all supported models and exit.", + ) + + # --- Model identification and input paths --- + parser.add_argument( + "--checkpoint", + type=str, + default=None, + help="Checkpoint to load the weights from. Required for the models whose " + "registry row carries no local weights.", + ) + parser.add_argument( + "--params", + type=str, + default=None, + help="Params JSON describing the model's shapes. Required for the models " + "whose registry row carries no params file (the llama family).", + ) + parser.add_argument( + "--tokenizer-model", + type=str, + default=None, + help="Tokenizer model to load instead of the model's own.", + ) + parser.add_argument( + "--tokenizer-bin", + type=str, + default=None, + help="Tokenizer binary, for Llama2-era models that ship one.", + ) + parser.add_argument( + "--artifact-dir", + type=str, + default=DEFAULT_ARTIFACT_DIR, + help=f"Directory for compiled artifacts (default: {DEFAULT_ARTIFACT_DIR}).", + ) + + # --- Quantization --- + parser.add_argument( + "--use-fp16", + action="store_true", + help="Use FP16 precision (skip quantization).", + ) + parser.add_argument( + "--embedding-quantize", + type=str, + default=None, + help="Fall back to the CPU embedding operator and quantize it, as " + "',' -- e.g. '4,32'.", + ) + parser.add_argument( + "--quant-recipe-suggestion", + action="store_true", + help="Emit a per-layer mixed-precision recipe suggestion during PTQ.", + ) + parser.add_argument( + "--batch-size", + type=int, + default=DEFAULT_BATCH_SIZE, + help=f"Batch size for decoder quantization. Larger values raise " + f"throughput at the cost of host memory, and affect only the " + f"calibration graph (default: {DEFAULT_BATCH_SIZE}).", + ) + + # --- Calibration data selection --- + parser.add_argument( + "--calib-tasks", + type=str, + nargs="+", + default=None, + help="lm-eval tasks to draw calibration samples from, e.g. " + "--calib-tasks wikitext.", + ) + parser.add_argument( + "--calib-samples", + type=str, + nargs="+", + default=None, + help="One or more JSON files of calibration samples, each a flat list of " + "objects with a 'messages' list (and 'files' for multimodal models). " + "Multiple files are merged.", + ) + parser.add_argument( + "--calib-limit", + type=int, + default=DEFAULT_CALIB_LIMIT, + help=f"How many samples to calibrate on (default: {DEFAULT_CALIB_LIMIT}).", + ) + parser.add_argument( + "--calib-num-fewshot", + type=int, + default=None, + metavar="N", + help="Number of few-shot examples in each calibration sample.", + ) + parser.add_argument( + "--calib-hf-dataset", + type=str, + default=None, + help="HuggingFace chat dataset for additional calibration data " + "(e.g. 'HuggingFaceTB/smol-smoltalk').", + ) + parser.add_argument( + "--calib-hf-limit", + type=int, + default=DEFAULT_CALIB_HF_LIMIT, + help="Number of samples to load from --calib-hf-dataset " + f"(default: {DEFAULT_CALIB_HF_LIMIT}).", + ) + + # --- Quantization-aware training --- + parser.add_argument( + "--qat", + action="store_true", + help="Enable Quantization-Aware Training (QAT). If not set, defaults to PTQ.", + ) + parser.add_argument( + "--train-config", + type=str, + default=None, + help="(QAT) YAML file overriding training configuration defaults.", + ) + parser.add_argument( + "--lr-config", + type=str, + default=None, + help="(QAT) YAML file configuring optimizer parameter groups.", + ) + parser.add_argument( + "--train-tasks", + type=str, + nargs="+", + default=None, + help="(QAT) lm-eval tasks for training data. Use --calib-tasks to " + "specify calibration data separately.", + ) + parser.add_argument( + "--train-limit", + type=int, + default=DEFAULT_TRAIN_LIMIT, + help=f"(QAT) Number of samples for train tasks (default: {DEFAULT_TRAIN_LIMIT}).", + ) + parser.add_argument( + "--train-hf-dataset", + type=str, + default=None, + help="(QAT) HuggingFace instruct dataset for training " + "(e.g. 'HuggingFaceTB/smol-smoltalk').", + ) + parser.add_argument( + "--train-hf-limit", + type=int, + default=DEFAULT_TRAIN_HF_LIMIT, + help="(QAT) Number of samples to load from --train-hf-dataset " + f"(default: {DEFAULT_TRAIN_HF_LIMIT}).", + ) + parser.add_argument( + "--train-val-ratio", + type=float, + default=DEFAULT_TRAIN_VAL_RATIO, + help="(QAT) Fraction of non-calib samples used for training; the " + "remainder becomes validation. 1.0 disables validation " + f"(default: {DEFAULT_TRAIN_VAL_RATIO}).", + ) + parser.add_argument( + "--freeze-all-params", + action="store_true", + help="(QAT) Freeze model weights so only quantization parameters are updated.", + ) + + # --- Evaluation --- + parser.add_argument( + "--eval-tasks", + type=str, + nargs="+", + default=None, + help="lm-eval tasks to evaluate on, e.g. --eval-tasks wikitext.", + ) + parser.add_argument( + "--eval-methods", + type=str, + nargs="+", + default=[DEFAULT_EVAL_METHOD], + help=f"Evaluation methods to run (default: {DEFAULT_EVAL_METHOD}).", + ) + parser.add_argument( + "--eval-limit", + type=int, + default=DEFAULT_EVAL_LIMIT, + help=f"How many samples to evaluate on (default: {DEFAULT_EVAL_LIMIT}).", + ) + parser.add_argument( + "--eval-num-fewshot", + type=int, + default=None, + metavar="N", + help="Number of few-shot examples in each evaluation sample.", + ) + + # --- Backend and SoC selection --- + parser.add_argument( + "--soc", + type=str, + default=DEFAULT_SOC_MODEL, + help=f"Target SoC model (default: {DEFAULT_SOC_MODEL}).", + ) + parser.add_argument( + "--backend", + type=str, + default=DEFAULT_BACKEND, + choices=["htp", "gpu"], + help=f"QNN backend type (default: {DEFAULT_BACKEND}).", + ) + + # --- Runtime prompts and multimodal inputs --- + parser.add_argument( + "--prompt", + type=str, + nargs="+", + default=["Hello, how are you?"], + help="User prompt(s) for text generation.", + ) + parser.add_argument( + "--system-prompt", + type=str, + default="", + help="System prompt for models that support one.", + ) + parser.add_argument( + "--temperature", + type=float, + default=DEFAULT_TEMPERATURE, + help=f"Sampling temperature for text generation (default: {DEFAULT_TEMPERATURE}).", + ) + parser.add_argument( + "--audio-path", + type=str, + nargs="+", + default=[], + help="Audio file(s) for multimodal (ALM) models.", + ) + parser.add_argument( + "--image-path", + type=str, + nargs="+", + default=[], + help="Image file(s) for multimodal (VLM) models.", + ) + + # Graph shapes and modes + parser.add_argument( + "--model-mode", + type=str, + default=DEFAULT_MODEL_MODE, + choices=["kv", "hybrid", "lookahead"], + help=f"Decoder model mode (default: {DEFAULT_MODEL_MODE}).", + ) + parser.add_argument( + "--max-seq-len", + type=int, + default=DEFAULT_MAX_SEQ_LEN, + help=f"Maximum sequence length (default: {DEFAULT_MAX_SEQ_LEN}).", + ) + parser.add_argument( + "--max-context-len", + type=int, + default=DEFAULT_MAX_SEQ_LEN, # TODO: Add DEFAULT_MAX_CONTEXT_LEN once attention sink is introduced in GenAI Pipeline. + help=f"Maximum context length (default: {DEFAULT_MAX_SEQ_LEN}).", + ) + parser.add_argument( + "--prefill-ar-len", + type=int, + default=DEFAULT_PREFILL_AR_LEN, + help=f"Prefill auto regressive length (default: {DEFAULT_PREFILL_AR_LEN}).", + ) + parser.add_argument( + "--dtype-override", + type=str, + default=DEFAULT_DTYPE_OVERRIDE, + choices=["fp32", "fp16"], + help=f"Override the dtype the model is loaded at " + f"(default: {DEFAULT_DTYPE_OVERRIDE}).", + ) + + # --- Lookahead decoding shape parameters (--model-mode lookahead) --- + parser.add_argument( + "--ngram", + type=int, + default=DEFAULT_NGRAM, + help=f"(lookahead) N-gram size (default: {DEFAULT_NGRAM}).", + ) + parser.add_argument( + "--window", + type=int, + default=DEFAULT_WINDOW, + help=f"(lookahead) Window size (default: {DEFAULT_WINDOW}).", + ) + parser.add_argument( + "--gcap", + type=int, + default=DEFAULT_GCAP, + help=f"(lookahead) Guess capacity (default: {DEFAULT_GCAP}).", + ) + + # --- Flow control --- + parser.add_argument( + "--compile-only", + action="store_true", + help="Only compile the model (skip inference).", + ) + parser.add_argument( + "--pre-gen-pte", + type=str, + default=None, + help="Directory containing pre-generated .pte artifacts.", + ) + parser.add_argument( + "-v", + "--verbose", + action="store_true", + help="Enable verbose logging.", + ) + + # --- Features --- + # Long context feature: Attention Sink + parser.add_argument( + "--use-attention-sink", + default=None, + type=str, + help="Use the attention sink feature to have fluent multi-round conversations. Specify the settings as ',', for example, '4,32'." + "This setting is for compilation. Once you compile with a chosen and , they cannot be changed at runtime. If you need to update them, you can recompile the attention sink module along with llama.py.", + ) + + return parser + + +def list_models() -> None: + """Log all supported models and exit.""" + from executorch.backends.qualcomm.genai_pipeline.model_lookup import ( + get_supported_models, + ) + + models = get_supported_models() + logger.info("Supported models (%d):", len(models)) + for model in models: + logger.info(" - %s", model) + + +def _validate_args(args: argparse.Namespace) -> None: + """Validate CLI arguments for consistency and constraints. + + Args: + args: Parsed CLI arguments. + + Raises: + ValueError: If arguments violate constraints. + """ + # TODO: The legacy path already supports the QAT feature, + # but it is not yet introduced in GenAI Pipeline. + if args.qat: + raise ValueError("QAT is not yet supported.") + + # TODO: The legacy path supports embedding quantization, but GenAI Pipeline + # does not yet support it. + if args.embedding_quantize is not None: + raise ValueError( + "`embedding_quantize` is currently unsupported in GenAI Pipeline." + ) + + # TODO: The legacy path already supports the attention sink feature, + # but it is not yet introduced in GenAI Pipeline. + if args.use_attention_sink is not None: + raise ValueError( + "`use_attention_sink` is currently unsupported in GenAI Pipeline." + ) + + from executorch.backends.qualcomm.genai_pipeline.model_lookup import is_multimodal + + if is_multimodal(args.model) and args.batch_size != 1: + logger.warning( + "Multi-batch is not supported for multimodal LLMs yet; " + "forcing batch_size from %d to 1.", + args.batch_size, + ) + args.batch_size = 1 + + +def _prepare_configs(args: argparse.Namespace) -> dict: + """Prepare model configs, transforms, and dataset options. + + Args: + args: Parsed CLI arguments. + + Returns: + Dict with keys: model_config, model_arch, weight_transforms, + module_transforms, dataset_options, model_options, quantize_options. + """ + from executorch.backends.qualcomm.genai_pipeline.datasets.dataset_options import ( + DatasetOptions, + ) + from executorch.backends.qualcomm.genai_pipeline.model_lookup import ( + get_model_arch, + get_model_config, + get_model_num_sharding, + get_quant_dtype, + get_quant_recipe, + get_source_transform, + get_state_dict_loader, + ) + + # TODO: The legacy path already supports attention sink, but GenAI Pipeline + # does not yet. Until attention sink is supported, keep max_context_len + # equal to max_seq_len. + logger.info( + "Setting max_context_len=%s to match max_seq_len because " + "attention sink is not yet supported in GenAI Pipeline.", + args.max_seq_len, + ) + args.max_context_len = args.max_seq_len + + model_config = get_model_config(args.model) + logger.info("Model config loaded") + + model_arch = get_model_arch(args.model, args, model_config=model_config) + logger.info("Model arch loaded") + + weight_transforms, module_transforms = get_source_transform( + args.model, + control_args=args, + model_config=model_config, + ) + logger.info("Loaded source transforms for model '%s'", args.model) + + num_sharding = get_model_num_sharding(args.model) + state_dict_loaders = get_state_dict_loader( + args.model, + control_args=args, + model_config=model_config, + ) + model_options = { + "model_arch": model_arch, + "weight_transforms": weight_transforms, + "module_transforms": module_transforms, + "embedding_quantize": args.embedding_quantize, + "num_shardings": num_sharding, + "state_dict_loader": state_dict_loaders, + } + + quantize_options = { + "quant_dtype": get_quant_dtype(args.model), + "quant_recipe": get_quant_recipe(args.model), + } + + dataset_options = replace( + DatasetOptions.from_namespace(args), + llm_config=model_config, + ) + + return { + "dataset_options": dataset_options, + "model_options": model_options, + "quantize_options": quantize_options, + } + + +def get_model_preparation_stage(args: argparse.Namespace) -> "ModelPreparationStage": + from executorch.backends.qualcomm.genai_pipeline.model_lookup import ( + get_model_loader_adapter, + ) + from executorch.backends.qualcomm.genai_pipeline.strategies.model_preparation.executorch_model_preparation_strategy import ( + ExecuTorchModelPreparationStrategy, + ) + + model_loader_adapter = get_model_loader_adapter(args.model, args) + model_preparation_stage = ModelPreparationStage( + ExecuTorchModelPreparationStrategy(model_loader_adapter=model_loader_adapter) + ) + return model_preparation_stage + + +def get_quantization_stage( + args: argparse.Namespace, + dataset_options: Dict[str, Any], + quantize_options: Dict[str, Any], +) -> "QuantizationStage": + from executorch.backends.qualcomm.genai_pipeline.datasets import ( + get_calibration_dataset_adapter, + get_eval_dataset_adapter, + get_training_dataset_adapter, + ) + from executorch.backends.qualcomm.genai_pipeline.model_lookup import ( + get_quantizer_adapter, + is_multimodal, + ) + from executorch.backends.qualcomm.genai_pipeline.strategies.quantization.executorch_quantization_strategy import ( + ExecuTorchQuantizationStrategy, + ) + + model_is_multimodal = is_multimodal(args.model) + + calibration_data_adapter = get_calibration_dataset_adapter( + dataset_options, is_multimodal=model_is_multimodal + ) + training_data_adapter = get_training_dataset_adapter( + dataset_options, is_multimodal=model_is_multimodal + ) + evaluation_data_adapter = get_eval_dataset_adapter( + dataset_options, is_multimodal=model_is_multimodal + ) + + quantizer_adapter = get_quantizer_adapter(args.model) + recipe_names = { + component: ( + recipe.__name__ if isinstance(recipe, type) else type(recipe).__name__ + ) + for component, recipe in quantize_options["quant_recipe"].items() + if recipe is not None + } + logger.info( + "Quantizer adapter created with recipes: %s", + recipe_names or "None", + ) + quantization_stage = QuantizationStage( + ExecuTorchQuantizationStrategy( + quantizer_adapter=quantizer_adapter, + calibration_data_adapter=calibration_data_adapter, + training_data_adapter=training_data_adapter, + evaluation_data_adapter=evaluation_data_adapter, + ) + ) + return quantization_stage + + +def run_pipeline(args: argparse.Namespace) -> None: + """Run the GenAI Pipeline with the given arguments. + + Args: + args: Parsed CLI arguments. + """ + from executorch.backends.qualcomm.genai_pipeline.engine_proxy import EngineProxy + from executorch.backends.qualcomm.genai_pipeline.genai_pipeline import GenAIPipeline + from executorch.backends.qualcomm.genai_pipeline.pipeline_context import ( + PipelineContext, + ) + from executorch.backends.qualcomm.genai_pipeline.pipeline_types import ( + EngineType, + STAGE_MODEL_PREPARATION, + STAGE_QUANTIZATION, + ) + from executorch.backends.qualcomm.serialization.qc_schema import ( + QnnExecuTorchBackendType, + ) + + logger.info("Args initialized: model=%s, soc=%s", args.model, args.soc) + + # Validate CLI constraints and normalize unsupported model configurations. + _validate_args(args) + + # Prepare extra options required by each pipeline stage. + configs = _prepare_configs(args) + + # Build pipeline context + extra_options = { + "backend": args.backend, + "model_mode": args.model_mode, + "max_seq_len": args.max_seq_len, + "use_fp16": args.use_fp16, + "compile_only": args.compile_only, + "verbose": args.verbose, + "model_options": configs["model_options"], + "quantize_options": configs["quantize_options"], + "dataset_options": configs["dataset_options"], + } + context = ( + PipelineContext.builder() + .with_model(args.model) + .with_soc(args.soc) + .with_prompt(args.prompt) + .with_artifact_dir(args.artifact_dir) + .with_extra_options(extra_options) + .build() + ) + logger.info( + "Pipeline context: model=%s, soc=%s", context.model_name, context.soc_model + ) + logger.info("Model '%s' configured for %s", args.model, args.soc) + logger.info("Artifact dir: %s", args.artifact_dir) + logger.info("Mode: %s", args.model_mode) + + # Build GenAI Pipeline + pipeline = GenAIPipeline( + model_preparation_stage=get_model_preparation_stage(args), + quantization_stage=( + # FP16 skips quantization. + get_quantization_stage( + args, + dataset_options=configs["dataset_options"], + quantize_options=configs["quantize_options"], + ) + if not args.use_fp16 + else None + ), + compilation_stage=None, + inference_stage=None, + engine_proxy=EngineProxy( + stage_engines={ + STAGE_MODEL_PREPARATION: EngineType.EXECUTORCH, + STAGE_QUANTIZATION: EngineType.EXECUTORCH, + }, + backend_type=QnnExecuTorchBackendType[ + f"k{args.backend.capitalize()}Backend" + ], + ), + ) + + pipeline.invoke(context) + + # Currently, only model preparation and quantization are supported. + # TODO: Add support for compilation and inference. + if args.compile_only: + logger.info("Compile-only mode (no device inference)") + + +def main(argv: Optional[List[str]] = None) -> None: + """CLI entry point. + + Args: + argv: Command-line arguments. If None, reads from sys.argv. + """ + parser = build_parser() + args = parser.parse_args(argv) + + # Setup logging + log_level = logging.DEBUG if args.verbose else logging.INFO + logging.basicConfig( + level=log_level, + format="[%(levelname)s %(asctime)s %(name)s] %(message)s", + ) + + # Handle --list-models + if args.list_models: + list_models() + return + + # Validate --model is required + if not args.model: + parser.error("--model is required (use --list-models to see options)") + + # Run pipeline + try: + run_pipeline(args) + except KeyError as e: + logger.error("Error: %s", e) + sys.exit(1) + except Exception as e: + logger.error("Pipeline failed: %s", e, exc_info=True) + sys.exit(1) + + +if __name__ == "__main__": + main() 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 194616d3563..36ccd1ad3aa 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, 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 diff --git a/backends/qualcomm/genai_pipeline/configs/quantization_input_config.py b/backends/qualcomm/genai_pipeline/configs/quantization_input_config.py index 30fdedeb904..030c13971ac 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, Tuple, TYPE_CHECKING +from typing import Any, Dict, Optional, Tuple, TYPE_CHECKING if TYPE_CHECKING: from executorch.backends.qualcomm.serialization.qc_schema import ( @@ -21,11 +21,11 @@ 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})`` @@ -33,33 +33,41 @@ class QuantizationInputConfig: 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) diff --git a/backends/qualcomm/genai_pipeline/configs/quantization_output_config.py b/backends/qualcomm/genai_pipeline/configs/quantization_output_config.py index 7852a1b90fb..d2e02f5f2bc 100644 --- a/backends/qualcomm/genai_pipeline/configs/quantization_output_config.py +++ b/backends/qualcomm/genai_pipeline/configs/quantization_output_config.py @@ -5,7 +5,9 @@ # 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 @@ -13,7 +15,10 @@ 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 diff --git a/backends/qualcomm/genai_pipeline/datasets/__init__.py b/backends/qualcomm/genai_pipeline/datasets/__init__.py index ecfa032f00d..aae8c3a2e04 100644 --- a/backends/qualcomm/genai_pipeline/datasets/__init__.py +++ b/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", ] diff --git a/backends/qualcomm/genai_pipeline/datasets/calibration/__init__.py b/backends/qualcomm/genai_pipeline/datasets/calibration/__init__.py new file mode 100644 index 00000000000..9dcad77fa46 --- /dev/null +++ b/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", +] diff --git a/backends/qualcomm/genai_pipeline/datasets/calibration/calibration_data_adapter.py b/backends/qualcomm/genai_pipeline/datasets/calibration/calibration_data_adapter.py new file mode 100644 index 00000000000..6c2f15c895a --- /dev/null +++ b/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. + """ + ... diff --git a/backends/qualcomm/genai_pipeline/datasets/calibration/default_calibration_data_adapter.py b/backends/qualcomm/genai_pipeline/datasets/calibration/default_calibration_data_adapter.py new file mode 100644 index 00000000000..7b37e86055b --- /dev/null +++ b/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} diff --git a/backends/qualcomm/genai_pipeline/datasets/calibration/llm_calibration_data_adapter.py b/backends/qualcomm/genai_pipeline/datasets/calibration/llm_calibration_data_adapter.py new file mode 100644 index 00000000000..4eaf28d5bbc --- /dev/null +++ b/backends/qualcomm/genai_pipeline/datasets/calibration/llm_calibration_data_adapter.py @@ -0,0 +1,122 @@ +# 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_BATCH_SIZE = 1 + + +class LLMCalibrationDataAdapter: + """Corpus-backed calibration provider for text-only LLMs. + + Composes the injected dataset loaders into a single ``ARTIFACT_TEXT_DECODER`` dataset + and wraps it in a DataLoader using the collator the collector supplies. + Unlike ``DefaultCalibrationDataAdapter``, this is the production path: + real token sequences, padded and masked by ``LLMCalibCollator``. + + An LLM has a single ``ARTIFACT_TEXT_DECODER`` component; the multi-component case + (encoders) belongs to ``MLLMCalibrationDataAdapter``. + + The returned object is always ``{ARTIFACT_TEXT_DECODER: DataLoader}``. + + Args: + dataset_loaders: Purpose-agnostic loaders, each exposing + ``load_dataset`` and returning a ``{ARTIFACT_TEXT_DECODER: dataset}`` dict. + collector: Provides the ``ARTIFACT_TEXT_DECODER`` collator for wrapping the + dataset in a DataLoader. + max_context_len: Sequence length for encoding and padding. The single + source of truth: it is pushed into ``extra_options`` before the + loaders run, so their encoding length and the collator's padding + length cannot diverge. + batch_size: Batch size for the ``ARTIFACT_TEXT_DECODER`` DataLoader. + """ + + def __init__( + self, + dataset_loaders: List[Any], + collector: Any, + max_context_len: int, + batch_size: int = DEFAULT_BATCH_SIZE, + ) -> None: + self._dataset_loaders = dataset_loaders + self._collector = collector + self._max_context_len = max_context_len + self._batch_size = batch_size + + 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 ``{ARTIFACT_TEXT_DECODER: DataLoader}`` built from the loaders. + + Args: + tokenizer: Tokenizer passed to each configured dataset loader. + example_inputs: Calibration-graph signatures required to construct + the decoder collator. + extra_options: Optional settings forwarded to each dataset loader. + + Returns: + A map containing a decoder calibration ``DataLoader``. + + Note: + ``max_context_len`` and ``batch_size`` are fixed at construction. + The context length is passed to both dataset loaders and the + collator so their sequence shapes remain aligned. + """ + from executorch.backends.qualcomm.genai_pipeline.artifact_keys import ( + ARTIFACT_TEXT_DECODER, + ) + from torch.utils.data import ConcatDataset, DataLoader + + extra_options = extra_options or {} + + if not self._dataset_loaders: + raise ValueError( + "LLMCalibrationDataAdapter requires at least one dataset loader; " + "use DefaultCalibrationDataAdapter for the random fallback." + ) + if self._collector is None: + raise ValueError("collector is required when dataset_loaders are provided") + + datasets: List[Any] = [] + for loader in self._dataset_loaders: + result = loader.load_dataset( + tokenizer=tokenizer, + max_context_len=self._max_context_len, + extra_options=extra_options, + ) + dataset = result.get(ARTIFACT_TEXT_DECODER) + if dataset is not None: + datasets.append(dataset) + + if not datasets: + raise ValueError( + "No calibration dataset produced: the configured loaders " + "returned no ARTIFACT_TEXT_DECODER data." + ) + + merged = datasets[0] if len(datasets) == 1 else ConcatDataset(datasets) + collators = self._collector.create_collators( + example_inputs, self._max_context_len + ) + + drop_last = self._batch_size > 1 + return { + ARTIFACT_TEXT_DECODER: DataLoader( + merged, + batch_size=self._batch_size, + shuffle=False, + drop_last=drop_last, + collate_fn=collators[ARTIFACT_TEXT_DECODER], + ) + } diff --git a/backends/qualcomm/genai_pipeline/datasets/calibration/mllm_calibration_data_adapter.py b/backends/qualcomm/genai_pipeline/datasets/calibration/mllm_calibration_data_adapter.py new file mode 100644 index 00000000000..b25453b7379 --- /dev/null +++ b/backends/qualcomm/genai_pipeline/datasets/calibration/mllm_calibration_data_adapter.py @@ -0,0 +1,128 @@ +# 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_BATCH_SIZE = 1 + + +class MLLMCalibrationDataAdapter: + """Corpus-backed calibration provider for multimodal models (VLM/ALM). + + Composes the injected dataset loaders into a single dataset per component + (ARTIFACT_TEXT_DECODER plus ARTIFACT_VISION_ENCODER / ARTIFACT_AUDIO_ENCODER) and wraps each in a + DataLoader using the collator the collector supplies. The text decoder is + batched; encoder components stay at batch size 1. + + The returned object is always ``{component: DataLoader}``. + + Args: + dataset_loaders: Purpose-agnostic loaders, each exposing + ``load_dataset`` and returning a component-keyed dataset dict. + collector: Provides per-component collators for wrapping datasets in + DataLoaders. + max_context_len: Sequence length for encoding and padding. The single + source of truth: it is pushed into ``extra_options`` before the + loaders run, so their encoding length and the collator's padding + length cannot diverge. + batch_size: Batch size for the ``ARTIFACT_TEXT_DECODER`` DataLoader; encoder + components stay at batch size 1. + """ + + def __init__( + self, + dataset_loaders: List[Any], + collector: Any, + max_context_len: int, + batch_size: int = DEFAULT_BATCH_SIZE, + ) -> None: + self._dataset_loaders = dataset_loaders + self._collector = collector + self._max_context_len = max_context_len + self._batch_size = batch_size + + 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 a component-keyed DataLoader map built from the loaders. + + Args: + tokenizer: Tokenizer passed to each configured dataset loader. + example_inputs: Calibration-graph signatures required to construct + the component collators. + extra_options: Optional settings forwarded to each dataset loader. + + Returns: + A map from artifact component keys to calibration ``DataLoader``s. + + Note: + ``max_context_len`` and ``batch_size`` are fixed at construction. + The context length is passed to both dataset loaders and the + collators so their sequence shapes remain aligned. + """ + from executorch.backends.qualcomm.genai_pipeline.artifact_keys import ( + ARTIFACT_TEXT_DECODER, + ) + from torch.utils.data import ConcatDataset, DataLoader + + 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} + + if not self._dataset_loaders: + raise ValueError( + "MLLMCalibrationDataAdapter requires at least one dataset loader; " + "use DefaultCalibrationDataAdapter for the random fallback." + ) + if self._collector is None: + raise ValueError("collector is required when dataset_loaders are provided") + + grouped: Dict[str, List[Any]] = {} + for loader in self._dataset_loaders: + result = loader.load_dataset( + tokenizer=tokenizer, + max_context_len=self._max_context_len, + extra_options=extra_options, + ) + for component, dataset in result.items(): + if dataset is not None: + grouped.setdefault(component, []).append(dataset) + + merged = { + component: (datasets[0] if len(datasets) == 1 else ConcatDataset(datasets)) + for component, datasets in grouped.items() + } + if not merged: + raise ValueError( + "No calibration dataset produced: the configured loaders " + "returned no data for any component." + ) + + collators = self._collector.create_collators( + example_inputs, self._max_context_len + ) + + dataloaders: Dict[str, Iterable[Any]] = {} + for component, component_dataset in merged.items(): + dataloaders[component] = DataLoader( + component_dataset, + batch_size=self._batch_size, + shuffle=False, + drop_last=self._batch_size > 1, + collate_fn=collators[component], + ) + return dataloaders diff --git a/backends/qualcomm/genai_pipeline/datasets/calibration_data_adapter.py b/backends/qualcomm/genai_pipeline/datasets/calibration_data_adapter.py deleted file mode 100644 index 7b14034167a..00000000000 --- a/backends/qualcomm/genai_pipeline/datasets/calibration_data_adapter.py +++ /dev/null @@ -1,54 +0,0 @@ -# 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 - -from typing import Any, Dict, Iterable, Optional, Protocol, runtime_checkable - - -@runtime_checkable -class CalibrationDataAdapter(Protocol): - """Protocol for calibration dataset construction. - - Kept separate from :class:`ModelLoaderAdapter` because data sources - (HuggingFace prompts, ``lm_eval`` tasks, JSON prompt files, multimodal file - inputs, synthetic tokens) vary independently of how the model is loaded and - mostly depend only on the tokenizer. Swapping the source is therefore a - matter of injecting a different adapter. - - The same calibration corpus is also consumed by the inference stage for - on-device result evaluation -- including pre-built ``.pte`` flows where model - preparation never runs -- which is why this lives in ``datasets/`` rather - than under a single stage. - - For quantization-aware training, see :class:`TrainingDataAdapter`: training - data yields ``(features, labels)`` pairs rather than model-argument tuples, - so it is a distinct contract. - """ - - def generate_calibration_data( - self, - tokenizer: Any, - num_samples: int = 128, - seq_length: int = 1024, - extra_options: Optional[Dict[str, Any]] = None, - ) -> Iterable[Any]: - """Build the calibration dataset used for post-training quantization. - - Args: - tokenizer: The tokenizer used to encode text samples. - num_samples: Number of calibration samples to produce. - seq_length: Sequence length for each sample. - extra_options: Implementation-specific options. Implementations may - support a ``dataset`` key accepting any - ``Iterable[Tuple[Tensor, ...]]`` (including a ``DataLoader``) - to use caller-supplied data directly. - - Returns: - An iterable of calibration input tuples, each suitable for - ``model(*sample)``. - """ - ... diff --git a/backends/qualcomm/genai_pipeline/datasets/collators/__init__.py b/backends/qualcomm/genai_pipeline/datasets/collators/__init__.py new file mode 100644 index 00000000000..b8943f3755d --- /dev/null +++ b/backends/qualcomm/genai_pipeline/datasets/collators/__init__.py @@ -0,0 +1,24 @@ +# 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. + +"""Component-aware collators, one per modality. + +Each collator maps a model component to its collate function (padding, mask +construction). Modality-specific: ``LLMDatasetCollector`` covers the text +decoder alone; ``MLLMDatasetCollector`` adds the vision / audio encoders. +""" + +from executorch.backends.qualcomm.genai_pipeline.datasets.collators.llm_collator import ( + LLMDatasetCollector, +) +from executorch.backends.qualcomm.genai_pipeline.datasets.collators.mllm_collator import ( + MLLMDatasetCollector, +) + +__all__ = [ + "LLMDatasetCollector", + "MLLMDatasetCollector", +] diff --git a/backends/qualcomm/genai_pipeline/datasets/collators/llm_collator.py b/backends/qualcomm/genai_pipeline/datasets/collators/llm_collator.py new file mode 100644 index 00000000000..fe6195bf57b --- /dev/null +++ b/backends/qualcomm/genai_pipeline/datasets/collators/llm_collator.py @@ -0,0 +1,64 @@ +# 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. + +"""LLM component-aware collator provider.""" + +from __future__ import annotations + +from typing import Any, Dict + + +class LLMDatasetCollector: + """Provides the collator for each LLM component. + + LLM models only have a ARTIFACT_TEXT_DECODER component. Building the DataLoaders is + the purpose adapter's responsibility; this collector only maps each + component to its collate function. + """ + + def create_collators( + self, + example_inputs: Dict[str, Any], + max_context_len: int, + ) -> Dict[str, Any]: + import torch + from executorch.backends.qualcomm.genai_pipeline.artifact_keys import ( + ARTIFACT_TEXT_DECODER, + ) + + from executorch.examples.qualcomm.oss_scripts.llama.dataset.collators import ( + LLMCalibCollator, + ) + + if ( + not isinstance(example_inputs, dict) + or ARTIFACT_TEXT_DECODER not in example_inputs + ): + raise ValueError( + f"LLM example_inputs must be a dict containing {ARTIFACT_TEXT_DECODER}; " + f"got {type(example_inputs)}" + ) + + decoder_inputs = example_inputs[ARTIFACT_TEXT_DECODER] + if not isinstance(decoder_inputs, tuple) or len(decoder_inputs) < 2: + raise ValueError( + f"{ARTIFACT_TEXT_DECODER} example_inputs must be a tuple of " + f"(tokens, attention_mask, ...); got {type(decoder_inputs)}" + ) + + attn_mask = decoder_inputs[1] + if attn_mask is None: + raise ValueError( + f"attention_mask not found at {ARTIFACT_TEXT_DECODER} example_inputs[1]" + ) + + return { + ARTIFACT_TEXT_DECODER: LLMCalibCollator( + attn_mask_template=attn_mask, + max_context_len=max_context_len, + token_dtype=torch.int32, + ) + } diff --git a/backends/qualcomm/genai_pipeline/datasets/collators/mllm_collator.py b/backends/qualcomm/genai_pipeline/datasets/collators/mllm_collator.py new file mode 100644 index 00000000000..f6508b7a9aa --- /dev/null +++ b/backends/qualcomm/genai_pipeline/datasets/collators/mllm_collator.py @@ -0,0 +1,71 @@ +# 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. + +"""MLLM component-aware collator provider.""" + +from __future__ import annotations + +from typing import Any, Dict + + +class MLLMDatasetCollector: + """Provides the collator for each MLLM component. + + ARTIFACT_TEXT_DECODER uses ``LLMCalibCollator``; encoder components + (ARTIFACT_VISION_ENCODER / ARTIFACT_AUDIO_ENCODER) use ``ModalityEncoderCollator``. Building + the DataLoaders is the purpose adapter's responsibility; this collector + only maps each component to its collate function. + """ + + def create_collators( + self, + example_inputs: Dict[str, Any], + max_context_len: int, + ) -> Dict[str, Any]: + import torch + from executorch.backends.qualcomm.genai_pipeline.artifact_keys import ( + ARTIFACT_AUDIO_ENCODER, + ARTIFACT_TEXT_DECODER, + ARTIFACT_VISION_ENCODER, + ) + + from executorch.examples.qualcomm.oss_scripts.llama.dataset.collators import ( + LLMCalibCollator, + ModalityEncoderCollator, + ) + + if ( + not isinstance(example_inputs, dict) + or ARTIFACT_TEXT_DECODER not in example_inputs + ): + raise ValueError( + f"MLLM example_inputs must be a dict containing {ARTIFACT_TEXT_DECODER}; " + f"got {type(example_inputs)}" + ) + + decoder_inputs = example_inputs[ARTIFACT_TEXT_DECODER] + if not isinstance(decoder_inputs, tuple) or len(decoder_inputs) < 2: + raise ValueError( + f"{ARTIFACT_TEXT_DECODER} example_inputs must be a tuple of " + f"(tokens, attention_mask, ...); got {type(decoder_inputs)}" + ) + + attn_mask = decoder_inputs[1] + if attn_mask is None: + raise ValueError( + f"attention_mask not found at {ARTIFACT_TEXT_DECODER} example_inputs[1]" + ) + + encoder_collator = ModalityEncoderCollator() + return { + ARTIFACT_TEXT_DECODER: LLMCalibCollator( + attn_mask_template=attn_mask, + max_context_len=max_context_len, + token_dtype=torch.int32, + ), + ARTIFACT_VISION_ENCODER: encoder_collator, + ARTIFACT_AUDIO_ENCODER: encoder_collator, + } diff --git a/backends/qualcomm/genai_pipeline/datasets/dataset_lookup.py b/backends/qualcomm/genai_pipeline/datasets/dataset_lookup.py new file mode 100644 index 00000000000..0a2c70eccba --- /dev/null +++ b/backends/qualcomm/genai_pipeline/datasets/dataset_lookup.py @@ -0,0 +1,239 @@ +# 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. + +"""Dataset loader, collector, and purpose-adapter lookup.""" + +from __future__ import annotations + + +def get_dataset_adapter(dataset_options, is_multimodal=False): + """Build purpose-agnostic dataset loaders selected by ``dataset_options``. + + Each loader exposes ``load_dataset(tokenizer, extra_options)`` returning a + component-keyed dataset dict. The purpose (calibration / training / eval) + lives in the purpose adapters, not here. Multimodal message samples also + need the model config to build encoder datasets, so it is bound into that + loader here instead of being routed through the quantization strategy. + """ + if dataset_options is None: + return [] + + loaders = [] + if is_multimodal: + from executorch.backends.qualcomm.genai_pipeline.datasets.loaders.mllm.message_sample_adapter import ( + MessageSampleAdapter as MLLMMessageSampleAdapter, + ) + + if dataset_options.calib_samples: + loaders.append( + MLLMMessageSampleAdapter( + dataset_options.calib_samples, + llm_config=dataset_options.llm_config, + ) + ) + return loaders + + from executorch.backends.qualcomm.genai_pipeline.datasets.loaders.llm.hf_dataset_adapter import ( + HFDatasetAdapter, + ) + from executorch.backends.qualcomm.genai_pipeline.datasets.loaders.llm.lm_eval_adapter import ( + LMEvalAdapter, + ) + from executorch.backends.qualcomm.genai_pipeline.datasets.loaders.llm.message_sample_adapter import ( + MessageSampleAdapter, + ) + + if dataset_options.calib_tasks: + loaders.append( + LMEvalAdapter( + dataset_options.calib_tasks, + dataset_options.calib_limit, + dataset_options.calib_num_fewshot, + ) + ) + if dataset_options.calib_samples: + loaders.append(MessageSampleAdapter(dataset_options.calib_samples)) + if dataset_options.calib_hf_dataset: + loaders.append( + HFDatasetAdapter( + dataset_options.calib_hf_dataset, + dataset_options.calib_hf_limit, + ) + ) + return loaders + + +def get_collector(is_multimodal=False): + """Return the component-aware collator provider for the model modality.""" + if is_multimodal: + from executorch.backends.qualcomm.genai_pipeline.datasets.collators.mllm_collator import ( + MLLMDatasetCollector, + ) + + return MLLMDatasetCollector() + + from executorch.backends.qualcomm.genai_pipeline.datasets.collators.llm_collator import ( + LLMDatasetCollector, + ) + + return LLMDatasetCollector() + + +def get_calibration_dataset_adapter(dataset_options, is_multimodal=False): + """Assemble the calibration purpose adapter. + + Selects source loaders from ``dataset_options``. With no source, returns + the modality-agnostic ``DefaultCalibrationDataAdapter`` (random pipeline + sanity data, sized from the options). With sources, wires them and the + modality's collector through :class:`CalibrationDatasetBuilder`. + """ + loaders = get_dataset_adapter(dataset_options, is_multimodal) + + if not loaders: + from executorch.backends.qualcomm.genai_pipeline.datasets.calibration.default_calibration_data_adapter import ( + DefaultCalibrationDataAdapter, + ) + + kwargs = {} + if dataset_options is not None: + if dataset_options.num_samples is not None: + kwargs["num_samples"] = dataset_options.num_samples + kwargs["batch_size"] = dataset_options.batch_size + if dataset_options.seed is not None: + kwargs["seed"] = dataset_options.seed + return DefaultCalibrationDataAdapter(**kwargs) + + builder = CalibrationDatasetBuilder(is_multimodal) + for loader in loaders: + builder.with_dataset(loader) + builder.with_collector(get_collector(is_multimodal)) + builder.with_context_len(dataset_options.max_context_len) + builder.with_batch_size(dataset_options.batch_size) + return builder.build() + + +class CalibrationDatasetBuilder: + """Assembler for a corpus-backed calibration adapter. + + The axes are modality (fixed at construction), data sources + (:meth:`with_dataset`), collator (:meth:`with_collector`), and shaping + (:meth:`with_context_len` / :meth:`with_batch_size`). Modality is decided + once and reused for adapter selection. :meth:`build` wires the modality's + corpus adapter with the accumulated sources, the supplied collector, and + the shaping; the random default is not this builder's concern (see + ``get_calibration_dataset_adapter``). + + Example usage: + + adapter = ( + CalibrationDatasetBuilder(is_multimodal) + .with_dataset(lm_eval_loader) + .with_dataset(message_sample_loader) + .with_collector(collector) + .with_context_len(1024) + .with_batch_size(1) + .build() + ) + + Args: + is_multimodal: Whether the model is multimodal (VLM/ALM). Selects the + MLLM adapter instead of the LLM one. + """ + + def __init__(self, is_multimodal: bool = False) -> None: + self._is_multimodal = is_multimodal + self._dataset_adapters: list = [] + self._collector = None + self._max_context_len = None + self._batch_size = None + + def with_dataset(self, dataset_adapter) -> "CalibrationDatasetBuilder": + """Accumulate one data-source loader. May be called more than once.""" + self._dataset_adapters.append(dataset_adapter) + return self + + def with_collector(self, collector) -> "CalibrationDatasetBuilder": + """Set the collator provider.""" + self._collector = collector + return self + + def with_context_len(self, max_context_len) -> "CalibrationDatasetBuilder": + """Set the sequence length for encoding and padding.""" + self._max_context_len = max_context_len + return self + + def with_batch_size(self, batch_size) -> "CalibrationDatasetBuilder": + """Set the text-decoder DataLoader batch size.""" + self._batch_size = batch_size + return self + + def build(self): + """Build and validate the calibration adapter. + + Returns: + An ``LLMCalibrationDataAdapter`` or ``MLLMCalibrationDataAdapter``. + + Raises: + ValueError: If required fields are missing. + """ + missing = [] + if not self._dataset_adapters: + missing.append("dataset sources (use .with_dataset())") + if self._collector is None: + missing.append("collector (use .with_collector())") + if self._max_context_len is None: + missing.append("max_context_len (use .with_context_len())") + + if missing: + raise ValueError( + f"Cannot build calibration adapter, missing required fields: " + f"{', '.join(missing)}" + ) + + kwargs = { + "dataset_loaders": self._dataset_adapters, + "collector": self._collector, + "max_context_len": self._max_context_len, + } + if self._batch_size is not None: + kwargs["batch_size"] = self._batch_size + + if self._is_multimodal: + from executorch.backends.qualcomm.genai_pipeline.datasets.calibration.mllm_calibration_data_adapter import ( + MLLMCalibrationDataAdapter, + ) + + return MLLMCalibrationDataAdapter(**kwargs) + + from executorch.backends.qualcomm.genai_pipeline.datasets.calibration.llm_calibration_data_adapter import ( + LLMCalibrationDataAdapter, + ) + + return LLMCalibrationDataAdapter(**kwargs) + + +def get_training_dataset_adapter(dataset_options, is_multimodal=False): + """Assemble the training purpose adapter.""" + from executorch.backends.qualcomm.genai_pipeline.datasets.training.default_training_data_adapter import ( + DefaultTrainingDataAdapter, + ) + + return DefaultTrainingDataAdapter( + dataset_loaders=get_dataset_adapter(dataset_options, is_multimodal), + collector=get_collector(is_multimodal), + ) + + +def get_eval_dataset_adapter(dataset_options, is_multimodal=False): + """Assemble the evaluation purpose adapter.""" + from executorch.backends.qualcomm.genai_pipeline.datasets.evaluation.default_evaluation_data_adapter import ( + DefaultEvaluationDataAdapter, + ) + + return DefaultEvaluationDataAdapter( + dataset_loaders=get_dataset_adapter(dataset_options, is_multimodal), + collector=get_collector(is_multimodal), + ) diff --git a/backends/qualcomm/genai_pipeline/datasets/dataset_options.py b/backends/qualcomm/genai_pipeline/datasets/dataset_options.py new file mode 100644 index 00000000000..d20dbdf6598 --- /dev/null +++ b/backends/qualcomm/genai_pipeline/datasets/dataset_options.py @@ -0,0 +1,115 @@ +# 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. + +"""DatasetOptions: the typed calibration / training data configuration. + +Datasets are a **cross-stage** concern -- the same corpus feeds PTQ calibration +during quantization and result evaluation during inference -- so the knobs that +select and shape that corpus are grouped here at package level rather than nested +inside any one stage's options. Quantization consumes them today; the evaluation +path will consume the same type. + +Field names match ``llama.py``'s / the GenAI CLI parser ``dest`` names, so +:meth:`from_namespace` can build a ``DatasetOptions`` straight from a parsed +argument namespace, dropping attributes without a matching field. +""" + +from __future__ import annotations + +import argparse +import logging +from dataclasses import dataclass, fields +from typing import Any, List, Optional + +logger = logging.getLogger(__name__) + + +@dataclass(frozen=True) +class DatasetOptions: + """Typed calibration / training / evaluation data configuration. + + Every field selects or shapes the corpus that feeds one part of the flow; + the grouping below follows which part consumes it. Field names match + ``llama.py``'s / the GenAI CLI parser ``dest`` names, so + :meth:`from_namespace` builds a ``DatasetOptions`` straight from a parsed + namespace, dropping attributes without a matching field. + """ + + # --- Calibration data selection --- + # + # ``calib_tasks`` (lm_eval task names) and ``calib_samples`` (message-sample + # JSON files) are the corpus sources; ``calib_hf_dataset`` adds a HuggingFace + # chat dataset (e.g. ``HuggingFaceTB/smol-smoltalk``). The ``*_limit`` and + # ``num_fewshot`` fields shape how much of each is drawn. + calib_tasks: Optional[List[str]] = None + calib_samples: Optional[List[str]] = None + calib_limit: Optional[int] = None + calib_num_fewshot: Optional[int] = None + calib_hf_dataset: Optional[str] = None + calib_hf_limit: int = 1 + + # --- Sample shape--- + max_context_len: int = 1024 + batch_size: int = 1 + model_mode: Optional[str] = None + + # --- Quantization mode --- + qat: bool = False + + # --- Evaluation --- + eval_tasks: Optional[List[str]] = None + eval_limit: int = 1 + eval_num_fewshot: Optional[int] = None + + # --- QAT training data --- + # + # ``train_val_ratio`` is the fraction of non-calib samples used for training; + # the remainder becomes validation, and ``1.0`` disables validation. + train_tasks: Optional[List[str]] = None + train_limit: int = 1 + train_hf_dataset: Optional[str] = None + train_hf_limit: int = 1000 + train_val_ratio: float = 1.0 + + # --- Random fallback (DefaultCalibrationDataAdapter) --- + num_samples: Optional[int] = None + seed: Optional[int] = None + + # --- Non-CLI --- + # + # ``llm_config`` is the model's ``LLMModelConfig``, forwarded to adapters + # that need model-specific tokenization details. It has no CLI argument and + # is set by the caller after :meth:`from_namespace` (e.g. via + # ``dataclasses.replace``). + llm_config: Any = None + + @classmethod + def field_names(cls) -> frozenset: + """The set of field names this dataclass defines.""" + return frozenset(f.name for f in fields(cls)) + + @classmethod + def from_namespace(cls, namespace: argparse.Namespace) -> "DatasetOptions": + """Build from an ``argparse.Namespace``, ignoring unknown attributes. + + The CLI parser produces a superset of these fields; attributes without a + matching field are dropped. ``llm_config`` has no CLI argument and is set + by the caller afterwards (e.g. via ``dataclasses.replace``). + + Args: + namespace: A parsed namespace, e.g. from the GenAI CLI parser. + + Returns: + A ``DatasetOptions`` carrying every recognised attribute. + """ + known = cls.field_names() + supplied = vars(namespace) + + ignored = sorted(set(supplied) - known) + if ignored: + logger.debug("Ignoring unrecognised arguments: %s", ignored) + + return cls(**{k: v for k, v in supplied.items() if k in known}) diff --git a/backends/qualcomm/genai_pipeline/datasets/default_calibration_data_adapter.py b/backends/qualcomm/genai_pipeline/datasets/default_calibration_data_adapter.py deleted file mode 100644 index bef81896066..00000000000 --- a/backends/qualcomm/genai_pipeline/datasets/default_calibration_data_adapter.py +++ /dev/null @@ -1,89 +0,0 @@ -# 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_SEED = 42 - - -class DefaultCalibrationDataAdapter: - """Default calibration data provider using random token sequences. - - Suitable for bring-up. For production quantization accuracy, pass a real - corpus via ``extra_options["dataset"]`` or inject a corpus-backed adapter - (HuggingFace prompts, ``lm_eval`` tasks, JSON prompt files, ...). - """ - - def generate_calibration_data( - self, - tokenizer: Any, - num_samples: int = DEFAULT_NUM_SAMPLES, - seq_length: int = DEFAULT_SEQ_LENGTH, - extra_options: Optional[Dict[str, Any]] = None, - ) -> Iterable[Any]: - """Generate calibration data as random token sequences. - - Args: - tokenizer: The tokenizer to use. Only ``vocab_size`` is consulted. - num_samples: Number of calibration samples to generate. - seq_length: Sequence length for each sample. - extra_options: Additional options. Supported keys: - - ``dataset``: Caller-supplied calibration data to use instead - of random data. Accepts any ``Iterable[Tuple[Tensor, ...]]``, - including a ``DataLoader`` with a custom ``collate_fn``, which - allows batched/streamed calibration without materializing the - whole dataset. - - ``seed``: Random seed for reproducibility (default: 42). - - Returns: - A list of ``(input_ids, attention_mask)`` tuples of shape - ``(1, seq_length)``. - - Raises: - ValueError: If the tokenizer has no usable ``vocab_size`` and no - ``dataset`` was supplied. - """ - import torch - - extra_options = extra_options or {} - - # A caller-supplied dataset takes precedence over random generation. - if "dataset" in extra_options: - logger.info("Using caller-supplied dataset for calibration") - return extra_options["dataset"] - - seed = extra_options.get("seed", DEFAULT_SEED) - torch.manual_seed(seed) - - logger.info( - "Generating %d random calibration samples (seq_length=%d, seed=%d)", - num_samples, - seq_length, - 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(num_samples): - input_ids = torch.randint(0, vocab_size, (1, seq_length)) - attention_mask = torch.ones(1, seq_length, dtype=torch.long) - calibration_data.append((input_ids, attention_mask)) - - return calibration_data diff --git a/backends/qualcomm/genai_pipeline/datasets/default_training_data_adapter.py b/backends/qualcomm/genai_pipeline/datasets/default_training_data_adapter.py deleted file mode 100644 index 575e31374b3..00000000000 --- a/backends/qualcomm/genai_pipeline/datasets/default_training_data_adapter.py +++ /dev/null @@ -1,65 +0,0 @@ -# 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, Optional, Tuple - -logger = logging.getLogger(__name__) - -DEFAULT_NUM_SAMPLES = 128 -DEFAULT_SEQ_LENGTH = 1024 - - -class DefaultTrainingDataAdapter: - """Pass-through QAT training data provider. - - Training data is **not** synthesized: random tokens carry no learning - signal, so there is no meaningful default corpus. Callers enabling QAT must - supply their data via ``extra_options["training_data"]``, or inject a - corpus-backed :class:`TrainingDataAdapter`. - """ - - def generate_training_data( - self, - tokenizer: Any, - num_samples: int = DEFAULT_NUM_SAMPLES, - seq_length: int = DEFAULT_SEQ_LENGTH, - extra_options: Optional[Dict[str, Any]] = None, - ) -> Iterable[Tuple[Any, Any]]: - """Return caller-supplied QAT training data. - - Args: - tokenizer: The tokenizer (unused by this implementation). - num_samples: Ignored; retained for protocol conformance. - seq_length: Ignored; retained for protocol conformance. - extra_options: Additional options. Supported keys: - - ``training_data``: ``(features, labels)`` pairs, or any - ``Iterable`` yielding them (including a ``DataLoader``). - - Returns: - The caller-supplied ``(features, labels)`` pairs. - - Raises: - ValueError: If no ``training_data`` was supplied. QAT cannot proceed - without labels, so failing loudly is preferable to silently - degrading to PTQ. - """ - extra_options = extra_options or {} - - training_data = extra_options.get("training_data") - if training_data is None: - raise ValueError( - "No training data supplied. DefaultTrainingDataAdapter does not " - "synthesize labelled data; pass it via " - "extra_options['training_data'] or inject a corpus-backed " - "TrainingDataAdapter. For PTQ, skip QAT entirely and use " - "CalibrationDataAdapter instead." - ) - - logger.info("Using caller-supplied training data for QAT") - return training_data diff --git a/backends/qualcomm/genai_pipeline/datasets/evaluation/__init__.py b/backends/qualcomm/genai_pipeline/datasets/evaluation/__init__.py new file mode 100644 index 00000000000..c75b30b362d --- /dev/null +++ b/backends/qualcomm/genai_pipeline/datasets/evaluation/__init__.py @@ -0,0 +1,19 @@ +# 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. + +"""Evaluation purpose adapters (protocol + implementation).""" + +from executorch.backends.qualcomm.genai_pipeline.datasets.evaluation.default_evaluation_data_adapter import ( + DefaultEvaluationDataAdapter, +) +from executorch.backends.qualcomm.genai_pipeline.datasets.evaluation.evaluation_data_adapter import ( + EvaluationDataAdapter, +) + +__all__ = [ + "EvaluationDataAdapter", + "DefaultEvaluationDataAdapter", +] diff --git a/backends/qualcomm/genai_pipeline/datasets/evaluation/default_evaluation_data_adapter.py b/backends/qualcomm/genai_pipeline/datasets/evaluation/default_evaluation_data_adapter.py new file mode 100644 index 00000000000..9163bd44c33 --- /dev/null +++ b/backends/qualcomm/genai_pipeline/datasets/evaluation/default_evaluation_data_adapter.py @@ -0,0 +1,112 @@ +# 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 + + +class DefaultEvaluationDataAdapter: + """Post-quantization evaluation data provider. + + Evaluation data is not synthesized. Uses injected dataset loaders when + present; otherwise callers must supply data via ``extra_options["eval_data"]``. + + The returned object is always ``{component: iterable}``. + + Args: + dataset_loaders: Purpose-agnostic loaders, each exposing + ``load_dataset`` and returning a component-keyed dataset dict. + collector: Provides per-component collators for wrapping datasets in + DataLoaders. + """ + + def __init__( + self, + dataset_loaders: Optional[List[Any]] = None, + collector: Optional[Any] = None, + ) -> None: + self._dataset_loaders = dataset_loaders or [] + self._collector = collector + + def generate_eval_data( + self, + tokenizer: Any, + num_samples: int = DEFAULT_NUM_SAMPLES, + seq_length: int = DEFAULT_SEQ_LENGTH, + extra_options: Optional[Dict[str, Any]] = None, + ) -> Dict[str, Iterable[Any]]: + """Return component-keyed post-quantization evaluation data.""" + from executorch.backends.qualcomm.genai_pipeline.artifact_keys import ( + ARTIFACT_TEXT_DECODER, + ) + from torch.utils.data import ConcatDataset, DataLoader + + extra_options = extra_options or {} + + eval_data = extra_options.get("eval_data") + if eval_data is not None: + logger.info("Using caller-supplied evaluation data") + return ( + eval_data + if isinstance(eval_data, dict) + else {ARTIFACT_TEXT_DECODER: eval_data} + ) + + if not self._dataset_loaders: + raise ValueError( + "No evaluation data supplied. DefaultEvaluationDataAdapter does not " + "synthesize evaluation data; pass it via " + "extra_options['eval_data'] or inject corpus sources." + ) + + if self._collector is None: + raise ValueError("collector is required when dataset_loaders are provided") + + max_context_len = extra_options.get("max_context_len", seq_length) + example_inputs = extra_options.get("example_inputs") + batch_size = extra_options.get("batch_size", 1) + + grouped: Dict[str, List[Any]] = {} + for loader in self._dataset_loaders: + result = loader.load_dataset( + tokenizer=tokenizer, + max_context_len=max_context_len, + extra_options=extra_options, + ) + for component, dataset in result.items(): + if dataset is not None: + grouped.setdefault(component, []).append(dataset) + + merged = { + component: (datasets[0] if len(datasets) == 1 else ConcatDataset(datasets)) + for component, datasets in grouped.items() + } + if not merged: + return {} + + collators = self._collector.create_collators(example_inputs, max_context_len) + + dataloaders: Dict[str, Iterable[Any]] = {} + for component, component_dataset in merged.items(): + component_batch_size = ( + batch_size if component == ARTIFACT_TEXT_DECODER else 1 + ) + drop_last = component == ARTIFACT_TEXT_DECODER and component_batch_size > 1 + dataloaders[component] = DataLoader( + component_dataset, + batch_size=component_batch_size, + shuffle=False, + drop_last=drop_last, + collate_fn=collators[component], + ) + return dataloaders diff --git a/backends/qualcomm/genai_pipeline/datasets/evaluation/evaluation_data_adapter.py b/backends/qualcomm/genai_pipeline/datasets/evaluation/evaluation_data_adapter.py new file mode 100644 index 00000000000..251489b3a99 --- /dev/null +++ b/backends/qualcomm/genai_pipeline/datasets/evaluation/evaluation_data_adapter.py @@ -0,0 +1,31 @@ +# 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. + +"""Evaluation data adapter protocol. + +A purpose adapter assembling evaluation data from raw dataset loaders. Its output +is always component-keyed so inference/evaluation strategies can consume one +LLM/MLLM contract. +""" + +from __future__ import annotations + +from typing import Any, Dict, Iterable, Optional, Protocol, runtime_checkable + + +@runtime_checkable +class EvaluationDataAdapter(Protocol): + """Protocol for assembling evaluation data.""" + + def generate_eval_data( + self, + tokenizer: Any, + num_samples: int = ..., + seq_length: int = ..., + extra_options: Optional[Dict[str, Any]] = None, + ) -> Dict[str, Iterable[Any]]: + """Return ``{component: iterable}`` evaluation data.""" + ... diff --git a/backends/qualcomm/genai_pipeline/datasets/loaders/__init__.py b/backends/qualcomm/genai_pipeline/datasets/loaders/__init__.py new file mode 100644 index 00000000000..412c564a29d --- /dev/null +++ b/backends/qualcomm/genai_pipeline/datasets/loaders/__init__.py @@ -0,0 +1,13 @@ +# 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. + +"""Purpose-agnostic dataset loaders and collators, organized by modality. + +``llm/`` and ``mllm/`` each provide the data-source loaders (lm_eval tasks, +message-sample JSON) and the component-aware collator for that modality. The +purpose (calibration / training / evaluation) is layered on top by the adapter +packages, not here. +""" diff --git a/backends/qualcomm/genai_pipeline/datasets/loaders/llm/__init__.py b/backends/qualcomm/genai_pipeline/datasets/loaders/llm/__init__.py new file mode 100644 index 00000000000..28ce8434eb7 --- /dev/null +++ b/backends/qualcomm/genai_pipeline/datasets/loaders/llm/__init__.py @@ -0,0 +1,23 @@ +# 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. + +"""LLM dataset loaders for single-modality models.""" + +from executorch.backends.qualcomm.genai_pipeline.datasets.loaders.llm.hf_dataset_adapter import ( + HFDatasetAdapter, +) +from executorch.backends.qualcomm.genai_pipeline.datasets.loaders.llm.lm_eval_adapter import ( + LMEvalAdapter, +) +from executorch.backends.qualcomm.genai_pipeline.datasets.loaders.llm.message_sample_adapter import ( + MessageSampleAdapter, +) + +__all__ = [ + "HFDatasetAdapter", + "LMEvalAdapter", + "MessageSampleAdapter", +] diff --git a/backends/qualcomm/genai_pipeline/datasets/loaders/llm/hf_dataset_adapter.py b/backends/qualcomm/genai_pipeline/datasets/loaders/llm/hf_dataset_adapter.py new file mode 100644 index 00000000000..24b5c097dd6 --- /dev/null +++ b/backends/qualcomm/genai_pipeline/datasets/loaders/llm/hf_dataset_adapter.py @@ -0,0 +1,67 @@ +# 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. + +"""LLM HuggingFace chat dataset loader.""" + +from __future__ import annotations + +import logging +from typing import Any, Dict, Optional + +logger = logging.getLogger(__name__) + + +class HFDatasetAdapter: + """Loads an LLM dataset from a HuggingFace chat dataset. + + Purpose-agnostic: it produces a component-keyed dataset dict and knows + nothing about calibration / training / eval. The dataset name and sample + count are fixed at construction; ``load_dataset`` only reads shaping + options. + """ + + def __init__(self, dataset_name: str, limit: int) -> None: + self._dataset_name = dataset_name + self._limit = limit + + def load_dataset( + self, + tokenizer: Any, + max_context_len: int, + extra_options: Optional[Dict[str, Any]] = None, + ) -> Dict[str, Any]: + """Load the dataset as ``{ARTIFACT_TEXT_DECODER: Dataset}``. + + Args: + tokenizer: TokenizerWrapper instance. + max_context_len: Sequence length the samples are encoded at. + extra_options: Reserved for adapter-specific options; unused here. + """ + from executorch.backends.qualcomm.genai_pipeline.artifact_keys import ( + ARTIFACT_TEXT_DECODER, + ) + from executorch.examples.qualcomm.oss_scripts.llama.dataset.builders import ( + DecoderDatasetBuilder, + ) + + logger.info( + "Loading LLM HuggingFace dataset: dataset=%s, limit=%s", + self._dataset_name, + self._limit, + ) + + decoder_builder = DecoderDatasetBuilder( + tokenizer_wrapper=tokenizer, + max_context_len=max_context_len, + is_multimodal=False, + ) + dataset = decoder_builder.from_hf_source( + self._dataset_name, + num_samples=self._limit, + ) + + logger.info("Loaded LLM HuggingFace dataset with %d samples", len(dataset)) + return {ARTIFACT_TEXT_DECODER: dataset} diff --git a/backends/qualcomm/genai_pipeline/datasets/loaders/llm/lm_eval_adapter.py b/backends/qualcomm/genai_pipeline/datasets/loaders/llm/lm_eval_adapter.py new file mode 100644 index 00000000000..64cf6432638 --- /dev/null +++ b/backends/qualcomm/genai_pipeline/datasets/loaders/llm/lm_eval_adapter.py @@ -0,0 +1,73 @@ +# 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. + +"""LLM lm_eval dataset loader.""" + +from __future__ import annotations + +import logging +from typing import Any, Dict, List, Optional, Union + +logger = logging.getLogger(__name__) + + +class LMEvalAdapter: + """Loads an LLM dataset from lm_eval tasks. + + Purpose-agnostic: it produces a component-keyed dataset dict and knows + nothing about calibration / training / eval. The task selection is fixed at + construction; ``load_dataset`` only reads shaping options. + """ + + def __init__( + self, + tasks: Union[str, List[str]], + limit: int, + num_fewshot: Optional[int] = None, + ) -> None: + self._tasks = tasks + self._limit = limit + self._num_fewshot = num_fewshot + + def load_dataset( + self, + tokenizer: Any, + max_context_len: int, + extra_options: Optional[Dict[str, Any]] = None, + ) -> Dict[str, Any]: + """Load the dataset as ``{ARTIFACT_TEXT_DECODER: Dataset}``. + + Args: + tokenizer: TokenizerWrapper instance. + max_context_len: Sequence length the samples are encoded at. + extra_options: Reserved for adapter-specific options; unused here. + """ + from executorch.backends.qualcomm.genai_pipeline.artifact_keys import ( + ARTIFACT_TEXT_DECODER, + ) + from executorch.examples.qualcomm.oss_scripts.llama.dataset.builders import ( + DecoderDatasetBuilder, + ) + + logger.info( + "Loading LLM lm_eval dataset: tasks=%s, limit=%s", + self._tasks, + self._limit, + ) + + decoder_builder = DecoderDatasetBuilder( + tokenizer_wrapper=tokenizer, + max_context_len=max_context_len, + is_multimodal=False, + ) + dataset = decoder_builder.from_lm_eval( + tasks=self._tasks, + limit=self._limit, + num_fewshot=self._num_fewshot, + ) + + logger.info("Loaded LLM lm_eval dataset with %d samples", len(dataset)) + return {ARTIFACT_TEXT_DECODER: dataset} diff --git a/backends/qualcomm/genai_pipeline/datasets/loaders/llm/message_sample_adapter.py b/backends/qualcomm/genai_pipeline/datasets/loaders/llm/message_sample_adapter.py new file mode 100644 index 00000000000..7a3ac461c2a --- /dev/null +++ b/backends/qualcomm/genai_pipeline/datasets/loaders/llm/message_sample_adapter.py @@ -0,0 +1,66 @@ +# 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. + +"""LLM message sample dataset loader.""" + +from __future__ import annotations + +import logging +from typing import Any, Dict, List, Optional + +logger = logging.getLogger(__name__) + + +class MessageSampleAdapter: + """Loads an LLM dataset from message-sample JSON files. + + Purpose-agnostic: it produces a component-keyed dataset dict and knows + nothing about calibration / training / eval. The sample paths are fixed at + construction; ``load_dataset`` only reads shaping options. + """ + + def __init__(self, samples_paths: List[str]) -> None: + self._samples_paths = samples_paths + + def load_dataset( + self, + tokenizer: Any, + max_context_len: int, + extra_options: Optional[Dict[str, Any]] = None, + ) -> Dict[str, Any]: + """Load the dataset as ``{ARTIFACT_TEXT_DECODER: Dataset}``. + + Args: + tokenizer: TokenizerWrapper instance. + max_context_len: Sequence length the samples are encoded at. + extra_options: Reserved for adapter-specific options; unused here. + """ + from executorch.backends.qualcomm.genai_pipeline.artifact_keys import ( + ARTIFACT_TEXT_DECODER, + ) + from executorch.examples.qualcomm.oss_scripts.llama.dataset.builders import ( + DecoderDatasetBuilder, + ) + from executorch.examples.qualcomm.oss_scripts.llama.dataset.loaders import ( + load_conversation_samples, + ) + + logger.info( + "Loading LLM message sample dataset: paths=%s", + self._samples_paths, + ) + + samples = load_conversation_samples(self._samples_paths) + + decoder_builder = DecoderDatasetBuilder( + tokenizer_wrapper=tokenizer, + max_context_len=max_context_len, + is_multimodal=False, + ) + dataset = decoder_builder.from_conversation(samples) + + logger.info("Loaded LLM message sample dataset with %d samples", len(dataset)) + return {ARTIFACT_TEXT_DECODER: dataset} diff --git a/backends/qualcomm/genai_pipeline/datasets/loaders/mllm/__init__.py b/backends/qualcomm/genai_pipeline/datasets/loaders/mllm/__init__.py new file mode 100644 index 00000000000..9fa423c6294 --- /dev/null +++ b/backends/qualcomm/genai_pipeline/datasets/loaders/mllm/__init__.py @@ -0,0 +1,15 @@ +# 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. + +"""MLLM dataset loaders for multi-modality models.""" + +from executorch.backends.qualcomm.genai_pipeline.datasets.loaders.mllm.message_sample_adapter import ( + MessageSampleAdapter, +) + +__all__ = [ + "MessageSampleAdapter", +] diff --git a/backends/qualcomm/genai_pipeline/datasets/loaders/mllm/message_sample_adapter.py b/backends/qualcomm/genai_pipeline/datasets/loaders/mllm/message_sample_adapter.py new file mode 100644 index 00000000000..8066aa0572b --- /dev/null +++ b/backends/qualcomm/genai_pipeline/datasets/loaders/mllm/message_sample_adapter.py @@ -0,0 +1,104 @@ +# 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. + +"""MLLM message sample dataset loader.""" + +from __future__ import annotations + +import logging +from typing import Any, Dict, List, Optional + +logger = logging.getLogger(__name__) + + +class MessageSampleAdapter: + """Loads a multimodal dataset from message-sample JSON files. + + Purpose-agnostic: it produces a component-keyed dataset dict + (ARTIFACT_TEXT_DECODER, ARTIFACT_AUDIO_ENCODER, ARTIFACT_VISION_ENCODER) and knows nothing about + calibration / training / eval. The sample paths and model config are fixed + at construction because encoder dataset construction depends on the model's + modality config, while ``load_dataset`` only receives call-time shaping + options. + """ + + def __init__(self, samples_paths: List[str], llm_config: Any) -> None: + if llm_config is None: + raise ValueError("llm_config is required for multimodal message samples") + self._samples_paths = samples_paths + self._llm_config = llm_config + + def load_dataset( + self, + tokenizer: Any, + max_context_len: int, + extra_options: Optional[Dict[str, Any]] = None, + ) -> Dict[str, Any]: + """Load the dataset as ``{component: Dataset}`` for all modalities. + + Args: + tokenizer: TokenizerWrapper instance. + max_context_len: Sequence length the decoder samples are encoded at. + extra_options: Reserved for adapter-specific options; unused here. + """ + from executorch.backends.qualcomm.genai_pipeline.artifact_keys import ( + ARTIFACT_AUDIO_ENCODER, + ARTIFACT_TEXT_DECODER, + ARTIFACT_VISION_ENCODER, + ) + from executorch.examples.qualcomm.oss_scripts.llama.dataset.builders import ( + DecoderDatasetBuilder, + EncoderDatasetBuilder, + ) + from executorch.examples.qualcomm.oss_scripts.llama.dataset.loaders import ( + load_conversation_samples, + ) + + logger.info( + "Loading MLLM message sample dataset: paths=%s", + self._samples_paths, + ) + + samples = load_conversation_samples(self._samples_paths) + + datasets: Dict[str, Any] = {} + + decoder_builder = DecoderDatasetBuilder( + tokenizer_wrapper=tokenizer, + max_context_len=max_context_len, + is_multimodal=True, + ) + datasets[ARTIFACT_TEXT_DECODER] = decoder_builder.from_conversation(samples) + + encoder_builder = EncoderDatasetBuilder( + llm_config=self._llm_config, + tokenizer_wrapper=tokenizer, + ) + datasets[ARTIFACT_AUDIO_ENCODER] = encoder_builder.from_message_samples( + samples, ARTIFACT_AUDIO_ENCODER + ) + datasets[ARTIFACT_VISION_ENCODER] = encoder_builder.from_message_samples( + samples, ARTIFACT_VISION_ENCODER + ) + + logger.info( + "Loaded MLLM message sample dataset: " "%s=%d, %s=%s, %s=%s", + ARTIFACT_TEXT_DECODER, + len(datasets[ARTIFACT_TEXT_DECODER]), + ARTIFACT_AUDIO_ENCODER, + ( + len(datasets[ARTIFACT_AUDIO_ENCODER]) + if datasets[ARTIFACT_AUDIO_ENCODER] + else None + ), + ARTIFACT_VISION_ENCODER, + ( + len(datasets[ARTIFACT_VISION_ENCODER]) + if datasets[ARTIFACT_VISION_ENCODER] + else None + ), + ) + return datasets diff --git a/backends/qualcomm/genai_pipeline/datasets/training/__init__.py b/backends/qualcomm/genai_pipeline/datasets/training/__init__.py new file mode 100644 index 00000000000..337fc88606c --- /dev/null +++ b/backends/qualcomm/genai_pipeline/datasets/training/__init__.py @@ -0,0 +1,19 @@ +# 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. + +"""Training purpose adapters (protocol + implementation).""" + +from executorch.backends.qualcomm.genai_pipeline.datasets.training.default_training_data_adapter import ( + DefaultTrainingDataAdapter, +) +from executorch.backends.qualcomm.genai_pipeline.datasets.training.training_data_adapter import ( + TrainingDataAdapter, +) + +__all__ = [ + "TrainingDataAdapter", + "DefaultTrainingDataAdapter", +] diff --git a/backends/qualcomm/genai_pipeline/datasets/training/default_training_data_adapter.py b/backends/qualcomm/genai_pipeline/datasets/training/default_training_data_adapter.py new file mode 100644 index 00000000000..b0845378573 --- /dev/null +++ b/backends/qualcomm/genai_pipeline/datasets/training/default_training_data_adapter.py @@ -0,0 +1,114 @@ +# 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 + + +class DefaultTrainingDataAdapter: + """QAT training data provider. + + Training data is not synthesized: random tokens carry no learning signal. + Uses injected dataset loaders when present; otherwise callers enabling QAT + must supply data via ``extra_options["training_data"]``. + + The returned object is always ``{component: iterable}``. + + Args: + dataset_loaders: Purpose-agnostic loaders, each exposing + ``load_dataset`` and returning a component-keyed dataset dict. + collector: Provides per-component collators for wrapping datasets in + DataLoaders. + """ + + def __init__( + self, + dataset_loaders: Optional[List[Any]] = None, + collector: Optional[Any] = None, + ) -> None: + self._dataset_loaders = dataset_loaders or [] + self._collector = collector + + def generate_training_data( + self, + tokenizer: Any, + num_samples: int = DEFAULT_NUM_SAMPLES, + seq_length: int = DEFAULT_SEQ_LENGTH, + extra_options: Optional[Dict[str, Any]] = None, + ) -> Dict[str, Iterable[Any]]: + """Return component-keyed QAT training data.""" + from executorch.backends.qualcomm.genai_pipeline.artifact_keys import ( + ARTIFACT_TEXT_DECODER, + ) + from torch.utils.data import ConcatDataset, DataLoader + + extra_options = extra_options or {} + + training_data = extra_options.get("training_data") + if training_data is not None: + logger.info("Using caller-supplied training data for QAT") + return ( + training_data + if isinstance(training_data, dict) + else {ARTIFACT_TEXT_DECODER: training_data} + ) + + if not self._dataset_loaders: + raise ValueError( + "No training data supplied. DefaultTrainingDataAdapter does not " + "synthesize labelled data; pass it via " + "extra_options['training_data'] or inject corpus sources. For " + "PTQ, skip QAT entirely and use the calibration adapter instead." + ) + + if self._collector is None: + raise ValueError("collector is required when dataset_loaders are provided") + + max_context_len = extra_options.get("max_context_len", seq_length) + example_inputs = extra_options.get("example_inputs") + batch_size = extra_options.get("batch_size", 1) + + grouped: Dict[str, List[Any]] = {} + for loader in self._dataset_loaders: + result = loader.load_dataset( + tokenizer=tokenizer, + max_context_len=max_context_len, + extra_options=extra_options, + ) + for component, dataset in result.items(): + if dataset is not None: + grouped.setdefault(component, []).append(dataset) + + merged = { + component: (datasets[0] if len(datasets) == 1 else ConcatDataset(datasets)) + for component, datasets in grouped.items() + } + if not merged: + return {} + + collators = self._collector.create_collators(example_inputs, max_context_len) + + dataloaders: Dict[str, Iterable[Any]] = {} + for component, component_dataset in merged.items(): + component_batch_size = ( + batch_size if component == ARTIFACT_TEXT_DECODER else 1 + ) + drop_last = component == ARTIFACT_TEXT_DECODER and component_batch_size > 1 + dataloaders[component] = DataLoader( + component_dataset, + batch_size=component_batch_size, + shuffle=False, + drop_last=drop_last, + collate_fn=collators[component], + ) + return dataloaders diff --git a/backends/qualcomm/genai_pipeline/datasets/training/training_data_adapter.py b/backends/qualcomm/genai_pipeline/datasets/training/training_data_adapter.py new file mode 100644 index 00000000000..cd544de85a8 --- /dev/null +++ b/backends/qualcomm/genai_pipeline/datasets/training/training_data_adapter.py @@ -0,0 +1,31 @@ +# 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. + +"""Training data adapter protocol. + +A purpose adapter assembling QAT training data from raw dataset loaders. Its +output is always component-keyed so strategy code can consume one LLM/MLLM +contract. +""" + +from __future__ import annotations + +from typing import Any, Dict, Iterable, Optional, Protocol, runtime_checkable + + +@runtime_checkable +class TrainingDataAdapter(Protocol): + """Protocol for assembling QAT training data.""" + + def generate_training_data( + self, + tokenizer: Any, + num_samples: int = ..., + seq_length: int = ..., + extra_options: Optional[Dict[str, Any]] = None, + ) -> Dict[str, Iterable[Any]]: + """Return ``{component: iterable}`` training data.""" + ... diff --git a/backends/qualcomm/genai_pipeline/datasets/training_data_adapter.py b/backends/qualcomm/genai_pipeline/datasets/training_data_adapter.py deleted file mode 100644 index 8c505039f12..00000000000 --- a/backends/qualcomm/genai_pipeline/datasets/training_data_adapter.py +++ /dev/null @@ -1,44 +0,0 @@ -# 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 - -from typing import Any, Dict, Iterable, Optional, Protocol, runtime_checkable, Tuple - - -@runtime_checkable -class TrainingDataAdapter(Protocol): - """Protocol for quantization-aware training (QAT) dataset construction. - - Separate from :class:`CalibrationDataAdapter` because the contracts differ: - calibration yields model-argument tuples that get splatted as - ``model(*sample)``, whereas training yields ``(features, labels)`` pairs -- - labels are required to compute a loss and have no analogue in PTQ. - """ - - def generate_training_data( - self, - tokenizer: Any, - num_samples: int = 128, - seq_length: int = 1024, - extra_options: Optional[Dict[str, Any]] = None, - ) -> Iterable[Tuple[Any, Any]]: - """Build the training dataset used for quantization-aware training. - - Args: - tokenizer: The tokenizer used to encode text samples. - num_samples: Number of training samples to produce. - seq_length: Sequence length for each sample. - extra_options: Implementation-specific options. Implementations may - support a ``training_data`` key accepting caller-supplied data - directly, including a ``DataLoader``. - - Returns: - An iterable of ``(features, labels)`` pairs. ``features`` is a tuple - of model inputs; ``labels`` holds the ground truth. Mirrors the - ``qat_training_data`` contract in ``build_executorch_binary``. - """ - ... diff --git a/backends/qualcomm/genai_pipeline/genai_pipeline.py b/backends/qualcomm/genai_pipeline/genai_pipeline.py index a11bbcc42d6..7bc8160592e 100644 --- a/backends/qualcomm/genai_pipeline/genai_pipeline.py +++ b/backends/qualcomm/genai_pipeline/genai_pipeline.py @@ -96,6 +96,15 @@ class GenAIPipeline: Assembles stages from EngineProxy, wires data flow between InputConfig → Stage → OutputConfig, and executes sequentially: model_preparation → quantization → compilation → inference. + + One decoder is exported several times, once per Mode, so the orchestrator + plans ``{graph_name: Mode}`` once per run and hands it to the stages that + fan out over it. Every other per-graph fact is derived from a graph's Mode + where it is needed, rather than carried between stages. + + Calibration data is built inside the quantization stage from its injected + dataset adapters, so there is no separate dataset stage: a run that skips + quantization never touches a corpus. """ def __init__( @@ -212,7 +221,8 @@ def invoke(self, context: PipelineContext) -> InferenceOutputConfig: return result def _run_model_preparation( - self, context: PipelineContext + self, + context: PipelineContext, ) -> ModelPreparationOutputConfig: if self._model_preparation_stage is not None: logger.info("[GenAIPipeline] ModelPreparationStage started") @@ -220,6 +230,7 @@ def _run_model_preparation( input_config = ModelPreparationInputConfig( model_name=context.model_name, soc_model=context.soc_model, + extra_options=context.extra_options, ) output = self._model_preparation_stage.invoke(context, input_config) elapsed = time.monotonic() - start @@ -237,13 +248,19 @@ def _run_quantization( if self._quantization_stage is not None: logger.info("[GenAIPipeline] QuantizationStage started") start = time.monotonic() + + # The quantization stage builds its own calibration data from the + # tokenizer via its injected dataset adapters. input_config = QuantizationInputConfig( 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. + # Export inputs come from the model, not from the dataset. example_inputs=model_prep_output.example_inputs, - calibration_data=model_prep_output.calibration_data, + tokenizer=model_prep_output.tokenizer, + meta=model_prep_output.meta, + inference=model_prep_output.inference, + extra_options=context.extra_options, ) output = self._quantization_stage.invoke(context, input_config) elapsed = time.monotonic() - start @@ -262,7 +279,7 @@ def _run_compilation( start = time.monotonic() # Fall back to the prepared module when quantization was skipped # (FP16 / GPU flows), so the model isn't silently dropped. - model = quant_output.quantized_model or model_prep_output.model_module + model = quant_output.graphs or model_prep_output.model_module input_config = CompilationInputConfig( # Note: context.soc_model is str; str→QcomChipset conversion is # deferred to the strategy/adapter layer (see PR5). diff --git a/backends/qualcomm/genai_pipeline/model_components/__init__.py b/backends/qualcomm/genai_pipeline/model_components/__init__.py new file mode 100644 index 00000000000..85b2cd68483 --- /dev/null +++ b/backends/qualcomm/genai_pipeline/model_components/__init__.py @@ -0,0 +1,32 @@ +# 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. + +"""GenAI Pipeline model components. + +Re-exports the exportable ``nn.Module`` building blocks from the existing llama +module, providing a stable import path within the genai_pipeline namespace: + +.. code-block:: text + + model_components.decoder <- llama.model.static_llama, .layernorm, + .feed_forward, .apply_rope + model_components.encoders <- vision_encoder, audio_encoder + model_components.embedding <- token embedding + +This package intentionally re-exports nothing itself. ``decoder`` pulls in +``static_llama`` and ``encoders`` pulls in the transformers vision/audio modeling +code, both of which are expensive to import; flattening them here would make +every consumer of any component pay for all of them. Import the submodule you +need instead: + +Usage: + from executorch.backends.qualcomm.genai_pipeline.model_components.decoder import ( + LlamaModel, + ) + from executorch.backends.qualcomm.genai_pipeline.model_components.embedding import ( + TokenEmbedding, + ) +""" diff --git a/backends/qualcomm/genai_pipeline/model_components/decoder/__init__.py b/backends/qualcomm/genai_pipeline/model_components/decoder/__init__.py new file mode 100644 index 00000000000..ce3226107ed --- /dev/null +++ b/backends/qualcomm/genai_pipeline/model_components/decoder/__init__.py @@ -0,0 +1,86 @@ +# 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. + +"""Transformer decoder architectures and their architectural primitives. + +The decoder variants differ in their input representation or attention pattern: + +* :class:`LlamaModel` consumes token IDs with a standard causal decoder. +* :class:`LlamaModelWithoutEmbedding` consumes precomputed token embeddings. +* :class:`MultiScopeAwareLlamaModel` interleaves local and global attention. + +The registries parameterize the shared decoder with architecture-specific +normalization, feed-forward, and rotary-position-embedding implementations. +This captures model-family variation without duplicating the decoder stack. + +""" + + +from executorch.examples.qualcomm.oss_scripts.llama.model.apply_rope import ( + apply_partial_rotary_emb_single, + apply_rotary_emb_single, + register_rotary_emb, + ROTARY_EMB_REGISTRY, +) +from executorch.examples.qualcomm.oss_scripts.llama.model.feed_forward import ( + CodegenFeedForward, + FeedForward_REGISTRY, + FeedForwardBase, + GLMFeedForward, + register_feed_forward, +) +from executorch.examples.qualcomm.oss_scripts.llama.model.layernorm import ( + LayerNorm, + Norm, + NORM_REGISTRY, + register_norm, + RMSNorm, +) +from executorch.examples.qualcomm.oss_scripts.llama.model.static_llama import ( + AttentionSinkRope, + FeedForward, + LlamaAttention, + LlamaDecoderLayer, + LlamaModel, + LlamaModelWithoutEmbedding, + MultiScopeAwareLlamaModel, + repeat_kv, +) + +__all__ = [ + # Decoder architectures. + "LlamaModel", + "LlamaModelWithoutEmbedding", + "MultiScopeAwareLlamaModel", + # Attention masks. + "AttentionMask", + "BaseAttentionMask", + "CausalAttentionMask", + "SlidingWindowAttentionMask", + # Attention, KV-head expansion, and rotary positional encoding. + "AttentionSinkRope", + "LlamaAttention", + "repeat_kv", + "apply_partial_rotary_emb_single", + "apply_rotary_emb_single", + "register_rotary_emb", + "ROTARY_EMB_REGISTRY", + # Feed-forward layers. + "FeedForward", + "CodegenFeedForward", + "FeedForward_REGISTRY", + "FeedForwardBase", + "GLMFeedForward", + "register_feed_forward", + # Normalization layers. + "LayerNorm", + "Norm", + "NORM_REGISTRY", + "register_norm", + "RMSNorm", + # A decoder layers. + "LlamaDecoderLayer", +] diff --git a/backends/qualcomm/genai_pipeline/model_components/embedding.py b/backends/qualcomm/genai_pipeline/model_components/embedding.py new file mode 100644 index 00000000000..21041f871e9 --- /dev/null +++ b/backends/qualcomm/genai_pipeline/model_components/embedding.py @@ -0,0 +1,25 @@ +# 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. + +"""GenAI Pipeline token embedding component. + +Re-exports the token embedding module from the existing llama module, providing +a stable import path within the genai_pipeline namespace. Models that export +their embedding table as a separate graph construct it from here. + +Usage: + from executorch.backends.qualcomm.genai_pipeline.model_components.embedding import ( + TokenEmbedding, + ) +""" + +from executorch.examples.qualcomm.oss_scripts.llama.model.embedding import ( + TokenEmbedding, +) + +__all__ = [ + "TokenEmbedding", +] diff --git a/backends/qualcomm/genai_pipeline/model_components/encoders/__init__.py b/backends/qualcomm/genai_pipeline/model_components/encoders/__init__.py new file mode 100644 index 00000000000..275401c557f --- /dev/null +++ b/backends/qualcomm/genai_pipeline/model_components/encoders/__init__.py @@ -0,0 +1,38 @@ +# 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. + +"""GenAI Pipeline multimodal encoder components. + +Re-exports the vision and audio encoders from the existing llama module, +providing a stable import path within the genai_pipeline namespace. Each class +is the encoder graph of one modality, paired with a text decoder by +``llama.encoder.encoder_config``. + +The ``Custom*`` classes those encoders swap in via +``replace_module_with_custom_class`` are export workarounds internal to their own +modules, and are deliberately not re-exported here. + +Usage: + from executorch.backends.qualcomm.genai_pipeline.model_components.encoders import ( + Idefics3VisionEncoder, + ) +""" + +from executorch.examples.qualcomm.oss_scripts.llama.model.audio_encoder import ( + GraniteSpeechCTCEncoderWrapper, +) +from executorch.examples.qualcomm.oss_scripts.llama.model.vision_encoder import ( + Idefics3VisionEncoder, + InternVL3VisionEncoder, +) + +__all__ = [ + # audio + "GraniteSpeechCTCEncoderWrapper", + # vision + "Idefics3VisionEncoder", + "InternVL3VisionEncoder", +] diff --git a/backends/qualcomm/genai_pipeline/model_lookup.py b/backends/qualcomm/genai_pipeline/model_lookup.py new file mode 100644 index 00000000000..09cfd23c68b --- /dev/null +++ b/backends/qualcomm/genai_pipeline/model_lookup.py @@ -0,0 +1,675 @@ +# 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. + +"""Model lookup utilities for the GenAI Pipeline. + +This module is the registry-facing entry point for model-specific pipeline +configuration. It keeps user-facing model names separate from the concrete +objects and options required by each pipeline stage. + +Most callers should resolve a model config once with :func:`get_model_config` +and pass it to the more specific lookup helpers when available, so all derived +configuration comes from the same registry entry. +""" + +from __future__ import annotations + +import json +import logging +from functools import partial +from typing import Any, Callable, Dict, List, Optional, Tuple + +from executorch.backends.qualcomm.genai_pipeline.artifact_keys import ( + ARTIFACT_AUDIO_ENCODER, + ARTIFACT_TEXT_DECODER, + ARTIFACT_TOK_EMBEDDING, + ARTIFACT_VISION_ENCODER, +) +from executorch.backends.qualcomm.genai_pipeline.graph_names import ( + DECODER_GRAPH_NAMES, + GRAPH_FORWARD, + TOK_EMBEDDING_GRAPH_NAMES, +) +from executorch.backends.qualcomm.genai_pipeline.model_components.decoder import ( + LlamaModel, +) +from executorch.backends.qualcomm.genai_pipeline.model_components.embedding import ( + TokenEmbedding, +) +from executorch.backends.qualcomm.genai_pipeline.models import LLM_VARIANT_ARCHS +from executorch.examples.qualcomm.oss_scripts.llama.model.static_llama import ModelArgs +from executorch.examples.qualcomm.oss_scripts.llama.wrappers.base_component import ( + get_model_specific_kwargs, + Mode, + next_power_of_two, + process_model_args, +) + +logger = logging.getLogger(__name__) + + +def get_supported_models() -> List[str]: + """Get list of all supported model names. + + Returns: + Sorted list of supported model name strings. + """ + from executorch.backends.qualcomm.genai_pipeline.models import SUPPORTED_LLM_MODELS + + return sorted(SUPPORTED_LLM_MODELS.keys()) + + +def get_model_config(model_name: str) -> Any: + """Look up a model configuration by name. + Args: + model_name: Model identifier (e.g., "llama3_2-1b_instruct", "qwen2_5-0_5b"). + Returns: + The LLMModelConfig instance for the given model. + Raises: + KeyError: If model_name is not found in the registry. + """ + from executorch.backends.qualcomm.genai_pipeline.models import SUPPORTED_LLM_MODELS + + model_name_lower = model_name.lower() + if model_name_lower not in SUPPORTED_LLM_MODELS: + available = get_supported_models() + raise KeyError( + f"Model '{model_name}' not found. " f"Available models: {available}" + ) + + config = SUPPORTED_LLM_MODELS[model_name_lower] + + return config + + +def get_model_arch_config( + model_name: str, control_args: Any, model_config: Optional[Any] = None +) -> Any: + """Resolve the decoder's config, for reading model shapes. + + All graphs (calibration / decode / prefill) share one set of weights and + differ only in per-mode fields (``ar_len``, ``max_batch_size``, + ``use_kv_cache``); the shapes source transforms read do not vary by graph. + This returns a single config -- the decode graph's -- which + :func:`get_source_transform` reads those shapes off. + + Args: + model_name: The model's key in ``SUPPORTED_LLM_MODELS``. + control_args: Carries paths and lengths. + + Returns: + One decoder config (``ModelArgs``, or ``Gemma4Config`` for gemma4). + + Raises: + ValueError: If the model needs a params file that ``control_args`` does + not name. + """ + if not model_config: + model_config = get_model_config(model_name) + + if model_name == "gemma4-e2b": + from executorch.examples.models.gemma4.text_decoder.gemma4_config import ( + Gemma4Config, + ) + + gemma4_config = Gemma4Config.from_json(model_config.params_path) + gemma4_config.use_kv_cache = True + gemma4_config.max_batch_size = 1 + gemma4_config.max_seq_len = control_args.max_seq_len + gemma4_config.max_context_len = control_args.max_context_len + return gemma4_config + + params_path = ( + model_config.params_path if control_args.params is None else control_args.params + ) + if params_path is None: + raise ValueError( + f"Model '{model_name}' carries no params file, so its shapes " + "cannot be resolved; pass one as control_args.params" + ) + # The recipe is held as a CLASS, while ``process_model_args`` reads + # ``get_kv_io_bit_width()`` off an instance. + quant_recipe = model_config.quant_recipe + if isinstance(quant_recipe, type): + quant_recipe = quant_recipe() + with open(params_path) as f: + base_args = json.load(f) + + return process_model_args( + control_args, + ModelArgs(**base_args), + quant_recipe, + model_config, + Mode.DECODE, + ) + + +def get_model_arch( + model_name: str, control_args: Any, model_config: Optional[Any] = None +) -> Dict[str, Dict[str, Callable]]: + """Build the model constructors exported for a model. + + Each constructor has its model config and graph-specific arguments + pre-bound, allowing the loader to instantiate it without handling model + modes or architecture details. + + The text decoder component is always included. The token embedding + component is included only for multimodal models. + + Args: + model_name: The model key in ``SUPPORTED_LLM_MODELS``. + control_args: Runtime arguments controlling graph modes, shapes. + model_config: Optional pre-resolved model configuration. + + Returns: + A ``{component_name: {graph_name: constructor}}`` mapping. + """ + + # Get model config + if not model_config: + model_config = get_model_config(model_name) + + # Decode and calibration graphs are always exported. Prefill is additionally + # exported for hybrid and lookahead modes. + decode_graph_name, prefill_graph_name = DECODER_GRAPH_NAMES + embedding_decode_name, embedding_prefill_name = TOK_EMBEDDING_GRAPH_NAMES + graph_modes = [ + (decode_graph_name, embedding_decode_name, Mode.DECODE), + (GRAPH_FORWARD, GRAPH_FORWARD, Mode.CALIBRATE), + *( + [(prefill_graph_name, embedding_prefill_name, Mode.PREFILL)] + if control_args.model_mode in {"hybrid", "lookahead"} + else [] + ), + ] + model_arch = { + ARTIFACT_TEXT_DECODER: {}, + **({ARTIFACT_TOK_EMBEDDING: {}} if is_multimodal(model_name) else {}), + } + + # Get any model-specific kwargs and quantization recipe to get the kv IO bit width. + use_i64_token = control_args.embedding_quantize is not None + model_specific_kwargs = get_model_specific_kwargs(control_args, model_config) + quant_recipe = model_config.quant_recipe + + if model_name == "gemma4-e2b": + from executorch.examples.models.gemma4.text_decoder.gemma4_config import ( + Gemma4Config, + ) + from executorch.examples.qualcomm.oss_scripts.gemma4.model_wrapper import ( + Gemma4TextModelWrapper, + ) + + for decoder_graph_name, embedding_graph_name, mode in graph_modes: + config = Gemma4Config.from_json(model_config.params_path) + config.use_kv_cache = True + config.max_batch_size = ( + control_args.batch_size if mode == Mode.CALIBRATE else 1 + ) + config.max_seq_len = control_args.max_seq_len + config.max_context_len = control_args.max_context_len + + # Gemma 4 reads ar_len off the constructor, not the config. + if mode == Mode.CALIBRATE: + ar_len = control_args.max_context_len + elif mode == Mode.PREFILL: + ar_len = control_args.prefill_ar_len + elif control_args.model_mode == "lookahead": + ar_len = next_power_of_two( + (control_args.window + control_args.gcap) * (control_args.ngram - 1) + ) + else: + ar_len = 1 + + extra_kwargs = { + # 32 is the sentinel for "unquantized KV IO"; get_kv_io_bit_width() + # returns it too when the recipe has no default_quant_dtype. + "kv_io_bit_width": ( + quant_recipe().get_kv_io_bit_width() if quant_recipe else 32 + ), + } + + # Get Text Decoder model architecture. + model_arch[ARTIFACT_TEXT_DECODER][decoder_graph_name] = partial( + Gemma4TextModelWrapper, + config, + ar_len=ar_len, + output_new_cache_only=True, + output_cache=True, + use_i64_token=use_i64_token, + enable_masked_softmax=False, + **extra_kwargs, + ) + + # Get Token Embedding model architecture if the model is multimodal LLM. + if is_multimodal(model_name): + model_arch[ARTIFACT_TOK_EMBEDDING][embedding_graph_name] = partial( + TokenEmbedding, + max_batch_size=config.max_batch_size, + ar_len=ar_len, + vocab_size=config.vocab_size, + dim=config.dim, + use_i64_token=use_i64_token, + ) + else: + params_path = ( + model_config.params_path + if control_args.params is None + else control_args.params + ) + if params_path is None: + raise ValueError( + f"Model '{model_name}' carries no params file, so its shapes " + "cannot be resolved; pass one as control_args.params" + ) + + with open(params_path) as f: + base_args = json.load(f) + + for decoder_graph_name, embedding_graph_name, mode in graph_modes: + config = process_model_args( + control_args, + ModelArgs(**base_args), + quant_recipe(mode == Mode.CALIBRATE), + model_config, + mode, + ) + # TODO: Decouple example-input generation from the model because + # input shapes are graph-specific rather than model-specific. + model_arch[ARTIFACT_TEXT_DECODER][decoder_graph_name] = partial( + LLM_VARIANT_ARCHS.get(model_name, LlamaModel), + config, + ar_len=config.ar_len, + output_new_cache_only=True, + output_cache=True, + use_i64_token=use_i64_token, + **model_specific_kwargs, + ) + + # Get Token Embedding model architecture if the model is multimodal LLM. + if is_multimodal(model_name): + model_arch[ARTIFACT_TOK_EMBEDDING][embedding_graph_name] = partial( + TokenEmbedding, + max_batch_size=config.max_batch_size, + ar_len=config.ar_len, + vocab_size=config.vocab_size, + dim=config.dim, + use_i64_token=use_i64_token, + ) + + return model_arch + + +def get_source_transform( + model_name: str, + *, + control_args: Any, + model_config: Optional[Any] = None, +) -> Tuple[Dict[str, List[Callable]], Dict[str, List[Callable]]]: + """Resolve a model's source transforms, mirroring ``LLMWrapper._prepare_model``. + + Args: + model_name: The model's key in ``SUPPORTED_LLM_MODELS``. + control_args: CLI arguments used to determine the checkpoint source and + resolve the values bound to individual transforms. + model_config: Optional resolved model configuration. When omitted, it + is looked up from ``model_name``. + + Returns: + A ``(weight_transforms, module_transforms)`` tuple of component-keyed + transform maps: + + - ``weight_transforms``: Each callable accepts a state dict and returns + the transformed state dict before model weights are loaded. + - ``module_transforms``: Each callable accepts a module and returns the + transformed module after model weights are loaded. + + Components without transforms are omitted from each map. + + Raises: + RuntimeError: If the model config requests SpinQuant (``r1``/``r2``), + which is no longer supported -- matching the reference's guard. + """ + from executorch.backends.qualcomm.genai_pipeline.source_transform import ( + apply_dtype_override, + convert_linear_to_conv2d, + gemma_rmsnorm_offset, + permute_partial_rope, + prepare_conv_submodules, + remap_gemma4_keys, + scale_token_embedding, + strip_orig_mod_prefix, + unwrap_model_key, + ) + + if model_config is None: + model_config = get_model_config(model_name) + config = model_config + name = model_name.lower() + is_hf_path = control_args.checkpoint is None + model_args = get_model_arch_config(model_name, control_args, config) + + if config.r1 or config.r2: + raise RuntimeError( + "SpinQuant (r1/r2) is no longer supported: the " + "torchao.prototype.spinquant module has been deleted." + ) + + # Weight chain, in ``_prepare_model`` order. + weight_transforms = [unwrap_model_key] + if is_hf_path: + # HF path: gemma4 key rename, then the Gemma RMSNorm +1 offset, then the + # embedding scale (self-guarding on the factor / key). + if name == "gemma4-e2b": + weight_transforms.append(remap_gemma4_keys) + else: + if name in ("gemma-2b", "gemma2-2b", "gemma3-1b"): + weight_transforms.append(gemma_rmsnorm_offset) + weight_transforms.append( + partial( + scale_token_embedding, + embedding_scale_factor=model_args.embedding_scale_factor, + ) + ) + else: + # Local checkpoint path: stories260k carries torch.compile's + # ``_orig_mod.`` prefix and must be renamed. + if name == "stories260k": + weight_transforms.append(strip_orig_mod_prefix) + # RoPE weight layout permutation, gated on the model config flag. + if config.transform_weight: + weight_transforms.append( + partial( + permute_partial_rope, + n_layers=model_args.n_layers, + n_heads=model_args.n_heads, + n_kv_heads=model_args.n_kv_heads, + partial_rotary_factor=model_args.partial_rotary_factor, + ) + ) + + # Module chain, in ``_prepare_model`` order: submodule conv prep, the + # linear-to-conv2d rewrite it feeds, then the dtype override last so it also + # covers the conv2d modules. + module_transforms = [ + prepare_conv_submodules, + convert_linear_to_conv2d, + partial( + apply_dtype_override, + dtype_override=control_args.dtype_override, + ), + ] + + return {ARTIFACT_TEXT_DECODER: weight_transforms}, { + ARTIFACT_TEXT_DECODER: module_transforms + } + + +def load_hf_checkpoint_state_dict( + repo_id: str, + *, + convert_weights: Callable, +) -> Dict[str, Any]: + """Download a standard HF checkpoint and return its converted state dict.""" + import torch + from executorch.examples.models.llama.hf_download import ( + download_and_convert_hf_checkpoint, + ) + + checkpoint = download_and_convert_hf_checkpoint(repo_id, convert_weights) + return torch.load(checkpoint, weights_only=True, map_location="cpu", mmap=True) + + +def load_gemma4_hf_checkpoint_state_dict( + repo_id: str, + *, + convert_weights: Callable, + gemma4_config: Any, + dtype: Any, +) -> Dict[str, Any]: + """Download Gemma4 HF weights and return the converter-produced state dict.""" + from huggingface_hub import snapshot_download + + return convert_weights( + snapshot_download(repo_id=repo_id), + gemma4_config, + dtype, + ) + + +def get_state_dict_loader( + model_name: str, + *, + control_args: Any, + model_config: Optional[Any] = None, +) -> Dict[str, Callable[[str], Dict[str, Any]]]: + """Return component-keyed loaders for remote Hugging Face checkpoints. + + Args: + model_name: Registered model identifier used to select the loader. + control_args: Runtime arguments. A configured local ``checkpoint`` + bypasses remote loading and returns no loader. + model_config: Optional resolved model configuration. When omitted, it + is looked up from ``model_name``. + + Returns: + A component-keyed loader map. Each loader accepts a Hugging Face + ``repo_id`` and returns a converted state dict: + + - When ``control_args.checkpoint`` is set: An empty map, because local + checkpoint loading does not need a remote loader. + - Otherwise: A map containing the text-decoder loader. The loader + has model-specific conversion details bound into the callable. + + Raises: + KeyError: If ``model_name`` is not registered and ``model_config`` is + not supplied. + """ + import torch + + if model_config is None: + model_config = get_model_config(model_name) + if control_args.checkpoint is not None: + return {} + + convert_weights = model_config.convert_weights + if hasattr(convert_weights, "__func__"): + convert_weights = convert_weights.__func__ + + name = model_name.lower() + if name == "gemma4-e2b": + return { + ARTIFACT_TEXT_DECODER: partial( + load_gemma4_hf_checkpoint_state_dict, + convert_weights=convert_weights, + gemma4_config=get_model_arch_config( + model_name, control_args, model_config + ), + dtype=torch.float32, + ) + } + + return { + ARTIFACT_TEXT_DECODER: partial( + load_hf_checkpoint_state_dict, + convert_weights=convert_weights, + ) + } + + +def get_model_num_sharding(model_name: str) -> Dict[str, int]: + """Get the number of graph shards for each component of a model. + + Args: + model_name: Model identifier. + + Returns: + ``{component_name: num_shards}`` (1 = no sharding). The text decoder + is always present; vision/audio encoders are present only if the + model declares them. + + Raises: + KeyError: If model_name is not found. + """ + config = get_model_config(model_name) + num_sharding = {ARTIFACT_TEXT_DECODER: getattr(config, "num_sharding", 1)} + if hasattr(config, "vision_encoder"): + num_sharding[ARTIFACT_VISION_ENCODER] = config.vision_encoder.num_sharding + if hasattr(config, "audio_encoder"): + num_sharding[ARTIFACT_AUDIO_ENCODER] = config.audio_encoder.num_sharding + return num_sharding + + +def get_quant_dtype(model_name: str) -> Any: + """Get component-keyed quantization dtypes for the model. + + Components without an explicit model configuration default to the higher + precision ``QuantDtype.use_16a8w`` (16A8W). + + Returns: + Component-keyed quantization dtypes. + + Raises: + KeyError: If model_name is not found. + """ + from executorch.backends.qualcomm.quantizer.quantizer import QuantDtype + + config = get_model_config(model_name) + quant_dtype = { + ARTIFACT_TEXT_DECODER: getattr(config, "quant_dtype", QuantDtype.use_16a8w) + } + if hasattr(config, "vision_encoder"): + quant_dtype[ARTIFACT_TOK_EMBEDDING] = QuantDtype.use_16a8w + quant_dtype[ARTIFACT_VISION_ENCODER] = getattr( + config.vision_encoder, "quant_dtype", QuantDtype.use_16a8w + ) + if hasattr(config, "audio_encoder"): + quant_dtype[ARTIFACT_TOK_EMBEDDING] = QuantDtype.use_16a8w + quant_dtype[ARTIFACT_AUDIO_ENCODER] = getattr( + config.audio_encoder, "quant_dtype", QuantDtype.use_16a8w + ) + + return quant_dtype + + +def get_quant_recipe(model_name: str) -> Any: + """Get the quantization recipe class for the model. + + Returns: + Component-keyed quantization recipe classes. + """ + + config = get_model_config(model_name) + + quant_recipe = {ARTIFACT_TEXT_DECODER: getattr(config, "quant_recipe", None)} + if hasattr(config, "vision_encoder"): + quant_recipe[ARTIFACT_TOK_EMBEDDING] = None + quant_recipe[ARTIFACT_VISION_ENCODER] = getattr( + config.vision_encoder, "quant_recipe", None + ) + if hasattr(config, "audio_encoder"): + quant_recipe[ARTIFACT_TOK_EMBEDDING] = None + quant_recipe[ARTIFACT_AUDIO_ENCODER] = getattr( + config.audio_encoder, "quant_recipe", None + ) + return quant_recipe + + +def is_multimodal(model_name: str) -> bool: + """Check if a model is multimodal (has vision/audio encoders). + + Args: + model_name: Model identifier. + + Returns: + True if the model has vision or audio encoder capabilities. + + Raises: + KeyError: If model_name is not found. + """ + config = get_model_config(model_name) + return hasattr(config, "audio_encoder") or hasattr(config, "vision_encoder") + + +def get_model_loader_adapter( + model_name: str, + control_args: Any, +) -> Any: + """Get the model loader adapter for ``model_name``. + + The returned adapter is selected from the model registry: multimodal models + use ``MLLMLoaderAdapter`` and text-only models use ``LLMLoaderAdapter``. + Model architecture and source transforms are resolved separately and passed + through ``ModelPreparationInputConfig.extra_options``. + + Args: + model_name: Model identifier (e.g., "llama3_2-1b_instruct"). + control_args: CLI namespace carrying paths and shape controls. + + Returns: + The loader adapter for the model. + + Raises: + KeyError: If model_name is not found in the registry. + """ + config = get_model_config(model_name) + + if is_multimodal(model_name): + from executorch.backends.qualcomm.genai_pipeline.strategies.model_preparation.mllm_loader_adapter import ( + MLLMLoaderAdapter, + ) + + adapter_cls = MLLMLoaderAdapter + logger.debug("Created multimodal model loader adapter for '%s'", model_name) + else: + from executorch.backends.qualcomm.genai_pipeline.strategies.model_preparation.llm_loader_adapter import ( + LLMLoaderAdapter, + ) + + adapter_cls = LLMLoaderAdapter + logger.debug("Created static LLM model loader adapter for '%s'", model_name) + + return adapter_cls.from_model_config( + config, + control_args=control_args, + ) + + +def get_quantizer_adapter(model_name: str) -> Any: + """Create a quantizer adapter for the given model. + + For LLM models, returns LLMQuantizerAdapter. + For multimodal models (VLM/ALM), returns MLLMQuantizerAdapter. + + Recipes and dtypes are routed by ``ExecuTorchQuantizationStrategy`` from the + component-keyed ``quantize_options`` dict. + + Args: + model_name: Model identifier (e.g., "llama3_2-1b_instruct", "smolvlm_500m"). + + Returns: + LLMQuantizerAdapter for LLM models, MLLMQuantizerAdapter for multimodal models. + + Raises: + KeyError: If model_name is not found in the registry. + """ + if is_multimodal(model_name): + from executorch.backends.qualcomm.genai_pipeline.strategies.quantization.mllm_quantizer_adapter import ( + MLLMQuantizerAdapter, + ) + + adapter = MLLMQuantizerAdapter() + logger.debug("Created MLLMQuantizerAdapter for '%s'", model_name) + return adapter + else: + from executorch.backends.qualcomm.genai_pipeline.strategies.quantization.llm_quantizer_adapter import ( + LLMQuantizerAdapter, + ) + + adapter = LLMQuantizerAdapter() + + logger.debug("Created LLMQuantizerAdapter for '%s'", model_name) + return adapter diff --git a/backends/qualcomm/genai_pipeline/models/__init__.py b/backends/qualcomm/genai_pipeline/models/__init__.py new file mode 100644 index 00000000000..ac31be07127 --- /dev/null +++ b/backends/qualcomm/genai_pipeline/models/__init__.py @@ -0,0 +1,34 @@ +# 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. + +"""GenAI Pipeline model registry. + +Re-exports LLM model configurations from the existing llama module, providing a +stable import path within the genai_pipeline namespace. + +Which transforms each model uses is owned by +``model_lookup.get_source_transform``, not declared here. + +Usage: + from executorch.backends.qualcomm.genai_pipeline.models import ( + LLMModelConfig, + SUPPORTED_LLM_MODELS, + ) +""" + +from executorch.examples.qualcomm.oss_scripts.llama import ( + LLM_VARIANT_ARCHS, + LLMModelConfig, + register_llm_model, + SUPPORTED_LLM_MODELS, +) + +__all__ = [ + "LLMModelConfig", + "LLM_VARIANT_ARCHS", + "register_llm_model", + "SUPPORTED_LLM_MODELS", +] diff --git a/backends/qualcomm/genai_pipeline/quant_utilities.py b/backends/qualcomm/genai_pipeline/quant_utilities.py new file mode 100644 index 00000000000..d81ceff016e --- /dev/null +++ b/backends/qualcomm/genai_pipeline/quant_utilities.py @@ -0,0 +1,355 @@ +# 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 logging +from pathlib import Path +from typing import Any, Dict, Optional + +import torch + +logger = logging.getLogger(__name__) + + +def is_node_src_start_with_name(node: torch.fx.Node, kv_cache_prefix: str) -> bool: + """ + Return True if any NodeSource in node.meta['from_node'] has a name + starting with `kv_cache_prefix`. Used to identify K/V cache nodes by their + "k_" or "v_" name prefix in the traced graph. + """ + + def has_source_name_prefix( + node_src: torch.fx.traceback.NodeSource, kv_cache_prefix: str + ) -> bool: + + name = getattr(node_src, "name", None) + if isinstance(name, str) and name.startswith(kv_cache_prefix): + return True + + children = getattr(node_src, "from_node", None) + if not children: + return False + + for src in children: + if has_source_name_prefix(src, kv_cache_prefix): + return True + + return False + + node_srcs = node.meta.get("from_node", None) + if not node_srcs: + return False + + return any( + has_source_name_prefix(node_src, kv_cache_prefix) for node_src in node_srcs + ) + + +def make_quantizer( + quant_dtype: Any = None, + backend: Any = None, + soc_model: Any = None, + quant_recipe: Any = None, + **kwargs: Any, +) -> Any: + """Create and configure a QNN quantizer.""" + from executorch.backends.qualcomm.export_utils import ( + make_quantizer as _make_quantizer, + ) + + soc_model_str = soc_model.name if hasattr(soc_model, "name") else str(soc_model) + 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: + recipe_config = ( + quant_recipe.recipe if hasattr(quant_recipe, "recipe") else quant_recipe + ) + quantizer.set_recipe(recipe_config) + return quantizer + + +def save_quantized_module( + quantized_module: Any, + example_inputs: Any, + artifact_dir: str, +) -> Path: + """Export the calibration graph used for decode SQNR evaluation.""" + from executorch.backends.qualcomm.genai_pipeline.artifact_keys import ( + DECODE_QDQ_FILENAME, + ) + + qdq_ep_path = Path(artifact_dir) / DECODE_QDQ_FILENAME + qdq_ep_path.parent.mkdir(parents=True, exist_ok=True) + qdq_ep = torch.export.export(quantized_module, example_inputs, strict=True) + torch.export.save(qdq_ep, qdq_ep_path) + logger.info("QDQ EP saved to %s", qdq_ep_path) + return qdq_ep_path + + +def _logits_io_shape(meta: Dict[str, Any]) -> set: + """The logits output shape, derived from the graph's metadata.""" + return { + # logit output + ( + meta["get_max_batch_size"], + meta["get_ar_len"], + meta["get_vocab_size"], + ), + } + + +def _kv_cache_shape(meta: Dict[str, Any]) -> set: + """The set of valid K/V cache shapes (last two dims), from the metadata. + + A model whose metadata carries ``get_global_head_dim`` has a per-layer + head_dim (e.g. Gemma 4: sliding=256, full=512), so every head_dim variant + contributes its own input/output shapes; otherwise a single head_dim is used. + """ + if "get_global_head_dim" in meta: + # Gemma 4 has per-layer head_dim: sliding=256, full=512 + kv_head_dims = { + meta["get_head_dim"], + meta["get_global_head_dim"], + } + kv_cache_shape = set() + for head_dim in kv_head_dims: + kv_cache_shape.add((head_dim, meta["get_max_context_len"])) + kv_cache_shape.add((meta["get_max_context_len"], head_dim)) + kv_cache_shape.add((head_dim, meta["get_ar_len"])) + kv_cache_shape.add((meta["get_ar_len"], head_dim)) + return kv_cache_shape + return { + # single head, kv input + (meta["get_head_dim"], meta["get_max_context_len"]), + (meta["get_max_context_len"], meta["get_head_dim"]), + # single head, kv output + (meta["get_head_dim"], meta["get_ar_len"]), + (meta["get_ar_len"], meta["get_head_dim"]), + } + + +def save_logits_quant_attrs( + graph_module: torch.fx.GraphModule, meta: Dict[str, Any] +) -> None: + """Record the quantized logits scale/zero-point into ``meta`` (in place). + + Scans the graph output for the ``dequantize_per_tensor`` node whose source + tensor matches the logits shape, and writes ``get_logits_scale`` / + ``get_logits_zero_point``. + """ + io_shape = _logits_io_shape(meta) + for node in graph_module.graph.nodes: + if node.op == "output": + for output_node in node.args[0]: + if ( + output_node.target + == torch.ops.quantized_decomposed.dequantize_per_tensor.default + ): + source_node = output_node.args[0].args[0] + if source_node.meta["val"].size() in io_shape: + meta["get_logits_scale"] = output_node.args[1] + meta["get_logits_zero_point"] = output_node.args[2] + break + + +def save_output_kv_cache_quant_attrs( + graph_module: torch.fx.GraphModule, meta: Dict[str, Any] +) -> None: + """Record each output KV-cache node's quant attributes into ``meta``. + + For every graph output whose last two dims match a KV-cache shape, writes + ``get_kv_output_{i}_quant_attr`` = ``[scale, zero_point, quant_min, + quant_max, dtype]`` (used by the attention-sink feature). + """ + from executorch.backends.qualcomm.builders.utils import is_graph_output + + kv_cache_shape = _kv_cache_shape(meta) + kv_idx = 0 + for node in graph_module.graph.nodes: + if not is_graph_output(node): + continue + cache_output_node = node.args[0].args[0] + if cache_output_node.meta["val"].size()[-2:] in kv_cache_shape: + # [QCOM_SCALE, QCOM_ZERO_POINT, QCOM_QUANT_MIN, QCOM_QUANT_MAX, QCOM_DTYPE] + meta[f"get_kv_output_{kv_idx}_quant_attr"] = [ + node.args[1], + node.args[2], + node.args[3], + node.args[4], + str(node.args[5]), + ] + kv_idx += 1 + + +def encoding_override( # noqa: C901 + quantized_model: torch.fx.GraphModule, + unquantized_model: torch.fx.GraphModule, + n_cache_layers: Optional[int] = None, +) -> None: + """Copy calibration encodings to a deployed graph. + + Activation and parameter encodings are always copied. Supplying + ``n_cache_layers`` additionally copies KV-cache output encodings onto the + deployed graph's cache inputs; ``None`` leaves KV-cache encodings unchanged. + """ + from executorch.backends.qualcomm.builders.utils import is_graph_output + + pbq_target = { + torch.ops.torchao.dequantize_affine, + torch.ops.torchao.quantize_affine, + } + pcq_target = { + torch.ops.quantized_decomposed.dequantize_per_channel.default, + torch.ops.quantized_decomposed.quantize_per_channel.default, + } + ptq_target = { + torch.ops.quantized_decomposed.dequantize_per_tensor.default, + torch.ops.quantized_decomposed.quantize_per_tensor.default, + } + qdq_target = pbq_target | pcq_target | ptq_target + + def compare_nodes(quantized_node, unquantized_node): + def info(node): + return node.name + ( + str(node.meta["nn_module_stack"].values()) + if node.op == "call_function" + else "" + ) + + assert info(quantized_node) == info( + unquantized_node + ), f"found unmatched order for ops: {quantized_node} vs {unquantized_node}" + + def resolve_param_target(node): + return ( + node + if node.op == "call_function" and node.target not in qdq_target + else resolve_param_target(list(node.users)[0]) + ) + + def activation_override(quantized_node, unquantized_node): + for quantized_user, unquantized_user in zip( + list(quantized_node.users), list(unquantized_node.users) + ): + if "output" == quantized_user.name: + continue + assert quantized_user.target == unquantized_user.target, ( + "found unmatched targets: " + f"{quantized_user.target} vs {unquantized_user.target}" + ) + if quantized_user.target in qdq_target: + unquantized_user.args = ( + unquantized_user.args[0], + *quantized_user.args[1:], + ) + activation_override(quantized_user, unquantized_user) + + def parameter_override(quantized_node, unquantized_node): + # Some parameters need to be iterated over to retrieve attributes such as static_llama.tok_embedding.weight + def _get_attr(graph_module: torch.fx.GraphModule, target: str) -> Any: + attr: Any = graph_module + for target_atom in target.split("."): + attr = getattr(attr, target_atom) + return attr + + def _set_attr( + graph_module: torch.fx.GraphModule, target: str, replacement: Any + ) -> Any: + attr: Any = graph_module + target_list = target.split(".") + for target_atom in target_list[:-1]: + attr = getattr(attr, target_atom) + setattr(attr, target_list[-1], replacement) + + _set_attr( + unquantized_model, + unquantized_node.target, + _get_attr(quantized_model, quantized_node.target), + ) + # scale / zero point are part of op's attributes + if list(quantized_node.users)[0].target in ptq_target: + activation_override(quantized_node, unquantized_node) + + # copy encoding for hybrid mode + parameters = [ + {n: resolve_param_target(n) for n in model.graph.nodes if n.op == "get_attr"} + for model in (quantized_model, unquantized_model) + ] + activations = [ + [ + n + for n in model.graph.nodes + if n.target not in qdq_target and n.op in {"call_function", "placeholder"} + ] + for model in (quantized_model, unquantized_model) + ] + # check topology order by node name & nn_module_stack + for act_quantized, act_unquantized in zip(*activations): + compare_nodes(act_quantized, act_unquantized) + + for op_quantized, op_unquantized in zip(*[p.values() for p in parameters]): + compare_nodes(op_quantized, op_unquantized) + # perform encoding override + for act_quantized, act_unquantized in zip(*activations): + activation_override(act_quantized, act_unquantized) + + for param_quantized, param_unquantized in zip(*[p.keys() for p in parameters]): + parameter_override(param_quantized, param_unquantized) + + if n_cache_layers is not None: + k_input_cache_nodes = [] + v_input_cache_nodes = [] + for node in unquantized_model.graph.nodes: + if node.op != "placeholder": + continue + + if "args_" in node.name: + args_idx = int(node.name.split("_")[-1]) + + if args_idx >= n_cache_layers: + v_input_cache_nodes.append(node) + else: + k_input_cache_nodes.append(node) + + if not k_input_cache_nodes or not v_input_cache_nodes: + raise RuntimeError( + "KV cache input detection failed. This likely means the model naming " + "does not match expected prefixes." + ) + + k_output_cache_nodes = [] + v_output_cache_nodes = [] + for node in quantized_model.graph.nodes: + if not is_graph_output(node): + continue + cache_output_node = node.args[0].args[0] + if is_node_src_start_with_name(cache_output_node, kv_cache_prefix="k_"): + k_output_cache_nodes.append(cache_output_node) + elif is_node_src_start_with_name(cache_output_node, kv_cache_prefix="v_"): + v_output_cache_nodes.append(cache_output_node) + + if not k_output_cache_nodes or not v_output_cache_nodes: + raise RuntimeError( + "KV cache detection failed. This likely means the model naming " + "does not match expected prefixes." + ) + + for input_k_cache_node, output_k_cache_node in zip( + k_input_cache_nodes, k_output_cache_nodes + ): + activation_override(output_k_cache_node, input_k_cache_node) + for input_v_cache_node, output_v_cache_node in zip( + v_input_cache_nodes, v_output_cache_nodes + ): + activation_override(output_v_cache_node, input_v_cache_node) + + unquantized_model.recompile() diff --git a/backends/qualcomm/genai_pipeline/recipes/__init__.py b/backends/qualcomm/genai_pipeline/recipes/__init__.py new file mode 100644 index 00000000000..287a6ec0ff4 --- /dev/null +++ b/backends/qualcomm/genai_pipeline/recipes/__init__.py @@ -0,0 +1,65 @@ +# 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. + +"""GenAI Pipeline quantization recipes. + +Re-exports quantization recipe classes from the existing llama module, +providing a stable import path within the genai_pipeline namespace. + +Usage: + from backends.qualcomm.genai_pipeline.recipes import ( + StaticLLMQuantRecipe, + Llama3_1BQuantRecipe, + ) +""" + +from executorch.examples.qualcomm.oss_scripts.llama.static_llm_quant_recipe import ( + CodegenQuantRecipe, + Gemma2QuantRecipe, + Gemma3QuantRecipe, + Gemma_2BQuantRecipe, + GLM_1_5B_InstructQuantRecipe, + Granite_3_3_2B_InstructQuantRecipe, + GraniteSpeech_3_3_2B_InstructQuantRecipe, + InternVL3_1B_QuantRecipe, + Llama3_1BQuantRecipe, + Llama3_3BQuantRecipe, + LlamaStories110MQuantRecipe, + LlamaStories260KQuantRecipe, + Phi4MiniQuantRecipe, + Qwen2_5_0_5BQuantRecipe, + Qwen2_5_1_5BQuantRecipe, + Qwen3_0_6BQuantRecipe, + Qwen3_1_7BQuantRecipe, + Smollm2QuantRecipe, + Smollm3QuantRecipe, + SmolVLMQuantRecipe, + StaticLLMQuantRecipe, +) + +__all__ = [ + "StaticLLMQuantRecipe", + "CodegenQuantRecipe", + "Gemma2QuantRecipe", + "Gemma3QuantRecipe", + "Gemma_2BQuantRecipe", + "GLM_1_5B_InstructQuantRecipe", + "Granite_3_3_2B_InstructQuantRecipe", + "GraniteSpeech_3_3_2B_InstructQuantRecipe", + "InternVL3_1B_QuantRecipe", + "Llama3_1BQuantRecipe", + "Llama3_3BQuantRecipe", + "LlamaStories110MQuantRecipe", + "LlamaStories260KQuantRecipe", + "Phi4MiniQuantRecipe", + "Qwen2_5_0_5BQuantRecipe", + "Qwen2_5_1_5BQuantRecipe", + "Qwen3_0_6BQuantRecipe", + "Qwen3_1_7BQuantRecipe", + "Smollm2QuantRecipe", + "Smollm3QuantRecipe", + "SmolVLMQuantRecipe", +] diff --git a/backends/qualcomm/genai_pipeline/source_transform/__init__.py b/backends/qualcomm/genai_pipeline/source_transform/__init__.py new file mode 100644 index 00000000000..f88ade5cb97 --- /dev/null +++ b/backends/qualcomm/genai_pipeline/source_transform/__init__.py @@ -0,0 +1,81 @@ +# 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. + +"""Model preparation transforms, declared per model by ``SOURCE_TRANSFORMS``. + +Model preparation is *declared*, not subclassed. The reference flow embeds it in +``LLMWrapper._prepare_model()``, where overlapping transform sets (Llama needing +``[A, B, C]`` while Gemma needs ``[B, C, D]``) cannot be referenced +individually. Here each transform is a named function implemented once, and +enabling a model is a row in ``models.model_registry`` rather than a new adapter +class. + +Ordering between the two transform stages is load-bearing, so a single +``transforms`` list is insufficient. The stages run: + +.. code-block:: text + + construct module (model_class_name) + | + load / convert checkpoint + | + state_dict_transforms + | (with model-specific values already bound) + load_state_dict(assign=True) + | + module_transforms + +Each transform receives only its operand at execution time. Model-specific +construction in :func:`model_lookup.get_source_transform` binds every additional +dependency as an explicit value: RoPE receives layer and head counts, embedding +scaling receives its scale factor, and dtype conversion receives the requested +dtype. A transform therefore never needs the full model configuration or CLI +namespace. + +* ``state_dict_transforms``: ``(state_dict) -> state_dict``. Transforms that only + rewrite values mutate and return the same dict -- these state dicts hold + multi-GB of weights and copying them is not free. Transforms that rename keys + necessarily build a new dict. +* ``module_transforms``: ``(module) -> module``. In-place mutations return the + same module; module-to-module replacements return the new one, so callers must + always use the return value. +""" + +from executorch.backends.qualcomm.genai_pipeline.source_transform.checkpoint_key_remap import ( + remap_gemma4_keys, + strip_orig_mod_prefix, + unwrap_model_key, +) +from executorch.backends.qualcomm.genai_pipeline.source_transform.dtype_override import ( + apply_dtype_override, +) +from executorch.backends.qualcomm.genai_pipeline.source_transform.embedding_scale import ( + scale_token_embedding, +) +from executorch.backends.qualcomm.genai_pipeline.source_transform.linear_to_conv2d import ( + convert_linear_to_conv2d, + prepare_conv_submodules, +) +from executorch.backends.qualcomm.genai_pipeline.source_transform.rms_norm_offset import ( + gemma_rmsnorm_offset, +) +from executorch.backends.qualcomm.genai_pipeline.source_transform.rope_layout import ( + permute_partial_rope, +) + +__all__ = [ + # state_dict transforms + "gemma_rmsnorm_offset", + "permute_partial_rope", + "remap_gemma4_keys", + "scale_token_embedding", + "strip_orig_mod_prefix", + "unwrap_model_key", + # module transforms + "apply_dtype_override", + "convert_linear_to_conv2d", + "prepare_conv_submodules", +] diff --git a/backends/qualcomm/genai_pipeline/source_transform/checkpoint_key_remap.py b/backends/qualcomm/genai_pipeline/source_transform/checkpoint_key_remap.py new file mode 100644 index 00000000000..4980fabf041 --- /dev/null +++ b/backends/qualcomm/genai_pipeline/source_transform/checkpoint_key_remap.py @@ -0,0 +1,41 @@ +# 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. + +"""State-dict transforms that rewrite the checkpoint's key layout. + +None of these can mutate in place: they return a different dict -- a rebuilt one +for the renames, the nested one for the unwrap -- while the transforms elsewhere +in this package rewrite values. +""" + +from __future__ import annotations + +from typing import Any, Dict + + +def unwrap_model_key(state_dict: Dict[str, Any]) -> Dict[str, Any]: + """Take the weights out of a checkpoint that nests them under ``model``. + + Local checkpoints wrap the weights alongside their training state; + HF-converted ones are already flat. Self-guarding on the key, so no model has + to opt out and this can lead every model's chain -- it has to, since every + rename after it looks keys up by their canonical name. + """ + return state_dict["model"] if "model" in state_dict else state_dict + + +def strip_orig_mod_prefix(state_dict: Dict[str, Any]) -> Dict[str, Any]: + """Drop the ``_orig_mod.`` prefix ``torch.compile`` leaves on checkpoints.""" + return {k.replace("_orig_mod.", ""): v for k, v in state_dict.items()} + + +def remap_gemma4_keys(state_dict: Dict[str, Any]) -> Dict[str, Any]: + """Rename Gemma4's converted weights to the static decoder's naming.""" + from executorch.examples.qualcomm.oss_scripts.gemma4.text_decoder.convert_weights import ( + remap_keys, + ) + + return remap_keys(state_dict) diff --git a/backends/qualcomm/genai_pipeline/source_transform/dtype_override.py b/backends/qualcomm/genai_pipeline/source_transform/dtype_override.py new file mode 100644 index 00000000000..3d5ee4d615c --- /dev/null +++ b/backends/qualcomm/genai_pipeline/source_transform/dtype_override.py @@ -0,0 +1,21 @@ +# 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. + +"""Module transforms casting the decoder's dtype.""" + +from __future__ import annotations + +from typing import Any, Optional + + +def apply_dtype_override(module: Any, *, dtype_override: Optional[str]) -> Any: + """Cast the module to ``--dtype-override``, if one was requested.""" + if dtype_override is None: + return module + + from executorch.extension.llm.export.builder import DType + + return module.to(DType[dtype_override].to_torch_dtype()) diff --git a/backends/qualcomm/genai_pipeline/source_transform/embedding_scale.py b/backends/qualcomm/genai_pipeline/source_transform/embedding_scale.py new file mode 100644 index 00000000000..c0a311544ce --- /dev/null +++ b/backends/qualcomm/genai_pipeline/source_transform/embedding_scale.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. + +"""State-dict transforms for the token embedding table.""" + +from __future__ import annotations + +from typing import Any, Dict + + +def scale_token_embedding( + state_dict: Dict[str, Any], *, embedding_scale_factor: float +) -> Dict[str, Any]: + """Bake ``embedding_scale_factor`` into the embedding table. + + A no-op -- dtype included -- for models whose factor is 1.0, and for models + that export their token embedding as a separate graph and so carry no such + key. Declared only by the rows whose params file sets a factor; the guard + means a row that declares it anyway costs nothing. + """ + key = "tok_embeddings.weight" + if embedding_scale_factor != 1.0 and key in state_dict: + state_dict[key] = state_dict[key].float() * embedding_scale_factor + return state_dict diff --git a/backends/qualcomm/genai_pipeline/source_transform/linear_to_conv2d.py b/backends/qualcomm/genai_pipeline/source_transform/linear_to_conv2d.py new file mode 100644 index 00000000000..9c9b1c86219 --- /dev/null +++ b/backends/qualcomm/genai_pipeline/source_transform/linear_to_conv2d.py @@ -0,0 +1,43 @@ +# 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. + +"""Module transforms migrating linear layers to conv2d. + +The two transforms here are the halves of one migration and are ordered: +:func:`prepare_conv_submodules` establishes the weight shapes that +:func:`convert_linear_to_conv2d` then reads. +""" + +from __future__ import annotations + +from typing import Any + + +def prepare_conv_submodules(module: Any) -> Any: + """Let attention and feed-forward blocks pre-shape their weights. + + Must run before :func:`convert_linear_to_conv2d`, which reads the shapes + these hooks establish. + """ + for layer in module.layers: + if getattr(layer.attention, "prepare_attention_conv", None): + layer.attention.prepare_attention_conv() + if getattr(layer.feed_forward, "prepare_feedforward_conv", None): + layer.feed_forward.prepare_feedforward_conv() + return module + + +def convert_linear_to_conv2d(module: Any) -> Any: + """Rewrite every ``nn.Linear`` as a 1x1 conv2d, which HTP runs faster. + + The backend util walks named attributes, so linears reachable only through + an ``nn.Sequential`` are left alone. + """ + from executorch.backends.qualcomm.utils.utils import ( + convert_linear_to_conv2d as _convert, + ) + + return _convert(module) diff --git a/backends/qualcomm/genai_pipeline/source_transform/rms_norm_offset.py b/backends/qualcomm/genai_pipeline/source_transform/rms_norm_offset.py new file mode 100644 index 00000000000..1c63e6ea1d5 --- /dev/null +++ b/backends/qualcomm/genai_pipeline/source_transform/rms_norm_offset.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. + +"""State-dict transforms for RMSNorm weights.""" + +from __future__ import annotations + +from typing import Any, Dict + +import torch + + +def gemma_rmsnorm_offset(state_dict: Dict[str, Any]) -> Dict[str, Any]: + """Fold Gemma's implicit ``+1`` into every RMSNorm weight. + + Gemma computes ``(x * w).to(fp16)`` where Llama computes ``x.to(fp16) * w``; + the static decoder implements the Llama form. See + https://github.com/huggingface/transformers/pull/29402. + """ + for k, v in state_dict.items(): + if "norm" not in k: + continue + state_dict[k] = v.float() + torch.ones(v.shape, dtype=torch.float32) + return state_dict diff --git a/backends/qualcomm/genai_pipeline/source_transform/rope_layout.py b/backends/qualcomm/genai_pipeline/source_transform/rope_layout.py new file mode 100644 index 00000000000..51d1efd4c20 --- /dev/null +++ b/backends/qualcomm/genai_pipeline/source_transform/rope_layout.py @@ -0,0 +1,48 @@ +# 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. + +"""State-dict transforms for RoPE weight layout.""" + +from __future__ import annotations + +from typing import Any, Dict + + +def permute_partial_rope( + state_dict: Dict[str, Any], + *, + n_layers: int, + n_heads: int, + n_kv_heads: int, + partial_rotary_factor: float, +) -> Dict[str, Any]: + """change to HF weight to improve the performance of RoPE in HTP backend.""" + + def permute(w, heads, partial_rotary_dim): + dim_0 = w.size(0) + dim_1 = w.size(1) + transformed_weight = ( + w.view(heads, -1, dim_0 // heads // 2 // partial_rotary_dim, 2, dim_1) + .transpose(2, 3) + .reshape(dim_0, dim_1) + ) + return transformed_weight + + # TODO: handle cases where input size isn't divisible. + partial_rotary_dim = int(1 // partial_rotary_factor) + for layer_i in range(n_layers): + state_dict[f"layers.{layer_i}.attention.wq.weight"] = permute( + state_dict[f"layers.{layer_i}.attention.wq.weight"], + n_heads, + partial_rotary_dim, + ) + state_dict[f"layers.{layer_i}.attention.wk.weight"] = permute( + state_dict[f"layers.{layer_i}.attention.wk.weight"], + n_kv_heads, + partial_rotary_dim, + ) + + return state_dict 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 ef96a17b98f..c6208cffb06 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, Tuple +from typing import Any, Callable, Dict, Optional, Sequence, Tuple logger = logging.getLogger(__name__) @@ -22,12 +22,8 @@ class DefaultModelLoaderAdapter: .. note:: This adapter is designed for **text-only causal LM** models. For multimodal models (vision/audio encoders), use a specialized adapter - (e.g., ``MultiModalModelLoaderAdapter``) that handles separate encoder/decoder + (e.g., ``MLLMLoaderAdapter``) that handles separate encoder/decoder loading via ``AutoModel`` or ``AutoModelForSpeechSeq2Seq``. - - .. note:: - Calibration data is **not** produced here -- see - ``CalibrationDataAdapter``. """ #: Batch size and sequence length of the generated example inputs. HTP has no @@ -48,6 +44,11 @@ def load_model( ) -> Any: """Load a causal LM model from HuggingFace hub or local path. + The model-preparation strategy consumes every loader through the same + component/graph map shape. A plain HuggingFace causal LM has no split + decoder graphs or extra components, so this adapter wraps the loaded + module as the single text-decoder ``GRAPH_FORWARD`` graph. + Args: model_name: HuggingFace model ID or local path. extra_options: Additional options. Supported keys: @@ -56,7 +57,9 @@ def load_model( - ``attn_implementation``: attention implementation to use. Returns: - The loaded nn.Module in eval mode. + ``{ARTIFACT_TEXT_DECODER: {GRAPH_FORWARD: module}}``. The + default HuggingFace path is the single-component, single-graph case + of the registry loaders' component/graph contract. """ import torch from transformers import AutoModelForCausalLM @@ -79,7 +82,14 @@ def load_model( model.eval() logger.info("Model loaded successfully") - return model + from executorch.backends.qualcomm.genai_pipeline.artifact_keys import ( + ARTIFACT_TEXT_DECODER, + ) + from executorch.backends.qualcomm.genai_pipeline.graph_names import ( + GRAPH_FORWARD, + ) + + return {ARTIFACT_TEXT_DECODER: {GRAPH_FORWARD: model}} def load_tokenizer( self, @@ -105,6 +115,28 @@ def load_tokenizer( logger.info("Tokenizer loaded successfully") return tokenizer + def apply_module_transforms( + self, + module: Any, + module_transforms: Optional[Sequence[Callable[[Any], Any]]] = None, + ) -> Any: + """Apply transforms to module. + + Args: + module: The module to be transformed. + module_transforms: Transforms to apply in order. + + Returns: + The transformed module. + """ + if not module_transforms: + return module + + logger.debug("Applying %d module transforms", len(module_transforms)) + for transform in module_transforms: + module = transform(module) + return module + def get_example_inputs( self, model: Any, @@ -150,6 +182,28 @@ def get_example_inputs( # int64 token ids: the embedding lookup indexes with them. return (torch.zeros((batch_size, ar_len), dtype=torch.int64),) + def get_metadata(self, module: Any) -> Dict[str, Any]: + """Read constant metadata from a graph module when it exposes it.""" + model_provided = getattr(module, "get_metadata", None) + if callable(model_provided): + logger.debug("Using metadata provided by the model") + return model_provided() + return {} + + def get_inference( + self, + meta: Any, + example_inputs: Any, + extra_options: Optional[Dict[str, Any]] = None, + ) -> Any: + """No inference helper is currently provided by this adapter. + + The HuggingFace default adapter does not currently use a model-specific + inference helper during the workflow, so this method always returns + ``None``. + """ + return None + def export_tokenizer( self, tokenizer: Any, 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 1d9ea05a975..fc52f09f66e 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 @@ -6,9 +6,10 @@ from __future__ import annotations +import gc import logging from pathlib import Path -from typing import Optional +from typing import Any, Dict, Optional from executorch.backends.qualcomm.genai_pipeline.configs.model_preparation_input_config import ( ModelPreparationInputConfig, @@ -16,13 +17,6 @@ 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 ( @@ -32,45 +26,71 @@ ModelPreparationStrategy, ) + logger = logging.getLogger(__name__) _STAGE_NAME = "model_preparation" class ExecuTorchModelPreparationStrategy(ModelPreparationStrategy): - """ExecuTorch-based model preparation using HuggingFace transformers. - - 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 + """Prepare model artifacts for subsequent ExecuTorch pipeline stages. + + This strategy owns the model-agnostic preparation flow and delegates + model-family-specific loading to ``ModelLoaderAdapter``. Graph variants share + one weighted module per component, while their metadata and example inputs stay + keyed by component and graph because export signatures may differ. + + Example:: + + # Text-only + { + ARTIFACT_TEXT_DECODER: { + GRAPH_FORWARD: decoder, + GRAPH_KV_FORWARD: decoder, + GRAPH_PREFILL_FORWARD: decoder, + }, + } + + # Multimodal + { + ARTIFACT_TEXT_DECODER: { + GRAPH_FORWARD: decoder, + GRAPH_KV_FORWARD: decoder, + GRAPH_PREFILL_FORWARD: decoder, + }, + ARTIFACT_TOK_EMBEDDING: { + GRAPH_FORWARD: tok_embedding, + ..., + }, + ARTIFACT_VISION_ENCODER: { + GRAPH_FORWARD: vision_encoder, + }, + ARTIFACT_AUDIO_ENCODER: { + GRAPH_FORWARD: audio_encoder, + }, + } + + The preparation flow mirrors :meth:`invoke`: + + 1. Load one weighted module per component and graph variant. + 2. Read per-graph metadata. + 3. Build model-native example inputs for every graph variant. + 4. Select one weight-sharing module per component. + 5. Apply component-level module transforms. + 6. Load the tokenizer. + 7. Build the model-specific inference helper used by quantization, if any. + 8. Optionally export the tokenizer for runtime. + 9. Extract the tokenizer chat template, falling back to ``extra_options`` + when the tokenizer has none. 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 ( @@ -80,54 +100,54 @@ def __init__( 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, input_config: ModelPreparationInputConfig, ) -> ModelPreparationOutputConfig: - """Prepare the model, tokenizer, and calibration data. + """Prepare component modules and runtime metadata for later pipeline stages. Args: - 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. + context: Pipeline context. Its artifact directory is used when + exporting the runtime tokenizer. + input_config: Model identity, target SoC, and preparation inputs. + + Configuration: + ``input_config.extra_options["model_options"]`` configures model + loading: + + - ``model_arch``: Component- and graph-keyed module constructors. + - ``state_dict_loader``: Component-keyed loaders. Each loader + accepts a Hugging Face ``repo_id`` and returns a state dict. + - ``weight_transforms``: Component-keyed transforms applied while + loading weights. + - ``module_transforms``: Component-keyed transforms applied after + one shared module is selected for each component. + - ``num_shardings``: Optional component shard counts. + + Root ``input_config.extra_options`` configures tokenizer and runtime + output: + + - ``tokenizer_options``: Options passed to ``load_tokenizer``. + - ``export_tokenizer``: Enables runtime tokenizer export. + - ``tokenizer_export_options``: Options passed when tokenizer + export is enabled. + - ``chat_template``: Fallback used only when the loaded tokenizer + has no chat template. Returns: - ModelPreparationOutputConfig with model, tokenizer, and calibration data. + A ``ModelPreparationOutputConfig`` with component-level modules, + component- and graph-keyed inputs and metadata, tokenizer state, + optional quantization inference support, and shard counts. Raises: - StageError: If model_name is missing or any loading step fails. + StageError: If required fields are missing or any preparation step + fails. """ logger.info( "Starting model preparation for '%s' on SoC=%s", @@ -139,44 +159,59 @@ def invoke( try: extra = dict(input_config.extra_options) + model_options = extra.get("model_options", {}) # Step 1: Load model logger.debug("Loading model") - model_module = self._adapter.load_model( + modules = self._adapter.load_model( model_name=input_config.model_name, - extra_options=extra.get("model_options"), + extra_options=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 2: Get per-graph constant metadata. + # Text decoder metadata determines whether its nested example inputs + # include KV-cache arguments. + meta = self._get_metadata(modules) - # Step 3: Build the export example inputs from the model itself. + # Step 3: Build per-graph export example inputs from the model. # 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"), + example_inputs = self._get_example_inputs(modules) + + # Select one module per component. + # TODO: Make module-wrapper ``get_example_inputs`` accept shape-related + # parameters so different-shaped inputs are derived from the shape request + # rather than from separate graph-wrapper modules. + model_module = self._get_component_module(modules) + del modules + gc.collect() + + # Step 4: Apply component-level transforms. + logger.debug("Applying module transforms") + model_module = self._apply_module_transforms( + model_module, + extra_options=model_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: Load tokenizer. + logger.debug("Loading tokenizer") + tokenizer = self._adapter.load_tokenizer( + model_name=input_config.model_name, + extra_options=extra.get("tokenizer_options"), + ) + + # Step 6: Create the model-specific inference instance or callable used + # by quantization. Returns None when not applicable. + inference = self._adapter.get_inference( + meta, + example_inputs, + extra_options=extra, ) - # Step 5: Optionally export tokenizer for runtime + # Step 7: Optionally export tokenizer for runtime runtime_tokenizer_path = None if extra.get("export_tokenizer", False): logger.debug("Exporting tokenizer for runtime use") @@ -187,7 +222,7 @@ def invoke( extra_options=extra.get("tokenizer_export_options"), ) - # Step 6: Extract chat_template from tokenizer (for instruct models). + # Step 8: 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. @@ -205,9 +240,14 @@ def invoke( model_module=model_module, tokenizer=tokenizer, example_inputs=example_inputs, - calibration_data=calibration_data, runtime_tokenizer_path=runtime_tokenizer_path, chat_template=chat_template, + meta=meta, + inference=inference, + num_shardings=model_options.get( + "num_shardings", + extra.get("num_shardings"), + ), ) except StageError: @@ -238,3 +278,105 @@ def _validate_input(self, input_config: ModelPreparationInputConfig) -> None: stage_name=_STAGE_NAME, message="soc_model is required for model preparation", ) + + def _get_metadata(self, modules: Dict[str, Dict[str, Any]]) -> Dict: + """Collect non-empty constant metadata for every component graph. + + Args: + modules: Component- and graph-keyed loaded modules. + + Returns: + Component- and graph-keyed metadata, omitting graphs and components + whose adapter metadata is empty. + """ + meta = {} + for component, graph_modules in modules.items(): + component_meta = { + graph_name: graph_meta + for graph_name, graph_module in graph_modules.items() + if (graph_meta := self._adapter.get_metadata(graph_module)) + } + if component_meta: + meta[component] = component_meta + return meta + + def _get_example_inputs( + self, + modules: Dict[str, Dict[str, Any]], + ) -> Dict[str, Dict[str, Any]]: + """Collect model-native example inputs for every component graph. + + Args: + modules: Component- and graph-keyed loaded modules. + + Returns: + Component- and graph-keyed example inputs that define each graph's + export signature. + """ + example_inputs = {} + for component, graph_modules in modules.items(): + example_inputs[component] = {} + for graph_name, graph_module in graph_modules.items(): + logger.debug( + "Building example inputs for component '%s' graph '%s'", + component, + graph_name, + ) + example_inputs[component][graph_name] = ( + self._adapter.get_example_inputs(graph_module) + ) + return example_inputs + + def _apply_module_transforms( + self, + modules: Dict[str, Any], + extra_options: Dict[str, Any], + ) -> Dict[str, Any]: + """Apply configured transforms to each selected component module. + + Args: + modules: Component-keyed modules selected from graph variants. + extra_options: The ``model_options`` map containing component-keyed + ``module_transforms``. + + Returns: + Component-keyed modules after their configured transforms run. + + Raises: + ValueError: If ``module_transforms`` is not component-keyed. + """ + module_transforms = extra_options.get("module_transforms", {}) + if not isinstance(module_transforms, dict): + raise ValueError("module_transforms must be component-keyed") + return { + component: self._adapter.apply_module_transforms( + module, + module_transforms=module_transforms.get(component, []), + ) + for component, module in modules.items() + } + + def _get_component_module( + self, modules: Dict[str, Dict[str, Any]] + ) -> Dict[str, Any]: + """Select one shared module per component. + + Graph variants share a component's module and weights, while their + graph-specific example inputs and metadata remain separate. + + Args: + modules: Component- and graph-keyed loaded modules. + + Returns: + Component-keyed representative modules, using the first graph + variant for each non-empty component. + """ + model_module = {} + for component, graph_modules in modules.items(): + if not graph_modules: + continue + + selected_graph_name = next(iter(graph_modules)) + model_module[component] = graph_modules[selected_graph_name] + + return model_module diff --git a/backends/qualcomm/genai_pipeline/strategies/model_preparation/llm_loader_adapter.py b/backends/qualcomm/genai_pipeline/strategies/model_preparation/llm_loader_adapter.py new file mode 100644 index 00000000000..3ae510cb0d0 --- /dev/null +++ b/backends/qualcomm/genai_pipeline/strategies/model_preparation/llm_loader_adapter.py @@ -0,0 +1,379 @@ +# 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. + +"""Registry-driven loading of a static LLM decoder. + +The reference flow does construction, checkpoint acquisition and preparation in +one ~100-line ``LLMWrapper._prepare_model()`` with every per-model branch inlined +(``if decoder_model in {"gemma-2b", "gemma2-2b", "gemma3-1b"}: ...``). Here the +mechanism is written once and the per-model part is data: +:meth:`LLMLoaderAdapter.from_model_config` reads the model's row in +``models.model_registry``, which names the decoder class to construct and the +transforms to run. + +All graph variants (calibration, decode, prefill) share a single checkpoint but +expose different ``get_example_inputs()`` signatures based on their ``ar_len`` and +KV-cache configuration. The adapter constructs all variants in :meth:`load_model` +and loads the same checkpoint into each graph wrapper. +""" + +from __future__ import annotations + +import logging +from pathlib import Path +from typing import Any, Callable, Dict, List, Optional + +import torch +from executorch.backends.qualcomm.genai_pipeline.strategies.model_preparation.default_model_loader_adapter import ( + DefaultModelLoaderAdapter, +) + +logger = logging.getLogger(__name__) + + +class LLMLoaderAdapter(DefaultModelLoaderAdapter): + """Loads a registry-declared static decoder from its checkpoint. + + Implementations of :class:`ModelLoaderAdapter` vary by *loading mechanism*, + not by model family: this one covers every model in ``SUPPORTED_LLM_MODELS``, + because what differs per model arrives as the registry row + :meth:`from_model_config` reads rather than as a subclass. + + :meth:`load_model` returns a ``{ARTIFACT_TEXT_DECODER: {graph_name: module}}`` map. + The graph variants share one checkpoint and differ only in AR length and + KV-cache use. + + Args: + model_config: The model's ``LLMModelConfig``, carrying ``repo_id``. + control_args: The CLI argument namespace, providing the optional + ``checkpoint`` override and consumed by ``TokenizerWrapper``. + """ + + RUNTIME_TOKENIZER_NAMES = ("tokenizer.json", "tokenizer.model") + + def __init__( + self, + model_config: Any, + control_args: Any, + ) -> None: + self.model_config = model_config + self.control_args = control_args + + def load_model( + self, + model_name: str, + extra_options: Optional[Dict[str, Any]] = None, + ) -> Dict[str, Any]: + """Build the decoder's graph modules and load their weights. + + Constructs one module per graph variant (calibration, decode, prefill), + each with its own ``ar_len`` and KV-cache configuration. All variants + load the same checkpoint. + + Args: + model_name: Registry model name (e.g. ``"llama3_2-1b_instruct"``). + extra_options: Model preparation options containing component-keyed + ``model_arch`` and ``weight_transforms`` entries, plus an + optional component-keyed ``state_dict_loader`` map for HF + checkpoints. + + Returns: + ``{ARTIFACT_TEXT_DECODER: {graph_name: module}}``, all modules weight-carrying. + A text-only model is the degenerate one-component case of the + component-keyed shape the multimodal loader also produces. + + Raises: + ValueError: If model_arch is missing or its text-decoder map is empty. + """ + from executorch.backends.qualcomm.genai_pipeline.artifact_keys import ( + ARTIFACT_TEXT_DECODER, + ) + + extra_options = extra_options or {} + + model_arch = extra_options.get("model_arch") + weight_transforms = (extra_options.get("weight_transforms") or {}).get( + ARTIFACT_TEXT_DECODER, [] + ) + state_dict_loaders = extra_options.get("state_dict_loader") or {} + if not isinstance(state_dict_loaders, dict): + raise ValueError("state_dict_loader must be component-keyed") + state_dict_loader = state_dict_loaders.get(ARTIFACT_TEXT_DECODER) + + if model_arch is None: + raise ValueError("model_arch is required in model_options") + + graph_modules = self._build_custom_model(model_arch) + state_dict = self._load_state_dict(state_dict_loader) + state_dict = self._apply_weight_transforms(state_dict, weight_transforms) + + # For the nn.Module itself, different input-shape graphs share the same + # module and weights. Today the wrapper constructors still produce one + # module per graph shape, so every wrapper must receive the same weights. + # TODO: Make module-wrapper get_example_inputs accept shape-related + # parameters so different-shaped inputs are derived from the shape + # request rather than from separate graph-wrapper modules. + for module in graph_modules.values(): + module.load_state_dict(state_dict, strict=True, assign=True) + + logger.info( + "Model loaded successfully with %d graph variants", len(graph_modules) + ) + return {ARTIFACT_TEXT_DECODER: graph_modules} + + def _build_custom_model( + self, + model_arch: Any, + ) -> Dict[str, Any]: + """Construct the decoder modules, without weights. + + Args: + model_arch: Component-keyed dict of per-graph constructors, i.e. + ``{ARTIFACT_TEXT_DECODER: {graph_name: constructor}}``. Each constructor is + a ``partial`` with its config and every arch-specific kwarg + pre-bound by ``get_model_arch``, so this method stays arch-agnostic + and calls each with no arguments. + + Returns: + Dict mapping graph names to uninitialized modules. + + Raises: + ValueError: If the text-decoder constructor map is empty. + """ + from executorch.backends.qualcomm.genai_pipeline.artifact_keys import ( + ARTIFACT_TEXT_DECODER, + ) + + graph_arch = model_arch[ARTIFACT_TEXT_DECODER] + if not graph_arch: + raise ValueError("model_arch[ARTIFACT_TEXT_DECODER] must not be empty") + + modules: Dict[str, Any] = {} + for name, build in graph_arch.items(): + modules[name] = build().eval() + + return modules + + def _load_state_dict( + self, + state_dict_loader: Optional[Callable[[str], Dict[str, Any]]] = None, + ) -> Dict[str, Any]: + """The raw checkpoint state dict, before any weight transform. + + Whatever ``control_args`` names, or the model's ``repo_id`` downloaded + and converted when it names none. Registry-specific HF loaders can be + passed through ``model_options["state_dict_loader"]`` and must accept + only ``repo_id``. + + Raises: + ValueError: If the model ships no ``repo_id`` and ``control_args`` + names no checkpoint. + """ + checkpoint = self.control_args.checkpoint + if checkpoint is None: + if self.model_config.repo_id is None: + raise ValueError( + f"'{self.control_args.model}' is supplied as a local " + "checkpoint; pass one via control_args.checkpoint." + ) + if state_dict_loader is not None: + return state_dict_loader(self.model_config.repo_id) + + from executorch.examples.models.llama.hf_download import ( + download_and_convert_hf_checkpoint, + ) + + # convert_weights is held as a class attribute, so reading it off the + # model_config instance binds it; __func__ recovers the plain function + # the downloader expects. + checkpoint = download_and_convert_hf_checkpoint( + self.model_config.repo_id, self.model_config.convert_weights.__func__ + ) + + return torch.load(checkpoint, weights_only=True, map_location="cpu", mmap=True) + + def _apply_weight_transforms( + self, + state_dict: Dict[str, Any], + weight_transforms: Optional[List[Callable]] = None, + ) -> Dict[str, Any]: + """Apply weight transforms to the state dict. + + Args: + state_dict: The raw state dict to transform. + weight_transforms: Optional list of transforms to apply in order. + + Returns: + The transformed state dict. + """ + if not weight_transforms: + return state_dict + + logger.debug("Applying %d weight transforms", len(weight_transforms)) + for transform in weight_transforms: + state_dict = transform(state_dict) + + return state_dict + + def get_inference( + self, + meta: Dict[str, Any], + example_inputs: Dict[str, Any], + extra_options: Optional[Dict[str, Any]] = None, + ) -> Any: + """Build the ``ModelInference`` bound to the calibration graph. + + The calibration graph is the non-deployed ``GRAPH_FORWARD`` graph; + this selects it from the full per-graph maps and drives a + ``DecoderInference`` from its metadata and export-input signature. + + Args: + meta: ``{ARTIFACT_TEXT_DECODER: {graph_name: meta}}`` from :meth:`get_metadata`. + example_inputs: ``{ARTIFACT_TEXT_DECODER: {graph_name: inputs}}`` from + :meth:`get_example_inputs`. + extra_options: Reads ``embedding_quantize`` to decide the token dtype. + + Returns: + A ``ModelInference`` wrapping the calibration graph's + ``DecoderInference``, or ``None`` when no calibration graph is present. + """ + from executorch.backends.qualcomm.genai_pipeline.artifact_keys import ( + ARTIFACT_TEXT_DECODER, + ) + from executorch.backends.qualcomm.genai_pipeline.graph_names import ( + GRAPH_FORWARD, + ) + from executorch.examples.qualcomm.oss_scripts.llama.inference import ( + DecoderInference, + ModelInference, + ) + + calibration_meta = meta.get(ARTIFACT_TEXT_DECODER, {}).get(GRAPH_FORWARD) + calibration_inputs = example_inputs.get(ARTIFACT_TEXT_DECODER, {}).get( + GRAPH_FORWARD + ) + if calibration_meta is None or calibration_inputs is None: + return None + + extra_options = extra_options or {} + use_i64_token = extra_options.get("embedding_quantize") is not None + + return ModelInference( + decoder=DecoderInference( + get_example_inputs=lambda: calibration_inputs, + max_context_len=calibration_meta["get_max_context_len"], + max_batch_size=calibration_meta["get_max_batch_size"], + use_i64_token=use_i64_token, + ) + ) + + def load_tokenizer( + self, + model_name: str, + extra_options: Optional[Dict[str, Any]] = None, + ) -> Any: + """Build the ``TokenizerWrapper`` for this model. + + Returns the wrapper, not the bare tokenizer: the dataset builders and + evaluators need its chat-template and multimodal-prompt helpers, and it + also resolves ``runtime_tokenizer_path``. The bare tokenizer remains + reachable as ``.tokenizer``. + """ + from executorch.examples.qualcomm.oss_scripts.llama.tokenizer import ( + TokenizerWrapper, + ) + + # TODO: Remove this compatibility bridge once TokenizerWrapper accepts + # the pipeline's ``model`` and ``artifact_dir`` field names directly. + self.control_args.decoder_model = getattr( + self.control_args, "decoder_model", self.control_args.model + ) + self.control_args.artifact = getattr( + self.control_args, "artifact", self.control_args.artifact_dir + ) + + return TokenizerWrapper(self.control_args, self.model_config) + + def get_example_inputs( + self, + model: Any, + extra_options: Optional[Dict[str, Any]] = None, + ) -> Dict[str, 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`. + + Returns: + A flat tuple positionally matching ``model.forward``. + """ + + 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()) + + raise ValueError( + "Module has no get_example_inputs() method and " + "no example_inputs provided in extra_options" + ) + + def export_tokenizer( + self, + tokenizer: Any, + output_dir: Path, + extra_options: Optional[Dict[str, Any]] = None, + ) -> Path: + """Export tokenizer to disk and return the runtime tokenizer file. + + ``TokenizerWrapper`` has already written the artifacts. Prefer the + runtime formats recognized by ``pytorch_tokenizers.get_tokenizer`` over + the wrapper's order-dependent fallback. + + Args: + tokenizer: The tokenizer instance to export. + output_dir: Directory to write the exported tokenizer artifacts to. + extra_options: Additional export options. + + Returns: + Path to the runtime tokenizer file (e.g. ``tokenizer.json``). + + Raises: + FileNotFoundError: If no tokenizer artifacts were written. + """ + runtime_tokenizer_path = Path(tokenizer.runtime_tokenizer_path) + artifact_dir = Path(tokenizer.artifact) + artifacts = list(artifact_dir.iterdir()) if artifact_dir.is_dir() else [] + if runtime_tokenizer_path not in artifacts: + artifacts.append(runtime_tokenizer_path) + return self._select_runtime_tokenizer(artifacts) + + @classmethod + def from_model_config( + cls, + model_config: Any, + *, + control_args: Any, + ) -> "LLMLoaderAdapter": + """Build the adapter, taking its source transforms as pre-bound partials. + + Which transforms this model uses is resolved by + ``model_lookup.get_source_transform`` and passed in here, so the adapter + neither reads a registry nor knows the per-model mapping. + """ + return cls( + model_config=model_config, + control_args=control_args, + ) diff --git a/backends/qualcomm/genai_pipeline/strategies/model_preparation/mllm_loader_adapter.py b/backends/qualcomm/genai_pipeline/strategies/model_preparation/mllm_loader_adapter.py new file mode 100644 index 00000000000..90800ee463f --- /dev/null +++ b/backends/qualcomm/genai_pipeline/strategies/model_preparation/mllm_loader_adapter.py @@ -0,0 +1,520 @@ +# 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. + +"""Model loading for a multimodal model, every component implemented here. + +A multimodal model is several modules -- one or more encoders, a token +embedding, and a text decoder. Loading them has no inter-dependency (unlike +quantization's ordered chain), so the whole multi-module load lives in this one +adapter. + +This is the multimodal implementation of ``ModelLoaderAdapter``. It uses the +same component/graph map contract as the text-only loader, with extra +components for modality encoders and token embedding. + +It implements every component itself rather than delegating the decoder to +``LLMLoaderAdapter``: + +- **text_decoder**: constructed per graph from model_arch, its + checkpoint acquired and rewritten through ``get_source_transform``'s partials, + then loaded strict/assign -- the same mechanism the single-module adapter uses, + written here so this adapter owns its whole load. +- **encoder(s)** and **tok_embedding**: both derived from a single shared + HuggingFace ``auto_model`` -- the encoder through the modality config's + ``create_encoder``, the token embedding from ``auto_model.get_input_embeddings()`` + sized by the decoder's ``model_args``. +""" + +from __future__ import annotations + +import logging +from pathlib import Path +from typing import Any, Callable, Dict, List, Optional + +from executorch.backends.qualcomm.genai_pipeline.artifact_keys import ( + ARTIFACT_AUDIO_ENCODER, + ARTIFACT_TEXT_DECODER, + ARTIFACT_TOK_EMBEDDING, + ARTIFACT_VISION_ENCODER, +) +from executorch.backends.qualcomm.genai_pipeline.strategies.model_preparation.default_model_loader_adapter import ( + DefaultModelLoaderAdapter, +) + +logger = logging.getLogger(__name__) + + +class MLLMLoaderAdapter(DefaultModelLoaderAdapter): + """Loads every component of a multimodal model. + + Orchestrates loading of encoders, token embedding, and text decoder from a + shared HuggingFace model. Each component is loaded independently with no + inter-dependencies, enabling flexible composition of multimodal architectures. + + The decoder and embedding are constructed with multiple graph variants + (calibration, decode, prefill) to support different inference modes. All + variants are derived from a single nn.Module via torch.export.export() with + different input shapes (ar_len, KV-cache config), so only one module per + component needs to hold the checkpoint weights. + + Args: + model_config: The model's ``LLMModelConfig``, carrying ``repo_id``, the + decoder class (``model_arch``) and per-modality attributes + (``vision_encoder`` / ``audio_encoder``). + control_args: The CLI argument namespace, providing the optional + ``checkpoint`` override and consumed by ``TokenizerWrapper``. + """ + + RUNTIME_TOKENIZER_NAMES = ("tokenizer.json", "tokenizer.model") + + def __init__( + self, + model_config: Any, + control_args: Any, + ) -> None: + self.model_config = model_config + self.control_args = control_args + + self._modalities = tuple( + m + for m in (ARTIFACT_VISION_ENCODER, ARTIFACT_AUDIO_ENCODER) + if hasattr(model_config, m) + ) + if not self._modalities: + raise ValueError( + "MLLMLoaderAdapter requires at least one modality encoder " + f"({ARTIFACT_VISION_ENCODER!r} or {ARTIFACT_AUDIO_ENCODER!r}) to be defined " + f"in {type(model_config).__name__}." + ) + + def load_model( + self, + model_name: str, + extra_options: Optional[Dict[str, Any]] = None, + ) -> Dict[str, Dict[str, Any]]: + """Build every component's graph modules. + + Orchestrates loading of all model components (encoders, token embedding, + text decoder) from a shared HuggingFace model. Each component is loaded + independently with no inter-dependencies. + + The decoder and embedding are constructed with multiple graph variants + (calibration, decode, prefill) to support different inference modes. + + Args: + model_name: Registry model name (e.g. ``"smolvlm_500m"``). Unused; + the model is identified by ``control_args``. + extra_options: Model preparation options containing component-keyed + ``model_arch`` and ``weight_transforms`` entries, plus an + optional component-keyed ``state_dict_loader`` map for HF + checkpoints. + + Returns: + ``{component: {graph_name: module}}``, where component is one of + ARTIFACT_TEXT_DECODER, ARTIFACT_TOK_EMBEDDING, ARTIFACT_VISION_ENCODER, or ARTIFACT_AUDIO_ENCODER. + + Raises: + ValueError: If required model_options are missing or invalid. + """ + extra_options = extra_options or {} + + # MLLM model options are component-keyed so each component can own its + # graph constructors and transforms. + decoder_model_arch = extra_options.get("model_arch") + if not isinstance(decoder_model_arch, dict): + raise ValueError( + f"For multimodal models, model_arch must be a component-keyed dict " + f"(e.g., {{ARTIFACT_TEXT_DECODER: value}}), got {type(decoder_model_arch).__name__}" + ) + embedding_model_arch = decoder_model_arch.get(ARTIFACT_TOK_EMBEDDING) + decoder_model_arch = decoder_model_arch.get(ARTIFACT_TEXT_DECODER) + if decoder_model_arch is None: + raise ValueError( + "model_arch dict must contain ARTIFACT_TEXT_DECODER key for multimodal models" + ) + + decoder_weight_transforms = extra_options.get("weight_transforms") or {} + if not isinstance(decoder_weight_transforms, dict): + raise ValueError( + f"For multimodal models, weight_transforms must be a component-keyed dict " + f"(e.g., {{ARTIFACT_TEXT_DECODER: value}}), got {type(decoder_weight_transforms).__name__}" + ) + decoder_weight_transforms = decoder_weight_transforms.get( + ARTIFACT_TEXT_DECODER, [] + ) + if decoder_weight_transforms is None: + decoder_weight_transforms = [] + state_dict_loaders = extra_options.get("state_dict_loader") or {} + if not isinstance(state_dict_loaders, dict): + raise ValueError("state_dict_loader must be component-keyed") + state_dict_loader = state_dict_loaders.get(ARTIFACT_TEXT_DECODER) + + # Load auto model + auto_model = self._load_auto_model() + + modules: Dict[str, Dict[str, Any]] = { + # Load encoder + **{ + modality: self._load_encoder(modality, auto_model) + for modality in self._modalities + }, + # Load token embedding (constructors carry their own shape kwargs) + ARTIFACT_TOK_EMBEDDING: self._load_embedding( + auto_model, + embedding_model_arch, + ), + # Load text decoder + ARTIFACT_TEXT_DECODER: self._load_decoder( + decoder_model_arch, + decoder_weight_transforms, + state_dict_loader, + ), + } + + return modules + + def _load_encoder(self, modality: str, auto_model: Any) -> Dict[str, Any]: + """Load one modality encoder from the shared HuggingFace model. + + Creates an encoder wrapper for the specified modality (vision or audio) + using the modality config's ``create_encoder`` method. The encoder weights + are loaded from the shared HuggingFace model with strict=False to allow + the wrapper to expose only the relevant modality weights. + + Args: + modality: The modality name (ARTIFACT_VISION_ENCODER or ARTIFACT_AUDIO_ENCODER). + auto_model: The shared HuggingFace model. + + Returns: + ``{GRAPH_FORWARD: encoder_module}``. + + Raises: + AttributeError: If the modality config is not found in model_config. + """ + from executorch.backends.qualcomm.genai_pipeline.graph_names import ( + GRAPH_FORWARD, + ) + + logger.info("Loading %s", modality) + modality_config = getattr(self.model_config, modality) + encoder = modality_config().create_encoder(auto_model.config).eval() + # strict=False: the encoder wrapper exposes only this modality's weights. + encoder.load_state_dict(auto_model.state_dict(), strict=False) + return {GRAPH_FORWARD: encoder} + + def _load_embedding( + self, + auto_model: Any, + embedding_arch: Any, + ) -> Dict[str, Any]: + """Load the token embedding for all its graph variants. + + Creates one embedding module per graph variant (calibration, decode, + prefill), keyed by the token-embedding graph names carried in + ``embedding_arch``. All variants share the same embedding weights + from the HuggingFace model. + + Args: + auto_model: The shared HuggingFace model. + embedding_arch: ``{graph_name: constructor}`` for the token embedding. + Each constructor is a ``partial`` with the embedding's shape kwargs + pre-bound; the shared weights are the sole call-time argument. + + Returns: + ``{graph_name: embedding_module}``, all modules weight-carrying. + """ + import torch + + if not embedding_arch: + raise RuntimeError( + "embedding_arch not available. Ensure load_model() is " + "called with model_options containing the token-embedding arch." + ) + + logger.info("Loading token embedding") + embedding_weights = auto_model.get_input_embeddings().to(torch.float32) + + embedding_modules: Dict[str, Any] = {} + + # Each constructor is a partial with the embedding's shape kwargs pre-bound; + # the shared HuggingFace weights are the sole + # call-time argument. + for graph_name, build in embedding_arch.items(): + embedding_modules[graph_name] = build(embedding_weights) + + logger.info( + "Token embedding loaded successfully with %d graph variants", + len(embedding_modules), + ) + return embedding_modules + + def _load_decoder( + self, + model_arch: Any = None, + weight_transforms: Optional[List[Callable]] = None, + state_dict_loader: Optional[Callable[[str], Dict[str, Any]]] = None, + ) -> Dict[str, Any]: + """Load the text decoder component's modules. + + Constructs one decoder module per graph variant (calibration, decode, prefill), + each with its own ``ar_len`` and KV-cache configuration. All variants + load the same checkpoint. + + Args: + model_arch: ``{graph_name: constructor}`` for the decoder. Each + constructor is a ``partial`` with its config and construction kwargs + pre-bound by ``get_model_arch``, so it is called with no arguments. + weight_transforms: Optional list of transforms to apply to the state dict + in order before loading. + state_dict_loader: Optional HF checkpoint loader that receives only + ``repo_id`` and returns a state dict. + + Returns: + Dict mapping graph names to decoder modules, all weight-carrying. + + Raises: + ValueError: If model_arch is missing or empty. + """ + if not model_arch: + raise ValueError("model_arch is required and must not be empty") + + logger.info("Loading text decoder") + decoder_modules = self._build_custom_decoder_model(model_arch) + state_dict = self._load_decoder_state_dict(state_dict_loader) + + if weight_transforms: + for transform in weight_transforms: + state_dict = transform(state_dict) + + # For the nn.Module itself, different input-shape graphs share the same + # module and weights. Today the wrapper constructors still produce one + # module per graph shape, so every wrapper must receive the same weights. + # TODO: Make module-wrapper get_example_inputs accept shape-related + # parameters so different-shaped inputs are derived from the shape + # request rather than from separate graph-wrapper modules. + for module in decoder_modules.values(): + module.load_state_dict(state_dict, strict=True, assign=True) + + logger.info( + "Text decoder loaded successfully with %d graph variants", + len(decoder_modules), + ) + + return decoder_modules + + def get_example_inputs( + self, + model: Any, + extra_options: Optional[Dict[str, Any]] = None, + ) -> Dict[str, 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`. + + Returns: + A flat tuple positionally matching ``model.forward``. + """ + + extra_options = extra_options or {} + + model_provided = getattr(model, "get_example_inputs", None) or getattr( + model, "get_example_input", None + ) + if callable(model_provided): + logger.info("Using example inputs provided by the model") + return tuple(model_provided()) + + raise ValueError( + "Module has no get_example_inputs() method and " + "no example_inputs provided in extra_options" + ) + + def get_inference( + self, + meta: Dict[str, Any], + example_inputs: Dict[str, Any], + extra_options: Optional[Dict[str, Any]] = None, + ) -> Any: + """Build the ``ModelInference`` bound to the calibration graph. + + The calibration graph is the non-deployed ``GRAPH_FORWARD`` decoder + graph; this selects it from the full per-graph maps and drives a + ``DecoderInference`` (plus an ``EncoderInference``) from its metadata and + export-input signature. + + Args: + meta: ``{component: {graph_name: meta}}`` from :meth:`get_metadata`. + example_inputs: ``{component: {graph_name: inputs}}`` from + :meth:`get_example_inputs`. + extra_options: Reads ``embedding_quantize`` to decide the token dtype. + + Returns: + A ``ModelInference`` wrapping the calibration graph's + ``DecoderInference`` plus an ``EncoderInference``, or ``None`` when no + calibration graph is present. + """ + from executorch.backends.qualcomm.genai_pipeline.artifact_keys import ( + ARTIFACT_TEXT_DECODER, + ) + from executorch.backends.qualcomm.genai_pipeline.graph_names import ( + GRAPH_FORWARD, + ) + from executorch.examples.qualcomm.oss_scripts.llama.inference import ( + DecoderInference, + EncoderInference, + ModelInference, + ) + + calibration_meta = meta.get(ARTIFACT_TEXT_DECODER, {}).get(GRAPH_FORWARD) + calibration_inputs = example_inputs.get(ARTIFACT_TEXT_DECODER, {}).get( + GRAPH_FORWARD + ) + if calibration_meta is None or calibration_inputs is None: + return None + + extra_options = extra_options or {} + use_i64_token = extra_options.get("embedding_quantize") is not None + + return ModelInference( + decoder=DecoderInference( + get_example_inputs=lambda: calibration_inputs, + audio_token_id=calibration_meta.get("audio_token_id"), + image_token_id=calibration_meta.get("image_token_id"), + max_context_len=calibration_meta["get_max_context_len"], + max_batch_size=calibration_meta["get_max_batch_size"], + use_i64_token=use_i64_token, + ), + encoder=EncoderInference(), + ) + + def load_tokenizer( + self, + model_name: str, + extra_options: Optional[Dict[str, Any]] = None, + ) -> Any: + """Build the model-level tokenizer wrapper.""" + from executorch.examples.qualcomm.oss_scripts.llama.tokenizer import ( + TokenizerWrapper, + ) + + # TODO: Remove this compatibility bridge once TokenizerWrapper accepts + # the pipeline's ``model`` and ``artifact_dir`` field names directly. + self.control_args.decoder_model = getattr( + self.control_args, "decoder_model", self.control_args.model + ) + self.control_args.artifact = getattr( + self.control_args, "artifact", self.control_args.artifact_dir + ) + + return TokenizerWrapper(self.control_args, self.model_config) + + def export_tokenizer( + self, + tokenizer: Any, + output_dir: Path, + extra_options: Optional[Dict[str, Any]] = None, + ) -> Path: + """Return the runtime tokenizer file exported by ``TokenizerWrapper``.""" + runtime_tokenizer_path = Path(tokenizer.runtime_tokenizer_path) + artifact_dir = Path(tokenizer.artifact) + artifacts = list(artifact_dir.iterdir()) if artifact_dir.is_dir() else [] + if runtime_tokenizer_path not in artifacts: + artifacts.append(runtime_tokenizer_path) + return self._select_runtime_tokenizer(artifacts) + + def _build_custom_decoder_model( + self, + model_arch: Any, + ) -> Dict[str, Any]: + """Construct the decoder modules, without weights. + + ``model_arch`` is the decoder's ``{graph_name: constructor}`` map (already + unwrapped from the component axis by :meth:`load_model`); each constructor + is a ``partial`` with its config and every arch-specific kwarg pre-bound by + ``get_model_arch``, so it is called with no arguments. + """ + modules: Dict[str, Any] = {} + for name, build in model_arch.items(): + modules[name] = build().eval() + return modules + + def _load_decoder_state_dict( + self, + state_dict_loader: Optional[Callable[[str], Dict[str, Any]]] = None, + ) -> Dict[str, Any]: + """Acquire the decoder's raw checkpoint state dict. + + The checkpoint is whatever ``control_args`` names, or the model's + ``repo_id`` downloaded and converted when it names none. Registry-specific + HF loaders can be passed through ``model_options["state_dict_loader"]`` + and must accept only ``repo_id``. + """ + import torch + + checkpoint = self.control_args.checkpoint + if checkpoint is None: + if self.model_config.repo_id is None: + raise ValueError( + f"'{self.control_args.model}' is supplied as a local " + "checkpoint; pass one via control_args.checkpoint." + ) + if state_dict_loader is not None: + return state_dict_loader(self.model_config.repo_id) + + from executorch.examples.models.llama.hf_download import ( + download_and_convert_hf_checkpoint, + ) + + checkpoint = download_and_convert_hf_checkpoint( + self.model_config.repo_id, self.model_config.convert_weights.__func__ + ) + + return torch.load(checkpoint, weights_only=True, map_location="cpu", mmap=True) + + def _load_auto_model(self) -> Any: + """Load the shared HuggingFace model backing the encoders and embedding.""" + import torch + from transformers import AutoModel, AutoModelForSpeechSeq2Seq + + repo_id = self.model_config.repo_id + if ARTIFACT_AUDIO_ENCODER in self._modalities: + auto_model = AutoModelForSpeechSeq2Seq.from_pretrained( + repo_id, _attn_implementation="eager" + ) + else: + auto_model = AutoModel.from_pretrained( + repo_id, _attn_implementation="eager" + ) + return auto_model.to(torch.float32).eval() + + @classmethod + def from_model_config( + cls, + model_config: Any, + *, + control_args: Any, + ) -> "MLLMLoaderAdapter": + """Build the adapter. + + Args: + model_config: The model's ``LLMModelConfig``. + control_args: The CLI argument namespace. + + Returns: + A new ``MLLMLoaderAdapter`` instance. + """ + return cls( + model_config=model_config, + control_args=control_args, + ) 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 fedefd63071..4b8ef89289c 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,16 @@ from __future__ import annotations from pathlib import Path -from typing import Any, Dict, Optional, Protocol, runtime_checkable, Tuple +from typing import ( + Any, + Callable, + Dict, + Optional, + Protocol, + runtime_checkable, + Sequence, + Tuple, +) @runtime_checkable @@ -44,10 +53,17 @@ class ModelLoaderAdapter(Protocol): contract, so the registry columns land with the transforms themselves once each is extracted into a named, shared function. - A single call returns **one** model. Multi-graph export configurations of the - same weights (hybrid prefill/decode, a separate token-embedding graph) are - expanded by the quantization and compilation *strategies*, and genuinely - multi-module models (multimodal encoders) belong in a specialized adapter. + ``load_model`` returns the full graph map: + ``{component: {graph_name: module}}``. Text-only models are the one-component + ``{ARTIFACT_TEXT_DECODER: {...}}`` case; multimodal adapters add encoder and + token-embedding components. Example inputs and metadata retain this shape + because they vary by graph. + + The strategy drops deployed graph modules and flattens the selected weight + holders to ``{component: module}``, for example ``{ARTIFACT_TEXT_DECODER: decoder}``. + The graph axis is absent because the selected module is shared when + exporting every graph variant for that component; only its inputs and + metadata differ. """ def load_model( @@ -55,14 +71,45 @@ def load_model( model_name: str, extra_options: Optional[Dict[str, Any]] = None, ) -> Any: - """Load a model by name or path. + """Load the model, weights and all. + + The whole weight stage lives here: acquiring the checkpoint, rewriting it + (key renames, value transforms) and loading it into the module. The + module this returns therefore already carries its final weights; + :meth:`apply_module_transforms` then handles the module-level preparation. Args: model_name: Model identifier (e.g., HuggingFace model ID or local path). extra_options: Additional model loading options (dtype, device_map, etc.). Returns: - The loaded nn.Module. + The weight-carrying ``{component: {graph_name: module}}`` map. + """ + ... + + def apply_module_transforms( + self, + module: Any, + module_transforms: Optional[Sequence[Callable[[Any], Any]]] = None, + ) -> Any: + """Apply transforms to one component module. + + The stage that owns the modules calls this, not :meth:`load_model`: a + module transform may *replace* a module (``convert_linear_to_conv2d`` + returns a new one) rather than mutate it, so the caller must own the + result. Implementations whose models need no preparation return the + modules unchanged. + + The weight stage is not here -- it runs inside :meth:`load_model`, the + only holder of the checkpoint -- so no caller ever sees a state dict. + + Args: + module: One value from the single-level ``{component: module}`` map + selected by the strategy. + module_transforms: Transforms to apply, in order. + + Returns: + The transformed module. Callers must use the return value. """ ... @@ -98,7 +145,7 @@ def get_example_inputs( (``LLMWrapper.attn_mask`` returns ``example_input[1]``). Args: - model: The module previously returned by :meth:`load_model`. + model: One graph module from ``load_model``'s nested map. extra_options: Additional options controlling the example shapes. Returns: @@ -107,6 +154,57 @@ def get_example_inputs( """ ... + def get_metadata( + self, + module: Any, + ) -> Any: + """Read one graph module's constant metadata for the ``.pte``. + + The metadata (``get_n_layers``, ``get_head_dim``, ``get_max_context_len``, + ...) is baked into the compiled binary as constant methods and is also + needed downstream to reconstruct logits / KV-cache shapes during + quantization encoding reconciliation. + + Args: + module: One graph module from the component/graph map returned by + :meth:`load_model`. + + Returns: + The graph metadata, or an empty value when the module exposes none. + The strategy owns routing this into ``{component: {graph_name: meta}}``. + """ + ... + + def get_inference( + self, + meta: Any, + example_inputs: Any, + extra_options: Optional[Dict[str, Any]] = None, + ) -> Any: + """Build the ``ModelInference`` that drives PTQ calibration. + + The inference object is bound to the calibration graph -- the + non-deployed, KV-cache graph whose observers calibration populates. The + strategy passes the full per-graph metadata and example inputs; the loader + selects its calibration graph and assembles the ``ModelInference``, since + the inference shape (decoder-only vs decoder + encoder) is model-family + specific. + + Args: + meta: The full per-graph metadata from :meth:`get_metadata` + (``{component: {graph_name: meta}}`` or ``{graph_name: meta}``). + example_inputs: The full per-graph example inputs from + :meth:`get_example_inputs`, mirroring ``meta``'s shape. + extra_options: Additional options; reads ``embedding_quantize`` to + decide the token dtype. + + Returns: + A ``ModelInference`` wrapping a ``DecoderInference`` for the + calibration graph, or ``None`` when the model has no calibration + driver. + """ + ... + 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 5ce3a5f2200..eeee1125d4a 100644 --- a/backends/qualcomm/genai_pipeline/strategies/quantization/default_quantizer_adapter.py +++ b/backends/qualcomm/genai_pipeline/strategies/quantization/default_quantizer_adapter.py @@ -17,72 +17,10 @@ class DefaultQuantizerAdapter: """Default adapter delegating to real ExecuTorch/QNN quantization APIs. - Wraps ``export_utils.make_quantizer``, ``torchao.quantization.pt2e.prepare_pt2e``, - and ``torchao.quantization.pt2e.convert_pt2e`` for production use. + Wraps ``torchao.quantization.pt2e.prepare_pt2e`` and + ``torchao.quantization.pt2e.convert_pt2e`` for production use. """ - def make_quantizer( - self, - 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. ``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: - A configured ``QnnQuantizer`` instance. - """ - from executorch.backends.qualcomm.export_utils import ( - make_quantizer as _make_quantizer, - ) - - # export_utils.make_quantizer expects soc_model as a string for - # getattr(QcomChipset, soc_model) lookup. Normalize enum → string. - 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, recipe=%s", - quant_dtype, - backend, - soc_model_str, - quant_recipe, - ) - 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, model: Any, @@ -98,7 +36,7 @@ def export_model( The exported module (``ExportedProgram.module()``). """ logger.debug("Exporting model via torch.export.export") - return torch.export.export(model, sample_input, strict=False).module() + return torch.export.export(model, sample_input, strict=True).module() def prepare_pt2e( self, @@ -121,27 +59,63 @@ def prepare_pt2e( logger.debug("Preparing model for PT2E quantization") return _prepare_pt2e(model, quantizer) + def init_encodings( + self, + module: Any, + example_inputs: Any, + ) -> Any: + """Initialize a graph's observers with one dummy forward. + + Args: + module: The annotated graph module. + example_inputs: This graph's positional example-input tuple. + + Returns: + The module after the dummy forward. + """ + logger.debug("Initializing encodings with a dummy forward") + with torch.no_grad(): + module(*example_inputs) + return module + def calibrate( self, model: Any, calibration_data: Iterable[Any], + **kwargs: Any, ) -> Any: """Run calibration data through the annotated model. - One forward pass per sample. Adapters needing a stateful procedure -- - e.g. autoregressive LLM calibration, where each step's input depends on - the previous step's output and the KV cache mutates across steps -- - should override this method rather than passing a callable as data. + One forward pass per sample. The default adapter supports both the + legacy single-module shape and the strategy's component-map shape: + ``{component: module}`` is paired with ``{component: iterable}``, and + components without calibration data are left untouched. Adapters needing + a stateful procedure -- e.g. autoregressive LLM calibration, where each + step's input depends on the previous step's output and the KV cache + mutates across steps -- should override this method. Args: - model: The annotated model with observers. - calibration_data: Any ``Iterable[Tuple[Tensor, ...]]``, including a - plain list or a ``DataLoader``. + model: The annotated model with observers, or a component-keyed map + of selected calibration graph modules. + calibration_data: Any ``Iterable[Tuple[Tensor, ...]]`` for a single + module, or a component-keyed map of such iterables. Plain lists + and ``DataLoader`` instances are both accepted. + **kwargs: Extra adapter-specific options. Ignored by the default + adapter; model-specific adapters may read keys such as + ``inference``. Returns: - The calibrated model. + The calibrated model or component map. """ logger.debug("Running calibration") + if isinstance(model, dict): + with torch.no_grad(): + for component, module in model.items(): + component_data = calibration_data.get(component, ()) + for data in component_data: + module(*data) + return model + with torch.no_grad(): for data in calibration_data: model(*data) 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 33102f89d3d..e226f81d546 100644 --- a/backends/qualcomm/genai_pipeline/strategies/quantization/executorch_quantization_strategy.py +++ b/backends/qualcomm/genai_pipeline/strategies/quantization/executorch_quantization_strategy.py @@ -6,8 +6,19 @@ from __future__ import annotations +import gc import logging -from typing import Optional +from enum import auto, Enum +from typing import Any, Optional, TYPE_CHECKING + +import torch +from executorch.backends.qualcomm.genai_pipeline.artifact_keys import ( + ARTIFACT_AUDIO_ENCODER, + ARTIFACT_TEXT_DECODER, + ARTIFACT_TEXT_ENCODER, + ARTIFACT_TOK_EMBEDDING, + ARTIFACT_VISION_ENCODER, +) from executorch.backends.qualcomm.genai_pipeline.configs.quantization_input_config import ( QuantizationInputConfig, @@ -16,6 +27,14 @@ QuantizationOutputConfig, ) from executorch.backends.qualcomm.genai_pipeline.exceptions import StageError +from executorch.backends.qualcomm.genai_pipeline.graph_bundle import GraphBundle +from executorch.backends.qualcomm.genai_pipeline.graph_names import ( + DECODER_GRAPH_NAMES, + GRAPH_FORWARD, + GRAPH_KV_FORWARD, + GRAPH_PREFILL_FORWARD, + TOK_EMBEDDING_GRAPH_NAMES, +) from executorch.backends.qualcomm.genai_pipeline.pipeline_context import PipelineContext from executorch.backends.qualcomm.genai_pipeline.strategies.quantization.quantization_strategy import ( QuantizationStrategy, @@ -24,55 +43,94 @@ QuantizerAdapter, ) +if TYPE_CHECKING: + from executorch.backends.qualcomm.genai_pipeline.datasets.calibration.calibration_data_adapter import ( + CalibrationDataAdapter, + ) + from executorch.backends.qualcomm.genai_pipeline.datasets.evaluation.evaluation_data_adapter import ( + EvaluationDataAdapter, + ) + from executorch.backends.qualcomm.genai_pipeline.datasets.training.training_data_adapter import ( + TrainingDataAdapter, + ) + logger = logging.getLogger(__name__) _STAGE_NAME = "quantization" +class _GraphRole(Enum): + # Uses real calibration data but is never lowered for deployment. + QUANTIZE = auto() + # Initializes encodings, receives overrides from a quantize graph, then lowers. + DEPLOY = auto() + # Uses calibration data and lowers as the same graph; no encoding initialization. + SHARED = auto() + + +_GRAPH_ROLES = { + # The quantize graph supplies encodings for the decoder's deployed variants. + ARTIFACT_TEXT_DECODER: { + GRAPH_FORWARD: _GraphRole.QUANTIZE, + **dict.fromkeys(DECODER_GRAPH_NAMES, _GraphRole.DEPLOY), + }, + # The quantize graph supplies encodings for deployed token-embedding variants. + ARTIFACT_TOK_EMBEDDING: { + GRAPH_FORWARD: _GraphRole.QUANTIZE, + **dict.fromkeys(TOK_EMBEDDING_GRAPH_NAMES, _GraphRole.DEPLOY), + }, + # Audio-encoder graph is quantized and lowered for deployment. + ARTIFACT_AUDIO_ENCODER: {GRAPH_FORWARD: _GraphRole.SHARED}, + # Text-encoder graph is quantized and lowered for deployment. + ARTIFACT_TEXT_ENCODER: {GRAPH_FORWARD: _GraphRole.SHARED}, + # Vision-encoder graph is quantized and lowered for deployment. + ARTIFACT_VISION_ENCODER: {GRAPH_FORWARD: _GraphRole.SHARED}, +} + + class ExecuTorchQuantizationStrategy(QuantizationStrategy): """ExecuTorch-based quantization using QNN quantizer annotator 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. + Delegates single-graph PT2E operations to ``QuantizerAdapter`` and owns the + component/graph routing. Each graph has one of three roles: + + - ``QUANTIZE``: Runs on real calibration data to produce encodings, then is + removed without being deployed. + - ``DEPLOY``: Runs once to initialize observers, receives encodings from its + component's quantize graph, then is lowered for deployment. + - ``SHARED``: Runs on calibration data and is lowered as the same graph. + + The quantization flow mirrors :meth:`invoke`: + + 1. Create QNN quantizers and recipe instances for every component graph. + 2. Export graph variants and prepare them with observers. + 3. Initialize observers for separate deploy graphs. + 4. Assemble real-data inputs for native PTQ calibration. + 5. Quantize ``QUANTIZE`` and ``SHARED`` graphs. + 6. Convert prepared graphs to QDQ modules. + 7. Save decoder QDQ output and override quantize-graph encodings to deploy + variants, then release quantize-only graphs. Args: - quantizer_adapter: Injectable adapter for quantization operations. - Defaults to ``DefaultQuantizerAdapter`` if not provided. + quantizer_adapter: Injectable adapter for single-graph quantization + operations. Defaults to ``DefaultQuantizerAdapter`` if not provided. + calibration_data_adapter: Injectable purpose adapter that assembles + calibration data. Defaults to ``DefaultCalibrationDataAdapter`` (random + fallback) when not provided. + training_data_adapter: Injectable adapter for QAT training data. Retained + for the QAT flow; the current native PTQ implementation does not + consume it. + evaluation_data_adapter: Injectable adapter for post-quantization + evaluation data. Retained for evaluation integration; the current + strategy does not run evaluation or consume it. """ def __init__( self, quantizer_adapter: Optional[QuantizerAdapter] = None, + calibration_data_adapter: Optional["CalibrationDataAdapter"] = None, + training_data_adapter: Optional["TrainingDataAdapter"] = None, + evaluation_data_adapter: Optional["EvaluationDataAdapter"] = None, ) -> None: if quantizer_adapter is None: from executorch.backends.qualcomm.genai_pipeline.strategies.quantization.default_quantizer_adapter import ( @@ -81,6 +139,27 @@ def __init__( quantizer_adapter = DefaultQuantizerAdapter() self._adapter = quantizer_adapter + if calibration_data_adapter is None: + from executorch.backends.qualcomm.genai_pipeline.datasets.calibration.default_calibration_data_adapter import ( + DefaultCalibrationDataAdapter, + ) + + calibration_data_adapter = DefaultCalibrationDataAdapter() + self._calibration_data_adapter = calibration_data_adapter + if training_data_adapter is None: + from executorch.backends.qualcomm.genai_pipeline.datasets.training.default_training_data_adapter import ( + DefaultTrainingDataAdapter, + ) + + training_data_adapter = DefaultTrainingDataAdapter() + self._training_data_adapter = training_data_adapter + if evaluation_data_adapter is None: + from executorch.backends.qualcomm.genai_pipeline.datasets.evaluation.default_evaluation_data_adapter import ( + DefaultEvaluationDataAdapter, + ) + + evaluation_data_adapter = DefaultEvaluationDataAdapter() + self._evaluation_data_adapter = evaluation_data_adapter @property def adapter(self) -> QuantizerAdapter: @@ -94,19 +173,23 @@ def invoke( ) -> QuantizationOutputConfig: """Quantize the model using ExecuTorch/QNN quantization. - Executes the full PT2E quantization pipeline: - export → make_quantizer → prepare_pt2e → calibrate → convert_pt2e. + The native PTQ flow creates graph quantizers, exports and prepares every + graph, initializes deploy-graph observers, quantizes graphs that consume + calibration data, converts to QDQ, and propagates encodings to deploy + variants. Args: context: The pipeline context with global settings. input_config: The quantization input configuration. Returns: - QuantizationOutputConfig with the quantized model. + Component- and graph-keyed bundles for deployment graphs. Shared + graphs remain in the output; quantize-only graphs are removed after + encoding propagation. Raises: - StageError: If the model, example inputs or calibration data is - missing, or if any quantization step fails. + StageError: If required input is missing or any quantization step + fails. """ logger.info( "Starting quantization for model '%s' on SoC=%s, backend=%s", @@ -117,70 +200,56 @@ def invoke( 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, - ) + extra = input_config.extra_options or {} + quant_options = dict(extra.get("quantize_options") or {}) - # 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 + # Step 1: Create QNN quantizers and per-graph recipe instances. + quantizers, quant_recipes = self._make_quantizer( + input_config, + quant_options, ) - # 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 + # Step 2: Export graph variants and prepare them with observers. + prepared_modules = self._export_and_prepare(input_config, quantizers) - quantizer = self._adapter.make_quantizer(**make_quantizer_kwargs) + # Step 3: Initialize deploy-graph observers for QDQ conversion. + self._initialize_encodings(prepared_modules, input_config) - # 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 4: Assemble real-data inputs for native PTQ calibration. + calibration_data = self._calibration_data_adapter.generate_calibration_data( + tokenizer=input_config.tokenizer, + example_inputs=self._flatten_calibration_example_inputs( + input_config.example_inputs + ), ) - # Step 5: Convert to quantized model - logger.debug("Converting to quantized model") - quantized_model = self._adapter.convert_pt2e(calibrated_model) + # Step 5: Quantize graphs that consume calibration data. + self._calibrate(prepared_modules, input_config, calibration_data) + + # Step 6: Convert prepared graphs to QDQ modules. + converted_modules = self._convert_pt2e(prepared_modules) + + # Step 7: Override quantize-graph encodings to deploy variants. + self._override_encodings(converted_modules, input_config, context) logger.info("Quantization completed successfully") - return QuantizationOutputConfig(quantized_model=quantized_model) + + # Build deployment graph bundles. + graphs = {} + for component, converted_graphs in converted_modules.items(): + graphs[component] = { + graph_name: GraphBundle( + module=module, + inputs=input_config.example_inputs[component][graph_name], + meta=input_config.meta.get(component, {}).get(graph_name, {}), + quant_io_dtypes=self._get_quant_io_dtypes( + quant_recipes[component][graph_name] + ), + ) + for graph_name, module in converted_graphs.items() + } + return QuantizationOutputConfig(graphs=graphs) except StageError: raise @@ -191,6 +260,363 @@ def invoke( original_exception=e, ) from e + def _make_quantizer( + self, + input_config: QuantizationInputConfig, + quant_options: dict, + ) -> tuple[dict, dict]: + """Build QNN quantizers and recipe instances for every component graph. + + Args: + input_config: Model and backend configuration for quantization. + quant_options: Component-keyed dtype and recipe configuration. + + Returns: + A ``(quantizers, recipes)`` tuple, each keyed by component then + graph name. Recipe classes are instantiated per graph; only the + ``GRAPH_FORWARD`` recipe instance enables verbose output. + """ + from executorch.backends.qualcomm.genai_pipeline.quant_utilities import ( + make_quantizer, + ) + + quant_dtype = quant_options.get("quant_dtype") or {} + quant_recipe = quant_options.get("quant_recipe") or {} + quantizers = {} + recipe_instances = {} + for component, component_example_inputs in input_config.example_inputs.items(): + quantizers[component] = {} + recipe_instances[component] = {} + recipe_class = quant_recipe.get(component) + for graph_name in component_example_inputs: + recipe = ( + recipe_class(verbose=graph_name == GRAPH_FORWARD) + if isinstance(recipe_class, type) + else recipe_class + ) + make_quantizer_kwargs = { + "backend": input_config.backend_type, + "soc_model": input_config.soc_model, + **( + {"quant_dtype": dtype} + if (dtype := quant_dtype.get(component)) is not None + else {} + ), + **({"quant_recipe": recipe} if recipe is not None else {}), + } + quantizers[component][graph_name] = make_quantizer( + **make_quantizer_kwargs + ) + recipe_instances[component][graph_name] = recipe + if recipe is not None: + logger.info( + "Created QNN quantizer for '%s' graph '%s' with quant " + "recipe %s", + component, + graph_name, + type(recipe).__name__, + ) + else: + logger.info( + "Quant recipe not set yet; created QNN quantizer for " + "'%s' graph '%s' with quant dtype %s", + component, + graph_name, + make_quantizer_kwargs.get("quant_dtype"), + ) + return quantizers, recipe_instances + + def _export_and_prepare( + self, + input_config: QuantizationInputConfig, + quantizers: dict, + ) -> dict: + """Export every graph and prepare it for PT2E quantization. + + Text-decoder example inputs are flattened into the positional export + signature before export. Other component inputs pass through unchanged. + + Args: + input_config: Model modules, graph inputs, and graph metadata. + quantizers: Component- and graph-keyed QNN quantizers. + + Returns: + Component- and graph-keyed prepared PT2E modules. + """ + prepared_modules = {} + for component, component_example_inputs in input_config.example_inputs.items(): + prepared_modules[component] = {} + for graph_name, graph_example_inputs in component_example_inputs.items(): + export_inputs = self._post_process_example_inputs( + component, + graph_name, + graph_example_inputs, + input_config.meta, + ) + prepared_modules[component][graph_name] = self._adapter.prepare_pt2e( + self._adapter.export_model( + input_config.model_module[component], export_inputs + ), + quantizers[component][graph_name], + ) + return prepared_modules + + @staticmethod + def _post_process_example_inputs( + component: str, + graph_name: str, + example_inputs: Any, + meta: dict, + ) -> Any: + """Flatten text-decoder inputs into their positional export signature. + + Args: + component: Artifact component key for the graph. + graph_name: Graph key used to retrieve metadata. + example_inputs: Model-preparation inputs in structured decoder form. + meta: Component- and graph-keyed model metadata. + + Returns: + Unchanged inputs for non-decoder components. For text decoder graphs, + positional inputs with attention masks expanded and position IDs/KV + caches included only when ``get_use_kv_cache`` is enabled. + + Raises: + StageError: If a text decoder graph lacks ``get_use_kv_cache`` + metadata. + """ + if component != ARTIFACT_TEXT_DECODER: + return example_inputs + + if "get_use_kv_cache" not in ( + graph_meta := meta.get(component, {}).get(graph_name, {}) + ): + raise StageError( + stage_name=_STAGE_NAME, + message=( + f"Component '{component}' graph '{graph_name}' is missing " + "get_use_kv_cache metadata" + ), + ) + + use_kv_cache = graph_meta["get_use_kv_cache"] + return ( + example_inputs[0], + *example_inputs[1], + *((example_inputs[2],) if use_kv_cache else []), + *(example_inputs[3] if use_kv_cache else []), + *(example_inputs[4] if use_kv_cache else []), + ) + + def _initialize_encodings( + self, + prepared_modules: dict, + input_config: QuantizationInputConfig, + ) -> None: + """Initialize deployed graph observers with their export signatures. + + Deploy graphs do not consume the corpus calibration data because their + encodings are later overridden from the component's quantize graph. They + still require one forward pass so ``convert_pt2e`` can produce QDQ + graphs. + + Args: + prepared_modules: Component- and graph-keyed prepared PT2E modules. + input_config: Graph inputs and metadata used for observer execution. + """ + for component, prepared_graphs in prepared_modules.items(): + for graph_name, graph_module in prepared_graphs.items(): + if self._is_deploy_graph(component, graph_name): + self._adapter.init_encodings( + graph_module, + self._post_process_example_inputs( + component, + graph_name, + input_config.example_inputs[component][graph_name], + input_config.meta, + ), + ) + + def _calibrate( + self, + prepared_modules: dict, + input_config: QuantizationInputConfig, + calibration_data: dict, + ) -> None: + """Run the quantization algorithm on graphs that consume real data. + + ``QUANTIZE`` graphs produce encodings for separate deploy variants; + ``SHARED`` graphs produce their own encodings and are also deployed. The + implementation currently supports only native PTQ calibration through + the model-specific adapter. + + TODO: Support legacy quantization algorithms such as QAT and SeqMSE. + + Args: + prepared_modules: Component- and graph-keyed prepared PT2E modules. + input_config: Model-specific inference support for calibration. + calibration_data: Component-keyed real-data calibration inputs. + """ + quantization_graphs = {} + for component, prepared_graphs in prepared_modules.items(): + for graph_name, graph_module in prepared_graphs.items(): + if not self._is_deploy_graph(component, graph_name): + quantization_graphs[component] = graph_module + self._adapter.calibrate( + quantization_graphs, + calibration_data, + inference=input_config.inference, + ) + + def _convert_pt2e(self, prepared_modules: dict) -> dict: + """Convert prepared PT2E modules to QDQ modules and release inputs. + + Args: + prepared_modules: Component- and graph-keyed prepared PT2E modules. + + Returns: + Component- and graph-keyed converted QDQ modules. + """ + converted_modules = {} + for component, prepared_graphs in prepared_modules.items(): + converted_modules[component] = { + graph_name: self._adapter.convert_pt2e(graph_module) + for graph_name, graph_module in prepared_graphs.items() + } + prepared_modules.clear() + gc.collect() + return converted_modules + + def _override_encodings( + self, + converted_modules: dict, + input_config: QuantizationInputConfig, + context: PipelineContext, + ) -> None: + """Override quantize-graph encodings to separate deploy variants. + + For text decoder, saves the converted ``GRAPH_FORWARD`` QDQ module, + then overrides decode and optional prefill graph encodings using cache + layer metadata from the decode graph. For token embedding, overrides + each configured deployment variant. Quantize-only graphs are removed + after they have served as encoding sources. + + Args: + converted_modules: Component- and graph-keyed converted QDQ modules. + input_config: Decoder inputs and metadata needed for QDQ export and + encoding propagation. + context: Pipeline context providing the QDQ artifact directory. + + Raises: + StageError: If a required quantize graph, deployed decoder graph, or + decoder cache-layer metadata is missing. + """ + from executorch.backends.qualcomm.genai_pipeline.quant_utilities import ( + encoding_override, + save_logits_quant_attrs, + save_output_kv_cache_quant_attrs, + save_quantized_module, + ) + + if (decoder_graphs := converted_modules.get(ARTIFACT_TEXT_DECODER)) is not None: + if (quantized_decoder := decoder_graphs.get(GRAPH_FORWARD)) is None: + raise StageError( + stage_name=_STAGE_NAME, + message=( + f"Component '{ARTIFACT_TEXT_DECODER}' is missing its " + f"quantization graph '{GRAPH_FORWARD}'" + ), + ) + try: + # Saving Decoder QDQ Model EP + save_quantized_module( + quantized_module=quantized_decoder, + example_inputs=self._post_process_example_inputs( + ARTIFACT_TEXT_DECODER, + GRAPH_FORWARD, + input_config.example_inputs[ARTIFACT_TEXT_DECODER][ + GRAPH_FORWARD + ], + input_config.meta, + ), + artifact_dir=context.artifact_dir, + ) + + # Override decoder quant encodings + if (decoder := decoder_graphs.get(GRAPH_KV_FORWARD)) is None: + raise StageError( + stage_name=_STAGE_NAME, + message=( + f"Component '{ARTIFACT_TEXT_DECODER}' is missing " + f"its deployed decoder graph '{GRAPH_KV_FORWARD}'" + ), + ) + decoder_meta = input_config.meta.get(ARTIFACT_TEXT_DECODER, {}).get( + GRAPH_KV_FORWARD, {} + ) + if ( + n_cache_layers := decoder_meta.get("get_n_self_layers") + or decoder_meta.get("get_n_layers") + ) is None: + raise StageError( + stage_name=_STAGE_NAME, + message=( + f"Component '{ARTIFACT_TEXT_DECODER}' graph " + f"'{GRAPH_KV_FORWARD}' requires n_cache_layers metadata " + "for encoding override" + ), + ) + encoding_override( + quantized_model=quantized_decoder, + unquantized_model=decoder, + n_cache_layers=n_cache_layers, + ) + save_logits_quant_attrs(decoder, decoder_meta) + save_output_kv_cache_quant_attrs(decoder, decoder_meta) + + # Override prefill quant encodings + if (prefill := decoder_graphs.get(GRAPH_PREFILL_FORWARD)) is not None: + prefill_meta = input_config.meta.get(ARTIFACT_TEXT_DECODER, {}).get( + GRAPH_PREFILL_FORWARD, {} + ) + encoding_override( + quantized_model=quantized_decoder, + unquantized_model=prefill, + n_cache_layers=n_cache_layers, + ) + save_logits_quant_attrs(prefill, prefill_meta) + save_output_kv_cache_quant_attrs(prefill, prefill_meta) + finally: + # The quantization graph only sources encodings; never deployed. + decoder_graphs.pop(GRAPH_FORWARD, None) + gc.collect() + + # Override token embedding quant encodings + if ( + tok_embedding_graphs := converted_modules.get(ARTIFACT_TOK_EMBEDDING) + ) is None: + return + if (quantized_tok_embedding := tok_embedding_graphs.get(GRAPH_FORWARD)) is None: + raise StageError( + stage_name=_STAGE_NAME, + message=( + f"Component '{ARTIFACT_TOK_EMBEDDING}' is missing its " + f"quantization graph '{GRAPH_FORWARD}'" + ), + ) + try: + for graph_name in TOK_EMBEDDING_GRAPH_NAMES: + if (tok_embedding := tok_embedding_graphs.get(graph_name)) is None: + continue + encoding_override( + quantized_model=quantized_tok_embedding, + unquantized_model=tok_embedding, + ) + finally: + # The quantization graph only sources encodings; never deployed. + tok_embedding_graphs.pop(GRAPH_FORWARD, None) + gc.collect() + def _validate_input(self, input_config: QuantizationInputConfig) -> None: """Validate required fields in the input configuration. @@ -213,14 +639,6 @@ def _validate_input(self, input_config: QuantizationInputConfig) -> None: "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, @@ -231,3 +649,109 @@ def _validate_input(self, input_config: QuantizationInputConfig) -> None: stage_name=_STAGE_NAME, message="backend_type is required for quantization", ) + if isinstance(input_config.backend_type, dict): + raise StageError( + stage_name=_STAGE_NAME, + message="backend_type must be one shared QNN backend, not a map", + ) + + def _flatten_calibration_example_inputs(self, example_inputs: dict) -> dict: + """Flatten ``{component: {graph_name: inputs}}`` to ``{component: inputs}``. + + Selects the first graph per component that participates in real-data + quantization (``QUANTIZE`` or ``SHARED``) and drops the graph-name + level. Separate deploy graphs are excluded because their KV-cache export + signatures must not shape the collator's attention-mask template. + + Args: + example_inputs: Component- and graph-keyed model example inputs. + + Returns: + Component-keyed inputs for calibration data construction. + """ + calibration_inputs = {} + for component, graphs in example_inputs.items(): + for graph_name, inputs in graphs.items(): + if not self._is_deploy_graph(component, graph_name): + calibration_inputs[component] = inputs + break + return calibration_inputs + + def _is_deploy_graph(self, component: str, graph_name: str) -> bool: + """Whether a graph is a separate deployment variant. + + Shared graphs participate in quantization and deployment, so only + ``_GraphRole.DEPLOY`` variants need observer initialization and encoding + override from a separate quantization graph. + + Args: + component: Artifact component key. + graph_name: Graph key within the component. + + Returns: + ``True`` when the graph is a separate deployment variant. + + Raises: + StageError: If the component or graph lacks a role definition. + """ + if (component_roles := _GRAPH_ROLES.get(component)) is None: + raise StageError( + stage_name=_STAGE_NAME, + message=f"Component '{component}' has no graph-role definition", + ) + if (role := component_roles.get(graph_name)) is None: + raise StageError( + stage_name=_STAGE_NAME, + message=f"Component '{component}' has no role for graph '{graph_name}'", + ) + return role is _GraphRole.DEPLOY + + def _get_quant_io_dtypes(self, quant_recipe: Any) -> Optional[dict]: + """The ``{"kv_type", "io_type"}`` dtypes a graph's IO is tagged with. + + KV width comes from ``get_kv_io_bit_width`` when available. Output width + prefers ``get_logits_output_bit_width`` and falls back to + ``get_act_bit_width`` for recipes without a logits-specific method. + Unsupported widths are rejected explicitly; 8-bit logits IO is not + implemented. + + Args: + quant_recipe: Per-graph recipe instance, if one is configured. + + Returns: + A map containing available ``kv_type`` and ``io_type`` dtypes, or + ``None`` when the recipe exposes neither width. + + Raises: + NotImplementedError: If logits output width is 8 bits. + RuntimeError: If a reported IO width is unsupported. + """ + kv_bit_width = ( + quant_recipe.get_kv_io_bit_width() + if quant_recipe is not None and hasattr(quant_recipe, "get_kv_io_bit_width") + else None + ) + io_bit_width = None + if quant_recipe is not None: + if hasattr(quant_recipe, "get_logits_output_bit_width"): + io_bit_width = quant_recipe.get_logits_output_bit_width() + elif hasattr(quant_recipe, "get_act_bit_width"): + io_bit_width = quant_recipe.get_act_bit_width() + width_to_dtype = {8: torch.uint8, 16: torch.uint16} + + if io_bit_width == 8: + raise NotImplementedError(f"unknown io bit width {io_bit_width}") + + quant_io_dtypes = {} + for bit_width, dtype_key in ( + (kv_bit_width, "kv_type"), + (io_bit_width, "io_type"), + ): + if bit_width in width_to_dtype: + quant_io_dtypes[dtype_key] = width_to_dtype[bit_width] + elif bit_width is not None: + raise RuntimeError( + f"Unsupported quantization IO bit width: {bit_width}" + ) + + return quant_io_dtypes or None diff --git a/backends/qualcomm/genai_pipeline/strategies/quantization/llm_quantizer_adapter.py b/backends/qualcomm/genai_pipeline/strategies/quantization/llm_quantizer_adapter.py new file mode 100644 index 00000000000..3b8b4fb391d --- /dev/null +++ b/backends/qualcomm/genai_pipeline/strategies/quantization/llm_quantizer_adapter.py @@ -0,0 +1,148 @@ +# 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. + +"""LLMQuantizerAdapter: quantization adapter for LLMs. + +The strategy selects the recipe and dtype for each component. This adapter +creates one quantizer from those scalar options and applies the supplied recipe. + +Most methods operate on a **single graph module**. ``calibrate`` receives the +single-level ``{component: module}`` map selected by the strategy so it can drive +the text decoder. Fanning out over graph variants remains the strategy's job. +""" + +from __future__ import annotations + +from typing import Any, Iterable + +import torch + + +class LLMQuantizerAdapter: + """Quantizer adapter for LLM models. + + Handles LLM-specific calibration and encoding propagation. + """ + + def export_model( + self, + module: Any, + example_inputs: Any, + ) -> Any: + """Export a single graph module using torch.export.export. + + Args: + module: The decoder module to export. + example_inputs: Positional example inputs for this graph. + + Returns: + The exported module. + """ + return torch.export.export(module, example_inputs, strict=True).module() + + def prepare_pt2e( + self, + module: Any, + quantizer: Any, + ) -> Any: + """Prepare a single exported module for PT2E quantization. + + Args: + module: The exported module. + quantizer: The QnnQuantizer instance. + + Returns: + The annotated module with observers inserted. + """ + from torchao.quantization.pt2e.quantize_pt2e import ( + prepare_pt2e as _prepare_pt2e, + ) + + return _prepare_pt2e(module, quantizer) + + def init_encodings( + self, + module: Any, + example_inputs: Any, + ) -> Any: + """Initialize a deployed graph's observers with one dummy forward. + + Args: + module: The annotated deployed graph module. + example_inputs: This graph's positional example-input tuple. + + Returns: + The module after the dummy forward. + """ + with torch.no_grad(): + module(*example_inputs) + return module + + def calibrate( + self, + modules: Any, + calibration_data: Iterable[Any], + **kwargs: Any, + ): + """Run true PTQ calibration over the selected decoder. + + ``modules`` is the single-level ``{component: module}`` map created by + the strategy. It has no graph axis because it contains only the selected + calibration graph for each component. + + Args: + modules: ``{ARTIFACT_TEXT_DECODER: decoder_module}``. + calibration_data: ``{component: DataLoader}`` corpus batches. + **kwargs: Extra adapter-specific options. Expects ``inference``, the + ``ModelInference`` bound to the calibration graph. + + Returns: + ``None``. The decoder module is calibrated in place. + """ + import torch + + from executorch.backends.qualcomm.genai_pipeline.artifact_keys import ( + ARTIFACT_TEXT_DECODER, + ) + from torch.utils.data import DataLoader + + inference = kwargs["inference"] + text_dataloader = calibration_data.get(ARTIFACT_TEXT_DECODER) + if not isinstance(text_dataloader, DataLoader): + raise ValueError( + "Calibration requires a corpus-backed DataLoader for " + f"{ARTIFACT_TEXT_DECODER};" + ) + + decoder_module = modules[ARTIFACT_TEXT_DECODER] + with torch.no_grad(): + for batch in text_dataloader: + inference.predict_step( + decoder_module, + input_ids=batch["input_ids"], + attn_mask=batch["attention_mask"], + ) + + def convert_pt2e( + self, + module: Any, + ) -> Any: + """Convert a single calibrated module to a quantized module. + + Cross-graph encoding override and the recording of quantized logits / + KV-cache attributes are performed by the strategy after this returns. + + Args: + module: The calibrated module. + + Returns: + The quantized module. + """ + from torchao.quantization.pt2e.quantize_pt2e import ( + convert_pt2e as _convert_pt2e, + ) + + return _convert_pt2e(module) diff --git a/backends/qualcomm/genai_pipeline/strategies/quantization/mllm_quantizer_adapter.py b/backends/qualcomm/genai_pipeline/strategies/quantization/mllm_quantizer_adapter.py new file mode 100644 index 00000000000..38a92f0f3f1 --- /dev/null +++ b/backends/qualcomm/genai_pipeline/strategies/quantization/mllm_quantizer_adapter.py @@ -0,0 +1,156 @@ +# 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. + +"""Quantizer adapter for multimodal models. + +Most methods operate on a **single graph module**. ``calibrate`` receives the +single-level ``{component: module}`` map selected by the strategy so it can +coordinate encoder, embedding, and decoder calibration. Graph fan-out and graph +selection remain the strategy's job, including choosing each component's recipe +and dtype before creating its quantizer. +""" + +from __future__ import annotations + +from typing import Any, Iterable + +import torch + + +class MLLMQuantizerAdapter: + """Quantizer adapter for multimodal models.""" + + def export_model( + self, + module: Any, + example_inputs: Any, + ) -> Any: + """Export a single graph module using torch.export.export. + + Args: + module: The component module to export. + example_inputs: Positional example inputs for this graph. + + Returns: + The exported module. + """ + return torch.export.export(module, example_inputs, strict=True).module() + + def prepare_pt2e( + self, + module: Any, + quantizer: Any, + ) -> Any: + """Prepare a single exported module for PT2E quantization. + + Args: + module: The exported module. + quantizer: The QnnQuantizer instance. + + Returns: + The annotated module with observers inserted. + """ + from torchao.quantization.pt2e.quantize_pt2e import ( + prepare_pt2e as _prepare_pt2e, + ) + + return _prepare_pt2e(module, quantizer) + + def init_encodings( + self, + module: Any, + example_inputs: Any, + ) -> Any: + """Initialize a deployed graph's observers with one dummy forward. + + Args: + module: The annotated deployed graph module. + example_inputs: This graph's positional example-input tuple. + + Returns: + The module after the dummy forward. + """ + with torch.no_grad(): + module(*example_inputs) + return module + + def calibrate( + self, + modules: Any, + calibration_data: Iterable[Any], + **kwargs: Any, + ) -> Any: + """Calibrate selected components in a ``{component: module}`` map. + + The graph axis has already been resolved by the strategy. This method + coordinates the encoder, embedding, and decoder modules for one + multimodal calibration path. + + Args: + modules: Selected encoder, embedding, and decoder modules. + calibration_data: Per-component calibration batches. + **kwargs: Extra adapter-specific options. Expects ``inference``, the + model-specific calibration driver. + """ + + from executorch.backends.qualcomm.genai_pipeline.artifact_keys import ( + ARTIFACT_AUDIO_ENCODER, + ARTIFACT_TEXT_DECODER, + ARTIFACT_TOK_EMBEDDING, + ARTIFACT_VISION_ENCODER, + ) + from executorch.examples.qualcomm.oss_scripts.llama.utils import ( + safe_dataloader_iter, + ) + + inference = kwargs["inference"] + audio_dataloader = calibration_data.get(ARTIFACT_AUDIO_ENCODER) + vision_dataloader = calibration_data.get(ARTIFACT_VISION_ENCODER) + text_dataloader = calibration_data[ARTIFACT_TEXT_DECODER] + + audio_encoder = modules.get(ARTIFACT_AUDIO_ENCODER) + vision_encoder = modules.get(ARTIFACT_VISION_ENCODER) + tok_embedding = modules.get(ARTIFACT_TOK_EMBEDDING) + text_decoder = modules.get(ARTIFACT_TEXT_DECODER) + encoder = audio_encoder or vision_encoder + + for _, (audio_batch, vision_batch, text_batch) in enumerate( + zip( + safe_dataloader_iter(audio_dataloader), + safe_dataloader_iter(vision_dataloader), + text_dataloader, + ) + ): + encoder_inputs = (audio_batch or vision_batch or {}).get("inputs") + inference.predict_step( + text_decoder, + input_ids=text_batch["input_ids"], + attn_mask=text_batch["attention_mask"], + tok_embedding=tok_embedding, + encoder_module=encoder, + encoder_inputs=encoder_inputs, + ) + + def convert_pt2e( + self, + module: Any, + ) -> Any: + """Convert a single calibrated module to a quantized module. + + Cross-graph encoding override and the recording of quantized logits / + KV-cache attributes are performed by the strategy after this returns. + + Args: + module: The calibrated module. + + Returns: + The quantized module. + """ + from torchao.quantization.pt2e.quantize_pt2e import ( + convert_pt2e as _convert_pt2e, + ) + + return _convert_pt2e(module) diff --git a/backends/qualcomm/genai_pipeline/strategies/quantization/quantizer_adapter.py b/backends/qualcomm/genai_pipeline/strategies/quantization/quantizer_adapter.py index c03b79554bf..d7257dd100e 100644 --- a/backends/qualcomm/genai_pipeline/strategies/quantization/quantizer_adapter.py +++ b/backends/qualcomm/genai_pipeline/strategies/quantization/quantizer_adapter.py @@ -13,118 +13,118 @@ class QuantizerAdapter(Protocol): """Protocol for quantization operations. - Wraps external quantization APIs (make_quantizer, prepare_pt2e, convert_pt2e) + Wraps external quantization APIs (prepare_pt2e, convert_pt2e) behind an injectable interface for testability. .. note:: - These methods operate on a **single graph**, mirroring the underlying - ``torchao`` PT2E APIs, and adapters stay a thin 1:1 wrapper over them. - - For models exported as several graphs from the same weights -- e.g. a - hybrid decoder -- only one graph is actually quantized: a dedicated - full-auto-regressive calibration graph, which yields the best activation - statistics but is never deployed. Its scales and zero points are then - propagated onto the deployed graphs (AR-N prefill, AR-1 decode), which do - not run PT2E themselves. That propagation is inherently cross-graph, so - it cannot live in a per-graph method here; the **quantization strategy** - sequences it via a separate reconciliation hook after this adapter - returns. + ``export_model``, ``prepare_pt2e``, ``init_encodings``, and + ``convert_pt2e`` each operate on one graph module. Their module arguments + are scalar values, never component or graph maps. + + All routing lives in the **quantization strategy**: it fans out over the + components and graph variants, decides which single graph is the one to + quantize (versus the deployed graphs that are only run once for their + observers), collects those graphs, drives their quantization, and then + propagates the resulting scales / zero points onto the deployed graphs. + ``calibrate`` is the exception: + it receives the single-level ``{component: module}`` map so model-family + inference can drive cross-component calibration. There is no graph axis + because the strategy has already selected one calibration graph per + component. """ - def make_quantizer( + def export_model( self, - quant_dtype: Any = None, - backend: Any = None, - soc_model: Any = None, - quant_recipe: Any = None, - **kwargs: Any, + module: Any, + example_inputs: 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. + """Export a single graph module using torch.export. 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.). + module: One graph module, taken from the value of the strategy's + ``{component: module}`` map. It is not a mapping. + example_inputs: Positional example inputs describing this graph's + export signature. Returns: - A configured quantizer instance. + The exported module (e.g., ExportedProgram.module()). """ ... - def export_model( + def prepare_pt2e( self, - model: Any, - sample_input: Any, + module: Any, + quantizer: Any, ) -> Any: - """Export the model using torch.export. + """Prepare a single exported module for PT2E quantization. Args: - model: The nn.Module to export. - sample_input: Sample input tuple for tracing. + module: One exported graph module, not a component or graph map. + quantizer: The configured quantizer. Returns: - The exported model (e.g., ExportedProgram.module()). + The annotated module with observers inserted. """ ... - def prepare_pt2e( + def init_encodings( self, - model: Any, - quantizer: Any, + module: Any, + example_inputs: Any, ) -> Any: - """Prepare the model for PT2E quantization (insert observers). + """Initialize a graph's observers with a single dummy forward. + + graphs (AR-1 decode / AR-N prefill) are not truly calibrated; + they run once on their own example inputs so their observers' + placeholders are populated before the encoding-override step copies the + real encodings in from the calibration graph. Args: - model: The exported model. - quantizer: The configured quantizer. + module: One annotated deployed graph module, not a map. + example_inputs: This graph's positional example-input tuple. Returns: - The annotated model with observers inserted. + The module after the dummy forward. """ ... def calibrate( self, - model: Any, + modules: Any, calibration_data: Iterable[Any], + **kwargs: Any, ) -> Any: - """Run calibration data through the annotated model. + """Run true calibration over the quantization graph. - Implementations needing a non-trivial procedure -- e.g. autoregressive - LLM calibration, where each step's input depends on the previous step's - output -- should override this method rather than encoding the procedure - in ``calibration_data``. + Drives ``calibration_data`` through ``module``. Model-specific adapters + may use extra kwargs such as ``inference`` (a ``ModelInference`` bound + to the calibration graph); generic adapters can directly call the + module. Args: - model: The annotated model with observers. - calibration_data: Any ``Iterable[Tuple[Tensor, ...]]``, including a - plain list or a ``DataLoader``. + modules: ``{component: module}`` for the selected calibration graphs. + The graph axis has already been removed because each component + contributes only its calibration graph. + calibration_data: ``{component: DataLoader}`` of corpus-backed + calibration batches. + **kwargs: Extra adapter-specific options, such as optional + ``inference`` for model-specific adapters. Returns: - The calibrated model. + The calibrated module. """ ... def convert_pt2e( self, - model: Any, + module: Any, ) -> Any: - """Convert the calibrated model to a quantized model. + """Convert a single calibrated module to a quantized module. Args: - model: The calibrated model. + module: One calibrated graph module, not a component or graph map. Returns: - The quantized model with fake quantize nodes replaced. + The quantized module with fake quantize nodes replaced. """ ... 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 86693c77ebe..0e555650d0a 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 @@ -30,18 +30,10 @@ def test_optional_fields_default_to_none(self): ) self.assertIsNone(config.model_module) self.assertIsNone(config.example_inputs) - self.assertIsNone(config.calibration_data) - self.assertIsNone(config.training_data) + self.assertIsNone(config.tokenizer) self.assertIsNone(config.quant_recipe) - - def test_training_data_carries_qat_dataset(self): - training_data = [("features", "labels")] - config = QuantizationInputConfig( - soc_model=MagicMock(), - backend_type=MagicMock(), - training_data=training_data, - ) - self.assertIs(config.training_data, training_data) + self.assertIsNone(config.meta) + self.assertIsNone(config.inference) if __name__ == "__main__": diff --git a/backends/qualcomm/genai_pipeline/tests/datasets/test_default_calibration_data_adapter.py b/backends/qualcomm/genai_pipeline/tests/datasets/test_default_calibration_data_adapter.py index afe27a12f02..67e4b99ca8f 100644 --- a/backends/qualcomm/genai_pipeline/tests/datasets/test_default_calibration_data_adapter.py +++ b/backends/qualcomm/genai_pipeline/tests/datasets/test_default_calibration_data_adapter.py @@ -8,8 +8,11 @@ from unittest.mock import MagicMock import torch +from executorch.backends.qualcomm.genai_pipeline.artifact_keys import ( + ARTIFACT_TEXT_DECODER, +) -from executorch.backends.qualcomm.genai_pipeline.datasets.default_calibration_data_adapter import ( +from executorch.backends.qualcomm.genai_pipeline.datasets.calibration.default_calibration_data_adapter import ( DefaultCalibrationDataAdapter, ) @@ -27,21 +30,19 @@ def _make_tokenizer(vocab_size=TEST_VOCAB_SIZE): class TestGenerateCalibrationData(unittest.TestCase): def setUp(self): - self.adapter = DefaultCalibrationDataAdapter() + self.adapter = DefaultCalibrationDataAdapter(num_samples=TEST_NUM_SAMPLES) def test_generates_requested_number_of_samples(self): data = self.adapter.generate_calibration_data( - _make_tokenizer(), - num_samples=TEST_NUM_SAMPLES, - seq_length=TEST_SEQ_LENGTH, + _make_tokenizer(), extra_options={"max_context_len": TEST_SEQ_LENGTH} ) - self.assertEqual(len(data), TEST_NUM_SAMPLES) + self.assertEqual(len(data[ARTIFACT_TEXT_DECODER]), TEST_NUM_SAMPLES) def test_sample_shape_and_dtype(self): data = self.adapter.generate_calibration_data( - _make_tokenizer(), num_samples=1, seq_length=TEST_SEQ_LENGTH + _make_tokenizer(), extra_options={"max_context_len": TEST_SEQ_LENGTH} ) - input_ids, attention_mask = data[0] + input_ids, attention_mask = data[ARTIFACT_TEXT_DECODER][0] self.assertEqual(input_ids.shape, (1, TEST_SEQ_LENGTH)) self.assertEqual(attention_mask.shape, (1, TEST_SEQ_LENGTH)) self.assertEqual(attention_mask.dtype, torch.long) @@ -50,22 +51,24 @@ def test_tokens_within_vocab_range(self): vocab_size = 128 data = self.adapter.generate_calibration_data( _make_tokenizer(vocab_size), - num_samples=TEST_NUM_SAMPLES, - seq_length=TEST_SEQ_LENGTH, + extra_options={"max_context_len": TEST_SEQ_LENGTH}, ) - for input_ids, _ in data: + for input_ids, _ in data[ARTIFACT_TEXT_DECODER]: self.assertGreaterEqual(int(input_ids.min()), 0) self.assertLess(int(input_ids.max()), vocab_size) def test_seed_makes_generation_reproducible(self): kwargs = { - "num_samples": TEST_NUM_SAMPLES, - "seq_length": TEST_SEQ_LENGTH, - "extra_options": {"seed": 7}, + "extra_options": { + "max_context_len": TEST_SEQ_LENGTH, + "seed": 7, + }, } first = self.adapter.generate_calibration_data(_make_tokenizer(), **kwargs) second = self.adapter.generate_calibration_data(_make_tokenizer(), **kwargs) - for (ids_a, _), (ids_b, _) in zip(first, second): + for (ids_a, _), (ids_b, _) in zip( + first[ARTIFACT_TEXT_DECODER], second[ARTIFACT_TEXT_DECODER] + ): self.assertTrue(torch.equal(ids_a, ids_b)) def test_provided_list_dataset_is_returned_unchanged(self): @@ -73,7 +76,7 @@ def test_provided_list_dataset_is_returned_unchanged(self): result = self.adapter.generate_calibration_data( _make_tokenizer(), extra_options={"dataset": dataset} ) - self.assertIs(result, dataset) + self.assertIs(result[ARTIFACT_TEXT_DECODER], dataset) def test_provided_dataloader_is_returned_unchanged(self): # A DataLoader is an Iterable, so it is accepted without materialization. @@ -83,7 +86,7 @@ def test_provided_dataloader_is_returned_unchanged(self): result = self.adapter.generate_calibration_data( _make_tokenizer(), extra_options={"dataset": dataloader} ) - self.assertIs(result, dataloader) + self.assertIs(result[ARTIFACT_TEXT_DECODER], dataloader) def test_unusable_vocab_size_raises(self): for vocab_size in (None, 0): diff --git a/backends/qualcomm/genai_pipeline/tests/datasets/test_default_training_data_adapter.py b/backends/qualcomm/genai_pipeline/tests/datasets/test_default_training_data_adapter.py index 823f3e6be28..feb1d49aceb 100644 --- a/backends/qualcomm/genai_pipeline/tests/datasets/test_default_training_data_adapter.py +++ b/backends/qualcomm/genai_pipeline/tests/datasets/test_default_training_data_adapter.py @@ -8,8 +8,11 @@ from unittest.mock import MagicMock import torch +from executorch.backends.qualcomm.genai_pipeline.artifact_keys import ( + ARTIFACT_TEXT_DECODER, +) -from executorch.backends.qualcomm.genai_pipeline.datasets.default_training_data_adapter import ( +from executorch.backends.qualcomm.genai_pipeline.datasets.training.default_training_data_adapter import ( DefaultTrainingDataAdapter, ) @@ -25,7 +28,7 @@ def test_returns_provided_feature_label_pairs(self): result = self.adapter.generate_training_data( self.tokenizer, extra_options={"training_data": training_data} ) - self.assertIs(result, training_data) + self.assertIs(result[ARTIFACT_TEXT_DECODER], training_data) def test_accepts_a_dataloader(self): dataloader = torch.utils.data.DataLoader( @@ -34,7 +37,7 @@ def test_accepts_a_dataloader(self): result = self.adapter.generate_training_data( self.tokenizer, extra_options={"training_data": dataloader} ) - self.assertIs(result, dataloader) + self.assertIs(result[ARTIFACT_TEXT_DECODER], dataloader) def test_raises_when_no_training_data_supplied(self): # Labelled data cannot be synthesized, so QAT must fail loudly rather diff --git a/backends/qualcomm/genai_pipeline/tests/source_transform/test_checkpoint_key_remap.py b/backends/qualcomm/genai_pipeline/tests/source_transform/test_checkpoint_key_remap.py new file mode 100644 index 00000000000..ac2f2b8cab6 --- /dev/null +++ b/backends/qualcomm/genai_pipeline/tests/source_transform/test_checkpoint_key_remap.py @@ -0,0 +1,56 @@ +# 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. + +"""Tests for the checkpoint key-renaming state-dict transforms.""" + +import unittest +from unittest.mock import MagicMock, patch + +import torch + +from executorch.backends.qualcomm.genai_pipeline.source_transform.checkpoint_key_remap import ( + remap_gemma4_keys, + strip_orig_mod_prefix, +) + + +class TestStripOrigModPrefix(unittest.TestCase): + """``_orig_mod.`` is left on keys by torch.compile checkpoints.""" + + def test_removes_prefix(self): + result = strip_orig_mod_prefix( + {"_orig_mod.layers.0.attention.wq.weight": torch.zeros(1)} + ) + self.assertEqual(list(result), ["layers.0.attention.wq.weight"]) + + def test_leaves_unprefixed_keys_alone(self): + result = strip_orig_mod_prefix({"tok_embeddings.weight": torch.zeros(1)}) + self.assertEqual(list(result), ["tok_embeddings.weight"]) + + +class TestRemapGemma4Keys(unittest.TestCase): + """Gemma4 arrives pre-converted and only needs its keys renamed.""" + + def test_delegates_to_gemma4_remap(self): + state_dict = {"before": torch.zeros(1)} + remapped = {"after": torch.zeros(1)} + module = MagicMock() + module.remap_keys.return_value = remapped + + with patch.dict( + "sys.modules", + { + "executorch.examples.qualcomm.oss_scripts.gemma4.text_decoder.convert_weights": module + }, + ): + result = remap_gemma4_keys(state_dict) + + module.remap_keys.assert_called_once_with(state_dict) + self.assertIs(result, remapped) + + +if __name__ == "__main__": + unittest.main() diff --git a/backends/qualcomm/genai_pipeline/tests/source_transform/test_dtype_override.py b/backends/qualcomm/genai_pipeline/tests/source_transform/test_dtype_override.py new file mode 100644 index 00000000000..b48bb4f1b36 --- /dev/null +++ b/backends/qualcomm/genai_pipeline/tests/source_transform/test_dtype_override.py @@ -0,0 +1,38 @@ +# 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. + +"""Tests for the dtype-override module transform.""" + +import unittest + +import torch + +from executorch.backends.qualcomm.genai_pipeline.source_transform.dtype_override import ( + apply_dtype_override, +) +from torch import nn + + +class TestApplyDtypeOverride(unittest.TestCase): + """``--dtype-override`` casts the whole module before export.""" + + def test_casts_to_requested_dtype(self): + module = nn.Linear(2, 2) + + result = apply_dtype_override(module, dtype_override="fp16") + + self.assertEqual(result.weight.dtype, torch.float16) + + def test_no_op_when_unset(self): + module = nn.Linear(2, 2) + + result = apply_dtype_override(module, dtype_override=None) + + self.assertEqual(result.weight.dtype, torch.float32) + + +if __name__ == "__main__": + unittest.main() diff --git a/backends/qualcomm/genai_pipeline/tests/source_transform/test_embedding_scale.py b/backends/qualcomm/genai_pipeline/tests/source_transform/test_embedding_scale.py new file mode 100644 index 00000000000..36a8003e06c --- /dev/null +++ b/backends/qualcomm/genai_pipeline/tests/source_transform/test_embedding_scale.py @@ -0,0 +1,51 @@ +# 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. + +"""Tests for the token-embedding state-dict transform.""" + +import unittest + +import torch + +from executorch.backends.qualcomm.genai_pipeline.source_transform.embedding_scale import ( + scale_token_embedding, +) + + +class TestScaleTokenEmbedding(unittest.TestCase): + """Gemma-family models scale embeddings by ``sqrt(hidden_size)``.""" + + def test_scales_by_factor(self): + result = scale_token_embedding( + {"tok_embeddings.weight": torch.ones(2)}, + embedding_scale_factor=2.0, + ) + torch.testing.assert_close( + result["tok_embeddings.weight"], torch.full((2,), 2.0) + ) + + def test_no_op_when_embedding_absent(self): + """Models with a separated token-embedding graph carry no such key.""" + result = scale_token_embedding( + {"norm.weight": torch.ones(2)}, + embedding_scale_factor=2.0, + ) + torch.testing.assert_close(result["norm.weight"], torch.ones(2)) + + def test_no_op_at_unit_factor(self): + """Runs for every model, so a factor of 1.0 must change nothing at all -- + including dtype, which the multiply would otherwise upcast.""" + weight = torch.ones(2, dtype=torch.bfloat16) + result = scale_token_embedding( + {"tok_embeddings.weight": weight}, + embedding_scale_factor=1.0, + ) + self.assertIs(result["tok_embeddings.weight"], weight) + self.assertEqual(result["tok_embeddings.weight"].dtype, torch.bfloat16) + + +if __name__ == "__main__": + unittest.main() diff --git a/backends/qualcomm/genai_pipeline/tests/source_transform/test_linear_to_conv2d.py b/backends/qualcomm/genai_pipeline/tests/source_transform/test_linear_to_conv2d.py new file mode 100644 index 00000000000..fb1fd4fe3eb --- /dev/null +++ b/backends/qualcomm/genai_pipeline/tests/source_transform/test_linear_to_conv2d.py @@ -0,0 +1,134 @@ +# 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. + +"""Tests for the linear-to-conv2d module transforms.""" + +import sys +import types +import unittest +from unittest.mock import MagicMock + +import torch + +from executorch.backends.qualcomm.genai_pipeline.source_transform.linear_to_conv2d import ( + convert_linear_to_conv2d, + prepare_conv_submodules, +) +from torch import nn + + +class _LinearAsConv2d(nn.Module): + def __init__(self, linear): + super().__init__() + self.conv = nn.Conv2d( + linear.in_features, + linear.out_features, + kernel_size=1, + bias=linear.bias is not None, + ) + self.conv.weight.data.copy_( + linear.weight.data.reshape(linear.out_features, linear.in_features, 1, 1) + ) + if linear.bias is not None: + self.conv.bias.data.copy_(linear.bias.data) + + def forward(self, x): + y = self.conv(x.transpose(-1, -2).unsqueeze(-1)) + return y.squeeze(-1).transpose(-1, -2) + + +def _convert_named_linears_to_conv2d(module): + for name, child in list(module.named_children()): + if isinstance(child, nn.Linear): + setattr(module, name, _LinearAsConv2d(child)) + return module + + +def _fake_backend_utils(): + module = types.ModuleType("executorch.backends.qualcomm.utils.utils") + module.convert_linear_to_conv2d = _convert_named_linears_to_conv2d + return module + + +class TestPrepareConvSubmodules(unittest.TestCase): + """Attention and feed-forward blocks pre-shape their weights for conv2d.""" + + def test_calls_both_prepare_hooks_on_every_layer(self): + layers = [] + for _ in range(2): + layer = MagicMock() + layer.attention.prepare_attention_conv = MagicMock() + layer.feed_forward.prepare_feedforward_conv = MagicMock() + layers.append(layer) + module = MagicMock(layers=layers) + + result = prepare_conv_submodules(module) + + for layer in layers: + layer.attention.prepare_attention_conv.assert_called_once_with() + layer.feed_forward.prepare_feedforward_conv.assert_called_once_with() + self.assertIs(result, module) + + def test_skips_layers_without_the_hooks(self): + layer = MagicMock(spec=[]) + layer.attention = MagicMock(spec=[]) + layer.feed_forward = MagicMock(spec=[]) + module = MagicMock(layers=[layer]) + + self.assertIs(prepare_conv_submodules(module), module) + + +class TestConvertLinearToConv2d(unittest.TestCase): + """HTP runs conv2d faster than linear; the backend util rewrites them. + + The util walks named attributes, so the module under test mirrors the real + decoder's shape (``self.wq = nn.Linear(...)``) rather than using + ``nn.Sequential``, whose children it does not reach. + """ + + def test_replaces_named_linear_attributes(self): + class Block(nn.Module): + def __init__(self): + super().__init__() + self.wq = nn.Linear(4, 4, bias=False) + + module = Block() + + with unittest.mock.patch.dict( + sys.modules, + {"executorch.backends.qualcomm.utils.utils": _fake_backend_utils()}, + ): + result = convert_linear_to_conv2d(module) + + self.assertFalse( + any(isinstance(m, nn.Linear) for m in result.modules()), + "nn.Linear should have been rewritten", + ) + + def test_preserves_forward_semantics(self): + class Block(nn.Module): + def __init__(self): + super().__init__() + self.wq = nn.Linear(4, 4, bias=False) + + def forward(self, x): + return self.wq(x) + + module = Block().eval() + x = torch.randn(1, 2, 4) + with torch.no_grad(): + expected = module(x) + with unittest.mock.patch.dict( + sys.modules, + {"executorch.backends.qualcomm.utils.utils": _fake_backend_utils()}, + ): + actual = convert_linear_to_conv2d(module)(x) + + torch.testing.assert_close(actual, expected) + + +if __name__ == "__main__": + unittest.main() diff --git a/backends/qualcomm/genai_pipeline/tests/source_transform/test_rms_norm_offset.py b/backends/qualcomm/genai_pipeline/tests/source_transform/test_rms_norm_offset.py new file mode 100644 index 00000000000..bafd99d4825 --- /dev/null +++ b/backends/qualcomm/genai_pipeline/tests/source_transform/test_rms_norm_offset.py @@ -0,0 +1,47 @@ +# 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. + +"""Tests for the RMSNorm state-dict transform.""" + +import unittest + +import torch + +from executorch.backends.qualcomm.genai_pipeline.source_transform.rms_norm_offset import ( + gemma_rmsnorm_offset, +) + + +class TestGemmaRmsnormOffset(unittest.TestCase): + """Gemma computes ``(x * w).to(fp16)``, Llama ``x.to(fp16) * w``. + + The static decoder implements the Llama form, so Gemma's norm weights carry + a ``+1`` offset folded in at load time. + """ + + def test_adds_one_to_norm_weights(self): + result = gemma_rmsnorm_offset( + {"layers.0.attention_norm.weight": torch.zeros(3)} + ) + torch.testing.assert_close( + result["layers.0.attention_norm.weight"], torch.ones(3) + ) + + def test_leaves_non_norm_weights_untouched(self): + result = gemma_rmsnorm_offset({"layers.0.attention.wq.weight": torch.zeros(3)}) + torch.testing.assert_close( + result["layers.0.attention.wq.weight"], torch.zeros(3) + ) + + def test_promotes_norm_weights_to_float32(self): + result = gemma_rmsnorm_offset( + {"norm.weight": torch.zeros(3, dtype=torch.bfloat16)} + ) + self.assertEqual(result["norm.weight"].dtype, torch.float32) + + +if __name__ == "__main__": + unittest.main() diff --git a/backends/qualcomm/genai_pipeline/tests/source_transform/test_rope_layout.py b/backends/qualcomm/genai_pipeline/tests/source_transform/test_rope_layout.py new file mode 100644 index 00000000000..59ab8244a06 --- /dev/null +++ b/backends/qualcomm/genai_pipeline/tests/source_transform/test_rope_layout.py @@ -0,0 +1,110 @@ +# 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. + +"""Tests for the RoPE weight-layout state-dict transform.""" + +import unittest + +import torch + +from executorch.backends.qualcomm.genai_pipeline.source_transform.rope_layout import ( + permute_partial_rope, +) + + +class TestPermutePartialRope(unittest.TestCase): + """Re-lays-out q/k weights from HF's interleaved RoPE to HTP's split form.""" + + def test_permutes_query_and_key_weights(self): + weight = torch.arange(8, dtype=torch.float32).reshape(8, 1) + state_dict = { + "layers.0.attention.wq.weight": weight.clone(), + "layers.0.attention.wk.weight": weight.clone(), + } + + result = permute_partial_rope( + state_dict, + partial_rotary_factor=1.0, + n_layers=1, + n_heads=2, + n_kv_heads=2, + ) + + expected = torch.tensor( + [[0.0], [2.0], [1.0], [3.0], [4.0], [6.0], [5.0], [7.0]] + ) + torch.testing.assert_close(result["layers.0.attention.wq.weight"], expected) + torch.testing.assert_close(result["layers.0.attention.wk.weight"], expected) + + def test_leaves_value_weights_untouched(self): + """Only q and k feed RoPE; v must keep HF's layout.""" + weight = torch.arange(8, dtype=torch.float32).reshape(8, 1) + state_dict = { + "layers.0.attention.wq.weight": weight.clone(), + "layers.0.attention.wk.weight": weight.clone(), + "layers.0.attention.wv.weight": weight.clone(), + } + + result = permute_partial_rope( + state_dict, + partial_rotary_factor=1.0, + n_layers=1, + n_heads=2, + n_kv_heads=2, + ) + + torch.testing.assert_close(result["layers.0.attention.wv.weight"], weight) + + def test_uses_separate_head_counts_for_q_and_k(self): + """GQA models have n_kv_heads < n_heads; each is permuted by its own count.""" + state_dict = { + "layers.0.attention.wq.weight": torch.arange( + 8, dtype=torch.float32 + ).reshape(8, 1), + "layers.0.attention.wk.weight": torch.arange( + 8, dtype=torch.float32 + ).reshape(8, 1), + } + + result = permute_partial_rope( + state_dict, + partial_rotary_factor=1.0, + n_layers=1, + n_heads=2, + n_kv_heads=1, + ) + + self.assertFalse( + torch.equal( + result["layers.0.attention.wq.weight"], + result["layers.0.attention.wk.weight"], + ) + ) + + def test_covers_every_layer(self): + weight = torch.arange(8, dtype=torch.float32).reshape(8, 1) + state_dict = { + f"layers.{i}.attention.w{p}.weight": weight.clone() + for i in range(3) + for p in ("q", "k") + } + + result = permute_partial_rope( + state_dict, + partial_rotary_factor=1.0, + n_layers=3, + n_heads=2, + n_kv_heads=2, + ) + + for i in range(3): + self.assertFalse( + torch.equal(result[f"layers.{i}.attention.wq.weight"], weight) + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/backends/qualcomm/genai_pipeline/tests/stages/test_quantization_stage.py b/backends/qualcomm/genai_pipeline/tests/stages/test_quantization_stage.py index 3fa6d3d60d6..59bfefe934f 100644 --- a/backends/qualcomm/genai_pipeline/tests/stages/test_quantization_stage.py +++ b/backends/qualcomm/genai_pipeline/tests/stages/test_quantization_stage.py @@ -36,9 +36,7 @@ def test_name(self): def test_invoke_delegates_to_strategy(self): mock_strategy = MagicMock(spec=QuantizationStrategy) - mock_strategy.invoke.return_value = QuantizationOutputConfig( - quantized_model="quantized" - ) + mock_strategy.invoke.return_value = QuantizationOutputConfig(graphs={}) stage = QuantizationStage(mock_strategy) context = make_test_context() input_config = QuantizationInputConfig( 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 bc9820ee5af..da93d369f1c 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 @@ -7,8 +7,12 @@ import tempfile import unittest from pathlib import Path -from unittest.mock import MagicMock +from unittest.mock import MagicMock, patch +from executorch.backends.qualcomm.genai_pipeline.artifact_keys import ( + ARTIFACT_TEXT_DECODER, +) +from executorch.backends.qualcomm.genai_pipeline.graph_names import GRAPH_FORWARD from executorch.backends.qualcomm.genai_pipeline.strategies.model_preparation.default_model_loader_adapter import ( DefaultModelLoaderAdapter, ) @@ -20,6 +24,27 @@ TEST_ADDED_TOKENS = "added_tokens.json" +class TestLoadModel(unittest.TestCase): + def setUp(self): + self.adapter = DefaultModelLoaderAdapter() + + def test_returns_component_graph_map(self): + model = MagicMock() + model.eval.return_value = None + + with patch( + "transformers.AutoModelForCausalLM.from_pretrained", + return_value=model, + ): + result = self.adapter.load_model("test-model") + + self.assertEqual( + result, + {ARTIFACT_TEXT_DECODER: {GRAPH_FORWARD: model}}, + ) + model.eval.assert_called_once_with() + + class TestExportTokenizer(unittest.TestCase): """The runtime (llm::load_tokenizer / pytorch_tokenizers.get_tokenizer) expects a single tokenizer *file*, never the containing directory.""" 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 index 9e02a6f5a7c..6a1411b7bd4 100644 --- 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 @@ -6,7 +6,13 @@ import unittest from pathlib import Path -from unittest.mock import MagicMock, patch +from unittest.mock import call, MagicMock, patch + +from executorch.backends.qualcomm.genai_pipeline.artifact_keys import ( + ARTIFACT_TEXT_DECODER, + ARTIFACT_TOK_EMBEDDING, + ARTIFACT_VISION_ENCODER, +) from executorch.backends.qualcomm.genai_pipeline.configs.model_preparation_input_config import ( ModelPreparationInputConfig, @@ -14,11 +20,12 @@ 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.graph_names import ( + DECODER_GRAPH_NAMES, + GRAPH_FORWARD, + TOK_EMBEDDING_GRAPH_NAMES, +) from executorch.backends.qualcomm.genai_pipeline.strategies.model_preparation.executorch_model_preparation_strategy import ( ExecuTorchModelPreparationStrategy, ) @@ -29,34 +36,67 @@ make_test_context, ) +DEFAULT_GRAPH_NAME = GRAPH_FORWARD +DECODE_GRAPH_NAME = DECODER_GRAPH_NAMES[0] +TOK_EMBEDDING_GRAPH_NAME = TOK_EMBEDDING_GRAPH_NAMES[0] + + +def _make_modules(): + return { + ARTIFACT_TEXT_DECODER: { + DEFAULT_GRAPH_NAME: MagicMock(name=DEFAULT_GRAPH_NAME), + DECODE_GRAPH_NAME: MagicMock(name=DECODE_GRAPH_NAME), + }, + ARTIFACT_TOK_EMBEDDING: { + DEFAULT_GRAPH_NAME: MagicMock(name=DEFAULT_GRAPH_NAME), + TOK_EMBEDDING_GRAPH_NAME: MagicMock(name=TOK_EMBEDDING_GRAPH_NAME), + }, + ARTIFACT_VISION_ENCODER: { + DEFAULT_GRAPH_NAME: MagicMock(name=DEFAULT_GRAPH_NAME), + }, + } + def _make_mock_adapter(): """Create a mock model loader adapter with sensible defaults.""" adapter = MagicMock() - adapter.load_model.return_value = MagicMock(name="model_module") + modules = _make_modules() tokenizer = MagicMock(name="tokenizer") # Default to "no chat template" so tests opt in explicitly. tokenizer.chat_template = None + + text_decoder_metadata = [ + ( + modules[ARTIFACT_TEXT_DECODER][DEFAULT_GRAPH_NAME], + {"get_max_context_len": 1024}, + ), + ( + modules[ARTIFACT_TEXT_DECODER][DECODE_GRAPH_NAME], + {"get_max_context_len": 1024}, + ), + ] + + def _get_metadata(model): + for graph_module, metadata in text_decoder_metadata: + if model is graph_module: + return metadata + return {} + + adapter.load_model.return_value = modules + adapter.get_example_inputs.side_effect = lambda model: (model,) + adapter.get_metadata.side_effect = _get_metadata + adapter.apply_module_transforms.side_effect = ( + lambda module, module_transforms: module + ) adapter.load_tokenizer.return_value = tokenizer - adapter.get_example_inputs.return_value = (MagicMock(name="example_input"),) + adapter.get_inference.return_value = MagicMock(name="inference") 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.""" +def _make_strategy(loader=None): 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() - ), + model_loader_adapter=loader if loader is not None else _make_mock_adapter() ) @@ -77,166 +117,249 @@ def test_is_model_preparation_strategy(self): self.assertIsInstance(_make_strategy(), ModelPreparationStrategy) def test_default_adapters_created_when_none_provided(self): - """Both adapters fall back to their default implementations.""" + """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: + ) as mock_loader_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.""" + def test_custom_adapter_injected(self): loader = _make_mock_adapter() - calibration = _make_mock_calibration_adapter() - strategy = _make_strategy(loader, calibration) + strategy = _make_strategy(loader) 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) + strategy = _make_strategy(loader) + modules = loader.load_model.return_value + text_decoder_forward = modules[ARTIFACT_TEXT_DECODER][DEFAULT_GRAPH_NAME] + text_decoder_kv_forward = modules[ARTIFACT_TEXT_DECODER][DECODE_GRAPH_NAME] + tok_embedding_forward = modules[ARTIFACT_TOK_EMBEDDING][DEFAULT_GRAPH_NAME] + tok_embedding_kv_forward = modules[ARTIFACT_TOK_EMBEDDING][ + TOK_EMBEDDING_GRAPH_NAME + ] + vision_encoder_forward = modules[ARTIFACT_VISION_ENCODER][DEFAULT_GRAPH_NAME] 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.assertEqual( + result.model_module, + { + ARTIFACT_TEXT_DECODER: text_decoder_forward, + ARTIFACT_TOK_EMBEDDING: tok_embedding_forward, + ARTIFACT_VISION_ENCODER: vision_encoder_forward, + }, + ) self.assertIs(result.tokenizer, loader.load_tokenizer.return_value) self.assertEqual( - result.calibration_data, - calibration.generate_calibration_data.return_value, + result.example_inputs, + { + ARTIFACT_TEXT_DECODER: { + DEFAULT_GRAPH_NAME: (text_decoder_forward,), + DECODE_GRAPH_NAME: (text_decoder_kv_forward,), + }, + ARTIFACT_TOK_EMBEDDING: { + DEFAULT_GRAPH_NAME: (tok_embedding_forward,), + TOK_EMBEDDING_GRAPH_NAME: (tok_embedding_kv_forward,), + }, + ARTIFACT_VISION_ENCODER: { + DEFAULT_GRAPH_NAME: (vision_encoder_forward,), + }, + }, ) + self.assertEqual( + result.meta, + { + ARTIFACT_TEXT_DECODER: { + DEFAULT_GRAPH_NAME: {"get_max_context_len": 1024}, + DECODE_GRAPH_NAME: {"get_max_context_len": 1024}, + } + }, + ) + self.assertIs(result.inference, loader.get_inference.return_value) - def test_invoke_calls_loader_in_correct_order(self): - """The loader is driven in order: load_model → load_tokenizer → example inputs.""" + def test_invoke_builds_example_inputs_for_each_component_graph(self): loader = _make_mock_adapter() strategy = _make_strategy(loader) + modules = loader.load_model.return_value + expected_modules = [ + modules[ARTIFACT_TEXT_DECODER][DEFAULT_GRAPH_NAME], + modules[ARTIFACT_TEXT_DECODER][DECODE_GRAPH_NAME], + modules[ARTIFACT_TOK_EMBEDDING][DEFAULT_GRAPH_NAME], + modules[ARTIFACT_TOK_EMBEDDING][TOK_EMBEDDING_GRAPH_NAME], + modules[ARTIFACT_VISION_ENCODER][DEFAULT_GRAPH_NAME], + ] 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"], + loader.get_example_inputs.assert_has_calls( + [call(module) for module in expected_modules], + any_order=False, ) - 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. - """ + def test_invoke_builds_metadata_for_each_component_graph(self): loader = _make_mock_adapter() - calibration = _make_mock_calibration_adapter() - strategy = _make_strategy(loader, calibration) + strategy = _make_strategy(loader) + modules = loader.load_model.return_value + expected_modules = [ + modules[ARTIFACT_TEXT_DECODER][DEFAULT_GRAPH_NAME], + modules[ARTIFACT_TEXT_DECODER][DECODE_GRAPH_NAME], + modules[ARTIFACT_TOK_EMBEDDING][DEFAULT_GRAPH_NAME], + modules[ARTIFACT_TOK_EMBEDDING][TOK_EMBEDDING_GRAPH_NAME], + modules[ARTIFACT_VISION_ENCODER][DEFAULT_GRAPH_NAME], + ] 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, + loader.get_metadata.assert_has_calls( + [call(module) for module in expected_modules], + any_order=False, ) - self.assertIs(result.example_inputs, loader.get_example_inputs.return_value) + self.assertNotIn(ARTIFACT_VISION_ENCODER, result.meta) - def test_invoke_example_input_options_forwarded(self): - """example_input_options from extra_options reach get_example_inputs.""" + def test_invoke_selects_first_graph_module_before_module_transforms(self): loader = _make_mock_adapter() strategy = _make_strategy(loader) - example_opts = {"ar_len": 128} + text_transform = MagicMock(name="text_transform") input_config = _make_valid_input_config( - extra_options={"example_input_options": example_opts} + extra_options={ + "model_options": { + "module_transforms": {ARTIFACT_TEXT_DECODER: [text_transform]} + } + } ) - strategy.invoke(make_test_context(), input_config) + result = 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, + loader.apply_module_transforms.assert_has_calls( + [ + call( + loader.load_model.return_value[ARTIFACT_TEXT_DECODER][ + DEFAULT_GRAPH_NAME + ], + module_transforms=[text_transform], + ), + call( + loader.load_model.return_value[ARTIFACT_TOK_EMBEDDING][ + DEFAULT_GRAPH_NAME + ], + module_transforms=[], + ), + call( + loader.load_model.return_value[ARTIFACT_VISION_ENCODER][ + DEFAULT_GRAPH_NAME + ], + module_transforms=[], + ), + ], + any_order=False, + ) + self.assertIs( + result.model_module[ARTIFACT_TEXT_DECODER], + loader.load_model.return_value[ARTIFACT_TEXT_DECODER][DEFAULT_GRAPH_NAME], + ) + self.assertIs( + result.model_module[ARTIFACT_TOK_EMBEDDING], + loader.load_model.return_value[ARTIFACT_TOK_EMBEDDING][DEFAULT_GRAPH_NAME], ) - 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) + def test_get_component_module_does_not_filter_by_deploy_graph_name(self): + strategy = _make_strategy() + text_decoder_kv_forward = MagicMock(name="text_decoder_kv_forward") + tok_embedding_kv_forward = MagicMock(name="tok_embedding_kv_forward") + + result = strategy._get_component_module( + { + ARTIFACT_TEXT_DECODER: {DECODE_GRAPH_NAME: text_decoder_kv_forward}, + ARTIFACT_TOK_EMBEDDING: { + TOK_EMBEDDING_GRAPH_NAME: tok_embedding_kv_forward + }, + ARTIFACT_VISION_ENCODER: {}, + } + ) - strategy.invoke(make_test_context(), _make_valid_input_config()) + self.assertEqual( + result, + { + ARTIFACT_TEXT_DECODER: text_decoder_kv_forward, + ARTIFACT_TOK_EMBEDDING: tok_embedding_kv_forward, + }, + ) - 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_get_component_module_does_not_mutate_graph_modules(self): + strategy = _make_strategy() + text_decoder_forward = MagicMock(name="text_decoder_forward") + text_decoder_kv_forward = MagicMock(name="text_decoder_kv_forward") + modules = { + ARTIFACT_TEXT_DECODER: { + DEFAULT_GRAPH_NAME: text_decoder_forward, + DECODE_GRAPH_NAME: text_decoder_kv_forward, + } + } + + result = strategy._get_component_module(modules) + + self.assertEqual(result, {ARTIFACT_TEXT_DECODER: text_decoder_forward}) + self.assertEqual( + modules, + { + ARTIFACT_TEXT_DECODER: { + DEFAULT_GRAPH_NAME: text_decoder_forward, + DECODE_GRAPH_NAME: text_decoder_kv_forward, + } + }, ) - def test_invoke_passes_model_name_to_load_model(self): - """load_model receives model_name from input config.""" + def test_invoke_passes_tokenizer_options_to_load_tokenizer(self): loader = _make_mock_adapter() strategy = _make_strategy(loader) + tokenizer_options = {"use_fast": False} strategy.invoke( - make_test_context(), _make_valid_input_config(model_name="llama3_2-1b") + make_test_context(), + _make_valid_input_config( + extra_options={"tokenizer_options": tokenizer_options} + ), ) - loader.load_model.assert_called_once_with( - model_name="llama3_2-1b", - extra_options=None, + loader.load_tokenizer.assert_called_once_with( + model_name="test_model", + extra_options=tokenizer_options, ) - def test_invoke_passes_model_name_to_load_tokenizer(self): - """load_tokenizer receives model_name from input config.""" + def test_invoke_builds_inference_from_meta_and_example_inputs(self): loader = _make_mock_adapter() strategy = _make_strategy(loader) + extra = {"embedding_quantize": "4a"} - strategy.invoke( - make_test_context(), _make_valid_input_config(model_name="llama3_2-1b") + result = strategy.invoke( + make_test_context(), _make_valid_input_config(extra_options=extra) ) - loader.load_tokenizer.assert_called_once_with( - model_name="llama3_2-1b", - extra_options=None, + loader.get_inference.assert_called_once_with( + result.meta, + result.example_inputs, + extra_options=extra, ) - 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"} + def test_invoke_rejects_unkeyed_module_transforms(self): + strategy = _make_strategy() input_config = _make_valid_input_config( - extra_options={ - "num_calibration_samples": 64, - "calibration_seq_length": 256, - "calibration_options": calibration_options, - } + extra_options={"model_options": {"module_transforms": []}} ) - 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()) + with self.assertRaises(StageError) as cm: + strategy.invoke(make_test_context(), input_config) - loader.export_tokenizer.assert_not_called() - self.assertIsNone(result.runtime_tokenizer_path) + self.assertEqual(cm.exception.stage_name, "model_preparation") + self.assertIsInstance(cm.exception.original_exception, ValueError) + self.assertIn("module_transforms", str(cm.exception.original_exception)) def test_invoke_exports_tokenizer_when_requested(self): """export_tokenizer is called when export_tokenizer=True in extra_options.""" @@ -259,33 +382,8 @@ def test_invoke_exports_tokenizer_when_requested(self): 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. - """ + """A chat template on the tokenizer is carried into the output config.""" loader = _make_mock_adapter() loader.load_tokenizer.return_value.chat_template = "tokenizer_template" strategy = _make_strategy(loader) @@ -329,19 +427,7 @@ def test_invoke_adapter_exception_wrapped_in_stage_error(self): 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" diff --git a/backends/qualcomm/genai_pipeline/tests/strategies/model_preparation/test_llm_loader_adapter.py b/backends/qualcomm/genai_pipeline/tests/strategies/model_preparation/test_llm_loader_adapter.py new file mode 100644 index 00000000000..3b205c23c2b --- /dev/null +++ b/backends/qualcomm/genai_pipeline/tests/strategies/model_preparation/test_llm_loader_adapter.py @@ -0,0 +1,278 @@ +# 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 tempfile +import unittest +from pathlib import Path +from unittest.mock import MagicMock, patch + +from executorch.backends.qualcomm.genai_pipeline.artifact_keys import ( + ARTIFACT_TEXT_DECODER, +) +from executorch.backends.qualcomm.genai_pipeline.graph_names import ( + DECODER_GRAPH_NAMES, + GRAPH_FORWARD, +) + +from executorch.backends.qualcomm.genai_pipeline.strategies.model_preparation.llm_loader_adapter import ( + LLMLoaderAdapter, +) + +DEFAULT_GRAPH_NAME = GRAPH_FORWARD +DECODE_GRAPH_NAME = DECODER_GRAPH_NAMES[0] + +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" + + +def _make_module(name): + module = MagicMock(name=name) + module.eval.return_value = module + return module + + +class TestLLMLoaderAdapter(unittest.TestCase): + def setUp(self): + self.model_config = MagicMock() + self.model_config.repo_id = None + self.control_args = MagicMock() + self.control_args.checkpoint = "/tmp/checkpoint.pt" + self.control_args.model = "test_model" + self.adapter = LLMLoaderAdapter(self.model_config, self.control_args) + + def test_load_model_builds_graphs_and_loads_checkpoint_into_each_graph(self): + calibration_module = _make_module("calibration") + decode_module = _make_module("decode") + state_dict = {"weight": MagicMock(name="weight")} + transform = MagicMock(return_value={"weight": MagicMock(name="rewritten")}) + model_arch = { + ARTIFACT_TEXT_DECODER: { + DEFAULT_GRAPH_NAME: MagicMock(return_value=calibration_module), + DECODE_GRAPH_NAME: MagicMock(return_value=decode_module), + } + } + + with patch( + "executorch.backends.qualcomm.genai_pipeline.strategies." + "model_preparation.llm_loader_adapter.torch.load", + return_value=state_dict, + ) as torch_load: + result = self.adapter.load_model( + "test_model", + extra_options={ + "model_arch": model_arch, + "weight_transforms": {ARTIFACT_TEXT_DECODER: [transform]}, + }, + ) + + torch_load.assert_called_once_with( + "/tmp/checkpoint.pt", weights_only=True, map_location="cpu", mmap=True + ) + transform.assert_called_once_with(state_dict) + calibration_module.load_state_dict.assert_called_once_with( + transform.return_value, strict=True, assign=True + ) + decode_module.load_state_dict.assert_called_once_with( + transform.return_value, strict=True, assign=True + ) + self.assertEqual( + result, + { + ARTIFACT_TEXT_DECODER: { + DEFAULT_GRAPH_NAME: calibration_module, + DECODE_GRAPH_NAME: decode_module, + } + }, + ) + + def test_load_model_requires_model_arch(self): + with self.assertRaises(ValueError): + self.adapter.load_model("test_model", extra_options={}) + + def test_load_state_dict_uses_injected_hf_loader(self): + state_dict = {"weight": MagicMock(name="weight")} + state_dict_loader = MagicMock(return_value=state_dict) + self.control_args.checkpoint = None + self.model_config.repo_id = "test/repo" + + with patch( + "executorch.backends.qualcomm.genai_pipeline.strategies." + "model_preparation.llm_loader_adapter.torch.load" + ) as torch_load: + result = self.adapter._load_state_dict(state_dict_loader) + + state_dict_loader.assert_called_once_with("test/repo") + torch_load.assert_not_called() + self.assertIs(result, state_dict) + + def test_load_model_uses_component_keyed_state_dict_loader(self): + module = _make_module("decoder") + state_dict = {"weight": MagicMock(name="weight")} + state_dict_loader = MagicMock(return_value=state_dict) + self.control_args.checkpoint = None + self.model_config.repo_id = "test/repo" + + result = self.adapter.load_model( + "test_model", + extra_options={ + "model_arch": { + ARTIFACT_TEXT_DECODER: { + DEFAULT_GRAPH_NAME: MagicMock(return_value=module), + } + }, + "state_dict_loader": {ARTIFACT_TEXT_DECODER: state_dict_loader}, + }, + ) + + state_dict_loader.assert_called_once_with("test/repo") + module.load_state_dict.assert_called_once_with( + state_dict, strict=True, assign=True + ) + self.assertEqual( + result, + {ARTIFACT_TEXT_DECODER: {DEFAULT_GRAPH_NAME: module}}, + ) + + def test_get_metadata_reads_module_when_it_exposes_metadata(self): + with_metadata = MagicMock() + with_metadata.get_metadata.return_value = {"get_n_layers": 2} + + result = self.adapter.get_metadata(with_metadata) + + self.assertEqual(result, {"get_n_layers": 2}) + with_metadata.get_metadata.assert_called_once_with() + + def test_get_metadata_returns_empty_when_module_exposes_none(self): + without_metadata = MagicMock(spec=[]) + + result = self.adapter.get_metadata(without_metadata) + + self.assertEqual(result, {}) + + +class TestExportTokenizer(unittest.TestCase): + """TokenizerWrapper has already written artifacts; the adapter selects the + runtime file using the same name-based priority as the default adapter.""" + + def setUp(self): + self.model_config = MagicMock() + self.control_args = MagicMock() + self.adapter = LLMLoaderAdapter(self.model_config, self.control_args) + self._tmp = tempfile.TemporaryDirectory() + self.artifact_dir = Path(self._tmp.name) / "tokenizer_artifacts" + self.output_dir = Path(self._tmp.name) / "tokenizer_out" + self.artifact_dir.mkdir() + + def tearDown(self): + self._tmp.cleanup() + + def _touch_artifacts(self, names): + paths = [] + for name in names: + path = self.artifact_dir / name + path.touch() + paths.append(path) + return paths + + def _make_tokenizer(self, runtime_name): + return MagicMock( + runtime_tokenizer_path=str(self.artifact_dir / runtime_name), + artifact=str(self.artifact_dir), + ) + + def test_returns_tokenizer_file_not_the_directory(self): + self._touch_artifacts( + [TEST_TOKENIZER_CONFIG, TEST_SPECIAL_TOKENS_MAP, TEST_TOKENIZER_JSON] + ) + + result = self.adapter.export_tokenizer( + self._make_tokenizer(TEST_TOKENIZER_JSON), self.output_dir + ) + + self.assertEqual(result, self.artifact_dir / TEST_TOKENIZER_JSON) + self.assertNotEqual(result, self.artifact_dir) + + def test_selects_by_name_not_by_position(self): + self._touch_artifacts( + [TEST_TOKENIZER_CONFIG, TEST_TOKENIZER_JSON, TEST_ADDED_TOKENS] + ) + + result = self.adapter.export_tokenizer( + self._make_tokenizer(TEST_ADDED_TOKENS), self.output_dir + ) + + self.assertEqual(result, self.artifact_dir / TEST_TOKENIZER_JSON) + + def test_prefers_tokenizer_json_over_tokenizer_model(self): + self._touch_artifacts([TEST_TOKENIZER_MODEL, TEST_TOKENIZER_JSON]) + + result = self.adapter.export_tokenizer( + self._make_tokenizer(TEST_TOKENIZER_MODEL), self.output_dir + ) + + self.assertEqual(result, self.artifact_dir / TEST_TOKENIZER_JSON) + + def test_falls_back_to_tokenizer_model_when_no_json(self): + self._touch_artifacts( + [TEST_TOKENIZER_CONFIG, TEST_TOKENIZER_MODEL, TEST_SPECIAL_TOKENS_MAP] + ) + + result = self.adapter.export_tokenizer( + self._make_tokenizer(TEST_TOKENIZER_MODEL), self.output_dir + ) + + self.assertEqual(result, self.artifact_dir / TEST_TOKENIZER_MODEL) + + def test_falls_back_to_runtime_tokenizer_when_no_known_name(self): + runtime_name = "tokenizer.bin" + tokenizer = self._make_tokenizer(runtime_name) + + result = self.adapter.export_tokenizer(tokenizer, self.output_dir) + + self.assertEqual(result, self.artifact_dir / runtime_name) + + def test_appends_runtime_tokenizer_when_artifact_dir_does_not_list_it(self): + self._touch_artifacts([TEST_TOKENIZER_CONFIG]) + runtime_path = Path(self._tmp.name) / TEST_TOKENIZER_MODEL + runtime_path.touch() + tokenizer = MagicMock( + runtime_tokenizer_path=str(runtime_path), + artifact=str(self.artifact_dir), + ) + + result = self.adapter.export_tokenizer(tokenizer, self.output_dir) + + self.assertEqual(result, runtime_path) + + +class TestGetExampleInputs(unittest.TestCase): + """The static LLM graph module owns its export signature.""" + + def setUp(self): + self.model_config = MagicMock() + self.control_args = MagicMock() + self.adapter = LLMLoaderAdapter(self.model_config, self.control_args) + + 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_requires_model_signature_instead_of_synthesizing(self): + with self.assertRaises(ValueError): + self.adapter.get_example_inputs(MagicMock(spec=[])) + + +if __name__ == "__main__": + unittest.main() diff --git a/backends/qualcomm/genai_pipeline/tests/strategies/model_preparation/test_mllm_loader_adapter.py b/backends/qualcomm/genai_pipeline/tests/strategies/model_preparation/test_mllm_loader_adapter.py new file mode 100644 index 00000000000..d8d6bc3a096 --- /dev/null +++ b/backends/qualcomm/genai_pipeline/tests/strategies/model_preparation/test_mllm_loader_adapter.py @@ -0,0 +1,372 @@ +# 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 tempfile +import unittest +from pathlib import Path +from unittest.mock import MagicMock, patch + +from executorch.backends.qualcomm.genai_pipeline.artifact_keys import ( + ARTIFACT_TEXT_DECODER, + ARTIFACT_TOK_EMBEDDING, + ARTIFACT_VISION_ENCODER, +) +from executorch.backends.qualcomm.genai_pipeline.graph_names import ( + DECODER_GRAPH_NAMES, + GRAPH_FORWARD, + TOK_EMBEDDING_GRAPH_NAMES, +) + +from executorch.backends.qualcomm.genai_pipeline.strategies.model_preparation.mllm_loader_adapter import ( + MLLMLoaderAdapter, +) + +DEFAULT_GRAPH_NAME = GRAPH_FORWARD +DECODE_GRAPH_NAME = DECODER_GRAPH_NAMES[0] +EMBED_DECODE_GRAPH_NAME = TOK_EMBEDDING_GRAPH_NAMES[0] +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" + + +def _make_model_config(): + model_config = MagicMock() + model_config.repo_id = "repo" + model_config.convert_weights = MagicMock() + vision_config = MagicMock() + vision_config.create_encoder.return_value = MagicMock(name="encoder") + vision_config.create_encoder.return_value.eval.return_value = ( + vision_config.create_encoder.return_value + ) + model_config.vision_encoder = MagicMock(return_value=vision_config) + return model_config + + +def _make_adapter(): + control_args = MagicMock() + control_args.checkpoint = "/tmp/checkpoint.pt" + control_args.model = "test_model" + return MLLMLoaderAdapter(_make_model_config(), control_args) + + +def _make_module(name): + module = MagicMock(name=name) + module.eval.return_value = module + return module + + +class TestMLLMLoaderAdapter(unittest.TestCase): + def test_requires_modality_encoder_config(self): + model_config = MagicMock(spec=[]) + + with self.assertRaises(ValueError): + MLLMLoaderAdapter(model_config, MagicMock()) + + def test_get_metadata_reads_module_when_it_exposes_metadata(self): + adapter = _make_adapter() + decoder = MagicMock() + decoder.get_metadata.return_value = {"get_n_layers": 2} + + result = adapter.get_metadata(decoder) + + self.assertEqual(result, {"get_n_layers": 2}) + decoder.get_metadata.assert_called_once_with() + + def test_get_metadata_returns_empty_when_module_exposes_none(self): + adapter = _make_adapter() + vision = MagicMock(spec=[]) + + result = adapter.get_metadata(vision) + + self.assertEqual(result, {}) + + def test_load_embedding_builds_each_graph_with_shared_weights(self): + adapter = _make_adapter() + weights = MagicMock(name="weights") + weights.to.return_value = weights + auto_model = MagicMock() + auto_model.get_input_embeddings.return_value = weights + build_calib = MagicMock(return_value=MagicMock(name="embed_calib")) + build_decode = MagicMock(return_value=MagicMock(name="embed_decode")) + + result = adapter._load_embedding( + auto_model, + { + DEFAULT_GRAPH_NAME: build_calib, + EMBED_DECODE_GRAPH_NAME: build_decode, + }, + ) + + build_calib.assert_called_once_with(weights) + build_decode.assert_called_once_with(weights) + self.assertEqual( + result, + { + DEFAULT_GRAPH_NAME: build_calib.return_value, + EMBED_DECODE_GRAPH_NAME: build_decode.return_value, + }, + ) + + def test_load_encoder_wraps_shared_auto_model_state(self): + adapter = _make_adapter() + auto_model = MagicMock() + auto_model.state_dict.return_value = {"encoder.weight": MagicMock()} + encoder = ( + adapter.model_config.vision_encoder.return_value.create_encoder.return_value + ) + + result = adapter._load_encoder(ARTIFACT_VISION_ENCODER, auto_model) + + adapter.model_config.vision_encoder.assert_called_once_with() + adapter.model_config.vision_encoder.return_value.create_encoder.assert_called_once_with( + auto_model.config + ) + encoder.load_state_dict.assert_called_once_with( + auto_model.state_dict.return_value, strict=False + ) + self.assertEqual(result, {DEFAULT_GRAPH_NAME: encoder}) + + def test_load_decoder_loads_checkpoint_into_each_graph(self): + adapter = _make_adapter() + calibration_module = _make_module("decoder_calibration") + decode_module = _make_module("decoder_decode") + state_dict = {"weight": MagicMock(name="weight")} + transform = MagicMock(return_value={"weight": MagicMock(name="rewritten")}) + + with patch("torch.load", return_value=state_dict) as torch_load: + result = adapter._load_decoder( + { + DEFAULT_GRAPH_NAME: MagicMock(return_value=calibration_module), + DECODE_GRAPH_NAME: MagicMock(return_value=decode_module), + }, + [transform], + ) + + torch_load.assert_called_once_with( + "/tmp/checkpoint.pt", weights_only=True, map_location="cpu", mmap=True + ) + transform.assert_called_once_with(state_dict) + calibration_module.load_state_dict.assert_called_once_with( + transform.return_value, strict=True, assign=True + ) + decode_module.load_state_dict.assert_called_once_with( + transform.return_value, strict=True, assign=True + ) + self.assertEqual( + result, + { + DEFAULT_GRAPH_NAME: calibration_module, + DECODE_GRAPH_NAME: decode_module, + }, + ) + + def test_load_decoder_state_dict_uses_injected_hf_loader(self): + adapter = _make_adapter() + state_dict = {"weight": MagicMock(name="weight")} + state_dict_loader = MagicMock(return_value=state_dict) + adapter.control_args.checkpoint = None + adapter.model_config.repo_id = "test/repo" + + with patch("torch.load") as torch_load: + result = adapter._load_decoder_state_dict(state_dict_loader) + + state_dict_loader.assert_called_once_with("test/repo") + torch_load.assert_not_called() + self.assertIs(result, state_dict) + + def test_load_decoder_passes_injected_hf_loader(self): + adapter = _make_adapter() + calibration_module = _make_module("decoder_calibration") + state_dict_loader = MagicMock(name="state_dict_loader") + state_dict = {"weight": MagicMock(name="weight")} + + with patch.object( + adapter, + "_load_decoder_state_dict", + return_value=state_dict, + ) as load_state_dict: + adapter._load_decoder( + { + DEFAULT_GRAPH_NAME: MagicMock(return_value=calibration_module), + }, + [], + state_dict_loader, + ) + + load_state_dict.assert_called_once_with(state_dict_loader) + calibration_module.load_state_dict.assert_called_once_with( + state_dict, strict=True, assign=True + ) + + def test_load_model_uses_component_keyed_state_dict_loader_for_decoder(self): + adapter = _make_adapter() + calibration_module = _make_module("decoder_calibration") + embed_module = MagicMock(name="embed_calibration") + state_dict = {"weight": MagicMock(name="weight")} + state_dict_loader = MagicMock(return_value=state_dict) + adapter.control_args.checkpoint = None + adapter.model_config.repo_id = "test/repo" + + with patch.object(adapter, "_load_auto_model") as load_auto_model: + auto_model = load_auto_model.return_value + auto_model.get_input_embeddings.return_value.to.return_value = MagicMock( + name="embedding_weights" + ) + result = adapter.load_model( + "test_model", + extra_options={ + "model_arch": { + ARTIFACT_TEXT_DECODER: { + DEFAULT_GRAPH_NAME: MagicMock( + return_value=calibration_module + ), + }, + ARTIFACT_TOK_EMBEDDING: { + DEFAULT_GRAPH_NAME: MagicMock(return_value=embed_module), + }, + }, + "state_dict_loader": {ARTIFACT_TEXT_DECODER: state_dict_loader}, + }, + ) + + state_dict_loader.assert_called_once_with("test/repo") + calibration_module.load_state_dict.assert_called_once_with( + state_dict, strict=True, assign=True + ) + self.assertIs( + result[ARTIFACT_TEXT_DECODER][DEFAULT_GRAPH_NAME], + calibration_module, + ) + + +class TestGetExampleInputs(unittest.TestCase): + """Multimodal graph modules own their export signatures.""" + + def setUp(self): + self.adapter = _make_adapter() + + 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_accepts_singular_model_method(self): + expected = (MagicMock(name="input"),) + model = MagicMock(spec=["get_example_input"]) + model.get_example_input.return_value = expected + + result = self.adapter.get_example_inputs(model) + + self.assertEqual(result, expected) + model.get_example_input.assert_called_once_with() + + def test_requires_model_signature_instead_of_synthesizing(self): + with self.assertRaises(ValueError): + self.adapter.get_example_inputs(MagicMock(spec=[])) + + +class TestExportTokenizer(unittest.TestCase): + """TokenizerWrapper has already written artifacts; the adapter selects the + runtime file using the same name-based priority as the default adapter.""" + + def setUp(self): + self.adapter = _make_adapter() + self._tmp = tempfile.TemporaryDirectory() + self.artifact_dir = Path(self._tmp.name) / "tokenizer_artifacts" + self.output_dir = Path(self._tmp.name) / "tokenizer_out" + self.artifact_dir.mkdir() + + def tearDown(self): + self._tmp.cleanup() + + def _touch_artifacts(self, names): + paths = [] + for name in names: + path = self.artifact_dir / name + path.touch() + paths.append(path) + return paths + + def _make_tokenizer(self, runtime_name): + return MagicMock( + runtime_tokenizer_path=str(self.artifact_dir / runtime_name), + artifact=str(self.artifact_dir), + ) + + def test_returns_tokenizer_file_not_the_directory(self): + self._touch_artifacts( + [TEST_TOKENIZER_CONFIG, TEST_SPECIAL_TOKENS_MAP, TEST_TOKENIZER_JSON] + ) + + result = self.adapter.export_tokenizer( + self._make_tokenizer(TEST_TOKENIZER_JSON), self.output_dir + ) + + self.assertEqual(result, self.artifact_dir / TEST_TOKENIZER_JSON) + self.assertNotEqual(result, self.artifact_dir) + + def test_selects_by_name_not_by_position(self): + self._touch_artifacts( + [TEST_TOKENIZER_CONFIG, TEST_TOKENIZER_JSON, TEST_ADDED_TOKENS] + ) + + result = self.adapter.export_tokenizer( + self._make_tokenizer(TEST_ADDED_TOKENS), self.output_dir + ) + + self.assertEqual(result, self.artifact_dir / TEST_TOKENIZER_JSON) + + def test_prefers_tokenizer_json_over_tokenizer_model(self): + self._touch_artifacts([TEST_TOKENIZER_MODEL, TEST_TOKENIZER_JSON]) + + result = self.adapter.export_tokenizer( + self._make_tokenizer(TEST_TOKENIZER_MODEL), self.output_dir + ) + + self.assertEqual(result, self.artifact_dir / TEST_TOKENIZER_JSON) + + def test_falls_back_to_tokenizer_model_when_no_json(self): + self._touch_artifacts( + [TEST_TOKENIZER_CONFIG, TEST_TOKENIZER_MODEL, TEST_SPECIAL_TOKENS_MAP] + ) + + result = self.adapter.export_tokenizer( + self._make_tokenizer(TEST_TOKENIZER_MODEL), self.output_dir + ) + + self.assertEqual(result, self.artifact_dir / TEST_TOKENIZER_MODEL) + + def test_falls_back_to_runtime_tokenizer_when_no_known_name(self): + runtime_name = "tokenizer.bin" + tokenizer = self._make_tokenizer(runtime_name) + + result = self.adapter.export_tokenizer(tokenizer, self.output_dir) + + self.assertEqual(result, self.artifact_dir / runtime_name) + + def test_appends_runtime_tokenizer_when_artifact_dir_does_not_list_it(self): + self._touch_artifacts([TEST_TOKENIZER_CONFIG]) + runtime_path = Path(self._tmp.name) / TEST_TOKENIZER_MODEL + runtime_path.touch() + tokenizer = MagicMock( + runtime_tokenizer_path=str(runtime_path), + artifact=str(self.artifact_dir), + ) + + result = self.adapter.export_tokenizer(tokenizer, self.output_dir) + + self.assertEqual(result, runtime_path) + + +if __name__ == "__main__": + unittest.main() diff --git a/backends/qualcomm/genai_pipeline/tests/strategies/quantization/test_default_quantizer_adapter.py b/backends/qualcomm/genai_pipeline/tests/strategies/quantization/test_default_quantizer_adapter.py index 587bfafde51..a0b298f5bc9 100644 --- a/backends/qualcomm/genai_pipeline/tests/strategies/quantization/test_default_quantizer_adapter.py +++ b/backends/qualcomm/genai_pipeline/tests/strategies/quantization/test_default_quantizer_adapter.py @@ -5,10 +5,17 @@ # LICENSE file in the root directory of this source tree. import unittest -from unittest.mock import MagicMock +from pathlib import Path +from tempfile import TemporaryDirectory +from unittest.mock import MagicMock, patch import torch +from executorch.backends.qualcomm.genai_pipeline.artifact_keys import ( + ARTIFACT_TEXT_DECODER, + ARTIFACT_VISION_ENCODER, + DECODE_QDQ_FILENAME, +) from executorch.backends.qualcomm.genai_pipeline.strategies.quantization.default_quantizer_adapter import ( DefaultQuantizerAdapter, ) @@ -53,6 +60,48 @@ def test_empty_dataset_is_a_no_op(self): self.adapter.calibrate(model, []) model.assert_not_called() + def test_component_map_calibrates_components_with_data(self): + text_model = MagicMock() + vision_model = MagicMock() + input_ids, attention_mask = torch.zeros(1, 2), torch.ones(1, 2) + + result = self.adapter.calibrate( + { + ARTIFACT_TEXT_DECODER: text_model, + ARTIFACT_VISION_ENCODER: vision_model, + }, + { + ARTIFACT_TEXT_DECODER: [(input_ids, attention_mask)], + }, + ) + + self.assertIs(result[ARTIFACT_TEXT_DECODER], text_model) + text_model.assert_called_once_with(input_ids, attention_mask) + vision_model.assert_not_called() + + def test_save_quantized_module_exports_strict_qdq_program(self): + model = MagicMock(name="quantized_model") + example_inputs = (torch.zeros(1, 2),) + exported_program = MagicMock(name="exported_program") + + with TemporaryDirectory() as artifact_dir, patch( + "torch.export.export", return_value=exported_program + ) as export, patch("torch.export.save") as save: + from executorch.backends.qualcomm.genai_pipeline.quant_utilities import ( + save_quantized_module, + ) + + result = save_quantized_module( + model, + example_inputs, + artifact_dir, + ) + + expected_path = Path(artifact_dir) / DECODE_QDQ_FILENAME + export.assert_called_once_with(model, example_inputs, strict=True) + save.assert_called_once_with(exported_program, expected_path) + self.assertEqual(result, expected_path) + 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 a1e1bbc49ae..25a233fbfd1 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,7 +5,14 @@ # LICENSE file in the root directory of this source tree. import unittest -from unittest.mock import create_autospec, MagicMock, patch +from unittest.mock import call, MagicMock, patch + +import torch +from executorch.backends.qualcomm.genai_pipeline.artifact_keys import ( + ARTIFACT_TEXT_DECODER, + ARTIFACT_TOK_EMBEDDING, + ARTIFACT_VISION_ENCODER, +) from executorch.backends.qualcomm.genai_pipeline.configs.quantization_input_config import ( QuantizationInputConfig, @@ -14,6 +21,12 @@ QuantizationOutputConfig, ) from executorch.backends.qualcomm.genai_pipeline.exceptions import StageError +from executorch.backends.qualcomm.genai_pipeline.graph_bundle import GraphBundle +from executorch.backends.qualcomm.genai_pipeline.graph_names import ( + DECODER_GRAPH_NAMES, + GRAPH_FORWARD, + TOK_EMBEDDING_GRAPH_NAMES, +) from executorch.backends.qualcomm.genai_pipeline.strategies.quantization.executorch_quantization_strategy import ( ExecuTorchQuantizationStrategy, ) @@ -29,248 +42,723 @@ TEST_SOC_CHIPSET, ) +DECODE_GRAPH_NAME = DECODER_GRAPH_NAMES[0] +PREFILL_GRAPH_NAME = DECODER_GRAPH_NAMES[1] +EMBED_DECODE_GRAPH_NAME = TOK_EMBEDDING_GRAPH_NAMES[0] + + +def _make_recipe(name, kv_bit_width=8, logits_bit_width=16): + recipe = MagicMock(name=name) + recipe.recipe = {name: "recipe"} + recipe.get_kv_io_bit_width.return_value = kv_bit_width + recipe.get_logits_output_bit_width.return_value = logits_bit_width + return recipe + 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") + """Create a mock quantizer adapter with per-step return values.""" + adapter = MagicMock(spec=QuantizerAdapter) + adapter.export_model.side_effect = lambda module, example_inputs: MagicMock( + name=f"exported_{module._mock_name}_{example_inputs[0]}" + ) + adapter.prepare_pt2e.side_effect = lambda module, quantizer: MagicMock( + name=f"prepared_{module._mock_name}" + ) + adapter.convert_pt2e.side_effect = lambda module: MagicMock( + name=f"converted_{module._mock_name}" + ) + return adapter + + +def _make_mock_calibration_adapter(): + """Create a mock calibration data adapter with component-keyed data.""" + adapter = MagicMock() + adapter.generate_calibration_data.return_value = { + ARTIFACT_TEXT_DECODER: MagicMock(name="text_calibration_data") + } return adapter +def _make_strategy(adapter=None, calibration=None): + """Build the strategy with quantization and calibration adapters mocked.""" + return ExecuTorchQuantizationStrategy( + quantizer_adapter=adapter if adapter 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 QuantizationInputConfig with defaults.""" + """Create a valid QuantizationInputConfig with component/graph defaults.""" defaults = { "soc_model": TEST_SOC_CHIPSET, "backend_type": TEST_BACKEND_TYPE, - "model_module": MagicMock(name="test_model"), - "example_inputs": (MagicMock(name="example_input"),), - "calibration_data": [(MagicMock(),)], + "model_module": { + ARTIFACT_TEXT_DECODER: MagicMock(name="decoder"), + ARTIFACT_TOK_EMBEDDING: MagicMock(name="tok_embedding"), + ARTIFACT_VISION_ENCODER: MagicMock(name="vision_encoder"), + }, + "example_inputs": { + ARTIFACT_TEXT_DECODER: { + GRAPH_FORWARD: ( + "text_calib", + ("text_calib_mask",), + "text_calib_pos_ids", + ("text_calib_k_cache",), + ("text_calib_v_cache",), + ), + DECODE_GRAPH_NAME: ( + "text_decode", + ("text_decode_mask",), + "text_decode_pos_ids", + ("text_decode_k_cache",), + ("text_decode_v_cache",), + ), + PREFILL_GRAPH_NAME: ( + "text_prefill", + ("text_prefill_mask",), + "text_prefill_pos_ids", + ("text_prefill_k_cache",), + ("text_prefill_v_cache",), + ), + }, + ARTIFACT_TOK_EMBEDDING: { + GRAPH_FORWARD: ("embed_calib",), + EMBED_DECODE_GRAPH_NAME: ("embed_decode",), + }, + ARTIFACT_VISION_ENCODER: { + GRAPH_FORWARD: ("vision",), + }, + }, + "tokenizer": MagicMock(name="tokenizer"), + "meta": { + ARTIFACT_TEXT_DECODER: { + GRAPH_FORWARD: { + "get_n_layers": 2, + "get_use_kv_cache": False, + }, + DECODE_GRAPH_NAME: {"get_n_layers": 2, "get_use_kv_cache": True}, + PREFILL_GRAPH_NAME: {"get_n_layers": 2, "get_use_kv_cache": False}, + }, + ARTIFACT_TOK_EMBEDDING: { + GRAPH_FORWARD: {"get_n_layers": 0}, + EMBED_DECODE_GRAPH_NAME: {"get_n_layers": 0}, + }, + ARTIFACT_VISION_ENCODER: { + GRAPH_FORWARD: {}, + }, + }, + "inference": MagicMock(name="inference"), + "extra_options": { + "quantize_options": { + "quant_recipe": { + ARTIFACT_TEXT_DECODER: _make_recipe("text_recipe"), + ARTIFACT_TOK_EMBEDDING: None, + ARTIFACT_VISION_ENCODER: _make_recipe("vision_recipe"), + } + } + }, } defaults.update(overrides) return QuantizationInputConfig(**defaults) class TestExecuTorchQuantizationStrategy(unittest.TestCase): + def setUp(self): + self.save_quantized_module = patch( + "executorch.backends.qualcomm.genai_pipeline.quant_utilities." + "save_quantized_module" + ).start() + self.encoding_override = patch( + "executorch.backends.qualcomm.genai_pipeline.quant_utilities." + "encoding_override" + ).start() + self.save_logits_quant_attrs = patch( + "executorch.backends.qualcomm.genai_pipeline.quant_utilities." + "save_logits_quant_attrs" + ).start() + self.save_output_kv_cache_quant_attrs = patch( + "executorch.backends.qualcomm.genai_pipeline.quant_utilities." + "save_output_kv_cache_quant_attrs" + ).start() + self.make_quantizer = patch( + "executorch.backends.qualcomm.genai_pipeline.quant_utilities." + "make_quantizer" + ).start() + self.addCleanup(patch.stopall) def test_is_quantization_strategy(self): """Strategy inherits from QuantizationStrategy ABC.""" - adapter = _make_mock_adapter() - strategy = ExecuTorchQuantizationStrategy(quantizer_adapter=adapter) + strategy = _make_strategy() self.assertIsInstance(strategy, QuantizationStrategy) - def test_default_adapter_created_when_none_provided(self): - """When no adapter is provided, DefaultQuantizerAdapter is created.""" + def test_default_adapters_created_when_none_provided(self): + """Default quantization and purpose adapters are created when omitted.""" 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.""" + ) as quantizer_cls, patch( + "executorch.backends.qualcomm.genai_pipeline.datasets.calibration." + "default_calibration_data_adapter.DefaultCalibrationDataAdapter" + ) as calibration_cls, patch( + "executorch.backends.qualcomm.genai_pipeline.datasets.training." + "default_training_data_adapter.DefaultTrainingDataAdapter" + ) as training_cls, patch( + "executorch.backends.qualcomm.genai_pipeline.datasets.evaluation." + "default_evaluation_data_adapter.DefaultEvaluationDataAdapter" + ) as evaluation_cls: + strategy = ExecuTorchQuantizationStrategy() + + quantizer_cls.assert_called_once() + calibration_cls.assert_called_once() + training_cls.assert_called_once() + evaluation_cls.assert_called_once() + self.assertIs(strategy.adapter, quantizer_cls.return_value) + + def test_custom_adapters_injected(self): + """Custom quantizer and calibration adapters are used when provided.""" adapter = _make_mock_adapter() - strategy = ExecuTorchQuantizationStrategy(quantizer_adapter=adapter) + calibration = _make_mock_calibration_adapter() + strategy = _make_strategy(adapter, calibration) self.assertIs(strategy.adapter, adapter) - def test_invoke_happy_path(self): - """Full quantization pipeline runs successfully end-to-end.""" + def test_invoke_returns_graph_bundles_for_deployable_graphs(self): + """Full quantization returns deployable graph bundles per component.""" adapter = _make_mock_adapter() - strategy = ExecuTorchQuantizationStrategy(quantizer_adapter=adapter) + strategy = _make_strategy(adapter) + input_config = _make_valid_input_config() - result = strategy.invoke(make_test_context(), _make_valid_input_config()) + result = strategy.invoke(make_test_context(), input_config) self.assertIsInstance(result, QuantizationOutputConfig) - self.assertIs(result.quantized_model, adapter.convert_pt2e.return_value) + self.assertNotIn(GRAPH_FORWARD, result.graphs[ARTIFACT_TEXT_DECODER]) + self.assertIn(DECODE_GRAPH_NAME, result.graphs[ARTIFACT_TEXT_DECODER]) + self.assertIn(PREFILL_GRAPH_NAME, result.graphs[ARTIFACT_TEXT_DECODER]) + self.assertNotIn(GRAPH_FORWARD, result.graphs[ARTIFACT_TOK_EMBEDDING]) + self.assertIn(EMBED_DECODE_GRAPH_NAME, result.graphs[ARTIFACT_TOK_EMBEDDING]) + self.assertIn(GRAPH_FORWARD, result.graphs[ARTIFACT_VISION_ENCODER]) + self.assertIsInstance( + result.graphs[ARTIFACT_TEXT_DECODER][DECODE_GRAPH_NAME], GraphBundle + ) + + def test_invoke_exports_each_graph_with_component_weight_holder(self): + """export_model receives each graph's inputs and shared component module. - def test_invoke_calls_adapter_in_correct_order(self): - """Adapter methods are called in the correct PT2E sequence.""" + The export signature comes from model preparation's example inputs + (zero-initialized KV caches, fixed AR length), not from a calibration + sample, which lacks the full export signature. + """ adapter = _make_mock_adapter() - strategy = ExecuTorchQuantizationStrategy(quantizer_adapter=adapter) + strategy = _make_strategy(adapter) + input_config = _make_valid_input_config() - strategy.invoke(make_test_context(), _make_valid_input_config()) + strategy.invoke(make_test_context(), input_config) - # Verify call order: export → make_quantizer → prepare → calibrate → convert - self.assertEqual( - [c[0] for c in adapter.method_calls], + adapter.export_model.assert_has_calls( [ - "export_model", - "make_quantizer", - "prepare_pt2e", - "calibrate", - "convert_pt2e", + call( + input_config.model_module[ARTIFACT_TEXT_DECODER], + ("text_calib", "text_calib_mask"), + ), + call( + input_config.model_module[ARTIFACT_TEXT_DECODER], + ( + "text_decode", + "text_decode_mask", + "text_decode_pos_ids", + "text_decode_k_cache", + "text_decode_v_cache", + ), + ), + call( + input_config.model_module[ARTIFACT_TEXT_DECODER], + ("text_prefill", "text_prefill_mask"), + ), + call( + input_config.model_module[ARTIFACT_TOK_EMBEDDING], + input_config.example_inputs[ARTIFACT_TOK_EMBEDDING][GRAPH_FORWARD], + ), + call( + input_config.model_module[ARTIFACT_TOK_EMBEDDING], + input_config.example_inputs[ARTIFACT_TOK_EMBEDDING][ + EMBED_DECODE_GRAPH_NAME + ], + ), + call( + input_config.model_module[ARTIFACT_VISION_ENCODER], + input_config.example_inputs[ARTIFACT_VISION_ENCODER][GRAPH_FORWARD], + ), ], + any_order=False, ) - def test_invoke_exports_with_example_inputs_not_calibration_data(self): - """export_model receives ``example_inputs``, never a calibration sample. + def test_export_and_prepare_flattens_text_decoder_inputs(self): + adapter = _make_mock_adapter() + strategy = _make_strategy(adapter) + decoder = MagicMock(name="decoder") + tokens = MagicMock(name="tokens") + attention_masks = (MagicMock(name="causal_mask"), MagicMock(name="local_mask")) + pos_ids = MagicMock(name="pos_ids") + k_caches = (MagicMock(name="k_cache_0"), MagicMock(name="k_cache_1")) + v_caches = (MagicMock(name="v_cache_0"), MagicMock(name="v_cache_1")) + input_config = _make_valid_input_config( + model_module={ARTIFACT_TEXT_DECODER: decoder}, + example_inputs={ + ARTIFACT_TEXT_DECODER: { + DECODE_GRAPH_NAME: ( + tokens, + attention_masks, + pos_ids, + k_caches, + v_caches, + ) + } + }, + meta={ + ARTIFACT_TEXT_DECODER: {DECODE_GRAPH_NAME: {"get_use_kv_cache": True}} + }, + ) + + strategy._export_and_prepare( + input_config, + {ARTIFACT_TEXT_DECODER: {DECODE_GRAPH_NAME: MagicMock(name="quantizer")}}, + ) - 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.export_model.assert_called_once_with( + decoder, + (tokens, *attention_masks, pos_ids, *k_caches, *v_caches), + ) + + def test_invoke_generates_calibration_data_from_adapter(self): + """Calibration data is generated with flattened example inputs. + + The calibration adapter receives only the non-deployed graph inputs for + each component, so collators derive masks from the calibration signature + rather than deployed decode/prefill signatures. """ + calibration = _make_mock_calibration_adapter() + strategy = _make_strategy(calibration=calibration) + extra_options = { + "model_options": {"model_arch": MagicMock(name="model_arch")}, + "quantize_options": { + "quant_dtype": { + ARTIFACT_TEXT_DECODER: "text_dtype", + ARTIFACT_TOK_EMBEDDING: "embed_dtype", + ARTIFACT_VISION_ENCODER: "vision_dtype", + } + }, + "dataset_options": MagicMock(name="dataset_options"), + } + input_config = _make_valid_input_config(extra_options=extra_options) + + strategy.invoke(make_test_context(), input_config) + + _, kwargs = calibration.generate_calibration_data.call_args + self.assertIs(kwargs["tokenizer"], input_config.tokenizer) + self.assertEqual( + kwargs["example_inputs"], + { + ARTIFACT_TEXT_DECODER: input_config.example_inputs[ + ARTIFACT_TEXT_DECODER + ][GRAPH_FORWARD], + ARTIFACT_TOK_EMBEDDING: input_config.example_inputs[ + ARTIFACT_TOK_EMBEDDING + ][GRAPH_FORWARD], + ARTIFACT_VISION_ENCODER: input_config.example_inputs[ + ARTIFACT_VISION_ENCODER + ][GRAPH_FORWARD], + }, + ) + self.assertNotIn("extra_options", kwargs) + + def test_invoke_calibrates_only_non_deployed_graphs(self): + """Only each component's non-deployed graph is truly calibrated.""" 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], + calibration = _make_mock_calibration_adapter() + strategy = _make_strategy(adapter, calibration) + input_config = _make_valid_input_config() + + strategy.invoke(make_test_context(), input_config) + + quantization_graphs = adapter.calibrate.call_args[0][0] + self.assertEqual( + set(quantization_graphs), + {ARTIFACT_TEXT_DECODER, ARTIFACT_TOK_EMBEDDING, ARTIFACT_VISION_ENCODER}, + ) + adapter.calibrate.assert_called_once_with( + quantization_graphs, + calibration.generate_calibration_data.return_value, + inference=input_config.inference, ) + def test_invoke_initializes_deployed_graph_observers(self): + """Deployed graph observers are initialized with their own inputs.""" + adapter = _make_mock_adapter() + strategy = _make_strategy(adapter) + input_config = _make_valid_input_config() + strategy.invoke(make_test_context(), input_config) - adapter.export_model.assert_called_once_with(model, example_inputs) + initialized_inputs = [ + args[0][1] for args in adapter.init_encodings.call_args_list + ] + self.assertEqual( + initialized_inputs, + [ + ( + "text_decode", + "text_decode_mask", + "text_decode_pos_ids", + "text_decode_k_cache", + "text_decode_v_cache", + ), + ("text_prefill", "text_prefill_mask"), + input_config.example_inputs[ARTIFACT_TOK_EMBEDDING][ + EMBED_DECODE_GRAPH_NAME + ], + ], + ) + + def test_invoke_overrides_only_text_and_embedding_deployed_graphs(self): + """Encoding override applies only where calibration/deploy graph split exists.""" + adapter = _make_mock_adapter() + strategy = _make_strategy(adapter) + input_config = _make_valid_input_config() - def test_invoke_passes_calibration_data_through_untouched(self): - """``calibration_data`` reaches ``calibrate`` as the very same object. + strategy.invoke(make_test_context(), input_config) - 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. - """ + self.assertEqual(self.encoding_override.call_count, 3) + n_cache_layers = [ + kwargs.get("n_cache_layers") + for _, kwargs in self.encoding_override.call_args_list + ] + self.assertEqual( + n_cache_layers, + [ + 2, + 2, + None, + ], + ) + text_override_calls = self.encoding_override.call_args_list[:2] + self.assertIs( + text_override_calls[1].kwargs["quantized_model"], + text_override_calls[0].kwargs["quantized_model"], + ) + self.assertEqual(self.save_logits_quant_attrs.call_count, 2) + self.assertEqual(self.save_output_kv_cache_quant_attrs.call_count, 2) + + def test_invoke_skips_missing_text_decoder_prefill_graph(self): 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 = _make_strategy(adapter) + input_config = _make_valid_input_config() + input_config.example_inputs[ARTIFACT_TEXT_DECODER].pop(PREFILL_GRAPH_NAME) + input_config.meta[ARTIFACT_TEXT_DECODER].pop(PREFILL_GRAPH_NAME) 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) + self.assertEqual(self.encoding_override.call_count, 2) + self.assertEqual(self.save_logits_quant_attrs.call_count, 1) + self.assertEqual(self.save_output_kv_cache_quant_attrs.call_count, 1) - def test_invoke_missing_example_inputs_raises_stage_error(self): - """StageError raised when example_inputs is None.""" + def test_invoke_saves_text_decoder_qdq_module(self): + """The text decoder calibration graph is saved for SQNR evaluation.""" adapter = _make_mock_adapter() - strategy = ExecuTorchQuantizationStrategy(quantizer_adapter=adapter) - input_config = _make_valid_input_config(example_inputs=None) + strategy = _make_strategy(adapter) + input_config = _make_valid_input_config() + context = make_test_context(artifact_dir="/tmp/qdq_artifacts") - 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() + strategy.invoke(context, input_config) + + self.save_quantized_module.assert_called_once() + _, kwargs = self.save_quantized_module.call_args + self.assertIs( + kwargs["quantized_module"], + self.encoding_override.call_args_list[0].kwargs["quantized_model"], + ) + self.assertEqual( + kwargs["example_inputs"], + ("text_calib", "text_calib_mask"), + ) + self.assertEqual(kwargs["artifact_dir"], context.artifact_dir) - def test_invoke_passes_correct_args_to_make_quantizer(self): - """make_quantizer receives backend_type and soc_model (no quant_dtype when not set).""" + def test_save_quantized_module_flattens_quantization_graph_inputs(self): 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 = _make_strategy(adapter) + input_config = _make_valid_input_config() + tokens = MagicMock(name="tokens") + attention_masks = (MagicMock(name="causal_mask"), MagicMock(name="local_mask")) + pos_ids = MagicMock(name="pos_ids") + k_caches = (MagicMock(name="k_cache_0"), MagicMock(name="k_cache_1")) + v_caches = (MagicMock(name="v_cache_0"), MagicMock(name="v_cache_1")) + input_config.example_inputs[ARTIFACT_TEXT_DECODER][GRAPH_FORWARD] = ( + tokens, + attention_masks, + pos_ids, + k_caches, + v_caches, + ) + input_config.meta[ARTIFACT_TEXT_DECODER][GRAPH_FORWARD] = { + "get_use_kv_cache": True + } 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, + _, kwargs = self.save_quantized_module.call_args + self.assertEqual( + kwargs["example_inputs"], + (tokens, *attention_masks, pos_ids, *k_caches, *v_caches), ) - def test_invoke_passes_quant_dtype_from_extra_options(self): - """quant_dtype extracted from extra_options and forwarded.""" + def test_invoke_routes_component_quant_options_to_each_graph_quantizer(self): + """make_quantizer receives each graph's component recipe and dtype.""" adapter = _make_mock_adapter() - strategy = ExecuTorchQuantizationStrategy(quantizer_adapter=adapter) - quant_dtype = MagicMock(name="quant_dtype") + strategy = _make_strategy(adapter) + text_recipe = _make_recipe("text_recipe") + embed_recipe = _make_recipe("embed_recipe") + vision_recipe = _make_recipe("vision_recipe") input_config = _make_valid_input_config( - extra_options={"quant_dtype": quant_dtype} + extra_options={ + "model_arch": MagicMock(name="model_arch"), + "model_options": {"model_arch": MagicMock(name="nested_model_arch")}, + "quantize_options": { + "quant_dtype": { + ARTIFACT_TEXT_DECODER: "text_dtype", + ARTIFACT_TOK_EMBEDDING: "embed_dtype", + ARTIFACT_VISION_ENCODER: "vision_dtype", + }, + "quant_recipe": { + ARTIFACT_TEXT_DECODER: text_recipe, + ARTIFACT_TOK_EMBEDDING: embed_recipe, + ARTIFACT_VISION_ENCODER: vision_recipe, + }, + "per_channel_conv": True, + }, + "dataset_options": {"max_context_len": 128}, + } ) 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, + self.make_quantizer.assert_has_calls( + [ + call( + backend=input_config.backend_type, + soc_model=input_config.soc_model, + quant_dtype="text_dtype", + quant_recipe=text_recipe, + ), + call( + backend=input_config.backend_type, + soc_model=input_config.soc_model, + quant_dtype="text_dtype", + quant_recipe=text_recipe, + ), + call( + backend=input_config.backend_type, + soc_model=input_config.soc_model, + quant_dtype="text_dtype", + quant_recipe=text_recipe, + ), + call( + backend=input_config.backend_type, + soc_model=input_config.soc_model, + quant_dtype="embed_dtype", + quant_recipe=embed_recipe, + ), + call( + backend=input_config.backend_type, + soc_model=input_config.soc_model, + quant_dtype="embed_dtype", + quant_recipe=embed_recipe, + ), + call( + backend=input_config.backend_type, + soc_model=input_config.soc_model, + quant_dtype="vision_dtype", + quant_recipe=vision_recipe, + ), + ], + any_order=False, + ) + + def test_make_quantizer_sets_recipe_verbose_only_for_calibration_graph(self): + adapter = _make_mock_adapter() + strategy = _make_strategy(adapter) + input_config = _make_valid_input_config() + verbose_values = [] + + class RecordingRecipe: + def __init__(self, *, verbose): + verbose_values.append(verbose) + + quantizers, _ = strategy._make_quantizer( + input_config, + { + "quant_recipe": {ARTIFACT_TEXT_DECODER: RecordingRecipe}, + }, ) - def test_invoke_passes_quant_recipe_from_config(self): - """quant_recipe on the config reaches the adapter as a declared kwarg. + self.assertEqual( + set(quantizers[ARTIFACT_TEXT_DECODER]), + {GRAPH_FORWARD, DECODE_GRAPH_NAME, PREFILL_GRAPH_NAME}, + ) + self.assertEqual(verbose_values, [True, False, False]) - The adapter consumes it (``QnnQuantizer.set_recipe``) rather than - forwarding it to ``export_utils.make_quantizer``, which takes no such - argument. - """ + def test_invoke_omits_quant_dtype_when_not_set(self): + """quant_dtype is omitted when not explicitly set in extra_options.""" adapter = _make_mock_adapter() - strategy = ExecuTorchQuantizationStrategy(quantizer_adapter=adapter) - recipe = MagicMock(name="quant_recipe") - input_config = _make_valid_input_config(quant_recipe=recipe) + strategy = _make_strategy(adapter) + input_config = _make_valid_input_config() strategy.invoke(make_test_context(), input_config) - _, kwargs = adapter.make_quantizer.call_args - self.assertIs(kwargs["quant_recipe"], recipe) + for _, kwargs in self.make_quantizer.call_args_list: + self.assertNotIn("quant_dtype", kwargs) - def test_invoke_training_data_does_not_switch_to_qat(self): - """``training_data`` is accepted but this strategy still performs PTQ.""" + def test_invoke_populates_graph_bundle_fields(self): + """GraphBundle carries the converted module, inputs, metadata and IO dtype.""" adapter = _make_mock_adapter() - strategy = ExecuTorchQuantizationStrategy(quantizer_adapter=adapter) - input_config = _make_valid_input_config(training_data=[(MagicMock(),)]) + strategy = _make_strategy(adapter) + input_config = _make_valid_input_config() - with self.assertLogs( - "executorch.backends.qualcomm.genai_pipeline.strategies.quantization." - "executorch_quantization_strategy", - level="WARNING", - ): - result = strategy.invoke(make_test_context(), input_config) + 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) + bundle = result.graphs[ARTIFACT_TEXT_DECODER][DECODE_GRAPH_NAME] + self.assertIs( + bundle.inputs, + input_config.example_inputs[ARTIFACT_TEXT_DECODER][DECODE_GRAPH_NAME], + ) + self.assertIs( + bundle.meta, input_config.meta[ARTIFACT_TEXT_DECODER][DECODE_GRAPH_NAME] + ) + self.assertEqual( + bundle.quant_io_dtypes, + {"kv_type": torch.uint8, "io_type": torch.uint16}, + ) + + def test_quant_io_dtypes_returns_none_for_missing_recipe(self): + """Components without an IO recipe leave GraphBundle IO dtypes unset.""" + strategy = _make_strategy() + + self.assertIsNone(strategy._get_quant_io_dtypes(None)) + + def test_quant_io_dtypes_rejects_unsupported_widths(self): + """Unsupported IO widths must not silently produce a partial result.""" + strategy = _make_strategy() + + with self.assertRaisesRegex(RuntimeError, "Unsupported quantization IO"): + strategy._get_quant_io_dtypes( + _make_recipe("unquantized_recipe", kv_bit_width=32) + ) + + def test_invoke_rejects_unsupported_quant_io_widths(self): + """An invalid recipe must surface as a quantization failure.""" + adapter = _make_mock_adapter() + strategy = _make_strategy(adapter) + input_config = _make_valid_input_config( + extra_options={ + "quantize_options": { + "quant_recipe": { + ARTIFACT_TEXT_DECODER: _make_recipe( + "text_recipe", + kv_bit_width=32, + logits_bit_width=32, + ), + } + } + } + ) + + with self.assertRaises(StageError) as cm: + strategy.invoke(make_test_context(), input_config) + self.assertIn("Unsupported quantization IO", str(cm.exception)) + + def test_invoke_allows_components_without_metadata(self): + """Graph bundles tolerate modules that do not expose get_metadata().""" + adapter = _make_mock_adapter() + strategy = _make_strategy(adapter) + input_config = _make_valid_input_config( + meta={ + ARTIFACT_TEXT_DECODER: { + GRAPH_FORWARD: { + "get_n_layers": 2, + "get_use_kv_cache": False, + }, + DECODE_GRAPH_NAME: { + "get_n_layers": 2, + "get_use_kv_cache": True, + }, + PREFILL_GRAPH_NAME: { + "get_n_layers": 2, + "get_use_kv_cache": False, + }, + } + } + ) + + result = strategy.invoke(make_test_context(), input_config) + + self.assertEqual( + result.graphs[ARTIFACT_TOK_EMBEDDING][EMBED_DECODE_GRAPH_NAME].meta, + {}, + ) + + def test_invoke_missing_example_inputs_raises_stage_error(self): + """StageError raised when example_inputs is None.""" + strategy = _make_strategy() + 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") 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) + strategy = _make_strategy() 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) + def test_invoke_missing_backend_type_raises_stage_error(self): + """StageError raised when backend_type is None.""" + strategy = _make_strategy() + 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_rejects_per_component_backend_map(self): + """backend_type must be one shared QNN backend, not a component map.""" + strategy = _make_strategy() + input_config = _make_valid_input_config( + backend_type={ARTIFACT_TEXT_DECODER: "htp"} + ) with self.assertRaises(StageError) as cm: strategy.invoke(make_test_context(), input_config) - self.assertIn("calibration_data", str(cm.exception)) + + self.assertIn("backend_type must be one shared QNN backend", 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) + strategy = _make_strategy(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)) @@ -280,55 +768,12 @@ def test_invoke_stage_error_not_double_wrapped(self): 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) + strategy = _make_strategy(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, - ) + self.assertIs(cm.exception, original_error) if __name__ == "__main__": diff --git a/backends/qualcomm/genai_pipeline/tests/strategies/quantization/test_llm_quantizer_adapter.py b/backends/qualcomm/genai_pipeline/tests/strategies/quantization/test_llm_quantizer_adapter.py new file mode 100644 index 00000000000..ed58972bcb4 --- /dev/null +++ b/backends/qualcomm/genai_pipeline/tests/strategies/quantization/test_llm_quantizer_adapter.py @@ -0,0 +1,121 @@ +# 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 sys +import types +import unittest +from unittest.mock import MagicMock, patch + +import torch +from executorch.backends.qualcomm.genai_pipeline.strategies.quantization.llm_quantizer_adapter import ( + LLMQuantizerAdapter, +) +from torch.utils.data import DataLoader + + +def _fake_quantize_pt2e(**functions): + module = types.ModuleType("torchao.quantization.pt2e.quantize_pt2e") + for name, function in functions.items(): + setattr(module, name, function) + return module + + +class TestLLMQuantizerAdapter(unittest.TestCase): + def setUp(self): + self.adapter = LLMQuantizerAdapter() + + def test_export_model_returns_exported_module(self): + module = MagicMock(name="decoder") + example_inputs = (MagicMock(name="tokens"),) + exported_module = MagicMock(name="exported_module") + exported_program = MagicMock(name="exported_program") + exported_program.module.return_value = exported_module + + with patch("torch.export.export", return_value=exported_program) as export: + result = self.adapter.export_model(module, example_inputs) + + self.assertIs(result, exported_module) + export.assert_called_once_with(module, example_inputs, strict=True) + + def test_prepare_pt2e_delegates_to_torchao(self): + module = MagicMock(name="exported_module") + quantizer = MagicMock(name="quantizer") + prepared_module = MagicMock(name="prepared_module") + prepare_pt2e = MagicMock(return_value=prepared_module) + + with patch.dict( + sys.modules, + { + "torchao.quantization.pt2e.quantize_pt2e": _fake_quantize_pt2e( + prepare_pt2e=prepare_pt2e + ) + }, + ): + result = self.adapter.prepare_pt2e(module, quantizer) + + self.assertIs(result, prepared_module) + prepare_pt2e.assert_called_once_with(module, quantizer) + + def test_init_encodings_runs_one_forward(self): + module = MagicMock(name="prepared_decoder") + example_inputs = (MagicMock(name="tokens"), MagicMock(name="attn_mask")) + + result = self.adapter.init_encodings(module, example_inputs) + + self.assertIs(result, module) + module.assert_called_once_with(*example_inputs) + + def test_calibrate_drives_decoder_with_corpus_batches(self): + decoder = MagicMock(name="decoder") + inference = MagicMock(name="inference") + batch = { + "input_ids": torch.tensor([[1, 2]]), + "attention_mask": torch.tensor([[1, 1]]), + } + dataloader = DataLoader([batch], batch_size=None) + + result = self.adapter.calibrate( + {"text_decoder": decoder}, + {"text_decoder": dataloader}, + inference=inference, + ) + + self.assertIsNone(result) + inference.predict_step.assert_called_once_with( + decoder, + input_ids=batch["input_ids"], + attn_mask=batch["attention_mask"], + ) + + def test_calibrate_rejects_non_dataloader_text_data(self): + with self.assertRaisesRegex(ValueError, "corpus-backed DataLoader"): + self.adapter.calibrate( + {"text_decoder": MagicMock(name="decoder")}, + {"text_decoder": []}, + inference=MagicMock(name="inference"), + ) + + def test_convert_pt2e_delegates_to_torchao(self): + module = MagicMock(name="prepared_module") + converted_module = MagicMock(name="converted_module") + convert_pt2e = MagicMock(return_value=converted_module) + + with patch.dict( + sys.modules, + { + "torchao.quantization.pt2e.quantize_pt2e": _fake_quantize_pt2e( + convert_pt2e=convert_pt2e + ) + }, + ): + result = self.adapter.convert_pt2e(module) + + self.assertIs(result, converted_module) + convert_pt2e.assert_called_once_with(module) + + +if __name__ == "__main__": + unittest.main() diff --git a/backends/qualcomm/genai_pipeline/tests/strategies/quantization/test_mllm_quantizer_adapter.py b/backends/qualcomm/genai_pipeline/tests/strategies/quantization/test_mllm_quantizer_adapter.py new file mode 100644 index 00000000000..a4b1148388d --- /dev/null +++ b/backends/qualcomm/genai_pipeline/tests/strategies/quantization/test_mllm_quantizer_adapter.py @@ -0,0 +1,187 @@ +# 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 itertools +import sys +import types +import unittest +from unittest.mock import MagicMock, patch + +from executorch.backends.qualcomm.genai_pipeline.artifact_keys import ( + ARTIFACT_AUDIO_ENCODER, + ARTIFACT_TEXT_DECODER, + ARTIFACT_TOK_EMBEDDING, + ARTIFACT_VISION_ENCODER, +) +from executorch.backends.qualcomm.genai_pipeline.strategies.quantization.mllm_quantizer_adapter import ( + MLLMQuantizerAdapter, +) + + +def _fake_quantize_pt2e(**functions): + module = types.ModuleType("torchao.quantization.pt2e.quantize_pt2e") + for name, function in functions.items(): + setattr(module, name, function) + return module + + +def _fake_llama_utils(): + module = types.ModuleType("executorch.examples.qualcomm.oss_scripts.llama.utils") + module.safe_dataloader_iter = lambda dataloader: ( + itertools.chain(dataloader, itertools.repeat([])) + if dataloader is not None + else itertools.repeat([]) + ) + return module + + +class TestMLLMQuantizerAdapter(unittest.TestCase): + def setUp(self): + self.adapter = MLLMQuantizerAdapter() + + def test_export_model_returns_exported_module(self): + module = MagicMock(name="encoder") + example_inputs = (MagicMock(name="inputs"),) + exported_module = MagicMock(name="exported_module") + exported_program = MagicMock(name="exported_program") + exported_program.module.return_value = exported_module + + with patch("torch.export.export", return_value=exported_program) as export: + result = self.adapter.export_model(module, example_inputs) + + self.assertIs(result, exported_module) + export.assert_called_once_with(module, example_inputs, strict=True) + + def test_prepare_pt2e_delegates_to_torchao(self): + module = MagicMock(name="exported_module") + quantizer = MagicMock(name="quantizer") + prepared_module = MagicMock(name="prepared_module") + prepare_pt2e = MagicMock(return_value=prepared_module) + + with patch.dict( + sys.modules, + { + "torchao.quantization.pt2e.quantize_pt2e": _fake_quantize_pt2e( + prepare_pt2e=prepare_pt2e + ) + }, + ): + result = self.adapter.prepare_pt2e(module, quantizer) + + self.assertIs(result, prepared_module) + prepare_pt2e.assert_called_once_with(module, quantizer) + + def test_init_encodings_runs_one_forward(self): + module = MagicMock(name="prepared_encoder") + example_inputs = (MagicMock(name="inputs"),) + + result = self.adapter.init_encodings(module, example_inputs) + + self.assertIs(result, module) + module.assert_called_once_with(*example_inputs) + + def test_calibrate_prefers_audio_encoder_and_inputs(self): + audio_encoder = MagicMock(name="audio_encoder") + vision_encoder = MagicMock(name="vision_encoder") + tok_embedding = MagicMock(name="tok_embedding") + text_decoder = MagicMock(name="text_decoder") + inference = MagicMock(name="inference") + audio_inputs = MagicMock(name="audio_inputs") + vision_inputs = MagicMock(name="vision_inputs") + text_batch = { + "input_ids": MagicMock(name="input_ids"), + "attention_mask": MagicMock(name="attention_mask"), + } + + with patch.dict( + sys.modules, + { + "executorch.examples.qualcomm.oss_scripts.llama.utils": _fake_llama_utils() + }, + ): + self.adapter.calibrate( + { + ARTIFACT_AUDIO_ENCODER: audio_encoder, + ARTIFACT_VISION_ENCODER: vision_encoder, + ARTIFACT_TOK_EMBEDDING: tok_embedding, + ARTIFACT_TEXT_DECODER: text_decoder, + }, + { + ARTIFACT_AUDIO_ENCODER: [{"inputs": audio_inputs}], + ARTIFACT_VISION_ENCODER: [{"inputs": vision_inputs}], + ARTIFACT_TEXT_DECODER: [text_batch], + }, + inference=inference, + ) + + inference.predict_step.assert_called_once_with( + text_decoder, + input_ids=text_batch["input_ids"], + attn_mask=text_batch["attention_mask"], + tok_embedding=tok_embedding, + encoder_module=audio_encoder, + encoder_inputs=audio_inputs, + ) + + def test_calibrate_uses_vision_when_audio_is_absent(self): + vision_encoder = MagicMock(name="vision_encoder") + text_decoder = MagicMock(name="text_decoder") + inference = MagicMock(name="inference") + vision_inputs = MagicMock(name="vision_inputs") + text_batch = { + "input_ids": MagicMock(name="input_ids"), + "attention_mask": MagicMock(name="attention_mask"), + } + + with patch.dict( + sys.modules, + { + "executorch.examples.qualcomm.oss_scripts.llama.utils": _fake_llama_utils() + }, + ): + self.adapter.calibrate( + { + ARTIFACT_VISION_ENCODER: vision_encoder, + ARTIFACT_TEXT_DECODER: text_decoder, + }, + { + ARTIFACT_AUDIO_ENCODER: None, + ARTIFACT_VISION_ENCODER: [{"inputs": vision_inputs}], + ARTIFACT_TEXT_DECODER: [text_batch], + }, + inference=inference, + ) + + inference.predict_step.assert_called_once_with( + text_decoder, + input_ids=text_batch["input_ids"], + attn_mask=text_batch["attention_mask"], + tok_embedding=None, + encoder_module=vision_encoder, + encoder_inputs=vision_inputs, + ) + + def test_convert_pt2e_delegates_to_torchao(self): + module = MagicMock(name="prepared_module") + converted_module = MagicMock(name="converted_module") + convert_pt2e = MagicMock(return_value=converted_module) + + with patch.dict( + sys.modules, + { + "torchao.quantization.pt2e.quantize_pt2e": _fake_quantize_pt2e( + convert_pt2e=convert_pt2e + ) + }, + ): + result = self.adapter.convert_pt2e(module) + + self.assertIs(result, converted_module) + convert_pt2e.assert_called_once_with(module) + + +if __name__ == "__main__": + unittest.main() diff --git a/backends/qualcomm/genai_pipeline/tests/test_genai_pipeline.py b/backends/qualcomm/genai_pipeline/tests/test_genai_pipeline.py index 34045047734..9e889e68d3b 100644 --- a/backends/qualcomm/genai_pipeline/tests/test_genai_pipeline.py +++ b/backends/qualcomm/genai_pipeline/tests/test_genai_pipeline.py @@ -13,6 +13,9 @@ from executorch.backends.qualcomm.genai_pipeline.configs.inference_output_config import ( InferenceOutputConfig, ) +from executorch.backends.qualcomm.genai_pipeline.configs.model_preparation_output_config import ( + ModelPreparationOutputConfig, +) from executorch.backends.qualcomm.genai_pipeline.configs.quantization_output_config import ( QuantizationOutputConfig, ) @@ -31,6 +34,9 @@ from executorch.backends.qualcomm.genai_pipeline.stages.inference_stage import ( InferenceStage, ) +from executorch.backends.qualcomm.genai_pipeline.stages.model_preparation_stage import ( + ModelPreparationStage, +) from executorch.backends.qualcomm.genai_pipeline.stages.quantization_stage import ( QuantizationStage, ) @@ -40,6 +46,9 @@ from executorch.backends.qualcomm.genai_pipeline.strategies.inference.inference_strategy import ( InferenceStrategy, ) +from executorch.backends.qualcomm.genai_pipeline.strategies.model_preparation.model_preparation_strategy import ( + ModelPreparationStrategy, +) from executorch.backends.qualcomm.genai_pipeline.strategies.quantization.quantization_strategy import ( QuantizationStrategy, ) @@ -55,7 +64,7 @@ class _MockQuantizationStrategy(QuantizationStrategy): def invoke(self, context, input_config): - return QuantizationOutputConfig(quantized_model="mock_quantized_model") + return QuantizationOutputConfig(graphs={"mock": "graphs"}) class _MockCompilationStrategy(CompilationStrategy): @@ -215,9 +224,7 @@ def test_invoke_no_stages(self): def test_quantization_receives_soc_model(self): mock_quant = MagicMock(spec=QuantizationStrategy) - mock_quant.invoke.return_value = QuantizationOutputConfig( - quantized_model="quantized" - ) + mock_quant.invoke.return_value = QuantizationOutputConfig(graphs={}) test_soc = "SM8650" proxy = EngineProxy( @@ -237,6 +244,58 @@ def test_quantization_receives_soc_model(self): input_config = args[1] self.assertEqual(input_config.soc_model, test_soc) + def test_model_preparation_receives_full_extra_options(self): + mock_model_prep = MagicMock(spec=ModelPreparationStrategy) + mock_model_prep.invoke.return_value = ModelPreparationOutputConfig() + extra_options = { + "model_options": {"model_arch": MagicMock(name="model_arch")}, + "quantize_options": {"quant_dtype": MagicMock(name="quant_dtype")}, + "dataset_options": MagicMock(name="dataset_options"), + } + proxy = EngineProxy( + {STAGE_MODEL_PREPARATION: EngineType.EXECUTORCH}, + backend_type=TEST_MOCK_BACKEND_TYPE, + ) + pipeline = GenAIPipeline( + model_preparation_stage=ModelPreparationStage(mock_model_prep), + quantization_stage=None, + compilation_stage=None, + inference_stage=None, + engine_proxy=proxy, + ) + + pipeline.invoke(make_test_context(extra_options=extra_options)) + + args, _ = mock_model_prep.invoke.call_args + input_config = args[1] + self.assertEqual(input_config.extra_options, extra_options) + + def test_quantization_receives_full_extra_options(self): + mock_quant = MagicMock(spec=QuantizationStrategy) + mock_quant.invoke.return_value = QuantizationOutputConfig(graphs={}) + extra_options = { + "model_options": {"model_arch": MagicMock(name="model_arch")}, + "quantize_options": {"quant_dtype": MagicMock(name="quant_dtype")}, + "dataset_options": MagicMock(name="dataset_options"), + } + proxy = EngineProxy( + {STAGE_QUANTIZATION: EngineType.EXECUTORCH}, + backend_type=TEST_MOCK_BACKEND_TYPE, + ) + pipeline = GenAIPipeline( + model_preparation_stage=None, + quantization_stage=QuantizationStage(mock_quant), + compilation_stage=None, + inference_stage=None, + engine_proxy=proxy, + ) + + pipeline.invoke(make_test_context(extra_options=extra_options)) + + args, _ = mock_quant.invoke.call_args + input_config = args[1] + self.assertEqual(input_config.extra_options, extra_options) + def test_compilation_receives_backend_type(self): mock_compile = MagicMock(spec=CompilationStrategy) mock_compile.invoke.return_value = CompilationOutputConfig( From 44676cf54e542d3229f219c643be438605a39fb4 Mon Sep 17 00:00:00 2001 From: yuyazhua Date: Tue, 6 Oct 2026 15:57:01 +0800 Subject: [PATCH 2/3] Fix quantization and calibration behavior Require explicit calibration sources in the CLI and dataset factory, normalize model names before tokenizer setup, and validate the target SoC before model preparation. Preserve adapter-agnostic calibration strategy behavior while forwarding resolved calibration options to adapters. Clarify quantization metadata and add focused regression coverage for model lookup and graph quantization behavior. Co-Authored-By: Claude --- backends/qualcomm/genai_pipeline/cli.py | 204 ++++-------------- .../model_preparation_output_config.py | 6 +- .../configs/quantization_input_config.py | 3 - .../llm_calibration_data_adapter.py | 3 +- .../mllm_calibration_data_adapter.py | 3 +- .../genai_pipeline/datasets/dataset_lookup.py | 25 +-- .../qualcomm/genai_pipeline/genai_pipeline.py | 9 - .../qualcomm/genai_pipeline/graph_bundle.py | 66 +++--- .../model_components/decoder/__init__.py | 5 - .../qualcomm/genai_pipeline/model_lookup.py | 90 ++++---- .../genai_pipeline/quant_utilities.py | 30 +-- .../executorch_model_preparation_strategy.py | 2 +- .../model_preparation/llm_loader_adapter.py | 26 +-- .../model_preparation/mllm_loader_adapter.py | 14 +- .../executorch_quantization_strategy.py | 34 +-- ...t_executorch_model_preparation_strategy.py | 20 +- .../test_llm_loader_adapter.py | 17 +- .../test_mllm_loader_adapter.py | 17 +- .../test_executorch_quantization_strategy.py | 63 ++++-- .../tests/test_genai_pipeline.py | 4 +- .../genai_pipeline/tests/test_graph_bundle.py | 18 +- .../genai_pipeline/tests/test_model_lookup.py | 42 ++++ 22 files changed, 301 insertions(+), 400 deletions(-) create mode 100644 backends/qualcomm/genai_pipeline/tests/test_model_lookup.py diff --git a/backends/qualcomm/genai_pipeline/cli.py b/backends/qualcomm/genai_pipeline/cli.py index 5a050df53d8..64c8d1fb000 100644 --- a/backends/qualcomm/genai_pipeline/cli.py +++ b/backends/qualcomm/genai_pipeline/cli.py @@ -37,24 +37,17 @@ from typing import Any, Dict, List, Optional from executorch.backends.qualcomm.genai_pipeline.control_args import ( - DEFAULT_ARTIFACT_DIR, + DEFAULT_ARTIFACT, DEFAULT_BACKEND, DEFAULT_BATCH_SIZE, DEFAULT_CALIB_HF_LIMIT, DEFAULT_CALIB_LIMIT, DEFAULT_DTYPE_OVERRIDE, - DEFAULT_EVAL_LIMIT, - DEFAULT_EVAL_METHOD, DEFAULT_GCAP, DEFAULT_MAX_SEQ_LEN, DEFAULT_MODEL_MODE, DEFAULT_NGRAM, DEFAULT_PREFILL_AR_LEN, - DEFAULT_SOC_MODEL, - DEFAULT_TEMPERATURE, - DEFAULT_TRAIN_HF_LIMIT, - DEFAULT_TRAIN_LIMIT, - DEFAULT_TRAIN_VAL_RATIO, DEFAULT_WINDOW, ) from executorch.backends.qualcomm.genai_pipeline.genai_pipeline import QuantizationStage @@ -124,8 +117,8 @@ def build_parser() -> argparse.ArgumentParser: parser.add_argument( "--artifact-dir", type=str, - default=DEFAULT_ARTIFACT_DIR, - help=f"Directory for compiled artifacts (default: {DEFAULT_ARTIFACT_DIR}).", + default=DEFAULT_ARTIFACT, + help=f"Directory for compiled artifacts (default: {DEFAULT_ARTIFACT}).", ) # --- Quantization --- @@ -134,6 +127,11 @@ def build_parser() -> argparse.ArgumentParser: action="store_true", help="Use FP16 precision (skip quantization).", ) + parser.add_argument( + "--qat", + action="store_true", + help="Enable Quantization-Aware Training (QAT). If not set, defaults to PTQ.", + ) parser.add_argument( "--embedding-quantize", type=str, @@ -141,11 +139,6 @@ def build_parser() -> argparse.ArgumentParser: help="Fall back to the CPU embedding operator and quantize it, as " "',' -- e.g. '4,32'.", ) - parser.add_argument( - "--quant-recipe-suggestion", - action="store_true", - help="Emit a per-layer mixed-precision recipe suggestion during PTQ.", - ) parser.add_argument( "--batch-size", type=int, @@ -201,101 +194,12 @@ def build_parser() -> argparse.ArgumentParser: f"(default: {DEFAULT_CALIB_HF_LIMIT}).", ) - # --- Quantization-aware training --- - parser.add_argument( - "--qat", - action="store_true", - help="Enable Quantization-Aware Training (QAT). If not set, defaults to PTQ.", - ) - parser.add_argument( - "--train-config", - type=str, - default=None, - help="(QAT) YAML file overriding training configuration defaults.", - ) - parser.add_argument( - "--lr-config", - type=str, - default=None, - help="(QAT) YAML file configuring optimizer parameter groups.", - ) - parser.add_argument( - "--train-tasks", - type=str, - nargs="+", - default=None, - help="(QAT) lm-eval tasks for training data. Use --calib-tasks to " - "specify calibration data separately.", - ) - parser.add_argument( - "--train-limit", - type=int, - default=DEFAULT_TRAIN_LIMIT, - help=f"(QAT) Number of samples for train tasks (default: {DEFAULT_TRAIN_LIMIT}).", - ) - parser.add_argument( - "--train-hf-dataset", - type=str, - default=None, - help="(QAT) HuggingFace instruct dataset for training " - "(e.g. 'HuggingFaceTB/smol-smoltalk').", - ) - parser.add_argument( - "--train-hf-limit", - type=int, - default=DEFAULT_TRAIN_HF_LIMIT, - help="(QAT) Number of samples to load from --train-hf-dataset " - f"(default: {DEFAULT_TRAIN_HF_LIMIT}).", - ) - parser.add_argument( - "--train-val-ratio", - type=float, - default=DEFAULT_TRAIN_VAL_RATIO, - help="(QAT) Fraction of non-calib samples used for training; the " - "remainder becomes validation. 1.0 disables validation " - f"(default: {DEFAULT_TRAIN_VAL_RATIO}).", - ) - parser.add_argument( - "--freeze-all-params", - action="store_true", - help="(QAT) Freeze model weights so only quantization parameters are updated.", - ) - - # --- Evaluation --- - parser.add_argument( - "--eval-tasks", - type=str, - nargs="+", - default=None, - help="lm-eval tasks to evaluate on, e.g. --eval-tasks wikitext.", - ) - parser.add_argument( - "--eval-methods", - type=str, - nargs="+", - default=[DEFAULT_EVAL_METHOD], - help=f"Evaluation methods to run (default: {DEFAULT_EVAL_METHOD}).", - ) - parser.add_argument( - "--eval-limit", - type=int, - default=DEFAULT_EVAL_LIMIT, - help=f"How many samples to evaluate on (default: {DEFAULT_EVAL_LIMIT}).", - ) - parser.add_argument( - "--eval-num-fewshot", - type=int, - default=None, - metavar="N", - help="Number of few-shot examples in each evaluation sample.", - ) - # --- Backend and SoC selection --- parser.add_argument( "--soc", type=str, - default=DEFAULT_SOC_MODEL, - help=f"Target SoC model (default: {DEFAULT_SOC_MODEL}).", + default=None, + help="Target SoC model.", ) parser.add_argument( "--backend", @@ -305,7 +209,7 @@ def build_parser() -> argparse.ArgumentParser: help=f"QNN backend type (default: {DEFAULT_BACKEND}).", ) - # --- Runtime prompts and multimodal inputs --- + # --- Runtime prompt --- parser.add_argument( "--prompt", type=str, @@ -313,32 +217,6 @@ def build_parser() -> argparse.ArgumentParser: default=["Hello, how are you?"], help="User prompt(s) for text generation.", ) - parser.add_argument( - "--system-prompt", - type=str, - default="", - help="System prompt for models that support one.", - ) - parser.add_argument( - "--temperature", - type=float, - default=DEFAULT_TEMPERATURE, - help=f"Sampling temperature for text generation (default: {DEFAULT_TEMPERATURE}).", - ) - parser.add_argument( - "--audio-path", - type=str, - nargs="+", - default=[], - help="Audio file(s) for multimodal (ALM) models.", - ) - parser.add_argument( - "--image-path", - type=str, - nargs="+", - default=[], - help="Image file(s) for multimodal (VLM) models.", - ) # Graph shapes and modes parser.add_argument( @@ -357,8 +235,11 @@ def build_parser() -> argparse.ArgumentParser: parser.add_argument( "--max-context-len", type=int, - default=DEFAULT_MAX_SEQ_LEN, # TODO: Add DEFAULT_MAX_CONTEXT_LEN once attention sink is introduced in GenAI Pipeline. - help=f"Maximum context length (default: {DEFAULT_MAX_SEQ_LEN}).", + default=None, + help=( + "Maximum context length. Overrides --max-seq-len until GenAI Pipeline " + f"supports attention sink (default: {DEFAULT_MAX_SEQ_LEN})." + ), ) parser.add_argument( "--prefill-ar-len", @@ -401,12 +282,6 @@ def build_parser() -> argparse.ArgumentParser: action="store_true", help="Only compile the model (skip inference).", ) - parser.add_argument( - "--pre-gen-pte", - type=str, - default=None, - help="Directory containing pre-generated .pte artifacts.", - ) parser.add_argument( "-v", "--verbose", @@ -448,6 +323,9 @@ def _validate_args(args: argparse.Namespace) -> None: Raises: ValueError: If arguments violate constraints. """ + if not args.soc: + raise ValueError("--soc is required to run the GenAI Pipeline.") + # TODO: The legacy path already supports the QAT feature, # but it is not yet introduced in GenAI Pipeline. if args.qat: @@ -477,6 +355,14 @@ def _validate_args(args: argparse.Namespace) -> None: ) args.batch_size = 1 + if not args.use_fp16 and not ( + args.calib_tasks or args.calib_samples or args.calib_hf_dataset + ): + raise ValueError( + "Quantization requires calibration data. Provide --calib-tasks, " + "--calib-samples, or --calib-hf-dataset." + ) + def _prepare_configs(args: argparse.Namespace) -> dict: """Prepare model configs, transforms, and dataset options. @@ -501,26 +387,22 @@ def _prepare_configs(args: argparse.Namespace) -> dict: get_state_dict_loader, ) - # TODO: The legacy path already supports attention sink, but GenAI Pipeline - # does not yet. Until attention sink is supported, keep max_context_len - # equal to max_seq_len. - logger.info( - "Setting max_context_len=%s to match max_seq_len because " - "attention sink is not yet supported in GenAI Pipeline.", - args.max_seq_len, - ) - args.max_context_len = args.max_seq_len + # TODO: Separate sequence and context lengths once attention sink is + # supported in GenAI Pipeline. Until then, use one resolved graph shape. + if args.max_context_len is None: + args.max_context_len = args.max_seq_len + else: + args.max_seq_len = args.max_context_len model_config = get_model_config(args.model) logger.info("Model config loaded") - model_arch = get_model_arch(args.model, args, model_config=model_config) + model_arch = get_model_arch(args.model, args) logger.info("Model arch loaded") weight_transforms, module_transforms = get_source_transform( args.model, control_args=args, - model_config=model_config, ) logger.info("Loaded source transforms for model '%s'", args.model) @@ -528,7 +410,6 @@ def _prepare_configs(args: argparse.Namespace) -> dict: state_dict_loaders = get_state_dict_loader( args.model, control_args=args, - model_config=model_config, ) model_options = { "model_arch": model_arch, @@ -578,8 +459,6 @@ def get_quantization_stage( ) -> "QuantizationStage": from executorch.backends.qualcomm.genai_pipeline.datasets import ( get_calibration_dataset_adapter, - get_eval_dataset_adapter, - get_training_dataset_adapter, ) from executorch.backends.qualcomm.genai_pipeline.model_lookup import ( get_quantizer_adapter, @@ -594,13 +473,6 @@ def get_quantization_stage( calibration_data_adapter = get_calibration_dataset_adapter( dataset_options, is_multimodal=model_is_multimodal ) - training_data_adapter = get_training_dataset_adapter( - dataset_options, is_multimodal=model_is_multimodal - ) - evaluation_data_adapter = get_eval_dataset_adapter( - dataset_options, is_multimodal=model_is_multimodal - ) - quantizer_adapter = get_quantizer_adapter(args.model) recipe_names = { component: ( @@ -617,8 +489,6 @@ def get_quantization_stage( ExecuTorchQuantizationStrategy( quantizer_adapter=quantizer_adapter, calibration_data_adapter=calibration_data_adapter, - training_data_adapter=training_data_adapter, - evaluation_data_adapter=evaluation_data_adapter, ) ) return quantization_stage @@ -656,6 +526,7 @@ def run_pipeline(args: argparse.Namespace) -> None: extra_options = { "backend": args.backend, "model_mode": args.model_mode, + "max_context_len": args.max_context_len, "max_seq_len": args.max_seq_len, "use_fp16": args.use_fp16, "compile_only": args.compile_only, @@ -711,7 +582,9 @@ def run_pipeline(args: argparse.Namespace) -> None: # Currently, only model preparation and quantization are supported. # TODO: Add support for compilation and inference. if args.compile_only: - logger.info("Compile-only mode (no device inference)") + logger.info( + "--compile-only is reserved until compilation and inference stages are added" + ) def main(argv: Optional[List[str]] = None) -> None: @@ -738,6 +611,7 @@ def main(argv: Optional[List[str]] = None) -> None: # Validate --model is required if not args.model: parser.error("--model is required (use --list-models to see options)") + args.model = args.model.lower() # Run pipeline try: 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 36ccd1ad3aa..6e9be8b792d 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, Dict, Optional, Tuple, TYPE_CHECKING +from typing import Any, Callable, Dict, Optional, Tuple, TYPE_CHECKING if TYPE_CHECKING: from torch import nn @@ -39,7 +39,7 @@ class ModelPreparationOutputConfig: 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. + chat_template: Optional tokenizer-bound callable 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``. @@ -54,7 +54,7 @@ class ModelPreparationOutputConfig: tokenizer: Any = None example_inputs: Optional[Dict[str, Dict[str, Tuple[Any, ...]]]] = None runtime_tokenizer_path: Optional[Path] = None - chat_template: Optional[str] = None + chat_template: Optional[Callable[..., str]] = None meta: Optional[Dict[str, Dict[str, Any]]] = None inference: Any = None num_shardings: Optional[Dict[str, int]] = None diff --git a/backends/qualcomm/genai_pipeline/configs/quantization_input_config.py b/backends/qualcomm/genai_pipeline/configs/quantization_input_config.py index 030c13971ac..bef8b15bd8f 100644 --- a/backends/qualcomm/genai_pipeline/configs/quantization_input_config.py +++ b/backends/qualcomm/genai_pipeline/configs/quantization_input_config.py @@ -51,9 +51,6 @@ class QuantizationInputConfig: 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, or a per-component map of them. meta: Per-graph constant metadata from model preparation. inference: Optional model-specific inference instance or callable used diff --git a/backends/qualcomm/genai_pipeline/datasets/calibration/llm_calibration_data_adapter.py b/backends/qualcomm/genai_pipeline/datasets/calibration/llm_calibration_data_adapter.py index 4eaf28d5bbc..bbb8665361a 100644 --- a/backends/qualcomm/genai_pipeline/datasets/calibration/llm_calibration_data_adapter.py +++ b/backends/qualcomm/genai_pipeline/datasets/calibration/llm_calibration_data_adapter.py @@ -82,8 +82,7 @@ def generate_calibration_data( if not self._dataset_loaders: raise ValueError( - "LLMCalibrationDataAdapter requires at least one dataset loader; " - "use DefaultCalibrationDataAdapter for the random fallback." + "LLMCalibrationDataAdapter requires at least one dataset loader." ) if self._collector is None: raise ValueError("collector is required when dataset_loaders are provided") diff --git a/backends/qualcomm/genai_pipeline/datasets/calibration/mllm_calibration_data_adapter.py b/backends/qualcomm/genai_pipeline/datasets/calibration/mllm_calibration_data_adapter.py index b25453b7379..c4a97246d92 100644 --- a/backends/qualcomm/genai_pipeline/datasets/calibration/mllm_calibration_data_adapter.py +++ b/backends/qualcomm/genai_pipeline/datasets/calibration/mllm_calibration_data_adapter.py @@ -85,8 +85,7 @@ def generate_calibration_data( if not self._dataset_loaders: raise ValueError( - "MLLMCalibrationDataAdapter requires at least one dataset loader; " - "use DefaultCalibrationDataAdapter for the random fallback." + "MLLMCalibrationDataAdapter requires at least one dataset loader." ) if self._collector is None: raise ValueError("collector is required when dataset_loaders are provided") diff --git a/backends/qualcomm/genai_pipeline/datasets/dataset_lookup.py b/backends/qualcomm/genai_pipeline/datasets/dataset_lookup.py index 0a2c70eccba..b4db29ef780 100644 --- a/backends/qualcomm/genai_pipeline/datasets/dataset_lookup.py +++ b/backends/qualcomm/genai_pipeline/datasets/dataset_lookup.py @@ -23,6 +23,12 @@ def get_dataset_adapter(dataset_options, is_multimodal=False): loaders = [] if is_multimodal: + if dataset_options.calib_tasks or dataset_options.calib_hf_dataset: + raise ValueError( + "MLLM calibration currently supports --calib-samples only; " + "--calib-tasks and --calib-hf-dataset are unsupported." + ) + from executorch.backends.qualcomm.genai_pipeline.datasets.loaders.mllm.message_sample_adapter import ( MessageSampleAdapter as MLLMMessageSampleAdapter, ) @@ -83,28 +89,15 @@ def get_collector(is_multimodal=False): def get_calibration_dataset_adapter(dataset_options, is_multimodal=False): - """Assemble the calibration purpose adapter. + """Assemble a corpus-backed calibration purpose adapter. - Selects source loaders from ``dataset_options``. With no source, returns - the modality-agnostic ``DefaultCalibrationDataAdapter`` (random pipeline - sanity data, sized from the options). With sources, wires them and the + Selects source loaders from ``dataset_options`` and wires them with the modality's collector through :class:`CalibrationDatasetBuilder`. """ loaders = get_dataset_adapter(dataset_options, is_multimodal) if not loaders: - from executorch.backends.qualcomm.genai_pipeline.datasets.calibration.default_calibration_data_adapter import ( - DefaultCalibrationDataAdapter, - ) - - kwargs = {} - if dataset_options is not None: - if dataset_options.num_samples is not None: - kwargs["num_samples"] = dataset_options.num_samples - kwargs["batch_size"] = dataset_options.batch_size - if dataset_options.seed is not None: - kwargs["seed"] = dataset_options.seed - return DefaultCalibrationDataAdapter(**kwargs) + raise ValueError("No calibration dataset sources were configured.") builder = CalibrationDatasetBuilder(is_multimodal) for loader in loaders: diff --git a/backends/qualcomm/genai_pipeline/genai_pipeline.py b/backends/qualcomm/genai_pipeline/genai_pipeline.py index 7bc8160592e..9e901e70b32 100644 --- a/backends/qualcomm/genai_pipeline/genai_pipeline.py +++ b/backends/qualcomm/genai_pipeline/genai_pipeline.py @@ -96,15 +96,6 @@ class GenAIPipeline: Assembles stages from EngineProxy, wires data flow between InputConfig → Stage → OutputConfig, and executes sequentially: model_preparation → quantization → compilation → inference. - - One decoder is exported several times, once per Mode, so the orchestrator - plans ``{graph_name: Mode}`` once per run and hands it to the stages that - fan out over it. Every other per-graph fact is derived from a graph's Mode - where it is needed, rather than carried between stages. - - Calibration data is built inside the quantization stage from its injected - dataset adapters, so there is no separate dataset stage: a run that skips - quantization never touches a corpus. """ def __init__( diff --git a/backends/qualcomm/genai_pipeline/graph_bundle.py b/backends/qualcomm/genai_pipeline/graph_bundle.py index 08437a3c9e7..bf5bdb642e1 100644 --- a/backends/qualcomm/genai_pipeline/graph_bundle.py +++ b/backends/qualcomm/genai_pipeline/graph_bundle.py @@ -20,10 +20,12 @@ from dataclasses import dataclass, field from typing import Any, Dict, Optional, Tuple -#: The exact key set ``GraphBundle.quant_io_dtypes`` carries when it is not -#: ``None``. The quantized-IO tagger indexes both, so a mapping missing either -#: one is rejected at construction rather than at lowering. -_QUANT_IO_DTYPE_KEYS = frozenset({"kv_type", "io_type"}) +#: The supported key sets ``GraphBundle.quant_io_dtypes`` carries when it is +#: not ``None``. Graphs without KV caches only carry ``io_type``. +_QUANT_IO_DTYPE_KEY_SETS = ( + frozenset({"kv_type", "io_type"}), + frozenset({"io_type"}), +) @dataclass(frozen=True) @@ -67,11 +69,18 @@ class GraphBundle: attributes written during ``convert_pt2e``. Becomes the ``.pte``'s constant methods, so it is complete only **after** quantization has run. - quant_io_dtypes: The graph-boundary dtypes quantization chose, as - ``{"kv_type": torch.dtype, "io_type": torch.dtype}``. Compilation - builds the ``TagQuantIO`` pass settings from these; it cannot derive - them, since they come from the quantization recipe's KV and logits - bit widths. + quant_io_dtypes: The graph-boundary dtypes quantization chose. + Compilation builds the ``TagQuantIO`` pass settings from these; it + cannot derive them, since they come from the quantization recipe's + KV and logits bit widths. + + A non-``None`` mapping has one of these shapes: + + * ``{"kv_type": torch.dtype, "io_type": torch.dtype}`` for a + graph with a KV cache, such as a decoder. + * ``{"io_type": torch.dtype}`` for a graph without a KV cache, + such as an encoder. Encoders derive ``io_type`` from their + quantization recipe, but have no ``kv_type`` to record. ``None`` means *this graph's boundary dtypes were not chosen by a recipe* -- which happens for two different reasons, and compilation @@ -79,24 +88,14 @@ class GraphBundle: * **Quantization was skipped** for the graph, so its IO stays float32 and ``TagQuantIO`` is left inactive. - * **The component never derives dtypes from a recipe.** Encoder - graphs are the case: the legacy flow tags an encoder's boundary - ``torch.float32`` from a literal and never consults the encoder - recipe, which exposes no KV or logits bit width at all. Such a - graph still needs ``TagQuantIO`` active when it is sharded, to - tag the ``llama.fallback.default`` boundary -- so ``None`` here - must not be read as "nothing to tag". - - The compilation stage's pass policy distinguishes the two by - component, since only the component knows whether a recipe was meant - to supply these at all. - - **Both keys or neither.** The tagger indexes both unconditionally, - once per node, so a mapping carrying only one of them is not a - partially-quantized boundary -- it is a ``KeyError`` during - lowering. A recipe whose bit width maps to no fixed-point dtype - makes the whole mapping ``None``; it does not drop a single key. - ``__post_init__`` rejects anything else. + * **No fixed-point boundary dtype was selected.** For example, a + recipe bit width may not map to a fixed-point dtype. The whole + mapping is then ``None`` rather than a mapping with one missing + key. + + ``io_type`` is required whenever the mapping is present. + ``kv_type`` is required only for graphs with a KV cache. + ``__post_init__`` rejects every other key combination. modality_inputs: Encoder inputs for a multimodal model, keyed by modality. ``None`` for a text-only model. executorch_config: Optional override for the ``to_executorch`` @@ -114,18 +113,21 @@ class GraphBundle: executorch_config: Optional[Any] = None def __post_init__(self) -> None: - """Reject a ``quant_io_dtypes`` mapping that is not both keys. + """Reject an unsupported ``quant_io_dtypes`` key combination. Raises: ValueError: If ``quant_io_dtypes`` is a mapping whose keys are not - exactly ``kv_type`` and ``io_type``. + ``io_type`` alone or ``kv_type`` and ``io_type`` together. """ if self.quant_io_dtypes is None: return - keys = set(self.quant_io_dtypes) - if keys != _QUANT_IO_DTYPE_KEYS: + keys = frozenset(self.quant_io_dtypes) + if keys not in _QUANT_IO_DTYPE_KEY_SETS: + valid_key_sets = " or ".join( + str(sorted(key_set)) for key_set in _QUANT_IO_DTYPE_KEY_SETS + ) raise ValueError( "quant_io_dtypes must carry exactly " - f"{sorted(_QUANT_IO_DTYPE_KEYS)} or be None; got {sorted(keys)}" + f"{valid_key_sets} or be None; got {sorted(keys)}" ) diff --git a/backends/qualcomm/genai_pipeline/model_components/decoder/__init__.py b/backends/qualcomm/genai_pipeline/model_components/decoder/__init__.py index ce3226107ed..96d9722adaf 100644 --- a/backends/qualcomm/genai_pipeline/model_components/decoder/__init__.py +++ b/backends/qualcomm/genai_pipeline/model_components/decoder/__init__.py @@ -55,11 +55,6 @@ "LlamaModel", "LlamaModelWithoutEmbedding", "MultiScopeAwareLlamaModel", - # Attention masks. - "AttentionMask", - "BaseAttentionMask", - "CausalAttentionMask", - "SlidingWindowAttentionMask", # Attention, KV-head expansion, and rotary positional encoding. "AttentionSinkRope", "LlamaAttention", diff --git a/backends/qualcomm/genai_pipeline/model_lookup.py b/backends/qualcomm/genai_pipeline/model_lookup.py index 09cfd23c68b..05a00f493ea 100644 --- a/backends/qualcomm/genai_pipeline/model_lookup.py +++ b/backends/qualcomm/genai_pipeline/model_lookup.py @@ -10,9 +10,9 @@ configuration. It keeps user-facing model names separate from the concrete objects and options required by each pipeline stage. -Most callers should resolve a model config once with :func:`get_model_config` -and pass it to the more specific lookup helpers when available, so all derived -configuration comes from the same registry entry. +Each lookup helper resolves its model configuration through +:func:`get_model_config`, so all derived configuration comes from the same +registry entry. """ from __future__ import annotations @@ -20,7 +20,7 @@ import json import logging from functools import partial -from typing import Any, Callable, Dict, List, Optional, Tuple +from typing import Any, Callable, Dict, List, Tuple from executorch.backends.qualcomm.genai_pipeline.artifact_keys import ( ARTIFACT_AUDIO_ENCODER, @@ -33,13 +33,6 @@ GRAPH_FORWARD, TOK_EMBEDDING_GRAPH_NAMES, ) -from executorch.backends.qualcomm.genai_pipeline.model_components.decoder import ( - LlamaModel, -) -from executorch.backends.qualcomm.genai_pipeline.model_components.embedding import ( - TokenEmbedding, -) -from executorch.backends.qualcomm.genai_pipeline.models import LLM_VARIANT_ARCHS from executorch.examples.qualcomm.oss_scripts.llama.model.static_llama import ModelArgs from executorch.examples.qualcomm.oss_scripts.llama.wrappers.base_component import ( get_model_specific_kwargs, @@ -85,9 +78,7 @@ def get_model_config(model_name: str) -> Any: return config -def get_model_arch_config( - model_name: str, control_args: Any, model_config: Optional[Any] = None -) -> Any: +def get_model_arch_config(model_name: str, control_args: Any) -> Any: """Resolve the decoder's config, for reading model shapes. All graphs (calibration / decode / prefill) share one set of weights and @@ -107,8 +98,14 @@ def get_model_arch_config( ValueError: If the model needs a params file that ``control_args`` does not name. """ - if not model_config: - model_config = get_model_config(model_name) + model_name = model_name.lower() + model_config = get_model_config(model_name) + + if model_config.quant_recipe is None: + raise ValueError( + f"Model '{model_name}' has no quantization recipe; the GenAI Pipeline " + "requires one to configure the KV-cache IO bit width" + ) if model_name == "gemma4-e2b": from executorch.examples.models.gemma4.text_decoder.gemma4_config import ( @@ -148,7 +145,7 @@ def get_model_arch_config( def get_model_arch( - model_name: str, control_args: Any, model_config: Optional[Any] = None + model_name: str, control_args: Any ) -> Dict[str, Dict[str, Callable]]: """Build the model constructors exported for a model. @@ -162,15 +159,21 @@ def get_model_arch( Args: model_name: The model key in ``SUPPORTED_LLM_MODELS``. control_args: Runtime arguments controlling graph modes, shapes. - model_config: Optional pre-resolved model configuration. Returns: A ``{component_name: {graph_name: constructor}}`` mapping. """ - # Get model config - if not model_config: - model_config = get_model_config(model_name) + from executorch.backends.qualcomm.genai_pipeline.model_components.decoder import ( + LlamaModel, + ) + from executorch.backends.qualcomm.genai_pipeline.model_components.embedding import ( + TokenEmbedding, + ) + from executorch.backends.qualcomm.genai_pipeline.models import LLM_VARIANT_ARCHS + + model_name = model_name.lower() + model_config = get_model_config(model_name) # Decode and calibration graphs are always exported. Prefill is additionally # exported for hybrid and lookahead modes. @@ -194,6 +197,11 @@ def get_model_arch( use_i64_token = control_args.embedding_quantize is not None model_specific_kwargs = get_model_specific_kwargs(control_args, model_config) quant_recipe = model_config.quant_recipe + if quant_recipe is None: + raise ValueError( + f"Model '{model_name}' has no quantization recipe; the GenAI Pipeline " + "requires one to configure the KV-cache IO bit width" + ) if model_name == "gemma4-e2b": from executorch.examples.models.gemma4.text_decoder.gemma4_config import ( @@ -225,11 +233,7 @@ def get_model_arch( ar_len = 1 extra_kwargs = { - # 32 is the sentinel for "unquantized KV IO"; get_kv_io_bit_width() - # returns it too when the recipe has no default_quant_dtype. - "kv_io_bit_width": ( - quant_recipe().get_kv_io_bit_width() if quant_recipe else 32 - ), + "kv_io_bit_width": quant_recipe().get_kv_io_bit_width(), } # Get Text Decoder model architecture. @@ -307,7 +311,6 @@ def get_source_transform( model_name: str, *, control_args: Any, - model_config: Optional[Any] = None, ) -> Tuple[Dict[str, List[Callable]], Dict[str, List[Callable]]]: """Resolve a model's source transforms, mirroring ``LLMWrapper._prepare_model``. @@ -315,8 +318,6 @@ def get_source_transform( model_name: The model's key in ``SUPPORTED_LLM_MODELS``. control_args: CLI arguments used to determine the checkpoint source and resolve the values bound to individual transforms. - model_config: Optional resolved model configuration. When omitted, it - is looked up from ``model_name``. Returns: A ``(weight_transforms, module_transforms)`` tuple of component-keyed @@ -345,12 +346,10 @@ def get_source_transform( unwrap_model_key, ) - if model_config is None: - model_config = get_model_config(model_name) - config = model_config - name = model_name.lower() + model_name = model_name.lower() + config = get_model_config(model_name) is_hf_path = control_args.checkpoint is None - model_args = get_model_arch_config(model_name, control_args, config) + model_args = get_model_arch_config(model_name, control_args) if config.r1 or config.r2: raise RuntimeError( @@ -363,10 +362,10 @@ def get_source_transform( if is_hf_path: # HF path: gemma4 key rename, then the Gemma RMSNorm +1 offset, then the # embedding scale (self-guarding on the factor / key). - if name == "gemma4-e2b": + if model_name == "gemma4-e2b": weight_transforms.append(remap_gemma4_keys) else: - if name in ("gemma-2b", "gemma2-2b", "gemma3-1b"): + if model_name in ("gemma-2b", "gemma2-2b", "gemma3-1b"): weight_transforms.append(gemma_rmsnorm_offset) weight_transforms.append( partial( @@ -377,7 +376,7 @@ def get_source_transform( else: # Local checkpoint path: stories260k carries torch.compile's # ``_orig_mod.`` prefix and must be renamed. - if name == "stories260k": + if model_name == "stories260k": weight_transforms.append(strip_orig_mod_prefix) # RoPE weight layout permutation, gated on the model config flag. if config.transform_weight: @@ -444,7 +443,6 @@ def get_state_dict_loader( model_name: str, *, control_args: Any, - model_config: Optional[Any] = None, ) -> Dict[str, Callable[[str], Dict[str, Any]]]: """Return component-keyed loaders for remote Hugging Face checkpoints. @@ -452,8 +450,6 @@ def get_state_dict_loader( model_name: Registered model identifier used to select the loader. control_args: Runtime arguments. A configured local ``checkpoint`` bypasses remote loading and returns no loader. - model_config: Optional resolved model configuration. When omitted, it - is looked up from ``model_name``. Returns: A component-keyed loader map. Each loader accepts a Hugging Face @@ -465,13 +461,12 @@ def get_state_dict_loader( has model-specific conversion details bound into the callable. Raises: - KeyError: If ``model_name`` is not registered and ``model_config`` is - not supplied. + KeyError: If ``model_name`` is not registered. """ import torch - if model_config is None: - model_config = get_model_config(model_name) + model_name = model_name.lower() + model_config = get_model_config(model_name) if control_args.checkpoint is not None: return {} @@ -479,15 +474,12 @@ def get_state_dict_loader( if hasattr(convert_weights, "__func__"): convert_weights = convert_weights.__func__ - name = model_name.lower() - if name == "gemma4-e2b": + if model_name == "gemma4-e2b": return { ARTIFACT_TEXT_DECODER: partial( load_gemma4_hf_checkpoint_state_dict, convert_weights=convert_weights, - gemma4_config=get_model_arch_config( - model_name, control_args, model_config - ), + gemma4_config=get_model_arch_config(model_name, control_args), dtype=torch.float32, ) } diff --git a/backends/qualcomm/genai_pipeline/quant_utilities.py b/backends/qualcomm/genai_pipeline/quant_utilities.py index d81ceff016e..0aa3cde74bc 100644 --- a/backends/qualcomm/genai_pipeline/quant_utilities.py +++ b/backends/qualcomm/genai_pipeline/quant_utilities.py @@ -58,11 +58,15 @@ def make_quantizer( from executorch.backends.qualcomm.export_utils import ( make_quantizer as _make_quantizer, ) + from torchao.quantization.pt2e import MinMaxObserver soc_model_str = soc_model.name if hasattr(soc_model, "name") else str(soc_model) make_quantizer_kwargs = { "backend": backend, "soc_model": soc_model_str, + "per_channel_conv": True, + "per_channel_linear": True, + "act_observer": MinMaxObserver, **kwargs, } if quant_dtype is not None: @@ -201,7 +205,11 @@ def encoding_override( # noqa: C901 ``n_cache_layers`` additionally copies KV-cache output encodings onto the deployed graph's cache inputs; ``None`` leaves KV-cache encodings unchanged. """ - from executorch.backends.qualcomm.builders.utils import is_graph_output + from executorch.backends.qualcomm.builders.utils import ( + get_attr_from_target, + is_graph_output, + set_attr_from_target, + ) pbq_target = { torch.ops.torchao.dequantize_affine, @@ -254,26 +262,10 @@ def activation_override(quantized_node, unquantized_node): activation_override(quantized_user, unquantized_user) def parameter_override(quantized_node, unquantized_node): - # Some parameters need to be iterated over to retrieve attributes such as static_llama.tok_embedding.weight - def _get_attr(graph_module: torch.fx.GraphModule, target: str) -> Any: - attr: Any = graph_module - for target_atom in target.split("."): - attr = getattr(attr, target_atom) - return attr - - def _set_attr( - graph_module: torch.fx.GraphModule, target: str, replacement: Any - ) -> Any: - attr: Any = graph_module - target_list = target.split(".") - for target_atom in target_list[:-1]: - attr = getattr(attr, target_atom) - setattr(attr, target_list[-1], replacement) - - _set_attr( + set_attr_from_target( unquantized_model, unquantized_node.target, - _get_attr(quantized_model, quantized_node.target), + get_attr_from_target(quantized_model, quantized_node.target), ) # scale / zero point are part of op's attributes if list(quantized_node.users)[0].target in ptq_target: 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 fc52f09f66e..cb10fd20740 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 @@ -208,7 +208,7 @@ def invoke( inference = self._adapter.get_inference( meta, example_inputs, - extra_options=extra, + extra_options=model_options, ) # Step 7: Optionally export tokenizer for runtime diff --git a/backends/qualcomm/genai_pipeline/strategies/model_preparation/llm_loader_adapter.py b/backends/qualcomm/genai_pipeline/strategies/model_preparation/llm_loader_adapter.py index 3ae510cb0d0..7dc59ce5400 100644 --- a/backends/qualcomm/genai_pipeline/strategies/model_preparation/llm_loader_adapter.py +++ b/backends/qualcomm/genai_pipeline/strategies/model_preparation/llm_loader_adapter.py @@ -336,29 +336,13 @@ def export_tokenizer( output_dir: Path, extra_options: Optional[Dict[str, Any]] = None, ) -> Path: - """Export tokenizer to disk and return the runtime tokenizer file. + """Return the runtime tokenizer path selected by ``TokenizerWrapper``. - ``TokenizerWrapper`` has already written the artifacts. Prefer the - runtime formats recognized by ``pytorch_tokenizers.get_tokenizer`` over - the wrapper's order-dependent fallback. - - Args: - tokenizer: The tokenizer instance to export. - output_dir: Directory to write the exported tokenizer artifacts to. - extra_options: Additional export options. - - Returns: - Path to the runtime tokenizer file (e.g. ``tokenizer.json``). - - Raises: - FileNotFoundError: If no tokenizer artifacts were written. + ``TokenizerWrapper`` creates its artifacts while loading the tokenizer + and owns the runtime file selection. ``output_dir`` is retained for + protocol compatibility and does not control the artifact location. """ - runtime_tokenizer_path = Path(tokenizer.runtime_tokenizer_path) - artifact_dir = Path(tokenizer.artifact) - artifacts = list(artifact_dir.iterdir()) if artifact_dir.is_dir() else [] - if runtime_tokenizer_path not in artifacts: - artifacts.append(runtime_tokenizer_path) - return self._select_runtime_tokenizer(artifacts) + return Path(tokenizer.runtime_tokenizer_path) @classmethod def from_model_config( diff --git a/backends/qualcomm/genai_pipeline/strategies/model_preparation/mllm_loader_adapter.py b/backends/qualcomm/genai_pipeline/strategies/model_preparation/mllm_loader_adapter.py index 90800ee463f..04ff58fccfc 100644 --- a/backends/qualcomm/genai_pipeline/strategies/model_preparation/mllm_loader_adapter.py +++ b/backends/qualcomm/genai_pipeline/strategies/model_preparation/mllm_loader_adapter.py @@ -425,13 +425,13 @@ def export_tokenizer( output_dir: Path, extra_options: Optional[Dict[str, Any]] = None, ) -> Path: - """Return the runtime tokenizer file exported by ``TokenizerWrapper``.""" - runtime_tokenizer_path = Path(tokenizer.runtime_tokenizer_path) - artifact_dir = Path(tokenizer.artifact) - artifacts = list(artifact_dir.iterdir()) if artifact_dir.is_dir() else [] - if runtime_tokenizer_path not in artifacts: - artifacts.append(runtime_tokenizer_path) - return self._select_runtime_tokenizer(artifacts) + """Return the runtime tokenizer path selected by ``TokenizerWrapper``. + + ``TokenizerWrapper`` creates its artifacts while loading the tokenizer + and owns the runtime file selection. ``output_dir`` is retained for + protocol compatibility and does not control the artifact location. + """ + return Path(tokenizer.runtime_tokenizer_path) def _build_custom_decoder_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 e226f81d546..a9462f27a0a 100644 --- a/backends/qualcomm/genai_pipeline/strategies/quantization/executorch_quantization_strategy.py +++ b/backends/qualcomm/genai_pipeline/strategies/quantization/executorch_quantization_strategy.py @@ -115,8 +115,8 @@ class ExecuTorchQuantizationStrategy(QuantizationStrategy): quantizer_adapter: Injectable adapter for single-graph quantization operations. Defaults to ``DefaultQuantizerAdapter`` if not provided. calibration_data_adapter: Injectable purpose adapter that assembles - calibration data. Defaults to ``DefaultCalibrationDataAdapter`` (random - fallback) when not provided. + calibration data. Defaults to ``DefaultCalibrationDataAdapter`` when + not provided. training_data_adapter: Injectable adapter for QAT training data. Retained for the QAT flow; the current native PTQ implementation does not consume it. @@ -222,6 +222,7 @@ def invoke( example_inputs=self._flatten_calibration_example_inputs( input_config.example_inputs ), + extra_options=dict(extra), ) # Step 5: Quantize graphs that consume calibration data. @@ -318,8 +319,8 @@ def _make_quantizer( ) else: logger.info( - "Quant recipe not set yet; created QNN quantizer for " - "'%s' graph '%s' with quant dtype %s", + "Created dtype-driven QNN quantizer for '%s' graph '%s' " + "with quant dtype %s", component, graph_name, make_quantizer_kwargs.get("quant_dtype"), @@ -462,6 +463,7 @@ def _calibrate( for graph_name, graph_module in prepared_graphs.items(): if not self._is_deploy_graph(component, graph_name): quantization_graphs[component] = graph_module + break self._adapter.calibrate( quantization_graphs, calibration_data, @@ -469,7 +471,7 @@ def _calibrate( ) def _convert_pt2e(self, prepared_modules: dict) -> dict: - """Convert prepared PT2E modules to QDQ modules and release inputs. + """Convert prepared PT2E modules to QDQ modules and release the prepared map. Args: prepared_modules: Component- and graph-keyed prepared PT2E modules. @@ -477,17 +479,18 @@ def _convert_pt2e(self, prepared_modules: dict) -> dict: Returns: Component- and graph-keyed converted QDQ modules. """ - converted_modules = {} - for component, prepared_graphs in prepared_modules.items(): - converted_modules[component] = { + converted_modules = { + component: { graph_name: self._adapter.convert_pt2e(graph_module) for graph_name, graph_module in prepared_graphs.items() } + for component, prepared_graphs in prepared_modules.items() + } prepared_modules.clear() gc.collect() return converted_modules - def _override_encodings( + def _override_encodings( # noqa: C901 self, converted_modules: dict, input_config: QuantizationInputConfig, @@ -579,16 +582,18 @@ def _override_encodings( prefill_meta = input_config.meta.get(ARTIFACT_TEXT_DECODER, {}).get( GRAPH_PREFILL_FORWARD, {} ) - encoding_override( - quantized_model=quantized_decoder, - unquantized_model=prefill, - n_cache_layers=n_cache_layers, - ) + if prefill_meta.get("get_use_kv_cache", False): + encoding_override( + quantized_model=quantized_decoder, + unquantized_model=prefill, + n_cache_layers=n_cache_layers, + ) save_logits_quant_attrs(prefill, prefill_meta) save_output_kv_cache_quant_attrs(prefill, prefill_meta) finally: # The quantization graph only sources encodings; never deployed. decoder_graphs.pop(GRAPH_FORWARD, None) + del quantized_decoder gc.collect() # Override token embedding quant encodings @@ -615,7 +620,6 @@ def _override_encodings( finally: # The quantization graph only sources encodings; never deployed. tok_embedding_graphs.pop(GRAPH_FORWARD, None) - gc.collect() def _validate_input(self, input_config: QuantizationInputConfig) -> None: """Validate required fields in the input configuration. 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 index 6a1411b7bd4..d2b7fe0f75e 100644 --- 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 @@ -336,16 +336,17 @@ def test_invoke_passes_tokenizer_options_to_load_tokenizer(self): def test_invoke_builds_inference_from_meta_and_example_inputs(self): loader = _make_mock_adapter() strategy = _make_strategy(loader) - extra = {"embedding_quantize": "4a"} + model_options = {"embedding_quantize": "4a"} result = strategy.invoke( - make_test_context(), _make_valid_input_config(extra_options=extra) + make_test_context(), + _make_valid_input_config(extra_options={"model_options": model_options}), ) loader.get_inference.assert_called_once_with( result.meta, result.example_inputs, - extra_options=extra, + extra_options=model_options, ) def test_invoke_rejects_unkeyed_module_transforms(self): @@ -384,16 +385,23 @@ def test_invoke_exports_tokenizer_when_requested(self): def test_invoke_chat_template_prefers_tokenizer_over_extra_options(self): """A chat template on the tokenizer is carried into the output config.""" + + def tokenizer_template(): + return "tokenizer_template" + + def fallback_template(): + return "fallback" + loader = _make_mock_adapter() - loader.load_tokenizer.return_value.chat_template = "tokenizer_template" + loader.load_tokenizer.return_value.chat_template = tokenizer_template strategy = _make_strategy(loader) input_config = _make_valid_input_config( - extra_options={"chat_template": "fallback"} + extra_options={"chat_template": fallback_template} ) result = strategy.invoke(make_test_context(), input_config) - self.assertEqual(result.chat_template, "tokenizer_template") + self.assertIs(result.chat_template, tokenizer_template) def test_invoke_missing_model_name_raises_stage_error(self): """StageError raised when model_name is empty.""" diff --git a/backends/qualcomm/genai_pipeline/tests/strategies/model_preparation/test_llm_loader_adapter.py b/backends/qualcomm/genai_pipeline/tests/strategies/model_preparation/test_llm_loader_adapter.py index 3b205c23c2b..115bf5e0950 100644 --- a/backends/qualcomm/genai_pipeline/tests/strategies/model_preparation/test_llm_loader_adapter.py +++ b/backends/qualcomm/genai_pipeline/tests/strategies/model_preparation/test_llm_loader_adapter.py @@ -198,7 +198,7 @@ def test_returns_tokenizer_file_not_the_directory(self): self.assertEqual(result, self.artifact_dir / TEST_TOKENIZER_JSON) self.assertNotEqual(result, self.artifact_dir) - def test_selects_by_name_not_by_position(self): + def test_returns_wrapper_runtime_path_over_other_artifacts(self): self._touch_artifacts( [TEST_TOKENIZER_CONFIG, TEST_TOKENIZER_JSON, TEST_ADDED_TOKENS] ) @@ -207,18 +207,9 @@ def test_selects_by_name_not_by_position(self): self._make_tokenizer(TEST_ADDED_TOKENS), self.output_dir ) - self.assertEqual(result, self.artifact_dir / TEST_TOKENIZER_JSON) - - def test_prefers_tokenizer_json_over_tokenizer_model(self): - self._touch_artifacts([TEST_TOKENIZER_MODEL, TEST_TOKENIZER_JSON]) - - result = self.adapter.export_tokenizer( - self._make_tokenizer(TEST_TOKENIZER_MODEL), self.output_dir - ) - - self.assertEqual(result, self.artifact_dir / TEST_TOKENIZER_JSON) + self.assertEqual(result, self.artifact_dir / TEST_ADDED_TOKENS) - def test_falls_back_to_tokenizer_model_when_no_json(self): + def test_returns_wrapper_runtime_path_for_known_name(self): self._touch_artifacts( [TEST_TOKENIZER_CONFIG, TEST_TOKENIZER_MODEL, TEST_SPECIAL_TOKENS_MAP] ) @@ -229,7 +220,7 @@ def test_falls_back_to_tokenizer_model_when_no_json(self): self.assertEqual(result, self.artifact_dir / TEST_TOKENIZER_MODEL) - def test_falls_back_to_runtime_tokenizer_when_no_known_name(self): + def test_returns_wrapper_runtime_path_when_no_known_name(self): runtime_name = "tokenizer.bin" tokenizer = self._make_tokenizer(runtime_name) diff --git a/backends/qualcomm/genai_pipeline/tests/strategies/model_preparation/test_mllm_loader_adapter.py b/backends/qualcomm/genai_pipeline/tests/strategies/model_preparation/test_mllm_loader_adapter.py index d8d6bc3a096..adab6e42935 100644 --- a/backends/qualcomm/genai_pipeline/tests/strategies/model_preparation/test_mllm_loader_adapter.py +++ b/backends/qualcomm/genai_pipeline/tests/strategies/model_preparation/test_mllm_loader_adapter.py @@ -315,7 +315,7 @@ def test_returns_tokenizer_file_not_the_directory(self): self.assertEqual(result, self.artifact_dir / TEST_TOKENIZER_JSON) self.assertNotEqual(result, self.artifact_dir) - def test_selects_by_name_not_by_position(self): + def test_returns_wrapper_runtime_path_over_other_artifacts(self): self._touch_artifacts( [TEST_TOKENIZER_CONFIG, TEST_TOKENIZER_JSON, TEST_ADDED_TOKENS] ) @@ -324,18 +324,9 @@ def test_selects_by_name_not_by_position(self): self._make_tokenizer(TEST_ADDED_TOKENS), self.output_dir ) - self.assertEqual(result, self.artifact_dir / TEST_TOKENIZER_JSON) - - def test_prefers_tokenizer_json_over_tokenizer_model(self): - self._touch_artifacts([TEST_TOKENIZER_MODEL, TEST_TOKENIZER_JSON]) - - result = self.adapter.export_tokenizer( - self._make_tokenizer(TEST_TOKENIZER_MODEL), self.output_dir - ) - - self.assertEqual(result, self.artifact_dir / TEST_TOKENIZER_JSON) + self.assertEqual(result, self.artifact_dir / TEST_ADDED_TOKENS) - def test_falls_back_to_tokenizer_model_when_no_json(self): + def test_returns_wrapper_runtime_path_for_known_name(self): self._touch_artifacts( [TEST_TOKENIZER_CONFIG, TEST_TOKENIZER_MODEL, TEST_SPECIAL_TOKENS_MAP] ) @@ -346,7 +337,7 @@ def test_falls_back_to_tokenizer_model_when_no_json(self): self.assertEqual(result, self.artifact_dir / TEST_TOKENIZER_MODEL) - def test_falls_back_to_runtime_tokenizer_when_no_known_name(self): + def test_returns_wrapper_runtime_path_when_no_known_name(self): runtime_name = "tokenizer.bin" tokenizer = self._make_tokenizer(runtime_name) 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 25a233fbfd1..be17f6ab6fb 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 @@ -45,6 +45,7 @@ DECODE_GRAPH_NAME = DECODER_GRAPH_NAMES[0] PREFILL_GRAPH_NAME = DECODER_GRAPH_NAMES[1] EMBED_DECODE_GRAPH_NAME = TOK_EMBEDDING_GRAPH_NAMES[0] +EMBED_PREFILL_GRAPH_NAME = TOK_EMBEDDING_GRAPH_NAMES[1] def _make_recipe(name, kv_bit_width=8, logits_bit_width=16): @@ -194,7 +195,7 @@ def test_is_quantization_strategy(self): self.assertIsInstance(strategy, QuantizationStrategy) def test_default_adapters_created_when_none_provided(self): - """Default quantization and purpose adapters are created when omitted.""" + """Default adapters are created when callers do not inject one.""" with patch( "executorch.backends.qualcomm.genai_pipeline.strategies.quantization." "default_quantizer_adapter.DefaultQuantizerAdapter" @@ -370,7 +371,7 @@ def test_invoke_generates_calibration_data_from_adapter(self): ][GRAPH_FORWARD], }, ) - self.assertNotIn("extra_options", kwargs) + self.assertEqual(kwargs["extra_options"], input_config.extra_options) def test_invoke_calibrates_only_non_deployed_graphs(self): """Only each component's non-deployed graph is truly calibrated.""" @@ -420,11 +421,17 @@ def test_invoke_initializes_deployed_graph_observers(self): ], ) - def test_invoke_overrides_only_text_and_embedding_deployed_graphs(self): - """Encoding override applies only where calibration/deploy graph split exists.""" + def test_invoke_skips_prefill_decoder_override_without_kv_cache(self): + """Only decoder cache reconciliation depends on prefill cache topology.""" adapter = _make_mock_adapter() strategy = _make_strategy(adapter) input_config = _make_valid_input_config() + input_config.example_inputs[ARTIFACT_TOK_EMBEDDING][ + EMBED_PREFILL_GRAPH_NAME + ] = ("embed_prefill",) + input_config.meta[ARTIFACT_TOK_EMBEDDING][EMBED_PREFILL_GRAPH_NAME] = { + "get_n_layers": 0 + } strategy.invoke(make_test_context(), input_config) @@ -433,21 +440,40 @@ def test_invoke_overrides_only_text_and_embedding_deployed_graphs(self): kwargs.get("n_cache_layers") for _, kwargs in self.encoding_override.call_args_list ] + self.assertEqual(n_cache_layers, [2, None, None]) + embedding_override_calls = self.encoding_override.call_args_list[1:] self.assertEqual( - n_cache_layers, - [ - 2, - 2, - None, - ], + embedding_override_calls[0].kwargs["unquantized_model"]._mock_name, + "converted_prepared_exported_tok_embedding_embed_decode", ) + self.assertEqual( + embedding_override_calls[1].kwargs["unquantized_model"]._mock_name, + "converted_prepared_exported_tok_embedding_embed_prefill", + ) + self.assertEqual(self.save_logits_quant_attrs.call_count, 2) + self.assertEqual(self.save_output_kv_cache_quant_attrs.call_count, 2) + + def test_invoke_overrides_prefill_with_kv_cache(self): + adapter = _make_mock_adapter() + strategy = _make_strategy(adapter) + input_config = _make_valid_input_config() + input_config.meta[ARTIFACT_TEXT_DECODER][PREFILL_GRAPH_NAME][ + "get_use_kv_cache" + ] = True + + strategy.invoke(make_test_context(), input_config) + + self.assertEqual(self.encoding_override.call_count, 3) + n_cache_layers = [ + kwargs.get("n_cache_layers") + for _, kwargs in self.encoding_override.call_args_list + ] + self.assertEqual(n_cache_layers, [2, 2, None]) text_override_calls = self.encoding_override.call_args_list[:2] self.assertIs( text_override_calls[1].kwargs["quantized_model"], text_override_calls[0].kwargs["quantized_model"], ) - self.assertEqual(self.save_logits_quant_attrs.call_count, 2) - self.assertEqual(self.save_output_kv_cache_quant_attrs.call_count, 2) def test_invoke_skips_missing_text_decoder_prefill_graph(self): adapter = _make_mock_adapter() @@ -644,6 +670,19 @@ def test_quant_io_dtypes_returns_none_for_missing_recipe(self): self.assertIsNone(strategy._get_quant_io_dtypes(None)) + def test_quant_io_dtypes_returns_io_dtype_for_encoder_recipe(self): + """A non-KV component recipe produces an io_type-only mapping.""" + strategy = _make_strategy() + recipe = MagicMock() + del recipe.get_kv_io_bit_width + del recipe.get_logits_output_bit_width + recipe.get_act_bit_width.return_value = 16 + + self.assertEqual( + strategy._get_quant_io_dtypes(recipe), + {"io_type": torch.uint16}, + ) + def test_quant_io_dtypes_rejects_unsupported_widths(self): """Unsupported IO widths must not silently produce a partial result.""" strategy = _make_strategy() diff --git a/backends/qualcomm/genai_pipeline/tests/test_genai_pipeline.py b/backends/qualcomm/genai_pipeline/tests/test_genai_pipeline.py index 9e889e68d3b..f2d4561e93b 100644 --- a/backends/qualcomm/genai_pipeline/tests/test_genai_pipeline.py +++ b/backends/qualcomm/genai_pipeline/tests/test_genai_pipeline.py @@ -254,7 +254,7 @@ def test_model_preparation_receives_full_extra_options(self): } proxy = EngineProxy( {STAGE_MODEL_PREPARATION: EngineType.EXECUTORCH}, - backend_type=TEST_MOCK_BACKEND_TYPE, + backend_type=TEST_BACKEND_TYPE, ) pipeline = GenAIPipeline( model_preparation_stage=ModelPreparationStage(mock_model_prep), @@ -280,7 +280,7 @@ def test_quantization_receives_full_extra_options(self): } proxy = EngineProxy( {STAGE_QUANTIZATION: EngineType.EXECUTORCH}, - backend_type=TEST_MOCK_BACKEND_TYPE, + backend_type=TEST_BACKEND_TYPE, ) pipeline = GenAIPipeline( model_preparation_stage=None, diff --git a/backends/qualcomm/genai_pipeline/tests/test_graph_bundle.py b/backends/qualcomm/genai_pipeline/tests/test_graph_bundle.py index 342840a3df5..46381db80a7 100644 --- a/backends/qualcomm/genai_pipeline/tests/test_graph_bundle.py +++ b/backends/qualcomm/genai_pipeline/tests/test_graph_bundle.py @@ -48,25 +48,33 @@ def test_meta_is_not_shared_between_bundles(self): class TestQuantIoDtypes(unittest.TestCase): - """The tagger indexes both keys per node, so a partial mapping is invalid.""" + """A graph records ``io_type`` and, when applicable, ``kv_type``.""" def test_both_dtypes_are_accepted(self): - """The mapping quantization publishes carries both boundary dtypes.""" + """A KV-cache graph carries both boundary dtypes.""" quant_io_dtypes = {"kv_type": torch.uint8, "io_type": torch.uint16} bundle = _make_bundle(quant_io_dtypes=quant_io_dtypes) self.assertIs(bundle.quant_io_dtypes, quant_io_dtypes) + def test_io_dtype_without_kv_dtype_is_accepted(self): + """A graph without a KV cache carries only its IO dtype.""" + quant_io_dtypes = {"io_type": torch.uint16} + + bundle = _make_bundle(quant_io_dtypes=quant_io_dtypes) + + self.assertIs(bundle.quant_io_dtypes, quant_io_dtypes) + def test_skipped_quantization_is_none_not_a_partial_mapping(self): """None is how a graph says its IO stays float32.""" bundle = _make_bundle(quant_io_dtypes=None) self.assertIsNone(bundle.quant_io_dtypes) - def test_partial_mapping_is_rejected(self): - """One key alone would be a KeyError at lowering, so reject it here.""" - for partial in ({"kv_type": torch.uint8}, {"io_type": torch.uint16}, {}): + def test_mapping_without_io_dtype_is_rejected(self): + """Every quantized graph needs an IO-boundary dtype.""" + for partial in ({"kv_type": torch.uint8}, {}): with self.subTest(quant_io_dtypes=partial): with self.assertRaises(ValueError): _make_bundle(quant_io_dtypes=partial) diff --git a/backends/qualcomm/genai_pipeline/tests/test_model_lookup.py b/backends/qualcomm/genai_pipeline/tests/test_model_lookup.py new file mode 100644 index 00000000000..b9a2e7e463a --- /dev/null +++ b/backends/qualcomm/genai_pipeline/tests/test_model_lookup.py @@ -0,0 +1,42 @@ +# 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 argparse +import unittest +from types import SimpleNamespace +from unittest.mock import patch + +from executorch.backends.qualcomm.genai_pipeline.model_lookup import ( + get_model_arch, + get_model_config, +) + + +class TestModelLookup(unittest.TestCase): + + def test_mixed_case_model_name_resolves_to_canonical_config(self): + self.assertIs( + get_model_config("StOrIeS260K"), + get_model_config("stories260k"), + ) + + def test_get_model_arch_rejects_decoder_without_recipe(self): + config = SimpleNamespace(quant_recipe=None) + control_args = argparse.Namespace( + embedding_quantize=None, + model_mode="kv", + ) + + with patch( + "executorch.backends.qualcomm.genai_pipeline.model_lookup.get_model_config", + return_value=config, + ): + with self.assertRaisesRegex(ValueError, "KV-cache IO bit width"): + get_model_arch("custom", control_args) + + +if __name__ == "__main__": + unittest.main() From 7987450a543bf920fd7443e4d6431650cce6f738 Mon Sep 17 00:00:00 2001 From: yuyazhua Date: Thu, 8 Oct 2026 18:13:37 +0800 Subject: [PATCH 3/3] Lazy-load legacy imports in model lookup --- .../qualcomm/genai_pipeline/model_lookup.py | 24 +++++++++++++------ 1 file changed, 17 insertions(+), 7 deletions(-) diff --git a/backends/qualcomm/genai_pipeline/model_lookup.py b/backends/qualcomm/genai_pipeline/model_lookup.py index 05a00f493ea..3ea3e3464fe 100644 --- a/backends/qualcomm/genai_pipeline/model_lookup.py +++ b/backends/qualcomm/genai_pipeline/model_lookup.py @@ -33,13 +33,6 @@ GRAPH_FORWARD, TOK_EMBEDDING_GRAPH_NAMES, ) -from executorch.examples.qualcomm.oss_scripts.llama.model.static_llama import ModelArgs -from executorch.examples.qualcomm.oss_scripts.llama.wrappers.base_component import ( - get_model_specific_kwargs, - Mode, - next_power_of_two, - process_model_args, -) logger = logging.getLogger(__name__) @@ -98,6 +91,14 @@ def get_model_arch_config(model_name: str, control_args: Any) -> Any: ValueError: If the model needs a params file that ``control_args`` does not name. """ + from executorch.examples.qualcomm.oss_scripts.llama.model.static_llama import ( + ModelArgs, + ) + from executorch.examples.qualcomm.oss_scripts.llama.wrappers.base_component import ( + Mode, + process_model_args, + ) + model_name = model_name.lower() model_config = get_model_config(model_name) @@ -171,6 +172,15 @@ def get_model_arch( TokenEmbedding, ) from executorch.backends.qualcomm.genai_pipeline.models import LLM_VARIANT_ARCHS + from executorch.examples.qualcomm.oss_scripts.llama.model.static_llama import ( + ModelArgs, + ) + from executorch.examples.qualcomm.oss_scripts.llama.wrappers.base_component import ( + get_model_specific_kwargs, + Mode, + next_power_of_two, + process_model_args, + ) model_name = model_name.lower() model_config = get_model_config(model_name)