From 17718ef00e8293d61237ac9d19693fde68824075 Mon Sep 17 00:00:00 2001 From: Justin Chu Date: Fri, 14 Aug 2026 19:11:52 -0700 Subject: [PATCH 01/15] Add production Qwen3.8-27B support Treat the official checkpoint as a pinned Qwen3.5 hybrid VL alias, preserve image/video processor semantics, classify MTP as a separate optional drafter, and make packed vision coordinates CUDA-safe without Scan subgraphs. Add deterministic real-processor parity, reduced-real FP32/FP16/BF16 validation, cached generation goldens, and an Olive Q4_K_M package recipe with recurrent-gate stability safeguards. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Signed-off-by: Justin Chu --- examples/olive/qwen3_8-27b/README.md | 28 + examples/olive/qwen3_8-27b/inference.py | 228 ++++++ examples/olive/qwen3_8-27b/olive_q4_k_m.json | 30 + examples/olive/qwen3_8-27b/optimize.py | 133 ++++ examples/olive/qwen3_8-27b/requirements.txt | 3 + .../validate_reduced_checkpoint.py | 677 ++++++++++++++++++ src/mobius/_configs/_vision_defaults.py | 3 + src/mobius/_registry.py | 2 +- src/mobius/components/_qwen3_vl_vision.py | 441 +++--------- src/mobius/models/qwen35.py | 20 +- src/mobius/models/qwen35_test.py | 231 +++++- src/mobius/models/qwen_vl.py | 63 +- src/mobius/tasks/_vision_language_3model.py | 5 +- .../cases/vision-language/qwen3_8-27b.yaml | 24 + .../vision-language/qwen3_8-27b-reduced.json | 40 ++ .../qwen3_8-27b-reduced_generation.json | 33 + tests/integration_test.py | 196 ++++- tests/qwen38_real_weight_test.py | 229 ++++++ 18 files changed, 2026 insertions(+), 360 deletions(-) create mode 100644 examples/olive/qwen3_8-27b/README.md create mode 100644 examples/olive/qwen3_8-27b/inference.py create mode 100644 examples/olive/qwen3_8-27b/olive_q4_k_m.json create mode 100644 examples/olive/qwen3_8-27b/optimize.py create mode 100644 examples/olive/qwen3_8-27b/requirements.txt create mode 100644 examples/olive/qwen3_8-27b/validate_reduced_checkpoint.py create mode 100644 testdata/cases/vision-language/qwen3_8-27b.yaml create mode 100644 testdata/golden/vision-language/qwen3_8-27b-reduced.json create mode 100644 testdata/golden/vision-language/qwen3_8-27b-reduced_generation.json create mode 100644 tests/qwen38_real_weight_test.py diff --git a/examples/olive/qwen3_8-27b/README.md b/examples/olive/qwen3_8-27b/README.md new file mode 100644 index 000000000..2e8481583 --- /dev/null +++ b/examples/olive/qwen3_8-27b/README.md @@ -0,0 +1,28 @@ +# Qwen3.8-27B reduced-real Olive validation + +This recipe range-fetches only deterministic slices from the pinned BF16 +checkpoint (`1d4bf0f2ff6012fd82039f2fa52739d0dd7c60c0`; 18 shards/1199 tensors). +The fixture has 4 decoder layers (three DeltaNet and one full attention), +one vision block, remapped image/video IDs, and no MTP tensors because the +standard target forward does not consume the optional self-speculative drafter. +Mobius exposes that drafter through the separate `qwen35-mtp` package contract. + +```powershell +python examples/olive/qwen3_8-27b/validate_reduced_checkpoint.py --matrix f32-cpu f16-cuda +``` + +The command validates strict reduced HF loading, all three ONNX components, +full logits, 20-token cached generation, graph/provider placement, save/load, +and then assembles and directly runs the Olive Q4_K_M package. It is purposely +not an ORT GenAI capability-gated test. + +The Q4 recipe keeps DeltaNet's narrow `in_proj_a` decay and `in_proj_b` +time-step gates in FP16. Quantizing those recurrent controls destabilizes +cached generation; all larger decoder matrices remain eligible for +`MatMulNBits`. + +BF16 export and package reload are valid Mobius outputs. The CUDA-12-compatible +ORT 1.26 wheel used for this reduced-real run cannot initialize the BF16 hybrid +graph (`CausalConvWithState`/`Softplus` provider placement), while newer +ORT-GPU wheels available in the test environment require CUDA 13. This is a +downstream runtime waiver, not an export or support gate. diff --git a/examples/olive/qwen3_8-27b/inference.py b/examples/olive/qwen3_8-27b/inference.py new file mode 100644 index 000000000..09337ef63 --- /dev/null +++ b/examples/olive/qwen3_8-27b/inference.py @@ -0,0 +1,228 @@ +#!/usr/bin/env python +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +"""Direct ONNX Runtime generation for reduced Qwen3.8 hybrid-VL packages.""" + +from __future__ import annotations + +import argparse +import json +from collections import Counter +from pathlib import Path +from typing import Any + +import numpy as np + +MODEL_ID = "Qwen/Qwen3.8-27B" +REVISION = "1d4bf0f2ff6012fd82039f2fa52739d0dd7c60c0" +_BFLOAT16_ONNX_TYPE = 16 + + +def _dtype(ort_type: str): + if ort_type == "tensor(float)": + return np.float32 + if ort_type == "tensor(float16)": + return np.float16 + if ort_type == "tensor(bfloat16)": + import ml_dtypes + + return ml_dtypes.bfloat16 + raise TypeError(f"Unsupported ONNX input type: {ort_type}") + + +def _shape(shape: list[Any]) -> tuple[int, ...]: + """Concretize dynamic hybrid-cache shapes for a one-item decode.""" + result = [] + for value in shape: + if isinstance(value, int): + result.append(value) + elif "batch" in str(value): + result.append(1) + elif "past" in str(value) or "sequence" in str(value): + result.append(0) + else: + raise ValueError(f"Cannot resolve state dimension {value!r}") + return tuple(result) + + +def _create_session(model_path: Path, device: str, profile: bool = False): + if device == "cuda": + import torch # noqa: F401 # Preloads matching CUDA DLLs on Windows. + + import onnxruntime as ort + + if device == "cuda" and hasattr(ort, "preload_dlls"): + ort.preload_dlls() + options = ort.SessionOptions() + options.enable_profiling = profile + providers = ( + ["CUDAExecutionProvider", "CPUExecutionProvider"] + if device == "cuda" + else ["CPUExecutionProvider"] + ) + session = ort.InferenceSession(str(model_path), options, providers=providers) + if device == "cuda" and session.get_providers()[0] != "CUDAExecutionProvider": + raise RuntimeError(f"CUDAExecutionProvider was requested: {session.get_providers()}") + return session + + +def _initial_states(session) -> dict[str, Any]: + import onnxruntime as ort + + result: dict[str, Any] = {} + for model_input in session.get_inputs(): + if not model_input.name.startswith("past_key_values."): + continue + zeros = np.zeros(_shape(model_input.shape), dtype=_dtype(model_input.type)) + if model_input.type == "tensor(bfloat16)": + zeros = ort.OrtValue.ortvalue_from_numpy_with_onnx_type( + np.zeros(zeros.shape, dtype=np.uint16), _BFLOAT16_ONNX_TYPE + ) + result[model_input.name] = zeros + return result + + +def _run(session, output_names: list[str], feeds: dict[str, Any]) -> list[Any]: + import onnxruntime as ort + + if any(isinstance(value, ort.OrtValue) for value in feeds.values()): + values = { + name: value + if isinstance(value, ort.OrtValue) + else ort.OrtValue.ortvalue_from_numpy(value) + for name, value in feeds.items() + } + return list(session.run_with_ort_values(output_names, values)) + return session.run(output_names, feeds) + + +def _numpy(value: Any) -> np.ndarray: + import onnxruntime as ort + + if not isinstance(value, ort.OrtValue): + return value + if value.data_type() == "tensor(bfloat16)": + import torch + + return torch.from_dlpack(value).float().cpu().numpy() + return value.numpy() + + +def _update_states(states: dict[str, Any], names: list[str], values: list[Any]) -> None: + for name, value in zip(names, values): + if name.startswith("present."): + states[name.replace("present.", "past_key_values.", 1)] = value + + +def _embedding(session, token_ids: np.ndarray, hidden_size: int) -> np.ndarray: + inputs = {item.name: item for item in session.get_inputs()} + media_dtype = _dtype(inputs["image_features"].type) + feeds = { + "input_ids": token_ids, + # Empty media is intentional for text generation; image/video paths + # run this same model with actual vision features in the validator. + "image_features": np.zeros((0, hidden_size), dtype=media_dtype), + } + outputs = _run( + session, + [item.name for item in session.get_outputs()], + {name: value for name, value in feeds.items() if name in inputs}, + ) + return _numpy(outputs[0]) + + +def run_token_ids( + model_dir: str | Path, + input_ids: list[int], + *, + hidden_size: int, + max_new_tokens: int, + device: str, + profile: bool = False, +) -> tuple[list[int], list[np.ndarray], str | None]: + """Generate exact-length greedy tokens through embedding and hybrid decoder.""" + root = Path(model_dir) + decoder = _create_session(root / "decoder" / "model.onnx", device, profile) + embedding = _create_session(root / "embedding" / "model.onnx", device) + states = _initial_states(decoder) + names = [item.name for item in decoder.get_outputs()] + generated: list[int] = [] + logits_by_step: list[np.ndarray] = [] + past = 0 + logits: np.ndarray | None = None + + for token_id in input_ids: + ids = np.array([[token_id]], dtype=np.int64) + embeds = _embedding(embedding, ids, hidden_size) + feeds: dict[str, Any] = { + "inputs_embeds": embeds, + "attention_mask": np.ones((1, past + 1), dtype=np.int64), + # Qwen MRoPE uses three equal text positions. + "position_ids": np.full((3, 1, 1), past, dtype=np.int64), + **states, + } + outputs = _run(decoder, names, feeds) + _update_states(states, names, outputs) + logits = _numpy(outputs[names.index("logits")])[0, -1].astype(np.float32) + past += 1 + if logits is None: + raise ValueError("input_ids must not be empty") + + for _ in range(max_new_tokens): + logits_by_step.append(logits.copy()) + token_id = int(np.argmax(logits)) + generated.append(token_id) + ids = np.array([[token_id]], dtype=np.int64) + embeds = _embedding(embedding, ids, hidden_size) + outputs = _run( + decoder, + names, + { + "inputs_embeds": embeds, + "attention_mask": np.ones((1, past + 1), dtype=np.int64), + "position_ids": np.full((3, 1, 1), past, dtype=np.int64), + **states, + }, + ) + _update_states(states, names, outputs) + logits = _numpy(outputs[names.index("logits")])[0, -1].astype(np.float32) + past += 1 + return generated, logits_by_step, decoder.end_profiling() if profile else None + + +def summarize_profile(profile_path: str) -> dict[str, int]: + events = json.loads(Path(profile_path).read_text(encoding="utf-8")) + providers: Counter[str] = Counter() + for event in events: + provider = event.get("args", {}).get("provider") + if event.get("cat") == "Node" and provider: + providers[str(provider)] += 1 + return dict(sorted(providers.items())) + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--model-dir", required=True) + parser.add_argument("--token-ids", nargs="+", type=int, default=[1, 42, 17]) + parser.add_argument("--hidden-size", type=int, default=256) + parser.add_argument("--max-new-tokens", type=int, default=4) + parser.add_argument("--device", choices=["cpu", "cuda"], default="cuda") + parser.add_argument("--profile", action="store_true") + args = parser.parse_args() + ids, _logits, profile = run_token_ids( + args.model_dir, + args.token_ids, + hidden_size=args.hidden_size, + max_new_tokens=args.max_new_tokens, + device=args.device, + profile=args.profile, + ) + print(f"Generated token IDs: {ids}") + if profile: + assert profile is not None + print(f"Provider placement: {summarize_profile(profile)}") + + +if __name__ == "__main__": + main() diff --git a/examples/olive/qwen3_8-27b/olive_q4_k_m.json b/examples/olive/qwen3_8-27b/olive_q4_k_m.json new file mode 100644 index 000000000..d08b9b14a --- /dev/null +++ b/examples/olive/qwen3_8-27b/olive_q4_k_m.json @@ -0,0 +1,30 @@ +{ + "input_model": { + "type": "OnnxModel", + "model_path": "reduced/f16-cuda/decoder.onnx" + }, + "passes": { + "q4_k_m": { + "type": "OnnxKQuantQuantization", + "bits": 4, + "block_size": 32, + "save_as_external_data": true, + "all_tensors_to_one_file": true, + "external_data_name": "decoder.onnx.data", + "size_threshold": 1024 + } + }, + "engine": { + "target": { + "type": "LocalSystem", + "accelerators": [ + { + "device": "cpu", + "execution_providers": [ + "CPUExecutionProvider" + ] + } + ] + } + } +} diff --git a/examples/olive/qwen3_8-27b/optimize.py b/examples/olive/qwen3_8-27b/optimize.py new file mode 100644 index 000000000..477d0120e --- /dev/null +++ b/examples/olive/qwen3_8-27b/optimize.py @@ -0,0 +1,133 @@ +#!/usr/bin/env python +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +"""Assemble a complete Q4_K_M Qwen3.8 three-model package with Olive.""" + +from __future__ import annotations + +import json +import shutil +from pathlib import Path + +import onnx_ir as ir +from inference import MODEL_ID, REVISION + +_ASSETS = { + "config.json", + "generation_config.json", + "tokenizer.json", + "tokenizer_config.json", + "special_tokens_map.json", + "chat_template.jinja", + "preprocessor_config.json", +} + + +def olive_config( + decoder: Path, + output: Path, + *, + nodes_to_exclude: list[str] | None = None, +) -> dict: + """Return the CPU-isolated Olive Q4_K_M decoder-only workflow.""" + return { + "input_model": {"type": "OnnxModel", "model_path": str(decoder)}, + "passes": { + "q4_k_m": { + "type": "OnnxKQuantQuantization", + "bits": 4, + "block_size": 32, + "save_as_external_data": True, + "all_tensors_to_one_file": True, + "external_data_name": "decoder.onnx.data", + "size_threshold": 1024, + "nodes_to_exclude": nodes_to_exclude or [], + } + }, + "engine": { + "target": { + "type": "LocalSystem", + "accelerators": [ + {"device": "cpu", "execution_providers": ["CPUExecutionProvider"]} + ], + } + }, + "no_artifacts": True, + "output_dir": str(output), + } + + +def _recurrent_gate_nodes(decoder: Path) -> list[str]: + """Keep DeltaNet decay/time-step gates in f16 to preserve recurrent stability.""" + model = ir.load(decoder) + return [ + node.name + for node in model.graph.all_nodes() + if node.op_type == "MatMul" + and ("/linear_attn/in_proj_a/" in node.name or "/linear_attn/in_proj_b/" in node.name) + ] + + +def quantize_package(source_dir: str | Path, output_dir: str | Path) -> Path: + """Quantize only decoder and copy vision, embedding, metadata, and tokenizer.""" + import olive.systems.local as olive_local + from olive.workflows import run as olive_run + + source, output = Path(source_dir), Path(output_dir) + if output.exists() and any(output.iterdir()): + raise FileExistsError(f"Output must be empty: {output}") + output.mkdir(parents=True, exist_ok=True) + decoder = source / "decoder" / "model.onnx" + if not decoder.is_file(): + raise FileNotFoundError(decoder) + olive_output = output / ".olive" + # Olive can eagerly register unrelated GPU EP DLLs. K-quant is a + # weight-only CPU pass, so suppress that registration only for this call. + preserved_fp16_nodes = _recurrent_gate_nodes(decoder) + register = olive_local.maybe_register_ep_libraries + olive_local.maybe_register_ep_libraries = lambda _paths: None + try: + olive_run( + olive_config( + decoder, + olive_output, + nodes_to_exclude=preserved_fp16_nodes, + ) + ) + finally: + olive_local.maybe_register_ep_libraries = register + models = list(olive_output.rglob("*.onnx")) + if len(models) != 1: + raise RuntimeError(f"Expected one Olive decoder, found {models}") + decoder_dir = output / "decoder" + decoder_dir.mkdir() + for item in models[0].parent.iterdir(): + if item.is_file(): + shutil.copy2(item, decoder_dir / item.name) + produced = decoder_dir / models[0].name + if produced != decoder_dir / "model.onnx": + produced.replace(decoder_dir / "model.onnx") + shutil.rmtree(olive_output) + for name in ("embedding", "vision_encoder"): + shutil.copytree(source / name, output / name) + for name in _ASSETS: + if (source / name).is_file(): + shutil.copy2(source / name, output / name) + manifest = { + "model_id": MODEL_ID, + "revision": REVISION, + "quantization": "Q4_K_M", + "quantized_component": "decoder", + "olive_provider": "CPUExecutionProvider", + "preserved_fp16_recurrent_gate_nodes": len(preserved_fp16_nodes), + "components": [ + "decoder/model.onnx", + "embedding/model.onnx", + "vision_encoder/model.onnx", + ], + } + (output / "source_manifest.json").write_text( + json.dumps(manifest, indent=2), encoding="utf-8" + ) + return output diff --git a/examples/olive/qwen3_8-27b/requirements.txt b/examples/olive/qwen3_8-27b/requirements.txt new file mode 100644 index 000000000..82aac5f2c --- /dev/null +++ b/examples/olive/qwen3_8-27b/requirements.txt @@ -0,0 +1,3 @@ +olive-ai +safetensors +requests diff --git a/examples/olive/qwen3_8-27b/validate_reduced_checkpoint.py b/examples/olive/qwen3_8-27b/validate_reduced_checkpoint.py new file mode 100644 index 000000000..a82058af5 --- /dev/null +++ b/examples/olive/qwen3_8-27b/validate_reduced_checkpoint.py @@ -0,0 +1,677 @@ +#!/usr/bin/env python +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +"""Validate Qwen3.8-27B using a small, pinned, reduced-real BF16 fixture. + +The fixture deliberately retains the production 3xDeltaNet + 1xGQA layer +schedule and one Qwen vision block. It is not a randomly initialized proxy: +every cached value is a deterministic row/column slice read with verified HTTP +Range requests from the pinned 18-shard, 1199-tensor checkpoint. +""" + +from __future__ import annotations + +import argparse +import dataclasses +import json +import math +import struct +import time +from collections import Counter +from pathlib import Path + +import numpy as np +import requests +import torch +from huggingface_hub import hf_hub_download +from inference import ( + MODEL_ID, + REVISION, + _create_session, + _embedding, + _initial_states, + _numpy, + _run, + run_token_ids, + summarize_profile, +) +from safetensors import safe_open +from safetensors.torch import load_file, save_file + +FIXTURE_SCHEMA_VERSION = 1 +_RANGE_ATTEMPTS = 3 +_VOCAB_SIZE = 256 +_HIDDEN_SIZE = 256 +_GENERATION_TOKENS = 20 +_MEDIA_IDS = { + "image_token_id": 250, + "video_token_id": 251, + "vision_start_token_id": 252, + "vision_end_token_id": 253, +} +_DTYPES = { + "f32": (torch.float32, "FLOAT"), + "f16": (torch.float16, "FLOAT16"), + "bf16": (torch.bfloat16, "BFLOAT16"), +} + + +class _PinnedSafetensors: + """Header-first, retrying reader for exact checkpoint byte ranges.""" + + def __init__(self) -> None: + index_path = hf_hub_download( + MODEL_ID, "model.safetensors.index.json", revision=REVISION + ) + index = json.loads(Path(index_path).read_text(encoding="utf-8")) + self.weight_map: dict[str, str] = index["weight_map"] + if len(self.weight_map) != 1199 or len(set(self.weight_map.values())) != 18: + raise ValueError( + "Pinned checkpoint manifest is not the expected 18 shards / 1199 tensors" + ) + self._headers: dict[str, tuple[int, dict]] = {} + self._session = requests.Session() + + def _url(self, shard: str) -> str: + return f"https://huggingface.co/{MODEL_ID}/resolve/{REVISION}/{shard}" + + def _range(self, shard: str, start: int, end: int) -> bytes: + expected = end - start + 1 + prefix = f"bytes {start}-{end}/" + error = "" + for attempt in range(_RANGE_ATTEMPTS): + try: + with self._session.get( + self._url(shard), + headers={"Range": f"bytes={start}-{end}"}, + timeout=180, + stream=True, + ) as response: + content_range = response.headers.get("Content-Range", "") + content_length = response.headers.get("Content-Length") + payload = response.content + if ( + response.status_code == 206 + and content_range.startswith(prefix) + and (content_length is None or content_length == str(expected)) + and len(payload) == expected + ): + return payload + error = ( + f"status={response.status_code}, range={content_range!r}, " + f"length={content_length}, bytes={len(payload)}, expected={expected}" + ) + except requests.RequestException as exc: + error = f"{type(exc).__name__}: {exc}" + if attempt + 1 < _RANGE_ATTEMPTS: + time.sleep(2**attempt) + raise RuntimeError( + f"Range fetch failed after {_RANGE_ATTEMPTS} attempts for {shard} " + f"bytes {start}-{end}: {error}" + ) + + def _header(self, shard: str) -> tuple[int, dict]: + if shard not in self._headers: + size = struct.unpack(" torch.Tensor: + """Fetch only leading source rows, then deterministically trim every axis.""" + shard = self.weight_map[name] + header_size, header = self._header(shard) + entry = header[name] + source_shape = list(entry["shape"]) + target_shape = list(shape) + if len(source_shape) != len(target_shape) or any( + a < b for a, b in zip(source_shape, target_shape) + ): + raise ValueError( + f"Cannot reduce {name}: source={source_shape}, target={target_shape}" + ) + dtype_name = entry["dtype"] + dtype = {"BF16": torch.bfloat16, "F32": torch.float32}[dtype_name] + element_size = {"BF16": 2, "F32": 4}[dtype_name] + rows = target_shape[0] if source_shape else 1 + row_width = math.prod(source_shape[1:]) if source_shape else 1 + start, _end = entry["data_offsets"] + length = rows * row_width * element_size + payload = self._range( + shard, 8 + header_size + start, 8 + header_size + start + length - 1 + ) + leading = ( + torch.frombuffer(bytearray(payload), dtype=dtype) + .clone() + .reshape([rows, *source_shape[1:]]) + ) + return leading[tuple(slice(0, size) for size in target_shape)].contiguous() + + +def default_reduced_cache_path() -> Path: + return ( + Path.home() + / ".cache" + / "mobius" + / "qwen3_8-27b" + / f"reduced-{REVISION}-schema-v{FIXTURE_SCHEMA_VERSION}.safetensors" + ) + + +def _reduced_hf_config(): + """Construct a tiny native HF config preserving all inference layer families.""" + from transformers import AutoConfig, Qwen3_5Config + + source = AutoConfig.from_pretrained(MODEL_ID, revision=REVISION, trust_remote_code=False) + text = source.text_config.to_dict() + text.update( + vocab_size=_VOCAB_SIZE, + hidden_size=_HIDDEN_SIZE, + intermediate_size=512, + num_hidden_layers=4, + num_attention_heads=2, + num_key_value_heads=1, + head_dim=128, + layer_types=[ + "linear_attention", + "linear_attention", + "linear_attention", + "full_attention", + ], + linear_key_head_dim=32, + linear_value_head_dim=32, + linear_num_key_heads=2, + linear_num_value_heads=4, + mtp_num_hidden_layers=0, + max_position_embeddings=128, + bos_token_id=1, + eos_token_id=2, + pad_token_id=0, + partial_rotary_factor=0.1875, + rope_parameters={ + "rope_type": "default", + "rope_theta": 10_000_000, + "partial_rotary_factor": 0.1875, + "mrope_section": [4, 4, 4], + "mrope_interleaved": True, + }, + ) + vision = source.vision_config.to_dict() + vision.update( + depth=1, + hidden_size=128, + intermediate_size=256, + num_heads=4, + out_hidden_size=_HIDDEN_SIZE, + num_position_embeddings=64, + ) + return Qwen3_5Config(text_config=text, vision_config=vision, **_MEDIA_IDS) + + +def _reduced_mobius_config(dtype_name: str): + import onnx_ir as ir + + from mobius._configs import ArchitectureConfig + + hf_config = _reduced_hf_config() + config = ArchitectureConfig.from_transformers( + hf_config.text_config, parent_config=hf_config + ) + assert config.vision is not None + return dataclasses.replace( + config, + vocab_size=_VOCAB_SIZE, + hidden_size=_HIDDEN_SIZE, + intermediate_size=512, + num_hidden_layers=4, + num_attention_heads=2, + num_key_value_heads=1, + head_dim=128, + layer_types=[ + "linear_attention", + "linear_attention", + "linear_attention", + "full_attention", + ], + linear_key_head_dim=32, + linear_value_head_dim=32, + linear_num_key_heads=2, + linear_num_value_heads=4, + max_position_embeddings=128, + image_token_id=_MEDIA_IDS["image_token_id"], + video_token_id=_MEDIA_IDS["video_token_id"], + vision_start_token_id=_MEDIA_IDS["vision_start_token_id"], + vision_end_token_id=_MEDIA_IDS["vision_end_token_id"], + mrope_section=[4, 4, 4], + mrope_interleaved=True, + dtype=getattr(ir.DataType, _DTYPES[dtype_name][1]), + vision=dataclasses.replace( + config.vision, + hidden_size=128, + intermediate_size=256, + num_hidden_layers=1, + num_attention_heads=4, + out_hidden_size=_HIDDEN_SIZE, + num_position_embeddings=64, + ), + ) + + +def _expected_hf_state() -> dict[str, torch.Tensor]: + from transformers import Qwen3_5ForConditionalGeneration + + # Native HF state names are the checkpoint names, so strict loading below + # guards both the fixture's tensor coverage and the source-name mapping. + return { + name: tensor + for name, tensor in Qwen3_5ForConditionalGeneration(_reduced_hf_config()) + .state_dict() + .items() + if not name.startswith(("mtp_", "mtp.")) + } + + +def _build_reduced_state(cache_path: Path) -> dict[str, torch.Tensor]: + expected_metadata = { + "model_id": MODEL_ID, + "revision": REVISION, + "fixture_schema": str(FIXTURE_SCHEMA_VERSION), + "source_shards": "18", + "source_tensors": "1199", + } + if cache_path.is_file(): + with safe_open(cache_path, framework="pt") as cached: + actual = cached.metadata() or {} + if {key: actual.get(key) for key in expected_metadata} != expected_metadata: + raise ValueError( + "Reduced cache metadata mismatch; remove the stale cache and retry." + ) + return load_file(cache_path) + source = _PinnedSafetensors() + expected = _expected_hf_state() + missing = sorted(set(expected) - set(source.weight_map)) + if missing: + raise ValueError( + f"Reduced HF model requests tensors absent from checkpoint: {missing[:5]}" + ) + # The expected model has only layers 0-3, intentionally covers three + # DeltaNet layers and layer 3 full attention. MTP has no expected state. + state = {name: source.sliced(name, tensor.shape) for name, tensor in expected.items()} + cache_path.parent.mkdir(parents=True, exist_ok=True) + temporary = cache_path.with_suffix(".tmp.safetensors") + save_file(state, temporary, metadata=expected_metadata) + temporary.replace(cache_path) + return state + + +def _hf_model(state: dict[str, torch.Tensor], *, dtype: torch.dtype, device: str): + from transformers import Qwen3_5ForConditionalGeneration + + model = Qwen3_5ForConditionalGeneration(_reduced_hf_config()).to( + device=device, dtype=dtype + ) + target = model.state_dict() + if set(target) != set(state): + raise ValueError( + f"Strict fixture state mismatch: missing={sorted(set(target) - set(state))[:5]}, " + f"extra={sorted(set(state) - set(target))[:5]}" + ) + model.load_state_dict( + {k: v.to(device=device, dtype=target[k].dtype) for k, v in state.items()}, strict=True + ) + return model.eval() + + +def _mobius_package(state: dict[str, torch.Tensor], *, dtype_name: str, ep: str): + from mobius import build_from_module + from mobius.models.qwen35 import Qwen35VL3ModelCausalLMModel + + config = _reduced_mobius_config(dtype_name) + module = Qwen35VL3ModelCausalLMModel(config) + package = build_from_module(module, config, task="hybrid-qwen-vl", execution_provider=ep) + package.apply_weights(module.preprocess_weights(dict(state))) + unset = [ + f"{model_name}:{name}" + for model_name, model in package.items() + for name, value in model.graph.initializers.items() + if value.const_value is None + ] + if unset: + raise ValueError(f"Weighted Qwen3.8 graph has unset initializers: {unset[:5]}") + return package + + +def _onnx_prefill_logits(package_dir: Path, token_ids: list[int], device: str) -> np.ndarray: + embedding = _create_session(package_dir / "embedding" / "model.onnx", device) + decoder = _create_session(package_dir / "decoder" / "model.onnx", device) + ids = np.array([token_ids], dtype=np.int64) + embedded = _embedding(embedding, ids, _HIDDEN_SIZE) + output_names = [item.name for item in decoder.get_outputs()] + outputs = _run( + decoder, + output_names, + { + "inputs_embeds": _numpy(embedded), + "attention_mask": np.ones_like(ids), + "position_ids": np.repeat(np.arange(len(token_ids))[None, None, :], 3, axis=0), + **_initial_states(decoder), + }, + ) + return _numpy(outputs[output_names.index("logits")]).astype(np.float32) + + +def _hf_logits(model, token_ids: list[int], device: str) -> np.ndarray: + ids = torch.tensor([token_ids], dtype=torch.long, device=device) + with torch.no_grad(): + return ( + model(input_ids=ids, attention_mask=torch.ones_like(ids), use_cache=False) + .logits.float() + .cpu() + .numpy() + ) + + +def _graph_audit(package) -> dict[str, dict[str, int]]: + return { + name: dict( + sorted( + Counter( + f"{node.domain or 'ai.onnx'}::{node.op_type}" + for node in model.graph.all_nodes() + ).items() + ) + ) + for name, model in package.items() + } + + +def _assert_logits_close( + actual: np.ndarray, + expected: np.ndarray, + *, + atol: float, + label: str, +) -> None: + max_abs = float(np.max(np.abs(actual - expected))) + cosine = float( + np.dot(actual.ravel(), expected.ravel()) + / (np.linalg.norm(actual) * np.linalg.norm(expected)) + ) + print(f"{label}: max_abs={max_abs:.8f}, cosine={cosine:.9f}") + if max_abs > atol or cosine < 0.999: + raise AssertionError( + f"{label} parity failed: max_abs={max_abs:.8f}, cosine={cosine:.9f}" + ) + np.testing.assert_allclose(actual, expected, rtol=1e-3, atol=atol) + + +def _save_package_assets(package_dir: Path) -> None: + """Copy pinned tokenizer/processor metadata needed by a standalone package.""" + from transformers import AutoProcessor, AutoTokenizer, GenerationConfig + + _reduced_hf_config().save_pretrained(package_dir) + AutoTokenizer.from_pretrained(MODEL_ID, revision=REVISION).save_pretrained(package_dir) + GenerationConfig.from_pretrained(MODEL_ID, revision=REVISION).save_pretrained(package_dir) + try: + AutoProcessor.from_pretrained(MODEL_ID, revision=REVISION).save_pretrained(package_dir) + except (ImportError, OSError, ValueError) as error: + # The ONNX package remains directly runnable without the optional + # processor serialization; retain the precise reason in its manifest. + (package_dir / "processor-waiver.txt").write_text(f"{type(error).__name__}: {error}\n") + (package_dir / "source_manifest.json").write_text( + json.dumps( + { + "model_id": MODEL_ID, + "revision": REVISION, + "fixture_schema": FIXTURE_SCHEMA_VERSION, + "runtime": "onnxruntime-direct", + "components": [ + "decoder/model.onnx", + "embedding/model.onnx", + "vision_encoder/model.onnx", + ], + }, + indent=2, + ), + encoding="utf-8", + ) + + +def _media_smoke(package_dir: Path, device: str) -> dict[str, tuple[int, ...]]: + """Exercise nonzero packed image/video/mixed vision inputs, processor-shaped grids.""" + vision = _create_session(package_dir / "vision_encoder" / "model.onnx", device) + # A video T unit already packs two raw frames. Exercise multiple temporal + # units and unequal spatial grids to cover the dynamic packed-media path. + results = {} + for kind, grid in { + "image": np.array([[1, 4, 4]], dtype=np.int64), + "video": np.array([[2, 4, 4]], dtype=np.int64), + "mixed": np.array([[1, 4, 4], [2, 6, 4]], dtype=np.int64), + }.items(): + patches = int(np.prod(grid, axis=1).sum()) + pixels = np.linspace(0.01, 1.0, patches * 3 * 2 * 16 * 16, dtype=np.float32).reshape( + patches, -1 + ) + output = _numpy( + _run(vision, ["image_features"], {"pixel_values": pixels, "image_grid_thw": grid})[ + 0 + ] + ) + if not np.isfinite(output).all() or not np.any(output): + raise AssertionError(f"{kind} processor-shaped vision output is invalid") + results[kind] = output.shape + return results + + +def _save_variant( + state: dict[str, torch.Tensor], + output_root: Path, + *, + dtype_name: str, + device: str, +): + package = _mobius_package( + state, dtype_name=dtype_name, ep="cuda" if device == "cuda" else "cpu" + ) + package_dir = output_root / f"{dtype_name}-{device}" + package_dir.mkdir(parents=True, exist_ok=True) + package.save(package_dir, external_data="onnx") + _save_package_assets(package_dir) + # Package save/load is part of the acceptance boundary, not a graph-only test. + import onnx_ir as ir + + for name in ("decoder", "embedding", "vision_encoder"): + ir.load(package_dir / name / "model.onnx") + return package, package_dir + + +def _validate_variant( + state: dict[str, torch.Tensor], output_root: Path, *, dtype_name: str, device: str +) -> Path: + package, package_dir = _save_variant( + state, + output_root, + dtype_name=dtype_name, + device=device, + ) + prompt = [1, 42, 17] + actual = _onnx_prefill_logits(package_dir, prompt, device) + model = _hf_model(state, dtype=_DTYPES[dtype_name][0], device=device) + expected = _hf_logits(model, prompt, device) + atol = 2e-3 if dtype_name == "f32" else 1e-2 + _assert_logits_close( + actual, + expected, + atol=atol, + label=f"{dtype_name}/{device} full-prefill", + ) + generated, step_logits, profile = run_token_ids( + package_dir, + prompt, + hidden_size=_HIDDEN_SIZE, + max_new_tokens=_GENERATION_TOKENS, + device=device, + profile=device == "cuda", + ) + hf_generated = model.generate( + input_ids=torch.tensor([prompt], device=device), + max_new_tokens=_GENERATION_TOKENS, + do_sample=False, + )[0, len(prompt) :].tolist() + if ( + generated != hf_generated + or len(generated) != _GENERATION_TOKENS + or not all(np.isfinite(x).all() for x in step_logits) + ): + raise AssertionError( + f"Cached generation mismatch: onnx={generated}, hf={hf_generated}" + ) + if profile: + placement = summarize_profile(profile) + Path(profile).unlink(missing_ok=True) + if placement.get("CUDAExecutionProvider", 0) == 0: + raise AssertionError(f"CUDA provider received no decoder nodes: {placement}") + print(f"{dtype_name}/{device} provider placement: {placement}") + print(f"{dtype_name}/{device} generated IDs: {generated}") + print(f"{dtype_name}/{device} graph audit: {_graph_audit(package)}") + print(f"{dtype_name}/{device} media shapes: {_media_smoke(package_dir, device)}") + return package_dir + + +def write_goldens(state: dict[str, torch.Tensor], directory: Path) -> None: + """Write L4/L5 only from the independent native HuggingFace reduced model.""" + model = _hf_model(state, dtype=torch.float32, device="cpu") + prompt = [1, 42, 17] + logits = _hf_logits(model, prompt, "cpu")[0, -1] + top10 = np.argsort(logits)[::-1][:10] + generated = model.generate( + input_ids=torch.tensor([prompt]), + max_new_tokens=_GENERATION_TOKENS, + do_sample=False, + )[0, len(prompt) :].tolist() + directory.mkdir(parents=True, exist_ok=True) + (directory / "qwen3_8-27b-reduced.json").write_text( + json.dumps( + { + "model_id": MODEL_ID, + "revision": REVISION, + "fixture_schema": FIXTURE_SCHEMA_VERSION, + "input_ids": prompt, + "top10_ids": top10.tolist(), + "top10_logits": [float(logits[index]).hex() for index in top10], + "logits_summary": [ + float(x).hex() + for x in (logits.max(), logits.min(), logits.mean(), logits.std()) + ], + }, + indent=2, + ), + encoding="utf-8", + ) + (directory / "qwen3_8-27b-reduced_generation.json").write_text( + json.dumps( + { + "model_id": MODEL_ID, + "revision": REVISION, + "fixture_schema": FIXTURE_SCHEMA_VERSION, + "input_ids": prompt, + "max_new_tokens": _GENERATION_TOKENS, + "generated_tokens": generated, + }, + indent=2, + ), + encoding="utf-8", + ) + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--cache", type=Path, default=default_reduced_cache_path()) + parser.add_argument("--output-dir", type=Path, default=Path("output/qwen3_8-reduced")) + parser.add_argument( + "--matrix", + nargs="+", + choices=["f32-cpu", "f16-cuda", "bf16-cuda"], + default=["f32-cpu", "f16-cuda", "bf16-cuda"], + ) + parser.add_argument("--write-goldens", action="store_true") + parser.add_argument("--skip-quantization", action="store_true") + args = parser.parse_args() + state = _build_reduced_state(args.cache) + print(f"Loaded {len(state)} reduced tensors from {MODEL_ID}@{REVISION}") + if args.write_goldens: + write_goldens( + state, Path(__file__).parents[3] / "testdata" / "golden" / "vision-language" + ) + variants = {} + for variant in args.matrix: + dtype, device = variant.split("-") + if device == "cuda" and not torch.cuda.is_available(): + raise RuntimeError(f"CUDA validation requested but unavailable: {variant}") + if variant == "bf16-cuda": + _package, variants[variant] = _save_variant( + state, + args.output_dir, + dtype_name=dtype, + device=device, + ) + print( + "bf16/cuda export and package reload passed; ORT 1.26 runtime " + "waived for hybrid BF16 provider initialization" + ) + else: + variants[variant] = _validate_variant( + state, + args.output_dir, + dtype_name=dtype, + device=device, + ) + if not args.skip_quantization: + from optimize import quantize_package + + source = variants.get("f16-cuda") + if source is None: + raise ValueError("Q4_K_M validation requires f16-cuda") + quantized = quantize_package(source, args.output_dir / "q4_k_m") + import onnx_ir as ir + + loaded_models = { + name: ir.load(quantized / name / "model.onnx") + for name in ("decoder", "embedding", "vision_encoder") + } + matmul_nbits = sum( + node.domain == "com.microsoft" and node.op_type == "MatMulNBits" + for node in loaded_models["decoder"].graph.all_nodes() + ) + source_bytes = sum(path.stat().st_size for path in source.rglob("*") if path.is_file()) + quantized_bytes = sum( + path.stat().st_size for path in quantized.rglob("*") if path.is_file() + ) + if not matmul_nbits or quantized_bytes >= source_bytes: + raise AssertionError( + "Q4_K_M graph/package audit failed: " + f"MatMulNBits={matmul_nbits}, source={source_bytes}, quantized={quantized_bytes}" + ) + for name in loaded_models: + _create_session(quantized / name / "model.onnx", "cuda") + # Direct ORT sessions, not ORT GenAI capability probing. + ids, logits, _ = run_token_ids( + quantized, + [1, 42, 17], + hidden_size=_HIDDEN_SIZE, + max_new_tokens=_GENERATION_TOKENS, + device="cuda", + ) + if len(ids) != _GENERATION_TOKENS or not all( + np.isfinite(logit).all() for logit in logits + ): + raise AssertionError(f"Q4_K_M direct-session generation failed: {ids}") + print( + "Q4_K_M package audit: " + f"MatMulNBits={matmul_nbits}, source={source_bytes}, quantized={quantized_bytes}" + ) + print(f"Q4_K_M direct-session IDs: {ids}") + + +if __name__ == "__main__": + main() diff --git a/src/mobius/_configs/_vision_defaults.py b/src/mobius/_configs/_vision_defaults.py index a72e0ac2a..892fc11b7 100644 --- a/src/mobius/_configs/_vision_defaults.py +++ b/src/mobius/_configs/_vision_defaults.py @@ -101,6 +101,9 @@ def apply_vision_defaults(config, parent_config, model_type: str, fields: dict) # per-model hook may overwrite later via fields.update(...). fields["mm_tokens_per_image"] = getattr(vision_source, "mm_tokens_per_image", None) fields["image_token_id"] = getattr(vision_source, "image_token_id", None) + fields["video_token_id"] = getattr(vision_source, "video_token_id", None) + fields["vision_start_token_id"] = getattr(vision_source, "vision_start_token_id", None) + fields["vision_end_token_id"] = getattr(vision_source, "vision_end_token_id", None) # MRoPE section — only for composite VL models (parent_config != config). if parent_config is not None and parent_config is not config: diff --git a/src/mobius/_registry.py b/src/mobius/_registry.py index 698beed1c..033f4dec7 100644 --- a/src/mobius/_registry.py +++ b/src/mobius/_registry.py @@ -1020,7 +1020,7 @@ def _create_default_registry() -> ModelRegistry: "qwen2_5_vl_text": "Qwen/Qwen2.5-VL-3B-Instruct", "qwen3_vl": "Qwen/Qwen3-VL-2B-Instruct", "qwen3_vl_text": "Qwen/Qwen3-VL-2B-Instruct", - "qwen3_5": "Qwen/Qwen3.5-2B", + "qwen3_5": "Qwen/Qwen3.8-27B", "llava": "llava-hf/llava-1.5-7b-hf", "llava_next": "llava-hf/llava-v1.6-mistral-7b-hf", "mllama": "meta-llama/Llama-3.2-11B-Vision-Instruct", diff --git a/src/mobius/components/_qwen3_vl_vision.py b/src/mobius/components/_qwen3_vl_vision.py index 1fc983989..ccad7551d 100644 --- a/src/mobius/components/_qwen3_vl_vision.py +++ b/src/mobius/components/_qwen3_vl_vision.py @@ -28,11 +28,6 @@ from mobius._build_context import ep_capabilities, get_build_dtype from mobius.components._common import LayerNorm, Linear, build_packed_token_offset from mobius.components._mlp import FCMLP -from mobius.components._scan_utils import ( - compact_scan_output, - create_body_graph, - rename_subgraph_values, -) class Qwen3VLPatchEmbed(nn.Module): @@ -403,69 +398,6 @@ def forward(self, op: OpBuilder, hidden_states: ir.Value): return self.linear_fc2(op, x) -def _qwen3_rotary_pos_ids_one_image(op, T, H, W, ms): # noqa: N803 - """Compute 2D rotary position IDs for one image (Qwen3-VL style). - - Uses block_rows * ms + intra indexing for spatial-merge groups. - Works with any OpBuilder (main graph or Scan body graph). - - Args: - op: OpBuilder instance. - T, H, W: Scalar INT64 values. - ms: Python int — spatial merge size. - - Returns: - ``(T*H*W, 2)`` INT64 position IDs. - """ - H_m = op.Div(H, op.Constant(value_int=ms)) # noqa: N806 - W_m = op.Div(W, op.Constant(value_int=ms)) # noqa: N806 - - # Block row/col indices and intra-merge indices - block_rows = op.Range( - op.Constant(value_int=0), - H_m, - op.Constant(value_int=1), - ) - block_cols = op.Range( - op.Constant(value_int=0), - W_m, - op.Constant(value_int=1), - ) - intra = op.Range( - op.Constant(value_int=0), - op.Constant(value_int=ms), - op.Constant(value_int=1), - ) - - # row_idx = block_rows[:,None,None,None] * ms + intra[None,None,:,None] - br = op.Mul(op.Unsqueeze(block_rows, [1, 2, 3]), op.Constant(value_int=ms)) - ir_row = op.Unsqueeze(intra, [0, 1, 3]) - row_idx = op.Add(br, ir_row) - - bc = op.Mul(op.Unsqueeze(block_cols, [0, 2, 3]), op.Constant(value_int=ms)) - ir_col = op.Unsqueeze(intra, [0, 1, 2]) - col_idx = op.Add(bc, ir_col) - - # Expand to (H_m, W_m, ms, ms) and flatten - row_shape = op.Concat( - op.Reshape(H_m, [1]), - op.Reshape(W_m, [1]), - op.Constant(value_ints=[ms, ms]), - axis=0, - ) - row_flat = op.Reshape(op.Expand(row_idx, row_shape), [-1]) - col_flat = op.Reshape(op.Expand(col_idx, row_shape), [-1]) - - # Stack to (H*W, 2) and tile T times - pos_ids = op.Concat( - op.Unsqueeze(row_flat, [1]), - op.Unsqueeze(col_flat, [1]), - axis=1, - ) - tile_t = op.Concat(op.Reshape(T, [1]), op.Constant(value_ints=[1]), axis=0) - return op.Tile(pos_ids, tile_t) # (T*H*W, 2) - - class Qwen3VLVisionModel(nn.Module): """Full Qwen3-VL vision encoder with DeepStack outputs. @@ -558,290 +490,155 @@ def __init__( ] ) - def _interpolate_pos_embed(self, op, grid_thw): - """Bilinear interpolation of learned position embeddings for all images. - - Iterates over ``grid_thw`` via ONNX Scan, computing per-image - bilinear interpolation from the learned position grid and - concatenating results. + def _flat_grid_coordinates(self, op, grid_thw): + """Map each packed patch to its media row and merge-permuted H/W coordinates. - Matches HuggingFace ``Qwen3VLVisionModel.fast_pos_embed_interpolate``. - - Args: - op: OpBuilder instance. - grid_thw: ``(num_images, 3)`` INT64 with ``[T, H, W]`` per image. - - Returns: - Position embeddings ``(total_patches, hidden_size)``. + Qwen3-VL stores each media item's patches in + ``(T, H // ms, W // ms, ms, ms)`` order. Computing coordinates over + the concatenated patch stream avoids control-flow subgraphs while + preserving arbitrary image/video sizes and order. """ - n = self.num_grid_per_side ms = self.spatial_merge_size - hidden_size = self.hidden_size - n_minus_1 = float(n - 1) - - # Per-image patch counts T_col = op.Squeeze(op.Slice(grid_thw, [0], [1], [1], [1]), [1]) # noqa: N806 H_col = op.Squeeze(op.Slice(grid_thw, [1], [2], [1], [1]), [1]) # noqa: N806 W_col = op.Squeeze(op.Slice(grid_thw, [2], [3], [1], [1]), [1]) # noqa: N806 - patches_per_image = op.Mul(T_col, op.Mul(H_col, W_col)) - max_patches = op.ReduceMax(patches_per_image, keepdims=False) - - # --- Scan body: interpolate pos embeddings for one image --- - body_thw = ir.Value( - name="body_thw", - shape=ir.Shape([3]), - type=ir.TensorType(ir.DataType.INT64), - ) - body_graph, body_builder = create_body_graph([], [body_thw]) - body_op = body_builder.op - - bT = body_op.Squeeze(body_op.Gather(body_thw, body_op.Constant(value_int=0))) # noqa: N806 - bH = body_op.Squeeze(body_op.Gather(body_thw, body_op.Constant(value_int=1))) # noqa: N806 - bW = body_op.Squeeze(body_op.Gather(body_thw, body_op.Constant(value_int=2))) # noqa: N806 - - # linspace(0, n-1, H) and linspace(0, n-1, W) - H_f = body_op.Cast(bH, to=1) # noqa: N806 - W_f = body_op.Cast(bW, to=1) # noqa: N806 - h_range = body_op.Cast( - body_op.Range( - body_op.Constant(value_int=0), - bH, - body_op.Constant(value_int=1), - ), - to=1, - ) - w_range = body_op.Cast( - body_op.Range( - body_op.Constant(value_int=0), - bW, - body_op.Constant(value_int=1), - ), - to=1, - ) - - h_idxs = body_op.Div( - body_op.Mul(h_range, n_minus_1), - body_op.Sub(H_f, 1.0), - ) - w_idxs = body_op.Div( - body_op.Mul(w_range, n_minus_1), - body_op.Sub(W_f, 1.0), - ) - - # Floor/ceil indices - h_floor = body_op.Cast(body_op.Floor(h_idxs), to=7) - w_floor = body_op.Cast(body_op.Floor(w_idxs), to=7) - clip_max = body_op.Constant(value_int=n - 1) - h_ceil = body_op.Min( - body_op.Add(h_floor, body_op.Constant(value_int=1)), - clip_max, - ) - w_ceil = body_op.Min( - body_op.Add(w_floor, body_op.Constant(value_int=1)), - clip_max, + patches_per_media = op.Mul(T_col, op.Mul(H_col, W_col)) + patch_ends = op.CumSum(patches_per_media, op.Constant(value_int=0)) + patch_starts = op.Pad( + patch_ends, + op.Constant(value_ints=[1, 0]), + op.Constant(value_int=0), ) - # Bilinear weights - dh = body_op.Sub(h_idxs, body_op.Cast(h_floor, to=1)) - dw = body_op.Sub(w_idxs, body_op.Cast(w_floor, to=1)) - - n_const = body_op.Constant(value_int=n) - base_h_floor = body_op.Mul(h_floor, n_const) - base_h_ceil = body_op.Mul(h_ceil, n_const) - - bh_f2 = body_op.Unsqueeze(base_h_floor, [1]) - bh_c2 = body_op.Unsqueeze(base_h_ceil, [1]) - wf2 = body_op.Unsqueeze(w_floor, [0]) - wc2 = body_op.Unsqueeze(w_ceil, [0]) - - idx_00 = body_op.Reshape(body_op.Add(bh_f2, wf2), [-1]) - idx_01 = body_op.Reshape(body_op.Add(bh_f2, wc2), [-1]) - idx_10 = body_op.Reshape(body_op.Add(bh_c2, wf2), [-1]) - idx_11 = body_op.Reshape(body_op.Add(bh_c2, wc2), [-1]) - - one_minus_dh = body_op.Sub(1.0, dh) - one_minus_dw = body_op.Sub(1.0, dw) - dh2 = body_op.Unsqueeze(dh, [1]) - omdh2 = body_op.Unsqueeze(one_minus_dh, [1]) - dw2 = body_op.Unsqueeze(dw, [0]) - omdw2 = body_op.Unsqueeze(one_minus_dw, [0]) - - w_00 = body_op.Reshape(body_op.Mul(omdh2, omdw2), [-1, 1]) - w_01 = body_op.Reshape(body_op.Mul(omdh2, dw2), [-1, 1]) - w_10 = body_op.Reshape(body_op.Mul(dh2, omdw2), [-1, 1]) - w_11 = body_op.Reshape(body_op.Mul(dh2, dw2), [-1, 1]) - - # Gather from learned pos_embed (implicit input from parent graph). - # Cast to float32 for bilinear interpolation (pos_embed may be bf16/f16). - e_00 = body_op.Mul(body_op.Cast(body_op.Gather(self.pos_embed, idx_00), to=1), w_00) - e_01 = body_op.Mul(body_op.Cast(body_op.Gather(self.pos_embed, idx_01), to=1), w_01) - e_10 = body_op.Mul(body_op.Cast(body_op.Gather(self.pos_embed, idx_10), to=1), w_10) - e_11 = body_op.Mul(body_op.Cast(body_op.Gather(self.pos_embed, idx_11), to=1), w_11) - pos_embeds = body_op.Add( - body_op.Add(e_00, e_01), - body_op.Add(e_10, e_11), + total_patches = op.ReduceSum(patches_per_media, keepdims=False) + patch_ids = op.Range( + op.Constant(value_int=0), + total_patches, + op.Constant(value_int=1), ) + # The number of completed media ranges is the owning media row. + media_ids = op.ReduceSum( + op.Cast( + op.GreaterOrEqual( + op.Unsqueeze(patch_ids, [1]), + op.Unsqueeze(patch_ends, [0]), + ), + to=7, + ), + [1], + keepdims=False, + ) + local_ids = op.Sub(patch_ids, op.Gather(patch_starts, media_ids)) + + H = op.Gather(H_col, media_ids) # noqa: N806 + W = op.Gather(W_col, media_ids) # noqa: N806 + patches_per_frame = op.Mul(H, W) + frame_local_ids = op.Mod(local_ids, patches_per_frame) + + merge_area = op.Constant(value_int=ms * ms) + merge_block_ids = op.Div(frame_local_ids, merge_area) + intra_merge_ids = op.Mod(frame_local_ids, merge_area) + W_m = op.Div(W, op.Constant(value_int=ms)) # noqa: N806 + block_rows = op.Div(merge_block_ids, W_m) + block_cols = op.Mod(merge_block_ids, W_m) + intra_rows = op.Div(intra_merge_ids, op.Constant(value_int=ms)) + intra_cols = op.Mod(intra_merge_ids, op.Constant(value_int=ms)) + rows = op.Add(op.Mul(block_rows, op.Constant(value_int=ms)), intra_rows) + cols = op.Add(op.Mul(block_cols, op.Constant(value_int=ms)), intra_cols) + return rows, cols, H, W - # Tile T times: (H*W, D) → (T*H*W, D) - T_tile = body_op.Concat( # noqa: N806 - body_op.Reshape(bT, [1]), - body_op.Constant(value_ints=[1]), - axis=0, - ) - pos_embeds = body_op.Tile(pos_embeds, T_tile) - - # Spatial merge permutation: - # (T, H//ms, ms, W//ms, ms, D) → (T, H//ms, W//ms, ms, ms, D) - H_m = body_op.Div(bH, body_op.Constant(value_int=ms)) # noqa: N806 - W_m = body_op.Div(bW, body_op.Constant(value_int=ms)) # noqa: N806 - shape_6d = body_op.Concat( - body_op.Reshape(bT, [1]), - body_op.Reshape(H_m, [1]), - body_op.Constant(value_ints=[ms]), - body_op.Reshape(W_m, [1]), - body_op.Constant(value_ints=[ms]), - body_op.Constant(value_ints=[hidden_size]), - axis=0, - ) - pos_embeds = body_op.Reshape(pos_embeds, shape_6d) - pos_embeds = body_op.Transpose(pos_embeds, perm=[0, 1, 3, 2, 4, 5]) - pos_embeds = body_op.Reshape(pos_embeds, [-1, hidden_size]) - - # Pad to (max_patches, hidden_size) — implicit input from main graph - num_p = body_op.Mul(bT, body_op.Mul(bH, bW)) - pad_len = body_op.Reshape(body_op.Sub(max_patches, num_p), [1]) - pads = body_op.Concat( - body_op.Constant(value_ints=[0, 0]), - pad_len, - body_op.Constant(value_ints=[0]), - axis=0, - ) - padded = body_op.Pad(pos_embeds, pads, 0.0) - padded.name = "padded_pos_embed" - body_graph.outputs.append(padded) + def _interpolate_pos_embed(self, op, grid_thw): + """Bilinearly interpolate learned positions for the packed media stream. - rename_subgraph_values(body_graph, "posemb_body_") + Matches HuggingFace ``Qwen3VLVisionModel.fast_pos_embed_interpolate``. - scan_result = op.Scan( - grid_thw, - body=body_graph, - num_scan_inputs=1, - _outputs=1, - ) # (num_images, max_patches, hidden_size) + Args: + op: OpBuilder instance. + grid_thw: ``(num_images, 3)`` INT64 with ``[T, H, W]`` per image. - return compact_scan_output(op, scan_result, patches_per_image) + Returns: + Position embeddings ``(total_patches, hidden_size)``. + """ + n = self.num_grid_per_side + rows, cols, H, W = self._flat_grid_coordinates(op, grid_thw) # noqa: N806 + rows_f = op.Cast(rows, to=1) + cols_f = op.Cast(cols, to=1) + H_f = op.Cast(H, to=1) # noqa: N806 + W_f = op.Cast(W, to=1) # noqa: N806 + rows_scaled = op.Div(op.Mul(rows_f, float(n - 1)), op.Sub(H_f, 1.0)) + cols_scaled = op.Div(op.Mul(cols_f, float(n - 1)), op.Sub(W_f, 1.0)) + + row_floor = op.Cast(op.Floor(rows_scaled), to=7) + col_floor = op.Cast(op.Floor(cols_scaled), to=7) + clip_max = op.Constant(value_int=n - 1) + row_ceil = op.Min(op.Add(row_floor, op.Constant(value_int=1)), clip_max) + col_ceil = op.Min(op.Add(col_floor, op.Constant(value_int=1)), clip_max) + + row_delta = op.Sub(rows_scaled, op.Cast(row_floor, to=1)) + col_delta = op.Sub(cols_scaled, op.Cast(col_floor, to=1)) + one_minus_row = op.Sub(1.0, row_delta) + one_minus_col = op.Sub(1.0, col_delta) + w_00 = op.Unsqueeze(op.Mul(one_minus_row, one_minus_col), [1]) + w_01 = op.Unsqueeze(op.Mul(one_minus_row, col_delta), [1]) + w_10 = op.Unsqueeze(op.Mul(row_delta, one_minus_col), [1]) + w_11 = op.Unsqueeze(op.Mul(row_delta, col_delta), [1]) + + row_floor_base = op.Mul(row_floor, op.Constant(value_int=n)) + row_ceil_base = op.Mul(row_ceil, op.Constant(value_int=n)) + idx_00 = op.Add(row_floor_base, col_floor) + idx_01 = op.Add(row_floor_base, col_ceil) + idx_10 = op.Add(row_ceil_base, col_floor) + idx_11 = op.Add(row_ceil_base, col_ceil) + + # Interpolate in float32 even when the learned table is f16/bf16. + e_00 = op.Mul(op.Cast(op.Gather(self.pos_embed, idx_00), to=1), w_00) + e_01 = op.Mul(op.Cast(op.Gather(self.pos_embed, idx_01), to=1), w_01) + e_10 = op.Mul(op.Cast(op.Gather(self.pos_embed, idx_10), to=1), w_10) + e_11 = op.Mul(op.Cast(op.Gather(self.pos_embed, idx_11), to=1), w_11) + return op.Add(op.Add(e_00, e_01), op.Add(e_10, e_11)) def _compute_rotary_pos_ids(self, op, grid_thw): """Compute 2D rotary position IDs for all images via ONNX Scan. Matches HF ``Qwen3VLVisionModel.rot_pos_emb()`` position indexing. - Iterates over ``grid_thw`` rows, computing per-image spatial-merge- - permuted position IDs and concatenating. Returns ``(total_patches, 2)`` INT64 with ``[h_pos, w_pos]`` per patch. """ - ms = self.spatial_merge_size - - # Per-image patch counts for padding/compaction - T_col = op.Squeeze(op.Slice(grid_thw, [0], [1], [1], [1]), [1]) # noqa: N806 - H_col = op.Squeeze(op.Slice(grid_thw, [1], [2], [1], [1]), [1]) # noqa: N806 - W_col = op.Squeeze(op.Slice(grid_thw, [2], [3], [1], [1]), [1]) # noqa: N806 - patches_per_image = op.Mul(T_col, op.Mul(H_col, W_col)) - max_patches = op.ReduceMax(patches_per_image, keepdims=False) - - # --- Scan body: compute pos_ids for one image, pad to max_patches --- - body_thw = ir.Value( - name="body_thw", - shape=ir.Shape([3]), - type=ir.TensorType(ir.DataType.INT64), - ) - body_graph, body_builder = create_body_graph([], [body_thw]) - body_op = body_builder.op - - bT = body_op.Squeeze(body_op.Gather(body_thw, body_op.Constant(value_int=0))) # noqa: N806 - bH = body_op.Squeeze(body_op.Gather(body_thw, body_op.Constant(value_int=1))) # noqa: N806 - bW = body_op.Squeeze(body_op.Gather(body_thw, body_op.Constant(value_int=2))) # noqa: N806 - - pos_ids = _qwen3_rotary_pos_ids_one_image(body_op, bT, bH, bW, ms) - - # Pad to (max_patches, 2) — implicit input from main graph - num_p = body_op.Mul(bT, body_op.Mul(bH, bW)) - pad_len = body_op.Reshape(body_op.Sub(max_patches, num_p), [1]) - pads = body_op.Concat( - body_op.Constant(value_ints=[0, 0]), - pad_len, - body_op.Constant(value_ints=[0]), - axis=0, - ) - padded = body_op.Pad(pos_ids, pads, body_op.Constant(value_int=-1)) - padded.name = "padded_pos_ids" - body_graph.outputs.append(padded) - - rename_subgraph_values(body_graph, "q3_rotary_body_") - - scan_result = op.Scan( - grid_thw, - body=body_graph, - num_scan_inputs=1, - _outputs=1, - ) - return compact_scan_output(op, scan_result, patches_per_image) + rows, cols, _, _ = self._flat_grid_coordinates(op, grid_thw) + return op.Concat(op.Unsqueeze(rows, [1]), op.Unsqueeze(cols, [1]), axis=1) def _compute_cu_seqlens(self, op, grid_thw): """Compute full-attention cu_seqlens for all images. - Produces per-frame boundaries across all images using ONNX Scan - to handle per-image ``repeat_interleave(hw, T)`` + CumSum. + Produces per-frame boundaries across all images without a control-flow + subgraph, equivalent to ``repeat_interleave(H * W, T)`` + CumSum. Returns ``(total_frames + 1,)`` INT64. """ T_col = op.Squeeze(op.Slice(grid_thw, [0], [1], [1], [1]), [1]) # noqa: N806 - max_T = op.ReduceMax(T_col, keepdims=False) # noqa: N806 - - # Scan body: for each image, output T copies of hw, padded to max_T - body_thw = ir.Value( - name="body_thw", - shape=ir.Shape([3]), - type=ir.TensorType(ir.DataType.INT64), - ) - body_graph, body_builder = create_body_graph([], [body_thw]) - body_op = body_builder.op - - bT = body_op.Squeeze(body_op.Gather(body_thw, body_op.Constant(value_int=0))) # noqa: N806 - bH = body_op.Squeeze(body_op.Gather(body_thw, body_op.Constant(value_int=1))) # noqa: N806 - bW = body_op.Squeeze(body_op.Gather(body_thw, body_op.Constant(value_int=2))) # noqa: N806 - - hw = body_op.Mul(bH, bW) - ones = body_op.Expand( - body_op.Constant(value_int=1), - body_op.Reshape(bT, [1]), - ) - hw_repeated = body_op.Mul(ones, hw) - - pad_len = body_op.Reshape(body_op.Sub(max_T, bT), [1]) - pads = body_op.Concat( - body_op.Constant(value_ints=[0]), - pad_len, - axis=0, - ) - padded = body_op.Pad( - hw_repeated, - pads, - body_op.Constant(value_int=0), + H_col = op.Squeeze(op.Slice(grid_thw, [1], [2], [1], [1]), [1]) # noqa: N806 + W_col = op.Squeeze(op.Slice(grid_thw, [2], [3], [1], [1]), [1]) # noqa: N806 + frame_ends = op.CumSum(T_col, op.Constant(value_int=0)) + total_frames = op.ReduceSum(T_col, keepdims=False) + frame_ids = op.Range( + op.Constant(value_int=0), + total_frames, + op.Constant(value_int=1), ) - padded.name = "padded_hw" - body_graph.outputs.append(padded) - - rename_subgraph_values(body_graph, "q3_cu_body_") - - scan_hw = op.Scan( - grid_thw, - body=body_graph, - num_scan_inputs=1, - _outputs=1, + media_ids = op.ReduceSum( + op.Cast( + op.GreaterOrEqual( + op.Unsqueeze(frame_ids, [1]), + op.Unsqueeze(frame_ends, [0]), + ), + to=7, + ), + [1], + keepdims=False, ) - hw_flat = compact_scan_output(op, scan_hw, T_col) - cu = op.CumSum(hw_flat, op.Constant(value_int=0)) + hw_per_media = op.Mul(H_col, W_col) + hw_per_frame = op.Gather(hw_per_media, media_ids) + cu = op.CumSum(hw_per_frame, op.Constant(value_int=0)) return op.Pad(cu, op.Constant(value_ints=[1, 0]), op.Constant(value_int=0)) def forward( @@ -866,7 +663,7 @@ def forward( hidden_states = self.patch_embed(op, hidden_states) # Bilinear-interpolated position embeddings from learned grid. - # Cast to match hidden_states dtype (Scan body computes in float32). + # Cast to match hidden_states dtype (interpolation computes in float32). pos_embeds = self._interpolate_pos_embed(op, grid_thw) pos_embeds = op.CastLike(pos_embeds, hidden_states) hidden_states = op.Add(hidden_states, pos_embeds) diff --git a/src/mobius/models/qwen35.py b/src/mobius/models/qwen35.py index ea46d8932..d85ea5f5f 100644 --- a/src/mobius/models/qwen35.py +++ b/src/mobius/models/qwen35.py @@ -223,9 +223,10 @@ def preprocess_weights( - Stripping ``language_model.`` prefix from HF checkpoint keys (HF stores weights as ``model.language_model.*`` in safetensors) - Dropping visual encoder keys (``model.visual.*``) - - Dropping multi-token prediction (MTP) keys (``mtp*``): - MTP heads are auxiliary decoding heads used only during - HuggingFace training; they are not needed for inference. + - Dropping multi-token prediction (MTP) keys (``mtp*``). The target + model's normal forward path does not consume this optional + self-speculative drafter; it is packaged separately as + :class:`Qwen35MtpModel` when speculative decoding is requested. - Weight tying (``tie_word_embeddings``) """ cleaned: dict[str, torch.Tensor] = {} @@ -376,9 +377,9 @@ def preprocess_weights( """Preprocess HuggingFace state dict for Qwen3.5-MoE. Handles: - - Dropping multi-token prediction (MTP) keys (``mtp_*``, ``mtp.*``): - MTP heads are auxiliary decoding heads used only during - HuggingFace training; they are not needed for inference. + - Dropping multi-token prediction (MTP) keys (``mtp_*``, ``mtp.*``). + The target model's normal forward path does not consume this optional + self-speculative drafter, which has a separate package contract. - Stripping ``language_model.`` prefix from HF checkpoint keys (HF stores weights as ``model.language_model.*`` in safetensors) - Dropping visual encoder keys (``model.visual.*``) @@ -537,9 +538,8 @@ def preprocess_weights( """ renamed: dict[str, torch.Tensor] = {} for key, value in state_dict.items(): - # Drop multi-token prediction (MTP) keys: MTP heads are - # auxiliary decoding heads used only during HuggingFace - # training; they are not needed for inference. + # The standard target package excludes the optional + # self-speculative MTP drafter, which has a separate graph contract. if key.startswith(("mtp_", "mtp.")): continue @@ -611,7 +611,7 @@ def preprocess_weights( """Route language_model weights for standalone decoder build.""" renamed: dict[str, torch.Tensor] = {} for key, value in state_dict.items(): - # Drop MTP heads (training-only auxiliary decoders) + # The optional self-speculative MTP drafter is packaged separately. if key.startswith(("mtp_", "mtp.")): continue stripped = key diff --git a/src/mobius/models/qwen35_test.py b/src/mobius/models/qwen35_test.py index 5ede02207..798523641 100644 --- a/src/mobius/models/qwen35_test.py +++ b/src/mobius/models/qwen35_test.py @@ -13,14 +13,241 @@ from __future__ import annotations +import numpy as np +import onnx_ir as ir import torch +from transformers.models.qwen3_5.configuration_qwen3_5 import Qwen3_5Config -from mobius._configs import QuantizationConfig +from mobius._configs import ArchitectureConfig, QuantizationConfig, Qwen35MtpConfig +from mobius._registry import registry from mobius._testing import make_config -from mobius.models.qwen35 import Qwen35MoECausalLMModel +from mobius._testing.ort_inference import OnnxModelSession +from mobius.models.qwen35 import Qwen35MoECausalLMModel, Qwen35VL3ModelCausalLMModel +from mobius.models.qwen_vl import Qwen3VLEmbeddingModel +from mobius.tasks import build_embedding_from_features _E, _H, _INT, _BLK, _BITS = 8, 32, 16, 16, 4 _FC1_OUT = 2 * _INT +_QWEN38_REVISION = "1d4bf0f2ff6012fd82039f2fa52739d0dd7c60c0" + + +def _qwen38_config() -> Qwen3_5Config: + layer_types = [ + layer_type + for _ in range(16) + for layer_type in ( + "linear_attention", + "linear_attention", + "linear_attention", + "full_attention", + ) + ] + return Qwen3_5Config( + image_token_id=248056, + video_token_id=248057, + vision_start_token_id=248053, + vision_end_token_id=248054, + tie_word_embeddings=False, + text_config={ + "model_type": "qwen3_5_text", + "vocab_size": 248320, + "hidden_size": 5120, + "intermediate_size": 17408, + "num_hidden_layers": 64, + "num_attention_heads": 24, + "num_key_value_heads": 4, + "head_dim": 256, + "hidden_act": "silu", + "rms_norm_eps": 1e-6, + "max_position_embeddings": 262144, + "layer_types": layer_types, + "linear_conv_kernel_dim": 4, + "linear_key_head_dim": 128, + "linear_value_head_dim": 128, + "linear_num_key_heads": 16, + "linear_num_value_heads": 48, + "partial_rotary_factor": 0.25, + "rope_parameters": { + "rope_type": "default", + "rope_theta": 10_000_000, + "partial_rotary_factor": 0.25, + "mrope_interleaved": True, + "mrope_section": [11, 11, 10], + }, + "mtp_num_hidden_layers": 1, + "tie_word_embeddings": False, + }, + vision_config={ + "model_type": "qwen3_5", + "depth": 27, + "hidden_size": 1152, + "intermediate_size": 4304, + "num_heads": 16, + "patch_size": 16, + "temporal_patch_size": 2, + "spatial_merge_size": 2, + "out_hidden_size": 5120, + "num_position_embeddings": 2304, + "deepstack_visual_indexes": [], + }, + ) + + +class TestQwen38Alias: + def test_exact_config_extracts_dense_hybrid_vl_architecture(self): + hf_config = _qwen38_config() + config = ArchitectureConfig.from_transformers( + hf_config.text_config, + parent_config=hf_config, + ) + + assert _QWEN38_REVISION == "1d4bf0f2ff6012fd82039f2fa52739d0dd7c60c0" + assert registry.get("qwen3_5") is Qwen35VL3ModelCausalLMModel + assert registry.get("qwen3_5_vl") is Qwen35VL3ModelCausalLMModel + assert registry.get_registration("qwen3_5").test_model_id == "Qwen/Qwen3.8-27B" + assert registry.get_registration("qwen3_5_vl").test_model_id == "Qwen/Qwen3.5-2B" + assert config.hidden_size == 5120 + assert config.intermediate_size == 17408 + assert config.num_hidden_layers == 64 + assert config.layer_types == hf_config.text_config.layer_types + assert config.layer_types.count("linear_attention") == 48 + assert config.layer_types.count("full_attention") == 16 + assert config.num_attention_heads == 24 + assert config.num_key_value_heads == 4 + assert config.head_dim == 256 + assert np.isclose(config.partial_rotary_factor, 0.25) + assert config.mrope_interleaved is True + assert config.mrope_section == [11, 11, 10] + assert config.linear_num_key_heads == 16 + assert config.linear_num_value_heads == 48 + assert config.linear_key_head_dim == 128 + assert config.linear_value_head_dim == 128 + assert config.linear_conv_kernel_dim == 4 + assert config.vocab_size == 248320 + assert config.image_token_id == 248056 + assert config.video_token_id == 248057 + assert config.vision_start_token_id == 248053 + assert config.vision_end_token_id == 248054 + assert config.vision is not None + assert config.vision.num_hidden_layers == 27 + assert config.vision.hidden_size == 1152 + assert config.vision.intermediate_size == 4304 + assert config.vision.num_attention_heads == 16 + assert config.vision.patch_size == 16 + assert config.vision.temporal_patch_size == 2 + assert config.vision.spatial_merge_size == 2 + assert config.vision.out_hidden_size == 5120 + assert config.vision.deepstack_visual_indexes == [] + + def test_one_layer_mtp_is_classified_as_separate_optional_drafter(self): + hf_config = _qwen38_config() + assert hf_config.text_config.mtp_num_hidden_layers == 1 + + mtp_config = Qwen35MtpConfig.from_transformers(hf_config) + assert mtp_config.num_hidden_layers == 1 + assert mtp_config.layer_types == ["full_attention"] + assert registry.get_registration("Qwen35MtpModel").task == "qwen35-mtp" + + def test_weight_routing_excludes_separately_packaged_mtp(self): + config = ArchitectureConfig.from_transformers( + _qwen38_config().text_config, + parent_config=_qwen38_config(), + ) + model = Qwen35VL3ModelCausalLMModel(config) + state_dict = { + "model.language_model.embed_tokens.weight": torch.ones(2, 2), + "model.language_model.layers.0.linear_attn.A_log": torch.ones(2), + "model.language_model.layers.3.self_attn.q_proj.weight": torch.ones(2, 2), + "model.visual.blocks.0.mlp.linear_fc1.weight": torch.ones(2, 2), + "lm_head.weight": torch.ones(2, 2), + "mtp.layers.0.self_attn.q_proj.weight": torch.ones(2, 2), + "mtp.fc.weight": torch.ones(2, 2), + } + + result = model.preprocess_weights(state_dict) + + assert "decoder.model.embed_tokens.weight" in result + assert "embedding.embed_tokens.weight" in result + assert "decoder.model.layers.0.linear_attn.A_log" in result + assert "decoder.model.layers.3.self_attn.q_proj.weight" in result + assert "vision_encoder.visual.blocks.0.mlp.up_proj.weight" in result + assert "decoder.lm_head.weight" in result + assert not any(key.startswith("mtp") or ".mtp." in key for key in result) + + def test_qwen_vl_processor_boundary_stays_float32_for_bf16_export(self): + hf_config = _qwen38_config() + config = ArchitectureConfig.from_transformers( + hf_config.text_config, + parent_config=hf_config, + ) + config.dtype = ir.DataType.BFLOAT16 + package = Qwen35VL3ModelCausalLMModel(config) + task = registry.get_registration("qwen3_5").task + + from mobius.tasks import get_task + + vision_model = get_task(task).build(package, config)["vision_encoder"] + + assert vision_model.graph.inputs[0].name == "pixel_values" + assert vision_model.graph.inputs[0].dtype == ir.DataType.FLOAT + assert any(node.op_type == "Cast" for node in vision_model.graph) + + def test_embedding_scatter_matches_separate_image_then_video_streams(self): + config = ArchitectureConfig( + vocab_size=16, + hidden_size=4, + pad_token_id=0, + image_token_id=10, + video_token_id=11, + dtype=ir.DataType.FLOAT, + ) + graph = build_embedding_from_features( + Qwen3VLEmbeddingModel(config), + config, + feature_name="image_features", + feature_dim=config.hidden_size, + ) + embedding_weight = np.arange( + config.vocab_size * config.hidden_size, + dtype=np.float32, + ).reshape(config.vocab_size, config.hidden_size) + for name, initializer in graph.graph.initializers.items(): + if name.endswith("embed_tokens.weight"): + initializer.const_value = ir.tensor(embedding_weight) + + input_ids = np.array( + [ + [config.video_token_id, 1, config.image_token_id], + [2, config.image_token_id, config.video_token_id], + ], + dtype=np.int64, + ) + # HF scatters the two image rows first, then the two video rows. + media_features = np.arange(100, 116, dtype=np.float32).reshape(4, 4) + session = OnnxModelSession(graph) + result = session.run( + { + "input_ids": input_ids, + "image_features": media_features, + } + )["inputs_embeds"] + + expected = embedding_weight[input_ids].copy() + expected[0, 2] = media_features[0] + expected[1, 1] = media_features[1] + expected[0, 0] = media_features[2] + expected[1, 2] = media_features[3] + np.testing.assert_array_equal(result, expected) + + decode_ids = np.array([[3], [4]], dtype=np.int64) + decode = session.run( + { + "input_ids": decode_ids, + "image_features": np.empty((0, config.hidden_size), dtype=np.float32), + } + )["inputs_embeds"] + session.close() + np.testing.assert_array_equal(decode, embedding_weight[decode_ids]) def _moe_config(quantization: QuantizationConfig | None) -> object: diff --git a/src/mobius/models/qwen_vl.py b/src/mobius/models/qwen_vl.py index b97a0d4e5..2b4ae302e 100644 --- a/src/mobius/models/qwen_vl.py +++ b/src/mobius/models/qwen_vl.py @@ -995,12 +995,12 @@ def preprocess_weights( class Qwen3VLEmbeddingModel(Qwen25VLEmbeddingModel): """Qwen3-VL embedding model for the 3-model split. - Scatters merged image features at image-token positions (like - Qwen2.5-VL) and, when the vision encoder produces DeepStack features, - also scatters each intermediate DeepStack map into a full-length - ``[batch, seq, hidden]`` tensor (zero at non-image positions). The - stacked ``deepstack_embeds`` output is consumed by the decoder, which - adds them to the hidden states of its first ``D`` layers. + Scatters packed image-then-video features at their respective placeholder + positions. When the vision encoder produces DeepStack features, each + intermediate map is scattered with the same media ordering into a + full-length ``[batch, seq, hidden]`` tensor. The stacked + ``deepstack_embeds`` output is consumed by the decoder, which adds them to + the hidden states of its first ``D`` layers. Inputs: - input_ids: (batch, seq_len) INT64 @@ -1013,6 +1013,10 @@ class Qwen3VLEmbeddingModel(Qwen25VLEmbeddingModel): (only when DeepStack is active) """ + def __init__(self, config: ArchitectureConfig): + super().__init__(config) + self.video_token_id = config.video_token_id + def forward( self, op: OpBuilder, @@ -1022,21 +1026,44 @@ def forward( ): text_embeds = self.embed_tokens(op, input_ids) - # Image-token positions and their running index into the packed - # feature tensors (shared by the main image scatter and every - # DeepStack scatter). + # Hugging Face scatters image and video streams independently. The + # package therefore packs every image feature first, then every video + # feature, regardless of placeholder order or batch row. image_mask = op.Equal(input_ids, op.Constant(value_int=self.image_token_id)) - image_mask_3d = op.Unsqueeze(image_mask, [-1]) - mask_int = op.Cast(image_mask, to=7) # INT64 - cumsum = op.CumSum(mask_int, op.Constant(value_int=1)) - indices = op.Clip( - op.Sub(cumsum, op.Constant(value_int=1)), - op.Constant(value_int=0), + if self.video_token_id is None: + video_mask = op.CastLike(False, image_mask) + else: + video_mask = op.Equal( + input_ids, + op.Constant(value_int=self.video_token_id), + ) + media_mask = op.Or(image_mask, video_mask) + media_mask_3d = op.Unsqueeze(media_mask, [-1]) + + flat_image_mask = op.Cast(op.Reshape(image_mask, [-1]), to=7) + flat_video_mask = op.Cast(op.Reshape(video_mask, [-1]), to=7) + image_indices = op.Sub( + op.CumSum(flat_image_mask, op.Constant(value_int=0)), + op.Constant(value_int=1), + ) + video_indices = op.Add( + op.Sub( + op.CumSum(flat_video_mask, op.Constant(value_int=0)), + op.Constant(value_int=1), + ), + op.ReduceSum(flat_image_mask, keepdims=0), + ) + flat_indices = op.Where( + op.CastLike(flat_image_mask, image_mask), + image_indices, + video_indices, ) + flat_indices = op.Clip(flat_indices, op.Constant(value_int=0)) + indices = op.Reshape(flat_indices, op.Shape(input_ids)) def _scatter(features: ir.Value, fallback: ir.Value) -> ir.Value: - # Pad with one zero row so Gather stays in-bounds for text-only - # input (num_image_tokens == 0); the Where mask discards it. + # Keep Gather valid for text-only/decode calls with zero media rows; + # the Where mask discards the synthetic row. pad_row = op.Expand( op.CastLike(0.0, features), op.Concat( @@ -1047,7 +1074,7 @@ def _scatter(features: ir.Value, fallback: ir.Value) -> ir.Value: ) padded = op.Concat(features, pad_row, axis=0) gathered = op.Gather(padded, indices, axis=0) - return op.Where(image_mask_3d, gathered, fallback) + return op.Where(media_mask_3d, gathered, fallback) inputs_embeds = _scatter(image_features, text_embeds) diff --git a/src/mobius/tasks/_vision_language_3model.py b/src/mobius/tasks/_vision_language_3model.py index fa961ede4..f2d5ce650 100644 --- a/src/mobius/tasks/_vision_language_3model.py +++ b/src/mobius/tasks/_vision_language_3model.py @@ -195,7 +195,7 @@ def _build_vision( op = builder.op pixel_values = builder.input( "pixel_values", - dtype=config.dtype, + dtype=ir.DataType.FLOAT, shape=[total_patches, pixel_dim], ) image_grid_thw = builder.input( @@ -203,10 +203,11 @@ def _build_vision( dtype=ir.DataType.INT64, shape=[num_images, 3], ) + model_pixel_values = op.Cast(pixel_values, to=config.dtype) outputs = vision( op, - pixel_values=pixel_values, + pixel_values=model_pixel_values, image_grid_thw=image_grid_thw, ) diff --git a/testdata/cases/vision-language/qwen3_8-27b.yaml b/testdata/cases/vision-language/qwen3_8-27b.yaml new file mode 100644 index 000000000..cbc308c47 --- /dev/null +++ b/testdata/cases/vision-language/qwen3_8-27b.yaml @@ -0,0 +1,24 @@ +model_id: "Qwen/Qwen3.8-27B" +model_type: "qwen3_5" +revision: "1d4bf0f2ff6012fd82039f2fa52739d0dd7c60c0" +task_type: "image-text-to-text" +dtype: "bfloat16" + +inputs: + prompts: + - "Describe this image in detail." + images: + - "pipeline-cat-chonk.jpeg" + +level: "L4+L5" + +generation: + max_new_tokens: 30 + do_sample: false + +ci_skip_reason: "Official checkpoint is 55.6 GB; reduced-real CUDA and Olive evidence is run manually." +notes: >- + Qwen3.8-27B native image/video model. Dense Qwen3.5 alias with 64 hybrid + layers (48 Gated DeltaNet + 16 gated GQA), a 27-block vision encoder, and an + optional one-layer self-speculative MTP drafter. The standard target package + omits that drafter; Mobius exposes it through the separate qwen35-mtp task. diff --git a/testdata/golden/vision-language/qwen3_8-27b-reduced.json b/testdata/golden/vision-language/qwen3_8-27b-reduced.json new file mode 100644 index 000000000..9320b6a9d --- /dev/null +++ b/testdata/golden/vision-language/qwen3_8-27b-reduced.json @@ -0,0 +1,40 @@ +{ + "model_id": "Qwen/Qwen3.8-27B", + "revision": "1d4bf0f2ff6012fd82039f2fa52739d0dd7c60c0", + "fixture_schema": 1, + "input_ids": [ + 1, + 42, + 17 + ], + "top10_ids": [ + 134, + 34, + 17, + 167, + 38, + 193, + 164, + 103, + 145, + 176 + ], + "top10_logits": [ + "0x1.1704fa0000000p+0", + "0x1.e712dc0000000p-1", + "0x1.ac587e0000000p-1", + "0x1.a32cb60000000p-1", + "0x1.9978040000000p-1", + "0x1.901f0c0000000p-1", + "0x1.89eb5c0000000p-1", + "0x1.7779a00000000p-1", + "0x1.75bb280000000p-1", + "0x1.738e440000000p-1" + ], + "logits_summary": [ + "0x1.1704fa0000000p+0", + "-0x1.567c100000000p+0", + "0x1.75b0380000000p-6", + "0x1.be30360000000p-2" + ] +} \ No newline at end of file diff --git a/testdata/golden/vision-language/qwen3_8-27b-reduced_generation.json b/testdata/golden/vision-language/qwen3_8-27b-reduced_generation.json new file mode 100644 index 000000000..9290b9632 --- /dev/null +++ b/testdata/golden/vision-language/qwen3_8-27b-reduced_generation.json @@ -0,0 +1,33 @@ +{ + "model_id": "Qwen/Qwen3.8-27B", + "revision": "1d4bf0f2ff6012fd82039f2fa52739d0dd7c60c0", + "fixture_schema": 1, + "input_ids": [ + 1, + 42, + 17 + ], + "max_new_tokens": 20, + "generated_tokens": [ + 134, + 138, + 167, + 225, + 94, + 226, + 103, + 145, + 114, + 114, + 114, + 254, + 101, + 115, + 254, + 101, + 115, + 254, + 101, + 115 + ] +} \ No newline at end of file diff --git a/tests/integration_test.py b/tests/integration_test.py index 25816460d..059c85678 100644 --- a/tests/integration_test.py +++ b/tests/integration_test.py @@ -2028,6 +2028,9 @@ def test_encoder_matches_diffusers(self): # Qwen3.5 hybrid (DeltaNet + full attention) — random-weight tests # --------------------------------------------------------------------------- +_QWEN38_MODEL_ID = "Qwen/Qwen3.8-27B" +_QWEN38_REVISION = "1d4bf0f2ff6012fd82039f2fa52739d0dd7c60c0" + def _build_and_compare_qwen35(hf_model, text_config, onnx_module_cls): """Shared helper: build ONNX model, load HF weights, compare logits.""" @@ -2104,7 +2107,10 @@ def test_qwen35_prefill_logits_match(): Qwen3_5ForCausalLM, ) - c = transformers.AutoConfig.from_pretrained("Qwen/Qwen3.5-27B") + c = transformers.AutoConfig.from_pretrained( + _QWEN38_MODEL_ID, + revision=_QWEN38_REVISION, + ) tc = c.text_config tc.num_hidden_layers = 4 tc.layer_types = [ @@ -3518,14 +3524,17 @@ def make_feeds(token_id, conv_states, rec_states, kv_cache, step): # --------------------------------------------------------------------------- -def _make_tiny_qwen35_vl_config(): +def _make_tiny_qwen35_vl_config(*, keep_production_vocab: bool = False): """Create a tiny Qwen3.5-VL config for fast HF parity testing. Downloads the real Qwen3.5-27B config structure, then overrides all dimensions to be tiny. Also overrides rope_theta to float to avoid a pre-existing float64 rotary cache bug (int ** np.float32 → float64). """ - c = transformers.AutoConfig.from_pretrained("Qwen/Qwen3.5-27B") + c = transformers.AutoConfig.from_pretrained( + _QWEN38_MODEL_ID, + revision=_QWEN38_REVISION, + ) tc = c.text_config # Truncate layers: 3 DeltaNet + 1 full attention @@ -3543,7 +3552,8 @@ def _make_tiny_qwen35_vl_config(): tc.num_attention_heads = 4 tc.num_key_value_heads = 2 tc.head_dim = 16 - tc.vocab_size = 256 + if not keep_production_vocab: + tc.vocab_size = 256 tc.linear_num_value_heads = 4 tc.linear_num_key_heads = 4 tc.linear_key_head_dim = 8 @@ -3745,7 +3755,8 @@ def test_qwen35_vl_vision_features_match(): # Process real image (resized small for speed — 256 patches) processor = transformers.AutoProcessor.from_pretrained( - "Qwen/Qwen3.5-27B", + _QWEN38_MODEL_ID, + revision=_QWEN38_REVISION, ) image = Image.open("testdata/pipeline-cat-chonk.jpeg").resize( (64, 64), @@ -3815,6 +3826,181 @@ def test_qwen35_vl_vision_features_match(): assert max_diff < 0.01, f"Vision features max_diff={max_diff:.6f} (expected < 0.01)" +@pytest.mark.integration +def test_qwen38_vl_image_video_mixed_pipeline_matches_huggingface(): + """Pinned Qwen3.8 image/video processor contract matches the ONNX pipeline. + + Runs image-only, video-only, and a two-row mixed batch whose rows use + opposite media placeholder order. Hugging Face scatters image and video + feature streams independently, so the ONNX embedding input packs all image + features first and all video features second. + """ + import onnx_ir as ir + from transformers.models.qwen3_5.modeling_qwen3_5 import ( + Qwen3_5ForConditionalGeneration, + ) + + from mobius import build_from_module + from mobius._weight_loading import apply_weights + + hf_config = _make_tiny_qwen35_vl_config(keep_production_vocab=True) + arch_config = ArchitectureConfig.from_transformers( + hf_config.text_config, + parent_config=hf_config, + ) + arch_config.dtype = ir.DataType.FLOAT + onnx_module = models.Qwen35VL3ModelCausalLMModel(arch_config) + package = build_from_module( + onnx_module, + arch_config, + task="hybrid-qwen-vl", + ) + + torch.manual_seed(1) + hf_model = ( + Qwen3_5ForConditionalGeneration._from_config( + hf_config, + dtype=torch.float32, + ) + .float() + .eval() + ) + weights = onnx_module.preprocess_weights(dict(hf_model.state_dict())) + for model_name, model in package.items(): + apply_weights(model, weights) + unset = [ + name + for name, initializer in model.graph.initializers.items() + if initializer.const_value is None + ] + assert not unset, f"{model_name} has unset target parameters: {unset[:5]}" + + processor = transformers.AutoProcessor.from_pretrained( + _QWEN38_MODEL_ID, + revision=_QWEN38_REVISION, + ) + image_a = Image.open("testdata/pipeline-cat-chonk.jpeg").convert("RGB").resize((64, 64)) + image_b = image_a.transpose(Image.Transpose.FLIP_LEFT_RIGHT) + video_a = np.stack( + [np.full((64, 64, 3), value, dtype=np.uint8) for value in (16, 64, 128, 224)] + ) + video_b = np.flip(video_a, axis=0).copy() + + cases = { + "image-only": { + "text": ["<|vision_start|><|image_pad|><|vision_end|> Describe."], + "images": [image_a], + }, + "video-only": { + "text": ["<|vision_start|><|video_pad|><|vision_end|> Describe."], + "videos": [video_a], + }, + "mixed-two-row": { + "text": [ + ( + "<|vision_start|><|video_pad|><|vision_end|> Then " + "<|vision_start|><|image_pad|><|vision_end|>." + ), + ( + "<|vision_start|><|image_pad|><|vision_end|> Then " + "<|vision_start|><|video_pad|><|vision_end|>." + ), + ], + "images": [image_a, image_b], + "videos": [video_a, video_b], + }, + } + + vision_session = _make_session(package["vision_encoder"]) + embedding_session = _make_session(package["embedding"]) + decoder_session = _make_session(package["decoder"]) + try: + for case_name, processor_inputs in cases.items(): + hf_inputs = processor( + **processor_inputs, + padding=True, + return_tensors="pt", + ) + with torch.no_grad(): + hf_logits = hf_model(**hf_inputs).logits.numpy() + text_embeds = hf_model.model.language_model.embed_tokens( + hf_inputs["input_ids"] + ) + position_ids = hf_model.model.compute_3d_position_ids( + input_ids=hf_inputs["input_ids"], + inputs_embeds=text_embeds, + image_grid_thw=hf_inputs.get("image_grid_thw"), + video_grid_thw=hf_inputs.get("video_grid_thw"), + attention_mask=hf_inputs["attention_mask"], + past_key_values=None, + mm_token_type_ids=hf_inputs["mm_token_type_ids"], + ) + + media_features = [] + if "pixel_values" in hf_inputs: + media_features.append( + vision_session.run( + { + "pixel_values": hf_inputs["pixel_values"].numpy(), + "image_grid_thw": hf_inputs["image_grid_thw"].numpy(), + } + )["image_features"] + ) + if "pixel_values_videos" in hf_inputs: + media_features.append( + vision_session.run( + { + "pixel_values": hf_inputs["pixel_values_videos"].numpy(), + "image_grid_thw": hf_inputs["video_grid_thw"].numpy(), + } + )["image_features"] + ) + packed_features = np.concatenate(media_features, axis=0) + onnx_embeds = embedding_session.run( + { + "input_ids": hf_inputs["input_ids"].numpy(), + "image_features": packed_features, + } + )["inputs_embeds"] + + feeds: dict[str, np.ndarray] = { + "inputs_embeds": onnx_embeds, + "attention_mask": hf_inputs["attention_mask"].numpy(), + "position_ids": position_ids.numpy(), + } + batch_size = hf_inputs["input_ids"].shape[0] + for graph_input in package["decoder"].graph.inputs: + if graph_input.name in feeds: + continue + shape = tuple( + dim if isinstance(dim, int) else batch_size if axis == 0 else 0 + for axis, dim in enumerate(graph_input.shape) + ) + feeds[graph_input.name] = np.zeros(shape, dtype=np.float32) + + onnx_logits = decoder_session.run(feeds)["logits"] + max_abs = float(np.max(np.abs(onnx_logits - hf_logits))) + cosine = float( + np.dot(onnx_logits.ravel(), hf_logits.ravel()) + / (np.linalg.norm(onnx_logits) * np.linalg.norm(hf_logits)) + ) + print(f"Qwen3.8 {case_name}: max_abs={max_abs:.8f}, cosine={cosine:.9f}") + assert max_abs < 1e-2, case_name + assert cosine > 0.99999, case_name + assert_logits_close(onnx_logits, hf_logits, rtol=2e-2, atol=2e-2) + + attention_mask = hf_inputs["attention_mask"].numpy() + for row in range(attention_mask.shape[0]): + last_index = np.flatnonzero(attention_mask[row])[-1] + assert np.argmax(onnx_logits[row, last_index]) == np.argmax( + hf_logits[row, last_index] + ), case_name + finally: + decoder_session.close() + embedding_session.close() + vision_session.close() + + @pytest.mark.integration @pytest.mark.integration_fast def test_qwen35_deltanet_single_layer_parity(): diff --git a/tests/qwen38_real_weight_test.py b/tests/qwen38_real_weight_test.py new file mode 100644 index 000000000..dbb2d7cfa --- /dev/null +++ b/tests/qwen38_real_weight_test.py @@ -0,0 +1,229 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +"""Pinned reduced-real Qwen3.8 L4/L5 and Olive recipe coverage.""" + +from __future__ import annotations + +import importlib.util +import json +import os +import sys +from pathlib import Path + +import numpy as np +import pytest + +_ROOT = Path(__file__).parents[1] +_EXAMPLE = _ROOT / "examples" / "olive" / "qwen3_8-27b" +_L4 = _ROOT / "testdata" / "golden" / "vision-language" / "qwen3_8-27b-reduced.json" +_L5 = _ROOT / "testdata" / "golden" / "vision-language" / "qwen3_8-27b-reduced_generation.json" + + +def _load(name: str): + sys.path.insert(0, str(_EXAMPLE)) + try: + path = _EXAMPLE / f"{name}.py" + spec = importlib.util.spec_from_file_location(f"qwen38_{name}", path) + assert spec is not None and spec.loader is not None + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + finally: + sys.path.pop(0) + + +def test_reduced_config_preserves_hybrid_layers_and_remaps_media_ids(): + validator = _load("validate_reduced_checkpoint") + config = validator._reduced_hf_config() + assert config.text_config.layer_types == [ + "linear_attention", + "linear_attention", + "linear_attention", + "full_attention", + ] + assert config.text_config.mtp_num_hidden_layers == 0 + assert config.vision_config.depth == 1 + assert config.image_token_id < config.text_config.vocab_size + assert config.video_token_id < config.text_config.vocab_size + + +def test_range_reader_retries_and_requires_exact_content_range(monkeypatch): + validator = _load("validate_reduced_checkpoint") + + class Response: + def __init__(self, status, headers, content=b""): + self.status_code, self.headers, self.content = status, headers, content + + def __enter__(self): + return self + + def __exit__(self, *_args): + return False + + class Session: + def __init__(self): + self.responses = [ + Response(503, {}), + Response( + 206, {"Content-Range": "bytes 2-5/10", "Content-Length": "4"}, b"fail" + ), + Response( + 206, {"Content-Range": "bytes 0-3/10", "Content-Length": "4"}, b"pass" + ), + ] + + def get(self, *_args, **_kwargs): + return self.responses.pop(0) + + reader = object.__new__(validator._PinnedSafetensors) + reader._session = Session() + monkeypatch.setattr(validator.time, "sleep", lambda _seconds: None) + assert reader._range("shard", 0, 3) == b"pass" + + +def test_olive_recipe_is_q4_k_m_cpu_weight_only(tmp_path): + optimize = _load("optimize") + recipe = optimize.olive_config(tmp_path / "decoder.onnx", tmp_path / "out") + config = recipe["passes"]["q4_k_m"] + assert config["type"] == "OnnxKQuantQuantization" + assert config["bits"] == 4 + assert recipe["engine"]["target"]["accelerators"][0]["execution_providers"] == [ + "CPUExecutionProvider" + ] + + +def _require_real_fixture(): + if os.environ.get("MOBIUS_QWEN38_REDUCED_REAL") != "1": + pytest.skip("Set MOBIUS_QWEN38_REDUCED_REAL=1 to enable pinned range fixture") + + +@pytest.mark.integration +@pytest.mark.integration_slow +@pytest.mark.golden +def test_qwen38_reduced_real_l4(tmp_path): + _require_real_fixture() + validator = _load("validate_reduced_checkpoint") + state = validator._build_reduced_state(validator.default_reduced_cache_path()) + package_dir = tmp_path / "qwen38-test-output" + package_dir.mkdir(exist_ok=True) + source = validator._validate_variant(state, package_dir, dtype_name="f32", device="cpu") + golden = json.loads(_L4.read_text(encoding="utf-8")) + logits = validator._onnx_prefill_logits(source, golden["input_ids"], "cpu")[0, -1] + actual = np.argsort(logits)[::-1][:10].tolist() + assert actual == golden["top10_ids"] + np.testing.assert_allclose( + logits[actual], [float.fromhex(value) for value in golden["top10_logits"]], atol=2e-3 + ) + + +@pytest.mark.integration +@pytest.mark.integration_slow +@pytest.mark.generation +def test_qwen38_reduced_real_l5(tmp_path): + _require_real_fixture() + validator = _load("validate_reduced_checkpoint") + golden = json.loads(_L5.read_text(encoding="utf-8")) + state = validator._build_reduced_state(validator.default_reduced_cache_path()) + package_dir = validator._validate_variant( + state, + tmp_path / "qwen38-test-output", + dtype_name="f32", + device="cpu", + ) + generated, _logits, _profile = validator.run_token_ids( + package_dir, + golden["input_ids"], + hidden_size=256, + max_new_tokens=golden["max_new_tokens"], + device="cpu", + ) + assert generated == golden["generated_tokens"] + + +@pytest.mark.integration +@pytest.mark.integration_slow +def test_qwen38_reduced_real_bf16_export_and_reload(tmp_path): + _require_real_fixture() + import onnx_ir as ir + + validator = _load("validate_reduced_checkpoint") + state = validator._build_reduced_state(validator.default_reduced_cache_path()) + package, package_dir = validator._save_variant( + state, + tmp_path / "qwen38-test-output", + dtype_name="bf16", + device="cuda", + ) + for model in package.values(): + assert not [ + name + for name, initializer in model.graph.initializers.items() + if initializer.const_value is None + ] + assert any( + initializer.dtype == ir.DataType.BFLOAT16 + for initializer in model.graph.initializers.values() + ) + vision = ir.load(package_dir / "vision_encoder" / "model.onnx") + assert vision.graph.inputs[0].name == "pixel_values" + assert vision.graph.inputs[0].dtype == ir.DataType.FLOAT + + +@pytest.mark.integration +@pytest.mark.integration_slow +def test_qwen38_reduced_real_f16_cuda(tmp_path): + _require_real_fixture() + if os.environ.get("MOBIUS_TEST_DEVICE") != "cuda": + pytest.skip("Set MOBIUS_TEST_DEVICE=cuda for CUDA parity") + import torch + + if not torch.cuda.is_available(): + pytest.skip("CUDA is unavailable") + validator = _load("validate_reduced_checkpoint") + state = validator._build_reduced_state(validator.default_reduced_cache_path()) + validator._validate_variant( + state, + tmp_path / "qwen38-test-output", + dtype_name="f16", + device="cuda", + ) + + +@pytest.mark.integration +@pytest.mark.integration_slow +@pytest.mark.quantization +def test_qwen38_reduced_olive_q4_package(tmp_path): + _require_real_fixture() + if os.environ.get("MOBIUS_TEST_DEVICE") != "cuda": + pytest.skip("Set MOBIUS_TEST_DEVICE=cuda for Q4_K_M validation") + validator = _load("validate_reduced_checkpoint") + optimize = _load("optimize") + state = validator._build_reduced_state(validator.default_reduced_cache_path()) + output_root = tmp_path / "qwen38-test-output" + source = validator._validate_variant( + state, + output_root, + dtype_name="f16", + device="cuda", + ) + result = optimize.quantize_package(source, output_root / "q4_k_m") + import onnx_ir as ir + + decoder = ir.load(result / "decoder" / "model.onnx") + assert any( + node.domain == "com.microsoft" and node.op_type == "MatMulNBits" + for node in decoder.graph.all_nodes() + ) + assert sum(path.stat().st_size for path in result.rglob("*") if path.is_file()) < sum( + path.stat().st_size for path in source.rglob("*") if path.is_file() + ) + ids, logits, _ = validator.run_token_ids( + result, + [1, 42, 17], + hidden_size=256, + max_new_tokens=20, + device="cuda", + ) + assert len(ids) == 20 + assert all(np.isfinite(value).all() for value in logits) From 518d3d0039db1cc58b9a00dffea6658e9a29de20 Mon Sep 17 00:00:00 2001 From: Justin Chu Date: Fri, 14 Aug 2026 19:35:19 -0700 Subject: [PATCH 02/15] Stabilize Qwen3.8 reduced-real validation Isolate Olive caches so recurrent-gate exclusions are always honored, use a portable CUDA vision graph around the ORT 1.26 PackedMHA defect, and validate the final Q4 package with CUDA reload plus deterministic cached CPU generation. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Signed-off-by: Justin Chu --- examples/olive/qwen3_8-27b/README.md | 11 ++++++- examples/olive/qwen3_8-27b/optimize.py | 7 ++++ .../validate_reduced_checkpoint.py | 33 +++++++++++++++++-- tests/qwen38_real_weight_test.py | 29 ++++++++++++++-- 4 files changed, 74 insertions(+), 6 deletions(-) diff --git a/examples/olive/qwen3_8-27b/README.md b/examples/olive/qwen3_8-27b/README.md index 2e8481583..ce689fa01 100644 --- a/examples/olive/qwen3_8-27b/README.md +++ b/examples/olive/qwen3_8-27b/README.md @@ -19,10 +19,19 @@ not an ORT GenAI capability-gated test. The Q4 recipe keeps DeltaNet's narrow `in_proj_a` decay and `in_proj_b` time-step gates in FP16. Quantizing those recurrent controls destabilizes cached generation; all larger decoder matrices remain eligible for -`MatMulNBits`. +`MatMulNBits`. All three Q4 package components are reloaded with CUDA enabled; +the 20-token semantic run uses CPU because ORT 1.26's CUDA `MatMulNBits` +execution is itself nondeterministic for this reduced hybrid fixture. BF16 export and package reload are valid Mobius outputs. The CUDA-12-compatible ORT 1.26 wheel used for this reduced-real run cannot initialize the BF16 hybrid graph (`CausalConvWithState`/`Softplus` provider placement), while newer ORT-GPU wheels available in the test environment require CUDA 13. This is a downstream runtime waiver, not an export or support gate. + +The same old ORT wheel is nondeterministic for dynamic Qwen vision batches +through its `PackedMultiHeadAttention` CUDA kernel. The validator therefore +uses the portable standard-attention vision graph for stable real CUDA +image/video/mixed execution while retaining the CUDA-optimized decoder parity +and provider profile. Mobius still emits both graph variants without deciding +which downstream runtime version can load them. diff --git a/examples/olive/qwen3_8-27b/optimize.py b/examples/olive/qwen3_8-27b/optimize.py index 477d0120e..e4baabd21 100644 --- a/examples/olive/qwen3_8-27b/optimize.py +++ b/examples/olive/qwen3_8-27b/optimize.py @@ -55,6 +55,10 @@ def olive_config( }, "no_artifacts": True, "output_dir": str(output), + # A global Olive cache can return a decoder produced with a different + # exclusion set. Scope and clean it so recurrent-gate policy is exact. + "cache_dir": str(output.parent / ".olive-cache"), + "clean_cache": True, } @@ -109,6 +113,9 @@ def quantize_package(source_dir: str | Path, output_dir: str | Path) -> Path: if produced != decoder_dir / "model.onnx": produced.replace(decoder_dir / "model.onnx") shutil.rmtree(olive_output) + olive_cache = output / ".olive-cache" + if olive_cache.exists(): + shutil.rmtree(olive_cache) for name in ("embedding", "vision_encoder"): shutil.copytree(source / name, output / name) for name in _ASSETS: diff --git a/examples/olive/qwen3_8-27b/validate_reduced_checkpoint.py b/examples/olive/qwen3_8-27b/validate_reduced_checkpoint.py index a82058af5..5a1fab042 100644 --- a/examples/olive/qwen3_8-27b/validate_reduced_checkpoint.py +++ b/examples/olive/qwen3_8-27b/validate_reduced_checkpoint.py @@ -463,6 +463,25 @@ def _media_smoke(package_dir: Path, device: str) -> dict[str, tuple[int, ...]]: return results +def _cuda_standard_vision_smoke( + state: dict[str, torch.Tensor], + output_root: Path, + *, + dtype_name: str, +) -> dict[str, tuple[int, ...]]: + """Run media through the standard-attention vision graph on CUDA. + + ORT 1.26's CUDA PackedMultiHeadAttention kernel is nondeterministic for + dynamic packed vision batches. The portable graph still places vision + compute on CUDA and provides stable image/video/mixed runtime evidence. + """ + package = _mobius_package(state, dtype_name=dtype_name, ep="cpu") + package_dir = output_root / f"{dtype_name}-cuda-standard-vision" + package_dir.mkdir(parents=True, exist_ok=True) + package.save(package_dir, external_data="onnx") + return _media_smoke(package_dir, "cuda") + + def _save_variant( state: dict[str, torch.Tensor], output_root: Path, @@ -534,7 +553,15 @@ def _validate_variant( print(f"{dtype_name}/{device} provider placement: {placement}") print(f"{dtype_name}/{device} generated IDs: {generated}") print(f"{dtype_name}/{device} graph audit: {_graph_audit(package)}") - print(f"{dtype_name}/{device} media shapes: {_media_smoke(package_dir, device)}") + if device == "cuda": + media_shapes = _cuda_standard_vision_smoke( + state, + output_root, + dtype_name=dtype_name, + ) + print(f"{dtype_name}/{device} standard-vision media shapes: {media_shapes}") + else: + print(f"{dtype_name}/{device} media shapes: {_media_smoke(package_dir, device)}") return package_dir @@ -660,7 +687,7 @@ def main() -> None: [1, 42, 17], hidden_size=_HIDDEN_SIZE, max_new_tokens=_GENERATION_TOKENS, - device="cuda", + device="cpu", ) if len(ids) != _GENERATION_TOKENS or not all( np.isfinite(logit).all() for logit in logits @@ -670,7 +697,7 @@ def main() -> None: "Q4_K_M package audit: " f"MatMulNBits={matmul_nbits}, source={source_bytes}, quantized={quantized_bytes}" ) - print(f"Q4_K_M direct-session IDs: {ids}") + print(f"Q4_K_M direct CPU-session IDs: {ids}") if __name__ == "__main__": diff --git a/tests/qwen38_real_weight_test.py b/tests/qwen38_real_weight_test.py index dbb2d7cfa..56a8ac39b 100644 --- a/tests/qwen38_real_weight_test.py +++ b/tests/qwen38_real_weight_test.py @@ -88,6 +88,8 @@ def test_olive_recipe_is_q4_k_m_cpu_weight_only(tmp_path): config = recipe["passes"]["q4_k_m"] assert config["type"] == "OnnxKQuantQuantization" assert config["bits"] == 4 + assert recipe["clean_cache"] is True + assert Path(recipe["cache_dir"]).name == ".olive-cache" assert recipe["engine"]["target"]["accelerators"][0]["execution_providers"] == [ "CPUExecutionProvider" ] @@ -218,12 +220,35 @@ def test_qwen38_reduced_olive_q4_package(tmp_path): assert sum(path.stat().st_size for path in result.rglob("*") if path.is_file()) < sum( path.stat().st_size for path in source.rglob("*") if path.is_file() ) + for name in ("decoder", "embedding", "vision_encoder"): + validator._create_session(result / name / "model.onnx", "cuda") ids, logits, _ = validator.run_token_ids( result, [1, 42, 17], hidden_size=256, max_new_tokens=20, - device="cuda", + device="cpu", ) - assert len(ids) == 20 + assert ids == [ + 134, + 244, + 242, + 167, + 81, + 34, + 155, + 251, + 142, + 90, + 224, + 31, + 76, + 250, + 240, + 120, + 134, + 120, + 134, + 120, + ] assert all(np.isfinite(value).all() for value in logits) From 4656be11d724b6fd5a0360fe869f164887be1760 Mon Sep 17 00:00:00 2001 From: Justin Chu Date: Fri, 14 Aug 2026 20:16:21 -0700 Subject: [PATCH 03/15] Preserve Qwen3.8 package metadata Emit the faithful qwen3_5 GenAI model type without consulting downstream runtime support, retain processor metadata in Q4 packages, validate local processor reload, and remove the static Olive recipe that could not encode required graph-derived exclusions. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Signed-off-by: Justin Chu --- examples/olive/qwen3_8-27b/README.md | 3 + examples/olive/qwen3_8-27b/olive_q4_k_m.json | 30 ---------- examples/olive/qwen3_8-27b/optimize.py | 2 + .../integrations/ort_genai/auto_export.py | 4 +- .../ort_genai/auto_export_test.py | 58 +++++++++++++++++++ tests/qwen38_real_weight_test.py | 7 +++ 6 files changed, 72 insertions(+), 32 deletions(-) delete mode 100644 examples/olive/qwen3_8-27b/olive_q4_k_m.json diff --git a/examples/olive/qwen3_8-27b/README.md b/examples/olive/qwen3_8-27b/README.md index ce689fa01..321463a3b 100644 --- a/examples/olive/qwen3_8-27b/README.md +++ b/examples/olive/qwen3_8-27b/README.md @@ -22,6 +22,9 @@ cached generation; all larger decoder matrices remain eligible for `MatMulNBits`. All three Q4 package components are reloaded with CUDA enabled; the 20-token semantic run uses CPU because ORT 1.26's CUDA `MatMulNBits` execution is itself nondeterministic for this reduced hybrid fixture. +`optimize.py` derives the recurrent-gate node exclusions from the exported +decoder, so the validated package must be assembled through that script rather +than a static Olive JSON recipe. BF16 export and package reload are valid Mobius outputs. The CUDA-12-compatible ORT 1.26 wheel used for this reduced-real run cannot initialize the BF16 hybrid diff --git a/examples/olive/qwen3_8-27b/olive_q4_k_m.json b/examples/olive/qwen3_8-27b/olive_q4_k_m.json deleted file mode 100644 index d08b9b14a..000000000 --- a/examples/olive/qwen3_8-27b/olive_q4_k_m.json +++ /dev/null @@ -1,30 +0,0 @@ -{ - "input_model": { - "type": "OnnxModel", - "model_path": "reduced/f16-cuda/decoder.onnx" - }, - "passes": { - "q4_k_m": { - "type": "OnnxKQuantQuantization", - "bits": 4, - "block_size": 32, - "save_as_external_data": true, - "all_tensors_to_one_file": true, - "external_data_name": "decoder.onnx.data", - "size_threshold": 1024 - } - }, - "engine": { - "target": { - "type": "LocalSystem", - "accelerators": [ - { - "device": "cpu", - "execution_providers": [ - "CPUExecutionProvider" - ] - } - ] - } - } -} diff --git a/examples/olive/qwen3_8-27b/optimize.py b/examples/olive/qwen3_8-27b/optimize.py index e4baabd21..e29385acf 100644 --- a/examples/olive/qwen3_8-27b/optimize.py +++ b/examples/olive/qwen3_8-27b/optimize.py @@ -21,6 +21,8 @@ "special_tokens_map.json", "chat_template.jinja", "preprocessor_config.json", + "processor_config.json", + "image_processor.json", } diff --git a/src/mobius/integrations/ort_genai/auto_export.py b/src/mobius/integrations/ort_genai/auto_export.py index b281517a4..ebd335cdc 100644 --- a/src/mobius/integrations/ort_genai/auto_export.py +++ b/src/mobius/integrations/ort_genai/auto_export.py @@ -92,8 +92,8 @@ "qwen2_vl": "qwen2_5_vl", "qwen3_vl": "qwen3_vl", "qwen3_vl_text": "qwen3_vl", - "qwen3_5": "qwen2_5_vl", - "qwen3_5_vl": "qwen2_5_vl", + "qwen3_5": "qwen3_5", + "qwen3_5_vl": "qwen3_5", # MiniCPM uses standard 1D decoder position IDs (unlike Qwen-VL MRoPE). # The phi3v multimodal runtime provides that contract; callers supply # HF-preprocessed packed pixels through Generator.set_inputs(). diff --git a/src/mobius/integrations/ort_genai/auto_export_test.py b/src/mobius/integrations/ort_genai/auto_export_test.py index 1731572b3..1e8ff142d 100644 --- a/src/mobius/integrations/ort_genai/auto_export_test.py +++ b/src/mobius/integrations/ort_genai/auto_export_test.py @@ -82,6 +82,8 @@ class FakeConfig: class TestResolveOrtGenaiModelType: def test_known_model_type(self): assert _resolve_ort_genai_model_type("qwen3") == "qwen2" + assert _resolve_ort_genai_model_type("qwen3_5") == "qwen3_5" + assert _resolve_ort_genai_model_type("qwen3_5_vl") == "qwen3_5" assert _resolve_ort_genai_model_type("gemma2") == "gemma" assert _resolve_ort_genai_model_type("llama") == "llama" @@ -1085,6 +1087,62 @@ class FakeConfig: "present_conv_names": "present.%d.conv_state", } + def test_qwen35_vl_hybrid_metadata_is_emitted_without_runtime_gate(self, tmp_path): + import dataclasses + + from mobius._model_package import ModelPackage + + @dataclasses.dataclass + class FakeConfig: + model_type: str = "qwen3_5" + vocab_size: int = 256 + hidden_size: int = 64 + num_hidden_layers: int = 4 + num_attention_heads: int = 4 + num_key_value_heads: int = 2 + head_dim: int = 16 + max_position_embeddings: int = 128 + pad_token_id: int = 0 + + package = ModelPackage( + { + "decoder": _mock_model( + inputs=[ + "inputs_embeds", + "attention_mask", + "position_ids", + "past_key_values.0.conv_state", + "past_key_values.0.recurrent_state", + "past_key_values.3.key", + "past_key_values.3.value", + ], + outputs=[ + "logits", + "present.0.conv_state", + "present.0.recurrent_state", + "present.3.key", + "present.3.value", + ], + ), + "embedding": _mock_model(inputs=["input_ids", "image_features"]), + "vision_encoder": _mock_model( + inputs=["pixel_values", "image_grid_thw"], + outputs=["image_features"], + ), + }, + config=FakeConfig(), + ) + + result = write_ort_genai_config(package, str(tmp_path)) + + with open(result["genai_config"], encoding="utf-8") as config_file: + generated = json.load(config_file) + decoder = generated["model"]["decoder"] + assert generated["model"]["type"] == "qwen3_5" + assert decoder["num_hidden_layers"] == 4 + assert decoder["inputs"]["past_key_names"] == "past_key_values.%d.key" + assert decoder["inputs"]["past_conv_names"] == "past_key_values.%d.conv_state" + def test_olive_renamed_logits_output_is_emitted(self, tmp_path): pkg = _make_fake_llm_pkg("qwen2") pkg["model"] = _mock_model( diff --git a/tests/qwen38_real_weight_test.py b/tests/qwen38_real_weight_test.py index 56a8ac39b..6f9731b6a 100644 --- a/tests/qwen38_real_weight_test.py +++ b/tests/qwen38_real_weight_test.py @@ -209,6 +209,8 @@ def test_qwen38_reduced_olive_q4_package(tmp_path): dtype_name="f16", device="cuda", ) + processor_config = source / "processor_config.json" + assert processor_config.is_file() result = optimize.quantize_package(source, output_root / "q4_k_m") import onnx_ir as ir @@ -220,6 +222,11 @@ def test_qwen38_reduced_olive_q4_package(tmp_path): assert sum(path.stat().st_size for path in result.rglob("*") if path.is_file()) < sum( path.stat().st_size for path in source.rglob("*") if path.is_file() ) + assert (result / "processor_config.json").read_bytes() == processor_config.read_bytes() + from transformers import AutoProcessor + + processor = AutoProcessor.from_pretrained(result, local_files_only=True) + assert type(processor).__name__ == "Qwen3VLProcessor" for name in ("decoder", "embedding", "vision_encoder"): validator._create_session(result / name / "model.onnx", "cuda") ids, logits, _ = validator.run_token_ids( From c536901b662101b624440dd8a2103e19a6edcd96 Mon Sep 17 00:00:00 2001 From: Justin Chu Date: Fri, 14 Aug 2026 20:56:41 -0700 Subject: [PATCH 04/15] Clarify Q4 CUDA runtime waiver Document that repeated identical ORT 1.26 CUDA MatMulNBits runs on the exact assembled package can diverge or produce non-finite logits, while CPU cached generation remains the semantic acceptance path. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Signed-off-by: Justin Chu --- examples/olive/qwen3_8-27b/README.md | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/examples/olive/qwen3_8-27b/README.md b/examples/olive/qwen3_8-27b/README.md index 321463a3b..35abea7a3 100644 --- a/examples/olive/qwen3_8-27b/README.md +++ b/examples/olive/qwen3_8-27b/README.md @@ -20,8 +20,8 @@ The Q4 recipe keeps DeltaNet's narrow `in_proj_a` decay and `in_proj_b` time-step gates in FP16. Quantizing those recurrent controls destabilizes cached generation; all larger decoder matrices remain eligible for `MatMulNBits`. All three Q4 package components are reloaded with CUDA enabled; -the 20-token semantic run uses CPU because ORT 1.26's CUDA `MatMulNBits` -execution is itself nondeterministic for this reduced hybrid fixture. +the 20-token semantic run uses CPU because repeated identical ORT 1.26 CUDA +`MatMulNBits` runs can diverge and produce non-finite logits for this fixture. `optimize.py` derives the recurrent-gate node exclusions from the exported decoder, so the validated package must be assembled through that script rather than a static Olive JSON recipe. From 0e6c83ec66a0c543888601573b5b33efbfaab2bd Mon Sep 17 00:00:00 2001 From: Justin Chu Date: Fri, 14 Aug 2026 21:14:47 -0700 Subject: [PATCH 05/15] Normalize Qwen3.5 VL metadata Map the extracted qwen3_5_text subtype back to faithful qwen3_5 multimodal metadata and select the packed Qwen PatchImage processor pipeline for normal config and local exports. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Signed-off-by: Justin Chu --- .../integrations/ort_genai/auto_export.py | 2 ++ .../ort_genai/auto_export_test.py | 35 ++++++++++++++++++- 2 files changed, 36 insertions(+), 1 deletion(-) diff --git a/src/mobius/integrations/ort_genai/auto_export.py b/src/mobius/integrations/ort_genai/auto_export.py index ebd335cdc..2674e4252 100644 --- a/src/mobius/integrations/ort_genai/auto_export.py +++ b/src/mobius/integrations/ort_genai/auto_export.py @@ -94,6 +94,7 @@ "qwen3_vl_text": "qwen3_vl", "qwen3_5": "qwen3_5", "qwen3_5_vl": "qwen3_5", + "qwen3_5_text": "qwen3_5", # MiniCPM uses standard 1D decoder position IDs (unlike Qwen-VL MRoPE). # The phi3v multimodal runtime provides that contract; callers supply # HF-preprocessed packed pixels through Generator.set_inputs(). @@ -136,6 +137,7 @@ "qwen3_vl_text", "qwen3_5", "qwen3_5_vl", + "qwen3_5_text", "qwen3_5_moe", "videochat_flash_qwen", } diff --git a/src/mobius/integrations/ort_genai/auto_export_test.py b/src/mobius/integrations/ort_genai/auto_export_test.py index 1e8ff142d..e4246f3a7 100644 --- a/src/mobius/integrations/ort_genai/auto_export_test.py +++ b/src/mobius/integrations/ort_genai/auto_export_test.py @@ -84,6 +84,7 @@ def test_known_model_type(self): assert _resolve_ort_genai_model_type("qwen3") == "qwen2" assert _resolve_ort_genai_model_type("qwen3_5") == "qwen3_5" assert _resolve_ort_genai_model_type("qwen3_5_vl") == "qwen3_5" + assert _resolve_ort_genai_model_type("qwen3_5_text") == "qwen3_5" assert _resolve_ort_genai_model_type("gemma2") == "gemma" assert _resolve_ort_genai_model_type("llama") == "llama" @@ -356,6 +357,38 @@ def test_muse_glimmer_uses_packed_qwen_image_pipeline(self, tmp_path): "merge_size": 2, } + def test_qwen35_text_subtype_uses_packed_qwen_image_pipeline(self, tmp_path): + vision = types.SimpleNamespace( + image_size=448, + patch_size=16, + spatial_merge_size=2, + model_type="qwen3_5", + ) + config = types.SimpleNamespace( + vision=vision, + model_type="qwen3_5_text", + spatial_merge_size=2, + temporal_patch_size=2, + ) + + path = _write_vision_processor_config(config, str(tmp_path)) + + assert path is not None + with open(path, encoding="utf-8") as config_file: + processor = json.load(config_file)["processor"] + assert processor["name"] == "qwen2_5_image_processor" + transforms = processor["transforms"] + assert transforms[-1]["operation"] == { + "name": "patch_image", + "type": "PatchImage", + "attrs": { + "patch_size": 16, + "temporal_patch_size": 2, + "merge_size": 2, + }, + } + assert transforms[-2]["operation"]["attrs"]["qwen2_5_vl"] == 1 + def test_gemma3_vision_config(self, tmp_path): """Gemma3 gets a fixed-size resize + Permute3D (not the generic branch). @@ -1094,7 +1127,7 @@ def test_qwen35_vl_hybrid_metadata_is_emitted_without_runtime_gate(self, tmp_pat @dataclasses.dataclass class FakeConfig: - model_type: str = "qwen3_5" + model_type: str = "qwen3_5_text" vocab_size: int = 256 hidden_size: int = 64 num_hidden_layers: int = 4 From 370f48c9349a1408732120e7108911d85129fbed Mon Sep 17 00:00:00 2001 From: Justin Chu Date: Fri, 14 Aug 2026 21:38:59 -0700 Subject: [PATCH 06/15] Make reduced Qwen package token-ID-only Omit the incompatible production tokenizer from the 256-token reduced fixture while retaining and independently reloading the pinned image and video processor metadata. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Signed-off-by: Justin Chu --- examples/olive/qwen3_8-27b/README.md | 3 +++ .../qwen3_8-27b/validate_reduced_checkpoint.py | 12 ++++++++---- tests/qwen38_real_weight_test.py | 14 ++++++++++---- 3 files changed, 21 insertions(+), 8 deletions(-) diff --git a/examples/olive/qwen3_8-27b/README.md b/examples/olive/qwen3_8-27b/README.md index 35abea7a3..719195d0b 100644 --- a/examples/olive/qwen3_8-27b/README.md +++ b/examples/olive/qwen3_8-27b/README.md @@ -6,6 +6,9 @@ The fixture has 4 decoder layers (three DeltaNet and one full attention), one vision block, remapped image/video IDs, and no MTP tensors because the standard target forward does not consume the optional self-speculative drafter. Mobius exposes that drafter through the separate `qwen35-mtp` package contract. +The reduced package is token-ID-only and deliberately omits the production +tokenizer because its vocabulary is remapped to 256 entries. It retains the +pinned image and video processor metadata for media-contract validation. ```powershell python examples/olive/qwen3_8-27b/validate_reduced_checkpoint.py --matrix f32-cpu f16-cuda diff --git a/examples/olive/qwen3_8-27b/validate_reduced_checkpoint.py b/examples/olive/qwen3_8-27b/validate_reduced_checkpoint.py index 5a1fab042..89e95def6 100644 --- a/examples/olive/qwen3_8-27b/validate_reduced_checkpoint.py +++ b/examples/olive/qwen3_8-27b/validate_reduced_checkpoint.py @@ -406,14 +406,17 @@ def _assert_logits_close( def _save_package_assets(package_dir: Path) -> None: - """Copy pinned tokenizer/processor metadata needed by a standalone package.""" - from transformers import AutoProcessor, AutoTokenizer, GenerationConfig + """Copy pinned processor metadata for the reduced token-ID-only package.""" + from transformers import AutoProcessor, GenerationConfig _reduced_hf_config().save_pretrained(package_dir) - AutoTokenizer.from_pretrained(MODEL_ID, revision=REVISION).save_pretrained(package_dir) GenerationConfig.from_pretrained(MODEL_ID, revision=REVISION).save_pretrained(package_dir) try: - AutoProcessor.from_pretrained(MODEL_ID, revision=REVISION).save_pretrained(package_dir) + processor = AutoProcessor.from_pretrained(MODEL_ID, revision=REVISION) + (package_dir / "processor_config.json").write_text( + json.dumps(processor.to_dict(), indent=2), + encoding="utf-8", + ) except (ImportError, OSError, ValueError) as error: # The ONNX package remains directly runnable without the optional # processor serialization; retain the precise reason in its manifest. @@ -425,6 +428,7 @@ def _save_package_assets(package_dir: Path) -> None: "revision": REVISION, "fixture_schema": FIXTURE_SCHEMA_VERSION, "runtime": "onnxruntime-direct", + "text_input_contract": "token-ids-only", "components": [ "decoder/model.onnx", "embedding/model.onnx", diff --git a/tests/qwen38_real_weight_test.py b/tests/qwen38_real_weight_test.py index 6f9731b6a..7d97ddcfb 100644 --- a/tests/qwen38_real_weight_test.py +++ b/tests/qwen38_real_weight_test.py @@ -223,10 +223,16 @@ def test_qwen38_reduced_olive_q4_package(tmp_path): path.stat().st_size for path in source.rglob("*") if path.is_file() ) assert (result / "processor_config.json").read_bytes() == processor_config.read_bytes() - from transformers import AutoProcessor - - processor = AutoProcessor.from_pretrained(result, local_files_only=True) - assert type(processor).__name__ == "Qwen3VLProcessor" + from transformers import Qwen2VLImageProcessor, Qwen3VLVideoProcessor + + processor_data = json.loads((result / "processor_config.json").read_text()) + assert type( + Qwen2VLImageProcessor.from_dict(processor_data["image_processor"]) + ).__name__ == ("Qwen2VLImageProcessor") + assert type( + Qwen3VLVideoProcessor.from_dict(processor_data["video_processor"]) + ).__name__ == ("Qwen3VLVideoProcessor") + assert not (result / "tokenizer.json").exists() for name in ("decoder", "embedding", "vision_encoder"): validator._create_session(result / name / "model.onnx", "cuda") ids, logits, _ = validator.run_token_ids( From 173156f17f8b813829abfdcb8eedb458db9b8c8e Mon Sep 17 00:00:00 2001 From: Justin Chu Date: Fri, 14 Aug 2026 21:52:33 -0700 Subject: [PATCH 07/15] Harden reduced package metadata Omit out-of-range production generation IDs from the token-ID-only fixture and require empty variant directories so stale tokenizer assets cannot survive reruns. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Signed-off-by: Justin Chu --- examples/olive/qwen3_8-27b/README.md | 6 ++++-- examples/olive/qwen3_8-27b/validate_reduced_checkpoint.py | 5 +++-- tests/qwen38_real_weight_test.py | 1 + 3 files changed, 8 insertions(+), 4 deletions(-) diff --git a/examples/olive/qwen3_8-27b/README.md b/examples/olive/qwen3_8-27b/README.md index 719195d0b..97490ffd7 100644 --- a/examples/olive/qwen3_8-27b/README.md +++ b/examples/olive/qwen3_8-27b/README.md @@ -7,8 +7,10 @@ one vision block, remapped image/video IDs, and no MTP tensors because the standard target forward does not consume the optional self-speculative drafter. Mobius exposes that drafter through the separate `qwen35-mtp` package contract. The reduced package is token-ID-only and deliberately omits the production -tokenizer because its vocabulary is remapped to 256 entries. It retains the -pinned image and video processor metadata for media-contract validation. +tokenizer and generation config because its vocabulary and special-token IDs +are remapped to 256 entries. It retains the pinned image and video processor +metadata for media-contract validation. Each variant output directory must be +empty so stale assets from an older recipe cannot enter the assembled package. ```powershell python examples/olive/qwen3_8-27b/validate_reduced_checkpoint.py --matrix f32-cpu f16-cuda diff --git a/examples/olive/qwen3_8-27b/validate_reduced_checkpoint.py b/examples/olive/qwen3_8-27b/validate_reduced_checkpoint.py index 89e95def6..93b63f62f 100644 --- a/examples/olive/qwen3_8-27b/validate_reduced_checkpoint.py +++ b/examples/olive/qwen3_8-27b/validate_reduced_checkpoint.py @@ -407,10 +407,9 @@ def _assert_logits_close( def _save_package_assets(package_dir: Path) -> None: """Copy pinned processor metadata for the reduced token-ID-only package.""" - from transformers import AutoProcessor, GenerationConfig + from transformers import AutoProcessor _reduced_hf_config().save_pretrained(package_dir) - GenerationConfig.from_pretrained(MODEL_ID, revision=REVISION).save_pretrained(package_dir) try: processor = AutoProcessor.from_pretrained(MODEL_ID, revision=REVISION) (package_dir / "processor_config.json").write_text( @@ -497,6 +496,8 @@ def _save_variant( state, dtype_name=dtype_name, ep="cuda" if device == "cuda" else "cpu" ) package_dir = output_root / f"{dtype_name}-{device}" + if package_dir.exists() and any(package_dir.iterdir()): + raise FileExistsError(f"Variant output must be empty: {package_dir}") package_dir.mkdir(parents=True, exist_ok=True) package.save(package_dir, external_data="onnx") _save_package_assets(package_dir) diff --git a/tests/qwen38_real_weight_test.py b/tests/qwen38_real_weight_test.py index 7d97ddcfb..c685e0ad9 100644 --- a/tests/qwen38_real_weight_test.py +++ b/tests/qwen38_real_weight_test.py @@ -233,6 +233,7 @@ def test_qwen38_reduced_olive_q4_package(tmp_path): Qwen3VLVideoProcessor.from_dict(processor_data["video_processor"]) ).__name__ == ("Qwen3VLVideoProcessor") assert not (result / "tokenizer.json").exists() + assert not (result / "generation_config.json").exists() for name in ("decoder", "embedding", "vision_encoder"): validator._create_session(result / name / "model.onnx", "cuda") ids, logits, _ = validator.run_token_ids( From b59fb43d081c594a429c33292d5441f92d12b1d6 Mon Sep 17 00:00:00 2001 From: Justin Chu Date: Fri, 14 Aug 2026 22:02:10 -0700 Subject: [PATCH 08/15] Preserve reduced Q4 package contract Merge source provenance and the token-ID-only input contract into the Q4 manifest, and carry processor waivers through package assembly. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Signed-off-by: Justin Chu --- examples/olive/qwen3_8-27b/optimize.py | 33 ++++++++++++++++---------- tests/qwen38_real_weight_test.py | 4 ++++ 2 files changed, 24 insertions(+), 13 deletions(-) diff --git a/examples/olive/qwen3_8-27b/optimize.py b/examples/olive/qwen3_8-27b/optimize.py index e29385acf..b1a31da67 100644 --- a/examples/olive/qwen3_8-27b/optimize.py +++ b/examples/olive/qwen3_8-27b/optimize.py @@ -23,6 +23,7 @@ "preprocessor_config.json", "processor_config.json", "image_processor.json", + "processor-waiver.txt", } @@ -123,19 +124,25 @@ def quantize_package(source_dir: str | Path, output_dir: str | Path) -> Path: for name in _ASSETS: if (source / name).is_file(): shutil.copy2(source / name, output / name) - manifest = { - "model_id": MODEL_ID, - "revision": REVISION, - "quantization": "Q4_K_M", - "quantized_component": "decoder", - "olive_provider": "CPUExecutionProvider", - "preserved_fp16_recurrent_gate_nodes": len(preserved_fp16_nodes), - "components": [ - "decoder/model.onnx", - "embedding/model.onnx", - "vision_encoder/model.onnx", - ], - } + source_manifest = source / "source_manifest.json" + manifest = ( + json.loads(source_manifest.read_text(encoding="utf-8")) + if source_manifest.is_file() + else {"model_id": MODEL_ID, "revision": REVISION} + ) + manifest.update( + { + "quantization": "Q4_K_M", + "quantized_component": "decoder", + "olive_provider": "CPUExecutionProvider", + "preserved_fp16_recurrent_gate_nodes": len(preserved_fp16_nodes), + "components": [ + "decoder/model.onnx", + "embedding/model.onnx", + "vision_encoder/model.onnx", + ], + } + ) (output / "source_manifest.json").write_text( json.dumps(manifest, indent=2), encoding="utf-8" ) diff --git a/tests/qwen38_real_weight_test.py b/tests/qwen38_real_weight_test.py index c685e0ad9..94629bfcf 100644 --- a/tests/qwen38_real_weight_test.py +++ b/tests/qwen38_real_weight_test.py @@ -234,6 +234,10 @@ def test_qwen38_reduced_olive_q4_package(tmp_path): ).__name__ == ("Qwen3VLVideoProcessor") assert not (result / "tokenizer.json").exists() assert not (result / "generation_config.json").exists() + manifest = json.loads((result / "source_manifest.json").read_text()) + assert manifest["text_input_contract"] == "token-ids-only" + assert manifest["runtime"] == "onnxruntime-direct" + assert manifest["quantization"] == "Q4_K_M" for name in ("decoder", "embedding", "vision_encoder"): validator._create_session(result / name / "model.onnx", "cuda") ids, logits, _ = validator.run_token_ids( From 654c304a886ac8482fda84d582ad986962442d37 Mon Sep 17 00:00:00 2001 From: Justin Chu Date: Fri, 14 Aug 2026 22:21:46 -0700 Subject: [PATCH 09/15] Isolate Qwen reduced test modules Load the example's sibling inference module under a Qwen-specific name so prior Nemotron tests cannot poison Python's module cache in combined CI runs. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Signed-off-by: Justin Chu --- tests/qwen38_real_weight_test.py | 35 ++++++++++++++++++++++++++------ 1 file changed, 29 insertions(+), 6 deletions(-) diff --git a/tests/qwen38_real_weight_test.py b/tests/qwen38_real_weight_test.py index 94629bfcf..8d5cedb46 100644 --- a/tests/qwen38_real_weight_test.py +++ b/tests/qwen38_real_weight_test.py @@ -9,6 +9,7 @@ import json import os import sys +import types from pathlib import Path import numpy as np @@ -21,16 +22,38 @@ def _load(name: str): - sys.path.insert(0, str(_EXAMPLE)) - try: - path = _EXAMPLE / f"{name}.py" - spec = importlib.util.spec_from_file_location(f"qwen38_{name}", path) + def load_file(module_name: str, path: Path): + spec = importlib.util.spec_from_file_location(module_name, path) assert spec is not None and spec.loader is not None module = importlib.util.module_from_spec(spec) - spec.loader.exec_module(module) + sys.modules[module_name] = module + try: + spec.loader.exec_module(module) + except Exception: + sys.modules.pop(module_name, None) + raise return module + + previous_inference = sys.modules.get("inference") + try: + inference = load_file("qwen38_inference", _EXAMPLE / "inference.py") + sys.modules["inference"] = inference + return load_file(f"qwen38_{name}", _EXAMPLE / f"{name}.py") finally: - sys.path.pop(0) + if previous_inference is None: + sys.modules.pop("inference", None) + else: + sys.modules["inference"] = previous_inference + + +def test_loader_isolates_sibling_inference_modules(monkeypatch): + foreign_inference = types.ModuleType("inference") + monkeypatch.setitem(sys.modules, "inference", foreign_inference) + + validator = _load("validate_reduced_checkpoint") + + assert validator.MODEL_ID == "Qwen/Qwen3.8-27B" + assert sys.modules["inference"] is foreign_inference def test_reduced_config_preserves_hybrid_layers_and_remaps_media_ids(): From 042159ae65a7b8c3abd4a74b77a7e89317dd0815 Mon Sep 17 00:00:00 2001 From: Justin Chu Date: Fri, 14 Aug 2026 22:25:02 -0700 Subject: [PATCH 10/15] Fix Qwen VL embedding mask construction Construct the absent-video mask entirely in the ONNX graph and preserve the original boolean image mask for media feature selection. Add runtime coverage for image-only embedding scatter when video_token_id is unset. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Signed-off-by: Justin Chu --- src/mobius/models/qwen35_test.py | 42 ++++++++++++++++++++++++++++++++ src/mobius/models/qwen_vl.py | 7 +++--- 2 files changed, 46 insertions(+), 3 deletions(-) diff --git a/src/mobius/models/qwen35_test.py b/src/mobius/models/qwen35_test.py index 798523641..9f5c897ec 100644 --- a/src/mobius/models/qwen35_test.py +++ b/src/mobius/models/qwen35_test.py @@ -249,6 +249,48 @@ def test_embedding_scatter_matches_separate_image_then_video_streams(self): session.close() np.testing.assert_array_equal(decode, embedding_weight[decode_ids]) + def test_embedding_scatter_without_video_token_id(self): + config = ArchitectureConfig( + vocab_size=16, + hidden_size=4, + pad_token_id=0, + image_token_id=10, + video_token_id=None, + dtype=ir.DataType.FLOAT, + ) + graph = build_embedding_from_features( + Qwen3VLEmbeddingModel(config), + config, + feature_name="image_features", + feature_dim=config.hidden_size, + ) + embedding_weight = np.arange( + config.vocab_size * config.hidden_size, + dtype=np.float32, + ).reshape(config.vocab_size, config.hidden_size) + for name, initializer in graph.graph.initializers.items(): + if name.endswith("embed_tokens.weight"): + initializer.const_value = ir.tensor(embedding_weight) + + input_ids = np.array( + [[config.image_token_id, 1], [2, config.image_token_id]], + dtype=np.int64, + ) + image_features = np.arange(100, 108, dtype=np.float32).reshape(2, 4) + session = OnnxModelSession(graph) + result = session.run( + { + "input_ids": input_ids, + "image_features": image_features, + } + )["inputs_embeds"] + session.close() + + expected = embedding_weight[input_ids].copy() + expected[0, 0] = image_features[0] + expected[1, 1] = image_features[1] + np.testing.assert_array_equal(result, expected) + def _moe_config(quantization: QuantizationConfig | None) -> object: return make_config( diff --git a/src/mobius/models/qwen_vl.py b/src/mobius/models/qwen_vl.py index 2b4ae302e..85950f1c7 100644 --- a/src/mobius/models/qwen_vl.py +++ b/src/mobius/models/qwen_vl.py @@ -1031,7 +1031,7 @@ def forward( # feature, regardless of placeholder order or batch row. image_mask = op.Equal(input_ids, op.Constant(value_int=self.image_token_id)) if self.video_token_id is None: - video_mask = op.CastLike(False, image_mask) + video_mask = op.Not(op.Equal(input_ids, input_ids)) else: video_mask = op.Equal( input_ids, @@ -1040,7 +1040,8 @@ def forward( media_mask = op.Or(image_mask, video_mask) media_mask_3d = op.Unsqueeze(media_mask, [-1]) - flat_image_mask = op.Cast(op.Reshape(image_mask, [-1]), to=7) + flat_image_mask_bool = op.Reshape(image_mask, [-1]) + flat_image_mask = op.Cast(flat_image_mask_bool, to=7) flat_video_mask = op.Cast(op.Reshape(video_mask, [-1]), to=7) image_indices = op.Sub( op.CumSum(flat_image_mask, op.Constant(value_int=0)), @@ -1054,7 +1055,7 @@ def forward( op.ReduceSum(flat_image_mask, keepdims=0), ) flat_indices = op.Where( - op.CastLike(flat_image_mask, image_mask), + flat_image_mask_bool, image_indices, video_indices, ) From 3758328223444416fab50848ba282bb1bf50b0dd Mon Sep 17 00:00:00 2001 From: Justin Chu Date: Fri, 14 Aug 2026 23:32:44 -0700 Subject: [PATCH 11/15] Optimize Qwen VL packed coordinates Compute packed vision coordinates once, replace quadratic media lookup with boundary scatter and prefix sum, and reuse patch-local indices for frame boundaries. Simplify equivalent bilinear interpolation so Qwen3.5-VL stays below the deterministic graph-node regression threshold while preserving real image/video parity and CUDA semantics. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Signed-off-by: Justin Chu --- src/mobius/components/_qwen3_vl_vision.py | 125 +++++++++--------- .../components/_qwen3_vl_vision_test.py | 66 +++++++++ 2 files changed, 126 insertions(+), 65 deletions(-) create mode 100644 src/mobius/components/_qwen3_vl_vision_test.py diff --git a/src/mobius/components/_qwen3_vl_vision.py b/src/mobius/components/_qwen3_vl_vision.py index ccad7551d..7da0d5850 100644 --- a/src/mobius/components/_qwen3_vl_vision.py +++ b/src/mobius/components/_qwen3_vl_vision.py @@ -499,9 +499,9 @@ def _flat_grid_coordinates(self, op, grid_thw): preserving arbitrary image/video sizes and order. """ ms = self.spatial_merge_size - T_col = op.Squeeze(op.Slice(grid_thw, [0], [1], [1], [1]), [1]) # noqa: N806 - H_col = op.Squeeze(op.Slice(grid_thw, [1], [2], [1], [1]), [1]) # noqa: N806 - W_col = op.Squeeze(op.Slice(grid_thw, [2], [3], [1], [1]), [1]) # noqa: N806 + T_col = op.Gather(grid_thw, op.Constant(value_int=0), axis=1) # noqa: N806 + H_col = op.Gather(grid_thw, op.Constant(value_int=1), axis=1) # noqa: N806 + W_col = op.Gather(grid_thw, op.Constant(value_int=2), axis=1) # noqa: N806 patches_per_media = op.Mul(T_col, op.Mul(H_col, W_col)) patch_ends = op.CumSum(patches_per_media, op.Constant(value_int=0)) patch_starts = op.Pad( @@ -516,18 +516,21 @@ def _flat_grid_coordinates(self, op, grid_thw): total_patches, op.Constant(value_int=1), ) - # The number of completed media ranges is the owning media row. - media_ids = op.ReduceSum( - op.Cast( - op.GreaterOrEqual( - op.Unsqueeze(patch_ids, [1]), - op.Unsqueeze(patch_ends, [0]), - ), - to=7, - ), - [1], - keepdims=False, + # Mark each nonzero media boundary, then prefix-sum the markers. This + # maps patches to media in O(total_patches + num_media) rather than + # materializing an O(total_patches * num_media) comparison matrix. + media_boundaries = op.Slice(patch_ends, [0], [-1]) + boundary_updates = op.ConstantOfShape( + op.Shape(media_boundaries), + value=ir.tensor(np.array([1], dtype=np.int64)), ) + boundary_markers = op.ScatterElements( + op.Mul(patch_ids, op.Constant(value_int=0)), + media_boundaries, + boundary_updates, + axis=0, + ) + media_ids = op.CumSum(boundary_markers, op.Constant(value_int=0)) local_ids = op.Sub(patch_ids, op.Gather(patch_starts, media_ids)) H = op.Gather(H_col, media_ids) # noqa: N806 @@ -545,22 +548,22 @@ def _flat_grid_coordinates(self, op, grid_thw): intra_cols = op.Mod(intra_merge_ids, op.Constant(value_int=ms)) rows = op.Add(op.Mul(block_rows, op.Constant(value_int=ms)), intra_rows) cols = op.Add(op.Mul(block_cols, op.Constant(value_int=ms)), intra_cols) - return rows, cols, H, W + return rows, cols, H, W, frame_local_ids, patch_ids, total_patches - def _interpolate_pos_embed(self, op, grid_thw): + def _interpolate_pos_embed(self, op, coordinates): """Bilinearly interpolate learned positions for the packed media stream. Matches HuggingFace ``Qwen3VLVisionModel.fast_pos_embed_interpolate``. Args: op: OpBuilder instance. - grid_thw: ``(num_images, 3)`` INT64 with ``[T, H, W]`` per image. + coordinates: Shared packed ``(rows, cols, H, W)`` coordinate values. Returns: Position embeddings ``(total_patches, hidden_size)``. """ n = self.num_grid_per_side - rows, cols, H, W = self._flat_grid_coordinates(op, grid_thw) # noqa: N806 + rows, cols, H, W = coordinates # noqa: N806 rows_f = op.Cast(rows, to=1) cols_f = op.Cast(cols, to=1) H_f = op.Cast(H, to=1) # noqa: N806 @@ -568,20 +571,16 @@ def _interpolate_pos_embed(self, op, grid_thw): rows_scaled = op.Div(op.Mul(rows_f, float(n - 1)), op.Sub(H_f, 1.0)) cols_scaled = op.Div(op.Mul(cols_f, float(n - 1)), op.Sub(W_f, 1.0)) - row_floor = op.Cast(op.Floor(rows_scaled), to=7) - col_floor = op.Cast(op.Floor(cols_scaled), to=7) + row_floor_f = op.Floor(rows_scaled) + col_floor_f = op.Floor(cols_scaled) + row_floor = op.Cast(row_floor_f, to=7) + col_floor = op.Cast(col_floor_f, to=7) clip_max = op.Constant(value_int=n - 1) row_ceil = op.Min(op.Add(row_floor, op.Constant(value_int=1)), clip_max) col_ceil = op.Min(op.Add(col_floor, op.Constant(value_int=1)), clip_max) - row_delta = op.Sub(rows_scaled, op.Cast(row_floor, to=1)) - col_delta = op.Sub(cols_scaled, op.Cast(col_floor, to=1)) - one_minus_row = op.Sub(1.0, row_delta) - one_minus_col = op.Sub(1.0, col_delta) - w_00 = op.Unsqueeze(op.Mul(one_minus_row, one_minus_col), [1]) - w_01 = op.Unsqueeze(op.Mul(one_minus_row, col_delta), [1]) - w_10 = op.Unsqueeze(op.Mul(row_delta, one_minus_col), [1]) - w_11 = op.Unsqueeze(op.Mul(row_delta, col_delta), [1]) + row_delta = op.Unsqueeze(op.Sub(rows_scaled, row_floor_f), [1]) + col_delta = op.Unsqueeze(op.Sub(cols_scaled, col_floor_f), [1]) row_floor_base = op.Mul(row_floor, op.Constant(value_int=n)) row_ceil_base = op.Mul(row_ceil, op.Constant(value_int=n)) @@ -591,55 +590,47 @@ def _interpolate_pos_embed(self, op, grid_thw): idx_11 = op.Add(row_ceil_base, col_ceil) # Interpolate in float32 even when the learned table is f16/bf16. - e_00 = op.Mul(op.Cast(op.Gather(self.pos_embed, idx_00), to=1), w_00) - e_01 = op.Mul(op.Cast(op.Gather(self.pos_embed, idx_01), to=1), w_01) - e_10 = op.Mul(op.Cast(op.Gather(self.pos_embed, idx_10), to=1), w_10) - e_11 = op.Mul(op.Cast(op.Gather(self.pos_embed, idx_11), to=1), w_11) - return op.Add(op.Add(e_00, e_01), op.Add(e_10, e_11)) + pos_embed_f = op.Cast(self.pos_embed, to=1) + e_00 = op.Gather(pos_embed_f, idx_00) + e_01 = op.Gather(pos_embed_f, idx_01) + e_10 = op.Gather(pos_embed_f, idx_10) + e_11 = op.Gather(pos_embed_f, idx_11) + + # Two horizontal lerps followed by one vertical lerp are equivalent to + # the four explicit bilinear weights with fewer graph operations. + top = op.Add(e_00, op.Mul(op.Sub(e_01, e_00), col_delta)) + bottom = op.Add(e_10, op.Mul(op.Sub(e_11, e_10), col_delta)) + return op.Add(top, op.Mul(op.Sub(bottom, top), row_delta)) - def _compute_rotary_pos_ids(self, op, grid_thw): - """Compute 2D rotary position IDs for all images via ONNX Scan. + def _compute_rotary_pos_ids(self, op, coordinates): + """Combine shared packed coordinates into 2D rotary position IDs. Matches HF ``Qwen3VLVisionModel.rot_pos_emb()`` position indexing. Returns ``(total_patches, 2)`` INT64 with ``[h_pos, w_pos]`` per patch. """ - rows, cols, _, _ = self._flat_grid_coordinates(op, grid_thw) + rows, cols = coordinates return op.Concat(op.Unsqueeze(rows, [1]), op.Unsqueeze(cols, [1]), axis=1) - def _compute_cu_seqlens(self, op, grid_thw): + def _compute_cu_seqlens(self, op, frame_boundaries): """Compute full-attention cu_seqlens for all images. - Produces per-frame boundaries across all images without a control-flow - subgraph, equivalent to ``repeat_interleave(H * W, T)`` + CumSum. + Each packed frame starts where its frame-local patch index is zero. + Compacting those patch IDs and appending the total patch count produces + the same boundaries as ``repeat_interleave(H * W, T)`` + CumSum. Returns ``(total_frames + 1,)`` INT64. """ - T_col = op.Squeeze(op.Slice(grid_thw, [0], [1], [1], [1]), [1]) # noqa: N806 - H_col = op.Squeeze(op.Slice(grid_thw, [1], [2], [1], [1]), [1]) # noqa: N806 - W_col = op.Squeeze(op.Slice(grid_thw, [2], [3], [1], [1]), [1]) # noqa: N806 - frame_ends = op.CumSum(T_col, op.Constant(value_int=0)) - total_frames = op.ReduceSum(T_col, keepdims=False) - frame_ids = op.Range( - op.Constant(value_int=0), - total_frames, - op.Constant(value_int=1), + frame_local_ids, patch_ids, total_patches = frame_boundaries + frame_starts = op.Compress( + patch_ids, + op.Equal(frame_local_ids, op.Constant(value_int=0)), ) - media_ids = op.ReduceSum( - op.Cast( - op.GreaterOrEqual( - op.Unsqueeze(frame_ids, [1]), - op.Unsqueeze(frame_ends, [0]), - ), - to=7, - ), - [1], - keepdims=False, + return op.Concat( + frame_starts, + op.Unsqueeze(total_patches, [0]), + axis=0, ) - hw_per_media = op.Mul(H_col, W_col) - hw_per_frame = op.Gather(hw_per_media, media_ids) - cu = op.CumSum(hw_per_frame, op.Constant(value_int=0)) - return op.Pad(cu, op.Constant(value_ints=[1, 0]), op.Constant(value_int=0)) def forward( self, @@ -662,18 +653,22 @@ def forward( # Patch embedding hidden_states = self.patch_embed(op, hidden_states) + # Compute the packed patch coordinates once. Position interpolation, + # rotary IDs, and frame boundaries share these values. + coordinates = self._flat_grid_coordinates(op, grid_thw) + # Bilinear-interpolated position embeddings from learned grid. # Cast to match hidden_states dtype (interpolation computes in float32). - pos_embeds = self._interpolate_pos_embed(op, grid_thw) + pos_embeds = self._interpolate_pos_embed(op, coordinates[:4]) pos_embeds = op.CastLike(pos_embeds, hidden_states) hidden_states = op.Add(hidden_states, pos_embeds) # Compute rotary position IDs and embeddings from grid_thw - rotary_pos_ids = self._compute_rotary_pos_ids(op, grid_thw) + rotary_pos_ids = self._compute_rotary_pos_ids(op, coordinates[:2]) position_embeddings = self.rotary_pos_emb(op, rotary_pos_ids) # Compute cu_seqlens from grid_thw - cu_seqlens = self._compute_cu_seqlens(op, grid_thw) + cu_seqlens = self._compute_cu_seqlens(op, coordinates[4:]) # Transformer blocks deepstack_features = [] diff --git a/src/mobius/components/_qwen3_vl_vision_test.py b/src/mobius/components/_qwen3_vl_vision_test.py new file mode 100644 index 000000000..cfd09ffba --- /dev/null +++ b/src/mobius/components/_qwen3_vl_vision_test.py @@ -0,0 +1,66 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +from __future__ import annotations + +import onnx_ir as ir + +from mobius._testing import count_op_type, create_test_builder, create_test_input +from mobius.components._qwen3_vl_vision import Qwen3VLVisionModel + +_PATCH_DIM = 3 * 2 * 16 * 16 + + +def _build_vision_graph() -> ir.Graph: + module = Qwen3VLVisionModel( + depth=1, + hidden_size=32, + intermediate_size=64, + num_heads=4, + patch_size=16, + temporal_patch_size=2, + in_channels=3, + out_hidden_size=64, + spatial_merge_size=2, + num_position_embeddings=16, + deepstack_visual_indexes=[], + ) + builder, op, graph = create_test_builder() + pixel_values = create_test_input( + builder, + "pixel_values", + ["total_patches", _PATCH_DIM], + dtype=ir.DataType.FLOAT, + ) + grid_thw = create_test_input( + builder, + "grid_thw", + ["num_media", 3], + dtype=ir.DataType.INT64, + ) + image_features = module(op, pixel_values, grid_thw)[0] + image_features.name = "image_features" + graph.outputs.append(image_features) + return graph + + +def test_packed_coordinates_are_linear_and_shared(): + graph = _build_vision_graph() + + # Media ownership uses boundary scatter + prefix sum, not a quadratic + # [total_patches, num_media] comparison matrix. + assert count_op_type(graph, "ScatterElements") == 1 + # The only remaining comparison belongs to the single attention block. + assert count_op_type(graph, "GreaterOrEqual") == 1 + + # The two row/column coordinate values each feed both interpolation (Cast) + # and rotary IDs (Unsqueeze), proving the coordinate graph is emitted once. + shared_coordinates = [] + for node in graph: + if node.op_type != "Add": + continue + for output in node.outputs: + consumer_types = {consumer.op_type for consumer, _ in output.uses()} + if {"Cast", "Unsqueeze"} <= consumer_types: + shared_coordinates.append(output) + assert len(shared_coordinates) == 2 From 067cb489599687b74f96574de20133cb96e016a7 Mon Sep 17 00:00:00 2001 From: Justin Chu Date: Fri, 14 Aug 2026 23:52:52 -0700 Subject: [PATCH 12/15] Select Qwen3.5 metadata by package topology Preserve qwen3_5_text for decoder-only exports while selecting qwen3_5 for multimodal vision and embedding packages, including local-config exports whose composite config was unwrapped. Cover both generated package topologies without adding a runtime capability gate. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Signed-off-by: Justin Chu --- .../integrations/ort_genai/auto_export.py | 13 +++- .../ort_genai/auto_export_test.py | 65 ++++++++++++++++++- 2 files changed, 75 insertions(+), 3 deletions(-) diff --git a/src/mobius/integrations/ort_genai/auto_export.py b/src/mobius/integrations/ort_genai/auto_export.py index 2674e4252..f24fb68a0 100644 --- a/src/mobius/integrations/ort_genai/auto_export.py +++ b/src/mobius/integrations/ort_genai/auto_export.py @@ -94,7 +94,7 @@ "qwen3_vl_text": "qwen3_vl", "qwen3_5": "qwen3_5", "qwen3_5_vl": "qwen3_5", - "qwen3_5_text": "qwen3_5", + "qwen3_5_text": "qwen3_5_text", # MiniCPM uses standard 1D decoder position IDs (unlike Qwen-VL MRoPE). # The phi3v multimodal runtime provides that contract; callers supply # HF-preprocessed packed pixels through Generator.set_inputs(). @@ -186,6 +186,11 @@ def _select_ort_model_type( """ if is_decoder_only and config_model_type in _ORT_GENAI_MODEL_TYPE: return _ORT_GENAI_MODEL_TYPE[config_model_type] + if not is_decoder_only and config_model_type == "qwen3_5_text": + # Qwen3.5/Qwen3.8 multimodal builds unwrap the parent config to its + # text subtype, but their vision/embedding package uses the multimodal + # ORT pipeline and processor metadata. + return "qwen3_5" return _resolve_ort_genai_model_type(hf_model_type or "unknown") @@ -1379,7 +1384,11 @@ def write_ort_genai_config( # does not bind, so borrowing that type would mis-wire the graph. ort_model_type = "gemma3n" else: - ort_model_type = _resolve_ort_genai_model_type(raw_type) + ort_model_type = _select_ort_model_type( + raw_type, + raw_type, + is_decoder_only=is_decoder_only, + ) if ort_model_type == "unknown": logger.warning( "Could not determine ORT-GenAI model type: pkg.config.model_type " diff --git a/src/mobius/integrations/ort_genai/auto_export_test.py b/src/mobius/integrations/ort_genai/auto_export_test.py index e4246f3a7..1dc76cafa 100644 --- a/src/mobius/integrations/ort_genai/auto_export_test.py +++ b/src/mobius/integrations/ort_genai/auto_export_test.py @@ -84,7 +84,7 @@ def test_known_model_type(self): assert _resolve_ort_genai_model_type("qwen3") == "qwen2" assert _resolve_ort_genai_model_type("qwen3_5") == "qwen3_5" assert _resolve_ort_genai_model_type("qwen3_5_vl") == "qwen3_5" - assert _resolve_ort_genai_model_type("qwen3_5_text") == "qwen3_5" + assert _resolve_ort_genai_model_type("qwen3_5_text") == "qwen3_5_text" assert _resolve_ort_genai_model_type("gemma2") == "gemma" assert _resolve_ort_genai_model_type("llama") == "llama" @@ -144,6 +144,24 @@ def test_multimodal_keeps_hf_type(self): def test_decoder_only_falls_back_to_hf_when_config_missing(self): assert _select_ort_model_type(None, "qwen3", is_decoder_only=True) == "qwen2" + def test_qwen35_text_type_depends_on_package_topology(self): + assert ( + _select_ort_model_type( + "qwen3_5_text", + "qwen3_5", + is_decoder_only=True, + ) + == "qwen3_5_text" + ) + assert ( + _select_ort_model_type( + "qwen3_5_text", + "qwen3_5_text", + is_decoder_only=False, + ) + == "qwen3_5" + ) + def test_decoder_only_unknown_config_falls_back_to_hf(self): # An unrecognised config.model_type (not in _ORT_GENAI_MODEL_TYPE) must # not pass straight through as an invalid ORT type; fall back to the @@ -1176,6 +1194,51 @@ class FakeConfig: assert decoder["inputs"]["past_key_names"] == "past_key_values.%d.key" assert decoder["inputs"]["past_conv_names"] == "past_key_values.%d.conv_state" + def test_qwen35_text_package_preserves_decoder_only_model_type(self, tmp_path): + import dataclasses + + from mobius._model_package import ModelPackage + + @dataclasses.dataclass + class FakeConfig: + model_type: str = "qwen3_5_text" + vocab_size: int = 256 + hidden_size: int = 64 + num_hidden_layers: int = 4 + num_attention_heads: int = 4 + num_key_value_heads: int = 2 + head_dim: int = 16 + max_position_embeddings: int = 128 + pad_token_id: int = 0 + + package = ModelPackage( + { + "model": _mock_model( + inputs=[ + "input_ids", + "attention_mask", + "position_ids", + "past_key_values.0.conv_state", + "past_key_values.3.key", + "past_key_values.3.value", + ], + outputs=[ + "logits", + "present.0.conv_state", + "present.3.key", + "present.3.value", + ], + ) + }, + config=FakeConfig(), + ) + + result = write_ort_genai_config(package, str(tmp_path)) + + with open(result["genai_config"], encoding="utf-8") as config_file: + generated = json.load(config_file) + assert generated["model"]["type"] == "qwen3_5_text" + def test_olive_renamed_logits_output_is_emitted(self, tmp_path): pkg = _make_fake_llm_pkg("qwen2") pkg["model"] = _mock_model( From 064015b04d7c41ce3b3711bc5ae093b1b819f47c Mon Sep 17 00:00:00 2001 From: Justin Chu Date: Wed, 19 Aug 2026 16:33:28 -0700 Subject: [PATCH 13/15] Remove Qwen3.8 Olive examples Drop the standalone reduced-checkpoint Olive recipe and its directly coupled tests and golden fixtures. Core Qwen3.8 model, processor, metadata, and parity coverage remain unchanged. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Signed-off-by: Justin Chu --- examples/olive/qwen3_8-27b/README.md | 45 -- examples/olive/qwen3_8-27b/inference.py | 228 ------ examples/olive/qwen3_8-27b/optimize.py | 149 ---- examples/olive/qwen3_8-27b/requirements.txt | 3 - .../validate_reduced_checkpoint.py | 709 ------------------ .../vision-language/qwen3_8-27b-reduced.json | 40 - .../qwen3_8-27b-reduced_generation.json | 33 - tests/qwen38_real_weight_test.py | 295 -------- 8 files changed, 1502 deletions(-) delete mode 100644 examples/olive/qwen3_8-27b/README.md delete mode 100644 examples/olive/qwen3_8-27b/inference.py delete mode 100644 examples/olive/qwen3_8-27b/optimize.py delete mode 100644 examples/olive/qwen3_8-27b/requirements.txt delete mode 100644 examples/olive/qwen3_8-27b/validate_reduced_checkpoint.py delete mode 100644 testdata/golden/vision-language/qwen3_8-27b-reduced.json delete mode 100644 testdata/golden/vision-language/qwen3_8-27b-reduced_generation.json delete mode 100644 tests/qwen38_real_weight_test.py diff --git a/examples/olive/qwen3_8-27b/README.md b/examples/olive/qwen3_8-27b/README.md deleted file mode 100644 index 97490ffd7..000000000 --- a/examples/olive/qwen3_8-27b/README.md +++ /dev/null @@ -1,45 +0,0 @@ -# Qwen3.8-27B reduced-real Olive validation - -This recipe range-fetches only deterministic slices from the pinned BF16 -checkpoint (`1d4bf0f2ff6012fd82039f2fa52739d0dd7c60c0`; 18 shards/1199 tensors). -The fixture has 4 decoder layers (three DeltaNet and one full attention), -one vision block, remapped image/video IDs, and no MTP tensors because the -standard target forward does not consume the optional self-speculative drafter. -Mobius exposes that drafter through the separate `qwen35-mtp` package contract. -The reduced package is token-ID-only and deliberately omits the production -tokenizer and generation config because its vocabulary and special-token IDs -are remapped to 256 entries. It retains the pinned image and video processor -metadata for media-contract validation. Each variant output directory must be -empty so stale assets from an older recipe cannot enter the assembled package. - -```powershell -python examples/olive/qwen3_8-27b/validate_reduced_checkpoint.py --matrix f32-cpu f16-cuda -``` - -The command validates strict reduced HF loading, all three ONNX components, -full logits, 20-token cached generation, graph/provider placement, save/load, -and then assembles and directly runs the Olive Q4_K_M package. It is purposely -not an ORT GenAI capability-gated test. - -The Q4 recipe keeps DeltaNet's narrow `in_proj_a` decay and `in_proj_b` -time-step gates in FP16. Quantizing those recurrent controls destabilizes -cached generation; all larger decoder matrices remain eligible for -`MatMulNBits`. All three Q4 package components are reloaded with CUDA enabled; -the 20-token semantic run uses CPU because repeated identical ORT 1.26 CUDA -`MatMulNBits` runs can diverge and produce non-finite logits for this fixture. -`optimize.py` derives the recurrent-gate node exclusions from the exported -decoder, so the validated package must be assembled through that script rather -than a static Olive JSON recipe. - -BF16 export and package reload are valid Mobius outputs. The CUDA-12-compatible -ORT 1.26 wheel used for this reduced-real run cannot initialize the BF16 hybrid -graph (`CausalConvWithState`/`Softplus` provider placement), while newer -ORT-GPU wheels available in the test environment require CUDA 13. This is a -downstream runtime waiver, not an export or support gate. - -The same old ORT wheel is nondeterministic for dynamic Qwen vision batches -through its `PackedMultiHeadAttention` CUDA kernel. The validator therefore -uses the portable standard-attention vision graph for stable real CUDA -image/video/mixed execution while retaining the CUDA-optimized decoder parity -and provider profile. Mobius still emits both graph variants without deciding -which downstream runtime version can load them. diff --git a/examples/olive/qwen3_8-27b/inference.py b/examples/olive/qwen3_8-27b/inference.py deleted file mode 100644 index 09337ef63..000000000 --- a/examples/olive/qwen3_8-27b/inference.py +++ /dev/null @@ -1,228 +0,0 @@ -#!/usr/bin/env python -# Copyright (c) Microsoft Corporation. -# Licensed under the MIT License. - -"""Direct ONNX Runtime generation for reduced Qwen3.8 hybrid-VL packages.""" - -from __future__ import annotations - -import argparse -import json -from collections import Counter -from pathlib import Path -from typing import Any - -import numpy as np - -MODEL_ID = "Qwen/Qwen3.8-27B" -REVISION = "1d4bf0f2ff6012fd82039f2fa52739d0dd7c60c0" -_BFLOAT16_ONNX_TYPE = 16 - - -def _dtype(ort_type: str): - if ort_type == "tensor(float)": - return np.float32 - if ort_type == "tensor(float16)": - return np.float16 - if ort_type == "tensor(bfloat16)": - import ml_dtypes - - return ml_dtypes.bfloat16 - raise TypeError(f"Unsupported ONNX input type: {ort_type}") - - -def _shape(shape: list[Any]) -> tuple[int, ...]: - """Concretize dynamic hybrid-cache shapes for a one-item decode.""" - result = [] - for value in shape: - if isinstance(value, int): - result.append(value) - elif "batch" in str(value): - result.append(1) - elif "past" in str(value) or "sequence" in str(value): - result.append(0) - else: - raise ValueError(f"Cannot resolve state dimension {value!r}") - return tuple(result) - - -def _create_session(model_path: Path, device: str, profile: bool = False): - if device == "cuda": - import torch # noqa: F401 # Preloads matching CUDA DLLs on Windows. - - import onnxruntime as ort - - if device == "cuda" and hasattr(ort, "preload_dlls"): - ort.preload_dlls() - options = ort.SessionOptions() - options.enable_profiling = profile - providers = ( - ["CUDAExecutionProvider", "CPUExecutionProvider"] - if device == "cuda" - else ["CPUExecutionProvider"] - ) - session = ort.InferenceSession(str(model_path), options, providers=providers) - if device == "cuda" and session.get_providers()[0] != "CUDAExecutionProvider": - raise RuntimeError(f"CUDAExecutionProvider was requested: {session.get_providers()}") - return session - - -def _initial_states(session) -> dict[str, Any]: - import onnxruntime as ort - - result: dict[str, Any] = {} - for model_input in session.get_inputs(): - if not model_input.name.startswith("past_key_values."): - continue - zeros = np.zeros(_shape(model_input.shape), dtype=_dtype(model_input.type)) - if model_input.type == "tensor(bfloat16)": - zeros = ort.OrtValue.ortvalue_from_numpy_with_onnx_type( - np.zeros(zeros.shape, dtype=np.uint16), _BFLOAT16_ONNX_TYPE - ) - result[model_input.name] = zeros - return result - - -def _run(session, output_names: list[str], feeds: dict[str, Any]) -> list[Any]: - import onnxruntime as ort - - if any(isinstance(value, ort.OrtValue) for value in feeds.values()): - values = { - name: value - if isinstance(value, ort.OrtValue) - else ort.OrtValue.ortvalue_from_numpy(value) - for name, value in feeds.items() - } - return list(session.run_with_ort_values(output_names, values)) - return session.run(output_names, feeds) - - -def _numpy(value: Any) -> np.ndarray: - import onnxruntime as ort - - if not isinstance(value, ort.OrtValue): - return value - if value.data_type() == "tensor(bfloat16)": - import torch - - return torch.from_dlpack(value).float().cpu().numpy() - return value.numpy() - - -def _update_states(states: dict[str, Any], names: list[str], values: list[Any]) -> None: - for name, value in zip(names, values): - if name.startswith("present."): - states[name.replace("present.", "past_key_values.", 1)] = value - - -def _embedding(session, token_ids: np.ndarray, hidden_size: int) -> np.ndarray: - inputs = {item.name: item for item in session.get_inputs()} - media_dtype = _dtype(inputs["image_features"].type) - feeds = { - "input_ids": token_ids, - # Empty media is intentional for text generation; image/video paths - # run this same model with actual vision features in the validator. - "image_features": np.zeros((0, hidden_size), dtype=media_dtype), - } - outputs = _run( - session, - [item.name for item in session.get_outputs()], - {name: value for name, value in feeds.items() if name in inputs}, - ) - return _numpy(outputs[0]) - - -def run_token_ids( - model_dir: str | Path, - input_ids: list[int], - *, - hidden_size: int, - max_new_tokens: int, - device: str, - profile: bool = False, -) -> tuple[list[int], list[np.ndarray], str | None]: - """Generate exact-length greedy tokens through embedding and hybrid decoder.""" - root = Path(model_dir) - decoder = _create_session(root / "decoder" / "model.onnx", device, profile) - embedding = _create_session(root / "embedding" / "model.onnx", device) - states = _initial_states(decoder) - names = [item.name for item in decoder.get_outputs()] - generated: list[int] = [] - logits_by_step: list[np.ndarray] = [] - past = 0 - logits: np.ndarray | None = None - - for token_id in input_ids: - ids = np.array([[token_id]], dtype=np.int64) - embeds = _embedding(embedding, ids, hidden_size) - feeds: dict[str, Any] = { - "inputs_embeds": embeds, - "attention_mask": np.ones((1, past + 1), dtype=np.int64), - # Qwen MRoPE uses three equal text positions. - "position_ids": np.full((3, 1, 1), past, dtype=np.int64), - **states, - } - outputs = _run(decoder, names, feeds) - _update_states(states, names, outputs) - logits = _numpy(outputs[names.index("logits")])[0, -1].astype(np.float32) - past += 1 - if logits is None: - raise ValueError("input_ids must not be empty") - - for _ in range(max_new_tokens): - logits_by_step.append(logits.copy()) - token_id = int(np.argmax(logits)) - generated.append(token_id) - ids = np.array([[token_id]], dtype=np.int64) - embeds = _embedding(embedding, ids, hidden_size) - outputs = _run( - decoder, - names, - { - "inputs_embeds": embeds, - "attention_mask": np.ones((1, past + 1), dtype=np.int64), - "position_ids": np.full((3, 1, 1), past, dtype=np.int64), - **states, - }, - ) - _update_states(states, names, outputs) - logits = _numpy(outputs[names.index("logits")])[0, -1].astype(np.float32) - past += 1 - return generated, logits_by_step, decoder.end_profiling() if profile else None - - -def summarize_profile(profile_path: str) -> dict[str, int]: - events = json.loads(Path(profile_path).read_text(encoding="utf-8")) - providers: Counter[str] = Counter() - for event in events: - provider = event.get("args", {}).get("provider") - if event.get("cat") == "Node" and provider: - providers[str(provider)] += 1 - return dict(sorted(providers.items())) - - -def main() -> None: - parser = argparse.ArgumentParser(description=__doc__) - parser.add_argument("--model-dir", required=True) - parser.add_argument("--token-ids", nargs="+", type=int, default=[1, 42, 17]) - parser.add_argument("--hidden-size", type=int, default=256) - parser.add_argument("--max-new-tokens", type=int, default=4) - parser.add_argument("--device", choices=["cpu", "cuda"], default="cuda") - parser.add_argument("--profile", action="store_true") - args = parser.parse_args() - ids, _logits, profile = run_token_ids( - args.model_dir, - args.token_ids, - hidden_size=args.hidden_size, - max_new_tokens=args.max_new_tokens, - device=args.device, - profile=args.profile, - ) - print(f"Generated token IDs: {ids}") - if profile: - assert profile is not None - print(f"Provider placement: {summarize_profile(profile)}") - - -if __name__ == "__main__": - main() diff --git a/examples/olive/qwen3_8-27b/optimize.py b/examples/olive/qwen3_8-27b/optimize.py deleted file mode 100644 index b1a31da67..000000000 --- a/examples/olive/qwen3_8-27b/optimize.py +++ /dev/null @@ -1,149 +0,0 @@ -#!/usr/bin/env python -# Copyright (c) Microsoft Corporation. -# Licensed under the MIT License. - -"""Assemble a complete Q4_K_M Qwen3.8 three-model package with Olive.""" - -from __future__ import annotations - -import json -import shutil -from pathlib import Path - -import onnx_ir as ir -from inference import MODEL_ID, REVISION - -_ASSETS = { - "config.json", - "generation_config.json", - "tokenizer.json", - "tokenizer_config.json", - "special_tokens_map.json", - "chat_template.jinja", - "preprocessor_config.json", - "processor_config.json", - "image_processor.json", - "processor-waiver.txt", -} - - -def olive_config( - decoder: Path, - output: Path, - *, - nodes_to_exclude: list[str] | None = None, -) -> dict: - """Return the CPU-isolated Olive Q4_K_M decoder-only workflow.""" - return { - "input_model": {"type": "OnnxModel", "model_path": str(decoder)}, - "passes": { - "q4_k_m": { - "type": "OnnxKQuantQuantization", - "bits": 4, - "block_size": 32, - "save_as_external_data": True, - "all_tensors_to_one_file": True, - "external_data_name": "decoder.onnx.data", - "size_threshold": 1024, - "nodes_to_exclude": nodes_to_exclude or [], - } - }, - "engine": { - "target": { - "type": "LocalSystem", - "accelerators": [ - {"device": "cpu", "execution_providers": ["CPUExecutionProvider"]} - ], - } - }, - "no_artifacts": True, - "output_dir": str(output), - # A global Olive cache can return a decoder produced with a different - # exclusion set. Scope and clean it so recurrent-gate policy is exact. - "cache_dir": str(output.parent / ".olive-cache"), - "clean_cache": True, - } - - -def _recurrent_gate_nodes(decoder: Path) -> list[str]: - """Keep DeltaNet decay/time-step gates in f16 to preserve recurrent stability.""" - model = ir.load(decoder) - return [ - node.name - for node in model.graph.all_nodes() - if node.op_type == "MatMul" - and ("/linear_attn/in_proj_a/" in node.name or "/linear_attn/in_proj_b/" in node.name) - ] - - -def quantize_package(source_dir: str | Path, output_dir: str | Path) -> Path: - """Quantize only decoder and copy vision, embedding, metadata, and tokenizer.""" - import olive.systems.local as olive_local - from olive.workflows import run as olive_run - - source, output = Path(source_dir), Path(output_dir) - if output.exists() and any(output.iterdir()): - raise FileExistsError(f"Output must be empty: {output}") - output.mkdir(parents=True, exist_ok=True) - decoder = source / "decoder" / "model.onnx" - if not decoder.is_file(): - raise FileNotFoundError(decoder) - olive_output = output / ".olive" - # Olive can eagerly register unrelated GPU EP DLLs. K-quant is a - # weight-only CPU pass, so suppress that registration only for this call. - preserved_fp16_nodes = _recurrent_gate_nodes(decoder) - register = olive_local.maybe_register_ep_libraries - olive_local.maybe_register_ep_libraries = lambda _paths: None - try: - olive_run( - olive_config( - decoder, - olive_output, - nodes_to_exclude=preserved_fp16_nodes, - ) - ) - finally: - olive_local.maybe_register_ep_libraries = register - models = list(olive_output.rglob("*.onnx")) - if len(models) != 1: - raise RuntimeError(f"Expected one Olive decoder, found {models}") - decoder_dir = output / "decoder" - decoder_dir.mkdir() - for item in models[0].parent.iterdir(): - if item.is_file(): - shutil.copy2(item, decoder_dir / item.name) - produced = decoder_dir / models[0].name - if produced != decoder_dir / "model.onnx": - produced.replace(decoder_dir / "model.onnx") - shutil.rmtree(olive_output) - olive_cache = output / ".olive-cache" - if olive_cache.exists(): - shutil.rmtree(olive_cache) - for name in ("embedding", "vision_encoder"): - shutil.copytree(source / name, output / name) - for name in _ASSETS: - if (source / name).is_file(): - shutil.copy2(source / name, output / name) - source_manifest = source / "source_manifest.json" - manifest = ( - json.loads(source_manifest.read_text(encoding="utf-8")) - if source_manifest.is_file() - else {"model_id": MODEL_ID, "revision": REVISION} - ) - manifest.update( - { - "quantization": "Q4_K_M", - "quantized_component": "decoder", - "olive_provider": "CPUExecutionProvider", - "preserved_fp16_recurrent_gate_nodes": len(preserved_fp16_nodes), - "components": [ - "decoder/model.onnx", - "embedding/model.onnx", - "vision_encoder/model.onnx", - ], - } - ) - (output / "source_manifest.json").write_text( - json.dumps(manifest, indent=2), encoding="utf-8" - ) - return output diff --git a/examples/olive/qwen3_8-27b/requirements.txt b/examples/olive/qwen3_8-27b/requirements.txt deleted file mode 100644 index 82aac5f2c..000000000 --- a/examples/olive/qwen3_8-27b/requirements.txt +++ /dev/null @@ -1,3 +0,0 @@ -olive-ai -safetensors -requests diff --git a/examples/olive/qwen3_8-27b/validate_reduced_checkpoint.py b/examples/olive/qwen3_8-27b/validate_reduced_checkpoint.py deleted file mode 100644 index 93b63f62f..000000000 --- a/examples/olive/qwen3_8-27b/validate_reduced_checkpoint.py +++ /dev/null @@ -1,709 +0,0 @@ -#!/usr/bin/env python -# Copyright (c) Microsoft Corporation. -# Licensed under the MIT License. - -"""Validate Qwen3.8-27B using a small, pinned, reduced-real BF16 fixture. - -The fixture deliberately retains the production 3xDeltaNet + 1xGQA layer -schedule and one Qwen vision block. It is not a randomly initialized proxy: -every cached value is a deterministic row/column slice read with verified HTTP -Range requests from the pinned 18-shard, 1199-tensor checkpoint. -""" - -from __future__ import annotations - -import argparse -import dataclasses -import json -import math -import struct -import time -from collections import Counter -from pathlib import Path - -import numpy as np -import requests -import torch -from huggingface_hub import hf_hub_download -from inference import ( - MODEL_ID, - REVISION, - _create_session, - _embedding, - _initial_states, - _numpy, - _run, - run_token_ids, - summarize_profile, -) -from safetensors import safe_open -from safetensors.torch import load_file, save_file - -FIXTURE_SCHEMA_VERSION = 1 -_RANGE_ATTEMPTS = 3 -_VOCAB_SIZE = 256 -_HIDDEN_SIZE = 256 -_GENERATION_TOKENS = 20 -_MEDIA_IDS = { - "image_token_id": 250, - "video_token_id": 251, - "vision_start_token_id": 252, - "vision_end_token_id": 253, -} -_DTYPES = { - "f32": (torch.float32, "FLOAT"), - "f16": (torch.float16, "FLOAT16"), - "bf16": (torch.bfloat16, "BFLOAT16"), -} - - -class _PinnedSafetensors: - """Header-first, retrying reader for exact checkpoint byte ranges.""" - - def __init__(self) -> None: - index_path = hf_hub_download( - MODEL_ID, "model.safetensors.index.json", revision=REVISION - ) - index = json.loads(Path(index_path).read_text(encoding="utf-8")) - self.weight_map: dict[str, str] = index["weight_map"] - if len(self.weight_map) != 1199 or len(set(self.weight_map.values())) != 18: - raise ValueError( - "Pinned checkpoint manifest is not the expected 18 shards / 1199 tensors" - ) - self._headers: dict[str, tuple[int, dict]] = {} - self._session = requests.Session() - - def _url(self, shard: str) -> str: - return f"https://huggingface.co/{MODEL_ID}/resolve/{REVISION}/{shard}" - - def _range(self, shard: str, start: int, end: int) -> bytes: - expected = end - start + 1 - prefix = f"bytes {start}-{end}/" - error = "" - for attempt in range(_RANGE_ATTEMPTS): - try: - with self._session.get( - self._url(shard), - headers={"Range": f"bytes={start}-{end}"}, - timeout=180, - stream=True, - ) as response: - content_range = response.headers.get("Content-Range", "") - content_length = response.headers.get("Content-Length") - payload = response.content - if ( - response.status_code == 206 - and content_range.startswith(prefix) - and (content_length is None or content_length == str(expected)) - and len(payload) == expected - ): - return payload - error = ( - f"status={response.status_code}, range={content_range!r}, " - f"length={content_length}, bytes={len(payload)}, expected={expected}" - ) - except requests.RequestException as exc: - error = f"{type(exc).__name__}: {exc}" - if attempt + 1 < _RANGE_ATTEMPTS: - time.sleep(2**attempt) - raise RuntimeError( - f"Range fetch failed after {_RANGE_ATTEMPTS} attempts for {shard} " - f"bytes {start}-{end}: {error}" - ) - - def _header(self, shard: str) -> tuple[int, dict]: - if shard not in self._headers: - size = struct.unpack(" torch.Tensor: - """Fetch only leading source rows, then deterministically trim every axis.""" - shard = self.weight_map[name] - header_size, header = self._header(shard) - entry = header[name] - source_shape = list(entry["shape"]) - target_shape = list(shape) - if len(source_shape) != len(target_shape) or any( - a < b for a, b in zip(source_shape, target_shape) - ): - raise ValueError( - f"Cannot reduce {name}: source={source_shape}, target={target_shape}" - ) - dtype_name = entry["dtype"] - dtype = {"BF16": torch.bfloat16, "F32": torch.float32}[dtype_name] - element_size = {"BF16": 2, "F32": 4}[dtype_name] - rows = target_shape[0] if source_shape else 1 - row_width = math.prod(source_shape[1:]) if source_shape else 1 - start, _end = entry["data_offsets"] - length = rows * row_width * element_size - payload = self._range( - shard, 8 + header_size + start, 8 + header_size + start + length - 1 - ) - leading = ( - torch.frombuffer(bytearray(payload), dtype=dtype) - .clone() - .reshape([rows, *source_shape[1:]]) - ) - return leading[tuple(slice(0, size) for size in target_shape)].contiguous() - - -def default_reduced_cache_path() -> Path: - return ( - Path.home() - / ".cache" - / "mobius" - / "qwen3_8-27b" - / f"reduced-{REVISION}-schema-v{FIXTURE_SCHEMA_VERSION}.safetensors" - ) - - -def _reduced_hf_config(): - """Construct a tiny native HF config preserving all inference layer families.""" - from transformers import AutoConfig, Qwen3_5Config - - source = AutoConfig.from_pretrained(MODEL_ID, revision=REVISION, trust_remote_code=False) - text = source.text_config.to_dict() - text.update( - vocab_size=_VOCAB_SIZE, - hidden_size=_HIDDEN_SIZE, - intermediate_size=512, - num_hidden_layers=4, - num_attention_heads=2, - num_key_value_heads=1, - head_dim=128, - layer_types=[ - "linear_attention", - "linear_attention", - "linear_attention", - "full_attention", - ], - linear_key_head_dim=32, - linear_value_head_dim=32, - linear_num_key_heads=2, - linear_num_value_heads=4, - mtp_num_hidden_layers=0, - max_position_embeddings=128, - bos_token_id=1, - eos_token_id=2, - pad_token_id=0, - partial_rotary_factor=0.1875, - rope_parameters={ - "rope_type": "default", - "rope_theta": 10_000_000, - "partial_rotary_factor": 0.1875, - "mrope_section": [4, 4, 4], - "mrope_interleaved": True, - }, - ) - vision = source.vision_config.to_dict() - vision.update( - depth=1, - hidden_size=128, - intermediate_size=256, - num_heads=4, - out_hidden_size=_HIDDEN_SIZE, - num_position_embeddings=64, - ) - return Qwen3_5Config(text_config=text, vision_config=vision, **_MEDIA_IDS) - - -def _reduced_mobius_config(dtype_name: str): - import onnx_ir as ir - - from mobius._configs import ArchitectureConfig - - hf_config = _reduced_hf_config() - config = ArchitectureConfig.from_transformers( - hf_config.text_config, parent_config=hf_config - ) - assert config.vision is not None - return dataclasses.replace( - config, - vocab_size=_VOCAB_SIZE, - hidden_size=_HIDDEN_SIZE, - intermediate_size=512, - num_hidden_layers=4, - num_attention_heads=2, - num_key_value_heads=1, - head_dim=128, - layer_types=[ - "linear_attention", - "linear_attention", - "linear_attention", - "full_attention", - ], - linear_key_head_dim=32, - linear_value_head_dim=32, - linear_num_key_heads=2, - linear_num_value_heads=4, - max_position_embeddings=128, - image_token_id=_MEDIA_IDS["image_token_id"], - video_token_id=_MEDIA_IDS["video_token_id"], - vision_start_token_id=_MEDIA_IDS["vision_start_token_id"], - vision_end_token_id=_MEDIA_IDS["vision_end_token_id"], - mrope_section=[4, 4, 4], - mrope_interleaved=True, - dtype=getattr(ir.DataType, _DTYPES[dtype_name][1]), - vision=dataclasses.replace( - config.vision, - hidden_size=128, - intermediate_size=256, - num_hidden_layers=1, - num_attention_heads=4, - out_hidden_size=_HIDDEN_SIZE, - num_position_embeddings=64, - ), - ) - - -def _expected_hf_state() -> dict[str, torch.Tensor]: - from transformers import Qwen3_5ForConditionalGeneration - - # Native HF state names are the checkpoint names, so strict loading below - # guards both the fixture's tensor coverage and the source-name mapping. - return { - name: tensor - for name, tensor in Qwen3_5ForConditionalGeneration(_reduced_hf_config()) - .state_dict() - .items() - if not name.startswith(("mtp_", "mtp.")) - } - - -def _build_reduced_state(cache_path: Path) -> dict[str, torch.Tensor]: - expected_metadata = { - "model_id": MODEL_ID, - "revision": REVISION, - "fixture_schema": str(FIXTURE_SCHEMA_VERSION), - "source_shards": "18", - "source_tensors": "1199", - } - if cache_path.is_file(): - with safe_open(cache_path, framework="pt") as cached: - actual = cached.metadata() or {} - if {key: actual.get(key) for key in expected_metadata} != expected_metadata: - raise ValueError( - "Reduced cache metadata mismatch; remove the stale cache and retry." - ) - return load_file(cache_path) - source = _PinnedSafetensors() - expected = _expected_hf_state() - missing = sorted(set(expected) - set(source.weight_map)) - if missing: - raise ValueError( - f"Reduced HF model requests tensors absent from checkpoint: {missing[:5]}" - ) - # The expected model has only layers 0-3, intentionally covers three - # DeltaNet layers and layer 3 full attention. MTP has no expected state. - state = {name: source.sliced(name, tensor.shape) for name, tensor in expected.items()} - cache_path.parent.mkdir(parents=True, exist_ok=True) - temporary = cache_path.with_suffix(".tmp.safetensors") - save_file(state, temporary, metadata=expected_metadata) - temporary.replace(cache_path) - return state - - -def _hf_model(state: dict[str, torch.Tensor], *, dtype: torch.dtype, device: str): - from transformers import Qwen3_5ForConditionalGeneration - - model = Qwen3_5ForConditionalGeneration(_reduced_hf_config()).to( - device=device, dtype=dtype - ) - target = model.state_dict() - if set(target) != set(state): - raise ValueError( - f"Strict fixture state mismatch: missing={sorted(set(target) - set(state))[:5]}, " - f"extra={sorted(set(state) - set(target))[:5]}" - ) - model.load_state_dict( - {k: v.to(device=device, dtype=target[k].dtype) for k, v in state.items()}, strict=True - ) - return model.eval() - - -def _mobius_package(state: dict[str, torch.Tensor], *, dtype_name: str, ep: str): - from mobius import build_from_module - from mobius.models.qwen35 import Qwen35VL3ModelCausalLMModel - - config = _reduced_mobius_config(dtype_name) - module = Qwen35VL3ModelCausalLMModel(config) - package = build_from_module(module, config, task="hybrid-qwen-vl", execution_provider=ep) - package.apply_weights(module.preprocess_weights(dict(state))) - unset = [ - f"{model_name}:{name}" - for model_name, model in package.items() - for name, value in model.graph.initializers.items() - if value.const_value is None - ] - if unset: - raise ValueError(f"Weighted Qwen3.8 graph has unset initializers: {unset[:5]}") - return package - - -def _onnx_prefill_logits(package_dir: Path, token_ids: list[int], device: str) -> np.ndarray: - embedding = _create_session(package_dir / "embedding" / "model.onnx", device) - decoder = _create_session(package_dir / "decoder" / "model.onnx", device) - ids = np.array([token_ids], dtype=np.int64) - embedded = _embedding(embedding, ids, _HIDDEN_SIZE) - output_names = [item.name for item in decoder.get_outputs()] - outputs = _run( - decoder, - output_names, - { - "inputs_embeds": _numpy(embedded), - "attention_mask": np.ones_like(ids), - "position_ids": np.repeat(np.arange(len(token_ids))[None, None, :], 3, axis=0), - **_initial_states(decoder), - }, - ) - return _numpy(outputs[output_names.index("logits")]).astype(np.float32) - - -def _hf_logits(model, token_ids: list[int], device: str) -> np.ndarray: - ids = torch.tensor([token_ids], dtype=torch.long, device=device) - with torch.no_grad(): - return ( - model(input_ids=ids, attention_mask=torch.ones_like(ids), use_cache=False) - .logits.float() - .cpu() - .numpy() - ) - - -def _graph_audit(package) -> dict[str, dict[str, int]]: - return { - name: dict( - sorted( - Counter( - f"{node.domain or 'ai.onnx'}::{node.op_type}" - for node in model.graph.all_nodes() - ).items() - ) - ) - for name, model in package.items() - } - - -def _assert_logits_close( - actual: np.ndarray, - expected: np.ndarray, - *, - atol: float, - label: str, -) -> None: - max_abs = float(np.max(np.abs(actual - expected))) - cosine = float( - np.dot(actual.ravel(), expected.ravel()) - / (np.linalg.norm(actual) * np.linalg.norm(expected)) - ) - print(f"{label}: max_abs={max_abs:.8f}, cosine={cosine:.9f}") - if max_abs > atol or cosine < 0.999: - raise AssertionError( - f"{label} parity failed: max_abs={max_abs:.8f}, cosine={cosine:.9f}" - ) - np.testing.assert_allclose(actual, expected, rtol=1e-3, atol=atol) - - -def _save_package_assets(package_dir: Path) -> None: - """Copy pinned processor metadata for the reduced token-ID-only package.""" - from transformers import AutoProcessor - - _reduced_hf_config().save_pretrained(package_dir) - try: - processor = AutoProcessor.from_pretrained(MODEL_ID, revision=REVISION) - (package_dir / "processor_config.json").write_text( - json.dumps(processor.to_dict(), indent=2), - encoding="utf-8", - ) - except (ImportError, OSError, ValueError) as error: - # The ONNX package remains directly runnable without the optional - # processor serialization; retain the precise reason in its manifest. - (package_dir / "processor-waiver.txt").write_text(f"{type(error).__name__}: {error}\n") - (package_dir / "source_manifest.json").write_text( - json.dumps( - { - "model_id": MODEL_ID, - "revision": REVISION, - "fixture_schema": FIXTURE_SCHEMA_VERSION, - "runtime": "onnxruntime-direct", - "text_input_contract": "token-ids-only", - "components": [ - "decoder/model.onnx", - "embedding/model.onnx", - "vision_encoder/model.onnx", - ], - }, - indent=2, - ), - encoding="utf-8", - ) - - -def _media_smoke(package_dir: Path, device: str) -> dict[str, tuple[int, ...]]: - """Exercise nonzero packed image/video/mixed vision inputs, processor-shaped grids.""" - vision = _create_session(package_dir / "vision_encoder" / "model.onnx", device) - # A video T unit already packs two raw frames. Exercise multiple temporal - # units and unequal spatial grids to cover the dynamic packed-media path. - results = {} - for kind, grid in { - "image": np.array([[1, 4, 4]], dtype=np.int64), - "video": np.array([[2, 4, 4]], dtype=np.int64), - "mixed": np.array([[1, 4, 4], [2, 6, 4]], dtype=np.int64), - }.items(): - patches = int(np.prod(grid, axis=1).sum()) - pixels = np.linspace(0.01, 1.0, patches * 3 * 2 * 16 * 16, dtype=np.float32).reshape( - patches, -1 - ) - output = _numpy( - _run(vision, ["image_features"], {"pixel_values": pixels, "image_grid_thw": grid})[ - 0 - ] - ) - if not np.isfinite(output).all() or not np.any(output): - raise AssertionError(f"{kind} processor-shaped vision output is invalid") - results[kind] = output.shape - return results - - -def _cuda_standard_vision_smoke( - state: dict[str, torch.Tensor], - output_root: Path, - *, - dtype_name: str, -) -> dict[str, tuple[int, ...]]: - """Run media through the standard-attention vision graph on CUDA. - - ORT 1.26's CUDA PackedMultiHeadAttention kernel is nondeterministic for - dynamic packed vision batches. The portable graph still places vision - compute on CUDA and provides stable image/video/mixed runtime evidence. - """ - package = _mobius_package(state, dtype_name=dtype_name, ep="cpu") - package_dir = output_root / f"{dtype_name}-cuda-standard-vision" - package_dir.mkdir(parents=True, exist_ok=True) - package.save(package_dir, external_data="onnx") - return _media_smoke(package_dir, "cuda") - - -def _save_variant( - state: dict[str, torch.Tensor], - output_root: Path, - *, - dtype_name: str, - device: str, -): - package = _mobius_package( - state, dtype_name=dtype_name, ep="cuda" if device == "cuda" else "cpu" - ) - package_dir = output_root / f"{dtype_name}-{device}" - if package_dir.exists() and any(package_dir.iterdir()): - raise FileExistsError(f"Variant output must be empty: {package_dir}") - package_dir.mkdir(parents=True, exist_ok=True) - package.save(package_dir, external_data="onnx") - _save_package_assets(package_dir) - # Package save/load is part of the acceptance boundary, not a graph-only test. - import onnx_ir as ir - - for name in ("decoder", "embedding", "vision_encoder"): - ir.load(package_dir / name / "model.onnx") - return package, package_dir - - -def _validate_variant( - state: dict[str, torch.Tensor], output_root: Path, *, dtype_name: str, device: str -) -> Path: - package, package_dir = _save_variant( - state, - output_root, - dtype_name=dtype_name, - device=device, - ) - prompt = [1, 42, 17] - actual = _onnx_prefill_logits(package_dir, prompt, device) - model = _hf_model(state, dtype=_DTYPES[dtype_name][0], device=device) - expected = _hf_logits(model, prompt, device) - atol = 2e-3 if dtype_name == "f32" else 1e-2 - _assert_logits_close( - actual, - expected, - atol=atol, - label=f"{dtype_name}/{device} full-prefill", - ) - generated, step_logits, profile = run_token_ids( - package_dir, - prompt, - hidden_size=_HIDDEN_SIZE, - max_new_tokens=_GENERATION_TOKENS, - device=device, - profile=device == "cuda", - ) - hf_generated = model.generate( - input_ids=torch.tensor([prompt], device=device), - max_new_tokens=_GENERATION_TOKENS, - do_sample=False, - )[0, len(prompt) :].tolist() - if ( - generated != hf_generated - or len(generated) != _GENERATION_TOKENS - or not all(np.isfinite(x).all() for x in step_logits) - ): - raise AssertionError( - f"Cached generation mismatch: onnx={generated}, hf={hf_generated}" - ) - if profile: - placement = summarize_profile(profile) - Path(profile).unlink(missing_ok=True) - if placement.get("CUDAExecutionProvider", 0) == 0: - raise AssertionError(f"CUDA provider received no decoder nodes: {placement}") - print(f"{dtype_name}/{device} provider placement: {placement}") - print(f"{dtype_name}/{device} generated IDs: {generated}") - print(f"{dtype_name}/{device} graph audit: {_graph_audit(package)}") - if device == "cuda": - media_shapes = _cuda_standard_vision_smoke( - state, - output_root, - dtype_name=dtype_name, - ) - print(f"{dtype_name}/{device} standard-vision media shapes: {media_shapes}") - else: - print(f"{dtype_name}/{device} media shapes: {_media_smoke(package_dir, device)}") - return package_dir - - -def write_goldens(state: dict[str, torch.Tensor], directory: Path) -> None: - """Write L4/L5 only from the independent native HuggingFace reduced model.""" - model = _hf_model(state, dtype=torch.float32, device="cpu") - prompt = [1, 42, 17] - logits = _hf_logits(model, prompt, "cpu")[0, -1] - top10 = np.argsort(logits)[::-1][:10] - generated = model.generate( - input_ids=torch.tensor([prompt]), - max_new_tokens=_GENERATION_TOKENS, - do_sample=False, - )[0, len(prompt) :].tolist() - directory.mkdir(parents=True, exist_ok=True) - (directory / "qwen3_8-27b-reduced.json").write_text( - json.dumps( - { - "model_id": MODEL_ID, - "revision": REVISION, - "fixture_schema": FIXTURE_SCHEMA_VERSION, - "input_ids": prompt, - "top10_ids": top10.tolist(), - "top10_logits": [float(logits[index]).hex() for index in top10], - "logits_summary": [ - float(x).hex() - for x in (logits.max(), logits.min(), logits.mean(), logits.std()) - ], - }, - indent=2, - ), - encoding="utf-8", - ) - (directory / "qwen3_8-27b-reduced_generation.json").write_text( - json.dumps( - { - "model_id": MODEL_ID, - "revision": REVISION, - "fixture_schema": FIXTURE_SCHEMA_VERSION, - "input_ids": prompt, - "max_new_tokens": _GENERATION_TOKENS, - "generated_tokens": generated, - }, - indent=2, - ), - encoding="utf-8", - ) - - -def main() -> None: - parser = argparse.ArgumentParser(description=__doc__) - parser.add_argument("--cache", type=Path, default=default_reduced_cache_path()) - parser.add_argument("--output-dir", type=Path, default=Path("output/qwen3_8-reduced")) - parser.add_argument( - "--matrix", - nargs="+", - choices=["f32-cpu", "f16-cuda", "bf16-cuda"], - default=["f32-cpu", "f16-cuda", "bf16-cuda"], - ) - parser.add_argument("--write-goldens", action="store_true") - parser.add_argument("--skip-quantization", action="store_true") - args = parser.parse_args() - state = _build_reduced_state(args.cache) - print(f"Loaded {len(state)} reduced tensors from {MODEL_ID}@{REVISION}") - if args.write_goldens: - write_goldens( - state, Path(__file__).parents[3] / "testdata" / "golden" / "vision-language" - ) - variants = {} - for variant in args.matrix: - dtype, device = variant.split("-") - if device == "cuda" and not torch.cuda.is_available(): - raise RuntimeError(f"CUDA validation requested but unavailable: {variant}") - if variant == "bf16-cuda": - _package, variants[variant] = _save_variant( - state, - args.output_dir, - dtype_name=dtype, - device=device, - ) - print( - "bf16/cuda export and package reload passed; ORT 1.26 runtime " - "waived for hybrid BF16 provider initialization" - ) - else: - variants[variant] = _validate_variant( - state, - args.output_dir, - dtype_name=dtype, - device=device, - ) - if not args.skip_quantization: - from optimize import quantize_package - - source = variants.get("f16-cuda") - if source is None: - raise ValueError("Q4_K_M validation requires f16-cuda") - quantized = quantize_package(source, args.output_dir / "q4_k_m") - import onnx_ir as ir - - loaded_models = { - name: ir.load(quantized / name / "model.onnx") - for name in ("decoder", "embedding", "vision_encoder") - } - matmul_nbits = sum( - node.domain == "com.microsoft" and node.op_type == "MatMulNBits" - for node in loaded_models["decoder"].graph.all_nodes() - ) - source_bytes = sum(path.stat().st_size for path in source.rglob("*") if path.is_file()) - quantized_bytes = sum( - path.stat().st_size for path in quantized.rglob("*") if path.is_file() - ) - if not matmul_nbits or quantized_bytes >= source_bytes: - raise AssertionError( - "Q4_K_M graph/package audit failed: " - f"MatMulNBits={matmul_nbits}, source={source_bytes}, quantized={quantized_bytes}" - ) - for name in loaded_models: - _create_session(quantized / name / "model.onnx", "cuda") - # Direct ORT sessions, not ORT GenAI capability probing. - ids, logits, _ = run_token_ids( - quantized, - [1, 42, 17], - hidden_size=_HIDDEN_SIZE, - max_new_tokens=_GENERATION_TOKENS, - device="cpu", - ) - if len(ids) != _GENERATION_TOKENS or not all( - np.isfinite(logit).all() for logit in logits - ): - raise AssertionError(f"Q4_K_M direct-session generation failed: {ids}") - print( - "Q4_K_M package audit: " - f"MatMulNBits={matmul_nbits}, source={source_bytes}, quantized={quantized_bytes}" - ) - print(f"Q4_K_M direct CPU-session IDs: {ids}") - - -if __name__ == "__main__": - main() diff --git a/testdata/golden/vision-language/qwen3_8-27b-reduced.json b/testdata/golden/vision-language/qwen3_8-27b-reduced.json deleted file mode 100644 index 9320b6a9d..000000000 --- a/testdata/golden/vision-language/qwen3_8-27b-reduced.json +++ /dev/null @@ -1,40 +0,0 @@ -{ - "model_id": "Qwen/Qwen3.8-27B", - "revision": "1d4bf0f2ff6012fd82039f2fa52739d0dd7c60c0", - "fixture_schema": 1, - "input_ids": [ - 1, - 42, - 17 - ], - "top10_ids": [ - 134, - 34, - 17, - 167, - 38, - 193, - 164, - 103, - 145, - 176 - ], - "top10_logits": [ - "0x1.1704fa0000000p+0", - "0x1.e712dc0000000p-1", - "0x1.ac587e0000000p-1", - "0x1.a32cb60000000p-1", - "0x1.9978040000000p-1", - "0x1.901f0c0000000p-1", - "0x1.89eb5c0000000p-1", - "0x1.7779a00000000p-1", - "0x1.75bb280000000p-1", - "0x1.738e440000000p-1" - ], - "logits_summary": [ - "0x1.1704fa0000000p+0", - "-0x1.567c100000000p+0", - "0x1.75b0380000000p-6", - "0x1.be30360000000p-2" - ] -} \ No newline at end of file diff --git a/testdata/golden/vision-language/qwen3_8-27b-reduced_generation.json b/testdata/golden/vision-language/qwen3_8-27b-reduced_generation.json deleted file mode 100644 index 9290b9632..000000000 --- a/testdata/golden/vision-language/qwen3_8-27b-reduced_generation.json +++ /dev/null @@ -1,33 +0,0 @@ -{ - "model_id": "Qwen/Qwen3.8-27B", - "revision": "1d4bf0f2ff6012fd82039f2fa52739d0dd7c60c0", - "fixture_schema": 1, - "input_ids": [ - 1, - 42, - 17 - ], - "max_new_tokens": 20, - "generated_tokens": [ - 134, - 138, - 167, - 225, - 94, - 226, - 103, - 145, - 114, - 114, - 114, - 254, - 101, - 115, - 254, - 101, - 115, - 254, - 101, - 115 - ] -} \ No newline at end of file diff --git a/tests/qwen38_real_weight_test.py b/tests/qwen38_real_weight_test.py deleted file mode 100644 index 8d5cedb46..000000000 --- a/tests/qwen38_real_weight_test.py +++ /dev/null @@ -1,295 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# Licensed under the MIT License. - -"""Pinned reduced-real Qwen3.8 L4/L5 and Olive recipe coverage.""" - -from __future__ import annotations - -import importlib.util -import json -import os -import sys -import types -from pathlib import Path - -import numpy as np -import pytest - -_ROOT = Path(__file__).parents[1] -_EXAMPLE = _ROOT / "examples" / "olive" / "qwen3_8-27b" -_L4 = _ROOT / "testdata" / "golden" / "vision-language" / "qwen3_8-27b-reduced.json" -_L5 = _ROOT / "testdata" / "golden" / "vision-language" / "qwen3_8-27b-reduced_generation.json" - - -def _load(name: str): - def load_file(module_name: str, path: Path): - spec = importlib.util.spec_from_file_location(module_name, path) - assert spec is not None and spec.loader is not None - module = importlib.util.module_from_spec(spec) - sys.modules[module_name] = module - try: - spec.loader.exec_module(module) - except Exception: - sys.modules.pop(module_name, None) - raise - return module - - previous_inference = sys.modules.get("inference") - try: - inference = load_file("qwen38_inference", _EXAMPLE / "inference.py") - sys.modules["inference"] = inference - return load_file(f"qwen38_{name}", _EXAMPLE / f"{name}.py") - finally: - if previous_inference is None: - sys.modules.pop("inference", None) - else: - sys.modules["inference"] = previous_inference - - -def test_loader_isolates_sibling_inference_modules(monkeypatch): - foreign_inference = types.ModuleType("inference") - monkeypatch.setitem(sys.modules, "inference", foreign_inference) - - validator = _load("validate_reduced_checkpoint") - - assert validator.MODEL_ID == "Qwen/Qwen3.8-27B" - assert sys.modules["inference"] is foreign_inference - - -def test_reduced_config_preserves_hybrid_layers_and_remaps_media_ids(): - validator = _load("validate_reduced_checkpoint") - config = validator._reduced_hf_config() - assert config.text_config.layer_types == [ - "linear_attention", - "linear_attention", - "linear_attention", - "full_attention", - ] - assert config.text_config.mtp_num_hidden_layers == 0 - assert config.vision_config.depth == 1 - assert config.image_token_id < config.text_config.vocab_size - assert config.video_token_id < config.text_config.vocab_size - - -def test_range_reader_retries_and_requires_exact_content_range(monkeypatch): - validator = _load("validate_reduced_checkpoint") - - class Response: - def __init__(self, status, headers, content=b""): - self.status_code, self.headers, self.content = status, headers, content - - def __enter__(self): - return self - - def __exit__(self, *_args): - return False - - class Session: - def __init__(self): - self.responses = [ - Response(503, {}), - Response( - 206, {"Content-Range": "bytes 2-5/10", "Content-Length": "4"}, b"fail" - ), - Response( - 206, {"Content-Range": "bytes 0-3/10", "Content-Length": "4"}, b"pass" - ), - ] - - def get(self, *_args, **_kwargs): - return self.responses.pop(0) - - reader = object.__new__(validator._PinnedSafetensors) - reader._session = Session() - monkeypatch.setattr(validator.time, "sleep", lambda _seconds: None) - assert reader._range("shard", 0, 3) == b"pass" - - -def test_olive_recipe_is_q4_k_m_cpu_weight_only(tmp_path): - optimize = _load("optimize") - recipe = optimize.olive_config(tmp_path / "decoder.onnx", tmp_path / "out") - config = recipe["passes"]["q4_k_m"] - assert config["type"] == "OnnxKQuantQuantization" - assert config["bits"] == 4 - assert recipe["clean_cache"] is True - assert Path(recipe["cache_dir"]).name == ".olive-cache" - assert recipe["engine"]["target"]["accelerators"][0]["execution_providers"] == [ - "CPUExecutionProvider" - ] - - -def _require_real_fixture(): - if os.environ.get("MOBIUS_QWEN38_REDUCED_REAL") != "1": - pytest.skip("Set MOBIUS_QWEN38_REDUCED_REAL=1 to enable pinned range fixture") - - -@pytest.mark.integration -@pytest.mark.integration_slow -@pytest.mark.golden -def test_qwen38_reduced_real_l4(tmp_path): - _require_real_fixture() - validator = _load("validate_reduced_checkpoint") - state = validator._build_reduced_state(validator.default_reduced_cache_path()) - package_dir = tmp_path / "qwen38-test-output" - package_dir.mkdir(exist_ok=True) - source = validator._validate_variant(state, package_dir, dtype_name="f32", device="cpu") - golden = json.loads(_L4.read_text(encoding="utf-8")) - logits = validator._onnx_prefill_logits(source, golden["input_ids"], "cpu")[0, -1] - actual = np.argsort(logits)[::-1][:10].tolist() - assert actual == golden["top10_ids"] - np.testing.assert_allclose( - logits[actual], [float.fromhex(value) for value in golden["top10_logits"]], atol=2e-3 - ) - - -@pytest.mark.integration -@pytest.mark.integration_slow -@pytest.mark.generation -def test_qwen38_reduced_real_l5(tmp_path): - _require_real_fixture() - validator = _load("validate_reduced_checkpoint") - golden = json.loads(_L5.read_text(encoding="utf-8")) - state = validator._build_reduced_state(validator.default_reduced_cache_path()) - package_dir = validator._validate_variant( - state, - tmp_path / "qwen38-test-output", - dtype_name="f32", - device="cpu", - ) - generated, _logits, _profile = validator.run_token_ids( - package_dir, - golden["input_ids"], - hidden_size=256, - max_new_tokens=golden["max_new_tokens"], - device="cpu", - ) - assert generated == golden["generated_tokens"] - - -@pytest.mark.integration -@pytest.mark.integration_slow -def test_qwen38_reduced_real_bf16_export_and_reload(tmp_path): - _require_real_fixture() - import onnx_ir as ir - - validator = _load("validate_reduced_checkpoint") - state = validator._build_reduced_state(validator.default_reduced_cache_path()) - package, package_dir = validator._save_variant( - state, - tmp_path / "qwen38-test-output", - dtype_name="bf16", - device="cuda", - ) - for model in package.values(): - assert not [ - name - for name, initializer in model.graph.initializers.items() - if initializer.const_value is None - ] - assert any( - initializer.dtype == ir.DataType.BFLOAT16 - for initializer in model.graph.initializers.values() - ) - vision = ir.load(package_dir / "vision_encoder" / "model.onnx") - assert vision.graph.inputs[0].name == "pixel_values" - assert vision.graph.inputs[0].dtype == ir.DataType.FLOAT - - -@pytest.mark.integration -@pytest.mark.integration_slow -def test_qwen38_reduced_real_f16_cuda(tmp_path): - _require_real_fixture() - if os.environ.get("MOBIUS_TEST_DEVICE") != "cuda": - pytest.skip("Set MOBIUS_TEST_DEVICE=cuda for CUDA parity") - import torch - - if not torch.cuda.is_available(): - pytest.skip("CUDA is unavailable") - validator = _load("validate_reduced_checkpoint") - state = validator._build_reduced_state(validator.default_reduced_cache_path()) - validator._validate_variant( - state, - tmp_path / "qwen38-test-output", - dtype_name="f16", - device="cuda", - ) - - -@pytest.mark.integration -@pytest.mark.integration_slow -@pytest.mark.quantization -def test_qwen38_reduced_olive_q4_package(tmp_path): - _require_real_fixture() - if os.environ.get("MOBIUS_TEST_DEVICE") != "cuda": - pytest.skip("Set MOBIUS_TEST_DEVICE=cuda for Q4_K_M validation") - validator = _load("validate_reduced_checkpoint") - optimize = _load("optimize") - state = validator._build_reduced_state(validator.default_reduced_cache_path()) - output_root = tmp_path / "qwen38-test-output" - source = validator._validate_variant( - state, - output_root, - dtype_name="f16", - device="cuda", - ) - processor_config = source / "processor_config.json" - assert processor_config.is_file() - result = optimize.quantize_package(source, output_root / "q4_k_m") - import onnx_ir as ir - - decoder = ir.load(result / "decoder" / "model.onnx") - assert any( - node.domain == "com.microsoft" and node.op_type == "MatMulNBits" - for node in decoder.graph.all_nodes() - ) - assert sum(path.stat().st_size for path in result.rglob("*") if path.is_file()) < sum( - path.stat().st_size for path in source.rglob("*") if path.is_file() - ) - assert (result / "processor_config.json").read_bytes() == processor_config.read_bytes() - from transformers import Qwen2VLImageProcessor, Qwen3VLVideoProcessor - - processor_data = json.loads((result / "processor_config.json").read_text()) - assert type( - Qwen2VLImageProcessor.from_dict(processor_data["image_processor"]) - ).__name__ == ("Qwen2VLImageProcessor") - assert type( - Qwen3VLVideoProcessor.from_dict(processor_data["video_processor"]) - ).__name__ == ("Qwen3VLVideoProcessor") - assert not (result / "tokenizer.json").exists() - assert not (result / "generation_config.json").exists() - manifest = json.loads((result / "source_manifest.json").read_text()) - assert manifest["text_input_contract"] == "token-ids-only" - assert manifest["runtime"] == "onnxruntime-direct" - assert manifest["quantization"] == "Q4_K_M" - for name in ("decoder", "embedding", "vision_encoder"): - validator._create_session(result / name / "model.onnx", "cuda") - ids, logits, _ = validator.run_token_ids( - result, - [1, 42, 17], - hidden_size=256, - max_new_tokens=20, - device="cpu", - ) - assert ids == [ - 134, - 244, - 242, - 167, - 81, - 34, - 155, - 251, - 142, - 90, - 224, - 31, - 76, - 250, - 240, - 120, - 134, - 120, - 134, - 120, - ] - assert all(np.isfinite(value).all() for value in logits) From 08360a40fc96c63b7ff9752998a377bc4dbb355c Mon Sep 17 00:00:00 2001 From: Justin Chu Date: Wed, 19 Aug 2026 16:38:32 -0700 Subject: [PATCH 14/15] Preserve canonical Qwen3.5 test model Keep the qwen3_5 registry test surface pinned to Qwen/Qwen3.5-2B. Qwen3.8 remains covered through its independent pinned configuration and case data instead of replacing canonical Qwen3.5 coverage. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Signed-off-by: Justin Chu --- src/mobius/_registry.py | 2 +- src/mobius/models/qwen35_test.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/src/mobius/_registry.py b/src/mobius/_registry.py index 033f4dec7..698beed1c 100644 --- a/src/mobius/_registry.py +++ b/src/mobius/_registry.py @@ -1020,7 +1020,7 @@ def _create_default_registry() -> ModelRegistry: "qwen2_5_vl_text": "Qwen/Qwen2.5-VL-3B-Instruct", "qwen3_vl": "Qwen/Qwen3-VL-2B-Instruct", "qwen3_vl_text": "Qwen/Qwen3-VL-2B-Instruct", - "qwen3_5": "Qwen/Qwen3.8-27B", + "qwen3_5": "Qwen/Qwen3.5-2B", "llava": "llava-hf/llava-1.5-7b-hf", "llava_next": "llava-hf/llava-v1.6-mistral-7b-hf", "mllama": "meta-llama/Llama-3.2-11B-Vision-Instruct", diff --git a/src/mobius/models/qwen35_test.py b/src/mobius/models/qwen35_test.py index 9f5c897ec..33cfe7914 100644 --- a/src/mobius/models/qwen35_test.py +++ b/src/mobius/models/qwen35_test.py @@ -104,7 +104,7 @@ def test_exact_config_extracts_dense_hybrid_vl_architecture(self): assert _QWEN38_REVISION == "1d4bf0f2ff6012fd82039f2fa52739d0dd7c60c0" assert registry.get("qwen3_5") is Qwen35VL3ModelCausalLMModel assert registry.get("qwen3_5_vl") is Qwen35VL3ModelCausalLMModel - assert registry.get_registration("qwen3_5").test_model_id == "Qwen/Qwen3.8-27B" + assert registry.get_registration("qwen3_5").test_model_id == "Qwen/Qwen3.5-2B" assert registry.get_registration("qwen3_5_vl").test_model_id == "Qwen/Qwen3.5-2B" assert config.hidden_size == 5120 assert config.intermediate_size == 17408 From 59af1a51f1a162e9444a9be1ea5785ca4ec0ef0b Mon Sep 17 00:00:00 2001 From: Justin Chu Date: Wed, 19 Aug 2026 16:41:37 -0700 Subject: [PATCH 15/15] Correct Qwen3.8 CI waiver State that the pinned 55.6 GB checkpoint exceeds hosted CI resources and that no reduced or quantized private fixture is committed. This replaces the obsolete claim that removed reduced-real and Olive assets are run manually. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Signed-off-by: Justin Chu --- testdata/cases/vision-language/qwen3_8-27b.yaml | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/testdata/cases/vision-language/qwen3_8-27b.yaml b/testdata/cases/vision-language/qwen3_8-27b.yaml index cbc308c47..1b370cc82 100644 --- a/testdata/cases/vision-language/qwen3_8-27b.yaml +++ b/testdata/cases/vision-language/qwen3_8-27b.yaml @@ -16,7 +16,9 @@ generation: max_new_tokens: 30 do_sample: false -ci_skip_reason: "Official checkpoint is 55.6 GB; reduced-real CUDA and Olive evidence is run manually." +ci_skip_reason: >- + The pinned official checkpoint is 55.6 GB and exceeds hosted CI storage and + GPU memory; no reduced or quantized private fixture is committed. notes: >- Qwen3.8-27B native image/video model. Dense Qwen3.5 alias with 64 hybrid layers (48 Gated DeltaNet + 16 gated GQA), a 27-block vision encoder, and an