From 5ac6fae8ed950a5ceda48cb925bde844ab481429 Mon Sep 17 00:00:00 2001 From: Justin Chu Date: Thu, 13 Aug 2026 20:31:50 -0700 Subject: [PATCH 01/13] Add Nemotron 3.5 Lightning export support Normalize current and legacy NemotronH configs, preserve routing and recurrent-state precision, and filter only auxiliary MTP weights. Add pinned reduced-real L4/L5 evidence, CUDA/ORT GenAI guards, and an executable Olive Q4 direct-runtime recipe. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Signed-off-by: Justin Chu --- .agents/skills/moe-models/SKILL.md | 19 + .agents/skills/ort-genai-config/SKILL.md | 6 + .github/workflows/gpu_l4_golden_parity.yml | 4 +- .github/workflows/gpu_l5_generation_e2e.yml | 4 +- .../nemotron-3_5-lightning-30b/.gitignore | 2 + .../nemotron-3_5-lightning-30b/README.md | 153 ++++++ .../nemotron-3_5-lightning-30b/inference.py | 286 +++++++++++ .../nemotron-3_5-lightning-30b/olive_q4.json | 32 ++ .../nemotron-3_5-lightning-30b/optimize.py | 240 +++++++++ .../requirements.txt | 3 + .../validate_reduced_checkpoint.py | 480 ++++++++++++++++++ src/mobius/__main__.py | 23 + src/mobius/_configs/_base.py | 35 +- src/mobius/_configs_test.py | 55 ++ src/mobius/_registry.py | 2 +- .../integrations/ort_genai/auto_export.py | 13 + .../ort_genai/auto_export_test.py | 17 + src/mobius/models/nemotron_h.py | 34 +- src/mobius/tasks/_cache_utils.py | 2 +- .../causal-lm/nemotron-3_5-lightning-30b.yaml | 20 + .../nemotron-3_5-lightning-30b-reduced.json | 43 ++ ...-3_5-lightning-30b-reduced_generation.json | 18 + tests/_test_configs.py | 2 + tests/arch_validation_test.py | 44 +- tests/build_graph_test.py | 55 ++ tests/cli_test.py | 23 + tests/model_coverage_test.py | 1 - tests/nemotron_h_real_weight_test.py | 117 +++++ tests/synthetic_parity_test.py | 18 +- tests/weight_alignment_test.py | 52 ++ 30 files changed, 1776 insertions(+), 27 deletions(-) create mode 100644 examples/olive/nemotron-3_5-lightning-30b/.gitignore create mode 100644 examples/olive/nemotron-3_5-lightning-30b/README.md create mode 100644 examples/olive/nemotron-3_5-lightning-30b/inference.py create mode 100644 examples/olive/nemotron-3_5-lightning-30b/olive_q4.json create mode 100644 examples/olive/nemotron-3_5-lightning-30b/optimize.py create mode 100644 examples/olive/nemotron-3_5-lightning-30b/requirements.txt create mode 100644 examples/olive/nemotron-3_5-lightning-30b/validate_reduced_checkpoint.py create mode 100644 testdata/cases/causal-lm/nemotron-3_5-lightning-30b.yaml create mode 100644 testdata/golden/causal-lm/nemotron-3_5-lightning-30b-reduced.json create mode 100644 testdata/golden/causal-lm/nemotron-3_5-lightning-30b-reduced_generation.json create mode 100644 tests/nemotron_h_real_weight_test.py diff --git a/.agents/skills/moe-models/SKILL.md b/.agents/skills/moe-models/SKILL.md index 76622ab96..aed597972 100644 --- a/.agents/skills/moe-models/SKILL.md +++ b/.agents/skills/moe-models/SKILL.md @@ -476,6 +476,25 @@ config.norm_topk_prob # Whether to normalize routing weights config.routed_scaling_factor # Post-normalization scale ``` +Current Transformers exposes NemotronH layer types as +`linear_attention` / `full_attention`; older configs use +`mamba` / `attention`. Normalize both vocabularies to Mobius +`mamba2` / `full_attention`, preserve `mlp` and `moe` distinctly, and reject +unknown values. Never map `mlp` to `moe` in parity fixtures. + +### Reduced-precision routing + +ONNX has no implicit mixed-float type promotion. NemotronH routing computes in +fp32, so keep the correction-bias initializer in fp32 and explicitly cast the +gate weight to fp32. Cast each expert output up, multiply and accumulate all +routed contributions in fp32, then cast the completed routed tensor back once. +Graph construction alone may miss this; execute fp16 and bf16 MoE paths. + +Official Nemotron 3.5 checkpoints also contain auxiliary `mtp.*` tensors. +The base `NemotronHForCausalLM` generation graph does not instantiate them and +marks them unexpected. Filter only that prefix and prove weight alignment still +populates every base-decoder initializer. + ### com.microsoft.MoE compatibility **Not compatible with NemotronH.** Three blockers: diff --git a/.agents/skills/ort-genai-config/SKILL.md b/.agents/skills/ort-genai-config/SKILL.md index 6ab46aa10..d644c6a7b 100644 --- a/.agents/skills/ort-genai-config/SKILL.md +++ b/.agents/skills/ort-genai-config/SKILL.md @@ -291,6 +291,12 @@ global cache-slot count. Preserve intrinsic schema/config validation, but do not gate or reject export based on the current GenAI model registry, runtime version, topology support, or cache executor capability. +NemotronH is one such unsupported contract in ORT GenAI 0.15.2: sparse global +layer indices mix attention key/value caches with Mamba `conv_state` and +`ssm_state`. The generic decoder schema cannot bind all three cache families. +Reject `--runtime ort-genai` before creating output and direct users to a +token-by-token ONNX Runtime loop until the runtime adds a dedicated model type. + ### "input_ids size exceeds max length" For image prompts, the tokenized input_ids (including image_pad tokens) can diff --git a/.github/workflows/gpu_l4_golden_parity.yml b/.github/workflows/gpu_l4_golden_parity.yml index 91dfb2c8e..d05499c50 100644 --- a/.github/workflows/gpu_l4_golden_parity.yml +++ b/.github/workflows/gpu_l4_golden_parity.yml @@ -86,7 +86,7 @@ jobs: AFFECTED='${{ inputs.affected_models }}' if [ -z "$AFFECTED" ] || [ "$AFFECTED" = "[]" ]; then echo "Running all L4 golden comparison tests" - pytest tests/e2e_golden_test.py \ + pytest tests/e2e_golden_test.py tests/nemotron_h_real_weight_test.py \ -m golden \ -v \ --timeout=300 \ @@ -98,7 +98,7 @@ jobs: # Convert JSON array to comma-separated list for --models MODELS=$(echo "$AFFECTED" | python -c "import json, sys; print(','.join(json.load(sys.stdin)))") if [ -n "$MODELS" ]; then - pytest tests/e2e_golden_test.py \ + pytest tests/e2e_golden_test.py tests/nemotron_h_real_weight_test.py \ -m golden \ -v \ --models "$MODELS" \ diff --git a/.github/workflows/gpu_l5_generation_e2e.yml b/.github/workflows/gpu_l5_generation_e2e.yml index 23808927b..889de86d1 100644 --- a/.github/workflows/gpu_l5_generation_e2e.yml +++ b/.github/workflows/gpu_l5_generation_e2e.yml @@ -74,7 +74,7 @@ jobs: AFFECTED='${{ inputs.affected_models }}' if [ -z "$AFFECTED" ] || [ "$AFFECTED" = "[]" ]; then echo "Running all L5 generation E2E tests" - pytest tests/e2e_golden_test.py \ + pytest tests/e2e_golden_test.py tests/nemotron_h_real_weight_test.py \ -m generation \ -v \ --timeout=300 \ @@ -86,7 +86,7 @@ jobs: # Convert JSON array to comma-separated list for --models MODELS=$(echo "$AFFECTED" | python -c "import json, sys; print(','.join(json.load(sys.stdin)))") if [ -n "$MODELS" ]; then - pytest tests/e2e_golden_test.py \ + pytest tests/e2e_golden_test.py tests/nemotron_h_real_weight_test.py \ -m generation \ -v \ --models "$MODELS" \ diff --git a/examples/olive/nemotron-3_5-lightning-30b/.gitignore b/examples/olive/nemotron-3_5-lightning-30b/.gitignore new file mode 100644 index 000000000..ab3bb1740 --- /dev/null +++ b/examples/olive/nemotron-3_5-lightning-30b/.gitignore @@ -0,0 +1,2 @@ +cache/ +output/ diff --git a/examples/olive/nemotron-3_5-lightning-30b/README.md b/examples/olive/nemotron-3_5-lightning-30b/README.md new file mode 100644 index 000000000..78ab3c8e7 --- /dev/null +++ b/examples/olive/nemotron-3_5-lightning-30b/README.md @@ -0,0 +1,153 @@ +# Nemotron 3.5 Lightning: BF16 checkpoint + Olive + +This is **Option A** for +[`nvidia/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-BF16`](https://huggingface.co/nvidia/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-BF16): +export the official BF16 checkpoint to supported FP16 ONNX, quantize the model +with Olive, assemble a direct ONNX Runtime package, then run cached generation. + +Every Hub access is pinned to revision +`d468880b6ad3c6e0d21377ce7242adaea4cc884d`. + +## Architecture and runtime contract + +The checkpoint is a real `nemotron_h` model, not an alias: + +- 52 base-decoder layers mixing Mamba2, sigmoid-routed MoE, and full GQA. +- 128 routed experts with top-6 selection and one shared expert. +- `mtp.*` contains 270 auxiliary multi-token-prediction tensors. This export + intentionally targets the base `NemotronHForCausalLM` decoder; its forward + graph does not instantiate MTP, and upstream marks those keys unexpected. + No base-decoder generation input, cache, logit, or weight depends on them. + +ONNX Runtime GenAI 0.15.2 cannot bind this model's mixed cache: Mamba layers +need `conv_state` plus `ssm_state`, while sparse full-attention layers need +key/value caches at global layer indices. Mobius therefore rejects +`--runtime ort-genai` before writing artifacts. The package uses direct ONNX +Runtime generation through `inference.py`. + +## Install + +From the repository root: + +```powershell +python -m pip install -e ".[transformers,testing]" ` + --index-url https://packagefeedproxy.microsoft.io/pypi/simple +python -m pip install -r examples\olive\nemotron-3_5-lightning-30b\requirements.txt ` + --index-url https://packagefeedproxy.microsoft.io/pypi/simple +``` + +Use an ONNX Runtime GPU build with CUDA 12 and cuDNN 9 for CUDA inference. + +## Full export, quantization, and smoke test + +```powershell +cd examples\olive\nemotron-3_5-lightning-30b +python optimize.py ` + --source-dir output\f16\cuda ` + --output-dir output\Q4_K_M\cuda ` + --ep cuda ` + --precision q4_k_m +``` + +The script performs four gated steps: + +1. Downloads the exact 14-shard BF16 checkpoint revision and exports FP16 ONNX. + BF16 execution is rejected explicitly because corrected reduced-real parity + reaches `0.8594` max logit error, above the `1e-2` reduced-precision gate. +2. Applies CUDA GQA/LinearAttention fusion and the grouped-RMSNorm workaround. +3. Runs Olive Q4 K-quant with an explicitly CPU-only target. It also suppresses + Olive 0.13's unrelated GPU-EP DLL auto-registration, so a missing TensorRT + installation cannot abort CPU weight-only quantization. +4. Reloads the assembled quantized package and generates four cached tokens. + +To reuse an existing source package: + +```powershell +python optimize.py --skip-export ` + --source-dir output\f16\cuda ` + --output-dir output\Q4_K_M\cuda ` + --ep cuda +``` + +`olive_q4.json` records the equivalent pass and provider-isolated target. Use +`optimize.py` rather than invoking the JSON directly when the installed ORT +wheel bundles unconfigured providers; the script contains the verified Olive +0.13 registration isolation. + +## Package layout + +```text +output/ +├── f16/cuda/ +│ ├── model.onnx +│ ├── model.onnx.data +│ ├── config.json +│ ├── generation_config.json +│ ├── tokenizer.json +│ ├── tokenizer_config.json +│ └── source_manifest.json +└── Q4_K_M/cuda/ + ├── model.onnx + ├── model.onnx.data + ├── config.json + ├── generation_config.json + ├── tokenizer.json + ├── tokenizer_config.json + └── source_manifest.json +``` + +No `genai_config.json` is emitted because that would claim unsupported ORT +GenAI runtime compatibility. + +## Direct generation and profiling + +```powershell +python inference.py ` + --model-dir output\Q4_K_M\cuda ` + --device cuda ` + --prompt "What is 84 * 3 / 2?" ` + --max-new-tokens 20 ` + --profile +``` + +The script initializes every cache from the saved graph, processes the prompt +token by token, carries Mamba and KV state independently, and fails if CUDA was +requested but not registered. + +## Reduced real-checkpoint validation + +The full checkpoint is 65.8 GB and cannot execute on the validation host's +8 GB RTX A1000. The reproducible reduced check range-downloads 236 MiB of real +weights while retaining production dimensions: + +- checkpoint layer 0: complete Mamba2 block; +- layer 1: router, shared expert, and four complete routed experts; +- layer 5: complete full-attention block; +- sliced real embedding and LM-head rows plus final norm. + +```powershell +python validate_reduced_checkpoint.py +``` + +Validated results on ORT 1.28.0 / Olive 0.13.0: + +| Variant | Full-logit max abs | Generated IDs | Placement | +|---|---:|---|---| +| FP32 CPU | `9.54e-6` | `12, 13, 12, 12` | CPU | +| FP16 CUDA | `0.00977` | `12, 13, 12, 12` | 833 CUDA / 14 CPU events | +| BF16 CUDA | rejected (`0.8594`) | N/A | fails numerical gate | +| Olive Q4 | quantized | `12, 13, 12, 12` | 833 CUDA / 14 CPU events | + +The reduced FP16 package is 247,380,256 bytes; Q4 is 73,190,965 bytes +(`0.296x`). Its weighted graph contains 15 `com.microsoft::MatMulNBits` +nodes and reloads successfully for multi-token generation. + +## Evidence-based waivers + +- Full-checkpoint L4/L5 coherent-text generation: requires roughly 66 GB just + for checkpoint storage and substantially more than 8 GB accelerator memory. +- Full 30B Olive run: the recipe and reduced production-dimension pass are + validated; completing all 2,944 expert subgraphs requires a large-memory + host. +- Foundry Local: not available on this host, and its ORT GenAI-based model + contract cannot represent NemotronH hybrid state today. diff --git a/examples/olive/nemotron-3_5-lightning-30b/inference.py b/examples/olive/nemotron-3_5-lightning-30b/inference.py new file mode 100644 index 000000000..e5bc0f105 --- /dev/null +++ b/examples/olive/nemotron-3_5-lightning-30b/inference.py @@ -0,0 +1,286 @@ +#!/usr/bin/env python +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +"""Direct ONNX Runtime generation for NemotronH hybrid-cache 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 = "nvidia/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-BF16" +REVISION = "d468880b6ad3c6e0d21377ce7242adaea4cc884d" +_BFLOAT16_ONNX_TYPE = 16 + + +def _numpy_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 state input type: {ort_type}") + + +def _concrete_state_shape(shape: list[Any]) -> tuple[int, ...]: + concrete: list[int] = [] + for dim in shape: + if isinstance(dim, int): + concrete.append(dim) + elif "batch" in str(dim): + concrete.append(1) + elif "past" in str(dim): + concrete.append(0) + else: + raise ValueError(f"Cannot resolve hybrid-cache dimension {dim!r}") + return tuple(concrete) + + +def _initial_states(session) -> dict[str, Any]: + import onnxruntime as ort + + states: dict[str, Any] = {} + for model_input in session.get_inputs(): + if not model_input.name.startswith("past_key_values."): + continue + shape = _concrete_state_shape(model_input.shape) + if model_input.type == "tensor(bfloat16)": + states[model_input.name] = ort.OrtValue.ortvalue_from_numpy_with_onnx_type( + np.zeros(shape, dtype=np.uint16), + _BFLOAT16_ONNX_TYPE, + ) + else: + states[model_input.name] = np.zeros( + shape, + dtype=_numpy_dtype(model_input.type), + ) + return states + + +def _update_states( + states: dict[str, Any], + output_names: list[str], + output_values: list[Any], +) -> None: + for name, value in zip(output_names, output_values): + if not name.startswith("present."): + continue + input_name = name.replace("present.", "past_key_values.", 1) + if input_name in states: + states[input_name] = value + + +def _token_feeds( + session, + token_ids: np.ndarray, + *, + total_length: int, + position_ids: np.ndarray, + states: dict[str, Any], +) -> dict[str, Any]: + available = {model_input.name for model_input in session.get_inputs()} + candidates = { + "input_ids": token_ids, + "attention_mask": np.ones((1, total_length), dtype=np.int64), + "position_ids": position_ids, + **states, + } + return {name: value for name, value in candidates.items() if name in available} + + +def _run_session(session, output_names: list[str], feeds: dict[str, Any]) -> list[Any]: + import onnxruntime as ort + + if not any(isinstance(value, ort.OrtValue) for value in feeds.values()): + return session.run(output_names, feeds) + ort_feeds = { + 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, ort_feeds)) + + +def _as_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 _create_session(model_path: Path, device: str, profile: bool): + if device == "cuda": + # Importing PyTorch first preloads its matching CUDA/cuDNN DLLs on Windows. + import torch # noqa: F401 + + 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), + sess_options=options, + providers=providers, + ) + if device == "cuda" and session.get_providers()[0] != "CUDAExecutionProvider": + raise RuntimeError( + f"CUDAExecutionProvider was requested but providers are {session.get_providers()}" + ) + return session + + +def run_token_ids( + model_dir: str | Path, + input_ids: list[int], + *, + max_new_tokens: int, + device: str, + profile: bool = False, +) -> tuple[list[int], list[np.ndarray], str | None]: + """Run token-by-token hybrid-cache generation and return IDs plus logits.""" + model_path = Path(model_dir) / "model.onnx" + if not model_path.is_file(): + raise FileNotFoundError(f"Missing ONNX model: {model_path}") + if not input_ids: + raise ValueError("input_ids must not be empty") + + session = _create_session(model_path, device, profile) + states = _initial_states(session) + output_names = [output.name for output in session.get_outputs()] + generated: list[int] = [] + logits_by_step: list[np.ndarray] = [] + past_length = 0 + outputs: list[Any] | None = None + + for token_id in input_ids: + feeds = _token_feeds( + session, + np.array([[token_id]], dtype=np.int64), + total_length=past_length + 1, + position_ids=np.array([[past_length]], dtype=np.int64), + states=states, + ) + outputs = _run_session(session, output_names, feeds) + _update_states(states, output_names, outputs) + past_length += 1 + + assert outputs is not None + logits = _as_numpy(outputs[output_names.index("logits")])[0, -1].astype(np.float32) + for _ in range(max_new_tokens): + logits_by_step.append(logits.copy()) + token_id = int(np.argmax(logits)) + generated.append(token_id) + feeds = _token_feeds( + session, + np.array([[token_id]], dtype=np.int64), + total_length=past_length + 1, + position_ids=np.array([[past_length]], dtype=np.int64), + states=states, + ) + outputs = _run_session(session, output_names, feeds) + _update_states(states, output_names, outputs) + past_length += 1 + logits = _as_numpy(outputs[output_names.index("logits")])[0, -1].astype(np.float32) + + profile_path = session.end_profiling() if profile else None + return generated, logits_by_step, profile_path + + +def summarize_profile(profile_path: str) -> dict[str, int]: + """Summarize actual node placement from an ORT profiling JSON file.""" + events = json.loads(Path(profile_path).read_text(encoding="utf-8")) + providers: Counter[str] = Counter() + for event in events: + args = event.get("args", {}) + provider = args.get("provider") + if event.get("cat") == "Node" and provider: + providers[str(provider)] += 1 + return dict(sorted(providers.items())) + + +def _tokenize_prompt(model_dir: Path, prompt: str, use_chat: bool): + from transformers import AutoTokenizer + + tokenizer = AutoTokenizer.from_pretrained( + model_dir, + revision=None, + local_files_only=True, + ) + if use_chat: + ids = tokenizer.apply_chat_template( + [{"role": "user", "content": prompt}], + tokenize=True, + add_generation_prompt=True, + ) + else: + ids = tokenizer.encode(prompt) + return tokenizer, list(ids) + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--model-dir", required=True) + parser.add_argument("--prompt", default="What is 84 * 3 / 2?") + parser.add_argument("--token-ids", nargs="+", type=int) + parser.add_argument("--max-new-tokens", type=int, default=20) + parser.add_argument("--device", choices=["cpu", "cuda"], default="cuda") + parser.add_argument("--no-chat", action="store_true") + parser.add_argument("--profile", action="store_true") + args = parser.parse_args() + + model_dir = Path(args.model_dir) + tokenizer = None + if args.token_ids: + input_ids = args.token_ids + else: + tokenizer, input_ids = _tokenize_prompt(model_dir, args.prompt, not args.no_chat) + + generated, _logits, profile_path = run_token_ids( + model_dir, + input_ids, + max_new_tokens=args.max_new_tokens, + device=args.device, + profile=args.profile, + ) + if not generated: + raise RuntimeError("Generation produced no tokens") + + if tokenizer is None: + print("Generated token IDs:", generated) + else: + text = tokenizer.decode(generated, skip_special_tokens=True) + if not text.strip(): + raise RuntimeError("Generation produced only empty/special-token text") + print(text) + + if profile_path is not None: + print(f"ORT profile: {profile_path}") + print("Node placement:", summarize_profile(profile_path)) + + +if __name__ == "__main__": + main() diff --git a/examples/olive/nemotron-3_5-lightning-30b/olive_q4.json b/examples/olive/nemotron-3_5-lightning-30b/olive_q4.json new file mode 100644 index 000000000..e4d316cde --- /dev/null +++ b/examples/olive/nemotron-3_5-lightning-30b/olive_q4.json @@ -0,0 +1,32 @@ +{ + "input_model": { + "type": "OnnxModel", + "model_path": "output/f16/cuda/model.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": "model.onnx.data", + "size_threshold": 1024 + } + }, + "engine": { + "target": { + "type": "LocalSystem", + "accelerators": [ + { + "device": "cpu", + "execution_providers": [ + "CPUExecutionProvider" + ] + } + ] + } + }, + "no_artifacts": true, + "output_dir": "output/Q4_K_M/cuda" +} diff --git a/examples/olive/nemotron-3_5-lightning-30b/optimize.py b/examples/olive/nemotron-3_5-lightning-30b/optimize.py new file mode 100644 index 000000000..e5bfd7236 --- /dev/null +++ b/examples/olive/nemotron-3_5-lightning-30b/optimize.py @@ -0,0 +1,240 @@ +#!/usr/bin/env python +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +"""Pinned BF16 export and Olive INT4 packaging for Nemotron 3.5 Lightning.""" + +from __future__ import annotations + +import argparse +import json +import shutil +import tempfile +from pathlib import Path + +from inference import MODEL_ID, REVISION, run_token_ids + +_METADATA_FILES = { + "added_tokens.json", + "chat_template.jinja", + "config.json", + "generation_config.json", + "merges.txt", + "special_tokens_map.json", + "tokenizer.json", + "tokenizer_config.json", + "tokenizer.model", + "vocab.json", +} + + +def _require_empty_output(path: Path) -> None: + if path.exists() and any(path.iterdir()): + raise FileExistsError(f"Output directory must be empty: {path}") + path.mkdir(parents=True, exist_ok=True) + + +def _save_pinned_metadata(output_dir: Path) -> None: + from transformers import AutoConfig, AutoTokenizer, GenerationConfig + + config = AutoConfig.from_pretrained( + MODEL_ID, + revision=REVISION, + trust_remote_code=False, + ) + config.save_pretrained(output_dir) + tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, revision=REVISION) + tokenizer.save_pretrained(output_dir) + generation = GenerationConfig.from_pretrained(MODEL_ID, revision=REVISION) + generation.save_pretrained(output_dir) + + (output_dir / "source_manifest.json").write_text( + json.dumps( + { + "model_id": MODEL_ID, + "revision": REVISION, + "runtime": "onnxruntime-direct", + "ort_genai_supported": False, + }, + indent=2, + ), + encoding="utf-8", + ) + + +def export_checkpoint(output_dir: str | Path, *, ep: str) -> Path: + """Export the pinned BF16 checkpoint as a supported FP16 ONNX package.""" + from mobius import build + from mobius._flags import override_flags + + output = Path(output_dir) + _require_empty_output(output) + with override_flags(ort_cuda_grouped_rmsnorm_workaround=ep == "cuda"): + package = build( + MODEL_ID, + revision=REVISION, + dtype="f16", + load_weights=True, + trust_remote_code=False, + execution_provider=ep, + ) + package.save(output, external_data="onnx") + _save_pinned_metadata(output) + manifest_path = output / "source_manifest.json" + manifest = json.loads(manifest_path.read_text(encoding="utf-8")) + manifest.update({"source_dtype": "bf16", "dtype": "f16", "target_ep": ep}) + manifest_path.write_text(json.dumps(manifest, indent=2), encoding="utf-8") + return output + + +def _olive_config(source_model: Path, output_dir: Path, precision: str) -> dict: + if precision == "q4_k_m": + pass_config = { + "type": "OnnxKQuantQuantization", + "bits": 4, + "block_size": 32, + } + elif precision == "nf4": + pass_config = { + "type": "OnnxBnb4Quantization", + "precision": "nf4", + } + else: + raise ValueError(f"Unsupported quantization precision: {precision}") + + pass_config.update( + { + "save_as_external_data": True, + "all_tensors_to_one_file": True, + "external_data_name": "model.onnx.data", + "size_threshold": 1024, + } + ) + return { + "input_model": { + "type": "OnnxModel", + "model_path": str(source_model), + }, + "passes": {precision: pass_config}, + "engine": { + "target": { + "type": "LocalSystem", + "accelerators": [ + { + "device": "cpu", + "execution_providers": ["CPUExecutionProvider"], + } + ], + } + }, + "no_artifacts": True, + "output_dir": str(output_dir), + } + + +def _find_olive_model(output_dir: Path) -> Path: + candidates = list(output_dir.rglob("*.onnx")) + if len(candidates) != 1: + raise RuntimeError( + f"Expected exactly one Olive ONNX output under {output_dir}, got {candidates}" + ) + return candidates[0] + + +def _copy_olive_model(model_path: Path, destination: Path) -> None: + for child in model_path.parent.iterdir(): + if child.is_file(): + shutil.copy2(child, destination / child.name) + copied_model = destination / model_path.name + canonical_model = destination / "model.onnx" + if copied_model != canonical_model: + copied_model.replace(canonical_model) + + +def quantize_package( + source_dir: str | Path, + output_dir: str | Path, + *, + precision: str = "q4_k_m", +) -> Path: + """Quantize model.onnx with a CPU-isolated Olive workflow.""" + from olive.workflows import run as olive_run + import olive.systems.local as olive_local + + source = Path(source_dir) + source_model = source / "model.onnx" + if not source_model.is_file(): + raise FileNotFoundError(f"Missing source model: {source_model}") + output = Path(output_dir) + _require_empty_output(output) + + with tempfile.TemporaryDirectory(prefix="olive-nemotron-") as temp: + olive_output = Path(temp) / "output" + # Olive 0.13 auto-registers every DLL bundled in a GPU ORT wheel, + # even for a CPU-only target. That makes an unrelated TensorRT DLL + # failure abort weight-only quantization. Suppress registration for + # this pass; the explicit workflow target remains CPU-only. + register_ep_libraries = olive_local.maybe_register_ep_libraries + olive_local.maybe_register_ep_libraries = lambda _paths: None + try: + olive_run(_olive_config(source_model, olive_output, precision)) + finally: + olive_local.maybe_register_ep_libraries = register_ep_libraries + _copy_olive_model(_find_olive_model(olive_output), output) + + for name in _METADATA_FILES | {"source_manifest.json"}: + path = source / name + if path.is_file(): + shutil.copy2(path, output / name) + manifest_path = output / "source_manifest.json" + manifest = ( + json.loads(manifest_path.read_text(encoding="utf-8")) + if manifest_path.is_file() + else {"model_id": MODEL_ID, "revision": REVISION} + ) + manifest.update({"quantization": precision, "olive_provider": "CPUExecutionProvider"}) + manifest_path.write_text(json.dumps(manifest, indent=2), encoding="utf-8") + return output + + +def smoke_test(model_dir: str | Path, *, device: str) -> list[int]: + """Load the assembled package and perform cached multi-token generation.""" + import numpy as np + + generated, logits, _profile = run_token_ids( + model_dir, + [1, 42, 17], + max_new_tokens=4, + device=device, + ) + if len(generated) != 4 or any(not np.isfinite(step).all() for step in logits): + raise RuntimeError(f"Quantized generation smoke test failed: {generated}") + return generated + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--source-dir", default="output/f16/cuda") + parser.add_argument("--output-dir", default="output/Q4_K_M/cuda") + parser.add_argument("--ep", choices=["cpu", "cuda"], default="cuda") + parser.add_argument("--precision", choices=["q4_k_m", "nf4"], default="q4_k_m") + parser.add_argument("--skip-export", action="store_true") + parser.add_argument("--skip-quantization", action="store_true") + parser.add_argument("--skip-smoke", action="store_true") + args = parser.parse_args() + + if not args.skip_export: + export_checkpoint(args.source_dir, ep=args.ep) + result_dir = Path(args.source_dir) + if not args.skip_quantization: + result_dir = quantize_package( + args.source_dir, + args.output_dir, + precision=args.precision, + ) + if not args.skip_smoke: + print("Generated token IDs:", smoke_test(result_dir, device=args.ep)) + + +if __name__ == "__main__": + main() diff --git a/examples/olive/nemotron-3_5-lightning-30b/requirements.txt b/examples/olive/nemotron-3_5-lightning-30b/requirements.txt new file mode 100644 index 000000000..b29382a5f --- /dev/null +++ b/examples/olive/nemotron-3_5-lightning-30b/requirements.txt @@ -0,0 +1,3 @@ +olive-ai>=0.13.0 +requests>=2.25 +safetensors>=0.4 diff --git a/examples/olive/nemotron-3_5-lightning-30b/validate_reduced_checkpoint.py b/examples/olive/nemotron-3_5-lightning-30b/validate_reduced_checkpoint.py new file mode 100644 index 000000000..b5fc1df49 --- /dev/null +++ b/examples/olive/nemotron-3_5-lightning-30b/validate_reduced_checkpoint.py @@ -0,0 +1,480 @@ +#!/usr/bin/env python +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +"""Validate a reduced model assembled from byte ranges of the pinned checkpoint.""" + +from __future__ import annotations + +import argparse +import json +import math +import struct +from collections import Counter +from pathlib import Path + +import numpy as np +import torch +from huggingface_hub import hf_hub_download +from safetensors import safe_open +from safetensors.torch import load_file, save_file + +from inference import ( + MODEL_ID, + REVISION, + _create_session, + _initial_states, + _as_numpy, + _run_session, + _token_feeds, + _update_states, + run_token_ids, + summarize_profile, +) +from optimize import quantize_package + +_VOCAB_SIZE = 256 +_NUM_EXPERTS = 4 +_LAYER_REMAP = {0: 0, 1: 1, 5: 2} +_DTYPES = { + "f32": (torch.float32, "FLOAT"), + "f16": (torch.float16, "FLOAT16"), + "bf16": (torch.bfloat16, "BFLOAT16"), +} + + +class _PinnedSafetensors: + """Read selected tensors via HTTP Range without downloading 65.8 GB.""" + + def __init__(self) -> None: + import requests + + 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"] + 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: + response = self._session.get( + self._url(shard), + headers={"Range": f"bytes={start}-{end}"}, + timeout=180, + ) + expected = end - start + 1 + if response.status_code != 206 or len(response.content) != expected: + raise RuntimeError( + f"Range fetch failed for {shard} bytes {start}-{end}: " + f"status={response.status_code}, bytes={len(response.content)}" + ) + return response.content + + def _header(self, shard: str) -> tuple[int, dict]: + if shard not in self._headers: + header_length = struct.unpack(" torch.Tensor: + shard = self.weight_map[name] + header_length, header = self._header(shard) + entry = header[name] + shape = list(entry["shape"]) + dtype_name = entry["dtype"] + dtype = {"BF16": torch.bfloat16, "F32": torch.float32}[dtype_name] + element_size = {"BF16": 2, "F32": 4}[dtype_name] + start, end = entry["data_offsets"] + if rows is not None: + if not shape or rows > shape[0]: + raise ValueError(f"Invalid row slice {rows} for {name}: {shape}") + row_elements = math.prod(shape[1:]) + end = start + rows * row_elements * element_size + shape[0] = rows + + data_start = 8 + header_length + payload = self._range(shard, data_start + start, data_start + end - 1) + tensor = torch.frombuffer(bytearray(payload), dtype=dtype).clone() + return tensor.reshape(shape) + + +def _source_to_target(name: str) -> str: + if name == "backbone.embeddings.weight": + return "model.embeddings.weight" + if name == "backbone.norm_f.weight": + return "model.norm_f.weight" + if not name.startswith("backbone.layers."): + return name + parts = name.split(".") + source_layer = int(parts[2]) + parts[2] = str(_LAYER_REMAP[source_layer]) + parts[0] = "model" + return ".".join(parts) + + +def _build_reduced_state(cache_path: Path) -> dict[str, torch.Tensor]: + if cache_path.is_file(): + with safe_open(cache_path, framework="pt") as cached: + metadata = cached.metadata() or {} + if metadata.get("revision") != REVISION: + raise ValueError( + f"Reduced cache revision mismatch: {metadata.get('revision')} != {REVISION}" + ) + return load_file(cache_path) + + source = _PinnedSafetensors() + state: dict[str, torch.Tensor] = { + "model.embeddings.weight": source.tensor( + "backbone.embeddings.weight", + rows=_VOCAB_SIZE, + ), + "model.norm_f.weight": source.tensor("backbone.norm_f.weight"), + "lm_head.weight": source.tensor("lm_head.weight", rows=_VOCAB_SIZE), + } + + for source_layer in (0, 5): + prefix = f"backbone.layers.{source_layer}." + for name in sorted(source.weight_map): + if name.startswith(prefix): + state[_source_to_target(name)] = source.tensor(name) + + moe_prefix = "backbone.layers.1.mixer" + state["model.layers.1.norm.weight"] = source.tensor("backbone.layers.1.norm.weight") + state["model.layers.1.mixer.gate.weight"] = source.tensor( + f"{moe_prefix}.gate.weight", + rows=_NUM_EXPERTS, + ) + state["model.layers.1.mixer.gate.e_score_correction_bias"] = source.tensor( + f"{moe_prefix}.gate.e_score_correction_bias", + rows=_NUM_EXPERTS, + ) + for projection in ("up_proj", "down_proj"): + state[f"model.layers.1.mixer.experts.{projection}"] = torch.stack( + [ + source.tensor(f"{moe_prefix}.experts.{expert}.{projection}.weight") + for expert in range(_NUM_EXPERTS) + ] + ) + state[f"model.layers.1.mixer.shared_experts.{projection}.weight"] = source.tensor( + f"{moe_prefix}.shared_experts.{projection}.weight" + ) + + cache_path.parent.mkdir(parents=True, exist_ok=True) + save_file( + {name: tensor.contiguous() for name, tensor in state.items()}, + cache_path, + metadata={"model_id": MODEL_ID, "revision": REVISION}, + ) + return state + + +def _hf_config(): + from transformers import NemotronHConfig + + return NemotronHConfig( + vocab_size=_VOCAB_SIZE, + hidden_size=2688, + layers_block_type=["linear_attention", "moe", "full_attention"], + num_attention_heads=32, + num_key_value_heads=2, + head_dim=128, + intermediate_size=1856, + mamba_num_heads=64, + mamba_head_dim=64, + ssm_state_size=128, + n_groups=8, + conv_kernel=4, + expand=2, + use_mamba_kernels=False, + moe_intermediate_size=1856, + moe_shared_expert_intermediate_size=3712, + n_routed_experts=_NUM_EXPERTS, + num_experts_per_tok=2, + routed_scaling_factor=2.5, + n_group=1, + topk_group=1, + norm_topk_prob=True, + layer_norm_epsilon=1e-5, + rescale_prenorm_residual=False, + max_position_embeddings=262144, + ) + + +def _hf_model( + state: dict[str, torch.Tensor], + *, + dtype: torch.dtype, + device: str, +): + from transformers import NemotronHForCausalLM + + model = NemotronHForCausalLM(_hf_config()).to(device=device, dtype=dtype) + target = model.state_dict() + if set(target) != set(state): + missing = sorted(set(target) - set(state)) + extra = sorted(set(state) - set(target)) + raise ValueError(f"Reduced state mismatch; missing={missing}, extra={extra}") + converted = { + name: tensor.to(device=device, dtype=target[name].dtype) + for name, tensor in state.items() + } + model.load_state_dict(converted, strict=True) + + # The production loader keeps this selection-only bias in fp32. + gate = model.model.layers[1].mixer.gate + gate.e_score_correction_bias = state[ + "model.layers.1.mixer.gate.e_score_correction_bias" + ].to(device=device, dtype=torch.float32) + return model.eval() + + +def _mobius_package( + state: dict[str, torch.Tensor], + *, + dtype_name: str, + ep: str, +): + import onnx_ir as ir + + from mobius import build_from_module + from mobius._configs import NemotronHConfig + from mobius._flags import override_flags + from mobius.models.nemotron_h import NemotronHCausalLMModel + + config = NemotronHConfig.from_transformers(_hf_config()) + config.dtype = getattr(ir.DataType, _DTYPES[dtype_name][1]) + module = NemotronHCausalLMModel(config) + with override_flags(ort_cuda_grouped_rmsnorm_workaround=ep == "cuda"): + package = build_from_module( + module, + config, + task="hybrid-text-generation", + execution_provider=ep, + trace_optimization=True, + ) + package.apply_weights(module.preprocess_weights(dict(state))) + unset = [ + name + for name, value in package["model"].graph.initializers.items() + if value.const_value is None + ] + if unset: + raise ValueError(f"Weighted graph still has {len(unset)} unset parameters: {unset[:5]}") + return package + + +def _full_prefill(session, token_ids: list[int]) -> np.ndarray: + states = _initial_states(session) + output_names = [output.name for output in session.get_outputs()] + feeds = _token_feeds( + session, + np.array([token_ids], dtype=np.int64), + total_length=len(token_ids), + position_ids=np.arange(len(token_ids), dtype=np.int64)[None, :], + states=states, + ) + outputs = _run_session(session, output_names, feeds) + return _as_numpy(outputs[output_names.index("logits")]).astype(np.float32) + + +def _hf_full_prefill(model, token_ids: list[int], device: str) -> np.ndarray: + ids = torch.tensor([token_ids], dtype=torch.long, device=device) + with torch.no_grad(): + logits = model( + input_ids=ids, + attention_mask=torch.ones_like(ids), + position_ids=torch.arange(len(token_ids), device=device)[None, :], + use_cache=False, + ).logits + return logits.float().cpu().numpy() + + +def _hf_generate( + model, + token_ids: list[int], + device: str, + max_new_tokens: int, +) -> tuple[list[int], list[np.ndarray]]: + from transformers import DynamicCache + + cache = DynamicCache(config=model.config) + past_length = 0 + outputs = None + with torch.no_grad(): + for token_id in token_ids: + ids = torch.tensor([[token_id]], dtype=torch.long, device=device) + outputs = model( + input_ids=ids, + attention_mask=torch.ones((1, past_length + 1), dtype=torch.long, device=device), + position_ids=torch.tensor([[past_length]], dtype=torch.long, device=device), + past_key_values=cache, + use_cache=True, + ) + cache = outputs.past_key_values + past_length += 1 + + assert outputs is not None + generated: list[int] = [] + logits_by_step: list[np.ndarray] = [] + for _ in range(max_new_tokens): + logits_by_step.append(outputs.logits[0, -1].float().cpu().numpy()) + token_id = int(outputs.logits[0, -1].argmax()) + generated.append(token_id) + ids = torch.tensor([[token_id]], dtype=torch.long, device=device) + outputs = model( + input_ids=ids, + attention_mask=torch.ones((1, past_length + 1), dtype=torch.long, device=device), + position_ids=torch.tensor([[past_length]], dtype=torch.long, device=device), + past_key_values=cache, + use_cache=True, + ) + cache = outputs.past_key_values + past_length += 1 + return generated, logits_by_step + + +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:.6g}, cosine={cosine:.9f}") + np.testing.assert_allclose(actual, expected, rtol=1e-3, atol=atol) + + +def _graph_audit(model) -> dict[str, int]: + counts = Counter( + f"{node.domain or 'ai.onnx'}::{node.op_type}" for node in model.graph.all_nodes() + ) + return dict(sorted(counts.items())) + + +def _validate_variant( + state: dict[str, torch.Tensor], + output_root: Path, + *, + dtype_name: str, + device: str, +) -> Path: + torch_dtype = _DTYPES[dtype_name][0] + ep = "cuda" if device == "cuda" else "cpu" + package = _mobius_package(state, dtype_name=dtype_name, ep=ep) + variant_dir = output_root / f"{dtype_name}-{ep}" + variant_dir.mkdir(parents=True, exist_ok=True) + package.save(variant_dir, external_data="onnx") + _hf_config().save_pretrained(variant_dir) + + profile = device == "cuda" + session = _create_session(variant_dir / "model.onnx", device, profile) + prompt_ids = [1, 42, 17] + actual = _full_prefill(session, prompt_ids) + hf_model = _hf_model(state, dtype=torch_dtype, device=device) + expected = _hf_full_prefill(hf_model, prompt_ids, device) + atol = 2e-3 if dtype_name == "f32" else 1e-2 + _assert_logits_close(actual, expected, atol=atol, label=f"{dtype_name}/{ep} prefill") + + generated, logits, profile_path = run_token_ids( + variant_dir, + prompt_ids, + max_new_tokens=4, + device=device, + profile=profile, + ) + if len(generated) != 4 or any(not np.isfinite(step).all() for step in logits): + raise AssertionError(f"Invalid generation for {dtype_name}/{ep}: {generated}") + hf_generated, hf_step_logits = _hf_generate( + hf_model, + prompt_ids, + device, + max_new_tokens=4, + ) + if generated != hf_generated: + raise AssertionError( + f"{dtype_name}/{ep} generation mismatch: ONNX={generated}, HF={hf_generated}" + ) + for index, (actual_step, expected_step) in enumerate(zip(logits, hf_step_logits)): + _assert_logits_close( + actual_step, + expected_step, + atol=atol, + label=f"{dtype_name}/{ep} cached step {index}", + ) + print(f"{dtype_name}/{ep} generated IDs: {generated}") + print(f"{dtype_name}/{ep} weighted graph ops: {_graph_audit(package['model'])}") + + if profile_path is not None: + placement = summarize_profile(profile_path) + print(f"{dtype_name}/{ep} provider placement: {placement}") + if placement.get("CUDAExecutionProvider", 0) == 0: + raise AssertionError(f"No CUDA nodes found in profile: {placement}") + del hf_model + if device == "cuda": + torch.cuda.empty_cache() + return variant_dir + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument( + "--cache", + default="cache/nemotron-3_5-reduced-real.safetensors", + ) + parser.add_argument("--output-dir", default="output/reduced-validation") + parser.add_argument( + "--matrix", + nargs="+", + choices=["f32-cpu", "f16-cuda", "bf16-cuda"], + default=["f32-cpu", "f16-cuda"], + ) + parser.add_argument("--skip-quantization", action="store_true") + args = parser.parse_args() + + state = _build_reduced_state(Path(args.cache)) + print(f"Loaded {len(state)} reduced real-weight tensors from revision {REVISION}") + output_root = Path(args.output_dir) + output_root.mkdir(parents=True, exist_ok=True) + variants: dict[str, Path] = {} + for variant in args.matrix: + dtype_name, device = variant.split("-") + if device == "cuda" and not torch.cuda.is_available(): + raise RuntimeError(f"CUDA validation requested but unavailable: {variant}") + variants[variant] = _validate_variant( + state, + output_root, + dtype_name=dtype_name, + device=device, + ) + + if not args.skip_quantization: + source = variants.get("f16-cuda") + if source is None: + raise ValueError("Olive validation requires f16-cuda in --matrix") + quantized = quantize_package(source, output_root / "q4_k_m-cuda") + generated, logits, _profile = run_token_ids( + quantized, + [1, 42, 17], + max_new_tokens=4, + device="cuda", + profile=False, + ) + if len(generated) != 4 or any(not np.isfinite(step).all() for step in logits): + raise AssertionError(f"Quantized generation failed: {generated}") + print(f"Olive Q4_K_M package loaded and generated IDs: {generated}") + + +if __name__ == "__main__": + main() diff --git a/src/mobius/__main__.py b/src/mobius/__main__.py index bed089bcb..4ceb6cf68 100644 --- a/src/mobius/__main__.py +++ b/src/mobius/__main__.py @@ -307,6 +307,7 @@ def _resolve_static_cache_task(model_type: str) -> ModelTask: config_path, trust_remote_code=trust_remote_code ) model_type = hf_config.model_type + _preflight_ort_genai_runtime(args, model_type) parent_config = hf_config if hasattr(hf_config, "text_config"): hf_config = hf_config.text_config @@ -337,6 +338,14 @@ def _resolve_static_cache_task(model_type: str) -> ModelTask: pkg.apply_weights(state_dict) else: model_id_or_path = args.model + if getattr(args, "runtime", None) == "ort-genai": + import transformers + + runtime_config = transformers.AutoConfig.from_pretrained( + model_id_or_path, + trust_remote_code=trust_remote_code, + ) + _preflight_ort_genai_runtime(args, runtime_config.model_type) if static_cache_params is not None: # Detect model type to resolve the correct static cache task. import transformers @@ -362,6 +371,20 @@ def _resolve_static_cache_task(model_type: str) -> ModelTask: _save_package(pkg, output_dir, args, optimize, component_filter) +def _preflight_ort_genai_runtime(args, model_type: str) -> None: + """Reject known unsupported ORT GenAI types before weights are downloaded.""" + if getattr(args, "runtime", None) != "ort-genai": + return + from mobius.integrations.ort_genai.auto_export import ( + _validate_ort_genai_model_type, + ) + + try: + _validate_ort_genai_model_type(model_type) + except ValueError as error: + raise SystemExit(f"Error: {error}") from error + + def _save_package( pkg, output_dir: str, args, optimize: str | None, component_filter: str | None ) -> None: diff --git a/src/mobius/_configs/_base.py b/src/mobius/_configs/_base.py index 78dcb4e7e..9c31eb49b 100644 --- a/src/mobius/_configs/_base.py +++ b/src/mobius/_configs/_base.py @@ -40,6 +40,16 @@ def _resolve_dtype(config) -> ir.DataType | None: return None +def _resolve_explicit_dtype(value, *, field_name: str) -> ir.DataType: + """Normalize a required dtype string/torch/IR value.""" + if isinstance(value, ir.DataType): + return value + torch_dtype = getattr(torch, value, None) if isinstance(value, str) else value + if isinstance(torch_dtype, torch.dtype): + return tensor_adapters.from_torch_dtype(torch_dtype) + raise ValueError(f"Unsupported {field_name}: {value!r}") + + def _resolve_hidden_act(config, model_type: str) -> str | None: """Resolve the hidden activation function from HF config patterns. @@ -2643,6 +2653,7 @@ class NemotronHConfig(ArchitectureConfig): mamba_conv_bias: bool = True mamba_proj_bias: bool = False mamba_time_step_min: float = 0.001 + mamba_ssm_cache_dtype: ir.DataType = ir.DataType.FLOAT moe_latent_size: int | None = None @classmethod @@ -2660,15 +2671,29 @@ def from_transformers(cls, config, parent_config=None) -> NemotronHConfig: "-": "mlp", "E": "moe", } - layers_block_type = [char_map.get(c, "mamba2") for c in pattern] + invalid_chars = sorted(set(pattern) - set(char_map)) + if invalid_chars: + raise ValueError( + "Unsupported NemotronH hybrid_override_pattern character(s): " + f"{invalid_chars}" + ) + layers_block_type = [char_map[c] for c in pattern] else: - # Convert HF names to mobius names + # Transformers 5.x uses ``linear_attention``/``full_attention``; + # older configs use ``mamba``/``attention``. Normalize both + # vocabularies to Mobius cache-layer names. type_map = { "mamba": "mamba2", + "linear_attention": "mamba2", "attention": "full_attention", + "full_attention": "full_attention", "moe": "moe", + "mlp": "mlp", } - layers_block_type = [type_map.get(t, t) for t in layers_block_type] + invalid_types = sorted(set(layers_block_type) - set(type_map)) + if invalid_types: + raise ValueError(f"Unsupported NemotronH layer type(s): {invalid_types}") + layers_block_type = [type_map[t] for t in layers_block_type] # Override num_hidden_layers based on actual pattern length n = len(layers_block_type) if layers_block_type else base.num_hidden_layers @@ -2717,6 +2742,10 @@ def from_transformers(cls, config, parent_config=None) -> NemotronHConfig: mamba_conv_bias=getattr(config, "use_conv_bias", True), mamba_proj_bias=getattr(config, "mamba_proj_bias", False), mamba_time_step_min=getattr(config, "time_step_min", 0.001), + mamba_ssm_cache_dtype=_resolve_explicit_dtype( + getattr(config, "mamba_ssm_cache_dtype", "float32"), + field_name="mamba_ssm_cache_dtype", + ), moe_latent_size=getattr(config, "moe_latent_size", None), shared_expert_intermediate_size=shared_expert_intermediate_size, ) diff --git a/src/mobius/_configs_test.py b/src/mobius/_configs_test.py index f353b9084..a242485be 100644 --- a/src/mobius/_configs_test.py +++ b/src/mobius/_configs_test.py @@ -16,6 +16,7 @@ ArchitectureConfig, AudioConfig, MuseGlimmerConfig, + NemotronHConfig, QuantizationConfig, VisionConfig, _extract_audio_config, @@ -402,6 +403,60 @@ class FakeNemotronH: # rope_interleave stays at its inert False default. assert config.rope_interleave is False + +class TestNemotronHConfig: + @staticmethod + def _fake_config(layer_types: list[str]) -> SimpleNamespace: + return SimpleNamespace( + model_type="nemotron_h", + layers_block_type=layer_types, + num_hidden_layers=len(layer_types), + vocab_size=131072, + hidden_size=2688, + intermediate_size=1856, + num_attention_heads=32, + num_key_value_heads=2, + head_dim=128, + max_position_embeddings=262144, + pad_token_id=0, + layer_norm_epsilon=1e-5, + mamba_num_heads=64, + mamba_head_dim=64, + ssm_state_size=128, + n_groups=8, + conv_kernel=4, + expand=2, + n_routed_experts=128, + num_experts_per_tok=6, + moe_intermediate_size=1856, + moe_shared_expert_intermediate_size=3712, + routed_scaling_factor=2.5, + ) + + @pytest.mark.parametrize( + "hf_layer_types", + [ + ["mamba", "attention", "moe", "mlp"], + ["linear_attention", "full_attention", "moe", "mlp"], + ], + ids=["legacy-transformers", "current-transformers"], + ) + def test_normalizes_transformers_layer_type_vocabularies( + self, hf_layer_types: list[str] + ) -> None: + config = NemotronHConfig.from_transformers(self._fake_config(hf_layer_types)) + + assert config.layer_types == ["mamba2", "full_attention", "moe", "mlp"] + assert config.num_hidden_layers == 4 + assert config.num_local_experts == 128 + assert config.num_experts_per_tok == 6 + assert config.shared_expert_intermediate_size == 3712 + assert config.mamba_ssm_cache_dtype.name == "FLOAT" + + def test_rejects_unknown_layer_type(self) -> None: + with pytest.raises(ValueError, match="Unsupported NemotronH layer type"): + NemotronHConfig.from_transformers(self._fake_config(["linear_attention", "bogus"])) + def test_from_transformers_legacy_rotary_dim_enables_rope(self): """GPT-J / CodeGen-style legacy configs use ``rotary_dim``.""" diff --git a/src/mobius/_registry.py b/src/mobius/_registry.py index bbe9a68a3..698beed1c 100644 --- a/src/mobius/_registry.py +++ b/src/mobius/_registry.py @@ -919,7 +919,7 @@ def _create_default_registry() -> ModelRegistry: "command_r": "CohereForAI/c4ai-command-r-v01", "csm": "sesame/csm-1b", "evolla": "westlake-repl/Evolla-10B-hf", - "nemotron_h": "nvidia/NVIDIA-Nemotron-3-Nano-4B-BF16", + "nemotron_h": "nvidia/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-BF16", "nemotron_parse": "nvidia/NVIDIA-Nemotron-Parse-2.0", "open-llama": "openlm-research/open_llama_3b", "persimmon": "adept/persimmon-8b-base", diff --git a/src/mobius/integrations/ort_genai/auto_export.py b/src/mobius/integrations/ort_genai/auto_export.py index 35a611c9c..a84b7f8e2 100644 --- a/src/mobius/integrations/ort_genai/auto_export.py +++ b/src/mobius/integrations/ort_genai/auto_export.py @@ -1237,9 +1237,22 @@ def _write_genai_config( return generator.write(output_dir) +def _validate_ort_genai_model_type(model_type: str | None) -> None: + """Reject model types whose runtime contract is known before graph build.""" + if model_type == "nemotron_h": + raise ValueError( + "onnxruntime-genai does not currently support NemotronH's mixed cache " + "contract: Mamba layers require conv_state + ssm_state while full-attention " + "layers require key/value caches at sparse global layer indices. Export " + "without --runtime ort-genai and run the saved model directly with ONNX " + "Runtime (see examples/nemotron_3_nano_text_generation.py)." + ) + + def _validate_ort_genai_compatibility(pkg: ModelPackage) -> None: """Reject packages whose required inputs cannot be supplied by ORT GenAI.""" config = getattr(pkg, "config", None) + _validate_ort_genai_model_type(getattr(config, "model_type", None)) if getattr(config, "model_type", None) == "parakeet_ctc": raise ValueError( "ORT GenAI does not define a feature-input CTC ASR pipeline; " diff --git a/src/mobius/integrations/ort_genai/auto_export_test.py b/src/mobius/integrations/ort_genai/auto_export_test.py index d9ca4fe2a..a74f6d1e3 100644 --- a/src/mobius/integrations/ort_genai/auto_export_test.py +++ b/src/mobius/integrations/ort_genai/auto_export_test.py @@ -1103,6 +1103,23 @@ class FakeConfig: write_ort_genai_config(pkg, str(tmp_path)) assert not (tmp_path / "genai_config.json").exists() + def test_rejects_nemotron_h_mixed_cache_before_writing(self, tmp_path): + import dataclasses + + from mobius._model_package import ModelPackage + + @dataclasses.dataclass + class FakeConfig: + model_type: str = "nemotron_h" + + pkg = ModelPackage(config=FakeConfig()) + output_dir = tmp_path / "ort-genai" + + with pytest.raises(ValueError, match=r"NemotronH.*conv_state \+ ssm_state.*key/value"): + write_ort_genai_config(pkg, str(output_dir)) + + assert not output_dir.exists() + def test_processor_config_written_with_vision(self, tmp_path): """image_processor.json is written when pkg.config.vision is set.""" import dataclasses diff --git a/src/mobius/models/nemotron_h.py b/src/mobius/models/nemotron_h.py index d56b4f87c..ff7cc7529 100644 --- a/src/mobius/models/nemotron_h.py +++ b/src/mobius/models/nemotron_h.py @@ -30,8 +30,8 @@ from __future__ import annotations import re -from typing import TYPE_CHECKING +import onnx_ir as ir import torch from onnxscript import OpBuilder, nn @@ -47,9 +47,6 @@ create_padding_mask, ) -if TYPE_CHECKING: - import onnx_ir as ir - # --------------------------------------------------------------------------- # Decoder layers # --------------------------------------------------------------------------- @@ -233,6 +230,7 @@ def __init__( self.weight = nn.Parameter([num_experts, hidden_size]) # Correction bias for expert selection (loaded from checkpoint) self.e_score_correction_bias = nn.Parameter([num_experts]) + self.e_score_correction_bias._keep_float32 = True def forward(self, op: OpBuilder, hidden_states: ir.Value): # Cast to float32 for numerical stability (eps=1e-20 underflows @@ -240,14 +238,17 @@ def forward(self, op: OpBuilder, hidden_states: ir.Value): # in NemotronHTopkRouter.forward and never casts back. hidden_states = op.Cast(hidden_states, to=1) # FLOAT32 - weight_t = op.Transpose(self.weight, perm=[1, 0]) + weight_t = op.Transpose(op.Cast(self.weight, to=1), perm=[1, 0]) router_logits = op.MatMul(hidden_states, weight_t) # Sigmoid probabilities (these become the final routing weights) probs = op.Sigmoid(router_logits) # Add correction bias for expert selection only - choice_scores = op.Add(probs, self.e_score_correction_bias) + choice_scores = op.Add( + probs, + op.Cast(self.e_score_correction_bias, to=1), + ) # Select top-k experts based on biased scores k = op.Constant(value_ints=[self.top_k]) @@ -374,12 +375,16 @@ def forward(self, op: OpBuilder, hidden_states: ir.Value): weighted = op.Mul(routing_weights, match_float) # Sum matched routing weights across top_k dim → per-token weight weight = op.ReduceSum(weighted, [-1], keepdims=True) - contribution = op.Mul(expert_output, weight) + # Match HF: accumulate routed expert contributions in the fp32 + # routing-weight dtype, then cast the completed routed result once. + contribution = op.Mul(op.Cast(expert_output, to=1), weight) if result is None: result = contribution else: result = op.Add(result, contribution) + result = op.CastLike(result, hidden_states) + # Optional latent projection back to hidden_size if self._has_latent: result = self.fc2_latent_proj(op, result) @@ -504,6 +509,9 @@ class NemotronHCausalLMModel(nn.Module): Uses ``HybridCausalLMTask`` with mixed ``"mamba2"``, ``"full_attention"``, and ``"mlp"`` layer types for the cache. + The exported task is the base decoder used by + ``NemotronHForCausalLM.forward``; auxiliary ``mtp.*`` training heads are + outside that generation graph and are intentionally not loaded. HuggingFace reference: ``NemotronHForCausalLM``. """ @@ -514,6 +522,12 @@ class NemotronHCausalLMModel(nn.Module): def __init__(self, config: NemotronHConfig): super().__init__() + if config.dtype == ir.DataType.BFLOAT16: + raise ValueError( + "NemotronH BF16 execution is not numerically supported: reduced real-weight " + "CUDA parity exceeds the 1e-2 logit tolerance. Build the BF16 checkpoint " + "with dtype='f16' instead." + ) self.config = config self.model = _NemotronHTextModel(config) self.lm_head = Linear(config.hidden_size, config.vocab_size, bias=False) @@ -554,6 +568,7 @@ def preprocess_weights( - mlp: ``mixer.`` → ``mlp.`` - moe: ``mixer.`` → ``moe.`` 6. MoE stacked 3D expert tensors split into per-expert 2D weights + 7. Auxiliary ``mtp.*`` training heads omitted by NemotronHForCausalLM """ layer_types = self.config.layer_types or [] @@ -572,6 +587,11 @@ def preprocess_weights( new_state_dict: dict[str, torch.Tensor] = {} for key, value in state_dict.items(): + # The official Nemotron 3.5 checkpoint includes multi-token + # prediction heads, but the trusted NemotronHForCausalLM forward + # does not instantiate them and marks ``mtp.*`` as unexpected. + if key.startswith("mtp."): + continue new_key = _rename_nemotron_h_weight(key, layer_types) # Split stacked 3D expert tensors into per-expert 2D weights. # HF stores experts.up_proj as (num_experts, inter, input) and diff --git a/src/mobius/tasks/_cache_utils.py b/src/mobius/tasks/_cache_utils.py index 68f994407..e2b09663d 100644 --- a/src/mobius/tasks/_cache_utils.py +++ b/src/mobius/tasks/_cache_utils.py @@ -310,7 +310,7 @@ def _make_hybrid_cache_inputs( ) ssm_state = builder.input( f"{prefix}.{i}.ssm_state", - dtype=dtype, + dtype=getattr(config, "mamba_ssm_cache_dtype", dtype), shape=[batch, mamba2_n_heads, mamba2_d_state, mamba2_d_head], ) pairs.append((conv_state, ssm_state)) diff --git a/testdata/cases/causal-lm/nemotron-3_5-lightning-30b.yaml b/testdata/cases/causal-lm/nemotron-3_5-lightning-30b.yaml new file mode 100644 index 000000000..a92660814 --- /dev/null +++ b/testdata/cases/causal-lm/nemotron-3_5-lightning-30b.yaml @@ -0,0 +1,20 @@ +model_id: "nvidia/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-BF16" +model_type: "nemotron_h" +revision: "d468880b6ad3c6e0d21377ce7242adaea4cc884d" +task_type: "text-generation" +dtype: "float16" +trust_remote_code: false + +inputs: + prompts: + - "Here is my poem:" + +level: "L4+L5" + +generation: + max_new_tokens: 20 + do_sample: false + +skip_reason: "The pinned checkpoint is 65.8 GB with 30B total parameters; L4/L5 reference generation and weighted ONNX export exceed the storage and accelerator memory available to standard CI." +ci_skip_reason: "30B MoE checkpoint requires a large-memory CUDA host and about 66 GB of checkpoint storage." +notes: "NVIDIA Nemotron 3.5 Lightning 30B-A3B. Hybrid Mamba2 + sigmoid-routed MoE + full attention; 128 routed experts, top-6. The 270 auxiliary mtp.* checkpoint tensors are intentionally ignored by the trusted NemotronHForCausalLM inference model." diff --git a/testdata/golden/causal-lm/nemotron-3_5-lightning-30b-reduced.json b/testdata/golden/causal-lm/nemotron-3_5-lightning-30b-reduced.json new file mode 100644 index 000000000..bbe28a8bd --- /dev/null +++ b/testdata/golden/causal-lm/nemotron-3_5-lightning-30b-reduced.json @@ -0,0 +1,43 @@ +{ + "model_id": "nvidia/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-BF16", + "revision": "d468880b6ad3c6e0d21377ce7242adaea4cc884d", + "fixture": "Production dimensions; complete checkpoint layers 0 (Mamba2), 1 (MoE with experts 0-3), and 5 (attention); embedding/lm_head rows 0-255.", + "input_ids": [ + 1, + 42, + 17 + ], + "top1_id": 12, + "top2_id": 13, + "top10_ids": [ + 12, + 13, + 14, + 1, + 11, + 17, + 16, + 10, + 170, + 201 + ], + "top10_logits": [ + "0x1.503e9a0000000p+1", + "0x1.1e09ce0000000p+1", + "0x1.0927420000000p+1", + "0x1.a5d06c0000000p-1", + "0x1.5fe3a60000000p-1", + "0x1.46611e0000000p-1", + "0x1.4523960000000p-1", + "0x1.3cc3800000000p-1", + "0x1.3bab820000000p-1", + "0x1.3b882c0000000p-1" + ], + "logits_summary": [ + "0x1.503e9a0000000p+1", + "-0x1.b120180000000p+0", + "0x1.3e00480000000p-1", + "0x1.e850aa0000000p-3" + ], + "full_logits_sha256": "7500c2ed86137e20dddf43a85b016746b8b38c4a932277c29d2efa926cc4a23d" +} diff --git a/testdata/golden/causal-lm/nemotron-3_5-lightning-30b-reduced_generation.json b/testdata/golden/causal-lm/nemotron-3_5-lightning-30b-reduced_generation.json new file mode 100644 index 000000000..51769b52f --- /dev/null +++ b/testdata/golden/causal-lm/nemotron-3_5-lightning-30b-reduced_generation.json @@ -0,0 +1,18 @@ +{ + "model_id": "nvidia/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-BF16", + "revision": "d468880b6ad3c6e0d21377ce7242adaea4cc884d", + "fixture": "Independent HuggingFace NemotronHForCausalLM greedy decode using the pinned reduced-real-weight fixture.", + "input_ids": [ + 1, + 42, + 17 + ], + "max_new_tokens": 4, + "do_sample": false, + "generated_tokens": [ + 12, + 13, + 12, + 12 + ] +} diff --git a/tests/_test_configs.py b/tests/_test_configs.py index 402e20dec..6fdb52fe3 100644 --- a/tests/_test_configs.py +++ b/tests/_test_configs.py @@ -1303,6 +1303,7 @@ def _base_config(config_cls=None, **overrides) -> ArchitectureConfig: "nemotron_h", { "hidden_act": "relu2", + "rms_norm_eps": 1e-5, "layer_types": ["mamba2", "mlp", "full_attention", "mamba2"], "_config_cls": NemotronHConfig, "num_hidden_layers": 4, @@ -1320,6 +1321,7 @@ def _base_config(config_cls=None, **overrides) -> ArchitectureConfig: "nemotron_h", { "hidden_act": "relu2", + "rms_norm_eps": 1e-5, "layer_types": [ "mamba2", "moe", diff --git a/tests/arch_validation_test.py b/tests/arch_validation_test.py index d0abc65f7..35ca29c3a 100644 --- a/tests/arch_validation_test.py +++ b/tests/arch_validation_test.py @@ -26,11 +26,14 @@ from __future__ import annotations +import dataclasses import logging import pytest +from mobius._builder import resolve_dtype from mobius._registry import registry +from mobius._testing.golden import discover_test_cases from mobius.integrations.transformers._config_resolver import ( _config_from_hf, _default_task_for_model, @@ -40,6 +43,9 @@ logger = logging.getLogger(__name__) +_PINNED_REVISIONS = {case.model_id: case.revision for case in discover_test_cases()} +_DECLARED_DTYPES = {case.model_id: case.dtype for case in discover_test_cases()} + # Build parametrized test cases from registry entries that have a test_model_id. # # We split known failures by which subset of tests they apply to: @@ -103,10 +109,39 @@ def _load_hf_config(model_id: str): """ import transformers + revision = _PINNED_REVISIONS.get(model_id) try: - return transformers.AutoConfig.from_pretrained(model_id, trust_remote_code=False) + return transformers.AutoConfig.from_pretrained( + model_id, + revision=revision, + trust_remote_code=False, + ) except (ValueError, OSError): - return _try_load_config_json(model_id) + return _try_load_config_json(model_id, revision=revision) + + +def test_load_hf_config_forwards_yaml_revision(monkeypatch): + model_id = "nvidia/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-BF16" + calls = [] + + def _from_pretrained(received_model_id, **kwargs): + calls.append((received_model_id, kwargs)) + return object() + + import transformers + + monkeypatch.setattr(transformers.AutoConfig, "from_pretrained", _from_pretrained) + _load_hf_config(model_id) + + assert calls == [ + ( + model_id, + { + "revision": "d468880b6ad3c6e0d21377ce7242adaea4cc884d", + "trust_remote_code": False, + }, + ) + ] def _resolve_hf_config(hf_config): @@ -160,6 +195,11 @@ def _build_graph(model_type: str, model_id: str): parent_config=parent_config, module_class=registration.module_class, ) + if model_type == "nemotron_h": + config = dataclasses.replace( + config, + dtype=resolve_dtype(_DECLARED_DTYPES[model_id]), + ) module = registration.module_class(config) task_name = registration.task or _default_task_for_model(model_type) diff --git a/tests/build_graph_test.py b/tests/build_graph_test.py index 201b97a4f..cf0e19830 100644 --- a/tests/build_graph_test.py +++ b/tests/build_graph_test.py @@ -5240,6 +5240,61 @@ def test_nemotron_h_moe_preprocess_weights(self): for key in result: assert not key.startswith("backbone."), f"Unrenamed key: {key}" + def test_reduced_precision_keeps_router_bias_and_ssm_cache_float32(self): + import onnx_ir as ir + + from mobius import build_from_module + from mobius._configs import NemotronHConfig + from mobius.models.nemotron_h import NemotronHCausalLMModel + + config = NemotronHConfig( + vocab_size=TINY_VOCAB, + hidden_size=TINY_HIDDEN, + intermediate_size=TINY_INTERMEDIATE, + num_hidden_layers=2, + num_attention_heads=TINY_HEADS, + num_key_value_heads=TINY_KV_HEADS, + rms_norm_eps=1e-5, + layer_types=["mamba2", "moe"], + mamba_n_heads=TINY_KV_HEADS, + mamba_d_head=TINY_HEAD_DIM, + mamba_d_state=16, + mamba_n_groups=1, + mamba_d_conv=4, + hidden_act="relu2", + head_dim=TINY_HEAD_DIM, + num_local_experts=4, + num_experts_per_tok=2, + moe_intermediate_size=TINY_INTERMEDIATE, + dtype=ir.DataType.FLOAT16, + ) + package = build_from_module( + NemotronHCausalLMModel(config), + config, + task="hybrid-text-generation", + execution_provider="cuda", + ) + model = package["model"] + + assert ( + model.graph.initializers["model.layers.1.moe.gate.e_score_correction_bias"].dtype + == ir.DataType.FLOAT + ) + inputs = {value.name: value for value in model.graph.inputs} + assert inputs["past_key_values.0.conv_state"].dtype == ir.DataType.FLOAT16 + assert inputs["past_key_values.0.ssm_state"].dtype == ir.DataType.FLOAT + + def test_bfloat16_is_rejected_with_actionable_error(self): + import onnx_ir as ir + + config = self._nemotron_h_config() + config.dtype = ir.DataType.BFLOAT16 + + with pytest.raises(ValueError, match=r"BF16.*dtype='f16'"): + from mobius.models.nemotron_h import NemotronHCausalLMModel + + NemotronHCausalLMModel(config) + # =========================================================================== # Hybrid SSM+Attention (Jamba) model tests diff --git a/tests/cli_test.py b/tests/cli_test.py index 07c288249..152fe0932 100644 --- a/tests/cli_test.py +++ b/tests/cli_test.py @@ -557,6 +557,29 @@ def test_runtime_ort_genai_rejects_mage_vl_before_saving(self): save.assert_not_called() config_writer.assert_not_called() + def test_runtime_ort_genai_rejects_nemotron_h_before_build(self): + with ( + tempfile.TemporaryDirectory() as tmpdir, + mock.patch( + "transformers.AutoConfig.from_pretrained", + return_value=SimpleNamespace(model_type="nemotron_h"), + ), + mock.patch("mobius.__main__.build") as build_model, + pytest.raises(SystemExit, match=r"NemotronH.*mixed cache"), + ): + main( + [ + "build", + "--model", + "nvidia/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-BF16", + tmpdir, + "--runtime", + "ort-genai", + ] + ) + + build_model.assert_not_called() + def test_runtime_onnx_genai_uses_native_vlm_emitter(self): pkg = mock.MagicMock() pkg.items.return_value = [] diff --git a/tests/model_coverage_test.py b/tests/model_coverage_test.py index 4df4bf723..049c8878d 100644 --- a/tests/model_coverage_test.py +++ b/tests/model_coverage_test.py @@ -241,7 +241,6 @@ def _all_registered_with_test_id() -> dict[str, str]: "mctct": "Audio model — no test_model_id yet", "megatron-bert": "Encoder — no test_model_id yet", "modernbert-decoder": "Decoder variant — no test_model_id yet", - "nemotron_h": "No test_model_id — no suitable public checkpoint", "nllb-moe": "Seq2seq MoE — no test_model_id yet", "nllb_moe": "Seq2seq MoE — no test_model_id yet", "ovis2": "VL model — no test_model_id yet", diff --git a/tests/nemotron_h_real_weight_test.py b/tests/nemotron_h_real_weight_test.py new file mode 100644 index 000000000..02ace1451 --- /dev/null +++ b/tests/nemotron_h_real_weight_test.py @@ -0,0 +1,117 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +"""Pinned reduced-real-weight L4/L5 tests for Nemotron 3.5 Lightning.""" + +from __future__ import annotations + +import importlib.util +import json +import os +import sys +from pathlib import Path + +import numpy as np +import pytest +import torch + +_ROOT = Path(__file__).parents[1] +_EXAMPLE_DIR = _ROOT / "examples" / "olive" / "nemotron-3_5-lightning-30b" +_VALIDATOR_PATH = _EXAMPLE_DIR / "validate_reduced_checkpoint.py" +_L4_PATH = ( + _ROOT / "testdata" / "golden" / "causal-lm" / "nemotron-3_5-lightning-30b-reduced.json" +) +_L5_PATH = ( + _ROOT + / "testdata" + / "golden" + / "causal-lm" + / "nemotron-3_5-lightning-30b-reduced_generation.json" +) + + +def _load_validator(): + sys.path.insert(0, str(_EXAMPLE_DIR)) + try: + spec = importlib.util.spec_from_file_location( + "nemotron_reduced_validator", _VALIDATOR_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) + + +@pytest.fixture(scope="module") +def reduced_real_outputs(tmp_path_factory): + validator = _load_validator() + configured_cache = os.environ.get("MOBIUS_NEMOTRON_REDUCED_CACHE") + cache = ( + Path(configured_cache) + if configured_cache + else tmp_path_factory.mktemp("nemotron-real") / "reduced.safetensors" + ) + state = validator._build_reduced_state(cache) + package = validator._mobius_package(state, dtype_name="f32", ep="cpu") + output_dir = tmp_path_factory.mktemp("nemotron-onnx") + package.save(output_dir, external_data="onnx") + session = validator._create_session(output_dir / "model.onnx", "cpu", False) + prompt_ids = [1, 42, 17] + onnx_logits = validator._full_prefill(session, prompt_ids) + onnx_tokens, _step_logits, _profile = validator.run_token_ids( + output_dir, + prompt_ids, + max_new_tokens=4, + device="cpu", + ) + + hf_model = validator._hf_model(state, dtype=torch.float32, device="cpu") + hf_logits = validator._hf_full_prefill(hf_model, prompt_ids, "cpu") + hf_tokens, _hf_step_logits = validator._hf_generate(hf_model, prompt_ids, "cpu", 4) + return onnx_logits, onnx_tokens, hf_logits, hf_tokens + + +@pytest.mark.integration +@pytest.mark.integration_slow +@pytest.mark.golden +@pytest.mark.parametrize("model_type", ["nemotron_h"]) +def test_nemotron_h_3_5_reduced_real_l4(reduced_real_outputs, model_type): + del model_type + onnx_logits, _onnx_tokens, hf_logits, _hf_tokens = reduced_real_outputs + golden = json.loads(_L4_PATH.read_text(encoding="utf-8")) + + np.testing.assert_allclose(onnx_logits, hf_logits, rtol=1e-3, atol=2e-3) + last_logits = hf_logits[0, -1] + top10 = np.argsort(last_logits)[::-1][:10] + summary = [last_logits.max(), last_logits.min(), last_logits.mean(), last_logits.std()] + + assert top10.tolist() == golden["top10_ids"] + np.testing.assert_allclose( + last_logits[top10], + [float.fromhex(value) for value in golden["top10_logits"]], + rtol=1e-5, + atol=1e-5, + ) + np.testing.assert_allclose( + summary, + [float.fromhex(value) for value in golden["logits_summary"]], + rtol=1e-5, + atol=1e-5, + ) + + +@pytest.mark.integration +@pytest.mark.integration_slow +@pytest.mark.generation +@pytest.mark.parametrize("model_type", ["nemotron_h"]) +def test_nemotron_h_3_5_reduced_real_l5(reduced_real_outputs, model_type): + del model_type + _onnx_logits, onnx_tokens, _hf_logits, hf_tokens = reduced_real_outputs + golden = json.loads(_L5_PATH.read_text(encoding="utf-8")) + + assert len(onnx_tokens) == golden["max_new_tokens"] + assert len(hf_tokens) == golden["max_new_tokens"] + assert hf_tokens == golden["generated_tokens"] + assert onnx_tokens == golden["generated_tokens"] diff --git a/tests/synthetic_parity_test.py b/tests/synthetic_parity_test.py index 2dc92f45a..8548956a0 100644 --- a/tests/synthetic_parity_test.py +++ b/tests/synthetic_parity_test.py @@ -208,11 +208,6 @@ # DeepSeek MLA: deepseek_v2_0 uses group_limited_greedy routing which hits a # HF transformers 5.3.0 bug (DeepseekV2Moe missing num_experts attr). "deepseek_v2_0": "HF transformers 5.3.0 bug: DeepseekV2Moe missing num_experts attr", - # Additional divergences (newly registered models) - # NemotronH Mamba2 layers diverge (cos=0.65): LinearAttention gated-SSM - # recurrence on CPU produces different results than HF's naive Mamba2. - # Attention-only layers match perfectly (cos=0.9999). - "nemotron_h": "Mamba2 SSM recurrence diverges on CPU (LinearAttention vs HF naive)", } # Fields that are properties in HF configs and cannot be set directly, @@ -581,12 +576,18 @@ def _create_hf_config(model_type: str, config_overrides: dict): for lt in layer_types ] - # NemotronH uses layers_block_type with HF values {"mamba", "attention", "moe"}. - # Convert our internal layer_types names (mamba2, full_attention, mlp) to HF names. + # NemotronH uses layers_block_type with current HF values + # {"linear_attention", "full_attention", "moe", "mlp"}. + # Convert our internal layer_types names to that vocabulary. # Also translate mobius Mamba field names to HF NemotronHConfig field names. if hf_model_type in ("nemotron_h",) and "layer_types" in hf_kwargs: layer_types = hf_kwargs.pop("layer_types") - _nemotron_type_map = {"mamba2": "mamba", "full_attention": "attention", "mlp": "moe"} + _nemotron_type_map = { + "mamba2": "linear_attention", + "full_attention": "full_attention", + "moe": "moe", + "mlp": "mlp", + } hf_kwargs["layers_block_type"] = [_nemotron_type_map.get(lt, lt) for lt in layer_types] # Mobius NemotronHConfig → HF NemotronHConfig field name mapping _nemotron_field_map = { @@ -596,6 +597,7 @@ def _create_hf_config(model_type: str, config_overrides: dict): "mamba_n_groups": "n_groups", "mamba_d_conv": "conv_kernel", "mamba_expand": "expand", + "shared_expert_intermediate_size": "moe_shared_expert_intermediate_size", } for old_name, new_name in _nemotron_field_map.items(): if old_name in hf_kwargs: diff --git a/tests/weight_alignment_test.py b/tests/weight_alignment_test.py index 03763b342..2e1dcc351 100644 --- a/tests/weight_alignment_test.py +++ b/tests/weight_alignment_test.py @@ -159,6 +159,58 @@ def test_identity_state_dict_roundtrip(self, model_type: str, config_overrides: _assert_identity_roundtrip(model_type, config_overrides) +def test_nemotron_h_filters_only_auxiliary_mtp_weights() -> None: + """Official MTP tensors are dropped without losing any decoder parameter.""" + config_overrides = next( + overrides + for model_type, overrides, _ in ALL_CAUSAL_LM_CONFIGS + if model_type == "nemotron_h" and "moe" in overrides.get("layer_types", []) + ) + config = _base_config(**config_overrides) + module = registry.get("nemotron_h")(config) + pkg = get_task(_default_task_for_model("nemotron_h")).build(module, config) + parameter_names = _collect_parameter_names(pkg) + state_dict = _build_identity_state_dict(pkg, parameter_names) + + # The 3.5 checkpoint carries these auxiliary training heads, while the + # trusted NemotronHForCausalLM implementation intentionally ignores them. + state_dict["mtp.layers.0.eh_proj.weight"] = torch.ones(1) + state_dict["mtp.layers.1.mixer.experts.0.up_proj.weight"] = torch.ones(1) + + aligned = module.preprocess_weights(state_dict) + + assert not any(name.startswith("mtp.") for name in aligned) + assert parameter_names <= set(aligned) + + +def test_nemotron_h_maps_per_expert_checkpoint_weights() -> None: + """The official 3.5 per-expert safetensor names map to MoE initializers.""" + config_overrides = next( + overrides + for model_type, overrides, _ in ALL_CAUSAL_LM_CONFIGS + if model_type == "nemotron_h" and "moe" in overrides.get("layer_types", []) + ) + config = _base_config(**config_overrides) + module = registry.get("nemotron_h")(config) + moe_layer = config.layer_types.index("moe") + state_dict = { + f"backbone.layers.{moe_layer}.mixer.experts.0.up_proj.weight": torch.ones(1), + f"backbone.layers.{moe_layer}.mixer.experts.0.down_proj.weight": torch.ones(1), + f"backbone.layers.{moe_layer}.mixer.gate.weight": torch.ones(1), + f"backbone.layers.{moe_layer}.mixer.gate.e_score_correction_bias": torch.ones(1), + } + + aligned = module.preprocess_weights(state_dict) + + prefix = f"model.layers.{moe_layer}.moe" + assert set(aligned) == { + f"{prefix}.experts.0.up_proj.weight", + f"{prefix}.experts.0.down_proj.weight", + f"{prefix}.gate.weight", + f"{prefix}.gate.e_score_correction_bias", + } + + # --------------------------------------------------------------------------- # Encoder-only weight alignment # --------------------------------------------------------------------------- From 4e1260fc8228eb662da0b97325b9793df74525f0 Mon Sep 17 00:00:00 2001 From: Justin Chu Date: Thu, 13 Aug 2026 20:36:03 -0700 Subject: [PATCH 02/13] Fix CI lint for Nemotron recipe Apply the Linux formatter/import order to the Olive scripts and include the pending ModelPackage formatting fixes required by the repository-wide lint job. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Signed-off-by: Justin Chu --- .../nemotron-3_5-lightning-30b/optimize.py | 2 +- .../validate_reduced_checkpoint.py | 20 +++++++++++-------- 2 files changed, 13 insertions(+), 9 deletions(-) diff --git a/examples/olive/nemotron-3_5-lightning-30b/optimize.py b/examples/olive/nemotron-3_5-lightning-30b/optimize.py index e5bfd7236..a8df08dd9 100644 --- a/examples/olive/nemotron-3_5-lightning-30b/optimize.py +++ b/examples/olive/nemotron-3_5-lightning-30b/optimize.py @@ -158,8 +158,8 @@ def quantize_package( precision: str = "q4_k_m", ) -> Path: """Quantize model.onnx with a CPU-isolated Olive workflow.""" - from olive.workflows import run as olive_run import olive.systems.local as olive_local + from olive.workflows import run as olive_run source = Path(source_dir) source_model = source / "model.onnx" diff --git a/examples/olive/nemotron-3_5-lightning-30b/validate_reduced_checkpoint.py b/examples/olive/nemotron-3_5-lightning-30b/validate_reduced_checkpoint.py index b5fc1df49..726e5ad39 100644 --- a/examples/olive/nemotron-3_5-lightning-30b/validate_reduced_checkpoint.py +++ b/examples/olive/nemotron-3_5-lightning-30b/validate_reduced_checkpoint.py @@ -16,22 +16,20 @@ import numpy as np import torch from huggingface_hub import hf_hub_download -from safetensors import safe_open -from safetensors.torch import load_file, save_file - from inference import ( MODEL_ID, REVISION, + _as_numpy, _create_session, _initial_states, - _as_numpy, _run_session, _token_feeds, - _update_states, run_token_ids, summarize_profile, ) from optimize import quantize_package +from safetensors import safe_open +from safetensors.torch import load_file, save_file _VOCAB_SIZE = 256 _NUM_EXPERTS = 4 @@ -266,7 +264,9 @@ def _mobius_package( if value.const_value is None ] if unset: - raise ValueError(f"Weighted graph still has {len(unset)} unset parameters: {unset[:5]}") + raise ValueError( + f"Weighted graph still has {len(unset)} unset parameters: {unset[:5]}" + ) return package @@ -312,7 +312,9 @@ def _hf_generate( ids = torch.tensor([[token_id]], dtype=torch.long, device=device) outputs = model( input_ids=ids, - attention_mask=torch.ones((1, past_length + 1), dtype=torch.long, device=device), + attention_mask=torch.ones( + (1, past_length + 1), dtype=torch.long, device=device + ), position_ids=torch.tensor([[past_length]], dtype=torch.long, device=device), past_key_values=cache, use_cache=True, @@ -330,7 +332,9 @@ def _hf_generate( ids = torch.tensor([[token_id]], dtype=torch.long, device=device) outputs = model( input_ids=ids, - attention_mask=torch.ones((1, past_length + 1), dtype=torch.long, device=device), + attention_mask=torch.ones( + (1, past_length + 1), dtype=torch.long, device=device + ), position_ids=torch.tensor([[past_length]], dtype=torch.long, device=device), past_key_values=cache, use_cache=True, From dfb5e8393b077650a5233e3efc4f089a58ebb24a Mon Sep 17 00:00:00 2001 From: Justin Chu Date: Thu, 13 Aug 2026 21:07:34 -0700 Subject: [PATCH 03/13] Add GPU and Olive integration gates Separate unsupported BF16 evidence from the supported validation matrix, add automated FP16 CUDA cached-logit parity, and verify fresh Olive Q4 quantization through final-package CUDA reload and generation. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Signed-off-by: Justin Chu --- .github/workflows/gpu_l5_generation_e2e.yml | 1 + .../nemotron-3_5-lightning-30b/README.md | 8 ++ .../nemotron-3_5-lightning-30b/optimize.py | 5 +- .../validate_reduced_checkpoint.py | 124 +++++++++++++++++- pyproject.toml | 1 + tests/nemotron_h_real_weight_test.py | 112 +++++++++++++++- 6 files changed, 247 insertions(+), 4 deletions(-) diff --git a/.github/workflows/gpu_l5_generation_e2e.yml b/.github/workflows/gpu_l5_generation_e2e.yml index 889de86d1..9183faae7 100644 --- a/.github/workflows/gpu_l5_generation_e2e.yml +++ b/.github/workflows/gpu_l5_generation_e2e.yml @@ -64,6 +64,7 @@ jobs: pip install -r requirements/ci/requirements.txt pip install soundfile librosa ml_dtypes flatbuffers numpy packaging protobuf sympy coloredlogs pip install -e '.[testing,transformers]' + pip install -r examples/olive/nemotron-3_5-lightning-30b/requirements.txt --index-url https://packagefeedproxy.microsoft.io/pypi/simple pip install --pre --extra-index-url https://aiinfra.pkgs.visualstudio.com/PublicPackages/_packaging/ort-cuda-12-nightly/pypi/simple/ onnxruntime-gpu onnxruntime-genai-cuda - name: Run L5 generation E2E tests diff --git a/examples/olive/nemotron-3_5-lightning-30b/README.md b/examples/olive/nemotron-3_5-lightning-30b/README.md index 78ab3c8e7..c390d1d58 100644 --- a/examples/olive/nemotron-3_5-lightning-30b/README.md +++ b/examples/olive/nemotron-3_5-lightning-30b/README.md @@ -129,6 +129,14 @@ weights while retaining production dimensions: python validate_reduced_checkpoint.py ``` +The supported matrix is intentionally limited to FP32/CPU and FP16/CUDA. +Reproduce the BF16 rejection evidence separately without creating a supported +package or weakening the production guard: + +```powershell +python validate_reduced_checkpoint.py --bf16-rejection-evidence +``` + Validated results on ORT 1.28.0 / Olive 0.13.0: | Variant | Full-logit max abs | Generated IDs | Placement | diff --git a/examples/olive/nemotron-3_5-lightning-30b/optimize.py b/examples/olive/nemotron-3_5-lightning-30b/optimize.py index a8df08dd9..4bba51c35 100644 --- a/examples/olive/nemotron-3_5-lightning-30b/optimize.py +++ b/examples/olive/nemotron-3_5-lightning-30b/optimize.py @@ -170,6 +170,9 @@ def quantize_package( with tempfile.TemporaryDirectory(prefix="olive-nemotron-") as temp: olive_output = Path(temp) / "output" + config = _olive_config(source_model, olive_output, precision) + config["cache_dir"] = str(Path(temp) / "cache") + config["clean_cache"] = True # Olive 0.13 auto-registers every DLL bundled in a GPU ORT wheel, # even for a CPU-only target. That makes an unrelated TensorRT DLL # failure abort weight-only quantization. Suppress registration for @@ -177,7 +180,7 @@ def quantize_package( register_ep_libraries = olive_local.maybe_register_ep_libraries olive_local.maybe_register_ep_libraries = lambda _paths: None try: - olive_run(_olive_config(source_model, olive_output, precision)) + olive_run(config) finally: olive_local.maybe_register_ep_libraries = register_ep_libraries _copy_olive_model(_find_olive_model(olive_output), output) diff --git a/examples/olive/nemotron-3_5-lightning-30b/validate_reduced_checkpoint.py b/examples/olive/nemotron-3_5-lightning-30b/validate_reduced_checkpoint.py index 726e5ad39..7bfad5b0b 100644 --- a/examples/olive/nemotron-3_5-lightning-30b/validate_reduced_checkpoint.py +++ b/examples/olive/nemotron-3_5-lightning-30b/validate_reduced_checkpoint.py @@ -37,7 +37,6 @@ _DTYPES = { "f32": (torch.float32, "FLOAT"), "f16": (torch.float16, "FLOAT16"), - "bf16": (torch.bfloat16, "BFLOAT16"), } @@ -270,6 +269,110 @@ def _mobius_package( return package +def _bf16_rejection_evidence_package(state: dict[str, torch.Tensor]): + """Build a test-only BF16 graph without weakening the production guard.""" + import dataclasses + + import onnx_ir as ir + from onnxscript import OpBuilder, nn + + from mobius import build_from_module + from mobius._configs import NemotronHConfig + from mobius._flags import override_flags + from mobius.components import Linear + from mobius.models.nemotron_h import ( + NemotronHCausalLMModel, + _NemotronHTextModel, + ) + + config = NemotronHConfig.from_transformers(_hf_config()) + config.dtype = ir.DataType.BFLOAT16 + + # Prove this evidence path has not weakened or bypassed the production API. + try: + NemotronHCausalLMModel(config) + except ValueError as error: + if "BF16 execution is not numerically supported" not in str(error): + raise + else: + raise AssertionError("Production NemotronH BF16 guard did not reject the model") + + class _Bf16EvidenceCausalLM(nn.Module): + """Test-only wrapper around the production components.""" + + def __init__(self, evidence_config: NemotronHConfig): + super().__init__() + self.model = _NemotronHTextModel(evidence_config) + self.lm_head = Linear( + evidence_config.hidden_size, + evidence_config.vocab_size, + bias=False, + ) + + def forward( + self, + op: OpBuilder, + input_ids, + attention_mask, + position_ids, + past_key_values=None, + ): + hidden_states, present_key_values = self.model( + op, + input_ids=input_ids, + attention_mask=attention_mask, + position_ids=position_ids, + past_key_values=past_key_values, + ) + return self.lm_head(op, hidden_states), present_key_values + + module = _Bf16EvidenceCausalLM(config) + with override_flags(ort_cuda_grouped_rmsnorm_workaround=True): + package = build_from_module( + module, + config, + task="hybrid-text-generation", + execution_provider="onnx-standard", + ) + + preprocessing_config = dataclasses.replace(config, dtype=ir.DataType.FLOAT) + preprocessor = NemotronHCausalLMModel(preprocessing_config) + package.apply_weights(preprocessor.preprocess_weights(dict(state))) + unset = [ + name + for name, value in package["model"].graph.initializers.items() + if value.const_value is None + ] + if unset: + raise ValueError(f"BF16 evidence graph has {len(unset)} unset parameters: {unset[:5]}") + return package + + +def _measure_bf16_rejection( + state: dict[str, torch.Tensor], + output_root: Path, +) -> float: + """Measure the rejected BF16 path on CUDA and return its maximum logit error.""" + if not torch.cuda.is_available(): + raise RuntimeError("BF16 rejection evidence requires CUDA") + + package = _bf16_rejection_evidence_package(state) + evidence_dir = output_root / "bf16-rejection-evidence" + evidence_dir.mkdir(parents=True, exist_ok=True) + package.save(evidence_dir, external_data="onnx") + + session = _create_session(evidence_dir / "model.onnx", "cuda", False) + prompt_ids = [1, 42, 17] + actual = _full_prefill(session, prompt_ids) + hf_model = _hf_model(state, dtype=torch.bfloat16, device="cuda") + expected = _hf_full_prefill(hf_model, prompt_ids, "cuda") + max_abs = float(np.max(np.abs(actual - expected))) + if not np.isfinite(max_abs): + raise AssertionError(f"BF16 rejection evidence is non-finite: {max_abs}") + print(f"BF16 rejection evidence: max_abs={max_abs:.6g} (limit=0.01)") + return max_abs + + def _full_prefill(session, token_ids: list[int]) -> np.ndarray: states = _initial_states(session) output_names = [output.name for output in session.get_outputs()] @@ -425,6 +528,7 @@ def _validate_variant( print(f"{dtype_name}/{ep} provider placement: {placement}") if placement.get("CUDAExecutionProvider", 0) == 0: raise AssertionError(f"No CUDA nodes found in profile: {placement}") + Path(profile_path).unlink(missing_ok=True) del hf_model if device == "cuda": torch.cuda.empty_cache() @@ -441,9 +545,17 @@ def main() -> None: parser.add_argument( "--matrix", nargs="+", - choices=["f32-cpu", "f16-cuda", "bf16-cuda"], + choices=["f32-cpu", "f16-cuda"], default=["f32-cpu", "f16-cuda"], ) + parser.add_argument( + "--bf16-rejection-evidence", + action="store_true", + help=( + "Measure the rejected BF16 CUDA path with a test-only component wrapper; " + "does not produce a supported package." + ), + ) parser.add_argument("--skip-quantization", action="store_true") args = parser.parse_args() @@ -451,6 +563,14 @@ def main() -> None: print(f"Loaded {len(state)} reduced real-weight tensors from revision {REVISION}") output_root = Path(args.output_dir) output_root.mkdir(parents=True, exist_ok=True) + if args.bf16_rejection_evidence: + max_abs = _measure_bf16_rejection(state, output_root) + if max_abs <= 1e-2: + raise AssertionError( + f"BF16 now meets the 1e-2 gate ({max_abs}); revisit the production rejection" + ) + return + variants: dict[str, Path] = {} for variant in args.matrix: dtype_name, device = variant.split("-") diff --git a/pyproject.toml b/pyproject.toml index 756b03aa5..59e91a385 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -75,6 +75,7 @@ markers = [ "arch_validation: L2 architecture validation tests that download real HF configs (no weights) and build full-size ONNX graphs (deselect with '-m \"not arch_validation\"')", "golden: L4 checkpoint-verified golden comparison tests (deselect with '-m \"not golden\"')", "generation: L5 generation end-to-end golden tests (deselect with '-m \"not generation\"')", + "quantization: quantized-package integration tests", ] [tool.mypy] diff --git a/tests/nemotron_h_real_weight_test.py b/tests/nemotron_h_real_weight_test.py index 02ace1451..83641c676 100644 --- a/tests/nemotron_h_real_weight_test.py +++ b/tests/nemotron_h_real_weight_test.py @@ -45,7 +45,7 @@ def _load_validator(): @pytest.fixture(scope="module") -def reduced_real_outputs(tmp_path_factory): +def reduced_real_state(tmp_path_factory): validator = _load_validator() configured_cache = os.environ.get("MOBIUS_NEMOTRON_REDUCED_CACHE") cache = ( @@ -54,6 +54,12 @@ def reduced_real_outputs(tmp_path_factory): else tmp_path_factory.mktemp("nemotron-real") / "reduced.safetensors" ) state = validator._build_reduced_state(cache) + return validator, state + + +@pytest.fixture(scope="module") +def reduced_real_outputs(reduced_real_state, tmp_path_factory): + validator, state = reduced_real_state package = validator._mobius_package(state, dtype_name="f32", ep="cpu") output_dir = tmp_path_factory.mktemp("nemotron-onnx") package.save(output_dir, external_data="onnx") @@ -73,6 +79,33 @@ def reduced_real_outputs(tmp_path_factory): return onnx_logits, onnx_tokens, hf_logits, hf_tokens +def _require_cuda() -> None: + if os.environ.get("MOBIUS_TEST_DEVICE") != "cuda": + pytest.skip("Set MOBIUS_TEST_DEVICE=cuda to run reduced-real CUDA coverage") + if not torch.cuda.is_available(): + pytest.skip("PyTorch CUDA is unavailable") + import onnxruntime as ort + + if hasattr(ort, "preload_dlls"): + ort.preload_dlls() + if "CUDAExecutionProvider" not in ort.get_available_providers(): + pytest.skip("ONNX Runtime CUDAExecutionProvider is unavailable") + + +@pytest.fixture(scope="module") +def reduced_real_fp16_cuda(reduced_real_state, tmp_path_factory): + validator, state = reduced_real_state + _require_cuda() + output_root = tmp_path_factory.mktemp("nemotron-fp16-cuda") + package_dir = validator._validate_variant( + state, + output_root, + dtype_name="f16", + device="cuda", + ) + return validator, state, package_dir + + @pytest.mark.integration @pytest.mark.integration_slow @pytest.mark.golden @@ -115,3 +148,80 @@ def test_nemotron_h_3_5_reduced_real_l5(reduced_real_outputs, model_type): assert len(hf_tokens) == golden["max_new_tokens"] assert hf_tokens == golden["generated_tokens"] assert onnx_tokens == golden["generated_tokens"] + + +@pytest.mark.integration +@pytest.mark.integration_slow +@pytest.mark.golden +@pytest.mark.parametrize("model_type", ["nemotron_h"]) +def test_nemotron_h_3_5_reduced_real_fp16_cuda(reduced_real_fp16_cuda, model_type): + del model_type + _validator, _state, package_dir = reduced_real_fp16_cuda + + assert (package_dir / "model.onnx").is_file() + assert (package_dir / "model.onnx.data").is_file() + + +@pytest.mark.integration +@pytest.mark.integration_slow +@pytest.mark.golden +@pytest.mark.parametrize("model_type", ["nemotron_h"]) +def test_nemotron_h_3_5_bf16_rejection_evidence( + reduced_real_state, + tmp_path, + model_type, +): + del model_type + validator, state = reduced_real_state + _require_cuda() + + max_abs = validator._measure_bf16_rejection(state, tmp_path) + + assert max_abs > 1e-2 + + +@pytest.mark.integration +@pytest.mark.integration_slow +@pytest.mark.generation +@pytest.mark.quantization +@pytest.mark.parametrize("model_type", ["nemotron_h"]) +def test_nemotron_h_3_5_olive_q4_final_package( + reduced_real_fp16_cuda, + tmp_path, + model_type, +): + del model_type + import onnx_ir as ir + + validator, _state, source_dir = reduced_real_fp16_cuda + quantized_dir = validator.quantize_package(source_dir, tmp_path / "q4_k_m-cuda") + + assert (quantized_dir / "model.onnx").is_file() + assert (quantized_dir / "model.onnx.data").is_file() + assert (quantized_dir / "config.json").is_file() + assert sum( + path.stat().st_size for path in quantized_dir.iterdir() if path.is_file() + ) < sum(path.stat().st_size for path in source_dir.iterdir() if path.is_file()) + + quantized_model = ir.load(quantized_dir / "model.onnx") + assert ( + sum( + node.domain == "com.microsoft" and node.op_type == "MatMulNBits" + for node in quantized_model.graph.all_nodes() + ) + == 15 + ) + + generated, logits, profile_path = validator.run_token_ids( + quantized_dir, + [1, 42, 17], + max_new_tokens=4, + device="cuda", + profile=True, + ) + + assert generated == [12, 13, 12, 12] + assert all(np.isfinite(step).all() for step in logits) + assert profile_path is not None + assert validator.summarize_profile(profile_path).get("CUDAExecutionProvider", 0) > 0 + Path(profile_path).unlink(missing_ok=True) From d7df5c8a356177aa0fed421f93db5e73ac6444fd Mon Sep 17 00:00:00 2001 From: Justin Chu Date: Thu, 13 Aug 2026 21:24:50 -0700 Subject: [PATCH 04/13] Enforce strict FP16 logit error Fail reduced-real validation whenever the maximum absolute prefill or cached-step error exceeds the advertised 1e-2 gate, independent of relative tolerance. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Signed-off-by: Justin Chu --- .../nemotron-3_5-lightning-30b/validate_reduced_checkpoint.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/examples/olive/nemotron-3_5-lightning-30b/validate_reduced_checkpoint.py b/examples/olive/nemotron-3_5-lightning-30b/validate_reduced_checkpoint.py index 7bfad5b0b..bef395fb2 100644 --- a/examples/olive/nemotron-3_5-lightning-30b/validate_reduced_checkpoint.py +++ b/examples/olive/nemotron-3_5-lightning-30b/validate_reduced_checkpoint.py @@ -460,6 +460,8 @@ def _assert_logits_close( / (np.linalg.norm(actual) * np.linalg.norm(expected)) ) print(f"{label}: max_abs={max_abs:.6g}, cosine={cosine:.9f}") + if max_abs > atol: + raise AssertionError(f"{label}: max_abs={max_abs:.6g} exceeds {atol=}") np.testing.assert_allclose(actual, expected, rtol=1e-3, atol=atol) From 812895a9f81bcc14095f93289c520d22b5a1b09e Mon Sep 17 00:00:00 2001 From: Justin Chu Date: Thu, 13 Aug 2026 21:31:39 -0700 Subject: [PATCH 05/13] Address Nemotron export review Cache test-case discovery, use generic reduced-real GPU workflow discovery, and emit an honest NemotronH GenAI config without a runtime-version-specific hard rejection. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Signed-off-by: Justin Chu --- .agents/skills/ort-genai-config/SKILL.md | 11 +++++---- .github/workflows/gpu_l4_golden_parity.yml | 4 ++-- .github/workflows/gpu_l5_generation_e2e.yml | 4 ++-- .../nemotron-3_5-lightning-30b/README.md | 12 ++++++---- src/mobius/__main__.py | 23 ------------------- .../integrations/ort_genai/auto_export.py | 13 ----------- .../ort_genai/auto_export_test.py | 20 +++++----------- tests/arch_validation_test.py | 5 ++-- tests/cli_test.py | 23 ------------------- 9 files changed, 26 insertions(+), 89 deletions(-) diff --git a/.agents/skills/ort-genai-config/SKILL.md b/.agents/skills/ort-genai-config/SKILL.md index d644c6a7b..7cd011aed 100644 --- a/.agents/skills/ort-genai-config/SKILL.md +++ b/.agents/skills/ort-genai-config/SKILL.md @@ -291,11 +291,12 @@ global cache-slot count. Preserve intrinsic schema/config validation, but do not gate or reject export based on the current GenAI model registry, runtime version, topology support, or cache executor capability. -NemotronH is one such unsupported contract in ORT GenAI 0.15.2: sparse global -layer indices mix attention key/value caches with Mamba `conv_state` and -`ssm_state`. The generic decoder schema cannot bind all three cache families. -Reject `--runtime ort-genai` before creating output and direct users to a -token-by-token ONNX Runtime loop until the runtime adds a dedicated model type. +NemotronH is unsupported in ORT GenAI 0.15.2 because sparse global layer +indices mix attention key/value caches with Mamba `conv_state` and +`ssm_state`. Do not hard-code a model-type rejection based only on that runtime +version: emit a structurally honest config so downstream releases can evolve, +record the tested-version waiver, and provide a direct ONNX Runtime loop for +current users. ### "input_ids size exceeds max length" diff --git a/.github/workflows/gpu_l4_golden_parity.yml b/.github/workflows/gpu_l4_golden_parity.yml index d05499c50..041a5975f 100644 --- a/.github/workflows/gpu_l4_golden_parity.yml +++ b/.github/workflows/gpu_l4_golden_parity.yml @@ -86,7 +86,7 @@ jobs: AFFECTED='${{ inputs.affected_models }}' if [ -z "$AFFECTED" ] || [ "$AFFECTED" = "[]" ]; then echo "Running all L4 golden comparison tests" - pytest tests/e2e_golden_test.py tests/nemotron_h_real_weight_test.py \ + pytest tests/e2e_golden_test.py tests/*_real_weight_test.py \ -m golden \ -v \ --timeout=300 \ @@ -98,7 +98,7 @@ jobs: # Convert JSON array to comma-separated list for --models MODELS=$(echo "$AFFECTED" | python -c "import json, sys; print(','.join(json.load(sys.stdin)))") if [ -n "$MODELS" ]; then - pytest tests/e2e_golden_test.py tests/nemotron_h_real_weight_test.py \ + pytest tests/e2e_golden_test.py tests/*_real_weight_test.py \ -m golden \ -v \ --models "$MODELS" \ diff --git a/.github/workflows/gpu_l5_generation_e2e.yml b/.github/workflows/gpu_l5_generation_e2e.yml index 9183faae7..bd1cd6211 100644 --- a/.github/workflows/gpu_l5_generation_e2e.yml +++ b/.github/workflows/gpu_l5_generation_e2e.yml @@ -75,7 +75,7 @@ jobs: AFFECTED='${{ inputs.affected_models }}' if [ -z "$AFFECTED" ] || [ "$AFFECTED" = "[]" ]; then echo "Running all L5 generation E2E tests" - pytest tests/e2e_golden_test.py tests/nemotron_h_real_weight_test.py \ + pytest tests/e2e_golden_test.py tests/*_real_weight_test.py \ -m generation \ -v \ --timeout=300 \ @@ -87,7 +87,7 @@ jobs: # Convert JSON array to comma-separated list for --models MODELS=$(echo "$AFFECTED" | python -c "import json, sys; print(','.join(json.load(sys.stdin)))") if [ -n "$MODELS" ]; then - pytest tests/e2e_golden_test.py tests/nemotron_h_real_weight_test.py \ + pytest tests/e2e_golden_test.py tests/*_real_weight_test.py \ -m generation \ -v \ --models "$MODELS" \ diff --git a/examples/olive/nemotron-3_5-lightning-30b/README.md b/examples/olive/nemotron-3_5-lightning-30b/README.md index c390d1d58..7abbb4f40 100644 --- a/examples/olive/nemotron-3_5-lightning-30b/README.md +++ b/examples/olive/nemotron-3_5-lightning-30b/README.md @@ -21,9 +21,10 @@ The checkpoint is a real `nemotron_h` model, not an alias: ONNX Runtime GenAI 0.15.2 cannot bind this model's mixed cache: Mamba layers need `conv_state` plus `ssm_state`, while sparse full-attention layers need -key/value caches at global layer indices. Mobius therefore rejects -`--runtime ort-genai` before writing artifacts. The package uses direct ONNX -Runtime generation through `inference.py`. +key/value caches at global layer indices. Mobius still emits a structurally +honest config rather than hard-coding a version-specific rejection, allowing +future runtime releases to add support. This validated package uses direct +ONNX Runtime generation through `inference.py`. ## Install @@ -96,8 +97,9 @@ output/ └── source_manifest.json ``` -No `genai_config.json` is emitted because that would claim unsupported ORT -GenAI runtime compatibility. +This recipe intentionally omits `genai_config.json`: direct ONNX Runtime is +the validated runtime for ORT GenAI 0.15.2. Core Mobius config emission remains +available for testing future runtime releases. ## Direct generation and profiling diff --git a/src/mobius/__main__.py b/src/mobius/__main__.py index 4ceb6cf68..bed089bcb 100644 --- a/src/mobius/__main__.py +++ b/src/mobius/__main__.py @@ -307,7 +307,6 @@ def _resolve_static_cache_task(model_type: str) -> ModelTask: config_path, trust_remote_code=trust_remote_code ) model_type = hf_config.model_type - _preflight_ort_genai_runtime(args, model_type) parent_config = hf_config if hasattr(hf_config, "text_config"): hf_config = hf_config.text_config @@ -338,14 +337,6 @@ def _resolve_static_cache_task(model_type: str) -> ModelTask: pkg.apply_weights(state_dict) else: model_id_or_path = args.model - if getattr(args, "runtime", None) == "ort-genai": - import transformers - - runtime_config = transformers.AutoConfig.from_pretrained( - model_id_or_path, - trust_remote_code=trust_remote_code, - ) - _preflight_ort_genai_runtime(args, runtime_config.model_type) if static_cache_params is not None: # Detect model type to resolve the correct static cache task. import transformers @@ -371,20 +362,6 @@ def _resolve_static_cache_task(model_type: str) -> ModelTask: _save_package(pkg, output_dir, args, optimize, component_filter) -def _preflight_ort_genai_runtime(args, model_type: str) -> None: - """Reject known unsupported ORT GenAI types before weights are downloaded.""" - if getattr(args, "runtime", None) != "ort-genai": - return - from mobius.integrations.ort_genai.auto_export import ( - _validate_ort_genai_model_type, - ) - - try: - _validate_ort_genai_model_type(model_type) - except ValueError as error: - raise SystemExit(f"Error: {error}") from error - - def _save_package( pkg, output_dir: str, args, optimize: str | None, component_filter: str | None ) -> None: diff --git a/src/mobius/integrations/ort_genai/auto_export.py b/src/mobius/integrations/ort_genai/auto_export.py index a84b7f8e2..35a611c9c 100644 --- a/src/mobius/integrations/ort_genai/auto_export.py +++ b/src/mobius/integrations/ort_genai/auto_export.py @@ -1237,22 +1237,9 @@ def _write_genai_config( return generator.write(output_dir) -def _validate_ort_genai_model_type(model_type: str | None) -> None: - """Reject model types whose runtime contract is known before graph build.""" - if model_type == "nemotron_h": - raise ValueError( - "onnxruntime-genai does not currently support NemotronH's mixed cache " - "contract: Mamba layers require conv_state + ssm_state while full-attention " - "layers require key/value caches at sparse global layer indices. Export " - "without --runtime ort-genai and run the saved model directly with ONNX " - "Runtime (see examples/nemotron_3_nano_text_generation.py)." - ) - - def _validate_ort_genai_compatibility(pkg: ModelPackage) -> None: """Reject packages whose required inputs cannot be supplied by ORT GenAI.""" config = getattr(pkg, "config", None) - _validate_ort_genai_model_type(getattr(config, "model_type", None)) if getattr(config, "model_type", None) == "parakeet_ctc": raise ValueError( "ORT GenAI does not define a feature-input CTC ASR pipeline; " diff --git a/src/mobius/integrations/ort_genai/auto_export_test.py b/src/mobius/integrations/ort_genai/auto_export_test.py index a74f6d1e3..6886b93f8 100644 --- a/src/mobius/integrations/ort_genai/auto_export_test.py +++ b/src/mobius/integrations/ort_genai/auto_export_test.py @@ -1103,22 +1103,14 @@ class FakeConfig: write_ort_genai_config(pkg, str(tmp_path)) assert not (tmp_path / "genai_config.json").exists() - def test_rejects_nemotron_h_mixed_cache_before_writing(self, tmp_path): - import dataclasses - - from mobius._model_package import ModelPackage - - @dataclasses.dataclass - class FakeConfig: - model_type: str = "nemotron_h" + def test_nemotron_h_config_is_emitted_for_future_runtime_support(self, tmp_path): + pkg = _make_fake_llm_pkg("nemotron_h") - pkg = ModelPackage(config=FakeConfig()) - output_dir = tmp_path / "ort-genai" - - with pytest.raises(ValueError, match=r"NemotronH.*conv_state \+ ssm_state.*key/value"): - write_ort_genai_config(pkg, str(output_dir)) + result = write_ort_genai_config(pkg, str(tmp_path)) - assert not output_dir.exists() + with open(result["genai_config"], encoding="utf-8") as config_file: + generated = json.load(config_file) + assert generated["model"]["type"] == "nemotron_h" def test_processor_config_written_with_vision(self, tmp_path): """image_processor.json is written when pkg.config.vision is set.""" diff --git a/tests/arch_validation_test.py b/tests/arch_validation_test.py index 35ca29c3a..81558b93c 100644 --- a/tests/arch_validation_test.py +++ b/tests/arch_validation_test.py @@ -43,8 +43,9 @@ logger = logging.getLogger(__name__) -_PINNED_REVISIONS = {case.model_id: case.revision for case in discover_test_cases()} -_DECLARED_DTYPES = {case.model_id: case.dtype for case in discover_test_cases()} +_TEST_CASES = discover_test_cases() +_PINNED_REVISIONS = {case.model_id: case.revision for case in _TEST_CASES} +_DECLARED_DTYPES = {case.model_id: case.dtype for case in _TEST_CASES} # Build parametrized test cases from registry entries that have a test_model_id. # diff --git a/tests/cli_test.py b/tests/cli_test.py index 152fe0932..07c288249 100644 --- a/tests/cli_test.py +++ b/tests/cli_test.py @@ -557,29 +557,6 @@ def test_runtime_ort_genai_rejects_mage_vl_before_saving(self): save.assert_not_called() config_writer.assert_not_called() - def test_runtime_ort_genai_rejects_nemotron_h_before_build(self): - with ( - tempfile.TemporaryDirectory() as tmpdir, - mock.patch( - "transformers.AutoConfig.from_pretrained", - return_value=SimpleNamespace(model_type="nemotron_h"), - ), - mock.patch("mobius.__main__.build") as build_model, - pytest.raises(SystemExit, match=r"NemotronH.*mixed cache"), - ): - main( - [ - "build", - "--model", - "nvidia/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-BF16", - tmpdir, - "--runtime", - "ort-genai", - ] - ) - - build_model.assert_not_called() - def test_runtime_onnx_genai_uses_native_vlm_emitter(self): pkg = mock.MagicMock() pkg.items.return_value = [] From 96635c938f9562b8e6c5cc897be4aad050019c2e Mon Sep 17 00:00:00 2001 From: Justin Chu Date: Thu, 13 Aug 2026 21:54:14 -0700 Subject: [PATCH 06/13] Target real-weight GPU coverage Teach affected-model detection about reduced-real tests and example assets, and stop direct generation on configured EOS without an unused final decoder invocation. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Signed-off-by: Justin Chu --- .../nemotron-3_5-lightning-30b/inference.py | 24 +++++++- .../nemotron-3_5-lightning-30b/optimize.py | 3 +- .../validate_reduced_checkpoint.py | 3 + scripts/detect_affected_models.py | 22 ++++++- scripts/detect_affected_models_test.py | 21 +++++++ tests/nemotron_h_real_weight_test.py | 59 +++++++++++++++++++ 6 files changed, 129 insertions(+), 3 deletions(-) diff --git a/examples/olive/nemotron-3_5-lightning-30b/inference.py b/examples/olive/nemotron-3_5-lightning-30b/inference.py index e5bc0f105..f4b83337e 100644 --- a/examples/olive/nemotron-3_5-lightning-30b/inference.py +++ b/examples/olive/nemotron-3_5-lightning-30b/inference.py @@ -9,6 +9,7 @@ import argparse import json from collections import Counter +from collections.abc import Collection from pathlib import Path from typing import Any @@ -153,6 +154,22 @@ def _create_session(model_path: Path, device: str, profile: bool): return session +def load_eos_token_ids(model_dir: str | Path) -> set[int]: + """Load scalar or list EOS IDs from the assembled package metadata.""" + model_dir = Path(model_dir) + eos_ids: set[int] = set() + for filename in ("generation_config.json", "config.json"): + path = model_dir / filename + if not path.is_file(): + continue + raw_eos = json.loads(path.read_text(encoding="utf-8")).get("eos_token_id") + if isinstance(raw_eos, int): + eos_ids.add(raw_eos) + elif isinstance(raw_eos, list): + eos_ids.update(value for value in raw_eos if isinstance(value, int)) + return eos_ids + + def run_token_ids( model_dir: str | Path, input_ids: list[int], @@ -160,6 +177,7 @@ def run_token_ids( max_new_tokens: int, device: str, profile: bool = False, + eos_token_ids: Collection[int] | None = None, ) -> tuple[list[int], list[np.ndarray], str | None]: """Run token-by-token hybrid-cache generation and return IDs plus logits.""" model_path = Path(model_dir) / "model.onnx" @@ -190,10 +208,13 @@ def run_token_ids( assert outputs is not None logits = _as_numpy(outputs[output_names.index("logits")])[0, -1].astype(np.float32) - for _ in range(max_new_tokens): + eos_ids = set(eos_token_ids or ()) + for token_index in range(max_new_tokens): logits_by_step.append(logits.copy()) token_id = int(np.argmax(logits)) generated.append(token_id) + if token_id in eos_ids or token_index + 1 == max_new_tokens: + break feeds = _token_feeds( session, np.array([[token_id]], dtype=np.int64), @@ -265,6 +286,7 @@ def main() -> None: max_new_tokens=args.max_new_tokens, device=args.device, profile=args.profile, + eos_token_ids=load_eos_token_ids(model_dir), ) if not generated: raise RuntimeError("Generation produced no tokens") diff --git a/examples/olive/nemotron-3_5-lightning-30b/optimize.py b/examples/olive/nemotron-3_5-lightning-30b/optimize.py index 4bba51c35..38001e09c 100644 --- a/examples/olive/nemotron-3_5-lightning-30b/optimize.py +++ b/examples/olive/nemotron-3_5-lightning-30b/optimize.py @@ -12,7 +12,7 @@ import tempfile from pathlib import Path -from inference import MODEL_ID, REVISION, run_token_ids +from inference import MODEL_ID, REVISION, load_eos_token_ids, run_token_ids _METADATA_FILES = { "added_tokens.json", @@ -209,6 +209,7 @@ def smoke_test(model_dir: str | Path, *, device: str) -> list[int]: [1, 42, 17], max_new_tokens=4, device=device, + eos_token_ids=load_eos_token_ids(model_dir), ) if len(generated) != 4 or any(not np.isfinite(step).all() for step in logits): raise RuntimeError(f"Quantized generation smoke test failed: {generated}") diff --git a/examples/olive/nemotron-3_5-lightning-30b/validate_reduced_checkpoint.py b/examples/olive/nemotron-3_5-lightning-30b/validate_reduced_checkpoint.py index bef395fb2..8b0f92324 100644 --- a/examples/olive/nemotron-3_5-lightning-30b/validate_reduced_checkpoint.py +++ b/examples/olive/nemotron-3_5-lightning-30b/validate_reduced_checkpoint.py @@ -24,6 +24,7 @@ _initial_states, _run_session, _token_feeds, + load_eos_token_ids, run_token_ids, summarize_profile, ) @@ -502,6 +503,7 @@ def _validate_variant( max_new_tokens=4, device=device, profile=profile, + eos_token_ids=load_eos_token_ids(variant_dir), ) if len(generated) != 4 or any(not np.isfinite(step).all() for step in logits): raise AssertionError(f"Invalid generation for {dtype_name}/{ep}: {generated}") @@ -596,6 +598,7 @@ def main() -> None: max_new_tokens=4, device="cuda", profile=False, + eos_token_ids=load_eos_token_ids(quantized), ) if len(generated) != 4 or any(not np.isfinite(step).all() for step in logits): raise AssertionError(f"Quantized generation failed: {generated}") diff --git a/scripts/detect_affected_models.py b/scripts/detect_affected_models.py index 55494f0cf..779b25317 100644 --- a/scripts/detect_affected_models.py +++ b/scripts/detect_affected_models.py @@ -53,6 +53,25 @@ "src/mobius/tasks/", ) +# Model-specific examples whose integration tests depend on files outside src/. +_MODEL_PATH_HINTS: tuple[tuple[str, tuple[str, ...]], ...] = ( + ("examples/olive/nemotron-3_5-lightning-30b/", ("nemotron_h",)), + ("testdata/cases/causal-lm/nemotron-3_5-lightning-30b.yaml", ("nemotron_h",)), + ("testdata/golden/causal-lm/nemotron-3_5-lightning-30b-", ("nemotron_h",)), +) + + +def _model_type_hints(path: str) -> set[str]: + """Infer targeted model types from real-weight tests and example assets.""" + normalized = path.replace("\\", "/") + if normalized.startswith("tests/") and normalized.endswith("_real_weight_test.py"): + filename = normalized.rsplit("/", 1)[-1] + return {filename[: -len("_real_weight_test.py")]} + for prefix, model_types in _MODEL_PATH_HINTS: + if normalized.startswith(prefix): + return set(model_types) + return set() + def classify_file(path: str) -> str: """Classify a changed file path. @@ -435,6 +454,7 @@ def detect_affected_models( model_files: list[str] = [] traceable_files: list[str] = [] for path in changed_files: + affected.update(_model_type_hints(path)) category = classify_file(path) if category == "shared_infra": run_all = True @@ -462,7 +482,7 @@ def detect_affected_models( return {"affected": [], "run_all": True} if not model_files and not traceable_files: - return {"affected": [], "run_all": False} + return {"affected": sorted(affected), "run_all": False} # Build the registry map: source_module → [model_types] registry_map = _build_source_module_to_types() diff --git a/scripts/detect_affected_models_test.py b/scripts/detect_affected_models_test.py index a472f056f..1444621c2 100644 --- a/scripts/detect_affected_models_test.py +++ b/scripts/detect_affected_models_test.py @@ -26,6 +26,7 @@ _build_registry_class_to_types, _build_source_module_to_types, _find_reverse_dependents, + _model_type_hints, _parse_imports, classify_file, detect_affected_models, @@ -84,6 +85,16 @@ def test_windows_paths(self): assert classify_file("src\\mobius\\models\\falcon.py") == "model" +class TestModelTypeHints: + def test_real_weight_test_infers_model_type(self): + assert _model_type_hints("tests/nemotron_h_real_weight_test.py") == {"nemotron_h"} + + def test_model_example_uses_explicit_mapping(self): + assert _model_type_hints("examples/olive/nemotron-3_5-lightning-30b/optimize.py") == { + "nemotron_h" + } + + # ---------------------------------------------------------------- # AST registry parsing tests # ---------------------------------------------------------------- @@ -233,6 +244,16 @@ def test_test_file_no_affected(self): assert result["run_all"] is False assert result["affected"] == [] + def test_real_weight_test_targets_its_model(self): + result = detect_affected_models(["tests/nemotron_h_real_weight_test.py"]) + assert result == {"affected": ["nemotron_h"], "run_all": False} + + def test_model_example_targets_its_integration_tests(self): + result = detect_affected_models( + ["examples/olive/nemotron-3_5-lightning-30b/inference.py"] + ) + assert result == {"affected": ["nemotron_h"], "run_all": False} + def test_falcon_model_file(self): result = detect_affected_models(["src/mobius/models/falcon.py"]) assert result["run_all"] is False diff --git a/tests/nemotron_h_real_weight_test.py b/tests/nemotron_h_real_weight_test.py index 83641c676..9b3691cee 100644 --- a/tests/nemotron_h_real_weight_test.py +++ b/tests/nemotron_h_real_weight_test.py @@ -44,6 +44,63 @@ def _load_validator(): sys.path.pop(0) +def test_load_eos_token_ids_unions_generation_and_model_config(tmp_path): + validator = _load_validator() + (tmp_path / "generation_config.json").write_text( + json.dumps({"eos_token_id": [2, 11]}), + encoding="utf-8", + ) + (tmp_path / "config.json").write_text( + json.dumps({"eos_token_id": 2}), + encoding="utf-8", + ) + + assert validator.load_eos_token_ids(tmp_path) == {2, 11} + + +def test_run_token_ids_stops_on_eos_without_unused_forward(tmp_path, monkeypatch): + validator = _load_validator() + inference_globals = validator.run_token_ids.__globals__ + (tmp_path / "model.onnx").touch() + + class _Output: + name = "logits" + + class _Session: + @staticmethod + def get_inputs(): + return [] + + @staticmethod + def get_outputs(): + return [_Output()] + + calls = [] + + def _run_session(_session, _output_names, _feeds): + calls.append(1) + logits = np.zeros((1, 1, 16), dtype=np.float32) + logits[0, 0, 2] = 1.0 + return [logits] + + monkeypatch.setitem(inference_globals, "_create_session", lambda *_args: _Session()) + monkeypatch.setitem(inference_globals, "_initial_states", lambda _session: {}) + monkeypatch.setitem(inference_globals, "_run_session", _run_session) + + generated, logits, profile = validator.run_token_ids( + tmp_path, + [1], + max_new_tokens=4, + device="cpu", + eos_token_ids={2}, + ) + + assert generated == [2] + assert len(logits) == 1 + assert len(calls) == 1 + assert profile is None + + @pytest.fixture(scope="module") def reduced_real_state(tmp_path_factory): validator = _load_validator() @@ -71,6 +128,7 @@ def reduced_real_outputs(reduced_real_state, tmp_path_factory): prompt_ids, max_new_tokens=4, device="cpu", + eos_token_ids=validator.load_eos_token_ids(output_dir), ) hf_model = validator._hf_model(state, dtype=torch.float32, device="cpu") @@ -218,6 +276,7 @@ def test_nemotron_h_3_5_olive_q4_final_package( max_new_tokens=4, device="cuda", profile=True, + eos_token_ids=validator.load_eos_token_ids(quantized_dir), ) assert generated == [12, 13, 12, 12] From 31b46b37ec6b3c3de37e27534e110f2a793197d2 Mon Sep 17 00:00:00 2001 From: Justin Chu Date: Thu, 13 Aug 2026 21:58:04 -0700 Subject: [PATCH 07/13] Clean up validation profiles Profile only the cached-generation session used for placement evidence so prefill session teardown does not leave untracked ORT profile files. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Signed-off-by: Justin Chu --- .../nemotron-3_5-lightning-30b/validate_reduced_checkpoint.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/examples/olive/nemotron-3_5-lightning-30b/validate_reduced_checkpoint.py b/examples/olive/nemotron-3_5-lightning-30b/validate_reduced_checkpoint.py index 8b0f92324..35a5df02e 100644 --- a/examples/olive/nemotron-3_5-lightning-30b/validate_reduced_checkpoint.py +++ b/examples/olive/nemotron-3_5-lightning-30b/validate_reduced_checkpoint.py @@ -489,7 +489,7 @@ def _validate_variant( _hf_config().save_pretrained(variant_dir) profile = device == "cuda" - session = _create_session(variant_dir / "model.onnx", device, profile) + session = _create_session(variant_dir / "model.onnx", device, False) prompt_ids = [1, 42, 17] actual = _full_prefill(session, prompt_ids) hf_model = _hf_model(state, dtype=torch_dtype, device=device) From 2b1a3c30435b41d6dd669769b28fccab55115601 Mon Sep 17 00:00:00 2001 From: Justin Chu Date: Thu, 13 Aug 2026 22:16:08 -0700 Subject: [PATCH 08/13] Harden runtime and fixture contracts Reject structurally unrepresentable SSM cache graphs before artifact creation and persist the pinned reduced-real fixture with validated range retries, atomic schema-versioned writes, and shared L4/L5 CI caching. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Signed-off-by: Justin Chu --- .agents/skills/ort-genai-config/SKILL.md | 10 +- .github/workflows/gpu_l4_golden_parity.yml | 8 ++ .github/workflows/gpu_l5_generation_e2e.yml | 8 ++ .../nemotron-3_5-lightning-30b/README.md | 19 ++-- .../validate_reduced_checkpoint.py | 88 +++++++++++++---- .../integrations/ort_genai/auto_export.py | 18 ++++ .../ort_genai/auto_export_test.py | 34 +++++-- tests/nemotron_h_real_weight_test.py | 96 ++++++++++++++++++- 8 files changed, 243 insertions(+), 38 deletions(-) diff --git a/.agents/skills/ort-genai-config/SKILL.md b/.agents/skills/ort-genai-config/SKILL.md index 7cd011aed..7319f4617 100644 --- a/.agents/skills/ort-genai-config/SKILL.md +++ b/.agents/skills/ort-genai-config/SKILL.md @@ -293,10 +293,12 @@ version, topology support, or cache executor capability. NemotronH is unsupported in ORT GenAI 0.15.2 because sparse global layer indices mix attention key/value caches with Mamba `conv_state` and -`ssm_state`. Do not hard-code a model-type rejection based only on that runtime -version: emit a structurally honest config so downstream releases can evolve, -record the tested-version waiver, and provide a direct ONNX Runtime loop for -current users. +`ssm_state`. Do not hard-code a model-type or runtime-version rejection. +Inspect the actual decoder graph and reject only when it requires cache inputs +the generated config has no template for (currently `ssm_state`). This keeps +the guard structural: remove it when config emission can represent that graph +contract. Record the tested-version waiver and provide a direct ONNX Runtime +loop for current users. ### "input_ids size exceeds max length" diff --git a/.github/workflows/gpu_l4_golden_parity.yml b/.github/workflows/gpu_l4_golden_parity.yml index 041a5975f..164214838 100644 --- a/.github/workflows/gpu_l4_golden_parity.yml +++ b/.github/workflows/gpu_l4_golden_parity.yml @@ -42,6 +42,8 @@ jobs: - "1ES.Pool=onnxruntime-ep-mobius-github-linux-a10" - "JobId=gpu-golden-${{ github.run_id }}-${{ github.run_number }}-${{ github.run_attempt }}" timeout-minutes: 60 + env: + MOBIUS_NEMOTRON_REDUCED_CACHE: ${{ runner.temp }}/mobius-nemotron-cache/reduced-d468880b6ad3c6e0d21377ce7242adaea4cc884d-schema-v1.safetensors steps: - uses: actions/checkout@v7 @@ -56,6 +58,12 @@ jobs: path: ~/.cache/huggingface key: hf-gpu-${{ hashFiles('testdata/cases/**/*.yaml') }} + - name: Cache reduced Nemotron fixture + uses: actions/cache@v6 + with: + path: ${{ runner.temp }}/mobius-nemotron-cache + key: nemotron-reduced-d468880b6ad3c6e0d21377ce7242adaea4cc884d-schema-v1 + - name: Install PyTorch (CUDA) run: pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu128 diff --git a/.github/workflows/gpu_l5_generation_e2e.yml b/.github/workflows/gpu_l5_generation_e2e.yml index bd1cd6211..47be44cad 100644 --- a/.github/workflows/gpu_l5_generation_e2e.yml +++ b/.github/workflows/gpu_l5_generation_e2e.yml @@ -42,6 +42,8 @@ jobs: - "1ES.Pool=onnxruntime-ep-mobius-github-linux-a10" - "JobId=gpu-l5-${{ github.run_id }}-${{ github.run_number }}-${{ github.run_attempt }}" timeout-minutes: 60 + env: + MOBIUS_NEMOTRON_REDUCED_CACHE: ${{ runner.temp }}/mobius-nemotron-cache/reduced-d468880b6ad3c6e0d21377ce7242adaea4cc884d-schema-v1.safetensors steps: - uses: actions/checkout@v7 @@ -56,6 +58,12 @@ jobs: path: ~/.cache/huggingface key: hf-gpu-${{ hashFiles('testdata/cases/**/*.yaml') }} + - name: Cache reduced Nemotron fixture + uses: actions/cache@v6 + with: + path: ${{ runner.temp }}/mobius-nemotron-cache + key: nemotron-reduced-d468880b6ad3c6e0d21377ce7242adaea4cc884d-schema-v1 + - name: Install PyTorch (CUDA) run: pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu128 diff --git a/examples/olive/nemotron-3_5-lightning-30b/README.md b/examples/olive/nemotron-3_5-lightning-30b/README.md index 7abbb4f40..8a9771442 100644 --- a/examples/olive/nemotron-3_5-lightning-30b/README.md +++ b/examples/olive/nemotron-3_5-lightning-30b/README.md @@ -21,10 +21,11 @@ The checkpoint is a real `nemotron_h` model, not an alias: ONNX Runtime GenAI 0.15.2 cannot bind this model's mixed cache: Mamba layers need `conv_state` plus `ssm_state`, while sparse full-attention layers need -key/value caches at global layer indices. Mobius still emits a structurally -honest config rather than hard-coding a version-specific rejection, allowing -future runtime releases to add support. This validated package uses direct -ONNX Runtime generation through `inference.py`. +key/value caches at global layer indices. Mobius rejects this graph contract +structurally because `genai_config.json` has no `ssm_state` input/output +template; it does not hard-code a NemotronH or runtime-version check. The guard +can be removed when config emission represents that state. This validated +package uses direct ONNX Runtime generation through `inference.py`. ## Install @@ -98,8 +99,8 @@ output/ ``` This recipe intentionally omits `genai_config.json`: direct ONNX Runtime is -the validated runtime for ORT GenAI 0.15.2. Core Mobius config emission remains -available for testing future runtime releases. +the validated runtime, and core Mobius rejects the currently unrepresentable +SSM cache contract before creating runtime artifacts. ## Direct generation and profiling @@ -131,6 +132,12 @@ weights while retaining production dimensions: python validate_reduced_checkpoint.py ``` +The fixture is stored persistently under `~/.cache/mobius/` by default. Each +range request validates status, `Content-Range`, declared length, and payload +length, with three bounded attempts (1s then 2s backoff). The cache metadata +must match the pinned model, revision, and fixture schema; writes are atomic. +GPU CI restores the same revision/schema-keyed cache for L4 and L5. + The supported matrix is intentionally limited to FP32/CPU and FP16/CUDA. Reproduce the BF16 rejection evidence separately without creating a supported package or weakening the production guard: diff --git a/examples/olive/nemotron-3_5-lightning-30b/validate_reduced_checkpoint.py b/examples/olive/nemotron-3_5-lightning-30b/validate_reduced_checkpoint.py index 35a5df02e..e51049b32 100644 --- a/examples/olive/nemotron-3_5-lightning-30b/validate_reduced_checkpoint.py +++ b/examples/olive/nemotron-3_5-lightning-30b/validate_reduced_checkpoint.py @@ -10,10 +10,12 @@ 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 ( @@ -35,6 +37,8 @@ _VOCAB_SIZE = 256 _NUM_EXPERTS = 4 _LAYER_REMAP = {0: 0, 1: 1, 5: 2} +FIXTURE_SCHEMA_VERSION = 1 +_RANGE_ATTEMPTS = 3 _DTYPES = { "f32": (torch.float32, "FLOAT"), "f16": (torch.float16, "FLOAT16"), @@ -45,8 +49,6 @@ class _PinnedSafetensors: """Read selected tensors via HTTP Range without downloading 65.8 GB.""" def __init__(self) -> None: - import requests - index_path = hf_hub_download( MODEL_ID, "model.safetensors.index.json", @@ -61,18 +63,42 @@ 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: - response = self._session.get( - self._url(shard), - headers={"Range": f"bytes={start}-{end}"}, - timeout=180, - ) expected = end - start + 1 - if response.status_code != 206 or len(response.content) != expected: - raise RuntimeError( - f"Range fetch failed for {shard} bytes {start}-{end}: " - f"status={response.status_code}, bytes={len(response.content)}" - ) - return response.content + expected_range_prefix = f"bytes {start}-{end}/" + last_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") + if response.status_code != 206: + last_error = f"status={response.status_code}" + elif not content_range.startswith(expected_range_prefix): + last_error = f"content-range={content_range!r}" + elif content_length is not None and ( + not content_length.isdecimal() or int(content_length) != expected + ): + last_error = f"content-length={content_length}, expected={expected}" + else: + payload = response.content + if len(payload) == expected: + return payload + last_error = f"bytes={len(payload)}, expected={expected}" + except requests.RequestException as error: + last_error = f"{type(error).__name__}: {error}" + + if attempt + 1 < _RANGE_ATTEMPTS: + time.sleep(2**attempt) + + raise RuntimeError( + f"Range fetch failed after {_RANGE_ATTEMPTS} attempts for " + f"{shard} bytes {start}-{end}: {last_error}" + ) def _header(self, shard: str) -> tuple[int, dict]: if shard not in self._headers: @@ -117,13 +143,32 @@ def _source_to_target(name: str) -> str: return ".".join(parts) +def default_reduced_cache_path() -> Path: + """Return the persistent, revision-and-schema-keyed fixture cache path.""" + return ( + Path.home() + / ".cache" + / "mobius" + / "nemotron-3_5-lightning" + / f"reduced-{REVISION}-schema-v{FIXTURE_SCHEMA_VERSION}.safetensors" + ) + + def _build_reduced_state(cache_path: Path) -> dict[str, torch.Tensor]: if cache_path.is_file(): with safe_open(cache_path, framework="pt") as cached: metadata = cached.metadata() or {} - if metadata.get("revision") != REVISION: + expected_metadata = { + "model_id": MODEL_ID, + "revision": REVISION, + "fixture_schema": str(FIXTURE_SCHEMA_VERSION), + } + actual_metadata = {key: metadata.get(key) for key in expected_metadata} + if actual_metadata != expected_metadata: raise ValueError( - f"Reduced cache revision mismatch: {metadata.get('revision')} != {REVISION}" + "Reduced cache metadata mismatch: " + f"expected={expected_metadata}, actual={actual_metadata}. " + "Remove the stale cache file and retry." ) return load_file(cache_path) @@ -165,11 +210,18 @@ def _build_reduced_state(cache_path: Path) -> dict[str, torch.Tensor]: ) cache_path.parent.mkdir(parents=True, exist_ok=True) + temporary_path = cache_path.with_name(f"{cache_path.name}.tmp") + temporary_path.unlink(missing_ok=True) save_file( {name: tensor.contiguous() for name, tensor in state.items()}, - cache_path, - metadata={"model_id": MODEL_ID, "revision": REVISION}, + temporary_path, + metadata={ + "model_id": MODEL_ID, + "revision": REVISION, + "fixture_schema": str(FIXTURE_SCHEMA_VERSION), + }, ) + temporary_path.replace(cache_path) return state @@ -543,7 +595,7 @@ def main() -> None: parser = argparse.ArgumentParser(description=__doc__) parser.add_argument( "--cache", - default="cache/nemotron-3_5-reduced-real.safetensors", + default=default_reduced_cache_path(), ) parser.add_argument("--output-dir", default="output/reduced-validation") parser.add_argument( diff --git a/src/mobius/integrations/ort_genai/auto_export.py b/src/mobius/integrations/ort_genai/auto_export.py index 35a611c9c..ba757f2d5 100644 --- a/src/mobius/integrations/ort_genai/auto_export.py +++ b/src/mobius/integrations/ort_genai/auto_export.py @@ -1240,6 +1240,24 @@ def _write_genai_config( def _validate_ort_genai_compatibility(pkg: ModelPackage) -> None: """Reject packages whose required inputs cannot be supplied by ORT GenAI.""" config = getattr(pkg, "config", None) + decoder_key = "decoder" if "decoder" in pkg else "model" + decoder_model = pkg.get(decoder_key) + if decoder_model is not None: + cache_suffixes = { + parts[2] + for model_input in decoder_model.graph.inputs + if model_input.name is not None + and len(parts := model_input.name.split(".")) == 3 + and parts[0] == "past_key_values" + and parts[1].isdigit() + } + if "ssm_state" in cache_suffixes: + raise ValueError( + "ORT GenAI config generation cannot represent this decoder's cache " + f"inputs {sorted(cache_suffixes)}: no input/output template exists for " + "ssm_state. Export without --runtime ort-genai and run the ONNX model " + "directly until the config schema and runtime support this graph contract." + ) if getattr(config, "model_type", None) == "parakeet_ctc": raise ValueError( "ORT GenAI does not define a feature-input CTC ASR pipeline; " diff --git a/src/mobius/integrations/ort_genai/auto_export_test.py b/src/mobius/integrations/ort_genai/auto_export_test.py index 6886b93f8..467debd1e 100644 --- a/src/mobius/integrations/ort_genai/auto_export_test.py +++ b/src/mobius/integrations/ort_genai/auto_export_test.py @@ -1103,14 +1103,36 @@ class FakeConfig: write_ort_genai_config(pkg, str(tmp_path)) assert not (tmp_path / "genai_config.json").exists() - def test_nemotron_h_config_is_emitted_for_future_runtime_support(self, tmp_path): - pkg = _make_fake_llm_pkg("nemotron_h") + def test_rejects_unrepresentable_mixed_cache_graph_before_writing(self, tmp_path): + import dataclasses - result = write_ort_genai_config(pkg, str(tmp_path)) + from mobius._model_package import ModelPackage + + @dataclasses.dataclass + class FakeConfig: + model_type: str = "future_hybrid" + + pkg = ModelPackage( + { + "model": _mock_model( + inputs=[ + "input_ids", + "attention_mask", + "past_key_values.0.conv_state", + "past_key_values.0.ssm_state", + "past_key_values.2.key", + "past_key_values.2.value", + ] + ) + }, + config=FakeConfig(), + ) + output_dir = tmp_path / "ort-genai" + + with pytest.raises(ValueError, match=r"no input/output template.*ssm_state"): + write_ort_genai_config(pkg, str(output_dir)) - with open(result["genai_config"], encoding="utf-8") as config_file: - generated = json.load(config_file) - assert generated["model"]["type"] == "nemotron_h" + assert not output_dir.exists() def test_processor_config_written_with_vision(self, tmp_path): """image_processor.json is written when pkg.config.vision is set.""" diff --git a/tests/nemotron_h_real_weight_test.py b/tests/nemotron_h_real_weight_test.py index 9b3691cee..cc17abdda 100644 --- a/tests/nemotron_h_real_weight_test.py +++ b/tests/nemotron_h_real_weight_test.py @@ -101,14 +101,102 @@ def _run_session(_session, _output_names, _feeds): assert profile is None +def test_pinned_range_fetch_retries_and_validates_headers(monkeypatch): + validator = _load_validator() + + class _Response: + def __init__(self, status, headers, content=b""): + self.status_code = status + self.headers = headers + self.content = 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 1-4/10", "Content-Length": "4"}), + _Response( + 206, {"Content-Range": "bytes 0-3/10", "Content-Length": "4"}, b"data" + ), + ] + self.calls = 0 + + def get(self, *_args, **_kwargs): + self.calls += 1 + return self.responses.pop(0) + + reader = object.__new__(validator._PinnedSafetensors) + reader._session = _Session() + sleeps = [] + monkeypatch.setattr(validator.time, "sleep", sleeps.append) + + assert reader._range("model.safetensors", 0, 3) == b"data" + assert reader._session.calls == 3 + assert sleeps == [1, 2] + + +def test_pinned_range_fetch_fails_explicitly_after_retries(monkeypatch): + validator = _load_validator() + + class _Response: + def __init__(self): + self.status_code = 503 + self.headers = {} + self.content = b"" + + def __enter__(self): + return self + + def __exit__(self, *_args): + return False + + class _Session: + calls = 0 + + def get(self, *_args, **_kwargs): + self.calls += 1 + return _Response() + + reader = object.__new__(validator._PinnedSafetensors) + reader._session = _Session() + monkeypatch.setattr(validator.time, "sleep", lambda _delay: None) + + with pytest.raises(RuntimeError, match=r"failed after 3 attempts.*status=503"): + reader._range("model.safetensors", 0, 3) + assert reader._session.calls == 3 + + +def test_reduced_cache_rejects_stale_fixture_schema(tmp_path): + validator = _load_validator() + from safetensors.torch import save_file + + cache_path = tmp_path / "stale.safetensors" + save_file( + {"placeholder": torch.zeros(1)}, + cache_path, + metadata={ + "model_id": validator.MODEL_ID, + "revision": validator.REVISION, + "fixture_schema": "0", + }, + ) + + with pytest.raises(ValueError, match=r"fixture_schema.*Remove the stale cache"): + validator._build_reduced_state(cache_path) + + @pytest.fixture(scope="module") -def reduced_real_state(tmp_path_factory): +def reduced_real_state(): validator = _load_validator() configured_cache = os.environ.get("MOBIUS_NEMOTRON_REDUCED_CACHE") cache = ( - Path(configured_cache) - if configured_cache - else tmp_path_factory.mktemp("nemotron-real") / "reduced.safetensors" + Path(configured_cache) if configured_cache else validator.default_reduced_cache_path() ) state = validator._build_reduced_state(cache) return validator, state From 592129ca7d44c496485a2dde9b76f47c2917ae05 Mon Sep 17 00:00:00 2001 From: Justin Chu Date: Thu, 13 Aug 2026 22:24:52 -0700 Subject: [PATCH 09/13] Fix GPU cache environment scope Set the runner-temp fixture path on the L4/L5 execution steps, where the runner context is available, while retaining the shared cache action path and key. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Signed-off-by: Justin Chu --- .github/workflows/gpu_l4_golden_parity.yml | 3 +-- .github/workflows/gpu_l5_generation_e2e.yml | 3 +-- 2 files changed, 2 insertions(+), 4 deletions(-) diff --git a/.github/workflows/gpu_l4_golden_parity.yml b/.github/workflows/gpu_l4_golden_parity.yml index 164214838..d0af55c4d 100644 --- a/.github/workflows/gpu_l4_golden_parity.yml +++ b/.github/workflows/gpu_l4_golden_parity.yml @@ -42,8 +42,6 @@ jobs: - "1ES.Pool=onnxruntime-ep-mobius-github-linux-a10" - "JobId=gpu-golden-${{ github.run_id }}-${{ github.run_number }}-${{ github.run_attempt }}" timeout-minutes: 60 - env: - MOBIUS_NEMOTRON_REDUCED_CACHE: ${{ runner.temp }}/mobius-nemotron-cache/reduced-d468880b6ad3c6e0d21377ce7242adaea4cc884d-schema-v1.safetensors steps: - uses: actions/checkout@v7 @@ -90,6 +88,7 @@ jobs: env: HF_TOKEN: ${{ secrets.HF_TOKEN }} MOBIUS_TEST_DEVICE: cuda + MOBIUS_NEMOTRON_REDUCED_CACHE: ${{ runner.temp }}/mobius-nemotron-cache/reduced-d468880b6ad3c6e0d21377ce7242adaea4cc884d-schema-v1.safetensors run: | AFFECTED='${{ inputs.affected_models }}' if [ -z "$AFFECTED" ] || [ "$AFFECTED" = "[]" ]; then diff --git a/.github/workflows/gpu_l5_generation_e2e.yml b/.github/workflows/gpu_l5_generation_e2e.yml index 47be44cad..255e756ad 100644 --- a/.github/workflows/gpu_l5_generation_e2e.yml +++ b/.github/workflows/gpu_l5_generation_e2e.yml @@ -42,8 +42,6 @@ jobs: - "1ES.Pool=onnxruntime-ep-mobius-github-linux-a10" - "JobId=gpu-l5-${{ github.run_id }}-${{ github.run_number }}-${{ github.run_attempt }}" timeout-minutes: 60 - env: - MOBIUS_NEMOTRON_REDUCED_CACHE: ${{ runner.temp }}/mobius-nemotron-cache/reduced-d468880b6ad3c6e0d21377ce7242adaea4cc884d-schema-v1.safetensors steps: - uses: actions/checkout@v7 @@ -79,6 +77,7 @@ jobs: env: HF_TOKEN: ${{ secrets.HF_TOKEN }} MOBIUS_TEST_DEVICE: cuda + MOBIUS_NEMOTRON_REDUCED_CACHE: ${{ runner.temp }}/mobius-nemotron-cache/reduced-d468880b6ad3c6e0d21377ce7242adaea4cc884d-schema-v1.safetensors run: | AFFECTED='${{ inputs.affected_models }}' if [ -z "$AFFECTED" ] || [ "$AFFECTED" = "[]" ]; then From f22819e734307512b1377f9756b9e739735c8ece Mon Sep 17 00:00:00 2001 From: Justin Chu Date: Fri, 14 Aug 2026 19:20:40 -0700 Subject: [PATCH 10/13] Emit faithful GenAI metadata for Nemotron Remove downstream ORT GenAI capability gates and derive decoder cache/output metadata directly from exported graphs. Use portable standard-ONNX cache operators for validated Nemotron CUDA and Olive generation, including quantized logits handling. Preserve GGUF special-token metadata and update model-development guidance to keep runtime acceptance downstream. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Signed-off-by: Justin Chu --- .agents/skills/ort-genai-config/SKILL.md | 9 - .../nemotron-3_5-lightning-30b/README.md | 39 +-- .../nemotron-3_5-lightning-30b/inference.py | 17 +- .../nemotron-3_5-lightning-30b/optimize.py | 15 +- .../validate_reduced_checkpoint.py | 5 +- src/mobius/__main__.py | 47 ++- .../integrations/gguf/_config_mapping.py | 26 ++ src/mobius/integrations/gguf/_reader_test.py | 16 ++ src/mobius/integrations/nemo/_genai_config.py | 9 +- .../integrations/ort_genai/auto_export.py | 187 ++++++------ .../ort_genai/auto_export_test.py | 267 ++++++++---------- .../integrations/ort_genai/genai_config.py | 13 +- src/mobius/models/parakeet_ctc_test.py | 11 +- tests/cli_test.py | 28 -- tests/gguf_test.py | 37 ++- tests/nemotron_h_real_weight_test.py | 17 +- 16 files changed, 362 insertions(+), 381 deletions(-) diff --git a/.agents/skills/ort-genai-config/SKILL.md b/.agents/skills/ort-genai-config/SKILL.md index 7319f4617..6ab46aa10 100644 --- a/.agents/skills/ort-genai-config/SKILL.md +++ b/.agents/skills/ort-genai-config/SKILL.md @@ -291,15 +291,6 @@ global cache-slot count. Preserve intrinsic schema/config validation, but do not gate or reject export based on the current GenAI model registry, runtime version, topology support, or cache executor capability. -NemotronH is unsupported in ORT GenAI 0.15.2 because sparse global layer -indices mix attention key/value caches with Mamba `conv_state` and -`ssm_state`. Do not hard-code a model-type or runtime-version rejection. -Inspect the actual decoder graph and reject only when it requires cache inputs -the generated config has no template for (currently `ssm_state`). This keeps -the guard structural: remove it when config emission can represent that graph -contract. Record the tested-version waiver and provide a direct ONNX Runtime -loop for current users. - ### "input_ids size exceeds max length" For image prompts, the tokenized input_ids (including image_pad tokens) can diff --git a/examples/olive/nemotron-3_5-lightning-30b/README.md b/examples/olive/nemotron-3_5-lightning-30b/README.md index 8a9771442..46ed93113 100644 --- a/examples/olive/nemotron-3_5-lightning-30b/README.md +++ b/examples/olive/nemotron-3_5-lightning-30b/README.md @@ -19,13 +19,13 @@ The checkpoint is a real `nemotron_h` model, not an alias: graph does not instantiate MTP, and upstream marks those keys unexpected. No base-decoder generation input, cache, logit, or weight depends on them. -ONNX Runtime GenAI 0.15.2 cannot bind this model's mixed cache: Mamba layers -need `conv_state` plus `ssm_state`, while sparse full-attention layers need -key/value caches at global layer indices. Mobius rejects this graph contract -structurally because `genai_config.json` has no `ssm_state` input/output -template; it does not hard-code a NemotronH or runtime-version check. The guard -can be removed when config emission represents that state. This validated -package uses direct ONNX Runtime generation through `inference.py`. +This model's graph mixes `conv_state` plus `ssm_state` with sparse +full-attention key/value caches. Mobius emits every field the current +`genai_config.json` schema can represent (semantic inputs, key/value and +convolution templates, and global cache-slot count) and leaves runtime +acceptance to ORT GenAI. The schema currently has no `ssm_state` template. The +validated recipe uses direct ONNX Runtime generation through `inference.py`; +downstream load/generation outcomes are informational and do not gate export. ## Install @@ -56,7 +56,9 @@ The script performs four gated steps: 1. Downloads the exact 14-shard BF16 checkpoint revision and exports FP16 ONNX. BF16 execution is rejected explicitly because corrected reduced-real parity reaches `0.8594` max logit error, above the `1e-2` reduced-precision gate. -2. Applies CUDA GQA/LinearAttention fusion and the grouped-RMSNorm workaround. +2. Emits standard ONNX cache operations and applies the grouped-RMSNorm CUDA + workaround. CUDA still executes supported compute nodes; portable cache + operations avoid provider-dependent fused decode drift. 3. Runs Olive Q4 K-quant with an explicitly CPU-only target. It also suppresses Olive 0.13's unrelated GPU-EP DLL auto-registration, so a missing TensorRT installation cannot abort CPU weight-only quantization. @@ -98,9 +100,9 @@ output/ └── source_manifest.json ``` -This recipe intentionally omits `genai_config.json`: direct ONNX Runtime is -the validated runtime, and core Mobius rejects the currently unrepresentable -SSM cache contract before creating runtime artifacts. +This recipe intentionally uses direct ONNX Runtime, while +`mobius build --runtime ort-genai` remains allowed and emits the best current +schema metadata for downstream testing. ## Direct generation and profiling @@ -151,13 +153,14 @@ Validated results on ORT 1.28.0 / Olive 0.13.0: | Variant | Full-logit max abs | Generated IDs | Placement | |---|---:|---|---| | FP32 CPU | `9.54e-6` | `12, 13, 12, 12` | CPU | -| FP16 CUDA | `0.00977` | `12, 13, 12, 12` | 833 CUDA / 14 CPU events | +| FP16 CUDA | `<= 0.0078125` (prefill + every cached step) | `12, 13, 12, 12` | portable ONNX graph on CUDA | | BF16 CUDA | rejected (`0.8594`) | N/A | fails numerical gate | -| Olive Q4 | quantized | `12, 13, 12, 12` | 833 CUDA / 14 CPU events | +| Olive Q4 | quantized | `12, 13, 12, 12` | portable ONNX graph on CUDA | -The reduced FP16 package is 247,380,256 bytes; Q4 is 73,190,965 bytes -(`0.296x`). Its weighted graph contains 15 `com.microsoft::MatMulNBits` -nodes and reloads successfully for multi-token generation. +The reduced package's portable weighted graph quantizes 17 matrix +multiplications to `com.microsoft::MatMulNBits` and reloads successfully for +multi-token generation. Record size/compression from the produced package; +it varies with external-data serialization and Olive version. ## Evidence-based waivers @@ -166,5 +169,5 @@ nodes and reloads successfully for multi-token generation. - Full 30B Olive run: the recipe and reduced production-dimension pass are validated; completing all 2,944 expert subgraphs requires a large-memory host. -- Foundry Local: not available on this host, and its ORT GenAI-based model - contract cannot represent NemotronH hybrid state today. +- Foundry Local was not available on this host. That downstream validation + remains informational and does not block Mobius export. diff --git a/examples/olive/nemotron-3_5-lightning-30b/inference.py b/examples/olive/nemotron-3_5-lightning-30b/inference.py index f4b83337e..272580f43 100644 --- a/examples/olive/nemotron-3_5-lightning-30b/inference.py +++ b/examples/olive/nemotron-3_5-lightning-30b/inference.py @@ -126,6 +126,14 @@ def _as_numpy(value: Any) -> np.ndarray: return value.numpy() +def _logits_output_name(output_names: list[str]) -> str: + """Resolve the full-precision or Olive-renamed logits output.""" + for candidate in ("logits", "logits_Q4"): + if candidate in output_names: + return candidate + raise ValueError(f"Model has no logits output; found {output_names}") + + def _create_session(model_path: Path, device: str, profile: bool): if device == "cuda": # Importing PyTorch first preloads its matching CUDA/cuDNN DLLs on Windows. @@ -189,6 +197,7 @@ def run_token_ids( session = _create_session(model_path, device, profile) states = _initial_states(session) output_names = [output.name for output in session.get_outputs()] + logits_output_name = _logits_output_name(output_names) generated: list[int] = [] logits_by_step: list[np.ndarray] = [] past_length = 0 @@ -207,7 +216,9 @@ def run_token_ids( past_length += 1 assert outputs is not None - logits = _as_numpy(outputs[output_names.index("logits")])[0, -1].astype(np.float32) + logits = _as_numpy(outputs[output_names.index(logits_output_name)])[0, -1].astype( + np.float32 + ) eos_ids = set(eos_token_ids or ()) for token_index in range(max_new_tokens): logits_by_step.append(logits.copy()) @@ -225,7 +236,9 @@ def run_token_ids( outputs = _run_session(session, output_names, feeds) _update_states(states, output_names, outputs) past_length += 1 - logits = _as_numpy(outputs[output_names.index("logits")])[0, -1].astype(np.float32) + logits = _as_numpy(outputs[output_names.index(logits_output_name)])[0, -1].astype( + np.float32 + ) profile_path = session.end_profiling() if profile else None return generated, logits_by_step, profile_path diff --git a/examples/olive/nemotron-3_5-lightning-30b/optimize.py b/examples/olive/nemotron-3_5-lightning-30b/optimize.py index 38001e09c..4ac40d470 100644 --- a/examples/olive/nemotron-3_5-lightning-30b/optimize.py +++ b/examples/olive/nemotron-3_5-lightning-30b/optimize.py @@ -69,6 +69,10 @@ def export_checkpoint(output_dir: str | Path, *, ep: str) -> Path: output = Path(output_dir) _require_empty_output(output) + # CUDA fused cache kernels diverge during multi-step decode for this hybrid + # architecture. Keep a portable graph and let CUDA place supported standard + # ops; reduced-real validation enforces <=1e-2 at every cached step. + build_ep = "onnx-standard" if ep == "cuda" else ep with override_flags(ort_cuda_grouped_rmsnorm_workaround=ep == "cuda"): package = build( MODEL_ID, @@ -76,13 +80,20 @@ def export_checkpoint(output_dir: str | Path, *, ep: str) -> Path: dtype="f16", load_weights=True, trust_remote_code=False, - execution_provider=ep, + execution_provider=build_ep, ) package.save(output, external_data="onnx") _save_pinned_metadata(output) manifest_path = output / "source_manifest.json" manifest = json.loads(manifest_path.read_text(encoding="utf-8")) - manifest.update({"source_dtype": "bf16", "dtype": "f16", "target_ep": ep}) + manifest.update( + { + "source_dtype": "bf16", + "dtype": "f16", + "target_ep": ep, + "build_ep": build_ep, + } + ) manifest_path.write_text(json.dumps(manifest, indent=2), encoding="utf-8") return output diff --git a/examples/olive/nemotron-3_5-lightning-30b/validate_reduced_checkpoint.py b/examples/olive/nemotron-3_5-lightning-30b/validate_reduced_checkpoint.py index e51049b32..839ae4a68 100644 --- a/examples/olive/nemotron-3_5-lightning-30b/validate_reduced_checkpoint.py +++ b/examples/olive/nemotron-3_5-lightning-30b/validate_reduced_checkpoint.py @@ -301,12 +301,15 @@ def _mobius_package( config = NemotronHConfig.from_transformers(_hf_config()) config.dtype = getattr(ir.DataType, _DTYPES[dtype_name][1]) module = NemotronHCausalLMModel(config) + # Keep CUDA execution portable: fused hybrid cache kernels show + # provider-dependent multi-step drift even when prefill matches. + build_ep = "onnx-standard" if ep == "cuda" else ep with override_flags(ort_cuda_grouped_rmsnorm_workaround=ep == "cuda"): package = build_from_module( module, config, task="hybrid-text-generation", - execution_provider=ep, + execution_provider=build_ep, trace_optimization=True, ) package.apply_weights(module.preprocess_weights(dict(state))) diff --git a/src/mobius/__main__.py b/src/mobius/__main__.py index bed089bcb..697223bef 100644 --- a/src/mobius/__main__.py +++ b/src/mobius/__main__.py @@ -367,16 +367,6 @@ def _save_package( ) -> None: """Save a ModelPackage to disk, applying optimizations and runtime configs.""" runtime = getattr(args, "runtime", None) - if runtime == "ort-genai": - from mobius.integrations.ort_genai.auto_export import ( - _validate_ort_genai_compatibility, - ) - - try: - _validate_ort_genai_compatibility(pkg) - except ValueError as error: - raise SystemExit(f"Error: {error}") from error - components = (lambda name: name == component_filter) if component_filter else None for name, model in pkg.items(): if components is not None and not components(name): @@ -507,15 +497,6 @@ def _cmd_build_gguf(args: argparse.Namespace) -> None: ) raise SystemExit(1) - if getattr(args, "runtime", None) == "ort-genai": - raise SystemExit( - "Error: mobius build-gguf does not yet support --runtime ort-genai. " - "The command cannot emit a valid genai_config.json until the selected " - "GGUF architecture's cache and tokenizer contracts have passed real " - "ORT GenAI generation. Use --runtime onnx-genai where supported, or " - "omit --runtime and run the ONNX model directly." - ) - mmproj_path = getattr(args, "mmproj", None) keep_quantized = not args.dequantize @@ -560,9 +541,9 @@ def _cmd_build_gguf(args: argparse.Namespace) -> None: path = os.path.join(output_dir, "model.onnx") print(f"Saved {name} to {path}") - if getattr(args, "runtime", None) == "onnx-genai": + runtime = getattr(args, "runtime", None) + if runtime in ("ort-genai", "onnx-genai"): from mobius.integrations.gguf import write_gguf_tokenizer_json - from mobius.integrations.onnx_genai import write_onnx_genai_config # A GGUF checkpoint has no Hugging Face source directory, so the # tokenizer is reconstructed from the file's embedded ggml metadata @@ -570,9 +551,20 @@ def _cmd_build_gguf(args: argparse.Namespace) -> None: tokenizer_path = write_gguf_tokenizer_json(gguf_path, output_dir) if tokenizer_path is not None: print(f" tokenizer: {tokenizer_path}") - artifacts = write_onnx_genai_config( - pkg, output_dir, config=getattr(pkg, "config", None), source=None - ) + if runtime == "onnx-genai": + from mobius.integrations.onnx_genai import write_onnx_genai_config + + artifacts = write_onnx_genai_config( + pkg, output_dir, config=getattr(pkg, "config", None), source=None + ) + else: + from mobius.integrations.ort_genai import write_ort_genai_config + + artifacts = write_ort_genai_config( + pkg, + output_dir, + ep=args.execution_provider, + ) for name, path in artifacts.items(): print(f" {name}: {path}") @@ -896,10 +888,9 @@ def main(argv: list[str] | None = None) -> None: default=None, help=( "Generate runtime-specific config files after building. " - "'onnx-genai' writes inference_metadata.yaml plus a tokenizer.json " - "reconstructed from the GGUF's embedded tokenizer metadata; " - "'ort-genai' is currently rejected until GGUF cache/tokenizer " - "contracts have runtime generation coverage." + "Both modes reconstruct tokenizer.json from GGUF metadata. " + "'onnx-genai' writes inference_metadata.yaml; 'ort-genai' writes " + "the best graph-derived genai_config.json metadata." ), ) gguf_parser.add_argument( diff --git a/src/mobius/integrations/gguf/_config_mapping.py b/src/mobius/integrations/gguf/_config_mapping.py index f1752b505..a52f498c4 100644 --- a/src/mobius/integrations/gguf/_config_mapping.py +++ b/src/mobius/integrations/gguf/_config_mapping.py @@ -22,6 +22,7 @@ __all__ = ["gguf_to_config"] +import contextlib import dataclasses import logging from typing import Any @@ -105,6 +106,14 @@ "attention.sliding_window": "sliding_window", } +_TOKENIZER_KEY_MAP: dict[str, str] = { + "tokenizer.ggml.bos_token_id": "bos_token_id", + "tokenizer.ggml.eos_token_id": "eos_token_id", + "tokenizer.ggml.padding_token_id": "pad_token_id", + # Non-standard but emitted by some multimodal converters. + "tokenizer.ggml.image_token_id": "image_token_id", +} + _ARCH_KEY_MAPS: dict[str, dict[str, str]] = { # Both spellings are accepted wherever muse-glimmer is recognized. "muse-glimmer": _MUSE_GLIMMER_KEY_MAP, @@ -212,6 +221,17 @@ def _extract_config_fields( if isinstance(tokens, list): hf_fields["vocab_size"] = len(tokens) + for gguf_key, config_key in _TOKENIZER_KEY_MAP.items(): + if gguf_key in metadata: + hf_fields[config_key] = int(metadata[gguf_key]) + + # Standard GGUF has no dedicated image-token key. Preserve the canonical + # HuggingFace placeholder when it is embedded in the tokenizer vocabulary. + tokens = metadata.get("tokenizer.ggml.tokens") + if "image_token_id" not in hf_fields and isinstance(tokens, list): + with contextlib.suppress(ValueError): + hf_fields["image_token_id"] = tokens.index("") + return hf_fields @@ -387,6 +407,11 @@ def gguf_to_config( if isinstance(swiglu_limit, (list, np.ndarray)): swiglu_limit = swiglu_limit[0] if len(swiglu_limit) else 0.0 + special_token_fields = { + name: int(hf_fields[name]) + for name in ("bos_token_id", "eos_token_id", "pad_token_id", "image_token_id") + if hf_fields.get(name) is not None + } config = ArchitectureConfig( hidden_size=hidden_size, intermediate_size=hf_fields.get("intermediate_size", 4 * hidden_size), @@ -452,6 +477,7 @@ def gguf_to_config( linear_key_head_dim=linear_key_head_dim, linear_value_head_dim=linear_value_head_dim, linear_conv_kernel_dim=(hf_fields.get("linear_conv_kernel_dim") or 4), + **special_token_fields, ) # Store model_type for registry lookup and tensor processor dispatch. diff --git a/src/mobius/integrations/gguf/_reader_test.py b/src/mobius/integrations/gguf/_reader_test.py index 0ec0e1089..99426b6ee 100644 --- a/src/mobius/integrations/gguf/_reader_test.py +++ b/src/mobius/integrations/gguf/_reader_test.py @@ -559,6 +559,22 @@ def test_extract_config_fields_with_prefix(self): assert fields["hidden_size"] == 4096 assert fields["num_hidden_layers"] == 32 + def test_extract_config_fields_preserves_tokenizer_metadata(self): + fields = _extract_config_fields( + "llama", + { + "tokenizer.ggml.bos_token_id": 1, + "tokenizer.ggml.eos_token_id": 2, + "tokenizer.ggml.padding_token_id": 3, + "tokenizer.ggml.tokens": ["a", "", "b"], + }, + ) + + assert fields["bos_token_id"] == 1 + assert fields["eos_token_id"] == 2 + assert fields["pad_token_id"] == 3 + assert fields["image_token_id"] == 1 + def test_infer_tie_embeddings_true(self, tied_gguf: Path): model = GGUFModel(tied_gguf) assert _infer_tie_embeddings(model) is True diff --git a/src/mobius/integrations/nemo/_genai_config.py b/src/mobius/integrations/nemo/_genai_config.py index 2b919cf3f..b284c72a5 100644 --- a/src/mobius/integrations/nemo/_genai_config.py +++ b/src/mobius/integrations/nemo/_genai_config.py @@ -269,8 +269,7 @@ def write_genai_bundle( Args: pkg: The :class:`~mobius._model_package.ModelPackage` produced by :func:`~mobius.integrations.nemo.build_from_nemo`. Must contain the - ``encoder_streaming``, ``decoder`` and ``joint`` graphs and be built - in float32 (the GenAI pipeline only supports float32 encoder I/O). + ``encoder_streaming``, ``decoder`` and ``joint`` graphs. archive: The source :class:`NeMoArchive` (for preprocessor parameters and the SentencePiece tokenizer). dest_dir: Output directory (created if needed). @@ -283,12 +282,6 @@ def write_genai_bundle( Returns: The resolved output directory path. """ - dtype = getattr(pkg.config, "dtype", None) - if dtype is not None and dtype != ir.DataType.FLOAT: - raise ValueError( - "ORT GenAI nemotron_speech only supports float32 encoder I/O; build " - f"the package with dtype='f32' (got {dtype})." - ) for key in ("encoder_streaming", "decoder", "joint"): if key not in pkg: raise KeyError(f"ModelPackage is missing required model {key!r}") diff --git a/src/mobius/integrations/ort_genai/auto_export.py b/src/mobius/integrations/ort_genai/auto_export.py index ba757f2d5..987f0dd3f 100644 --- a/src/mobius/integrations/ort_genai/auto_export.py +++ b/src/mobius/integrations/ort_genai/auto_export.py @@ -243,6 +243,71 @@ def _count_cache_layer_slots(model: ir.Model | None) -> int | None: return max(layer_indices) + 1 if layer_indices else None +_CACHE_INPUT_FIELDS = { + (None, "key"): "past_key_names", + (None, "value"): "past_value_names", + ("self", "key"): "past_key_names", + ("self", "value"): "past_value_names", + ("cross", "key"): "cross_past_key_names", + ("cross", "value"): "cross_past_value_names", + (None, "conv_state"): "past_conv_names", +} +_CACHE_OUTPUT_FIELDS = { + (None, "key"): "present_key_names", + (None, "value"): "present_value_names", + ("self", "key"): "present_key_names", + ("self", "value"): "present_value_names", + (None, "conv_state"): "present_conv_names", +} + + +def _cache_name_parts(name: str, prefix: str) -> tuple[str | None, str, str] | None: + """Return ``(scope, suffix, template)`` for indexed cache names.""" + parts = name.split(".") + if len(parts) not in (3, 4) or parts[0] != prefix or not parts[1].isdigit(): + return None + scope = parts[2] if len(parts) == 4 else None + suffix = parts[-1] + parts[1] = "%d" + return scope, suffix, ".".join(parts) + + +def _decoder_cache_templates(model: ir.Model) -> tuple[dict[str, str], dict[str, str]]: + """Map graph cache suffixes to every template the current config schema supports.""" + inputs: dict[str, str] = {} + outputs: dict[str, str] = {} + for value in model.graph.inputs: + if ( + value.name is None + or (parts := _cache_name_parts(value.name, "past_key_values")) is None + ): + continue + scope, suffix, template = parts + if (config_name := _CACHE_INPUT_FIELDS.get((scope, suffix))) is not None: + inputs[config_name] = template + for value in model.graph.outputs: + if value.name is None or (parts := _cache_name_parts(value.name, "present")) is None: + continue + scope, suffix, template = parts + if (config_name := _CACHE_OUTPUT_FIELDS.get((scope, suffix))) is not None: + outputs[config_name] = template + return inputs, outputs + + +def _decoder_output_mapping(model: ir.Model) -> dict[str, str] | None: + """Return semantic decoder outputs and graph-derived cache templates.""" + output_names = [value.name for value in model.graph.outputs if value.name is not None] + logits_name = next( + (name for name in output_names if name == "logits"), + next((name for name in output_names if name.startswith("logits_")), None), + ) + _cache_inputs, cache_outputs = _decoder_cache_templates(model) + outputs = dict(cache_outputs) + if logits_name is not None: + outputs["logits"] = logits_name + return outputs or None + + def _introspect_inputs(pkg: ModelPackage, key: str) -> dict[str, str] | None: """Return ``{name: name}`` identity mapping for a sub-model's inputs. @@ -1007,45 +1072,23 @@ def _write_genai_config( # --- Discover decoder inputs from the ONNX graph --- decoder_key = "decoder" if "decoder" in pkg else "model" decoder_inputs = _introspect_inputs(pkg, decoder_key) - if decoder_inputs is not None: - # KV cache entries are template-based, not per-input - decoder_inputs["past_key_names"] = "past_key_values.%d.key" - decoder_inputs["past_value_names"] = "past_key_values.%d.value" + decoder_model = pkg.get(decoder_key) + decoder_outputs = None + if decoder_model is not None: + cache_inputs, _cache_outputs = _decoder_cache_templates(decoder_model) + if decoder_inputs is not None: + decoder_inputs.update(cache_inputs) + decoder_outputs = _decoder_output_mapping(decoder_model) # Derive decoder filename from the actual package key decoder_filename = ( f"{decoder_key}/model.onnx" if len(pkg) > 1 or decoder_key != "model" else "model.onnx" ) - # ORT GenAI's ``past_present_share_buffer`` mode requires the decoder - # graph to write the KV cache in place. Only ``com.microsoft. - # GroupQueryAttention`` does that; the standard ONNX ``Attention`` op - # concatenates ``past_key`` with the new ``K`` and returns a dynamic- - # shape ``present_key``, which is incompatible with the pre-allocated - # shared buffer. Introspect the graph: if there is at least one GQA - # node, the model supports shared-buffer mode; otherwise force it off - # regardless of the EP capability flag. - # - # ``com.microsoft.LinearAttention`` (linear/recurrent-attention layers, - # e.g. Qwen3.5's GatedDeltaNet) is a separate, *mandatory* case: its - # recurrent state requires ``past_present_share_buffer=True`` regardless - # of whether any other layer uses GQA (ORT GenAI raises "RecurrentState - # requires past_present_share_buffer=true" otherwise). - # - # Hybrid models mix LinearAttention layers with full-attention layers, - # which may lower to GQA *or* to the standard (non-GQA) ``Attention`` op - # depending on EP/dtype (e.g. the CPU EP only lowers to GQA for fp32; - # fp16 falls back to standard Attention -- see ``_execution_providers.py`` - # ``gqa_dtypes``). If a hybrid graph has LinearAttention but its - # full-attention layers are still standard (non-GQA) Attention, forcing - # ``past_present_share_buffer=True`` produces an unrunnable config: the - # recurrent state requires it, but standard Attention's dynamic-shape KV - # concat cannot honor a pre-allocated shared buffer, which fails at - # generation time with an ``attn_mask``/``total_sequence_length`` - # mismatch rather than at load time. Rather than silently emit a broken - # config, raise a clear error so the caller picks an EP/dtype combination - # (e.g. fp32 on CPU) that lowers full attention to GQA. - decoder_model = pkg.get(decoder_key) + # Derive shared-buffer metadata from the graph. GQA supports in-place KV, + # while LinearAttention requires a shared recurrent-state buffer. Mixed + # topologies are still emitted faithfully; downstream runtime acceptance + # must not become a Mobius export capability gate. supports_in_place_kv_cache: bool | None = None if decoder_model is not None: has_gqa = any( @@ -1056,26 +1099,6 @@ def _write_genai_config( node.op_type == "LinearAttention" and node.domain == "com.microsoft" for node in decoder_model.graph ) - has_standard_attention = any( - node.op_type == "Attention" and node.domain in ("", "ai.onnx") - for node in decoder_model.graph - ) - if has_recurrent_state and has_standard_attention: - # A GQA node elsewhere in the graph does NOT make a co-existing - # standard Attention node compatible with a shared buffer -- - # each op instance is independently (in)compatible, so this - # must reject on the mere presence of standard Attention, not - # only when GQA is completely absent (partial GQA fusion still - # leaves the unfused standard Attention layers broken). - raise ValueError( - "This decoder graph mixes com.microsoft.LinearAttention " - "(recurrent state, requires past_present_share_buffer=True) " - "with standard (non-GQA) Attention (incompatible with " - "past_present_share_buffer=True). This EP/dtype combination " - "cannot produce a runnable genai_config -- pick an EP/dtype " - "that lowers *all* full-attention layers to " - "GroupQueryAttention instead (e.g. fp32 on the CPU EP)." - ) supports_in_place_kv_cache = has_gqa or has_recurrent_state generator = GenaiConfigGenerator.from_config( @@ -1087,6 +1110,7 @@ def _write_genai_config( eos_token_id=eos_token_id, pad_token_id=pad_token_id, decoder_inputs=decoder_inputs, + decoder_outputs=decoder_outputs, decoder_filename=decoder_filename, supports_in_place_kv_cache=supports_in_place_kv_cache, num_cache_layer_slots=_count_cache_layer_slots(decoder_model), @@ -1237,48 +1261,6 @@ def _write_genai_config( return generator.write(output_dir) -def _validate_ort_genai_compatibility(pkg: ModelPackage) -> None: - """Reject packages whose required inputs cannot be supplied by ORT GenAI.""" - config = getattr(pkg, "config", None) - decoder_key = "decoder" if "decoder" in pkg else "model" - decoder_model = pkg.get(decoder_key) - if decoder_model is not None: - cache_suffixes = { - parts[2] - for model_input in decoder_model.graph.inputs - if model_input.name is not None - and len(parts := model_input.name.split(".")) == 3 - and parts[0] == "past_key_values" - and parts[1].isdigit() - } - if "ssm_state" in cache_suffixes: - raise ValueError( - "ORT GenAI config generation cannot represent this decoder's cache " - f"inputs {sorted(cache_suffixes)}: no input/output template exists for " - "ssm_state. Export without --runtime ort-genai and run the ONNX model " - "directly until the config schema and runtime support this graph contract." - ) - if getattr(config, "model_type", None) == "parakeet_ctc": - raise ValueError( - "ORT GenAI does not define a feature-input CTC ASR pipeline; " - "export Parakeet CTC as ONNX and run it directly with ONNX Runtime." - ) - if {"vision_encoder", "decoder"}.issubset(pkg) and "embedding" not in pkg: - model_type = getattr(config, "model_type", "unknown") - raise NotImplementedError( - "onnxruntime-genai does not support generic vision encoder-decoder " - f"packages such as {model_type!r}. Run the vision_encoder and decoder " - "ONNX sessions directly; emitting genai_config.json would create an " - "artifact that the runtime cannot load." - ) - if getattr(config, "model_type", None) == "mage_vl": - raise ValueError( - "ORT GenAI does not support Mage-VL's required patch_positions vision " - "input or its 1D decoder position_ids contract. Export without " - "--runtime ort-genai to save the runnable direct three-model ONNX package." - ) - - def write_ort_genai_config( pkg: ModelPackage, directory: str, @@ -1343,15 +1325,6 @@ def write_ort_genai_config( "This is set automatically when building with mobius.build(). " "Diffusion models (which have no config) are not supported." ) - _validate_ort_genai_compatibility(pkg) - - if getattr(config, "model_type", None) == "moonshine": - raise NotImplementedError( - "onnxruntime-genai does not support Moonshine's variable-length raw-waveform " - "encoder. Run the exported encoder and cached decoder directly with " - "ONNX Runtime." - ) - os.makedirs(directory, exist_ok=True) # Normalize EP: 'default' and 'onnx-standard' are portable-ONNX modes @@ -1411,6 +1384,10 @@ def write_ort_genai_config( # Gemma3 multimodal configs are unwrapped to the text sub-config # during build, but ORT GenAI needs the multimodal parent type. ort_model_type = "gemma3" + elif is_vlm and raw_type == "gemma4_text": + # GGUF multimodal builds retain the text checkpoint's model type + # after attaching the companion vision projector. + ort_model_type = "gemma4" elif is_vlm and raw_type == "gemma3n_text": # Same unwrapping for Gemma3n, whose parent type is "gemma3n". # Deliberately *not* aliased to "gemma3": the package threads @@ -1637,8 +1614,6 @@ def export_package( "Diffusion models (which have no config) are not supported — " "use ModelPackage.save() directly for those." ) - _validate_ort_genai_compatibility(pkg) - os.makedirs(output_dir, exist_ok=True) # 1. Save ONNX models + weights diff --git a/src/mobius/integrations/ort_genai/auto_export_test.py b/src/mobius/integrations/ort_genai/auto_export_test.py index 467debd1e..f65ea7a44 100644 --- a/src/mobius/integrations/ort_genai/auto_export_test.py +++ b/src/mobius/integrations/ort_genai/auto_export_test.py @@ -56,25 +56,6 @@ def _mock_model_with_outputs(names: list[str]) -> ir.Model: return _mock_model(outputs=names) -def test_moonshine_native_runtime_is_rejected(tmp_path): - from mobius._model_package import ModelPackage - - config = mock.MagicMock() - config.model_type = "moonshine" - package = ModelPackage( - {"encoder": _mock_model(), "decoder": _mock_model()}, - config=config, - ) - - with pytest.raises( - NotImplementedError, - match="variable-length raw-waveform encoder", - ) as error: - write_ort_genai_config(package, str(tmp_path)) - assert "onnx-genai" not in str(error.value) - assert "ONNX Runtime" in str(error.value) - - def _make_fake_llm_pkg(model_type: str = "qwen2"): """Build a minimal LLM-only ModelPackage with a fake config.""" import dataclasses @@ -1080,37 +1061,22 @@ def test_genai_config_json_is_written(self, tmp_path): assert "model" in data assert data["model"]["type"] == "qwen2" - def test_rejects_generic_vision_encoder_decoder_package(self, tmp_path): + def test_nemotron_h_mixed_cache_metadata_is_emitted(self, tmp_path): import dataclasses from mobius._model_package import ModelPackage @dataclasses.dataclass class FakeConfig: - model_type: str = "nemotron_parse" - - pkg = ModelPackage( - { - "vision_encoder": _mock_model(), - "decoder": _mock_model(), - }, - config=FakeConfig(), - ) - with pytest.raises( - NotImplementedError, - match="does not support generic vision encoder-decoder", - ): - write_ort_genai_config(pkg, str(tmp_path)) - assert not (tmp_path / "genai_config.json").exists() - - def test_rejects_unrepresentable_mixed_cache_graph_before_writing(self, tmp_path): - import dataclasses - - from mobius._model_package import ModelPackage - - @dataclasses.dataclass - class FakeConfig: - model_type: str = "future_hybrid" + model_type: str = "nemotron_h" + vocab_size: int = 256 + hidden_size: int = 64 + num_hidden_layers: int = 3 + 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 pkg = ModelPackage( { @@ -1122,17 +1088,89 @@ class FakeConfig: "past_key_values.0.ssm_state", "past_key_values.2.key", "past_key_values.2.value", - ] + ], + outputs=[ + "logits", + "present.0.conv_state", + "present.0.ssm_state", + "present.2.key", + "present.2.value", + ], ) }, config=FakeConfig(), ) output_dir = tmp_path / "ort-genai" - with pytest.raises(ValueError, match=r"no input/output template.*ssm_state"): - write_ort_genai_config(pkg, str(output_dir)) + result = write_ort_genai_config(pkg, str(output_dir)) + + with open(result["genai_config"], encoding="utf-8") as config_file: + generated = json.load(config_file) + decoder = generated["model"]["decoder"] + assert generated["model"]["type"] == "nemotron_h" + assert decoder["num_hidden_layers"] == 3 + assert decoder["inputs"] == { + "input_ids": "input_ids", + "attention_mask": "attention_mask", + "past_key_names": "past_key_values.%d.key", + "past_value_names": "past_key_values.%d.value", + "past_conv_names": "past_key_values.%d.conv_state", + } + assert decoder["outputs"] == { + "logits": "logits", + "present_key_names": "present.%d.key", + "present_value_names": "present.%d.value", + "present_conv_names": "present.%d.conv_state", + } + + def test_olive_renamed_logits_output_is_emitted(self, tmp_path): + pkg = _make_fake_llm_pkg("qwen2") + pkg["model"] = _mock_model( + inputs=["input_ids", "past_key_values.0.key", "past_key_values.0.value"], + outputs=["logits_Q4", "present.0.key", "present.0.value"], + ) + + result = write_ort_genai_config(pkg, str(tmp_path)) + + with open(result["genai_config"], encoding="utf-8") as config_file: + generated = json.load(config_file) + assert generated["model"]["decoder"]["outputs"]["logits"] == "logits_Q4" - assert not output_dir.exists() + def test_nested_self_and_cross_cache_templates_are_emitted(self, tmp_path): + pkg = _make_fake_llm_pkg("decoder") + pkg["model"] = _mock_model( + inputs=[ + "input_ids", + "encoder_hidden_states", + "past_key_values.0.self.key", + "past_key_values.0.self.value", + "past_key_values.0.cross.key", + "past_key_values.0.cross.value", + ], + outputs=[ + "logits", + "present.0.self.key", + "present.0.self.value", + ], + ) + + result = write_ort_genai_config(pkg, str(tmp_path)) + + with open(result["genai_config"], encoding="utf-8") as config_file: + decoder = json.load(config_file)["model"]["decoder"] + assert decoder["inputs"] == { + "input_ids": "input_ids", + "encoder_hidden_states": "encoder_hidden_states", + "past_key_names": "past_key_values.%d.self.key", + "past_value_names": "past_key_values.%d.self.value", + "cross_past_key_names": "past_key_values.%d.cross.key", + "cross_past_value_names": "past_key_values.%d.cross.value", + } + assert decoder["outputs"] == { + "logits": "logits", + "present_key_names": "present.%d.self.key", + "present_value_names": "present.%d.self.value", + } def test_processor_config_written_with_vision(self, tmp_path): """image_processor.json is written when pkg.config.vision is set.""" @@ -1233,48 +1271,6 @@ class FakeConfig: assert model["vision"]["patch_size"] == 16 assert model["vision"]["window_size"] == 64 - def test_mage_vl_is_rejected_before_writing_runtime_artifacts(self, tmp_path): - import dataclasses - - from mobius._model_package import ModelPackage - from mobius.integrations.ort_genai.auto_export import write_ort_genai_config - - @dataclasses.dataclass - class FakeVision: - image_size: int = 448 - patch_size: int = 16 - spatial_merge_size: int = 2 - - @dataclasses.dataclass - class FakeConfig: - model_type: str = "mage_vl" - vocab_size: int = 151936 - hidden_size: int = 2560 - num_hidden_layers: int = 1 - num_attention_heads: int = 32 - num_key_value_heads: int = 8 - head_dim: int = 128 - image_token_id: int = 151655 - temporal_patch_size: int = 1 - vision: FakeVision = dataclasses.field(default_factory=FakeVision) - - pkg = ModelPackage( - { - "decoder": _mock_model(), - "vision_encoder": _mock_model(), - "embedding": _mock_model(), - }, - config=FakeConfig(), - ) - - output_dir = tmp_path / "ort-genai" - with pytest.raises( - ValueError, - match=r"Mage-VL.*patch_positions.*1D decoder position_ids", - ): - write_ort_genai_config(pkg, str(output_dir)) - assert not output_dir.exists() - def test_processor_config_not_written_without_vision(self, tmp_path): """image_processor.json is NOT written when pkg.config has no vision attr.""" from mobius.integrations.ort_genai.auto_export import write_ort_genai_config @@ -1670,8 +1666,17 @@ class FakeConfig: # "gemma2" maps to "gemma" in _ORT_GENAI_MODEL_TYPE assert data["model"]["type"] == "gemma" - def test_config_mode_gemma3_text_vlm_uses_multimodal_model_type(self, tmp_path): - """Gemma3 VLM --config exports use ORT's multimodal gemma3 type.""" + @pytest.mark.parametrize( + ("text_model_type", "multimodal_model_type"), + [("gemma3_text", "gemma3"), ("gemma4_text", "gemma4")], + ) + def test_config_mode_text_vlm_uses_multimodal_model_type( + self, + tmp_path, + text_model_type, + multimodal_model_type, + ): + """Unwrapped text configs retain their multimodal runtime type.""" import dataclasses from mobius._model_package import ModelPackage @@ -1686,8 +1691,7 @@ class FakeVision: @dataclasses.dataclass class FakeConfig: - # build() stores the unwrapped text sub-config type on Gemma3 VLMs. - model_type: str = "gemma3_text" + model_type: str vocab_size: int = 262144 hidden_size: int = 64 num_hidden_layers: int = 2 @@ -1704,13 +1708,13 @@ class FakeConfig: "vision_encoder": _mock_model_with_inputs(["pixel_values"]), "embedding": _mock_model_with_inputs(["input_ids", "image_features"]), }, - config=FakeConfig(), + config=FakeConfig(model_type=text_model_type), ) result = write_ort_genai_config(pkg, str(tmp_path), hf_model_id=None) with open(result["genai_config"]) as f: data = json.load(f) - assert data["model"]["type"] == "gemma3" + assert data["model"]["type"] == multimodal_model_type def test_config_mode_gemma3n_text_vlm_uses_multimodal_model_type(self, tmp_path): """Gemma3n unwraps to "gemma3n_text" too, and must not alias to gemma3. @@ -2030,8 +2034,6 @@ def _make_pkg(): def test_writes_both_onnx_and_genai_config(self, tmp_path, monkeypatch): """export_package calls pkg.save AND writes genai_config.json.""" - from mobius.integrations.ort_genai.auto_export import export_package - pkg = self._make_pkg() save_calls = [] @@ -2051,22 +2053,8 @@ def fake_save(self, directory, **kwargs): # ONNX path is in the manifest (single-component package) assert result["model"] == os.path.join(str(tmp_path), "model.onnx") - def test_mage_vl_is_rejected_before_saving_onnx(self, tmp_path): - pkg = self._make_pkg() - pkg.config.model_type = "mage_vl" - - with ( - mock.patch.object(pkg, "save") as save, - pytest.raises(ValueError, match=r"Mage-VL.*patch_positions"), - ): - export_package(pkg, str(tmp_path)) - - save.assert_not_called() - def test_propagates_save_kwargs(self, tmp_path, monkeypatch): """external_data and progress_bar are forwarded to pkg.save.""" - from mobius.integrations.ort_genai.auto_export import export_package - pkg = self._make_pkg() save_calls = [] @@ -2087,8 +2075,6 @@ def fake_save(self, directory, **kwargs): def test_propagates_genai_config_kwargs(self, tmp_path, monkeypatch): """The ep and context_length kwargs reach the generated genai_config.json.""" - from mobius.integrations.ort_genai.auto_export import export_package - pkg = self._make_pkg() monkeypatch.setattr(pkg.__class__, "save", lambda self, d, **kw: None) @@ -2114,7 +2100,6 @@ def test_preflights_missing_config(self, tmp_path, monkeypatch): no genai_config.json. """ from mobius._model_package import ModelPackage - from mobius.integrations.ort_genai.auto_export import export_package pkg = ModelPackage({"model": _mock_model()}, config=None) save_called = [] @@ -2132,8 +2117,6 @@ def fake_save(self, *a, **kw): def test_returns_manifest_with_all_artifacts(self, tmp_path, monkeypatch): """Returned manifest contains ONNX paths AND config artifacts.""" - from mobius.integrations.ort_genai.auto_export import export_package - pkg = self._make_pkg() monkeypatch.setattr(pkg.__class__, "save", lambda self, d, **kw: None) @@ -2349,16 +2332,8 @@ def test_pixtral_config_filename_is_processor_config(self, tmp_path): assert data["model"]["image_token_id"] == 10 -class TestHybridAttentionShareBufferGuard: - """Tests for the LinearAttention/GQA past_present_share_buffer guard. - - See the comment above ``supports_in_place_kv_cache`` in - ``_write_genai_config``: recurrent-state layers (LinearAttention) - mandate ``past_present_share_buffer=True``, but standard (non-GQA) - Attention is incompatible with it. A hybrid graph with both, and no - GQA node to lower the standard Attention layers to, must raise a clear - build-time error rather than silently emit a broken config. - """ +class TestHybridAttentionShareBufferMetadata: + """Tests graph-derived shared-buffer metadata without runtime gating.""" @staticmethod def _make_pkg(node_op_types: list[tuple[str, str]]): @@ -2406,13 +2381,14 @@ def _write(self, pkg, tmp_path): has_speech=False, ) - def test_recurrent_state_with_standard_attention_and_no_gqa_raises(self, tmp_path): - """LinearAttention + standard Attention + no GQA is an unrunnable config.""" + def test_recurrent_state_with_standard_attention_is_emitted(self, tmp_path): pkg = self._make_pkg( [("LinearAttention", "com.microsoft"), ("Attention", "")], ) - with pytest.raises(ValueError, match="past_present_share_buffer"): - self._write(pkg, tmp_path) + path = self._write(pkg, tmp_path) + with open(path) as f: + data = json.load(f) + assert data["search"]["past_present_share_buffer"] is True def test_recurrent_state_with_gqa_does_not_raise(self, tmp_path): """LinearAttention + GQA (no standard Attention) is a valid hybrid config.""" @@ -2427,16 +2403,7 @@ def test_recurrent_state_with_gqa_does_not_raise(self, tmp_path): data = json.load(f) assert data["search"]["past_present_share_buffer"] is True - def test_recurrent_state_with_standard_attention_and_gqa_raises(self, tmp_path): - """Partial GQA fusion still leaves an incompatible standard Attention node. - - Regression test: the guard previously read - ``has_recurrent_state and has_standard_attention and not has_gqa``, so - a GQA node present *anywhere* in the graph would short-circuit the - check even though a separate, unfused standard Attention node - coexists. A GQA node on one layer doesn't make a standard Attention - node on another layer safe for ``past_present_share_buffer=True``. - """ + def test_recurrent_state_with_standard_attention_and_gqa_is_emitted(self, tmp_path): pkg = self._make_pkg( [ ("LinearAttention", "com.microsoft"), @@ -2444,8 +2411,10 @@ def test_recurrent_state_with_standard_attention_and_gqa_raises(self, tmp_path): ("Attention", ""), ], ) - with pytest.raises(ValueError, match="past_present_share_buffer"): - self._write(pkg, tmp_path) + path = self._write(pkg, tmp_path) + with open(path) as f: + data = json.load(f) + assert data["search"]["past_present_share_buffer"] is True def test_recurrent_state_only_does_not_raise(self, tmp_path): """LinearAttention with no full-attention layers at all is unaffected.""" @@ -2799,18 +2768,6 @@ def fake_export_package(pkg, output_dir, **kwargs): assert captured["execution_provider"] == "default" assert captured["text_only"] is False - def test_auto_export_rejects_mage_vl_before_saving(self, tmp_path): - pkg = _make_fake_llm_pkg("mage_vl") - - with ( - mock.patch("mobius.integrations.transformers.build", return_value=pkg), - mock.patch.object(pkg, "save") as save, - pytest.raises(ValueError, match=r"Mage-VL.*patch_positions"), - ): - auto_export("microsoft/Mage-VL", str(tmp_path)) - - save.assert_not_called() - def test_auto_export_produces_genai_config(self, tmp_path): """Mock build() to return a tiny package, verify genai_config.""" import onnx_ir as ir diff --git a/src/mobius/integrations/ort_genai/genai_config.py b/src/mobius/integrations/ort_genai/genai_config.py index 532b92a58..a3f0027ed 100644 --- a/src/mobius/integrations/ort_genai/genai_config.py +++ b/src/mobius/integrations/ort_genai/genai_config.py @@ -156,6 +156,8 @@ class GenaiConfigGenerator: :func:`_default_decoder_inputs`. Must already include KV cache template entries (``past_key_names``, ``past_value_names``). + decoder_outputs: Explicit decoder output name mapping. When + provided, used instead of :func:`_default_decoder_outputs`. """ def __init__( @@ -174,6 +176,7 @@ def __init__( eos_token_id: int | list[int] | None = None, pad_token_id: int | None = None, decoder_inputs: dict[str, str] | None = None, + decoder_outputs: dict[str, str] | None = None, decoder_filename: str | None = None, supports_in_place_kv_cache: bool | None = None, decoder_graph_capture: bool | None = None, @@ -195,6 +198,8 @@ def __init__( # Explicit decoder inputs (from graph introspection); None -> use defaults self._decoder_inputs = decoder_inputs + # Explicit decoder outputs (from graph introspection); None -> use defaults + self._decoder_outputs = decoder_outputs # Explicit decoder filename; None -> use "model.onnx" self._decoder_filename = decoder_filename # Whether the exported decoder ONNX graph supports in-place KV-cache @@ -230,6 +235,7 @@ def from_config( eos_token_id: int | list[int] | None = None, pad_token_id: int | None = None, decoder_inputs: dict[str, str] | None = None, + decoder_outputs: dict[str, str] | None = None, decoder_filename: str | None = None, supports_in_place_kv_cache: bool | None = None, num_cache_layer_slots: int | None = None, @@ -278,6 +284,7 @@ def from_config( eos_token_id=eos_token_id, pad_token_id=pad, decoder_inputs=decoder_inputs, + decoder_outputs=decoder_outputs, decoder_filename=decoder_filename, supports_in_place_kv_cache=supports_in_place_kv_cache, layer_types=getattr(config, "layer_types", None), @@ -459,6 +466,10 @@ def generate(self) -> dict[str, Any]: decoder_inputs = dict(self._decoder_inputs) else: decoder_inputs = _default_decoder_inputs(is_vlm=is_multimodal) + if self._decoder_outputs is not None: + decoder_outputs = dict(self._decoder_outputs) + else: + decoder_outputs = _default_decoder_outputs() decoder_filename = "decoder/model.onnx" if is_multimodal else "model.onnx" decoder: dict[str, Any] = { "session_options": _make_session_options( @@ -469,7 +480,7 @@ def generate(self) -> dict[str, Any]: "head_size": self.head_dim, "hidden_size": self.hidden_size, "inputs": decoder_inputs, - "outputs": _default_decoder_outputs(), + "outputs": decoder_outputs, "num_attention_heads": self.num_attention_heads, "num_hidden_layers": self.num_hidden_layers, "num_key_value_heads": self.num_key_value_heads, diff --git a/src/mobius/models/parakeet_ctc_test.py b/src/mobius/models/parakeet_ctc_test.py index 0c86c2084..ac90923c1 100644 --- a/src/mobius/models/parakeet_ctc_test.py +++ b/src/mobius/models/parakeet_ctc_test.py @@ -168,11 +168,10 @@ def test_parakeet_synthetic_parity_with_padding(): np.testing.assert_allclose(actual, expected, atol=1e-5, rtol=1e-5) -def test_parakeet_rejects_unsupported_ort_genai_export(tmp_path): +def test_parakeet_emits_ort_genai_metadata(tmp_path): _, _, _, package = _build_tiny() - with pytest.raises( - ValueError, - match="does not define a feature-input CTC ASR pipeline", - ): - write_ort_genai_config(package, str(tmp_path)) + result = write_ort_genai_config(package, str(tmp_path)) + + assert "genai_config" in result + assert (tmp_path / "genai_config.json").is_file() diff --git a/tests/cli_test.py b/tests/cli_test.py index 07c288249..dd7a6d8eb 100644 --- a/tests/cli_test.py +++ b/tests/cli_test.py @@ -529,34 +529,6 @@ def test_runtime_ort_genai_propagates_trust_remote_code(self): assert mock_export.call_args.kwargs["trust_remote_code"] is True - def test_runtime_ort_genai_rejects_mage_vl_before_saving(self): - with ( - tempfile.TemporaryDirectory() as tmpdir, - mock.patch("mobius._model_package.ModelPackage.save") as save, - mock.patch( - "mobius.integrations.ort_genai.write_ort_genai_config" - ) as config_writer, - pytest.raises( - SystemExit, - match=r"Mage-VL.*patch_positions.*1D decoder position_ids", - ), - ): - main( - [ - "build", - "--model", - "microsoft/Mage-VL", - tmpdir, - "--no-weights", - "--trust-remote-code", - "--runtime", - "ort-genai", - ] - ) - - save.assert_not_called() - config_writer.assert_not_called() - def test_runtime_onnx_genai_uses_native_vlm_emitter(self): pkg = mock.MagicMock() pkg.items.return_value = [] diff --git a/tests/gguf_test.py b/tests/gguf_test.py index b651254d6..8dcb7c527 100644 --- a/tests/gguf_test.py +++ b/tests/gguf_test.py @@ -9,6 +9,7 @@ from __future__ import annotations +import json from unittest import mock import numpy as np @@ -52,6 +53,9 @@ def _create_tiny_gguf( writer.add_context_length(128) writer.add_feed_forward_length(ffn_size) writer.add_vocab_size(vocab) + writer.add_bos_token_id(1) + writer.add_eos_token_id(2) + writer.add_pad_token_id(3) # Tensors — unquantized random weights rng = np.random.default_rng(42) @@ -544,20 +548,25 @@ def test_contradictory_quantization_flags_error(self, tmp_path): ) assert exc_info.value.code == 2 - def test_ort_genai_runtime_is_rejected_before_artifacts(self, tmp_path): - """build-gguf must not silently ignore an ORT GenAI runtime request.""" + def test_ort_genai_runtime_writes_graph_derived_config(self, tmp_path): from mobius.__main__ import main + path = _create_tiny_gguf(tmp_path / "test.gguf") output_dir = tmp_path / "output" - with pytest.raises(SystemExit, match="does not yet support --runtime ort-genai"): - main( - [ - "build-gguf", - str(tmp_path / "not-downloaded.gguf"), - "--runtime", - "ort-genai", - "--output", - str(output_dir), - ] - ) - assert not output_dir.exists() + main( + [ + "build-gguf", + path, + "--runtime", + "ort-genai", + "--output", + str(output_dir), + ] + ) + + assert (output_dir / "model.onnx").is_file() + assert (output_dir / "genai_config.json").is_file() + config = json.loads((output_dir / "genai_config.json").read_text(encoding="utf-8")) + assert config["model"]["bos_token_id"] == 1 + assert config["model"]["eos_token_id"] == 2 + assert config["model"]["pad_token_id"] == 3 diff --git a/tests/nemotron_h_real_weight_test.py b/tests/nemotron_h_real_weight_test.py index cc17abdda..d2dd232d9 100644 --- a/tests/nemotron_h_real_weight_test.py +++ b/tests/nemotron_h_real_weight_test.py @@ -58,13 +58,13 @@ def test_load_eos_token_ids_unions_generation_and_model_config(tmp_path): assert validator.load_eos_token_ids(tmp_path) == {2, 11} -def test_run_token_ids_stops_on_eos_without_unused_forward(tmp_path, monkeypatch): +def test_run_token_ids_accepts_olive_logits_and_stops_on_eos(tmp_path, monkeypatch): validator = _load_validator() inference_globals = validator.run_token_ids.__globals__ (tmp_path / "model.onnx").touch() class _Output: - name = "logits" + name = "logits_Q4" class _Session: @staticmethod @@ -303,9 +303,20 @@ def test_nemotron_h_3_5_reduced_real_l5(reduced_real_outputs, model_type): def test_nemotron_h_3_5_reduced_real_fp16_cuda(reduced_real_fp16_cuda, model_type): del model_type _validator, _state, package_dir = reduced_real_fp16_cuda + import onnx_ir as ir assert (package_dir / "model.onnx").is_file() assert (package_dir / "model.onnx.data").is_file() + model = ir.load(package_dir / "model.onnx") + op_types = {(node.domain, node.op_type) for node in model.graph.all_nodes()} + assert ("", "Attention") in op_types + assert ("", "Scan") in op_types + assert ("", "Conv") in op_types + assert not any( + domain == "com.microsoft" + and op_type in {"GroupQueryAttention", "LinearAttention", "CausalConvWithState"} + for domain, op_type in op_types + ) @pytest.mark.integration @@ -355,7 +366,7 @@ def test_nemotron_h_3_5_olive_q4_final_package( node.domain == "com.microsoft" and node.op_type == "MatMulNBits" for node in quantized_model.graph.all_nodes() ) - == 15 + == 17 ) generated, logits, profile_path = validator.run_token_ids( From e02aa8a399ccec41442ab8aa2584fe30f5694a9b Mon Sep 17 00:00:00 2001 From: Justin Chu Date: Fri, 14 Aug 2026 19:54:10 -0700 Subject: [PATCH 11/13] Validate GGUF generation token metadata Drop serialized null-token sentinels and out-of-vocabulary special-token IDs while preserving valid BOS and padding metadata. Combine embedded EOS, EOT, and EOM IDs into the generation stop-token contract. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Signed-off-by: Justin Chu --- .../integrations/gguf/_config_mapping.py | 41 ++++++++++++++++--- src/mobius/integrations/gguf/_reader_test.py | 28 ++++++++++++- 2 files changed, 62 insertions(+), 7 deletions(-) diff --git a/src/mobius/integrations/gguf/_config_mapping.py b/src/mobius/integrations/gguf/_config_mapping.py index a52f498c4..77a508306 100644 --- a/src/mobius/integrations/gguf/_config_mapping.py +++ b/src/mobius/integrations/gguf/_config_mapping.py @@ -108,11 +108,15 @@ _TOKENIZER_KEY_MAP: dict[str, str] = { "tokenizer.ggml.bos_token_id": "bos_token_id", - "tokenizer.ggml.eos_token_id": "eos_token_id", "tokenizer.ggml.padding_token_id": "pad_token_id", # Non-standard but emitted by some multimodal converters. "tokenizer.ggml.image_token_id": "image_token_id", } +_STOP_TOKEN_KEYS = ( + "tokenizer.ggml.eos_token_id", + "tokenizer.ggml.eot_token_id", + "tokenizer.ggml.eom_token_id", +) _ARCH_KEY_MAPS: dict[str, dict[str, str]] = { # Both spellings are accepted wherever muse-glimmer is recognized. @@ -221,9 +225,34 @@ def _extract_config_fields( if isinstance(tokens, list): hf_fields["vocab_size"] = len(tokens) + vocab_size = hf_fields.get("vocab_size") + + def _valid_token_id(value: Any) -> int | None: + token_id = int(value) + if token_id < 0 or token_id == 0xFFFFFFFF: + return None + if vocab_size is not None and token_id >= int(vocab_size): + return None + return token_id + for gguf_key, config_key in _TOKENIZER_KEY_MAP.items(): - if gguf_key in metadata: - hf_fields[config_key] = int(metadata[gguf_key]) + if ( + gguf_key in metadata + and (token_id := _valid_token_id(metadata[gguf_key])) is not None + ): + hf_fields[config_key] = token_id + + stop_token_ids: list[int] = [] + for gguf_key in _STOP_TOKEN_KEYS: + if gguf_key not in metadata: + continue + token_id = _valid_token_id(metadata[gguf_key]) + if token_id is not None and token_id not in stop_token_ids: + stop_token_ids.append(token_id) + if stop_token_ids: + hf_fields["eos_token_id"] = ( + stop_token_ids[0] if len(stop_token_ids) == 1 else stop_token_ids + ) # Standard GGUF has no dedicated image-token key. Preserve the canonical # HuggingFace placeholder when it is embedded in the tokenizer vocabulary. @@ -407,11 +436,13 @@ def gguf_to_config( if isinstance(swiglu_limit, (list, np.ndarray)): swiglu_limit = swiglu_limit[0] if len(swiglu_limit) else 0.0 - special_token_fields = { + special_token_fields: dict[str, Any] = { name: int(hf_fields[name]) - for name in ("bos_token_id", "eos_token_id", "pad_token_id", "image_token_id") + for name in ("bos_token_id", "pad_token_id", "image_token_id") if hf_fields.get(name) is not None } + if (eos_token_id := hf_fields.get("eos_token_id")) is not None: + special_token_fields["eos_token_id"] = eos_token_id config = ArchitectureConfig( hidden_size=hidden_size, intermediate_size=hf_fields.get("intermediate_size", 4 * hidden_size), diff --git a/src/mobius/integrations/gguf/_reader_test.py b/src/mobius/integrations/gguf/_reader_test.py index 99426b6ee..c3cde77b5 100644 --- a/src/mobius/integrations/gguf/_reader_test.py +++ b/src/mobius/integrations/gguf/_reader_test.py @@ -566,15 +566,39 @@ def test_extract_config_fields_preserves_tokenizer_metadata(self): "tokenizer.ggml.bos_token_id": 1, "tokenizer.ggml.eos_token_id": 2, "tokenizer.ggml.padding_token_id": 3, - "tokenizer.ggml.tokens": ["a", "", "b"], + "tokenizer.ggml.eot_token_id": 4, + "tokenizer.ggml.eom_token_id": 2, + "tokenizer.ggml.tokens": [ + "a", + "", + "b", + "c", + "d", + ], }, ) assert fields["bos_token_id"] == 1 - assert fields["eos_token_id"] == 2 + assert fields["eos_token_id"] == [2, 4] assert fields["pad_token_id"] == 3 assert fields["image_token_id"] == 1 + def test_extract_config_fields_omits_invalid_token_sentinels(self): + fields = _extract_config_fields( + "llama", + { + "tokenizer.ggml.tokens": ["a"], + "tokenizer.ggml.bos_token_id": 0xFFFFFFFF, + "tokenizer.ggml.eos_token_id": 0xFFFFFFFF, + "tokenizer.ggml.eot_token_id": -1, + "tokenizer.ggml.padding_token_id": 999, + }, + ) + + assert "bos_token_id" not in fields + assert "eos_token_id" not in fields + assert "pad_token_id" not in fields + def test_infer_tie_embeddings_true(self, tied_gguf: Path): model = GGUFModel(tied_gguf) assert _infer_tie_embeddings(model) is True From eb1b69d7a4da84b495b9d112b41193d880427403 Mon Sep 17 00:00:00 2001 From: Justin Chu Date: Wed, 19 Aug 2026 16:51:28 -0700 Subject: [PATCH 12/13] Add production Qwen3.8-27B support (#498) - add production support for pinned `Qwen/Qwen3.8-27B@1d4bf0f2ff6012fd82039f2fa52739d0dd7c60c0` as a checkpoint alias of the existing dense Qwen3.5 hybrid vision-language architecture - preserve the exact 64-layer 48-DeltaNet/16-full-attention text schedule, 27-layer vision encoder, image/video processor contracts, mixed-batch scatter order, dtype boundaries, and optional separately packaged MTP drafter - preserve canonical `qwen3_5` registry regression coverage on `Qwen/Qwen3.5-2B`; Qwen3.8 uses its own pinned case/config assertions instead of replacing that mapping - emit topology-faithful ORT GenAI metadata without a runtime capability gate: decoder-only packages keep `qwen3_5_text`, while multimodal Qwen3.5/Qwen3.8 packages emit `qwen3_5` - remove the standalone Qwen3.8 Olive example, its coupled reduced-real test module, and its private reduced golden fixtures; core model, processor, metadata, and parity coverage remain unchanged Stacked on #487 final skill-split head `a8cd77570bca980861eeadab9e6fa077464f0138`. Exact PR head: `59af1a51f1a162e9444a9be1ea5785ca4ec0ef0b`. Tracks #483. - the real processor from the pinned official revision drives deterministic tiny-model HF/ONNX pipeline parity for image-only, video-only, and a mixed two-row batch with opposite placeholder order: image max abs `0.00697723`, video `0.00411959`, mixed `0.00677243`; cosine is above `0.99999` and final argmax matches - exact raw-config assertions cover the 48/16 hybrid schedule, text/vision dimensions, processor token contracts, and intentional separate MTP packaging - packed-vision coordinates are emitted once and shared by interpolation, rotary, and frame-boundary consumers; media ownership uses boundary `ScatterElements` + `CumSum` in O(patches + media) - deterministic Qwen3.5-VL benchmark: `428 -> 467` top-level nodes (`+9.1%`, model size unchanged), below the unchanged 10% blocker threshold - topology-specific metadata regression coverage generates both package shapes: decoder-only `qwen3_5_text` and multimodal `qwen3_5` - post-removal YAML/coverage/focused suite: `987 passed, 227 skipped`; diff lintrunner clean; independent review found no stale reduced/Olive claims or dangling references - canonical non-integration suite before the deletion-only follow-up: `3962 passed, 61 skipped` The pinned official checkpoint is 55.6 GB and exceeds hosted CI storage and GPU memory; no reduced or quantized private fixture is committed. The Qwen3.8 case records this exact `ci_skip_reason` rather than claiming retained executable evidence. Earlier reduced-real and Olive development artifacts were removed and are not part of the final PR validation surface. Current head `59af1a51f1a162e9444a9be1ea5785ca4ec0ef0b`: - Benchmark run `32314325378` is complete and successful (base, head, and comparison). - CI run `32314325553` is still active. Lint, affected-model detection, L1, L3, and lintrunner have completed successfully; the Linux/Windows test matrices are in progress, while Integration (fast), L4, and L5 are queued. - Architecture Diff run `32314325423` is in progress. Historical run `31870831200` belongs to pre-removal head `067cb489599687b74f96574de20133cb96e016a7`, not the final head. It is retained only as historical context: its three red jobs reproduced on the exact #487 base run `31869640946`, but it is not cited as current-head CI evidence. PR #498 remains ready for maintainer review; auto-merge is off. --------- Signed-off-by: Justin Chu Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- src/mobius/_configs/_vision_defaults.py | 3 + src/mobius/components/_qwen3_vl_vision.py | 458 +++++------------- .../components/_qwen3_vl_vision_test.py | 66 +++ .../integrations/ort_genai/auto_export.py | 17 +- .../ort_genai/auto_export_test.py | 154 ++++++ src/mobius/models/qwen35.py | 20 +- src/mobius/models/qwen35_test.py | 275 ++++++++++- src/mobius/models/qwen_vl.py | 64 ++- src/mobius/tasks/_vision_language_3model.py | 5 +- .../cases/vision-language/qwen3_8-27b.yaml | 26 + tests/integration_test.py | 196 +++++++- 11 files changed, 912 insertions(+), 372 deletions(-) create mode 100644 src/mobius/components/_qwen3_vl_vision_test.py create mode 100644 testdata/cases/vision-language/qwen3_8-27b.yaml 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/components/_qwen3_vl_vision.py b/src/mobius/components/_qwen3_vl_vision.py index 1fc983989..7da0d5850 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,291 +490,147 @@ 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, - ) - - # 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), + 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( + patch_ends, + op.Constant(value_ints=[1, 0]), + op.Constant(value_int=0), ) - # 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, + 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), ) - 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]), + # 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, ) - padded = body_op.Pad(pos_embeds, pads, 0.0) - padded.name = "padded_pos_embed" - body_graph.outputs.append(padded) + 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 + 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, frame_local_ids, patch_ids, total_patches + + def _interpolate_pos_embed(self, op, coordinates): + """Bilinearly interpolate learned positions for the packed media stream. - rename_subgraph_values(body_graph, "posemb_body_") - - scan_result = op.Scan( - grid_thw, - body=body_graph, - num_scan_inputs=1, - _outputs=1, - ) # (num_images, max_patches, hidden_size) + Matches HuggingFace ``Qwen3VLVisionModel.fast_pos_embed_interpolate``. - return compact_scan_output(op, scan_result, patches_per_image) + Args: + op: OpBuilder instance. + coordinates: Shared packed ``(rows, cols, H, W)`` coordinate values. - def _compute_rotary_pos_ids(self, op, grid_thw): - """Compute 2D rotary position IDs for all images via ONNX Scan. + Returns: + Position embeddings ``(total_patches, hidden_size)``. + """ + n = self.num_grid_per_side + 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 + 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_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.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)) + 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. + 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, coordinates): + """Combine shared packed coordinates into 2D rotary position IDs. 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) + rows, cols = coordinates + return op.Concat(op.Unsqueeze(rows, [1]), op.Unsqueeze(cols, [1]), axis=1) - 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) - - 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 using ONNX Scan - to handle per-image ``repeat_interleave(hw, 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 - 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, + 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)), + ) + return op.Concat( + frame_starts, + op.Unsqueeze(total_patches, [0]), axis=0, ) - padded = body_op.Pad( - hw_repeated, - pads, - body_op.Constant(value_int=0), - ) - 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, - ) - hw_flat = compact_scan_output(op, scan_hw, T_col) - cu = op.CumSum(hw_flat, op.Constant(value_int=0)) - return op.Pad(cu, op.Constant(value_ints=[1, 0]), op.Constant(value_int=0)) def forward( self, @@ -865,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 (Scan body computes in float32). - pos_embeds = self._interpolate_pos_embed(op, grid_thw) + # Cast to match hidden_states dtype (interpolation computes in float32). + 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 diff --git a/src/mobius/integrations/ort_genai/auto_export.py b/src/mobius/integrations/ort_genai/auto_export.py index 987f0dd3f..88143284a 100644 --- a/src/mobius/integrations/ort_genai/auto_export.py +++ b/src/mobius/integrations/ort_genai/auto_export.py @@ -94,8 +94,9 @@ "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", + "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(). @@ -138,6 +139,7 @@ "qwen3_vl_text", "qwen3_5", "qwen3_5_vl", + "qwen3_5_text", "qwen3_5_moe", "qwen3_5_moe_text", "videochat_flash_qwen", @@ -187,6 +189,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") @@ -1395,7 +1402,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 f65ea7a44..445e2ff78 100644 --- a/src/mobius/integrations/ort_genai/auto_export_test.py +++ b/src/mobius/integrations/ort_genai/auto_export_test.py @@ -82,6 +82,9 @@ 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("qwen3_5_text") == "qwen3_5_text" assert _resolve_ort_genai_model_type("gemma2") == "gemma" assert _resolve_ort_genai_model_type("llama") == "llama" @@ -141,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 @@ -388,6 +409,38 @@ def test_qwen35_moe_text_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). @@ -1123,6 +1176,107 @@ 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_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( + { + "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_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( diff --git a/src/mobius/models/qwen35.py b/src/mobius/models/qwen35.py index 359f5d1ba..d3746b828 100644 --- a/src/mobius/models/qwen35.py +++ b/src/mobius/models/qwen35.py @@ -282,9 +282,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] = {} @@ -471,9 +472,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.*``) @@ -599,9 +600,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 @@ -699,7 +699,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 adbfc0311..9bc1ab49f 100644 --- a/src/mobius/models/qwen35_test.py +++ b/src/mobius/models/qwen35_test.py @@ -15,21 +15,294 @@ import dataclasses +import numpy as np import onnx_ir as ir import pytest import torch +from transformers.models.qwen3_5.configuration_qwen3_5 import Qwen3_5Config -from mobius._configs import QuantizationConfig, VisionConfig +from mobius._configs import ( + ArchitectureConfig, + QuantizationConfig, + Qwen35MtpConfig, + VisionConfig, +) +from mobius._registry import registry from mobius._testing import make_config +from mobius._testing.ort_inference import OnnxModelSession from mobius.models.qwen35 import ( Qwen35CausalLMModel, Qwen35MoECausalLMModel, Qwen35MoEVL3ModelCausalLMModel, 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.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 + 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 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: diff --git a/src/mobius/models/qwen_vl.py b/src/mobius/models/qwen_vl.py index b97a0d4e5..85950f1c7 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,45 @@ 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.Not(op.Equal(input_ids, input_ids)) + 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_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)), + 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( + flat_image_mask_bool, + 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 +1075,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..1b370cc82 --- /dev/null +++ b/testdata/cases/vision-language/qwen3_8-27b.yaml @@ -0,0 +1,26 @@ +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: >- + 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 + 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/tests/integration_test.py b/tests/integration_test.py index d6e8127c6..0cbe5d002 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(): From 8ccfc25977d9551c7dfc07088eea1827a5fcb8b8 Mon Sep 17 00:00:00 2001 From: Justin Chu Date: Wed, 19 Aug 2026 17:14:22 -0700 Subject: [PATCH 13/13] Update Qwen3.8 weight-loading import Use the integrations namespace introduced on main so the rebased Qwen3.8 parity test remains executable after the ecosystem refactor. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Signed-off-by: Justin Chu --- tests/integration_test.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/integration_test.py b/tests/integration_test.py index 0cbe5d002..3e7ffc1b9 100644 --- a/tests/integration_test.py +++ b/tests/integration_test.py @@ -3841,7 +3841,7 @@ def test_qwen38_vl_image_video_mixed_pipeline_matches_huggingface(): ) from mobius import build_from_module - from mobius._weight_loading import apply_weights + from mobius.integrations._weight_loading import apply_weights hf_config = _make_tiny_qwen35_vl_config(keep_production_vocab=True) arch_config = ArchitectureConfig.from_transformers(