From 7fdb147190528eff7667dc0d9349617178439fba Mon Sep 17 00:00:00 2001 From: yuyazhua Date: Mon, 17 Aug 2026 17:16:57 +0800 Subject: [PATCH] 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(