Gemma3-4b QNN example fixes

examples/gemma3/qnn/README.md

@@ -0,0 +1,122 @@
+# Gemma-3-4B Model Optimization
+
+This repository demonstrates the optimization of the [Google Gemma-3-4B](https://huggingface.co/google/gemma-3-4b-it) model using **post-training quantization (PTQ)** techniques for QNN (Qualcomm Neural Network) execution. The optimization process utilizes an environment based heavily upon the [PTQ tutorial for Phi-3.5](https://github.com/CodeLinaro/Olive/blob/main/examples/phi3_5/README.md)
+
+## File Overview
+
+This example contains the following key files:
+
+- **`env_setup.sh`** - Automated environment setup script (Linux only)
+- **`gemma3-4b-text-qnn-config.json`** - Olive configuration for optimizing the text component
+- **`gemma3-4b-vision-qnn-config.json`** - Olive configuration for optimizing the vision component
+- **`user_script.py`** - Dataset handling and preprocessing utilities
+- **`custom_gemma3_4b_it_vision.py`** - Vision model loader for the optimization pipeline
+
+## Prerequisites
+
+### System Requirements
+- **Operating System**: Linux (automated setup script is Linux-only)
+- **Python**: 3.10
+- **Package Manager**: [uv](https://docs.astral.sh/uv/getting-started/installation/#installation-methods)
+- **Storage**: ~13GB for COCO train2017 dataset (downloaded automatically)
+
+### Dependencies Installed by Setup Script
+The `env_setup.sh` script installs the following components:
+- setuptools (for building Olive from source)
+- Olive requirements and dependencies
+- AutoGPTQ (from source)
+- GPTQModel (specific commit: `558449bed3ef2653c36041650d30da6bbbca440d`)
+- onnxruntime-qnn (pre-release version)
+
+## Setup Instructions
+
+### Automated Setup (Recommended)
+```bash
+source env_setup.sh
+```
+
+### Manual Setup (Alternative)
+If you prefer to set up manually or need to troubleshoot:
+
+1. Install setuptools:
+   ```bash
+   uv pip install setuptools
+   ```
+
+2. Install requirements:
+   ```bash
+   uv pip install -r ../requirements.txt
+   uv pip install -r ../../../requirements.txt
+   ```
+
+3. Install AutoGPTQ from source:
+   ```bash
+   export BUILD_CUDA_EXT=0
+   uv pip install --no-build-isolation git+https://github.com/PanQiWei/AutoGPTQ.git
+   ```
+
+4. Install GPTQModel with Gemma3 fix:
+   ```bash
+   uv pip install --no-build-isolation git+https://github.com/ModelCloud/GPTQModel.git@558449bed3ef2653c36041650d30da6bbbca440d
+   ```
+
+5. Install onnxruntime-qnn:
+   ```bash
+   uv pip install -r https://raw.githubusercontent.com/microsoft/onnxruntime/refs/heads/main/requirements.txt
+   uv pip install -U --pre --extra-index-url https://aiinfra.pkgs.visualstudio.com/PublicPackages/_packaging/ORT-Nightly/pypi/simple onnxruntime-qnn --no-deps
+   ```
+
+> **Important:** The setup uses a specific commit hash for GPTQModel (`558449bed3ef2653c36041650d30da6bbbca440d`) to address a [memory leak issue](https://github.com/ModelCloud/GPTQModel/commit/558449bed3ef2653c36041650d30da6bbbca440d) with Gemma3 models.
+
+## Optimization Process
+
+Since Gemma-3-4B is a multi-modal model composed of both vision and text components, the strategy for optimizing it through Olive is to operate on the constituent models separately before configuring them to work together at the onnxruntime-genai stage.
+
+### Configuration Differences
+
+**Text Configuration (`gemma3-4b-text-qnn-config.json`)**:
+- Uses HuggingFace model directly (`google/gemma-3-4b-it`)
+- Applies comprehensive optimization pipeline: QuaRot → GptqModel → ModelBuilder → Quantization
+- Outputs to: `models/gemma-3-4b-it-text/`
+
+**Vision Configuration (`gemma3-4b-vision-qnn-config.json`)**:
+- Uses custom PyTorch model loader (`custom_gemma3_4b_it_vision.py`)
+- Simpler pipeline: ONNX Conversion → Graph Surgery → Quantization
+- Outputs to: `models/gemma-3-4b-it-vision/`
+
+### Running Optimization
+
+Execute the following commands to separately produce optimized binaries for each component:
+
+```bash
+olive run --config gemma3-4b-text-qnn-config.json
+```
+
+```bash
+olive run --config gemma3-4b-vision-qnn-config.json
+```
+
+## Expected Outputs
+
+After successful optimization, you will find:
+
+- **Text model outputs**: `models/gemma-3-4b-it-text/`
+- **Vision model outputs**: `models/gemma-3-4b-it-vision/`
+- **Cache directory**: `cache/` (intermediate files and downloaded datasets)
+- **Dataset**: `.cache/train2017/` (COCO train2017 images, ~13GB)
+
+Both configurations use `"no_artifacts": true`, meaning only the final optimized models are retained.
+
+## Troubleshooting
+
+### Common Issues
+
+**Insufficient Storage**: The COCO train2017 dataset requires ~13GB of storage and is downloaded automatically to `.cache/train2017/`.
+
+**Memory Requirements**: The optimization process, particularly for the text model with its comprehensive pipeline, requires substantial memory.
+
+**QNN Provider**: Ensure the QNNExecutionProvider is properly installed and configured in your environment.
+
+**Platform Limitation**: The current setup script is designed for Linux only. Windows/macOS users will need to adapt the manual setup steps.
+
+**Dataset Download**: If the COCO dataset download fails, check your internet connection and available storage. The script uses `wget` which must be available on your system.

examples/gemma3/qnn/app.py

@@ -0,0 +1,170 @@
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License
+
+import argparse
+import glob
+import json
+import os
+import time
+from pathlib import Path
+
+import onnxruntime_genai as og
+
+# og.set_log_options(enabled=True, model_input_values=True, model_output_values=True)
+
+
+def _find_dir_contains_sub_dir(current_dir: Path, target_dir_name):
+    curr_path = Path(current_dir).absolute()
+    target_dir = glob.glob(target_dir_name, root_dir=curr_path)
+    if target_dir:
+        return Path(curr_path / target_dir[0]).absolute()
+    else:
+        if curr_path.parent == curr_path:
+            # Root dir
+            return None
+        return _find_dir_contains_sub_dir(curr_path / "..", target_dir_name)
+
+
+def _complete(text, state):
+    return (glob.glob(text + "*") + [None])[state]
+
+
+def run(args: argparse.Namespace):
+    print("Loading model...")
+    config = og.Config(args.model_path)
+    if args.execution_provider != "follow_config":
+        config.clear_providers()
+        if args.execution_provider != "cpu":
+            print(f"Setting model to {args.execution_provider}...")
+            config.append_provider(args.execution_provider)
+    model = og.Model(config)
+    print("Model loaded")
+
+    tokenizer = og.Tokenizer(model)
+    processor = model.create_multimodal_processor()
+    stream = processor.create_stream()
+
+    interactive = not args.non_interactive
+
+    while True:
+        if interactive:
+            try:
+                import readline
+
+                readline.set_completer_delims(" \t\n;")
+                readline.parse_and_bind("tab: complete")
+                readline.set_completer(_complete)
+            except ImportError:
+                # Not available on some platforms. Ignore it.
+                pass
+            image_paths = [
+                image_path.strip()
+                for image_path in input("Image Path (comma separated; leave empty if no image): ").split(",")
+            ]
+        else:
+            if args.image_paths:
+                image_paths = args.image_paths
+            else:
+                image_paths = [str(Path(__file__).parent / "images" / "dog.jpg")]
+
+        image_paths = [image_path for image_path in image_paths if image_path]
+        print(image_paths)
+
+        images = None
+        if len(image_paths) == 0:
+            print("No image provided")
+        else:
+            for i, image_path in enumerate(image_paths):
+                if not os.path.exists(image_path):
+                    raise FileNotFoundError(f"Image file not found: {image_path}")
+                print(f"Using image: {image_path}")
+
+            images = og.Images.open(*image_paths)
+
+        if interactive:
+            text = input("Prompt: ")
+        else:
+            if args.prompt:
+                text = args.prompt
+            else:
+                text = "What is shown in this image?"
+
+        # Construct the "messages" argument passed to apply_chat_template
+        messages = []
+        if model.type == "phi3v":
+            # Combine all image tags and text into one user message
+            content = "".join([f"<|image_{i + 1}|>\n" for i in range(len(image_paths))]) + text
+            messages.append({"role": "user", "content": content})
+        else:
+            # Gemma3-style multimodal: structured content
+            content_list = [{"type": "image"} for _ in image_paths]
+            content_list.append({"type": "text", "text": text})
+            messages.append({"role": "user", "content": content_list})
+
+        # Apply the chat template using the tokenizer
+        message_json = json.dumps(messages)
+        print(message_json)
+        prompt = tokenizer.apply_chat_template(message_json, add_generation_prompt=True)
+
+        print("Processing images and prompt...")
+        inputs = processor(prompt, images=images)
+
+        print("Generating response...")
+        params = og.GeneratorParams(model)
+        params.set_search_options(max_length=1024)
+
+        print(inputs)
+
+        generator = og.Generator(model, params)
+        generator.set_inputs(inputs)
+        start_time = time.time()
+
+        while not generator.is_done():
+            generator.generate_next_token()
+
+            new_token = generator.get_next_tokens()[0]
+            print(stream.decode(new_token), end="", flush=True)
+
+        print()
+        total_run_time = time.time() - start_time
+        print(f"Total Time : {total_run_time:.2f}")
+
+        for _ in range(3):
+            print()
+
+        # Delete the generator to free the captured graph before creating another one
+        del generator
+
+        if not interactive:
+            break
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "-m", "--model_path", type=str, default="", required=True, help="Path to the folder containing the model"
+    )
+    parser.add_argument(
+        "-e",
+        "--execution_provider",
+        type=str,
+        required=False,
+        default="follow_config",
+        choices=["cpu", "cuda", "dml", "follow_config"],
+        help="Execution provider to run the ONNX Runtime session with. Defaults to follow_config that uses the execution provider listed in the genai_config.json instead.",
+    )
+    parser.add_argument(
+        "--image_paths", nargs="*", type=str, required=False, help="Path to the images, mainly for CI usage"
+    )
+    parser.add_argument(
+        "-pr", "--prompt", required=False, help="Input prompts to generate tokens from, mainly for CI usage"
+    )
+    parser.add_argument(
+        "--non-interactive",
+        action=argparse.BooleanOptionalAction,
+        default=True,
+        required=False,
+        help="Non-interactive mode, mainly for CI usage",
+    )
+    args = parser.parse_args()
+    run(args)