From 03a9b370e39dae9905c04b044d6c7ef1bdbb03b7 Mon Sep 17 00:00:00 2001 From: hualxie Date: Tue, 16 Jun 2026 18:04:24 +0800 Subject: [PATCH 1/8] 1st iter --- .aitk/configs/checks.json | 26 +- .aitk/configs/model_list.json | 20 +- .aitk/docs/guide/ModelList.md | 1 + Qwen-Qwen3.5-2B/aitk/.gitignore | 5 + Qwen-Qwen3.5-2B/aitk/README.md | 37 ++ Qwen-Qwen3.5-2B/aitk/_copy.json.config | 8 + Qwen-Qwen3.5-2B/aitk/codes/__init__.py | 0 .../aitk/codes/modeling_qwen3_5.py | 511 ++++++++++++++++++ Qwen-Qwen3.5-2B/aitk/embedding.json | 58 ++ Qwen-Qwen3.5-2B/aitk/inference_model.json | 31 ++ Qwen-Qwen3.5-2B/aitk/inference_sample.ipynb | 118 ++++ Qwen-Qwen3.5-2B/aitk/info.yml | 13 + Qwen-Qwen3.5-2B/aitk/model_project.config | 12 + .../aitk/qwen_trtrtx_workflow.json | 31 ++ .../aitk/qwen_trtrtx_workflow.json.config | 84 +++ Qwen-Qwen3.5-2B/aitk/qwen_trtrtx_workflow.py | 227 ++++++++ Qwen-Qwen3.5-2B/aitk/requirements.txt | 5 + Qwen-Qwen3.5-2B/aitk/text.json | 43 ++ Qwen-Qwen3.5-2B/aitk/user_script.py | 133 +++++ Qwen-Qwen3.5-2B/aitk/vision.json | 89 +++ Qwen-Qwen3.5-2B/aitk/winml.py | 58 ++ 21 files changed, 1495 insertions(+), 15 deletions(-) create mode 100644 Qwen-Qwen3.5-2B/aitk/.gitignore create mode 100644 Qwen-Qwen3.5-2B/aitk/README.md create mode 100644 Qwen-Qwen3.5-2B/aitk/_copy.json.config create mode 100644 Qwen-Qwen3.5-2B/aitk/codes/__init__.py create mode 100644 Qwen-Qwen3.5-2B/aitk/codes/modeling_qwen3_5.py create mode 100644 Qwen-Qwen3.5-2B/aitk/embedding.json create mode 100644 Qwen-Qwen3.5-2B/aitk/inference_model.json create mode 100644 Qwen-Qwen3.5-2B/aitk/inference_sample.ipynb create mode 100644 Qwen-Qwen3.5-2B/aitk/info.yml create mode 100644 Qwen-Qwen3.5-2B/aitk/model_project.config create mode 100644 Qwen-Qwen3.5-2B/aitk/qwen_trtrtx_workflow.json create mode 100644 Qwen-Qwen3.5-2B/aitk/qwen_trtrtx_workflow.json.config create mode 100644 Qwen-Qwen3.5-2B/aitk/qwen_trtrtx_workflow.py create mode 100644 Qwen-Qwen3.5-2B/aitk/requirements.txt create mode 100644 Qwen-Qwen3.5-2B/aitk/text.json create mode 100644 Qwen-Qwen3.5-2B/aitk/user_script.py create mode 100644 Qwen-Qwen3.5-2B/aitk/vision.json create mode 100644 Qwen-Qwen3.5-2B/aitk/winml.py diff --git a/.aitk/configs/checks.json b/.aitk/configs/checks.json index 4c4fc46f4..544545274 100644 --- a/.aitk/configs/checks.json +++ b/.aitk/configs/checks.json @@ -1,19 +1,19 @@ { - "configCheck": 173, - "copyCheck": 183, + "configCheck": 174, + "copyCheck": 184, "executePatchPyCheck": 0, - "executeRuntimeCheck": 106, + "executeRuntimeCheck": 107, "extensionCheck": 2, - "gitignoreCheck": 44, - "inferenceModelCheck": 25, - "ipynbCheck": 51, - "licenseCheck": 41, - "modelProjectCheck": 46, - "oliveCheck": 82, - "oliveJsonCheck": 173, - "pathCheck": 1463, - "requirementsCheck": 44, + "gitignoreCheck": 45, + "inferenceModelCheck": 26, + "ipynbCheck": 52, + "licenseCheck": 42, + "modelProjectCheck": 47, + "oliveCheck": 85, + "oliveJsonCheck": 174, + "pathCheck": 1468, + "requirementsCheck": 45, "templateCheck": 3, "venvRequirementsCheck": 23, - "winmlCopyCheck": 38 + "winmlCopyCheck": 39 } diff --git a/.aitk/configs/model_list.json b/.aitk/configs/model_list.json index 672d1cd26..858d7cf5b 100644 --- a/.aitk/configs/model_list.json +++ b/.aitk/configs/model_list.json @@ -485,7 +485,7 @@ "relativePath": "sam-vit-base/aitk", "version": 2, "pipeline_tags": [ - "fill-mask" + "mask-generation" ] }, { @@ -501,7 +501,7 @@ "relativePath": "sam2.1-hiera-small/aitk", "version": 2, "pipeline_tags": [ - "fill-mask" + "mask-generation" ] }, { @@ -837,6 +837,22 @@ "text-generation" ] }, + { + "displayName": "Qwen/Qwen3.5-2B", + "icon": "qwen", + "modelLink": "https://huggingface.co/Qwen/Qwen3.5-2B", + "id": "huggingface/Qwen/Qwen3.5-2B", + "runtimes": [ + "NvidiaTRTRTX" + ], + "architecture": "Transformer", + "status": "Ready", + "relativePath": "Qwen-Qwen3.5-2B/aitk", + "version": 1, + "pipeline_tags": [ + "image-text-to-text" + ] + }, { "displayName": "sd2-community/stable-diffusion-2-1", "icon": "HuggingFace", diff --git a/.aitk/docs/guide/ModelList.md b/.aitk/docs/guide/ModelList.md index 77fabca47..169eac43c 100644 --- a/.aitk/docs/guide/ModelList.md +++ b/.aitk/docs/guide/ModelList.md @@ -30,6 +30,7 @@ | [Qwen2.5 Coder 14B Instruct](https://huggingface.co/Qwen/Qwen2.5-Coder-14B-Instruct) | [NVIDIA TensorRT for RTX](../../../Qwen-Qwen2.5-Coder-14B-Instruct/aitk/qwen2_5_trtrtx.json), [Intel CPU](../../../Qwen-Qwen2.5-Coder-14B-Instruct/aitk/qwen2_5_ov_config.json), [Intel GPU](../../../Qwen-Qwen2.5-Coder-14B-Instruct/aitk/qwen2_5_ov_config.json), [Intel NPU](../../../Qwen-Qwen2.5-Coder-14B-Instruct/aitk/qwen2_5_ov_npu_config.json) | | [Qwen2.5 Coder 3B Instruct](https://huggingface.co/Qwen/Qwen2.5-Coder-3B-Instruct) | [Intel CPU](../../../Qwen-Qwen2.5-Coder-3B-Instruct/aitk/qwen2_5_ov_config.json), [Intel GPU](../../../Qwen-Qwen2.5-Coder-3B-Instruct/aitk/qwen2_5_ov_config.json), [Intel NPU](../../../Qwen-Qwen2.5-Coder-3B-Instruct/aitk/qwen2_5_ov_npu_config.json) | | [Qwen2.5 Coder 7B Instruct](https://huggingface.co/Qwen/Qwen2.5-Coder-7B-Instruct) | [AMD NPU](../../../Qwen-Qwen2.5-Coder-7B-Instruct/aitk/qwen2_5_vitis_ai_config.json), [NVIDIA TensorRT for RTX](../../../Qwen-Qwen2.5-Coder-7B-Instruct/aitk/qwen2_5_trtrtx.json), [Intel CPU](../../../Qwen-Qwen2.5-Coder-7B-Instruct/aitk/qwen2_5_ov_config.json), [Intel GPU](../../../Qwen-Qwen2.5-Coder-7B-Instruct/aitk/qwen2_5_ov_config.json), [Intel NPU](../../../Qwen-Qwen2.5-Coder-7B-Instruct/aitk/qwen2_5_ov_npu_config.json) | +| [Qwen3.5 2B](https://huggingface.co/Qwen/Qwen3.5-2B) | [NVIDIA TensorRT for RTX](../../../Qwen-Qwen3.5-2B/aitk/qwen_trtrtx_workflow.json) | ## Non-LLM Models | Model Name | Supported Runtimes | diff --git a/Qwen-Qwen3.5-2B/aitk/.gitignore b/Qwen-Qwen3.5-2B/aitk/.gitignore new file mode 100644 index 000000000..48c038823 --- /dev/null +++ b/Qwen-Qwen3.5-2B/aitk/.gitignore @@ -0,0 +1,5 @@ +__pycache__ +/cache +/history/*/* +!/history/*/history.config +!/history/*/olive_config.json diff --git a/Qwen-Qwen3.5-2B/aitk/README.md b/Qwen-Qwen3.5-2B/aitk/README.md new file mode 100644 index 000000000..7c8202328 --- /dev/null +++ b/Qwen-Qwen3.5-2B/aitk/README.md @@ -0,0 +1,37 @@ +# Qwen3.5-2B Model Optimization — NVIDIA TRT for RTX + +This recipe converts the [Qwen3.5-2B](https://huggingface.co/Qwen/Qwen3.5-2B) +vision-language model to ONNX for the **NVIDIA TensorRT for RTX** execution +provider (`NvTensorRTRTXExecutionProvider`) and runs it with ONNX Runtime +GenAI. + +Qwen3.5 is a hybrid architecture combining GatedDeltaNet linear attention +layers with standard full attention layers. The pipeline exports three +sub-models and assembles them into a single ONNX Runtime GenAI model folder: + +- **embedding.json** — token embedding + image feature fusion (FP16) +- **vision.json** — vision encoder, packed patches → image features (FP16) +- **text.json** — text decoder via ModelBuilder (INT4, hybrid GatedDeltaNet + full attention) + +Because AITK runs a single Olive workflow per recipe, the three inner Olive +configs are wrapped behind one `AitkPython` pass +(`qwen_trtrtx_workflow.py`). The script runs each inner config, then patches +`genai_config.json` / `processor_config.json` and the tokenizer for the GenAI +runtime. + +## Optimization + +| Sub-model | Precision | +|-----------|-----------| +| Vision encoder | FP16 | +| Text embedding | FP16 | +| Text decoder | INT4 (block size 128, accuracy level 4) | + +## Inference + +Run the provided `inference_sample.ipynb`. It loads the optimized model from +`./model`, registers the NVIDIA TRT for RTX execution provider, and streams a +response for a text (and optional image) prompt. + +> Metrics (latency / accuracy on a specific device) to be added after a +> benchmark run on target RTX hardware. diff --git a/Qwen-Qwen3.5-2B/aitk/_copy.json.config b/Qwen-Qwen3.5-2B/aitk/_copy.json.config new file mode 100644 index 000000000..28f2a7c4c --- /dev/null +++ b/Qwen-Qwen3.5-2B/aitk/_copy.json.config @@ -0,0 +1,8 @@ +{ + "copies": [ + { + "src": "../../deepseek-ai-DeepSeek-R1-Distill-Qwen-1.5B/aitk/winml.py", + "dst": "winml.py" + } + ] +} diff --git a/Qwen-Qwen3.5-2B/aitk/codes/__init__.py b/Qwen-Qwen3.5-2B/aitk/codes/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/Qwen-Qwen3.5-2B/aitk/codes/modeling_qwen3_5.py b/Qwen-Qwen3.5-2B/aitk/codes/modeling_qwen3_5.py new file mode 100644 index 000000000..45c61b217 --- /dev/null +++ b/Qwen-Qwen3.5-2B/aitk/codes/modeling_qwen3_5.py @@ -0,0 +1,511 @@ +# Copyright 2025 The Qwen Team and The HuggingFace Inc. team. All rights reserved. +# +# 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 typing import Optional + +import torch +import torch.nn as nn +import torch.nn.functional as F +from transformers.activations import ACT2FN +from transformers.modeling_layers import GradientCheckpointingLayer +from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel +from transformers.models.qwen3_5.configuration_qwen3_5 import Qwen3_5Config, Qwen3_5VisionConfig +from transformers.utils import auto_docstring, logging + +logger = logging.get_logger(__name__) + + +# ═══════════════════════════════════════════════════════════════ +# Vision encoder components (identical architecture to Qwen3-VL) +# ═══════════════════════════════════════════════════════════════ + + +class VisionMLP(nn.Module): + def __init__(self, config): + super().__init__() + self.hidden_size = config.hidden_size + self.intermediate_size = config.intermediate_size + self.linear_fc1 = nn.Linear(self.hidden_size, self.intermediate_size, bias=True) + self.linear_fc2 = nn.Linear(self.intermediate_size, self.hidden_size, bias=True) + self.act_fn = ACT2FN[config.hidden_act] + + def forward(self, hidden_state): + return self.linear_fc2(self.act_fn(self.linear_fc1(hidden_state))) + + +class VisionPatchEmbed(nn.Module): + def __init__(self, config) -> None: + super().__init__() + self.patch_size = config.patch_size + self.temporal_patch_size = config.temporal_patch_size + self.in_channels = config.in_channels + self.embed_dim = config.hidden_size + + kernel_size = [self.temporal_patch_size, self.patch_size, self.patch_size] + self.proj = nn.Conv3d(self.in_channels, self.embed_dim, kernel_size=kernel_size, stride=kernel_size, bias=True) + + 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 VisionRotaryEmbedding(nn.Module): + inv_freq: torch.Tensor + + 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) -> 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 VisionPatchMerger(nn.Module): + def __init__(self, config: Qwen3_5VisionConfig, use_postshuffle_norm=False) -> None: + super().__init__() + self.hidden_size = config.hidden_size * (config.spatial_merge_size**2) + self.use_postshuffle_norm = use_postshuffle_norm + self.norm = nn.LayerNorm(self.hidden_size if use_postshuffle_norm else config.hidden_size, eps=1e-6) + self.linear_fc1 = nn.Linear(self.hidden_size, self.hidden_size) + self.act_fn = nn.GELU() + self.linear_fc2 = nn.Linear(self.hidden_size, config.out_hidden_size) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + x = self.norm(x.view(-1, self.hidden_size) if self.use_postshuffle_norm else x).view(-1, self.hidden_size) + x = self.linear_fc2(self.act_fn(self.linear_fc1(x))) + return x + + +def rotate_half(x): + 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) + return q_embed.to(orig_q_dtype), k_embed.to(orig_k_dtype) + + +def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor: + 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, query, key, value, attention_mask, scaling, dropout=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 VisionAttention(nn.Module): + def __init__(self, config: Qwen3_5VisionConfig) -> 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 + 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 = 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": + 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 getattr(torch.compiler, "is_exporting", lambda: False)(): + # ONNX export: use PackedAttention custom op + 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], + query_states.shape[2], + query_states.shape[1], + query_states.shape[3], + ), + version=1, + ) + attn_output = attn_output.to(self.proj.weight.device) + else: + # SDPA fallback: 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 = [] + for q, k, v in zip(*splits): + out = torch.nn.functional.scaled_dot_product_attention( + q, + k, + v, + attn_mask=None, + dropout_p=0.0 if not self.training else self.attention_dropout, + scale=self.scaling, + is_causal=False, + ) + attn_outputs.append(out.transpose(1, 2)) + 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 VisionBlock(GradientCheckpointingLayer): + def __init__(self, config, attn_implementation: str = "sdpa") -> None: + super().__init__() + self.norm1 = nn.LayerNorm(config.hidden_size, eps=1e-6) + self.norm2 = nn.LayerNorm(config.hidden_size, eps=1e-6) + self.attn = VisionAttention(config=config) + self.mlp = VisionMLP(config=config) + + def forward(self, hidden_states, cu_seqlens, rotary_pos_emb=None, position_embeddings=None, **kwargs): + 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 + + +# ═══════════════════════════════════════════════════════════════ +# Vision model (PreTrainedModel wrapper) +# ═══════════════════════════════════════════════════════════════ + + +@auto_docstring +class Qwen3_5PreTrainedModel(PreTrainedModel): + config: Qwen3_5Config + base_model_prefix = "model" + supports_gradient_checkpointing = True + _no_split_modules = ["VisionBlock"] + _skip_keys_device_placement = "past_key_values" + _supports_sdpa = True + _supports_attention_backend = True + + +class Qwen3_5VisionModel(Qwen3_5PreTrainedModel): + config: Qwen3_5VisionConfig + _no_split_modules = ["VisionBlock"] + + 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.spatial_merge_unit = self.spatial_merge_size * self.spatial_merge_size + + self.patch_embed = VisionPatchEmbed(config=config) + + self.pos_embed = nn.Embedding(config.num_position_embeddings, config.hidden_size) + self.num_grid_per_side = int(config.num_position_embeddings**0.5) + + head_dim = config.hidden_size // config.num_heads + self.rotary_pos_emb = VisionRotaryEmbedding(head_dim // 2) + + self.blocks = nn.ModuleList([VisionBlock(config) for _ in range(config.depth)]) + self.merger = VisionPatchMerger(config=config, use_postshuffle_norm=False) + + self.deepstack_visual_indexes = config.deepstack_visual_indexes + self.deepstack_merger_list = nn.ModuleList( + [ + VisionPatchMerger(config=config, use_postshuffle_norm=True) + for _ in range(len(config.deepstack_visual_indexes)) + ] + ) + + self.gradient_checkpointing = False + + def rot_pos_emb(self, grid_thw: torch.Tensor) -> torch.Tensor: + merge_size = self.spatial_merge_size + max_hw = grid_thw[:, 1:].max() + freq_table = self.rotary_pos_emb(max_hw) + device = freq_table.device + + all_embeddings = [] + for num_frames, height, width in grid_thw: + merged_h, merged_w = height // merge_size, width // merge_size + + 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) + coords = coords.repeat(num_frames, 1) + all_embeddings.append(freq_table[coords].flatten(1)) + + return torch.cat(all_embeddings, dim=0) + + def fast_pos_embed_interpolate(self, grid_thw): + """Bilinear interpolation of learnable 2D position embeddings.""" + merge_size = self.config.spatial_merge_size + dev = self.pos_embed.weight.device + dtype = self.pos_embed.weight.dtype + n = self.num_grid_per_side + + all_pos_embeds = [] + for t, h, w in zip(grid_thw[:, 0], grid_thw[:, 1], grid_thw[:, 2]): + torch._check(t.item() >= 1) + torch._check(h.item() >= 2) + torch._check(w.item() >= 2) + + h_idxs = torch.arange(h, dtype=torch.float32, device=dev) * ((n - 1) / (h - 1)) + w_idxs = torch.arange(w, dtype=torch.float32, device=dev) * ((n - 1) / (w - 1)) + + h_floor = h_idxs.int() + w_floor = w_idxs.int() + h_ceil = (h_floor + 1).clamp(max=n - 1) + w_ceil = (w_floor + 1).clamp(max=n - 1) + + dh = (h_idxs - h_floor.float()).to(dtype) + dw = (w_idxs - w_floor.float()).to(dtype) + + base_h = h_floor.long() * n + base_hc = h_ceil.long() * n + + idx_00 = (base_h[:, None] + w_floor.long()[None]).reshape(-1) + idx_01 = (base_h[:, None] + w_ceil.long()[None]).reshape(-1) + idx_10 = (base_hc[:, None] + w_floor.long()[None]).reshape(-1) + idx_11 = (base_hc[:, None] + w_ceil.long()[None]).reshape(-1) + + wt_00 = ((1.0 - dh)[:, None] * (1.0 - dw)[None]).reshape(-1) + wt_01 = ((1.0 - dh)[:, None] * dw[None]).reshape(-1) + wt_10 = (dh[:, None] * (1.0 - dw)[None]).reshape(-1) + wt_11 = (dh[:, None] * dw[None]).reshape(-1) + + pos = ( + self.pos_embed(idx_00.to(dev)) * wt_00[:, None] + + self.pos_embed(idx_01.to(dev)) * wt_01[:, None] + + self.pos_embed(idx_10.to(dev)) * wt_10[:, None] + + self.pos_embed(idx_11.to(dev)) * wt_11[:, None] + ) + + pos = pos.repeat(t, 1) + pos = ( + pos.reshape(t, h // merge_size, merge_size, w // merge_size, merge_size, -1) + .permute(0, 1, 3, 2, 4, 5) + .flatten(0, 4) + ) + all_pos_embeds.append(pos) + + return torch.cat(all_pos_embeds) + + def forward(self, hidden_states: torch.Tensor, grid_thw: torch.Tensor, **kwargs) -> torch.Tensor: + hidden_states = self.patch_embed(hidden_states) + pos_embeds = self.fast_pos_embed_interpolate(grid_thw) + hidden_states = hidden_states + pos_embeds + + rotary_pos_emb = self.rot_pos_emb(grid_thw) + seq_len, _ = hidden_states.size() + hidden_states = hidden_states.reshape(seq_len, -1) + rotary_pos_emb = rotary_pos_emb.reshape(seq_len, -1) + emb = torch.cat((rotary_pos_emb, rotary_pos_emb), dim=-1) + position_embeddings = (emb.cos(), emb.sin()) + + cu_seqlens = torch.repeat_interleave(grid_thw[:, 1] * grid_thw[:, 2], grid_thw[:, 0]).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 blk in self.blocks: + hidden_states = blk(hidden_states, cu_seqlens=cu_seqlens, position_embeddings=position_embeddings, **kwargs) + + hidden_states = self.merger(hidden_states) + return hidden_states + + +# ═══════════════════════════════════════════════════════════════ +# Top-level Qwen3.5 model (vision + text embedding shell) +# ═══════════════════════════════════════════════════════════════ + + +class Qwen3_5Model(Qwen3_5PreTrainedModel): + """Qwen3.5 composite model for vision + embedding ONNX export. + + The text decoder is exported separately via ModelBuilder. + This class provides: + - get_image_features(): vision encoder export + - get_fused_input_embeddings(): embedding fusion export + """ + + base_model_prefix = "" + _checkpoint_conversion_mapping = {} + config: Qwen3_5Config + _no_split_modules = ["VisionBlock"] + + def __init__(self, config): + super().__init__(config) + self.visual = Qwen3_5VisionModel._from_config(config.vision_config) + # Minimal text model: only need embed_tokens for embedding export + text_config = config.text_config + self.embed_tokens = nn.Embedding(text_config.vocab_size, text_config.hidden_size) + self.post_init() + + def get_input_embeddings(self): + return self.embed_tokens + + def get_image_features(self, pixel_values: torch.FloatTensor, image_grid_thw: Optional[torch.LongTensor] = None): + """Vision encoder: pixel_values + grid_thw -> image features. + + Returns list of per-image feature tensors (split by grid sizes). + """ + pixel_values = pixel_values.type(self.visual.dtype) + image_embeds = self.visual(pixel_values, grid_thw=image_grid_thw) + split_sizes = (image_grid_thw.prod(-1) // self.visual.spatial_merge_size**2).tolist() + image_embeds = torch.split(image_embeds, split_sizes) + return image_embeds + + def get_fused_input_embeddings(self, input_ids, image_features=None): + """Embedding fusion: input_ids + image_features -> inputs_embeds. + + Scatters image features at image_token_id positions in the text embedding. + """ + + 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 llm_input_ids + + def false_fn_for_input_ids(input_ids): + return input_ids + + 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.embed_tokens(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): + 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 + + 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 forward(self, *args, **kwargs): + raise NotImplementedError( + "Qwen3_5Model.forward() should not be called directly. " + "Use get_image_features() or get_fused_input_embeddings() via method swap." + ) + + +__all__ = [ + "Qwen3_5Model", + "Qwen3_5PreTrainedModel", + "Qwen3_5VisionModel", +] diff --git a/Qwen-Qwen3.5-2B/aitk/embedding.json b/Qwen-Qwen3.5-2B/aitk/embedding.json new file mode 100644 index 000000000..464931487 --- /dev/null +++ b/Qwen-Qwen3.5-2B/aitk/embedding.json @@ -0,0 +1,58 @@ +{ + "input_model": { + "type": "PyTorchModel", + "model_path": "Qwen/Qwen3.5-2B", + "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": [ + "NvTensorRTRTXExecutionProvider" + ] + } + ] + } + }, + "no_artifacts": true, + "output_dir": "model" +} diff --git a/Qwen-Qwen3.5-2B/aitk/inference_model.json b/Qwen-Qwen3.5-2B/aitk/inference_model.json new file mode 100644 index 000000000..b662d56fd --- /dev/null +++ b/Qwen-Qwen3.5-2B/aitk/inference_model.json @@ -0,0 +1,31 @@ +{ + "Name": "Qwen3.5-2B", + "PromptTemplate": { + "assistant": "{Content}", + "prompt": "<|im_start|>user\n{Content}<|im_end|>\n<|im_start|>assistant\n" + }, + "ParameterSchema": { + "enabled": [ + { + "name": "max_tokens", + "default": 512 + }, + { + "name": "temperature", + "default": 0.6 + }, + { + "name": "top_p", + "default": 0.95 + }, + { + "name": "top_k", + "default": 5 + }, + { + "name": "random_seed", + "default": 3328 + } + ] + } +} diff --git a/Qwen-Qwen3.5-2B/aitk/inference_sample.ipynb b/Qwen-Qwen3.5-2B/aitk/inference_sample.ipynb new file mode 100644 index 000000000..cf2b53781 --- /dev/null +++ b/Qwen-Qwen3.5-2B/aitk/inference_sample.ipynb @@ -0,0 +1,118 @@ +{ + "cells": [ + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "text=\"Describe this image in detail.\"\n", + "image=None # set to an image path (e.g. \"./photo.jpg\") for image + text\n", + "ExecutionProvider=\"NvTensorRTRTXExecutionProvider\"\n", + "model_folder = \"./model\"" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "from winml import register_execution_providers\n", + "register_execution_providers(ExecutionProvider)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "import json\n", + "import time\n", + "\n", + "import onnxruntime_genai as og\n", + "\n", + "# Load the model and the multimodal processor\n", + "model = og.Model(model_folder)\n", + "processor = model.create_multimodal_processor()\n", + "tokenizer = og.Tokenizer(model)\n", + "tokenizer_stream = processor.create_stream()\n", + "\n", + "# Build the chat prompt (optionally with an image)\n", + "images = None\n", + "if image:\n", + " images = og.Images.open(image)\n", + " messages = [\n", + " {\n", + " \"role\": \"user\",\n", + " \"content\": [\n", + " {\"type\": \"image\"},\n", + " {\"type\": \"text\", \"text\": text},\n", + " ],\n", + " }\n", + " ]\n", + "else:\n", + " messages = [{\"role\": \"user\", \"content\": text}]\n", + "\n", + "full_prompt = tokenizer.apply_chat_template(json.dumps(messages), add_generation_prompt=True)\n", + "inputs = processor(full_prompt, images=images)\n", + "\n", + "# Create params and generator\n", + "params = og.GeneratorParams(model)\n", + "params.set_search_options(max_length=4096)\n", + "\n", + "generator = og.Generator(model, params)\n", + "generator.set_inputs(inputs)\n", + "\n", + "print(\"\")\n", + "print(\"Output: \", end=\"\", flush=True)\n", + "\n", + "token_times = []\n", + "\n", + "# Stream the output\n", + "while not generator.is_done():\n", + " start_time = time.time()\n", + " generator.generate_next_token()\n", + " new_token = generator.get_next_tokens()[0]\n", + " token_times.append(time.time() - start_time)\n", + "\n", + " print(tokenizer_stream.decode(new_token), end=\"\", flush=True)\n", + "\n", + "print()\n", + "\n", + "# Calculate and display timing statistics\n", + "if token_times:\n", + " total_tokens = len(token_times)\n", + " avg_time = sum(token_times) / total_tokens\n", + "\n", + " print(f\"Total tokens generated: {total_tokens}\")\n", + " print(f\"Average time per token: {avg_time:.4f} seconds\")\n", + " print(f\"Tokens per second: {total_tokens / sum(token_times):.2f}\")\n", + "\n", + "del generator" + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.11.9" + } + }, + "nbformat": 4, + "nbformat_minor": 2 +} diff --git a/Qwen-Qwen3.5-2B/aitk/info.yml b/Qwen-Qwen3.5-2B/aitk/info.yml new file mode 100644 index 000000000..932f06b1d --- /dev/null +++ b/Qwen-Qwen3.5-2B/aitk/info.yml @@ -0,0 +1,13 @@ +keywords: + aitk +arch: qwen3_5 +recipes: + - file: "qwen_trtrtx_workflow.json" + device: gpu + ep: NvTensorRTRTXExecutionProvider + aitk: + auto: false +aitk: + modelInfo: + id: "huggingface/Qwen/Qwen3.5-2B" + version: 1 diff --git a/Qwen-Qwen3.5-2B/aitk/model_project.config b/Qwen-Qwen3.5-2B/aitk/model_project.config new file mode 100644 index 000000000..02af5b0aa --- /dev/null +++ b/Qwen-Qwen3.5-2B/aitk/model_project.config @@ -0,0 +1,12 @@ +{ + "workflows": [ + { + "file": "qwen_trtrtx_workflow.json", + "templateName": "qwen_trtrtx_workflow" + } + ], + "modelInfo": { + "id": "huggingface/Qwen/Qwen3.5-2B", + "version": 1 + } +} diff --git a/Qwen-Qwen3.5-2B/aitk/qwen_trtrtx_workflow.json b/Qwen-Qwen3.5-2B/aitk/qwen_trtrtx_workflow.json new file mode 100644 index 000000000..d190068d7 --- /dev/null +++ b/Qwen-Qwen3.5-2B/aitk/qwen_trtrtx_workflow.json @@ -0,0 +1,31 @@ +{ + "input_model": { + "type": "HfModel", + "model_path": "Qwen/Qwen3.5-2B" + }, + "systems": { + "local_system": { + "type": "LocalSystem", + "accelerators": [ + { + "device": "gpu", + "execution_providers": [ + "NvTensorRTRTXExecutionProvider" + ] + } + ] + } + }, + "passes": { + "aitkpython": { + "type": "AitkPython", + "precision": "int4", + "user_script": "qwen_trtrtx_workflow.py" + } + }, + "target": "local_system", + "output_dir": "model/qwen3_5_trtrtx", + "cache_dir": "cache", + "no_artifacts": true, + "evaluate_input_model": false +} diff --git a/Qwen-Qwen3.5-2B/aitk/qwen_trtrtx_workflow.json.config b/Qwen-Qwen3.5-2B/aitk/qwen_trtrtx_workflow.json.config new file mode 100644 index 000000000..c749d67af --- /dev/null +++ b/Qwen-Qwen3.5-2B/aitk/qwen_trtrtx_workflow.json.config @@ -0,0 +1,84 @@ +{ + "$schema": "https://github.com/microsoft/olive-recipes/raw/refs/heads/main/.aitk/configs/config_schema.json", + "name": "Convert to NVIDIA TRT for RTX", + "oliveFile": { + "embedding.json": "builtin/cuda/embedding.json", + "text.json": "builtin/cuda/text.json", + "vision.json": "builtin/cuda/vision.json" + }, + "isLLM": true, + "needHFLogin": true, + "runtimeOverwrite": { + "autoGenerated": true, + "executeRequirement": "General/CPU_py3.12.9" + }, + "runtime": { + "autoGenerated": true, + "name": "Evaluate on", + "type": "enum", + "displayNames": [ + "NVIDIA TensorRT for RTX" + ], + "path": "systems.local_system.accelerators.0.execution_providers.0", + "values": [ + "NvTensorRTRTXExecutionProvider" + ], + "readOnly": false + }, + "optimizationPaths": [ + { + "path": "passes.aitkpython.precision" + } + ], + "optimizationDefault": "int4", + "aitkPython": "qwen_trtrtx_workflow.py", + "sections": [ + { + "autoGenerated": true, + "name": "Convert", + "phase": "Conversion", + "parameters": [], + "toggle": { + "autoGenerated": true, + "name": "Convert to ONNX format", + "type": "bool", + "path": "passes.aitkpython", + "actions": [ + [], + [] + ], + "readOnly": true + } + }, + { + "name": "Optimization", + "phase": "Quantization", + "parameters": [ + { + "name": "Precision", + "description": "Precision of the text decoder", + "type": "enum", + "displayNames": [ + "Int4" + ], + "displayType": "RadioGroup", + "path": "passes.aitkpython.precision", + "values": [ + "int4" + ] + } + ], + "disableToggleGeneration": true, + "toggle": { + "name": "Optimize model", + "type": "bool", + "path": "passes.aitkpython", + "actions": [ + [], + [] + ], + "readOnly": true + } + } + ] +} diff --git a/Qwen-Qwen3.5-2B/aitk/qwen_trtrtx_workflow.py b/Qwen-Qwen3.5-2B/aitk/qwen_trtrtx_workflow.py new file mode 100644 index 000000000..195129422 --- /dev/null +++ b/Qwen-Qwen3.5-2B/aitk/qwen_trtrtx_workflow.py @@ -0,0 +1,227 @@ +"""AITK aitkpython driver for the Qwen3.5-2B NVIDIA TRT for RTX recipe. + +The Qwen3.5 vision-language model is exported as three ONNX sub-models +(vision encoder, text embedding, text decoder). AITK only runs a single +Olive workflow per recipe, so this script wraps the three inner Olive +configs behind one ``AitkPython`` pass: it runs each inner config, then +patches ``genai_config.json``/``processor_config.json`` and the tokenizer +for the ONNX Runtime GenAI runtime. + +Adapted from ``../builtin/optimize.py`` (CUDA recipe) using the +``NvTensorRTRTXExecutionProvider`` execution provider. +""" + +import argparse +import json +import logging +import os + +import olive.workflows + +logger = logging.getLogger(os.path.basename(__file__)) +logging.basicConfig(level=logging.INFO) +logging.getLogger("onnxscript").setLevel(logging.WARNING) +logging.getLogger("onnx_ir").setLevel(logging.WARNING) + +SCRIPT_DIR = os.path.dirname(os.path.abspath(__file__)) +INNER_CONFIGS = ("embedding.json", "text.json", "vision.json") + + +def parse_arguments(): + parser = argparse.ArgumentParser() + parser.add_argument("--config", required=True, help="path to input config file") + parser.add_argument("--model_config", help="path to input model config file") + parser.add_argument("--runtime", required=True, help="runtime") + return parser.parse_args() + + +def export_models(models_dir: str, cache_dir: str): + """Run Olive for all 3 sub-models (embedding, text, vision) into models_dir.""" + for config_name in INNER_CONFIGS: + logger.info(f"Running {config_name}...") + with open(os.path.join(SCRIPT_DIR, config_name), "r", encoding="utf-8") as f: + inner = json.load(f) + + # all sub-models land in the single genai model folder + inner["output_dir"] = models_dir + inner["cache_dir"] = cache_dir + # resolve the model_script relative to this script's folder + model_script = inner.get("input_model", {}).get("model_script") + if model_script: + inner["input_model"]["model_script"] = os.path.join(SCRIPT_DIR, model_script) + + output = olive.workflows.run(inner) + if output is None or (hasattr(output, "has_output_model") and not output.has_output_model()): + raise Exception(f"Model file is not generated for {config_name}") + + +def update_genai_config(models_dir: str, device: str = "gpu"): + """Patch genai_config.json with embedding/vision sections and processor_config.""" + config_path = os.path.join(models_dir, "genai_config.json") + + with open(config_path, "r", encoding="utf-8") as f: + config = json.load(f) + + if device == "gpu": + provider_options = [{"cuda": {"enable_cuda_graph": "1", "enable_skip_layer_norm_strict_mode": "1"}}] + # Vision model has Loop nodes (one per ViT block) which are incompatible + # with CUDA graph capture, so disable it for vision and embedding only. + vision_provider_options = [{"cuda": {"enable_cuda_graph": "0", "enable_skip_layer_norm_strict_mode": "1"}}] + else: + provider_options = [] + vision_provider_options = [] + + session_options = {"log_id": "onnxruntime-genai", "provider_options": provider_options} + vision_session_options = {"log_id": "onnxruntime-genai", "provider_options": vision_provider_options} + + config["model"]["decoder"]["session_options"] = session_options + + config["model"]["embedding"] = { + "filename": "embedding.onnx", + "inputs": {"input_ids": "input_ids", "image_features": "image_features"}, + "outputs": {"inputs_embeds": "inputs_embeds"}, + "session_options": vision_session_options, + } + + config["model"]["vision"] = { + "filename": "vision.onnx", + "config_filename": "processor_config.json", + "spatial_merge_size": 2, + "tokens_per_second": 2.0, + "patch_size": 16, + "inputs": {"pixel_values": "pixel_values", "image_grid_thw": "image_grid_thw"}, + "outputs": {"image_features": "image_features"}, + "session_options": vision_session_options, + } + + config["model"]["bos_token_id"] = 248044 + config["model"]["eos_token_id"] = [248044] + config["model"]["pad_token_id"] = 248044 + config["model"]["image_token_id"] = 248056 + config["model"]["video_token_id"] = 248057 + config["model"]["vision_start_token_id"] = 248053 + + config["search"]["top_k"] = 1 + if config["search"].get("top_p") is None: + config["search"]["top_p"] = 1.0 + + with open(config_path, "w", encoding="utf-8") as f: + json.dump(config, f, indent=4) + logger.info(f"Updated {config_path}") + + 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": 960, + "height": 672, + "smart_resize": 1, + "min_pixels": 65536, + "max_pixels": 16777216, + "patch_size": 16, + "merge_size": 2, + }, + } + }, + { + "operation": { + "name": "rescale", + "type": "Rescale", + "attrs": { + "rescale_factor": 0.00392156862745098, + }, + } + }, + { + "operation": { + "name": "normalize", + "type": "Normalize", + "attrs": { + "mean": [0.5, 0.5, 0.5], + "std": [0.5, 0.5, 0.5], + "qwen2_5_vl": 1, + }, + } + }, + { + "operation": { + "name": "patch_image", + "type": "PatchImage", + "attrs": { + "patch_size": 16, + "temporal_patch_size": 2, + "merge_size": 2, + }, + } + }, + ], + } + } + + processor_path = os.path.join(models_dir, "processor_config.json") + with open(processor_path, "w", encoding="utf-8") as f: + json.dump(processor_config, f, indent=2) + logger.info(f"Created {processor_path}") + + +def fix_tokenizer(models_dir: str): + """Fix tokenizer.json for C++ std::regex compatibility. + + Qwen3.5's tokenizer uses Unicode property escapes (\\p{L}, \\p{N}) in its + Split pre-tokenizer, which aren't supported by std::regex in onnxruntime-genai. + Remove the Split and keep only ByteLevel with use_regex=True. + """ + tk_path = os.path.join(models_dir, "tokenizer.json") + if not os.path.exists(tk_path): + return + with open(tk_path, "r", encoding="utf-8") as f: + tk = json.load(f) + pt = tk.get("pre_tokenizer", {}) + if pt.get("type") == "Sequence": + pt["pretokenizers"] = [s for s in pt["pretokenizers"] if s.get("type") == "ByteLevel"] + for s in pt["pretokenizers"]: + s["use_regex"] = True + with open(tk_path, "w", encoding="utf-8") as f: + json.dump(tk, f, ensure_ascii=False) + + tc_path = os.path.join(models_dir, "tokenizer_config.json") + if os.path.exists(tc_path): + with open(tc_path, "r", encoding="utf-8") as f: + tc = json.load(f) + tc["tokenizer_class"] = "Qwen2Tokenizer" + with open(tc_path, "w", encoding="utf-8") as f: + json.dump(tc, f, indent=2, ensure_ascii=False) + logger.info("Fixed tokenizer for C++ std::regex compatibility") + + +def main(): + args = parse_arguments() + + with open(args.config, "r", encoding="utf-8") as f: + olive_json = json.load(f) + + # evaluation entrypoint — this recipe has no evaluation. + if args.model_config: + return + + models_dir = olive_json["output_dir"] + cache_dir = olive_json["cache_dir"] + os.makedirs(models_dir, exist_ok=True) + + export_models(models_dir, cache_dir) + + logger.info("=== Generating genai configs ===") + update_genai_config(models_dir, device="gpu") + fix_tokenizer(models_dir) + logger.info("Done.") + + +if __name__ == "__main__": + main() diff --git a/Qwen-Qwen3.5-2B/aitk/requirements.txt b/Qwen-Qwen3.5-2B/aitk/requirements.txt new file mode 100644 index 000000000..98af38216 --- /dev/null +++ b/Qwen-Qwen3.5-2B/aitk/requirements.txt @@ -0,0 +1,5 @@ +# This file will be installed together with Foundry Toolkit runtime requirements +# For the full requirements, see Foundry Toolkit +git+https://github.com/microsoft/Olive.git@main +torch>=2.10.0,<2.11.0 +transformers>=4.57.0 diff --git a/Qwen-Qwen3.5-2B/aitk/text.json b/Qwen-Qwen3.5-2B/aitk/text.json new file mode 100644 index 000000000..9e353e6a8 --- /dev/null +++ b/Qwen-Qwen3.5-2B/aitk/text.json @@ -0,0 +1,43 @@ +{ + "input_model": { + "type": "HfModel", + "model_path": "Qwen/Qwen3.5-2B" + }, + "passes": { + "m": { + "type": "ModelBuilder", + "precision": "int4", + "int4_accuracy_level": 4, + "int4_algo_config": "k_quant_mixed", + "int4_block_size": 128, + "extra_options": { + "filename": "text.onnx", + "quant_mode": "int4", + "prune_lm_head": true + } + }, + "t": { + "type": "GraphSurgeries", + "surgeries": [ + { + "surgeon": "TieWordEmbeddings" + } + ] + } + }, + "engine": { + "target": { + "type": "LocalSystem", + "accelerators": [ + { + "device": "gpu", + "execution_providers": [ + "NvTensorRTRTXExecutionProvider" + ] + } + ] + } + }, + "no_artifacts": true, + "output_dir": "model" +} diff --git a/Qwen-Qwen3.5-2B/aitk/user_script.py b/Qwen-Qwen3.5-2B/aitk/user_script.py new file mode 100644 index 000000000..e86d6bb2b --- /dev/null +++ b/Qwen-Qwen3.5-2B/aitk/user_script.py @@ -0,0 +1,133 @@ +import os +import sys + +import torch +from transformers import Qwen3_5Config + +_script_dir = os.path.dirname(os.path.abspath(__file__)) +if _script_dir not in sys.path: + sys.path.insert(0, _script_dir) + +from codes.modeling_qwen3_5 import Qwen3_5Model + +model_name = "Qwen/Qwen3.5-2B" +config = Qwen3_5Config.from_pretrained(model_name) + + +def _load_base_model(model_path): + """Load weights from safetensors into custom Qwen3_5Model.""" + import glob + + from huggingface_hub import hf_hub_download + from safetensors.torch import load_file + + config_path = hf_hub_download(model_path, "config.json") + model_dir = os.path.dirname(config_path) + st_files = sorted(glob.glob(os.path.join(model_dir, "*.safetensors"))) + + state_dict = {} + for sf in st_files: + tensors = load_file(sf) + for k, v in tensors.items(): + if k.startswith("model."): + stripped = k[6:] + state_dict[stripped] = v + # Map language_model.embed_tokens -> embed_tokens for our flat model + if stripped.startswith("language_model.embed_tokens."): + state_dict[stripped[len("language_model.") :]] = v + + custom_model = Qwen3_5Model(config) + result = custom_model.load_state_dict(state_dict, strict=False) + if result.missing_keys: + print(f"Warning: {len(result.missing_keys)} missing keys") + custom_model = custom_model.to(torch.bfloat16) + custom_model.eval() + + del state_dict + return custom_model + + +# ── Embedding ──────────────────────────────────────────────────────────── + + +def get_embedding_model(model_path=None): + model = _load_base_model(model_path) + model = model.to(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): + return { + "input_names": ["input_ids", "image_features"], + "output_names": ["inputs_embeds"], + "dynamic_axes": { + "input_ids": {0: "batch_size", 1: "sequence_length"}, + "image_features": {0: "num_logical_patches"}, + "inputs_embeds": {0: "batch_size", 1: "sequence_length"}, + }, + } + + +def get_embedding_dummy_inputs(model=None): + # Qwen3.5-2B: out_hidden_size=2048, patch_size=16 + batch_size, sequence_length, patches_per_image, out_hidden_size = ( + 2, + 216, + 187, + 2048, + ) + num_logical_patches = batch_size * patches_per_image + + vision_start_token_id = config.vision_start_token_id + vision_end_token_id = config.vision_end_token_id + image_token_id = config.image_token_id + + inputs = { + "input_ids": torch.randint(0, image_token_id, (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 + + for b in range(batch_size): + inputs["input_ids"][b][2] = vision_start_token_id + inputs["input_ids"][b][img_start_index:img_end_index] = image_token_id + inputs["input_ids"][b][img_end_index] = vision_end_token_id + + return inputs + + +# ── Vision ─────────────────────────────────────────────────────────────── + + +def get_vision_model(model_path=None): + model = _load_base_model(model_path) + model = model.to(torch.float32) + model.forward, model.get_image_features = model.get_image_features, model.forward + return model + + +def get_vision_io_config(model_path=None): + return { + "input_names": ["pixel_values", "image_grid_thw"], + "output_names": ["image_features"], + "dynamic_shapes": { + "pixel_values": {0: "num_patches"}, + "image_grid_thw": None, + }, + } + + +def get_vision_dummy_inputs(model=None): + # patch_size=16, temporal_patch_size=2, in_channels=3 + # patch dim: 3 * 2 * 16 * 16 = 1536 + # For 544x352 image: grid=(1, 22, 34), 748 raw patches + patches = 22 * 34 + pixel_values = torch.randn((patches, 1536), dtype=torch.float32) + grid_thw = torch.tensor([[1, 22, 34]], dtype=torch.int64) + return {"pixel_values": pixel_values, "image_grid_thw": grid_thw} diff --git a/Qwen-Qwen3.5-2B/aitk/vision.json b/Qwen-Qwen3.5-2B/aitk/vision.json new file mode 100644 index 000000000..f82991d8c --- /dev/null +++ b/Qwen-Qwen3.5-2B/aitk/vision.json @@ -0,0 +1,89 @@ +{ + "input_model": { + "type": "PyTorchModel", + "model_path": "Qwen/Qwen3.5-2B", + "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" + } + ] + }, + "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" + } + ], + "save_as_external_data": true, + "external_data_name": "vision.onnx.data" + } + }, + "engine": { + "target": { + "type": "LocalSystem", + "accelerators": [ + { + "device": "gpu", + "execution_providers": [ + "NvTensorRTRTXExecutionProvider" + ] + } + ] + } + }, + "no_artifacts": true, + "output_dir": "model" +} diff --git a/Qwen-Qwen3.5-2B/aitk/winml.py b/Qwen-Qwen3.5-2B/aitk/winml.py new file mode 100644 index 000000000..5a26031d6 --- /dev/null +++ b/Qwen-Qwen3.5-2B/aitk/winml.py @@ -0,0 +1,58 @@ +# https://learn.microsoft.com/en-us/windows/ai/new-windows-ml/initialize-execution-providers?tabs=python#production-app-example + + +def _get_ep_paths(ep: str | None = None) -> dict[str, str]: + import ctypes + import importlib.util + from pathlib import Path + + # Locate onnxruntime package path without importing it first + ort_spec = importlib.util.find_spec("onnxruntime") + assert ort_spec is not None and ort_spec.origin is not None + ort_package_path = Path(ort_spec.origin).parent + ort_capi_dir = ort_package_path / "capi" + ort_dll_path = ort_capi_dir / "onnxruntime.dll" + + # Load the onnxruntime DLL because "C:\Windows\System32\onnxruntime.dll" may be exist and loaded first + ctypes.WinDLL(str(ort_dll_path)) + + # remove the msvcp140.dll from the winrt-runtime package. + # So it does not cause issues with other libraries. + from importlib import metadata + + site_packages_path = Path(str(metadata.distribution("winrt-runtime").locate_file(""))) + dll_path = site_packages_path / "winrt" / "msvcp140.dll" + if dll_path.exists(): + dll_path.unlink() + + import winui3.microsoft.windows.ai.machinelearning as winml + from winui3.microsoft.windows.applicationmodel.dynamicdependency.bootstrap import InitializeOptions, initialize + + eps = {} + with initialize(options=InitializeOptions.ON_NO_MATCH_SHOW_UI): + catalog = winml.ExecutionProviderCatalog.get_default() + providers = catalog.find_all_providers() + for provider in providers: + if ep is not None and provider.name != ep: + continue + result = provider.ensure_ready_async().get() + if result.status == winml.ExecutionProviderReadyResultState.SUCCESS: + eps[provider.name] = provider.library_path + else: + print( + f"Execution provider '{provider.name}' is unavailable. Status: {result.status}; reason: {result.diagnostic_text}; error code: {result.extended_error.value}" + ) + return eps + + +def register_execution_providers(ep: str | None = None): + paths = _get_ep_paths(ep) + + import onnxruntime_genai as og + + for item in paths.items(): + try: + og.register_execution_provider_library(item[0], item[1]) # pyright: ignore[reportAttributeAccessIssue] + print(f"Successfully registered execution provider {item[0]} from {item[1]}") + except Exception as e: + print(f"Failed to register execution provider {item[0]} from {item[1]}: {e}") From 4a2e6329038f1584cc4e9972fb91014338ac6cd6 Mon Sep 17 00:00:00 2001 From: hualxie Date: Wed, 17 Jun 2026 11:38:50 +0800 Subject: [PATCH 2/8] add feature --- .aitk/configs/checks.json | 2 +- .aitk/requirements/General/CPU_py3.12.9-QWEN3.txt | 1 + Qwen-Qwen3.5-2B/aitk/qwen_trtrtx_workflow.json.config | 3 +++ 3 files changed, 5 insertions(+), 1 deletion(-) create mode 100644 .aitk/requirements/General/CPU_py3.12.9-QWEN3.txt diff --git a/.aitk/configs/checks.json b/.aitk/configs/checks.json index 544545274..829c6f87d 100644 --- a/.aitk/configs/checks.json +++ b/.aitk/configs/checks.json @@ -14,6 +14,6 @@ "pathCheck": 1468, "requirementsCheck": 45, "templateCheck": 3, - "venvRequirementsCheck": 23, + "venvRequirementsCheck": 24, "winmlCopyCheck": 39 } diff --git a/.aitk/requirements/General/CPU_py3.12.9-QWEN3.txt b/.aitk/requirements/General/CPU_py3.12.9-QWEN3.txt new file mode 100644 index 000000000..0aa984f18 --- /dev/null +++ b/.aitk/requirements/General/CPU_py3.12.9-QWEN3.txt @@ -0,0 +1 @@ +onnxscript==0.7.0 diff --git a/Qwen-Qwen3.5-2B/aitk/qwen_trtrtx_workflow.json.config b/Qwen-Qwen3.5-2B/aitk/qwen_trtrtx_workflow.json.config index c749d67af..b6f5c4bb3 100644 --- a/Qwen-Qwen3.5-2B/aitk/qwen_trtrtx_workflow.json.config +++ b/Qwen-Qwen3.5-2B/aitk/qwen_trtrtx_workflow.json.config @@ -12,6 +12,9 @@ "autoGenerated": true, "executeRequirement": "General/CPU_py3.12.9" }, + "executeRuntimeFeatures": [ + "QWEN3" + ], "runtime": { "autoGenerated": true, "name": "Evaluate on", From 5ac1f6f9518ccb811c2b9a86029edeb90d6e70ca Mon Sep 17 00:00:00 2001 From: hualxie Date: Wed, 17 Jun 2026 13:58:17 +0800 Subject: [PATCH 3/8] update more --- .aitk/requirements/General/CPU_py3.12.9-QWEN3.txt | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/.aitk/requirements/General/CPU_py3.12.9-QWEN3.txt b/.aitk/requirements/General/CPU_py3.12.9-QWEN3.txt index 0aa984f18..7cb96b0fd 100644 --- a/.aitk/requirements/General/CPU_py3.12.9-QWEN3.txt +++ b/.aitk/requirements/General/CPU_py3.12.9-QWEN3.txt @@ -1 +1,5 @@ +olive-ai==0.13.0 +onnx-ir==0.2.1 onnxscript==0.7.0 +transformers==5.2.0 +onnxruntime-genai-winml==0.14.1 From 15047685b26157d06175f8170fdc5020b55633c8 Mon Sep 17 00:00:00 2001 From: hualxie Date: Wed, 17 Jun 2026 14:21:30 +0800 Subject: [PATCH 4/8] update --- Qwen-Qwen3.5-2B/aitk/inference_sample.ipynb | 18 +++++++- Qwen-Qwen3.5-2B/aitk/qwen_trtrtx_workflow.py | 45 ++++++++++++++------ 2 files changed, 48 insertions(+), 15 deletions(-) diff --git a/Qwen-Qwen3.5-2B/aitk/inference_sample.ipynb b/Qwen-Qwen3.5-2B/aitk/inference_sample.ipynb index cf2b53781..15cc60713 100644 --- a/Qwen-Qwen3.5-2B/aitk/inference_sample.ipynb +++ b/Qwen-Qwen3.5-2B/aitk/inference_sample.ipynb @@ -7,11 +7,27 @@ "outputs": [], "source": [ "text=\"Describe this image in detail.\"\n", - "image=None # set to an image path (e.g. \"./photo.jpg\") for image + text\n", + "image=\"dog.jpeg\"\n", "ExecutionProvider=\"NvTensorRTRTXExecutionProvider\"\n", "model_folder = \"./model\"" ] }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "import os\n", + "import urllib.request\n", + "\n", + "# Download a sample image (https://onnxruntime.ai/images/dog.jpeg)\n", + "image_url = \"https://onnxruntime.ai/images/dog.jpeg\"\n", + "if not os.path.exists(image):\n", + " urllib.request.urlretrieve(image_url, image)\n", + " print(f\"Downloaded {image}\")" + ] + }, { "cell_type": "code", "execution_count": null, diff --git a/Qwen-Qwen3.5-2B/aitk/qwen_trtrtx_workflow.py b/Qwen-Qwen3.5-2B/aitk/qwen_trtrtx_workflow.py index 195129422..9e1c02921 100644 --- a/Qwen-Qwen3.5-2B/aitk/qwen_trtrtx_workflow.py +++ b/Qwen-Qwen3.5-2B/aitk/qwen_trtrtx_workflow.py @@ -35,22 +35,38 @@ def parse_arguments(): return parser.parse_args() -def export_models(models_dir: str, cache_dir: str): +def load_update_config(config_path: str, cache_dir: str, output_dir: str) -> dict: + """Load an inner Olive config and rewire its cache/output dirs and model_script.""" + with open(config_path, "r", encoding="utf-8") as f: + oliveJson = json.load(f) + + oliveJson["cache_dir"] = cache_dir + # all sub-models land in the single genai model folder + oliveJson["output_dir"] = output_dir + # resolve the model_script relative to this script's folder + model_script = oliveJson.get("input_model", {}).get("model_script") + if model_script: + oliveJson["input_model"]["model_script"] = os.path.join(SCRIPT_DIR, model_script) + + return oliveJson + + +def copy_olive_config(history_folder: str, config_name: str, cache_dir: str, output_dir: str) -> dict: + """Save the resolved inner config into the history folder for record, and return it.""" + logger.info(f"Copying {config_name} to {history_folder}...") + oliveJson = load_update_config(os.path.join(SCRIPT_DIR, config_name), cache_dir, output_dir) + os.makedirs(history_folder, exist_ok=True) + with open(os.path.join(history_folder, config_name), "w", encoding="utf-8") as f: + json.dump(oliveJson, f, indent=4) + return oliveJson + + +def export_models(history_folder: str, models_dir: str, cache_dir: str): """Run Olive for all 3 sub-models (embedding, text, vision) into models_dir.""" for config_name in INNER_CONFIGS: + oliveJson = copy_olive_config(history_folder, config_name, cache_dir, models_dir) logger.info(f"Running {config_name}...") - with open(os.path.join(SCRIPT_DIR, config_name), "r", encoding="utf-8") as f: - inner = json.load(f) - - # all sub-models land in the single genai model folder - inner["output_dir"] = models_dir - inner["cache_dir"] = cache_dir - # resolve the model_script relative to this script's folder - model_script = inner.get("input_model", {}).get("model_script") - if model_script: - inner["input_model"]["model_script"] = os.path.join(SCRIPT_DIR, model_script) - - output = olive.workflows.run(inner) + output = olive.workflows.run(oliveJson) if output is None or (hasattr(output, "has_output_model") and not output.has_output_model()): raise Exception(f"Model file is not generated for {config_name}") @@ -211,11 +227,12 @@ def main(): if args.model_config: return + history_folder = os.path.dirname(args.config) models_dir = olive_json["output_dir"] cache_dir = olive_json["cache_dir"] os.makedirs(models_dir, exist_ok=True) - export_models(models_dir, cache_dir) + export_models(history_folder, models_dir, cache_dir) logger.info("=== Generating genai configs ===") update_genai_config(models_dir, device="gpu") From 443c4a37cbe37aa9b9f2f095963544d5963edcdf Mon Sep 17 00:00:00 2001 From: hualxie Date: Wed, 17 Jun 2026 14:25:22 +0800 Subject: [PATCH 5/8] rename --- Qwen-Qwen3.5-2B/aitk/qwen_trtrtx_workflow.json.config | 6 +++--- Qwen-Qwen3.5-2B/aitk/qwen_trtrtx_workflow.py | 2 +- Qwen-Qwen3.5-2B/aitk/requirements.txt | 3 --- .../aitk/{embedding.json => trtrtx_embedding.json} | 0 Qwen-Qwen3.5-2B/aitk/{text.json => trtrtx_text.json} | 0 Qwen-Qwen3.5-2B/aitk/{vision.json => trtrtx_vision.json} | 0 6 files changed, 4 insertions(+), 7 deletions(-) rename Qwen-Qwen3.5-2B/aitk/{embedding.json => trtrtx_embedding.json} (100%) rename Qwen-Qwen3.5-2B/aitk/{text.json => trtrtx_text.json} (100%) rename Qwen-Qwen3.5-2B/aitk/{vision.json => trtrtx_vision.json} (100%) diff --git a/Qwen-Qwen3.5-2B/aitk/qwen_trtrtx_workflow.json.config b/Qwen-Qwen3.5-2B/aitk/qwen_trtrtx_workflow.json.config index b6f5c4bb3..5262c69af 100644 --- a/Qwen-Qwen3.5-2B/aitk/qwen_trtrtx_workflow.json.config +++ b/Qwen-Qwen3.5-2B/aitk/qwen_trtrtx_workflow.json.config @@ -2,9 +2,9 @@ "$schema": "https://github.com/microsoft/olive-recipes/raw/refs/heads/main/.aitk/configs/config_schema.json", "name": "Convert to NVIDIA TRT for RTX", "oliveFile": { - "embedding.json": "builtin/cuda/embedding.json", - "text.json": "builtin/cuda/text.json", - "vision.json": "builtin/cuda/vision.json" + "trtrtx_embedding.json": "builtin/cuda/embedding.json", + "trtrtx_text.json": "builtin/cuda/text.json", + "trtrtx_vision.json": "builtin/cuda/vision.json" }, "isLLM": true, "needHFLogin": true, diff --git a/Qwen-Qwen3.5-2B/aitk/qwen_trtrtx_workflow.py b/Qwen-Qwen3.5-2B/aitk/qwen_trtrtx_workflow.py index 9e1c02921..0f37e3b84 100644 --- a/Qwen-Qwen3.5-2B/aitk/qwen_trtrtx_workflow.py +++ b/Qwen-Qwen3.5-2B/aitk/qwen_trtrtx_workflow.py @@ -24,7 +24,7 @@ logging.getLogger("onnx_ir").setLevel(logging.WARNING) SCRIPT_DIR = os.path.dirname(os.path.abspath(__file__)) -INNER_CONFIGS = ("embedding.json", "text.json", "vision.json") +INNER_CONFIGS = ("trtrtx_embedding.json", "trtrtx_text.json", "trtrtx_vision.json") def parse_arguments(): diff --git a/Qwen-Qwen3.5-2B/aitk/requirements.txt b/Qwen-Qwen3.5-2B/aitk/requirements.txt index 98af38216..e3a05fd76 100644 --- a/Qwen-Qwen3.5-2B/aitk/requirements.txt +++ b/Qwen-Qwen3.5-2B/aitk/requirements.txt @@ -1,5 +1,2 @@ # This file will be installed together with Foundry Toolkit runtime requirements # For the full requirements, see Foundry Toolkit -git+https://github.com/microsoft/Olive.git@main -torch>=2.10.0,<2.11.0 -transformers>=4.57.0 diff --git a/Qwen-Qwen3.5-2B/aitk/embedding.json b/Qwen-Qwen3.5-2B/aitk/trtrtx_embedding.json similarity index 100% rename from Qwen-Qwen3.5-2B/aitk/embedding.json rename to Qwen-Qwen3.5-2B/aitk/trtrtx_embedding.json diff --git a/Qwen-Qwen3.5-2B/aitk/text.json b/Qwen-Qwen3.5-2B/aitk/trtrtx_text.json similarity index 100% rename from Qwen-Qwen3.5-2B/aitk/text.json rename to Qwen-Qwen3.5-2B/aitk/trtrtx_text.json diff --git a/Qwen-Qwen3.5-2B/aitk/vision.json b/Qwen-Qwen3.5-2B/aitk/trtrtx_vision.json similarity index 100% rename from Qwen-Qwen3.5-2B/aitk/vision.json rename to Qwen-Qwen3.5-2B/aitk/trtrtx_vision.json From 00466b0f5e78db05b85e8acde1f40bd3c57ff260 Mon Sep 17 00:00:00 2001 From: hualxie Date: Wed, 17 Jun 2026 14:59:50 +0800 Subject: [PATCH 6/8] update torch --- .aitk/requirements/General/CPU_py3.12.9-QWEN3.txt | 3 +++ 1 file changed, 3 insertions(+) diff --git a/.aitk/requirements/General/CPU_py3.12.9-QWEN3.txt b/.aitk/requirements/General/CPU_py3.12.9-QWEN3.txt index 7cb96b0fd..55e5ad794 100644 --- a/.aitk/requirements/General/CPU_py3.12.9-QWEN3.txt +++ b/.aitk/requirements/General/CPU_py3.12.9-QWEN3.txt @@ -3,3 +3,6 @@ onnx-ir==0.2.1 onnxscript==0.7.0 transformers==5.2.0 onnxruntime-genai-winml==0.14.1 +torch==2.10.0 +torchmetrics==1.9.0 +torchvision==0.25.0 From 6a57de65eee766e2d8d3c43f6973f8ac85901880 Mon Sep 17 00:00:00 2001 From: hualxie Date: Wed, 17 Jun 2026 15:05:36 +0800 Subject: [PATCH 7/8] remove update model_script --- Qwen-Qwen3.5-2B/aitk/qwen_trtrtx_workflow.py | 6 +----- 1 file changed, 1 insertion(+), 5 deletions(-) diff --git a/Qwen-Qwen3.5-2B/aitk/qwen_trtrtx_workflow.py b/Qwen-Qwen3.5-2B/aitk/qwen_trtrtx_workflow.py index 0f37e3b84..da4f6476f 100644 --- a/Qwen-Qwen3.5-2B/aitk/qwen_trtrtx_workflow.py +++ b/Qwen-Qwen3.5-2B/aitk/qwen_trtrtx_workflow.py @@ -36,17 +36,13 @@ def parse_arguments(): def load_update_config(config_path: str, cache_dir: str, output_dir: str) -> dict: - """Load an inner Olive config and rewire its cache/output dirs and model_script.""" + """Load an inner Olive config and rewire its cache/output dirs.""" with open(config_path, "r", encoding="utf-8") as f: oliveJson = json.load(f) oliveJson["cache_dir"] = cache_dir # all sub-models land in the single genai model folder oliveJson["output_dir"] = output_dir - # resolve the model_script relative to this script's folder - model_script = oliveJson.get("input_model", {}).get("model_script") - if model_script: - oliveJson["input_model"]["model_script"] = os.path.join(SCRIPT_DIR, model_script) return oliveJson From 399d55e9996cd89a1660de8a38f765da5332df26 Mon Sep 17 00:00:00 2001 From: hualxie Date: Wed, 17 Jun 2026 15:07:23 +0800 Subject: [PATCH 8/8] update output_name --- Qwen-Qwen3.5-2B/aitk/qwen_trtrtx_workflow.py | 23 ++++++++++++++------ 1 file changed, 16 insertions(+), 7 deletions(-) diff --git a/Qwen-Qwen3.5-2B/aitk/qwen_trtrtx_workflow.py b/Qwen-Qwen3.5-2B/aitk/qwen_trtrtx_workflow.py index da4f6476f..75afa472d 100644 --- a/Qwen-Qwen3.5-2B/aitk/qwen_trtrtx_workflow.py +++ b/Qwen-Qwen3.5-2B/aitk/qwen_trtrtx_workflow.py @@ -24,7 +24,11 @@ logging.getLogger("onnx_ir").setLevel(logging.WARNING) SCRIPT_DIR = os.path.dirname(os.path.abspath(__file__)) -INNER_CONFIGS = ("trtrtx_embedding.json", "trtrtx_text.json", "trtrtx_vision.json") +INNER_CONFIGS = { + "trtrtx_embedding.json": "embedding.onnx", + "trtrtx_text.json": "text.onnx", + "trtrtx_vision.json": "vision.onnx", +} def parse_arguments(): @@ -35,22 +39,21 @@ def parse_arguments(): return parser.parse_args() -def load_update_config(config_path: str, cache_dir: str, output_dir: str) -> dict: +def load_update_config(config_path: str, cache_dir: str, output_dir: str, output_name: str) -> dict: """Load an inner Olive config and rewire its cache/output dirs.""" with open(config_path, "r", encoding="utf-8") as f: oliveJson = json.load(f) oliveJson["cache_dir"] = cache_dir # all sub-models land in the single genai model folder - oliveJson["output_dir"] = output_dir - + oliveJson["output_dir"] = os.path.join(output_dir, output_name) if output_name else output_dir return oliveJson -def copy_olive_config(history_folder: str, config_name: str, cache_dir: str, output_dir: str) -> dict: +def copy_olive_config(history_folder: str, config_name: str, cache_dir: str, output_dir: str, output_name: str) -> dict: """Save the resolved inner config into the history folder for record, and return it.""" logger.info(f"Copying {config_name} to {history_folder}...") - oliveJson = load_update_config(os.path.join(SCRIPT_DIR, config_name), cache_dir, output_dir) + oliveJson = load_update_config(os.path.join(SCRIPT_DIR, config_name), cache_dir, output_dir, output_name) os.makedirs(history_folder, exist_ok=True) with open(os.path.join(history_folder, config_name), "w", encoding="utf-8") as f: json.dump(oliveJson, f, indent=4) @@ -60,7 +63,7 @@ def copy_olive_config(history_folder: str, config_name: str, cache_dir: str, out def export_models(history_folder: str, models_dir: str, cache_dir: str): """Run Olive for all 3 sub-models (embedding, text, vision) into models_dir.""" for config_name in INNER_CONFIGS: - oliveJson = copy_olive_config(history_folder, config_name, cache_dir, models_dir) + oliveJson = copy_olive_config(history_folder, config_name, cache_dir, models_dir, INNER_CONFIGS[config_name]) logger.info(f"Running {config_name}...") output = olive.workflows.run(oliveJson) if output is None or (hasattr(output, "has_output_model") and not output.has_output_model()): @@ -228,6 +231,12 @@ def main(): cache_dir = olive_json["cache_dir"] os.makedirs(models_dir, exist_ok=True) + import sys + + # A fix for Windows console output encoding issues (e.g., encodings\cp1252.py) + sys.stdout.reconfigure(encoding="utf-8") + sys.stderr.reconfigure(encoding="utf-8") + export_models(history_folder, models_dir, cache_dir) logger.info("=== Generating genai configs ===")