From 142e9500461e7bbc6b5156b29d5e7e8f8661d606 Mon Sep 17 00:00:00 2001 From: Xiaoyu Date: Mon, 29 Jun 2026 15:15:10 -0700 Subject: [PATCH 01/11] Add multi components recipes --- .../multi_comp/.gitignore | 12 ++ .../multi_comp/README.md | 178 ++++++++++++++++++ Qwen-Qwen3-VL-2B-Instruct/multi_comp/info.yml | 15 ++ .../multi_comp/vlm_inference.py | 113 +++++++++++ .../multi_comp/vlm_optimize_components.json | 20 ++ .../multi_comp/vlm_quantize_then_export.json | 25 +++ .../LICENSE | 14 ++ .../multi_comp/.gitignore | 12 ++ .../multi_comp/README.md | 107 +++++++++++ .../multi_comp/info.yml | 9 + .../multi_comp/sd3_inference.py | 169 +++++++++++++++++ .../multi_comp/sd3_optimize_components.json | 90 +++++++++ 12 files changed, 764 insertions(+) create mode 100644 Qwen-Qwen3-VL-2B-Instruct/multi_comp/.gitignore create mode 100644 Qwen-Qwen3-VL-2B-Instruct/multi_comp/README.md create mode 100644 Qwen-Qwen3-VL-2B-Instruct/multi_comp/info.yml create mode 100644 Qwen-Qwen3-VL-2B-Instruct/multi_comp/vlm_inference.py create mode 100644 Qwen-Qwen3-VL-2B-Instruct/multi_comp/vlm_optimize_components.json create mode 100644 Qwen-Qwen3-VL-2B-Instruct/multi_comp/vlm_quantize_then_export.json create mode 100644 stabilityai-stable-diffusion-3-medium-diffusers/LICENSE create mode 100644 stabilityai-stable-diffusion-3-medium-diffusers/multi_comp/.gitignore create mode 100644 stabilityai-stable-diffusion-3-medium-diffusers/multi_comp/README.md create mode 100644 stabilityai-stable-diffusion-3-medium-diffusers/multi_comp/info.yml create mode 100644 stabilityai-stable-diffusion-3-medium-diffusers/multi_comp/sd3_inference.py create mode 100644 stabilityai-stable-diffusion-3-medium-diffusers/multi_comp/sd3_optimize_components.json diff --git a/Qwen-Qwen3-VL-2B-Instruct/multi_comp/.gitignore b/Qwen-Qwen3-VL-2B-Instruct/multi_comp/.gitignore new file mode 100644 index 000000000..ccb2ee32b --- /dev/null +++ b/Qwen-Qwen3-VL-2B-Instruct/multi_comp/.gitignore @@ -0,0 +1,12 @@ +# Exported ONNX packages +exported_vlm_pkg/ +exported_vlm_gptq_pkg/ + +# Quantized HF checkpoint +vlm_decoder_gptq_hf/ + +# Optimized components +out/ + +# Olive cache +cache/ diff --git a/Qwen-Qwen3-VL-2B-Instruct/multi_comp/README.md b/Qwen-Qwen3-VL-2B-Instruct/multi_comp/README.md new file mode 100644 index 000000000..2723a6b71 --- /dev/null +++ b/Qwen-Qwen3-VL-2B-Instruct/multi_comp/README.md @@ -0,0 +1,178 @@ +# Qwen3-VL-2B-Instruct — Multi-Component Optimization + +These recipes demonstrate two multi-component flows for +[Qwen3-VL-2B-Instruct](https://huggingface.co/Qwen/Qwen3-VL-2B-Instruct): + +- **Flow A — export first, then per-component optimization** + (`vlm_optimize_components.json`): export the VLM to ONNX once with the Mobius builder, then run a + single Olive config whose `builds` apply a **different pipeline to each component**. +- **Flow B — optimize a Torch component first, then export** + (`vlm_quantize_then_export.json`): run a Torch-stage GPTQ pass on the decoder component while + saving a complete HF directory, then export that directory with + `olive capture-onnx-graph --use_mobius_builder`. + +Olive loads an exported directory as a `CompositeModel` whose **component names are the subfolder +names**, so there is no need to memorize component names. + +## Prerequisites + +``` +pip install olive-ai +pip install mobius-ai +``` + +Exporting also needs `transformers` and access to the model on Hugging Face. + +--- + +## Recipe 1 — Export then per-component optimize (`vlm_optimize_components.json`) + +### Step 1 — Export + +``` +olive capture-onnx-graph --model_name_or_path Qwen/Qwen3-VL-2B-Instruct --use_mobius_builder --output_path exported_vlm_pkg +``` + +Mobius exports this model as three components, each in its own subfolder: + +``` +exported_vlm_pkg/ + decoder/model.onnx + vision_encoder/model.onnx + embedding/model.onnx +``` + +### Step 2 — Optimize + +``` +olive run --config vlm_optimize_components.json +``` + +| component | pipeline | intent | +|------------------|-----------------|-------------------------------------| +| `decoder` | `dynamic_quant` | INT8-quantize the language decoder | +| `vision_encoder` | `to_fp16` | keep the vision tower in FP16 | +| `embedding` | `to_fp16` | keep the embedding in FP16 | + +> The three component names (`decoder`, `vision_encoder`, `embedding`) are exactly what Mobius +> produces for `Qwen/Qwen3-VL-2B-Instruct`. For a different VLM, adjust the component names in the +> config to match the subfolder names your export actually produced. + +### Step 3 — Inference with ORT GenAI + +Run text generation with the exported ONNX models using **onnxruntime-genai**: + +```bash +# Text-only +python vlm_inference.py --prompt "The capital of France is" + +# With image input +python vlm_inference.py --prompt "Describe this image." --image photo.jpg + +# Custom settings +python vlm_inference.py --model_dir exported_vlm_pkg --max_new_tokens 256 +``` + +The inference script (`vlm_inference.py`) uses ORT GenAI which handles: +- **Tokenization**: built-in tokenizer from saved HF tokenizer files +- **Embedding**: ONNX `embedding/model.onnx` (token embed + image feature mixing) +- **Vision encoding**: ONNX `vision_encoder/model.onnx` (when `--image` is provided) +- **Decoding**: ONNX `decoder/model.onnx` with KV cache (autoregressive generation) + +Options: +``` +--prompt TEXT Text prompt +--image PATH Optional image file for multimodal input +--max_new_tokens N Maximum tokens to generate (default: 128) +--model_dir DIR Path to exported model directory (default: exported_vlm_pkg) +``` + +#### Setup requirements + +The export directory needs these files alongside the ONNX models: + +``` +exported_vlm_pkg/ + genai_config.json # Model type, I/O mappings, search config + tokenizer.json # HF tokenizer + tokenizer_config.json + vision_processor.json # Vision preprocessing config + decoder/model.onnx + vision_encoder/model.onnx + embedding/model.onnx +``` + +To create the tokenizer files after export: + +```python +from transformers import AutoTokenizer +tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen3-VL-2B-Instruct", trust_remote_code=True) +tokenizer.save_pretrained("exported_vlm_pkg") +``` + +For the `genai_config.json` structure, see the +[Mobius ORT GenAI examples](https://github.com/microsoft/mobius/tree/main/examples) which write the +config automatically. + +> **Note.** Install `onnxruntime-genai` (`pip install onnxruntime-genai`) to use this script. + +--- + +## Recipe 2 — Torch decoder quantization, then Mobius export (`vlm_quantize_then_export.json`) + +This is **Flow B**: quantize only the Torch decoder component first, then export the resulting +complete HF directory with the Olive capture CLI using the Mobius builder. + +### Step 1 — Quantize the decoder component + +``` +olive run --config vlm_quantize_then_export.json +``` + +The config uses `builds.components: ["decoder"]`, so Olive asks Mobius for the VLM component plan, +scopes the Torch GPTQ pass to the decoder submodule, and saves the original HF folder layout with +the decoder quantized in place. This output is **not** a standalone decoder checkpoint; it is a +complete HF model directory: + +``` +out/vlm_decoder_gptq_hf/ +``` + +The recipe uses the GPTQ pass defaults for calibration data (`Salesforce/wikitext`). For production, +add a `data_config` to `decoder_gptq` with your own text or multimodal calibration set. + +### Step 2 — Export the quantized HF directory with the Mobius builder + +``` +olive capture-onnx-graph \ + --model_name_or_path vlm_decoder_gptq_hf \ + --use_mobius_builder \ + --trust_remote_code \ + --precision fp16 \ + --output_path exported_vlm_gptq_pkg +``` + +Output: + +``` +exported_vlm_gptq_pkg/ + decoder/model.onnx + vision_encoder/model.onnx + embedding/model.onnx +``` + +> **Note.** The Torch GPTQ pass saves Olive-packed weights (`quant_method="olive"`). Use this export +> step with a Mobius builder version that supports Olive-packed quantized HF checkpoints. + +--- + +## Notes + +- The passes in Recipe 1 (`OnnxFloatToFloat16`, `OnnxDynamicQuantization`) are **illustrative** and + chosen to run without calibration data. Swap in `OrtTransformersOptimization`, + `OnnxStaticQuantization` (with a `data_config`), or other ONNX passes for production-quality + optimization. +- The ONNX component recipe runs on the EP declared in its `systems` section. The Torch GPTQ recipe + targets CUDA because VLM decoder GPTQ is GPU-oriented. +- `builds.components` selects which exported components to optimize. Only the components with a build + are touched; the rest remain as exported. diff --git a/Qwen-Qwen3-VL-2B-Instruct/multi_comp/info.yml b/Qwen-Qwen3-VL-2B-Instruct/multi_comp/info.yml new file mode 100644 index 000000000..26a0c37da --- /dev/null +++ b/Qwen-Qwen3-VL-2B-Instruct/multi_comp/info.yml @@ -0,0 +1,15 @@ +keywords: + - olive-ai +recipes: + - name: qwen3vl-2B-Instruct + file: vlm_optimize_components.json + eps: + - CPUExecutionProvider + devices: + - cpu + - name: qwen3vl-2B-Instruct + file: vlm_quantize_then_export.json + eps: + - CUDAExecutionProvider + devices: + - gpu diff --git a/Qwen-Qwen3-VL-2B-Instruct/multi_comp/vlm_inference.py b/Qwen-Qwen3-VL-2B-Instruct/multi_comp/vlm_inference.py new file mode 100644 index 000000000..15e11aeb7 --- /dev/null +++ b/Qwen-Qwen3-VL-2B-Instruct/multi_comp/vlm_inference.py @@ -0,0 +1,113 @@ +#!/usr/bin/env python +"""VLM (Qwen3-VL-2B-Instruct) inference using ORT GenAI with exported ONNX models. + +Usage: + # Text-only + python vlm_inference.py --prompt "The capital of France is" + + # With image + python vlm_inference.py --prompt "Describe this image." --image photo.jpg + + # Custom model directory + python vlm_inference.py --model_dir exported_vlm_pkg --prompt "What is 2+2?" +""" + +import argparse +import os + +import onnxruntime_genai as og + + +def generate_text(model_dir: str, prompt: str, max_new_tokens: int = 128) -> str: + """Run text-only generation.""" + model = og.Model(model_dir) + tokenizer = og.Tokenizer(model) + + input_ids = tokenizer.encode(prompt) + params = og.GeneratorParams(model) + params.set_search_options(max_length=len(input_ids) + max_new_tokens) + + generator = og.Generator(model, params) + generator.append_tokens(input_ids) + + tokenizer_stream = tokenizer.create_stream() + generated = [] + while not generator.is_done(): + generator.generate_next_token() + token = generator.get_next_tokens()[0] + generated.append(token) + print(tokenizer_stream.decode(token), end="", flush=True) + if len(generated) >= max_new_tokens: + break + + print() + del generator + return tokenizer.decode(generated) + + +def generate_with_image(model_dir: str, prompt: str, image_path: str, max_new_tokens: int = 128) -> str: + """Run multimodal generation with image input.""" + model = og.Model(model_dir) + tokenizer = og.Tokenizer(model) + processor = model.create_multimodal_processor() + + images = og.Images.open(image_path) + inputs = processor(prompt, images=images) + + params = og.GeneratorParams(model) + params.set_search_options(max_length=4096) + + generator = og.Generator(model, params) + generator.set_inputs(inputs) + + tokenizer_stream = tokenizer.create_stream() + generated = [] + while not generator.is_done(): + generator.generate_next_token() + token = generator.get_next_tokens()[0] + generated.append(token) + print(tokenizer_stream.decode(token), end="", flush=True) + if len(generated) >= max_new_tokens: + break + + print() + del generator + return tokenizer.decode(generated) + + +def main(): + parser = argparse.ArgumentParser(description="VLM inference with ORT GenAI") + parser.add_argument("--prompt", default="The capital of France is") + parser.add_argument("--image", default=None, help="Path to an image file for vision input") + parser.add_argument("--max_new_tokens", type=int, default=128) + parser.add_argument("--model_dir", default="exported_vlm_pkg") + args = parser.parse_args() + + genai_config = os.path.join(args.model_dir, "genai_config.json") + if not os.path.exists(genai_config): + print(f"Error: genai_config.json not found in {args.model_dir}") + print("Run export first:") + print( + " olive capture-onnx-graph --model_name_or_path Qwen/Qwen3-VL-2B-Instruct " + "--use_mobius_builder --output_path exported_vlm_pkg" + ) + print("Then create genai_config.json and save tokenizer (see README.md).") + return + + print(f"Model: {args.model_dir}") + print(f"Prompt: {args.prompt}") + if args.image: + print(f"Image: {args.image}") + print("-" * 50) + + if args.image: + output = generate_with_image(args.model_dir, args.prompt, args.image, args.max_new_tokens) + else: + output = generate_text(args.model_dir, args.prompt, args.max_new_tokens) + + print("-" * 50) + print(f"Output: {output}") + + +if __name__ == "__main__": + main() diff --git a/Qwen-Qwen3-VL-2B-Instruct/multi_comp/vlm_optimize_components.json b/Qwen-Qwen3-VL-2B-Instruct/multi_comp/vlm_optimize_components.json new file mode 100644 index 000000000..89e4d47d5 --- /dev/null +++ b/Qwen-Qwen3-VL-2B-Instruct/multi_comp/vlm_optimize_components.json @@ -0,0 +1,20 @@ +{ + "input_model": { "type": "CompositeModel", "config": { "model_path": "exported_vlm_pkg" } }, + "systems": { + "local_system": { + "type": "LocalSystem", + "accelerators": [ { "device": "cpu", "execution_providers": [ "CPUExecutionProvider" ] } ] + } + }, + "passes": { "to_fp16": { "type": "OnnxFloatToFloat16" }, "dynamic_quant": { "type": "OnnxDynamicQuantization" } }, + "engine": { "host": "local_system", "target": "local_system", "evaluate_input_model": false, "cache_dir": "cache" }, + "builds": { + "decoder": { "components": [ "decoder" ], "pipeline": [ "dynamic_quant" ], "output_dir": "out/decoder" }, + "vision_encoder": { + "components": [ "vision_encoder" ], + "pipeline": [ "to_fp16" ], + "output_dir": "out/vision_encoder" + }, + "embedding": { "components": [ "embedding" ], "pipeline": [ "to_fp16" ], "output_dir": "out/embedding" } + } +} diff --git a/Qwen-Qwen3-VL-2B-Instruct/multi_comp/vlm_quantize_then_export.json b/Qwen-Qwen3-VL-2B-Instruct/multi_comp/vlm_quantize_then_export.json new file mode 100644 index 000000000..cd1855016 --- /dev/null +++ b/Qwen-Qwen3-VL-2B-Instruct/multi_comp/vlm_quantize_then_export.json @@ -0,0 +1,25 @@ +{ + "input_model": { + "type": "HfModel", + "config": { + "model_path": "Qwen/Qwen3-VL-2B-Instruct", + "task": "image-text-to-text", + "load_kwargs": { "torch_dtype": "float16", "trust_remote_code": true } + } + }, + "systems": { + "local_gpu": { + "type": "LocalSystem", + "accelerators": [ { "device": "gpu", "execution_providers": [ "CUDAExecutionProvider" ] } ] + } + }, + "passes": { "decoder_gptq": { "type": "Gptq", "bits": 4, "group_size": 128, "sym": true, "lm_head": false } }, + "engine": { "host": "local_gpu", "target": "local_gpu", "evaluate_input_model": false, "cache_dir": "cache" }, + "builds": { + "decoder_gptq": { + "components": [ "decoder" ], + "pipeline": [ "decoder_gptq" ], + "output_dir": "vlm_decoder_gptq_hf" + } + } +} diff --git a/stabilityai-stable-diffusion-3-medium-diffusers/LICENSE b/stabilityai-stable-diffusion-3-medium-diffusers/LICENSE new file mode 100644 index 000000000..d7a6188fb --- /dev/null +++ b/stabilityai-stable-diffusion-3-medium-diffusers/LICENSE @@ -0,0 +1,14 @@ +Stable Diffusion 3 Medium is released under the Stability AI Community License. + +The model "stabilityai/stable-diffusion-3-medium-diffusers" is a gated model. You must +review and accept its license on Hugging Face before downloading or using the weights: + + https://huggingface.co/stabilityai/stable-diffusion-3-medium-diffusers + +Full license text: + + https://stability.ai/license + +The recipes in this directory are provided by the olive-recipes project under the +repository's root LICENSE. The license referenced above governs use of the Stable +Diffusion 3 model weights only. diff --git a/stabilityai-stable-diffusion-3-medium-diffusers/multi_comp/.gitignore b/stabilityai-stable-diffusion-3-medium-diffusers/multi_comp/.gitignore new file mode 100644 index 000000000..0bdc513bc --- /dev/null +++ b/stabilityai-stable-diffusion-3-medium-diffusers/multi_comp/.gitignore @@ -0,0 +1,12 @@ +# Exported ONNX packages +exported_pkg/ +exported_sd3_full2/ + +# Optimized components +out/ + +# Olive cache +cache/ + +# Generated images +*.png diff --git a/stabilityai-stable-diffusion-3-medium-diffusers/multi_comp/README.md b/stabilityai-stable-diffusion-3-medium-diffusers/multi_comp/README.md new file mode 100644 index 000000000..4ba1ab3fc --- /dev/null +++ b/stabilityai-stable-diffusion-3-medium-diffusers/multi_comp/README.md @@ -0,0 +1,107 @@ +# Stable Diffusion 3 Medium — Multi-Component Optimization + +This recipe demonstrates a **multi-component flow** for +[Stable Diffusion 3 Medium](https://huggingface.co/stabilityai/stable-diffusion-3-medium-diffusers): +export the pipeline to ONNX once with the Mobius builder, then run a single Olive config whose +`builds` apply a **different pipeline to each component**. + +The flow is two explicit steps: + +1. **Export** the model to a directory of per-component ONNX subfolders using the Olive CLI with the + Mobius builder. +2. **Optimize** by pointing an Olive config at that directory; each component subfolder becomes a + selectable component that a `build` can target. + +There is no need to memorize component names: each exported component lives in its own folder, and +Olive loads the export directory as a `CompositeModel` whose **component names are the subfolder +names**. + +## Prerequisites + +``` +pip install olive-ai +pip install mobius-ai +``` + +Exporting a diffusion pipeline also needs `diffusers`/`transformers` and access to the model on +Hugging Face (Stable Diffusion 3 is a gated model — accept its license and `huggingface-cli login` +first). + +## Step 1 — Export with the CLI + +``` +olive capture-onnx-graph --model_name_or_path stabilityai/stable-diffusion-3-medium-diffusers --use_mobius_builder --output_path exported_pkg +``` + +Mobius exports each neural-network component to its own subfolder: + +``` +exported_pkg/ + text_encoder/model.onnx # CLIP-L text encoder + text_encoder_2/model.onnx # CLIP-G text encoder + text_encoder_3/model.onnx # T5-XXL text encoder + transformer/model.onnx # MMDiT denoising backbone + vae_encoder/model.onnx + vae_decoder/model.onnx +``` + +> **Note.** The exact subfolders depend on the pipeline; the optimize config below only +> needs `builds` for the components you actually want to optimize. + +## Step 2 — Optimize each component + +Run from the directory that contains `exported_pkg/`: + +``` +olive run --config sd3_optimize_components.json +``` + +This applies a different pipeline per component: + +| component | pipeline | intent | +|------------------|------------------------------|--------------------------------------------| +| `transformer` | `OrtTransformersOptimization`| FP16-optimize the heavy denoising backbone | +| `vae_encoder` | `OrtTransformersOptimization`| FP16-optimize the VAE encoder | +| `vae_decoder` | `OrtTransformersOptimization`| FP16-optimize the VAE decoder | + +Output: + +``` +out/transformer/ # optimized transformer +out/vae_encoder/ # optimized VAE encoder +out/vae_decoder/ # optimized VAE decoder +``` + +Each build writes one optimized component; components without a build stay as exported. + +## Step 3 — Inference + +Run end-to-end image generation with the exported ONNX models: + +``` +python sd3_inference.py --prompt "A photo of a cat sitting on a windowsill" --steps 28 --output result.png +``` + +The inference script (`sd3_inference.py`) uses: +- **Text encoding**: ONNX Runtime with exported CLIP-L, CLIP-G, and T5-XXL encoders (run once) +- **Denoising**: ONNX Runtime with the exported SD3 transformer (28 steps) +- **VAE decoding**: ONNX Runtime with the exported VAE decoder + +Options: +``` +--prompt TEXT Text prompt for image generation +--steps N Number of denoising steps (default: 28) +--seed N Random seed (default: 42) +--output PATH Output image path (default: sd3_output.png) +--onnx_dir DIR Path to exported model directory (default: exported_sd3_full2) +``` + +> **Note.** SD3 is a gated model — you need `huggingface-cli login` or set `HF_TOKEN` to export. +> The tokenizers (CLIP and T5) still run via the `transformers` library. + +## Notes + +- The passes here are **illustrative**. Swap in `OnnxStaticQuantization` (with a `data_config`), + `OnnxDynamicQuantization`, or other ONNX passes for production-quality optimization. +- `builds.components` selects which exported components to optimize. Only the components with a build + are touched; the rest remain as exported. diff --git a/stabilityai-stable-diffusion-3-medium-diffusers/multi_comp/info.yml b/stabilityai-stable-diffusion-3-medium-diffusers/multi_comp/info.yml new file mode 100644 index 000000000..bae3dc157 --- /dev/null +++ b/stabilityai-stable-diffusion-3-medium-diffusers/multi_comp/info.yml @@ -0,0 +1,9 @@ +keywords: + - olive-ai +recipes: + - name: stable-diffusion-3-medium-multi-component + file: sd3_optimize_components.json + eps: + - CUDAExecutionProvider + devices: + - gpu diff --git a/stabilityai-stable-diffusion-3-medium-diffusers/multi_comp/sd3_inference.py b/stabilityai-stable-diffusion-3-medium-diffusers/multi_comp/sd3_inference.py new file mode 100644 index 000000000..0c1e79fad --- /dev/null +++ b/stabilityai-stable-diffusion-3-medium-diffusers/multi_comp/sd3_inference.py @@ -0,0 +1,169 @@ +#!/usr/bin/env python +"""SD3 end-to-end inference using all ONNX components (text encoders + transformer + VAE). + +Usage: + python sd3_inference.py --prompt "A photo of a cat sitting on a windowsill" + python sd3_inference.py --prompt "A futuristic city" --steps 50 --output city.png +""" + +import argparse +import os + +import numpy as np +import onnxruntime as ort +import torch +from diffusers import FlowMatchEulerDiscreteScheduler +from PIL import Image +from transformers import CLIPTokenizer, T5TokenizerFast + +MODEL_ID = "stabilityai/stable-diffusion-3-medium-diffusers" +ONNX_DIR = "exported_sd3_full2" + + +def encode_text(prompt: str, onnx_dir: str, model_id: str) -> tuple[np.ndarray, np.ndarray]: + """Encode prompt using ONNX CLIP-L, CLIP-G, and T5-XXL text encoders. + + Returns: + encoder_hidden_states: [1, 410, 4096] + pooled_projections: [1, 2048] + + """ + # Load tokenizers (lightweight, no model weights) + tokenizer_l = CLIPTokenizer.from_pretrained(model_id, subfolder="tokenizer") + tokenizer_g = CLIPTokenizer.from_pretrained(model_id, subfolder="tokenizer_2") + tokenizer_t5 = T5TokenizerFast.from_pretrained(model_id, subfolder="tokenizer_3") + + # Load ONNX sessions + sess_l = ort.InferenceSession(os.path.join(onnx_dir, "text_encoder", "model.onnx")) + sess_g = ort.InferenceSession(os.path.join(onnx_dir, "text_encoder_2", "model.onnx")) + sess_t5 = ort.InferenceSession(os.path.join(onnx_dir, "text_encoder_3", "model.onnx")) + + # CLIP-L + tokens_l = tokenizer_l(prompt, padding="max_length", max_length=77, return_tensors="np", truncation=True) + out_l = sess_l.run( + None, + { + "input_ids": tokens_l["input_ids"].astype(np.int64), + "attention_mask": tokens_l["attention_mask"].astype(np.int64), + }, + ) + clip_l_hidden = out_l[0] # last_hidden_state [1, 77, 768] + clip_l_pooled = out_l[1] # text_embeds [1, 768] + + # CLIP-G + tokens_g = tokenizer_g(prompt, padding="max_length", max_length=77, return_tensors="np", truncation=True) + out_g = sess_g.run( + None, + { + "input_ids": tokens_g["input_ids"].astype(np.int64), + "attention_mask": tokens_g["attention_mask"].astype(np.int64), + }, + ) + clip_g_hidden = out_g[0] # last_hidden_state [1, 77, 1280] + clip_g_pooled = out_g[1] # text_embeds [1, 1280] + + # T5-XXL + tokens_t5 = tokenizer_t5(prompt, padding="max_length", max_length=256, return_tensors="np", truncation=True) + out_t5 = sess_t5.run(None, {"input_ids": tokens_t5["input_ids"].astype(np.int64)}) + t5_hidden = out_t5[0] # last_hidden_state [1, 256, 4096] + + # Pad CLIP outputs to 4096 and concatenate + clip_l_padded = np.pad(clip_l_hidden, ((0, 0), (0, 0), (0, 4096 - 768))) # [1, 77, 4096] + clip_g_padded = np.pad(clip_g_hidden, ((0, 0), (0, 0), (0, 4096 - 1280))) # [1, 77, 4096] + encoder_hidden_states = np.concatenate([clip_l_padded, clip_g_padded, t5_hidden], axis=1) # [1, 410, 4096] + pooled_projections = np.concatenate([clip_l_pooled, clip_g_pooled], axis=-1) # [1, 2048] + + return encoder_hidden_states.astype(np.float32), pooled_projections.astype(np.float32) + + +def denoise( + onnx_dir: str, + encoder_hidden_states: np.ndarray, + pooled_projections: np.ndarray, + scheduler: FlowMatchEulerDiscreteScheduler, + latent_shape: tuple = (1, 16, 64, 64), + seed: int = 42, +) -> torch.Tensor: + """Run the denoising loop using the ONNX transformer.""" + sess = ort.InferenceSession(os.path.join(onnx_dir, "transformer", "model.onnx")) + + torch.manual_seed(seed) + latents = torch.randn(latent_shape) + + for i, t in enumerate(scheduler.timesteps): + noise_pred = sess.run( + None, + { + "sample": latents.numpy(), + "timestep": np.array([t.item()], dtype=np.int64), + "encoder_hidden_states": encoder_hidden_states, + "pooled_projections": pooled_projections, + }, + )[0] + latents = scheduler.step(torch.from_numpy(noise_pred), t, latents, return_dict=False)[0] + if i % 7 == 0: + print(f" Step {i}/{len(scheduler.timesteps)}, t={t.item():.1f}") + + return latents + + +def decode_latents(latents: torch.Tensor, onnx_dir: str) -> np.ndarray: + """Decode latents to image using the ONNX VAE decoder.""" + sess = ort.InferenceSession(os.path.join(onnx_dir, "vae_decoder", "model.onnx")) + + # SD3 VAE scaling: latents / scaling_factor + shift_factor + # SD3 defaults: scaling_factor=1.5305, shift_factor=0.0609 + scaling_factor = 1.5305 + shift_factor = 0.0609 + latents_scaled = latents / scaling_factor + shift_factor + + output = sess.run(None, {"latent_sample": latents_scaled.numpy()})[0] + # output: [1, 3, H, W] in [-1, 1] + image = (output / 2 + 0.5).clip(0, 1) + image = np.transpose(image[0], (1, 2, 0)) # [H, W, 3] + return (image * 255).astype(np.uint8) + + +def main(): + parser = argparse.ArgumentParser(description="SD3 all-ONNX inference") + parser.add_argument("--prompt", default="A photo of a cat sitting on a windowsill") + parser.add_argument("--steps", type=int, default=28) + parser.add_argument("--seed", type=int, default=42) + parser.add_argument("--output", default="sd3_output.png") + parser.add_argument("--model_id", default=MODEL_ID) + parser.add_argument("--onnx_dir", default=ONNX_DIR) + args = parser.parse_args() + + # Verify exported model exists + transformer_path = os.path.join(args.onnx_dir, "transformer", "model.onnx") + if not os.path.exists(transformer_path): + print(f"Error: ONNX model not found at {args.onnx_dir}/") + print( + "Run: olive capture-onnx-graph --model_name_or_path " + "stabilityai/stable-diffusion-3-medium-diffusers " + "--use_mobius_builder --output_path exported_sd3_full2" + ) + return + + print(f"Prompt: {args.prompt}") + print(f"Steps: {args.steps}, Seed: {args.seed}") + print(f"ONNX dir: {args.onnx_dir}") + + print("\n1. Encoding text (ONNX CLIP-L + CLIP-G + T5-XXL)...") + encoder_hidden_states, pooled_projections = encode_text(args.prompt, args.onnx_dir, args.model_id) + print(f" encoder_hidden_states: {encoder_hidden_states.shape}") + print(f" pooled_projections: {pooled_projections.shape}") + + print("\n2. Denoising (ONNX SD3 transformer)...") + scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained(args.model_id, subfolder="scheduler") + scheduler.set_timesteps(args.steps) + latents = denoise(args.onnx_dir, encoder_hidden_states, pooled_projections, scheduler, seed=args.seed) + + print("\n3. Decoding latents (ONNX VAE decoder)...") + image = decode_latents(latents, args.onnx_dir) + Image.fromarray(image).save(args.output) + print(f"\nSaved: {args.output} ({image.shape[1]}x{image.shape[0]})") + + +if __name__ == "__main__": + main() diff --git a/stabilityai-stable-diffusion-3-medium-diffusers/multi_comp/sd3_optimize_components.json b/stabilityai-stable-diffusion-3-medium-diffusers/multi_comp/sd3_optimize_components.json new file mode 100644 index 000000000..8551d5325 --- /dev/null +++ b/stabilityai-stable-diffusion-3-medium-diffusers/multi_comp/sd3_optimize_components.json @@ -0,0 +1,90 @@ +{ + "input_model": { "type": "CompositeModel", "config": { "model_path": "exported_pkg" } }, + "systems": { + "local_system": { + "type": "LocalSystem", + "accelerators": [ { "device": "gpu", "execution_providers": [ "CUDAExecutionProvider" ] } ] + } + }, + "passes": { + "optimize_transformer": { + "type": "OrtTransformersOptimization", + "model_type": "unet", + "opt_level": 0, + "float16": true, + "use_gpu": true, + "keep_io_types": false, + "optimization_options": { + "enable_gelu": true, + "enable_layer_norm": true, + "enable_attention": true, + "use_multi_head_attention": true, + "enable_skip_layer_norm": false, + "enable_embed_layer_norm": true, + "enable_bias_skip_layer_norm": false, + "enable_bias_gelu": true, + "enable_gelu_approximation": false, + "enable_qordered_matmul": false, + "enable_shape_inference": true, + "enable_gemm_fast_gelu": false, + "enable_nhwc_conv": false, + "enable_group_norm": true, + "enable_bias_splitgelu": false, + "enable_packed_qkv": true, + "enable_packed_kv": true, + "enable_bias_add": false, + "group_norm_channels_last": false + }, + "force_fp32_ops": [ "RandomNormalLike" ] + }, + "optimize_vae": { + "type": "OrtTransformersOptimization", + "model_type": "vae", + "opt_level": 0, + "float16": true, + "use_gpu": true, + "keep_io_types": false, + "optimization_options": { + "enable_gelu": true, + "enable_layer_norm": true, + "enable_attention": true, + "use_multi_head_attention": true, + "enable_skip_layer_norm": false, + "enable_embed_layer_norm": true, + "enable_bias_skip_layer_norm": false, + "enable_bias_gelu": true, + "enable_gelu_approximation": false, + "enable_qordered_matmul": false, + "enable_shape_inference": true, + "enable_gemm_fast_gelu": false, + "enable_nhwc_conv": false, + "enable_group_norm": true, + "enable_bias_splitgelu": false, + "enable_packed_qkv": true, + "enable_packed_kv": true, + "enable_bias_add": false, + "group_norm_channels_last": false + }, + "force_fp32_ops": [ "RandomNormalLike" ], + "force_fp16_inputs": { "GroupNorm": [ 0, 1, 2 ] } + } + }, + "engine": { "host": "local_system", "target": "local_system", "evaluate_input_model": false, "cache_dir": "cache" }, + "builds": { + "transformer": { + "components": [ "transformer" ], + "pipeline": [ "optimize_transformer" ], + "output_dir": "out/transformer" + }, + "vae_encoder": { + "components": [ "vae_encoder" ], + "pipeline": [ "optimize_vae" ], + "output_dir": "out/vae_encoder" + }, + "vae_decoder": { + "components": [ "vae_decoder" ], + "pipeline": [ "optimize_vae" ], + "output_dir": "out/vae_decoder" + } + } +} From ac3cb34b2d7f0606836c04196890843de622965c Mon Sep 17 00:00:00 2001 From: Xiaoyu Zhang Date: Tue, 25 Aug 2026 17:22:28 -0700 Subject: [PATCH 02/11] Add Gemma4 quantize-then-export recipe --- google-gemma-4-E2B-it/README.md | 11 +- google-gemma-4-E2B-it/inference.py | 16 ++- google-gemma-4-E2B-it/multi_comp/.gitignore | 11 ++ google-gemma-4-E2B-it/multi_comp/README.md | 120 ++++++++++++++++++ .../gemma4_quantize_then_export.json | 35 +++++ google-gemma-4-E2B-it/multi_comp/info.yml | 15 +++ google-gemma-4-E2B-it/requirements.txt | 3 +- 7 files changed, 207 insertions(+), 4 deletions(-) create mode 100644 google-gemma-4-E2B-it/multi_comp/.gitignore create mode 100644 google-gemma-4-E2B-it/multi_comp/README.md create mode 100644 google-gemma-4-E2B-it/multi_comp/gemma4_quantize_then_export.json create mode 100644 google-gemma-4-E2B-it/multi_comp/info.yml diff --git a/google-gemma-4-E2B-it/README.md b/google-gemma-4-E2B-it/README.md index 50ccd7c9d..8349e62d2 100644 --- a/google-gemma-4-E2B-it/README.md +++ b/google-gemma-4-E2B-it/README.md @@ -16,10 +16,14 @@ post-processing required. ## Prerequisites ```bash -pip install olive-ai mobius-ai +pip install olive-ai +pip install "git+https://github.com/onnxruntime/mobius.git@b9b4ef4" pip install -r requirements.txt ``` +The pinned Mobius revision includes support for the heterogeneous per-layer +Gemma 4 configuration used by current Transformers releases. + Install ONNX Runtime GenAI: | Device | Install Command | @@ -39,6 +43,11 @@ Install ONNX Runtime GenAI: K-Quant (Q4_K_M) is significantly faster with GPU acceleration — install `cupy-cuda12x` for a 19–51× speedup during quantization. +For a Torch-stage quantize-then-export flow, see +[`multi_comp/README.md`](multi_comp/README.md). It applies INT4 RTN only to +the decoder, saves a complete Hugging Face checkpoint, and then exports all +four components with Mobius. + ## Build ```bash diff --git a/google-gemma-4-E2B-it/inference.py b/google-gemma-4-E2B-it/inference.py index 18eb60dc9..fbc7695fe 100644 --- a/google-gemma-4-E2B-it/inference.py +++ b/google-gemma-4-E2B-it/inference.py @@ -38,8 +38,20 @@ def format_chat_prompt(tokenizer, prompt: str, system_prompt: str | None = None) messages.append({"role": "system", "content": system_prompt}) messages.append({"role": "user", "content": prompt}) - # ORT GenAI tokenizer expects the messages as a JSON string - return tokenizer.apply_chat_template(json.dumps(messages)) + # ORT GenAI tokenizer expects the messages as a JSON string. Older + # releases cannot execute Gemma 4's newer Jinja template, so fall back to + # the equivalent canonical turn markers for simple text generation. + try: + return tokenizer.apply_chat_template(json.dumps(messages)) + except RuntimeError as exc: + if "Invalid or unsupported chat template" not in str(exc): + raise + + turns = [""] + if system_prompt: + turns.append(f"<|turn>system\n{system_prompt.strip()}\n") + turns.append(f"<|turn>user\n{prompt.strip()}\n<|turn>model\n") + return "".join(turns) def generate( diff --git a/google-gemma-4-E2B-it/multi_comp/.gitignore b/google-gemma-4-E2B-it/multi_comp/.gitignore new file mode 100644 index 000000000..3ed0b3c9c --- /dev/null +++ b/google-gemma-4-E2B-it/multi_comp/.gitignore @@ -0,0 +1,11 @@ +# Generated Hugging Face and ONNX packages +gemma4_decoder_int4_hf/ +exported_gemma4_int4_pkg/ + +# Olive caches and generated dependency lists +.olive-cache/ +cache/ +olive_requirements.txt + +# Local run logs +*.log diff --git a/google-gemma-4-E2B-it/multi_comp/README.md b/google-gemma-4-E2B-it/multi_comp/README.md new file mode 100644 index 000000000..6bcdd467f --- /dev/null +++ b/google-gemma-4-E2B-it/multi_comp/README.md @@ -0,0 +1,120 @@ +# Gemma 4 E2B — Quantize Then Export + +This recipe quantizes the Torch decoder component of +[`google/gemma-4-E2B-it`](https://huggingface.co/google/gemma-4-E2B-it) +before exporting the complete multimodal model with Mobius. + +The flow has two explicit stages: + +1. Olive selects the Gemma 4 `decoder` component, applies INT4 RTN, and saves a + complete Hugging Face directory. Vision, audio, and embedding weights remain + available for the later export. +2. `olive capture-onnx-graph --use_mobius_builder` loads that quantized + directory and exports the four-component ORT GenAI package. + +## Prerequisites + +This recipe requires Olive multi-build support and the current Mobius +component/quantized-checkpoint integration. Until those changes are included in +published releases, install the tested source revisions and runtime +dependencies: + +```bash +pip install "git+https://github.com/microsoft/Olive.git@4081c1bb" +pip install "git+https://github.com/onnxruntime/mobius.git@b9b4ef4" +pip install transformers torch onnxruntime-genai requests +``` + +Gemma 4 is gated. Accept the model license, then authenticate after installing +`huggingface_hub` through the dependencies above: + +```bash +hf auth login +``` + +Run the commands below from this `multi_comp` directory. + +## Step 1 — Quantize the decoder + +```bash +olive run --config gemma4_quantize_then_export.json +``` + +The build selects only the decoder: + +```json +{ + "components": ["decoder"], + "pipeline": ["decoder_rtn"] +} +``` + +`Rtn` performs calibration-free INT4 weight quantization with group size 128. +The embedding table, LM head, and Gemma 4's runtime-specific +`per_layer_input_gate` / `per_layer_projection` modules remain floating point; +Mobius currently represents those two modules as ordinary Linear operators. +Olive saves a complete Hugging Face checkpoint, not a standalone decoder: + +```text +gemma4_decoder_int4_hf/ + model/ + config.json + generation_config.json + model*.safetensors + tokenizer and processor files + model_config.json + footprint.json +``` + +The complete directory is required because Mobius still needs the unquantized +vision encoder, audio encoder, and multimodal embedding components. + +## Step 2 — Export all components with Mobius + +```bash +olive capture-onnx-graph \ + --model_name_or_path gemma4_decoder_int4_hf/model \ + --use_mobius_builder \ + --trust_remote_code \ + --precision fp32 \ + --output_path exported_gemma4_int4_pkg +``` + +Mobius preserves the Olive-packed INT4 decoder weights and exports: + +```text +exported_gemma4_int4_pkg/ + decoder/model.onnx + vision_encoder/model.onnx + audio_encoder/model.onnx + embedding/model.onnx + genai_config.json + tokenizer.json + processor and audio feature-extraction files +``` + +## Step 3 — Inference + +Use the inference entry point in the parent Gemma 4 recipe: + +```bash +python ../inference.py \ + --model-path exported_gemma4_int4_pkg \ + --prompt "What is the capital of France?" \ + --verbose +``` + +For CUDA inference, install `onnxruntime-genai-cuda` and change the export +precision to `fp16` on a CUDA-capable machine. The RTN stage itself can run on +CPU or CUDA. + +## Notes + +- `builds.components: ["decoder"]` scopes RTN to the language decoder while + preserving the full Hugging Face checkpoint layout. +- `lm_head: false` and `embeds: false` avoid quantizing the tied/output tables. +- `modules_to_not_convert` keeps Gemma 4's per-layer input gate/projection in + the floating-point format expected by the current Mobius graph. +- This is intentionally a quantize-then-export flow. The existing sibling + recipes under `cpu/` and `cuda/` demonstrate export-then-ONNX-quantize flows. +- The model download and full quantization require substantial disk and memory. diff --git a/google-gemma-4-E2B-it/multi_comp/gemma4_quantize_then_export.json b/google-gemma-4-E2B-it/multi_comp/gemma4_quantize_then_export.json new file mode 100644 index 000000000..240ee60ce --- /dev/null +++ b/google-gemma-4-E2B-it/multi_comp/gemma4_quantize_then_export.json @@ -0,0 +1,35 @@ +{ + "input_model": { + "type": "HfModel", + "config": { + "model_path": "google/gemma-4-E2B-it", + "task": "image-text-to-text", + "load_kwargs": { "torch_dtype": "float16", "trust_remote_code": true } + } + }, + "systems": { + "local_cpu": { + "type": "LocalSystem", + "accelerators": [ { "device": "cpu", "execution_providers": [ "CPUExecutionProvider" ] } ] + } + }, + "passes": { + "decoder_rtn": { + "type": "Rtn", + "bits": 4, + "group_size": 128, + "sym": true, + "lm_head": false, + "embeds": false, + "modules_to_not_convert": [ "per_layer_input_gate", "per_layer_projection" ] + } + }, + "engine": { "host": "local_cpu", "target": "local_cpu", "evaluate_input_model": false, "cache_dir": "cache" }, + "builds": { + "decoder_int4": { + "components": [ "decoder" ], + "pipeline": [ "decoder_rtn" ], + "output_dir": "gemma4_decoder_int4_hf" + } + } +} diff --git a/google-gemma-4-E2B-it/multi_comp/info.yml b/google-gemma-4-E2B-it/multi_comp/info.yml new file mode 100644 index 000000000..48de183c9 --- /dev/null +++ b/google-gemma-4-E2B-it/multi_comp/info.yml @@ -0,0 +1,15 @@ +keywords: + - olive-ai + - gemma4 + - multimodal + - multi-component + - rtn + - int4 + - mobius +recipes: + - name: gemma4-e2b-decoder-rtn-then-mobius + file: gemma4_quantize_then_export.json + eps: + - CPUExecutionProvider + devices: + - cpu diff --git a/google-gemma-4-E2B-it/requirements.txt b/google-gemma-4-E2B-it/requirements.txt index a72e46373..5fb784554 100644 --- a/google-gemma-4-E2B-it/requirements.txt +++ b/google-gemma-4-E2B-it/requirements.txt @@ -1,3 +1,4 @@ lm-eval -mobius-ai +mobius-onnx @ git+https://github.com/onnxruntime/mobius.git@b9b4ef4 olive-ai[gpu] +requests From 4c9e01b1aa2f566517e7043b4936b9542aa5171c Mon Sep 17 00:00:00 2001 From: Xiaoyu Zhang Date: Tue, 25 Aug 2026 18:30:38 -0700 Subject: [PATCH 03/11] Quantize Gemma4 decoder and vision together --- google-gemma-4-E2B-it/README.md | 6 +- google-gemma-4-E2B-it/inference.py | 74 +++++++++++++++---- google-gemma-4-E2B-it/multi_comp/.gitignore | 2 + google-gemma-4-E2B-it/multi_comp/README.md | 63 ++++++++++------ .../gemma4_quantize_then_export.json | 11 +-- google-gemma-4-E2B-it/multi_comp/info.yml | 2 +- 6 files changed, 110 insertions(+), 48 deletions(-) diff --git a/google-gemma-4-E2B-it/README.md b/google-gemma-4-E2B-it/README.md index 8349e62d2..35bf6cea4 100644 --- a/google-gemma-4-E2B-it/README.md +++ b/google-gemma-4-E2B-it/README.md @@ -44,9 +44,9 @@ K-Quant (Q4_K_M) is significantly faster with GPU acceleration — install `cupy-cuda12x` for a 19–51× speedup during quantization. For a Torch-stage quantize-then-export flow, see -[`multi_comp/README.md`](multi_comp/README.md). It applies INT4 RTN only to -the decoder, saves a complete Hugging Face checkpoint, and then exports all -four components with Mobius. +[`multi_comp/README.md`](multi_comp/README.md). It applies INT4 RTN to the +decoder and vision encoder in one multi-component build, saves a complete +Hugging Face checkpoint, and then exports all four components with Mobius. ## Build diff --git a/google-gemma-4-E2B-it/inference.py b/google-gemma-4-E2B-it/inference.py index fbc7695fe..ef2e66eac 100644 --- a/google-gemma-4-E2B-it/inference.py +++ b/google-gemma-4-E2B-it/inference.py @@ -1,9 +1,10 @@ """ONNX Runtime GenAI inference for Gemma 4 models. -Supports text-only inference with chat template formatting. +Supports text and image inference with chat template formatting. Usage: python inference.py --prompt "What is the capital of France?" + python inference.py --image photo.jpg --prompt "Describe this image." python inference.py --device gpu --variant int4 --prompt "Explain quantum computing" python inference.py --interactive python inference.py --model-path /path/to/models --prompt "Hello" @@ -27,7 +28,12 @@ def resolve_model_path(device: str, variant: str | None) -> str: return f"cuda/{variant}/models" -def format_chat_prompt(tokenizer, prompt: str, system_prompt: str | None = None) -> str: +def format_chat_prompt( + tokenizer, + prompt: str, + system_prompt: str | None = None, + has_image: bool = False, +) -> str: """Format a prompt using Gemma4's chat template. Gemma4 instruction-tuned models require chat formatting for best results. @@ -36,7 +42,10 @@ def format_chat_prompt(tokenizer, prompt: str, system_prompt: str | None = None) messages = [] if system_prompt: messages.append({"role": "system", "content": system_prompt}) - messages.append({"role": "user", "content": prompt}) + content = ( + [{"type": "image"}, {"type": "text", "text": prompt}] if has_image else prompt + ) + messages.append({"role": "user", "content": content}) # ORT GenAI tokenizer expects the messages as a JSON string. Older # releases cannot execute Gemma 4's newer Jinja template, so fall back to @@ -50,7 +59,8 @@ def format_chat_prompt(tokenizer, prompt: str, system_prompt: str | None = None) turns = [""] if system_prompt: turns.append(f"<|turn>system\n{system_prompt.strip()}\n") - turns.append(f"<|turn>user\n{prompt.strip()}\n<|turn>model\n") + image_token = "<|image|>" if has_image else "" + turns.append(f"<|turn>user\n{image_token}{prompt.strip()}\n<|turn>model\n") return "".join(turns) @@ -60,26 +70,41 @@ def generate( prompt: str, max_length: int = 2048, system_prompt: str | None = None, + image_path: str | None = None, verbose: bool = False, ) -> str: """Generate text from a prompt.""" - formatted = format_chat_prompt(tokenizer, prompt, system_prompt) - input_ids = tokenizer.encode(formatted) + formatted = format_chat_prompt( + tokenizer, + prompt, + system_prompt, + has_image=image_path is not None, + ) + input_ids = None + multimodal_inputs = None + if image_path is not None: + processor = model.create_multimodal_processor() + images = og.Images.open(image_path) + multimodal_inputs = processor(formatted, images=images) + else: + input_ids = tokenizer.encode(formatted) - if verbose: + if verbose and input_ids is not None: print(f" Input tokens: {len(input_ids)}") params = og.GeneratorParams(model) params.set_search_options( max_length=max_length, - past_present_share_buffer=False, do_sample=False, top_k=1, ) start = time.time() generator = og.Generator(model, params) - generator.append_tokens([input_ids]) + if multimodal_inputs is not None: + generator.set_inputs(multimodal_inputs) + else: + generator.append_tokens([input_ids]) output_tokens = [] tokenizer_stream = tokenizer.create_stream() @@ -99,7 +124,9 @@ def generate( if verbose: print() tps = len(output_tokens) / elapsed if elapsed > 0 else 0 - print(f" Output tokens: {len(output_tokens)}, Time: {elapsed:.2f}s, Speed: {tps:.1f} tok/s") + print( + f" Output tokens: {len(output_tokens)}, Time: {elapsed:.2f}s, Speed: {tps:.1f} tok/s" + ) return output_text @@ -131,16 +158,28 @@ def main(): parser.add_argument("--model-path", default=None, help="Override model directory") parser.add_argument("--prompt", type=str, default=None, help="Text prompt") parser.add_argument("--system-prompt", type=str, default=None, help="System prompt") - parser.add_argument("--max-length", type=int, default=2048, help="Max generation length") + parser.add_argument( + "--image", type=str, default=None, help="Path to an input image" + ) + parser.add_argument( + "--max-length", type=int, default=2048, help="Max generation length" + ) parser.add_argument("--interactive", action="store_true", help="Interactive mode") - parser.add_argument("--verbose", action="store_true", help="Show token-by-token output") + parser.add_argument( + "--verbose", action="store_true", help="Show token-by-token output" + ) args = parser.parse_args() model_path = args.model_path or resolve_model_path(args.device, args.variant) if not Path(model_path).exists(): print(f"ERROR: Model directory not found: {model_path}") - print("Run `olive run --config //config.json` first to generate the models.") + print( + "Run `olive run --config //config.json` first to generate the models." + ) + sys.exit(1) + if args.image and not Path(args.image).is_file(): + print(f"ERROR: Image file not found: {args.image}") sys.exit(1) print(f"Loading model from {model_path}...") @@ -154,9 +193,12 @@ def main(): interactive_mode(model, tokenizer, args.max_length) elif args.prompt: response = generate( - model, tokenizer, args.prompt, + model, + tokenizer, + args.prompt, max_length=args.max_length, system_prompt=args.system_prompt, + image_path=args.image, verbose=args.verbose, ) if not args.verbose: @@ -167,7 +209,9 @@ def main(): print(f"Demo prompt: {demo_prompt}") print() generate( - model, tokenizer, demo_prompt, + model, + tokenizer, + demo_prompt, max_length=args.max_length, verbose=True, ) diff --git a/google-gemma-4-E2B-it/multi_comp/.gitignore b/google-gemma-4-E2B-it/multi_comp/.gitignore index 3ed0b3c9c..551fc302d 100644 --- a/google-gemma-4-E2B-it/multi_comp/.gitignore +++ b/google-gemma-4-E2B-it/multi_comp/.gitignore @@ -1,6 +1,8 @@ # Generated Hugging Face and ONNX packages gemma4_decoder_int4_hf/ +gemma4_decoder_vision_int4_hf/ exported_gemma4_int4_pkg/ +exported_gemma4_decoder_vision_int4_pkg/ # Olive caches and generated dependency lists .olive-cache/ diff --git a/google-gemma-4-E2B-it/multi_comp/README.md b/google-gemma-4-E2B-it/multi_comp/README.md index 6bcdd467f..8469b0acf 100644 --- a/google-gemma-4-E2B-it/multi_comp/README.md +++ b/google-gemma-4-E2B-it/multi_comp/README.md @@ -1,14 +1,15 @@ # Gemma 4 E2B — Quantize Then Export -This recipe quantizes the Torch decoder component of +This recipe quantizes the Torch decoder and vision components of [`google/gemma-4-E2B-it`](https://huggingface.co/google/gemma-4-E2B-it) before exporting the complete multimodal model with Mobius. The flow has two explicit stages: -1. Olive selects the Gemma 4 `decoder` component, applies INT4 RTN, and saves a - complete Hugging Face directory. Vision, audio, and embedding weights remain - available for the later export. +1. Olive selects the Gemma 4 `decoder` and `vision_encoder` components in one + build, applies INT4 RTN to both, and saves one complete Hugging Face + directory. Audio and embedding weights remain available for the later + export. 2. `olive capture-onnx-graph --use_mobius_builder` loads that quantized directory and exports the four-component ORT GenAI package. @@ -20,8 +21,8 @@ published releases, install the tested source revisions and runtime dependencies: ```bash -pip install "git+https://github.com/microsoft/Olive.git@4081c1bb" -pip install "git+https://github.com/onnxruntime/mobius.git@b9b4ef4" +pip install "git+https://github.com/microsoft/Olive.git@faa15641" +pip install "git+https://github.com/onnxruntime/mobius.git@d048028" pip install transformers torch onnxruntime-genai requests ``` @@ -34,29 +35,31 @@ hf auth login Run the commands below from this `multi_comp` directory. -## Step 1 — Quantize the decoder +## Step 1 — Quantize the decoder and vision encoder ```bash olive run --config gemma4_quantize_then_export.json ``` -The build selects only the decoder: +The build selects both components so the two sets of packed weights are saved +in the same Hugging Face checkpoint: ```json { - "components": ["decoder"], - "pipeline": ["decoder_rtn"] + "components": ["decoder", "vision_encoder"], + "pipeline": ["decoder_vision_rtn"] } ``` `Rtn` performs calibration-free INT4 weight quantization with group size 128. -The embedding table, LM head, and Gemma 4's runtime-specific -`per_layer_input_gate` / `per_layer_projection` modules remain floating point; -Mobius currently represents those two modules as ordinary Linear operators. -Olive saves a complete Hugging Face checkpoint, not a standalone decoder: +`quantize_vision: true` includes the vision tower and its vision-to-text +projector. The embedding table, LM head, audio encoder, and Gemma 4's +runtime-specific `per_layer_input_gate` / `per_layer_projection` modules remain +floating point. Olive saves a complete Hugging Face checkpoint, not standalone +component fragments: ```text -gemma4_decoder_int4_hf/ +gemma4_decoder_vision_int4_hf/ model/ config.json generation_config.json @@ -66,24 +69,24 @@ gemma4_decoder_int4_hf/ footprint.json ``` -The complete directory is required because Mobius still needs the unquantized -vision encoder, audio encoder, and multimodal embedding components. +The complete directory lets Mobius load the decoder and vision INT4 sidecars +together with the unquantized audio and multimodal embedding components. ## Step 2 — Export all components with Mobius ```bash olive capture-onnx-graph \ - --model_name_or_path gemma4_decoder_int4_hf/model \ + --model_name_or_path gemma4_decoder_vision_int4_hf/model \ --use_mobius_builder \ --trust_remote_code \ --precision fp32 \ - --output_path exported_gemma4_int4_pkg + --output_path exported_gemma4_decoder_vision_int4_pkg ``` -Mobius preserves the Olive-packed INT4 decoder weights and exports: +Mobius preserves the Olive-packed INT4 decoder and vision weights and exports: ```text -exported_gemma4_int4_pkg/ +exported_gemma4_decoder_vision_int4_pkg/ decoder/model.onnx vision_encoder/model.onnx audio_encoder/model.onnx @@ -99,19 +102,31 @@ Use the inference entry point in the parent Gemma 4 recipe: ```bash python ../inference.py \ - --model-path exported_gemma4_int4_pkg \ + --model-path exported_gemma4_decoder_vision_int4_pkg \ --prompt "What is the capital of France?" \ --verbose ``` +To execute the quantized vision encoder, provide an image: + +```bash +python ../inference.py \ + --model-path exported_gemma4_decoder_vision_int4_pkg \ + --image path/to/image.jpg \ + --prompt "Describe this image." \ + --verbose +``` + For CUDA inference, install `onnxruntime-genai-cuda` and change the export precision to `fp16` on a CUDA-capable machine. The RTN stage itself can run on CPU or CUDA. ## Notes -- `builds.components: ["decoder"]` scopes RTN to the language decoder while - preserving the full Hugging Face checkpoint layout. +- `builds.components: ["decoder", "vision_encoder"]` scopes one RTN pass to + both selected subtrees while preserving the full Hugging Face checkpoint. +- `quantize_vision: true` quantizes the vision tower and vision-to-text + projector instead of applying RTN to the decoder only. - `lm_head: false` and `embeds: false` avoid quantizing the tied/output tables. - `modules_to_not_convert` keeps Gemma 4's per-layer input gate/projection in the floating-point format expected by the current Mobius graph. diff --git a/google-gemma-4-E2B-it/multi_comp/gemma4_quantize_then_export.json b/google-gemma-4-E2B-it/multi_comp/gemma4_quantize_then_export.json index 240ee60ce..dc5fc603c 100644 --- a/google-gemma-4-E2B-it/multi_comp/gemma4_quantize_then_export.json +++ b/google-gemma-4-E2B-it/multi_comp/gemma4_quantize_then_export.json @@ -14,22 +14,23 @@ } }, "passes": { - "decoder_rtn": { + "decoder_vision_rtn": { "type": "Rtn", "bits": 4, "group_size": 128, "sym": true, "lm_head": false, "embeds": false, + "quantize_vision": true, "modules_to_not_convert": [ "per_layer_input_gate", "per_layer_projection" ] } }, "engine": { "host": "local_cpu", "target": "local_cpu", "evaluate_input_model": false, "cache_dir": "cache" }, "builds": { - "decoder_int4": { - "components": [ "decoder" ], - "pipeline": [ "decoder_rtn" ], - "output_dir": "gemma4_decoder_int4_hf" + "decoder_vision_int4": { + "components": [ "decoder", "vision_encoder" ], + "pipeline": [ "decoder_vision_rtn" ], + "output_dir": "gemma4_decoder_vision_int4_hf" } } } diff --git a/google-gemma-4-E2B-it/multi_comp/info.yml b/google-gemma-4-E2B-it/multi_comp/info.yml index 48de183c9..c45d077f4 100644 --- a/google-gemma-4-E2B-it/multi_comp/info.yml +++ b/google-gemma-4-E2B-it/multi_comp/info.yml @@ -7,7 +7,7 @@ keywords: - int4 - mobius recipes: - - name: gemma4-e2b-decoder-rtn-then-mobius + - name: gemma4-e2b-decoder-vision-rtn-then-mobius file: gemma4_quantize_then_export.json eps: - CPUExecutionProvider From 4f2f9025461823276e909ef6b6096b1f16662ba8 Mon Sep 17 00:00:00 2001 From: Xiaoyu Zhang Date: Tue, 25 Aug 2026 18:38:19 -0700 Subject: [PATCH 04/11] Update Gemma4 Olive revision --- google-gemma-4-E2B-it/multi_comp/README.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/google-gemma-4-E2B-it/multi_comp/README.md b/google-gemma-4-E2B-it/multi_comp/README.md index 8469b0acf..c72e76083 100644 --- a/google-gemma-4-E2B-it/multi_comp/README.md +++ b/google-gemma-4-E2B-it/multi_comp/README.md @@ -21,7 +21,7 @@ published releases, install the tested source revisions and runtime dependencies: ```bash -pip install "git+https://github.com/microsoft/Olive.git@faa15641" +pip install "git+https://github.com/microsoft/Olive.git@6e2fe601" pip install "git+https://github.com/onnxruntime/mobius.git@d048028" pip install transformers torch onnxruntime-genai requests ``` From 933a4625895735445cc2eff038cc5026de0d2e7a Mon Sep 17 00:00:00 2001 From: Xiaoyu Zhang Date: Tue, 25 Aug 2026 19:46:45 -0700 Subject: [PATCH 05/11] Consolidate Qwen2.5 VL component builds --- .../builtin/.gitignore | 3 + Qwen-Qwen2.5-VL-3B-Instruct/builtin/README.md | 184 +- .../builtin/codes/__init__.py | 0 .../builtin/codes/modeling_qwen2_5_vl.py | 1888 ----------------- .../builtin/cpu_and_mobile/config.json | 67 + .../builtin/cpu_and_mobile/embedding.json | 48 - .../builtin/cpu_and_mobile/text.json | 11 - .../builtin/cpu_and_mobile/vision.json | 71 - .../builtin/cuda/config.json | 77 + .../builtin/cuda/embedding.json | 58 - .../builtin/cuda/text.json | 17 - .../builtin/cuda/vision.json | 98 - .../builtin/optimize.py | 277 ++- .../builtin/requirements.txt | 3 +- .../builtin/user_script.py | 166 -- 15 files changed, 386 insertions(+), 2582 deletions(-) delete mode 100644 Qwen-Qwen2.5-VL-3B-Instruct/builtin/codes/__init__.py delete mode 100644 Qwen-Qwen2.5-VL-3B-Instruct/builtin/codes/modeling_qwen2_5_vl.py create mode 100644 Qwen-Qwen2.5-VL-3B-Instruct/builtin/cpu_and_mobile/config.json delete mode 100644 Qwen-Qwen2.5-VL-3B-Instruct/builtin/cpu_and_mobile/embedding.json delete mode 100644 Qwen-Qwen2.5-VL-3B-Instruct/builtin/cpu_and_mobile/text.json delete mode 100644 Qwen-Qwen2.5-VL-3B-Instruct/builtin/cpu_and_mobile/vision.json create mode 100644 Qwen-Qwen2.5-VL-3B-Instruct/builtin/cuda/config.json delete mode 100644 Qwen-Qwen2.5-VL-3B-Instruct/builtin/cuda/embedding.json delete mode 100644 Qwen-Qwen2.5-VL-3B-Instruct/builtin/cuda/text.json delete mode 100644 Qwen-Qwen2.5-VL-3B-Instruct/builtin/cuda/vision.json delete mode 100644 Qwen-Qwen2.5-VL-3B-Instruct/builtin/user_script.py diff --git a/Qwen-Qwen2.5-VL-3B-Instruct/builtin/.gitignore b/Qwen-Qwen2.5-VL-3B-Instruct/builtin/.gitignore index ec6a5ba42..686569bdb 100644 --- a/Qwen-Qwen2.5-VL-3B-Instruct/builtin/.gitignore +++ b/Qwen-Qwen2.5-VL-3B-Instruct/builtin/.gitignore @@ -1,5 +1,6 @@ # Generated model artifacts models/ +mobius_base/ # Python bytecode __pycache__/ @@ -7,6 +8,8 @@ __pycache__/ # Olive cache .olive-cache/ +cache/ +mobius_cache/ # Temp and log files *.temp diff --git a/Qwen-Qwen2.5-VL-3B-Instruct/builtin/README.md b/Qwen-Qwen2.5-VL-3B-Instruct/builtin/README.md index b61ddfe1c..78c1419e3 100644 --- a/Qwen-Qwen2.5-VL-3B-Instruct/builtin/README.md +++ b/Qwen-Qwen2.5-VL-3B-Instruct/builtin/README.md @@ -1,8 +1,17 @@ -# Qwen2.5-VL-3B-Instruct ONNX Runtime GenAI Example +# Qwen2.5-VL-3B-Instruct — Olive + Mobius Multi-Component Recipe -This example demonstrates how to convert [Qwen2.5-VL-3B-Instruct](https://huggingface.co/Qwen/Qwen2.5-VL-3B-Instruct) vision-language model to ONNX format using Olive and run inference with ONNX Runtime GenAI. +This recipe exports +[`Qwen/Qwen2.5-VL-3B-Instruct`](https://huggingface.co/Qwen/Qwen2.5-VL-3B-Instruct) +with Olive's `MobiusBuilder`, then optimizes all three ONNX components from one +Olive multi-build config: -The pipeline exports three sub-models (vision encoder, text embedding, text decoder), applies graph optimizations (Cast chain elimination, Gemm→MatMul conversion), and quantizes all three sub-models to INT4. +- `decoder` +- `vision_encoder` +- `embedding` + +Mobius owns the model graph, weight mapping, ORT GenAI configuration, tokenizer, +and image processor generation. The previous custom PyTorch model and three +independent component configs are no longer required. ## Prerequisites @@ -10,109 +19,122 @@ The pipeline exports three sub-models (vision encoder, text embedding, text deco pip install -r requirements.txt ``` -Install ONNX Runtime GenAI based on your target device: +Install ONNX Runtime GenAI for the target: -| Device | Install Command | -|--------|-----------------| -| GPU (CUDA) | `pip install onnxruntime-genai-cuda --index-url https://aiinfra.pkgs.visualstudio.com/PublicPackages/_packaging/ORT-Nightly/pypi/simple` | -| CPU | `pip install onnxruntime-genai --index-url https://aiinfra.pkgs.visualstudio.com/PublicPackages/_packaging/ORT-Nightly/pypi/simple` | +| Target | Install command | +|---|---| +| CPU | `pip install onnxruntime-genai` | +| CUDA | `pip install onnxruntime-genai-cuda` | -## Steps +Run commands from this `builtin` directory. -### 1. Export & Optimize Models (CPU) +## Export and optimize -All graph transformations and quantization are declared in the JSON config files inside `cpu_and_mobile/` and `cuda/`. The top-level `optimize.py` script orchestrates the three Olive runs and generates the GenAI runtime configs. +### CPU and mobile -| Command | Description | -|---------|-------------| -| `python optimize.py --config-dir cpu_and_mobile --device cpu` | Full pipeline: export, optimize, INT4 quantize (CPU) | -| `python optimize.py --config-dir cuda --device gpu` | Full pipeline with FP16 + INT4 (CUDA) | -| `python optimize.py --config-dir cpu_and_mobile --skip-export` | Regenerate configs only (models already exported) | +```bash +python optimize.py --config-dir cpu_and_mobile --device cpu +``` -> **Note:** The text model is always exported as INT4 via ModelBuilder. The vision encoder is graph-optimized and quantized to INT4 by Olive passes. The embedding model's Gather-based embedding table is quantized to INT4 using GatherBlockQuantized. -> -> The vision encoder is exported for a single image using the Dynamo exporter. At runtime, ONNX Runtime GenAI handles multiple images by calling the vision encoder once per image and concatenating the results — so there is no upper bound on the number of images passed to the model. +`cpu_and_mobile/config.json` is one Olive config containing three named builds. +All three components use block-wise INT4 RTN: -### 2. Run Inference +| Build | Pipeline | +|---|---| +| `decoder` | `OnnxBlockWiseRtnQuantization` | +| `vision_encoder` | `OnnxBlockWiseRtnQuantization` | +| `embedding` | `OnnxBlockWiseRtnQuantization` | -From the top-level model directory: +### CUDA ```bash -# Text-only (CPU models, default) -python inference.py --prompt "What is the capital of France?" +python optimize.py --config-dir cuda --device gpu +``` -# With a single image -python inference.py --prompt "Describe this image" --image cat.jpeg +`cuda/config.json` preserves the previous target intent: -# CUDA models -python inference.py --model_path cuda/models --prompt "Describe this image" --image cat.jpeg +| Build | Pipeline | +|---|---| +| `decoder` | INT4 RTN | +| `vision_encoder` | Mobius FP16, Olive resave | +| `embedding` | Mobius FP16, Olive resave | -# Interactive mode -python inference.py --interactive -``` +Both flows have two stages: + +1. Olive runs `MobiusBuilder` once and saves the complete package under + `/mobius_base/`. +2. Olive runs the target's single `config.json`; its `builds` select and + optimize the three Mobius components into `/models/`. -**Multi-image inference** is supported via `model-mm.py` from the `onnxruntime-genai` examples: +To reuse an existing Mobius export while rerunning the three component builds: ```bash -# Two images — compare or reason across multiple images -# Adjust paths to your onnxruntime-genai checkout and model directory -python /examples/python/model-mm.py \ - -m /cpu_and_mobile/models \ - -up "Are these two images the same?" \ - --image_paths image1.jpeg image2.jpeg \ - --non_interactive +python optimize.py --config-dir cpu_and_mobile --device cpu --skip-export ``` -## Evaluation +The final ORT GenAI package uses Mobius's native component layout: + +```text +cpu_and_mobile/models/ + decoder/model.onnx + vision_encoder/model.onnx + embedding/model.onnx + genai_config.json + processor_config.json + tokenizer.json + tokenizer_config.json +``` -`eval.py` measures model quality on [AI2D](https://huggingface.co/datasets/lmms-lab/ai2d) — a multiple-choice visual QA benchmark on scientific diagrams. It supports side-by-side comparison of the quantized ONNX model against the PyTorch FP32 baseline. +## Inference ```bash -# ONNX only (fastest) -python eval.py --num_samples 100 - -# ONNX + PyTorch comparison -python eval.py --num_samples 100 --pytorch_model Qwen/Qwen2.5-VL-3B-Instruct - -# Evaluate CUDA models -python eval.py --model_path cuda/models --num_samples 100 +# Text only +python inference.py \ + --model_path cpu_and_mobile/models \ + --prompt "What is the capital of France?" + +# Image + text +python inference.py \ + --model_path cpu_and_mobile/models \ + --image cat.jpeg \ + --prompt "Describe this image." + +# CUDA package +python inference.py \ + --model_path cuda/models \ + --image cat.jpeg \ + --prompt "Describe this image." ``` -### Results (AI2D, 100 samples) - -| Model | Accuracy | Avg latency | -|-------|----------|-------------| -| PyTorch FP32 (baseline) | 81.00% | 11.87 s/sample | -| **ONNX INT4 (CPU)** | **82.00%** | **9.52 s/sample** | -| **ONNX FP16 (CUDA)** | **85.00%** | **0.31 s/sample** | -| Random chance | 25.00% | — | - -- **CPU INT4 accuracy delta: −1 pp** (81% → 82%) -- **CUDA FP16 accuracy delta: +4 pp** (81% → 85%) -- **CPU speedup: 1.25×** vs PyTorch FP32 +ORT GenAI executes the vision encoder only when images are present, fuses its +features in the embedding component, and runs autoregressive generation through +the decoder. -> Results measured with `--num_samples 100` from the AI2D test split. +## Evaluation -## Directory Structure +`eval.py` evaluates the final package on AI2D: +```bash +python eval.py --model_path cpu_and_mobile/models --num_samples 100 +python eval.py \ + --model_path cpu_and_mobile/models \ + --num_samples 100 \ + --pytorch_model Qwen/Qwen2.5-VL-3B-Instruct ``` -Qwen-Qwen2.5-VL-3B-Instruct/ -├── LICENSE -└── builtin/ - ├── optimize.py # End-to-end Olive pipeline + GenAI config generation - ├── user_script.py # Olive callbacks: model loading, dummy inputs, IO configs - ├── eval.py # AI2D accuracy evaluation (ONNX vs PyTorch) - ├── inference.py # ONNX Runtime GenAI inference - ├── cat.jpeg # Sample test image - ├── codes/ # Custom Qwen2.5-VL PyTorch model adapted for ONNX export - ├── cpu_and_mobile/ - │ ├── embedding.json # Olive config: export → optimize → INT4 - │ ├── vision.json # Olive config: Dynamo export → graph surgeries → INT4 - │ ├── text.json # Olive config: ModelBuilder INT4 - │ └── models/ # Exported ONNX models (generated) - └── cuda/ - ├── embedding.json # Olive config with FP16 + INT4 + CUDA EP - ├── vision.json # Olive config with FP16 + INT4 + CUDA EP - ├── text.json # ModelBuilder INT4 with CUDA EP - └── models/ # Exported CUDA ONNX models (generated) + +Re-run evaluation when changing Mobius, quantization settings, or runtime +versions; results from the previous custom export graph are not comparable. + +## Directory structure + +```text +builtin/ + optimize.py + inference.py + eval.py + cat.jpeg + cpu_and_mobile/ + config.json + cuda/ + config.json ``` diff --git a/Qwen-Qwen2.5-VL-3B-Instruct/builtin/codes/__init__.py b/Qwen-Qwen2.5-VL-3B-Instruct/builtin/codes/__init__.py deleted file mode 100644 index e69de29bb..000000000 diff --git a/Qwen-Qwen2.5-VL-3B-Instruct/builtin/codes/modeling_qwen2_5_vl.py b/Qwen-Qwen2.5-VL-3B-Instruct/builtin/codes/modeling_qwen2_5_vl.py deleted file mode 100644 index 5a107662c..000000000 --- a/Qwen-Qwen2.5-VL-3B-Instruct/builtin/codes/modeling_qwen2_5_vl.py +++ /dev/null @@ -1,1888 +0,0 @@ -# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨 -# This file was automatically generated from src/transformers/models/qwen2_5_vl/modular_qwen2_5_vl.py. -# Do NOT edit this file manually as any edits will be overwritten by the generation of -# the file from the modular. If any change should be done, please apply the change to the -# modular_qwen2_5_vl.py file directly. One of our CI enforces this. -# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨 -# coding=utf-8 -# Copyright 2025 The Qwen Team and The HuggingFace Inc. team. All rights reserved. -# -# This code is based on EleutherAI's GPT-NeoX library and the GPT-NeoX -# and OPT implementations in this library. It has been modified from its -# original forms to accommodate minor architectural differences compared -# to GPT-NeoX and OPT used by the Meta AI team that trained the model. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -from collections.abc import Callable -from dataclasses import dataclass -from typing import Any, Optional, Union - -import torch -import torch.nn as nn -import torch.nn.functional as F - -from transformers.activations import ACT2FN -from transformers.cache_utils import Cache, DynamicCache -from transformers.generation import GenerationMixin -from transformers.masking_utils import create_causal_mask, create_sliding_window_causal_mask -from transformers.modeling_flash_attention_utils import FlashAttentionKwargs -from transformers.modeling_layers import GradientCheckpointingLayer -from transformers.modeling_outputs import BaseModelOutputWithPast, ModelOutput -from transformers.modeling_rope_utils import ROPE_INIT_FUNCTIONS -from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel -from transformers.processing_utils import Unpack -from transformers.utils import TransformersKwargs, auto_docstring, can_return_tuple, logging -from transformers.models.qwen2.modeling_qwen2 import Qwen2RMSNorm -from transformers.models.qwen2_5_vl.configuration_qwen2_5_vl import Qwen2_5_VLConfig, Qwen2_5_VLTextConfig, Qwen2_5_VLVisionConfig - - -logger = logging.get_logger(__name__) - - -def _get_rope_parameters(config) -> dict: - """Build a rope_parameters-compatible dict. - - Handles both the old format (config.rope_parameters dict) and the newer - transformers 4.57+ format (config.rope_theta + config.rope_scaling). - """ - if hasattr(config, "rope_parameters"): - return config.rope_parameters - rope_scaling = getattr(config, "rope_scaling", {}) or {} - return { - "rope_type": rope_scaling.get("rope_type", "default"), - "rope_theta": getattr(config, "rope_theta", 1000000.0), - "mrope_section": rope_scaling.get("mrope_section", [16, 24, 24]), - } - - -class Qwen2_5_VLMLP(nn.Module): - def __init__(self, config, bias: bool = False): - super().__init__() - self.hidden_size = config.hidden_size - self.intermediate_size = config.intermediate_size - self.gate_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=bias) - self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=bias) - self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=bias) - self.act_fn = ACT2FN[config.hidden_act] - - def forward(self, hidden_state): - return self.down_proj(self.act_fn(self.gate_proj(hidden_state)) * self.up_proj(hidden_state)) - - -class Qwen2_5_VisionPatchEmbed(nn.Module): - def __init__( - self, - patch_size: int = 14, - temporal_patch_size: int = 2, - in_channels: int = 3, - embed_dim: int = 1152, - ) -> None: - super().__init__() - self.patch_size = patch_size - self.temporal_patch_size = temporal_patch_size - self.in_channels = in_channels - self.embed_dim = embed_dim - - kernel_size = [temporal_patch_size, patch_size, patch_size] - self.proj = nn.Conv3d(in_channels, embed_dim, kernel_size=kernel_size, stride=kernel_size, bias=False) - - def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: - target_dtype = self.proj.weight.dtype - hidden_states = hidden_states.view( - -1, self.in_channels, self.temporal_patch_size, self.patch_size, self.patch_size - ) - hidden_states = self.proj(hidden_states.to(dtype=target_dtype)).view(-1, self.embed_dim) - return hidden_states - - -class Qwen2_5_VisionRotaryEmbedding(nn.Module): - inv_freq: torch.Tensor # fix linting for `register_buffer` - - def __init__(self, dim: int, theta: float = 10000.0) -> None: - super().__init__() - inv_freq = 1.0 / (theta ** (torch.arange(0, dim, 2, dtype=torch.float) / dim)) - self.register_buffer("inv_freq", inv_freq, persistent=True) - - def forward(self, seqlen: int) -> torch.Tensor: - seq = torch.arange(seqlen, device=self.inv_freq.device, dtype=self.inv_freq.dtype) - freqs = torch.outer(seq, self.inv_freq) - return freqs - - -class Qwen2_5_VLPatchMerger(nn.Module): - def __init__(self, dim: int, context_dim: int, spatial_merge_size: int = 2) -> None: - super().__init__() - self.hidden_size = context_dim * (spatial_merge_size**2) - self.ln_q = Qwen2RMSNorm(context_dim, eps=1e-6) - self.mlp = nn.Sequential( - nn.Linear(self.hidden_size, self.hidden_size), - nn.GELU(), - nn.Linear(self.hidden_size, dim), - ) - - def forward(self, x: torch.Tensor) -> torch.Tensor: - x = self.mlp(self.ln_q(x).view(-1, self.hidden_size)) - return x - - -def rotate_half(x): - """Rotates half the hidden dims of the input.""" - x1 = x[..., : x.shape[-1] // 2] - x2 = x[..., x.shape[-1] // 2 :] - return torch.cat((-x2, x1), dim=-1) - - -def apply_rotary_pos_emb_vision( - q: torch.Tensor, k: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor -) -> tuple[torch.Tensor, torch.Tensor]: - orig_q_dtype = q.dtype - orig_k_dtype = k.dtype - q, k = q.float(), k.float() - cos, sin = cos.unsqueeze(-2).float(), sin.unsqueeze(-2).float() - q_embed = (q * cos) + (rotate_half(q) * sin) - k_embed = (k * cos) + (rotate_half(k) * sin) - q_embed = q_embed.to(orig_q_dtype) - k_embed = k_embed.to(orig_k_dtype) - return q_embed, k_embed - - -def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor: - """ - This is the equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep). The hidden states go from (batch, - num_key_value_heads, seqlen, head_dim) to (batch, num_attention_heads, seqlen, head_dim) - """ - batch, num_key_value_heads, slen, head_dim = hidden_states.shape - if n_rep == 1: - return hidden_states - hidden_states = hidden_states[:, :, None, :, :].expand(batch, num_key_value_heads, n_rep, slen, head_dim) - return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim) - - -def eager_attention_forward( - module: nn.Module, - query: torch.Tensor, - key: torch.Tensor, - value: torch.Tensor, - attention_mask: Optional[torch.Tensor], - scaling: float, - dropout: float = 0.0, - **kwargs, -): - key_states = repeat_kv(key, module.num_key_value_groups) - value_states = repeat_kv(value, module.num_key_value_groups) - - attn_weights = torch.matmul(query, key_states.transpose(2, 3)) * scaling - if attention_mask is not None: - causal_mask = attention_mask[:, :, :, : key_states.shape[-2]] - attn_weights = attn_weights + causal_mask - - attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query.dtype) - attn_weights = nn.functional.dropout(attn_weights, p=dropout, training=module.training) - attn_output = torch.matmul(attn_weights, value_states) - attn_output = attn_output.transpose(1, 2).contiguous() - - return attn_output, attn_weights - - -class Qwen2_5_VLVisionAttention(nn.Module): - def __init__(self, config: Qwen2_5_VLVisionConfig) -> None: - super().__init__() - self.dim = config.hidden_size - self.num_heads = config.num_heads - self.head_dim = self.dim // self.num_heads - self.num_key_value_groups = 1 # needed for eager attention - self.qkv = nn.Linear(self.dim, self.dim * 3, bias=True) - self.proj = nn.Linear(self.dim, self.dim) - self.scaling = self.head_dim**-0.5 - self.config = config - self.attention_dropout = 0.0 - self.is_causal = False - - def forward( - self, - hidden_states: torch.Tensor, - cu_seqlens: torch.Tensor, - rotary_pos_emb: Optional[torch.Tensor] = None, - position_embeddings: Optional[tuple[torch.Tensor, torch.Tensor]] = None, - **kwargs, - ) -> torch.Tensor: - seq_length = hidden_states.shape[0] - query_states, key_states, value_states = ( - self.qkv(hidden_states).reshape(seq_length, 3, self.num_heads, -1).permute(1, 0, 2, 3).unbind(0) - ) - cos, sin = position_embeddings - query_states, key_states = apply_rotary_pos_emb_vision(query_states, key_states, cos, sin) - - query_states = query_states.transpose(0, 1).unsqueeze(0) - key_states = key_states.transpose(0, 1).unsqueeze(0) - value_states = value_states.transpose(0, 1).unsqueeze(0) - - attention_interface: Callable = eager_attention_forward - if self.config._attn_implementation != "eager": - attention_interface = ALL_ATTENTION_FUNCTIONS[self.config._attn_implementation] - - if self.config._attn_implementation == "flash_attention_2": - # Flash Attention 2: Use cu_seqlens for variable length attention - max_seqlen = (cu_seqlens[1:] - cu_seqlens[:-1]).max() - attn_output, _ = attention_interface( - self, - query_states, - key_states, - value_states, - attention_mask=None, - scaling=self.scaling, - dropout=0.0 if not self.training else self.attention_dropout, - cu_seq_lens_q=cu_seqlens, - cu_seq_lens_k=cu_seqlens, - max_length_q=max_seqlen, - max_length_k=max_seqlen, - is_causal=False, - **kwargs, - ) - elif torch.compiler.is_exporting(): - # ONNX export: emit a custom PackedAttention op node directly. - # Olive's PackedAttentionToLoopMHA surgery will replace this in the ONNX graph. - attn_output = torch.onnx.ops.symbolic( - "custom::PackedAttention", - ( - query_states, - key_states, - value_states, - cu_seqlens, - ), - dict( - scale=self.scaling, - num_heads=self.num_heads, - ), - dtype=query_states.dtype, - shape=( - query_states.shape[0], # batch_size - query_states.shape[2], # sequence_length - query_states.shape[1], # num_heads - query_states.shape[3], # head_size - ), - version=1, - ) - # ONNX symbolic outputs default to CPU; move to projection weight device - attn_output = attn_output.to(self.proj.weight.device) - else: - # Other implementations: Process each chunk separately - lengths = cu_seqlens[1:] - cu_seqlens[:-1] - splits = [ - torch.split(tensor, lengths.tolist(), dim=2) for tensor in (query_states, key_states, value_states) - ] - - attn_outputs = [ - attention_interface( - self, - q, - k, - v, - attention_mask=None, - scaling=self.scaling, - dropout=0.0 if not self.training else self.attention_dropout, - is_causal=False, - **kwargs, - )[0] - for q, k, v in zip(*splits) - ] - attn_output = torch.cat(attn_outputs, dim=1) - - attn_output = attn_output.reshape(seq_length, -1).contiguous() - attn_output = self.proj(attn_output) - return attn_output - - -class Qwen2_5_VLVisionBlock(GradientCheckpointingLayer): - def __init__(self, config, attn_implementation: str = "sdpa") -> None: - super().__init__() - self.norm1 = Qwen2RMSNorm(config.hidden_size, eps=1e-6) - self.norm2 = Qwen2RMSNorm(config.hidden_size, eps=1e-6) - self.attn = Qwen2_5_VLVisionAttention(config=config) - self.mlp = Qwen2_5_VLMLP(config, bias=True) - - def forward( - self, - hidden_states: torch.Tensor, - cu_seqlens: torch.Tensor, - rotary_pos_emb: Optional[torch.Tensor] = None, - position_embeddings: Optional[tuple[torch.Tensor, torch.Tensor]] = None, - **kwargs, - ) -> torch.Tensor: - hidden_states = hidden_states + self.attn( - self.norm1(hidden_states), - cu_seqlens=cu_seqlens, - rotary_pos_emb=rotary_pos_emb, - position_embeddings=position_embeddings, - **kwargs, - ) - hidden_states = hidden_states + self.mlp(self.norm2(hidden_states)) - return hidden_states - - -@auto_docstring -class Qwen2_5_VLPreTrainedModel(PreTrainedModel): - config: Qwen2_5_VLConfig - base_model_prefix = "model" - input_modalities = ["image", "video", "text"] - supports_gradient_checkpointing = True - _no_split_modules = ["Qwen2_5_VLDecoderLayer", "Qwen2_5_VLVisionBlock"] - _skip_keys_device_placement = "past_key_values" - _supports_flash_attn = True - _supports_sdpa = True - - _can_compile_fullgraph = True - _supports_attention_backend = True - - -class Qwen2_5_VisionTransformerPretrainedModel(Qwen2_5_VLPreTrainedModel): - config: Qwen2_5_VLVisionConfig - _no_split_modules = ["Qwen2_5_VLVisionBlock"] - - def __init__(self, config, *inputs, **kwargs) -> None: - super().__init__(config, *inputs, **kwargs) - self.spatial_merge_size = config.spatial_merge_size - self.patch_size = config.patch_size - self.fullatt_block_indexes = config.fullatt_block_indexes - self.window_size = config.window_size - self.spatial_merge_unit = self.spatial_merge_size * self.spatial_merge_size - - self.patch_embed = Qwen2_5_VisionPatchEmbed( - patch_size=config.patch_size, - temporal_patch_size=config.temporal_patch_size, - in_channels=config.in_channels, - embed_dim=config.hidden_size, - ) - - head_dim = config.hidden_size // config.num_heads - self.rotary_pos_emb = Qwen2_5_VisionRotaryEmbedding(head_dim // 2) - - self.blocks = nn.ModuleList([Qwen2_5_VLVisionBlock(config) for _ in range(config.depth)]) - self.merger = Qwen2_5_VLPatchMerger( - dim=config.out_hidden_size, - context_dim=config.hidden_size, - spatial_merge_size=config.spatial_merge_size, - ) - self.gradient_checkpointing = False - - def rot_pos_emb(self, grid_thw): - # Vectorized — no Python loop over image batch — to support a dynamic - # num_images dimension in image_grid_thw. - # - # Uniform-grid assumption: all images in one call share the same (t, h, w). - # This holds in practice because inference engines resize images to a common - # resolution before batching. Under this assumption we compute rotary - # embeddings for one representative image and tile for num_images. - merge_size = self.spatial_merge_size - - max_grid_size = grid_thw[:, 1:].max() - freq_table = self.rotary_pos_emb(max_grid_size) - device = freq_table.device - - num_images = grid_thw.shape[0] # symbolic when exporting dynamic shape - num_frames = grid_thw[0, 0] - height = grid_thw[0, 1] - width = grid_thw[0, 2] - merged_h, merged_w = height // merge_size, width // merge_size - - # Shape constraints for torch.export - torch._check(merged_h.item() >= 1) - torch._check(merged_w.item() >= 1) - torch._check(num_frames.item() >= 1) - - block_rows = torch.arange(merged_h, device=device) - block_cols = torch.arange(merged_w, device=device) - intra_row = torch.arange(merge_size, device=device) - intra_col = torch.arange(merge_size, device=device) - - row_idx = ( - block_rows[:, None, None, None] * merge_size + intra_row[None, None, :, None] - ).expand(merged_h, merged_w, merge_size, merge_size).reshape(-1) - - col_idx = ( - block_cols[None, :, None, None] * merge_size + intra_col[None, None, None, :] - ).expand(merged_h, merged_w, merge_size, merge_size).reshape(-1) - - coords = torch.stack((row_idx, col_idx), dim=-1) - # Repeat across temporal frames for one image - coords = coords.repeat(num_frames, 1) - single_emb = freq_table[coords].flatten(1) # [t*h*w, pos_dim] - - # Tile for all images (uniform grid → identical per-image embeddings) - return single_emb.repeat(num_images, 1) # [num_images*t*h*w, pos_dim] - - def get_window_index(self, grid_thw): - # Extended to support dynamic num_images (uniform-grid assumption). - # Computes permutation indices for one image, then tiles across all images. - ws = self.window_size // self.spatial_merge_size // self.patch_size - grid_t = grid_thw[0, 0] - grid_h = grid_thw[0, 1] - grid_w = grid_thw[0, 2] - llm_h = grid_h // self.spatial_merge_size - llm_w = grid_w // self.spatial_merge_size - num_images = grid_thw.shape[0] # symbolic when exporting dynamic shape - - # Shape constraints for torch.export - torch._check(grid_t.item() >= 1) - torch._check(llm_h.item() >= 1) - torch._check(llm_w.item() >= 1) - - total = grid_t * llm_h * llm_w # logical patches per image - - # Window-order permutation for a single image. - flat_idx = torch.arange(total) - rows = flat_idx // llm_w - cols = flat_idx % llm_w - num_win_w = (llm_w + ws - 1) // ws - win_ids = (rows // ws) * num_win_w + cols // ws - sort_key = win_ids * (ws * ws) + (rows % ws) * ws + (cols % ws) - window_index_single = torch.argsort(sort_key) # [total] - - # Tile for num_images: image i → window_index_single + i*total. - offsets = torch.arange(num_images) * total # [N] - window_index = (offsets.unsqueeze(1) + window_index_single.unsqueeze(0)).reshape(-1) # [N*total] - - # Per-window sizes for one image, then tiled across all images. - num_win_h = (llm_h + ws - 1) // ws - total_windows = grid_t * num_win_h * num_win_w - win_flat = torch.arange(total_windows) - wr = (win_flat % (num_win_h * num_win_w)) // num_win_w - wc = (win_flat % (num_win_h * num_win_w)) % num_win_w - valid_h = torch.clamp(llm_h - wr * ws, min=0, max=ws) - valid_w = torch.clamp(llm_w - wc * ws, min=0, max=ws) - win_sizes_single = valid_h * valid_w * self.spatial_merge_unit # [total_windows] - - # Tile win_sizes for N images and compute cumulative seqlens. - win_sizes_all = win_sizes_single.unsqueeze(0).expand(num_images, -1).reshape(-1) - cu_seqlens = F.pad(win_sizes_all.cumsum(0), (1, 0), value=0).to(torch.int32) - - return window_index, cu_seqlens - - def forward(self, hidden_states: torch.Tensor, grid_thw: torch.Tensor, **kwargs) -> torch.Tensor: - """ - Args: - hidden_states (`torch.Tensor` of shape `(seq_len, hidden_size)`): - The final hidden states of the model. - grid_thw (`torch.Tensor` of shape `(num_images_or_videos, 3)`): - The temporal, height and width of feature shape of each image in LLM. - - Returns: - `torch.Tensor`: hidden_states. - """ - # 1. Transpose to [-1, channels, temporal, patch_size, patch_size] - # Conv3D to get patch embeddings. - # eg: hidden_states (input): [14308, 1176] - # eg: hidden_states (output): [14308, 1280] - hidden_states = self.patch_embed(hidden_states) - # Calculate 2D rotary positional embeddings for each patch (14308 patches) - # eg: grid_thw: [1, 3] - rotary_pos_emb = self.rot_pos_emb(grid_thw) - # eg: rotary_pos_emb: [14308, 40] - # Calculate windows for windowed attention - window_index, cu_window_seqlens = self.get_window_index(grid_thw) - # NOTE: The rewritten get_window_index never produces zero-size windows, - # so the unique_consecutive dedup is no longer needed. - - # Constrain window_index size (u7) so torch.export knows it's nonzero. - torch._check_is_size(window_index.shape[0]) - torch._check(window_index.shape[0] > 0) - - seq_len, embed_dim = hidden_states.size() - smu = self.spatial_merge_unit - # Constrain seq_len: divisible by spatial_merge_unit, and nonzero. - torch._check(hidden_states.shape[0] > 0) - torch._check(seq_len % smu == 0) - - hidden_states = hidden_states.reshape(seq_len // smu, smu, embed_dim) - hidden_states = hidden_states[window_index, :, :] - # Use window_index.shape[0]*smu (not seq_len) to avoid Eq(s0, u7) guard - n_win = window_index.shape[0] - hidden_states = hidden_states.reshape(n_win * smu, embed_dim) - - pos_dim = rotary_pos_emb.shape[-1] - rpe_len = rotary_pos_emb.shape[0] - torch._check(rpe_len % smu == 0) - rotary_pos_emb = rotary_pos_emb.reshape(rpe_len // smu, smu, pos_dim) - rotary_pos_emb = rotary_pos_emb[window_index, :, :] - rotary_pos_emb = rotary_pos_emb.reshape(n_win * smu, pos_dim) - emb = torch.cat((rotary_pos_emb, rotary_pos_emb), dim=-1) - position_embeddings = (emb.cos(), emb.sin()) - - # cu_seqlens: boundaries of each image's patches in the packed sequence. - # Original: repeat_interleave(h*w, t) for each image — output size is data- - # dependent and can't be traced with symbolic num_images. - # For uniform grids (all images share t, h, w), this is equivalent to - # tiling [h*w] exactly t*num_images times, which avoids data-dependent shapes. - hw0 = grid_thw[0, 1] * grid_thw[0, 2] - t0 = grid_thw[0, 0] - num_images_fwd = grid_thw.shape[0] # symbolic - cu_vals = hw0.unsqueeze(0).expand(t0 * num_images_fwd) - cu_seqlens = cu_vals.cumsum( - dim=0, - dtype=grid_thw.dtype if torch.jit.is_tracing() else torch.int32, - ) - cu_seqlens = F.pad(cu_seqlens, (1, 0), value=0) - for layer_num, blk in enumerate(self.blocks): - # NOTE: Decide whether to use full attention or windowed attention - if layer_num in self.fullatt_block_indexes: - cu_seqlens_now = cu_seqlens - else: - cu_seqlens_now = cu_window_seqlens - - hidden_states = blk( - hidden_states, - cu_seqlens=cu_seqlens_now, - position_embeddings=position_embeddings, - **kwargs, - ) - - hidden_states = self.merger(hidden_states) - reverse_indices = torch.argsort(window_index) - hidden_states = hidden_states[reverse_indices, :] - - return hidden_states - - -@dataclass -@auto_docstring( - custom_intro=""" - Base class for Llava outputs, with hidden states and attentions. - """ -) -class Qwen2_5_VLModelOutputWithPast(ModelOutput): - r""" - past_key_values (`Cache`, *optional*, returned when `use_cache=True` is passed or when `config.use_cache=True`): - It is a [`~cache_utils.Cache`] instance. For more details, see our [kv cache guide](https://huggingface.co/docs/transformers/en/kv_cache). - - Contains pre-computed hidden-states (key and values in the self-attention blocks) that can be used (see - `past_key_values` input) to speed up sequential decoding. - rope_deltas (`torch.LongTensor` of shape `(batch_size, )`, *optional*): - The rope index difference between sequence length and multimodal rope. - """ - - last_hidden_state: Optional[torch.FloatTensor] = None - past_key_values: Optional[Cache] = None - hidden_states: Optional[tuple[torch.FloatTensor]] = None - attentions: Optional[tuple[torch.FloatTensor]] = None - rope_deltas: Optional[torch.LongTensor] = None - - -class Qwen2_5_VLRotaryEmbedding(nn.Module): - inv_freq: torch.Tensor # fix linting for `register_buffer` - - def __init__(self, config: Qwen2_5_VLConfig, device=None): - super().__init__() - self.max_seq_len_cached = config.max_position_embeddings - self.original_max_seq_len = config.max_position_embeddings - - self.config = config - - self.rope_type = _get_rope_parameters(self.config)["rope_type"] - rope_init_fn: Callable = self.compute_default_rope_parameters - if self.rope_type != "default": - rope_init_fn = ROPE_INIT_FUNCTIONS[self.rope_type] - inv_freq, self.attention_scaling = rope_init_fn(self.config, device) - - self.register_buffer("inv_freq", inv_freq, persistent=True) - self.original_inv_freq = inv_freq - - @staticmethod - def compute_default_rope_parameters( - config: Optional[Qwen2_5_VLConfig] = None, - device: Optional["torch.device"] = None, - seq_len: Optional[int] = None, - ) -> tuple["torch.Tensor", float]: - """ - Computes the inverse frequencies according to the original RoPE implementation - Args: - config ([`~transformers.PreTrainedConfig`]): - The model configuration. - device (`torch.device`): - The device to use for initialization of the inverse frequencies. - seq_len (`int`, *optional*): - The current sequence length. Unused for this type of RoPE. - Returns: - Tuple of (`torch.Tensor`, `float`), containing the inverse frequencies for the RoPE embeddings and the - post-processing scaling factor applied to the computed cos/sin (unused in this type of RoPE). - """ - base = _get_rope_parameters(config)["rope_theta"] - dim = getattr(config, "head_dim", None) or config.hidden_size // config.num_attention_heads - - attention_factor = 1.0 # Unused in this type of RoPE - - # Compute the inverse frequencies - inv_freq = 1.0 / ( - base ** (torch.arange(0, dim, 2, dtype=torch.int64).to(device=device, dtype=torch.float) / dim) - ) - return inv_freq, attention_factor - - # Ignore copy - def forward(self, x, position_ids): - # In contrast to other models, Qwen2_5_VL has different position ids for the grids - # So we expand the inv_freq to shape (3, ...) - inv_freq_expanded = self.inv_freq[None, None, :, None].float().expand(3, position_ids.shape[1], -1, 1) - position_ids_expanded = position_ids[:, :, None, :].float() # shape (3, bs, 1, positions) - - device_type = x.device.type if isinstance(x.device.type, str) and x.device.type != "mps" else "cpu" - with torch.autocast(device_type=device_type, enabled=False): # Force float32 - freqs = (inv_freq_expanded.float() @ position_ids_expanded.float()).transpose(2, 3) - emb = torch.cat((freqs, freqs), dim=-1) - cos = emb.cos() * self.attention_scaling - sin = emb.sin() * self.attention_scaling - - return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype) - - -class Qwen2MLP(nn.Module): - def __init__(self, config): - super().__init__() - self.config = config - self.hidden_size = config.hidden_size - self.intermediate_size = config.intermediate_size - self.gate_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False) - self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False) - self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=False) - self.act_fn = ACT2FN[config.hidden_act] - - def forward(self, x): - down_proj = self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x)) - return down_proj - - -def apply_multimodal_rotary_pos_emb(q, k, cos, sin, mrope_section, unsqueeze_dim=1): - """Applies Rotary Position Embedding with Multimodal Sections to the query and key tensors (https://qwenlm.github.io/blog/qwen2-vl/). - - Explanation: - Multimodal 3D rotary position embedding is an extension to 1D rotary position embedding. The input embedding - sequence contains vision (images / videos) embedding and text embedding or just contains text embedding. For - vision embedding part, we apply rotary position embedding on temporal, height and width dimension separately. - Here we split the channel dimension to 3 chunks for the temporal, height and width rotary position embedding. - For text embedding part, we just apply 1D rotary position embedding. The three rotary position index (temporal, - height and width) of text embedding is always the same, so the text embedding rotary position embedding has no - difference with modern LLMs. - - Args: - q (`torch.Tensor`): The query tensor. - k (`torch.Tensor`): The key tensor. - cos (`torch.Tensor`): The cosine part of the rotary embedding. - sin (`torch.Tensor`): The sine part of the rotary embedding. - position_ids (`torch.Tensor`): - The position indices of the tokens corresponding to the query and key tensors. For example, this can be - used to pass offsetted position ids when working with a KV-cache. - mrope_section(`List(int)`): - Multimodal rope section is for channel dimension of temporal, height and width in rope calculation. - unsqueeze_dim (`int`, *optional*, defaults to 1): - The 'unsqueeze_dim' argument specifies the dimension along which to unsqueeze cos[position_ids] and - sin[position_ids] so that they can be properly broadcasted to the dimensions of q and k. For example, note - that cos[position_ids] and sin[position_ids] have the shape [batch_size, seq_len, head_dim]. Then, if q and - k have the shape [batch_size, heads, seq_len, head_dim], then setting unsqueeze_dim=1 makes - cos[position_ids] and sin[position_ids] broadcastable to the shapes of q and k. Similarly, if q and k have - the shape [batch_size, seq_len, heads, head_dim], then set unsqueeze_dim=2. - Returns: - `tuple(torch.Tensor)` comprising of the query and key tensors rotated using the Rotary Position Embedding. - """ - mrope_section = mrope_section * 2 - cos = torch.cat([m[i % 3] for i, m in enumerate(cos.split(mrope_section, dim=-1))], dim=-1).unsqueeze( - unsqueeze_dim - ) - sin = torch.cat([m[i % 3] for i, m in enumerate(sin.split(mrope_section, dim=-1))], dim=-1).unsqueeze( - unsqueeze_dim - ) - - q_embed = (q * cos) + (rotate_half(q) * sin) - k_embed = (k * cos) + (rotate_half(k) * sin) - return q_embed, k_embed - - -class Qwen2_5_VLAttention(nn.Module): - """ - Multi-headed attention from 'Attention Is All You Need' paper. Modified to use sliding window attention: Longformer - and "Generating Long Sequences with Sparse Transformers". - """ - - def __init__(self, config: Qwen2_5_VLTextConfig, layer_idx: Optional[int] = None): - super().__init__() - self.config = config - self.layer_idx = layer_idx - if layer_idx is None: - logger.warning_once( - f"Instantiating {self.__class__.__name__} without passing `layer_idx` is not recommended and will " - "to errors during the forward call, if caching is used. Please make sure to provide a `layer_idx` " - "when creating this class." - ) - - self.hidden_size = config.hidden_size - self.num_heads = config.num_attention_heads - self.head_dim = self.hidden_size // self.num_heads - self.num_key_value_heads = config.num_key_value_heads - self.num_key_value_groups = self.num_heads // self.num_key_value_heads - self.is_causal = True - self.attention_dropout = config.attention_dropout - self.rope_parameters = _get_rope_parameters(config) - self.scaling = self.head_dim**-0.5 - - if (self.head_dim * self.num_heads) != self.hidden_size: - raise ValueError( - f"hidden_size must be divisible by num_heads (got `hidden_size`: {self.hidden_size}" - f" and `num_heads`: {self.num_heads})." - ) - self.q_proj = nn.Linear(self.hidden_size, self.num_heads * self.head_dim, bias=True) - self.k_proj = nn.Linear(self.hidden_size, self.num_key_value_heads * self.head_dim, bias=True) - self.v_proj = nn.Linear(self.hidden_size, self.num_key_value_heads * self.head_dim, bias=True) - self.o_proj = nn.Linear(self.num_heads * self.head_dim, self.hidden_size, bias=False) - self.layer_type = config.layer_types[layer_idx] if hasattr(config, "layer_types") else None - self.sliding_window = config.sliding_window if self.layer_type == "sliding_attention" else None - - def forward( - self, - hidden_states: torch.Tensor, - attention_mask: Optional[torch.Tensor] = None, - position_ids: Optional[torch.LongTensor] = None, - past_key_values: Optional[Cache] = None, - output_attentions: bool = False, - use_cache: bool = False, - cache_position: Optional[torch.LongTensor] = None, - position_embeddings: Optional[tuple[torch.Tensor, torch.Tensor]] = None, - **kwargs: Unpack[FlashAttentionKwargs], - ) -> tuple[torch.Tensor, Optional[torch.Tensor], Optional[tuple[torch.Tensor]]]: - bsz, q_len, _ = hidden_states.size() - - query_states = self.q_proj(hidden_states) - key_states = self.k_proj(hidden_states) - value_states = self.v_proj(hidden_states) - - query_states = query_states.view(bsz, q_len, -1, self.head_dim).transpose(1, 2) - key_states = key_states.view(bsz, q_len, -1, self.head_dim).transpose(1, 2) - value_states = value_states.view(bsz, q_len, -1, self.head_dim).transpose(1, 2) - - cos, sin = position_embeddings - query_states, key_states = apply_multimodal_rotary_pos_emb( - query_states, key_states, cos, sin, self.rope_parameters["mrope_section"] - ) - - if past_key_values is not None: - cache_kwargs = {"sin": sin, "cos": cos, "cache_position": cache_position} # Specific to RoPE models - key_states, value_states = past_key_values.update(key_states, value_states, self.layer_idx, cache_kwargs) - - attention_interface: Callable = eager_attention_forward - if self.config._attn_implementation != "eager": - attention_interface = ALL_ATTENTION_FUNCTIONS[self.config._attn_implementation] - - attn_output, attn_weights = attention_interface( - self, - query_states, - key_states, - value_states, - attention_mask, - dropout=0.0 if not self.training else self.attention_dropout, - scaling=self.scaling, - sliding_window=self.sliding_window, - position_ids=position_ids, # pass positions for FA2 - **kwargs, - ) - - attn_output = attn_output.reshape(bsz, q_len, -1).contiguous() - attn_output = self.o_proj(attn_output) - return attn_output, attn_weights - - -class Qwen2_5_VLDecoderLayer(GradientCheckpointingLayer): - def __init__(self, config: Qwen2_5_VLTextConfig, layer_idx: int): - super().__init__() - self.hidden_size = config.hidden_size - - if config.use_sliding_window and config._attn_implementation != "flash_attention_2": - logger.warning_once( - f"Sliding Window Attention is enabled but not implemented for `{config._attn_implementation}`; " - "unexpected results may be encountered." - ) - self.self_attn = Qwen2_5_VLAttention(config, layer_idx) - - self.mlp = Qwen2MLP(config) - self.input_layernorm = Qwen2RMSNorm(config.hidden_size, eps=config.rms_norm_eps) - self.post_attention_layernorm = Qwen2RMSNorm(config.hidden_size, eps=config.rms_norm_eps) - self.attention_type = config.layer_types[layer_idx] - - def forward( - self, - hidden_states: torch.Tensor, - attention_mask: Optional[torch.Tensor] = None, - position_ids: Optional[torch.LongTensor] = None, - past_key_values: Optional[Cache] = None, - output_attentions: Optional[bool] = False, - use_cache: Optional[bool] = False, - cache_position: Optional[torch.LongTensor] = None, - position_embeddings: Optional[tuple[torch.Tensor, torch.Tensor]] = None, - **kwargs: Unpack[FlashAttentionKwargs], - ) -> tuple[torch.FloatTensor, Optional[tuple[torch.FloatTensor, torch.FloatTensor]]]: - """ - Args: - hidden_states (`torch.FloatTensor`): input to the layer of shape `(batch, seq_len, embed_dim)` - attention_mask (`torch.FloatTensor`, *optional*): attention mask of size - `(batch, sequence_length)` where padding elements are indicated by 0. - output_attentions (`bool`, *optional*): - Whether or not to return the attentions tensors of all attention layers. See `attentions` under - returned tensors for more detail. - use_cache (`bool`, *optional*): - If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding - (see `past_key_values`). - past_key_values (`Cache`, *optional*): cached past key and value projection states - cache_position (`torch.LongTensor` of shape `(sequence_length)`, *optional*): - Indices depicting the position of the input sequence tokens in the sequence. - position_embeddings (`tuple[torch.FloatTensor, torch.FloatTensor]`, *optional*): - Tuple containing the cosine and sine positional embeddings of shape `(batch_size, seq_len, head_dim)`, - with `head_dim` being the embedding dimension of each attention head. - kwargs (`dict`, *optional*): - Arbitrary kwargs to be ignored, used for FSDP and other methods that injects code - into the model - """ - - residual = hidden_states - - hidden_states = self.input_layernorm(hidden_states) - - # Self Attention - hidden_states, self_attn_weights = self.self_attn( - hidden_states=hidden_states, - attention_mask=attention_mask, - position_ids=position_ids, - past_key_values=past_key_values, - output_attentions=output_attentions, - use_cache=use_cache, - cache_position=cache_position, - position_embeddings=position_embeddings, - **kwargs, - ) - hidden_states = residual + hidden_states - - # Fully Connected - residual = hidden_states - hidden_states = self.post_attention_layernorm(hidden_states) - hidden_states = self.mlp(hidden_states) - hidden_states = residual + hidden_states - - outputs = (hidden_states,) - - if output_attentions: - outputs += (self_attn_weights,) - - return outputs - - -@auto_docstring -class Qwen2_5_VLTextModel(Qwen2_5_VLPreTrainedModel): - config: Qwen2_5_VLTextConfig - input_modalities = "text" - - def __init__(self, config: Qwen2_5_VLTextConfig): - super().__init__(config) - self.padding_idx = config.pad_token_id - self.vocab_size = config.vocab_size - - self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx) - self.layers = nn.ModuleList( - [Qwen2_5_VLDecoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)] - ) - self._attn_implementation = config._attn_implementation - self.norm = Qwen2RMSNorm(config.hidden_size, eps=config.rms_norm_eps) - self.has_sliding_layers = "sliding_attention" in self.config.layer_types - self.rotary_emb = Qwen2_5_VLRotaryEmbedding(config=config) - - self.gradient_checkpointing = False - # Initialize weights and apply final processing - self.post_init() - - @auto_docstring - def forward( - self, - input_ids: Optional[torch.LongTensor] = None, - attention_mask: Optional[torch.Tensor] = None, - position_ids: Optional[torch.LongTensor] = None, - past_key_values: Optional[Cache] = None, - inputs_embeds: Optional[torch.FloatTensor] = None, - use_cache: Optional[bool] = None, - output_attentions: Optional[bool] = None, - output_hidden_states: Optional[bool] = None, - return_dict: Optional[bool] = None, - cache_position: Optional[torch.LongTensor] = None, - **kwargs: Unpack[FlashAttentionKwargs], - ) -> Union[tuple, BaseModelOutputWithPast]: - output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions - output_hidden_states = ( - output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states - ) - use_cache = use_cache if use_cache is not None else self.config.use_cache - - return_dict = return_dict if return_dict is not None else self.config.use_return_dict - - if (input_ids is None) ^ (inputs_embeds is not None): - raise ValueError("You must specify exactly one of input_ids or inputs_embeds") - - if self.gradient_checkpointing and self.training: - if use_cache: - logger.warning_once( - "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`..." - ) - use_cache = False - - # torch.jit.trace() doesn't support cache objects in the output - if use_cache and past_key_values is None and not torch.jit.is_tracing(): - past_key_values = DynamicCache(config=self.config) - - if inputs_embeds is None: - inputs_embeds = self.embed_tokens(input_ids) - - if cache_position is None: - past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0 - cache_position = torch.arange( - past_seen_tokens, past_seen_tokens + inputs_embeds.shape[1], device=inputs_embeds.device - ) - - # the hard coded `3` is for temporal, height and width. - if position_ids is None: - position_ids = cache_position.view(1, 1, -1).expand(3, inputs_embeds.shape[0], -1) - elif position_ids.ndim == 2: - position_ids = position_ids[None, ...].expand(3, position_ids.shape[0], -1) - - # NOTE: we need to pass text position ids for packing. Qwen2-VL uses 3D positions - # where each dim indicates visual spatial positions for temporal/height/width grids. - # There are two scenarios when FA2-like packed masking might be activated. - # 1. User specifically passed packed `position_ids` and no attention mask. - # In this case we expect the useer to create correct position ids for all 3 grids - # and prepend text-only position ids to it. The final tensor will be [4, bs, seq-len] - # 2. User runs forward with no attention mask and no position ids. In this case, position ids - # are prepared by the model (`get_rope_index`) as `[4, bs, seq-len]` tensor. Text-only positions are - # prepended by us when creating positions so that the mask is constructed correctly. NOTE: failing to pass - # text-only positions will cause incorrect mask construction, do not change `prepare_input_for_generation` - if position_ids.ndim == 3 and position_ids.shape[0] == 4: - text_position_ids = position_ids[0] - position_ids = position_ids[1:] - else: - # If inputs are not packed (usual 3D positions), do not prepare mask from position_ids - text_position_ids = None - - # It may already have been prepared by e.g. `generate` - if not isinstance(causal_mask_mapping := attention_mask, dict): - # Prepare mask arguments - mask_kwargs = { - "config": self.config, - "input_embeds": inputs_embeds, - "attention_mask": attention_mask, - "cache_position": cache_position, - "past_key_values": past_key_values, - "position_ids": text_position_ids, - } - # Create the masks - causal_mask_mapping = { - "full_attention": create_causal_mask(**mask_kwargs), - } - # The sliding window alternating layers are not always activated depending on the config - if self.has_sliding_layers: - causal_mask_mapping["sliding_attention"] = create_sliding_window_causal_mask(**mask_kwargs) - - hidden_states = inputs_embeds - position_embeddings = self.rotary_emb(hidden_states, position_ids) - - # decoder layers - all_hidden_states = () if output_hidden_states else None - all_self_attns = () if output_attentions else None - - for decoder_layer in self.layers: - if output_hidden_states: - all_hidden_states += (hidden_states,) - - layer_outputs = decoder_layer( - hidden_states, - attention_mask=causal_mask_mapping[decoder_layer.attention_type], - position_embeddings=position_embeddings, - position_ids=text_position_ids, - past_key_values=past_key_values, - output_attentions=output_attentions, - use_cache=use_cache, - cache_position=cache_position, - **kwargs, - ) - - hidden_states = layer_outputs[0] - - if output_attentions: - all_self_attns += (layer_outputs[1],) - - hidden_states = self.norm(hidden_states) - - # add hidden states from the last decoder layer - if output_hidden_states: - all_hidden_states += (hidden_states,) - - if not return_dict: - return tuple( - v for v in [hidden_states, past_key_values, all_hidden_states, all_self_attns] if v is not None - ) - return BaseModelOutputWithPast( - last_hidden_state=hidden_states, - past_key_values=past_key_values, - hidden_states=all_hidden_states, - attentions=all_self_attns, - ) - - -@auto_docstring -class Qwen2_5_VLModel(Qwen2_5_VLPreTrainedModel): - base_model_prefix = "model" - _checkpoint_conversion_mapping = {} - # Reference: fix gemma3 grad acc #37208 - accepts_loss_kwargs = False - config: Qwen2_5_VLConfig - _no_split_modules = ["Qwen2_5_VLDecoderLayer", "Qwen2_5_VLVisionBlock"] - - def __init__(self, config): - super().__init__(config) - self.visual = Qwen2_5_VisionTransformerPretrainedModel._from_config(config.vision_config) - self.language_model = Qwen2_5_VLTextModel._from_config(config.text_config) - self.rope_deltas = None # cache rope_deltas here - - # Initialize weights and apply final processing - self.post_init() - - def get_input_embeddings(self): - return self.language_model.get_input_embeddings() - - def set_input_embeddings(self, value): - self.language_model.set_input_embeddings(value) - - def set_decoder(self, decoder): - self.language_model = decoder - - def get_decoder(self): - return self.language_model - - def get_rope_index( - self, - input_ids: Optional[torch.LongTensor] = None, - image_grid_thw: Optional[torch.LongTensor] = None, - video_grid_thw: Optional[torch.LongTensor] = None, - second_per_grid_ts: Optional[torch.Tensor] = None, - attention_mask: Optional[torch.Tensor] = None, - ) -> tuple[torch.Tensor, torch.Tensor]: - """ - Calculate the 3D rope index based on image and video's temporal, height and width in LLM. - - Explanation: - Each embedding sequence contains vision embedding and text embedding or just contains text embedding. - - For pure text embedding sequence, the rotary position embedding has no difference with modern LLMs. - Examples: - input_ids: [T T T T T], here T is for text. - temporal position_ids: [0, 1, 2, 3, 4] - height position_ids: [0, 1, 2, 3, 4] - width position_ids: [0, 1, 2, 3, 4] - - For vision and text embedding sequence, we calculate 3D rotary position embedding for vision part - and 1D rotary position embedding for text part. - Examples: - Temporal (Time): 3 patches, representing different segments of the video in time. - Height: 2 patches, dividing each frame vertically. - Width: 2 patches, dividing each frame horizontally. - We also have some important parameters: - fps (Frames Per Second): The video's frame rate, set to 1. This means one frame is processed each second. - tokens_per_second: This is a crucial parameter. It dictates how many "time-steps" or "temporal tokens" are conceptually packed into a one-second interval of the video. In this case, we have 25 tokens per second. So each second of the video will be represented with 25 separate time points. It essentially defines the temporal granularity. - temporal_patch_size: The number of frames that compose one temporal patch. Here, it's 2 frames. - interval: The step size for the temporal position IDs, calculated as tokens_per_second * temporal_patch_size / fps. In this case, 25 * 2 / 1 = 50. This means that each temporal patch will be have a difference of 50 in the temporal position IDs. - input_ids: [V V V V V V V V V V V V T T T T T], here V is for vision. - vision temporal position_ids: [0, 0, 0, 0, 50, 50, 50, 50, 100, 100, 100, 100] - vision height position_ids: [0, 0, 1, 1, 0, 0, 1, 1, 0, 0, 1, 1] - vision width position_ids: [0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1] - text temporal position_ids: [101, 102, 103, 104, 105] - text height position_ids: [101, 102, 103, 104, 105] - text width position_ids: [101, 102, 103, 104, 105] - Here we calculate the text start position_ids as the max vision position_ids plus 1. - - Args: - input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`): - Indices of input sequence tokens in the vocabulary. Padding will be ignored by default should you provide - it. - image_grid_thw (`torch.LongTensor` of shape `(num_images, 3)`, *optional*): - The temporal, height and width of feature shape of each image in LLM. - video_grid_thw (`torch.LongTensor` of shape `(num_videos, 3)`, *optional*): - The temporal, height and width of feature shape of each video in LLM. - second_per_grid_ts (`torch.Tensor` of shape `(num_videos)`, *optional*): - The time interval (in seconds) for each grid along the temporal dimension in the 3D position IDs. - attention_mask (`torch.Tensor` of shape `(batch_size, sequence_length)`, *optional*): - Mask to avoid performing attention on padding token indices. Mask values selected in `[0, 1]`: - - - 1 for tokens that are **not masked**, - - 0 for tokens that are **masked**. - - Returns: - position_ids (`torch.LongTensor` of shape `(3, batch_size, sequence_length)`) - mrope_position_deltas (`torch.Tensor` of shape `(batch_size)`) - """ - spatial_merge_size = self.config.vision_config.spatial_merge_size - image_token_id = self.config.image_token_id - video_token_id = self.config.video_token_id - vision_start_token_id = self.config.vision_start_token_id - mrope_position_deltas = [] - if input_ids is not None and (image_grid_thw is not None or video_grid_thw is not None): - total_input_ids = input_ids - if attention_mask is not None: - attention_mask = attention_mask == 1 - position_ids = torch.ones( - 3, - input_ids.shape[0], - input_ids.shape[1], - dtype=input_ids.dtype, - device=input_ids.device, - ) - image_index, video_index = 0, 0 - for i, input_ids in enumerate(total_input_ids): - if attention_mask is not None: - input_ids = input_ids[attention_mask[i]] - image_nums, video_nums = 0, 0 - vision_start_indices = torch.argwhere(input_ids == vision_start_token_id).squeeze(1) - vision_tokens = input_ids[vision_start_indices + 1] - image_nums = (vision_tokens == image_token_id).sum() - video_nums = (vision_tokens == video_token_id).sum() - input_tokens = input_ids.tolist() - llm_pos_ids_list: list = [] - st = 0 - remain_images, remain_videos = image_nums, video_nums - for _ in range(image_nums + video_nums): - if image_token_id in input_tokens and remain_images > 0: - ed_image = input_tokens.index(image_token_id, st) - else: - ed_image = len(input_tokens) + 1 - if video_token_id in input_tokens and remain_videos > 0: - ed_video = input_tokens.index(video_token_id, st) - else: - ed_video = len(input_tokens) + 1 - if ed_image < ed_video: - t, h, w = ( - image_grid_thw[image_index][0], - image_grid_thw[image_index][1], - image_grid_thw[image_index][2], - ) - second_per_grid_t = 0 - image_index += 1 - remain_images -= 1 - ed = ed_image - - else: - t, h, w = ( - video_grid_thw[video_index][0], - video_grid_thw[video_index][1], - video_grid_thw[video_index][2], - ) - if second_per_grid_ts is not None: - second_per_grid_t = second_per_grid_ts[video_index] - else: - second_per_grid_t = 1.0 - video_index += 1 - remain_videos -= 1 - ed = ed_video - llm_grid_t, llm_grid_h, llm_grid_w = ( - t.item(), - h.item() // spatial_merge_size, - w.item() // spatial_merge_size, - ) - text_len = ed - st - - st_idx = llm_pos_ids_list[-1].max() + 1 if len(llm_pos_ids_list) > 0 else 0 - llm_pos_ids_list.append(torch.arange(text_len).view(1, -1).expand(3, -1) + st_idx) - - range_tensor = torch.arange(llm_grid_t).view(-1, 1) - expanded_range = range_tensor.expand(-1, llm_grid_h * llm_grid_w) - - ## normalize type, send to device. - second_per_grid_t = torch.as_tensor( - second_per_grid_t, dtype=range_tensor.dtype, device=range_tensor.device - ) - - time_tensor = expanded_range * second_per_grid_t * self.config.vision_config.tokens_per_second - - time_tensor_long = time_tensor.long() - t_index = time_tensor_long.flatten() - - h_index = torch.arange(llm_grid_h).view(1, -1, 1).expand(llm_grid_t, -1, llm_grid_w).flatten() - w_index = torch.arange(llm_grid_w).view(1, 1, -1).expand(llm_grid_t, llm_grid_h, -1).flatten() - llm_pos_ids_list.append(torch.stack([t_index, h_index, w_index]) + text_len + st_idx) - st = ed + llm_grid_t * llm_grid_h * llm_grid_w - - if st < len(input_tokens): - st_idx = llm_pos_ids_list[-1].max() + 1 if len(llm_pos_ids_list) > 0 else 0 - text_len = len(input_tokens) - st - llm_pos_ids_list.append(torch.arange(text_len).view(1, -1).expand(3, -1) + st_idx) - - llm_positions = torch.cat(llm_pos_ids_list, dim=1).reshape(3, -1) - if attention_mask is not None: - position_ids[..., i, attention_mask[i]] = llm_positions.to(position_ids.device) - else: - position_ids[..., i, :] = llm_positions.to(position_ids.device) - mrope_position_deltas.append(llm_positions.max() + 1 - len(total_input_ids[i])) - mrope_position_deltas = torch.tensor(mrope_position_deltas).unsqueeze(1).to(device=input_ids.device) - return position_ids, mrope_position_deltas - else: - if attention_mask is not None: - position_ids = attention_mask.long().cumsum(-1) - 1 - position_ids.masked_fill_(attention_mask == 0, 1) - position_ids = position_ids.unsqueeze(0).expand(3, -1, -1).to(attention_mask.device) - max_position_ids = position_ids.max(0, keepdim=False)[0].max(-1, keepdim=True)[0] - mrope_position_deltas = max_position_ids + 1 - attention_mask.shape[-1] - else: - position_ids = ( - torch.arange(input_ids.shape[1], device=input_ids.device) - .view(1, 1, -1) - .expand(3, input_ids.shape[0], -1) - ) - mrope_position_deltas = torch.zeros( - [input_ids.shape[0], 1], - device=input_ids.device, - dtype=input_ids.dtype, - ) - - return position_ids, mrope_position_deltas - - def get_video_features( - self, pixel_values_videos: torch.FloatTensor, video_grid_thw: Optional[torch.LongTensor] = None - ): - """ - Encodes videos into continuous embeddings that can be forwarded to the language model. - - Args: - pixel_values_videos (`torch.FloatTensor` of shape `(batch_size, num_channels, image_size, image_size)`): - The tensors corresponding to the input videos. - video_grid_thw (`torch.LongTensor` of shape `(num_videos, 3)`, *optional*): - The temporal, height and width of feature shape of each video in LLM. - """ - pixel_values_videos = pixel_values_videos.type(self.visual.dtype) - video_embeds = self.visual(pixel_values_videos, grid_thw=video_grid_thw) - split_sizes = (video_grid_thw.prod(-1) // self.visual.spatial_merge_size**2).tolist() - video_embeds = torch.split(video_embeds, split_sizes) - return video_embeds - - def get_image_features(self, pixel_values: torch.FloatTensor, image_grid_thw: Optional[torch.LongTensor] = None): - """ - Encodes images into continuous embeddings that can be forwarded to the language model. - - Args: - pixel_values (`torch.FloatTensor` of shape `(batch_size, num_channels, image_size, image_size)`): - The tensors corresponding to the input images. - image_grid_thw (`torch.LongTensor` of shape `(num_images, 3)`, *optional*): - The temporal, height and width of feature shape of each image in LLM. - """ - pixel_values = pixel_values.type(self.visual.dtype) - image_embeds = self.visual(pixel_values, grid_thw=image_grid_thw) - # Return the full concatenated feature tensor. Callers that need per-image - # slices can split offline using: - # sizes = (image_grid_thw.prod(-1) // spatial_merge_size**2).tolist() - # The ONNX model outputs a single tensor; the genai runtime handles per-image - # splitting using image_grid_thw. Removing .tolist() and torch.split here - # allows image_grid_thw.shape[0] to remain symbolic during export. - return image_embeds - - def get_fused_input_embeddings(self, input_ids, image_features=None): - """ - Fuses the input embeddings from the language model with the image features. - - Args: - input_ids (`torch.LongTensor`): The input IDs for the language model. - image_features (`torch.FloatTensor`, optional): The image features to fuse with the input embeddings. - - Returns: - `torch.FloatTensor`: The fused input embeddings. - """ - def true_fn_for_input_ids(input_ids): - special_image_mask = input_ids == self.config.image_token_id - llm_input_ids = input_ids.clone() - llm_input_ids[special_image_mask] = 0 - return input_ids - def false_fn_for_input_ids(input_ids): - return input_ids - - # condition 1 on the image token index - llm_input_ids = torch.cond( - input_ids is not None and self.config.image_token_id >= self.config.text_config.vocab_size, - true_fn_for_input_ids, - false_fn_for_input_ids, - (input_ids, ) - ) - - inputs_embeds = self.language_model.get_input_embeddings()(llm_input_ids) - - def image_features_is_none(inputs_embeds, image_features=None): - return inputs_embeds - - def image_features_is_not_none(inputs_embeds, image_features=None): - # input_ids: [batch_size, seq_len] - # input_embeds: [batch_size, seq_len, 2560 (hidden_size)] - special_image_mask = (llm_input_ids == self.config.image_token_id).unsqueeze(-1) - special_image_mask = special_image_mask.expand_as(inputs_embeds).to(inputs_embeds.device) - - image_features = image_features.to(inputs_embeds.device, inputs_embeds.dtype) - inputs_embeds = inputs_embeds.masked_scatter(special_image_mask, image_features) - return inputs_embeds - - # condition 2 on the image features - inputs_embeds = torch.cond( - image_features is None, - image_features_is_none, - image_features_is_not_none, - (inputs_embeds, image_features,) - ) - - return inputs_embeds - - def get_placeholder_mask( - self, - input_ids: torch.LongTensor, - inputs_embeds: torch.FloatTensor, - image_features: Optional[torch.FloatTensor] = None, - video_features: Optional[torch.FloatTensor] = None, - ): - """ - Obtains multimodal placeholder mask from `input_ids` or `inputs_embeds`, and checks that the placeholder token count is - equal to the length of multimodal features. If the lengths are different, an error is raised. - """ - if input_ids is None: - special_image_mask = inputs_embeds == self.get_input_embeddings()( - torch.tensor(self.config.image_token_id, dtype=torch.long, device=inputs_embeds.device) - ) - special_image_mask = special_image_mask.all(-1) - special_video_mask = inputs_embeds == self.get_input_embeddings()( - torch.tensor(self.config.video_token_id, dtype=torch.long, device=inputs_embeds.device) - ) - special_video_mask = special_video_mask.all(-1) - else: - special_image_mask = input_ids == self.config.image_token_id - special_video_mask = input_ids == self.config.video_token_id - - n_image_tokens = special_image_mask.sum() - special_image_mask = special_image_mask.unsqueeze(-1).expand_as(inputs_embeds).to(inputs_embeds.device) - if image_features is not None and inputs_embeds[special_image_mask].numel() != image_features.numel(): - raise ValueError( - f"Image features and image tokens do not match: tokens: {n_image_tokens}, features {image_features.shape[0]}" - ) - - n_video_tokens = special_video_mask.sum() - special_video_mask = special_video_mask.unsqueeze(-1).expand_as(inputs_embeds).to(inputs_embeds.device) - if video_features is not None and inputs_embeds[special_video_mask].numel() != video_features.numel(): - raise ValueError( - f"Videos features and video tokens do not match: tokens: {n_video_tokens}, features {video_features.shape[0]}" - ) - - return special_image_mask, special_video_mask - - @auto_docstring - def forward( - self, - input_ids: Optional[torch.LongTensor] = None, - attention_mask: Optional[torch.Tensor] = None, - position_ids: Optional[torch.LongTensor] = None, - past_key_values: Optional[Cache] = None, - inputs_embeds: Optional[torch.FloatTensor] = None, - use_cache: Optional[bool] = None, - output_attentions: Optional[bool] = None, - output_hidden_states: Optional[bool] = None, - return_dict: Optional[bool] = None, - pixel_values: Optional[torch.Tensor] = None, - pixel_values_videos: Optional[torch.FloatTensor] = None, - image_grid_thw: Optional[torch.LongTensor] = None, - video_grid_thw: Optional[torch.LongTensor] = None, - rope_deltas: Optional[torch.LongTensor] = None, - cache_position: Optional[torch.LongTensor] = None, - second_per_grid_ts: Optional[torch.Tensor] = None, - **kwargs: Unpack[TransformersKwargs], - ) -> Union[tuple, Qwen2_5_VLModelOutputWithPast]: - r""" - image_grid_thw (`torch.LongTensor` of shape `(num_images, 3)`, *optional*): - The temporal, height and width of feature shape of each image in LLM. - video_grid_thw (`torch.LongTensor` of shape `(num_videos, 3)`, *optional*): - The temporal, height and width of feature shape of each video in LLM. - rope_deltas (`torch.LongTensor` of shape `(batch_size, )`, *optional*): - The rope index difference between sequence length and multimodal rope. - second_per_grid_ts (`torch.Tensor` of shape `(num_videos)`, *optional*): - The time interval (in seconds) for each grid along the temporal dimension in the 3D position IDs. - """ - - output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions - output_hidden_states = ( - output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states - ) - return_dict = return_dict if return_dict is not None else self.config.use_return_dict - - if inputs_embeds is None: - inputs_embeds = self.get_input_embeddings()(input_ids) - - if pixel_values is not None: - image_embeds = self.get_image_features(pixel_values, image_grid_thw) - image_embeds = torch.cat(image_embeds, dim=0).to(inputs_embeds.device, inputs_embeds.dtype) - image_mask, _ = self.get_placeholder_mask( - input_ids, inputs_embeds=inputs_embeds, image_features=image_embeds - ) - inputs_embeds = inputs_embeds.masked_scatter(image_mask, image_embeds) - - if pixel_values_videos is not None: - video_embeds = self.get_video_features(pixel_values_videos, video_grid_thw) - video_embeds = torch.cat(video_embeds, dim=0).to(inputs_embeds.device, inputs_embeds.dtype) - _, video_mask = self.get_placeholder_mask( - input_ids, inputs_embeds=inputs_embeds, video_features=video_embeds - ) - inputs_embeds = inputs_embeds.masked_scatter(video_mask, video_embeds) - - if position_ids is None: - if self.rope_deltas is None or cache_position is None or cache_position[0] == 0: - position_ids, rope_deltas = self.get_rope_index( - input_ids, - image_grid_thw, - video_grid_thw, - second_per_grid_ts=second_per_grid_ts, - attention_mask=attention_mask, - ) - self.rope_deltas = rope_deltas - else: - batch_size, seq_length, _ = inputs_embeds.shape - position_ids = torch.arange(seq_length, device=inputs_embeds.device) - position_ids = position_ids.view(1, 1, -1).expand(3, batch_size, -1) - if cache_position is not None: - delta = (cache_position[0] + self.rope_deltas).to(inputs_embeds.device) - else: - delta = torch.zeros((batch_size, seq_length), device=inputs_embeds.device) - delta = delta.repeat_interleave(batch_size // delta.shape[0], dim=1) - position_ids = position_ids + delta.to(position_ids.device) - - outputs = self.language_model( - input_ids=None, - position_ids=position_ids, - attention_mask=attention_mask, - past_key_values=past_key_values, - inputs_embeds=inputs_embeds, - use_cache=use_cache, - output_attentions=output_attentions, - output_hidden_states=output_hidden_states, - return_dict=True, - cache_position=cache_position, - **kwargs, - ) - - output = Qwen2_5_VLModelOutputWithPast( - last_hidden_state=outputs.last_hidden_state, - past_key_values=outputs.past_key_values, - hidden_states=outputs.hidden_states, - attentions=outputs.attentions, - rope_deltas=self.rope_deltas, - ) - return output if return_dict else output.to_tuple() - - -@dataclass -@auto_docstring( - custom_intro=""" - Base class for Qwen2_5_VL causal language model (or autoregressive) outputs. - """ -) -class Qwen2_5_VLCausalLMOutputWithPast(ModelOutput): - r""" - loss (`torch.FloatTensor` of shape `(1,)`, *optional*, returned when `labels` is provided): - Language modeling loss (for next-token prediction). - logits (`torch.FloatTensor` of shape `(batch_size, sequence_length, config.vocab_size)`): - Prediction scores of the language modeling head (scores for each vocabulary token before SoftMax). - past_key_values (`Cache`, *optional*, returned when `use_cache=True` is passed or when `config.use_cache=True`): - It is a [`~cache_utils.Cache`] instance. For more details, see our [kv cache guide](https://huggingface.co/docs/transformers/en/kv_cache). - - Contains pre-computed hidden-states (key and values in the self-attention blocks) that can be used (see - `past_key_values` input) to speed up sequential decoding. - rope_deltas (`torch.LongTensor` of shape `(batch_size, )`, *optional*): - The rope index difference between sequence length and multimodal rope. - """ - - loss: Optional[torch.FloatTensor] = None - logits: Optional[torch.FloatTensor] = None - past_key_values: Optional[Cache] = None - hidden_states: Optional[tuple[torch.FloatTensor]] = None - attentions: Optional[tuple[torch.FloatTensor]] = None - rope_deltas: Optional[torch.LongTensor] = None - - -class Qwen2_5_VLForConditionalGeneration(Qwen2_5_VLPreTrainedModel, GenerationMixin): - _checkpoint_conversion_mapping = { - "^visual": "model.visual", - r"^model(?!\.(language_model|visual))": "model.language_model", - } - _tied_weights_keys = ["lm_head.weight"] - # Reference: fix gemma3 grad acc #37208 - accepts_loss_kwargs = False - - def __init__(self, config): - super().__init__(config) - self.model = Qwen2_5_VLModel(config) - self.lm_head = nn.Linear(config.text_config.hidden_size, config.text_config.vocab_size, bias=False) - - self.post_init() - - def get_input_embeddings(self): - return self.model.get_input_embeddings() - - def set_input_embeddings(self, value): - self.model.set_input_embeddings(value) - - def set_decoder(self, decoder): - self.model.set_decoder(decoder) - - def get_decoder(self): - return self.model.get_decoder() - - def get_video_features( - self, pixel_values_videos: torch.FloatTensor, video_grid_thw: Optional[torch.LongTensor] = None - ): - return self.model.get_video_features(pixel_values_videos, video_grid_thw) - - def get_image_features(self, pixel_values: torch.FloatTensor, image_grid_thw: Optional[torch.LongTensor] = None): - return self.model.get_image_features(pixel_values, image_grid_thw) - - # Make modules available through conditional class for BC - @property - def language_model(self): - return self.model.language_model - - @property - def visual(self): - return self.model.visual - - @can_return_tuple - @auto_docstring - def forward( - self, - input_ids: Optional[torch.LongTensor] = None, - attention_mask: Optional[torch.Tensor] = None, - position_ids: Optional[torch.LongTensor] = None, - past_key_values: Optional[Cache] = None, - inputs_embeds: Optional[torch.FloatTensor] = None, - labels: Optional[torch.LongTensor] = None, - use_cache: Optional[bool] = None, - output_attentions: Optional[bool] = None, - output_hidden_states: Optional[bool] = None, - pixel_values: Optional[torch.Tensor] = None, - pixel_values_videos: Optional[torch.FloatTensor] = None, - image_grid_thw: Optional[torch.LongTensor] = None, - video_grid_thw: Optional[torch.LongTensor] = None, - rope_deltas: Optional[torch.LongTensor] = None, - cache_position: Optional[torch.LongTensor] = None, - second_per_grid_ts: Optional[torch.Tensor] = None, - logits_to_keep: Union[int, torch.Tensor] = 0, - **kwargs: Unpack[TransformersKwargs], - ) -> Union[tuple, Qwen2_5_VLCausalLMOutputWithPast]: - r""" - labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*): - Labels for computing the masked language modeling loss. Indices should either be in `[0, ..., - config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored - (masked), the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`. - image_grid_thw (`torch.LongTensor` of shape `(num_images, 3)`, *optional*): - The temporal, height and width of feature shape of each image in LLM. - video_grid_thw (`torch.LongTensor` of shape `(num_videos, 3)`, *optional*): - The temporal, height and width of feature shape of each video in LLM. - rope_deltas (`torch.LongTensor` of shape `(batch_size, )`, *optional*): - The rope index difference between sequence length and multimodal rope. - second_per_grid_ts (`torch.Tensor` of shape `(num_videos)`, *optional*): - The time interval (in seconds) for each grid along the temporal dimension in the 3D position IDs. - - Example: - - ```python - >>> from transformers import AutoProcessor, Qwen2_5_VLForConditionalGeneration - - >>> model = Qwen2_5_VLForConditionalGeneration.from_pretrained("Qwen/Qwen2.5-VL-7B-Instruct") - >>> processor = AutoProcessor.from_pretrained("Qwen/Qwen2.5-VL-7B-Instruct") - - >>> messages = [ - { - "role": "user", - "content": [ - { - "type": "image", - "image": "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/pipeline-cat-chonk.jpeg", - }, - {"type": "text", "text": "Describe the image."}, - ], - } - ] - - >>> inputs = processor.apply_chat_template( - messages, - tokenize=True, - add_generation_prompt=True, - return_dict=True, - return_tensors="pt" - ) - - >>> # Generate - >>> generated_ids = model.generate(**inputs, max_new_tokens=1024) - >>> generated_ids_trimmed = [out_ids[len(in_ids) :] for in_ids, out_ids in zip(inputs.input_ids, generated_ids)] - >>> output_text = processor.batch_decode(generated_ids_trimmed, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0] - >>> print(output_text) - ``` - """ - - output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions - output_hidden_states = ( - output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states - ) - - outputs = self.model( - input_ids=input_ids, - pixel_values=pixel_values, - pixel_values_videos=pixel_values_videos, - image_grid_thw=image_grid_thw, - video_grid_thw=video_grid_thw, - second_per_grid_ts=second_per_grid_ts, - position_ids=position_ids, - attention_mask=attention_mask, - past_key_values=past_key_values, - inputs_embeds=inputs_embeds, - use_cache=use_cache, - output_attentions=output_attentions, - output_hidden_states=output_hidden_states, - return_dict=True, - cache_position=cache_position, - **kwargs, - ) - - hidden_states = outputs[0] - - # Only compute necessary logits, and do not upcast them to float if we are not computing the loss - slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep - logits = self.lm_head(hidden_states[:, slice_indices, :]) - - loss = None - if labels is not None: - loss = self.loss_function( - logits=logits, labels=labels, vocab_size=self.config.text_config.vocab_size, **kwargs - ) - - return Qwen2_5_VLCausalLMOutputWithPast( - loss=loss, - logits=logits, - past_key_values=outputs.past_key_values, - hidden_states=outputs.hidden_states, - attentions=outputs.attentions, - rope_deltas=outputs.rope_deltas, - ) - - def prepare_inputs_for_generation( - self, - input_ids, - past_key_values=None, - attention_mask=None, - inputs_embeds=None, - cache_position=None, - position_ids=None, - use_cache=True, - pixel_values=None, - pixel_values_videos=None, - image_grid_thw=None, - video_grid_thw=None, - second_per_grid_ts=None, - **kwargs, - ): - # Overwritten -- in specific circumstances we don't want to forward image inputs to the model - - model_inputs = super().prepare_inputs_for_generation( - input_ids, - past_key_values=past_key_values, - attention_mask=attention_mask, - inputs_embeds=inputs_embeds, - cache_position=cache_position, - position_ids=position_ids, - pixel_values=pixel_values, - pixel_values_videos=pixel_values_videos, - image_grid_thw=image_grid_thw, - video_grid_thw=video_grid_thw, - second_per_grid_ts=second_per_grid_ts, - use_cache=use_cache, - **kwargs, - ) - - # Qwen2-5-VL position_ids are prepared with rope_deltas - if position_ids is None: - # Calculate RoPE index once per generation in the pre-fill stage only. - # When compiling, we can't check tensor values thus we check only input length - # It is safe to assume that `length!=1` means we're in pre-fill because compiled - # models currently cannot do assisted decoding - if cache_position[0] == 0 or self.model.rope_deltas is None: - vision_positions, rope_deltas = self.model.get_rope_index( - model_inputs.get("input_ids", None), - image_grid_thw=image_grid_thw, - video_grid_thw=video_grid_thw, - second_per_grid_ts=second_per_grid_ts, - attention_mask=attention_mask, - ) - self.model.rope_deltas = rope_deltas - # then use the prev pre-calculated rope-deltas to get the correct position ids - elif "position_ids" in model_inputs: - batch_size, seq_length = model_inputs["position_ids"].shape - device = model_inputs["position_ids"].device - position_ids = torch.arange(seq_length, device=device) - position_ids = position_ids.view(1, 1, -1).expand(3, batch_size, -1) - delta = cache_position[0] + self.model.rope_deltas - delta = delta.repeat_interleave(batch_size // delta.shape[0], dim=0) - vision_positions = position_ids + delta.expand_as(position_ids) - - # Concatenate "text + vision" positions into [4, bs, seq-len] - text_positions = model_inputs["position_ids"][None, ...] - model_inputs["position_ids"] = torch.cat([text_positions, vision_positions], dim=0) - - if cache_position[0] != 0: - model_inputs["pixel_values"] = None - model_inputs["pixel_values_videos"] = None - - return model_inputs - - def _get_image_nums_and_video_nums( - self, - input_ids: Optional[torch.LongTensor], - inputs_embeds: Optional[torch.Tensor] = None, - ) -> tuple[torch.Tensor, torch.Tensor]: - """ - Get the number of images and videos for each sample to calculate the separation length of the sample tensor. - These parameters are not passed through the processor to avoid unpredictable impacts from interface modifications. - - Args: - input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`): - Indices of input sequence tokens in the vocabulary. - - Returns: - image_nums (`torch.LongTensor` of shape `(batch_size, num_images_sample)`) - video_nums (`torch.LongTensor` of shape `(batch_size, num_videos_sample)`) - """ - image_token_id = self.config.image_token_id - video_token_id = self.config.video_token_id - vision_start_token_id = self.config.vision_start_token_id - - if inputs_embeds is not None: - vision_start_mask = ( - inputs_embeds - == self.get_input_embeddings()( - torch.tensor(vision_start_token_id, dtype=torch.long, device=inputs_embeds.device) - ) - )[..., 0] - image_mask = ( - inputs_embeds - == self.get_input_embeddings()( - torch.tensor(image_token_id, dtype=torch.long, device=inputs_embeds.device) - ) - )[..., 0] - video_mask = ( - inputs_embeds - == self.get_input_embeddings()( - torch.tensor(video_token_id, dtype=torch.long, device=inputs_embeds.device) - ) - )[..., 0] - else: - vision_start_mask = input_ids == vision_start_token_id - image_mask = input_ids == image_token_id - video_mask = input_ids == video_token_id - - vision_first_mask = torch.roll(vision_start_mask, shifts=1, dims=1) - image_nums = torch.sum(vision_first_mask & image_mask, dim=1) - video_nums = torch.sum(vision_first_mask & video_mask, dim=1) - - return image_nums, video_nums - - def _expand_inputs_for_generation( - self, - expand_size: int = 1, - is_encoder_decoder: bool = False, - input_ids: Optional[torch.LongTensor] = None, - **model_kwargs, - ) -> tuple[torch.LongTensor, dict[str, Any]]: - # Overwritten -- Support for expanding tensors without a batch size dimension - # e.g., pixel_values, image_grid_thw, pixel_values_videos, video_grid_thw, second_per_grid_t - # pixel_values.shape[0] is sum(seqlen_images for samples) - # image_grid_thw.shape[0] is sum(num_images for samples) - - if expand_size == 1: - return input_ids, model_kwargs - - visual_keys = ["pixel_values", "image_grid_thw", "pixel_values_videos", "video_grid_thw", "second_per_grid_ts"] - - def _expand_dict_for_generation_visual(dict_to_expand): - image_grid_thw = model_kwargs.get("image_grid_thw", None) - video_grid_thw = model_kwargs.get("video_grid_thw", None) - image_nums, video_nums = self._get_image_nums_and_video_nums( - input_ids, inputs_embeds=model_kwargs.get("inputs_embeds", None) - ) - - def _repeat_interleave_samples(x, lengths, repeat_times): - samples = torch.split(x, lengths) - repeat_args = [repeat_times] + [1] * (x.dim() - 1) - result = torch.cat([sample.repeat(*repeat_args) for sample in samples], dim=0) - return result - - for key in dict_to_expand: - if key == "pixel_values": - # split images into samples - samples = torch.split(image_grid_thw, list(image_nums)) - # compute the sequence length of images for each sample - lengths = [torch.prod(sample, dim=1).sum() for sample in samples] - dict_to_expand[key] = _repeat_interleave_samples( - dict_to_expand[key], lengths=lengths, repeat_times=expand_size - ) - elif key == "image_grid_thw": - # get the num of images for each sample - lengths = list(image_nums) - dict_to_expand[key] = _repeat_interleave_samples( - dict_to_expand[key], lengths=lengths, repeat_times=expand_size - ) - elif key == "pixel_values_videos": - samples = torch.split(video_grid_thw, list(video_nums)) - lengths = [torch.prod(sample, dim=1).sum() for sample in samples] - dict_to_expand[key] = _repeat_interleave_samples( - dict_to_expand[key], lengths=lengths, repeat_times=expand_size - ) - elif key == "video_grid_thw": - lengths = list(video_nums) - dict_to_expand[key] = _repeat_interleave_samples( - dict_to_expand[key], lengths=lengths, repeat_times=expand_size - ) - elif key == "second_per_grid_ts": - dict_to_expand[key] = _repeat_interleave_samples( - dict_to_expand[key], lengths=list(video_nums), repeat_times=expand_size - ) - return dict_to_expand - - def _expand_dict_for_generation(dict_to_expand): - for key in dict_to_expand: - if ( - key != "cache_position" - and dict_to_expand[key] is not None - and isinstance(dict_to_expand[key], torch.Tensor) - and key not in visual_keys - ): - dict_to_expand[key] = dict_to_expand[key].repeat_interleave(expand_size, dim=0) - return dict_to_expand - - model_kwargs = _expand_dict_for_generation_visual(model_kwargs) - - if input_ids is not None: - input_ids = input_ids.repeat_interleave(expand_size, dim=0) - - model_kwargs = _expand_dict_for_generation(model_kwargs) - - if is_encoder_decoder: - if model_kwargs.get("encoder_outputs") is None: - raise ValueError("If `is_encoder_decoder` is True, make sure that `encoder_outputs` is defined.") - model_kwargs["encoder_outputs"] = _expand_dict_for_generation(model_kwargs["encoder_outputs"]) - - return input_ids, model_kwargs - - -__all__ = ["Qwen2_5_VLForConditionalGeneration", "Qwen2_5_VLModel", "Qwen2_5_VLPreTrainedModel", "Qwen2_5_VLTextModel"] diff --git a/Qwen-Qwen2.5-VL-3B-Instruct/builtin/cpu_and_mobile/config.json b/Qwen-Qwen2.5-VL-3B-Instruct/builtin/cpu_and_mobile/config.json new file mode 100644 index 000000000..67476cf4d --- /dev/null +++ b/Qwen-Qwen2.5-VL-3B-Instruct/builtin/cpu_and_mobile/config.json @@ -0,0 +1,67 @@ +{ + "input_model": { + "type": "CompositeModel", + "config": { + "model_path": "cpu_and_mobile/mobius_base" + } + }, + "systems": { + "local_cpu": { + "type": "LocalSystem", + "accelerators": [ + { + "device": "cpu", + "execution_providers": [ + "CPUExecutionProvider" + ] + } + ] + } + }, + "passes": { + "int4": { + "type": "OnnxBlockWiseRtnQuantization", + "block_size": 128, + "is_symmetric": true, + "accuracy_level": 4, + "save_as_external_data": true, + "external_data_name": "model.onnx.data" + } + }, + "engine": { + "host": "local_cpu", + "target": "local_cpu", + "evaluate_input_model": false, + "cache_dir": "cpu_and_mobile/cache" + }, + "max_concurrent_builds": 1, + "builds": { + "_default": { + "output_dir": "cpu_and_mobile/models" + }, + "decoder": { + "components": [ + "decoder" + ], + "pipeline": [ + "int4" + ] + }, + "vision_encoder": { + "components": [ + "vision_encoder" + ], + "pipeline": [ + "int4" + ] + }, + "embedding": { + "components": [ + "embedding" + ], + "pipeline": [ + "int4" + ] + } + } +} diff --git a/Qwen-Qwen2.5-VL-3B-Instruct/builtin/cpu_and_mobile/embedding.json b/Qwen-Qwen2.5-VL-3B-Instruct/builtin/cpu_and_mobile/embedding.json deleted file mode 100644 index 661478576..000000000 --- a/Qwen-Qwen2.5-VL-3B-Instruct/builtin/cpu_and_mobile/embedding.json +++ /dev/null @@ -1,48 +0,0 @@ -{ - "input_model": { - "type": "PyTorchModel", - "model_path": "Qwen/Qwen2.5-VL-3B-Instruct", - "model_loader": "get_embedding_model", - "model_script": "user_script.py", - "io_config": "get_embedding_io_config", - "dummy_inputs_func": "get_embedding_dummy_inputs" - }, - "passes": { - "convert": { - "type": "OnnxConversion", - "use_dynamo_exporter": false - }, - "ort": { - "type": "OrtTransformersOptimization", - "model_type": "", - "opt_level": 1, - "only_onnxruntime": true - }, - "cast": { - "type": "OnnxPeepholeOptimizer", - "onnxscript_optimize": false, - "onnxoptimizer_optimize": false, - "fuse_reshape_operations": false, - "fix_com_microsoft_opset": true, - "cast_chain_elimination": true - }, - "gemm2mm": { - "type": "GraphSurgeries", - "surgeries": [ - { - "surgeon": "GemmToMatMulAdd" - } - ] - }, - "int4": { - "type": "OnnxBlockWiseRtnQuantization", - "block_size": 128, - "is_symmetric": true, - "accuracy_level": 4, - "save_as_external_data": true, - "external_data_name": "embedding.onnx.data" - } - }, - "no_artifacts": true, - "output_dir": "cpu_and_mobile/models/embedding.onnx" -} diff --git a/Qwen-Qwen2.5-VL-3B-Instruct/builtin/cpu_and_mobile/text.json b/Qwen-Qwen2.5-VL-3B-Instruct/builtin/cpu_and_mobile/text.json deleted file mode 100644 index 812a96cca..000000000 --- a/Qwen-Qwen2.5-VL-3B-Instruct/builtin/cpu_and_mobile/text.json +++ /dev/null @@ -1,11 +0,0 @@ -{ - "input_model": { - "type": "HfModel", - "model_path": "Qwen/Qwen2.5-VL-3B-Instruct" - }, - "passes": { - "convert": { "type": "ModelBuilder", "precision": "int4", "extra_options": { "filename": "text.onnx" } } - }, - "no_artifacts": true, - "output_dir": "cpu_and_mobile/models/text.onnx" -} diff --git a/Qwen-Qwen2.5-VL-3B-Instruct/builtin/cpu_and_mobile/vision.json b/Qwen-Qwen2.5-VL-3B-Instruct/builtin/cpu_and_mobile/vision.json deleted file mode 100644 index 25b331f36..000000000 --- a/Qwen-Qwen2.5-VL-3B-Instruct/builtin/cpu_and_mobile/vision.json +++ /dev/null @@ -1,71 +0,0 @@ -{ - "input_model": { - "type": "PyTorchModel", - "model_path": "Qwen/Qwen2.5-VL-3B-Instruct", - "model_loader": "get_vision_model", - "model_script": "user_script.py", - "io_config": "get_vision_io_config", - "dummy_inputs_func": "get_vision_dummy_inputs" - }, - "passes": { - "c": { - "type": "OnnxConversion", - "use_dynamo_exporter": true - }, - "gs": { - "type": "GraphSurgeries", - "surgeries": [ - { - "surgeon": "PackedAttentionToLoopMHA" - }, - { - "surgeon": "ReciprocalMulToDiv" - }, - { - "surgeon": "RenameOutputDims", - "output_idx": 0, - "dim_idx": 0, - "dim_name": "num_logical_patches" - }, - { - "surgeon": "RenameInputDims", - "input_name": "image_grid_thw", - "dim_idx": 0, - "dim_name": "num_images" - } - ] - }, - "ort": { - "type": "OrtTransformersOptimization", - "model_type": "", - "opt_level": 1, - "only_onnxruntime": true - }, - "cast": { - "type": "OnnxPeepholeOptimizer", - "onnxscript_optimize": false, - "onnxoptimizer_optimize": false, - "fuse_reshape_operations": false, - "fix_com_microsoft_opset": true, - "cast_chain_elimination": true - }, - "gs2": { - "type": "GraphSurgeries", - "surgeries": [ - { - "surgeon": "GemmToMatMulAdd" - } - ] - }, - "int4": { - "type": "OnnxBlockWiseRtnQuantization", - "block_size": 128, - "is_symmetric": true, - "accuracy_level": 4, - "save_as_external_data": true, - "external_data_name": "vision.onnx.data" - } - }, - "no_artifacts": true, - "output_dir": "cpu_and_mobile/models/vision.onnx" -} diff --git a/Qwen-Qwen2.5-VL-3B-Instruct/builtin/cuda/config.json b/Qwen-Qwen2.5-VL-3B-Instruct/builtin/cuda/config.json new file mode 100644 index 000000000..3f4f23c67 --- /dev/null +++ b/Qwen-Qwen2.5-VL-3B-Instruct/builtin/cuda/config.json @@ -0,0 +1,77 @@ +{ + "input_model": { + "type": "CompositeModel", + "config": { + "model_path": "cuda/mobius_base" + } + }, + "systems": { + "local_gpu": { + "type": "LocalSystem", + "accelerators": [ + { + "device": "gpu", + "execution_providers": [ + "CUDAExecutionProvider" + ] + } + ] + } + }, + "passes": { + "int4": { + "type": "OnnxBlockWiseRtnQuantization", + "block_size": 128, + "is_symmetric": true, + "accuracy_level": 4, + "save_as_external_data": true, + "external_data_name": "model.onnx.data" + }, + "resave": { + "type": "OnnxPeepholeOptimizer", + "onnxscript_optimize": false, + "onnxoptimizer_optimize": false, + "fuse_reshape_operations": false, + "fix_com_microsoft_opset": false, + "cast_chain_elimination": false, + "save_as_external_data": true, + "external_data_name": "model.onnx.data" + } + }, + "engine": { + "host": "local_gpu", + "target": "local_gpu", + "evaluate_input_model": false, + "cache_dir": "cuda/cache" + }, + "max_concurrent_builds": 1, + "builds": { + "_default": { + "output_dir": "cuda/models" + }, + "decoder": { + "components": [ + "decoder" + ], + "pipeline": [ + "int4" + ] + }, + "vision_encoder": { + "components": [ + "vision_encoder" + ], + "pipeline": [ + "resave" + ] + }, + "embedding": { + "components": [ + "embedding" + ], + "pipeline": [ + "resave" + ] + } + } +} diff --git a/Qwen-Qwen2.5-VL-3B-Instruct/builtin/cuda/embedding.json b/Qwen-Qwen2.5-VL-3B-Instruct/builtin/cuda/embedding.json deleted file mode 100644 index 887bfeec2..000000000 --- a/Qwen-Qwen2.5-VL-3B-Instruct/builtin/cuda/embedding.json +++ /dev/null @@ -1,58 +0,0 @@ -{ - "input_model": { - "type": "PyTorchModel", - "model_path": "Qwen/Qwen2.5-VL-3B-Instruct", - "model_loader": "get_embedding_model", - "model_script": "user_script.py", - "io_config": "get_embedding_io_config", - "dummy_inputs_func": "get_embedding_dummy_inputs" - }, - "passes": { - "convert": { - "type": "OnnxConversion", - "use_dynamo_exporter": false - }, - "ort": { - "type": "OrtTransformersOptimization", - "model_type": "", - "opt_level": 1, - "only_onnxruntime": true - }, - "cast": { - "type": "OnnxPeepholeOptimizer", - "onnxscript_optimize": false, - "onnxoptimizer_optimize": false, - "fuse_reshape_operations": false, - "fix_com_microsoft_opset": true, - "cast_chain_elimination": true - }, - "gemm2mm": { - "type": "GraphSurgeries", - "surgeries": [ - { - "surgeon": "GemmToMatMulAdd" - } - ] - }, - "fp16": { - "type": "OnnxFloatToFloat16", - "save_as_external_data": true, - "external_data_name": "embedding.onnx.data" - } - }, - "engine": { - "target": { - "type": "LocalSystem", - "accelerators": [ - { - "device": "gpu", - "execution_providers": [ - "CUDAExecutionProvider" - ] - } - ] - } - }, - "no_artifacts": true, - "output_dir": "cuda/models/embedding.onnx" -} diff --git a/Qwen-Qwen2.5-VL-3B-Instruct/builtin/cuda/text.json b/Qwen-Qwen2.5-VL-3B-Instruct/builtin/cuda/text.json deleted file mode 100644 index d865d7f40..000000000 --- a/Qwen-Qwen2.5-VL-3B-Instruct/builtin/cuda/text.json +++ /dev/null @@ -1,17 +0,0 @@ -{ - "input_model": { - "type": "HfModel", - "model_path": "Qwen/Qwen2.5-VL-3B-Instruct" - }, - "passes": { - "convert": { "type": "ModelBuilder", "precision": "int4", "extra_options": { "filename": "text.onnx" } } - }, - "engine": { - "target": { - "type": "LocalSystem", - "accelerators": [{ "device": "gpu", "execution_providers": ["CUDAExecutionProvider"] }] - } - }, - "no_artifacts": true, - "output_dir": "cuda/models/text.onnx" -} diff --git a/Qwen-Qwen2.5-VL-3B-Instruct/builtin/cuda/vision.json b/Qwen-Qwen2.5-VL-3B-Instruct/builtin/cuda/vision.json deleted file mode 100644 index 66cfa467e..000000000 --- a/Qwen-Qwen2.5-VL-3B-Instruct/builtin/cuda/vision.json +++ /dev/null @@ -1,98 +0,0 @@ -{ - "input_model": { - "type": "PyTorchModel", - "model_path": "Qwen/Qwen2.5-VL-3B-Instruct", - "model_loader": "get_vision_model", - "model_script": "user_script.py", - "io_config": "get_vision_io_config", - "dummy_inputs_func": "get_vision_dummy_inputs" - }, - "passes": { - "c": { - "type": "OnnxConversion", - "use_dynamo_exporter": true - }, - "gs": { - "type": "GraphSurgeries", - "surgeries": [ - { - "surgeon": "PackedAttentionToLoopMHA" - }, - { - "surgeon": "ReciprocalMulToDiv" - }, - { - "surgeon": "RenameOutputDims", - "output_idx": 0, - "dim_idx": 0, - "dim_name": "num_logical_patches" - }, - { - "surgeon": "RenameInputDims", - "input_name": "image_grid_thw", - "dim_idx": 0, - "dim_name": "num_images" - } - ] - }, - "ort": { - "type": "OrtTransformersOptimization", - "model_type": "vit", - "opt_level": 2, - "only_onnxruntime": true - }, - "dedup": { - "type": "GraphSurgeries", - "surgeries": [ - { - "surgeon": "DeduplicateSubgraphInitializers" - } - ] - }, - "cast": { - "type": "OnnxPeepholeOptimizer", - "onnxscript_optimize": false, - "onnxoptimizer_optimize": false, - "fuse_reshape_operations": false, - "fix_com_microsoft_opset": true, - "cast_chain_elimination": true - }, - "fp16": { - "type": "OnnxFloatToFloat16", - "op_block_list": [ - "LayerNormalization", - "Range" - ], - "save_as_external_data": true, - "external_data_name": "vision.onnx.data" - }, - "cleanup": { - "type": "GraphSurgeries", - "surgeries": [ - { - "surgeon": "DeduplicateNodes" - }, - { - "surgeon": "RemoveMemcpy" - } - ], - "save_as_external_data": true, - "external_data_name": "vision.onnx.data" - } - }, - "engine": { - "target": { - "type": "LocalSystem", - "accelerators": [ - { - "device": "gpu", - "execution_providers": [ - "CUDAExecutionProvider" - ] - } - ] - } - }, - "no_artifacts": true, - "output_dir": "cuda/models/vision.onnx" -} diff --git a/Qwen-Qwen2.5-VL-3B-Instruct/builtin/optimize.py b/Qwen-Qwen2.5-VL-3B-Instruct/builtin/optimize.py index 22b21ed4b..e5eaf72ca 100644 --- a/Qwen-Qwen2.5-VL-3B-Instruct/builtin/optimize.py +++ b/Qwen-Qwen2.5-VL-3B-Instruct/builtin/optimize.py @@ -1,168 +1,159 @@ -"""End-to-end optimization pipeline for Qwen2.5-VL ONNX models. +"""Export and optimize Qwen2.5-VL with Olive and Mobius. -All ONNX graph transformations (Gemm→MatMul, Cast chain elimination, -INT4 quantization) are now handled by Olive passes declared in the JSON -configs. This script orchestrates the three Olive runs and writes the -GenAI runtime configuration files. - -Usage: - # Full pipeline: export + optimize + INT4 quantize (CPU) - python optimize.py --config-dir cpu_and_mobile --device cpu - - # CUDA pipeline - python optimize.py --config-dir cuda --device gpu - - # Skip export (models already exist, just regenerate configs) - python optimize.py --config-dir cpu_and_mobile --device cpu --skip-export +Each target directory contains one Olive multi-build config. Mobius first +exports the complete three-component package once; the config then applies +the target-specific pipeline to decoder, vision_encoder, and embedding. """ + import argparse import json -import logging +import shutil from pathlib import Path -logging.getLogger("onnxscript").setLevel(logging.WARNING) -logging.getLogger("onnx_ir").setLevel(logging.WARNING) -MODELS_DIR = "models" +MODEL_ID = "Qwen/Qwen2.5-VL-3B-Instruct" +COMPONENTS = ("decoder", "vision_encoder", "embedding") -# ============================================================================= -# 1. Olive Export + Optimization + Quantization (all driven by JSON configs) -# ============================================================================= +def _system_config(device: str) -> tuple[str, dict]: + if device == "gpu": + name = "local_gpu" + accelerator = { + "device": "gpu", + "execution_providers": ["CUDAExecutionProvider"], + } + else: + name = "local_cpu" + accelerator = { + "device": "cpu", + "execution_providers": ["CPUExecutionProvider"], + } + return name, { + "type": "LocalSystem", + "accelerators": [accelerator], + } -def export_models(config_dir: str): - """Run Olive for all 3 sub-models (embedding, text, vision). - The JSON configs define the full pipeline: export → graph surgeries - → ORT optimization → Cast chain elimination → Gemm→MatMul → INT4 - quantization. - """ +def export_with_mobius(config_dir: Path, device: str, output_dir: Path) -> None: + """Export one complete Mobius package through Olive.""" from olive import run - config_path = Path(config_dir) - print(f"=== Running Olive pipelines (configs from {config_path}) ===") - for config in ("embedding.json", "text.json", "vision.json"): - print(f" Running {config}...") - run(str(config_path / config)) - print() + system_name, system = _system_config(device) + precision = "fp16" if device == "gpu" else "fp32" + export_config = { + "input_model": { + "type": "HfModel", + "config": { + "model_path": MODEL_ID, + "task": "image-text-to-text", + "load_kwargs": {"trust_remote_code": True}, + }, + }, + "systems": {system_name: system}, + "passes": { + "mobius": { + "type": "MobiusBuilder", + "precision": precision, + } + }, + "engine": { + "host": system_name, + "target": system_name, + "evaluate_input_model": False, + "cache_dir": str(config_dir / "mobius_cache"), + "output_dir": str(output_dir), + }, + } + if output_dir.exists(): + shutil.rmtree(output_dir) + print(f"=== Exporting all components with Olive + Mobius ({precision}) ===") + run(export_config) -# ============================================================================= -# 2. GenAI Runtime Config Generation -# ============================================================================= -def update_genai_config(output_dir: str = MODELS_DIR, device: str = "gpu"): - """Patch genai_config.json with embedding/vision sections and processor_config.""" - config_path = Path(output_dir) / "genai_config.json" +def run_component_builds(config_path: Path, base_dir: Path, models_dir: Path) -> None: + """Run all three component builds from one Olive config.""" + from olive import run - with open(config_path) as f: + with config_path.open(encoding="utf-8") as f: config = json.load(f) - - # Provider options - if device == "gpu": - provider_options = [ - {"cuda": {"enable_cuda_graph": "0", "enable_skip_layer_norm_strict_mode": "1"}} - ] - else: - provider_options = [] - - session_options = {"log_id": "onnxruntime-genai", "provider_options": provider_options} - - # Embedding configuration - config["model"]["embedding"] = { - "filename": "embedding.onnx", - "inputs": {"input_ids": "input_ids", "image_features": "image_features"}, - "outputs": {"inputs_embeds": "inputs_embeds"}, - "session_options": session_options, - } - - # Vision configuration (Qwen2.5-VL: patch_size=14) - config["model"]["vision"] = { - "filename": "vision.onnx", - "config_filename": "processor_config.json", - "spatial_merge_size": 2, - "tokens_per_second": 2.0, - "patch_size": 14, - "window_size": 56, - "inputs": {"pixel_values": "pixel_values", "image_grid_thw": "image_grid_thw"}, - "outputs": {"image_features": "image_features"}, - "session_options": session_options, - } - - config["model"]["image_token_id"] = 151655 - config["model"]["video_token_id"] = 151656 - config["model"]["vision_start_token_id"] = 151652 - - # Fix null search params - if config["search"].get("top_k") is None: - config["search"]["top_k"] = 50 - if config["search"].get("top_p") is None: - config["search"]["top_p"] = 1.0 - - with open(config_path, "w") as f: - json.dump(config, f, indent=4) - print(f" Updated {config_path}") - - # Create processor_config.json (Qwen2.5-VL: patch_size=14, CLIP normalization) - processor_config = { - "processor": { - "name": "qwen2_5_image_processor", - "transforms": [ - {"operation": {"name": "decode_image", "type": "DecodeImage", "attrs": {"color_space": "RGB"}}}, - {"operation": {"name": "convert_to_rgb", "type": "ConvertRGB"}}, - {"operation": {"name": "resize", "type": "Resize", "attrs": { - "width": 540, "height": 360, "smart_resize": 1, - "min_pixels": 3136, "max_pixels": 12845056, "patch_size": 14, "merge_size": 2, - }}}, - {"operation": {"name": "rescale", "type": "Rescale", "attrs": { - "rescale_factor": 0.00392156862745098, - }}}, - {"operation": {"name": "normalize", "type": "Normalize", "attrs": { - "mean": [0.48145466, 0.4578275, 0.40821073], - "std": [0.26862954, 0.26130258, 0.27577711], - "qwen2_5_vl": 1, - }}}, - {"operation": {"name": "patch_image", "type": "PatchImage", "attrs": { - "patch_size": 14, "temporal_patch_size": 2, "merge_size": 2, - }}}, - ], - } - } - - processor_path = Path(output_dir) / "processor_config.json" - with open(processor_path, "w") as f: - json.dump(processor_config, f, indent=2) - print(f" Created {processor_path}") - - -# ============================================================================= -# Main -# ============================================================================= - -def main(): - parser = argparse.ArgumentParser(description="Optimize Qwen2.5-VL ONNX models") - parser.add_argument("--device", choices=["gpu", "cpu"], default="cpu", - help="Target device (default: cpu)") - parser.add_argument("--config-dir", default="cpu_and_mobile", - help="Directory containing Olive JSON configs (default: cpu_and_mobile)") - parser.add_argument("--skip-export", action="store_true", - help="Skip Olive export (models already exist)") - parser.add_argument("--models-dir", default=None, - help="Models directory (default: /models)") + config["input_model"]["config"]["model_path"] = str(base_dir) + config["builds"]["_default"]["output_dir"] = str(models_dir) + + if models_dir.exists(): + shutil.rmtree(models_dir) + print( + f"=== Running decoder, vision_encoder, and embedding builds from {config_path} ===" + ) + run(config) + + +def copy_runtime_artifacts(base_dir: Path, models_dir: Path) -> None: + """Copy Mobius-generated GenAI, tokenizer, and processor metadata.""" + models_dir.mkdir(parents=True, exist_ok=True) + for path in base_dir.iterdir(): + if path.is_file(): + shutil.copy2(path, models_dir / path.name) + + +def _validate_base_package(base_dir: Path) -> None: + missing = [ + str(base_dir / component / "model.onnx") + for component in COMPONENTS + if not (base_dir / component / "model.onnx").is_file() + ] + missing.extend( + str(base_dir / name) + for name in ("genai_config.json", "processor_config.json", "tokenizer.json") + if not (base_dir / name).is_file() + ) + if missing: + raise FileNotFoundError( + "Mobius base package is incomplete; missing " + ", ".join(missing) + ) + + +def main() -> None: + parser = argparse.ArgumentParser( + description="Optimize Qwen2.5-VL with Olive and Mobius" + ) + parser.add_argument("--device", choices=["gpu", "cpu"], default="cpu") + parser.add_argument( + "--config-dir", + default="cpu_and_mobile", + help="Target directory containing config.json", + ) + parser.add_argument( + "--skip-export", + action="store_true", + help="Reuse /mobius_base and rerun the component builds", + ) + parser.add_argument( + "--models-dir", + default=None, + help="Final package directory (default: /models)", + ) args = parser.parse_args() - models_dir = args.models_dir or str(Path(args.config_dir) / MODELS_DIR) - - # Step 1: Export + optimize + quantize (all in Olive JSON pipelines) + config_dir = Path(args.config_dir) + config_path = config_dir / "config.json" + base_dir = config_dir / "mobius_base" + models_dir = Path(args.models_dir) if args.models_dir else config_dir / "models" + + expected_device = "gpu" if config_dir.name == "cuda" else "cpu" + if args.device != expected_device: + raise ValueError( + f"{config_dir} targets {expected_device}; got --device {args.device}" + ) + if not config_path.is_file(): + raise FileNotFoundError(f"Olive multi-build config not found: {config_path}") if not args.skip_export: - export_models(args.config_dir) - - # Step 2: Generate GenAI runtime configs - print("=== Generating configs ===") - update_genai_config(output_dir=models_dir, device=args.device) - print() + export_with_mobius(config_dir, args.device, base_dir) + _validate_base_package(base_dir) + run_component_builds(config_path, base_dir, models_dir) + copy_runtime_artifacts(base_dir, models_dir) - print("Done.") + print(f"Done. ORT GenAI package: {models_dir}") if __name__ == "__main__": diff --git a/Qwen-Qwen2.5-VL-3B-Instruct/builtin/requirements.txt b/Qwen-Qwen2.5-VL-3B-Instruct/builtin/requirements.txt index 14a0a2b86..a9b489e07 100644 --- a/Qwen-Qwen2.5-VL-3B-Instruct/builtin/requirements.txt +++ b/Qwen-Qwen2.5-VL-3B-Instruct/builtin/requirements.txt @@ -1,3 +1,4 @@ -git+https://github.com/microsoft/Olive.git@main +git+https://github.com/microsoft/Olive.git@c1d19c25 +git+https://github.com/onnxruntime/mobius.git@8dd2b4e torch>=2.10.0 transformers>=4.57.0,<6.0 diff --git a/Qwen-Qwen2.5-VL-3B-Instruct/builtin/user_script.py b/Qwen-Qwen2.5-VL-3B-Instruct/builtin/user_script.py deleted file mode 100644 index e265057bc..000000000 --- a/Qwen-Qwen2.5-VL-3B-Instruct/builtin/user_script.py +++ /dev/null @@ -1,166 +0,0 @@ -import os -import sys - -# Fix Windows cp1252 encoding crash when PyTorch prints emoji in error messages -if sys.stdout.encoding and sys.stdout.encoding.lower() != "utf-8": - sys.stdout.reconfigure(encoding="utf-8", errors="replace") -if sys.stderr.encoding and sys.stderr.encoding.lower() != "utf-8": - sys.stderr.reconfigure(encoding="utf-8", errors="replace") - -import torch - -from transformers import Qwen2_5_VLConfig - -# Add current directory to sys.path to import codes module -_this_dir = os.path.dirname(os.path.abspath(__file__)) -if _this_dir not in sys.path: - sys.path.insert(0, _this_dir) - -# Import custom model from codes directory -from codes.modeling_qwen2_5_vl import Qwen2_5_VLModel - -model_name = "Qwen/Qwen2.5-VL-3B-Instruct" -config = Qwen2_5_VLConfig.from_pretrained(model_name) - - -### Embedding -# Dynamo export - -def get_embedding_model(model_path=None): - model = Qwen2_5_VLModel.from_pretrained( - model_path, - attn_implementation="sdpa", - trust_remote_code=True, - torch_dtype=torch.float32, - ) - - model.get_fused_input_embeddings, model.forward = ( - model.forward, - model.get_fused_input_embeddings, - ) - return model - -def get_embedding_io_config(model_path=None): - dynamic_axes = { - "input_ids": {0: "batch_size", 1: "sequence_length"}, - "image_features": {0: "num_logical_patches"}, - "inputs_embeds": {0: "batch_size", 1: "sequence_length"}, - } - return { - "input_names": ["input_ids", "image_features"], - "output_names": ["inputs_embeds"], - "dynamic_axes": dynamic_axes, - } - - -def get_embedding_dummy_inputs(model=None): - # assume 2 batches, each with 1 image input (3577 logical patches) - # out_hidden_size: 2048 for 3B, 3584 for 7B - batch_size, sequence_length, patches_per_image, out_hidden_size = ( - 2, - 3606, - 3577, - 2048, # 3B model hidden_size - ) - num_logical_patches = batch_size * patches_per_image - - # Qwen2.5-VL special token IDs - vision_start_token_id = config.vision_start_token_id # 151652 - vision_end_token_id = config.vision_end_token_id # 151653 - image_token_id = config.image_token_id # 151655 - - inputs = { - "input_ids": torch.randint( - low=0, - high=image_token_id, - size=(batch_size, sequence_length), - dtype=torch.int64, - ), - "image_features": torch.randn( - num_logical_patches, - out_hidden_size, - dtype=torch.float32, - ), - } - - img_start_index = 3 - img_end_index = img_start_index + patches_per_image # 3 + 3577 = 3580 - - # Fill in with image token index - inputs["input_ids"][0][2] = vision_start_token_id # <|vision_start|> - inputs["input_ids"][0][ - img_start_index:img_end_index - ] = image_token_id # <|image_pad|> - inputs["input_ids"][0][img_end_index] = vision_end_token_id # <|vision_end|> - - inputs["input_ids"][1][2] = vision_start_token_id # <|vision_start|> - inputs["input_ids"][1][ - img_start_index:img_end_index - ] = image_token_id # <|image_pad|> - inputs["input_ids"][1][img_end_index] = vision_end_token_id # <|vision_end|> - - return { - "input_ids": inputs["input_ids"], # input_ids: torch.LongTensor - "image_features": inputs["image_features"], # image_features: Optional[torch.FloatTensor] = None, - } - - -### Vision -def _reinit_inv_freq(model): - """Recompute inv_freq buffers that are missing from the HF checkpoint. - - The upstream Qwen code registers inv_freq with persistent=False, so - the buffer is never saved in the checkpoint. Our local modeling code - uses persistent=True so that torch.export captures the buffer, but - from_pretrained's fast-init (meta device) leaves it uninitialized. - Re-derive the correct values from the same formula used in __init__. - """ - rope = model.visual.rotary_pos_emb - dim = rope.inv_freq.shape[0] * 2 # original dim passed to __init__ - theta = 10000.0 - inv_freq = 1.0 / (theta ** (torch.arange(0, dim, 2, dtype=torch.float32) / dim)) - rope.inv_freq.data.copy_(inv_freq) - - -def get_vision_model(model_path=None): - model = Qwen2_5_VLModel.from_pretrained( - model_path, - attn_implementation="sdpa", - trust_remote_code=True, - torch_dtype=torch.float32, - ) - _reinit_inv_freq(model) - model.forward, model.get_image_features = model.get_image_features, model.forward - return model - -def get_vision_io_config(model_path=None): - """Vision model IO config with dynamic shapes. - - Both pixel_values and image_grid_thw have symbolic dim-0 so the model - accepts any number of patches (any image resolution) and any number of - images in a single call. The RenameInputDims graph surgery in the Olive - config labels dim-0 of image_grid_thw as 'num_images' in the final ONNX. - - Requires torch >= 2.10 for reliable dynamo export with dynamic_shapes. - """ - return { - "input_names": ["pixel_values", "image_grid_thw"], - "output_names": ["image_features"], - "dynamic_shapes": { - "pixel_values": {0: "num_patches"}, - "image_grid_thw": {0: "num_images"}, - }, - } - - -def get_vision_dummy_inputs(model=None): - """Dummy inputs for vision model export. - - Two images with the same 14x14 grid (196 patches each, 392 total) - to exercise the dynamic num_images dimension during torch.export tracing. - Qwen2.5-VL: patch_size=14, temporal_patch_size=2 → 1176 channels/patch. - """ - pixel_values = torch.randn((2 * 196, 1176), dtype=torch.float32) - pixel_values = pixel_values * (0.95 - (-1)) + (-1) - grid_thw = torch.tensor([[1, 14, 14], [1, 14, 14]], dtype=torch.int64) - return {"pixel_values": pixel_values, "image_grid_thw": grid_thw} From c7a20099f056e7c81b158b1a67803604214388e5 Mon Sep 17 00:00:00 2001 From: Xiaoyu Zhang Date: Thu, 27 Aug 2026 11:23:59 -0700 Subject: [PATCH 06/11] Add mixed-quantization Gemma4 component builds --- google-gemma-4-E2B-it/README.md | 7 +- google-gemma-4-E2B-it/multi_comp/.gitignore | 4 + google-gemma-4-E2B-it/multi_comp/README.md | 128 ++++++++---------- .../gemma4_quantize_then_export.json | 26 +++- google-gemma-4-E2B-it/multi_comp/info.yml | 3 +- 5 files changed, 88 insertions(+), 80 deletions(-) diff --git a/google-gemma-4-E2B-it/README.md b/google-gemma-4-E2B-it/README.md index 35bf6cea4..379275147 100644 --- a/google-gemma-4-E2B-it/README.md +++ b/google-gemma-4-E2B-it/README.md @@ -44,9 +44,10 @@ K-Quant (Q4_K_M) is significantly faster with GPU acceleration — install `cupy-cuda12x` for a 19–51× speedup during quantization. For a Torch-stage quantize-then-export flow, see -[`multi_comp/README.md`](multi_comp/README.md). It applies INT4 RTN to the -decoder and vision encoder in one multi-component build, saves a complete -Hugging Face checkpoint, and then exports all four components with Mobius. +[`multi_comp/README.md`](multi_comp/README.md). It applies INT4 KQuant to the +decoder and INT4 RTN to the vision encoder in two independent component builds, +automatically assembles a standard Hugging Face checkpoint, and then exports +all four components with Mobius. ## Build diff --git a/google-gemma-4-E2B-it/multi_comp/.gitignore b/google-gemma-4-E2B-it/multi_comp/.gitignore index 551fc302d..e2e7752bf 100644 --- a/google-gemma-4-E2B-it/multi_comp/.gitignore +++ b/google-gemma-4-E2B-it/multi_comp/.gitignore @@ -1,8 +1,12 @@ # Generated Hugging Face and ONNX packages gemma4_decoder_int4_hf/ gemma4_decoder_vision_int4_hf/ +gemma4_decoder_kquant_vision_rtn_hf/ +gemma4_mixed_hf/ exported_gemma4_int4_pkg/ exported_gemma4_decoder_vision_int4_pkg/ +exported_gemma4_decoder_kquant_vision_rtn_pkg/ +exported_gemma4_mixed_pkg/ # Olive caches and generated dependency lists .olive-cache/ diff --git a/google-gemma-4-E2B-it/multi_comp/README.md b/google-gemma-4-E2B-it/multi_comp/README.md index c72e76083..697c93169 100644 --- a/google-gemma-4-E2B-it/multi_comp/README.md +++ b/google-gemma-4-E2B-it/multi_comp/README.md @@ -1,92 +1,90 @@ -# Gemma 4 E2B — Quantize Then Export +# Gemma 4 E2B — Decoder KQuant + Vision RTN -This recipe quantizes the Torch decoder and vision components of -[`google/gemma-4-E2B-it`](https://huggingface.co/google/gemma-4-E2B-it) -before exporting the complete multimodal model with Mobius. +This recipe uses two independent Olive component builds for +[`google/gemma-4-E2B-it`](https://huggingface.co/google/gemma-4-E2B-it): -The flow has two explicit stages: +- `decoder`: PyTorch KQuant, asymmetric INT4, group size 32 +- `vision_encoder`: PyTorch RTN, symmetric INT4, group size 128 -1. Olive selects the Gemma 4 `decoder` and `vision_encoder` components in one - build, applies INT4 RTN to both, and saves one complete Hugging Face - directory. Audio and embedding weights remain available for the later - export. -2. `olive capture-onnx-graph --use_mobius_builder` loads that quantized - directory and exports the four-component ORT GenAI package. +Olive automatically assembles the component-only artifacts with the unchanged +audio and embedding weights into one standard Hugging Face checkpoint. Mobius +then loads that checkpoint through the ordinary +`olive capture-onnx-graph --use_mobius_builder` CLI. ## Prerequisites -This recipe requires Olive multi-build support and the current Mobius -component/quantized-checkpoint integration. Until those changes are included in -published releases, install the tested source revisions and runtime -dependencies: - ```bash -pip install "git+https://github.com/microsoft/Olive.git@6e2fe601" -pip install "git+https://github.com/onnxruntime/mobius.git@d048028" +pip install "git+https://github.com/microsoft/Olive.git@14bb7a6c" +pip install "git+https://github.com/onnxruntime/mobius.git@ea293cb" pip install transformers torch onnxruntime-genai requests -``` - -Gemma 4 is gated. Accept the model license, then authenticate after installing -`huggingface_hub` through the dependencies above: - -```bash hf auth login ``` Run the commands below from this `multi_comp` directory. -## Step 1 — Quantize the decoder and vision encoder +## Step 1 — Run and assemble both component builds ```bash olive run --config gemma4_quantize_then_export.json ``` -The build selects both components so the two sets of packed weights are saved -in the same Hugging Face checkpoint: +The config contains two disjoint builds under one shared output parent: ```json { - "components": ["decoder", "vision_encoder"], - "pipeline": ["decoder_vision_rtn"] + "builds": { + "_default": { + "output_dir": "gemma4_mixed_hf" + }, + "decoder": { + "components": ["decoder"], + "pipeline": ["decoder_kquant"] + }, + "vision": { + "components": ["vision_encoder"], + "pipeline": ["vision_rtn"] + } + } } ``` -`Rtn` performs calibration-free INT4 weight quantization with group size 128. -`quantize_vision: true` includes the vision tower and its vision-to-text -projector. The embedding table, LM head, audio encoder, and Gemma 4's -runtime-specific `per_layer_input_gate` / `per_layer_projection` modules remain -floating point. Olive saves a complete Hugging Face checkpoint, not standalone -component fragments: +Olive writes component-only shards for the optimized components and retains all +unbuilt tensors from the source checkpoint: ```text -gemma4_decoder_vision_int4_hf/ - model/ - config.json - generation_config.json - model*.safetensors - tokenizer and processor files +gemma4_mixed_hf/ + config.json + model.safetensors.index.json + model-unoptimized-*.safetensors model_config.json - footprint.json + decoder/ + component.json + model-*.safetensors + vision/ + component.json + model-*.safetensors ``` -The complete directory lets Mobius load the decoder and vision INT4 sidecars -together with the unquantized audio and multimodal embedding components. +The root is a standard HF checkpoint. Its `component_quantization` mapping +records the independent decoder and vision layouts. The LM head, embeddings, +audio encoder, and Gemma 4 `per_layer_input_gate` / +`per_layer_projection` modules remain floating point. -## Step 2 — Export all components with Mobius +## Step 2 — Export with Mobius ```bash olive capture-onnx-graph \ - --model_name_or_path gemma4_decoder_vision_int4_hf/model \ + --model_name_or_path gemma4_mixed_hf \ --use_mobius_builder \ --trust_remote_code \ --precision fp32 \ - --output_path exported_gemma4_decoder_vision_int4_pkg + --output_path exported_gemma4_mixed_pkg ``` -Mobius preserves the Olive-packed INT4 decoder and vision weights and exports: +Output: ```text -exported_gemma4_decoder_vision_int4_pkg/ +exported_gemma4_mixed_pkg/ decoder/model.onnx vision_encoder/model.onnx audio_encoder/model.onnx @@ -96,40 +94,28 @@ exported_gemma4_decoder_vision_int4_pkg/ processor and audio feature-extraction files ``` +The exported decoder contains 205 asymmetric group-32 `MatMulNBits` nodes. The +vision encoder contains 114 symmetric group-128 `MatMulNBits` nodes. Audio and +embedding remain floating point, and all 70 runtime-specific per-layer +gate/projection operations remain ordinary `MatMul`. + ## Step 3 — Inference -Use the inference entry point in the parent Gemma 4 recipe: +Text: ```bash python ../inference.py \ - --model-path exported_gemma4_decoder_vision_int4_pkg \ + --model-path exported_gemma4_mixed_pkg \ --prompt "What is the capital of France?" \ --verbose ``` -To execute the quantized vision encoder, provide an image: +Image: ```bash python ../inference.py \ - --model-path exported_gemma4_decoder_vision_int4_pkg \ + --model-path exported_gemma4_mixed_pkg \ --image path/to/image.jpg \ - --prompt "Describe this image." \ + --prompt "What animal is shown? Answer in one short sentence." \ --verbose ``` - -For CUDA inference, install `onnxruntime-genai-cuda` and change the export -precision to `fp16` on a CUDA-capable machine. The RTN stage itself can run on -CPU or CUDA. - -## Notes - -- `builds.components: ["decoder", "vision_encoder"]` scopes one RTN pass to - both selected subtrees while preserving the full Hugging Face checkpoint. -- `quantize_vision: true` quantizes the vision tower and vision-to-text - projector instead of applying RTN to the decoder only. -- `lm_head: false` and `embeds: false` avoid quantizing the tied/output tables. -- `modules_to_not_convert` keeps Gemma 4's per-layer input gate/projection in - the floating-point format expected by the current Mobius graph. -- This is intentionally a quantize-then-export flow. The existing sibling - recipes under `cpu/` and `cuda/` demonstrate export-then-ONNX-quantize flows. -- The model download and full quantization require substantial disk and memory. diff --git a/google-gemma-4-E2B-it/multi_comp/gemma4_quantize_then_export.json b/google-gemma-4-E2B-it/multi_comp/gemma4_quantize_then_export.json index dc5fc603c..a9c7f0dbd 100644 --- a/google-gemma-4-E2B-it/multi_comp/gemma4_quantize_then_export.json +++ b/google-gemma-4-E2B-it/multi_comp/gemma4_quantize_then_export.json @@ -14,7 +14,17 @@ } }, "passes": { - "decoder_vision_rtn": { + "decoder_kquant": { + "type": "KQuant", + "bits": 4, + "group_size": 32, + "sym": false, + "lm_head": false, + "embeds": false, + "quantize_vision": false, + "modules_to_not_convert": [ "per_layer_input_gate", "per_layer_projection" ] + }, + "vision_rtn": { "type": "Rtn", "bits": 4, "group_size": 128, @@ -27,10 +37,16 @@ }, "engine": { "host": "local_cpu", "target": "local_cpu", "evaluate_input_model": false, "cache_dir": "cache" }, "builds": { - "decoder_vision_int4": { - "components": [ "decoder", "vision_encoder" ], - "pipeline": [ "decoder_vision_rtn" ], - "output_dir": "gemma4_decoder_vision_int4_hf" + "_default": { + "output_dir": "gemma4_mixed_hf" + }, + "decoder": { + "components": [ "decoder" ], + "pipeline": [ "decoder_kquant" ] + }, + "vision": { + "components": [ "vision_encoder" ], + "pipeline": [ "vision_rtn" ] } } } diff --git a/google-gemma-4-E2B-it/multi_comp/info.yml b/google-gemma-4-E2B-it/multi_comp/info.yml index c45d077f4..442394cb3 100644 --- a/google-gemma-4-E2B-it/multi_comp/info.yml +++ b/google-gemma-4-E2B-it/multi_comp/info.yml @@ -3,11 +3,12 @@ keywords: - gemma4 - multimodal - multi-component + - kquant - rtn - int4 - mobius recipes: - - name: gemma4-e2b-decoder-vision-rtn-then-mobius + - name: gemma4-e2b-decoder-kquant-vision-rtn file: gemma4_quantize_then_export.json eps: - CPUExecutionProvider From eac8ad667ad08da9cfa21063e2bf9f3db704ec81 Mon Sep 17 00:00:00 2001 From: Xiaoyu Zhang Date: Thu, 27 Aug 2026 14:14:38 -0700 Subject: [PATCH 07/11] Pin tested Mobius component stack --- google-gemma-4-E2B-it/multi_comp/README.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/google-gemma-4-E2B-it/multi_comp/README.md b/google-gemma-4-E2B-it/multi_comp/README.md index 697c93169..abd806e7d 100644 --- a/google-gemma-4-E2B-it/multi_comp/README.md +++ b/google-gemma-4-E2B-it/multi_comp/README.md @@ -15,7 +15,7 @@ then loads that checkpoint through the ordinary ```bash pip install "git+https://github.com/microsoft/Olive.git@14bb7a6c" -pip install "git+https://github.com/onnxruntime/mobius.git@ea293cb" +pip install "git+https://github.com/onnxruntime/mobius.git@459fc68" pip install transformers torch onnxruntime-genai requests hf auth login ``` From ac0f75abde32677c23c2b06236a3444cdd536359 Mon Sep 17 00:00:00 2001 From: Xiaoyu Zhang Date: Thu, 27 Aug 2026 15:27:33 -0700 Subject: [PATCH 08/11] Fix Qwen3 multi-component workflow --- .../multi_comp/.gitignore | 1 + .../multi_comp/README.md | 45 +++-- Qwen-Qwen3-VL-2B-Instruct/multi_comp/info.yml | 6 +- .../multi_comp/optimize.py | 32 ++++ .../multi_comp/requirements.txt | 6 + .../multi_comp/vlm_inference.py | 30 +++- .../multi_comp/vlm_optimize_components.json | 13 +- .../LICENSE | 14 -- .../multi_comp/.gitignore | 12 -- .../multi_comp/README.md | 107 ----------- .../multi_comp/info.yml | 9 - .../multi_comp/sd3_inference.py | 169 ------------------ .../multi_comp/sd3_optimize_components.json | 90 ---------- 13 files changed, 96 insertions(+), 438 deletions(-) create mode 100644 Qwen-Qwen3-VL-2B-Instruct/multi_comp/optimize.py create mode 100644 Qwen-Qwen3-VL-2B-Instruct/multi_comp/requirements.txt delete mode 100644 stabilityai-stable-diffusion-3-medium-diffusers/LICENSE delete mode 100644 stabilityai-stable-diffusion-3-medium-diffusers/multi_comp/.gitignore delete mode 100644 stabilityai-stable-diffusion-3-medium-diffusers/multi_comp/README.md delete mode 100644 stabilityai-stable-diffusion-3-medium-diffusers/multi_comp/info.yml delete mode 100644 stabilityai-stable-diffusion-3-medium-diffusers/multi_comp/sd3_inference.py delete mode 100644 stabilityai-stable-diffusion-3-medium-diffusers/multi_comp/sd3_optimize_components.json diff --git a/Qwen-Qwen3-VL-2B-Instruct/multi_comp/.gitignore b/Qwen-Qwen3-VL-2B-Instruct/multi_comp/.gitignore index ccb2ee32b..83b45d6fa 100644 --- a/Qwen-Qwen3-VL-2B-Instruct/multi_comp/.gitignore +++ b/Qwen-Qwen3-VL-2B-Instruct/multi_comp/.gitignore @@ -1,6 +1,7 @@ # Exported ONNX packages exported_vlm_pkg/ exported_vlm_gptq_pkg/ +optimized_vlm_pkg/ # Quantized HF checkpoint vlm_decoder_gptq_hf/ diff --git a/Qwen-Qwen3-VL-2B-Instruct/multi_comp/README.md b/Qwen-Qwen3-VL-2B-Instruct/multi_comp/README.md index 2723a6b71..058c012a3 100644 --- a/Qwen-Qwen3-VL-2B-Instruct/multi_comp/README.md +++ b/Qwen-Qwen3-VL-2B-Instruct/multi_comp/README.md @@ -16,9 +16,8 @@ names**, so there is no need to memorize component names. ## Prerequisites -``` -pip install olive-ai -pip install mobius-ai +```bash +pip install -r requirements.txt ``` Exporting also needs `transformers` and access to the model on Hugging Face. @@ -30,7 +29,11 @@ Exporting also needs `transformers` and access to the model on Hugging Face. ### Step 1 — Export ``` -olive capture-onnx-graph --model_name_or_path Qwen/Qwen3-VL-2B-Instruct --use_mobius_builder --output_path exported_vlm_pkg +olive capture-onnx-graph \ + --model_name_or_path Qwen/Qwen3-VL-2B-Instruct \ + --use_mobius_builder \ + --precision fp32 \ + --output_path exported_vlm_pkg ``` Mobius exports this model as three components, each in its own subfolder: @@ -45,14 +48,19 @@ exported_vlm_pkg/ ### Step 2 — Optimize ``` -olive run --config vlm_optimize_components.json +python optimize.py ``` | component | pipeline | intent | |------------------|-----------------|-------------------------------------| | `decoder` | `dynamic_quant` | INT8-quantize the language decoder | -| `vision_encoder` | `to_fp16` | keep the vision tower in FP16 | -| `embedding` | `to_fp16` | keep the embedding in FP16 | +| `vision_encoder` | `to_fp16` | convert the vision tower to FP16 | +| `embedding` | `to_fp16` | convert the embedding to FP16 | + +The FP16 builds preserve FP32 model inputs and outputs so the embedding output +remains compatible with the dynamically quantized FP32 decoder. `optimize.py` +copies Mobius's tokenizer, processor, and `genai_config.json` artifacts into +`optimized_vlm_pkg/`, producing a directly loadable ORT GenAI package. > The three component names (`decoder`, `vision_encoder`, `embedding`) are exactly what Mobius > produces for `Qwen/Qwen3-VL-2B-Instruct`. For a different VLM, adjust the component names in the @@ -64,13 +72,13 @@ Run text generation with the exported ONNX models using **onnxruntime-genai**: ```bash # Text-only -python vlm_inference.py --prompt "The capital of France is" +python vlm_inference.py --prompt "What is the capital of France? Answer in one sentence." # With image input python vlm_inference.py --prompt "Describe this image." --image photo.jpg # Custom settings -python vlm_inference.py --model_dir exported_vlm_pkg --max_new_tokens 256 +python vlm_inference.py --model_dir optimized_vlm_pkg --max_new_tokens 256 ``` The inference script (`vlm_inference.py`) uses ORT GenAI which handles: @@ -84,7 +92,7 @@ Options: --prompt TEXT Text prompt --image PATH Optional image file for multimodal input --max_new_tokens N Maximum tokens to generate (default: 128) ---model_dir DIR Path to exported model directory (default: exported_vlm_pkg) +--model_dir DIR Path to optimized model directory (default: optimized_vlm_pkg) ``` #### Setup requirements @@ -92,27 +100,18 @@ Options: The export directory needs these files alongside the ONNX models: ``` -exported_vlm_pkg/ +optimized_vlm_pkg/ genai_config.json # Model type, I/O mappings, search config tokenizer.json # HF tokenizer tokenizer_config.json - vision_processor.json # Vision preprocessing config + processor_config.json # Vision preprocessing config decoder/model.onnx vision_encoder/model.onnx embedding/model.onnx ``` -To create the tokenizer files after export: - -```python -from transformers import AutoTokenizer -tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen3-VL-2B-Instruct", trust_remote_code=True) -tokenizer.save_pretrained("exported_vlm_pkg") -``` - -For the `genai_config.json` structure, see the -[Mobius ORT GenAI examples](https://github.com/microsoft/mobius/tree/main/examples) which write the -config automatically. +Mobius writes the tokenizer, processor, and `genai_config.json` during Step 1; +`optimize.py` carries them into the optimized package. > **Note.** Install `onnxruntime-genai` (`pip install onnxruntime-genai`) to use this script. diff --git a/Qwen-Qwen3-VL-2B-Instruct/multi_comp/info.yml b/Qwen-Qwen3-VL-2B-Instruct/multi_comp/info.yml index 26a0c37da..c121ecef3 100644 --- a/Qwen-Qwen3-VL-2B-Instruct/multi_comp/info.yml +++ b/Qwen-Qwen3-VL-2B-Instruct/multi_comp/info.yml @@ -1,13 +1,13 @@ keywords: - olive-ai recipes: - - name: qwen3vl-2B-Instruct - file: vlm_optimize_components.json + - name: qwen3vl-2B-Instruct-components + file: optimize.py eps: - CPUExecutionProvider devices: - cpu - - name: qwen3vl-2B-Instruct + - name: qwen3vl-2B-Instruct-gptq file: vlm_quantize_then_export.json eps: - CUDAExecutionProvider diff --git a/Qwen-Qwen3-VL-2B-Instruct/multi_comp/optimize.py b/Qwen-Qwen3-VL-2B-Instruct/multi_comp/optimize.py new file mode 100644 index 000000000..cd7ad6b47 --- /dev/null +++ b/Qwen-Qwen3-VL-2B-Instruct/multi_comp/optimize.py @@ -0,0 +1,32 @@ +"""Optimize the three Qwen3-VL ONNX components and assemble a runtime package.""" + +import shutil +from pathlib import Path + +from olive import run + + +BASE_PACKAGE = Path("exported_vlm_pkg") +OUTPUT_PACKAGE = Path("optimized_vlm_pkg") +CONFIG = Path("vlm_optimize_components.json") + + +def main() -> None: + if not (BASE_PACKAGE / "genai_config.json").is_file(): + raise FileNotFoundError( + "exported_vlm_pkg is missing. Export the FP32 Mobius package first." + ) + if OUTPUT_PACKAGE.exists(): + shutil.rmtree(OUTPUT_PACKAGE) + + run(str(CONFIG)) + + for path in BASE_PACKAGE.iterdir(): + if path.is_file() and path.name != "model_config.json": + shutil.copy2(path, OUTPUT_PACKAGE / path.name) + + print(f"Optimized ORT GenAI package: {OUTPUT_PACKAGE}") + + +if __name__ == "__main__": + main() diff --git a/Qwen-Qwen3-VL-2B-Instruct/multi_comp/requirements.txt b/Qwen-Qwen3-VL-2B-Instruct/multi_comp/requirements.txt new file mode 100644 index 000000000..757734eb4 --- /dev/null +++ b/Qwen-Qwen3-VL-2B-Instruct/multi_comp/requirements.txt @@ -0,0 +1,6 @@ +git+https://github.com/microsoft/Olive.git@74a029ee +git+https://github.com/onnxruntime/mobius.git@459fc68 +onnxruntime-genai +requests +torch>=2.10.0 +transformers>=4.57.0,<6.0 diff --git a/Qwen-Qwen3-VL-2B-Instruct/multi_comp/vlm_inference.py b/Qwen-Qwen3-VL-2B-Instruct/multi_comp/vlm_inference.py index 15e11aeb7..abac72191 100644 --- a/Qwen-Qwen3-VL-2B-Instruct/multi_comp/vlm_inference.py +++ b/Qwen-Qwen3-VL-2B-Instruct/multi_comp/vlm_inference.py @@ -13,17 +13,29 @@ """ import argparse +import json import os import onnxruntime_genai as og +def format_prompt(tokenizer, prompt: str, has_image: bool = False) -> str: + content = ( + [{"type": "image"}, {"type": "text", "text": prompt}] if has_image else prompt + ) + messages = [{"role": "user", "content": content}] + return tokenizer.apply_chat_template( + json.dumps(messages), + add_generation_prompt=True, + ) + + def generate_text(model_dir: str, prompt: str, max_new_tokens: int = 128) -> str: """Run text-only generation.""" model = og.Model(model_dir) tokenizer = og.Tokenizer(model) - input_ids = tokenizer.encode(prompt) + input_ids = tokenizer.encode(format_prompt(tokenizer, prompt)) params = og.GeneratorParams(model) params.set_search_options(max_length=len(input_ids) + max_new_tokens) @@ -45,14 +57,16 @@ def generate_text(model_dir: str, prompt: str, max_new_tokens: int = 128) -> str return tokenizer.decode(generated) -def generate_with_image(model_dir: str, prompt: str, image_path: str, max_new_tokens: int = 128) -> str: +def generate_with_image( + model_dir: str, prompt: str, image_path: str, max_new_tokens: int = 128 +) -> str: """Run multimodal generation with image input.""" model = og.Model(model_dir) tokenizer = og.Tokenizer(model) processor = model.create_multimodal_processor() images = og.Images.open(image_path) - inputs = processor(prompt, images=images) + inputs = processor(format_prompt(tokenizer, prompt, has_image=True), images=images) params = og.GeneratorParams(model) params.set_search_options(max_length=4096) @@ -78,9 +92,11 @@ def generate_with_image(model_dir: str, prompt: str, image_path: str, max_new_to def main(): parser = argparse.ArgumentParser(description="VLM inference with ORT GenAI") parser.add_argument("--prompt", default="The capital of France is") - parser.add_argument("--image", default=None, help="Path to an image file for vision input") + parser.add_argument( + "--image", default=None, help="Path to an image file for vision input" + ) parser.add_argument("--max_new_tokens", type=int, default=128) - parser.add_argument("--model_dir", default="exported_vlm_pkg") + parser.add_argument("--model_dir", default="optimized_vlm_pkg") args = parser.parse_args() genai_config = os.path.join(args.model_dir, "genai_config.json") @@ -101,7 +117,9 @@ def main(): print("-" * 50) if args.image: - output = generate_with_image(args.model_dir, args.prompt, args.image, args.max_new_tokens) + output = generate_with_image( + args.model_dir, args.prompt, args.image, args.max_new_tokens + ) else: output = generate_text(args.model_dir, args.prompt, args.max_new_tokens) diff --git a/Qwen-Qwen3-VL-2B-Instruct/multi_comp/vlm_optimize_components.json b/Qwen-Qwen3-VL-2B-Instruct/multi_comp/vlm_optimize_components.json index 89e4d47d5..bd2f62a02 100644 --- a/Qwen-Qwen3-VL-2B-Instruct/multi_comp/vlm_optimize_components.json +++ b/Qwen-Qwen3-VL-2B-Instruct/multi_comp/vlm_optimize_components.json @@ -6,15 +6,18 @@ "accelerators": [ { "device": "cpu", "execution_providers": [ "CPUExecutionProvider" ] } ] } }, - "passes": { "to_fp16": { "type": "OnnxFloatToFloat16" }, "dynamic_quant": { "type": "OnnxDynamicQuantization" } }, + "passes": { + "to_fp16": { "type": "OnnxFloatToFloat16", "keep_io_types": true }, + "dynamic_quant": { "type": "OnnxDynamicQuantization" } + }, "engine": { "host": "local_system", "target": "local_system", "evaluate_input_model": false, "cache_dir": "cache" }, "builds": { - "decoder": { "components": [ "decoder" ], "pipeline": [ "dynamic_quant" ], "output_dir": "out/decoder" }, + "_default": { "output_dir": "optimized_vlm_pkg" }, + "decoder": { "components": [ "decoder" ], "pipeline": [ "dynamic_quant" ] }, "vision_encoder": { "components": [ "vision_encoder" ], - "pipeline": [ "to_fp16" ], - "output_dir": "out/vision_encoder" + "pipeline": [ "to_fp16" ] }, - "embedding": { "components": [ "embedding" ], "pipeline": [ "to_fp16" ], "output_dir": "out/embedding" } + "embedding": { "components": [ "embedding" ], "pipeline": [ "to_fp16" ] } } } diff --git a/stabilityai-stable-diffusion-3-medium-diffusers/LICENSE b/stabilityai-stable-diffusion-3-medium-diffusers/LICENSE deleted file mode 100644 index d7a6188fb..000000000 --- a/stabilityai-stable-diffusion-3-medium-diffusers/LICENSE +++ /dev/null @@ -1,14 +0,0 @@ -Stable Diffusion 3 Medium is released under the Stability AI Community License. - -The model "stabilityai/stable-diffusion-3-medium-diffusers" is a gated model. You must -review and accept its license on Hugging Face before downloading or using the weights: - - https://huggingface.co/stabilityai/stable-diffusion-3-medium-diffusers - -Full license text: - - https://stability.ai/license - -The recipes in this directory are provided by the olive-recipes project under the -repository's root LICENSE. The license referenced above governs use of the Stable -Diffusion 3 model weights only. diff --git a/stabilityai-stable-diffusion-3-medium-diffusers/multi_comp/.gitignore b/stabilityai-stable-diffusion-3-medium-diffusers/multi_comp/.gitignore deleted file mode 100644 index 0bdc513bc..000000000 --- a/stabilityai-stable-diffusion-3-medium-diffusers/multi_comp/.gitignore +++ /dev/null @@ -1,12 +0,0 @@ -# Exported ONNX packages -exported_pkg/ -exported_sd3_full2/ - -# Optimized components -out/ - -# Olive cache -cache/ - -# Generated images -*.png diff --git a/stabilityai-stable-diffusion-3-medium-diffusers/multi_comp/README.md b/stabilityai-stable-diffusion-3-medium-diffusers/multi_comp/README.md deleted file mode 100644 index 4ba1ab3fc..000000000 --- a/stabilityai-stable-diffusion-3-medium-diffusers/multi_comp/README.md +++ /dev/null @@ -1,107 +0,0 @@ -# Stable Diffusion 3 Medium — Multi-Component Optimization - -This recipe demonstrates a **multi-component flow** for -[Stable Diffusion 3 Medium](https://huggingface.co/stabilityai/stable-diffusion-3-medium-diffusers): -export the pipeline to ONNX once with the Mobius builder, then run a single Olive config whose -`builds` apply a **different pipeline to each component**. - -The flow is two explicit steps: - -1. **Export** the model to a directory of per-component ONNX subfolders using the Olive CLI with the - Mobius builder. -2. **Optimize** by pointing an Olive config at that directory; each component subfolder becomes a - selectable component that a `build` can target. - -There is no need to memorize component names: each exported component lives in its own folder, and -Olive loads the export directory as a `CompositeModel` whose **component names are the subfolder -names**. - -## Prerequisites - -``` -pip install olive-ai -pip install mobius-ai -``` - -Exporting a diffusion pipeline also needs `diffusers`/`transformers` and access to the model on -Hugging Face (Stable Diffusion 3 is a gated model — accept its license and `huggingface-cli login` -first). - -## Step 1 — Export with the CLI - -``` -olive capture-onnx-graph --model_name_or_path stabilityai/stable-diffusion-3-medium-diffusers --use_mobius_builder --output_path exported_pkg -``` - -Mobius exports each neural-network component to its own subfolder: - -``` -exported_pkg/ - text_encoder/model.onnx # CLIP-L text encoder - text_encoder_2/model.onnx # CLIP-G text encoder - text_encoder_3/model.onnx # T5-XXL text encoder - transformer/model.onnx # MMDiT denoising backbone - vae_encoder/model.onnx - vae_decoder/model.onnx -``` - -> **Note.** The exact subfolders depend on the pipeline; the optimize config below only -> needs `builds` for the components you actually want to optimize. - -## Step 2 — Optimize each component - -Run from the directory that contains `exported_pkg/`: - -``` -olive run --config sd3_optimize_components.json -``` - -This applies a different pipeline per component: - -| component | pipeline | intent | -|------------------|------------------------------|--------------------------------------------| -| `transformer` | `OrtTransformersOptimization`| FP16-optimize the heavy denoising backbone | -| `vae_encoder` | `OrtTransformersOptimization`| FP16-optimize the VAE encoder | -| `vae_decoder` | `OrtTransformersOptimization`| FP16-optimize the VAE decoder | - -Output: - -``` -out/transformer/ # optimized transformer -out/vae_encoder/ # optimized VAE encoder -out/vae_decoder/ # optimized VAE decoder -``` - -Each build writes one optimized component; components without a build stay as exported. - -## Step 3 — Inference - -Run end-to-end image generation with the exported ONNX models: - -``` -python sd3_inference.py --prompt "A photo of a cat sitting on a windowsill" --steps 28 --output result.png -``` - -The inference script (`sd3_inference.py`) uses: -- **Text encoding**: ONNX Runtime with exported CLIP-L, CLIP-G, and T5-XXL encoders (run once) -- **Denoising**: ONNX Runtime with the exported SD3 transformer (28 steps) -- **VAE decoding**: ONNX Runtime with the exported VAE decoder - -Options: -``` ---prompt TEXT Text prompt for image generation ---steps N Number of denoising steps (default: 28) ---seed N Random seed (default: 42) ---output PATH Output image path (default: sd3_output.png) ---onnx_dir DIR Path to exported model directory (default: exported_sd3_full2) -``` - -> **Note.** SD3 is a gated model — you need `huggingface-cli login` or set `HF_TOKEN` to export. -> The tokenizers (CLIP and T5) still run via the `transformers` library. - -## Notes - -- The passes here are **illustrative**. Swap in `OnnxStaticQuantization` (with a `data_config`), - `OnnxDynamicQuantization`, or other ONNX passes for production-quality optimization. -- `builds.components` selects which exported components to optimize. Only the components with a build - are touched; the rest remain as exported. diff --git a/stabilityai-stable-diffusion-3-medium-diffusers/multi_comp/info.yml b/stabilityai-stable-diffusion-3-medium-diffusers/multi_comp/info.yml deleted file mode 100644 index bae3dc157..000000000 --- a/stabilityai-stable-diffusion-3-medium-diffusers/multi_comp/info.yml +++ /dev/null @@ -1,9 +0,0 @@ -keywords: - - olive-ai -recipes: - - name: stable-diffusion-3-medium-multi-component - file: sd3_optimize_components.json - eps: - - CUDAExecutionProvider - devices: - - gpu diff --git a/stabilityai-stable-diffusion-3-medium-diffusers/multi_comp/sd3_inference.py b/stabilityai-stable-diffusion-3-medium-diffusers/multi_comp/sd3_inference.py deleted file mode 100644 index 0c1e79fad..000000000 --- a/stabilityai-stable-diffusion-3-medium-diffusers/multi_comp/sd3_inference.py +++ /dev/null @@ -1,169 +0,0 @@ -#!/usr/bin/env python -"""SD3 end-to-end inference using all ONNX components (text encoders + transformer + VAE). - -Usage: - python sd3_inference.py --prompt "A photo of a cat sitting on a windowsill" - python sd3_inference.py --prompt "A futuristic city" --steps 50 --output city.png -""" - -import argparse -import os - -import numpy as np -import onnxruntime as ort -import torch -from diffusers import FlowMatchEulerDiscreteScheduler -from PIL import Image -from transformers import CLIPTokenizer, T5TokenizerFast - -MODEL_ID = "stabilityai/stable-diffusion-3-medium-diffusers" -ONNX_DIR = "exported_sd3_full2" - - -def encode_text(prompt: str, onnx_dir: str, model_id: str) -> tuple[np.ndarray, np.ndarray]: - """Encode prompt using ONNX CLIP-L, CLIP-G, and T5-XXL text encoders. - - Returns: - encoder_hidden_states: [1, 410, 4096] - pooled_projections: [1, 2048] - - """ - # Load tokenizers (lightweight, no model weights) - tokenizer_l = CLIPTokenizer.from_pretrained(model_id, subfolder="tokenizer") - tokenizer_g = CLIPTokenizer.from_pretrained(model_id, subfolder="tokenizer_2") - tokenizer_t5 = T5TokenizerFast.from_pretrained(model_id, subfolder="tokenizer_3") - - # Load ONNX sessions - sess_l = ort.InferenceSession(os.path.join(onnx_dir, "text_encoder", "model.onnx")) - sess_g = ort.InferenceSession(os.path.join(onnx_dir, "text_encoder_2", "model.onnx")) - sess_t5 = ort.InferenceSession(os.path.join(onnx_dir, "text_encoder_3", "model.onnx")) - - # CLIP-L - tokens_l = tokenizer_l(prompt, padding="max_length", max_length=77, return_tensors="np", truncation=True) - out_l = sess_l.run( - None, - { - "input_ids": tokens_l["input_ids"].astype(np.int64), - "attention_mask": tokens_l["attention_mask"].astype(np.int64), - }, - ) - clip_l_hidden = out_l[0] # last_hidden_state [1, 77, 768] - clip_l_pooled = out_l[1] # text_embeds [1, 768] - - # CLIP-G - tokens_g = tokenizer_g(prompt, padding="max_length", max_length=77, return_tensors="np", truncation=True) - out_g = sess_g.run( - None, - { - "input_ids": tokens_g["input_ids"].astype(np.int64), - "attention_mask": tokens_g["attention_mask"].astype(np.int64), - }, - ) - clip_g_hidden = out_g[0] # last_hidden_state [1, 77, 1280] - clip_g_pooled = out_g[1] # text_embeds [1, 1280] - - # T5-XXL - tokens_t5 = tokenizer_t5(prompt, padding="max_length", max_length=256, return_tensors="np", truncation=True) - out_t5 = sess_t5.run(None, {"input_ids": tokens_t5["input_ids"].astype(np.int64)}) - t5_hidden = out_t5[0] # last_hidden_state [1, 256, 4096] - - # Pad CLIP outputs to 4096 and concatenate - clip_l_padded = np.pad(clip_l_hidden, ((0, 0), (0, 0), (0, 4096 - 768))) # [1, 77, 4096] - clip_g_padded = np.pad(clip_g_hidden, ((0, 0), (0, 0), (0, 4096 - 1280))) # [1, 77, 4096] - encoder_hidden_states = np.concatenate([clip_l_padded, clip_g_padded, t5_hidden], axis=1) # [1, 410, 4096] - pooled_projections = np.concatenate([clip_l_pooled, clip_g_pooled], axis=-1) # [1, 2048] - - return encoder_hidden_states.astype(np.float32), pooled_projections.astype(np.float32) - - -def denoise( - onnx_dir: str, - encoder_hidden_states: np.ndarray, - pooled_projections: np.ndarray, - scheduler: FlowMatchEulerDiscreteScheduler, - latent_shape: tuple = (1, 16, 64, 64), - seed: int = 42, -) -> torch.Tensor: - """Run the denoising loop using the ONNX transformer.""" - sess = ort.InferenceSession(os.path.join(onnx_dir, "transformer", "model.onnx")) - - torch.manual_seed(seed) - latents = torch.randn(latent_shape) - - for i, t in enumerate(scheduler.timesteps): - noise_pred = sess.run( - None, - { - "sample": latents.numpy(), - "timestep": np.array([t.item()], dtype=np.int64), - "encoder_hidden_states": encoder_hidden_states, - "pooled_projections": pooled_projections, - }, - )[0] - latents = scheduler.step(torch.from_numpy(noise_pred), t, latents, return_dict=False)[0] - if i % 7 == 0: - print(f" Step {i}/{len(scheduler.timesteps)}, t={t.item():.1f}") - - return latents - - -def decode_latents(latents: torch.Tensor, onnx_dir: str) -> np.ndarray: - """Decode latents to image using the ONNX VAE decoder.""" - sess = ort.InferenceSession(os.path.join(onnx_dir, "vae_decoder", "model.onnx")) - - # SD3 VAE scaling: latents / scaling_factor + shift_factor - # SD3 defaults: scaling_factor=1.5305, shift_factor=0.0609 - scaling_factor = 1.5305 - shift_factor = 0.0609 - latents_scaled = latents / scaling_factor + shift_factor - - output = sess.run(None, {"latent_sample": latents_scaled.numpy()})[0] - # output: [1, 3, H, W] in [-1, 1] - image = (output / 2 + 0.5).clip(0, 1) - image = np.transpose(image[0], (1, 2, 0)) # [H, W, 3] - return (image * 255).astype(np.uint8) - - -def main(): - parser = argparse.ArgumentParser(description="SD3 all-ONNX inference") - parser.add_argument("--prompt", default="A photo of a cat sitting on a windowsill") - parser.add_argument("--steps", type=int, default=28) - parser.add_argument("--seed", type=int, default=42) - parser.add_argument("--output", default="sd3_output.png") - parser.add_argument("--model_id", default=MODEL_ID) - parser.add_argument("--onnx_dir", default=ONNX_DIR) - args = parser.parse_args() - - # Verify exported model exists - transformer_path = os.path.join(args.onnx_dir, "transformer", "model.onnx") - if not os.path.exists(transformer_path): - print(f"Error: ONNX model not found at {args.onnx_dir}/") - print( - "Run: olive capture-onnx-graph --model_name_or_path " - "stabilityai/stable-diffusion-3-medium-diffusers " - "--use_mobius_builder --output_path exported_sd3_full2" - ) - return - - print(f"Prompt: {args.prompt}") - print(f"Steps: {args.steps}, Seed: {args.seed}") - print(f"ONNX dir: {args.onnx_dir}") - - print("\n1. Encoding text (ONNX CLIP-L + CLIP-G + T5-XXL)...") - encoder_hidden_states, pooled_projections = encode_text(args.prompt, args.onnx_dir, args.model_id) - print(f" encoder_hidden_states: {encoder_hidden_states.shape}") - print(f" pooled_projections: {pooled_projections.shape}") - - print("\n2. Denoising (ONNX SD3 transformer)...") - scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained(args.model_id, subfolder="scheduler") - scheduler.set_timesteps(args.steps) - latents = denoise(args.onnx_dir, encoder_hidden_states, pooled_projections, scheduler, seed=args.seed) - - print("\n3. Decoding latents (ONNX VAE decoder)...") - image = decode_latents(latents, args.onnx_dir) - Image.fromarray(image).save(args.output) - print(f"\nSaved: {args.output} ({image.shape[1]}x{image.shape[0]})") - - -if __name__ == "__main__": - main() diff --git a/stabilityai-stable-diffusion-3-medium-diffusers/multi_comp/sd3_optimize_components.json b/stabilityai-stable-diffusion-3-medium-diffusers/multi_comp/sd3_optimize_components.json deleted file mode 100644 index 8551d5325..000000000 --- a/stabilityai-stable-diffusion-3-medium-diffusers/multi_comp/sd3_optimize_components.json +++ /dev/null @@ -1,90 +0,0 @@ -{ - "input_model": { "type": "CompositeModel", "config": { "model_path": "exported_pkg" } }, - "systems": { - "local_system": { - "type": "LocalSystem", - "accelerators": [ { "device": "gpu", "execution_providers": [ "CUDAExecutionProvider" ] } ] - } - }, - "passes": { - "optimize_transformer": { - "type": "OrtTransformersOptimization", - "model_type": "unet", - "opt_level": 0, - "float16": true, - "use_gpu": true, - "keep_io_types": false, - "optimization_options": { - "enable_gelu": true, - "enable_layer_norm": true, - "enable_attention": true, - "use_multi_head_attention": true, - "enable_skip_layer_norm": false, - "enable_embed_layer_norm": true, - "enable_bias_skip_layer_norm": false, - "enable_bias_gelu": true, - "enable_gelu_approximation": false, - "enable_qordered_matmul": false, - "enable_shape_inference": true, - "enable_gemm_fast_gelu": false, - "enable_nhwc_conv": false, - "enable_group_norm": true, - "enable_bias_splitgelu": false, - "enable_packed_qkv": true, - "enable_packed_kv": true, - "enable_bias_add": false, - "group_norm_channels_last": false - }, - "force_fp32_ops": [ "RandomNormalLike" ] - }, - "optimize_vae": { - "type": "OrtTransformersOptimization", - "model_type": "vae", - "opt_level": 0, - "float16": true, - "use_gpu": true, - "keep_io_types": false, - "optimization_options": { - "enable_gelu": true, - "enable_layer_norm": true, - "enable_attention": true, - "use_multi_head_attention": true, - "enable_skip_layer_norm": false, - "enable_embed_layer_norm": true, - "enable_bias_skip_layer_norm": false, - "enable_bias_gelu": true, - "enable_gelu_approximation": false, - "enable_qordered_matmul": false, - "enable_shape_inference": true, - "enable_gemm_fast_gelu": false, - "enable_nhwc_conv": false, - "enable_group_norm": true, - "enable_bias_splitgelu": false, - "enable_packed_qkv": true, - "enable_packed_kv": true, - "enable_bias_add": false, - "group_norm_channels_last": false - }, - "force_fp32_ops": [ "RandomNormalLike" ], - "force_fp16_inputs": { "GroupNorm": [ 0, 1, 2 ] } - } - }, - "engine": { "host": "local_system", "target": "local_system", "evaluate_input_model": false, "cache_dir": "cache" }, - "builds": { - "transformer": { - "components": [ "transformer" ], - "pipeline": [ "optimize_transformer" ], - "output_dir": "out/transformer" - }, - "vae_encoder": { - "components": [ "vae_encoder" ], - "pipeline": [ "optimize_vae" ], - "output_dir": "out/vae_encoder" - }, - "vae_decoder": { - "components": [ "vae_decoder" ], - "pipeline": [ "optimize_vae" ], - "output_dir": "out/vae_decoder" - } - } -} From 8062c910334be2bc5505c7a8f46bc489f13eeed8 Mon Sep 17 00:00:00 2001 From: Xiaoyu Zhang Date: Thu, 27 Aug 2026 15:48:45 -0700 Subject: [PATCH 09/11] update recipe --- google-gemma-4-E2B-it/multi_comp/README.md | 10 ---------- .../multi_comp/gemma4_quantize_then_export.json | 7 ------- 2 files changed, 17 deletions(-) diff --git a/google-gemma-4-E2B-it/multi_comp/README.md b/google-gemma-4-E2B-it/multi_comp/README.md index abd806e7d..6ba343f1b 100644 --- a/google-gemma-4-E2B-it/multi_comp/README.md +++ b/google-gemma-4-E2B-it/multi_comp/README.md @@ -65,11 +65,6 @@ gemma4_mixed_hf/ model-*.safetensors ``` -The root is a standard HF checkpoint. Its `component_quantization` mapping -records the independent decoder and vision layouts. The LM head, embeddings, -audio encoder, and Gemma 4 `per_layer_input_gate` / -`per_layer_projection` modules remain floating point. - ## Step 2 — Export with Mobius ```bash @@ -94,11 +89,6 @@ exported_gemma4_mixed_pkg/ processor and audio feature-extraction files ``` -The exported decoder contains 205 asymmetric group-32 `MatMulNBits` nodes. The -vision encoder contains 114 symmetric group-128 `MatMulNBits` nodes. Audio and -embedding remain floating point, and all 70 runtime-specific per-layer -gate/projection operations remain ordinary `MatMul`. - ## Step 3 — Inference Text: diff --git a/google-gemma-4-E2B-it/multi_comp/gemma4_quantize_then_export.json b/google-gemma-4-E2B-it/multi_comp/gemma4_quantize_then_export.json index a9c7f0dbd..d0ac2c95e 100644 --- a/google-gemma-4-E2B-it/multi_comp/gemma4_quantize_then_export.json +++ b/google-gemma-4-E2B-it/multi_comp/gemma4_quantize_then_export.json @@ -7,12 +7,6 @@ "load_kwargs": { "torch_dtype": "float16", "trust_remote_code": true } } }, - "systems": { - "local_cpu": { - "type": "LocalSystem", - "accelerators": [ { "device": "cpu", "execution_providers": [ "CPUExecutionProvider" ] } ] - } - }, "passes": { "decoder_kquant": { "type": "KQuant", @@ -35,7 +29,6 @@ "modules_to_not_convert": [ "per_layer_input_gate", "per_layer_projection" ] } }, - "engine": { "host": "local_cpu", "target": "local_cpu", "evaluate_input_model": false, "cache_dir": "cache" }, "builds": { "_default": { "output_dir": "gemma4_mixed_hf" From a9e0cdf0ccbe9b99a7342dc3f41f6689aaeb6c2d Mon Sep 17 00:00:00 2001 From: Xiaoyu Zhang Date: Tue, 1 Sep 2026 13:54:28 -0700 Subject: [PATCH 10/11] Opt in to Gemma4 component assembly --- .../multi_comp/gemma4_quantize_then_export.json | 1 + 1 file changed, 1 insertion(+) diff --git a/google-gemma-4-E2B-it/multi_comp/gemma4_quantize_then_export.json b/google-gemma-4-E2B-it/multi_comp/gemma4_quantize_then_export.json index d0ac2c95e..0ea283fb8 100644 --- a/google-gemma-4-E2B-it/multi_comp/gemma4_quantize_then_export.json +++ b/google-gemma-4-E2B-it/multi_comp/gemma4_quantize_then_export.json @@ -29,6 +29,7 @@ "modules_to_not_convert": [ "per_layer_input_gate", "per_layer_projection" ] } }, + "assemble_components": true, "builds": { "_default": { "output_dir": "gemma4_mixed_hf" From 09852fb2ecc91f0d656bd83cd4816e67be6b3be7 Mon Sep 17 00:00:00 2001 From: Xiaoyu Zhang Date: Wed, 2 Sep 2026 15:28:07 -0700 Subject: [PATCH 11/11] Use workflow output for Gemma4 assembly --- .../multi_comp/gemma4_quantize_then_export.json | 7 +++---- 1 file changed, 3 insertions(+), 4 deletions(-) diff --git a/google-gemma-4-E2B-it/multi_comp/gemma4_quantize_then_export.json b/google-gemma-4-E2B-it/multi_comp/gemma4_quantize_then_export.json index 0ea283fb8..2dd189ab1 100644 --- a/google-gemma-4-E2B-it/multi_comp/gemma4_quantize_then_export.json +++ b/google-gemma-4-E2B-it/multi_comp/gemma4_quantize_then_export.json @@ -29,11 +29,10 @@ "modules_to_not_convert": [ "per_layer_input_gate", "per_layer_projection" ] } }, - "assemble_components": true, + "engine": { + "output_dir": "gemma4_mixed_hf" + }, "builds": { - "_default": { - "output_dir": "gemma4_mixed_hf" - }, "decoder": { "components": [ "decoder" ], "pipeline": [ "decoder_kquant" ]