From 12426484dcde7cd3bfe6f612b6e072ce0a723143 Mon Sep 17 00:00:00 2001 From: Justin Chu Date: Tue, 21 Apr 2026 15:54:05 +0000 Subject: [PATCH 1/8] Restructure ort-genai-config skill for progressive disclosure Split the 799-line SKILL.md into a 304-line overview + 3 reference files: - references/genai-config-fields.md: complete field tables for all sections - references/processor-config-fields.md: processor_config.json full reference - references/multimodal-pipeline.md: VLM pipeline architecture and helpers No content removed; all fields, examples, and code snippets preserved. SKILL.md now contains overview, model type registry, minimal valid config, processor overview, troubleshooting, and cross-reference directives. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Justin Chu --- .agents/skills/ort-genai-config/SKILL.md | 304 ++++++++++++++++++ .../references/genai-config-fields.md | 301 +++++++++++++++++ .../references/multimodal-pipeline.md | 158 +++++++++ .../references/processor-config-fields.md | 176 ++++++++++ 4 files changed, 939 insertions(+) create mode 100644 .agents/skills/ort-genai-config/SKILL.md create mode 100644 .agents/skills/ort-genai-config/references/genai-config-fields.md create mode 100644 .agents/skills/ort-genai-config/references/multimodal-pipeline.md create mode 100644 .agents/skills/ort-genai-config/references/processor-config-fields.md diff --git a/.agents/skills/ort-genai-config/SKILL.md b/.agents/skills/ort-genai-config/SKILL.md new file mode 100644 index 00000000..fc2b7bfe --- /dev/null +++ b/.agents/skills/ort-genai-config/SKILL.md @@ -0,0 +1,304 @@ +--- +name: ort-genai-config +description: > + Use this skill when generating genai_config.json or processor_config.json + for onnxruntime-genai model exports, debugging ORT GenAI model loading + errors, understanding the model type registry, or integrating ONNX models + with the onnxruntime-genai runtime. Covers the full config format, the + MultiModal pipeline architecture (vision/audio/embedding/decoder), and + the ort-extensions processor_config.json format. +--- + +# Skill: ORT GenAI Config Format + +## When to use + +Use this skill when: + +- Writing `genai_config.json` for a new model export +- Writing `processor_config.json` for image/audio preprocessing +- Debugging ORT GenAI model loading errors (protobuf parsing, missing keys) +- Understanding how the ORT GenAI pipeline feeds inputs to vision, embedding, + and decoder models +- Adding support for a new model type in ORT GenAI + +## Detailed references + +Read these companion documents for exhaustive field-by-field details: + +- Read **[`references/genai-config-fields.md`](references/genai-config-fields.md)** + when you need the complete field table for any `genai_config.json` section + (model, decoder, vision, embedding, speech, encoder, search, engine, + session_options). +- Read **[`references/processor-config-fields.md`](references/processor-config-fields.md)** + when you need the full `processor_config.json` transform reference, the + HuggingFace-to-ort-extensions conversion code, or the Qwen2.5-VL example. +- Read **[`references/multimodal-pipeline.md`](references/multimodal-pipeline.md)** + when you need the VLM 3-model prompt/generation flow, input routing + details, QwenImageProcessor output tensors, the multimodal processor + factory table, or the full `_write_genai_config` helper code. + +--- + +## Overview + +onnxruntime-genai loads models from a directory containing: + +``` +model_dir/ +├── genai_config.json # Required — model config + search params +├── model.onnx # Decoder model +├── model.onnx.data # External weights (optional) +├── vision.onnx # Vision encoder (multimodal only) +├── embedding.onnx # Embedding model (multimodal only) +├── tokenizer.json # Tokenizer (HuggingFace format) +├── tokenizer_config.json # Tokenizer config +├── chat_template.jinja # Chat template (optional) +└── processor_config.json # Image processor (multimodal only) +``` + +## genai_config.json — Structure + +The config has three top-level sections: + +```json +{ + "model": { ... }, + "search": { ... }, + "engine": { ... } +} +``` + +- **`model`**: Model architecture — `type`, token IDs, and sub-sections + `decoder`, `vision`, `embedding`, `speech`, `encoder`. +- **`search`**: Generation parameters — sampling, beam search, max length. +- **`engine`** *(optional)*: Batched serving (dynamic or static batching). + +Key model-level fields: `type` (required), `vocab_size`, `context_length` +(required), `eos_token_id`, `pad_token_id`. VLMs also need `image_token_id` +and `vision_start_token_id`. + +Key decoder fields: `filename`, `hidden_size`, `head_size`, +`num_attention_heads`, `num_key_value_heads`, `num_hidden_layers`, plus +`inputs`/`outputs` name mappings. + +> For the complete field-by-field tables, see +> [`references/genai-config-fields.md`](references/genai-config-fields.md). + +--- + +## Model type registry + +### LLM (decoder-only → `DecoderOnly_Model`) + +``` +chatglm, decoder, ernie4_5, gemma, gemma2, gemma3_text, gpt2, +gptoss, granite, internlm2, llama, mistral, nemotron, olmo, +phi, phimoe, phi3, phi3small, qwen2, qwen3, smollm3 +``` + +### VLM (vision-language → `MultiModalLanguageModel`) + +``` +fara, gemma3, phi3v, qwen2_5_vl +``` + +### MMM (vision + audio → `MultiModalLanguageModel`) + +``` +phi4mm +``` + +### ALM (audio-language → `WhisperModel`) + +``` +whisper +``` + +### Pipeline models → `DecoderOnlyPipelineModel` + +``` +phi3small_pipeline, qwen2_5_vl_pipeline +``` + +### Special handling + +- `fara` / `qwen2_5_vl` with non-empty `model.decoder.pipeline` → + `Qwen2_5_VL_PipelineModel` +- `IsQwen25VL()` (type == `"fara"` or `"qwen2_5_vl"`) enables 3D MRoPE + position ID handling +- `gpt2` has a special code path (`Gpt_Model`) but is also in the LLM list + +--- + +## Minimal valid config — decoder-only LLM + +```json +{ + "model": { + "type": "llama", + "vocab_size": 32000, + "context_length": 4096, + "eos_token_id": 2, + "pad_token_id": 0, + "decoder": { + "filename": "model.onnx", + "hidden_size": 4096, + "head_size": 128, + "num_attention_heads": 32, + "num_key_value_heads": 8, + "num_hidden_layers": 32, + "inputs": { + "input_ids": "input_ids", + "attention_mask": "attention_mask", + "position_ids": "position_ids", + "past_key_names": "past_key_values.%d.key", + "past_value_names": "past_key_values.%d.value" + }, + "outputs": { + "logits": "logits", + "present_key_names": "present.%d.key", + "present_value_names": "present.%d.value" + } + } + }, + "search": { + "do_sample": false, + "max_length": 4096, + "num_beams": 1, + "past_present_share_buffer": false + } +} +``` + +VLM models additionally require `model.vision`, `model.embedding`, +`image_token_id`, and `vision_start_token_id`. See +[`references/genai-config-fields.md`](references/genai-config-fields.md) for +the full vision/embedding/speech/encoder schemas. + +--- + +## processor_config.json overview + +> **Critical:** ORT GenAI expects the **ort-extensions** format — NOT the +> HuggingFace `processor_config.json` format. HF uses `"image_processor"` +> as the top key; ort-extensions uses `"processor"` with an ordered +> transform pipeline. + +Structure: + +```json +{ + "processor": { + "name": "", + "transforms": [ + { "operation": { "name": "...", "type": "...", "attrs": { ... } } } + ] + } +} +``` + +Transform types: `DecodeImage`, `ConvertRGB`, `Resize`, `Rescale`, +`Normalize`, `PatchImage`. + +> For the full Qwen2.5-VL example, transform field reference, and +> HuggingFace conversion code, see +> [`references/processor-config-fields.md`](references/processor-config-fields.md). + +--- + +## MultiModal pipeline overview + +VLM models use a 3-model split: + +``` +pixel_values + image_grid_thw → [vision.onnx] → image_features + │ +input_ids + image_features → [embedding.onnx] → inputs_embeds + │ +inputs_embeds + position_ids → [model.onnx] → logits + + past_kv +``` + +During generation, the vision model runs once at prompt time. The embedding +and decoder models run each token step. + +> For the full generation flow, input routing, QwenImageProcessor output +> tensors, and the multimodal processor factory, see +> [`references/multimodal-pipeline.md`](references/multimodal-pipeline.md). + +--- + +## Troubleshooting + +### "Protobuf parsing failed" + +Missing `model.vision` and/or `model.embedding` sections in genai_config.json. +VLM models require all three model sections. + +### "key 'processor' not found" + +The `processor_config.json` is in HuggingFace format instead of ort-extensions +format. The HF format has `"image_processor"` as the top key; ORT extensions +needs `"processor"` with a transforms pipeline. + +### "Missing Input: cu_window_seqlens" + +The vision ONNX model expects packed-attention inputs that the ORT GenAI +processor doesn't provide. Either: +1. Compute them externally and inject via NamedTensors, or +2. Modify the vision model to compute them from `image_grid_thw` internally + +### "input_ids size exceeds max length" + +For image prompts, the tokenized input_ids (including image_pad tokens) can +be much longer than the default `max_length` in search options. Use +`params.set_search_options(max_length=4096)` or a sufficiently large value. + +### "OrtValue shape verification failed" + +Mismatch between `num_image_tokens` (computed by the processor) and the +actual vision model output shape. Ensure the same image processor is used +consistently — don't mix ORT GenAI processor output with HF processor +pixel_values. + +### Image not recognized despite being processed + +If the model generates coherent text but fails to describe image contents: + +1. **Missing `image_token_id` or `spatial_merge_size`:** Without these, + ORT GenAI cannot compute 3D M-RoPE position IDs. Add `image_token_id`, + `vision_start_token_id` at model level and `spatial_merge_size` under + model.vision. + +2. **processor_config.json resize mismatch:** The Resize transform uses + `width`/`height` as direct target dimensions. If too small, image loses + detail. Compute correct dimensions: + ```python + factor = patch_size * merge_size # 28 + new_h = max(factor, round(orig_h / factor) * factor) + new_w = max(factor, round(orig_w / factor) * factor) + ``` + +3. **ONNX model numerical accuracy:** Logits may differ from HF (typical + max_diff ~8 for VLMs), causing greedy decoding to diverge after 3-4 + tokens. + +--- + +## Source reference files + +- **ORT GenAI config structs:** + `/home/justinchu/dev/onnxruntime-genai/src/config.h` +- **ORT GenAI config parsing:** + `/home/justinchu/dev/onnxruntime-genai/src/config.cpp` +- **Model type registry:** + `/home/justinchu/dev/onnxruntime-genai/src/model_type.h` +- **VLM pipeline:** + `/home/justinchu/dev/onnxruntime-genai/src/models/multi_modal.cpp` +- **Qwen image processor:** + `/home/justinchu/dev/onnxruntime-genai/src/models/qwen2_5_vl_image_processor.cpp` +- **Reference processor_config.json:** + `/home/justinchu/dev/onnxruntime-genai/test/test_models/qwen-vision-preprocessing/processor_config.json` +- **Example genai_config generation:** + `examples/qwen25_vl_ort_genai.py` diff --git a/.agents/skills/ort-genai-config/references/genai-config-fields.md b/.agents/skills/ort-genai-config/references/genai-config-fields.md new file mode 100644 index 00000000..c9c79263 --- /dev/null +++ b/.agents/skills/ort-genai-config/references/genai-config-fields.md @@ -0,0 +1,301 @@ +# genai_config.json — Complete Field Reference + +This document is the exhaustive field-by-field reference for every section of +`genai_config.json`. For an overview and minimal config example, see the +parent [SKILL.md](../SKILL.md). + +--- + +## model section — Model-level fields + +| Field | Type | Required | Description | +|---|---|---|---| +| `type` | string | **yes** | Model type identifier (see registry in SKILL.md) | +| `vocab_size` | int | yes | Vocabulary size | +| `context_length` | int | **yes** | Maximum context length; must be > 0 | +| `bos_token_id` | int | no | Beginning-of-sequence token | +| `eos_token_id` | int \| int[] | no | End-of-sequence token(s); defaults to `pad_token_id` | +| `pad_token_id` | int | no | Padding token | +| `sep_token_id` | int | no | Separator token | +| `decoder_start_token_id` | int | no | Decoder start token (encoder-decoder models) | +| `image_token_id` | int | VLM | Token ID for image placeholders (e.g. 151655 for Qwen2.5-VL). **Required** for 3D M-RoPE position ID computation. | +| `video_token_id` | int | no | Token ID for video placeholders (e.g. 151656) | +| `vision_start_token_id` | int | VLM | Token ID for `<\|vision_start\|>` (e.g. 151652). Used to locate image/video regions in input_ids. | + +--- + +## model.decoder + +The decoder (text model) configuration. + +### Core fields + +| Field | Type | Required | Description | +|---|---|---|---| +| `filename` | string | **yes** | ONNX model filename (e.g. `"model.onnx"`) | +| `hidden_size` | int | yes | Hidden dimension | +| `head_size` | int | yes | Size per attention head | +| `num_attention_heads` | int | yes | Number of query attention heads | +| `num_key_value_heads` | int | yes | Number of KV heads (for GQA) | +| `num_hidden_layers` | int | yes | Number of transformer layers | +| `session_options` | object | no | ORT session configuration | +| `run_options` | object | no | ORT run options | + +### Decoder inputs + +```json +"inputs": { + "input_ids": "input_ids", + "inputs_embeds": "inputs_embeds", + "attention_mask": "attention_mask", + "position_ids": "position_ids", + "past_key_names": "past_key_values.%d.key", + "past_value_names": "past_key_values.%d.value" +} +``` + +The `%d` in `past_key_names` / `past_value_names` is replaced with the layer +index (0 to num_hidden_layers-1) at load time. + +Additional optional inputs for advanced scenarios: + +```json +"past_names": "", +"cross_past_key_names": "", +"cross_past_value_names": "", +"past_key_values_length": "past_key_values_length", +"past_sequence_length": "past_sequence_length", +"current_sequence_length": "current_sequence_length", +"total_sequence_length": "total_sequence_length", +"cache_indirection": "cache_indirection", +"encoder_hidden_states": "encoder_hidden_states", +"encoder_attention_mask": "encoder_attention_mask", +"cumulative_sequence_lengths": "cumulative_sequence_lengths", +"past_sequence_lengths": "past_sequence_lengths", +"block_table": "block_table" +``` + +### Decoder outputs + +```json +"outputs": { + "logits": "logits", + "present_key_names": "present.%d.key", + "present_value_names": "present.%d.value" +} +``` + +### Sliding window (optional) + +```json +"sliding_window": { + "window_size": 4096, + "pad_value": -1, + "alignment": "right", + "slide_key_value_cache": true, + "slide_inputs": true, + "layers": [0, 2, 4] +} +``` + +--- + +## model.embedding + +Required for VLM and MMM models. Merges text token embeddings with vision/audio +features. + +```json +"embedding": { + "filename": "embedding.onnx", + "inputs": { + "input_ids": "input_ids", + "image_features": "image_features", + "audio_features": "audio_features" + }, + "outputs": { + "inputs_embeds": "inputs_embeds" + } +} +``` + +--- + +## model.vision + +Required for VLM and MMM models. + +| Field | Type | Default | Description | +|---|---|---|---| +| `filename` | string | — | Vision ONNX model | +| `config_filename` | string | `"processor_config.json"` | Processor config file | +| `adapter_filename` | string | — | Optional adapter model | +| `spatial_merge_size` | int | 2 | **Required for Qwen2.5-VL.** Controls how many vision patches are merged into one token. Used to compute grid dimensions for 3D M-RoPE position IDs (h/merge × w/merge). | +| `tokens_per_second` | float | 2.0 | Video tokens/second | + +### Vision inputs + +```json +"inputs": { + "pixel_values": "pixel_values", + "image_sizes": "image_sizes", + "image_grid_thw": "image_grid_thw", + "attention_mask": "image_attention_mask" +} +``` + +### Vision outputs + +```json +"outputs": { + "image_features": "image_features" +} +``` + +### Vision pipeline (optional) + +For models that split vision into stages (e.g. patch_embed → attention → +merger): + +```json +"pipeline": [ + { + "filename": "patch_embed.onnx", + "model_id": "patch_embed", + "inputs": ["pixel_values"], + "outputs": ["patch_embeddings"], + "run_on_cpu": false, + "session_options": {} + } +] +``` + +--- + +## model.speech + +For audio-language models (whisper, phi4mm). + +```json +"speech": { + "filename": "speech.onnx", + "config_filename": "audio_processor_config.json", + "inputs": { + "audio_embeds": "audio_embeds", + "attention_mask": "audio_attention_mask", + "audio_sizes": "audio_sizes", + "audio_projection_mode": "audio_projection_mode" + }, + "outputs": { + "audio_features": "audio_features" + } +} +``` + +--- + +## model.encoder + +For encoder-decoder models (whisper). + +```json +"encoder": { + "filename": "encoder.onnx", + "hidden_size": 1280, + "num_attention_heads": 20, + "num_hidden_layers": 32, + "head_size": 64, + "inputs": { + "input_ids": "input_ids", + "attention_mask": "attention_mask" + }, + "outputs": { + "encoder_hidden_states": "encoder_hidden_states" + } +} +``` + +--- + +## search section + +Controls generation behavior. + +| Field | Type | Default | Description | +|---|---|---|---| +| `do_sample` | bool | false | Sampling vs greedy | +| `min_length` | int | 0 | Minimum output length | +| `max_length` | int | context_length | Maximum total length (prompt + output) | +| `batch_size` | int | 1 | Batch size | +| `num_beams` | int | 1 | Beam width (1 = greedy) | +| `num_return_sequences` | int | 1 | Sequences to return | +| `top_k` | int | 50 | Top-K sampling | +| `top_p` | float | 0.0 | Nucleus sampling | +| `temperature` | float | 1.0 | Sampling temperature | +| `repetition_penalty` | float | 1.0 | Repetition penalty (1.0 = none) | +| `length_penalty` | float | 1.0 | Beam search length penalty | +| `early_stopping` | bool | true | Stop beam search early | +| `past_present_share_buffer` | bool | false | Share KV cache buffer (CUDA) | +| `random_seed` | int | -1 | RNG seed (-1 = random) | +| `chunk_size` | int | — | Prefill chunking size | + +--- + +## engine section (optional) + +For batched serving. + +```json +"engine": { + "dynamic_batching": { + "block_size": 256, + "num_blocks": 16, + "gpu_utilization_factor": 0.9, + "max_batch_size": 16 + } +} +``` + +Or static batching: + +```json +"engine": { + "static_batching": { + "max_batch_size": 4 + } +} +``` + +Dynamic and static batching are mutually exclusive. + +--- + +## session_options + +Nested inside `decoder`, `encoder`, `vision`, `speech`, or `embedding`. + +```json +"session_options": { + "intra_op_num_threads": 8, + "inter_op_num_threads": 1, + "log_id": "onnxruntime-genai", + "log_severity_level": 2, + "enable_cpu_mem_arena": true, + "enable_mem_pattern": true, + "enable_profiling": "profile.json", + "graph_optimization_level": "ORT_ENABLE_EXTENDED", + "provider_options": [ + { + "cuda": { + "device_id": "0" + } + } + ] +} +``` + +Graph optimization levels: `ORT_DISABLE_ALL`, `ORT_ENABLE_BASIC`, +`ORT_ENABLE_EXTENDED`, `ORT_ENABLE_ALL`. + +Provider names are normalized: `"qnn"` → `"QNN"`, `"dml"` → `"DML"`, +`"webgpu"` → `"WebGPU"`, `"openvino"` → `"OpenVINO"`. diff --git a/.agents/skills/ort-genai-config/references/multimodal-pipeline.md b/.agents/skills/ort-genai-config/references/multimodal-pipeline.md new file mode 100644 index 00000000..6ec3174e --- /dev/null +++ b/.agents/skills/ort-genai-config/references/multimodal-pipeline.md @@ -0,0 +1,158 @@ +# MultiModal Pipeline Architecture + +This document describes the detailed ORT GenAI MultiModal pipeline +architecture, including the VLM prompt flow, generation loop, input routing, +and processor outputs. For an overview, see the parent [SKILL.md](../SKILL.md). + +--- + +## VLM prompt flow (3-model split) + +``` +pixel_values + image_grid_thw → [vision.onnx] → image_features + │ +input_ids + image_features → [embedding.onnx] → inputs_embeds + │ +inputs_embeds + position_ids → [model.onnx] → logits + + past_kv +``` + +--- + +## VLM generation flow + +``` +Prompt stage: + 1. VisionState.Run() → image_features + 2. EmbeddingState.ReuseFeaturesBuffer(image_features) + 3. EmbeddingState.Run() → inputs_embeds + 4. DecoderState.Run() → logits + present_kv + 5. VisionState destroyed (no longer needed) + +Token generation stage (loop): + 1. EmbeddingState.Run() → inputs_embeds (from single token) + 2. DecoderState.Run() → logits + present_kv +``` + +--- + +## Input flow + +When `generator.set_inputs(named_tensors)` is called: + +1. Tensors matching vision model input names → fed to VisionState +2. Tensors matching embedding model input names → fed to EmbeddingState +3. `input_ids` → used for token counting and embedding lookup +4. `num_image_tokens` → used to allocate image_features buffer size + +--- + +## QwenImageProcessor outputs + +| Tensor | Shape | Description | +|---|---|---| +| `input_ids` | (1, seq_len) | Tokenized prompt with image_pad tokens | +| `pixel_values` | (total_patches, C×T×P×P) | Flattened image patches | +| `image_grid_thw` | (num_images, 3) | Grid dimensions per image | +| `num_image_tokens` | (1,) | Total merged image tokens | + +> **Important:** The ORT GenAI QwenImageProcessor does NOT produce +> `cu_seqlens`, `cu_window_seqlens`, or `rotary_pos_ids`. If the vision +> ONNX model requires these, they must be computed externally and injected +> into the NamedTensors. + +--- + +## Multimodal processor factory + +When `model.create_multimodal_processor()` is called: + +| model.type | Processor class | +|---|---| +| `phi3v` | PhiImageProcessor | +| `whisper` | WhisperProcessor | +| `phi4mm` | PhiMultiModalProcessor | +| `gemma3` | GemmaImageProcessor | +| `fara` | QwenImageProcessor | +| `qwen2_5_vl` | QwenImageProcessor | + +> Models not in this table cannot use `create_multimodal_processor()`. + +--- + +## Writing genai_config.json from ArchitectureConfig + +```python +def _write_genai_config(config, output_dir, model_type="qwen2_5_vl"): + genai_config = { + "model": { + "bos_token_id": config.bos_token_id or 151643, + "context_length": 4096, + "decoder": { + "session_options": { + "log_id": "onnxruntime-genai", + "provider_options": [], + }, + "filename": "model.onnx", + "head_size": config.head_dim, + "hidden_size": config.hidden_size, + "inputs": { + "inputs_embeds": "inputs_embeds", + "attention_mask": "attention_mask", + "position_ids": "position_ids", + "past_key_names": "past_key_values.%d.key", + "past_value_names": "past_key_values.%d.value", + }, + "outputs": { + "logits": "logits", + "present_key_names": "present.%d.key", + "present_value_names": "present.%d.value", + }, + "num_attention_heads": config.num_attention_heads, + "num_hidden_layers": config.num_hidden_layers, + "num_key_value_heads": config.num_key_value_heads, + }, + "embedding": { + "filename": "embedding.onnx", + "inputs": { + "input_ids": "input_ids", + "image_features": "image_features", + }, + "outputs": { + "inputs_embeds": "inputs_embeds", + }, + }, + "vision": { + "filename": "vision.onnx", + "spatial_merge_size": 2, + "inputs": { + "pixel_values": "pixel_values", + "image_grid_thw": "image_grid_thw", + }, + "outputs": { + "image_features": "image_features", + }, + }, + "eos_token_id": config.eos_token_id or [151645, 151643], + "pad_token_id": config.pad_token_id or 151643, + "image_token_id": 151655, + "vision_start_token_id": 151652, + "type": model_type, + "vocab_size": config.vocab_size, + }, + "search": { + "do_sample": False, + "early_stopping": True, + "max_length": 4096, + "num_beams": 1, + "num_return_sequences": 1, + "past_present_share_buffer": False, + "repetition_penalty": 1.0, + "temperature": 1.0, + "top_k": 1, + "top_p": 1.0, + }, + } + with open(os.path.join(output_dir, "genai_config.json"), "w") as f: + json.dump(genai_config, f, indent=4) +``` diff --git a/.agents/skills/ort-genai-config/references/processor-config-fields.md b/.agents/skills/ort-genai-config/references/processor-config-fields.md new file mode 100644 index 00000000..616d1dd4 --- /dev/null +++ b/.agents/skills/ort-genai-config/references/processor-config-fields.md @@ -0,0 +1,176 @@ +# processor_config.json — Complete Field Reference + +This document is the exhaustive reference for the `processor_config.json` file +used by ort-extensions image processors. For an overview, see the parent +[SKILL.md](../SKILL.md). + +--- + +## Format: ort-extensions vs HuggingFace + +> **Critical:** ORT GenAI expects the ort-extensions format — NOT the +> HuggingFace `processor_config.json` format. The HF format wraps data under +> `"image_processor"` with different keys; ORT extensions expects a `"processor"` +> key with an ordered transform pipeline. + +--- + +## Qwen2.5-VL full example + +```json +{ + "processor": { + "name": "qwen2_5_image_processor", + "transforms": [ + { + "operation": { + "name": "decode_image", + "type": "DecodeImage", + "attrs": { "color_space": "RGB" } + } + }, + { + "operation": { + "name": "convert_to_rgb", + "type": "ConvertRGB" + } + }, + { + "operation": { + "name": "resize", + "type": "Resize", + "attrs": { + "width": 540, + "height": 360, + "smart_resize": 1, + "min_pixels": 3136, + "max_pixels": 12845056, + "patch_size": 14, + "merge_size": 2 + } + } + }, + { + "operation": { + "name": "rescale", + "type": "Rescale", + "attrs": { "rescale_factor": 0.00392156862745098 } + } + }, + { + "operation": { + "name": "normalize", + "type": "Normalize", + "attrs": { + "mean": [0.48145466, 0.4578275, 0.40821073], + "std": [0.26862954, 0.26130258, 0.27577711], + "qwen2_5_vl": 1 + } + } + }, + { + "operation": { + "name": "patch_image", + "type": "PatchImage", + "attrs": { + "patch_size": 14, + "temporal_patch_size": 2, + "merge_size": 2 + } + } + } + ] + } +} +``` + +--- + +## Transform types + +| Type | Purpose | Key attrs | +|---|---|---| +| `DecodeImage` | Decode from bytes | `color_space` | +| `ConvertRGB` | Ensure RGB | — | +| `Resize` | Smart resize | `width`, `height`, `smart_resize`, `min_pixels`, `max_pixels`, `patch_size`, `merge_size` | +| `Rescale` | Scale pixel values | `rescale_factor` | +| `Normalize` | Mean/std normalization | `mean`, `std` | +| `PatchImage` | Extract patches | `patch_size`, `temporal_patch_size`, `merge_size` | + +--- + +## Generating from HuggingFace config + +```python +from transformers import AutoProcessor + +processor = AutoProcessor.from_pretrained(model_id) +ip = processor.image_processor + +processor_config = { + "processor": { + "name": "qwen2_5_image_processor", + "transforms": [ + {"operation": {"name": "decode_image", "type": "DecodeImage", + "attrs": {"color_space": "RGB"}}}, + {"operation": {"name": "convert_to_rgb", "type": "ConvertRGB"}}, + {"operation": {"name": "resize", "type": "Resize", + "attrs": { + "width": 540, "height": 360, "smart_resize": 1, + "min_pixels": ip.size.get("shortest_edge", 3136), + "max_pixels": ip.size.get("longest_edge", 12845056), + "patch_size": ip.patch_size, "merge_size": ip.merge_size, + }}}, + {"operation": {"name": "rescale", "type": "Rescale", + "attrs": {"rescale_factor": ip.rescale_factor}}}, + {"operation": {"name": "normalize", "type": "Normalize", "attrs": { + "mean": list(ip.image_mean), "std": list(ip.image_std), + "qwen2_5_vl": 1, + }}}, + {"operation": {"name": "patch_image", "type": "PatchImage", "attrs": { + "patch_size": ip.patch_size, + "temporal_patch_size": ip.temporal_patch_size, + "merge_size": ip.merge_size, + }}}, + ], + } +} +``` + +--- + +## Writing processor_config.json from HuggingFace + +```python +def _write_processor_config(processor, output_dir): + ip = processor.image_processor + config = { + "processor": { + "name": "qwen2_5_image_processor", + "transforms": [ + {"operation": {"name": "decode_image", "type": "DecodeImage", + "attrs": {"color_space": "RGB"}}}, + {"operation": {"name": "convert_to_rgb", "type": "ConvertRGB"}}, + {"operation": {"name": "resize", "type": "Resize", "attrs": { + "width": 540, "height": 360, "smart_resize": 1, + "min_pixels": ip.size.get("shortest_edge", 3136), + "max_pixels": ip.size.get("longest_edge", 12845056), + "patch_size": ip.patch_size, "merge_size": ip.merge_size, + }}}, + {"operation": {"name": "rescale", "type": "Rescale", + "attrs": {"rescale_factor": ip.rescale_factor}}}, + {"operation": {"name": "normalize", "type": "Normalize", "attrs": { + "mean": list(ip.image_mean), "std": list(ip.image_std), + "qwen2_5_vl": 1, + }}}, + {"operation": {"name": "patch_image", "type": "PatchImage", "attrs": { + "patch_size": ip.patch_size, + "temporal_patch_size": ip.temporal_patch_size, + "merge_size": ip.merge_size, + }}}, + ], + } + } + with open(os.path.join(output_dir, "processor_config.json"), "w") as f: + json.dump(config, f, indent=2) +``` From 5d761f712f7a44ad4eca98b2722f1ba4fee54326 Mon Sep 17 00:00:00 2001 From: Justin Chu Date: Tue, 21 Apr 2026 15:54:41 +0000 Subject: [PATCH 2/8] Restructure writing-tests skill for progressive disclosure Reduce SKILL.md from 699 to 273 lines by extracting detailed reference material into references/ subdirectory: - references/test-examples.md: Full code examples, YAML format, golden file format, step-by-step coverage instructions - references/tolerance-guidelines.md: Detailed tolerance tables, failure checklist, debugging scripts, dtype-specific guidance - references/test-utilities.md: API reference for OnnxModelSession, OnnxGenerator, comparison functions, feed creation patterns SKILL.md retains: L1-L5 summary table, test commands, file layout, shared config overview, quick patterns, tolerance quick-reference, and gotchas section with clear directives to reference files. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Justin Chu --- .agents/skills/writing-tests/SKILL.md | 273 ++++++++++++++++ .../writing-tests/references/test-examples.md | 304 ++++++++++++++++++ .../references/test-utilities.md | 86 +++++ .../references/tolerance-guidelines.md | 108 +++++++ 4 files changed, 771 insertions(+) create mode 100644 .agents/skills/writing-tests/SKILL.md create mode 100644 .agents/skills/writing-tests/references/test-examples.md create mode 100644 .agents/skills/writing-tests/references/test-utilities.md create mode 100644 .agents/skills/writing-tests/references/tolerance-guidelines.md diff --git a/.agents/skills/writing-tests/SKILL.md b/.agents/skills/writing-tests/SKILL.md new file mode 100644 index 00000000..742bbe2c --- /dev/null +++ b/.agents/skills/writing-tests/SKILL.md @@ -0,0 +1,273 @@ +--- +name: writing-tests +description: > + Use this skill when writing or modifying tests for mobius models and + components. Covers the L1–L5 confidence system: unit tests (graph + construction from tiny configs), integration tests (numerical parity + with HuggingFace), golden tests (pre-computed reference comparison), + and generation tests (multi-token output verification). Includes test + commands, shared config infrastructure, testing utilities, tolerance + guidelines, and common pitfalls. +--- + +# Skill: Writing Tests + +## When to use + +Use this skill whenever you need to: +- Add tests for a new model or component +- Write integration tests comparing ONNX output to HuggingFace +- Debug numerical parity failures +- Add golden test coverage (L4/L5) +- Understand the test infrastructure and conventions + +## References + +Detailed material is extracted into reference files: + +- **Read [`references/test-examples.md`](references/test-examples.md)** when + you need full code examples for any test type, YAML test case format, + golden file format, or step-by-step coverage instructions. +- **Read [`references/tolerance-guidelines.md`](references/tolerance-guidelines.md)** + when debugging numerical mismatches, choosing tolerance values, or + investigating dtype-specific bugs. +- **Read [`references/test-utilities.md`](references/test-utilities.md)** when + using `OnnxModelSession`, `OnnxGenerator`, comparison functions, or + dealing with test feed creation for symbolic dimensions. + +--- + +## Confidence levels (L1–L5) + +Each level is detected and counted **independently**. A model can pass L3 +without passing L2, or have L4 golden data without passing L3. + +| Level | Name | What it verifies | Data source | +|-------|------|-----------------|-------------| +| **L1** | Graph builds | ONNX graph builds from a tiny synthetic config | `tests/_test_configs.py` + `tests/build_graph_test.py` | +| **L2** | Config compatible | Full-size HF config produces a valid graph | `test_model_id` in YAML test case (`testdata/cases/`) | +| **L3** | Synthetic parity | Random-weight forward pass matches HF numerically | `tests/integration_test.py` parametrized tests | +| **L4** | Golden match | Real-weight prefill logits match pre-computed reference | `testdata/golden//.json` | +| **L5** | Generation verified | Full multi-token generation matches golden output | `testdata/golden//_generation.json` | + +The dashboard shows **per-flag counts** — a model is counted at every level +it passes, not just the highest. L1 equals the total number of registered +models. + +--- + +## Test commands + +```bash +# All non-integration tests (fast, no downloads) +python -m pytest tests/build_graph_test.py tests/cli_test.py src/ -q \ + -k "not phi4mm and not apply_weights_unknown" --tb=short + +# Representative models only (~5 seconds) +python -m pytest tests/build_graph_test.py --fast + +# Single model type +python -m pytest tests/build_graph_test.py -k "phi4mm" + +# Integration tests (slow, downloads models) +python -m pytest tests/integration_test.py -m integration -k "qwen2.5-0.5b" + +# L4/L5 golden tests +python -m pytest tests/e2e_golden_test.py -m golden --level L4 -v +python -m pytest tests/e2e_golden_test.py -m golden --level L5 -v + +# Generate golden data +python scripts/generate_golden.py --level L4 --filter 'my-model*' +``` + +--- + +## Test file layout + +``` +tests/ +├── build_graph_test.py # L1: graph construction (no weights) +├── _test_configs.py # shared model configs for all tests +├── integration_test.py # L3: real-weight numerical parity +├── e2e_golden_test.py # L4 + L5: golden file comparison +├── yaml_schema_test.py # YAML test case schema validation +├── weight_alignment_test.py # L1: preprocess_weights correctness +└── arch_validation_test.py # L2: full HF config graph build + +testdata/ +├── cases/ # YAML test case definitions (L2, L4, L5) +└── golden/ # Pre-computed reference outputs +``` + +--- + +## Shared test configuration + +All model configs live in `tests/_test_configs.py`, organized by category: + +| List | Test class | Task type | +|------|-----------|-----------| +| `CAUSAL_LM_CONFIGS` | `TestBuildGraph` | text-generation | +| `ENCODER_CONFIGS` | `TestBuildEncoderGraph` | feature-extraction | +| `SEQ2SEQ_CONFIGS` | `TestBuildSeq2SeqGraph` | seq2seq | +| `VISION_CONFIGS` | `TestBuildVisionGraph` | image-classification | +| `DETECTION_CONFIGS` | `TestBuildDetectionGraph` | object-detection | + +Each entry is a 3-tuple: `(model_type, config_overrides, is_representative)`. + +- **`is_representative=True`**: Models with unique behaviour (custom class, + softcapping, MoE, ALiBi, etc.). Always tested. +- **`is_representative=False`**: Simple aliases. Skipped with `--fast`. +- **Auto-generation**: Text-generation models with no explicit entry get + `(model_type, {}, False)` automatically from the registry. + +To add a new model: +```python +CAUSAL_LM_CONFIGS: list[tuple[str, dict, bool]] = [ + ("my_model", {"hidden_act": "gelu", "attn_qkv_bias": True}, True), +] +``` + +--- + +## L1: Graph build tests + +Located in `tests/build_graph_test.py`. Uses tiny synthetic configs +(64 hidden, 2 layers, 256 vocab) — no weights, no network. + +The framework checks: inputs exist (`input_ids`, `attention_mask`, +`position_ids`), outputs exist (`logits`, KV cache), and initializers +are present. + +VLM/audio models use dedicated test methods tracked in +`_SPECIALIZED_TEST_MODEL_TYPES`. + +### Weight alignment tests + +`tests/weight_alignment_test.py` verifies `preprocess_weights()` maps +HF state dict keys to ONNX initializer names correctly. Catches bugs +like prefix replacement corrupting names or fused weight names being dropped. + +### Rewrite rule unit tests + +Place rewrite rule tests **next to** the source file: +- Source: `src/mobius/rewrite_rules/_packed_attention.py` +- Test: `src/mobius/rewrite_rules/_packed_attention_test.py` + +--- + +## L2: Config compatibility + +Detected from the `test_model_id` field in YAML test cases. To add L2: +create `testdata/cases//my-model.yaml` with `test_model_id` +set to a real HF model ID. + +--- + +## L3: Integration tests + +Located in `tests/integration_test.py`. Parametrized with +`(model_id, trust_remote_code)`. Prefer models ≤ 1B, publicly accessible, +one per distinct model class. + +> Read [`references/test-examples.md`](references/test-examples.md) for +> full prefill/decode/generation code patterns. + +--- + +## L4 + L5: Golden tests + +Compare ONNX outputs against pre-computed golden files in `testdata/golden/`. + +| Level | File pattern | Contents | +|-------|-------------|----------| +| L4 | `.json` | Prefill top-1/top-2 token IDs + logit summary | +| L5 | `_generation.json` | Prompt + generated token IDs + text | + +> Read [`references/test-examples.md`](references/test-examples.md) for +> YAML format, golden file format, and step-by-step coverage instructions. + +--- + +## Testing utilities + +| Utility | Purpose | +|---------|---------| +| `OnnxModelSession(model)` | Save + load + run ONNX model | +| `OnnxGenerator(session, config)` | Multi-step greedy decoding | +| `load_torch_model(id)` | Load HF model + tokenizer | +| `torch_forward(model, ...)` | Single forward pass | +| `torch_generate_greedy(...)` | Multi-token HF generation | +| `assert_logits_close(a, b)` | Logit comparison with diagnostics | +| `assert_generation_match(a, b)` | Token-ID exact match | + +> Read [`references/test-utilities.md`](references/test-utilities.md) for +> detailed API, feed creation patterns, and ONNX function registration. + +--- + +## Tolerances (quick reference) + +| Model type | rtol / atol | +|------------|-------------| +| Standard text, encoder, seq2seq, diffusion, audio | `1e-3` / `1e-3` | +| Multimodal (vision pipeline) | `1e-2` / `1e-2` | +| Generation (token IDs) | Exact match | +| fp16/bf16 logits | `1e-2` / `1e-2` | + +Key rules: +- `assert_logits_close` checks shape + dtype match (`strict=True`) +- If max abs diff > 0.5 → likely a norm or scaling bug +- If max abs diff > 10 → weights loaded to wrong parameters + +> Read [`references/tolerance-guidelines.md`](references/tolerance-guidelines.md) +> for the full failure checklist, debugging scripts, and dtype-specific guidance. + +--- + +## Gotchas and common mistakes + +### L1 tests are necessary but not sufficient + +L1 verifies graph construction, not execution. A Scan body MatMul shape +mismatch can pass all L1 tests but crash at runtime. **Always write an +integration test alongside any new custom function or Scan op.** + +### Integration tests must exercise all code paths + +- **Text-only first** — verify logit parity before adding other modalities +- **Vision with real pixel values** — zeros don't exercise the encoder +- **All dtypes** (f32, f16, bf16) — each can expose different bugs +- **GPU when available** — different kernels on CUDA + +### Recurrent state ≠ KV cache + +Recurrent state batch dim must match the actual input batch size. Do not +copy the KV cache `batch=0` initialization pattern: +```python +# WRONG — collapses Scan output +past_state = np.zeros((0, num_heads, d_k, d_v), dtype=np.float32) +# CORRECT +past_state = np.zeros((batch_size, num_heads, d_k, d_v), dtype=np.float32) +``` + +### fp16 Exp overflow + +`exp(x)` overflows to `inf` for `x > ~11.09` in fp16. Upcast to float32 +for Exp/Softplus. bf16 does NOT need this workaround (same exponent range +as fp32). + +### Compare full logits, not just tokens + +Generated tokens hide logit divergence — two different logit vectors can +agree on top-1. Always use `assert_logits_close` on the full tensor. + +### Use `--compare-hf` in example scripts + +The `--compare-hf` flag is the gold-standard correctness check. Run it +for all supported dtypes as part of every significant model change. + +### Enable automated code review + +Code review catches bugs that unit tests cannot (fp16 overflow risks, +missing input validation). Enable it on every PR modifying model code. diff --git a/.agents/skills/writing-tests/references/test-examples.md b/.agents/skills/writing-tests/references/test-examples.md new file mode 100644 index 00000000..6de2d3b2 --- /dev/null +++ b/.agents/skills/writing-tests/references/test-examples.md @@ -0,0 +1,304 @@ +# Test Examples Reference + +Full code examples for each test type in mobius. See the main +[SKILL.md](../SKILL.md) for the overview and decision framework. + +## L1: Graph build test patterns + +### Adding a new model type + +Add an entry to the appropriate list in `tests/_test_configs.py` with the +model type, config overrides, and `is_representative` flag: + +```python +# In tests/_test_configs.py: +CAUSAL_LM_CONFIGS: list[tuple[str, dict, bool]] = [ + ("llama", {}, True), + ("my_model", {"attn_qkv_bias": True, "hidden_act": "gelu"}, True), + # ... +] +``` + +The test framework automatically creates a tiny config, builds the graph, +and checks: +- Graph has inputs (`input_ids`, `attention_mask`, `position_ids`) +- Graph has outputs (`logits`, `present.{i}.key`, `present.{i}.value`) +- Graph has initializers (embedding, attention, MLP/expert parameters) + +### Model-specific structure tests + +For models with unique structure (e.g. LoRA), add a dedicated test class: + +```python +class TestBuildGraphLoRA: + def test_lora_initializers_present(self): + config = _base_config( + vision_lora={"r": 4, "lora_alpha": 8}, + speech_lora={"r": 8, "lora_alpha": 16}, + ) + model_cls = registry.get("phi4mm") + module = model_cls(config) + task = CausalLMTask() + model = task.build_graph(module, config, opset_version=23) + + init_names = list(model.graph.initializers) + lora_names = [n for n in init_names if "lora" in n] + assert len(lora_names) > 0 +``` + +## L3: Integration test patterns + +### Prefill + decode numerical comparison + +```python +@pytest.mark.integration +@pytest.mark.parametrize("model_id,trust_remote_code", _TEXT_MODELS) +class TestForwardNumerical: + def test_prefill_logits_match(self, model_id, trust_remote_code): + onnx_model = build(model_id, load_weights=True) + torch_model, tokenizer = load_torch_model(model_id) + config = _get_config(model_id, trust_remote_code) + + # Tokenize, run both models, compare + feeds = _make_prefill_feeds(config, input_ids, attention_mask, position_ids) + onnx_outputs = session.run(feeds) + assert_logits_close(onnx_outputs["logits"], torch_logits, rtol=1e-3, atol=1e-3) + + def test_decode_step_logits_match(self, model_id, trust_remote_code): + # Prefill first, then feed next token + KV cache + decode_feeds = _make_decode_feeds(config, ...) + onnx_out_2 = session.run(decode_feeds) + assert_logits_close(onnx_out_2["logits"], torch_logits_2, rtol=1e-3, atol=1e-3) +``` + +### Greedy generation + +```python +@pytest.mark.integration +class TestGreedyGeneration: + def test_generate_tokens_match(self, model_id, trust_remote_code): + session = OnnxModelSession(onnx_model) + generator = OnnxGenerator(session, config) + onnx_ids = generator.generate(input_ids, max_new_tokens=10, eos_token_id=...) + + torch_ids = torch_generate_greedy(torch_model, input_ids, max_new_tokens=10, eos_token_id=...) + assert_generation_match(onnx_ids[0].tolist(), torch_ids[0].tolist()) +``` + +### Adding a new model to integration tests + +Add a `pytest.param` to `_TEXT_MODELS`: + +```python +_TEXT_MODELS = [ + pytest.param("Qwen/Qwen2.5-0.5B", False, id="qwen2.5-0.5b"), + pytest.param("my-org/my-small-model", False, id="my-model"), + # ... +] +``` + +Guidelines for choosing models: +- Prefer models ≤ 1B parameters for CI speed +- Models must be publicly accessible (no gated/private repos) +- One representative model per distinct model class + +## L4 + L5: Golden test patterns + +### YAML test case format + +**Location:** `testdata/cases//.yaml` + +Categories match task types: `causal-lm`, `encoder`, `seq2seq`, `audio`, +`vision`, `vision-language`, `diffusion`. + +**Required fields:** + +```yaml +model_id: "Qwen/Qwen2.5-1.5B-Instruct" # HuggingFace model ID +revision: "main" # Git revision / commit SHA +task_type: "text-generation" # Task type string +dtype: "float32" # "float32", "float16", or "bfloat16" +level: "L4+L5" # "L4", "L5", or "L4+L5" + +inputs: + prompts: + - "Here is my poem:" # Text prompt(s); use this default +``` + +For image models, use `images:` instead of (or alongside) `prompts:`: + +```yaml +inputs: + images: + - "pipeline-cat-chonk.jpeg" # Path relative to testdata/ +``` + +For audio models: + +```yaml +inputs: + audio: + - "652-129742-0006.flac" +``` + +**Optional fields:** + +```yaml +# Identifier for the test model used in L2 config compatibility check. +# If set, the dashboard counts this model as L2 (full HF config valid). +test_model_id: "Qwen/Qwen2.5-1.5B-Instruct" + +# Skip this test case entirely (model too large, gated repo, etc.). +# Dashboard shows the model as 'skipped' rather than counting it toward coverage. +skip_reason: "Model too large (47B MoE) for CPU golden generation." + +# Pass trust_remote_code=True when loading HuggingFace model (default: false). +trust_remote_code: true + +# Minimum fraction of generated tokens that must match the golden reference. +# Use for VL/audio pipelines where floating-point variance causes later tokens +# to diverge. A value of 0.25 means at least 25% of tokens must match exactly. +# Green (≥0.9) / Yellow (0.5–0.9) / Red (<0.5) on dashboard. +min_token_match_ratio: 0.25 + +# Human-readable notes about this model. +notes: "GPT-2 124M. Absolute positional embeddings, no RoPE." + +generation: + max_new_tokens: 20 # Override token generation limit + do_sample: false +``` + +**`skip_reason` vs `_SKIP_REASONS` dict:** Always use the YAML `skip_reason` +field for new cases. The legacy `_SKIP_REASONS` dict in `e2e_golden_test.py` +has been removed — YAML is the canonical location. + +### Golden file format + +**L4 golden file** (`testdata/golden//.json`): +Generated automatically by `generate_golden.py`. Contains `top1_id`, +`top2_id`, `top10_ids`, `top10_logits`, and `logits_summary` from the last +token position of the prefill pass. + +**L5 generation file** (`testdata/golden//_generation.json`): +Contains `model_id`, `prompt`, `generated_tokens` (list of token IDs), and +`generated_text`. This is the authoritative source for L5 tests — the main +golden JSON does **not** contain generation data. + +### Generating golden data + +```bash +# Generate for all test cases at a given level +python scripts/generate_golden.py --level L4 + +# Generate for a specific task type +python scripts/generate_golden.py --level L4 --task-type causal-lm + +# Generate for a specific model (glob filter on model name) +python scripts/generate_golden.py --level L4 --filter 'llama*' +``` + +Golden files must be committed alongside new test case YAML files. + +### Step-by-step: adding coverage for a new model + +**L1 — Graph builds:** +1. Add `("my_model", {config_overrides}, True)` to the appropriate list in + `tests/_test_configs.py` (or add a dedicated method if the model is a VLM/audio). +2. Run `python -m pytest tests/build_graph_test.py -k "my_model"`. + +**L2 — Config compatible:** +1. Create `testdata/cases//my-model.yaml`. +2. Set `test_model_id: "org/my-model-id"`. +3. Run schema validation: `python -m pytest tests/yaml_schema_test.py`. + +**L3 — Synthetic parity:** +1. Add `pytest.param("org/my-model", False, id="my-model")` to the + appropriate parametrized list in `tests/integration_test.py`. +2. Run `python -m pytest tests/integration_test.py -m integration -k "my-model"`. + +**L4 — Golden match:** +1. Create/update `testdata/cases//my-model.yaml` with `level: "L4"`. +2. Set `inputs.prompts: ["Here is my poem:"]` (standard default prompt). +3. Run `python scripts/generate_golden.py --level L4 --filter 'my-model*'`. +4. Commit the generated `testdata/golden//my-model.json`. +5. Run `python -m pytest tests/e2e_golden_test.py -m golden --level L4 -k "my-model"`. + +**L5 — Generation verified:** +1. Update YAML to `level: "L5"` or `"L4+L5"`. +2. Add a `generation:` block with `max_new_tokens` and `do_sample: false`. +3. Optionally set `min_token_match_ratio` if you expect partial divergence + (VL pipelines, long generation sequences). +4. Run `python scripts/generate_golden.py --level L5 --filter 'my-model*'`. +5. Commit `testdata/golden//my-model_generation.json`. +6. Run `python -m pytest tests/e2e_golden_test.py -m golden --level L5 -k "my-model"`. + +## Debugging multi-model pipelines (TTS, VLM) + +When a multi-model pipeline produces wrong output but individual +model prefill logits look correct, isolate each model boundary: + +1. **Compare each model's output against HF at the boundary** — e.g. + `last_hidden_state` from the talker, `codec_sum` from embeddings, + `inputs_embeds` constructed for the code predictor. + +2. **Check pre-norm vs post-norm** — `outputs.last_hidden_state` in HF + is typically post-norm. If your ONNX model returns pre-norm hidden + states, downstream models receive wrong values. + +3. **Verify external construction matches HF** — for models where the + generation loop constructs inputs externally (e.g. concatenating + hidden states with embeddings), write a comparison script that + checks the constructed input matches HF token-by-token: + ```python + # Compare inputs_embeds at each generation step + for step in range(num_steps): + onnx_input = construct_inputs_embeds(step, ...) + hf_input = hf_model.get_inputs_embeds(step, ...) + diff = np.abs(onnx_input - hf_input).max() + print(f"Step {step}: max diff = {diff:.6f}") + ``` + +4. **Embedding weight vs lookup mismatch** — if embedding weights are + identical but lookups differ, the issue is usually which code index + or embedding table is being used (off-by-one errors). + +## Parity testing methodology + +Compare full logit tensors, not just generated tokens. Generated tokens +hide logit divergence (two very-different logit vectors can agree on the +top-1 token): + +```python +# Always compare full logits at every position +assert_logits_close(onnx_logits, hf_logits, atol=1e-3, rtol=1e-3) # fp32 +assert_logits_close(onnx_logits, hf_logits, atol=1e-2, rtol=1e-2) # fp16/bf16 + +# Also check last-position argmax matches (quick sanity check) +assert onnx_logits[0, -1].argmax() == hf_logits[0, -1].argmax() +``` + +If argmax matches but full logit tolerance fails, the model is numerically +correct but some intermediate accumulation differs — this is usually +acceptable for fp16/bf16 and worth a brief comment in the test. + +## Examples as QA tools + +The `--compare-hf` flag in example scripts is the gold-standard correctness +check for a model. Run it as part of every significant change: + +```bash +# Primary correctness check +python examples/qwen35_text_generation.py --compare-hf + +# Test all supported dtypes +python examples/qwen35_text_generation.py --compare-hf --dtype f16 +python examples/qwen35_text_generation.py --compare-hf --dtype bf16 + +# Test on GPU (if available) +python examples/qwen35_text_generation.py --compare-hf --device cuda +``` + +Target: **100% token match** in fp32 greedy generation. fp16/bf16 may +diverge after the first few tokens due to floating-point accumulation, which +is acceptable if logit parity holds at `atol=rtol=1e-2`. diff --git a/.agents/skills/writing-tests/references/test-utilities.md b/.agents/skills/writing-tests/references/test-utilities.md new file mode 100644 index 00000000..9dd18db4 --- /dev/null +++ b/.agents/skills/writing-tests/references/test-utilities.md @@ -0,0 +1,86 @@ +# Test Utilities Reference + +Detailed API reference for mobius testing utilities, fixture patterns, and +test feed creation. See the main [SKILL.md](../SKILL.md) for the overview. + +## Utility API reference + +| Utility | Import path | Purpose | +|---------|-------------|---------| +| `OnnxModelSession(model)` | `_testing.ort_inference` | Save + load + run ONNX model via ONNX Runtime | +| `OnnxGenerator(session, config)` | `_testing.generation` | Multi-step greedy decoding loop | +| `load_torch_model(id)` | `_testing.torch_reference` | Load HuggingFace model + tokenizer | +| `torch_forward(model, ...)` | `_testing.torch_reference` | Single forward pass through HF model | +| `torch_generate_greedy(...)` | `_testing.generation` | Multi-token HF greedy generation | +| `assert_logits_close(a, b)` | `_testing.comparison` | Logit comparison with diagnostics | +| `assert_generation_match(a, b)` | `_testing.comparison` | Token-ID exact match assertion | + +## `OnnxModelSession` + +Wraps the build → save → load → run cycle for integration tests: + +```python +from mobius._testing.ort_inference import OnnxModelSession + +onnx_model = build(model_id, load_weights=True) +session = OnnxModelSession(onnx_model) +outputs = session.run(feed_dict) +``` + +## `OnnxGenerator` + +Implements multi-step greedy decoding over an `OnnxModelSession`: + +```python +from mobius._testing.generation import OnnxGenerator + +generator = OnnxGenerator(session, config) +token_ids = generator.generate(input_ids, max_new_tokens=10, eos_token_id=...) +``` + +## Comparison functions + +### `assert_logits_close(actual, expected, rtol, atol)` + +Uses `np.testing.assert_allclose` with `strict=True` (checks shape + dtype). +On failure, prints diagnostic info including max/mean abs diff. + +### `assert_generation_match(actual_ids, expected_ids)` + +Exact match on token ID lists. Fails with a clear diff showing the first +divergent position. + +## Test feed creation: symbolic dimensions + +ONNX models export symbolic batch/sequence dimensions. When feeding the +model for ORT inference: + +- **Recurrent state batch dim must match input batch dim** — unlike KV + cache (which initialises to zeros and grows), recurrent state tensors + have a fixed `(B, ...)` shape. Feeding batch=0 produces a zero-sized + carry state that collapses the Scan output. +- **Scan carry state is not KV cache** — do not copy the KV cache + zero-initialisation pattern for recurrent state; the batch dimension + must be the actual inference batch size. + +```python +# WRONG — batch=0 zeros out Scan carry +past_state = np.zeros((0, num_heads, d_k, d_v), dtype=np.float32) + +# CORRECT — must match actual batch size +batch_size = input_ids.shape[0] +past_state = np.zeros((batch_size, num_heads, d_k, d_v), dtype=np.float32) +``` + +## ONNX function registration + +When renaming a custom function's `op_type` (e.g. `CausalConvNdWithState` +→ `CausalConvWithState`), the function must be re-registered under the new +name in ORT's function decomposition list. ORT needs the function embedded +in `model.functions` to decompose the custom op before execution. + +**Checklist when renaming a custom function:** +1. Rename the Python factory function and the `ir.Function.name` +2. Update all call sites that reference the old op_type string +3. Update any integration tests that check the op_type name +4. Verify the function appears in `onnx_model.functions` after build diff --git a/.agents/skills/writing-tests/references/tolerance-guidelines.md b/.agents/skills/writing-tests/references/tolerance-guidelines.md new file mode 100644 index 00000000..336508ae --- /dev/null +++ b/.agents/skills/writing-tests/references/tolerance-guidelines.md @@ -0,0 +1,108 @@ +# Tolerance Guidelines Reference + +Detailed tolerance tables, debugging strategies, and per-dtype guidance for +numerical parity testing. See the main [SKILL.md](../SKILL.md) for the +overview. + +## Tolerance table by model type + +| Test type | Recommended rtol/atol | +|-----------|----------------------| +| Standard text models | `1e-3` / `1e-3` | +| Encoder-only (BERT) | `1e-3` / `1e-3` | +| Encoder-decoder (Whisper, BART, T5) | `1e-3` / `1e-3` | +| Multimodal models | `1e-2` / `1e-2` | +| Diffusion models (UNet, DiT, VAE) | `1e-3` / `1e-3` | +| Audio encoder models | `1e-3` / `1e-3` | +| Generation (token IDs) | Exact match | + +Multimodal models use looser tolerances because the vision pipeline +introduces additional floating-point variance. + +## `assert_logits_close` behavior + +`assert_logits_close` uses `strict=True` in `np.testing.assert_allclose`, +which also checks shape and dtype match. + +## Tolerance failure checklist + +If tolerances fail, verify these in order: + +1. **Norm epsilon** — LayerNorm/RMSNorm eps must match HF config exactly + (e.g., Whisper uses `1e-5`, not the default `1e-6`) +2. **Norm type** — Check if the model uses RMSNorm or LayerNorm. OLMo-1B + uses weight-free LayerNorm (not RMSNorm). Using the wrong type causes + max abs diff > 1.0. +3. **Q scaling order** — some models (Whisper) pre-scale Q before attention + and pass `scale=1.0` to the op, which is numerically different from + passing `scale=head_dim**-0.5` +4. **Attention scale** — some models (Granite) replace `1/sqrt(head_dim)` with + a custom `attention_multiplier` from the config +5. **Scaling multipliers** — check HF config for `embedding_multiplier`, + `logits_scaling`, `residual_multiplier` that aren't in standard Llama +6. **Residual pattern** — verify `residual + output * scale` vs + `residual * scale + output` by reading HF source +7. **Weight loading** — compare ONNX initializers against HF state_dict to + rule out name mapping bugs +8. **Float64 contamination** — numpy arrays created from config values default + to float64; always use `dtype=np.float32` + +## Debugging large logit differences + +When max abs diff is large (> 0.5), run this diagnostic: + +```python +import numpy as np +diff = np.abs(onnx_logits[0, -1] - hf_logits) +print(f"Max abs diff: {diff.max():.4f}") +print(f"Mean abs diff: {diff.mean():.4f}") +# If max > 0.5, it's likely a norm or scaling bug, not just floating-point +# If max > 10, weights are probably loaded to wrong parameters +``` + +Check the HF norm class directly: +```python +import inspect +from transformers.models.olmo.modeling_olmo import OlmoLayerNorm +print(inspect.getsource(OlmoLayerNorm)) +``` + +Check for unextracted config fields: +```python +config = AutoConfig.from_pretrained("model-id") +for k, v in config.to_dict().items(): + if any(s in k for s in ("multiplier", "scaling", "factor", "epsilon")): + print(f"{k}: {v}") +``` + +## Dtype-specific tolerance guidance + +### fp32 + +Standard tolerance: `atol=1e-3, rtol=1e-3`. Target **100% token match** in +greedy generation. + +### fp16 + +Looser tolerance: `atol=1e-2, rtol=1e-2`. May diverge after the first few +tokens in generation due to floating-point accumulation. Acceptable if logit +parity holds. + +**fp16 Exp overflow:** `exp(x)` overflows to `inf` for `x > ~11.09` in +fp16. The Softplus activation (`log(1 + exp(x))`) and decay computation +`exp(-softplus(x))` are common overflow sites. Always upcast to float32 +for Exp/Softplus in fp16 models: + +```python +x_f32 = op.Cast(x, to=ir.DataType.FLOAT) +result = op.Exp(x_f32) +result = op.Cast(result, to=x.dtype) # cast back +``` + +### bf16 + +Same tolerance as fp16: `atol=1e-2, rtol=1e-2`. bf16 has the same exponent +range as fp32 (no overflow at 11.09) but much less precision (7-bit mantissa +vs 10-bit). bf16 does NOT need the Exp upcast workaround. + +If a computation works in bf16 but not fp16, check for Exp overflow first. From 8f86411b268872bcaf78ed1b539518babe0caa26 Mon Sep 17 00:00:00 2001 From: Justin Chu Date: Tue, 21 Apr 2026 16:10:37 +0000 Subject: [PATCH 3/8] Restructure 3 skills for progressive disclosure with references/ MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Extract detailed reference material from multimodal-models, reusable-components, and phi4mm-component-parity SKILL.md files into per-skill references/ subdirectories. Each SKILL.md is now under 500 lines (177, 221, 146) with 'Read references/X.md when...' directives for progressive disclosure. multimodal-models: - references/projector-variants.md: Detailed projector code, Qwen-VL specifics - references/vision-encoder-details.md: VisionModel construction, model template - references/weight-mappings.md: Full weight mapping tables, shape fixes reusable-components: - references/onnx-op-patterns.md: Scalar constants, CastLike, fp32 upcast, shapes - references/component-examples.md: All component variants, adapter patterns phi4mm-component-parity: - references/common-failures.md: 9 detailed failure modes with code - references/debugging-cookbook.md: Step-by-step procedures, test patterns All description fields updated with imperative phrasing and user intent focus. No content was lost — all material preserved in reference files. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Justin Chu --- .agents/skills/multimodal-models/SKILL.md | 177 +++++++++++ .../references/projector-variants.md | 192 +++++++++++ .../references/vision-encoder-details.md | 166 ++++++++++ .../references/weight-mappings.md | 104 ++++++ .../skills/phi4mm-component-parity/SKILL.md | 146 +++++++++ .../references/common-failures.md | 244 ++++++++++++++ .../references/debugging-cookbook.md | 187 +++++++++++ .agents/skills/reusable-components/SKILL.md | 221 +++++++++++++ .../references/component-examples.md | 297 ++++++++++++++++++ .../references/onnx-op-patterns.md | 218 +++++++++++++ 10 files changed, 1952 insertions(+) create mode 100644 .agents/skills/multimodal-models/SKILL.md create mode 100644 .agents/skills/multimodal-models/references/projector-variants.md create mode 100644 .agents/skills/multimodal-models/references/vision-encoder-details.md create mode 100644 .agents/skills/multimodal-models/references/weight-mappings.md create mode 100644 .agents/skills/phi4mm-component-parity/SKILL.md create mode 100644 .agents/skills/phi4mm-component-parity/references/common-failures.md create mode 100644 .agents/skills/phi4mm-component-parity/references/debugging-cookbook.md create mode 100644 .agents/skills/reusable-components/SKILL.md create mode 100644 .agents/skills/reusable-components/references/component-examples.md create mode 100644 .agents/skills/reusable-components/references/onnx-op-patterns.md diff --git a/.agents/skills/multimodal-models/SKILL.md b/.agents/skills/multimodal-models/SKILL.md new file mode 100644 index 00000000..26b69f33 --- /dev/null +++ b/.agents/skills/multimodal-models/SKILL.md @@ -0,0 +1,177 @@ +--- +name: multimodal-models +description: > + Add or modify multimodal (vision + language + audio) models in mobius. + Use when wiring a VisionModel, projector, InputMixer, or + VisionLanguageTask; handling image/audio token placeholders; choosing + a projector variant; or splitting a model into 3-or-4 ONNX sub-models + for ORT GenAI deployment. +--- + +# Skill: Multimodal (Vision + Language + Audio) Models + +## When to use + +Use this skill when adding a model that processes both images and text — such +as Gemma3, Gemma4, LLaVA, LLaVA-NeXT, Phi-3-Vision, PaliGemma, InternVL2, +Pixtral, Idefics2/3, Molmo, Florence2, or Video-LLaVA — or a model that +also processes audio (e.g. Gemma4 with speech/audio inputs). + +## Architecture overview + +``` +pixel_values ──► VisionModel ──► MultiModalProjector ──► InputMixer ──┐ + ├──► TextDecoder ──► logits +input_ids ──────► Embedding ──────────────────────────────────────────┘ +``` + +### Key components + +| Component | File | Purpose | +|-----------|------|---------| +| `VisionModel` | `components/_vision.py` | SigLIP-style patch embedding + transformer encoder | +| `PixtralVisionTower` | `components/_pixtral_vision.py` | Pixtral 2D RoPE vision encoder | +| `Gemma3MultiModalProjector` | `components/_multimodal.py` | AvgPool2d → RMSNorm → MatMul | +| `MLPMultiModalProjector` | `components/_multimodal.py` | Linear → GELU → Linear | +| `Mistral3MultiModalProjector` | `components/_pixtral_vision.py` | RMSNorm → spatial merge → Linear → GELU → Linear | +| `LinearMultiModalProjector` | `components/_multimodal.py` | Single Linear | +| `InputMixer` | `components/_multimodal.py` | Scatter vision embeddings at image-token positions | +| `VisionLanguageTask` | `tasks/__init__.py` | ONNX I/O contract with `pixel_values` input | + +## Projector variants + +Choose the projector that matches the HuggingFace implementation: + +| Projector | Architecture | Models | +|-----------|-------------|--------| +| `Gemma3MultiModalProjector` | AvgPool2d → RMSNorm → MatMul | Gemma3 | +| `MLPMultiModalProjector` | Linear → GELU → Linear | LLaVA, LLaVA-NeXT, VipLLaVA, Phi-4-MM, InternVL2, Molmo | +| `Mistral3MultiModalProjector` | RMSNorm → spatial merge → Linear → GELU → Linear | Mistral-3, Pixtral | +| `LinearMultiModalProjector` | Single Linear | PaliGemma, Qwen2-Audio, Idefics2, Florence2 | + +> Read `references/projector-variants.md` when you need detailed constructor +> arguments, Qwen-VL vision encoder specifics (Conv3d, 2D rotary, windowed +> attention, spatial merge), or Qwen3.5-VL architecture details. + +## InputMixer pattern + +`InputMixer` replaces placeholder tokens in the text embedding sequence +with projected vision (or audio) features using `GatherElements` + `Where`: + +1. Find special token positions in `input_ids` +2. Scatter vision embeddings at those positions +3. Return fused `hidden_states` for the decoder + +**Always invoke child modules through `__call__`** so `onnxscript.nn.Module` +pushes the correct naming context. Direct access like +`self.language_model.model.embed_tokens(op, x)` skips intermediate naming +scopes and produces wrong initializer names. + +## VisionLanguageTask I/O contract + +For ORT GenAI deployment, multimodal models split into 3 (or 4) ONNX models: + +| Model | Inputs | Outputs | +|-------|--------|---------| +| Vision | `pixel_values: float32`, `grid_thw: int64` | `image_features: float32` | +| Embedding | `input_ids: int64`, `image_features: float32` | `inputs_embeds: float32` | +| Decoder | `inputs_embeds: float32`, `attention_mask: int64`, `position_ids: int64`, `past_key_values.*` | `logits: float32`, `present.*` | +| Speech (optional) | `input_features: float32` | `audio_features: float32` | + +The embedding model must handle `num_image_tokens=0` (text-only input) by +zero-padding `image_features` before Gather so indices stay in-bounds. + +### Conditional 3-or-4-model task + +Some models come in two tiers (e.g. Gemma4): small variants include a +speech encoder (4 models), large variants are vision-only (3 models). A +single task class checks `config.audio is not None` to decide whether to +include the speech encoder. Reference: `Gemma4Task` in +`src/mobius/tasks/_gemma4.py`. + +## Image and audio token handling + +### Image tokens + +Insert `mm_tokens_per_image` placeholder tokens (not 1!) to match the +number of vision features the projector produces: + +```python +mm_tokens = config.mm_tokens_per_image or 1 +img_tokens = np.full((1, mm_tokens), image_token_id, dtype=np.int64) +input_ids = np.concatenate([input_ids[:, :1], img_tokens, input_ids[:, 1:]], axis=1) +``` + +### Audio boundary markers (Gemma4) + +HuggingFace wraps audio tokens with boundary markers: +``` +<|audio> (256000) + N × <|audio|> (258881) + (258883) +``` +**Missing boundary markers** cause garbled transcription even when the +audio encoder output is numerically correct. + +### Testing tolerances + +Use `rtol=1e-2, atol=1e-2` for multimodal tests. After prefill with image, +the decode step still needs `pixel_values` as input (use zeros). + +## ClippableLinear (critical for Gemma4) + +Gemma4's vision and audio encoders use `ClippableLinear` — a `Linear` with +learned finite input/output activation clamping. Using plain `Linear` +causes max diff 52.68 (audio) / 3.92 (vision) → 0.0003 / 0.00007 after +fix. See `reusable-components` skill for full API reference. + +## Weight name mapping overview + +Multimodal HF models often prefix text weights differently. Implement +`preprocess_weights()` to strip prefixes and rename keys: + +| HF key | Our key | +|--------|---------| +| `language_model.model.layers.0.…` | `layers.0.…` | +| `vision_tower.vision_model.encoder.…` | `vision_tower.encoder.…` | + +> Read `references/weight-mappings.md` when you need full weight mapping +> tables, shape mismatch fixes, ClippableLinear weight conventions, or +> per-layer embedding splitting details. + +> Read `references/vision-encoder-details.md` when you need step-by-step +> instructions for adding a new multimodal model, vision config extraction +> code, or the full model class template with InputMixer wiring. + +## genai_config.json required fields for VLMs + +**Required fields** (without these, VLM output is wrong): +- `image_token_id`: Token ID for `<|image_pad|>` — needed for 3D M-RoPE +- `vision_start_token_id`: Token ID for `<|vision_start|>` — marks boundaries +- `spatial_merge_size`: Grid merge factor (2 for Qwen2.5-VL) + +See `.agents/skills/ort-genai-config/SKILL.md` for the complete reference +and `.agents/skills/debugging-vl-pipeline/SKILL.md` for troubleshooting. + +### processor_config.json for image preprocessing + +ORT GenAI uses ort-extensions for image preprocessing (not HuggingFace). +The `processor_config.json` must use `qwen2_5_image_processor` format with +DecodeImage → ConvertRGB → Resize → Rescale → Normalize → PatchImage +transforms. The `width`/`height` in Resize are direct target dimensions; +compute them as +`round(original_dim / (patch_size * merge_size)) * (patch_size * merge_size)`. + +## Gemma4: per-layer embeddings (CUDA ORT workaround) + +Gemma4 uses `embed_tokens_per_layer` with shape `[V, L*D]`. For large +models this overflows ORT's CUDA Gather kernel. **Workaround:** Split into +L separate `Embedding([V, D])` tables via `nn.ModuleList`, and use `Slice` +instead of `Gather` for per-layer projection indexing. + +## Cross-references + +- **VL debugging:** `.agents/skills/debugging-vl-pipeline/SKILL.md` +- **ORT GenAI config:** `.agents/skills/ort-genai-config/SKILL.md` +- **Weight name alignment:** `.agents/skills/weight-name-alignment/SKILL.md` +- **Multi-image Scan pattern:** `.agents/skills/scan-and-multi-image/SKILL.md` +- **Component parity debugging:** `.agents/skills/phi4mm-component-parity/SKILL.md` +- **Reusable components (ClippableLinear):** `.agents/skills/reusable-components/SKILL.md` diff --git a/.agents/skills/multimodal-models/references/projector-variants.md b/.agents/skills/multimodal-models/references/projector-variants.md new file mode 100644 index 00000000..1adaf648 --- /dev/null +++ b/.agents/skills/multimodal-models/references/projector-variants.md @@ -0,0 +1,192 @@ +# Projector Variants — Detailed Reference + +## Gemma3MultiModalProjector + +```python +Gemma3MultiModalProjector( + vision_hidden_size=1152, # SigLIP hidden dim + text_hidden_size=2560, # Text model hidden dim + patches_per_image=64, # sqrt(num_patches) per side + tokens_per_image=256, # mm_tokens_per_image from config + norm=Gemma3RMSNorm(1152), # Gemma3-specific RMSNorm with +1 offset +) +``` + +The pooling kernel is computed as `patches_per_image / sqrt(tokens_per_image)`. +For Gemma3-4B: `64 / 16 = 4`, so `AvgPool2d(kernel_size=4, stride=4)`. + +## MLPMultiModalProjector + +```python +MLPMultiModalProjector( + vision_hidden_size=1024, + text_hidden_size=4096, + bias=True, +) +``` + +Two-layer MLP with GELU activation. The most common projector pattern. + +## LinearMultiModalProjector + +```python +LinearMultiModalProjector( + vision_hidden_size=1024, + text_hidden_size=4096, + bias=True, +) +``` + +Simple single linear layer. + +## Mistral3MultiModalProjector + +RMSNorm → spatial merge → Linear → GELU → Linear. Used by Mistral-3 and +Pixtral. Defined in `components/_pixtral_vision.py`. + +## Qwen2.5-VL / Qwen3-VL vision encoder specifics + +These models use a **custom vision encoder** (not SigLIP) with unique +architectural features. The encoder is in +`components/_qwen25_vl_vision.py` and `components/_qwen3_vl_vision.py`. + +### Architecture differences from standard VisionModel + +| Feature | Standard (SigLIP) | Qwen2.5-VL / Qwen3-VL | +|---------|-------------------|----------------------| +| Patch embedding | Conv2d | **Conv3d** (temporal + spatial) | +| Position encoding | Learnable embedding | **2D rotary** (height, width) | +| Attention | Standard self-attention | **Windowed + full attention** alternating | +| Normalization | LayerNorm | **RMSNorm** | +| Output merging | CLS token or mean pool | **Spatial merge** (2×2 → 1) | +| MLP | fc1/fc2 | **Gated MLP** (gate_proj/up_proj/down_proj + SiLU) | + +### Critical: 2D rotary embedding dimension + +The vision encoder computes separate rotary frequencies for height and +width positions. The rotary embedding dimension must be `head_dim // 2` +(not `head_dim`): + +```python +# CORRECT: each spatial dimension gets head_dim//4 frequencies +self.rotary_pos_emb = Qwen25VLVisionRotaryEmbedding(head_dim // 2) + +# WRONG: produces 2× too many frequencies with wrong values +self.rotary_pos_emb = Qwen25VLVisionRotaryEmbedding(head_dim) +``` + +The frequency table has shape `(num_patches, head_dim//2)`. Each half +(`head_dim//4` values) covers one spatial dimension. The `forward` method +concatenates `cos(h_freqs)` and `cos(w_freqs)` to produce the final +`(num_patches, head_dim)` rotary embeddings. + +### Critical: fullatt_block_indexes config + +Qwen2.5-VL uses a **hybrid attention pattern**: most blocks use windowed +attention (local windows for efficiency), but certain blocks use full +attention (all patches attend to all patches): + +```python +# Must be extracted from HF vision_config +fullatt_block_indexes = [7, 15, 23, 31] # For 32-block encoder +window_size = 112 # Window size in patches for windowed blocks +``` + +If `fullatt_block_indexes` is missing, ALL blocks use windowed attention, +causing massive feature divergence (cos ≈ 0.25). The first few blocks may +appear correct since they happen to be windowed blocks. + +**Config extraction** — these must be in `_configs.py` VisionConfig: + +```python +@dataclasses.dataclass +class VisionConfig: + ... + fullatt_block_indexes: list[int] | None = None + window_size: int | None = None +``` + +### Window index and attention bias + +- **Windowed blocks**: Patches are grouped into windows of `window_size`. + Each window attends only within itself. The attention bias is block-diagonal. +- **Full attention blocks**: Use `cu_seqlens` (not `cu_window_seqlens`) to + attend across all patches in each image. +- `window_index` permutes patches into window-ordered layout before the + transformer blocks, then `reverse_indices = argsort(window_index)` restores + the original order after. + +### Multi-image support + +Both vision encoders support multiple images via the ONNX `Scan` op. +Per-image values (position IDs, window indices, cu_seqlens) are computed +in a Scan body and concatenated. See `.agents/skills/scan-and-multi-image/SKILL.md`. + +### Spatial merge (post-encoder) + +After the transformer blocks, a spatial merge layer combines 2×2 patches +into 1 token: +``` +(num_patches, hidden_size) → reshape to (num_merged, 4*hidden_size) → MLP → (num_merged, text_hidden_size) +``` +The merge reduces token count by 4× and projects to text model dimension. + +## Qwen3.5-VL + +Qwen3.5-VL uses the same **3-model split** as Qwen3-VL (decoder + vision + +embedding), but swaps the text decoder for the **Qwen3.5 architecture** +which uses hybrid DeltaNet + full attention instead of standard GQA. + +### Architecture + +The vision encoder is **identical to Qwen3-VL** — it reuses +`Qwen3VLVisionModel` (patch_size=16, hidden=1152, depth=27). Only the +text decoder changes. + +| Component | Class | Notes | +|-----------|-------|-------| +| 3-model composite | `Qwen35VL3ModelCausalLMModel` | Splits into decoder + vision + embedding | +| Decoder (standalone) | `Qwen35VLDecoderModel` | Uses `Qwen35TextModel` internally | +| Text model | `Qwen35VLTextModel` | Text-only decoder; strips VL weight prefixes | + +### Registration + +| `model_type` | Variant | Description | +|--------------|---------|-------------| +| `qwen3_5_vl` | 3-model split | Full VLM with vision encoder | +| `qwen3_5_vl_text` | Text-only | Decoder without vision | + +### Task + +Reuses `Qwen3VLVisionLanguage3ModelTask` (task name: `qwen35-vl`). + +### Config + +The HF config is VL-style with a nested `text_config`: + +``` +config.json → model_type: "qwen3_5_vl" +config.text_config → model_type: "qwen3_5" (or "qwen3_5_text") +``` + +### Token IDs + +| Token | ID | +|-------|----| +| `image` | 248056 | +| `video` | 248057 | +| `vision_start` | 248053 | +| `vision_end` | 248054 | + +### Interleaved MRoPE + +Uses `InterleavedMRope` (not `ChunkedMRope`) with: + +- `partial_rotary_factor=0.25` +- `mrope_section=[11, 11, 10]` + +### Key insight + +The vision pipeline is completely shared with Qwen3-VL — only the text +decoder differs (hybrid DeltaNet + full attention). This means vision +encoder bugs/fixes apply to both models equally. diff --git a/.agents/skills/multimodal-models/references/vision-encoder-details.md b/.agents/skills/multimodal-models/references/vision-encoder-details.md new file mode 100644 index 00000000..08388482 --- /dev/null +++ b/.agents/skills/multimodal-models/references/vision-encoder-details.md @@ -0,0 +1,166 @@ +# Vision Encoder Details + +## VisionModel / VisionEncoder construction + +The standard vision encoder (`components/_vision.py`) follows a SigLIP-style +architecture: + +| Component | File | Purpose | +|-----------|------|---------| +| `VisionModel` | `components/_vision.py` | SigLIP-style patch embedding + transformer encoder | +| `PixtralVisionTower` | `components/_pixtral_vision.py` | Pixtral 2D RoPE vision encoder (bidirectional attention) | +| `PatchEmbedding` | `components/_vision.py` | Conv2d → positional embedding | + +### PatchEmbedding naming + +`PatchEmbedding` has three parameters. These use explicit `name=` because the +attribute names don't match the desired ONNX names (e.g. `patch_embedding` +needs to map to `patch_embedding.weight`): + +```python +self.patch_embedding = nn.Parameter([...], name="patch_embedding.weight") +self.patch_embedding_bias = nn.Parameter([...], name="patch_embedding.bias") +self.position_embedding = nn.Parameter([...], name="position_embedding.weight") +``` + +In most cases, `name=` is **not needed** because `nn.Module.__setattr__` +automatically sets the parameter name from the attribute name. Only use +`name=` when the attribute name differs from the desired ONNX initializer name. + +## Step-by-step: adding a new multimodal model + +### 1. Identify the projector architecture + +Look at the HuggingFace source in `modeling_.py`: + +```bash +grep -n "class.*Projector\|class.*projector" \ + transformers/models//modeling_.py +``` + +Match it to one of the projector variants, or create a new one. + +### 2. Extract vision config + +Multimodal HF configs have a `vision_config` sub-object. Extract vision +fields in the test or integration code: + +```python +hf_config = transformers.AutoConfig.from_pretrained(model_id) +text_config = hf_config.text_config +vision_config = hf_config.vision_config + +config = ArchitectureConfig.from_transformers(text_config) +# Add vision fields +config.vision_hidden_size = vision_config.hidden_size +config.vision_intermediate_size = vision_config.intermediate_size +config.vision_num_hidden_layers = vision_config.num_hidden_layers +config.vision_num_attention_heads = vision_config.num_attention_heads +config.vision_image_size = vision_config.image_size +config.vision_patch_size = vision_config.patch_size +config.vision_norm_eps = getattr(vision_config, "layer_norm_eps", 1e-6) +config.mm_tokens_per_image = getattr(hf_config, "mm_tokens_per_image", None) +config.image_token_id = getattr(hf_config, "image_token_id", None) +``` + +### 3. Create the model class + +**Important:** Always invoke child modules through `__call__` (not by +accessing their sub-modules directly) so that `onnxscript.nn.Module` pushes +the correct naming context. Direct access like +`self.language_model.model.embed_tokens(op, x)` skips intermediate naming +scopes and produces wrong initializer names. + +The recommended pattern is to pass vision embeddings as a kwarg through the +`__call__` chain, and have the text model perform the mixing internally: + +```python +class _MyTextModelForMultimodal(MyTextModel): + """Text model that mixes vision embeddings into the input.""" + + def __init__(self, config): + super().__init__(config) + self.input_mixer = InputMixer(image_token_id=config.image_token_id or 0) + + def forward(self, op, input_ids, attention_mask, position_ids, + past_key_values=None, vision_embeddings=None): + hidden_states = self.embed_tokens(op, input_ids) + if vision_embeddings is not None: + hidden_states = self.input_mixer( + op, hidden_states, vision_embeddings, input_ids + ) + return super().forward( + op, input_ids, attention_mask, position_ids, + past_key_values=past_key_values, inputs_embeds=hidden_states, + ) + + +class _MyForMultimodalLM(MyCausalLMModel): + """CausalLM that passes vision_embeddings to the text model.""" + + def __init__(self, config): + nn.Module.__init__(self) + self.config = config + self.model = _MyTextModelForMultimodal(config) + self.lm_head = Linear(config.hidden_size, config.vocab_size, bias=False) + + +class MyMultiModalModel(nn.Module): + def __init__(self, config): + super().__init__() + self.config = config + self.vision_tower = VisionModel(config) + self.multi_modal_projector = MLPMultiModalProjector( + vision_hidden_size=config.vision_hidden_size, + text_hidden_size=config.hidden_size, + ) + self.language_model = _MyForMultimodalLM(config) + + def forward(self, op, input_ids, attention_mask, position_ids, pixel_values, + past_key_values=None): + # 1. Encode vision + vision_features = self.vision_tower(op, pixel_values) + vision_embeddings = self.multi_modal_projector(op, vision_features) + + # 2. Pass through __call__ chain — naming is correct automatically + return self.language_model( + op, input_ids, attention_mask, position_ids, + past_key_values=past_key_values, + vision_embeddings=vision_embeddings, + ) +``` + +See `models/gemma3.py` for the full working example. + +### 4. Handle weight name mismatches + +Multimodal HF models often prefix text weights differently: + +| HF key | Our key | +|--------|---------| +| `language_model.model.layers.0.…` | `layers.0.…` | +| `vision_tower.vision_model.encoder.…` | `vision_tower.encoder.…` | +| `multi_modal_projector.mm_input_projection_weight` | `multi_modal_projector.weight` | + +Implement `preprocess_weights` to strip prefixes and rename keys. + +### 5. Handle weight tying + +If `tie_word_embeddings=True`, the HF checkpoint may not include +`lm_head.weight`. Copy it from `embed_tokens.weight`: + +```python +if self.config.tie_word_embeddings: + if "lm_head.weight" not in renamed and "embed_tokens.weight" in renamed: + renamed["lm_head.weight"] = renamed["embed_tokens.weight"] +``` + +### 6. Use VisionLanguageTask + +Build with the `VisionLanguageTask` to add `pixel_values` to graph inputs: + +```python +from mobius.tasks import VisionLanguageTask + +onnx_model = build_from_module(module, config, task=VisionLanguageTask()) +``` diff --git a/.agents/skills/multimodal-models/references/weight-mappings.md b/.agents/skills/multimodal-models/references/weight-mappings.md new file mode 100644 index 00000000..7fbf897b --- /dev/null +++ b/.agents/skills/multimodal-models/references/weight-mappings.md @@ -0,0 +1,104 @@ +# Weight Name Mappings — Full Reference + +## Common multimodal weight prefixes + +Multimodal HF models often prefix text weights differently from the ONNX +model structure. These must be handled in `preprocess_weights()`. + +| HF key | ONNX key | +|--------|---------| +| `language_model.model.layers.0.…` | `layers.0.…` | +| `vision_tower.vision_model.encoder.…` | `vision_tower.encoder.…` | +| `multi_modal_projector.mm_input_projection_weight` | `multi_modal_projector.weight` | + +## Weight tying + +If `tie_word_embeddings=True`, the HF checkpoint may not include +`lm_head.weight`. Copy it from `embed_tokens.weight`: + +```python +if self.config.tie_word_embeddings: + if "lm_head.weight" not in renamed and "embed_tokens.weight" in renamed: + renamed["lm_head.weight"] = renamed["embed_tokens.weight"] +``` + +## ClippableLinear weight mapping + +HuggingFace stores `.linear.weight` for the actual weight (the +`.linear.` segment is stripped by `preprocess_weights`), and +`.input_min`, `.input_max`, `.output_min`, +`.output_max` as direct scalar buffers. + +## Per-layer embedding weight splitting (Gemma4) + +Gemma4 uses per-layer embedding: `embed_tokens_per_layer` with shape +`[V, L*D]` where V=vocab_size, L=num_layers, D=per_layer_dim. In +`preprocess_weights`, split the HF weight column-wise: + +```python +for i in range(num_layers): + renamed[f"embed_tokens_per_layer.{i}.weight"] = value[:, i*D:(i+1)*D] +``` + +## PatchEmbedding parameter names + +`PatchEmbedding` uses explicit `name=` because the attribute names don't +match the desired ONNX names: + +```python +self.patch_embedding = nn.Parameter([...], name="patch_embedding.weight") +self.patch_embedding_bias = nn.Parameter([...], name="patch_embedding.bias") +self.position_embedding = nn.Parameter([...], name="position_embedding.weight") +``` + +## Shape mismatches requiring preprocess_weights transforms + +Some HF weights have different shapes from the ONNX parameter declaration: + +```python +# Squeeze extra batch dimension from position embeddings +# HF: [1, num_patches, hidden_size] (3D) → ONNX: [num_patches, hidden_size] (2D) +if "position_embedding.weight" in key and state_dict[key].dim() == 3: + state_dict[key] = state_dict[key].squeeze(0) +``` + +**General rule:** Check whether the preprocess_weights transform goes +in the correct direction (squeeze vs unsqueeze). A common mistake is +writing the transform backwards. + +## Testing multimodal models + +### Image token count + +Insert `mm_tokens_per_image` image tokens (not 1!) into the input to match +the number of vision features the projector produces: + +```python +mm_tokens = config.mm_tokens_per_image or 1 +img_tokens = np.full((1, mm_tokens), image_token_id, dtype=np.int64) +input_ids = np.concatenate([input_ids[:, :1], img_tokens, input_ids[:, 1:]], axis=1) +``` + +### Dummy pixel values + +Use random pixel values for testing (we only need numerical parity, not +meaningful images): + +```python +rng = np.random.default_rng(42) +pixel_values = rng.standard_normal((1, 3, image_size, image_size)).astype(np.float32) +``` + +### Tolerances + +Use `rtol=1e-2, atol=1e-2` for multimodal tests — the vision pipeline +introduces more floating-point variance than text-only models. + +### Decode step + +After prefill with image, the decode step is text-only but still needs +`pixel_values` as a graph input (use zeros): + +```python +decode_pixel_values = np.zeros_like(pixel_values) +``` diff --git a/.agents/skills/phi4mm-component-parity/SKILL.md b/.agents/skills/phi4mm-component-parity/SKILL.md new file mode 100644 index 00000000..172e25c2 --- /dev/null +++ b/.agents/skills/phi4mm-component-parity/SKILL.md @@ -0,0 +1,146 @@ +--- +name: phi4mm-component-parity +description: > + Debug multimodal ONNX model output that diverges from HuggingFace. + Use when isolating which component (vision encoder, speech encoder, + embedding, or decoder) causes numerical divergence; when integration + tests fail with large differences; or when adding a new multimodal + model and verifying each stage independently. Applicable to Phi4MM, + Gemma4, and any multi-encoder architecture. +--- + +# Skill: Multimodal Component Parity Debugging + +## When to use + +Use this skill when: + +- A multimodal ONNX model's logits diverge systematically from HuggingFace +- You're adding a new multimodal model and need to verify each component +- Integration tests fail with large numerical differences (not just tolerance) +- Weights appear to load but the model produces wrong outputs +- You need to isolate which component (vision, speech, embedding, decoder) + is causing divergence + +For vision-language-only models (no speech), see also the +`debugging-vl-pipeline` skill which covers VLM-specific issues like 3D M-RoPE. + +## Pipeline isolation methodology + +Multimodal models with N encoders have N+2 stages (encoders + embedding + +decoder). Debug by comparing each stage independently against HuggingFace +at every boundary. + +### 4-model multimodal pipeline (e.g., Phi4MM) + +``` +pixel_values ──► [1. Vision Encoder] ──► image_features ──┐ + │ +audio_embeds ──► [2. Speech Encoder] ──► speech_features ──┤ + │ +input_ids ──► [3. Embedding/Fusion] ◄───────────────────┘ + │ + ▼ + inputs_embeds + │ + ▼ + [4. Decoder + LoRA] ──► logits +``` + +**Golden rule:** Start from the simplest case (text-only, no encoders), +verify it matches HF, then add one modality at a time. + +### Stage-by-stage comparison + +**Stage 1 — Vision encoder:** Compare SigLIP/ViT output. Check output +shape `(num_image_tokens, text_hidden_size)`, projection MLP correctness, +and position embeddings (2D vs 3D shape). Target: cos_sim > 0.99. + +**Stage 2 — Speech encoder:** Compare Conformer output. Check compression +rate (typically 8× time reduction), projection branch selection, and conv +subsampling output length. + +**Stage 3 — Embedding/fusion:** Compare token embeddings. Text-only should +match HF `embed_tokens` exactly (< 1e-5). With features, verify token +replacement at correct positions. InputMixer must handle zero-length tensors. + +**Stage 4 — Decoder:** Compare logits. Acceptable float32 metrics: +max_diff 5-10, mean_diff 0.5-1.5, cos_sim > 0.98, argmax match exact. + +## Quick-start 4-step process + +1. **Text-only baseline:** Build ONNX → run embedding + decoder (skip + encoders) → compare against HF `embed_tokens` and full forward logits. + If this diverges, fix weight loading / decoder before touching encoders. + +2. **Add vision:** Run vision encoder → feed features to embedding → + compare. If newly divergent, isolate vision encoder output vs HF. + +3. **Add audio:** Same as above for speech encoder. + +4. **Combined:** All modalities together. If divergent only in combined + mode, suspect LoRA mode mismatch or embedding fusion ordering. + +## Gotchas + +### Module forward() bypass + +The #1 source of missing weights. Directly accessing nested sub-module +parameters (`self.glu.ext_pw_conv_1d.weight`) instead of calling +`self.glu(op, x)` makes onnxscript unable to resolve the full module path. +**Detection:** Conv/MatMul nodes where weight inputs have +`is_initializer=False` and generic names. + +### LoRA mode mismatch + +Some models (Phi4MM) apply LoRA conditionally per modality. If ONNX applies +all adapters unconditionally, run HF reference with `input_mode=3` to match. +**Detection:** Text-only inference diverges but output is reasonable (not +garbage). + +### ClippableLinear omission (Gemma4) + +Using plain `Linear` instead of `ClippableLinear` in Gemma4 vision/audio +encoders causes max_diff 52.68 (audio) / 3.92 (vision). Check HF source +for `ClippableLinear` usage. + +### Empty tensor handling + +Text-only inference crashes when no image/audio features are present. +**Fix:** Zero-pad `image_features` before Gather, then mask with Where. + +### Missing boundary tokens + +Audio/image boundary markers (`<|audio>`, ``) are required for +correct modality region identification. Missing markers → garbled output +even with correct encoder output. + +> Read `references/common-failures.md` when you need detailed code examples +> for each failure mode, including weight name alignment (ModuleList subclass +> name doubling, setattr with dotted names), shape mismatches, dtype +> mismatches (float64 vs float32), and HD transform format issues. + +> Read `references/debugging-cookbook.md` when you need step-by-step +> debugging procedures with code for each phase, intermediate value +> extraction methods, integration test patterns (text-only, audio, vision), +> tolerance guidelines, weight loading verification, and HD multi-crop +> verification. + +## Reference files + +- **Integration tests:** `tests/phi4mm_integration_test.py`, + `tests/integration_test.py` +- **VL debugging skill:** `.agents/skills/debugging-vl-pipeline/SKILL.md` +- **Model implementation:** `src/mobius/models/phi.py`, + `src/mobius/models/gemma4.py` +- **Audio components:** `src/mobius/components/_audio.py`, + `src/mobius/components/_gemma4_audio.py` +- **Vision components:** `src/mobius/components/_vision.py` +- **ClippableLinear:** `src/mobius/components/_gemma4_audio.py` +- **LoRA component:** `src/mobius/components/_lora.py` +- **Weight loading:** `src/mobius/_weight_loading.py` +- **Feature flags:** `src/mobius/_flags.py` +- **ORT GenAI config skill:** `.agents/skills/ort-genai-config/SKILL.md` +- **Weight name alignment skill:** + `.agents/skills/weight-name-alignment/SKILL.md` +- **Gemma4 example:** `examples/gemma4_multimodal.py` diff --git a/.agents/skills/phi4mm-component-parity/references/common-failures.md b/.agents/skills/phi4mm-component-parity/references/common-failures.md new file mode 100644 index 00000000..ed82efa0 --- /dev/null +++ b/.agents/skills/phi4mm-component-parity/references/common-failures.md @@ -0,0 +1,244 @@ +# Common Failure Modes — Detailed Reference + +## 1. Weight name alignment (missing weights) + +**Symptoms:** Hundreds or thousands of weights reported as "unmatched" by +`apply_weights`. Model runs but produces garbage output. + +**Root causes encountered:** + +### a. Module forward() bypass (240 missing weights in Phi4MM) + +The most insidious bug. When a component's `forward()` method directly +accesses nested sub-module parameters (e.g., `self.glu.ext_pw_conv_1d.weight`) +instead of calling the sub-module's `forward()` method, onnxscript cannot +resolve the full module path for the parameter. The weight ends up as an +unnamed, non-initializer constant in the graph. + +```python +# BAD — weights become unnamed +def forward(self, op, x): + return op.Conv(x, self.glu.ext_pw_conv_1d.weight, + self.glu.ext_pw_conv_1d.bias, ...) + +# GOOD — onnxscript resolves full module path +def forward(self, op, x): + return self.glu(op, x) # GLU.forward() calls op.Conv internally +``` + +**Detection:** Check the ONNX graph for Conv/MatMul nodes where weight +inputs have `is_initializer=False` and generic names like "weight"/"bias". + +**Fix:** Add `forward()` methods to sub-modules and call them instead of +directly accessing their parameters. + +### b. ModuleList subclass causing name doubling (8 weights) + +Subclassing `nn.ModuleList` causes the module's own name to appear twice +in the parameter path: `img_projection.img_projection.0.weight` instead +of `img_projection.0.weight`. + +```python +# BAD — name doubling +class ProjectionMLP(nn.ModuleList): + def __init__(self): + super().__init__() + self.append(nn.Linear(1152, 3072)) + self.append(nn.Linear(3072, 3072)) + +# GOOD — use nn.Module with indexed children +class ProjectionMLP(nn.Module): + def __init__(self): + super().__init__() + layers = [nn.Linear(1152, 3072), nn.Linear(3072, 3072)] + for i, layer in enumerate(layers): + setattr(self, str(i), layer) +``` + +### c. setattr with dotted names + +Using `setattr(self, "audio_projection.speech", module)` creates a single +attribute with a dot in its name, rather than a nested module. The resulting +ONNX parameter names won't match HuggingFace's `ModuleDict`-style naming. + +**Fix:** Use `nn.ModuleDict` or create proper nested attributes. + +## 2. Shape mismatches (position embedding 2D vs 3D) + +**Symptoms:** `RuntimeError: shape mismatch` during weight loading. + +**Root cause:** The ONNX component declares a parameter with a different +number of dimensions than the HuggingFace weight. Example: PatchEmbedding +declares `position_embedding.weight` as `[num_patches, hidden_size]` (2D), +but HF stores `[1, num_patches, hidden_size]` (3D). + +**Fix in `preprocess_weights()`:** +```python +# Squeeze the extra batch dimension to match ONNX declaration +if "position_embedding.weight" in key and state_dict[key].dim() == 3: + state_dict[key] = state_dict[key].squeeze(0) # [1,N,H] → [N,H] +``` + +**General rule:** Check whether the preprocess_weights transform goes +in the correct direction (squeeze vs unsqueeze). A common mistake is +writing the transform backwards. + +## 3. Dtype mismatches (float64 vs float32) + +**Symptoms:** ONNX Runtime error: "type mismatch in Mul/Add node" during +inference. + +**Root causes:** + +### a. NumPy default float64 + +`numpy.array(python_float)` defaults to float64. Any constant created +from a Python scalar without explicit dtype will be float64 in the graph. + +```python +# BAD — float64 constant +scale = numpy.array(alpha / rank) # defaults to float64 +op.Mul(x, scale) # Mul(float32, float64) → type error + +# GOOD — explicit float32 +scale = numpy.array(alpha / rank, dtype=numpy.float32) +op.Mul(x, scale) +``` + +### b. Python int auto-promotion + +When passing a Python `int` to an op that expects a tensor, onnxscript +may auto-promote to float64 (implementation-dependent). + +```python +# RISKY — Python int may become float64 +op.Mul(int64_tensor, self.max_position_embeddings) + +# SAFE — explicit constant +op.Mul(int64_tensor, + op.Constant(value_int=self.max_position_embeddings)) +``` + +**Detection:** Run the ONNX model and look for type mismatch errors. +The error message includes the node name — trace it back to the source. + +## 4. LoRA application mismatch (conditional vs unconditional) + +**Symptoms:** Systematic divergence (> 80% logits mismatch) across ALL +test cases, but the model structurally runs correctly. + +**Root cause:** Some models apply LoRA adapters conditionally based on +input modality. For example, Phi4MM applies: +- `input_mode=0` (text): no adapters +- `input_mode=1` (vision): vision LoRA only +- `input_mode=2` (speech): speech LoRA only +- `input_mode=3` (combined): both adapters + +If the ONNX model unconditionally applies all adapters (both vision and +speech LoRA always active), it diverges from HF when HF uses a different +input mode. + +**Quick fix for integration tests:** Run the HF reference with the mode +that matches the ONNX model's behavior (e.g., `input_mode=3` to match +unconditional application of both adapters). + +**Proper fix:** Add an `input_mode` input to the decoder model and use +conditional logic to selectively apply adapters. + +**Detection:** If text-only inference diverges but the model generates +reasonable (not garbage) output, suspect LoRA mode mismatch. Temporarily +zero out all LoRA weights — if base model matches HF perfectly, the +LoRA application mode is the issue. + +## 5. Empty tensor handling (zero-length features) + +**Symptoms:** Crash during text-only inference when no image/audio +features are present. + +**Root cause:** The embedding model's `InputMixer` uses `GatherElements` +to place features at special token positions. With zero features, the +gather indices are empty but the operation may still execute on the +padded dimension, causing shape errors. + +**Fix pattern:** Zero-pad before Gather, then use Where to mask results: +```python +# Pad with one zero row so Gather never accesses out-of-bounds +padded = op.Concat( + op.ConstantOfShape(op.Constant(value_ints=[1, hidden_size])), + features, # may be [0, hidden_size] + axis=0, +) +# After Gather, mask out the padding positions with Where +result = op.Where(feature_mask, gathered, text_embeddings) +``` + +## 6. HD transform image format (5D vs 4D) + +**Symptoms:** Vision model crashes or produces wrong output with multi-crop +HD images. + +**Root cause:** HD-capable vision models expect images in different formats: +- Some expect `[batch, channels, height, width]` (4D, single crop per batch) +- Others expect `[num_images, num_crops, channels, height, width]` (5D) + +The HF processor output format must match the ONNX model's input format. +If using the HF processor for test input preparation, verify it produces +the expected format. + +**Fix:** Check the HF model's preprocessing code for the expected format, +and ensure the ONNX model's input signature matches. For tests, either: +- Use the HF processor: `processor(images=image, return_tensors="np")` +- Or manually construct the correct format for simple test cases + +## 7. Causal mask construction (inputs_embeds vs input_ids) + +**Symptoms:** Attention mask has wrong length, causing decoder crash or +wrong output. + +**Root cause:** When the decoder receives `inputs_embeds` instead of +`input_ids`, the sequence length must be derived from the embeds tensor +shape, not from input_ids. If the mask is built from input_ids length but +the actual sequence includes fused image/audio tokens, the lengths diverge. + +**Fix:** Always derive `seq_len` from `inputs_embeds.shape[1]` when the +model uses inputs_embeds as input. + +## 8. ClippableLinear not used in encoder (Gemma4) + +**Symptoms:** Encoder output has large numerical divergence (max diff > 1.0) +from HuggingFace, despite all weights loading correctly. + +**Root cause:** Some HuggingFace models (e.g. Gemma4) use +`ClippableLinear` — a linear layer with learned finite input/output +activation clipping — for ALL linear layers in their vision and audio +encoders (attention q/k/v/o projections AND MLP gate/up/down projections). +Using plain `Linear` misses the clamping and causes divergence. + +**Detection:** Check HF source for `ClippableLinear`: +```bash +grep -n "ClippableLinear" transformers/models//modeling_.py +``` + +**Fix:** Use `ClippableLinear` from `mobius.components`: +- For attention: use `ClippableLinear` for q/k/v/o_proj +- For MLP: pass `linear_class=ClippableLinear` parameter + +**Impact (Gemma4):** +- Audio: max diff 52.68 → 0.0003 +- Vision: max diff 3.92 → 0.00007 + +## 9. Missing audio/image boundary tokens + +**Symptoms:** Audio transcription or vision description is garbled, but +encoder output matches HuggingFace. + +**Root cause:** HuggingFace wraps modality placeholder tokens with +boundary markers: +- Image: `<|image>` (open) + N × `<|image|>` (pad) + `` (close) +- Audio: `<|audio>` (open) + N × `<|audio|>` (pad) + `` (close) + +Missing boundary markers prevent the model from correctly identifying +modality regions in the input sequence. + +**Fix:** Ensure `build_input_ids` wraps placeholder tokens with the +correct open/close marker token IDs. diff --git a/.agents/skills/phi4mm-component-parity/references/debugging-cookbook.md b/.agents/skills/phi4mm-component-parity/references/debugging-cookbook.md new file mode 100644 index 00000000..202832f9 --- /dev/null +++ b/.agents/skills/phi4mm-component-parity/references/debugging-cookbook.md @@ -0,0 +1,187 @@ +# Debugging Cookbook — Step-by-Step Procedures + +## Step-by-step debugging process + +### Phase 1: Text-only baseline + +1. **Build ONNX model** with tiny config (for fast iteration) or full + weights (for accuracy). +2. **Run text-only** through embedding → decoder (skip encoders). +3. **Compare embedding output** against `hf_model.model.embed_tokens(ids)`. + If this diverges, the issue is in weight loading or embedding model. +4. **Compare decoder logits** against HF full forward. + If embedding matches but logits diverge, issue is in decoder. + +### Phase 2: Isolate decoder issues + +5. **Check weight count** — verify all expected weights are loaded: + ```python + pkg = build(model_id, load_weights=True) + # apply_weights prints statistics: applied, skipped, unmatched + ``` +6. **Disable LoRA** — if the model uses LoRA, zero out adapter weights and + compare base model output against HF with adapters disabled. +7. **Layer-by-layer** — add intermediate outputs to the ONNX graph (see + `debugging-vl-pipeline` skill) to find which decoder layer first diverges. + +### Phase 3: Add modalities + +8. **Vision only** — run vision encoder, feed features to embedding, compare. +9. **Audio only** — run speech encoder, feed features to embedding, compare. +10. **Combined** — all modalities together. + +At each step, if a newly added component causes divergence, isolate that +component's output against HF. + +### Phase 4: LoRA verification + +11. **Match input modes** — ensure HF reference uses the same adapter + activation mode as ONNX (e.g., `input_mode=3` for both adapters). +12. **Compare with LoRA** — verify LoRA scaling factor: `alpha / rank`. +13. **Check adapter routing** — for multi-adapter models, verify the correct + adapter set is active for each modality combination. + +## Integration test patterns + +### Test configuration + +```python +# Always use for HF reference: +hf_model = AutoModelForCausalLM.from_pretrained( + model_id, + trust_remote_code=True, + attn_implementation="eager", # No flash_attn dependency + torch_dtype=torch.float32, # Match ONNX precision +) + +# For models with conditional LoRA: +hf_model.input_mode = 3 # Match ONNX unconditional LoRA +# Or: pass input_mode=3 to forward() if supported +``` + +### Text-only test + +```python +def test_text_only_prefill_logits_match(self): + input_ids = tokenizer.encode("Hello world", return_tensors="np") + empty_image = np.zeros((0, hidden_size), dtype=np.float32) + empty_audio = np.zeros((0, hidden_size), dtype=np.float32) + + embeds = embedding_session.run({ + "input_ids": input_ids, + "image_features": empty_image, + "audio_features": empty_audio, + })["inputs_embeds"] + + onnx_logits = decoder_session.run({ + "inputs_embeds": embeds, + "attention_mask": np.ones((1, seq_len), dtype=np.int64), + "position_ids": np.arange(seq_len).reshape(1, -1), + # ... zero KV cache + })["logits"] + + hf_logits = hf_model(input_ids=..., input_mode=3).logits.numpy() + assert_logits_close(onnx_logits, hf_logits) +``` + +### Audio test + +```python +def test_audio_prefill_logits_match(self): + # Prepare mel spectrogram input + mel = load_audio_as_mel(audio_path) # [1, n_mel, time] + + speech_out = speech_session.run({ + "audio_embeds": mel, + "audio_sizes": np.array([[mel.shape[-1]]], dtype=np.int64), + "audio_projection_mode": np.array(0, dtype=np.int64), + }) + speech_features = speech_out["audio_features"] + + # Build input_ids with audio placeholder tokens + input_ids = build_audio_input_ids(prompt, num_audio_tokens) + + embeds = embedding_session.run({ + "input_ids": input_ids, + "image_features": np.zeros((0, hidden_size), dtype=np.float32), + "audio_features": speech_features, + })["inputs_embeds"] + + onnx_logits = decoder_session.run(...)["logits"] + hf_logits = hf_forward_with_audio(...) + assert_logits_close(onnx_logits, hf_logits) +``` + +### Tolerance guidelines + +| Precision | atol | rtol | Notes | +|-----------|------|------|-------| +| float32 | 1e-4 | 2e-2 | Standard for single-forward-pass | +| float32 (deep model, 32+ layers) | 1e-3 | 5e-2 | Error compounds over layers | +| float16 / bfloat16 | 0.01 | 0.05 | Wider tolerance for mixed precision | +| Cosine similarity (last token) | > 0.98 | — | Primary correctness metric | +| Argmax match (first prediction) | exact | — | Should always match | + +### Weight loading verification + +After `apply_weights`, check the statistics: +```python +# Expected output: +# Applied: 485/485 weights +# Skipped: 0 (weights in state_dict but not in graph) +# Unmatched: 0 (weights in state_dict with no graph match) + +# If unmatched > 0, dump the names to find alignment issues: +pkg = build(model_id) +state_dict = download_weights(model_id) +state_dict = module.preprocess_weights(state_dict) +# Compare state_dict.keys() vs graph initializer names +``` + +## Vision-specific verification (HD multi-crop) + +For models with HD dynamic resolution (Phi4MM, Phi3-Vision): + +### Input preparation + +```python +from transformers import AutoProcessor + +processor = AutoProcessor.from_pretrained( + model_id, trust_remote_code=True +) +inputs = processor( + images=image, + text=prompt, + return_tensors="pt", +) +pixel_values = inputs["pixel_values"] # [num_crops, C, H, W] or 5D +image_sizes = inputs["image_sizes"] # [num_images, 2] +``` + +### HD transform verification + +The HD transform typically: +1. Splits image into base (global) + sub-image crops +2. Encodes each crop through vision encoder → `[num_patches, hidden_size]` +3. Applies spatial merge (e.g., AvgPool2d + reshape) → compressed tokens +4. Adds learned separators (glb_GN between global/sub, sub_GN between subs) +5. Projects to text dimension via MLP + +```python +# Verify token count matches expected: +# global: (image_size/patch_size)^2 / merge^2 tokens +# per sub-image: same count +# separators: 1 glb_GN + (num_subs - 1) sub_GN rows +total_expected = global_tokens + num_subs * sub_tokens + separator_count +assert image_features.shape[0] == total_expected +``` + +### Testing without HD (simpler) + +For initial validation, use base resolution (single crop, no HD): +```python +# Single image at base resolution — bypasses HD transform +pixel_values = np.random.randn(1, 3, 384, 384).astype(np.float32) +image_sizes = np.array([[384, 384]], dtype=np.int64) +``` diff --git a/.agents/skills/reusable-components/SKILL.md b/.agents/skills/reusable-components/SKILL.md new file mode 100644 index 00000000..f0265315 --- /dev/null +++ b/.agents/skills/reusable-components/SKILL.md @@ -0,0 +1,221 @@ +--- +name: reusable-components +description: > + Create or extend reusable ONNX building blocks in the mobius component + library. Use when adding Attention, MLP, norm, RoPE, or embedding + components; understanding parameter naming and nn.Module conventions; + applying design principles (subclass over flags, model-agnostic); + or wiring shared-weight / per-layer adapter patterns. +--- + +# Skill: Reusable Components + +## When to use + +Use this skill when creating or extending the building blocks that models are +composed from — attention layers, MLPs, normalisations, embeddings, RoPE +variants, and activations. + +## Component library overview + +All components live in `src/mobius/components/` and inherit from +`onnxscript.nn.Module`. Each component's `forward(op, ...)` method builds +ONNX nodes via the `OpBuilder`. + +``` +components/ +├── _activations.py # get_activation(), SiLU module +├── _attention.py # Multi-head / GQA attention with KV cache; Qwen35Attention (gated GQA) +├── _audio.py # ConformerEncoder (NeMo subsampling, T5 bias, Conformer layers) +├── _common.py # Embedding, Linear, LayerNorm, LayerNormNoAffine, GroupNorm, create_attention_bias +├── _gemma4_audio.py # Gemma4 audio encoder: ClippableLinear, ConvSubsampling, SlidingWindowAttention +├── _conv.py # Conv2d (2D convolution with bias and groups) +├── _decoder.py # DecoderLayer (pre-norm residual block) +├── _encoder.py # BertEmbeddings, EncoderAttention, EncoderLayer +├── _lora.py # LoRALinear (base + per-adapter A/B/scale) +├── _mlp.py # Gate-up-down MLP +├── _moe.py # MoELayer, TopKGate, SparseMixerGate +├── _multimodal.py # Projectors + InputMixer +├── _qwen3_vl_vision.py # Qwen3-VL block-diagonal vision encoder +├── _gated_deltanet.py # GatedDeltaNet (recurrent linear attention for Qwen3.5 hybrid) +├── _rms_norm.py # RMSNorm, OffsetRMSNorm (1+weight), GatedRMSNorm (norm * SiLU gate) +├── _rotary_embedding.py # RoPE variants (Default, Linear, Dynamic, Llama3, InterleavedMRope, ChunkedMRope) +├── _vision.py # PatchEmbedding, VisionEncoder, VisionModel +└── _whisper.py # Conv1d, WhisperAttention, WhisperDecoderLayer, WhisperEncoderLayer +``` + +Model files import shared primitives from `components/` and alias them with +an underscore prefix for local use: + +```python +from mobius.components import Conv2d as _Conv2d, SiLU as _SiLU +``` + +Model-specific compound blocks (e.g. `_TimestepEmbedding`, `_DiTBlock`, +`_ResNetBlock2D`) remain in the model files they belong to. + +## How to create a new component + +### 1. Define the class + +```python +from onnxscript import nn +from onnxscript._internal import builder + + +class MyComponent(nn.Module): + def __init__(self, hidden_size: int): + super().__init__() + self.weight = nn.Parameter([hidden_size]) + + def forward(self, op: builder.OpBuilder, hidden_states): + return op.Mul(hidden_states, self.weight) +``` + +### 2. Parameter naming + +Parameter names are **automatically set** from the attribute name by +`nn.Module.__setattr__`. You do **not** need to pass `name=` when the +attribute name matches the desired ONNX name: + +```python +# GOOD — name is automatically "weight" +self.weight = nn.Parameter([hidden_size]) + +# Only use name= when the attribute name differs from the desired ONNX name +self.patch_embedding = nn.Parameter( + [out_ch, in_ch, kH, kW], name="patch_embedding.weight" +) +``` + +When the component is nested in a module tree, names are automatically +prefixed by parent attribute names: + +```python +# In model: self.layer = MyComponent(...) +# Resulting ONNX name: "layer.weight" +``` + +**Critical:** Parameter names must be unique within a component. If two +parameters share the same attribute name at different levels, one will +silently overwrite the other. + +To create a parameter with precomputed data (e.g. frozen positional embeddings), +use the `data=` argument: + +```python +import onnx_ir as ir +self.embed_positions = nn.Parameter( + [max_positions, d_model], + name="embed_positions.weight", + data=ir.tensor(numpy_array), +) +``` + +Do **not** assign `_const_value` directly. + +### 3. Export from `__init__.py` + +Add to `src/mobius/components/__init__.py`: + +```python +__all__ = [..., "MyComponent"] +from mobius.components._my_component import MyComponent +``` + +### 4. Write unit tests + +Create `_my_component_test.py` alongside the source: + +```python +from mobius._testing import create_test_builder, create_test_input + +class TestMyComponent: + def test_forward(self): + comp = MyComponent(hidden_size=64) + b, op, graph = create_test_builder() + x = create_test_input(b, "x", [1, 10, 64]) + result = comp(op, x) + b._adapt_outputs([result]) + assert graph.num_nodes() > 0 + + def test_parameter_names(self): + comp = MyComponent(hidden_size=64) + names = [n for n, _ in comp.named_parameters()] + assert "weight" in names +``` + +## Key components (concise) + +| Component | Signature | Notes | +|-----------|-----------|-------| +| `Attention(config)` | MHA/GQA/MQA, optional QK norm, custom scale | Opset 23 `op.Attention` | +| `Qwen35Attention(config)` | Gated GQA with partial RoPE | `attn_output * sigmoid(gate)` | +| `MLP(config)` | gate_proj + up_proj + down_proj | Activation from `config.hidden_act` | +| `DecoderLayer(config)` | Pre-norm residual block | Subclass to customize norms | +| `GatedDeltaNet(config)` | Recurrent linear attention | Qwen3.5 hybrid; delta rule recurrence | +| `RMSNorm(h, eps)` | Opset 23 `RMSNormalization` | `OffsetRMSNorm` for `1+weight` variant | +| `LayerNorm(h, eps)` | `LayerNormalization` op | Check HF config for correct eps | +| `ClippableLinear(in, out)` | Linear + learned input/output clipping | Critical for Gemma4 encoders | +| `Embedding(V, D)` | Gather on weight matrix | | + +> Read `references/component-examples.md` when you need detailed constructor +> arguments, Attention/MLP/norm variant code, ClippableLinear weight mapping, +> RoPE factory usage, shared-weight + per-layer adapter patterns, or the +> `op.Identity` pattern for exposing parameters as graph outputs. + +## Design principles + +1. **Favour subclasses over flags.** When a model family has a unique variant + (e.g. Gemma's `weight + 1` norm), create a subclass rather than adding a + boolean flag to the base class. + +2. **Keep components model-agnostic.** A component should work for any model + that has the right config fields. Model-specific wiring belongs in the + model module. + +3. **One file per concern.** Attention in `_attention.py`, RoPE in + `_rotary_embedding.py`, etc. Tests co-located as `_*_test.py`. + +4. **Reuse across model families.** The same `Attention` component is used by + LLaMA, Mistral, Qwen, Phi, and others. Only override when the + architecture genuinely differs. + +5. **Multiple reusable variants, not one-size-fits-all.** When models need + different behaviour (e.g. MoE gates, projector types), create separate + classes rather than cramming everything into one class with many branches. + +6. **Comment generously with architecture context.** Annotate tensor shapes + after ops (e.g. `# (N, num_heads, head_dim)`), explain multi-step + computations (window reordering, RoPE, spatial merge), and document how + the ONNX graph maps to the HuggingFace reference implementation. + +7. **Match HuggingFace's precision behaviour.** Components must work with any + compute dtype (float32, float16, bfloat16). For numerically sensitive ops + (`exp`, `softplus`, RMSNorm variance), upcast to float32 with + `op.Cast(to=ir.DataType.FLOAT)`, compute, then cast back with `op.CastLike(result, input)`. + For dtype-adaptive parameters, use `op.CastLike(param, reference)`. + +## ONNX op patterns overview + +Key patterns for building components: + +- **Scalar constants:** Use `op.Constant(value_ints=[...])` for tensor inputs +- **CastLike:** Use `op.CastLike(param, activation)` for dtype-agnostic casting +- **fp32 upcast:** `op.Cast(to=FLOAT)` → compute → `op.CastLike(result, input)` + for numerically sensitive ops (Exp, Softplus, RMSNorm variance) +- **Shape extraction:** `op.Shape(x, start=i, end=i+1)` — never `Gather(Shape(x), ...)` +- **ModuleList vs Sequential:** Use `nn.Sequential` for fixed chains, + `nn.ModuleList` for custom iteration + +> Read `references/onnx-op-patterns.md` when you need full code examples for +> scalar constants, CastLike, fp32 upcast tables, shape manipulation, module +> containers, conditional ops, or the `op.Identity` graph output pattern. + +## Cross-references + +- **Weight name alignment:** `.agents/skills/weight-name-alignment/SKILL.md` +- **Multimodal components:** `.agents/skills/multimodal-models/SKILL.md` +- **MoE components:** `.agents/skills/moe-models/SKILL.md` +- **Writing tests:** `.agents/skills/writing-tests/SKILL.md` +- **Rewrite rules:** `.agents/skills/writing-rewrite-rules/SKILL.md` diff --git a/.agents/skills/reusable-components/references/component-examples.md b/.agents/skills/reusable-components/references/component-examples.md new file mode 100644 index 00000000..ddd44b6b --- /dev/null +++ b/.agents/skills/reusable-components/references/component-examples.md @@ -0,0 +1,297 @@ +# Component Examples — Detailed Reference + +## Attention + +```python +Attention(config) +Attention(config, scale=0.015625) # Override default 1/sqrt(head_dim) scale +# Inputs: hidden_states, attention_bias, position_embeddings, past_key_value +# Outputs: attn_output, (key_cache, value_cache) +``` + +Handles MHA, GQA, and MQA via `num_key_value_heads`. Supports optional QK +norm (`attn_qk_norm=True`) and bias on Q/K/V/O projections. + +The optional `scale` parameter overrides the default `head_dim**-0.5` attention +scale. Use this when a model specifies a custom attention multiplier (e.g. +Granite's `attention_multiplier`). When `None` (default), uses `1/sqrt(head_dim)`. + +The ONNX `Attention` op (opset 23) has an `is_causal` attribute. For +decoder self-attention in encoder-decoder models (e.g., Whisper), set +`is_causal=1` instead of building an explicit causal mask with +`create_attention_bias`. + +Some models (Whisper) require **Q pre-scaling** for numerical parity with +HuggingFace: multiply Q by `head_dim**-0.5` before passing to `op.Attention` +and set `scale=1.0`. This matches HF's order of operations and avoids +floating-point divergence in softmax. + +**Qwen35Attention** (`_attention.py`): Gated GQA variant for Qwen3.5. Doubles +the Q projection to produce both Q and a gate signal, applies per-head +`OffsetRMSNorm` to Q and K, supports partial RoPE, and gates the output with +`attn_output * sigmoid(gate)`. + +## MLP + +```python +MLP(config) +# Uses gate_proj + up_proj + down_proj with configurable activation +``` + +The activation function comes from `config.hidden_act` and is resolved by +`get_activation()`. + +## DecoderLayer + +```python +DecoderLayer(config) +# Pre-norm residual: LayerNorm → Attention → Add → LayerNorm → MLP → Add +``` + +To customise, subclass and override the components: + +```python +class MyDecoderLayer(DecoderLayer): + def __init__(self, config): + super().__init__(config) + # Replace norm with custom variant + self.input_layernorm = MyRMSNorm(config.hidden_size, eps=config.rms_norm_eps) +``` + +## GatedDeltaNet (Linear Attention) + +```python +GatedDeltaNet(config) +# Inputs: hidden_states, position_embeddings (unused), past_key_value (unused) +# Outputs: output, (conv_state, recurrent_state) +``` + +Recurrent linear attention mechanism from the Qwen3.5 hybrid architecture +(`_gated_deltanet.py`). Key operations: fused QKV projection, causal +depthwise Conv1D, L2-normalised Q/K, exponential decay gates, delta rule +recurrence, and gated output via `GatedRMSNorm`. Supports GQA-like key +head grouping (`num_k_heads` → repeat to `num_v_heads`). State is +`conv_state` + `recurrent_state` (currently zero-initialised for stateless +export). + +## RoPE variants + +Created via the factory function `initialize_rope(config)`: + +| `config.rope_type` | Class | Use case | +|--------------------|-------|----------| +| `"default"` | `DefaultRope` | Standard RoPE | +| `"linear"` | `LinearRope` | Linear scaling (factor in `rope_scaling`) | +| `"dynamic"` | `DynamicNTKRope` | Dynamic NTK scaling | +| `"llama3"` | `Llama3Rope` | LLaMA-3 piecewise scaling | + +**MRope (Multimodal RoPE):** Two variants share a `_MRopeBase` base class +that splits frequencies into temporal (T), height (H), and width (W) sections. +`ChunkedMRope` uses a chunked layout `[TTT...HHH...WWW]` (Qwen2-VL). +`InterleavedMRope` uses an interleaved layout `[T,H,W,T,H,W,...]` and +supports `partial_rotary_factor` for partial RoPE (Qwen3-VL, Qwen3.5). + +RoPE embeddings are precomputed as `cos_cache` / `sin_cache` initializers +and looked up at runtime via `Gather` on `position_ids`. + +## RMSNorm + +```python +RMSNorm(hidden_size, eps=1e-6) +``` + +Uses the ONNX `RMSNormalization` op from opset 23. The `eps` is a float +attribute (not a Parameter). + +For Gemma's `weight + 1` variant, subclass: + +```python +class GemmaRMSNorm(RMSNorm): + def forward(self, op, hidden_states): + weight_plus_one = op.Add(self.weight, 1.0) + return apply_rms_norm(op, hidden_states, weight_plus_one, self.variance_epsilon) +``` + +**OffsetRMSNorm** (`_rms_norm.py`): `output * (1 + weight)` variant where +HuggingFace stores weights initialised to 0, so the effective multiplier is +`1 + weight`. Used by Qwen3.5 for per-head Q/K normalisation. + +**GatedRMSNorm** (`_rms_norm.py`): `RMSNorm(x) * SiLU(gate)` — applies +RMS normalisation then element-wise gates the result with a SiLU activation +on a separate gate input. Used by GatedDeltaNet output projection. + +## LayerNorm + +```python +LayerNorm(hidden_size, eps=1e-6) +``` + +Uses the ONNX `LayerNormalization` op. **Always check the model's HF +config for the correct epsilon** — the default `1e-6` does not match all +models. For example, Whisper uses `1e-5`. A wrong epsilon causes large +numerical drift that amplifies through the network. + +## LayerNormNoAffine + +```python +LayerNormNoAffine(dim, eps=1e-5) +``` + +Layer normalization **without learnable parameters** (`elementwise_affine=False` +in PyTorch). Used in AdaLayerNorm blocks where scale/shift come from a +separate modulation projection. Calls `op.LayerNormalization` with no +`Scale` or `Bias` inputs. + +For weight-free LayerNorm that still needs frozen ones/zeros (e.g. OLMo-1B), +create constant parameters with `data=ir.tensor(...)` instead. + +**Key:** RMSNorm vs LayerNorm is NOT interchangeable. LayerNorm subtracts +the mean; RMSNorm does not. Using the wrong type causes max abs diff > 1.0 +that grows through layers. + +## GroupNorm + +```python +GroupNorm(num_groups, num_channels, eps=1e-5) +``` + +Group normalization with learnable `weight` and `bias`. Uses the ONNX +`GroupNormalization` op. Commonly used in diffusion models (UNet, VAE). + +## Conv2d + +```python +Conv2d(in_channels, out_channels, kernel_size=3, stride=1, padding=0, groups=1) +``` + +2D convolution with bias, matching `torch.nn.Conv2d(bias=True)`. Used in +diffusion models (VAE, UNet, ControlNet) and vision patch embeddings. +Parameters: `weight` (`[out, in/groups, kH, kW]`) and `bias` (`[out]`). + +## SiLU + +```python +SiLU() +# SiLU (Swish) activation as a module: x * sigmoid(x) +``` + +Useful in `nn.Sequential` containers where an activation needs to be a +module with a `forward()` method. For functional use, call +`get_activation("silu")` instead. + +## Linear + +```python +Linear(in_features, out_features, bias=False) +# Uses MatMul (+ optional Add for bias) +``` + +## ClippableLinear + +```python +ClippableLinear(in_features, out_features, bias=False) +# Linear with learned input/output activation clamping +# Matches HuggingFace Gemma4ClippableLinear +``` + +Wraps a standard `Linear` with 4 learned scalar parameters: +`input_min`, `input_max`, `output_min`, `output_max`. Inputs are clamped +before the linear projection and outputs are clamped after: + +```python +x = Clip(x, input_min, input_max) +x = MatMul(x, weight.T) [+ bias] +x = Clip(x, output_min, output_max) +``` + +**Critical:** HuggingFace `Gemma4ClippableLinear` stores *finite* learned +bounds (not ±inf). Using plain `Linear` instead of `ClippableLinear` causes +large numerical divergence in Gemma4 vision and audio encoders: + +- Audio encoder: max diff 52.68 → 0.0003 after fix +- Vision encoder: max diff 3.92 → 0.00007 after fix + +The component is exported from the public API: `from mobius.components +import ClippableLinear`. + +**Weight mapping:** HF stores `.linear.weight` for the actual +weight (the `.linear.` segment is stripped by `preprocess_weights`), and +`.input_min`, `.input_max`, `.output_min`, +`.output_max` as direct scalar buffers. + +The MLP component accepts a `linear_class` parameter, so you can pass +`linear_class=ClippableLinear` to use it for all projections in the MLP. + +## Embedding + +```python +Embedding(num_embeddings, embedding_dim, padding_idx=0) +# Uses Gather on weight matrix +``` + +## Shared weights with per-layer adapters + +Some architectures reuse the same transformer block across multiple layers, +with per-layer low-rank adapters that differentiate each usage (e.g. Zamba2, +which shares one transformer across 6 hybrid layers). + +### The scope challenge + +`onnxscript.nn` determines ONNX initializer names from the module call stack. +A single module instance called multiple times produces the **same** initializer +names each time — which is exactly what we want for shared weights. But +per-layer adapters need **different** names for each layer. + +### Pattern: split shared + per-instance modules + +```python +class _TextModel(nn.Module): + def __init__(self, config): + super().__init__() + # Shared weights: ONE instance → single set of ONNX initializers + self.shared_transformer = SharedAttentionLayer(config) + + # Per-layer adapters: ModuleList → "adapters.0.*", "adapters.1.*" + self.adapters = nn.ModuleList([ + AdapterModule(config) for _ in range(num_layers) + ]) + + # Shared MLP at model scope if adapter output must mix with MLP + self.gate_proj = Linear(hidden, intermediate) +``` + +### Critical: use `__call__` for per-index scope + +When iterating over adapter ModuleList elements, you **must** call the +element (triggering `__call__`) rather than accessing its sub-attributes: + +```python +# ❌ Broken: adapter_out gets "q_adapter.weight" scope (same for all idx!) +adapter_out = self.adapters[idx].q_adapter(op, x) + +# ✅ Correct: adapter_out gets "adapters.{idx}.q_adapter.weight" scope +adapter_out = self.adapters[idx](op, x) +``` + +### Handling the MLP circular dependency + +When per-layer adapter output must be combined with shared MLP weights, and +the adapter input comes from inside the shared module, split the shared +module into phases: + +1. **Phase 1 — Shared attention:** `shared_transformer(op, x)` → returns + `mlp_input` (pre-MLP hidden states) + KV cache +2. **Phase 2 — Per-layer adapter:** `self.adapters[idx](op, mlp_input)` → + per-layer contribution (correct `adapters.{idx}` scope) +3. **Phase 3 — Shared MLP at caller scope:** Apply gate/up/down projections + registered on the caller module, combining with adapter output + +```python +# In _TextModel.forward(): +mlp_input, kv = self.shared_transformer(op, x) # shared scope +adapter_out = self.mlp_adapters[idx](op, mlp_input) # per-layer scope +gate = op.Add(self.gate_proj(op, mlp_input), adapter_out) # model scope +``` + +**Reference implementation:** `models/zamba2.py` — Zamba2 hybrid Mamba2 + +shared attention with Q/K/V/MLP low-rank adapters. diff --git a/.agents/skills/reusable-components/references/onnx-op-patterns.md b/.agents/skills/reusable-components/references/onnx-op-patterns.md new file mode 100644 index 00000000..10eeabc5 --- /dev/null +++ b/.agents/skills/reusable-components/references/onnx-op-patterns.md @@ -0,0 +1,218 @@ +# Common ONNX Op Patterns — Detailed Reference + +## Scalar constants + +Many ONNX ops require tensor inputs, not Python scalars: + +```python +# K for TopK must be a 1-D tensor +k = op.Constant(value_ints=[2]) +values, indices = op.TopK(logits, k, axis=-1) + +# Integer constants +one = op.Constant(value_int=1) + +# Float constants +eps = op.Constant(value_float=1e-6) +``` + +## Dtype-agnostic casting with `CastLike` + +When a parameter or constant needs to match an activation tensor's dtype +without knowing what it is at graph-build time, use `op.CastLike`: + +```python +# GOOD — adapts to whatever dtype hidden_states has +scale = op.CastLike(op.Constant(value_float=1e-6), hidden_states) +``` + +**When `op.Cast(to=...)` IS appropriate:** converting between fundamentally +different types (e.g. int64 position_ids to float for arithmetic, or float +timesteps to the model's compute type), and for the fp32 upcast pattern +below. + +## Precision-sensitive ops: fp32 upcast pattern + +Some operations are numerically unstable in float16/bfloat16 and must run +in float32 to match HuggingFace's behaviour. The pattern is: +**upcast → compute → cast back**. + +```python +# Upcast inputs to fp32 for numerically sensitive exp/softplus +dt_f32 = op.Cast(dt, to=ir.DataType.FLOAT) +dt_f32 = op.Softplus(dt_f32) +a_neg = op.Neg(op.Exp(op.Cast(self.A_log, to=ir.DataType.FLOAT))) +... +# Cast output back to input dtype +y = op.CastLike(y_f32, x) +``` + +**Operations that need fp32 (based on HuggingFace source):** + +| Op | Why | HF pattern | +|----|-----|-----------| +| `Exp` on A_log/decay | Overflow/underflow in fp16 range | `self.A_log.float()` | +| `Softplus` (dt) | Uses exp internally | `softplus(dt + dt_bias)` stays in fp32 context | +| `Exp(dt * A)` (discretisation) | Exponential of product | `A.to(dtype=torch.float32)` | +| SSM state update | Accumulates over many steps | `hidden_states.float()`, `B.float()`, `C.float()` | +| RMSNorm variance | Small values squared then averaged | `hidden_states.to(torch.float32)` | +| GatedRMSNorm (SiLU + norm) | Both gate and variance need fp32 | `gate.to(torch.float32)` | + +**When fp32 upcast is NOT needed:** + +- Linear projections (`MatMul`) — handled by the runtime +- SiLU activation on conv output — stays in model dtype in HF +- Standard attention — ONNX `Attention` op handles precision internally +- `RMSNormalization` op — has `stash_type=1` (default) which auto-upcasts + the variance computation to fp32 + +**Rule of thumb:** Check the HuggingFace source for `.float()` or +`.to(torch.float32)` calls. Every such call indicates an fp32 upcast +region that the ONNX component must replicate with explicit +`op.Cast(to=ir.DataType.FLOAT)` ... `op.CastLike(result, input)` bracketing. + +## Shape manipulation + +Use `op.Shape` with `start` and `end` attributes to extract specific +dimensions directly — do **not** use `Gather(Shape(x), index)`: + +```python +# GOOD — single Shape node with start/end +batch_size = op.Shape(x, start=0, end=1) # 1-D [1]-element tensor +seq_len = op.Shape(x, start=1, end=2) +hidden_dim = op.Shape(x, start=2, end=3) + +# BAD — unnecessary Gather +batch_size = op.Gather(op.Shape(x), [0], axis=0) +``` + +Building dynamic shapes for Reshape/Concat: + +```python +new_shape = op.Concat(batch_size, hidden_dim, op.Constant(value_ints=[-1]), axis=0) +reshaped = op.Reshape(x, new_shape) +``` + +Since `Shape(start, end)` returns a 1-D tensor, it can be passed directly +to ops expecting 1-D shape inputs (e.g. `Slice` starts/ends, `Reshape`, +`Concat` for shape building) without intermediate `Reshape` calls. + +## Conditional operations + +```python +mask = op.Equal(input_ids, op.Constant(value_int=token_id)) +result = op.Where(mask, true_value, false_value) +``` + +## Module lists and sequential containers + +Use `nn.ModuleList` to register a list of child modules. It automatically +registers children with numeric keys (`"0"`, `"1"`, ...) and supports +iteration, indexing, and `len()`: + +```python +# GOOD — nn.ModuleList +self.layers = nn.ModuleList( + [DecoderLayer(config) for _ in range(config.num_hidden_layers)] +) + +# BAD — manual setattr loop +self.layers = [DecoderLayer(config) for _ in range(config.num_hidden_layers)] +for i, layer in enumerate(self.layers): + setattr(self, f"layers.{i}", layer) +``` + +For sequential containers where children should be called in order (e.g. +matching HF `nn.Sequential`), use `nn.Sequential`. It subclasses +`nn.ModuleList` and adds automatic forward chaining: + +```python +from mobius.components import Linear, SiLU + +# nn.Sequential chains forward calls: SiLU → Linear +self.img_mod = nn.Sequential(SiLU(), Linear(dim, 6 * dim)) + +# Clean call — output chains through each child +result = self.img_mod(op, temb) # equivalent to Linear(SiLU(temb)) +``` + +`nn.Sequential` produces the same parameter names as `nn.ModuleList` +(`img_mod.0.weight`, `img_mod.1.weight`). The key implementation detail: +it overrides `_set_name` to keep children with simple "0", "1" names +(not fully-qualified), because `__call__` already pushes the parent name +onto the scope stack. + +**When to use which:** +- `nn.Sequential` — children are called in a fixed chain (e.g. `to_out`, + modulation layers, FFN with activation gaps) +- `nn.ModuleList` — children need custom iteration logic (e.g. decoder + layers with residual connections, down/up blocks with skip connections) + +For non-consecutive indices (e.g. matching HF `nn.Sequential` with +activation/dropout layers at skipped positions), include parameter-free +placeholder modules to fill the gaps: + +```python +class _NoOpModule(nn.Module): + """Placeholder for HF Dropout (no params, identity at inference).""" + def forward(self, op, x): + return x + +# Matches HF net.0.proj.weight, net.2.weight (Dropout at index 1) +self.net = nn.Sequential( + _GELUGate(dim, inner_dim * 2), # index 0 + _NoOpModule(), # index 1 (Dropout placeholder) + Linear(inner_dim, dim), # index 2 +) +result = self.net(op, x) # chains: GELUGate → NoOp → Linear +``` + +If `nn.Sequential` is not available, fall back to `nn.ModuleList` with +explicit indexing: + +```python +self.img_mod = nn.ModuleList([SiLU(), Linear(dim, 6 * dim)]) +# Manual chaining: +result = self.img_mod[1](op, self.img_mod[0](op, temb)) +``` + +## Exposing parameters as graph outputs + +Sometimes a generation loop needs access to model weights for external +computation (e.g. embedding lookups in numpy). Use `op.Identity()` to +expose a parameter as a graph output without affecting the initializer +name used for weight loading: + +```python +class MyModel(nn.Module): + def __init__(self, config): + super().__init__() + # Stacked weight exposed for external lookup + self.stacked_embedding = nn.Parameter([num_groups, vocab, hidden]) + + def forward(self, op, ...): + # Use Identity to create a separate output value. + # This prevents the optimizer from renaming the initializer + # when the task sets a custom output name. + embeddings_out = op.Identity(self.stacked_embedding) + return logits, present_key_values, embeddings_out +``` + +In the task, you can safely rename the Identity output: + +```python +# Safe — Identity separates the output name from the initializer name +embeddings_out.name = "codec_embeddings" +graph.outputs.append(embeddings_out) +``` + +**Important:** Without the Identity node, the optimizer may fold the +reference and setting `output.name = "..."` would rename the +initializer itself, breaking `preprocess_weights` name mapping. + +The generation loop extracts the weights once via a dummy inference: + +```python +weights = session.run(dummy_inputs)["codec_embeddings"] # (N, vocab, H) +# Use as numpy lookup: embed = weights[step, code_id, :] +``` From 43fba1d04d2e9f13ee120128559ba745ca2626a3 Mon Sep 17 00:00:00 2001 From: Justin Chu Date: Tue, 21 Apr 2026 16:16:44 +0000 Subject: [PATCH 4/8] Restructure skills: trim debugging-vl-pipeline, improve descriptions, add references - Trim debugging-vl-pipeline SKILL.md from 533 to 222 lines by extracting 10 failure modes to references/failure-modes.md and 3 intermediate value extraction methods to references/extraction-methods.md - Rescope debugging-vl-pipeline for existing pipeline debugging vs phi4mm-component-parity for building new multi-encoder models - Update descriptions for 8 skills to use imperative phrasing with clear trigger keywords - Verify no .github/skills/ references remain (all use .agents/skills/) Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Justin Chu --- .agents/skills/adding-a-new-model/SKILL.md | 364 +++++++++++++++ .../references/architecture-patterns.md | 259 +++++++++++ .../references/weight-preprocessing.md | 277 ++++++++++++ .agents/skills/debugging-vl-pipeline/SKILL.md | 222 ++++++++++ .../references/extraction-methods.md | 144 ++++++ .../references/failure-modes.md | 206 +++++++++ .agents/skills/diffusion-models/SKILL.md | 389 ++++++++++++++++ .agents/skills/moe-models/SKILL.md | 384 ++++++++++++++++ .../skills/multi-agent-coordination/SKILL.md | 215 +++++++++ .../skills/phi4mm-component-parity/SKILL.md | 14 +- .agents/skills/quality-checklist/SKILL.md | 244 ++++++++++ .agents/skills/scan-and-multi-image/SKILL.md | 392 ++++++++++++++++ .agents/skills/weight-name-alignment/SKILL.md | 417 ++++++++++++++++++ .agents/skills/writing-rewrite-rules/SKILL.md | 302 +++++++++++++ 14 files changed, 3823 insertions(+), 6 deletions(-) create mode 100644 .agents/skills/adding-a-new-model/SKILL.md create mode 100644 .agents/skills/adding-a-new-model/references/architecture-patterns.md create mode 100644 .agents/skills/adding-a-new-model/references/weight-preprocessing.md create mode 100644 .agents/skills/debugging-vl-pipeline/SKILL.md create mode 100644 .agents/skills/debugging-vl-pipeline/references/extraction-methods.md create mode 100644 .agents/skills/debugging-vl-pipeline/references/failure-modes.md create mode 100644 .agents/skills/diffusion-models/SKILL.md create mode 100644 .agents/skills/moe-models/SKILL.md create mode 100644 .agents/skills/multi-agent-coordination/SKILL.md create mode 100644 .agents/skills/quality-checklist/SKILL.md create mode 100644 .agents/skills/scan-and-multi-image/SKILL.md create mode 100644 .agents/skills/weight-name-alignment/SKILL.md create mode 100644 .agents/skills/writing-rewrite-rules/SKILL.md diff --git a/.agents/skills/adding-a-new-model/SKILL.md b/.agents/skills/adding-a-new-model/SKILL.md new file mode 100644 index 00000000..0992a2f3 --- /dev/null +++ b/.agents/skills/adding-a-new-model/SKILL.md @@ -0,0 +1,364 @@ +--- +name: adding-a-new-model +description: > + Use this skill when adding a new HuggingFace model architecture to + mobius — including LLM, encoder-only, encoder-decoder, vision, audio, + diffusion, or multimodal models. Covers the full workflow: config + extraction, model class creation, registry registration, weight + preprocessing, and testing. Also covers MoE and hybrid architectures. +--- + +# Skill: Adding a New Model + +## When to use + +Use this skill when adding support for a new HuggingFace model architecture +(e.g. a new LLM family, vision model, encoder-decoder, audio model, or +diffusion component) to the `mobius` package. + +## Reference files + +Read these when you need deeper detail on a specific topic: + +- Read [`references/architecture-patterns.md`](references/architecture-patterns.md) + when implementing a **non-LLM model** (encoder-only, encoder-decoder, vision, + audio, diffusion, multimodal), when dealing with **KV sharing across layers**, + or when checking **false compatibility pitfalls** for registry aliases. +- Read [`references/weight-preprocessing.md`](references/weight-preprocessing.md) + when handling **weight name mismatches**, **fused weight splitting**, + **precision/dtype issues**, or debugging **logit mismatches** traced to + weight loading, identity folding, or fp32 upcast problems. + +## Prerequisites + +- Identify the HuggingFace `model_type` string (from the model's `config.json`) +- Find a small checkpoint on HuggingFace Hub for testing +- Have the HuggingFace `transformers` source available to reference the + PyTorch implementation + +## Step-by-step + +### 1. Check if the base `CausalLMModel` already works + +Many models (LLaMA, Mistral, Qwen2, DeepSeek) use the standard decoder-only +architecture with no special components. Before writing a custom class, +check whether `CausalLMModel` from `models/base.py` produces correct results: + +```python +from mobius._registry import registry +from mobius.models.base import CausalLMModel + +registry.register("my_model_type", CausalLMModel) +model = build("org/my-model-id", load_weights=True) +``` + +If the logits match HuggingFace, you only need the registry entry. + +### 2. Create the model file + +Create `src/mobius/models/.py`. The minimal template: + +```python +from __future__ import annotations + +import torch +from onnxscript import nn + +from mobius._configs import ArchitectureConfig +from mobius.components import ( + Attention, DecoderLayer, Embedding, Linear, MLP, RMSNorm, + create_attention_bias, initialize_rope, +) +from mobius.models.base import CausalLMModel + + +class MyTextModel(nn.Module): + """Text model for MyArchitecture.""" + + def __init__(self, config: ArchitectureConfig): + super().__init__() + self.embed_tokens = Embedding(config.vocab_size, config.hidden_size) + self.layers = [MyDecoderLayer(config) for _ in range(config.num_hidden_layers)] + self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.rotary_emb = initialize_rope(config) + + def forward(self, op, input_ids, attention_mask, position_ids, past_key_values=None): + hidden_states = self.embed_tokens(op, input_ids) + position_embeddings = self.rotary_emb(op, position_ids) + attention_bias = create_attention_bias( + op, input_ids=input_ids, attention_mask=attention_mask, + ) + present_key_values = [] + past_kvs = past_key_values or [None] * len(self.layers) + for layer, past_kv in zip(self.layers, past_kvs): + hidden_states, present_kv = layer( + op, hidden_states=hidden_states, + attention_bias=attention_bias, + position_embeddings=position_embeddings, + past_key_value=past_kv, + ) + present_key_values.append(present_kv) + hidden_states = self.norm(op, hidden_states) + return hidden_states, present_key_values + + +class MyCausalLMModel(CausalLMModel): + """Causal LM wrapper for MyArchitecture.""" + + def __init__(self, config: ArchitectureConfig): + nn.Module.__init__(self) + self.config = config + self.model = MyTextModel(config) + self.lm_head = Linear(config.hidden_size, config.vocab_size, bias=False) +``` + +#### Class metadata attributes + +Every registered model class should set two class-level attributes: + +| Attribute | Type | Default | Description | +|-----------|------|---------|-------------| +| `default_task` | `str` | `"text-generation"` | Task auto-selected by `build()` / CLI. | +| `category` | `str` | `"Text Generation"` | Grouping label in generated docs. | + +Override these when the model isn't a standard text-generation model: + +```python +class MyMultiModalModel(nn.Module): + default_task: str = "vision-language" + category: str = "Multimodal" +``` + +Standard categories: `"Text Generation"`, `"Mixture of Experts"`, +`"Multimodal"`, `"Speech-to-Text"`, `"Audio"`, `"Diffusion"`, +`"autoencoder"`, `"encoder-only"`, `"encoder"`, `"encoder-decoder"`, +`"vision"`, `"causal-lm"`. + +### 3. Identify what's different + +Compare the HuggingFace PyTorch source against `CausalLMModel` / `DecoderLayer`. +Common variations to look for: + +| Variation | Example | Solution | +|-----------|---------|----------| +| Custom norm (weight + 1) | Gemma | Subclass `RMSNorm` | +| Embedding scaling | Gemma (`* sqrt(d)`) | Subclass `Embedding` | +| Extra norms (pre/post feedforward) | Gemma2, Gemma3 | Custom `DecoderLayer` | +| QK normalization | Gemma3, Qwen3 | Set `attn_qk_norm=True` in config | +| Sliding window attention | Gemma2, Gemma3 | Alternating layer types + `sliding_window` config | +| Different activation | Various | Set `hidden_act` in config (handled by `MLP`) | +| Biased attention projections | Phi, PhiMoE | Set `attn_qkv_bias=True`, `attn_o_bias=True` | +| LayerNorm epsilon | Whisper (`1e-5`) | Pass eps from config — default `1e-6` is wrong for many models | +| Custom attention scale | Granite | Pass `scale=config.attention_multiplier` to `Attention` | +| Embedding/logits/residual multipliers | Granite | Apply in `forward` — see troubleshooting §3 | +| MoE layers | PhiMoE, GPTOSS | See the **moe-models** skill | +| Vision encoder | Gemma3 | See the **multimodal-models** skill | +| Gated attention output | Qwen3.5 | Subclass `Attention` with doubled q_proj → Q+gate split | +| Hybrid layer types | Qwen3.5 | Use `config.layer_types` list to dispatch per-layer | +| Fused QKV / gate+up | ModernBERT | Split in `preprocess_weights` | +| Subclass-only (weight rename) | BLIP, TrOCR | Override only `preprocess_weights` | + +### 4. Handle weight name mismatches (`preprocess_weights`) + +If HuggingFace uses different weight names than your component tree, override +`preprocess_weights`: + +```python +class MyCausalLMModel(CausalLMModel): + def preprocess_weights(self, state_dict): + renamed = {} + for key, value in state_dict.items(): + new_key = key.replace("old_prefix.", "new_prefix.") + renamed[new_key] = value + return super().preprocess_weights(renamed) +``` + +> For detailed examples (fused QKV splitting, expert renames, weight-free +> norms, identity folding), read +> [`references/weight-preprocessing.md`](references/weight-preprocessing.md). + +### 5. Register the model + +Add to `_create_default_registry()` in `src/mobius/_registry.py`: + +```python +from mobius.models import MyCausalLMModel +reg.register("my_model_type", MyCausalLMModel) +``` + +Also export from `src/mobius/models/__init__.py`. + +### 6. Update `ArchitectureConfig.from_transformers` if needed + +If the model has unusual config fields, update `from_transformers()` in +`_configs.py`. Use safe defaults (1.0 for multipliers, None for optional +features) so existing models are unaffected. + +### 7. Write tests + +See the **writing-tests** skill for full details. At minimum: + +1. **Add a config entry to `tests/_test_configs.py`** in the appropriate + group (`CAUSAL_LM_CONFIGS`, `ENCODER_CONFIGS`, `SEQ2SEQ_CONFIGS`, + `VISION_CONFIGS`, or `DETECTION_CONFIGS`): + ```python + CAUSAL_LM_CONFIGS: list[tuple[str, dict, bool]] = [ + ("my_model_type", {"hidden_act": "gelu"}, True), + ] + ``` + - `is_representative=True` if the model has unique behaviour + - `is_representative=False` if it's an alias with no special config + + Then verify: `pytest tests/build_graph_test.py -k "my_model_type"` + +2. **Add a small model to `tests/integration_test.py`** if a small + checkpoint exists (< 1B parameters preferred). + +3. **Testing large models with random weights:** Create a reduced HF model: + ```python + c = AutoConfig.from_pretrained("Qwen/Qwen3.5-27B") + tc = c.text_config + tc.num_hidden_layers = 4 + hf_model = Qwen3_5ForCausalLM._from_config(tc, dtype=torch.float32) + ``` + Then use `build_from_module` with the HF state dict to compare logits. + +4. **Test with the CLI:** + ```bash + mobius build --model org/my-small-model mymodel/output + ``` + +### 8. Documentation + +Model documentation is **auto-generated** from class metadata by +`docs/_generate_models.py`. No manual doc update is needed if you set +`default_task`, `category`, and a good class docstring. + +## Checklist + +This is the **implementation** checklist. For the full **definition-of-done** +quality checklist (L1–L5 tests, ORT GenAI, Foundry Local, Olive, multi-dtype), +see the [quality-checklist skill](../quality-checklist/SKILL.md). + +- [ ] Model file in `src/mobius/models/` with Microsoft MIT copyright header +- [ ] Class has `default_task` and `category` attributes (if not standard text-generation) +- [ ] Class has a descriptive docstring (first paragraph used in generated docs) +- [ ] `preprocess_weights` handles any key mismatches +- [ ] Registered in `_create_default_registry()` +- [ ] Exported from `models/__init__.py` +- [ ] Config extraction works (`ArchitectureConfig.from_transformers`) +- [ ] Tiny config in `tests/_test_configs.py` (with `is_representative` flag) +- [ ] L2 YAML test case in `testdata/cases/` with `test_model_id` +- [ ] L3 synthetic parity passes (`tests/synthetic_parity_test.py -k ""`) +- [ ] Integration test in `tests/integration_test.py` (if small checkpoint available) +- [ ] L4 golden file generated and committed (`testdata/golden/`) +- [ ] L5 generation golden file generated and committed +- [ ] ORT GenAI test added to `tests/ort_genai_test.py` (text-generation and VLM models) +- [ ] CLI build works (`mobius build --model ...`) +- [ ] Multi-dtype correctness verified (fp32, fp16, bf16) + +**Note:** Default optimizer passes (CSE, deduplicate initializers, identity +elimination, remove unused nodes/opsets) are applied automatically. + +## Example: minimal diff for a LLaMA-compatible model + +If the new model is fully LLaMA-compatible, the entire change is: + +```python +# _registry.py +reg.register("my_llama_variant", CausalLMModel) +``` + +No new model file needed. + +## Example: adding a non-LLM model + +For non-LLM architectures, use the appropriate base class and task. +See [`references/architecture-patterns.md`](references/architecture-patterns.md) +for the full model type → base class → task mapping table. + +Quick reference: + +| Model type | Base class | Task | +|------------|-----------|------| +| Encoder-only (BERT-like) | `BertModel` | `feature-extraction` | +| Encoder-decoder (BART/T5) | `BartForConditionalGeneration` | `seq2seq` | +| Vision (ViT-like) | `ViTModel` | `image-classification` | +| Multimodal (LLaVA-like) | `LLaVAModel` | `vision-language` | +| Diffusion denoiser | Custom | `denoising` | + +Many models can be registered as aliases of existing classes if the +architecture matches. + +## Troubleshooting: common pitfalls + +### 1. ORT rejects graph with `tensor(double)` / wrong dtype + +NumPy creates float64 arrays by default. Always pass `dtype=np.float32`: +```python +# BAD — creates float64 +long_factor = np.array(config.rope_scaling["long_factor"]) +# GOOD +long_factor = np.array(config.rope_scaling["long_factor"], dtype=np.float32) +``` + +### 2. Normalization type mismatch (RMSNorm vs LayerNorm) + +**Symptom:** Large max abs diff (> 0.5) from HuggingFace. Check what norm +class HF actually uses — LayerNorm and RMSNorm are NOT interchangeable. +Key cases: OLMo-1B uses weight-free LayerNorm, Whisper uses eps=1e-5, +Gemma adds 1 to RMSNorm weight. + +### 3. Missing scaling multipliers + +**Symptom:** Generation diverges after a few tokens. Check for multiplier +config fields (`embedding_multiplier`, `attention_multiplier`, +`logits_scaling`, `residual_multiplier`). + +**Critical:** Residual scaling direction matters: +```python +# CORRECT: residual + output * multiplier +# WRONG: residual * multiplier + output +``` + +### 4. Config fields not extracted + +**Symptom:** Model builds but multipliers default to 1.0. Add extraction +to `ArchitectureConfig.from_transformers()` in `_configs.py` with safe +defaults. + +### 5. Debugging workflow for logit mismatches + +1. Check max abs diff on prefill (> 0.01 suspicious, > 0.5 is a bug) +2. Check the HuggingFace norm class and epsilon +3. Check for model-specific config fields (multiplier/scaling/factor/epsilon) +4. Check weight dtype (numpy arrays must be float32) +5. Compare layer by layer (moderate diff → norm/residual issue; huge → wrong weights) + +> For additional troubleshooting (gated attention split ordering, DeltaNet +> scaling, identity node folding, fp32 upcast patterns, multi-token prefill, +> embedding table off-by-one), read +> [`references/weight-preprocessing.md`](references/weight-preprocessing.md). + +## Reference examples + +| Complexity | File | Why | +|---|---|---| +| **Minimal** | `models/phi3.py` | Only overrides `preprocess_weights()` | +| **Minimal encoder** | `models/layoutlmv3.py` | Encoder subclass, weight rename only | +| **Moderate** | `models/gemma.py` | Custom attention, MLP, normalization | +| **Complex** | `models/qwen3_tts.py` | 4-model TTS architecture | + +> For the full reference implementation table (20+ models) and KV sharing +> patterns, read +> [`references/architecture-patterns.md`](references/architecture-patterns.md). + +## Cross-references + +- **[moe-models](../moe-models/SKILL.md)** — MoE gate variants, expert weight naming +- **[multimodal-models](../multimodal-models/SKILL.md)** — Vision encoders, projectors, VisionLanguageTask +- **[diffusion-models](../diffusion-models/SKILL.md)** — UNet, VAE, DiT, Flux, SD3 +- **[weight-name-alignment](../weight-name-alignment/SKILL.md)** — Aligning ONNX parameter names with HF +- **[writing-tests](../writing-tests/SKILL.md)** — Unit, integration, and generation test patterns +- **[quality-checklist](../quality-checklist/SKILL.md)** — L1–L5 definition of done +- **[reusable-components](../reusable-components/SKILL.md)** — Component library and design principles diff --git a/.agents/skills/adding-a-new-model/references/architecture-patterns.md b/.agents/skills/adding-a-new-model/references/architecture-patterns.md new file mode 100644 index 00000000..35163b66 --- /dev/null +++ b/.agents/skills/adding-a-new-model/references/architecture-patterns.md @@ -0,0 +1,259 @@ +# Architecture Patterns Reference + +Detailed code templates, compatibility rules, and advanced patterns for +non-standard model architectures. Read this when implementing a model that +is **not** a standard decoder-only causal LM. + +## Non-LLM model type table + +For models that aren't causal LMs, use the appropriate base class and task: + +| Model type | Base class / pattern | Task | Config | +|------------|---------------------|------|--------| +| Encoder-only (BERT-like) | `BertModel` | `feature-extraction` | `ArchitectureConfig` | +| Encoder-only (ModernBERT) | `ModernBertModel` | `feature-extraction` | `ArchitectureConfig` | +| Encoder-decoder (BART/T5-like) | `BartForConditionalGeneration` or `T5ForConditionalGeneration` | `seq2seq` | `ArchitectureConfig` | +| Vision (ViT-like) | `ViTModel` or `CLIPVisionModel` | `image-classification` | `ArchitectureConfig` | +| Object detection | `YolosForObjectDetection` | `object-detection` | `ArchitectureConfig` | +| Depth estimation | `DepthAnythingForDepthEstimation` | `image-classification` | `ArchitectureConfig` | +| Segmentation | `SegformerForSemanticSegmentation` or `Sam2VisionModel` | `image-classification` | `ArchitectureConfig` | +| Audio encoder (Wav2Vec2-like) | `Wav2Vec2Model` | `audio-feature-extraction` | `ArchitectureConfig` | +| Multimodal (LLaVA-like) | `LLaVAModel` | `vision-language` | `ArchitectureConfig` | +| Document AI | `LayoutLMv3Model` | `feature-extraction` | `ArchitectureConfig` | +| OCR decoder | `TrOCRForConditionalGeneration` | `seq2seq` | `ArchitectureConfig` | +| Diffusion denoiser | Custom (`UNet2DConditionModel`, etc.) | `denoising` | Custom config (e.g. `UNet2DConfig`) | +| VAE | `AutoencoderKLModel` | `vae` | `VAEConfig` | +| Adapter | `T2IAdapterModel` / `IPAdapterModel` | `adapter` | Custom config | + +Many new models can be registered as aliases of existing classes (e.g. +`reg.register("my_bert_variant", BertModel)`) if the architecture matches. + +## False Compatibility Pitfalls + +When registering models as aliases of existing base classes, **tests passing +does not mean the mapping is correct.** Graph-build tests only check that an +ONNX graph can be constructed — they do NOT verify that the graph matches +the model's actual computation. + +### Safe approximate mappings + +The project accepts "approximate" registry aliases when the model uses +similar-but-not-identical attention. These produce structurally correct ONNX +graphs; weight-loading may need minor adjustments: + +| Model | Maps to | Why it works | +|-------|---------|-------------| +| DeBERTa | `BertModel` | Disentangled attention is a variant of standard attention | +| Swin | `ViTModel` | Shifted window attention is still self-attention over patches | +| SqueezeBERT | `BertModel` | Grouped convolution replaces dense attention, but same I/O shape | + +### NEVER safe as registry aliases + +These model families have fundamentally different computation that **cannot** +be represented by standard base classes, even though `build_graph_test` passes: + +| Category | Models | Why it fails | +|----------|--------|-------------| +| Pure CNNs | ConvNeXt, ResNet, MobileNet, EfficientNet, RegNet | No attention at all — base ViT/BERT classes produce attention-based graphs | +| Spatial pooling | PoolFormer | Uses spatial average pooling instead of attention — structurally incompatible | +| SSM / state-space models | Mamba, Mamba2, FalconMamba, RWKV, RecurrentGemma | Sequential scan / linear recurrence, not attention | +| Fundamentally different attention | Longformer (sparse), BigBird (block sparse), Funnel (downsampling) | Attention pattern differs from dense self-attention at a structural level | +| Custom tokenization | CANINE (character-level) | Byte-level input, hash embeddings — not a standard vocab embedding | + +**Rule of thumb:** If the HuggingFace model's `forward()` method doesn't call +`self_attn(query, key, value)` in a standard way, it is NOT a safe alias. + +### Future work + +CI currently only runs graph-build tests (shape inference, op validity). To +catch false compatibility in approximate mappings, we need **weight-loading +tests** that: +1. Load real HuggingFace weights into the ONNX graph +2. Run inference on a test input +3. Compare output against HuggingFace PyTorch output +4. Fail if max abs diff exceeds a threshold (e.g. 0.01) + +This would catch shape mismatches, wrong norm types, and missing scaling +factors that graph-build tests cannot detect. + +## KV sharing across layers (num_kv_shared_layers) + +Some models (e.g. Gemma 4) reduce parameter count by having the last N +decoder layers **borrow** Key and Value states from an earlier "source" layer +of the same type instead of projecting their own K,V. This is controlled by +`num_kv_shared_layers` in the HuggingFace config. + +### What it means + +``` +first_kv_shared_idx = num_hidden_layers - num_kv_shared_layers + +Layers [0 .. first_kv_shared_idx - 1]: normal — own k_proj, v_proj, k_norm +Layers [first_kv_shared_idx .. end]: shared — NO k_proj/v_proj weights +``` + +Each shared layer reuses K,V from the **last non-shared layer of the same +attention type** (e.g. sliding vs. full attention). Only Q is computed fresh. + +### Impact on the checkpoint + +Shared layers have **no `k_proj`, `v_proj`, `k_norm`** keys in the +HuggingFace checkpoint. `preprocess_weights` must not assert these keys +exist for shared-layer indices — they simply won't be present. + +```python +def preprocess_weights(self, state_dict): + # shared layers have no k/v proj — remove them silently if accidentally present + first_shared = self.config.num_hidden_layers - self.config.num_kv_shared_layers + for i in range(first_shared, self.config.num_hidden_layers): + for suffix in ("k_proj.weight", "v_proj.weight", "k_norm.weight"): + state_dict.pop(f"model.layers.{i}.self_attn.{suffix}", None) + return super().preprocess_weights(state_dict) +``` + +### Attention module: is_kv_shared_layer flag + +The attention class detects at `__init__` time whether it is a shared layer: + +```python +class Gemma4Attention(nn.Module): + def __init__(self, config, layer_idx, layer_types, first_kv_shared_idx, ...): + self.is_kv_shared_layer = layer_idx >= first_kv_shared_idx > 0 + prev_layers = layer_types[:first_kv_shared_idx] + + if self.is_kv_shared_layer: + # Index of the source layer whose K,V this layer borrows + self.kv_shared_layer_index = ( + len(prev_layers) - 1 - prev_layers[::-1].index(layer_types[layer_idx]) + ) + self.store_full_length_kv = False + else: + self.kv_shared_layer_index = None + # True for the last non-shared layer of each type that has downstream + # KV-shared layers depending on it — it stores K,V for reuse. + self.store_full_length_kv = first_kv_shared_idx > 0 and ( + layer_idx + == len(prev_layers) - 1 - prev_layers[::-1].index(layer_types[layer_idx]) + ) + + # All layers have Q projection + self.q_proj = Linear(config.hidden_size, num_heads * head_dim) + self.q_norm = RMSNorm(head_dim) + self.o_proj = Linear(num_heads * head_dim, config.hidden_size) + + # Only non-shared layers have K/V projections + if not self.is_kv_shared_layer: + self.k_proj = Linear(config.hidden_size, num_kv_heads * head_dim) + self.v_proj = Linear(config.hidden_size, num_kv_heads * head_dim) + self.k_norm = RMSNorm(head_dim) +``` + +### forward(): shared layers consume shared_kv_states dict + +Pass a mutable `shared_kv_states` dict through the forward call. Source +layers populate it; shared layers read from it: + +```python +def forward(self, op, hidden_states, ..., shared_kv_states, past_key_value): + # Q projection (all layers) + query_states = self.q_proj(op, hidden_states) + ... + + if self.is_kv_shared_layer: + # Borrow K,V from source layer (already in shared_kv_states) + src_key, src_value = shared_kv_states[self.kv_shared_layer_index] + # Reshape from present_kv 4D [B, kv_heads, total_seq, head_dim] + # to Attention input 3D [B, total_seq, kv_heads * head_dim] + src_key = op.Transpose(src_key, perm=[0, 2, 1, 3]) + key_states = op.Reshape(src_key, ...) + value_states = ... + else: + # Normal K/V projection + norm + key_states = self.k_proj(op, hidden_states) + value_states = self.v_proj(op, hidden_states) + ... + + hidden_out, present_kv = _apply_attention(op, query_states, key_states, ...) + + if self.store_full_length_kv: + # Store present_kv [B, kv_heads, total_seq, head_dim] for downstream shared layers + shared_kv_states[self.layer_idx] = (present_kv_key, present_kv_value) + + return hidden_out, present_kv +``` + +### Text model: KV cache has only num_kv_layers entries + +KV-shared layers do **not** append to `present_key_values`. The output list +has `num_hidden_layers - num_kv_shared_layers` entries, not `num_hidden_layers`: + +```python +# In Gemma4TextModel.forward(): +shared_kv_states: dict = {} +present_key_values = [] + +# past_key_values has only num_kv_layers entries (no entry for KV-shared layers). +# Expand it to a full per-layer list so we can zip cleanly over all layers. +if past_key_values is not None: + kv_iter = iter(past_key_values) + past_kvs: list = [ + None if layer.self_attn.is_kv_shared_layer else next(kv_iter) + for layer in self.layers + ] +else: + past_kvs = [None] * len(self.layers) + +for i, (layer, layer_type, past_kv) in enumerate( + zip(self.layers, self.layer_types, past_kvs) +): + hidden_states, present_kv = layer( + op, + hidden_states=hidden_states, + attention_bias=attention_bias_dict[layer_type], + position_embeddings=position_embeddings_dict[layer_type], + shared_kv_states=shared_kv_states, + past_key_value=past_kv, + ) + # KV-shared layers borrow K,V — exclude from present_key_values so the + # output has exactly num_kv_layers (not num_hidden_layers) entries. + if not layer.self_attn.is_kv_shared_layer: + present_key_values.append(present_kv) +``` + +The task's KV cache inputs/outputs must use the correct count: +`num_kv_layers = config.num_hidden_layers - config.num_kv_shared_layers`. + +## Reference implementations + +| Model | File | Key differences from base | +|-------|------|--------------------------| +| Granite | `models/granite.py` | 4 scaling multipliers, custom attention scale | +| OLMo-1B | `models/olmo.py` | Weight-free LayerNorm (not RMSNorm), eps=1e-5 | +| OLMo-2 | `models/olmo.py` | Post-norm decoder layers, QK full norm | +| Gemma | `models/gemma.py` | RMSNorm weight+1, embedding scaling | +| Whisper | `components/_whisper.py` | Q pre-scaling, LayerNorm eps=1e-5, is_causal attr | +| Phi3.5 | `components/_rotary_embedding.py` | LongRope with float32 factors | +| Qwen3.5 | `models/qwen.py` | Hybrid DeltaNet + full attention, gated GQA, OffsetRMSNorm, interleaved MRoPE | +| Qwen3.5-MoE | `models/qwen.py` | Same hybrid attention + MoE FFN with shared expert (sigmoid gate) | +| Qwen3-TTS | `models/qwen3_tts.py` | 4-model TTS split, 2-token code predictor prefill, small_to_mtp projection, Identity-exposed weights | +| **BLIP** | `models/blip.py` | Subclass of ViTModel — only `preprocess_weights` (fused QKV split, renaming) | +| **YOLOS** | `models/yolos.py` | ViT + detection tokens + DETR-style MLP heads. New `object-detection` task | +| **Depth Anything** | `models/depth_anything.py` | ViT backbone + DPT decoder (reassemble + fusion + depth head). Uses `ConvTranspose2d` | +| **Segformer** | `models/segformer.py` | Hierarchical 4-stage encoder, efficient attention (strided Conv2d on K/V), Mix-FFN with depthwise conv | +| **SAM2** | `models/sam2.py` | Hiera backbone (per-stage dim transitions, fused QKV attention) + FPN neck with top-down fusion | +| **LayoutLMv3** | `models/layoutlmv3.py` | Subclass of BertModel — only `preprocess_weights` (spatial embedding filtering) | +| **TrOCR** | `models/trocr.py` | Subclass of BartForConditionalGeneration — only `preprocess_weights` (`output_projection` rename) | +| **ModernBERT** | `models/modernbert.py` | Pre-norm encoder with RoPE + GeGLU + bidirectional attention. Fused QKV/Wi splitting. Both encoder and decoder variants | +| **Gemma3n** | `models/gemma3n.py` | AltUp predict/correct, Laurel low-rank, per-layer input gating, hybrid local/global attention | +| **Mllama** | `models/mllama.py` | Interleaved cross-attention decoder, tanh-gated residual, manual QK-norm | + +## Reference examples by complexity + +When adding a new model, use these files as canonical references: + +| Complexity | File | Why | +|---|---|---| +| **Minimal** — base class works, only weight mapping needed | `models/phi3.py` (38 lines) | Extends `CausalLMModel`, only overrides `preprocess_weights()` to split fused QKV and gate-up projections. Shows the simplest possible model addition. | +| **Minimal** — encoder subclass | `models/layoutlmv3.py` | Extends `BertModel`, only overrides `preprocess_weights()`. Same pattern for encoder-only models. | +| **Moderate** — custom components | `models/gemma.py` | Adds custom attention (soft-capping), custom MLP (GeGLU), and custom normalization. Good example of component subclassing. | +| **Complex** — multi-model architecture | `models/qwen3_tts.py` | 4-model TTS split with talker, code predictor, embedding, and speaker encoder sub-modules. Shows how to structure multi-model architectures. | diff --git a/.agents/skills/adding-a-new-model/references/weight-preprocessing.md b/.agents/skills/adding-a-new-model/references/weight-preprocessing.md new file mode 100644 index 00000000..60f74032 --- /dev/null +++ b/.agents/skills/adding-a-new-model/references/weight-preprocessing.md @@ -0,0 +1,277 @@ +# Weight Preprocessing Reference + +Detailed `preprocess_weights` examples, troubleshooting for weight loading, +and advanced debugging patterns. Read this when you need to handle weight +name mismatches, fused weight splitting, or diagnose numerical issues traced +to weight loading or dtype problems. + +## Common preprocess_weights operations + +- **Strip prefixes:** `language_model.model.X` → `X` (multimodal models) +- **Rename expert weights:** `w1` → `gate_proj` (MoE models) +- **Weight tying:** copy `embed_tokens.weight` → `lm_head.weight` +- **Split fused QKV:** `Wqkv.weight` [3H, H] → `q_proj`, `k_proj`, `v_proj` (ModernBERT) +- **Split fused gate+up:** `Wi.weight` [2I, H] → `gate_proj`, `up_proj` (ModernBERT) +- **Split fused BLIP QKV:** `in_proj_weight` [3H, H] → separate Q/K/V projections + +## Debugging mismatches: initializer comparison + +Build the ONNX model, list its initializer names, and compare against +the HuggingFace state dict keys to find mismatches: + +```python +model_names = set(onnx_model.graph.initializers.keys()) +hf_names = set(state_dict.keys()) +print("In HF but not model:", hf_names - model_names) +print("In model but not HF:", model_names - hf_names) +``` + +## Attention scale override + +**Symptom:** Attention scores are wrong, causing gradual drift in generation. + +**Diagnosis:** Some models override the default `1/sqrt(head_dim)` scale: + +```python +# Check HuggingFace attention class +from transformers.models.granite.modeling_granite import GraniteAttention +print(GraniteAttention.__init__) # Look for self.scaling = ... +``` + +**Fix:** Use the `scale` parameter on the `Attention` component: + +```python +# In your custom DecoderLayer: +self.self_attn = Attention(config, scale=config.attention_multiplier) +``` + +The `Attention.__init__` accepts an optional `scale: float | None` parameter. +When `None`, it defaults to `head_dim**-0.5`. + +## Weight-free norms (no learnable parameters) + +**Symptom:** Weight keys like `model.norm.weight` appear in the model but not +in the HuggingFace state dict, and `preprocess_weights` fills them with ones. +The norm still uses the wrong algorithm (e.g. RMSNorm instead of LayerNorm). + +**Fix:** Use `_WeightFreeLayerNorm` from `models/olmo.py` which creates +`nn.Parameter` with constant data (ones for scale, zeros for bias) that +the ONNX `LayerNormalization` op requires, but the HuggingFace model has no +corresponding weights for: + +```python +class _WeightFreeLayerNorm(nn.Module): + def __init__(self, hidden_size: int, eps: float = 1e-5): + super().__init__() + self.scale = nn.Parameter( + [hidden_size], data=ir.tensor(np.ones(hidden_size, dtype=np.float32)) + ) + self.bias = nn.Parameter( + [hidden_size], data=ir.tensor(np.zeros(hidden_size, dtype=np.float32)) + ) + self.eps = eps + + def forward(self, op, hidden_states): + return op.LayerNormalization( + hidden_states, self.scale, self.bias, epsilon=self.eps, axis=-1 + ) +``` + +## Gated attention Q/gate split ordering + +**Symptom:** Large logit diff (~2.0) on Qwen3.5 or similar gated attention models. + +**Root cause:** HF splits Q and gate *within each head*: reshapes to +`[B, S, num_heads, 2*head_dim]` then chunks on last dim. A naive midpoint +split of the flat tensor gives wrong results. + +**Fix:** Reshape to per-head layout before splitting: +```python +# WRONG: split flat tensor at midpoint +q, gate = op.Split(qg_proj, num_outputs=2, axis=-1) + +# CORRECT: reshape to per-head, then split +qg = op.Reshape(qg_proj, [0, 0, num_heads, 2 * head_dim]) +q, gate = op.Split(qg, [head_dim, head_dim], axis=-1) +``` + +## DeltaNet missing query scaling + +**Symptom:** Linear attention output is orders of magnitude too large. + +**Root cause:** After L2-normalizing Q and K, you still need a +`1/sqrt(key_head_dim)` scaling factor on the query, similar to standard +attention. + +## Extracting `last_hidden_state` before vs after norm + +**Symptom:** Downstream model (e.g. code predictor, projection layer) +receives wrong hidden states. Prefill logits match HF exactly, but +generation diverges immediately. + +**Root cause:** HuggingFace's `outputs.last_hidden_state` is the +**post-norm** hidden state (after RMSNorm/LayerNorm). If you extract +the hidden state before the final norm, downstream consumers get +pre-norm values. In single-model LLMs this doesn't matter (the lm_head +is after the norm). In multi-model pipelines (TTS, VLM), the +hidden state is passed to another model, so norm ordering is critical. + +**Fix:** Always extract hidden state *after* the model's final norm: +```python +# WRONG: hidden_states before norm +hidden_states = decoder_output # pre-norm +logits = lm_head(norm(hidden_states)) # logits correct, but... +return logits, hidden_states # hidden_states is WRONG for downstream + +# CORRECT: apply norm first, then use for both logits and output +hidden_states = norm(decoder_output) # post-norm +logits = lm_head(hidden_states) +return logits, hidden_states # hidden_states matches HF +``` + +## Identity node folding renames initializers + +**Symptom:** Weight loading fails — `preprocess_weights` maps to the +original parameter name (e.g. `code_predictor.stacked_codec_embedding`) +but the initializer in the ONNX graph has been renamed to something +like `v_code_predictor.Identity_174`. + +**Root cause:** The IR optimizer folds `Identity(initializer)` by +removing the Identity node and renaming the initializer to the +output name. If you then set a custom name on the output (e.g. +`codec_embeddings.name = "codec_embeddings"`), the initializer gets +renamed to `codec_embeddings` — breaking weight loading. + +**Fix:** Use `op.Identity()` to create a *real* Identity node between +the initializer and the graph output. Ensure the elimination pass +retains Identity nodes that feed graph outputs. This creates a separate +output value, so renaming the output doesn't affect the initializer: +```python +# In forward(): +codec_embeddings = op.Identity(self.stacked_codec_embedding) +return logits, present_key_values, codec_embeddings + +# In task (safe to rename — Identity separates the names): +codec_embeddings.name = "codec_embeddings" +graph.outputs.append(codec_embeddings) +``` + +## `np.ascontiguousarray` promotes 0-d arrays to 1-d + +**Symptom:** ONNX `Gather` axis-reducing semantics break — the output +has an extra dimension (e.g. `(1, vocab)` instead of `(vocab,)`). + +**Root cause:** `np.ascontiguousarray(scalar_array)` promotes shape +`()` to `(1,)`. This changes `Gather(axis=0)` from axis-reducing +(scalar index) to axis-preserving (1-d index). + +**Fix:** Guard against 0-d arrays: +```python +if v.ndim > 0: + v = np.ascontiguousarray(v) +``` + +## Multi-token prefill in code predictors + +**Symptom:** Code predictor generates garbage. Prefill logits are +slightly off compared to HF. + +**Root cause:** Some architectures (e.g. Qwen3-TTS code predictor) use +a **2-token prefill**: `concat(projected_hidden, embed(code_0))` as two +separate tokens through the transformer. Summing them into 1 token +changes attention patterns and all subsequent hidden states. + +**Diagnosis:** Compare the inputs_embeds shape at step 0. If HF passes +`(batch, 2, hidden)` but your model uses `(batch, 1, hidden)`, the +attention context window is wrong. + +**Fix:** Construct inputs_embeds externally to match HF's exact flow: +```python +# Step 0 (prefill): 2 tokens +inputs = np.concatenate([talker_hidden, embed(code_0)], axis=1) # (1, 2, H) +# Steps 1+: 1 token +inputs = cp_embed[step-1, code_i, :].reshape(1, 1, -1) # (1, 1, H) +``` + +## Embedding table index off-by-one in multi-step generation + +**Symptom:** Codes are plausible but audio quality is wrong. Codec sum +doesn't match HF. + +**Root cause:** In multi-step code prediction, HF uses +`embed[step-1](code)` at generation step `step`, not `embed[step]`. +The off-by-one means every embedding lookup uses the wrong table. + +**Fix:** Carefully trace HF's generation loop to determine which +embedding table index corresponds to which generation step. Write a +comparison script that checks individual embedding lookups match. + +## Codec sum uses output codes, not input codes + +**Symptom:** codec_sum diverges from HF even though individual +embeddings weights are identical. + +**Root cause:** The codec sum `Σ embed[i](code_{i+1})` uses the +*generated* (output) code at each step, not the input code. If the +model returns embeddings of the input codes, the sum is wrong. + +**Fix:** Compute codec_sum externally using the codes actually +generated at each step: +```python +codec_sum = talker_embed(code_0) +for i in range(num_groups - 1): + # codes[i+1] is the OUTPUT of code predictor step i + codec_sum += cp_embed[i, codes[i + 1], :] +``` + +## Precision-sensitive ops need fp32 upcast + +**Symptom:** Type mismatch errors (`tensor(float) vs tensor(bfloat16)`) +when loading a model built with `--dtype bf16`, or numerical drift compared +to HuggingFace when running in fp16/bf16. + +**Root cause:** Operations like `exp`, `softplus`, `sigmoid` (in gated norms), +and RMSNorm variance are numerically sensitive and must run in float32 to +match HuggingFace, which explicitly upcasts with `.float()` / +`.to(torch.float32)`. + +**Two distinct problems:** + +1. **Naive `CastLike` everywhere** — keeps everything in the model dtype + (e.g. bf16), but `exp` overflows and the SSM state diverges. +2. **Naive `Cast(to=ir.DataType.FLOAT)` everywhere** — computes in fp32 but forgets to cast + back, producing type mismatches with downstream bf16 ops. + +**Correct pattern — upcast → compute → cast back:** +```python +# 1. Upcast to fp32 for the sensitive region +dt_f32 = op.Cast(dt, to=ir.DataType.FLOAT) +dt_f32 = op.Softplus(dt_f32) +a_neg = op.Neg(op.Exp(op.Cast(self.A_log, to=ir.DataType.FLOAT))) +da = op.Exp(op.Mul(dt_4d, a_4d)) # all fp32 here +... +# 2. Cast back to input dtype at the boundary +y = op.CastLike(y_f32, x) +new_state = op.CastLike(new_state_f32, ssm_state) +``` + +**How to identify which ops need fp32:** Check the HuggingFace source for +`.float()` or `.to(torch.float32)` calls. Each one marks an fp32 region +that the ONNX graph must replicate. + +**Known fp32-required regions:** + +| Region | HF evidence | ONNX pattern | +|--------|-------------|-------------| +| SSM recurrence (A, dt, exp, state) | `self.A_log.float()`, `hidden_states.float()`, `B.float()`, `C.float()` | `Cast(to=ir.DataType.FLOAT)` all inputs, `CastLike` output | +| GatedRMSNorm (SiLU + variance) | `hidden_states.to(torch.float32)`, `gate.to(torch.float32)` | Explicit fp32 for both, `CastLike` output | +| RMSNorm variance | `hidden_states.to(torch.float32)` | ONNX `RMSNormalization` handles via `stash_type=1` (default) | + +**When fp32 upcast is NOT needed:** +- Linear projections (`MatMul`) — runtime handles mixed precision +- SiLU on conv output — HF keeps in model dtype +- Standard attention — ONNX `Attention` op handles precision internally + +**Use `CastLike` for** parameters/constants that should match the *current* +compute dtype (which is fp32 inside an upcast region, or the model dtype +outside). Use `Cast(to=ir.DataType.FLOAT)` to explicitly enter an fp32 region. diff --git a/.agents/skills/debugging-vl-pipeline/SKILL.md b/.agents/skills/debugging-vl-pipeline/SKILL.md new file mode 100644 index 00000000..e41cc004 --- /dev/null +++ b/.agents/skills/debugging-vl-pipeline/SKILL.md @@ -0,0 +1,222 @@ +--- +name: debugging-vl-pipeline +description: > + Use this skill when debugging wrong, garbled, or divergent output from an + existing ORT GenAI multimodal pipeline (vision-language or vision+audio). + Covers the 3-stage pipeline isolation methodology, quick diagnostic flow, + 3D M-RoPE position ID issues, CUDA EP gotchas, and numerical tolerance + expectations. For building or adding a new multi-encoder model, use the + phi4mm-component-parity skill instead. +--- + +# Skill: Debugging VL Pipeline Issues + +> **Scope boundary:** This skill is for debugging runtime/inference issues +> in an **existing** ORT GenAI multimodal pipeline. If you are **building +> or adding** a new multi-encoder model and need component-by-component +> parity verification, use the `phi4mm-component-parity` skill instead. + +## When to use + +Use this skill when: + +- ORT GenAI produces wrong or irrelevant output for image inputs +- ONNX model logits diverge significantly from HuggingFace +- The model generates text-only descriptions ignoring the image +- Image features appear correct but decoder output is wrong +- Audio transcription is garbled or wrong despite correct encoder output +- CUDA EP crashes or produces different results than CPU + +## Quick diagnostic flow + +1. **Text-only works?** If yes → problem is in vision/audio path or + position IDs. If no → decoder or weight loading issue. +2. **Vision cos_sim > 0.99?** If no → vision encoder issue (see Stage 1). +3. **Embedding text positions match?** If no → token replacement bug. +4. **3D M-RoPE fields set?** Missing `image_token_id`, + `vision_start_token_id`, or `spatial_merge_size` causes 1D fallback. +5. **CUDA-only failure?** See CUDA EP section below. + +## Debugging methodology: isolate each stage + +VL models have 3 stages. Debug by isolating and validating each stage +independently, comparing against HuggingFace at every boundary. + +``` +pixel_values ──► [1. Vision] ──► image_features + │ +input_ids ──► [2. Embedding] ◄──────┘ + │ + ▼ + inputs_embeds + position_ids + attention_mask + │ + ▼ + [3. Decoder] ──► logits +``` + +### Stage 1: Vision model + +**What to check:** +- Output shape: `(num_patches, hidden_size)` +- Expected patches = `t * (h / merge) * (w / merge)` from `grid_thw` +- Compare features against HF vision encoder output + +```python +# HF reference +with torch.no_grad(): + hf_vision_out = hf_model.model.visual( + pixel_values, grid_thw=grid_thw + ) +# ONNX +session = OnnxModelSession(pkg["vision"]) +onnx_out = session.run({"pixel_values": pv, "grid_thw": grid_thw}) + +# Compare +cos_sim = np.dot(hf_flat, onnx_flat) / (norm_hf * norm_onnx) +print(f"Vision cos_sim: {cos_sim:.6f}") # Should be > 0.99 +``` + +**Common issues:** +- Wrong pixel value normalization (mean/std mismatch) +- `grid_thw` shape or values don't match HF processor output +- Missing `temporal_patch_size` in patch embedding + +### Stage 2: Embedding model + +**What to check:** +- Image features injected at correct token positions +- Non-image positions have correct text embeddings +- Output shape: `(1, seq_len, hidden_size)` + +```python +# Verify image token positions +image_mask = (input_ids[0] == image_token_id) # 151655 +num_image_positions = image_mask.sum() +assert num_image_positions == image_features.shape[0] + +# Compare embeddings at text positions (should match HF exactly) +text_mask = ~image_mask +cos_sim_text = cosine_similarity( + onnx_embeds[0, text_mask], hf_embeds[0, text_mask] +) +print(f"Text embedding cos_sim: {cos_sim_text:.6f}") # Should be 1.0 +``` + +**Common issues:** +- Image token count mismatch between processor and vision model +- Missing zero-padding row in embedding model (for text-only inputs) +- Wrong `image_token_id` used for Gather/Where mask + +### Stage 3: Decoder + +**What to check:** +- Logits shape matches HF: `(1, seq_len, vocab_size)` +- First token prediction matches HF (argmax of last position) +- Cosine similarity of logit vectors + +```python +# With HF-computed position_ids (ground truth): +onnx_logits = decoder_session.run(feeds)["logits"] +hf_logits = hf_model(**hf_inputs).logits.numpy() + +max_diff = np.abs(onnx_logits - hf_logits).max() +cos_sim = cosine_similarity(onnx_logits[0, -1], hf_logits[0, -1]) +print(f"max_diff={max_diff:.2f}, cos_sim={cos_sim:.4f}") +# Typical: max_diff=5-10, cos_sim>0.98 +``` + +**Common issues:** +- Wrong position_ids (see "3D M-RoPE" section below) +- Missing KV cache initialization +- Wrong attention_mask length + +## Critical: 3D M-RoPE position IDs + +Qwen2-VL / Qwen2.5-VL / Qwen3-VL use **3D Multimodal RoPE** where +`position_ids` has shape `(3, batch, seq_len)`: + +``` +position_ids[0] = temporal positions +position_ids[1] = height positions +position_ids[2] = width positions +``` + +**Text tokens:** all 3 dimensions have the same sequential value. + +**Image tokens:** temporal is constant, height/width vary over the +image grid `(h/merge, w/merge)`: +``` +temporal: [offset, offset, offset, ..., offset] +height: [offset, offset+1, offset+1, ..., offset+h/merge-1] +width: [offset, offset+1, offset, offset+1, ..., offset+w/merge-1] +``` + +**Text after image:** all 3 dimensions resume from +`max(temporal, height, width) + 1`. + +### ORT GenAI config requirements + +For ORT GenAI to compute 3D M-RoPE automatically, the following +`genai_config.json` fields are **required**: + +| Field | Level | Purpose | +|-------|-------|---------| +| `model.image_token_id` | model | Token ID for `<\|image_pad\|>` (e.g. 151655) | +| `model.vision_start_token_id` | model | Token ID for `<\|vision_start\|>` (e.g. 151652) | +| `model.vision.spatial_merge_size` | vision | Grid merge factor (typically 2) | + +**Without these fields**, ORT GenAI falls back to standard 1D positions, +which produces completely wrong output for image inputs (the model may +describe a "snowy landscape" instead of the actual image content). + +## CUDA EP gotchas + +### ORT Gather int32 overflow + +CUDA EP crashes or produces incorrect results for models with large +embedding tables (> 2^31 elements). CPU EP works correctly. +**Workaround:** Split large embeddings via `nn.ModuleList`. +ORT bug: microsoft/onnxruntime#28107 + +### Opset 24 kernel registration + +ORT ≤1.24.x CUDA/TRT EPs don't register kernels for opset 24. +**Fix:** Use the `ort_lower_opset_for_ep` feature flag (enabled by +default). See `src/mobius/_flags.py`. + +## Integration test patterns + +### Full VL forward test +```python +# Build model → process image with HF processor → run ONNX → compare logits +assert_logits_close(onnx_logits, hf_logits, rtol=2e-2, atol=2e-1) +``` + +### 3-model pipeline test +```python +# Vision → Embedding → Decoder, each stage uses OnnxModelSession +# Compare final decoder logits against HF single-model forward +``` + +### ORT GenAI end-to-end test +```python +# Build → save flat → write genai_config → load with ort_genai → generate +# Verify output length > input (basic sanity) +``` + +## Reference files + +- **Detailed failure modes (10):** `references/failure-modes.md` +- **Intermediate value extraction methods:** `references/extraction-methods.md` +- **Integration tests:** `tests/integration_test.py` + (`TestVLFullForward`, `TestQwen25VL3Model`) +- **ORT GenAI tests:** `tests/ort_genai_test.py` + (`TestOrtGenaiQwen25VL.test_multimodal_image_generation`) +- **Example scripts:** `examples/qwen25_vl_ort_genai.py`, + `examples/qwen3_vl_ort_genai.py`, `examples/gemma4_multimodal.py` +- **genai_config reference:** `.agents/skills/ort-genai-config/SKILL.md` +- **Component parity skill:** `.agents/skills/phi4mm-component-parity/SKILL.md` +- **ORT GenAI position_ids code (external):** + `onnxruntime-genai/src/models/position_inputs.cpp:617-814` +- **Feature flags:** `src/mobius/_flags.py` + (`ort_lower_opset_for_ep`, `ort_cuda_grouped_rmsnorm_workaround`) diff --git a/.agents/skills/debugging-vl-pipeline/references/extraction-methods.md b/.agents/skills/debugging-vl-pipeline/references/extraction-methods.md new file mode 100644 index 00000000..26dbe6fa --- /dev/null +++ b/.agents/skills/debugging-vl-pipeline/references/extraction-methods.md @@ -0,0 +1,144 @@ +# Extracting Intermediate ONNX Values + +Three methods for extracting intermediate values from ONNX models to +compare block-by-block against HuggingFace. Used when a pipeline stage +(e.g., vision encoder) diverges and you need to narrow down the root cause. + +--- + +## Method 1: Add intermediate outputs to the ONNX graph + +The most reliable approach — expose any internal node's output as a graph +output so ORT returns it alongside normal outputs. + +```python +import onnx + +model = onnx.load("vision.onnx") +graph = model.graph + +# Find the node whose output you want to inspect +for node in graph.node: + if node.op_type == "RMSNormalization" and "block_0" in node.output[0]: + # Name the output (if unnamed, give it a name) + target_output = node.output[0] + break + +# Add as a graph output +graph.output.append( + onnx.helper.make_tensor_value_info(target_output, onnx.TensorProto.FLOAT, None) +) +onnx.save(model, "vision_debug.onnx") + +# Now ORT will return this value alongside image_features +session = ort.InferenceSession("vision_debug.onnx") +results = session.run(None, feeds) +# results[-1] is the intermediate value +``` + +## Method 2: Use `ir.Model` graph manipulation (preferred for mobius) + +When working with `ir.Model` objects from the build pipeline, manipulate +the graph directly without saving/loading: + +```python +from mobius._testing.ort_inference import OnnxModelSession + +pkg = build(model_id, dtype="f32", load_weights=True) +vision_model = pkg["vision"] +graph = vision_model.graph + +# Find target nodes by op type or name pattern +target_nodes = [n for n in graph if n.op_type == "RMSNormalization"] + +# RMSNorm input nodes are natural block boundaries: +# rms_nodes[0] = block 0 norm1 (input = patch_embed output) +# rms_nodes[2] = block 1 norm1 (input = block 0 output) +# rms_nodes[2*i] = block i norm1 (input = block i-1 output) +block_0_output = target_nodes[2].inputs[0] # block 0's output +block_0_output.name = "block_0_output" +graph.outputs.append(block_0_output) + +session = OnnxModelSession(vision_model) +out = session.run({"pixel_values": pv, "grid_thw": grid_thw}) +block_0_out = out["block_0_output"] +session.close() +``` + +## Method 3: Hook HuggingFace model for reference values + +Use PyTorch hooks to extract intermediate values from HuggingFace at +the same points: + +```python +intermediates = {} + +def hook_fn(name): + def fn(module, input, output): + if isinstance(output, tuple): + intermediates[name] = output[0].detach().cpu().numpy() + else: + intermediates[name] = output.detach().cpu().numpy() + return fn + +# Register hooks on specific blocks +for i, block in enumerate(hf_model.model.visual.blocks): + block.register_forward_hook(hook_fn(f"block_{i}")) + +# Run forward pass — hooks capture all intermediate values +with torch.no_grad(): + hf_out = hf_model.model.visual(pixel_values, grid_thw=grid_thw) + +# Now compare block by block +for i in range(num_blocks): + hf_block_out = intermediates[f"block_{i}"] + cos = cosine_similarity(onnx_block_out, hf_block_out) + print(f"Block {i}: cos={cos:.6f}") +``` + +## Block-by-block comparison strategy + +When overall output diverges, narrow down by comparing each transformer +block's output sequentially: + +```python +for i in range(num_blocks): + onnx_out_i = extract_onnx_block_output(vision_model, i, feeds) + hf_out_i = intermediates[f"block_{i}"] + + cos = cosine_similarity(onnx_out_i.flatten(), hf_out_i.flatten()) + max_diff = np.max(np.abs(onnx_out_i - hf_out_i)) + print(f"Block {i:2d}: cos={cos:.6f} max_diff={max_diff:.4f}") +``` + +Typical pattern for a bug in block N: +``` +Block 0: cos=1.000000 max_diff=0.0001 ← perfect +Block 1: cos=1.000000 max_diff=0.0001 ← perfect +... +Block N: cos=0.961000 max_diff=4.8500 ← divergence starts! +Block N+1: cos=0.892000 max_diff=25.00 ← error compounds +``` + +Once you identify the divergent block, drill deeper into that block's +sub-operations (attention, MLP, normalization) to find the root cause. + +## Comparing specific weight values + +To verify weights loaded correctly, compare ONNX initializers against +HuggingFace state dict: + +```python +from safetensors import safe_open + +# Load HF weights +with safe_open(safetensors_path, framework="numpy") as f: + hf_weight = f.get_tensor("visual.blocks.0.attn.qkv.weight") + +# Get ONNX weight from ir.Model +onnx_weight = vision_model.graph.initializers["blocks.0.attn.qkv.weight"] +onnx_np = onnx_weight.const_value.numpy() + +max_diff = np.max(np.abs(hf_weight.astype(np.float32) - onnx_np)) +print(f"Weight diff: {max_diff}") # Should be 0.0 +``` diff --git a/.agents/skills/debugging-vl-pipeline/references/failure-modes.md b/.agents/skills/debugging-vl-pipeline/references/failure-modes.md new file mode 100644 index 00000000..cae059ed --- /dev/null +++ b/.agents/skills/debugging-vl-pipeline/references/failure-modes.md @@ -0,0 +1,206 @@ +# Common Failure Modes and Fixes + +Detailed failure mode reference for debugging existing ORT GenAI multimodal +pipelines. For the high-level debugging methodology, see the parent +`SKILL.md`. + +--- + +## 1. Image not recognized (wrong output for image inputs) + +**Symptoms:** Model produces generic or hallucinated descriptions that +don't match the input image. Text-only generation works correctly. + +**Root causes (in order of likelihood):** + +1. **Missing genai_config fields** — `image_token_id`, + `vision_start_token_id`, or `spatial_merge_size` not set. + Without these, position_ids are 1D instead of 3D M-RoPE. + +2. **Image resize mismatch** — ORT processor resizes image to different + dimensions than HF processor, producing different number of vision + tokens. ORT's `width`/`height` in `processor_config.json` are used + as direct resize targets, unlike HF's smart_resize which computes + target from original image dimensions. + +3. **Processor config format** — ORT GenAI expects ort-extensions format + `processor_config.json`, not HuggingFace format. The file must include + `DecodeImage`, `ConvertRGB`, `Resize`, `Rescale`, `Normalize`, and + `PatchImage` transforms with correct attributes. + +**Fix for resize mismatch:** +```python +def _update_resize_for_image(processor_config_path, image_path): + """Recompute resize dimensions from actual image like HF does.""" + from PIL import Image + img = Image.open(image_path) + w, h = img.size + factor = 14 * 2 # patch_size * merge_size + new_w = round(w / factor) * factor + new_h = round(h / factor) * factor + # Update width/height in processor_config.json +``` + +## 2. Numerical divergence in greedy decoding + +**Symptoms:** First 1-3 tokens match HF, then output diverges. + +**Expected behavior:** This is inherent to ONNX vs PyTorch numerical +differences. ONNX models use different operator implementations that +accumulate small floating-point errors. + +**Typical metrics for Qwen2.5-VL 3B:** +- max_diff in logits: 5-10 +- mean_diff in logits: 0.5-1.5 +- cosine similarity: 0.98-0.99 +- First token: matches HF +- Greedy decoding: diverges at token 3-5 + +**This is NOT a bug** if the metrics above are within range. Both models +produce semantically similar descriptions. + +## 3. Vision model output shape mismatch + +**Symptoms:** Vision model produces wrong number of patches. + +**Debug:** Check `grid_thw` values: +```python +# For Qwen2.5-VL with merge_size=2: +t, h, w = grid_thw[0] +expected_patches = t * (h // 2) * (w // 2) +actual_patches = vision_output.shape[0] +assert expected_patches == actual_patches +``` + +## 4. Embedding model text-only failure + +**Symptoms:** Error when running without images (num_image_tokens=0). + +**Fix:** Ensure embedding model pads `image_features` with a zero row +before Gather, then uses a Where mask to select only real features: +```python +# Pad with zero row so Gather with index 0 doesn't fail +padded = op.Concat( + op.ConstantOfShape(...), # (1, hidden_size) zeros + image_features, + axis=0, +) +``` + +## 5. Vision encoder internal divergence (cos < 0.5) + +**Symptoms:** Vision features have very low cosine similarity (< 0.5) +against HuggingFace, even though patch embedding and weights are correct. + +**Debug with block-by-block comparison** (see extraction methods in +`references/extraction-methods.md`). Common root causes: + +1. **Wrong rotary embedding dimension** — Qwen2.5-VL vision uses 2D + position encoding (height + width). The rotary dim must be + `head_dim // 2`, not `head_dim`. Each half (head_dim // 4 frequencies) + covers one spatial dimension. With full `head_dim`, you get 2× too + many frequencies with wrong values. **Result: cos ≈ 0.25.** + +2. **Missing `fullatt_block_indexes`** — Qwen2.5-VL alternates between + windowed attention (local windows of `window_size` patches) and full + attention (all patches attend to all). Blocks at indexes `[7, 15, 23, 31]` + use full attention. If `fullatt_block_indexes` is not extracted from + HF config, all blocks use windowed attention. **Result: blocks 0-6 + are perfect (they're windowed anyway), but block 7+ diverges.** + +3. **Wrong attention bias construction** — Full-attention blocks should + have an all-zeros bias (everything attends to everything). Windowed + blocks have a block-diagonal bias. Check the bias by inspecting + sparsity: `(bias == -inf).float().mean()` should be ~0% for full + attention, ~98% for windowed. + +**Config extraction checklist for vision encoders:** +```python +# These fields MUST be extracted from HF vision_config: +fullatt_block_indexes = getattr(vc, "fullatt_block_indexes", None) +window_size = getattr(vc, "window_size", None) +spatial_merge_size = getattr(vc, "spatial_merge_size", 2) +temporal_patch_size = getattr(vc, "temporal_patch_size", 2) +``` + +## 6. ClippableLinear divergence (Gemma4) + +**Symptoms:** Vision or audio encoder output has large max diff (> 1.0) +against HuggingFace, even though weights are loaded correctly. + +**Root cause:** Gemma4 uses `Gemma4ClippableLinear` with learned finite +input/output activation clamping for ALL linear layers in its vision and +audio encoders. Using plain `Linear` misses the clamping. + +**Detection:** Check if the HuggingFace model uses `ClippableLinear`: +```bash +grep -n "ClippableLinear" transformers/models//modeling_.py +``` + +**Fix:** Use `ClippableLinear` (from `mobius.components`) for all +affected linear layers. For vision attention: q/k/v/o_proj. For MLP: +pass `linear_class=ClippableLinear` to the MLP component. + +**Impact:** +- Audio: max diff 52.68 → 0.0003 after fix +- Vision: max diff 3.92 → 0.00007 after fix + +## 7. Missing audio boundary markers + +**Symptoms:** Audio transcription is garbled or completely wrong, even +though the audio encoder output matches HuggingFace. + +**Root cause:** HuggingFace wraps audio placeholder tokens with boundary +markers in `input_ids`: +``` +<|audio> (256000) + N × <|audio|> (258881) + (258883) +``` +If boundary markers are missing, the model cannot distinguish audio +regions from text, producing wrong output. + +**Fix:** Add boundary markers to `build_input_ids()` when constructing +audio inputs, matching HuggingFace's token wrapping. + +## 8. CUDA EP: ORT Gather int32 overflow + +**Symptoms:** CUDA EP crashes or produces incorrect results for models +with large embedding tables. CPU EP works correctly. + +**Root cause:** ORT CUDA `gather_impl.cu` uses `int32` for element +offset computation: `input_index = idx * cols + col_offset`. For +tensors with > 2^31 elements (e.g. Gemma4 per-layer embedding +[262144, 8960] = 2.35B elements), this overflows. + +**ORT bug:** microsoft/onnxruntime#28107 + +**Workaround:** Split large embeddings into smaller tables via +`nn.ModuleList` so each individual Gather stays under the int32 limit. +Use Slice instead of Gather for column-wise indexing on large tensors. + +## 9. CUDA EP: opset 24 kernel registration + +**Symptoms:** ORT CUDA EP fails to find kernels for standard ops +(Squeeze, Reshape, etc.) even though they work on CPU. + +**Root cause:** ORT ≤1.24.x CUDA/TRT EPs don't register kernels for +opset 24, even though the op semantics are unchanged from opset 23. + +**Fix:** Use the `ort_lower_opset_for_ep` feature flag (enabled by +default) which lowers the declared opset import from 24 to 23 for +non-CPU EPs. See `src/mobius/_flags.py` and +`src/mobius/_testing/ort_inference.py`. + +## 10. Wrong audio feature extractor + +**Symptoms:** Audio model produces completely wrong output. Audio +encoder features don't match HuggingFace at all. + +**Root cause:** Using `WhisperFeatureExtractor` instead of the +model-specific feature extractor (e.g. `Gemma4AudioFeatureExtractor`). +Different extractors produce different mel spectrograms. + +**Fix:** Always use the correct feature extractor for the model: +```python +from transformers import AutoFeatureExtractor +feature_extractor = AutoFeatureExtractor.from_pretrained(model_id) +``` diff --git a/.agents/skills/diffusion-models/SKILL.md b/.agents/skills/diffusion-models/SKILL.md new file mode 100644 index 00000000..6a03d10c --- /dev/null +++ b/.agents/skills/diffusion-models/SKILL.md @@ -0,0 +1,389 @@ +--- +name: diffusion-models +description: > + Use this skill when adding or modifying a diffusion or image-generation + model in mobius. Covers UNet, VAE, DiT, Flux, SD3, ControlNet, adapters, + and QwenImage architectures. Includes pipeline detection from diffusers + configs, DiffusionTask and DiffusionModel classes, building blocks + (timestep embeddings, cross-attention, downsampling/upsampling), and + weight loading conventions for diffusers checkpoints. +--- + +# Skill: Diffusion Models + +## When to use + +Use this skill when: +- Adding a new diffusion model (denoiser, VAE, ControlNet, adapter) +- Working with `build_diffusers_pipeline()` or the `_DIFFUSERS_CLASS_MAP` +- Creating diffusers config classes or task types +- Debugging diffusers weight loading or pipeline detection + +## Architecture overview + +Diffusion models generate images/video by iteratively denoising latent +representations. The key components are: + +| Component | Role | Examples | +|-----------|------|----------| +| **Denoiser** | Predicts noise to remove at each step | UNet2D, DiT, Flux, SD3, QwenImage | +| **VAE** | Encodes images to latent / decodes latent to images | AutoencoderKL, QwenImage 3D VAE, Video VAE | +| **ControlNet** | Provides spatial conditioning (edges, depth, etc.) | ControlNetModel | +| **Adapter** | Adds image/conditioning features to denoiser | T2I-Adapter, IP-Adapter | + +## Pipeline detection and building + +### How it works + +When the CLI or `build()` encounters a model that isn't a transformers model, +it checks for `model_index.json` (the diffusers pipeline descriptor): + +```python +# In _diffusers_builder.py +# 1. Try transformers AutoConfig → if fails: +# 2. Try loading model_index.json → if found: +# 3. Parse components and build each via _DIFFUSERS_CLASS_MAP +``` + +### Registering a new diffusers model + +Add an entry to `_init_diffusers_class_map()` in `_diffusers_builder.py`: + +```python +def _init_diffusers_class_map(): + _DIFFUSERS_CLASS_MAP["MyTransformer2DModel"] = ( + MyTransformer2DModel, # Module class + MyConfig, # Config dataclass + "denoising", # Task name + ) +``` + +The key must match the class name in `model_index.json`: +```json +{ + "transformer": ["diffusers", "MyTransformer2DModel"], + "vae": ["diffusers", "AutoencoderKL"] +} +``` + +### Weight loading for diffusers + +Diffusers uses different weight file naming than transformers: + +| Convention | Filename | +|------------|----------| +| Single file | `diffusion_pytorch_model.safetensors` | +| Sharded index | `diffusion_pytorch_model.safetensors.index.json` | +| Fallback | `model.safetensors` (some models) | + +Weights live in component subdirectories: `transformer/`, `vae/`, etc. +The `_download_diffusers_component_weights()` function tries both naming +conventions. + +## Config classes + +Diffusers configs are separate from `ArchitectureConfig`. Each is a +`@dataclasses.dataclass` with a `from_diffusers(config: dict)` classmethod +that parses the JSON config. + +```python +@dataclasses.dataclass +class MyDiffuserConfig: + in_channels: int = 4 + out_channels: int = 4 + block_out_channels: tuple[int, ...] = (320, 640, 1280, 1280) + num_layers: int = 2 + attention_head_dim: int = 8 + + @classmethod + def from_diffusers(cls, config: dict) -> "MyDiffuserConfig": + return cls( + in_channels=config.get("in_channels", cls.in_channels), + out_channels=config.get("out_channels", cls.out_channels), + block_out_channels=tuple(config.get("block_out_channels", cls.block_out_channels)), + num_layers=config.get("num_layers", cls.num_layers), + attention_head_dim=config.get("attention_head_dim", cls.attention_head_dim), + ) +``` + +Existing config classes: +- `VAEConfig` — Standard SD VAE (2D) +- `UNet2DConfig` — UNet-based denoisers +- `DiTConfig` — Diffusion Transformer +- `SD3Config` — Stable Diffusion 3 / MMDiT +- `FluxConfig` — Flux transformer +- `ControlNetConfig` — ControlNet conditioning +- `T2IAdapterConfig` / `IPAdapterConfig` — Adapters +- `QwenImageConfig` — QwenImage transformer +- `QwenImageVAEConfig` — QwenImage 3D causal VAE +- `VideoVAEConfig` — 3D video VAE + +All defined in `src/mobius/_diffusers_configs.py`. + +## Task classes + +Each diffusion model type has a task class that defines I/O signatures: + +### DenoisingTask + +```python +# Inputs: +# sample: [B, in_channels, H, W] — noisy latent +# timestep: [B] — diffusion timestep +# encoder_hidden_states: [B, seq, dim] — text conditioning +# Output: +# noise_pred: [B, in_channels, H, W] — predicted noise +``` + +### VAETask + +Returns `ModelPackage` with two sub-models: +- `encoder`: `sample [B, 3, H, W]` → `latent_dist [B, 2*latent_ch, H/f, W/f]` +- `decoder`: `latent [B, latent_ch, H/f, W/f]` → `sample [B, 3, H, W]` + +### QwenImageVAETask + +3D causal VAE for video/images: +- `encoder`: `[B, 3, T, H, W]` → `[B, 2*z_dim, T', H', W']` +- `decoder`: `[B, z_dim, T', H', W']` → `[B, 3, T, H, W]` + +### ControlNetTask + +```python +# Inputs: sample, timestep, encoder_hidden_states, controlnet_cond +# Outputs: down_outputs (list of residuals), mid_output +``` + +### AdapterTask + +```python +# T2I: condition [B, in_channels, H, W] → feature_list +# IP: image_embeds [B, image_dim] → adapter_output [B, num_tokens, cross_dim] +``` + +## Building blocks + +### Shared components (from `components/`) + +Diffusion models import shared primitives from the component library: + +| Component | Import | Purpose | +|-----------|--------|---------| +| `Conv2d` | `components/_conv.py` | 2D convolution with bias | +| `GroupNorm` | `components/_common.py` | Group normalization | +| `LayerNormNoAffine` | `components/_common.py` | Norm without learnable params (AdaLN) | +| `Linear` | `components/_common.py` | Linear projection | +| `SiLU` | `components/_activations.py` | SiLU activation module | + +Model files import these with underscore-prefixed aliases: + +```python +from mobius.components import Conv2d as _Conv2d, GroupNorm as _GroupNorm, SiLU as _SiLU +``` + +### Common model-specific blocks + +| Block | Purpose | Used by | +|-------|---------|---------| +| `_TimestepEmbedding` | Sinusoidal → MLP time embedding | UNet, DiT, Flux, SD3, QwenImage | +| `nn.ModuleList` | Layer lists, skip connections | All | +| `nn.Sequential` | Callable container that chains forward calls | Diffusion `to_out`, `img_mod`, `net` | + +> **`nn.Sequential` vs `nn.ModuleList`**: `nn.Sequential` chains children's +> `forward()` calls automatically — prefer it for sequential containers like +> `to_out` and modulation layers. Use `nn.ModuleList` for layer lists that +> need custom iteration (e.g., `down_blocks`, `transformer_blocks`). + +### UNet-specific + +| Block | Purpose | +|-------|---------| +| `_ResNetBlock2DWithTime` | ResNet block with time embedding injection | +| `_CrossAttentionBlock` | Self-attention + cross-attention + FFN | +| `_BasicAttention` | Scaled dot-product attention (spatial) | +| `_DownBlock2D` | ResNet + optional attn + spatial downsample | +| `_UpBlock2D` | ResNet + optional attn + spatial upsample + skip concat | +| `_UNetMidBlock2DCrossAttn` | Mid block: ResNet → cross-attn → ResNet | + +### DiT / Transformer denoiser patterns + +| Block | Purpose | +|-------|---------| +| `_PatchEmbed` | Conv2d with stride=patch_size (patchify input) | +| `_AdaLayerNormZero` | Adaptive LayerNorm with scale/shift/gate from timestep | +| `_DiTBlock` | AdaLN-Zero → self-attn + cross-attn + FFN | +| `_JointAttentionBlock` | Concatenate text+image → joint attention → split | +| `_FluxSingleBlock` | Unified self-attention (single-stream) | + +### VAE-specific + +| Block | Purpose | +|-------|---------| +| `_ResNetBlock2D` | GroupNorm → SiLU → Conv → skip | +| `_AttentionBlock` | Self-attention in latent space | +| `_Downsample2D` / `_Upsample2D` | Spatial resampling | +| `_MidBlock2D` | ResNet + optional attention | + +### 3D / Video-specific + +| Block | Purpose | +|-------|---------| +| `_CausalConv3d` | 3D conv with temporal-causal padding | +| `_RMSNorm3d` | Channel-wise RMS normalization for 3D | +| `_ResidualBlock` (3D) | 3D ResNet with RMSNorm | +| `_Resample` | 2D spatial + optional 3D temporal resampling | + +## Denoiser architecture comparison + +| Model | Stream | Attention | Modulation | Normalization | +|-------|--------|-----------|-----------|---------------| +| **UNet2D** | Single | Cross-attn only | Time embed inject | GroupNorm | +| **DiT** | Single | Self + cross | AdaLN-Zero (6 params) | LayerNorm | +| **SD3 (MMDiT)** | Double | Joint (concat text+img) | AdaLN-Zero | LayerNorm | +| **Flux** | Double→Single | Joint then unified | AdaLN-Zero | LayerNorm | +| **QwenImage** | Double | Joint (separate proj) | AdaLN (2×6 params) | RMSNorm | + +### Double-stream pattern + +SD3, Flux, and QwenImage use a double-stream architecture: +1. **Image stream**: modulation → attention → FFN → residual +2. **Text stream**: separate modulation → attention → FFN → residual +3. **Joint attention**: concatenate Q/K/V from both streams, attend together + +```python +# In _JointAttentionBlock.forward(): +img_q, img_k, img_v = self.to_q(op, img), self.to_k(op, img), self.to_v(op, img) +txt_q, txt_k, txt_v = self.add_q(op, txt), self.add_k(op, txt), self.add_v(op, txt) +q = op.Concat(img_q, txt_q, axis=1) +k = op.Concat(img_k, txt_k, axis=1) +v = op.Concat(img_v, txt_v, axis=1) +attn_out = op.Attention(q, k, v, ...) +img_out, txt_out = op.Split(attn_out, axis=1, ...) +``` + +### AdaLN-Zero modulation + +Most transformer denoisers use AdaLN-Zero: project timestep embedding to +6 parameters (shift, scale, gate for attention and FFN): + +```python +# Modulation: timestep → SiLU → Linear → split into 6 chunks +mod = self.mod(op, temb) # [B, 6 * dim] +shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = op.Split( + mod, num_outputs=6, axis=-1 +) + +# Apply: norm → scale/shift → attention → gate → residual +normed = self.norm(op, x) +modulated = op.Add(op.Mul(normed, op.Add(scale_msa, ONE)), shift_msa) +attn_out = self.attn(op, modulated) +x = op.Add(x, op.Mul(attn_out, gate_msa)) +``` + +## Adding a new diffusion model — step by step + +### 1. Create config class + +In `_diffusers_configs.py`: + +```python +@dataclasses.dataclass +class MyDenoiserConfig: + # Parse from HF diffusers config.json + in_channels: int = 4 + ... + + @classmethod + def from_diffusers(cls, config: dict) -> "MyDenoiserConfig": + return cls(...) +``` + +### 2. Create model class + +In `models/my_model.py`: + +```python +class MyDenoiser2DModel(nn.Module): + default_task = "denoising" + category = "Diffusion" + config_class = MyDenoiserConfig + + def __init__(self, config: MyDenoiserConfig): + super().__init__() + # Build architecture matching HF naming + ... + + def forward(self, op, sample, timestep, encoder_hidden_states): + # Denoising forward pass + ... + + def preprocess_weights(self, state_dict): + # Ideally a no-op if naming matches HF + return state_dict +``` + +### 3. Register in _DIFFUSERS_CLASS_MAP + +In `_diffusers_builder.py`, add to `_init_diffusers_class_map()`: + +```python +_DIFFUSERS_CLASS_MAP["MyDenoiser2DModel"] = ( + MyDenoiser2DModel, MyDenoiserConfig, "denoising" +) +``` + +### 4. Add task class (if needed) + +If the existing `DenoisingTask` doesn't match the I/O signature, create a +new task in `tasks/`. Most denoisers use the standard `DenoisingTask`. + +### 5. Add unit test + +In `tests/build_graph_test.py`, add a tiny config: + +```python +("my_denoiser", MyDenoiser2DModel, MyDenoiserConfig( + in_channels=4, out_channels=4, ... # Tiny values +), "denoising"), +``` + +### 6. Match HF weight names + +Compare `named_parameters()` output with HF weight names. Use the +techniques from the **weight-name-alignment** skill to minimize +`preprocess_weights`. + +## Naming conventions for diffusers models + +Diffusers uses different conventions than transformers: + +| Diffusers | Transformers | +|-----------|-------------| +| `GroupNorm` | `LayerNorm` / `RMSNorm` | +| `ResNet` blocks | Decoder layers | +| `down_blocks` / `up_blocks` | `layers` | +| `resnets` / `attentions` (within blocks) | `self_attn` / `mlp` | +| `to_q` / `to_k` / `to_v` / `to_out` | `q_proj` / `k_proj` / `v_proj` / `o_proj` | +| `conv_in` / `conv_out` | `embed_tokens` / `lm_head` | +| `time_embedding` | N/A | +| `encoder_hid_proj` | N/A | + +## Reference files + +| File | Contains | +|------|----------| +| `_diffusers_configs.py` | All diffusers config dataclasses | +| `_diffusers_builder.py` | Pipeline detection, `_DIFFUSERS_CLASS_MAP`, build functions | +| `models/unet.py` | UNet2DConditionModel + all UNet blocks | +| `models/vae.py` | AutoencoderKLModel + encoder/decoder blocks | +| `models/dit.py` | DiTTransformer2DModel + AdaLN blocks | +| `models/flux_sd3.py` | FluxTransformer2DModel + SD3Transformer2DModel | +| `models/controlnet.py` | ControlNetModel | +| `models/adapters.py` | T2IAdapterModel + IPAdapterModel | +| `models/qwen_image.py` | QwenImageTransformer2DModel | +| `models/qwen_image_vae.py` | AutoencoderKLQwenImageModel (3D causal VAE) | +| `models/video_vae.py` | VideoAutoencoderModel (3D video VAE) | +| `tasks/_denoising.py` | DenoisingTask | +| `tasks/_vae.py` | VAETask, QwenImageVAETask | +| `tasks/_controlnet.py` | ControlNetTask | +| `tasks/_adapter.py` | AdapterTask | diff --git a/.agents/skills/moe-models/SKILL.md b/.agents/skills/moe-models/SKILL.md new file mode 100644 index 00000000..662cb582 --- /dev/null +++ b/.agents/skills/moe-models/SKILL.md @@ -0,0 +1,384 @@ +--- +name: moe-models +description: > + Use this skill when adding or modifying a model that uses Mixture-of-Experts + (MoE) layers. Covers gate variants (TopKGate, SparseMixerGate), MoELayer + composition and expert routing, expert weight naming conventions for + HuggingFace alignment, and preprocess_weights mappings for stacked + expert tensors. Applicable to models like Mixtral, DeepSeek, and Qwen-MoE. +--- + +# Skill: Mixture-of-Experts (MoE) Models + +## When to use + +Use this skill when adding or modifying a model that uses Mixture-of-Experts +layers — where each token is routed to a subset of expert MLPs. + +## Architecture overview + +``` +MoEDecoderLayer + ├── RMSNorm (input_layernorm) + ├── Attention (self_attn) + ├── RMSNorm (post_attention_layernorm) + └── MoELayer + ├── Gate (routing: input → expert selection + weights) + └── Experts[0..N-1] (each is a standard MLP) +``` + +### Key components + +| Component | File | Purpose | +|-----------|------|---------| +| `MoELayer` | `components/_moe.py` | Routes tokens to experts, combines outputs | +| `TopKGate` | `components/_moe.py` | Standard softmax + top-k routing | +| `SparseMixerGate` | `components/_moe.py` | Sequential selection with threshold masking | +| `MoEDecoderLayer` | `models/moe.py` | Decoder layer that uses `MoELayer` instead of `MLP` | +| `MoETextModel` | `models/moe.py` | Text model with MoE decoder layers | + +## How routing gates work + +### TopKGate (default) + +Standard routing used by most MoE models (Mixtral, GPTOSS): + +1. Compute router logits: `logits = MatMul(hidden_states, gate_weight)` +2. Top-k selection: `values, indices = TopK(logits, k=num_experts_per_tok)` +3. Softmax over selected experts: `weights = Softmax(values)` + +### SparseMixerGate (PhiMoE) + +Sequential expert selection with threshold-based masking: + +1. Compute router logits via `MatMul` +2. For each of `top_k` rounds: + - Find max score (`ReduceMax`) + - Threshold mask: experts whose scores are far from the max (relative to + `jitter_eps`) are masked with `-inf` + - Softmax over non-masked experts + - `TopK(k=1)` to select best expert + - `ScatterElements` to mask out the selected expert for the next round +3. Concatenate all selected expert indices and weights + +## Adding a new MoE model + +### 1. Determine the gate type + +Check the HuggingFace implementation for the routing logic. Look for: + +- `router_type` or `routing_type` in the config +- How `router_logits` are computed and processed +- Whether top-k is applied before or after softmax + +If neither `TopKGate` nor `SparseMixerGate` fits, create a new gate class +in `components/_moe.py`. + +### 2. Check expert MLP naming + +HuggingFace MoE models often use different weight names for expert MLPs: + +| HF name | Our name | Description | +|---------|----------|-------------| +| `w1` | `gate_proj.weight` | Gate projection | +| `w2` | `down_proj.weight` | Down projection | +| `w3` | `up_proj.weight` | Up projection | + +Implement a `_rename_moe_expert_weights()` function if the naming differs: + +```python +def _rename_moe_expert_weights(state_dict): + renamed = {} + for key, value in state_dict.items(): + new_key = key + if ".experts." in key: + new_key = new_key.replace(".w1.", ".gate_proj.") + new_key = new_key.replace(".w2.", ".down_proj.") + new_key = new_key.replace(".w3.", ".up_proj.") + renamed[new_key] = value + return renamed +``` + +### 3. Check the normalization + +Some MoE models use different norms than standard models: + +- **PhiMoE**: Uses `LayerNorm` (with bias), not `RMSNorm` +- **Mixtral**: Uses standard `RMSNorm` + +### 4. Create the model class + +Use `MoETextModel` with the correct gate factory: + +```python +class MyMoECausalLMModel(CausalLMModel): + def __init__(self, config): + nn.Module.__init__(self) + self.config = config + # Pass gate_factory for custom routing + self.model = MoETextModel(config, gate_factory=SparseMixerGate) + self.lm_head = Linear(config.hidden_size, config.vocab_size, bias=True) + + def preprocess_weights(self, state_dict): + state_dict = _rename_moe_expert_weights(state_dict) + return super().preprocess_weights(state_dict) +``` + +### 5. Inject a custom gate + +`MoELayer` accepts an optional `gate` parameter. `MoETextModel` accepts +`gate_factory` — a callable `(config) -> gate_instance` that creates a gate +per layer: + +```python +# Default (TopKGate) +MoETextModel(config) + +# Custom gate +MoETextModel(config, gate_factory=SparseMixerGate) +``` + +To create a new gate, implement a class with this interface: + +```python +class MyGate(nn.Module): + def __init__(self, config): + super().__init__() + self.weight = nn.Parameter([config.num_local_experts, config.hidden_size]) + # ... other params + + def forward(self, op, hidden_states): + # Returns: (expert_weights, expert_indices) + # expert_weights: [batch, seq, num_experts_per_tok] + # expert_indices: [batch, seq, num_experts_per_tok] (INT64) + ... + return weights, indices +``` + +## Config fields for MoE + +```python +config.num_local_experts # Total number of experts (e.g. 8, 16) +config.num_experts_per_tok # Experts activated per token (e.g. 2) +``` + +These are extracted automatically from HuggingFace configs. + +## TopK ONNX op gotcha + +The ONNX `TopK` op requires `K` as a **1-D int64 tensor**, not a Python int: + +```python +# WRONG +op.TopK(logits, self.top_k, axis=-1) + +# CORRECT +k_tensor = op.Constant(value_ints=[self.top_k]) +op.TopK(logits, k_tensor, axis=-1) +``` + +## Qwen3.5-MoE: Hybrid Attention + MoE + +Qwen3.5-MoE combines the hybrid DeltaNet/full-attention architecture of +Qwen3.5 dense with MoE FFN layers instead of a dense MLP. + +### Architecture + +``` +Qwen35MoEDecoderLayer + ├── OffsetRMSNorm (input_layernorm) + ├── GatedDeltaNet / Qwen35Attention (per layer_types) + ├── OffsetRMSNorm (post_attention_layernorm) + └── Qwen35MoEBlock + ├── TopKGate (router) + ├── Experts[0..N-1] (standard MLP: gate/up/down with SiLU) + ├── SharedExpert (MLP: gate/up/down with SiLU) + └── shared_expert_gate → sigmoid scalar +``` + +### Classes + +| Class | File | Purpose | +|-------|------|---------| +| `Qwen35MoEBlock` | `models/qwen.py` | MoE block with routed + shared experts | +| `Qwen35MoEDecoderLayer` | `models/qwen.py` | Hybrid attention + MoE FFN layer | +| `Qwen35MoETextModel` | `models/qwen.py` | Stacks decoder layers with RoPE | +| `Qwen35MoECausalLMModel` | `models/qwen.py` | Top-level causal LM model | + +### MoE block (`Qwen35MoEBlock`) + +- **TopKGate routing**: 256 experts, top-8 in the full model (configurable + via `num_local_experts` / `num_experts_per_tok`) +- **Expert MLPs**: Standard `MLP` (gate/up/down projections, SiLU activation), + each with `moe_intermediate_size` as the intermediate dim +- **Shared expert**: A separate `MLP` that runs on **all** tokens (not routed), + sized by `shared_expert_intermediate_size` +- **Shared expert gating**: `sigmoid(shared_expert_gate(x)) * shared_expert(x)`, + where `shared_expert_gate` is `Linear(hidden_size, 1, bias=False)` + +The key difference from standard MoE is the shared expert: its output is +gated by a learned sigmoid scalar and added to the routed expert output. + +### Config fields + +```python +config.moe_intermediate_size # Intermediate size per expert MLP +config.shared_expert_intermediate_size # Intermediate size for the shared expert +config.num_local_experts # Total number of routed experts (e.g. 256) +config.num_experts_per_tok # Experts activated per token (e.g. 8) +config.layer_types # Per-layer attention type list +``` + +### Weight naming + +HuggingFace weights map directly (no renames needed for expert names): + +``` +mlp.gate.weight → router logits +mlp.experts.N.{gate,up,down}_proj.weight → per-expert MLP +mlp.shared_expert.{gate,up,down}_proj.weight → shared expert MLP +mlp.shared_expert_gate.weight → sigmoid gate (Linear, no bias) +``` + +Note: HF checkpoints store experts as fused tensors +(`experts.gate_up_proj`, `experts.down_proj`). `preprocess_weights()` +unpacks these into per-expert tensors and also renames +`linear_attn.conv1d.weight` → `linear_attn.conv1d_weight`. + +### Testing + +Integration tests use a random-weight HF model with reduced layers and +experts (e.g. 4 layers, 4 experts, top-2 routing). See +`test_qwen35_moe_prefill_logits_match` in `tests/integration_test.py`. + +## Testing MoE models + +MoE integration tests require a model with MoE layers. A good test model +should be small enough for CI (~1-4B params). The test pattern: + +1. Build ONNX model with weights +2. Run prefill + decode against HuggingFace reference +3. Optionally test greedy generation (token-ID matching) + +See `tests/moe_integration_test.py` for the complete pattern. + +## Direct MoE op emission (com.microsoft.MoE) + +OnnxRuntime ships a fused `com.microsoft.MoE` contrib op (CUDA float32/fp16/bf16, +CPU float32/fp16). For new model architectures, **emit it directly** — like +`com.microsoft.GroupQueryAttention` — rather than relying on a rewrite rule. + +### When to use + +Use `com.microsoft.MoE` when: +- The model uses top-k MoE routing (standard softmax gate → TopK) +- All expert weights are the same shape (no dynamic expert counts) +- The EP's `caps.supports_fused_moe` is `True` + +Fall back to the loop-over-experts path when `supports_fused_moe` is `False` +(CPU EP without contrib ops, or EPs that don't support the custom op). + +### Gate output: full pre-topk router_probs + +The op takes the **full** `(num_tokens, num_experts)` probability tensor and +performs top-k selection internally via the `k` attribute. The gate must +produce the full softmax distribution — not already-selected top-k indices. + +```python +# In your gate forward(), return shape [num_tokens, num_experts] +router_probs = op.Softmax(op.MatMul(hidden_states, self.weight), axis=-1) +``` + +### Emission pattern (from Gemma 4 implementation) + +```python +from mobius._build_context import ep_capabilities + +caps = ep_capabilities() +if caps.supports_fused_moe: + moe_out = op.CastLike( + op.MoE( # type: ignore[attr-defined] + normed_hidden, # [num_tokens, hidden_size] + router_probs, # [num_tokens, num_experts] — full pre-topk + self.fc1_experts_weights, # [E, inter_size, hidden_size] + self.fc2_experts_weights, # [E, hidden_size, inter_size] + activation_type="silu", + k=self._top_k, + normalize_routing_weights=1, + _domain="com.microsoft", + ), + normed_hidden, # CastLike: preserve bf16/fp16/fp32 — NOT hardcoded float32 + ) +else: + moe_out = self._dispatch_moe_fallback(op, normed_hidden, router_probs) +``` + +**Critical: use `CastLike` after the MoE op.** The `com.microsoft.MoE` custom +op has `type=None` on its output — ONNX type propagation cannot infer the +output dtype. `op.CastLike(moe_output, target=input)` restores the correct +dtype (bf16/fp16/fp32), which allows downstream ops to share scalar +initializers and avoids hard-coded `Cast` to float32. + +### preprocess_weights: expert weight mapping + +HuggingFace Gemma 4 stores experts as a 3D tensor per projection: +`layers.N.experts.gate_up_proj [E, 2*inter, H]` +`layers.N.experts.down_proj [E, H, inter]` + +Map these to the parameter names used by the ONNX MoE op: + +```python +def preprocess_weights(self, state_dict): + for key in list(state_dict.keys()): + if ".experts.gate_up_proj" in key: + new_key = key.replace(".experts.gate_up_proj", ".fc1_experts_weights") + state_dict[new_key] = state_dict.pop(key) + elif ".experts.down_proj" in key: + new_key = key.replace(".experts.down_proj", ".fc2_experts_weights") + state_dict[new_key] = state_dict.pop(key) + return super().preprocess_weights(state_dict) +``` + +For models that store per-expert weights separately (one matrix per expert), +stack them into 3D tensors in `preprocess_weights`: + +```python +n = config.num_local_experts +gate = torch.stack([state_dict.pop(f"experts.{i}.gate_proj.weight") for i in range(n)]) +down = torch.stack([state_dict.pop(f"experts.{i}.down_proj.weight") for i in range(n)]) +state_dict["fc1_experts_weights"] = gate # [E, inter, hidden] +state_dict["fc2_experts_weights"] = down # [E, hidden, inter] +``` + +### EP capability check (matches GQA pattern) + +```python +# In _execution_providers.py EpCapabilities: +supports_fused_moe: bool = True # set False for EPs without com.microsoft.MoE + +# In model forward(): +from mobius._execution_providers import ep_capabilities +caps = ep_capabilities() +if caps.supports_fused_moe: + # emit com.microsoft.MoE +else: + # fallback loop +``` + +### Fallback: TopKGate + loop dispatch + +When `supports_fused_moe` is False, implement a static unroll: + +```python +def _dispatch_moe_fallback(self, op, hidden, router_probs): + k_tensor = op.Constant(value_ints=[self._top_k]) + top_weights, top_indices = op.TopK(router_probs, k_tensor, axis=-1) + top_weights = op.Softmax(top_weights, axis=-1) # renormalize + output = op.CastLike(op.ConstantOfShape(op.Shape(hidden), value=0.0), hidden) + for e_idx in range(self._num_experts): + w1 = op.Squeeze(op.Gather(self.fc1_experts_weights, [e_idx], axis=0), [0]) + w2 = op.Squeeze(op.Gather(self.fc2_experts_weights, [e_idx], axis=0), [0]) + # expert output, gated by routing weight + ... + return output +``` diff --git a/.agents/skills/multi-agent-coordination/SKILL.md b/.agents/skills/multi-agent-coordination/SKILL.md new file mode 100644 index 00000000..eb13d385 --- /dev/null +++ b/.agents/skills/multi-agent-coordination/SKILL.md @@ -0,0 +1,215 @@ +--- +name: multi-agent-coordination +description: > + Use this skill when managing parallel workstreams across multiple AI agents + on a shared Git repository. Covers worktree isolation to prevent merge + conflicts, commit coordination protocols, verification gates between + dependent tasks, context management for stateless agents, and failure + recovery patterns. Based on real experience coordinating 10–17 agents + simultaneously. +--- + +# Multi-Agent Coordination + +Practical guide for a project lead coordinating 10–17 AI agents working on a single codebase simultaneously. Each section documents a real failure mode encountered in practice and the pattern that prevents it. + +--- + +## 1. Worktree Isolation (CRITICAL) + +**Problem**: Multiple agents sharing a single working directory caused file contamination. One agent's uncommitted changes were overwritten by another checking out a different branch. Work was lost and had to be redone. + +**Solution**: Each agent gets its own git worktree. + +```bash +# Setup — run this before delegating work to an agent +git worktree add .worktrees/- 2>/dev/null || true +cd .worktrees/- +git fetch origin && git checkout && git pull +``` + +**Rules**: +- NEVER let two agents work in the same directory. +- Even "read-only" exploration can cause issues if an agent runs tests that generate artifacts or cache files. +- The main working directory (`/home/.../repo`) is reserved for the lead. Agents always use their worktree. +- Worktrees live at `.worktrees/-` — predictable names make auditing easy. + +**Cleanup**: Remove worktrees after agents are done with `git worktree remove .worktrees/-`. Do NOT clean up while agents are still active — you will destroy their workspace. + +--- + +## 2. Commit Coordination Protocol + +Multiple agents committing to the same branch requires discipline to avoid divergence. + +**Rules**: +1. **Pull before commit**: Always `git pull origin --rebase` immediately before committing. +2. **Push immediately after commit**: Don't let commits sit unpushed. Other agents may be waiting on your changes or about to conflict with yours. +3. **Atomic commits**: Each commit must be self-contained and leave tests passing. Never commit half-finished work. +4. **Descriptive commit messages**: Include what was fixed and the relevant test outcome. Other agents (and the lead) need to understand the diff at a glance. +5. **Conflict resolution**: If `pull --rebase` fails, the agent should report the conflict back to the lead rather than force-pushing or making arbitrary merge decisions. + +**Pattern for each agent commit**: +```bash +git pull origin --rebase # sync first +git add # never git add -A (picks up others' work) +git commit -m "fix: ..." +git push origin # push immediately +``` + +--- + +## 3. Verification Gates + +Checks at each stage prevent problems from compounding across agents. + +### Before starting work +- Verify the branch exists and is up to date. +- Run the relevant tests to establish a baseline — know what was already failing before you touched anything. +- Check for uncommitted changes left by previous agents in your worktree. + +### After each commit +- Run the affected tests immediately (don't batch; catch regressions early). +- Verify the commit landed on the correct branch (`git log --oneline -3`). +- Confirm the push succeeded. + +### Before final review (lead audit) +- Check ALL worktrees for unpushed commits or uncommitted changes. +- Verify the total commit count matches expectations. +- Run the full test suite from a clean checkout (not a worktree). +- Confirm no worktree has diverged from the remote branch. + +### Before PR creation +- Triple review: code correctness, critical logic, readability — ideally by three separate agents. +- All review findings addressed. +- Final test run passes. +- PR description updated with an accurate scorecard of what changed. + +--- + +## 4. Task Dependency Management (DAG) + +Track task dependencies as a directed acyclic graph (DAG). + +**State machine**: `pending → ready → running → done` + +A task becomes `ready` only when ALL its dependencies are `done`. + +**Dependency rules**: +- Review tasks depend on ALL implementation tasks (can't review code that isn't written). +- PR creation depends on ALL reviews passing. +- Integration tests depend on implementation being complete. +- Exploration/investigation tasks have no dependencies — launch them first. + +**Anti-pattern**: Circular dependencies. Reviews depend on implementation, not vice versa. If a review finds a bug, create a new implementation task — don't loop the dependency graph. + +--- + +## 5. Parallel Execution Patterns + +### Safe to parallelize +- Independent model/file fixes (different files, no shared state) +- Code review + readability review + critical review (all read-only) +- Investigation and exploration tasks +- Golden data generation for different models +- Any tasks touching completely separate files + +### Must be serialized +- Two agents editing the same file +- Commit + push sequences within the same agent (not cross-agent) +- Tasks with explicit DAG dependencies (review after implementation) +- Branch cleanup before final review + +### Optimal parallelism +4–6 agents working simultaneously is the sweet spot. Beyond that, coordination overhead increases and merge conflicts become more likely. Scale back if you see frequent push failures or agents blocking each other. + +--- + +## 6. Agent Role Specialization + +Match the agent role to the task type. Mismatched roles waste agent capacity. + +| Role | Best For | Avoid | +|------|----------|-------| +| **Architect** | Investigation, analysis, design decisions, understanding existing code | Simple mechanical fixes | +| **Developer** | Implementation, bug fixes, refactoring | Open-ended exploration | +| **QA Tester** | Verification, auditing, smoke tests, checking other agents' work | New implementation | +| **Code Reviewer** | Correctness, logic bugs, API contracts | Style opinions | +| **Critical Reviewer** | Security, edge cases, invariant violations | Routine fixes | +| **Readability Reviewer** | Naming, clarity, documentation | Implementation work | + +**Delegation tip**: Developers should receive clear instructions — specific files to change, specific test commands to run, and a concrete definition of "done." Architects are better suited for ambiguous investigation tasks. + +--- + +## 7. Communication Patterns + +**Agent → Lead**: Status updates after significant steps, completion summaries, blocker notifications. Don't wait until fully done — send progress updates so the lead can coordinate. + +**Lead → Agent**: Task delegation with full context (worktree path, setup commands, what to change, test commands, commit message format, definition of done). + +**Agent → Agent**: Coordinate through the lead. Agents should not directly orchestrate other agents unless explicitly set up as a sub-lead. + +**Anti-patterns**: +- Sending new tasks to a context-saturated agent. If they don't respond correctly, use a different agent. +- Vague task prompts ("fix the MLP stuff"). Be specific about files, methods, and expected outcomes. +- Assuming agents share context. Each agent starts fresh — always include relevant background in the task prompt. + +--- + +## 8. Failure Recovery + +| Failure | Recovery | +|---------|----------| +| **Lost work** | Check worktrees, `git stash list`, `git reflog`. Most work is recoverable. | +| **Broken tests** | Have the responsible agent fix it immediately. Don't leave broken tests for later. | +| **Merge conflict** | Have the agent that caused the conflict resolve it — they have the most context. | +| **Agent stuck/looping** | Terminate and create a fresh agent. Don't try to unstick a saturated agent. | +| **Wrong branch** | Cherry-pick commits to the correct branch (`git cherry-pick `). Don't redo work. | +| **Agent pushed to wrong branch** | Revert the wrong-branch commit, cherry-pick to the right branch. Coordinate with lead before force-pushing. | +| **Diverged worktree** | `git fetch origin && git rebase origin/` from the worktree. | + +--- + +## 9. Single-Branch Strategy + +For large multi-agent efforts, use **one shared feature branch** rather than one branch per agent. + +**Why**: Avoids complex multi-branch merge scenarios at the end. All agents commit to the same branch via their isolated worktrees. + +**Trade-offs**: +- More `pull --rebase` cycles per agent. +- Occasional push conflicts (recoverable with rebase). +- BUT: much simpler final state, single PR, linear history. + +**Alternative**: Separate branches per agent merged via separate PRs — only viable when features are truly independent and don't share files. + +--- + +## Checklist: Delegating a Task to an Agent + +Before sending a task, verify your prompt includes all of the following: + +1. ☐ **Worktree path** — where the agent should work (`.worktrees/-`) +2. ☐ **Setup commands** — `conda activate`, `git fetch`, `git checkout`, `git pull` +3. ☐ **What to change and why** — specific files, functions, or test cases +4. ☐ **Test commands** — what to run after changes to verify correctness +5. ☐ **Commit message format** — so the history is readable +6. ☐ **Push reminder** — explicit instruction to push after committing +7. ☐ **Definition of done** — what "finished" looks like (tests passing, specific output, etc.) + +**Example delegation prompt**: +``` +Work in .worktrees/developer-. Setup: + git fetch origin && git checkout justinchu/fix-model-parity && git pull + +Fix the Phi-2 MLP weight mismatch: replace the gated MLP (gate_proj/up_proj/down_proj) +with FCMLP (fc1/fc2) in src/mobius/models/phi.py. The HF weights use fc1/fc2. + +After changes, run: + python -m pytest tests/build_graph_test.py -k "phi" -q + +Commit message: "fix: use FCMLP for Phi-1/2 model (HF uses fc1/fc2 not gated MLP)" +Push to justinchu/fix-model-parity immediately after committing. +Done = tests pass, push confirmed. +``` diff --git a/.agents/skills/phi4mm-component-parity/SKILL.md b/.agents/skills/phi4mm-component-parity/SKILL.md index 172e25c2..50624db5 100644 --- a/.agents/skills/phi4mm-component-parity/SKILL.md +++ b/.agents/skills/phi4mm-component-parity/SKILL.md @@ -1,12 +1,14 @@ --- name: phi4mm-component-parity description: > - Debug multimodal ONNX model output that diverges from HuggingFace. - Use when isolating which component (vision encoder, speech encoder, - embedding, or decoder) causes numerical divergence; when integration - tests fail with large differences; or when adding a new multimodal - model and verifying each stage independently. Applicable to Phi4MM, - Gemma4, and any multi-encoder architecture. + Use this skill when building or adding a new multi-encoder multimodal + model (vision + language + audio) and verifying component-by-component + parity with HuggingFace. Covers pipeline isolation methodology, common + failure modes from real debugging experience (Phi4MM, Gemma4), the + step-by-step process for isolating which component (vision encoder, + speech encoder, embedding, or decoder) causes numerical divergence, + and integration test patterns. For debugging an existing deployed ORT + GenAI pipeline, use the debugging-vl-pipeline skill instead. --- # Skill: Multimodal Component Parity Debugging diff --git a/.agents/skills/quality-checklist/SKILL.md b/.agents/skills/quality-checklist/SKILL.md new file mode 100644 index 00000000..284ce6ff --- /dev/null +++ b/.agents/skills/quality-checklist/SKILL.md @@ -0,0 +1,244 @@ +--- +name: quality-checklist +description: > + Use this skill when verifying that a new model is truly done and ready to + merge. Provides a Definition-of-Done checklist covering all five test + confidence levels (L1 graph build through L5 Foundry Local smoke-test), + ORT GenAI runtime validation, Olive quantization compatibility, multi-dtype + (f32/f16/bf16) and multi-EP (CPU/CUDA/DML) correctness, documentation + requirements, and code review criteria. +--- + +# Skill: Quality Checklist + +## When to use + +Use this checklist before marking a new model addition as **done**. +Every item must be checked — or explicitly waived with a written reason — +before the PR is merged. + +--- + +## The Checklist + +### 1. Code quality + +- [ ] Model file is in `src/mobius/models/` with the Microsoft MIT copyright + header (`# Copyright (c) Microsoft Corporation. / # Licensed under the + MIT License.`) +- [ ] Class has a descriptive one-paragraph docstring (first paragraph is + used in generated docs) +- [ ] Class has `default_task` and `category` class-level attributes if + the model is not a standard text-generation model +- [ ] `preprocess_weights()` correctly maps every HuggingFace state-dict key + to the ONNX initializer name (verified by the weight-alignment test) +- [ ] All new components use `from mobius.components import ...` (public API), + not private submodule paths +- [ ] No explicit protobuf operations anywhere in new code + (`onnx.helper`, `onnx.TensorProto`, etc. are forbidden) +- [ ] Tensor shapes are annotated in comments after non-trivial operations +- [ ] Automated code review (Copilot/PR review) has been run and all + findings are resolved or explicitly dismissed with a reason +- [ ] `lintrunner -a` is clean — zero lint errors before merging + (`lintrunner f --output oneline --all-files` to auto-fix, then re-run to confirm) + +### 2. L1 — Graph builds + +- [ ] Entry exists in `tests/_test_configs.py` (or a dedicated test method + for VLM / audio models) +- [ ] `is_representative=True` if the model has unique behaviour (custom + class, special attention, MoE, hybrid layers, etc.) +- [ ] `python -m pytest tests/build_graph_test.py -k ""` passes +- [ ] Weight-alignment test passes: + `python -m pytest tests/weight_alignment_test.py -k ""` + +### 3. L2 — Config compatible + +- [ ] YAML test case created at `testdata/cases//.yaml` +- [ ] `test_model_id` field set to a real HuggingFace model ID +- [ ] Schema validates: `python -m pytest tests/yaml_schema_test.py` + +### 4. L3 — Synthetic parity + +- [ ] Model type is covered in `tests/synthetic_parity_test.py` (driven by + `_test_configs.py`; added automatically for text-generation models) +- [ ] `python -m pytest tests/synthetic_parity_test.py -k ""` passes + with `atol=1e-3` / `rtol=1e-3` (or `1e-2` for multimodal) +- [ ] Real-weight parity also checked via `tests/integration_test.py` (add + model to `_TEXT_MODELS` or equivalent if a small checkpoint is available) + +### 5. L4 — Golden match + +- [ ] YAML test case has `level: "L4"` or `"L4+L5"` +- [ ] `inputs.prompts: ["Here is my poem:"]` (standard default prompt unless + audio/image model) +- [ ] Golden file generated: + `python scripts/generate_golden.py --level L4 --filter '*'` +- [ ] Golden file committed to `testdata/golden//.json` +- [ ] `python -m pytest tests/e2e_golden_test.py -m golden -k ""` passes + +### 6. L5 — Generation verified + +- [ ] YAML test case has `level: "L5"` or `"L4+L5"` +- [ ] `generation.max_new_tokens` set (≥ 20 recommended) +- [ ] `generation.do_sample: false` (deterministic greedy decode) +- [ ] Generation golden file generated: + `python scripts/generate_golden.py --level L5 --filter '*'` +- [ ] Generation golden file committed to + `testdata/golden//_generation.json` +- [ ] `python -m pytest tests/e2e_golden_test.py -m generation -k ""` passes + +> **Speech-language models:** The golden generation script supports the +> `speech-language` task type for models that process audio inputs (e.g., +> Gemma4). The `_generate_speech_language()` function in +> `scripts/generate_golden.py` handles audio feature extraction and input +> construction for these models. Ensure the correct feature extractor +> (e.g. `Gemma4AudioFeatureExtractor`, not `WhisperFeatureExtractor`) is +> auto-detected via `AutoFeatureExtractor.from_pretrained()`. + +> **Why L4/L5 matter:** Graph-build tests (L1) only verify ONNX graph +> construction; they never execute the graph with real data. A MatMul shape +> mismatch that crashes at runtime, a wrong normalisation type, or a missing +> scaling multiplier can all pass L1 while producing completely wrong output. +> L4/L5 are the only tests that catch these classes of bugs. + +### 7. Multi-dtype correctness + +- [ ] fp32 tests pass (target: exact token match in greedy generation) +- [ ] fp16 tests pass (target: logit parity `atol=1e-2`; token match for + first N tokens) +- [ ] bf16 tests pass (target: logit parity `atol=1e-2`) + +Use the example `--compare-hf --dtype f16/bf16` flag if a comparison script +exists: + +```bash +python examples/_text_generation.py --compare-hf --dtype f16 +python examples/_text_generation.py --compare-hf --dtype bf16 +``` + +### 8. CLI build + +- [ ] `mobius build --model /tmp/out` completes without error +- [ ] Output directory contains the expected ONNX files and `genai_config.json` + +### 8a. Multi-EP correctness (CUDA) + +- [ ] Model runs correctly with `--ep cuda` (or `--device cuda`) +- [ ] CUDA results match CPU results (compare generation output) +- [ ] No crashes from large tensor operations (see ORT Gather int32 + overflow: microsoft/onnxruntime#28107) +- [ ] `ort_lower_opset_for_ep` flag handles opset 24→23 lowering for + CUDA EP (enabled by default in `src/mobius/_flags.py`) + +### 9. ORT GenAI runtime + +- [ ] Model can be loaded with `ort_genai.Model(output_dir)` without error +- [ ] Greedy text generation produces non-empty, coherent output +- [ ] ORT GenAI test added to `tests/ort_genai_test.py` + (or confirmed covered by an existing parametrized test) + +Run the ORT GenAI integration test: + +```bash +python -m pytest tests/ort_genai_test.py -m integration_slow -k "" -sv +``` + +### 10. Foundry Local smoke test + +- [ ] Model exported package can be loaded and run in Foundry Local +- [ ] At minimum, verify that the `genai_config.json` and all ONNX files + are present and the model responds to a short prompt + +> If Foundry Local is not available in the current environment, document the +> skip with a `# TODO: verify with Foundry Local` comment in the PR. + +### 11. Olive quantization compatibility + +- [ ] Model can be loaded from the exported ONNX package by Olive +- [ ] INT4 / INT8 quantization runs to completion without errors +- [ ] Quantized model produces non-degenerate output (coherent text) +- [ ] If quantization changes the graph structure (e.g. MatMulNBits), verify + the `genai_config.json` still loads correctly in ORT GenAI + +Run the quantization integration test suite to confirm existing patterns +are not broken: + +```bash +python -m pytest tests/quantization_integration_test.py -v +``` + +For new architectures, add a quantized variant test if the architecture has +novel weight layouts (e.g. fused QKV, non-standard expert routing). + +### 12. Documentation + +- [ ] Class docstring (first paragraph) clearly describes the model family + and the HuggingFace class it replicates +- [ ] `default_task` and `category` are set correctly (auto-generates the + model page and index entry) +- [ ] If the model has a notable architectural difference from the base class, + a comment in the source file or the skill notes explains it +- [ ] README model table updated if this is a significant new addition + +--- + +## Waiver policy + +Any item that cannot be completed must be waived explicitly in the PR +description: + +``` +**Waivers:** +- L5 golden: Model is 70B — generating golden data exceeds CI resources. + skip_reason added to YAML. +- Foundry Local: Not available in this environment. Tracked in issue #NNN. +``` + +Unchecked items without a waiver are grounds to request changes before merge. + +--- + +## Quick reference commands + +```bash +# Lint (auto-fix then verify clean) +lintrunner f --output oneline --all-files +lintrunner -a + +# L1 – graph build +python -m pytest tests/build_graph_test.py -k "" + +# L1 – weight alignment +python -m pytest tests/weight_alignment_test.py -k "" + +# L2 – YAML schema +python -m pytest tests/yaml_schema_test.py + +# L3 – synthetic parity +python -m pytest tests/synthetic_parity_test.py -k "" -sv + +# L3 – real-weight integration (if small checkpoint available) +python -m pytest tests/integration_test.py -m integration -k "" -sv + +# L4 – generate golden +python scripts/generate_golden.py --level L4 --filter '*' + +# L4 – run golden test +python -m pytest tests/e2e_golden_test.py -m golden -k "" -v + +# L5 – generate generation golden +python scripts/generate_golden.py --level L5 --filter '*' + +# L5 – run generation golden test +python -m pytest tests/e2e_golden_test.py -m generation -k "" -v + +# ORT GenAI runtime +python -m pytest tests/ort_genai_test.py -m integration_slow -k "" -sv + +# Quantization +python -m pytest tests/quantization_integration_test.py -v + +# CLI build +mobius build --model /tmp/out +``` diff --git a/.agents/skills/scan-and-multi-image/SKILL.md b/.agents/skills/scan-and-multi-image/SKILL.md new file mode 100644 index 00000000..347fc058 --- /dev/null +++ b/.agents/skills/scan-and-multi-image/SKILL.md @@ -0,0 +1,392 @@ +--- +name: scan-and-multi-image +description: > + Use this skill when building ONNX Scan or Loop subgraphs, or adding + multi-image support to a vision model. Covers the Scan + Padding + + Compaction pattern for variable-length per-image outputs, building Scan + body subgraphs with implicit inputs and carry states, the + rename-to-avoid-SSA-violations workaround, and the compact_scan_output + helper. Primary use case: multi-image vision models where per-image + computations produce variable output sizes. +--- + +# Skill: ONNX Scan Op & Multi-Image Vision + +## When to use + +Use this skill when: + +- Adding multi-image support to a vision encoder (iterating over + `image_grid_thw` rows). +- Building any ONNX subgraph that iterates over a variable-length sequence + with per-element computation that has dynamic output sizes. +- Working with the `Scan` or `Loop` ONNX ops through the `onnxscript` + builder API. + +## Background: the multi-image problem + +ORT GenAI calls the vision model **once** with all images packed together: + +- `pixel_values`: `(total_patches, pixel_dim)` — all images concatenated. +- `image_grid_thw`: `(num_images, 3)` INT64 — one `[T, H, W]` per image. + +HuggingFace iterates with Python `for t, h, w in grid_thw.tolist()` loops, +accumulating per-image results. ONNX has no Python loops, so we use the +**Scan** op. + +### Why not vectorize? + +Per-image computations like rotary position IDs and windowed attention +indices involve `arange(H)`, `Reshape(…, H_m, ms, W_m, ms)`, etc., where +H and W differ across images. These cannot be batched into a single tensor +operation — we need an explicit loop. + +## The Scan + Padding + Compaction pattern + +Since ONNX Scan requires **fixed-size outputs per iteration** but each +image produces variable-size results, we: + +1. **Pre-compute** `max_size = ReduceMax(per_image_sizes)` in the main + graph. +2. **Pad** each iteration's output to `max_size` inside the Scan body. +3. **Compact** the concatenated Scan output by removing padding using a + boolean mask + `Compress`. + +``` +Main graph: + max_patches = ReduceMax(T * H * W for each image) + │ +Scan body (per image): + pos_ids = compute(T_i, H_i, W_i) → (T_i*H_i*W_i, 2) + padded = Pad(pos_ids, max_patches) → (max_patches, 2) + │ +Scan output: → (num_images, max_patches, 2) + │ +Compact: + mask[i,j] = (j < patches_per_image[i]) + result = Compress(flatten(scan_output), flatten(mask)) + → (total_patches, 2) +``` + +## Helpers: `_scan_utils.py` + +Location: `src/mobius/components/_scan_utils.py` + +### `create_body_graph(state_inputs, scan_inputs, name)` + +Creates a Scan body `ir.Graph` and its `GraphBuilder`. + +```python +from mobius.components._scan_utils import ( + compact_scan_output, + create_body_graph, + rename_subgraph_values, +) + +# No carry state, one scan input per iteration +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 +``` + +### `rename_subgraph_values(graph, prefix)` + +**Critical step.** ONNX Scan body graphs share a value namespace with the +parent graph in ORT. Without renaming, node outputs like `v_Constant_11` +in the body collide with identically named values in the main graph, +causing an "SSA form violation" error. + +Call this **after** building all body graph ops and **before** calling +`op.Scan(...)` on the main graph: + +```python +rename_subgraph_values(body_graph, "rotary_body_") +``` + +The prefix must be unique per Scan in the model (e.g. `"rotary_body_"`, +`"win_body_"`, `"cu_body_"`). Graph input/output names are NOT renamed +— they define the Scan interface. + +### `compact_scan_output(op, scan_result, lengths_per_iter)` + +Removes padding from a `(num_iters, max_len, ...)` Scan output using a +boolean mask built from actual per-iteration lengths. + +```python +# After Scan +result = compact_scan_output(op, scan_result, patches_per_image) +# → (total_patches, ...) +``` + +## Step-by-step: building a Scan + +### 1. Compute per-image sizes in the main graph + +Extract column vectors from `grid_thw` using Slice + Squeeze. + +**Important:** Always specify the squeeze axis to avoid collapsing the +batch dimension when `num_images == 1`: + +```python +# GOOD — squeeze only the column axis +T_col = op.Squeeze(op.Slice(grid_thw, [0], [1], [1], [1]), [1]) + +# BAD — squeezes ALL size-1 dims; scalar when num_images=1 +T_col = op.Squeeze(op.Slice(grid_thw, [0], [1], [1], [1])) +``` + +Compute derived values: + +```python +H_col = op.Squeeze(op.Slice(grid_thw, [1], [2], [1], [1]), [1]) +W_col = op.Squeeze(op.Slice(grid_thw, [2], [3], [1], [1]), [1]) +patches_per_image = op.Mul(T_col, op.Mul(H_col, W_col)) # (N,) +max_patches = op.ReduceMax(patches_per_image, keepdims=False) # scalar +``` + +### 2. Create the body graph + +```python +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 +``` + +### 3. Build per-image computation in the body + +Extract T, H, W from the scan input and compute: + +```python +bT = body_op.Squeeze(body_op.Gather(body_thw, body_op.Constant(value_int=0))) +bH = body_op.Squeeze(body_op.Gather(body_thw, body_op.Constant(value_int=1))) +bW = body_op.Squeeze(body_op.Gather(body_thw, body_op.Constant(value_int=2))) + +result = some_per_image_computation(body_op, bT, bH, bW) +``` + +**Helper function pattern:** Extract the per-image logic into a +standalone function that takes `op` as a parameter. This function works +with either the main graph's `op` or the Scan body's `body_op`: + +```python +def _compute_one_image(op, T, H, W, ms): + """Works with any OpBuilder.""" + H_m = op.Div(H, op.Constant(value_int=ms)) + # ... computation ... + return result + +# In main graph (single-image fast path): +result = _compute_one_image(op, T, H, W, ms) + +# In Scan body: +result = _compute_one_image(body_op, bT, bH, bW, ms) +``` + +### 4. Pad the output to max_size + +```python +num_p = body_op.Mul(bT, body_op.Mul(bH, bW)) +pad_len = body_op.Reshape(body_op.Sub(max_patches, num_p), [1]) + +# For 2D output (patches, D): pads = [0, 0, pad_len, 0] +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(result, pads, body_op.Constant(value_int=-1)) +``` + +`max_patches` is an **implicit input** from the main graph — the Scan +body references it directly. The ONNX Scan spec allows body graphs to +reference outer-scope values. This is handled correctly by +`onnxscript`'s builder and `onnx_ir`'s serializer. + +### 5. Set body outputs and rename + +```python +padded.name = "padded_output" +body_graph.outputs.append(padded) + +rename_subgraph_values(body_graph, "my_scan_body_") +``` + +### 6. Call Scan on the main graph + +```python +scan_result = op.Scan( + grid_thw, # scan input (iterated over axis 0) + body=body_graph, # body subgraph + num_scan_inputs=1, # number of scan inputs + _outputs=1, # number of outputs +) +# scan_result: (num_images, max_patches, D) +``` + +### 7. Compact the result + +```python +result = compact_scan_output(op, scan_result, patches_per_image) +# → (total_patches, D) +``` + +## Advanced: carry states + +Carry states (also called "state variables") persist across Scan +iterations. Use them for accumulating offsets. + +**Body graph inputs:** `[state_1, state_2, ..., scan_input_1, ...]` +**Body graph outputs:** `[new_state_1, new_state_2, ..., scan_out_1, ...]` + +```python +# State inputs +body_offset = ir.Value( + name="offset", shape=ir.Shape([]), + type=ir.TensorType(ir.DataType.INT64), +) +body_thw = ir.Value( + name="body_thw", shape=ir.Shape([3]), + type=ir.TensorType(ir.DataType.INT64), +) + +# 1 carry state + 1 scan input +body_graph, body_builder = create_body_graph( + [body_offset], [body_thw], name="window_body", +) +body_op = body_builder.op + +# ... compute per-image result ... +win_idx = body_op.Add(local_indices, body_offset) # add offset +new_offset = body_op.Add(body_offset, total_merged) # update carry + +# Outputs: carry states FIRST, then scan outputs +new_offset.name = "new_offset" +padded_idx.name = "padded_window_index" +body_graph.outputs.extend([new_offset, padded_idx]) + +rename_subgraph_values(body_graph, "win_body_") + +# Call with initial state values +init_offset = op.Constant(value_int=0) +final_offset, scan_idx = op.Scan( + init_offset, # initial carry state + grid_thw, # scan input + body=body_graph, + num_scan_inputs=1, + _outputs=2, # 1 carry output + 1 scan output +) +``` + +### Carry state use cases + +| Use case | Carry state | Updated as | +|----------|-------------|------------| +| Window index global offsets | `merged_offset` (INT64 scalar) | `+= T * llm_h * llm_w` per image | +| cu_window_seqlens offsets | `cu_offset` (INT64 scalar) | `= last cu_window value` per image | +| Running patch count | `patch_offset` (INT64 scalar) | `+= T * H * W` per image | + +## Pitfalls and gotchas + +### 1. SSA violations from name collisions + +**Problem:** ORT rejects models where a body graph value name matches a +main graph value name (e.g. both have `v_Constant_11`). + +**Fix:** Always call `rename_subgraph_values(body_graph, "unique_prefix_")` +before `op.Scan(...)`. + +### 2. Squeeze removes batch dim when N=1 + +**Problem:** `op.Squeeze(tensor)` removes ALL size-1 dims. When +`num_images == 1`, a `(1,)` tensor becomes a scalar, causing downstream +`Unsqueeze` or `Reshape` failures. + +**Fix:** Always specify the axis: `op.Squeeze(tensor, [1])`. + +### 3. Pad format for multi-dim outputs + +ONNX `Pad` pads format for an N-dim tensor: +`[d0_begin, d1_begin, ..., dN_begin, d0_end, d1_end, ..., dN_end]` + +For a 2D `(rows, cols)` tensor padded only on rows: +`pads = [0, 0, pad_rows, 0]` + +For a 1D `(len,)` tensor: +`pads = [0, pad_len]` + +### 4. Implicit inputs from parent graph + +Scan body graphs can reference values from the parent graph (implicit +inputs). This is supported by the ONNX spec, `onnxscript`, and ORT. +Use this for `max_patches`, `max_merged`, learned parameters like +`self.pos_embed`, etc. + +**No special syntax needed:** just use the main-graph `ir.Value` directly +in `body_op` operations. + +### 5. Body graph needs opset imports + +`create_body_graph()` already handles this (sets `opset_imports={"": 23}`). +If building manually, ensure the body graph has opset imports. + +## Reference files + +| File | Content | +|------|---------| +| `src/mobius/components/_scan_utils.py` | `create_body_graph`, `rename_subgraph_values`, `compact_scan_output` | +| `src/mobius/components/_qwen25_vl_vision.py` | Qwen2.5-VL multi-image: rotary, window index, cu_seqlens via Scan | +| `src/mobius/components/_qwen3_vl_vision.py` | Qwen3-VL multi-image: rotary, cu_seqlens, pos embed interpolation via Scan | + +### Qwen2.5-VL examples (3 Scans) + +1. **`_compute_rotary_pos_ids`** — No carry state. Pads `(T*H*W, 2)` to + `(max_patches, 2)`. +2. **`_compute_window_index`** — Two carry states (`merged_offset`, + `cu_offset`). Two scan outputs (window index, cu_window). +3. **`_compute_cu_seqlens`** — No carry state. Pads `(T,)` of hw values + to `(max_T,)`. Post-Scan: compact → CumSum → Pad with leading 0. + +### Qwen3-VL examples (3 Scans) + +1. **`_compute_rotary_pos_ids`** — Same pattern as Qwen2.5-VL but with + block-row/col indexing. +2. **`_compute_cu_seqlens`** — Same as Qwen2.5-VL. +3. **`_interpolate_pos_embed`** — No carry state. Bilinear interpolation + of learned embeddings per image. References `self.pos_embed` as + implicit input. Pads `(T*H*W, hidden_size)` to `(max_patches, D)`. + +## Testing Scan-based code + +Unit tests (`build_graph_test.py`) verify graph construction only. To +verify Scan correctness at runtime, build the vision model, fill +initializers with random weights, and run with ORT: + +```python +import numpy as np +import onnx_ir as ir +import onnxruntime as ort + +# Build and fill weights... +sess = ort.InferenceSession(model_path) + +# Single image +r1 = sess.run(None, { + "pixel_values": np.random.randn(4, pd).astype(np.float32), + "image_grid_thw": np.array([[1, 2, 2]], dtype=np.int64), +}) +assert r1[0].shape[0] == 1 # 4 patches / smu(4) = 1 merged + +# Two different-size images +r2 = sess.run(None, { + "pixel_values": np.random.randn(12, pd).astype(np.float32), + "image_grid_thw": np.array([[1, 2, 2], [1, 2, 4]], dtype=np.int64), +}) +assert r2[0].shape[0] == 3 # (4+8) / 4 = 3 merged +``` diff --git a/.agents/skills/weight-name-alignment/SKILL.md b/.agents/skills/weight-name-alignment/SKILL.md new file mode 100644 index 00000000..422e74ad --- /dev/null +++ b/.agents/skills/weight-name-alignment/SKILL.md @@ -0,0 +1,417 @@ +--- +name: weight-name-alignment +description: > + Use this skill when adding or modifying a model's preprocess_weights method + to align ONNX parameter names with HuggingFace weight names. Covers + nn.ModuleList for Sequential patterns, wrapper modules for nesting, + placeholder modules, non-consecutive indices, and which rename categories + cannot be eliminated. Reduces or eliminates weight name renames by + structuring nn.Module attributes to match HuggingFace naming conventions. +--- + +# Skill: Weight Name Alignment + +## When to use + +Use this skill when: +- Adding a new model and designing `preprocess_weights` +- Simplifying an existing model's `preprocess_weights` method +- Debugging weight loading failures (mismatched parameter names) +- Deciding whether to restructure model construction vs. rename in + `preprocess_weights` + +## Core principle + +**The best `preprocess_weights` is a no-op.** Most renames exist because the +ONNX module hierarchy doesn't match HuggingFace's. By restructuring +`nn.Module` construction to produce parameter names that match HF directly, +you can eliminate renames entirely. + +## How parameter names are formed + +In `onnxscript.nn`, parameter names are built from the Python attribute chain: + +```python +class MyModel(nn.Module): + def __init__(self): + self.layers = nn.ModuleList([MyLayer()]) + # layers[0].weight → "layers.0.weight" + +class MyLayer(nn.Module): + def __init__(self): + self.linear = _Linear(4, 4) + # linear.weight → "linear.weight" +``` + +The full name is `"layers.0.linear.weight"`. + +## Categories of renames + +### ✅ Can be eliminated (restructure model construction) + +#### 1. Sequential index patterns (nn.Sequential / nn.ModuleList) + +**HF pattern:** `nn.Sequential(SiLU(), Linear(...))` → weights at `mod.1.weight` + +**Problem:** Using a plain `_Linear(...)` produces `mod.weight` (no index). + +**Preferred solution — `nn.Sequential`:** + +`nn.Sequential` (from `onnxscript.nn`) registers children with numeric keys +like PyTorch's `nn.Sequential`, AND chains `forward()` calls automatically. +This gives both correct naming and clean call sites: + +```python +from mobius.components import Linear, SiLU + +# Produces "img_mod.1.weight" — matching HF +self.img_mod = nn.Sequential(SiLU(), Linear(dim, 6 * dim)) + +# Forward: output chains through each child automatically +result = self.img_mod(op, temb) +``` + +`nn.Sequential` subclasses `nn.ModuleList`. Key implementation detail: it +overrides `_set_name` to keep children with simple "0", "1" names (not +fully-qualified), because `__call__` already pushes the parent name onto the +scope stack. Without this override, children would be double-prefixed. + +**Fallback — `nn.ModuleList` with manual indexing:** + +If `nn.Sequential` is not yet available, use `nn.ModuleList` with explicit +`[i]` indexing: + +```python +self.img_mod = nn.ModuleList([SiLU(), Linear(dim, 6 * dim)]) + +# Forward: manual chaining +result = self.img_mod[1](op, self.img_mod[0](op, temb)) +``` + +This produces the same parameter names but requires manual forward logic. + +#### 2. Non-consecutive indices with placeholder modules + +**HF pattern:** `nn.Sequential(Linear, GELU, Linear)` → weights at `0.weight` +and `2.weight` (GELU at index 1 has no params). + +**Problem:** `nn.ModuleList([linear1, linear2])` produces indices 0, 1. + +**Solution:** Include activation modules to fill gaps: + +```python +class _NoOpModule(nn.Module): + """Placeholder for HF Dropout (no params, identity at inference).""" + def forward(self, op, x): + return x + +class _GELUGate(nn.Module): + """Matches HF GEGLU wrapper with .proj sub-attribute.""" + def __init__(self, in_features, out_features): + super().__init__() + self.proj = _Linear(in_features, out_features) + +# Matches HF: net.0.proj.weight, net.2.weight +self.net = nn.ModuleList([ + _GELUGate(dim, inner_dim * 2), # index 0 + _NoOpModule(), # index 1 (Dropout placeholder) + _Linear(inner_dim, dim), # index 2 +]) +``` + +#### 3. Wrapper modules for extra nesting + +**HF pattern:** `time_text_embed.timestep_embedder.linear_1.weight` + +**Problem:** Flat structure produces `linear_1.weight` (missing prefix). + +**Solution:** Create wrapper module matching HF nesting: + +```python +class _TimestepMLP(nn.Module): + def __init__(self, in_channels, time_embed_dim): + super().__init__() + self.linear_1 = _Linear(in_channels, time_embed_dim) + self.linear_2 = _Linear(time_embed_dim, time_embed_dim) + +class _TimestepEmbedding(nn.Module): + def __init__(self, in_channels, time_embed_dim): + super().__init__() + self.timestep_embedder = _TimestepMLP(in_channels, time_embed_dim) +``` + +#### 4. Bare Parameter → Module wrapper + +**HF pattern:** `txt_norm.weight` (from `RMSNorm` module) + +**Problem:** Using `nn.Parameter` produces `txt_norm` (no `.weight` suffix). + +**Solution:** Use a proper module: + +```python +# BAD — produces "txt_norm" as a bare parameter name +self.txt_norm = nn.Parameter((dim,)) + +# GOOD — produces "txt_norm.weight" +class _RMSNorm(nn.Module): + def __init__(self, dim, eps=1e-6): + super().__init__() + self.weight = nn.Parameter((dim,)) + self._eps = eps + def forward(self, op, x): + return op.RMSNormalization(x, self.weight, epsilon=self._eps) + +self.txt_norm = _RMSNorm(dim) +``` + +#### 5. Inner model wrapper for prefix nesting + +**HF pattern:** `model.layers.0.self_attn.q_proj.weight` + +**Problem:** Without a `model` wrapper, you get `layers.0.self_attn...`. + +**Solution:** Create inner model class: + +```python +class _TextModel(nn.Module): + def __init__(self, config): + super().__init__() + self.layers = nn.ModuleList([...]) + self.norm = _RMSNorm(config.hidden_size) + +class MyCausalLMModel(nn.Module): + def __init__(self, config): + super().__init__() + self.model = _TextModel(config) # Creates "model." prefix + self.lm_head = _Linear(config.hidden_size, config.vocab_size) +``` + +### ❌ Cannot be eliminated (must stay in preprocess_weights) + +#### 1. QKV splitting + +HuggingFace fuses Q, K, V into a single tensor (`query_key_value`, +`c_attn`, `qkv_proj`), but ONNX uses separate `q_proj`, `k_proj`, `v_proj`. + +```python +def preprocess_weights(self, state_dict): + new_state = {} + for key, tensor in state_dict.items(): + if "query_key_value" in key: + q, k, v = self._split_qkv(tensor, self.config) + new_state[key.replace("query_key_value", "q_proj")] = q + new_state[key.replace("query_key_value", "k_proj")] = k + new_state[key.replace("query_key_value", "v_proj")] = v + else: + new_state[key] = tensor + return new_state +``` + +**Models affected:** GPT-2, Falcon, InternLM2, ChatGLM, Phi3/Phi3Small + +#### 2. Conv1D → Linear transpose + +GPT-2 uses Conv1D `[in, out]` layout; ONNX Linear needs `[out, in]`. + +```python +if key.endswith(".weight") and tensor.ndim == 2: + tensor = tensor.t() +``` + +**Models affected:** GPT-2 + +#### 3. Deep structural naming differences + +BERT, T5, BART have deeply different naming conventions that would require +rewriting fundamental component classes to match. + +```python +# BERT: "encoder.layer.0.attention.self.query.weight" +# Ours: "encoder.layer.0.self_attn.q_proj.weight" +``` + +Changing this would require BERT-specific Attention, MLP components — not +worth the complexity for a simple rename. + +**Models affected:** BERT, DistilBERT, RoBERTa, ALBERT, T5, BART, mBART, +Marian, CLIP, SigLIP + +#### 4. MoE expert weight remapping + +MoE models have mixed naming across architectures (Mixtral: `w1/w2/w3`, +Qwen2-MoE: `gate_proj/up_proj/down_proj`). The current `_rename_moe_expert_weights` +handles both conventions optimally. + +#### 5. Weight tying + +Always needed when `tie_word_embeddings=True`: + +```python +if self.config.tie_word_embeddings: + if "lm_head.weight" in state_dict: + state_dict["model.embed_tokens.weight"] = state_dict["lm_head.weight"] + elif "model.embed_tokens.weight" in state_dict: + state_dict["lm_head.weight"] = state_dict["model.embed_tokens.weight"] +``` + +#### 6. Weight deletion + +Some HF weights are not needed (e.g. `rotary_emb.inv_freq` — RoPE +frequencies are computed at runtime). + +## ⚠️ Scope mechanism pitfalls + +Understanding how `onnxscript.nn` builds parameter names is critical. Names +come from the **call-stack of `__call__` invocations**, not from the Python +attribute chain used to *reach* a module. + +### Rule: only `__call__` pushes scope + +When you call `module(op, x)`, the base `nn.Module.__call__` method: +1. Pushes `module._name` onto the scope stack +2. Calls `module.forward(op, x)` +3. Pops the scope + +**Accessing a child and calling its method directly bypasses the parent scope:** + +```python +# ❌ BAD — self.shared.child(op, x) only pushes "child", NOT "shared" +result = self.shared.child(op, x) +# Parameter name: "child.weight" (missing "shared." prefix!) + +# ✅ GOOD — call shared's forward, which internally calls child +result = self.shared(op, x) +# Parameter name: "shared.child.weight" ✓ +``` + +### ModuleList indexing requires `__call__` + +The same rule applies to `nn.ModuleList`. `self.items[i]` returns the +module object — its `_name` is `"items.{i}"` — but that scope is only +pushed when you call it: + +```python +# ❌ BAD — sub_module's scope is pushed, but items.0 scope is NOT +result = self.items[0].sub_module(op, x) +# Parameter name: "sub_module.weight" (missing "items.0." prefix!) + +# ✅ GOOD — items[0].__call__ pushes "items.0", then forward calls sub_module +result = self.items[0](op, x) +# Parameter name: "items.0.sub_module.weight" ✓ +``` + +**Key takeaway:** If you need per-index distinct parameter names (e.g. +per-layer adapters), you **must** call the ModuleList element via +`self.items[idx](op, ...)`, not reach into its sub-attributes. + +### Shared weights via single module instance + +When multiple layers reuse the same weights (e.g. Zamba2's shared +transformer), register ONE module instance. Calling it multiple times +produces the same initializer names — ONNX uses a single initializer: + +```python +class _TextModel(nn.Module): + def __init__(self, config): + super().__init__() + # ONE shared module → ONE set of initializers + self.shared_transformer = SharedLayer(config) + + def forward(self, op, x): + for i in range(num_uses): + # Same scope "shared_transformer.*" each time → same initializer + x = self.shared_transformer(op, x) +``` + +### Circular dependency: shared weights + per-instance data + +When shared weights and per-instance data (e.g. adapters) must interact in +the same computation, the scope model creates a tension: + +- **Shared weights** must be inside a shared module (for correct scope) +- **Per-instance data** must be outside (different scope per use) + +**Solution: split the computation.** Have the shared module return an +intermediate value. The caller computes per-instance contributions at its +scope, then continues the computation with shared weights at its level: + +```python +class _TextModel(nn.Module): + def __init__(self, config): + super().__init__() + self.shared_attn = SharedAttention(config) # shared weights + self.adapters = nn.ModuleList([...]) # per-layer + self.gate_proj = Linear(...) # shared MLP at model scope + + def forward(self, op, hidden): + for idx in range(num_layers): + # Phase 1: shared attention (inside shared module scope) + intermediate = self.shared_attn(op, hidden) + # Phase 2: per-layer adapter (at "adapters.{idx}" scope) + adapter_out = self.adapters[idx](op, intermediate) + # Phase 3: shared MLP at model scope (gate_proj is model attr) + hidden = self.gate_proj(op, op.Add(intermediate, adapter_out)) +``` + +**Reference implementation:** `models/zamba2.py` — Zamba2 hybrid model with +shared transformer + per-layer Q/K/V/MLP low-rank adapters. + +## How to analyze a model's preprocess_weights + +1. **Compare HF names to ONNX names:** + ```python + # Print ONNX parameter names + module = MyModel(config) + for name, _ in module.named_parameters(): + print(name) + + # Print HF weight names (from safetensors) + from safetensors import safe_open + with safe_open("model.safetensors", framework="pt") as f: + for key in f.keys(): + print(key) + ``` + +2. **Categorize each rename** as one of the types above. + +3. **For eliminable renames**, restructure the module constructor. + +4. **For non-eliminable renames**, keep them in `preprocess_weights`. + +## Non-consecutive index patterns (setattr fallback) + +When HF uses `nn.Sequential` with non-consecutive parameter indices AND the +gap modules have no natural implementation: + +```python +# HF: Sequential(Conv2d, SiLU, Conv2d, SiLU, Conv2d, SiLU, Conv2d) +# Weights at indices: 0, 2, 4, 6 (SiLU at 1, 3, 5 has no params) +# But we process differently in forward, so ModuleList doesn't work + +from mobius.components import Conv2d + +# Fallback: manual setattr +class _SequentialConv2d(nn.Module): + def __init__(self, in_channels, out_channels, **kwargs): + super().__init__() + conv = Conv2d(in_channels, out_channels, **kwargs) + setattr(self, "1", conv) # Matches HF Sequential index +``` + +Use this only when `nn.ModuleList` with activation placeholders doesn't +work (e.g., different forward logic, or HF Sequential wraps padding + conv). + +## Reference implementations + +| Pattern | Model | File | +|---------|-------|------| +| No-op (fully aligned) | QwenImage transformer | `models/qwen_image.py` | +| Weight tying only | CausalLMModel (base) | `models/base.py` | +| Sequential index (ModuleList) | UNet, DiT, VAE | `models/unet.py`, `models/dit.py`, `models/vae.py` | +| Wrapper + placeholder modules | QwenImage (all patterns) | `models/qwen_image.py` | +| QKV splitting | Falcon, GPT-2 | `models/falcon.py`, `models/gpt2.py` | +| Conv1D transpose | GPT-2 | `models/gpt2.py` | +| MoE expert remapping | MoE models | `models/moe.py` | +| Deep structural renames | BERT, T5 | `models/bert.py`, `models/t5.py` | +| Shared weights + per-layer adapters | Zamba2 | `models/zamba2.py` | +| Scope-aware ModuleList adapters | Zamba2 | `models/zamba2.py` | diff --git a/.agents/skills/writing-rewrite-rules/SKILL.md b/.agents/skills/writing-rewrite-rules/SKILL.md new file mode 100644 index 00000000..14f3b96a --- /dev/null +++ b/.agents/skills/writing-rewrite-rules/SKILL.md @@ -0,0 +1,302 @@ +--- +name: writing-rewrite-rules +description: > + Use this skill when creating ONNX rewrite rules that transform parts of an + ONNX model graph, such as replacing standard ops with custom or fused ops. + Covers the onnxscript.rewriter RewriteRuleClassBase API, pattern matching + with op functions, check/rewrite method conventions, file organization + under src/mobius/rewrite_rules/, and testing patterns for verifying rule + correctness. +--- + +# Skill: Writing Rewrite Rules + +## When to use + +Use this skill when creating rules that transform parts of an ONNX model graph +— for example, replacing standard ops with custom or fused ops for better +runtime performance. Rewrite rules live in `src/mobius/rewrite_rules/` +and are applied **after** model export. + +## API overview + +Rewrite rules use `RewriteRuleClassBase` from `onnxscript.rewriter`: + +```python +from onnxscript.rewriter._basics import MatchResult +from onnxscript.rewriter._rewrite_rule import RewriteRuleClassBase, RewriteRuleSet +from onnxscript.rewriter import rewrite +``` + +1. **Subclass** `RewriteRuleClassBase` with `pattern()`, `check()`, and + `rewrite()` methods. +2. Call `.rule()` on an instance to create a `RewriteRule`. +3. Wrap one or more rules in a `RewriteRuleSet`. +4. Apply via `rewrite(model, pattern_rewrite_rules=rule_set)`. + +> **Important:** The keyword argument is `pattern_rewrite_rules`, **not** `rules`. + +```python +class MyRule(RewriteRuleClassBase): + def pattern(self, op, ...): ... + def check(self, context, ...): ... + def rewrite(self, op, ...): ... + +def my_rules() -> RewriteRuleSet: + return RewriteRuleSet([MyRule().rule()]) + +# Apply +rewrite(model, pattern_rewrite_rules=my_rules()) +``` + +## Pattern function + +`pattern(self, op, ...)` defines the ONNX subgraph to match. The positional +parameters after `op` become the matched inputs. + +```python +def pattern(self, op, q, k, v, attn_bias_2d): + q_4d = op.Unsqueeze(q, [0]) + k_4d = op.Unsqueeze(k, [0]) + v_4d = op.Unsqueeze(v, [0]) + attn_bias_4d = op.Unsqueeze(attn_bias_2d, [0, 1]) + attn_out = op.Attention( + q_4d, k_4d, v_4d, attn_bias_4d, + _allow_other_attributes=True, + _outputs=["attn_out"], + ) + return op.Squeeze(attn_out, [0]) +``` + +Key options: + +- **`_outputs=["name"]`** — Capture an intermediate value so it can be + referenced in `check()` and `rewrite()` by that keyword name. +- **`_allow_other_attributes=True`** — Match an op even if it has extra + attributes not listed in the pattern (e.g. `scale`, `num_heads`). + +## Check function + +`check(self, context, **kwargs)` validates structural requirements that the +pattern alone cannot express. Matched inputs and captured outputs arrive as +keyword arguments. + +Return `MatchResult()` (no arguments) for success. Call `.fail("reason")` to +reject the match: + +```python +def check(self, context, attn_bias_2d, attn_out, **_): + result = MatchResult() + + # Walk the producer chain to verify structure + where = attn_bias_2d.producer() + if where is None or where.op_type != "Where": + return result.fail("Expected Where producing attention bias") + + # Access attributes on matched nodes + attn = attn_out.producer() + if attn.attributes.get_float("scale", None) is None: + return result.fail("Missing scale attribute on Attention") + if attn.attributes.get_int("q_num_heads", None) is None: + return result.fail("Missing q_num_heads attribute on Attention") + + return result +``` + +Common attribute accessors: + +- `node.attributes.get_float("name")` / `node.attributes.get_float("name", default)` +- `node.attributes.get_int("name")` / `node.attributes.get_int("name", default)` + +## Rewrite function + +`rewrite(self, op, **kwargs)` builds the replacement subgraph. The `op` +parameter is an **IR tape builder** (not the same as `onnxscript`'s +`OpBuilder`). Matched inputs and captured outputs arrive as keyword arguments. + +```python +def rewrite(self, op, q, k, v, attn_bias_2d, attn_out, **_): + attn = attn_out.producer() + scale = attn.attributes.get_float("scale") + num_heads = attn.attributes.get_int("q_num_heads") + + cu_seqlens = self._trace_cu_seqlens(attn_bias_2d) + cu_seqlens_i32 = op.Cast(cu_seqlens, to=6) + + return op.op( + "PackedMultiHeadAttention", + inputs=[q, k, v, None, token_offset, cu_seqlens_i32], + domain="com.microsoft", + attributes={"scale": scale, "num_heads": num_heads}, + ) +``` + +### Critical: constant tensors in rewrite + +Raw Python lists **cannot** be used as inputs in the rewrite function. +Always create constants explicitly: + +```python +# GOOD — explicit Constant node +axes_0 = op.Constant(value_ints=[0]) +neg_one = op.Constant(value_ints=[-1]) +result = op.Squeeze(x, axes_0) + +# BAD — raw list (will fail) +result = op.Squeeze(x, [0]) +``` + +### Custom / domain-specific ops + +Use `op.op(...)` to emit single-output ops from non-default domains: + +```python +op.op( + "PackedMultiHeadAttention", + inputs=[q, k, v, None, token_offset, cu_seqlens], + domain="com.microsoft", + attributes={"scale": scale, "num_heads": num_heads}, +) +``` + +Pass `None` in the inputs list for optional inputs that should be left empty. + +### Multi-output custom ops + +Use `op.op_multi_out(...)` for ops with multiple outputs. +**`op.op()` returns a single `ir.Value`; `op.op_multi_out()` returns +`Sequence[ir.Value]`.** + +```python +outputs = op.op_multi_out( + "GroupQueryAttention", + inputs=[q, k, v, past_key, past_value, seqlens_k, total_seq_len], + domain="com.microsoft", + attributes={"num_heads": num_heads, "kv_num_heads": kv_num_heads}, + num_outputs=3, +) +attn_out, present_key, present_value = outputs[0], outputs[1], outputs[2] +``` + +### Matching patterns with shared intermediate values + +The rewriter will **not** match a pattern if an intermediate node's output +has consumers outside the matched subgraph. For example, matching +`Add → RMSNorm` will fail if the `Add` output is also used by a downstream +residual connection. + +**Workaround:** Match only the end node, then trace back in `check()`: + +```python +def pattern(self, op, add_out, weight): + # Only match RMSNorm — don't include Add in the pattern + return op.RMSNormalization(add_out, weight, _allow_other_attributes=True) + +def check(self, context, add_out, **_): + result = MatchResult() + producer = add_out.producer() + if producer is None or producer.op_type != "Add": + return result.fail("Input is not from Add") + if len(list(add_out.uses())) < 2: + return result.fail("Add has only 1 consumer") + return result + +def rewrite(self, op, add_out, weight, **_): + add_node = add_out.producer() + input_a, input_b = add_node.inputs[0], add_node.inputs[1] + # Create fused op and reroute the shared output + outputs = op.op_multi_out("FusedOp", inputs=[input_a, input_b, weight], ...) + add_out.replace_all_uses_with(outputs[1]) # reroute skip connection + return outputs[0] # return the primary output +``` + +## File organization + +``` +src/mobius/rewrite_rules/ +├── __init__.py # Public exports +├── _packed_attention.py # Rule implementation (private module) +├── _packed_attention_test.py # Unit tests (next to source) +├── _group_query_attention.py # Attention → GQA rule +├── _group_query_attention_test.py # GQA rule tests +├── _skip_norm.py # Add+RMSNorm → SkipNorm rule +└── _skip_norm_test.py # SkipNorm rule tests +``` + +### Conventions + +- Rule files are **private modules**: `_rule_name.py`. +- Unit tests go **next to the source file**: `_rule_name_test.py`. +- Export the public factory function from `__init__.py`: + +```python +# __init__.py +__all__ = ["packed_attention_rules"] +from mobius.rewrite_rules._packed_attention import packed_attention_rules +``` + +- Each rule module should provide a factory function (e.g. + `packed_attention_rules()`) that returns a `RewriteRuleSet`. + +## Testing + +Write unit tests that: + +1. Build a model containing the target pattern (either a tiny model from + the model library or a synthetic graph). +2. Count ops before applying the rule. +3. Apply the rule set via `rewrite(model, pattern_rewrite_rules=rules)`. +4. Count ops after and assert the expected replacements occurred. +5. Verify that non-matching subgraphs are **not** affected. + +```python +from collections import Counter +from onnxscript.rewriter import rewrite +from mobius.rewrite_rules import packed_attention_rules + + +def _count_ops(model) -> Counter: + return Counter(node.op_type for node in model.graph) + + +class TestPackedAttentionRules: + def test_rule_replaces_vision_attention(self): + model = build_model_with_pattern(...) + counts_before = _count_ops(model) + assert counts_before["Attention"] == 4 + + rewrite(model, pattern_rewrite_rules=packed_attention_rules()) + + counts_after = _count_ops(model) + assert counts_after["PackedMultiHeadAttention"] == 2 + assert counts_after["Attention"] == 2 # text decoder untouched + + def test_rule_preserves_non_matching_model(self): + """Models without the pattern are not affected.""" + model = build_text_only_model(...) + counts_before = _count_ops(model) + + rewrite(model, pattern_rewrite_rules=packed_attention_rules()) + + counts_after = _count_ops(model) + assert counts_after["Attention"] == counts_before["Attention"] + assert counts_after.get("PackedMultiHeadAttention", 0) == 0 + + def test_rules_returns_rule_set(self): + from onnxscript.rewriter._rewrite_rule import RewriteRuleSet + rules = packed_attention_rules() + assert isinstance(rules, RewriteRuleSet) +``` + +## Reference files + +- **Full rule implementations:** + - `src/mobius/rewrite_rules/_packed_attention.py` — Block-diagonal → PackedMHA + - `src/mobius/rewrite_rules/_group_query_attention.py` — Attention → GQA + - `src/mobius/rewrite_rules/_skip_norm.py` — Add+RMSNorm → SkipNorm +- **Test examples:** + - `src/mobius/rewrite_rules/_packed_attention_test.py` + - `src/mobius/rewrite_rules/_group_query_attention_test.py` + - `src/mobius/rewrite_rules/_skip_norm_test.py` +- **Exports:** + `src/mobius/rewrite_rules/__init__.py` From 616cf378838c1bd55d51a1d58d6f53e03c7810d5 Mon Sep 17 00:00:00 2001 From: Justin Chu Date: Tue, 21 Apr 2026 16:23:26 +0000 Subject: [PATCH 5/8] Merge model-specific skills into general ones MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Merge 1: debugging-vl-pipeline + phi4mm-component-parity → debugging-multimodal - Combined 3-stage VL pipeline and 4-model multi-encoder methodologies - Merged failure modes (deduplicated overlapping content like ClippableLinear, boundary tokens, empty tensors) into references/failure-modes.md - Preserved extraction-methods.md and debugging-cookbook.md as references - Removed old debugging-vl-pipeline and phi4mm-component-parity directories Merge 2: scan-and-multi-image → multimodal-models/references/scan-pattern.md - Moved scan skill content as a reference file under multimodal-models - Added reference directive in multimodal-models SKILL.md - Removed old scan-and-multi-image directory Updated all cross-references in multimodal-models SKILL.md and projector-variants.md to point to new locations. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Justin Chu --- .agents/skills/debugging-multimodal/SKILL.md | 316 ++++++++++++++ .../references/debugging-cookbook.md | 20 +- .../references/extraction-methods.md | 0 .../references/failure-modes.md | 389 ++++++++++++++++++ .agents/skills/debugging-vl-pipeline/SKILL.md | 222 ---------- .../references/failure-modes.md | 206 ---------- .agents/skills/multimodal-models/SKILL.md | 10 +- .../references/projector-variants.md | 2 +- .../references/scan-pattern.md} | 20 +- .../skills/phi4mm-component-parity/SKILL.md | 148 ------- .../references/common-failures.md | 244 ----------- 11 files changed, 727 insertions(+), 850 deletions(-) create mode 100644 .agents/skills/debugging-multimodal/SKILL.md rename .agents/skills/{phi4mm-component-parity => debugging-multimodal}/references/debugging-cookbook.md (90%) rename .agents/skills/{debugging-vl-pipeline => debugging-multimodal}/references/extraction-methods.md (100%) create mode 100644 .agents/skills/debugging-multimodal/references/failure-modes.md delete mode 100644 .agents/skills/debugging-vl-pipeline/SKILL.md delete mode 100644 .agents/skills/debugging-vl-pipeline/references/failure-modes.md rename .agents/skills/{scan-and-multi-image/SKILL.md => multimodal-models/references/scan-pattern.md} (95%) delete mode 100644 .agents/skills/phi4mm-component-parity/SKILL.md delete mode 100644 .agents/skills/phi4mm-component-parity/references/common-failures.md diff --git a/.agents/skills/debugging-multimodal/SKILL.md b/.agents/skills/debugging-multimodal/SKILL.md new file mode 100644 index 00000000..2ab4bd4d --- /dev/null +++ b/.agents/skills/debugging-multimodal/SKILL.md @@ -0,0 +1,316 @@ +--- +name: debugging-multimodal +description: > + Debug wrong, garbled, or divergent output from multimodal ONNX models + (vision-language, vision+audio, multi-encoder). Use when ORT GenAI + multimodal output doesn't match HuggingFace, when building a new + multimodal model and verifying component-by-component parity (vision + encoder, speech encoder, embedding/projector, text decoder), when + integration tests fail with large numerical differences, or when CUDA + EP produces different results than CPU. Covers the 3-stage VL pipeline + and 4-model multi-encoder pipeline isolation, 3D M-RoPE position IDs, + CUDA EP gotchas, and systematic stage-by-stage comparison methodology. +--- + +# Skill: Debugging Multimodal Pipeline Issues + +## When to use + +Use this skill when: + +- ORT GenAI produces wrong or irrelevant output for image/audio inputs +- ONNX model logits diverge significantly from HuggingFace +- The model generates text-only descriptions ignoring the image +- Audio transcription is garbled or wrong despite correct encoder output +- You're adding a new multimodal model and need to verify each component +- Integration tests fail with large numerical differences (not just tolerance) +- Weights appear to load but the model produces wrong outputs +- CUDA EP crashes or produces different results than CPU + +## Quick diagnostic flow + +1. **Text-only works?** If yes → problem is in vision/audio path or + position IDs. If no → decoder or weight loading issue. +2. **Vision cos_sim > 0.99?** If no → vision encoder issue (see Stage 1). +3. **Embedding text positions match?** If no → token replacement bug. +4. **3D M-RoPE fields set?** Missing `image_token_id`, + `vision_start_token_id`, or `spatial_merge_size` causes 1D fallback. +5. **CUDA-only failure?** See CUDA EP section below. + +## Debugging methodology: isolate each stage + +### 3-stage VL pipeline (vision-language only) + +``` +pixel_values ──► [1. Vision] ──► image_features + │ +input_ids ──► [2. Embedding] ◄──────┘ + │ + ▼ + inputs_embeds + position_ids + attention_mask + │ + ▼ + [3. Decoder] ──► logits +``` + +### 4-model multi-encoder pipeline (e.g. Phi4MM, Gemma4) + +``` +pixel_values ──► [1. Vision Encoder] ──► image_features ──┐ + │ +audio_embeds ──► [2. Speech Encoder] ──► speech_features ──┤ + │ +input_ids ──► [3. Embedding/Fusion] ◄───────────────────┘ + │ + ▼ + inputs_embeds + │ + ▼ + [4. Decoder + LoRA] ──► logits +``` + +**Golden rule:** Start from the simplest case (text-only, no encoders), +verify it matches HF, then add one modality at a time. + +### Stage 1: Vision encoder + +**What to check:** +- Output shape: `(num_patches, hidden_size)` or + `(num_image_tokens, text_hidden_size)` after projection +- Expected patches = `t * (h / merge) * (w / merge)` from `grid_thw` +- Compare features against HF vision encoder output + +```python +# HF reference +with torch.no_grad(): + hf_vision_out = hf_model.model.visual( + pixel_values, grid_thw=grid_thw + ) +# ONNX +session = OnnxModelSession(pkg["vision"]) +onnx_out = session.run({"pixel_values": pv, "grid_thw": grid_thw}) + +# Compare +cos_sim = np.dot(hf_flat, onnx_flat) / (norm_hf * norm_onnx) +print(f"Vision cos_sim: {cos_sim:.6f}") # Should be > 0.99 +``` + +**Common issues:** +- Wrong pixel value normalization (mean/std mismatch) +- `grid_thw` shape or values don't match HF processor output +- Missing `temporal_patch_size` in patch embedding +- Wrong rotary embedding dimension (must be `head_dim // 2` for 2D) +- Missing `fullatt_block_indexes` (windowed vs full attention) + +### Stage 2: Speech encoder (multi-encoder models) + +**What to check:** +- Compression rate (typically 8× time reduction) +- Projection branch selection +- Conv subsampling output length + +Compare Conformer output against HF. Target: cos_sim > 0.99. + +### Stage 3: Embedding/fusion + +**What to check:** +- Image/audio features injected at correct token positions +- Non-image positions have correct text embeddings +- Output shape: `(1, seq_len, hidden_size)` + +```python +# Verify image token positions +image_mask = (input_ids[0] == image_token_id) +num_image_positions = image_mask.sum() +assert num_image_positions == image_features.shape[0] + +# Compare embeddings at text positions (should match HF exactly) +text_mask = ~image_mask +cos_sim_text = cosine_similarity( + onnx_embeds[0, text_mask], hf_embeds[0, text_mask] +) +print(f"Text embedding cos_sim: {cos_sim_text:.6f}") # Should be 1.0 +``` + +**Common issues:** +- Image token count mismatch between processor and vision model +- Missing zero-padding row in embedding model (for text-only inputs) +- Wrong `image_token_id` used for Gather/Where mask +- InputMixer must handle zero-length tensors + +### Stage 4: Decoder + +**What to check:** +- Logits shape matches HF: `(1, seq_len, vocab_size)` +- First token prediction matches HF (argmax of last position) +- Cosine similarity of logit vectors + +```python +onnx_logits = decoder_session.run(feeds)["logits"] +hf_logits = hf_model(**hf_inputs).logits.numpy() + +max_diff = np.abs(onnx_logits - hf_logits).max() +cos_sim = cosine_similarity(onnx_logits[0, -1], hf_logits[0, -1]) +print(f"max_diff={max_diff:.2f}, cos_sim={cos_sim:.4f}") +# Typical: max_diff=5-10, cos_sim>0.98 +``` + +## Quick-start 4-step process (new models) + +1. **Text-only baseline:** Build ONNX → run embedding + decoder (skip + encoders) → compare against HF `embed_tokens` and full forward logits. + If this diverges, fix weight loading / decoder before touching encoders. + +2. **Add vision:** Run vision encoder → feed features to embedding → + compare. If newly divergent, isolate vision encoder output vs HF. + +3. **Add audio:** Same as above for speech encoder. + +4. **Combined:** All modalities together. If divergent only in combined + mode, suspect LoRA mode mismatch or embedding fusion ordering. + +## Critical: 3D M-RoPE position IDs + +Qwen2-VL / Qwen2.5-VL / Qwen3-VL use **3D Multimodal RoPE** where +`position_ids` has shape `(3, batch, seq_len)`: + +``` +position_ids[0] = temporal positions +position_ids[1] = height positions +position_ids[2] = width positions +``` + +**Text tokens:** all 3 dimensions have the same sequential value. + +**Image tokens:** temporal is constant, height/width vary over the +image grid `(h/merge, w/merge)`. + +**Text after image:** all 3 dimensions resume from +`max(temporal, height, width) + 1`. + +### ORT GenAI config requirements + +For ORT GenAI to compute 3D M-RoPE automatically, these +`genai_config.json` fields are **required**: + +| Field | Level | Purpose | +|-------|-------|---------| +| `model.image_token_id` | model | Token ID for `<\|image_pad\|>` (e.g. 151655) | +| `model.vision_start_token_id` | model | Token ID for `<\|vision_start\|>` (e.g. 151652) | +| `model.vision.spatial_merge_size` | vision | Grid merge factor (typically 2) | + +**Without these fields**, ORT GenAI falls back to standard 1D positions, +which produces completely wrong output for image inputs. + +## Gotchas + +### Module forward() bypass + +The #1 source of missing weights. Directly accessing nested sub-module +parameters (`self.glu.ext_pw_conv_1d.weight`) instead of calling +`self.glu(op, x)` makes onnxscript unable to resolve the full module path. +**Detection:** Conv/MatMul nodes where weight inputs have +`is_initializer=False` and generic names. + +### LoRA mode mismatch + +Some models (Phi4MM) apply LoRA conditionally per modality. If ONNX applies +all adapters unconditionally, run HF reference with `input_mode=3` to match. +**Detection:** Text-only inference diverges but output is reasonable (not +garbage). + +### ClippableLinear omission (Gemma4) + +Using plain `Linear` instead of `ClippableLinear` in Gemma4 vision/audio +encoders causes max_diff 52.68 (audio) / 3.92 (vision). Check HF source +for `ClippableLinear` usage. + +### Empty tensor handling + +Text-only inference crashes when no image/audio features are present. +**Fix:** Zero-pad `image_features` before Gather, then mask with Where. + +### Missing boundary tokens + +Audio/image boundary markers (`<|audio>`, ``, `<|image>`, +``) are required for correct modality region identification. +Missing markers → garbled output even with correct encoder output. + +## CUDA EP gotchas + +### ORT Gather int32 overflow + +CUDA EP crashes or produces incorrect results for models with large +embedding tables (> 2^31 elements). CPU EP works correctly. +**Workaround:** Split large embeddings via `nn.ModuleList`. +ORT bug: microsoft/onnxruntime#28107 + +### Opset 24 kernel registration + +ORT ≤1.24.x CUDA/TRT EPs don't register kernels for opset 24. +**Fix:** Use the `ort_lower_opset_for_ep` feature flag (enabled by +default). See `src/mobius/_flags.py`. + +## Tolerance guidelines + +| Precision | atol | rtol | Notes | +|-----------|------|------|-------| +| float32 | 1e-4 | 2e-2 | Standard for single-forward-pass | +| float32 (deep model, 32+ layers) | 1e-3 | 5e-2 | Error compounds over layers | +| float16 / bfloat16 | 0.01 | 0.05 | Wider tolerance for mixed precision | +| Cosine similarity (last token) | > 0.98 | — | Primary correctness metric | +| Argmax match (first prediction) | exact | — | Should always match | + +## Integration test patterns + +### Full VL forward test +```python +assert_logits_close(onnx_logits, hf_logits, rtol=2e-2, atol=2e-1) +``` + +### 3-model pipeline test +```python +# Vision → Embedding → Decoder, each stage uses OnnxModelSession +# Compare final decoder logits against HF single-model forward +``` + +### ORT GenAI end-to-end test +```python +# Build → save flat → write genai_config → load with ort_genai → generate +# Verify output length > input (basic sanity) +``` + +## Reference files + +> Read `references/failure-modes.md` when you need detailed code examples +> for each failure mode, including weight name alignment, shape mismatches, +> dtype mismatches, ClippableLinear, HD transform format, and CUDA EP issues. + +> Read `references/extraction-methods.md` when you need to extract +> intermediate ONNX values for block-by-block comparison against HuggingFace. + +> Read `references/debugging-cookbook.md` when you need step-by-step +> debugging procedures with code for each phase, intermediate value +> extraction, integration test patterns, tolerance guidelines, weight +> loading verification, and HD multi-crop verification. + +- **Integration tests:** `tests/integration_test.py` + (`TestVLFullForward`, `TestQwen25VL3Model`), + `tests/phi4mm_integration_test.py` +- **ORT GenAI tests:** `tests/ort_genai_test.py` +- **Example scripts:** `examples/qwen25_vl_ort_genai.py`, + `examples/qwen3_vl_ort_genai.py`, `examples/gemma4_multimodal.py` +- **genai_config reference:** `.agents/skills/ort-genai-config/SKILL.md` +- **ORT GenAI position_ids code (external):** + `onnxruntime-genai/src/models/position_inputs.cpp:617-814` +- **Feature flags:** `src/mobius/_flags.py` +- **Model implementations:** `src/mobius/models/phi.py`, + `src/mobius/models/gemma4.py` +- **Audio components:** `src/mobius/components/_audio.py`, + `src/mobius/components/_gemma4_audio.py` +- **Vision components:** `src/mobius/components/_vision.py` +- **ClippableLinear:** `src/mobius/components/_gemma4_audio.py` +- **LoRA component:** `src/mobius/components/_lora.py` +- **Weight loading:** `src/mobius/_weight_loading.py` +- **Weight name alignment skill:** + `.agents/skills/weight-name-alignment/SKILL.md` diff --git a/.agents/skills/phi4mm-component-parity/references/debugging-cookbook.md b/.agents/skills/debugging-multimodal/references/debugging-cookbook.md similarity index 90% rename from .agents/skills/phi4mm-component-parity/references/debugging-cookbook.md rename to .agents/skills/debugging-multimodal/references/debugging-cookbook.md index 202832f9..33e0ab19 100644 --- a/.agents/skills/phi4mm-component-parity/references/debugging-cookbook.md +++ b/.agents/skills/debugging-multimodal/references/debugging-cookbook.md @@ -1,5 +1,10 @@ # Debugging Cookbook — Step-by-Step Procedures +Detailed step-by-step debugging procedures for multimodal model parity +issues. For the high-level methodology, see the parent `SKILL.md`. + +--- + ## Step-by-step debugging process ### Phase 1: Text-only baseline @@ -22,7 +27,8 @@ 6. **Disable LoRA** — if the model uses LoRA, zero out adapter weights and compare base model output against HF with adapters disabled. 7. **Layer-by-layer** — add intermediate outputs to the ONNX graph (see - `debugging-vl-pipeline` skill) to find which decoder layer first diverges. + `references/extraction-methods.md`) to find which decoder layer first + diverges. ### Phase 3: Add modalities @@ -112,17 +118,7 @@ def test_audio_prefill_logits_match(self): assert_logits_close(onnx_logits, hf_logits) ``` -### Tolerance guidelines - -| Precision | atol | rtol | Notes | -|-----------|------|------|-------| -| float32 | 1e-4 | 2e-2 | Standard for single-forward-pass | -| float32 (deep model, 32+ layers) | 1e-3 | 5e-2 | Error compounds over layers | -| float16 / bfloat16 | 0.01 | 0.05 | Wider tolerance for mixed precision | -| Cosine similarity (last token) | > 0.98 | — | Primary correctness metric | -| Argmax match (first prediction) | exact | — | Should always match | - -### Weight loading verification +## Weight loading verification After `apply_weights`, check the statistics: ```python diff --git a/.agents/skills/debugging-vl-pipeline/references/extraction-methods.md b/.agents/skills/debugging-multimodal/references/extraction-methods.md similarity index 100% rename from .agents/skills/debugging-vl-pipeline/references/extraction-methods.md rename to .agents/skills/debugging-multimodal/references/extraction-methods.md diff --git a/.agents/skills/debugging-multimodal/references/failure-modes.md b/.agents/skills/debugging-multimodal/references/failure-modes.md new file mode 100644 index 00000000..dfdda389 --- /dev/null +++ b/.agents/skills/debugging-multimodal/references/failure-modes.md @@ -0,0 +1,389 @@ +# Failure Modes — Detailed Reference + +Comprehensive failure mode reference for debugging multimodal ONNX pipelines. +Merges failure modes from both VL pipeline debugging and multi-encoder +component parity debugging. For the high-level methodology, see the parent +`SKILL.md`. + +--- + +## 1. Image not recognized (wrong output for image inputs) + +**Symptoms:** Model produces generic or hallucinated descriptions that +don't match the input image. Text-only generation works correctly. + +**Root causes (in order of likelihood):** + +1. **Missing genai_config fields** — `image_token_id`, + `vision_start_token_id`, or `spatial_merge_size` not set. + Without these, position_ids are 1D instead of 3D M-RoPE. + +2. **Image resize mismatch** — ORT processor resizes image to different + dimensions than HF processor, producing different number of vision + tokens. ORT's `width`/`height` in `processor_config.json` are used + as direct resize targets, unlike HF's smart_resize which computes + target from original image dimensions. + +3. **Processor config format** — ORT GenAI expects ort-extensions format + `processor_config.json`, not HuggingFace format. The file must include + `DecodeImage`, `ConvertRGB`, `Resize`, `Rescale`, `Normalize`, and + `PatchImage` transforms with correct attributes. + +**Fix for resize mismatch:** +```python +def _update_resize_for_image(processor_config_path, image_path): + """Recompute resize dimensions from actual image like HF does.""" + from PIL import Image + img = Image.open(image_path) + w, h = img.size + factor = 14 * 2 # patch_size * merge_size + new_w = round(w / factor) * factor + new_h = round(h / factor) * factor + # Update width/height in processor_config.json +``` + +## 2. Numerical divergence in greedy decoding + +**Symptoms:** First 1-3 tokens match HF, then output diverges. + +**Expected behavior:** This is inherent to ONNX vs PyTorch numerical +differences. ONNX models use different operator implementations that +accumulate small floating-point errors. + +**Typical metrics for Qwen2.5-VL 3B:** +- max_diff in logits: 5-10 +- mean_diff in logits: 0.5-1.5 +- cosine similarity: 0.98-0.99 +- First token: matches HF +- Greedy decoding: diverges at token 3-5 + +**This is NOT a bug** if the metrics above are within range. Both models +produce semantically similar descriptions. + +## 3. Vision model output shape mismatch + +**Symptoms:** Vision model produces wrong number of patches. + +**Debug:** Check `grid_thw` values: +```python +# For Qwen2.5-VL with merge_size=2: +t, h, w = grid_thw[0] +expected_patches = t * (h // 2) * (w // 2) +actual_patches = vision_output.shape[0] +assert expected_patches == actual_patches +``` + +## 4. Embedding model text-only failure + +**Symptoms:** Error when running without images (num_image_tokens=0). + +**Fix:** Ensure embedding model pads `image_features` with a zero row +before Gather, then uses a Where mask to select only real features: +```python +# Pad with zero row so Gather with index 0 doesn't fail +padded = op.Concat( + op.ConstantOfShape(...), # (1, hidden_size) zeros + image_features, + axis=0, +) +``` + +## 5. Vision encoder internal divergence (cos < 0.5) + +**Symptoms:** Vision features have very low cosine similarity (< 0.5) +against HuggingFace, even though patch embedding and weights are correct. + +**Debug with block-by-block comparison** (see `references/extraction-methods.md`). +Common root causes: + +1. **Wrong rotary embedding dimension** — Qwen2.5-VL vision uses 2D + position encoding (height + width). The rotary dim must be + `head_dim // 2`, not `head_dim`. Each half (head_dim // 4 frequencies) + covers one spatial dimension. With full `head_dim`, you get 2× too + many frequencies with wrong values. **Result: cos ≈ 0.25.** + +2. **Missing `fullatt_block_indexes`** — Qwen2.5-VL alternates between + windowed attention (local windows of `window_size` patches) and full + attention (all patches attend to all). Blocks at indexes `[7, 15, 23, 31]` + use full attention. If `fullatt_block_indexes` is not extracted from + HF config, all blocks use windowed attention. **Result: blocks 0-6 + are perfect (they're windowed anyway), but block 7+ diverges.** + +3. **Wrong attention bias construction** — Full-attention blocks should + have an all-zeros bias (everything attends to everything). Windowed + blocks have a block-diagonal bias. Check the bias by inspecting + sparsity: `(bias == -inf).float().mean()` should be ~0% for full + attention, ~98% for windowed. + +**Config extraction checklist for vision encoders:** +```python +# These fields MUST be extracted from HF vision_config: +fullatt_block_indexes = getattr(vc, "fullatt_block_indexes", None) +window_size = getattr(vc, "window_size", None) +spatial_merge_size = getattr(vc, "spatial_merge_size", 2) +temporal_patch_size = getattr(vc, "temporal_patch_size", 2) +``` + +## 6. ClippableLinear divergence (Gemma4) + +**Symptoms:** Vision or audio encoder output has large max diff (> 1.0) +against HuggingFace, even though weights are loaded correctly. + +**Root cause:** Gemma4 uses `Gemma4ClippableLinear` with learned finite +input/output activation clamping for ALL linear layers in its vision and +audio encoders (attention q/k/v/o projections AND MLP gate/up/down +projections). Using plain `Linear` misses the clamping. + +**Detection:** Check if the HuggingFace model uses `ClippableLinear`: +```bash +grep -n "ClippableLinear" transformers/models//modeling_.py +``` + +**Fix:** Use `ClippableLinear` (from `mobius.components`) for all +affected linear layers. For vision attention: q/k/v/o_proj. For MLP: +pass `linear_class=ClippableLinear` to the MLP component. + +**Impact:** +- Audio: max diff 52.68 → 0.0003 after fix +- Vision: max diff 3.92 → 0.00007 after fix + +## 7. Missing audio/image boundary markers + +**Symptoms:** Audio transcription or vision description is garbled, but +encoder output matches HuggingFace. + +**Root cause:** HuggingFace wraps modality placeholder tokens with +boundary markers: +- Image: `<|image>` (open) + N × `<|image|>` (pad) + `` (close) +- Audio: `<|audio>` (open) + N × `<|audio|>` (pad) + `` (close) + +Missing boundary markers prevent the model from correctly identifying +modality regions in the input sequence. + +**Fix:** Ensure `build_input_ids` wraps placeholder tokens with the +correct open/close marker token IDs. + +## 8. CUDA EP: ORT Gather int32 overflow + +**Symptoms:** CUDA EP crashes or produces incorrect results for models +with large embedding tables. CPU EP works correctly. + +**Root cause:** ORT CUDA `gather_impl.cu` uses `int32` for element +offset computation: `input_index = idx * cols + col_offset`. For +tensors with > 2^31 elements (e.g. Gemma4 per-layer embedding +[262144, 8960] = 2.35B elements), this overflows. + +**ORT bug:** microsoft/onnxruntime#28107 + +**Workaround:** Split large embeddings into smaller tables via +`nn.ModuleList` so each individual Gather stays under the int32 limit. +Use Slice instead of Gather for column-wise indexing on large tensors. + +## 9. CUDA EP: opset 24 kernel registration + +**Symptoms:** ORT CUDA EP fails to find kernels for standard ops +(Squeeze, Reshape, etc.) even though they work on CPU. + +**Root cause:** ORT ≤1.24.x CUDA/TRT EPs don't register kernels for +opset 24, even though the op semantics are unchanged from opset 23. + +**Fix:** Use the `ort_lower_opset_for_ep` feature flag (enabled by +default) which lowers the declared opset import from 24 to 23 for +non-CPU EPs. See `src/mobius/_flags.py` and +`src/mobius/_testing/ort_inference.py`. + +## 10. Wrong audio feature extractor + +**Symptoms:** Audio model produces completely wrong output. Audio +encoder features don't match HuggingFace at all. + +**Root cause:** Using `WhisperFeatureExtractor` instead of the +model-specific feature extractor (e.g. `Gemma4AudioFeatureExtractor`). +Different extractors produce different mel spectrograms. + +**Fix:** Always use the correct feature extractor for the model: +```python +from transformers import AutoFeatureExtractor +feature_extractor = AutoFeatureExtractor.from_pretrained(model_id) +``` + +## 11. Weight name alignment (missing weights) + +**Symptoms:** Hundreds or thousands of weights reported as "unmatched" by +`apply_weights`. Model runs but produces garbage output. + +**Root causes encountered:** + +### a. Module forward() bypass (240 missing weights in Phi4MM) + +The most insidious bug. When a component's `forward()` method directly +accesses nested sub-module parameters (e.g., `self.glu.ext_pw_conv_1d.weight`) +instead of calling the sub-module's `forward()` method, onnxscript cannot +resolve the full module path for the parameter. The weight ends up as an +unnamed, non-initializer constant in the graph. + +```python +# BAD — weights become unnamed +def forward(self, op, x): + return op.Conv(x, self.glu.ext_pw_conv_1d.weight, + self.glu.ext_pw_conv_1d.bias, ...) + +# GOOD — onnxscript resolves full module path +def forward(self, op, x): + return self.glu(op, x) # GLU.forward() calls op.Conv internally +``` + +**Detection:** Check the ONNX graph for Conv/MatMul nodes where weight +inputs have `is_initializer=False` and generic names like "weight"/"bias". + +**Fix:** Add `forward()` methods to sub-modules and call them instead of +directly accessing their parameters. + +### b. ModuleList subclass causing name doubling (8 weights) + +Subclassing `nn.ModuleList` causes the module's own name to appear twice +in the parameter path: `img_projection.img_projection.0.weight` instead +of `img_projection.0.weight`. + +```python +# BAD — name doubling +class ProjectionMLP(nn.ModuleList): + def __init__(self): + super().__init__() + self.append(nn.Linear(1152, 3072)) + self.append(nn.Linear(3072, 3072)) + +# GOOD — use nn.Module with indexed children +class ProjectionMLP(nn.Module): + def __init__(self): + super().__init__() + layers = [nn.Linear(1152, 3072), nn.Linear(3072, 3072)] + for i, layer in enumerate(layers): + setattr(self, str(i), layer) +``` + +### c. setattr with dotted names + +Using `setattr(self, "audio_projection.speech", module)` creates a single +attribute with a dot in its name, rather than a nested module. The resulting +ONNX parameter names won't match HuggingFace's `ModuleDict`-style naming. + +**Fix:** Use `nn.ModuleDict` or create proper nested attributes. + +## 12. Shape mismatches (position embedding 2D vs 3D) + +**Symptoms:** `RuntimeError: shape mismatch` during weight loading. + +**Root cause:** The ONNX component declares a parameter with a different +number of dimensions than the HuggingFace weight. Example: PatchEmbedding +declares `position_embedding.weight` as `[num_patches, hidden_size]` (2D), +but HF stores `[1, num_patches, hidden_size]` (3D). + +**Fix in `preprocess_weights()`:** +```python +if "position_embedding.weight" in key and state_dict[key].dim() == 3: + state_dict[key] = state_dict[key].squeeze(0) # [1,N,H] → [N,H] +``` + +## 13. Dtype mismatches (float64 vs float32) + +**Symptoms:** ONNX Runtime error: "type mismatch in Mul/Add node" during +inference. + +**Root causes:** + +### a. NumPy default float64 + +`numpy.array(python_float)` defaults to float64. Any constant created +from a Python scalar without explicit dtype will be float64 in the graph. + +```python +# BAD — float64 constant +scale = numpy.array(alpha / rank) # defaults to float64 +op.Mul(x, scale) # Mul(float32, float64) → type error + +# GOOD — explicit float32 +scale = numpy.array(alpha / rank, dtype=numpy.float32) +op.Mul(x, scale) +``` + +### b. Python int auto-promotion + +When passing a Python `int` to an op that expects a tensor, onnxscript +may auto-promote to float64 (implementation-dependent). + +```python +# RISKY — Python int may become float64 +op.Mul(int64_tensor, self.max_position_embeddings) + +# SAFE — explicit constant +op.Mul(int64_tensor, + op.Constant(value_int=self.max_position_embeddings)) +``` + +## 14. LoRA application mismatch (conditional vs unconditional) + +**Symptoms:** Systematic divergence (> 80% logits mismatch) across ALL +test cases, but the model structurally runs correctly. + +**Root cause:** Some models apply LoRA adapters conditionally based on +input modality. For example, Phi4MM applies: +- `input_mode=0` (text): no adapters +- `input_mode=1` (vision): vision LoRA only +- `input_mode=2` (speech): speech LoRA only +- `input_mode=3` (combined): both adapters + +**Quick fix for integration tests:** Run the HF reference with the mode +that matches the ONNX model's behavior (e.g., `input_mode=3`). + +**Proper fix:** Add an `input_mode` input to the decoder model and use +conditional logic to selectively apply adapters. + +**Detection:** If text-only inference diverges but the model generates +reasonable (not garbage) output, suspect LoRA mode mismatch. Temporarily +zero out all LoRA weights — if base model matches HF perfectly, the +LoRA application mode is the issue. + +## 15. Empty tensor handling (zero-length features) + +**Symptoms:** Crash during text-only inference when no image/audio +features are present. + +**Root cause:** The embedding model's `InputMixer` uses `GatherElements` +to place features at special token positions. With zero features, the +gather indices are empty but the operation may still execute on the +padded dimension, causing shape errors. + +**Fix pattern:** Zero-pad before Gather, then use Where to mask results: +```python +padded = op.Concat( + op.ConstantOfShape(op.Constant(value_ints=[1, hidden_size])), + features, # may be [0, hidden_size] + axis=0, +) +result = op.Where(feature_mask, gathered, text_embeddings) +``` + +## 16. HD transform image format (5D vs 4D) + +**Symptoms:** Vision model crashes or produces wrong output with multi-crop +HD images. + +**Root cause:** HD-capable vision models expect images in different formats: +- Some expect `[batch, channels, height, width]` (4D, single crop per batch) +- Others expect `[num_images, num_crops, channels, height, width]` (5D) + +**Fix:** Check the HF model's preprocessing code for the expected format, +and ensure the ONNX model's input signature matches. + +## 17. Causal mask construction (inputs_embeds vs input_ids) + +**Symptoms:** Attention mask has wrong length, causing decoder crash or +wrong output. + +**Root cause:** When the decoder receives `inputs_embeds` instead of +`input_ids`, the sequence length must be derived from the embeds tensor +shape, not from input_ids. + +**Fix:** Always derive `seq_len` from `inputs_embeds.shape[1]` when the +model uses inputs_embeds as input. diff --git a/.agents/skills/debugging-vl-pipeline/SKILL.md b/.agents/skills/debugging-vl-pipeline/SKILL.md deleted file mode 100644 index e41cc004..00000000 --- a/.agents/skills/debugging-vl-pipeline/SKILL.md +++ /dev/null @@ -1,222 +0,0 @@ ---- -name: debugging-vl-pipeline -description: > - Use this skill when debugging wrong, garbled, or divergent output from an - existing ORT GenAI multimodal pipeline (vision-language or vision+audio). - Covers the 3-stage pipeline isolation methodology, quick diagnostic flow, - 3D M-RoPE position ID issues, CUDA EP gotchas, and numerical tolerance - expectations. For building or adding a new multi-encoder model, use the - phi4mm-component-parity skill instead. ---- - -# Skill: Debugging VL Pipeline Issues - -> **Scope boundary:** This skill is for debugging runtime/inference issues -> in an **existing** ORT GenAI multimodal pipeline. If you are **building -> or adding** a new multi-encoder model and need component-by-component -> parity verification, use the `phi4mm-component-parity` skill instead. - -## When to use - -Use this skill when: - -- ORT GenAI produces wrong or irrelevant output for image inputs -- ONNX model logits diverge significantly from HuggingFace -- The model generates text-only descriptions ignoring the image -- Image features appear correct but decoder output is wrong -- Audio transcription is garbled or wrong despite correct encoder output -- CUDA EP crashes or produces different results than CPU - -## Quick diagnostic flow - -1. **Text-only works?** If yes → problem is in vision/audio path or - position IDs. If no → decoder or weight loading issue. -2. **Vision cos_sim > 0.99?** If no → vision encoder issue (see Stage 1). -3. **Embedding text positions match?** If no → token replacement bug. -4. **3D M-RoPE fields set?** Missing `image_token_id`, - `vision_start_token_id`, or `spatial_merge_size` causes 1D fallback. -5. **CUDA-only failure?** See CUDA EP section below. - -## Debugging methodology: isolate each stage - -VL models have 3 stages. Debug by isolating and validating each stage -independently, comparing against HuggingFace at every boundary. - -``` -pixel_values ──► [1. Vision] ──► image_features - │ -input_ids ──► [2. Embedding] ◄──────┘ - │ - ▼ - inputs_embeds + position_ids + attention_mask - │ - ▼ - [3. Decoder] ──► logits -``` - -### Stage 1: Vision model - -**What to check:** -- Output shape: `(num_patches, hidden_size)` -- Expected patches = `t * (h / merge) * (w / merge)` from `grid_thw` -- Compare features against HF vision encoder output - -```python -# HF reference -with torch.no_grad(): - hf_vision_out = hf_model.model.visual( - pixel_values, grid_thw=grid_thw - ) -# ONNX -session = OnnxModelSession(pkg["vision"]) -onnx_out = session.run({"pixel_values": pv, "grid_thw": grid_thw}) - -# Compare -cos_sim = np.dot(hf_flat, onnx_flat) / (norm_hf * norm_onnx) -print(f"Vision cos_sim: {cos_sim:.6f}") # Should be > 0.99 -``` - -**Common issues:** -- Wrong pixel value normalization (mean/std mismatch) -- `grid_thw` shape or values don't match HF processor output -- Missing `temporal_patch_size` in patch embedding - -### Stage 2: Embedding model - -**What to check:** -- Image features injected at correct token positions -- Non-image positions have correct text embeddings -- Output shape: `(1, seq_len, hidden_size)` - -```python -# Verify image token positions -image_mask = (input_ids[0] == image_token_id) # 151655 -num_image_positions = image_mask.sum() -assert num_image_positions == image_features.shape[0] - -# Compare embeddings at text positions (should match HF exactly) -text_mask = ~image_mask -cos_sim_text = cosine_similarity( - onnx_embeds[0, text_mask], hf_embeds[0, text_mask] -) -print(f"Text embedding cos_sim: {cos_sim_text:.6f}") # Should be 1.0 -``` - -**Common issues:** -- Image token count mismatch between processor and vision model -- Missing zero-padding row in embedding model (for text-only inputs) -- Wrong `image_token_id` used for Gather/Where mask - -### Stage 3: Decoder - -**What to check:** -- Logits shape matches HF: `(1, seq_len, vocab_size)` -- First token prediction matches HF (argmax of last position) -- Cosine similarity of logit vectors - -```python -# With HF-computed position_ids (ground truth): -onnx_logits = decoder_session.run(feeds)["logits"] -hf_logits = hf_model(**hf_inputs).logits.numpy() - -max_diff = np.abs(onnx_logits - hf_logits).max() -cos_sim = cosine_similarity(onnx_logits[0, -1], hf_logits[0, -1]) -print(f"max_diff={max_diff:.2f}, cos_sim={cos_sim:.4f}") -# Typical: max_diff=5-10, cos_sim>0.98 -``` - -**Common issues:** -- Wrong position_ids (see "3D M-RoPE" section below) -- Missing KV cache initialization -- Wrong attention_mask length - -## Critical: 3D M-RoPE position IDs - -Qwen2-VL / Qwen2.5-VL / Qwen3-VL use **3D Multimodal RoPE** where -`position_ids` has shape `(3, batch, seq_len)`: - -``` -position_ids[0] = temporal positions -position_ids[1] = height positions -position_ids[2] = width positions -``` - -**Text tokens:** all 3 dimensions have the same sequential value. - -**Image tokens:** temporal is constant, height/width vary over the -image grid `(h/merge, w/merge)`: -``` -temporal: [offset, offset, offset, ..., offset] -height: [offset, offset+1, offset+1, ..., offset+h/merge-1] -width: [offset, offset+1, offset, offset+1, ..., offset+w/merge-1] -``` - -**Text after image:** all 3 dimensions resume from -`max(temporal, height, width) + 1`. - -### ORT GenAI config requirements - -For ORT GenAI to compute 3D M-RoPE automatically, the following -`genai_config.json` fields are **required**: - -| Field | Level | Purpose | -|-------|-------|---------| -| `model.image_token_id` | model | Token ID for `<\|image_pad\|>` (e.g. 151655) | -| `model.vision_start_token_id` | model | Token ID for `<\|vision_start\|>` (e.g. 151652) | -| `model.vision.spatial_merge_size` | vision | Grid merge factor (typically 2) | - -**Without these fields**, ORT GenAI falls back to standard 1D positions, -which produces completely wrong output for image inputs (the model may -describe a "snowy landscape" instead of the actual image content). - -## CUDA EP gotchas - -### ORT Gather int32 overflow - -CUDA EP crashes or produces incorrect results for models with large -embedding tables (> 2^31 elements). CPU EP works correctly. -**Workaround:** Split large embeddings via `nn.ModuleList`. -ORT bug: microsoft/onnxruntime#28107 - -### Opset 24 kernel registration - -ORT ≤1.24.x CUDA/TRT EPs don't register kernels for opset 24. -**Fix:** Use the `ort_lower_opset_for_ep` feature flag (enabled by -default). See `src/mobius/_flags.py`. - -## Integration test patterns - -### Full VL forward test -```python -# Build model → process image with HF processor → run ONNX → compare logits -assert_logits_close(onnx_logits, hf_logits, rtol=2e-2, atol=2e-1) -``` - -### 3-model pipeline test -```python -# Vision → Embedding → Decoder, each stage uses OnnxModelSession -# Compare final decoder logits against HF single-model forward -``` - -### ORT GenAI end-to-end test -```python -# Build → save flat → write genai_config → load with ort_genai → generate -# Verify output length > input (basic sanity) -``` - -## Reference files - -- **Detailed failure modes (10):** `references/failure-modes.md` -- **Intermediate value extraction methods:** `references/extraction-methods.md` -- **Integration tests:** `tests/integration_test.py` - (`TestVLFullForward`, `TestQwen25VL3Model`) -- **ORT GenAI tests:** `tests/ort_genai_test.py` - (`TestOrtGenaiQwen25VL.test_multimodal_image_generation`) -- **Example scripts:** `examples/qwen25_vl_ort_genai.py`, - `examples/qwen3_vl_ort_genai.py`, `examples/gemma4_multimodal.py` -- **genai_config reference:** `.agents/skills/ort-genai-config/SKILL.md` -- **Component parity skill:** `.agents/skills/phi4mm-component-parity/SKILL.md` -- **ORT GenAI position_ids code (external):** - `onnxruntime-genai/src/models/position_inputs.cpp:617-814` -- **Feature flags:** `src/mobius/_flags.py` - (`ort_lower_opset_for_ep`, `ort_cuda_grouped_rmsnorm_workaround`) diff --git a/.agents/skills/debugging-vl-pipeline/references/failure-modes.md b/.agents/skills/debugging-vl-pipeline/references/failure-modes.md deleted file mode 100644 index cae059ed..00000000 --- a/.agents/skills/debugging-vl-pipeline/references/failure-modes.md +++ /dev/null @@ -1,206 +0,0 @@ -# Common Failure Modes and Fixes - -Detailed failure mode reference for debugging existing ORT GenAI multimodal -pipelines. For the high-level debugging methodology, see the parent -`SKILL.md`. - ---- - -## 1. Image not recognized (wrong output for image inputs) - -**Symptoms:** Model produces generic or hallucinated descriptions that -don't match the input image. Text-only generation works correctly. - -**Root causes (in order of likelihood):** - -1. **Missing genai_config fields** — `image_token_id`, - `vision_start_token_id`, or `spatial_merge_size` not set. - Without these, position_ids are 1D instead of 3D M-RoPE. - -2. **Image resize mismatch** — ORT processor resizes image to different - dimensions than HF processor, producing different number of vision - tokens. ORT's `width`/`height` in `processor_config.json` are used - as direct resize targets, unlike HF's smart_resize which computes - target from original image dimensions. - -3. **Processor config format** — ORT GenAI expects ort-extensions format - `processor_config.json`, not HuggingFace format. The file must include - `DecodeImage`, `ConvertRGB`, `Resize`, `Rescale`, `Normalize`, and - `PatchImage` transforms with correct attributes. - -**Fix for resize mismatch:** -```python -def _update_resize_for_image(processor_config_path, image_path): - """Recompute resize dimensions from actual image like HF does.""" - from PIL import Image - img = Image.open(image_path) - w, h = img.size - factor = 14 * 2 # patch_size * merge_size - new_w = round(w / factor) * factor - new_h = round(h / factor) * factor - # Update width/height in processor_config.json -``` - -## 2. Numerical divergence in greedy decoding - -**Symptoms:** First 1-3 tokens match HF, then output diverges. - -**Expected behavior:** This is inherent to ONNX vs PyTorch numerical -differences. ONNX models use different operator implementations that -accumulate small floating-point errors. - -**Typical metrics for Qwen2.5-VL 3B:** -- max_diff in logits: 5-10 -- mean_diff in logits: 0.5-1.5 -- cosine similarity: 0.98-0.99 -- First token: matches HF -- Greedy decoding: diverges at token 3-5 - -**This is NOT a bug** if the metrics above are within range. Both models -produce semantically similar descriptions. - -## 3. Vision model output shape mismatch - -**Symptoms:** Vision model produces wrong number of patches. - -**Debug:** Check `grid_thw` values: -```python -# For Qwen2.5-VL with merge_size=2: -t, h, w = grid_thw[0] -expected_patches = t * (h // 2) * (w // 2) -actual_patches = vision_output.shape[0] -assert expected_patches == actual_patches -``` - -## 4. Embedding model text-only failure - -**Symptoms:** Error when running without images (num_image_tokens=0). - -**Fix:** Ensure embedding model pads `image_features` with a zero row -before Gather, then uses a Where mask to select only real features: -```python -# Pad with zero row so Gather with index 0 doesn't fail -padded = op.Concat( - op.ConstantOfShape(...), # (1, hidden_size) zeros - image_features, - axis=0, -) -``` - -## 5. Vision encoder internal divergence (cos < 0.5) - -**Symptoms:** Vision features have very low cosine similarity (< 0.5) -against HuggingFace, even though patch embedding and weights are correct. - -**Debug with block-by-block comparison** (see extraction methods in -`references/extraction-methods.md`). Common root causes: - -1. **Wrong rotary embedding dimension** — Qwen2.5-VL vision uses 2D - position encoding (height + width). The rotary dim must be - `head_dim // 2`, not `head_dim`. Each half (head_dim // 4 frequencies) - covers one spatial dimension. With full `head_dim`, you get 2× too - many frequencies with wrong values. **Result: cos ≈ 0.25.** - -2. **Missing `fullatt_block_indexes`** — Qwen2.5-VL alternates between - windowed attention (local windows of `window_size` patches) and full - attention (all patches attend to all). Blocks at indexes `[7, 15, 23, 31]` - use full attention. If `fullatt_block_indexes` is not extracted from - HF config, all blocks use windowed attention. **Result: blocks 0-6 - are perfect (they're windowed anyway), but block 7+ diverges.** - -3. **Wrong attention bias construction** — Full-attention blocks should - have an all-zeros bias (everything attends to everything). Windowed - blocks have a block-diagonal bias. Check the bias by inspecting - sparsity: `(bias == -inf).float().mean()` should be ~0% for full - attention, ~98% for windowed. - -**Config extraction checklist for vision encoders:** -```python -# These fields MUST be extracted from HF vision_config: -fullatt_block_indexes = getattr(vc, "fullatt_block_indexes", None) -window_size = getattr(vc, "window_size", None) -spatial_merge_size = getattr(vc, "spatial_merge_size", 2) -temporal_patch_size = getattr(vc, "temporal_patch_size", 2) -``` - -## 6. ClippableLinear divergence (Gemma4) - -**Symptoms:** Vision or audio encoder output has large max diff (> 1.0) -against HuggingFace, even though weights are loaded correctly. - -**Root cause:** Gemma4 uses `Gemma4ClippableLinear` with learned finite -input/output activation clamping for ALL linear layers in its vision and -audio encoders. Using plain `Linear` misses the clamping. - -**Detection:** Check if the HuggingFace model uses `ClippableLinear`: -```bash -grep -n "ClippableLinear" transformers/models//modeling_.py -``` - -**Fix:** Use `ClippableLinear` (from `mobius.components`) for all -affected linear layers. For vision attention: q/k/v/o_proj. For MLP: -pass `linear_class=ClippableLinear` to the MLP component. - -**Impact:** -- Audio: max diff 52.68 → 0.0003 after fix -- Vision: max diff 3.92 → 0.00007 after fix - -## 7. Missing audio boundary markers - -**Symptoms:** Audio transcription is garbled or completely wrong, even -though the audio encoder output matches HuggingFace. - -**Root cause:** HuggingFace wraps audio placeholder tokens with boundary -markers in `input_ids`: -``` -<|audio> (256000) + N × <|audio|> (258881) + (258883) -``` -If boundary markers are missing, the model cannot distinguish audio -regions from text, producing wrong output. - -**Fix:** Add boundary markers to `build_input_ids()` when constructing -audio inputs, matching HuggingFace's token wrapping. - -## 8. CUDA EP: ORT Gather int32 overflow - -**Symptoms:** CUDA EP crashes or produces incorrect results for models -with large embedding tables. CPU EP works correctly. - -**Root cause:** ORT CUDA `gather_impl.cu` uses `int32` for element -offset computation: `input_index = idx * cols + col_offset`. For -tensors with > 2^31 elements (e.g. Gemma4 per-layer embedding -[262144, 8960] = 2.35B elements), this overflows. - -**ORT bug:** microsoft/onnxruntime#28107 - -**Workaround:** Split large embeddings into smaller tables via -`nn.ModuleList` so each individual Gather stays under the int32 limit. -Use Slice instead of Gather for column-wise indexing on large tensors. - -## 9. CUDA EP: opset 24 kernel registration - -**Symptoms:** ORT CUDA EP fails to find kernels for standard ops -(Squeeze, Reshape, etc.) even though they work on CPU. - -**Root cause:** ORT ≤1.24.x CUDA/TRT EPs don't register kernels for -opset 24, even though the op semantics are unchanged from opset 23. - -**Fix:** Use the `ort_lower_opset_for_ep` feature flag (enabled by -default) which lowers the declared opset import from 24 to 23 for -non-CPU EPs. See `src/mobius/_flags.py` and -`src/mobius/_testing/ort_inference.py`. - -## 10. Wrong audio feature extractor - -**Symptoms:** Audio model produces completely wrong output. Audio -encoder features don't match HuggingFace at all. - -**Root cause:** Using `WhisperFeatureExtractor` instead of the -model-specific feature extractor (e.g. `Gemma4AudioFeatureExtractor`). -Different extractors produce different mel spectrograms. - -**Fix:** Always use the correct feature extractor for the model: -```python -from transformers import AutoFeatureExtractor -feature_extractor = AutoFeatureExtractor.from_pretrained(model_id) -``` diff --git a/.agents/skills/multimodal-models/SKILL.md b/.agents/skills/multimodal-models/SKILL.md index 26b69f33..f3241135 100644 --- a/.agents/skills/multimodal-models/SKILL.md +++ b/.agents/skills/multimodal-models/SKILL.md @@ -141,6 +141,10 @@ Multimodal HF models often prefix text weights differently. Implement > instructions for adding a new multimodal model, vision config extraction > code, or the full model class template with InputMixer wiring. +> Read `references/scan-pattern.md` when building multi-image support using +> the ONNX Scan op, especially for variable-length per-image outputs that +> need the Scan + Padding + Compaction pattern. + ## genai_config.json required fields for VLMs **Required fields** (without these, VLM output is wrong): @@ -149,7 +153,7 @@ Multimodal HF models often prefix text weights differently. Implement - `spatial_merge_size`: Grid merge factor (2 for Qwen2.5-VL) See `.agents/skills/ort-genai-config/SKILL.md` for the complete reference -and `.agents/skills/debugging-vl-pipeline/SKILL.md` for troubleshooting. +and `.agents/skills/debugging-multimodal/SKILL.md` for troubleshooting. ### processor_config.json for image preprocessing @@ -169,9 +173,7 @@ instead of `Gather` for per-layer projection indexing. ## Cross-references -- **VL debugging:** `.agents/skills/debugging-vl-pipeline/SKILL.md` +- **Multimodal debugging:** `.agents/skills/debugging-multimodal/SKILL.md` - **ORT GenAI config:** `.agents/skills/ort-genai-config/SKILL.md` - **Weight name alignment:** `.agents/skills/weight-name-alignment/SKILL.md` -- **Multi-image Scan pattern:** `.agents/skills/scan-and-multi-image/SKILL.md` -- **Component parity debugging:** `.agents/skills/phi4mm-component-parity/SKILL.md` - **Reusable components (ClippableLinear):** `.agents/skills/reusable-components/SKILL.md` diff --git a/.agents/skills/multimodal-models/references/projector-variants.md b/.agents/skills/multimodal-models/references/projector-variants.md index 1adaf648..18e623fe 100644 --- a/.agents/skills/multimodal-models/references/projector-variants.md +++ b/.agents/skills/multimodal-models/references/projector-variants.md @@ -120,7 +120,7 @@ class VisionConfig: Both vision encoders support multiple images via the ONNX `Scan` op. Per-image values (position IDs, window indices, cu_seqlens) are computed -in a Scan body and concatenated. See `.agents/skills/scan-and-multi-image/SKILL.md`. +in a Scan body and concatenated. See `references/scan-pattern.md`. ### Spatial merge (post-encoder) diff --git a/.agents/skills/scan-and-multi-image/SKILL.md b/.agents/skills/multimodal-models/references/scan-pattern.md similarity index 95% rename from .agents/skills/scan-and-multi-image/SKILL.md rename to .agents/skills/multimodal-models/references/scan-pattern.md index 347fc058..52e83335 100644 --- a/.agents/skills/scan-and-multi-image/SKILL.md +++ b/.agents/skills/multimodal-models/references/scan-pattern.md @@ -1,16 +1,10 @@ ---- -name: scan-and-multi-image -description: > - Use this skill when building ONNX Scan or Loop subgraphs, or adding - multi-image support to a vision model. Covers the Scan + Padding + - Compaction pattern for variable-length per-image outputs, building Scan - body subgraphs with implicit inputs and carry states, the - rename-to-avoid-SSA-violations workaround, and the compact_scan_output - helper. Primary use case: multi-image vision models where per-image - computations produce variable output sizes. ---- - -# Skill: ONNX Scan Op & Multi-Image Vision +# ONNX Scan Op & Multi-Image Vision + +Reference for using the ONNX Scan op in multi-image vision models. +Covers the Scan + Padding + Compaction pattern for variable-length +per-image outputs, building Scan body subgraphs with implicit inputs +and carry states, the rename-to-avoid-SSA-violations workaround, and +the `compact_scan_output` helper. ## When to use diff --git a/.agents/skills/phi4mm-component-parity/SKILL.md b/.agents/skills/phi4mm-component-parity/SKILL.md deleted file mode 100644 index 50624db5..00000000 --- a/.agents/skills/phi4mm-component-parity/SKILL.md +++ /dev/null @@ -1,148 +0,0 @@ ---- -name: phi4mm-component-parity -description: > - Use this skill when building or adding a new multi-encoder multimodal - model (vision + language + audio) and verifying component-by-component - parity with HuggingFace. Covers pipeline isolation methodology, common - failure modes from real debugging experience (Phi4MM, Gemma4), the - step-by-step process for isolating which component (vision encoder, - speech encoder, embedding, or decoder) causes numerical divergence, - and integration test patterns. For debugging an existing deployed ORT - GenAI pipeline, use the debugging-vl-pipeline skill instead. ---- - -# Skill: Multimodal Component Parity Debugging - -## When to use - -Use this skill when: - -- A multimodal ONNX model's logits diverge systematically from HuggingFace -- You're adding a new multimodal model and need to verify each component -- Integration tests fail with large numerical differences (not just tolerance) -- Weights appear to load but the model produces wrong outputs -- You need to isolate which component (vision, speech, embedding, decoder) - is causing divergence - -For vision-language-only models (no speech), see also the -`debugging-vl-pipeline` skill which covers VLM-specific issues like 3D M-RoPE. - -## Pipeline isolation methodology - -Multimodal models with N encoders have N+2 stages (encoders + embedding + -decoder). Debug by comparing each stage independently against HuggingFace -at every boundary. - -### 4-model multimodal pipeline (e.g., Phi4MM) - -``` -pixel_values ──► [1. Vision Encoder] ──► image_features ──┐ - │ -audio_embeds ──► [2. Speech Encoder] ──► speech_features ──┤ - │ -input_ids ──► [3. Embedding/Fusion] ◄───────────────────┘ - │ - ▼ - inputs_embeds - │ - ▼ - [4. Decoder + LoRA] ──► logits -``` - -**Golden rule:** Start from the simplest case (text-only, no encoders), -verify it matches HF, then add one modality at a time. - -### Stage-by-stage comparison - -**Stage 1 — Vision encoder:** Compare SigLIP/ViT output. Check output -shape `(num_image_tokens, text_hidden_size)`, projection MLP correctness, -and position embeddings (2D vs 3D shape). Target: cos_sim > 0.99. - -**Stage 2 — Speech encoder:** Compare Conformer output. Check compression -rate (typically 8× time reduction), projection branch selection, and conv -subsampling output length. - -**Stage 3 — Embedding/fusion:** Compare token embeddings. Text-only should -match HF `embed_tokens` exactly (< 1e-5). With features, verify token -replacement at correct positions. InputMixer must handle zero-length tensors. - -**Stage 4 — Decoder:** Compare logits. Acceptable float32 metrics: -max_diff 5-10, mean_diff 0.5-1.5, cos_sim > 0.98, argmax match exact. - -## Quick-start 4-step process - -1. **Text-only baseline:** Build ONNX → run embedding + decoder (skip - encoders) → compare against HF `embed_tokens` and full forward logits. - If this diverges, fix weight loading / decoder before touching encoders. - -2. **Add vision:** Run vision encoder → feed features to embedding → - compare. If newly divergent, isolate vision encoder output vs HF. - -3. **Add audio:** Same as above for speech encoder. - -4. **Combined:** All modalities together. If divergent only in combined - mode, suspect LoRA mode mismatch or embedding fusion ordering. - -## Gotchas - -### Module forward() bypass - -The #1 source of missing weights. Directly accessing nested sub-module -parameters (`self.glu.ext_pw_conv_1d.weight`) instead of calling -`self.glu(op, x)` makes onnxscript unable to resolve the full module path. -**Detection:** Conv/MatMul nodes where weight inputs have -`is_initializer=False` and generic names. - -### LoRA mode mismatch - -Some models (Phi4MM) apply LoRA conditionally per modality. If ONNX applies -all adapters unconditionally, run HF reference with `input_mode=3` to match. -**Detection:** Text-only inference diverges but output is reasonable (not -garbage). - -### ClippableLinear omission (Gemma4) - -Using plain `Linear` instead of `ClippableLinear` in Gemma4 vision/audio -encoders causes max_diff 52.68 (audio) / 3.92 (vision). Check HF source -for `ClippableLinear` usage. - -### Empty tensor handling - -Text-only inference crashes when no image/audio features are present. -**Fix:** Zero-pad `image_features` before Gather, then mask with Where. - -### Missing boundary tokens - -Audio/image boundary markers (`<|audio>`, ``) are required for -correct modality region identification. Missing markers → garbled output -even with correct encoder output. - -> Read `references/common-failures.md` when you need detailed code examples -> for each failure mode, including weight name alignment (ModuleList subclass -> name doubling, setattr with dotted names), shape mismatches, dtype -> mismatches (float64 vs float32), and HD transform format issues. - -> Read `references/debugging-cookbook.md` when you need step-by-step -> debugging procedures with code for each phase, intermediate value -> extraction methods, integration test patterns (text-only, audio, vision), -> tolerance guidelines, weight loading verification, and HD multi-crop -> verification. - -## Reference files - -- **Integration tests:** `tests/phi4mm_integration_test.py`, - `tests/integration_test.py` -- **VL debugging skill:** `.agents/skills/debugging-vl-pipeline/SKILL.md` -- **Model implementation:** `src/mobius/models/phi.py`, - `src/mobius/models/gemma4.py` -- **Audio components:** `src/mobius/components/_audio.py`, - `src/mobius/components/_gemma4_audio.py` -- **Vision components:** `src/mobius/components/_vision.py` -- **ClippableLinear:** `src/mobius/components/_gemma4_audio.py` -- **LoRA component:** `src/mobius/components/_lora.py` -- **Weight loading:** `src/mobius/_weight_loading.py` -- **Feature flags:** `src/mobius/_flags.py` -- **ORT GenAI config skill:** `.agents/skills/ort-genai-config/SKILL.md` -- **Weight name alignment skill:** - `.agents/skills/weight-name-alignment/SKILL.md` -- **Gemma4 example:** `examples/gemma4_multimodal.py` diff --git a/.agents/skills/phi4mm-component-parity/references/common-failures.md b/.agents/skills/phi4mm-component-parity/references/common-failures.md deleted file mode 100644 index ed82efa0..00000000 --- a/.agents/skills/phi4mm-component-parity/references/common-failures.md +++ /dev/null @@ -1,244 +0,0 @@ -# Common Failure Modes — Detailed Reference - -## 1. Weight name alignment (missing weights) - -**Symptoms:** Hundreds or thousands of weights reported as "unmatched" by -`apply_weights`. Model runs but produces garbage output. - -**Root causes encountered:** - -### a. Module forward() bypass (240 missing weights in Phi4MM) - -The most insidious bug. When a component's `forward()` method directly -accesses nested sub-module parameters (e.g., `self.glu.ext_pw_conv_1d.weight`) -instead of calling the sub-module's `forward()` method, onnxscript cannot -resolve the full module path for the parameter. The weight ends up as an -unnamed, non-initializer constant in the graph. - -```python -# BAD — weights become unnamed -def forward(self, op, x): - return op.Conv(x, self.glu.ext_pw_conv_1d.weight, - self.glu.ext_pw_conv_1d.bias, ...) - -# GOOD — onnxscript resolves full module path -def forward(self, op, x): - return self.glu(op, x) # GLU.forward() calls op.Conv internally -``` - -**Detection:** Check the ONNX graph for Conv/MatMul nodes where weight -inputs have `is_initializer=False` and generic names like "weight"/"bias". - -**Fix:** Add `forward()` methods to sub-modules and call them instead of -directly accessing their parameters. - -### b. ModuleList subclass causing name doubling (8 weights) - -Subclassing `nn.ModuleList` causes the module's own name to appear twice -in the parameter path: `img_projection.img_projection.0.weight` instead -of `img_projection.0.weight`. - -```python -# BAD — name doubling -class ProjectionMLP(nn.ModuleList): - def __init__(self): - super().__init__() - self.append(nn.Linear(1152, 3072)) - self.append(nn.Linear(3072, 3072)) - -# GOOD — use nn.Module with indexed children -class ProjectionMLP(nn.Module): - def __init__(self): - super().__init__() - layers = [nn.Linear(1152, 3072), nn.Linear(3072, 3072)] - for i, layer in enumerate(layers): - setattr(self, str(i), layer) -``` - -### c. setattr with dotted names - -Using `setattr(self, "audio_projection.speech", module)` creates a single -attribute with a dot in its name, rather than a nested module. The resulting -ONNX parameter names won't match HuggingFace's `ModuleDict`-style naming. - -**Fix:** Use `nn.ModuleDict` or create proper nested attributes. - -## 2. Shape mismatches (position embedding 2D vs 3D) - -**Symptoms:** `RuntimeError: shape mismatch` during weight loading. - -**Root cause:** The ONNX component declares a parameter with a different -number of dimensions than the HuggingFace weight. Example: PatchEmbedding -declares `position_embedding.weight` as `[num_patches, hidden_size]` (2D), -but HF stores `[1, num_patches, hidden_size]` (3D). - -**Fix in `preprocess_weights()`:** -```python -# Squeeze the extra batch dimension to match ONNX declaration -if "position_embedding.weight" in key and state_dict[key].dim() == 3: - state_dict[key] = state_dict[key].squeeze(0) # [1,N,H] → [N,H] -``` - -**General rule:** Check whether the preprocess_weights transform goes -in the correct direction (squeeze vs unsqueeze). A common mistake is -writing the transform backwards. - -## 3. Dtype mismatches (float64 vs float32) - -**Symptoms:** ONNX Runtime error: "type mismatch in Mul/Add node" during -inference. - -**Root causes:** - -### a. NumPy default float64 - -`numpy.array(python_float)` defaults to float64. Any constant created -from a Python scalar without explicit dtype will be float64 in the graph. - -```python -# BAD — float64 constant -scale = numpy.array(alpha / rank) # defaults to float64 -op.Mul(x, scale) # Mul(float32, float64) → type error - -# GOOD — explicit float32 -scale = numpy.array(alpha / rank, dtype=numpy.float32) -op.Mul(x, scale) -``` - -### b. Python int auto-promotion - -When passing a Python `int` to an op that expects a tensor, onnxscript -may auto-promote to float64 (implementation-dependent). - -```python -# RISKY — Python int may become float64 -op.Mul(int64_tensor, self.max_position_embeddings) - -# SAFE — explicit constant -op.Mul(int64_tensor, - op.Constant(value_int=self.max_position_embeddings)) -``` - -**Detection:** Run the ONNX model and look for type mismatch errors. -The error message includes the node name — trace it back to the source. - -## 4. LoRA application mismatch (conditional vs unconditional) - -**Symptoms:** Systematic divergence (> 80% logits mismatch) across ALL -test cases, but the model structurally runs correctly. - -**Root cause:** Some models apply LoRA adapters conditionally based on -input modality. For example, Phi4MM applies: -- `input_mode=0` (text): no adapters -- `input_mode=1` (vision): vision LoRA only -- `input_mode=2` (speech): speech LoRA only -- `input_mode=3` (combined): both adapters - -If the ONNX model unconditionally applies all adapters (both vision and -speech LoRA always active), it diverges from HF when HF uses a different -input mode. - -**Quick fix for integration tests:** Run the HF reference with the mode -that matches the ONNX model's behavior (e.g., `input_mode=3` to match -unconditional application of both adapters). - -**Proper fix:** Add an `input_mode` input to the decoder model and use -conditional logic to selectively apply adapters. - -**Detection:** If text-only inference diverges but the model generates -reasonable (not garbage) output, suspect LoRA mode mismatch. Temporarily -zero out all LoRA weights — if base model matches HF perfectly, the -LoRA application mode is the issue. - -## 5. Empty tensor handling (zero-length features) - -**Symptoms:** Crash during text-only inference when no image/audio -features are present. - -**Root cause:** The embedding model's `InputMixer` uses `GatherElements` -to place features at special token positions. With zero features, the -gather indices are empty but the operation may still execute on the -padded dimension, causing shape errors. - -**Fix pattern:** Zero-pad before Gather, then use Where to mask results: -```python -# Pad with one zero row so Gather never accesses out-of-bounds -padded = op.Concat( - op.ConstantOfShape(op.Constant(value_ints=[1, hidden_size])), - features, # may be [0, hidden_size] - axis=0, -) -# After Gather, mask out the padding positions with Where -result = op.Where(feature_mask, gathered, text_embeddings) -``` - -## 6. HD transform image format (5D vs 4D) - -**Symptoms:** Vision model crashes or produces wrong output with multi-crop -HD images. - -**Root cause:** HD-capable vision models expect images in different formats: -- Some expect `[batch, channels, height, width]` (4D, single crop per batch) -- Others expect `[num_images, num_crops, channels, height, width]` (5D) - -The HF processor output format must match the ONNX model's input format. -If using the HF processor for test input preparation, verify it produces -the expected format. - -**Fix:** Check the HF model's preprocessing code for the expected format, -and ensure the ONNX model's input signature matches. For tests, either: -- Use the HF processor: `processor(images=image, return_tensors="np")` -- Or manually construct the correct format for simple test cases - -## 7. Causal mask construction (inputs_embeds vs input_ids) - -**Symptoms:** Attention mask has wrong length, causing decoder crash or -wrong output. - -**Root cause:** When the decoder receives `inputs_embeds` instead of -`input_ids`, the sequence length must be derived from the embeds tensor -shape, not from input_ids. If the mask is built from input_ids length but -the actual sequence includes fused image/audio tokens, the lengths diverge. - -**Fix:** Always derive `seq_len` from `inputs_embeds.shape[1]` when the -model uses inputs_embeds as input. - -## 8. ClippableLinear not used in encoder (Gemma4) - -**Symptoms:** Encoder output has large numerical divergence (max diff > 1.0) -from HuggingFace, despite all weights loading correctly. - -**Root cause:** Some HuggingFace models (e.g. Gemma4) use -`ClippableLinear` — a linear layer with learned finite input/output -activation clipping — for ALL linear layers in their vision and audio -encoders (attention q/k/v/o projections AND MLP gate/up/down projections). -Using plain `Linear` misses the clamping and causes divergence. - -**Detection:** Check HF source for `ClippableLinear`: -```bash -grep -n "ClippableLinear" transformers/models//modeling_.py -``` - -**Fix:** Use `ClippableLinear` from `mobius.components`: -- For attention: use `ClippableLinear` for q/k/v/o_proj -- For MLP: pass `linear_class=ClippableLinear` parameter - -**Impact (Gemma4):** -- Audio: max diff 52.68 → 0.0003 -- Vision: max diff 3.92 → 0.00007 - -## 9. Missing audio/image boundary tokens - -**Symptoms:** Audio transcription or vision description is garbled, but -encoder output matches HuggingFace. - -**Root cause:** HuggingFace wraps modality placeholder tokens with -boundary markers: -- Image: `<|image>` (open) + N × `<|image|>` (pad) + `` (close) -- Audio: `<|audio>` (open) + N × `<|audio|>` (pad) + `` (close) - -Missing boundary markers prevent the model from correctly identifying -modality regions in the input sequence. - -**Fix:** Ensure `build_input_ids` wraps placeholder tokens with the -correct open/close marker token IDs. From eb656a25dafb518d5ac707ba98ff963b8fa16851 Mon Sep 17 00:00:00 2001 From: Justin Chu Date: Tue, 21 Apr 2026 16:26:16 +0000 Subject: [PATCH 6/8] Remove old .github/skills/ and update path references The skills were moved to .agents/skills/ in prior commits but the old directory was not cleaned up. This commit: - Deletes the 14 original .github/skills/ SKILL.md files - Updates path references in copilot-instructions, README, CONTRIBUTING, CHANGELOG, docs, and tests to point to .agents/skills/ Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Justin Chu --- .github/copilot-instructions.md | 2 +- .github/skills/adding-a-new-model/SKILL.md | 980 ------------------ .github/skills/debugging-vl-pipeline/SKILL.md | 533 ---------- .github/skills/diffusion-models/SKILL.md | 388 ------- .github/skills/moe-models/SKILL.md | 383 ------- .../skills/multi-agent-coordination/SKILL.md | 209 ---- .github/skills/multimodal-models/SKILL.md | 707 ------------- .github/skills/ort-genai-config/SKILL.md | 799 -------------- .../skills/phi4mm-component-parity/SKILL.md | 618 ----------- .github/skills/quality-checklist/SKILL.md | 243 ----- .github/skills/reusable-components/SKILL.md | 698 ------------- .github/skills/scan-and-multi-image/SKILL.md | 392 ------- .github/skills/weight-name-alignment/SKILL.md | 416 -------- .github/skills/writing-rewrite-rules/SKILL.md | 301 ------ .github/skills/writing-tests/SKILL.md | 699 ------------- CHANGELOG.md | 2 +- CONTRIBUTING.md | 2 +- README.md | 2 +- docs/ai-model-support-strategy.md | 4 +- docs/getting-started.md | 2 +- tests/model_coverage_test.py | 2 +- 21 files changed, 8 insertions(+), 7374 deletions(-) delete mode 100644 .github/skills/adding-a-new-model/SKILL.md delete mode 100644 .github/skills/debugging-vl-pipeline/SKILL.md delete mode 100644 .github/skills/diffusion-models/SKILL.md delete mode 100644 .github/skills/moe-models/SKILL.md delete mode 100644 .github/skills/multi-agent-coordination/SKILL.md delete mode 100644 .github/skills/multimodal-models/SKILL.md delete mode 100644 .github/skills/ort-genai-config/SKILL.md delete mode 100644 .github/skills/phi4mm-component-parity/SKILL.md delete mode 100644 .github/skills/quality-checklist/SKILL.md delete mode 100644 .github/skills/reusable-components/SKILL.md delete mode 100644 .github/skills/scan-and-multi-image/SKILL.md delete mode 100644 .github/skills/weight-name-alignment/SKILL.md delete mode 100644 .github/skills/writing-rewrite-rules/SKILL.md delete mode 100644 .github/skills/writing-tests/SKILL.md diff --git a/.github/copilot-instructions.md b/.github/copilot-instructions.md index 700e0a9a..8c4e4a4c 100644 --- a/.github/copilot-instructions.md +++ b/.github/copilot-instructions.md @@ -90,7 +90,7 @@ and wrapper modules for nesting. See the `weight-name-alignment` skill. to source) - Top-level tests: `tests/build_graph_test.py` (graph construction), `tests/integration_test.py` (numerical accuracy vs HuggingFace) -- Skills: `.github/skills//SKILL.md` +- Skills: `.agents/skills//SKILL.md` ### Code style diff --git a/.github/skills/adding-a-new-model/SKILL.md b/.github/skills/adding-a-new-model/SKILL.md deleted file mode 100644 index 9fe78664..00000000 --- a/.github/skills/adding-a-new-model/SKILL.md +++ /dev/null @@ -1,980 +0,0 @@ ---- -name: adding-a-new-model -description: > - Step-by-step guide for adding a new HuggingFace model architecture to the - mobius package. Covers config extraction, model class creation, - registry registration, weight preprocessing, and testing. Use this skill - when the user wants to add support for a new model architecture (LLM, - encoder-only, encoder-decoder, vision, audio, diffusion, or multimodal). ---- - -# Skill: Adding a New Model - -## When to use - -Use this skill when adding support for a new HuggingFace model architecture -(e.g. a new LLM family, vision model, encoder-decoder, audio model, or -diffusion component) to the `mobius` package. - -## Prerequisites - -- Identify the HuggingFace `model_type` string (from the model's `config.json`) -- Find a small checkpoint on HuggingFace Hub for testing -- Have the HuggingFace `transformers` source available to reference the - PyTorch implementation - -## Step-by-step - -### 1. Check if the base `CausalLMModel` already works - -Many models (LLaMA, Mistral, Qwen2, DeepSeek) use the standard decoder-only -architecture with no special components. Before writing a custom class, -check whether `CausalLMModel` from `models/base.py` produces correct results: - -```python -from mobius._registry import registry -from mobius.models.base import CausalLMModel - -registry.register("my_model_type", CausalLMModel) -model = build("org/my-model-id", load_weights=True) -``` - -If the logits match HuggingFace, you only need the registry entry. - -### 2. Create the model file - -Create `src/mobius/models/.py`. The minimal template: - -```python -# Copyright (c) Microsoft Corporation. -# Licensed under the MIT License. - -from __future__ import annotations - -import torch -from onnxscript import nn -from onnxscript._internal import builder - -from mobius._configs import ArchitectureConfig -from mobius.components import ( - Attention, DecoderLayer, Embedding, Linear, MLP, RMSNorm, - create_attention_bias, initialize_rope, -) -from mobius.models.base import CausalLMModel - - -class MyTextModel(nn.Module): - """Text model for MyArchitecture.""" - - def __init__(self, config: ArchitectureConfig): - super().__init__() - self.embed_tokens = Embedding(config.vocab_size, config.hidden_size) - self.layers = [MyDecoderLayer(config) for _ in range(config.num_hidden_layers)] - self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) - self.rotary_emb = initialize_rope(config) - - def forward(self, op, input_ids, attention_mask, position_ids, past_key_values=None): - hidden_states = self.embed_tokens(op, input_ids) - position_embeddings = self.rotary_emb(op, position_ids) - attention_bias = create_attention_bias(op, input_ids=input_ids, attention_mask=attention_mask) - - present_key_values = [] - past_kvs = past_key_values or [None] * len(self.layers) - for layer, past_kv in zip(self.layers, past_kvs): - hidden_states, present_kv = layer( - op, hidden_states=hidden_states, - attention_bias=attention_bias, - position_embeddings=position_embeddings, - past_key_value=past_kv, - ) - present_key_values.append(present_kv) - - hidden_states = self.norm(op, hidden_states) - return hidden_states, present_key_values - - -class MyCausalLMModel(CausalLMModel): - """Causal LM wrapper for MyArchitecture.""" - - def __init__(self, config: ArchitectureConfig): - # Skip CausalLMModel.__init__ to use custom TextModel - nn.Module.__init__(self) - self.config = config - self.model = MyTextModel(config) - self.lm_head = Linear(config.hidden_size, config.vocab_size, bias=False) -``` - -#### Class metadata attributes - -Every registered model class should set two class-level attributes used for -**auto-task detection** and **documentation generation**: - -| Attribute | Type | Default (from `CausalLMModel`) | Description | -|-----------|------|------|-------------| -| `default_task` | `str` | `"text-generation"` | Task auto-selected by `build()` / CLI when no `--task` is specified. | -| `category` | `str` | `"Text Generation"` | Grouping label in the generated docs. | - -The `_default_task_for_model()` function reads `default_task` from the -registered class, so there is **no hardcoded task map** to maintain. -The doc generator (`docs/_generate_models.py`) reads both attributes to -produce per-model pages and the categorised index automatically. - -Override these when the model isn't a standard text-generation model: - -```python -class MyMultiModalModel(nn.Module): - """My vision-language model.""" - - default_task: str = "vision-language" - category: str = "Multimodal" -``` - -Standard categories: `"Text Generation"`, `"Mixture of Experts"`, -`"Multimodal"`, `"Speech-to-Text"`, `"Audio"`, `"Diffusion"`, -`"autoencoder"`, `"encoder-only"`, `"encoder"`, `"encoder-decoder"`, -`"vision"`, `"causal-lm"`. New categories are picked up -automatically by the doc generator. - -### 3. Identify what's different - -Compare the HuggingFace PyTorch source against `CausalLMModel` / `DecoderLayer`. -Common variations to look for: - -| Variation | Example | Solution | -|-----------|---------|----------| -| Custom norm (weight + 1) | Gemma | Subclass `RMSNorm` | -| Embedding scaling | Gemma (`* sqrt(d)`) | Subclass `Embedding` | -| Extra norms (pre/post feedforward) | Gemma2, Gemma3 | Custom `DecoderLayer` | -| QK normalization | Gemma3, Qwen3 | Set `attn_qk_norm=True` in config | -| Sliding window attention | Gemma2, Gemma3 | Alternating layer types + `sliding_window` config | -| Different activation | Various | Set `hidden_act` in config (handled by `MLP`) | -| Biased attention projections | Phi, PhiMoE | Set `attn_qkv_bias=True`, `attn_o_bias=True` | -| LayerNorm epsilon | Whisper (`1e-5`) | Pass eps from config to `LayerNorm(hidden_size, eps=...)` — default `1e-6` is wrong for many models | -| Q pre-scaling | Whisper | Multiply Q by `head_dim**-0.5` before Attention op, set `scale=1.0` in op | -| Causal self-attention (decoder) | Whisper | Set `is_causal=1` attribute on Attention op instead of explicit mask | -| Weight-free LayerNorm | OLMo-1B | Use `_WeightFreeLayerNorm` (constant scale=1, bias=0) — OLMo uses `F.layer_norm` with `weight=None, bias=None, eps=1e-5` | -| Custom attention scale | Granite | Pass `scale=config.attention_multiplier` to `Attention(config, scale=...)` instead of default `1/sqrt(head_dim)` | -| Embedding multiplier | Granite (`* 12.0`) | Multiply embeddings by `config.embedding_multiplier` in `TextModel.forward` after embed lookup | -| Logits scaling | Granite (`/ 8.0`) | Divide logits by `config.logits_scaling` in `CausalLMModel.forward` | -| Residual scaling | Granite (`* 0.22`) | Apply `residual + output * residual_multiplier` — **NOT** `residual * multiplier + output` | -| MoE layers | PhiMoE, GPTOSS | See the MoE skill | -| Vision encoder | Gemma3 | See the multimodal skill | -| Gated attention output | Qwen3.5 (`attn * sigmoid(gate)`) | Subclass `Attention` with doubled q_proj → Q+gate split | -| OffsetRMSNorm (1+weight) | Qwen3.5 | Use `OffsetRMSNorm` from `components/_rms_norm.py` | -| Hybrid layer types | Qwen3.5 (DeltaNet + full attention) | Use `config.layer_types` list to dispatch per-layer | -| Linear attention (DeltaNet) | Qwen3.5 | Use `GatedDeltaNet` component, stateless export | -| Shared expert with sigmoid gate | Qwen3.5-MoE | Custom `Qwen35MoEBlock` with `shared_expert_gate` Linear(hidden, 1) | -| Interleaved M-RoPE | Qwen3.5 | Use `InterleavedMRope` (not `ChunkedMRope`) | -| Fused QKV projection | ModernBERT | Split `Wqkv` [3H, H] into q/k/v in `preprocess_weights` | -| Fused gate+up (GeGLU/SwiGLU) | ModernBERT | Split `Wi` [2I, H] into gate/up in `preprocess_weights` | -| Pre-norm encoder with RoPE | ModernBERT | Custom encoder model — BERT is post-norm, CausalLM is causal | -| Hierarchical multi-scale | SAM2, Segformer | Per-stage embed dims/heads, dimension projection at stage boundaries | -| Efficient attention (seq reduction) | Segformer | Strided Conv2d on K/V to reduce sequence length before attention | -| DPT decoder (dense prediction) | Depth Anything | Multi-index backbone extraction + reassemble + fusion + prediction head | -| Detection tokens + heads | YOLOS | Learnable tokens appended to ViT sequence, MLP class/bbox heads | -| Subclass-only (weight rename) | BLIP, TrOCR, LayoutLMv3 | Override only `preprocess_weights` — lowest effort for compatible architectures | - -### 4. Handle weight name mismatches (`preprocess_weights`) - -If HuggingFace uses different weight names than your component tree, override -`preprocess_weights`: - -```python -class MyCausalLMModel(CausalLMModel): - def preprocess_weights(self, state_dict: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]: - # Example: rename keys - renamed = {} - for key, value in state_dict.items(): - new_key = key.replace("old_prefix.", "new_prefix.") - renamed[new_key] = value - # Call parent for weight tying - return super().preprocess_weights(renamed) -``` - -**Common operations:** - -- Strip prefixes: `language_model.model.X` → `X` (multimodal models) -- Rename expert weights: `w1` → `gate_proj` (MoE models) -- Weight tying: copy `embed_tokens.weight` → `lm_head.weight` -- Split fused QKV: `Wqkv.weight` [3H, H] → `q_proj`, `k_proj`, `v_proj` (ModernBERT) -- Split fused gate+up: `Wi.weight` [2I, H] → `gate_proj`, `up_proj` (ModernBERT) -- Split fused BLIP QKV: `in_proj_weight` [3H, H] → separate Q/K/V projections - -**Tip:** Build the ONNX model, list its initializer names, and compare against -the HuggingFace state dict keys to find mismatches: - -```python -model_names = set(onnx_model.graph.initializers.keys()) -hf_names = set(state_dict.keys()) -print("In HF but not model:", hf_names - model_names) -print("In model but not HF:", model_names - hf_names) -``` - -### 5. Register the model - -Add to `_create_default_registry()` in `src/mobius/_registry.py`: - -```python -from mobius.models import MyCausalLMModel - -# In _create_default_registry(): -reg.register("my_model_type", MyCausalLMModel) -``` - -Also export from `src/mobius/models/__init__.py`. - -### 6. Update `ArchitectureConfig.from_transformers` if needed - -If the model has unusual config fields, update `from_transformers()` in -`_configs.py`. For example, nested RoPE configs or custom head-dim formulas. - -### 7. Write tests - -See the **writing-tests** skill for full details. At minimum: - -1. **Add a config entry to `tests/_test_configs.py`** in the appropriate - group (`CAUSAL_LM_CONFIGS`, `ENCODER_CONFIGS`, `SEQ2SEQ_CONFIGS`, - `VISION_CONFIGS`, or `DETECTION_CONFIGS`): - ```python - # In tests/_test_configs.py: - CAUSAL_LM_CONFIGS: list[tuple[str, dict, bool]] = [ - # ... - ("my_model_type", {"hidden_act": "gelu"}, True), - ] - ``` - - Set `is_representative=True` if the model has unique behaviour - (custom model class, special config like softcapping, partial rotary, - ALiBi, MoE, etc.) - - Set `is_representative=False` if it's an alias of an existing base - class with no special config overrides - - Text-generation models with no special config are auto-generated from - the registry, but explicit entries are preferred for documentation - - Then verify: - ```bash - pytest tests/build_graph_test.py -k "my_model_type" - ``` - -2. **Add a small model to `tests/integration_test.py`** if a small HuggingFace - checkpoint exists (< 1 B parameters preferred): - ```python - # In _TEXT_MODELS list: - pytest.param("org/my-small-model", False, id="my-model"), - ``` - -3. **Testing large models with random weights:** - - When the smallest available checkpoint is too large for CI (e.g. Qwen3.5 - at 27B), create a HF model with random weights and reduced layers: - - ```python - from transformers import AutoConfig - from transformers.models.qwen3_5.modeling_qwen3_5 import Qwen3_5ForCausalLM - - c = AutoConfig.from_pretrained("Qwen/Qwen3.5-27B") - tc = c.text_config - tc.num_hidden_layers = 4 - tc.layer_types = tc.layer_types[:4] # Must match num_hidden_layers - hf_model = Qwen3_5ForCausalLM._from_config(tc, dtype=torch.float32) - ``` - - Then use `build_from_module` with the HF state dict to compare logits. - See `test_qwen35_prefill_logits_match` in `tests/integration_test.py`. - -4. **Test with the CLI** to verify end-to-end build works: - ```bash - # Single model - mobius build --model org/my-small-model mymodel/output - - # Multi-component (encoder-decoder) — produces separate encoder.onnx + decoder.onnx - mobius build --model org/my-small-model mymodel/output - ``` - The task is auto-detected from the model type. Use `--task` to override - if needed (e.g. `--task text-generation`). - -### 8. Documentation - -Model documentation is **auto-generated** from class metadata by -`docs/_generate_models.py` at doc-build time. No manual doc update is needed -if you set the `default_task`, `category`, and a good class docstring (the -first paragraph is used as the model description). - -The README model list is a curated highlight table — update it only if the -new model is a significant addition. - -## Checklist - -This is the **implementation** checklist — the steps needed to wire up a new -model in the codebase. For the full **definition-of-done** quality checklist -(L1–L5 tests, ORT GenAI, Foundry Local, Olive quantization, multi-dtype, -code review), see the -[quality-checklist skill](../quality-checklist/SKILL.md). - -- [ ] Model file in `src/mobius/models/` with Microsoft MIT copyright header -- [ ] Class has `default_task` and `category` attributes (if not standard text-generation) -- [ ] Class has a descriptive docstring (first paragraph used in generated docs) -- [ ] `preprocess_weights` handles any key mismatches -- [ ] Registered in `_create_default_registry()` -- [ ] Exported from `models/__init__.py` -- [ ] Config extraction works (`ArchitectureConfig.from_transformers`) -- [ ] Tiny config in `tests/_test_configs.py` (with `is_representative` flag) -- [ ] L2 YAML test case in `testdata/cases/` with `test_model_id` -- [ ] L3 synthetic parity passes (`tests/synthetic_parity_test.py -k ""`) -- [ ] Integration test model in `tests/integration_test.py` (real-weight integration suite, if small checkpoint available) -- [ ] L4 golden file generated and committed (`testdata/golden/`) -- [ ] L5 generation golden file generated and committed -- [ ] ORT GenAI test added to `tests/ort_genai_test.py` (text-generation and VLM models) -- [ ] CLI build works (`mobius build --model ...`) -- [ ] Multi-dtype correctness verified (fp32, fp16, bf16) - -**Note:** Default optimizer passes (CSE, deduplicate initializers, identity -elimination, remove unused nodes/opsets) are applied automatically by -`build_from_module`. - -## Example: minimal diff for a LLaMA-compatible model - -If the new model is fully LLaMA-compatible, the entire change is: - -```python -# _registry.py -reg.register("my_llama_variant", CausalLMModel) -``` - -No new model file needed. - -## Example: adding a non-LLM model - -For models that aren't causal LMs, follow the same steps but use the -appropriate base class and task: - -| Model type | Base class / pattern | Task | Config | -|------------|---------------------|------|--------| -| Encoder-only (BERT-like) | `BertModel` | `feature-extraction` | `ArchitectureConfig` | -| Encoder-only (ModernBERT) | `ModernBertModel` | `feature-extraction` | `ArchitectureConfig` | -| Encoder-decoder (BART/T5-like) | `BartForConditionalGeneration` or `T5ForConditionalGeneration` | `seq2seq` | `ArchitectureConfig` | -| Vision (ViT-like) | `ViTModel` or `CLIPVisionModel` | `image-classification` | `ArchitectureConfig` | -| Object detection | `YolosForObjectDetection` | `object-detection` | `ArchitectureConfig` | -| Depth estimation | `DepthAnythingForDepthEstimation` | `image-classification` | `ArchitectureConfig` | -| Segmentation | `SegformerForSemanticSegmentation` or `Sam2VisionModel` | `image-classification` | `ArchitectureConfig` | -| Audio encoder (Wav2Vec2-like) | `Wav2Vec2Model` | `audio-feature-extraction` | `ArchitectureConfig` | -| Multimodal (LLaVA-like) | `LLaVAModel` | `vision-language` | `ArchitectureConfig` | -| Document AI | `LayoutLMv3Model` | `feature-extraction` | `ArchitectureConfig` | -| OCR decoder | `TrOCRForConditionalGeneration` | `seq2seq` | `ArchitectureConfig` | -| Diffusion denoiser | Custom (`UNet2DConditionModel`, etc.) | `denoising` | Custom config (e.g. `UNet2DConfig`) | -| VAE | `AutoencoderKLModel` | `vae` | `VAEConfig` | -| Adapter | `T2IAdapterModel` / `IPAdapterModel` | `adapter` | Custom config | - -Many new models can be registered as aliases of existing classes (e.g. -`reg.register("my_bert_variant", BertModel)`) if the architecture matches. - -## False Compatibility Pitfalls - -When registering models as aliases of existing base classes, **tests passing -does not mean the mapping is correct.** Graph-build tests only check that an -ONNX graph can be constructed — they do NOT verify that the graph matches -the model's actual computation. - -### Safe approximate mappings - -The project accepts "approximate" registry aliases when the model uses -similar-but-not-identical attention. These produce structurally correct ONNX -graphs; weight-loading may need minor adjustments: - -| Model | Maps to | Why it works | -|-------|---------|-------------| -| DeBERTa | `BertModel` | Disentangled attention is a variant of standard attention | -| Swin | `ViTModel` | Shifted window attention is still self-attention over patches | -| SqueezeBERT | `BertModel` | Grouped convolution replaces dense attention, but same I/O shape | - -### NEVER safe as registry aliases - -These model families have fundamentally different computation that **cannot** -be represented by standard base classes, even though `build_graph_test` passes: - -| Category | Models | Why it fails | -|----------|--------|-------------| -| Pure CNNs | ConvNeXt, ResNet, MobileNet, EfficientNet, RegNet | No attention at all — base ViT/BERT classes produce attention-based graphs | -| Spatial pooling | PoolFormer | Uses spatial average pooling instead of attention — structurally incompatible | -| SSM / state-space models | Mamba, Mamba2, FalconMamba, RWKV, RecurrentGemma | Sequential scan / linear recurrence, not attention | -| Fundamentally different attention | Longformer (sparse), BigBird (block sparse), Funnel (downsampling) | Attention pattern differs from dense self-attention at a structural level | -| Custom tokenization | CANINE (character-level) | Byte-level input, hash embeddings — not a standard vocab embedding | - -**Rule of thumb:** If the HuggingFace model's `forward()` method doesn't call -`self_attn(query, key, value)` in a standard way, it is NOT a safe alias. - -### Future work - -CI currently only runs graph-build tests (shape inference, op validity). To -catch false compatibility in approximate mappings, we need **weight-loading -tests** that: -1. Load real HuggingFace weights into the ONNX graph -2. Run inference on a test input -3. Compare output against HuggingFace PyTorch output -4. Fail if max abs diff exceeds a threshold (e.g. 0.01) - -This would catch shape mismatches, wrong norm types, and missing scaling -factors that graph-build tests cannot detect. - -## Troubleshooting: common pitfalls - -This section documents real bugs found during model integration and how to -diagnose them. - -### 1. ORT rejects graph with `tensor(double)` / wrong dtype - -**Symptom:** `onnxruntime.capi.onnxruntime_pybind11_state.InvalidGraph` error -mentioning `Type 'tensor(double)'` for an input to a custom op like -`RotaryEmbedding`. - -**Root cause:** NumPy creates float64 arrays by default when given Python -floats/lists. If any intermediate computation uses float64, the result -propagates through the graph. - -**Common culprit:** `LongRope` in `_rotary_embedding.py`: -```python -# BAD — creates float64 array, poisons cos/sin caches -long_factor = np.array(config.rope_scaling["long_factor"]) - -# GOOD — explicit float32 -long_factor = np.array(config.rope_scaling["long_factor"], dtype=np.float32) -``` - -**Fix:** Always pass `dtype=np.float32` when creating numpy arrays from config -values. Search for `np.array(` without an explicit dtype to find other -instances. - -### 2. Normalization type mismatch (RMSNorm vs LayerNorm) - -**Symptom:** Logits have large max abs diff (> 0.5) from HuggingFace, but -weight loading succeeds and the greedy token may still match. - -**Diagnosis:** Check what norm class the HuggingFace model actually uses. -Many models have custom norm classes that differ from the standard: - -```python -import inspect -from transformers.models.olmo.modeling_olmo import OlmoLayerNorm -print(inspect.getsource(OlmoLayerNorm)) -``` - -**Known cases:** -- **OLMo-1B**: Uses `OlmoLayerNorm` which is `F.layer_norm(x, ..., weight=None, bias=None, eps=1e-5)` — a weight-free **LayerNorm**, NOT RMSNorm. Solution: use `_WeightFreeLayerNorm` (constant scale=1, bias=0) from `models/olmo.py`. -- **Whisper**: Uses `LayerNorm` with `eps=1e-5`, not the default `1e-6`. -- **Gemma**: Uses `RMSNorm` but adds 1 to the weight before applying. - -**Key difference between LayerNorm and RMSNorm:** -- **LayerNorm** subtracts the mean, then divides by std: `(x - mean) / sqrt(var + eps) * gamma + beta` -- **RMSNorm** does NOT subtract the mean: `x / sqrt(mean(x²) + eps) * gamma` - -Using the wrong type causes systematic error that grows through layers. - -### 3. Missing scaling multipliers - -**Symptom:** Token generation diverges after a few tokens, but the first few -tokens may match. - -**Diagnosis:** Check the HuggingFace model for multiplier/scaling attributes -that aren't in the standard Llama config: - -```python -config = AutoConfig.from_pretrained("model-id") -for k, v in config.to_dict().items(): - if "multiplier" in k or "scaling" in k: - print(f"{k}: {v}") -``` - -**Known cases:** -- **Granite**: Has four multipliers — `embedding_multiplier`, `attention_multiplier`, `logits_scaling`, `residual_multiplier`. All must be implemented for correct results. - -**Critical: residual scaling direction matters.** Granite uses: -```python -# CORRECT: multiplier on the output, not the residual -hidden_states = residual + attn_output * residual_multiplier - -# WRONG: multiplier on the residual -hidden_states = residual * residual_multiplier + attn_output -``` - -Always verify the direction by reading the HuggingFace source: -```python -from transformers.models.granite.modeling_granite import GraniteDecoderLayer -import inspect -print(inspect.getsource(GraniteDecoderLayer.forward)) -``` - -### 4. Config fields not extracted - -**Symptom:** Model builds without error but multipliers/special features are -not applied (they default to 1.0/None/False). - -**Fix:** Add extraction to `ArchitectureConfig.from_transformers()` in -`_configs.py`: - -```python -# In _configs.py ArchitectureConfig dataclass: -embedding_multiplier: float = 1.0 -attention_multiplier: float | None = None - -# In from_transformers() options dict: -embedding_multiplier=getattr(config, "embedding_multiplier", 1.0), -attention_multiplier=getattr(config, "attention_multiplier", None), -``` - -**Tip:** Use safe defaults (1.0 for multipliers, None for optional features) -so existing models are unaffected. - -### 5. Attention scale override - -**Symptom:** Attention scores are wrong, causing gradual drift in generation. - -**Diagnosis:** Some models override the default `1/sqrt(head_dim)` scale: - -```python -# Check HuggingFace attention class -from transformers.models.granite.modeling_granite import GraniteAttention -print(GraniteAttention.__init__) # Look for self.scaling = ... -``` - -**Fix:** Use the `scale` parameter on the `Attention` component: - -```python -# In your custom DecoderLayer: -self.self_attn = Attention(config, scale=config.attention_multiplier) -``` - -The `Attention.__init__` accepts an optional `scale: float | None` parameter. -When `None`, it defaults to `head_dim**-0.5`. - -### 6. Weight-free norms (no learnable parameters) - -**Symptom:** Weight keys like `model.norm.weight` appear in the model but not -in the HuggingFace state dict, and `preprocess_weights` fills them with ones. -The norm still uses the wrong algorithm (e.g. RMSNorm instead of LayerNorm). - -**Fix:** Use `_WeightFreeLayerNorm` from `models/olmo.py` which creates -`nn.Parameter` with constant data (ones for scale, zeros for bias) that -the ONNX `LayerNormalization` op requires, but the HuggingFace model has no -corresponding weights for: - -```python -class _WeightFreeLayerNorm(nn.Module): - def __init__(self, hidden_size: int, eps: float = 1e-5): - super().__init__() - self.scale = nn.Parameter( - [hidden_size], data=ir.tensor(np.ones(hidden_size, dtype=np.float32)) - ) - self.bias = nn.Parameter( - [hidden_size], data=ir.tensor(np.zeros(hidden_size, dtype=np.float32)) - ) - self.eps = eps - - def forward(self, op, hidden_states): - return op.LayerNormalization( - hidden_states, self.scale, self.bias, epsilon=self.eps, axis=-1 - ) -``` - -### 7. Debugging workflow for logit mismatches - -When ONNX logits don't match HuggingFace, follow this sequence: - -1. **Check max abs diff on prefill** — tells you if the issue is in the - model computation (not just autoregressive error accumulation): - ```python - diff = np.abs(onnx_logits[0, -1] - hf_logits) - print(f"Max abs diff: {diff.max()}") # > 0.01 is suspicious, > 0.5 is a bug - ``` - -2. **Check the HuggingFace norm class** — most mismatches come from using the - wrong normalization type or epsilon. - -3. **Check for model-specific config fields** — inspect the HuggingFace config - dict for any field with "multiplier", "scaling", "factor", "epsilon", or - "bias" that isn't being extracted. - -4. **Check weight dtype** — ensure numpy arrays in rope/embeddings use float32. - -5. **Compare layer by layer** — if the diff is moderate (0.01-1.0), the issue - is likely an architectural difference in the norm or residual pattern. If - the diff is huge (>10), check if weights are loaded to the wrong parameters. - -### 8. Gated attention Q/gate split ordering - -**Symptom:** Large logit diff (~2.0) on Qwen3.5 or similar gated attention models. - -**Root cause:** HF splits Q and gate *within each head*: reshapes to -`[B, S, num_heads, 2*head_dim]` then chunks on last dim. A naive midpoint -split of the flat tensor gives wrong results. - -**Fix:** Reshape to per-head layout before splitting: -```python -# WRONG: split flat tensor at midpoint -q, gate = op.Split(qg_proj, num_outputs=2, axis=-1) - -# CORRECT: reshape to per-head, then split -qg = op.Reshape(qg_proj, [0, 0, num_heads, 2 * head_dim]) -q, gate = op.Split(qg, [head_dim, head_dim], axis=-1) -``` - -### 9. DeltaNet missing query scaling - -**Symptom:** Linear attention output is orders of magnitude too large. - -**Root cause:** After L2-normalizing Q and K, you still need a -`1/sqrt(key_head_dim)` scaling factor on the query, similar to standard -attention. - -### 10. Extracting `last_hidden_state` before vs after norm - -**Symptom:** Downstream model (e.g. code predictor, projection layer) -receives wrong hidden states. Prefill logits match HF exactly, but -generation diverges immediately. - -**Root cause:** HuggingFace's `outputs.last_hidden_state` is the -**post-norm** hidden state (after RMSNorm/LayerNorm). If you extract -the hidden state before the final norm, downstream consumers get -pre-norm values. In single-model LLMs this doesn't matter (the lm_head -is after the norm). In multi-model pipelines (TTS, VLM), the -hidden state is passed to another model, so norm ordering is critical. - -**Fix:** Always extract hidden state *after* the model's final norm: -```python -# WRONG: hidden_states before norm -hidden_states = decoder_output # pre-norm -logits = lm_head(norm(hidden_states)) # logits correct, but... -return logits, hidden_states # hidden_states is WRONG for downstream - -# CORRECT: apply norm first, then use for both logits and output -hidden_states = norm(decoder_output) # post-norm -logits = lm_head(hidden_states) -return logits, hidden_states # hidden_states matches HF -``` - -### 11. Identity node folding renames initializers - -**Symptom:** Weight loading fails — `preprocess_weights` maps to the -original parameter name (e.g. `code_predictor.stacked_codec_embedding`) -but the initializer in the ONNX graph has been renamed to something -like `v_code_predictor.Identity_174`. - -**Root cause:** The IR optimizer folds `Identity(initializer)` by -removing the Identity node and renaming the initializer to the -output name. If you then set a custom name on the output (e.g. -`codec_embeddings.name = "codec_embeddings"`), the initializer gets -renamed to `codec_embeddings` — breaking weight loading. - -**Fix:** Use `op.Identity()` to create a *real* Identity node between -the initializer and the graph output. Ensure the elimination pass -retains Identity nodes that feed graph outputs. This creates a separate -output value, so renaming the output doesn't affect the initializer: -```python -# In forward(): -codec_embeddings = op.Identity(self.stacked_codec_embedding) -return logits, present_key_values, codec_embeddings - -# In task (safe to rename — Identity separates the names): -codec_embeddings.name = "codec_embeddings" -graph.outputs.append(codec_embeddings) -``` - -### 12. `np.ascontiguousarray` promotes 0-d arrays to 1-d - -**Symptom:** ONNX `Gather` axis-reducing semantics break — the output -has an extra dimension (e.g. `(1, vocab)` instead of `(vocab,)`). - -**Root cause:** `np.ascontiguousarray(scalar_array)` promotes shape -`()` to `(1,)`. This changes `Gather(axis=0)` from axis-reducing -(scalar index) to axis-preserving (1-d index). - -**Fix:** Guard against 0-d arrays: -```python -if v.ndim > 0: - v = np.ascontiguousarray(v) -``` - -### 13. Multi-token prefill in code predictors - -**Symptom:** Code predictor generates garbage. Prefill logits are -slightly off compared to HF. - -**Root cause:** Some architectures (e.g. Qwen3-TTS code predictor) use -a **2-token prefill**: `concat(projected_hidden, embed(code_0))` as two -separate tokens through the transformer. Summing them into 1 token -changes attention patterns and all subsequent hidden states. - -**Diagnosis:** Compare the inputs_embeds shape at step 0. If HF passes -`(batch, 2, hidden)` but your model uses `(batch, 1, hidden)`, the -attention context window is wrong. - -**Fix:** Construct inputs_embeds externally to match HF's exact flow: -```python -# Step 0 (prefill): 2 tokens -inputs = np.concatenate([talker_hidden, embed(code_0)], axis=1) # (1, 2, H) -# Steps 1+: 1 token -inputs = cp_embed[step-1, code_i, :].reshape(1, 1, -1) # (1, 1, H) -``` - -### 14. Embedding table index off-by-one in multi-step generation - -**Symptom:** Codes are plausible but audio quality is wrong. Codec sum -doesn't match HF. - -**Root cause:** In multi-step code prediction, HF uses -`embed[step-1](code)` at generation step `step`, not `embed[step]`. -The off-by-one means every embedding lookup uses the wrong table. - -**Fix:** Carefully trace HF's generation loop to determine which -embedding table index corresponds to which generation step. Write a -comparison script that checks individual embedding lookups match. - -### 15. Codec sum uses output codes, not input codes - -**Symptom:** codec_sum diverges from HF even though individual -embeddings weights are identical. - -**Root cause:** The codec sum `Σ embed[i](code_{i+1})` uses the -*generated* (output) code at each step, not the input code. If the -model returns embeddings of the input codes, the sum is wrong. - -**Fix:** Compute codec_sum externally using the codes actually -generated at each step: -```python -codec_sum = talker_embed(code_0) -for i in range(num_groups - 1): - # codes[i+1] is the OUTPUT of code predictor step i - codec_sum += cp_embed[i, codes[i + 1], :] -``` - -### 16. Precision-sensitive ops need fp32 upcast - -**Symptom:** Type mismatch errors (`tensor(float) vs tensor(bfloat16)`) -when loading a model built with `--dtype bf16`, or numerical drift compared -to HuggingFace when running in fp16/bf16. - -**Root cause:** Operations like `exp`, `softplus`, `sigmoid` (in gated norms), -and RMSNorm variance are numerically sensitive and must run in float32 to -match HuggingFace, which explicitly upcasts with `.float()` / -`.to(torch.float32)`. - -**Two distinct problems:** - -1. **Naive `CastLike` everywhere** — keeps everything in the model dtype - (e.g. bf16), but `exp` overflows and the SSM state diverges. -2. **Naive `Cast(to=ir.DataType.FLOAT)` everywhere** — computes in fp32 but forgets to cast - back, producing type mismatches with downstream bf16 ops. - -**Correct pattern — upcast → compute → cast back:** -```python -# 1. Upcast to fp32 for the sensitive region -dt_f32 = op.Cast(dt, to=ir.DataType.FLOAT) -dt_f32 = op.Softplus(dt_f32) -a_neg = op.Neg(op.Exp(op.Cast(self.A_log, to=ir.DataType.FLOAT))) -da = op.Exp(op.Mul(dt_4d, a_4d)) # all fp32 here -... -# 2. Cast back to input dtype at the boundary -y = op.CastLike(y_f32, x) -new_state = op.CastLike(new_state_f32, ssm_state) -``` - -**How to identify which ops need fp32:** Check the HuggingFace source for -`.float()` or `.to(torch.float32)` calls. Each one marks an fp32 region -that the ONNX graph must replicate. - -**Known fp32-required regions:** - -| Region | HF evidence | ONNX pattern | -|--------|-------------|-------------| -| SSM recurrence (A, dt, exp, state) | `self.A_log.float()`, `hidden_states.float()`, `B.float()`, `C.float()` | `Cast(to=ir.DataType.FLOAT)` all inputs, `CastLike` output | -| GatedRMSNorm (SiLU + variance) | `hidden_states.to(torch.float32)`, `gate.to(torch.float32)` | Explicit fp32 for both, `CastLike` output | -| RMSNorm variance | `hidden_states.to(torch.float32)` | ONNX `RMSNormalization` handles via `stash_type=1` (default) | - -**When fp32 upcast is NOT needed:** -- Linear projections (`MatMul`) — runtime handles mixed precision -- SiLU on conv output — HF keeps in model dtype -- Standard attention — ONNX `Attention` op handles precision internally - -**Use `CastLike` for** parameters/constants that should match the *current* -compute dtype (which is fp32 inside an upcast region, or the model dtype -outside). Use `Cast(to=ir.DataType.FLOAT)` to explicitly enter an fp32 region. - -## Reference implementations - -| Model | File | Key differences from base | -|-------|------|--------------------------| -| Granite | `models/granite.py` | 4 scaling multipliers, custom attention scale | -| OLMo-1B | `models/olmo.py` | Weight-free LayerNorm (not RMSNorm), eps=1e-5 | -| OLMo-2 | `models/olmo.py` | Post-norm decoder layers, QK full norm | -| Gemma | `models/gemma.py` | RMSNorm weight+1, embedding scaling | -| Whisper | `components/_whisper.py` | Q pre-scaling, LayerNorm eps=1e-5, is_causal attr | -| Phi3.5 | `components/_rotary_embedding.py` | LongRope with float32 factors | -| Qwen3.5 | `models/qwen.py` | Hybrid DeltaNet + full attention, gated GQA, OffsetRMSNorm, interleaved MRoPE | -| Qwen3.5-MoE | `models/qwen.py` | Same hybrid attention + MoE FFN with shared expert (sigmoid gate) | -| Qwen3-TTS | `models/qwen3_tts.py` | 4-model TTS split, 2-token code predictor prefill, small_to_mtp projection, Identity-exposed weights | -| **BLIP** | `models/blip.py` | Subclass of ViTModel — only `preprocess_weights` (fused QKV split, renaming) | -| **YOLOS** | `models/yolos.py` | ViT + detection tokens + DETR-style MLP heads. New `object-detection` task | -| **Depth Anything** | `models/depth_anything.py` | ViT backbone + DPT decoder (reassemble + fusion + depth head). Uses `ConvTranspose2d` | -| **Segformer** | `models/segformer.py` | Hierarchical 4-stage encoder, efficient attention (strided Conv2d on K/V), Mix-FFN with depthwise conv | -| **SAM2** | `models/sam2.py` | Hiera backbone (per-stage dim transitions, fused QKV attention) + FPN neck with top-down fusion | -| **LayoutLMv3** | `models/layoutlmv3.py` | Subclass of BertModel — only `preprocess_weights` (spatial embedding filtering) | -| **TrOCR** | `models/trocr.py` | Subclass of BartForConditionalGeneration — only `preprocess_weights` (`output_projection` rename) | -| **ModernBERT** | `models/modernbert.py` | Pre-norm encoder with RoPE + GeGLU + bidirectional attention. Fused QKV/Wi splitting. Both encoder and decoder variants | -| **Gemma3n** | `models/gemma3n.py` | AltUp predict/correct, Laurel low-rank, per-layer input gating, hybrid local/global attention | -| **Mllama** | `models/mllama.py` | Interleaved cross-attention decoder, tanh-gated residual, manual QK-norm | - -## Reference Examples - -When adding a new model, use these files as canonical references: - -| Complexity | File | Why | -|---|---|---| -| **Minimal** — base class works, only weight mapping needed | `models/phi3.py` (38 lines) | Extends `CausalLMModel`, only overrides `preprocess_weights()` to split fused QKV and gate-up projections. Shows the simplest possible model addition. | -| **Minimal** — encoder subclass | `models/layoutlmv3.py` | Extends `BertModel`, only overrides `preprocess_weights()`. Same pattern for encoder-only models. | -| **Moderate** — custom components | `models/gemma.py` | Adds custom attention (soft-capping), custom MLP (GeGLU), and custom normalization. Good example of component subclassing. | -| **Complex** — multi-model architecture | `models/qwen3_tts.py` | 4-model TTS split with talker, code predictor, embedding, and speaker encoder sub-modules. Shows how to structure multi-model architectures. | - -## KV sharing across layers (num_kv_shared_layers) - -Some models (e.g. Gemma 4) reduce parameter count by having the last N -decoder layers **borrow** Key and Value states from an earlier "source" layer -of the same type instead of projecting their own K,V. This is controlled by -`num_kv_shared_layers` in the HuggingFace config. - -### What it means - -``` -first_kv_shared_idx = num_hidden_layers - num_kv_shared_layers - -Layers [0 .. first_kv_shared_idx - 1]: normal — own k_proj, v_proj, k_norm -Layers [first_kv_shared_idx .. end]: shared — NO k_proj/v_proj weights -``` - -Each shared layer reuses K,V from the **last non-shared layer of the same -attention type** (e.g. sliding vs. full attention). Only Q is computed fresh. - -### Impact on the checkpoint - -Shared layers have **no `k_proj`, `v_proj`, `k_norm`** keys in the -HuggingFace checkpoint. `preprocess_weights` must not assert these keys -exist for shared-layer indices — they simply won't be present. - -```python -def preprocess_weights(self, state_dict): - # shared layers have no k/v proj — remove them silently if accidentally present - first_shared = self.config.num_hidden_layers - self.config.num_kv_shared_layers - for i in range(first_shared, self.config.num_hidden_layers): - for suffix in ("k_proj.weight", "v_proj.weight", "k_norm.weight"): - state_dict.pop(f"model.layers.{i}.self_attn.{suffix}", None) - return super().preprocess_weights(state_dict) -``` - -### Attention module: is_kv_shared_layer flag - -The attention class detects at `__init__` time whether it is a shared layer: - -```python -class Gemma4Attention(nn.Module): - def __init__(self, config, layer_idx, layer_types, first_kv_shared_idx, ...): - self.is_kv_shared_layer = layer_idx >= first_kv_shared_idx > 0 - prev_layers = layer_types[:first_kv_shared_idx] - - if self.is_kv_shared_layer: - # Index of the source layer whose K,V this layer borrows - self.kv_shared_layer_index = ( - len(prev_layers) - 1 - prev_layers[::-1].index(layer_types[layer_idx]) - ) - self.store_full_length_kv = False - else: - self.kv_shared_layer_index = None - # True for the last non-shared layer of each type that has downstream - # KV-shared layers depending on it — it stores K,V for reuse. - self.store_full_length_kv = first_kv_shared_idx > 0 and ( - layer_idx - == len(prev_layers) - 1 - prev_layers[::-1].index(layer_types[layer_idx]) - ) - - # All layers have Q projection - self.q_proj = Linear(config.hidden_size, num_heads * head_dim) - self.q_norm = RMSNorm(head_dim) - self.o_proj = Linear(num_heads * head_dim, config.hidden_size) - - # Only non-shared layers have K/V projections - if not self.is_kv_shared_layer: - self.k_proj = Linear(config.hidden_size, num_kv_heads * head_dim) - self.v_proj = Linear(config.hidden_size, num_kv_heads * head_dim) - self.k_norm = RMSNorm(head_dim) -``` - -### forward(): shared layers consume shared_kv_states dict - -Pass a mutable `shared_kv_states` dict through the forward call. Source -layers populate it; shared layers read from it: - -```python -def forward(self, op, hidden_states, ..., shared_kv_states, past_key_value): - # Q projection (all layers) - query_states = self.q_proj(op, hidden_states) - ... - - if self.is_kv_shared_layer: - # Borrow K,V from source layer (already in shared_kv_states) - src_key, src_value = shared_kv_states[self.kv_shared_layer_index] - # Reshape from present_kv 4D [B, kv_heads, total_seq, head_dim] - # to Attention input 3D [B, total_seq, kv_heads * head_dim] - src_key = op.Transpose(src_key, perm=[0, 2, 1, 3]) - key_states = op.Reshape(src_key, ...) - value_states = ... - else: - # Normal K/V projection + norm - key_states = self.k_proj(op, hidden_states) - value_states = self.v_proj(op, hidden_states) - ... - - hidden_out, present_kv = _apply_attention(op, query_states, key_states, ...) - - if self.store_full_length_kv: - # Store present_kv [B, kv_heads, total_seq, head_dim] for downstream shared layers - shared_kv_states[self.layer_idx] = (present_kv_key, present_kv_value) - - return hidden_out, present_kv -``` - -### Text model: KV cache has only num_kv_layers entries - -KV-shared layers do **not** append to `present_key_values`. The output list -has `num_hidden_layers - num_kv_shared_layers` entries, not `num_hidden_layers`: - -```python -# In Gemma4TextModel.forward(): -shared_kv_states: dict = {} -present_key_values = [] - -# past_key_values has only num_kv_layers entries (no entry for KV-shared layers). -# Expand it to a full per-layer list so we can zip cleanly over all layers. -if past_key_values is not None: - kv_iter = iter(past_key_values) - past_kvs: list = [ - None if layer.self_attn.is_kv_shared_layer else next(kv_iter) - for layer in self.layers - ] -else: - past_kvs = [None] * len(self.layers) - -for i, (layer, layer_type, past_kv) in enumerate( - zip(self.layers, self.layer_types, past_kvs) -): - hidden_states, present_kv = layer( - op, - hidden_states=hidden_states, - attention_bias=attention_bias_dict[layer_type], - position_embeddings=position_embeddings_dict[layer_type], - shared_kv_states=shared_kv_states, - past_key_value=past_kv, - ) - # KV-shared layers borrow K,V — exclude from present_key_values so the - # output has exactly num_kv_layers (not num_hidden_layers) entries. - if not layer.self_attn.is_kv_shared_layer: - present_key_values.append(present_kv) -``` - -The task's KV cache inputs/outputs must use the correct count: -`num_kv_layers = config.num_hidden_layers - config.num_kv_shared_layers`. diff --git a/.github/skills/debugging-vl-pipeline/SKILL.md b/.github/skills/debugging-vl-pipeline/SKILL.md deleted file mode 100644 index be7cf4e0..00000000 --- a/.github/skills/debugging-vl-pipeline/SKILL.md +++ /dev/null @@ -1,533 +0,0 @@ ---- -name: debugging-vl-pipeline -description: > - How to debug vision-language (VL) and multimodal (vision + audio) model - output issues in mobius. Covers the systematic pipeline isolation - methodology, common failure modes, stage-by-stage comparison with - HuggingFace, numerical tolerance expectations, and CUDA EP-specific - issues. Use this skill when ORT GenAI multimodal output is wrong, - garbled, or doesn't match HuggingFace. ---- - -# Skill: Debugging VL Pipeline Issues - -## When to use - -Use this skill when: - -- ORT GenAI produces wrong or irrelevant output for image inputs -- ONNX model logits diverge significantly from HuggingFace -- The model generates text-only descriptions ignoring the image -- Image features appear correct but decoder output is wrong -- Audio transcription is garbled or wrong despite correct encoder output -- CUDA EP crashes or produces different results than CPU - -## Debugging methodology: isolate each stage - -VL models have 3 stages. Debug by isolating and validating each stage -independently, comparing against HuggingFace at every boundary. - -``` -pixel_values ──► [1. Vision] ──► image_features - │ -input_ids ──► [2. Embedding] ◄──────┘ - │ - ▼ - inputs_embeds + position_ids + attention_mask - │ - ▼ - [3. Decoder] ──► logits -``` - -### Stage 1: Vision model - -**What to check:** -- Output shape: `(num_patches, hidden_size)` -- Expected patches = `t * (h / merge) * (w / merge)` from `grid_thw` -- Compare features against HF vision encoder output - -```python -# HF reference -with torch.no_grad(): - hf_vision_out = hf_model.model.visual( - pixel_values, grid_thw=grid_thw - ) -# ONNX -session = OnnxModelSession(pkg["vision"]) -onnx_out = session.run({"pixel_values": pv, "grid_thw": grid_thw}) - -# Compare -cos_sim = np.dot(hf_flat, onnx_flat) / (norm_hf * norm_onnx) -print(f"Vision cos_sim: {cos_sim:.6f}") # Should be > 0.99 -``` - -**Common issues:** -- Wrong pixel value normalization (mean/std mismatch) -- `grid_thw` shape or values don't match HF processor output -- Missing `temporal_patch_size` in patch embedding - -### Stage 2: Embedding model - -**What to check:** -- Image features injected at correct token positions -- Non-image positions have correct text embeddings -- Output shape: `(1, seq_len, hidden_size)` - -```python -# Verify image token positions -image_mask = (input_ids[0] == image_token_id) # 151655 -num_image_positions = image_mask.sum() -assert num_image_positions == image_features.shape[0] - -# Compare embeddings at text positions (should match HF exactly) -text_mask = ~image_mask -cos_sim_text = cosine_similarity( - onnx_embeds[0, text_mask], hf_embeds[0, text_mask] -) -print(f"Text embedding cos_sim: {cos_sim_text:.6f}") # Should be 1.0 -``` - -**Common issues:** -- Image token count mismatch between processor and vision model -- Missing zero-padding row in embedding model (for text-only inputs) -- Wrong `image_token_id` used for Gather/Where mask - -### Stage 3: Decoder - -**What to check:** -- Logits shape matches HF: `(1, seq_len, vocab_size)` -- First token prediction matches HF (argmax of last position) -- Cosine similarity of logit vectors - -```python -# With HF-computed position_ids (ground truth): -onnx_logits = decoder_session.run(feeds)["logits"] -hf_logits = hf_model(**hf_inputs).logits.numpy() - -max_diff = np.abs(onnx_logits - hf_logits).max() -cos_sim = cosine_similarity(onnx_logits[0, -1], hf_logits[0, -1]) -print(f"max_diff={max_diff:.2f}, cos_sim={cos_sim:.4f}") -# Typical: max_diff=5-10, cos_sim>0.98 -``` - -**Common issues:** -- Wrong position_ids (see "3D M-RoPE" section below) -- Missing KV cache initialization -- Wrong attention_mask length - -## Critical: 3D M-RoPE position IDs - -Qwen2-VL / Qwen2.5-VL / Qwen3-VL use **3D Multimodal RoPE** where -`position_ids` has shape `(3, batch, seq_len)`: - -``` -position_ids[0] = temporal positions -position_ids[1] = height positions -position_ids[2] = width positions -``` - -**Text tokens:** all 3 dimensions have the same sequential value. - -**Image tokens:** temporal is constant, height/width vary over the -image grid `(h/merge, w/merge)`: -``` -temporal: [offset, offset, offset, ..., offset] -height: [offset, offset+1, offset+1, ..., offset+h/merge-1] -width: [offset, offset+1, offset, offset+1, ..., offset+w/merge-1] -``` - -**Text after image:** all 3 dimensions resume from -`max(temporal, height, width) + 1`. - -### ORT GenAI config requirements - -For ORT GenAI to compute 3D M-RoPE automatically, the following -`genai_config.json` fields are **required**: - -| Field | Level | Purpose | -|-------|-------|---------| -| `model.image_token_id` | model | Token ID for `<\|image_pad\|>` (e.g. 151655) | -| `model.vision_start_token_id` | model | Token ID for `<\|vision_start\|>` (e.g. 151652) | -| `model.vision.spatial_merge_size` | vision | Grid merge factor (typically 2) | - -**Without these fields**, ORT GenAI falls back to standard 1D positions, -which produces completely wrong output for image inputs (the model may -describe a "snowy landscape" instead of the actual image content). - -## Common failure modes and fixes - -### 1. Image not recognized (wrong output for image inputs) - -**Symptoms:** Model produces generic or hallucinated descriptions that -don't match the input image. Text-only generation works correctly. - -**Root causes (in order of likelihood):** - -1. **Missing genai_config fields** — `image_token_id`, - `vision_start_token_id`, or `spatial_merge_size` not set. - Without these, position_ids are 1D instead of 3D M-RoPE. - -2. **Image resize mismatch** — ORT processor resizes image to different - dimensions than HF processor, producing different number of vision - tokens. ORT's `width`/`height` in `processor_config.json` are used - as direct resize targets, unlike HF's smart_resize which computes - target from original image dimensions. - -3. **Processor config format** — ORT GenAI expects ort-extensions format - `processor_config.json`, not HuggingFace format. The file must include - `DecodeImage`, `ConvertRGB`, `Resize`, `Rescale`, `Normalize`, and - `PatchImage` transforms with correct attributes. - -**Fix for resize mismatch:** -```python -def _update_resize_for_image(processor_config_path, image_path): - """Recompute resize dimensions from actual image like HF does.""" - from PIL import Image - img = Image.open(image_path) - w, h = img.size - factor = 14 * 2 # patch_size * merge_size - new_w = round(w / factor) * factor - new_h = round(h / factor) * factor - # Update width/height in processor_config.json -``` - -### 2. Numerical divergence in greedy decoding - -**Symptoms:** First 1-3 tokens match HF, then output diverges. - -**Expected behavior:** This is inherent to ONNX vs PyTorch numerical -differences. ONNX models use different operator implementations that -accumulate small floating-point errors. - -**Typical metrics for Qwen2.5-VL 3B:** -- max_diff in logits: 5-10 -- mean_diff in logits: 0.5-1.5 -- cosine similarity: 0.98-0.99 -- First token: matches HF -- Greedy decoding: diverges at token 3-5 - -**This is NOT a bug** if the metrics above are within range. Both models -produce semantically similar descriptions. - -### 3. Vision model output shape mismatch - -**Symptoms:** Vision model produces wrong number of patches. - -**Debug:** Check `grid_thw` values: -```python -# For Qwen2.5-VL with merge_size=2: -t, h, w = grid_thw[0] -expected_patches = t * (h // 2) * (w // 2) -actual_patches = vision_output.shape[0] -assert expected_patches == actual_patches -``` - -### 4. Embedding model text-only failure - -**Symptoms:** Error when running without images (num_image_tokens=0). - -**Fix:** Ensure embedding model pads `image_features` with a zero row -before Gather, then uses a Where mask to select only real features: -```python -# Pad with zero row so Gather with index 0 doesn't fail -padded = op.Concat( - op.ConstantOfShape(...), # (1, hidden_size) zeros - image_features, - axis=0, -) -``` - -## Extracting intermediate ONNX values - -When a stage (e.g., the vision encoder) diverges, drill down by extracting -intermediate values from the ONNX graph. This lets you compare block-by-block -or even op-by-op against HuggingFace. - -### Method 1: Add intermediate outputs to the ONNX graph - -The most reliable approach — expose any internal node's output as a graph -output so ORT returns it alongside normal outputs. - -```python -import onnx - -model = onnx.load("vision.onnx") -graph = model.graph - -# Find the node whose output you want to inspect -for node in graph.node: - if node.op_type == "RMSNormalization" and "block_0" in node.output[0]: - # Name the output (if unnamed, give it a name) - target_output = node.output[0] - break - -# Add as a graph output -graph.output.append( - onnx.helper.make_tensor_value_info(target_output, onnx.TensorProto.FLOAT, None) -) -onnx.save(model, "vision_debug.onnx") - -# Now ORT will return this value alongside image_features -session = ort.InferenceSession("vision_debug.onnx") -results = session.run(None, feeds) -# results[-1] is the intermediate value -``` - -### Method 2: Use `ir.Model` graph manipulation (preferred for mobius) - -When working with `ir.Model` objects from the build pipeline, manipulate -the graph directly without saving/loading: - -```python -from mobius._testing.ort_inference import OnnxModelSession - -pkg = build(model_id, dtype="f32", load_weights=True) -vision_model = pkg["vision"] -graph = vision_model.graph - -# Find target nodes by op type or name pattern -target_nodes = [n for n in graph if n.op_type == "RMSNormalization"] - -# RMSNorm input nodes are natural block boundaries: -# rms_nodes[0] = block 0 norm1 (input = patch_embed output) -# rms_nodes[2] = block 1 norm1 (input = block 0 output) -# rms_nodes[2*i] = block i norm1 (input = block i-1 output) -block_0_output = target_nodes[2].inputs[0] # block 0's output -block_0_output.name = "block_0_output" -graph.outputs.append(block_0_output) - -session = OnnxModelSession(vision_model) -out = session.run({"pixel_values": pv, "grid_thw": grid_thw}) -block_0_out = out["block_0_output"] -session.close() -``` - -### Method 3: Hook HuggingFace model for reference values - -Use PyTorch hooks to extract intermediate values from HuggingFace at -the same points: - -```python -intermediates = {} - -def hook_fn(name): - def fn(module, input, output): - if isinstance(output, tuple): - intermediates[name] = output[0].detach().cpu().numpy() - else: - intermediates[name] = output.detach().cpu().numpy() - return fn - -# Register hooks on specific blocks -for i, block in enumerate(hf_model.model.visual.blocks): - block.register_forward_hook(hook_fn(f"block_{i}")) - -# Run forward pass — hooks capture all intermediate values -with torch.no_grad(): - hf_out = hf_model.model.visual(pixel_values, grid_thw=grid_thw) - -# Now compare block by block -for i in range(num_blocks): - hf_block_out = intermediates[f"block_{i}"] - cos = cosine_similarity(onnx_block_out, hf_block_out) - print(f"Block {i}: cos={cos:.6f}") -``` - -### Block-by-block comparison strategy - -When overall output diverges, narrow down by comparing each transformer -block's output sequentially: - -```python -for i in range(num_blocks): - onnx_out_i = extract_onnx_block_output(vision_model, i, feeds) - hf_out_i = intermediates[f"block_{i}"] - - cos = cosine_similarity(onnx_out_i.flatten(), hf_out_i.flatten()) - max_diff = np.max(np.abs(onnx_out_i - hf_out_i)) - print(f"Block {i:2d}: cos={cos:.6f} max_diff={max_diff:.4f}") -``` - -Typical pattern for a bug in block N: -``` -Block 0: cos=1.000000 max_diff=0.0001 ← perfect -Block 1: cos=1.000000 max_diff=0.0001 ← perfect -... -Block N: cos=0.961000 max_diff=4.8500 ← divergence starts! -Block N+1: cos=0.892000 max_diff=25.00 ← error compounds -``` - -Once you identify the divergent block, drill deeper into that block's -sub-operations (attention, MLP, normalization) to find the root cause. - -### Comparing specific weight values - -To verify weights loaded correctly, compare ONNX initializers against -HuggingFace state dict: - -```python -from safetensors import safe_open - -# Load HF weights -with safe_open(safetensors_path, framework="numpy") as f: - hf_weight = f.get_tensor("visual.blocks.0.attn.qkv.weight") - -# Get ONNX weight from ir.Model -onnx_weight = vision_model.graph.initializers["blocks.0.attn.qkv.weight"] -onnx_np = onnx_weight.const_value.numpy() - -max_diff = np.max(np.abs(hf_weight.astype(np.float32) - onnx_np)) -print(f"Weight diff: {max_diff}") # Should be 0.0 -``` - -## Integration test patterns - -### Full VL forward test -```python -# Build model → process image with HF processor → run ONNX → compare logits -assert_logits_close(onnx_logits, hf_logits, rtol=2e-2, atol=2e-1) -``` - -### 3-model pipeline test -```python -# Vision → Embedding → Decoder, each stage uses OnnxModelSession -# Compare final decoder logits against HF single-model forward -``` - -### ORT GenAI end-to-end test -```python -# Build → save flat → write genai_config → load with ort_genai → generate -# Verify output length > input (basic sanity) -``` - -### 5. Vision encoder internal divergence (cos < 0.5) - -**Symptoms:** Vision features have very low cosine similarity (< 0.5) -against HuggingFace, even though patch embedding and weights are correct. - -**Debug with block-by-block comparison** (see "Extracting intermediate -ONNX values" above). Common root causes: - -1. **Wrong rotary embedding dimension** — Qwen2.5-VL vision uses 2D - position encoding (height + width). The rotary dim must be - `head_dim // 2`, not `head_dim`. Each half (head_dim // 4 frequencies) - covers one spatial dimension. With full `head_dim`, you get 2× too - many frequencies with wrong values. **Result: cos ≈ 0.25.** - -2. **Missing `fullatt_block_indexes`** — Qwen2.5-VL alternates between - windowed attention (local windows of `window_size` patches) and full - attention (all patches attend to all). Blocks at indexes `[7, 15, 23, 31]` - use full attention. If `fullatt_block_indexes` is not extracted from - HF config, all blocks use windowed attention. **Result: blocks 0-6 - are perfect (they're windowed anyway), but block 7+ diverges.** - -3. **Wrong attention bias construction** — Full-attention blocks should - have an all-zeros bias (everything attends to everything). Windowed - blocks have a block-diagonal bias. Check the bias by inspecting - sparsity: `(bias == -inf).float().mean()` should be ~0% for full - attention, ~98% for windowed. - -**Config extraction checklist for vision encoders:** -```python -# These fields MUST be extracted from HF vision_config: -fullatt_block_indexes = getattr(vc, "fullatt_block_indexes", None) -window_size = getattr(vc, "window_size", None) -spatial_merge_size = getattr(vc, "spatial_merge_size", 2) -temporal_patch_size = getattr(vc, "temporal_patch_size", 2) -``` - -### 6. ClippableLinear divergence (Gemma4) - -**Symptoms:** Vision or audio encoder output has large max diff (> 1.0) -against HuggingFace, even though weights are loaded correctly. - -**Root cause:** Gemma4 uses `Gemma4ClippableLinear` with learned finite -input/output activation clamping for ALL linear layers in its vision and -audio encoders. Using plain `Linear` misses the clamping. - -**Detection:** Check if the HuggingFace model uses `ClippableLinear`: -```bash -grep -n "ClippableLinear" transformers/models//modeling_.py -``` - -**Fix:** Use `ClippableLinear` (from `mobius.components`) for all -affected linear layers. For vision attention: q/k/v/o_proj. For MLP: -pass `linear_class=ClippableLinear` to the MLP component. - -**Impact:** -- Audio: max diff 52.68 → 0.0003 after fix -- Vision: max diff 3.92 → 0.00007 after fix - -### 7. Missing audio boundary markers - -**Symptoms:** Audio transcription is garbled or completely wrong, even -though the audio encoder output matches HuggingFace. - -**Root cause:** HuggingFace wraps audio placeholder tokens with boundary -markers in `input_ids`: -``` -<|audio> (256000) + N × <|audio|> (258881) + (258883) -``` -If boundary markers are missing, the model cannot distinguish audio -regions from text, producing wrong output. - -**Fix:** Add boundary markers to `build_input_ids()` when constructing -audio inputs, matching HuggingFace's token wrapping. - -### 8. CUDA EP: ORT Gather int32 overflow - -**Symptoms:** CUDA EP crashes or produces incorrect results for models -with large embedding tables. CPU EP works correctly. - -**Root cause:** ORT CUDA `gather_impl.cu` uses `int32` for element -offset computation: `input_index = idx * cols + col_offset`. For -tensors with > 2^31 elements (e.g. Gemma4 per-layer embedding -[262144, 8960] = 2.35B elements), this overflows. - -**ORT bug:** microsoft/onnxruntime#28107 - -**Workaround:** Split large embeddings into smaller tables via -`nn.ModuleList` so each individual Gather stays under the int32 limit. -Use Slice instead of Gather for column-wise indexing on large tensors. - -### 9. CUDA EP: opset 24 kernel registration - -**Symptoms:** ORT CUDA EP fails to find kernels for standard ops -(Squeeze, Reshape, etc.) even though they work on CPU. - -**Root cause:** ORT ≤1.24.x CUDA/TRT EPs don't register kernels for -opset 24, even though the op semantics are unchanged from opset 23. - -**Fix:** Use the `ort_lower_opset_for_ep` feature flag (enabled by -default) which lowers the declared opset import from 24 to 23 for -non-CPU EPs. See `src/mobius/_flags.py` and -`src/mobius/_testing/ort_inference.py`. - -### 10. Wrong audio feature extractor - -**Symptoms:** Audio model produces completely wrong output. Audio -encoder features don't match HuggingFace at all. - -**Root cause:** Using `WhisperFeatureExtractor` instead of the -model-specific feature extractor (e.g. `Gemma4AudioFeatureExtractor`). -Different extractors produce different mel spectrograms. - -**Fix:** Always use the correct feature extractor for the model: -```python -from transformers import AutoFeatureExtractor -feature_extractor = AutoFeatureExtractor.from_pretrained(model_id) -``` - -## Reference files - -- **Integration tests:** `tests/integration_test.py` - (`TestVLFullForward`, `TestQwen25VL3Model`) -- **ORT GenAI tests:** `tests/ort_genai_test.py` - (`TestOrtGenaiQwen25VL.test_multimodal_image_generation`) -- **Example scripts:** `examples/qwen25_vl_ort_genai.py`, - `examples/qwen3_vl_ort_genai.py`, `examples/gemma4_multimodal.py` -- **genai_config reference:** `.github/skills/ort-genai-config/SKILL.md` -- **ORT GenAI position_ids code (external):** - `onnxruntime-genai/src/models/position_inputs.cpp:617-814` -- **Feature flags:** `src/mobius/_flags.py` - (`ort_lower_opset_for_ep`, `ort_cuda_grouped_rmsnorm_workaround`) diff --git a/.github/skills/diffusion-models/SKILL.md b/.github/skills/diffusion-models/SKILL.md deleted file mode 100644 index 90cf20b3..00000000 --- a/.github/skills/diffusion-models/SKILL.md +++ /dev/null @@ -1,388 +0,0 @@ ---- -name: diffusion-models -description: > - How to add and work with diffusion / image-generation models in - mobius. Covers UNet, VAE, DiT, Flux, SD3, ControlNet, - adapters, and QwenImage architectures. Includes pipeline detection, - diffusers configs, task classes, building blocks, and weight loading - conventions. Use this skill when adding or modifying a diffusion model. ---- - -# Skill: Diffusion Models - -## When to use - -Use this skill when: -- Adding a new diffusion model (denoiser, VAE, ControlNet, adapter) -- Working with `build_diffusers_pipeline()` or the `_DIFFUSERS_CLASS_MAP` -- Creating diffusers config classes or task types -- Debugging diffusers weight loading or pipeline detection - -## Architecture overview - -Diffusion models generate images/video by iteratively denoising latent -representations. The key components are: - -| Component | Role | Examples | -|-----------|------|----------| -| **Denoiser** | Predicts noise to remove at each step | UNet2D, DiT, Flux, SD3, QwenImage | -| **VAE** | Encodes images to latent / decodes latent to images | AutoencoderKL, QwenImage 3D VAE, Video VAE | -| **ControlNet** | Provides spatial conditioning (edges, depth, etc.) | ControlNetModel | -| **Adapter** | Adds image/conditioning features to denoiser | T2I-Adapter, IP-Adapter | - -## Pipeline detection and building - -### How it works - -When the CLI or `build()` encounters a model that isn't a transformers model, -it checks for `model_index.json` (the diffusers pipeline descriptor): - -```python -# In _diffusers_builder.py -# 1. Try transformers AutoConfig → if fails: -# 2. Try loading model_index.json → if found: -# 3. Parse components and build each via _DIFFUSERS_CLASS_MAP -``` - -### Registering a new diffusers model - -Add an entry to `_init_diffusers_class_map()` in `_diffusers_builder.py`: - -```python -def _init_diffusers_class_map(): - _DIFFUSERS_CLASS_MAP["MyTransformer2DModel"] = ( - MyTransformer2DModel, # Module class - MyConfig, # Config dataclass - "denoising", # Task name - ) -``` - -The key must match the class name in `model_index.json`: -```json -{ - "transformer": ["diffusers", "MyTransformer2DModel"], - "vae": ["diffusers", "AutoencoderKL"] -} -``` - -### Weight loading for diffusers - -Diffusers uses different weight file naming than transformers: - -| Convention | Filename | -|------------|----------| -| Single file | `diffusion_pytorch_model.safetensors` | -| Sharded index | `diffusion_pytorch_model.safetensors.index.json` | -| Fallback | `model.safetensors` (some models) | - -Weights live in component subdirectories: `transformer/`, `vae/`, etc. -The `_download_diffusers_component_weights()` function tries both naming -conventions. - -## Config classes - -Diffusers configs are separate from `ArchitectureConfig`. Each is a -`@dataclasses.dataclass` with a `from_diffusers(config: dict)` classmethod -that parses the JSON config. - -```python -@dataclasses.dataclass -class MyDiffuserConfig: - in_channels: int = 4 - out_channels: int = 4 - block_out_channels: tuple[int, ...] = (320, 640, 1280, 1280) - num_layers: int = 2 - attention_head_dim: int = 8 - - @classmethod - def from_diffusers(cls, config: dict) -> "MyDiffuserConfig": - return cls( - in_channels=config.get("in_channels", cls.in_channels), - out_channels=config.get("out_channels", cls.out_channels), - block_out_channels=tuple(config.get("block_out_channels", cls.block_out_channels)), - num_layers=config.get("num_layers", cls.num_layers), - attention_head_dim=config.get("attention_head_dim", cls.attention_head_dim), - ) -``` - -Existing config classes: -- `VAEConfig` — Standard SD VAE (2D) -- `UNet2DConfig` — UNet-based denoisers -- `DiTConfig` — Diffusion Transformer -- `SD3Config` — Stable Diffusion 3 / MMDiT -- `FluxConfig` — Flux transformer -- `ControlNetConfig` — ControlNet conditioning -- `T2IAdapterConfig` / `IPAdapterConfig` — Adapters -- `QwenImageConfig` — QwenImage transformer -- `QwenImageVAEConfig` — QwenImage 3D causal VAE -- `VideoVAEConfig` — 3D video VAE - -All defined in `src/mobius/_diffusers_configs.py`. - -## Task classes - -Each diffusion model type has a task class that defines I/O signatures: - -### DenoisingTask - -```python -# Inputs: -# sample: [B, in_channels, H, W] — noisy latent -# timestep: [B] — diffusion timestep -# encoder_hidden_states: [B, seq, dim] — text conditioning -# Output: -# noise_pred: [B, in_channels, H, W] — predicted noise -``` - -### VAETask - -Returns `ModelPackage` with two sub-models: -- `encoder`: `sample [B, 3, H, W]` → `latent_dist [B, 2*latent_ch, H/f, W/f]` -- `decoder`: `latent [B, latent_ch, H/f, W/f]` → `sample [B, 3, H, W]` - -### QwenImageVAETask - -3D causal VAE for video/images: -- `encoder`: `[B, 3, T, H, W]` → `[B, 2*z_dim, T', H', W']` -- `decoder`: `[B, z_dim, T', H', W']` → `[B, 3, T, H, W]` - -### ControlNetTask - -```python -# Inputs: sample, timestep, encoder_hidden_states, controlnet_cond -# Outputs: down_outputs (list of residuals), mid_output -``` - -### AdapterTask - -```python -# T2I: condition [B, in_channels, H, W] → feature_list -# IP: image_embeds [B, image_dim] → adapter_output [B, num_tokens, cross_dim] -``` - -## Building blocks - -### Shared components (from `components/`) - -Diffusion models import shared primitives from the component library: - -| Component | Import | Purpose | -|-----------|--------|---------| -| `Conv2d` | `components/_conv.py` | 2D convolution with bias | -| `GroupNorm` | `components/_common.py` | Group normalization | -| `LayerNormNoAffine` | `components/_common.py` | Norm without learnable params (AdaLN) | -| `Linear` | `components/_common.py` | Linear projection | -| `SiLU` | `components/_activations.py` | SiLU activation module | - -Model files import these with underscore-prefixed aliases: - -```python -from mobius.components import Conv2d as _Conv2d, GroupNorm as _GroupNorm, SiLU as _SiLU -``` - -### Common model-specific blocks - -| Block | Purpose | Used by | -|-------|---------|---------| -| `_TimestepEmbedding` | Sinusoidal → MLP time embedding | UNet, DiT, Flux, SD3, QwenImage | -| `nn.ModuleList` | Layer lists, skip connections | All | -| `nn.Sequential` | Callable container that chains forward calls | Diffusion `to_out`, `img_mod`, `net` | - -> **`nn.Sequential` vs `nn.ModuleList`**: `nn.Sequential` chains children's -> `forward()` calls automatically — prefer it for sequential containers like -> `to_out` and modulation layers. Use `nn.ModuleList` for layer lists that -> need custom iteration (e.g., `down_blocks`, `transformer_blocks`). - -### UNet-specific - -| Block | Purpose | -|-------|---------| -| `_ResNetBlock2DWithTime` | ResNet block with time embedding injection | -| `_CrossAttentionBlock` | Self-attention + cross-attention + FFN | -| `_BasicAttention` | Scaled dot-product attention (spatial) | -| `_DownBlock2D` | ResNet + optional attn + spatial downsample | -| `_UpBlock2D` | ResNet + optional attn + spatial upsample + skip concat | -| `_UNetMidBlock2DCrossAttn` | Mid block: ResNet → cross-attn → ResNet | - -### DiT / Transformer denoiser patterns - -| Block | Purpose | -|-------|---------| -| `_PatchEmbed` | Conv2d with stride=patch_size (patchify input) | -| `_AdaLayerNormZero` | Adaptive LayerNorm with scale/shift/gate from timestep | -| `_DiTBlock` | AdaLN-Zero → self-attn + cross-attn + FFN | -| `_JointAttentionBlock` | Concatenate text+image → joint attention → split | -| `_FluxSingleBlock` | Unified self-attention (single-stream) | - -### VAE-specific - -| Block | Purpose | -|-------|---------| -| `_ResNetBlock2D` | GroupNorm → SiLU → Conv → skip | -| `_AttentionBlock` | Self-attention in latent space | -| `_Downsample2D` / `_Upsample2D` | Spatial resampling | -| `_MidBlock2D` | ResNet + optional attention | - -### 3D / Video-specific - -| Block | Purpose | -|-------|---------| -| `_CausalConv3d` | 3D conv with temporal-causal padding | -| `_RMSNorm3d` | Channel-wise RMS normalization for 3D | -| `_ResidualBlock` (3D) | 3D ResNet with RMSNorm | -| `_Resample` | 2D spatial + optional 3D temporal resampling | - -## Denoiser architecture comparison - -| Model | Stream | Attention | Modulation | Normalization | -|-------|--------|-----------|-----------|---------------| -| **UNet2D** | Single | Cross-attn only | Time embed inject | GroupNorm | -| **DiT** | Single | Self + cross | AdaLN-Zero (6 params) | LayerNorm | -| **SD3 (MMDiT)** | Double | Joint (concat text+img) | AdaLN-Zero | LayerNorm | -| **Flux** | Double→Single | Joint then unified | AdaLN-Zero | LayerNorm | -| **QwenImage** | Double | Joint (separate proj) | AdaLN (2×6 params) | RMSNorm | - -### Double-stream pattern - -SD3, Flux, and QwenImage use a double-stream architecture: -1. **Image stream**: modulation → attention → FFN → residual -2. **Text stream**: separate modulation → attention → FFN → residual -3. **Joint attention**: concatenate Q/K/V from both streams, attend together - -```python -# In _JointAttentionBlock.forward(): -img_q, img_k, img_v = self.to_q(op, img), self.to_k(op, img), self.to_v(op, img) -txt_q, txt_k, txt_v = self.add_q(op, txt), self.add_k(op, txt), self.add_v(op, txt) -q = op.Concat(img_q, txt_q, axis=1) -k = op.Concat(img_k, txt_k, axis=1) -v = op.Concat(img_v, txt_v, axis=1) -attn_out = op.Attention(q, k, v, ...) -img_out, txt_out = op.Split(attn_out, axis=1, ...) -``` - -### AdaLN-Zero modulation - -Most transformer denoisers use AdaLN-Zero: project timestep embedding to -6 parameters (shift, scale, gate for attention and FFN): - -```python -# Modulation: timestep → SiLU → Linear → split into 6 chunks -mod = self.mod(op, temb) # [B, 6 * dim] -shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = op.Split( - mod, num_outputs=6, axis=-1 -) - -# Apply: norm → scale/shift → attention → gate → residual -normed = self.norm(op, x) -modulated = op.Add(op.Mul(normed, op.Add(scale_msa, ONE)), shift_msa) -attn_out = self.attn(op, modulated) -x = op.Add(x, op.Mul(attn_out, gate_msa)) -``` - -## Adding a new diffusion model — step by step - -### 1. Create config class - -In `_diffusers_configs.py`: - -```python -@dataclasses.dataclass -class MyDenoiserConfig: - # Parse from HF diffusers config.json - in_channels: int = 4 - ... - - @classmethod - def from_diffusers(cls, config: dict) -> "MyDenoiserConfig": - return cls(...) -``` - -### 2. Create model class - -In `models/my_model.py`: - -```python -class MyDenoiser2DModel(nn.Module): - default_task = "denoising" - category = "Diffusion" - config_class = MyDenoiserConfig - - def __init__(self, config: MyDenoiserConfig): - super().__init__() - # Build architecture matching HF naming - ... - - def forward(self, op, sample, timestep, encoder_hidden_states): - # Denoising forward pass - ... - - def preprocess_weights(self, state_dict): - # Ideally a no-op if naming matches HF - return state_dict -``` - -### 3. Register in _DIFFUSERS_CLASS_MAP - -In `_diffusers_builder.py`, add to `_init_diffusers_class_map()`: - -```python -_DIFFUSERS_CLASS_MAP["MyDenoiser2DModel"] = ( - MyDenoiser2DModel, MyDenoiserConfig, "denoising" -) -``` - -### 4. Add task class (if needed) - -If the existing `DenoisingTask` doesn't match the I/O signature, create a -new task in `tasks/`. Most denoisers use the standard `DenoisingTask`. - -### 5. Add unit test - -In `tests/build_graph_test.py`, add a tiny config: - -```python -("my_denoiser", MyDenoiser2DModel, MyDenoiserConfig( - in_channels=4, out_channels=4, ... # Tiny values -), "denoising"), -``` - -### 6. Match HF weight names - -Compare `named_parameters()` output with HF weight names. Use the -techniques from the **weight-name-alignment** skill to minimize -`preprocess_weights`. - -## Naming conventions for diffusers models - -Diffusers uses different conventions than transformers: - -| Diffusers | Transformers | -|-----------|-------------| -| `GroupNorm` | `LayerNorm` / `RMSNorm` | -| `ResNet` blocks | Decoder layers | -| `down_blocks` / `up_blocks` | `layers` | -| `resnets` / `attentions` (within blocks) | `self_attn` / `mlp` | -| `to_q` / `to_k` / `to_v` / `to_out` | `q_proj` / `k_proj` / `v_proj` / `o_proj` | -| `conv_in` / `conv_out` | `embed_tokens` / `lm_head` | -| `time_embedding` | N/A | -| `encoder_hid_proj` | N/A | - -## Reference files - -| File | Contains | -|------|----------| -| `_diffusers_configs.py` | All diffusers config dataclasses | -| `_diffusers_builder.py` | Pipeline detection, `_DIFFUSERS_CLASS_MAP`, build functions | -| `models/unet.py` | UNet2DConditionModel + all UNet blocks | -| `models/vae.py` | AutoencoderKLModel + encoder/decoder blocks | -| `models/dit.py` | DiTTransformer2DModel + AdaLN blocks | -| `models/flux_sd3.py` | FluxTransformer2DModel + SD3Transformer2DModel | -| `models/controlnet.py` | ControlNetModel | -| `models/adapters.py` | T2IAdapterModel + IPAdapterModel | -| `models/qwen_image.py` | QwenImageTransformer2DModel | -| `models/qwen_image_vae.py` | AutoencoderKLQwenImageModel (3D causal VAE) | -| `models/video_vae.py` | VideoAutoencoderModel (3D video VAE) | -| `tasks/_denoising.py` | DenoisingTask | -| `tasks/_vae.py` | VAETask, QwenImageVAETask | -| `tasks/_controlnet.py` | ControlNetTask | -| `tasks/_adapter.py` | AdapterTask | diff --git a/.github/skills/moe-models/SKILL.md b/.github/skills/moe-models/SKILL.md deleted file mode 100644 index 8a03fb46..00000000 --- a/.github/skills/moe-models/SKILL.md +++ /dev/null @@ -1,383 +0,0 @@ ---- -name: moe-models -description: > - How to represent Mixture-of-Experts models in mobius. Covers gate - variants (TopKGate, SparseMixerGate), MoELayer composition, expert weight - naming conventions, and preprocess_weights mappings. Use this skill when - adding or modifying a model that uses MoE layers. ---- - -# Skill: Mixture-of-Experts (MoE) Models - -## When to use - -Use this skill when adding or modifying a model that uses Mixture-of-Experts -layers — where each token is routed to a subset of expert MLPs. - -## Architecture overview - -``` -MoEDecoderLayer - ├── RMSNorm (input_layernorm) - ├── Attention (self_attn) - ├── RMSNorm (post_attention_layernorm) - └── MoELayer - ├── Gate (routing: input → expert selection + weights) - └── Experts[0..N-1] (each is a standard MLP) -``` - -### Key components - -| Component | File | Purpose | -|-----------|------|---------| -| `MoELayer` | `components/_moe.py` | Routes tokens to experts, combines outputs | -| `TopKGate` | `components/_moe.py` | Standard softmax + top-k routing | -| `SparseMixerGate` | `components/_moe.py` | Sequential selection with threshold masking | -| `MoEDecoderLayer` | `models/moe.py` | Decoder layer that uses `MoELayer` instead of `MLP` | -| `MoETextModel` | `models/moe.py` | Text model with MoE decoder layers | - -## How routing gates work - -### TopKGate (default) - -Standard routing used by most MoE models (Mixtral, GPTOSS): - -1. Compute router logits: `logits = MatMul(hidden_states, gate_weight)` -2. Top-k selection: `values, indices = TopK(logits, k=num_experts_per_tok)` -3. Softmax over selected experts: `weights = Softmax(values)` - -### SparseMixerGate (PhiMoE) - -Sequential expert selection with threshold-based masking: - -1. Compute router logits via `MatMul` -2. For each of `top_k` rounds: - - Find max score (`ReduceMax`) - - Threshold mask: experts whose scores are far from the max (relative to - `jitter_eps`) are masked with `-inf` - - Softmax over non-masked experts - - `TopK(k=1)` to select best expert - - `ScatterElements` to mask out the selected expert for the next round -3. Concatenate all selected expert indices and weights - -## Adding a new MoE model - -### 1. Determine the gate type - -Check the HuggingFace implementation for the routing logic. Look for: - -- `router_type` or `routing_type` in the config -- How `router_logits` are computed and processed -- Whether top-k is applied before or after softmax - -If neither `TopKGate` nor `SparseMixerGate` fits, create a new gate class -in `components/_moe.py`. - -### 2. Check expert MLP naming - -HuggingFace MoE models often use different weight names for expert MLPs: - -| HF name | Our name | Description | -|---------|----------|-------------| -| `w1` | `gate_proj.weight` | Gate projection | -| `w2` | `down_proj.weight` | Down projection | -| `w3` | `up_proj.weight` | Up projection | - -Implement a `_rename_moe_expert_weights()` function if the naming differs: - -```python -def _rename_moe_expert_weights(state_dict): - renamed = {} - for key, value in state_dict.items(): - new_key = key - if ".experts." in key: - new_key = new_key.replace(".w1.", ".gate_proj.") - new_key = new_key.replace(".w2.", ".down_proj.") - new_key = new_key.replace(".w3.", ".up_proj.") - renamed[new_key] = value - return renamed -``` - -### 3. Check the normalization - -Some MoE models use different norms than standard models: - -- **PhiMoE**: Uses `LayerNorm` (with bias), not `RMSNorm` -- **Mixtral**: Uses standard `RMSNorm` - -### 4. Create the model class - -Use `MoETextModel` with the correct gate factory: - -```python -class MyMoECausalLMModel(CausalLMModel): - def __init__(self, config): - nn.Module.__init__(self) - self.config = config - # Pass gate_factory for custom routing - self.model = MoETextModel(config, gate_factory=SparseMixerGate) - self.lm_head = Linear(config.hidden_size, config.vocab_size, bias=True) - - def preprocess_weights(self, state_dict): - state_dict = _rename_moe_expert_weights(state_dict) - return super().preprocess_weights(state_dict) -``` - -### 5. Inject a custom gate - -`MoELayer` accepts an optional `gate` parameter. `MoETextModel` accepts -`gate_factory` — a callable `(config) -> gate_instance` that creates a gate -per layer: - -```python -# Default (TopKGate) -MoETextModel(config) - -# Custom gate -MoETextModel(config, gate_factory=SparseMixerGate) -``` - -To create a new gate, implement a class with this interface: - -```python -class MyGate(nn.Module): - def __init__(self, config): - super().__init__() - self.weight = nn.Parameter([config.num_local_experts, config.hidden_size]) - # ... other params - - def forward(self, op, hidden_states): - # Returns: (expert_weights, expert_indices) - # expert_weights: [batch, seq, num_experts_per_tok] - # expert_indices: [batch, seq, num_experts_per_tok] (INT64) - ... - return weights, indices -``` - -## Config fields for MoE - -```python -config.num_local_experts # Total number of experts (e.g. 8, 16) -config.num_experts_per_tok # Experts activated per token (e.g. 2) -``` - -These are extracted automatically from HuggingFace configs. - -## TopK ONNX op gotcha - -The ONNX `TopK` op requires `K` as a **1-D int64 tensor**, not a Python int: - -```python -# WRONG -op.TopK(logits, self.top_k, axis=-1) - -# CORRECT -k_tensor = op.Constant(value_ints=[self.top_k]) -op.TopK(logits, k_tensor, axis=-1) -``` - -## Qwen3.5-MoE: Hybrid Attention + MoE - -Qwen3.5-MoE combines the hybrid DeltaNet/full-attention architecture of -Qwen3.5 dense with MoE FFN layers instead of a dense MLP. - -### Architecture - -``` -Qwen35MoEDecoderLayer - ├── OffsetRMSNorm (input_layernorm) - ├── GatedDeltaNet / Qwen35Attention (per layer_types) - ├── OffsetRMSNorm (post_attention_layernorm) - └── Qwen35MoEBlock - ├── TopKGate (router) - ├── Experts[0..N-1] (standard MLP: gate/up/down with SiLU) - ├── SharedExpert (MLP: gate/up/down with SiLU) - └── shared_expert_gate → sigmoid scalar -``` - -### Classes - -| Class | File | Purpose | -|-------|------|---------| -| `Qwen35MoEBlock` | `models/qwen.py` | MoE block with routed + shared experts | -| `Qwen35MoEDecoderLayer` | `models/qwen.py` | Hybrid attention + MoE FFN layer | -| `Qwen35MoETextModel` | `models/qwen.py` | Stacks decoder layers with RoPE | -| `Qwen35MoECausalLMModel` | `models/qwen.py` | Top-level causal LM model | - -### MoE block (`Qwen35MoEBlock`) - -- **TopKGate routing**: 256 experts, top-8 in the full model (configurable - via `num_local_experts` / `num_experts_per_tok`) -- **Expert MLPs**: Standard `MLP` (gate/up/down projections, SiLU activation), - each with `moe_intermediate_size` as the intermediate dim -- **Shared expert**: A separate `MLP` that runs on **all** tokens (not routed), - sized by `shared_expert_intermediate_size` -- **Shared expert gating**: `sigmoid(shared_expert_gate(x)) * shared_expert(x)`, - where `shared_expert_gate` is `Linear(hidden_size, 1, bias=False)` - -The key difference from standard MoE is the shared expert: its output is -gated by a learned sigmoid scalar and added to the routed expert output. - -### Config fields - -```python -config.moe_intermediate_size # Intermediate size per expert MLP -config.shared_expert_intermediate_size # Intermediate size for the shared expert -config.num_local_experts # Total number of routed experts (e.g. 256) -config.num_experts_per_tok # Experts activated per token (e.g. 8) -config.layer_types # Per-layer attention type list -``` - -### Weight naming - -HuggingFace weights map directly (no renames needed for expert names): - -``` -mlp.gate.weight → router logits -mlp.experts.N.{gate,up,down}_proj.weight → per-expert MLP -mlp.shared_expert.{gate,up,down}_proj.weight → shared expert MLP -mlp.shared_expert_gate.weight → sigmoid gate (Linear, no bias) -``` - -Note: HF checkpoints store experts as fused tensors -(`experts.gate_up_proj`, `experts.down_proj`). `preprocess_weights()` -unpacks these into per-expert tensors and also renames -`linear_attn.conv1d.weight` → `linear_attn.conv1d_weight`. - -### Testing - -Integration tests use a random-weight HF model with reduced layers and -experts (e.g. 4 layers, 4 experts, top-2 routing). See -`test_qwen35_moe_prefill_logits_match` in `tests/integration_test.py`. - -## Testing MoE models - -MoE integration tests require a model with MoE layers. A good test model -should be small enough for CI (~1-4B params). The test pattern: - -1. Build ONNX model with weights -2. Run prefill + decode against HuggingFace reference -3. Optionally test greedy generation (token-ID matching) - -See `tests/moe_integration_test.py` for the complete pattern. - -## Direct MoE op emission (com.microsoft.MoE) - -OnnxRuntime ships a fused `com.microsoft.MoE` contrib op (CUDA float32/fp16/bf16, -CPU float32/fp16). For new model architectures, **emit it directly** — like -`com.microsoft.GroupQueryAttention` — rather than relying on a rewrite rule. - -### When to use - -Use `com.microsoft.MoE` when: -- The model uses top-k MoE routing (standard softmax gate → TopK) -- All expert weights are the same shape (no dynamic expert counts) -- The EP's `caps.supports_fused_moe` is `True` - -Fall back to the loop-over-experts path when `supports_fused_moe` is `False` -(CPU EP without contrib ops, or EPs that don't support the custom op). - -### Gate output: full pre-topk router_probs - -The op takes the **full** `(num_tokens, num_experts)` probability tensor and -performs top-k selection internally via the `k` attribute. The gate must -produce the full softmax distribution — not already-selected top-k indices. - -```python -# In your gate forward(), return shape [num_tokens, num_experts] -router_probs = op.Softmax(op.MatMul(hidden_states, self.weight), axis=-1) -``` - -### Emission pattern (from Gemma 4 implementation) - -```python -from mobius._build_context import ep_capabilities - -caps = ep_capabilities() -if caps.supports_fused_moe: - moe_out = op.CastLike( - op.MoE( # type: ignore[attr-defined] - normed_hidden, # [num_tokens, hidden_size] - router_probs, # [num_tokens, num_experts] — full pre-topk - self.fc1_experts_weights, # [E, inter_size, hidden_size] - self.fc2_experts_weights, # [E, hidden_size, inter_size] - activation_type="silu", - k=self._top_k, - normalize_routing_weights=1, - _domain="com.microsoft", - ), - normed_hidden, # CastLike: preserve bf16/fp16/fp32 — NOT hardcoded float32 - ) -else: - moe_out = self._dispatch_moe_fallback(op, normed_hidden, router_probs) -``` - -**Critical: use `CastLike` after the MoE op.** The `com.microsoft.MoE` custom -op has `type=None` on its output — ONNX type propagation cannot infer the -output dtype. `op.CastLike(moe_output, target=input)` restores the correct -dtype (bf16/fp16/fp32), which allows downstream ops to share scalar -initializers and avoids hard-coded `Cast` to float32. - -### preprocess_weights: expert weight mapping - -HuggingFace Gemma 4 stores experts as a 3D tensor per projection: -`layers.N.experts.gate_up_proj [E, 2*inter, H]` -`layers.N.experts.down_proj [E, H, inter]` - -Map these to the parameter names used by the ONNX MoE op: - -```python -def preprocess_weights(self, state_dict): - for key in list(state_dict.keys()): - if ".experts.gate_up_proj" in key: - new_key = key.replace(".experts.gate_up_proj", ".fc1_experts_weights") - state_dict[new_key] = state_dict.pop(key) - elif ".experts.down_proj" in key: - new_key = key.replace(".experts.down_proj", ".fc2_experts_weights") - state_dict[new_key] = state_dict.pop(key) - return super().preprocess_weights(state_dict) -``` - -For models that store per-expert weights separately (one matrix per expert), -stack them into 3D tensors in `preprocess_weights`: - -```python -n = config.num_local_experts -gate = torch.stack([state_dict.pop(f"experts.{i}.gate_proj.weight") for i in range(n)]) -down = torch.stack([state_dict.pop(f"experts.{i}.down_proj.weight") for i in range(n)]) -state_dict["fc1_experts_weights"] = gate # [E, inter, hidden] -state_dict["fc2_experts_weights"] = down # [E, hidden, inter] -``` - -### EP capability check (matches GQA pattern) - -```python -# In _execution_providers.py EpCapabilities: -supports_fused_moe: bool = True # set False for EPs without com.microsoft.MoE - -# In model forward(): -from mobius._execution_providers import ep_capabilities -caps = ep_capabilities() -if caps.supports_fused_moe: - # emit com.microsoft.MoE -else: - # fallback loop -``` - -### Fallback: TopKGate + loop dispatch - -When `supports_fused_moe` is False, implement a static unroll: - -```python -def _dispatch_moe_fallback(self, op, hidden, router_probs): - k_tensor = op.Constant(value_ints=[self._top_k]) - top_weights, top_indices = op.TopK(router_probs, k_tensor, axis=-1) - top_weights = op.Softmax(top_weights, axis=-1) # renormalize - output = op.CastLike(op.ConstantOfShape(op.Shape(hidden), value=0.0), hidden) - for e_idx in range(self._num_experts): - w1 = op.Squeeze(op.Gather(self.fc1_experts_weights, [e_idx], axis=0), [0]) - w2 = op.Squeeze(op.Gather(self.fc2_experts_weights, [e_idx], axis=0), [0]) - # expert output, gated by routing weight - ... - return output -``` diff --git a/.github/skills/multi-agent-coordination/SKILL.md b/.github/skills/multi-agent-coordination/SKILL.md deleted file mode 100644 index 5d31d210..00000000 --- a/.github/skills/multi-agent-coordination/SKILL.md +++ /dev/null @@ -1,209 +0,0 @@ ---- -name: multi-agent-coordination -description: How to coordinate multiple AI agents working on the same Git repository simultaneously. Covers worktree isolation, commit coordination, verification gates, context management, and failure recovery patterns. Use this skill when managing parallel workstreams across multiple agents on a shared codebase. ---- - -# Multi-Agent Coordination - -Practical guide for a project lead coordinating 10–17 AI agents working on a single codebase simultaneously. Each section documents a real failure mode encountered in practice and the pattern that prevents it. - ---- - -## 1. Worktree Isolation (CRITICAL) - -**Problem**: Multiple agents sharing a single working directory caused file contamination. One agent's uncommitted changes were overwritten by another checking out a different branch. Work was lost and had to be redone. - -**Solution**: Each agent gets its own git worktree. - -```bash -# Setup — run this before delegating work to an agent -git worktree add .worktrees/- 2>/dev/null || true -cd .worktrees/- -git fetch origin && git checkout && git pull -``` - -**Rules**: -- NEVER let two agents work in the same directory. -- Even "read-only" exploration can cause issues if an agent runs tests that generate artifacts or cache files. -- The main working directory (`/home/.../repo`) is reserved for the lead. Agents always use their worktree. -- Worktrees live at `.worktrees/-` — predictable names make auditing easy. - -**Cleanup**: Remove worktrees after agents are done with `git worktree remove .worktrees/-`. Do NOT clean up while agents are still active — you will destroy their workspace. - ---- - -## 2. Commit Coordination Protocol - -Multiple agents committing to the same branch requires discipline to avoid divergence. - -**Rules**: -1. **Pull before commit**: Always `git pull origin --rebase` immediately before committing. -2. **Push immediately after commit**: Don't let commits sit unpushed. Other agents may be waiting on your changes or about to conflict with yours. -3. **Atomic commits**: Each commit must be self-contained and leave tests passing. Never commit half-finished work. -4. **Descriptive commit messages**: Include what was fixed and the relevant test outcome. Other agents (and the lead) need to understand the diff at a glance. -5. **Conflict resolution**: If `pull --rebase` fails, the agent should report the conflict back to the lead rather than force-pushing or making arbitrary merge decisions. - -**Pattern for each agent commit**: -```bash -git pull origin --rebase # sync first -git add # never git add -A (picks up others' work) -git commit -m "fix: ..." -git push origin # push immediately -``` - ---- - -## 3. Verification Gates - -Checks at each stage prevent problems from compounding across agents. - -### Before starting work -- Verify the branch exists and is up to date. -- Run the relevant tests to establish a baseline — know what was already failing before you touched anything. -- Check for uncommitted changes left by previous agents in your worktree. - -### After each commit -- Run the affected tests immediately (don't batch; catch regressions early). -- Verify the commit landed on the correct branch (`git log --oneline -3`). -- Confirm the push succeeded. - -### Before final review (lead audit) -- Check ALL worktrees for unpushed commits or uncommitted changes. -- Verify the total commit count matches expectations. -- Run the full test suite from a clean checkout (not a worktree). -- Confirm no worktree has diverged from the remote branch. - -### Before PR creation -- Triple review: code correctness, critical logic, readability — ideally by three separate agents. -- All review findings addressed. -- Final test run passes. -- PR description updated with an accurate scorecard of what changed. - ---- - -## 4. Task Dependency Management (DAG) - -Track task dependencies as a directed acyclic graph (DAG). - -**State machine**: `pending → ready → running → done` - -A task becomes `ready` only when ALL its dependencies are `done`. - -**Dependency rules**: -- Review tasks depend on ALL implementation tasks (can't review code that isn't written). -- PR creation depends on ALL reviews passing. -- Integration tests depend on implementation being complete. -- Exploration/investigation tasks have no dependencies — launch them first. - -**Anti-pattern**: Circular dependencies. Reviews depend on implementation, not vice versa. If a review finds a bug, create a new implementation task — don't loop the dependency graph. - ---- - -## 5. Parallel Execution Patterns - -### Safe to parallelize -- Independent model/file fixes (different files, no shared state) -- Code review + readability review + critical review (all read-only) -- Investigation and exploration tasks -- Golden data generation for different models -- Any tasks touching completely separate files - -### Must be serialized -- Two agents editing the same file -- Commit + push sequences within the same agent (not cross-agent) -- Tasks with explicit DAG dependencies (review after implementation) -- Branch cleanup before final review - -### Optimal parallelism -4–6 agents working simultaneously is the sweet spot. Beyond that, coordination overhead increases and merge conflicts become more likely. Scale back if you see frequent push failures or agents blocking each other. - ---- - -## 6. Agent Role Specialization - -Match the agent role to the task type. Mismatched roles waste agent capacity. - -| Role | Best For | Avoid | -|------|----------|-------| -| **Architect** | Investigation, analysis, design decisions, understanding existing code | Simple mechanical fixes | -| **Developer** | Implementation, bug fixes, refactoring | Open-ended exploration | -| **QA Tester** | Verification, auditing, smoke tests, checking other agents' work | New implementation | -| **Code Reviewer** | Correctness, logic bugs, API contracts | Style opinions | -| **Critical Reviewer** | Security, edge cases, invariant violations | Routine fixes | -| **Readability Reviewer** | Naming, clarity, documentation | Implementation work | - -**Delegation tip**: Developers should receive clear instructions — specific files to change, specific test commands to run, and a concrete definition of "done." Architects are better suited for ambiguous investigation tasks. - ---- - -## 7. Communication Patterns - -**Agent → Lead**: Status updates after significant steps, completion summaries, blocker notifications. Don't wait until fully done — send progress updates so the lead can coordinate. - -**Lead → Agent**: Task delegation with full context (worktree path, setup commands, what to change, test commands, commit message format, definition of done). - -**Agent → Agent**: Coordinate through the lead. Agents should not directly orchestrate other agents unless explicitly set up as a sub-lead. - -**Anti-patterns**: -- Sending new tasks to a context-saturated agent. If they don't respond correctly, use a different agent. -- Vague task prompts ("fix the MLP stuff"). Be specific about files, methods, and expected outcomes. -- Assuming agents share context. Each agent starts fresh — always include relevant background in the task prompt. - ---- - -## 8. Failure Recovery - -| Failure | Recovery | -|---------|----------| -| **Lost work** | Check worktrees, `git stash list`, `git reflog`. Most work is recoverable. | -| **Broken tests** | Have the responsible agent fix it immediately. Don't leave broken tests for later. | -| **Merge conflict** | Have the agent that caused the conflict resolve it — they have the most context. | -| **Agent stuck/looping** | Terminate and create a fresh agent. Don't try to unstick a saturated agent. | -| **Wrong branch** | Cherry-pick commits to the correct branch (`git cherry-pick `). Don't redo work. | -| **Agent pushed to wrong branch** | Revert the wrong-branch commit, cherry-pick to the right branch. Coordinate with lead before force-pushing. | -| **Diverged worktree** | `git fetch origin && git rebase origin/` from the worktree. | - ---- - -## 9. Single-Branch Strategy - -For large multi-agent efforts, use **one shared feature branch** rather than one branch per agent. - -**Why**: Avoids complex multi-branch merge scenarios at the end. All agents commit to the same branch via their isolated worktrees. - -**Trade-offs**: -- More `pull --rebase` cycles per agent. -- Occasional push conflicts (recoverable with rebase). -- BUT: much simpler final state, single PR, linear history. - -**Alternative**: Separate branches per agent merged via separate PRs — only viable when features are truly independent and don't share files. - ---- - -## Checklist: Delegating a Task to an Agent - -Before sending a task, verify your prompt includes all of the following: - -1. ☐ **Worktree path** — where the agent should work (`.worktrees/-`) -2. ☐ **Setup commands** — `conda activate`, `git fetch`, `git checkout`, `git pull` -3. ☐ **What to change and why** — specific files, functions, or test cases -4. ☐ **Test commands** — what to run after changes to verify correctness -5. ☐ **Commit message format** — so the history is readable -6. ☐ **Push reminder** — explicit instruction to push after committing -7. ☐ **Definition of done** — what "finished" looks like (tests passing, specific output, etc.) - -**Example delegation prompt**: -``` -Work in .worktrees/developer-. Setup: - git fetch origin && git checkout justinchu/fix-model-parity && git pull - -Fix the Phi-2 MLP weight mismatch: replace the gated MLP (gate_proj/up_proj/down_proj) -with FCMLP (fc1/fc2) in src/mobius/models/phi.py. The HF weights use fc1/fc2. - -After changes, run: - python -m pytest tests/build_graph_test.py -k "phi" -q - -Commit message: "fix: use FCMLP for Phi-1/2 model (HF uses fc1/fc2 not gated MLP)" -Push to justinchu/fix-model-parity immediately after committing. -Done = tests pass, push confirmed. -``` diff --git a/.github/skills/multimodal-models/SKILL.md b/.github/skills/multimodal-models/SKILL.md deleted file mode 100644 index 3c64af5f..00000000 --- a/.github/skills/multimodal-models/SKILL.md +++ /dev/null @@ -1,707 +0,0 @@ ---- -name: multimodal-models -description: > - How to add multimodal (vision + language + audio) models to mobius. - Covers projector variants (Gemma3, MLP, Linear), the VisionModel encoder, - InputMixer, VisionLanguageTask, image/audio token handling, ClippableLinear, - and weight name mappings. Use this skill when adding a model that processes - images, audio, or both alongside text. ---- - -# Skill: Multimodal (Vision + Language + Audio) Models - -## When to use - -Use this skill when adding a model that processes both images and text — such -as Gemma3, Gemma4, LLaVA, LLaVA-NeXT, Phi-3-Vision, PaliGemma, InternVL2, -Pixtral, Idefics2/3, Molmo, Florence2, or Video-LLaVA — or a model that -also processes audio (e.g. Gemma4 with speech/audio inputs). - -## Architecture overview - -``` -pixel_values ──► VisionModel ──► MultiModalProjector ──► InputMixer ──┐ - ├──► TextDecoder ──► logits -input_ids ──────► Embedding ──────────────────────────────────────────┘ -``` - -### Key components - -| Component | File | Purpose | -|-----------|------|---------| -| `VisionModel` | `components/_vision.py` | SigLIP-style patch embedding + transformer encoder | -| `PixtralVisionTower` | `components/_pixtral_vision.py` | Pixtral 2D RoPE vision encoder (bidirectional attention) | -| `PatchEmbedding` | `components/_vision.py` | Conv2d → positional embedding | -| `Gemma3MultiModalProjector` | `components/_multimodal.py` | AvgPool2d → RMSNorm → MatMul | -| `MLPMultiModalProjector` | `components/_multimodal.py` | Linear → GELU → Linear | -| `Mistral3MultiModalProjector` | `components/_pixtral_vision.py` | RMSNorm → spatial merge → Linear → GELU → Linear | -| `LinearMultiModalProjector` | `components/_multimodal.py` | Single Linear | -| `InputMixer` | `components/_multimodal.py` | Scatter vision embeddings at image-token positions | -| `VisionLanguageTask` | `tasks/__init__.py` | ONNX I/O contract with `pixel_values` input | - -## Projector variants - -Different model families use different projectors to bridge vision and text -embedding spaces. Choose the one that matches the HuggingFace implementation: - -| Projector | Architecture | Models | -|-----------|-------------|--------| -| `Gemma3MultiModalProjector` | AvgPool2d → RMSNorm → MatMul | Gemma3 | -| `MLPMultiModalProjector` | Linear → GELU → Linear | LLaVA, LLaVA-NeXT, VipLLaVA, Phi-4-MM, InternVL2, Molmo | -| `Mistral3MultiModalProjector` | RMSNorm → spatial merge → Linear → GELU → Linear | Mistral-3, Pixtral | -| `LinearMultiModalProjector` | Single Linear | PaliGemma, Qwen2-Audio, Idefics2, Florence2 | - -### Gemma3MultiModalProjector - -```python -Gemma3MultiModalProjector( - vision_hidden_size=1152, # SigLIP hidden dim - text_hidden_size=2560, # Text model hidden dim - patches_per_image=64, # sqrt(num_patches) per side - tokens_per_image=256, # mm_tokens_per_image from config - norm=Gemma3RMSNorm(1152), # Gemma3-specific RMSNorm with +1 offset -) -``` - -The pooling kernel is computed as `patches_per_image / sqrt(tokens_per_image)`. -For Gemma3-4B: `64 / 16 = 4`, so `AvgPool2d(kernel_size=4, stride=4)`. - -### MLPMultiModalProjector - -```python -MLPMultiModalProjector( - vision_hidden_size=1024, - text_hidden_size=4096, - bias=True, -) -``` - -Two-layer MLP with GELU activation. The most common projector pattern. - -### LinearMultiModalProjector - -```python -LinearMultiModalProjector( - vision_hidden_size=1024, - text_hidden_size=4096, - bias=True, -) -``` - -Simple single linear layer. - -## Step-by-step: adding a new multimodal model - -### 1. Identify the projector architecture - -Look at the HuggingFace source in `modeling_.py`: - -```bash -grep -n "class.*Projector\|class.*projector" \ - transformers/models//modeling_.py -``` - -Match it to one of the three projector variants, or create a new one. - -### 2. Extract vision config - -Multimodal HF configs have a `vision_config` sub-object. Extract vision -fields in the test or integration code: - -```python -hf_config = transformers.AutoConfig.from_pretrained(model_id) -text_config = hf_config.text_config -vision_config = hf_config.vision_config - -config = ArchitectureConfig.from_transformers(text_config) -# Add vision fields -config.vision_hidden_size = vision_config.hidden_size -config.vision_intermediate_size = vision_config.intermediate_size -config.vision_num_hidden_layers = vision_config.num_hidden_layers -config.vision_num_attention_heads = vision_config.num_attention_heads -config.vision_image_size = vision_config.image_size -config.vision_patch_size = vision_config.patch_size -config.vision_norm_eps = getattr(vision_config, "layer_norm_eps", 1e-6) -config.mm_tokens_per_image = getattr(hf_config, "mm_tokens_per_image", None) -config.image_token_id = getattr(hf_config, "image_token_id", None) -``` - -### 3. Create the model class - -**Important:** Always invoke child modules through `__call__` (not by -accessing their sub-modules directly) so that `onnxscript.nn.Module` pushes -the correct naming context. Direct access like -`self.language_model.model.embed_tokens(op, x)` skips intermediate naming -scopes and produces wrong initializer names. - -The recommended pattern is to pass vision embeddings as a kwarg through the -`__call__` chain, and have the text model perform the mixing internally: - -```python -class _MyTextModelForMultimodal(MyTextModel): - """Text model that mixes vision embeddings into the input.""" - - def __init__(self, config): - super().__init__(config) - self.input_mixer = InputMixer(image_token_id=config.image_token_id or 0) - - def forward(self, op, input_ids, attention_mask, position_ids, - past_key_values=None, vision_embeddings=None): - hidden_states = self.embed_tokens(op, input_ids) - if vision_embeddings is not None: - hidden_states = self.input_mixer( - op, hidden_states, vision_embeddings, input_ids - ) - return super().forward( - op, input_ids, attention_mask, position_ids, - past_key_values=past_key_values, inputs_embeds=hidden_states, - ) - - -class _MyForMultimodalLM(MyCausalLMModel): - """CausalLM that passes vision_embeddings to the text model.""" - - def __init__(self, config): - nn.Module.__init__(self) - self.config = config - self.model = _MyTextModelForMultimodal(config) - self.lm_head = Linear(config.hidden_size, config.vocab_size, bias=False) - - -class MyMultiModalModel(nn.Module): - def __init__(self, config): - super().__init__() - self.config = config - self.vision_tower = VisionModel(config) - self.multi_modal_projector = MLPMultiModalProjector( - vision_hidden_size=config.vision_hidden_size, - text_hidden_size=config.hidden_size, - ) - self.language_model = _MyForMultimodalLM(config) - - def forward(self, op, input_ids, attention_mask, position_ids, pixel_values, - past_key_values=None): - # 1. Encode vision - vision_features = self.vision_tower(op, pixel_values) - vision_embeddings = self.multi_modal_projector(op, vision_features) - - # 2. Pass through __call__ chain — naming is correct automatically - return self.language_model( - op, input_ids, attention_mask, position_ids, - past_key_values=past_key_values, - vision_embeddings=vision_embeddings, - ) -``` - -See `models/gemma3.py` for the full working example. - -### 4. Handle weight name mismatches - -Multimodal HF models often prefix text weights differently: - -| HF key | Our key | -|--------|---------| -| `language_model.model.layers.0.…` | `layers.0.…` | -| `vision_tower.vision_model.encoder.…` | `vision_tower.encoder.…` | -| `multi_modal_projector.mm_input_projection_weight` | `multi_modal_projector.weight` | - -Implement `preprocess_weights` to strip prefixes and rename keys. - -### 5. Handle weight tying - -If `tie_word_embeddings=True`, the HF checkpoint may not include -`lm_head.weight`. Copy it from `embed_tokens.weight`: - -```python -if self.config.tie_word_embeddings: - if "lm_head.weight" not in renamed and "embed_tokens.weight" in renamed: - renamed["lm_head.weight"] = renamed["embed_tokens.weight"] -``` - -### 6. Use VisionLanguageTask - -Build with the `VisionLanguageTask` to add `pixel_values` to graph inputs: - -```python -from mobius.tasks import VisionLanguageTask - -onnx_model = build_from_module(module, config, task=VisionLanguageTask()) -``` - -### 7. PatchEmbedding naming - -`PatchEmbedding` has three parameters. These use explicit `name=` because the -attribute names don't match the desired ONNX names (e.g. `patch_embedding` -needs to map to `patch_embedding.weight`): - -```python -self.patch_embedding = nn.Parameter([...], name="patch_embedding.weight") -self.patch_embedding_bias = nn.Parameter([...], name="patch_embedding.bias") -self.position_embedding = nn.Parameter([...], name="position_embedding.weight") -``` - -In most cases, `name=` is **not needed** because `nn.Module.__setattr__` -automatically sets the parameter name from the attribute name. Only use -`name=` when the attribute name differs from the desired ONNX initializer name. - -## Qwen2.5-VL / Qwen3-VL vision encoder specifics - -These models use a **custom vision encoder** (not SigLIP) with unique -architectural features. The encoder is in -`components/_qwen25_vl_vision.py` and `components/_qwen3_vl_vision.py`. - -### Architecture differences from standard VisionModel - -| Feature | Standard (SigLIP) | Qwen2.5-VL / Qwen3-VL | -|---------|-------------------|----------------------| -| Patch embedding | Conv2d | **Conv3d** (temporal + spatial) | -| Position encoding | Learnable embedding | **2D rotary** (height, width) | -| Attention | Standard self-attention | **Windowed + full attention** alternating | -| Normalization | LayerNorm | **RMSNorm** | -| Output merging | CLS token or mean pool | **Spatial merge** (2×2 → 1) | -| MLP | fc1/fc2 | **Gated MLP** (gate_proj/up_proj/down_proj + SiLU) | - -### Critical: 2D rotary embedding dimension - -The vision encoder computes separate rotary frequencies for height and -width positions. The rotary embedding dimension must be `head_dim // 2` -(not `head_dim`): - -```python -# CORRECT: each spatial dimension gets head_dim//4 frequencies -self.rotary_pos_emb = Qwen25VLVisionRotaryEmbedding(head_dim // 2) - -# WRONG: produces 2× too many frequencies with wrong values -self.rotary_pos_emb = Qwen25VLVisionRotaryEmbedding(head_dim) -``` - -The frequency table has shape `(num_patches, head_dim//2)`. Each half -(`head_dim//4` values) covers one spatial dimension. The `forward` method -concatenates `cos(h_freqs)` and `cos(w_freqs)` to produce the final -`(num_patches, head_dim)` rotary embeddings. - -### Critical: fullatt_block_indexes config - -Qwen2.5-VL uses a **hybrid attention pattern**: most blocks use windowed -attention (local windows for efficiency), but certain blocks use full -attention (all patches attend to all patches): - -```python -# Must be extracted from HF vision_config -fullatt_block_indexes = [7, 15, 23, 31] # For 32-block encoder -window_size = 112 # Window size in patches for windowed blocks -``` - -If `fullatt_block_indexes` is missing, ALL blocks use windowed attention, -causing massive feature divergence (cos ≈ 0.25). The first few blocks may -appear correct since they happen to be windowed blocks. - -**Config extraction** — these must be in `_configs.py` VisionConfig: - -```python -@dataclasses.dataclass -class VisionConfig: - ... - fullatt_block_indexes: list[int] | None = None - window_size: int | None = None -``` - -### Window index and attention bias - -- **Windowed blocks**: Patches are grouped into windows of `window_size`. - Each window attends only within itself. The attention bias is block-diagonal. -- **Full attention blocks**: Use `cu_seqlens` (not `cu_window_seqlens`) to - attend across all patches in each image. -- `window_index` permutes patches into window-ordered layout before the - transformer blocks, then `reverse_indices = argsort(window_index)` restores - the original order after. - -### Multi-image support - -Both vision encoders support multiple images via the ONNX `Scan` op. -Per-image values (position IDs, window indices, cu_seqlens) are computed -in a Scan body and concatenated. See `.github/skills/scan-and-multi-image/SKILL.md`. - -### Spatial merge (post-encoder) - -After the transformer blocks, a spatial merge layer combines 2×2 patches -into 1 token: -``` -(num_patches, hidden_size) → reshape to (num_merged, 4*hidden_size) → MLP → (num_merged, text_hidden_size) -``` -The merge reduces token count by 4× and projects to text model dimension. - -## Qwen3.5-VL - -Qwen3.5-VL uses the same **3-model split** as Qwen3-VL (decoder + vision + -embedding), but swaps the text decoder for the **Qwen3.5 architecture** -which uses hybrid DeltaNet + full attention instead of standard GQA. - -### Architecture - -The vision encoder is **identical to Qwen3-VL** — it reuses -`Qwen3VLVisionModel` (patch_size=16, hidden=1152, depth=27). Only the -text decoder changes. - -| Component | Class | Notes | -|-----------|-------|-------| -| 3-model composite | `Qwen35VL3ModelCausalLMModel` | Splits into decoder + vision + embedding | -| Decoder (standalone) | `Qwen35VLDecoderModel` | Uses `Qwen35TextModel` internally | -| Text model | `Qwen35VLTextModel` | Text-only decoder; strips VL weight prefixes | - -### Registration - -| `model_type` | Variant | Description | -|--------------|---------|-------------| -| `qwen3_5_vl` | 3-model split | Full VLM with vision encoder | -| `qwen3_5_vl_text` | Text-only | Decoder without vision | - -### Task - -Reuses `Qwen3VLVisionLanguage3ModelTask` (task name: `qwen35-vl`). - -### Config - -The HF config is VL-style with a nested `text_config`: - -``` -config.json → model_type: "qwen3_5_vl" -config.text_config → model_type: "qwen3_5" (or "qwen3_5_text") -``` - -### Token IDs - -| Token | ID | -|-------|----| -| `image` | 248056 | -| `video` | 248057 | -| `vision_start` | 248053 | -| `vision_end` | 248054 | - -### Interleaved MRoPE - -Uses `InterleavedMRope` (not `ChunkedMRope`) with: - -- `partial_rotary_factor=0.25` -- `mrope_section=[11, 11, 10]` - -### Key insight - -The vision pipeline is completely shared with Qwen3-VL — only the text -decoder differs (hybrid DeltaNet + full attention). This means vision -encoder bugs/fixes apply to both models equally. - -## Gemma4: vision + audio multimodal - -Gemma4 models (E2B, E4B, 26B-A4B, 31B) support **both vision and audio** -inputs. The architecture has 4 sub-models: decoder, vision encoder, audio -encoder, and embedding. - -### Architecture - -``` -pixel_values ──► [Vision Encoder] ──► image_features ──┐ - │ -audio_features ─► [Audio Encoder] ──► audio_features ──┤ - │ -input_ids ──────► [Embedding] ◄────────────────────────┘ - │ - ▼ - [Text Decoder] ──► logits -``` - -### ClippableLinear (critical for Gemma4) - -Gemma4's vision and audio encoders use `Gemma4ClippableLinear` — a -`Linear` with learned finite input/output activation clamping: - -```python -x = Clip(x, input_min, input_max) -x = x @ weight.T [+ bias] -x = Clip(x, output_min, output_max) -``` - -**This is the single most common source of Gemma4 divergence.** Using -plain `Linear` instead of `ClippableLinear` causes: -- Audio encoder max diff: 52.68 → 0.0003 after fix -- Vision encoder max diff: 3.92 → 0.00007 after fix - -Vision encoder uses ClippableLinear for ALL linear layers: -- Q/K/V/O projections in `Gemma4VisionSelfAttention` -- gate/up/down projections in MLP (via `linear_class=ClippableLinear`) - -Audio encoder uses ClippableLinear for its linear layers as well. - -See `reusable-components` skill for full `ClippableLinear` API reference. - -### Audio boundary markers - -HuggingFace wraps audio tokens with boundary markers that must be present -in `input_ids` for correct generation: - -``` -<|audio> (256000) + N × <|audio|> (258881) + (258883) -``` - -This parallels the image token pattern: -``` -<|image> (255999) + N × <|image|> (258880) + (258882) -``` - -**Missing audio boundary markers** cause garbled audio transcription output -even when the audio encoder output is numerically correct. - -### Per-layer embeddings (CUDA ORT workaround) - -Gemma4 uses per-layer embedding: `embed_tokens_per_layer` with shape -`[V, L*D]` where V=vocab_size, L=num_layers, D=per_layer_dim. For large -models this creates a single Gather on a 2.35B-element tensor, which -**overflows ORT's CUDA Gather kernel** (int32 offset computation in -`gather_impl.cu`). - -**Workaround:** Split into L separate `Embedding([V, D])` tables via -`nn.ModuleList`. In `preprocess_weights`, split the HF weight column-wise: -```python -for i in range(num_layers): - renamed[f"embed_tokens_per_layer.{i}.weight"] = value[:, i*D:(i+1)*D] -``` - -Use `Slice` instead of `Gather` for per-layer projection indexing to -avoid the large-tensor issue entirely. - -### Audio feature extraction - -Gemma4 uses `Gemma4AudioFeatureExtractor` (not `WhisperFeatureExtractor`). -Using the wrong feature extractor produces completely different mel features -and the model fails silently (produces garbage transcription). - -## Testing multimodal models - -### Image token count - -Insert `mm_tokens_per_image` image tokens (not 1!) into the input to match -the number of vision features the projector produces: - -```python -mm_tokens = config.mm_tokens_per_image or 1 -img_tokens = np.full((1, mm_tokens), image_token_id, dtype=np.int64) -input_ids = np.concatenate([input_ids[:, :1], img_tokens, input_ids[:, 1:]], axis=1) -``` - -### Dummy pixel values - -Use random pixel values for testing (we only need numerical parity, not -meaningful images): - -```python -rng = np.random.default_rng(42) -pixel_values = rng.standard_normal((1, 3, image_size, image_size)).astype(np.float32) -``` - -### Tolerances - -Use `rtol=1e-2, atol=1e-2` for multimodal tests — the vision pipeline -introduces more floating-point variance than text-only models. - -### Decode step - -After prefill with image, the decode step is text-only but still needs -`pixel_values` as a graph input (use zeros): - -```python -decode_pixel_values = np.zeros_like(pixel_values) -``` - -## 3-model split for ORT GenAI - -For deployment with onnxruntime-genai, multimodal models are split into -3 separate ONNX models: - -``` -[vision.onnx] pixel_values, grid_thw → image_features -[embedding.onnx] input_ids, image_features → inputs_embeds -[model.onnx] inputs_embeds, attention_mask, position_ids, past_kv → logits, present_kv -``` - -### I/O contract - -| Model | Inputs | Outputs | -|-------|--------|---------| -| Vision | `pixel_values: float32`, `grid_thw: int64` | `image_features: float32` | -| Embedding | `input_ids: int64`, `image_features: float32` | `inputs_embeds: float32` | -| Decoder | `inputs_embeds: float32`, `attention_mask: int64`, `position_ids: int64`, `past_key_values.*` | `logits: float32`, `present.*` | - -### genai_config.json required fields for VLMs - -ORT GenAI needs these fields to compute 3D M-RoPE position_ids for -image tokens. **Without them, image inputs produce wrong output.** - -```json -{ - "model": { - "image_token_id": 151655, - "video_token_id": 151656, - "vision_start_token_id": 151652, - "vision": { - "filename": "vision.onnx", - "config_filename": "processor_config.json", - "spatial_merge_size": 2, - "tokens_per_second": 2.0, - "inputs": { "pixel_values": "pixel_values", "image_grid_thw": "image_grid_thw" }, - "outputs": { "image_features": "image_features" }, - "session_options": { "log_id": "onnxruntime-genai", "provider_options": [] } - }, - "embedding": { - "filename": "embedding.onnx", - "inputs": { "input_ids": "input_ids", "image_features": "image_features" }, - "outputs": { "inputs_embeds": "inputs_embeds" }, - "session_options": { "log_id": "onnxruntime-genai", "provider_options": [] } - } - } -} -``` - -**Required fields** (without these, VLM output is wrong): -- `image_token_id`: Token ID for `<|image_pad|>` — needed for 3D M-RoPE -- `vision_start_token_id`: Token ID for `<|vision_start|>` — marks image boundaries -- `spatial_merge_size`: Grid merge factor (2 for Qwen2.5-VL) — used in position computation - -**Recommended fields** (from olive reference): -- `video_token_id`: Token ID for video frames (151656) -- `config_filename`: Points to `processor_config.json` for image preprocessing -- `tokens_per_second`: Controls temporal position increment (2.0) -- `session_options`: ORT session configuration for each sub-model - -See `.github/skills/ort-genai-config/SKILL.md` for the complete reference -and `.github/skills/debugging-vl-pipeline/SKILL.md` for troubleshooting. - -### processor_config.json for image preprocessing - -ORT GenAI uses ort-extensions for image preprocessing (not HuggingFace). -The `processor_config.json` must use this format: - -```json -{ - "processor": { - "name": "qwen2_5_image_processor", - "transforms": [ - {"operation": {"name": "decode_image", "type": "DecodeImage", "attrs": {"color_space": "RGB"}}}, - {"operation": {"name": "convert_to_rgb", "type": "ConvertRGB"}}, - {"operation": {"name": "resize", "type": "Resize", "attrs": { - "width": 960, "height": 672, - "smart_resize": 1, "min_pixels": 3136, "max_pixels": 12845056, - "patch_size": 14, "merge_size": 2 - }}}, - {"operation": {"name": "rescale", "type": "Rescale", "attrs": {"rescale_factor": 0.00392156862745098}}}, - {"operation": {"name": "normalize", "type": "Normalize", "attrs": {"mean": [0.4814, 0.4578, 0.4082], "std": [0.2686, 0.2613, 0.2758]}}}, - {"operation": {"name": "patch_image", "type": "PatchImage", "attrs": {"patch_size": 14, "temporal_patch_size": 2, "merge_size": 2}}} - ] - } -} -``` - -**Important:** The `width`/`height` in the Resize transform are used as -direct target dimensions, unlike HF's smart_resize which computes targets -from original image dimensions. Compute them as -`round(original_dim / (patch_size * merge_size)) * (patch_size * merge_size)`. - -### Embedding model padding for text-only input - -The embedding model must handle the case where `num_image_tokens=0` -(text-only input). Pad `image_features` with a zero row before Gather -so that the Gather index doesn't go out-of-bounds: - -```python -# In embedding model forward: -padded = op.Concat(zero_row, image_features, axis=0) -gathered = op.Gather(padded, indices, axis=0) -# Where mask selects only real features; padding row is never used -result = op.Where(image_mask, gathered, text_embeddings) -``` - -## Conditional 3-or-4-model task (vision+audio+text) - -Some models come in two tiers: small variants support vision **and** audio -(Any-to-Any), while large variants support vision only (Image-Text-to-Text). -A **single unified task class** handles both tiers by checking whether -`config.audio is not None` to decide whether to include the speech encoder. - -### Tier split - -| Tier | Models | ONNX split | Example | -|------|--------|-----------|---------| -| Small Any-to-Any | E2B, E4B | 4 models: decoder + vision + **speech** + embedding | `google/gemma-4-E2B-it` | -| Large Image-Text-to-Text | 26B-A4B, 31B | 3 models: decoder + vision + embedding | `google/gemma-4-26B-A4B-it` | - -`ArchitectureConfig.from_transformers` populates the `audio` field when the -HuggingFace config contains an audio sub-config; otherwise it is `None`. - -### 4-model task structure (when audio is present) - -``` -decoder inputs_embeds [B, S, H] → logits + KV cache -vision pixel_values [B, N, 3*P^2] → image_features [B*N, H] -speech input_features [B, T, mel] → audio_features [num_audio_tokens, H] -embedding input_ids + image_features + audio_features → inputs_embeds [B, S, H] -``` - -Reference implementation: `Gemma4Task` in `src/mobius/tasks/_gemma4.py`. -This follows the same multi-model structural pattern as -`Phi4MMMultiModalTask` in `src/mobius/tasks/_phi4mm_multimodal.py` -(each modality is a separate ONNX model; embedding splices features at placeholder -positions), though the exact I/O shapes differ per architecture. - -`Gemma4VisionLanguageTask` and `Gemma4AnyToAnyTask` are backward-compatible -aliases pointing to `Gemma4Task`. - -### Speech encoder wiring - -The speech (audio) encoder takes raw mel-spectrogram frames and outputs -token-level features at the text hidden size: - -```python -# In Gemma4Task._build_speech(): -input_features = ir.Value( - name="input_features", - shape=ir.Shape([batch, time, input_size]), # [B, T, 128] - type=ir.TensorType(config.dtype), -) -audio_features = audio_encoder(op, input_features) -# audio_features: [B, T//4, text_hidden_size] -``` - -The audio encoder (`Gemma4AudioEncoder` / `_Gemma4AudioEncoderModel`) is -its own `nn.Module` subgraph exported as the `"speech"` key in the -`ModelPackage`. - -### Embedding model fuses all modalities - -The embedding model receives `input_ids`, `image_features`, and optionally -`audio_features` as inputs and splices them into the token embedding sequence -at the placeholder positions: - -```python -# In Gemma4Task._build_embedding(): -inputs_embeds = embedding( - op, - input_ids=input_ids, # [B, S] - image_features=image_features, # [num_image_tokens, H] - audio_features=audio_features, # [num_audio_tokens, H] — only when audio present -) -# returns inputs_embeds: [B, S, H] -``` - -### Task class conditional pattern - -```python -class MyConditionalTask(ModelTask): - def build(self, module, config): - models = {} - models["decoder"] = self._build_decoder(module.decoder, config) - models["vision"] = self._build_vision(module.vision_encoder, config) - # Build speech encoder only when audio config is present - if config.audio is not None: - models["speech"] = self._build_speech(module.audio_encoder, config) - models["embedding"] = self._build_embedding(module.embedding, config) - return ModelPackage(models, config=config) -``` diff --git a/.github/skills/ort-genai-config/SKILL.md b/.github/skills/ort-genai-config/SKILL.md deleted file mode 100644 index c5c51ca5..00000000 --- a/.github/skills/ort-genai-config/SKILL.md +++ /dev/null @@ -1,799 +0,0 @@ ---- -name: ort-genai-config -description: > - Complete reference for the onnxruntime-genai genai_config.json format, - processor_config.json format, model type registry, and the MultiModal - pipeline architecture. Use this skill when generating genai_config.json, - processor_config.json, debugging ORT GenAI model loading, or integrating - exported ONNX models with the onnxruntime-genai runtime. ---- - -# Skill: ORT GenAI Config Format - -## When to use - -Use this skill when: - -- Writing `genai_config.json` for a new model export -- Writing `processor_config.json` for image/audio preprocessing -- Debugging ORT GenAI model loading errors (protobuf parsing, missing keys) -- Understanding how the ORT GenAI pipeline feeds inputs to vision, embedding, - and decoder models -- Adding support for a new model type in ORT GenAI - -## Overview - -onnxruntime-genai loads models from a directory containing: - -``` -model_dir/ -├── genai_config.json # Required — model config + search params -├── model.onnx # Decoder model -├── model.onnx.data # External weights (optional) -├── vision.onnx # Vision encoder (multimodal only) -├── embedding.onnx # Embedding model (multimodal only) -├── tokenizer.json # Tokenizer (HuggingFace format) -├── tokenizer_config.json # Tokenizer config -├── chat_template.jinja # Chat template (optional) -└── processor_config.json # Image processor (multimodal only) -``` - -## genai_config.json — Top-level structure - -```json -{ - "model": { ... }, - "search": { ... }, - "engine": { ... } -} -``` - -The `engine` section is optional and only used for batched serving. - ---- - -## model section - -### Model-level fields - -| Field | Type | Required | Description | -|---|---|---|---| -| `type` | string | **yes** | Model type identifier (see registry below) | -| `vocab_size` | int | yes | Vocabulary size | -| `context_length` | int | **yes** | Maximum context length; must be > 0 | -| `bos_token_id` | int | no | Beginning-of-sequence token | -| `eos_token_id` | int \| int[] | no | End-of-sequence token(s); defaults to `pad_token_id` | -| `pad_token_id` | int | no | Padding token | -| `sep_token_id` | int | no | Separator token | -| `decoder_start_token_id` | int | no | Decoder start token (encoder-decoder models) | -| `image_token_id` | int | VLM | Token ID for image placeholders (e.g. 151655 for Qwen2.5-VL). **Required** for 3D M-RoPE position ID computation. | -| `video_token_id` | int | no | Token ID for video placeholders (e.g. 151656) | -| `vision_start_token_id` | int | VLM | Token ID for `<\|vision_start\|>` (e.g. 151652). Used to locate image/video regions in input_ids. | - -### Model type registry - -**LLM** (decoder-only, maps to `DecoderOnly_Model`): - -``` -chatglm, decoder, ernie4_5, gemma, gemma2, gemma3_text, gpt2, -gptoss, granite, internlm2, llama, mistral, nemotron, olmo, -phi, phimoe, phi3, phi3small, qwen2, qwen3, smollm3 -``` - -> `gpt2` has a special code path (`Gpt_Model`) but is also in the LLM list. - -**VLM** (vision-language, maps to `MultiModalLanguageModel`): - -``` -fara, gemma3, phi3v, qwen2_5_vl -``` - -**MMM** (multi-modal with vision + audio, maps to `MultiModalLanguageModel`): - -``` -phi4mm -``` - -**ALM** (audio-language, maps to `WhisperModel`): - -``` -whisper -``` - -**Pipeline models** (maps to `DecoderOnlyPipelineModel`): - -``` -phi3small_pipeline, qwen2_5_vl_pipeline -``` - -**Special handling:** - -- `fara` and `qwen2_5_vl` with non-empty `model.decoder.pipeline` → - `Qwen2_5_VL_PipelineModel` -- `IsQwen25VL()` check (type == `"fara"` or `"qwen2_5_vl"`) enables 3D - MRoPE position ID handling - -### Multimodal processor factory - -When `model.create_multimodal_processor()` is called: - -| model.type | Processor class | -|---|---| -| `phi3v` | PhiImageProcessor | -| `whisper` | WhisperProcessor | -| `phi4mm` | PhiMultiModalProcessor | -| `gemma3` | GemmaImageProcessor | -| `fara` | QwenImageProcessor | -| `qwen2_5_vl` | QwenImageProcessor | - -> Models not in this table cannot use `create_multimodal_processor()`. - ---- - -## model.decoder - -The decoder (text model) configuration. - -### Core fields - -| Field | Type | Required | Description | -|---|---|---|---| -| `filename` | string | **yes** | ONNX model filename (e.g. `"model.onnx"`) | -| `hidden_size` | int | yes | Hidden dimension | -| `head_size` | int | yes | Size per attention head | -| `num_attention_heads` | int | yes | Number of query attention heads | -| `num_key_value_heads` | int | yes | Number of KV heads (for GQA) | -| `num_hidden_layers` | int | yes | Number of transformer layers | -| `session_options` | object | no | ORT session configuration | -| `run_options` | object | no | ORT run options | - -### Decoder inputs - -```json -"inputs": { - "input_ids": "input_ids", - "inputs_embeds": "inputs_embeds", - "attention_mask": "attention_mask", - "position_ids": "position_ids", - "past_key_names": "past_key_values.%d.key", - "past_value_names": "past_key_values.%d.value" -} -``` - -The `%d` in `past_key_names` / `past_value_names` is replaced with the layer -index (0 to num_hidden_layers-1) at load time. - -Additional optional inputs for advanced scenarios: - -```json -"past_names": "", -"cross_past_key_names": "", -"cross_past_value_names": "", -"past_key_values_length": "past_key_values_length", -"past_sequence_length": "past_sequence_length", -"current_sequence_length": "current_sequence_length", -"total_sequence_length": "total_sequence_length", -"cache_indirection": "cache_indirection", -"encoder_hidden_states": "encoder_hidden_states", -"encoder_attention_mask": "encoder_attention_mask", -"cumulative_sequence_lengths": "cumulative_sequence_lengths", -"past_sequence_lengths": "past_sequence_lengths", -"block_table": "block_table" -``` - -### Decoder outputs - -```json -"outputs": { - "logits": "logits", - "present_key_names": "present.%d.key", - "present_value_names": "present.%d.value" -} -``` - -### Sliding window (optional) - -```json -"sliding_window": { - "window_size": 4096, - "pad_value": -1, - "alignment": "right", - "slide_key_value_cache": true, - "slide_inputs": true, - "layers": [0, 2, 4] -} -``` - ---- - -## model.embedding - -Required for VLM and MMM models. Merges text token embeddings with vision/audio -features. - -```json -"embedding": { - "filename": "embedding.onnx", - "inputs": { - "input_ids": "input_ids", - "image_features": "image_features", - "audio_features": "audio_features" - }, - "outputs": { - "inputs_embeds": "inputs_embeds" - } -} -``` - ---- - -## model.vision - -Required for VLM and MMM models. - -| Field | Type | Default | Description | -|---|---|---|---| -| `filename` | string | — | Vision ONNX model | -| `config_filename` | string | `"processor_config.json"` | Processor config file | -| `adapter_filename` | string | — | Optional adapter model | -| `spatial_merge_size` | int | 2 | **Required for Qwen2.5-VL.** Controls how many vision patches are merged into one token. Used to compute grid dimensions for 3D M-RoPE position IDs (h/merge × w/merge). | -| `tokens_per_second` | float | 2.0 | Video tokens/second | - -### Vision inputs - -```json -"inputs": { - "pixel_values": "pixel_values", - "image_sizes": "image_sizes", - "image_grid_thw": "image_grid_thw", - "attention_mask": "image_attention_mask" -} -``` - -### Vision outputs - -```json -"outputs": { - "image_features": "image_features" -} -``` - -### Vision pipeline (optional) - -For models that split vision into stages (e.g. patch_embed → attention → -merger): - -```json -"pipeline": [ - { - "filename": "patch_embed.onnx", - "model_id": "patch_embed", - "inputs": ["pixel_values"], - "outputs": ["patch_embeddings"], - "run_on_cpu": false, - "session_options": {} - } -] -``` - ---- - -## model.speech - -For audio-language models (whisper, phi4mm). - -```json -"speech": { - "filename": "speech.onnx", - "config_filename": "audio_processor_config.json", - "inputs": { - "audio_embeds": "audio_embeds", - "attention_mask": "audio_attention_mask", - "audio_sizes": "audio_sizes", - "audio_projection_mode": "audio_projection_mode" - }, - "outputs": { - "audio_features": "audio_features" - } -} -``` - ---- - -## model.encoder - -For encoder-decoder models (whisper). - -```json -"encoder": { - "filename": "encoder.onnx", - "hidden_size": 1280, - "num_attention_heads": 20, - "num_hidden_layers": 32, - "head_size": 64, - "inputs": { - "input_ids": "input_ids", - "attention_mask": "attention_mask" - }, - "outputs": { - "encoder_hidden_states": "encoder_hidden_states" - } -} -``` - ---- - -## search section - -Controls generation behavior. - -| Field | Type | Default | Description | -|---|---|---|---| -| `do_sample` | bool | false | Sampling vs greedy | -| `min_length` | int | 0 | Minimum output length | -| `max_length` | int | context_length | Maximum total length (prompt + output) | -| `batch_size` | int | 1 | Batch size | -| `num_beams` | int | 1 | Beam width (1 = greedy) | -| `num_return_sequences` | int | 1 | Sequences to return | -| `top_k` | int | 50 | Top-K sampling | -| `top_p` | float | 0.0 | Nucleus sampling | -| `temperature` | float | 1.0 | Sampling temperature | -| `repetition_penalty` | float | 1.0 | Repetition penalty (1.0 = none) | -| `length_penalty` | float | 1.0 | Beam search length penalty | -| `early_stopping` | bool | true | Stop beam search early | -| `past_present_share_buffer` | bool | false | Share KV cache buffer (CUDA) | -| `random_seed` | int | -1 | RNG seed (-1 = random) | -| `chunk_size` | int | — | Prefill chunking size | - -### Minimal search config - -```json -"search": { - "do_sample": false, - "max_length": 4096, - "num_beams": 1, - "past_present_share_buffer": false -} -``` - ---- - -## engine section (optional) - -For batched serving. - -```json -"engine": { - "dynamic_batching": { - "block_size": 256, - "num_blocks": 16, - "gpu_utilization_factor": 0.9, - "max_batch_size": 16 - } -} -``` - -Or static batching: - -```json -"engine": { - "static_batching": { - "max_batch_size": 4 - } -} -``` - -Dynamic and static batching are mutually exclusive. - ---- - -## session_options - -Nested inside `decoder`, `encoder`, `vision`, `speech`, or `embedding`. - -```json -"session_options": { - "intra_op_num_threads": 8, - "inter_op_num_threads": 1, - "log_id": "onnxruntime-genai", - "log_severity_level": 2, - "enable_cpu_mem_arena": true, - "enable_mem_pattern": true, - "enable_profiling": "profile.json", - "graph_optimization_level": "ORT_ENABLE_EXTENDED", - "provider_options": [ - { - "cuda": { - "device_id": "0" - } - } - ] -} -``` - -Graph optimization levels: `ORT_DISABLE_ALL`, `ORT_ENABLE_BASIC`, -`ORT_ENABLE_EXTENDED`, `ORT_ENABLE_ALL`. - -Provider names are normalized: `"qnn"` → `"QNN"`, `"dml"` → `"DML"`, -`"webgpu"` → `"WebGPU"`, `"openvino"` → `"OpenVINO"`. - ---- - -## processor_config.json (ort-extensions format) - -> **Critical:** ORT GenAI expects the ort-extensions format — NOT the -> HuggingFace `processor_config.json` format. The HF format wraps data under -> `"image_processor"` with different keys; ORT extensions expects a `"processor"` -> key with an ordered transform pipeline. - -### Qwen2.5-VL example - -```json -{ - "processor": { - "name": "qwen2_5_image_processor", - "transforms": [ - { - "operation": { - "name": "decode_image", - "type": "DecodeImage", - "attrs": { "color_space": "RGB" } - } - }, - { - "operation": { - "name": "convert_to_rgb", - "type": "ConvertRGB" - } - }, - { - "operation": { - "name": "resize", - "type": "Resize", - "attrs": { - "width": 540, - "height": 360, - "smart_resize": 1, - "min_pixels": 3136, - "max_pixels": 12845056, - "patch_size": 14, - "merge_size": 2 - } - } - }, - { - "operation": { - "name": "rescale", - "type": "Rescale", - "attrs": { "rescale_factor": 0.00392156862745098 } - } - }, - { - "operation": { - "name": "normalize", - "type": "Normalize", - "attrs": { - "mean": [0.48145466, 0.4578275, 0.40821073], - "std": [0.26862954, 0.26130258, 0.27577711], - "qwen2_5_vl": 1 - } - } - }, - { - "operation": { - "name": "patch_image", - "type": "PatchImage", - "attrs": { - "patch_size": 14, - "temporal_patch_size": 2, - "merge_size": 2 - } - } - } - ] - } -} -``` - -### Transform types - -| Type | Purpose | Key attrs | -|---|---|---| -| `DecodeImage` | Decode from bytes | `color_space` | -| `ConvertRGB` | Ensure RGB | — | -| `Resize` | Smart resize | `width`, `height`, `smart_resize`, `min_pixels`, `max_pixels`, `patch_size`, `merge_size` | -| `Rescale` | Scale pixel values | `rescale_factor` | -| `Normalize` | Mean/std normalization | `mean`, `std` | -| `PatchImage` | Extract patches | `patch_size`, `temporal_patch_size`, `merge_size` | - -### Generating from HuggingFace config - -```python -from transformers import AutoProcessor - -processor = AutoProcessor.from_pretrained(model_id) -ip = processor.image_processor - -processor_config = { - "processor": { - "name": "qwen2_5_image_processor", - "transforms": [ - {"operation": {"name": "decode_image", "type": "DecodeImage", - "attrs": {"color_space": "RGB"}}}, - {"operation": {"name": "convert_to_rgb", "type": "ConvertRGB"}}, - {"operation": {"name": "resize", "type": "Resize", - "attrs": { - "width": 540, "height": 360, "smart_resize": 1, - "min_pixels": ip.size.get("shortest_edge", 3136), - "max_pixels": ip.size.get("longest_edge", 12845056), - "patch_size": ip.patch_size, - "merge_size": ip.merge_size, - }}}, - {"operation": {"name": "rescale", "type": "Rescale", - "attrs": {"rescale_factor": ip.rescale_factor}}}, - {"operation": {"name": "normalize", "type": "Normalize", - "attrs": { - "mean": list(ip.image_mean), - "std": list(ip.image_std), - "qwen2_5_vl": 1, - }}}, - {"operation": {"name": "patch_image", "type": "PatchImage", - "attrs": { - "patch_size": ip.patch_size, - "temporal_patch_size": ip.temporal_patch_size, - "merge_size": ip.merge_size, - }}}, - ], - } -} -``` - ---- - -## MultiModal pipeline architecture - -### VLM prompt flow (3-model split) - -``` -pixel_values + image_grid_thw → [vision.onnx] → image_features - │ -input_ids + image_features → [embedding.onnx] → inputs_embeds - │ -inputs_embeds + position_ids → [model.onnx] → logits - + past_kv -``` - -### VLM generation flow - -``` -Prompt stage: - 1. VisionState.Run() → image_features - 2. EmbeddingState.ReuseFeaturesBuffer(image_features) - 3. EmbeddingState.Run() → inputs_embeds - 4. DecoderState.Run() → logits + present_kv - 5. VisionState destroyed (no longer needed) - -Token generation stage (loop): - 1. EmbeddingState.Run() → inputs_embeds (from single token) - 2. DecoderState.Run() → logits + present_kv -``` - -### Input flow - -When `generator.set_inputs(named_tensors)` is called: - -1. Tensors matching vision model input names → fed to VisionState -2. Tensors matching embedding model input names → fed to EmbeddingState -3. `input_ids` → used for token counting and embedding lookup -4. `num_image_tokens` → used to allocate image_features buffer size - -### The QwenImageProcessor produces - -| Tensor | Shape | Description | -|---|---|---| -| `input_ids` | (1, seq_len) | Tokenized prompt with image_pad tokens | -| `pixel_values` | (total_patches, C×T×P×P) | Flattened image patches | -| `image_grid_thw` | (num_images, 3) | Grid dimensions per image | -| `num_image_tokens` | (1,) | Total merged image tokens | - -> **Important:** The ORT GenAI QwenImageProcessor does NOT produce -> `cu_seqlens`, `cu_window_seqlens`, or `rotary_pos_ids`. If the vision -> ONNX model requires these, they must be computed externally and injected -> into the NamedTensors. - ---- - -## Common patterns in this package - -### Writing genai_config.json from ArchitectureConfig - -```python -def _write_genai_config(config, output_dir, model_type="qwen2_5_vl"): - genai_config = { - "model": { - "bos_token_id": config.bos_token_id or 151643, - "context_length": 4096, - "decoder": { - "session_options": { - "log_id": "onnxruntime-genai", - "provider_options": [], - }, - "filename": "model.onnx", - "head_size": config.head_dim, - "hidden_size": config.hidden_size, - "inputs": { - "inputs_embeds": "inputs_embeds", - "attention_mask": "attention_mask", - "position_ids": "position_ids", - "past_key_names": "past_key_values.%d.key", - "past_value_names": "past_key_values.%d.value", - }, - "outputs": { - "logits": "logits", - "present_key_names": "present.%d.key", - "present_value_names": "present.%d.value", - }, - "num_attention_heads": config.num_attention_heads, - "num_hidden_layers": config.num_hidden_layers, - "num_key_value_heads": config.num_key_value_heads, - }, - "embedding": { - "filename": "embedding.onnx", - "inputs": { - "input_ids": "input_ids", - "image_features": "image_features", - }, - "outputs": { - "inputs_embeds": "inputs_embeds", - }, - }, - "vision": { - "filename": "vision.onnx", - "spatial_merge_size": 2, - "inputs": { - "pixel_values": "pixel_values", - "image_grid_thw": "image_grid_thw", - }, - "outputs": { - "image_features": "image_features", - }, - }, - "eos_token_id": config.eos_token_id or [151645, 151643], - "pad_token_id": config.pad_token_id or 151643, - "image_token_id": 151655, - "vision_start_token_id": 151652, - "type": model_type, - "vocab_size": config.vocab_size, - }, - "search": { - "do_sample": False, - "early_stopping": True, - "max_length": 4096, - "num_beams": 1, - "num_return_sequences": 1, - "past_present_share_buffer": False, - "repetition_penalty": 1.0, - "temperature": 1.0, - "top_k": 1, - "top_p": 1.0, - }, - } - with open(os.path.join(output_dir, "genai_config.json"), "w") as f: - json.dump(genai_config, f, indent=4) -``` - -### Writing processor_config.json from HuggingFace - -```python -def _write_processor_config(processor, output_dir): - ip = processor.image_processor - config = { - "processor": { - "name": "qwen2_5_image_processor", - "transforms": [ - {"operation": {"name": "decode_image", "type": "DecodeImage", - "attrs": {"color_space": "RGB"}}}, - {"operation": {"name": "convert_to_rgb", "type": "ConvertRGB"}}, - {"operation": {"name": "resize", "type": "Resize", "attrs": { - "width": 540, "height": 360, "smart_resize": 1, - "min_pixels": ip.size.get("shortest_edge", 3136), - "max_pixels": ip.size.get("longest_edge", 12845056), - "patch_size": ip.patch_size, "merge_size": ip.merge_size, - }}}, - {"operation": {"name": "rescale", "type": "Rescale", - "attrs": {"rescale_factor": ip.rescale_factor}}}, - {"operation": {"name": "normalize", "type": "Normalize", "attrs": { - "mean": list(ip.image_mean), "std": list(ip.image_std), - "qwen2_5_vl": 1, - }}}, - {"operation": {"name": "patch_image", "type": "PatchImage", "attrs": { - "patch_size": ip.patch_size, - "temporal_patch_size": ip.temporal_patch_size, - "merge_size": ip.merge_size, - }}}, - ], - } - } - with open(os.path.join(output_dir, "processor_config.json"), "w") as f: - json.dump(config, f, indent=2) -``` - ---- - -## Common errors and fixes - -### "Protobuf parsing failed" - -Missing `model.vision` and/or `model.embedding` sections in genai_config.json. -VLM models require all three model sections. - -### "key 'processor' not found" - -The `processor_config.json` is in HuggingFace format instead of ort-extensions -format. The HF format has `"image_processor"` as the top key; ORT extensions -needs `"processor"` with a transforms pipeline. - -### "Missing Input: cu_window_seqlens" - -The vision ONNX model expects packed-attention inputs that the ORT GenAI -processor doesn't provide. Either: -1. Compute them externally and inject via NamedTensors, or -2. Modify the vision model to compute them from `image_grid_thw` internally - -### "input_ids size exceeds max length" - -For image prompts, the tokenized input_ids (including image_pad tokens) can -be much longer than the default `max_length` in search options. Use -`params.set_search_options(max_length=4096)` or a sufficiently large value. - -### "OrtValue shape verification failed" - -Mismatch between `num_image_tokens` (computed by the processor) and the -actual vision model output shape. Ensure the same image processor is used -consistently — don't mix ORT GenAI processor output with HF processor -pixel_values. - -### Image not recognized despite being processed - -If the model generates coherent text but fails to describe image contents -(e.g. describes "snow" but not the cat in the image): - -1. **Missing `image_token_id` or `spatial_merge_size`:** Without these - config fields, ORT GenAI cannot compute 3D M-RoPE position IDs for - image tokens. The model runs but has no spatial understanding. Add - `image_token_id`, `vision_start_token_id` at model level and - `spatial_merge_size` under model.vision. - -2. **processor_config.json resize mismatch:** The ORT GenAI processor - uses `width`/`height` in the Resize transform as direct target - dimensions (unlike HF which computes from original image size + - min/max pixels). If set too small, the image loses detail. Compute - correct dimensions per-image: - ```python - factor = patch_size * merge_size # 28 - new_h = max(factor, round(orig_h / factor) * factor) - new_w = max(factor, round(orig_w / factor) * factor) - ``` - -3. **ONNX model numerical accuracy:** The ONNX model's logits may differ - from HF (typical max_diff ~8 for VLMs). This causes greedy decoding - to diverge after 3-4 tokens even though the first tokens match. - ---- - -## Reference files - -- **ORT GenAI config structs:** - `/home/justinchu/dev/onnxruntime-genai/src/config.h` -- **ORT GenAI config parsing:** - `/home/justinchu/dev/onnxruntime-genai/src/config.cpp` -- **Model type registry:** - `/home/justinchu/dev/onnxruntime-genai/src/model_type.h` -- **VLM pipeline:** - `/home/justinchu/dev/onnxruntime-genai/src/models/multi_modal.cpp` -- **Qwen image processor:** - `/home/justinchu/dev/onnxruntime-genai/src/models/qwen2_5_vl_image_processor.cpp` -- **Reference processor_config.json:** - `/home/justinchu/dev/onnxruntime-genai/test/test_models/qwen-vision-preprocessing/processor_config.json` -- **Example genai_config generation:** - `examples/qwen25_vl_ort_genai.py` diff --git a/.github/skills/phi4mm-component-parity/SKILL.md b/.github/skills/phi4mm-component-parity/SKILL.md deleted file mode 100644 index 0b6b903c..00000000 --- a/.github/skills/phi4mm-component-parity/SKILL.md +++ /dev/null @@ -1,618 +0,0 @@ ---- -name: phi4mm-component-parity -description: > - How to ensure each multimodal model component (vision encoder, speech encoder, - embedding/projector, text decoder) matches HuggingFace output. Covers pipeline - isolation methodology, common failure modes from real debugging experience, - step-by-step debugging process, and integration test patterns. Applicable to - any multimodal model with similar architecture (Phi4MM, Gemma4, future - audio+vision models). Use this skill when multimodal ONNX model output - diverges from HuggingFace, or when adding a new multimodal model with - multiple encoders. ---- - -# Skill: Multimodal Component Parity Debugging - -## When to use - -Use this skill when: - -- A multimodal ONNX model's logits diverge systematically from HuggingFace -- You're adding a new multimodal model and need to verify each component -- Integration tests fail with large numerical differences (not just tolerance) -- Weights appear to load but the model produces wrong outputs -- You need to isolate which component (vision, speech, embedding, decoder) - is causing divergence - -For vision-language-only models (no speech), see also the -`debugging-vl-pipeline` skill which covers VLM-specific issues like 3D M-RoPE. - -## Pipeline isolation methodology - -Multimodal models with N encoders have N+2 stages (encoders + embedding + -decoder). Debug by comparing each stage independently against HuggingFace -at every boundary. - -### 4-model multimodal pipeline (e.g., Phi4MM) - -``` -pixel_values ──► [1. Vision Encoder] ──► image_features ──┐ - │ -audio_embeds ──► [2. Speech Encoder] ──► speech_features ──┤ - │ -input_ids ──► [3. Embedding/Fusion] ◄───────────────────┘ - │ - ▼ - inputs_embeds - │ - ▼ - [4. Decoder + LoRA] ──► logits -``` - -**Golden rule:** Start from the simplest case (text-only, no encoders), -verify it matches HF, then add one modality at a time. - -### Stage 1: Vision encoder - -Compare SigLIP/ViT encoder output between ONNX and HF. - -```python -# ONNX -vision_session = OnnxModelSession(pkg["vision"]) -onnx_out = vision_session.run({ - "pixel_values": pixel_values, - "image_sizes": image_sizes, # for HD multi-crop models -}) -onnx_features = onnx_out["image_features"] - -# HuggingFace reference -with torch.no_grad(): - hf_features = hf_model.model.embed_tokens_extend.image_embed( - pixel_values_tensor, image_sizes_tensor - ) - -# Compare -cos_sim = cosine_similarity(onnx_features.flatten(), hf_features.flatten()) -print(f"Vision cos_sim: {cos_sim:.6f}") # Should be > 0.99 -``` - -**What to check:** -- Output shape: `(num_image_tokens, text_hidden_size)` after projection -- For HD models: token count varies by image resolution -- Projection MLP maps vision hidden dim → text hidden dim -- Position embeddings added correctly (2D vs 3D shape) - -### Stage 2: Speech encoder - -Compare Conformer encoder output between ONNX and HF. - -```python -# ONNX -speech_session = OnnxModelSession(pkg["speech"]) -onnx_out = speech_session.run({ - "audio_embeds": mel_features, - "audio_sizes": audio_sizes, - "audio_projection_mode": np.array(0, dtype=np.int64), # 0=speech -}) -onnx_features = onnx_out["audio_features"] - -# HuggingFace reference -with torch.no_grad(): - hf_features = hf_model.model.embed_tokens_extend.audio_embed( - mel_tensor, audio_sizes_tensor, input_mode=2 # speech mode - ) -``` - -**What to check:** -- Output shape: `(num_audio_tokens, text_hidden_size)` after projection -- Compression rate: Conformer typically does 8× time reduction -- Projection branch selection: speech vs vision projection -- Conv subsampling produces correct output length - -### Stage 3: Embedding / fusion - -Compare token embedding + feature fusion between ONNX and HF. - -```python -# ONNX -emb_session = OnnxModelSession(pkg["embedding"]) -onnx_embeds = emb_session.run({ - "input_ids": input_ids, - "image_features": image_features, # or zeros([0, H]) if text-only - "audio_features": speech_features, # or zeros([0, H]) if no audio -})["inputs_embeds"] - -# HuggingFace reference (text-only baseline) -with torch.no_grad(): - hf_embeds = hf_model.model.embed_tokens(input_ids_tensor) - -# For text-only, these should match exactly -max_diff = np.abs(onnx_embeds - hf_embeds).max() -print(f"Embedding max_diff: {max_diff:.6f}") # Should be < 1e-5 -``` - -**What to check:** -- Text-only: should match HF `embed_tokens` exactly (no fusion) -- With features: verify token replacement at correct positions -- InputMixer handles zero-length feature tensors without crashing -- Feature positions align with special token positions in input_ids - -### Stage 4: Decoder - -Compare logits between ONNX and HF for the full forward pass. - -```python -# ONNX decoder with KV cache -decoder_session = OnnxModelSession(pkg["model"]) -onnx_logits = decoder_session.run({ - "inputs_embeds": inputs_embeds, - "attention_mask": attention_mask, - "position_ids": position_ids, - # ... past_key_values (zeros for prefill) -})["logits"] - -# HuggingFace full model forward -with torch.no_grad(): - hf_out = hf_model(input_ids=input_ids_tensor, ...) -hf_logits = hf_out.logits.numpy() - -cos_sim = cosine_similarity(onnx_logits[0, -1], hf_logits[0, -1]) -max_diff = np.abs(onnx_logits - hf_logits).max() -print(f"Decoder: cos_sim={cos_sim:.4f}, max_diff={max_diff:.2f}") -``` - -**Typical acceptable metrics (float32):** -- max_diff: 5-10 -- mean_diff: 0.5-1.5 -- cosine similarity: > 0.98 -- First token argmax: matches HF - -## Common failure modes and fixes - -### 1. Weight name alignment (missing weights) - -**Symptoms:** Hundreds or thousands of weights reported as "unmatched" by -`apply_weights`. Model runs but produces garbage output. - -**Root causes encountered:** - -#### a. Module forward() bypass (240 missing weights in Phi4MM) - -The most insidious bug. When a component's `forward()` method directly -accesses nested sub-module parameters (e.g., `self.glu.ext_pw_conv_1d.weight`) -instead of calling the sub-module's `forward()` method, onnxscript cannot -resolve the full module path for the parameter. The weight ends up as an -unnamed, non-initializer constant in the graph. - -```python -# BAD — weights become unnamed -def forward(self, op, x): - return op.Conv(x, self.glu.ext_pw_conv_1d.weight, - self.glu.ext_pw_conv_1d.bias, ...) - -# GOOD — onnxscript resolves full module path -def forward(self, op, x): - return self.glu(op, x) # GLU.forward() calls op.Conv internally -``` - -**Detection:** Check the ONNX graph for Conv/MatMul nodes where weight -inputs have `is_initializer=False` and generic names like "weight"/"bias". - -**Fix:** Add `forward()` methods to sub-modules and call them instead of -directly accessing their parameters. - -#### b. ModuleList subclass causing name doubling (8 weights) - -Subclassing `nn.ModuleList` causes the module's own name to appear twice -in the parameter path: `img_projection.img_projection.0.weight` instead -of `img_projection.0.weight`. - -```python -# BAD — name doubling -class ProjectionMLP(nn.ModuleList): - def __init__(self): - super().__init__() - self.append(nn.Linear(1152, 3072)) - self.append(nn.Linear(3072, 3072)) - -# GOOD — use nn.Module with indexed children -class ProjectionMLP(nn.Module): - def __init__(self): - super().__init__() - layers = [nn.Linear(1152, 3072), nn.Linear(3072, 3072)] - for i, layer in enumerate(layers): - setattr(self, str(i), layer) -``` - -#### c. setattr with dotted names - -Using `setattr(self, "audio_projection.speech", module)` creates a single -attribute with a dot in its name, rather than a nested module. The resulting -ONNX parameter names won't match HuggingFace's `ModuleDict`-style naming. - -**Fix:** Use `nn.ModuleDict` or create proper nested attributes. - -### 2. Shape mismatches (position embedding 2D vs 3D) - -**Symptoms:** `RuntimeError: shape mismatch` during weight loading. - -**Root cause:** The ONNX component declares a parameter with a different -number of dimensions than the HuggingFace weight. Example: PatchEmbedding -declares `position_embedding.weight` as `[num_patches, hidden_size]` (2D), -but HF stores `[1, num_patches, hidden_size]` (3D). - -**Fix in `preprocess_weights()`:** -```python -# Squeeze the extra batch dimension to match ONNX declaration -if "position_embedding.weight" in key and state_dict[key].dim() == 3: - state_dict[key] = state_dict[key].squeeze(0) # [1,N,H] → [N,H] -``` - -**General rule:** Check whether the preprocess_weights transform goes -in the correct direction (squeeze vs unsqueeze). A common mistake is -writing the transform backwards. - -### 3. Dtype mismatches (float64 vs float32) - -**Symptoms:** ONNX Runtime error: "type mismatch in Mul/Add node" during -inference. - -**Root causes:** - -#### a. NumPy default float64 - -`numpy.array(python_float)` defaults to float64. Any constant created -from a Python scalar without explicit dtype will be float64 in the graph. - -```python -# BAD — float64 constant -scale = numpy.array(alpha / rank) # defaults to float64 -op.Mul(x, scale) # Mul(float32, float64) → type error - -# GOOD — explicit float32 -scale = numpy.array(alpha / rank, dtype=numpy.float32) -op.Mul(x, scale) -``` - -#### b. Python int auto-promotion - -When passing a Python `int` to an op that expects a tensor, onnxscript -may auto-promote to float64 (implementation-dependent). - -```python -# RISKY — Python int may become float64 -op.Mul(int64_tensor, self.max_position_embeddings) - -# SAFE — explicit constant -op.Mul(int64_tensor, - op.Constant(value_int=self.max_position_embeddings)) -``` - -**Detection:** Run the ONNX model and look for type mismatch errors. -The error message includes the node name — trace it back to the source. - -### 4. LoRA application mismatch (conditional vs unconditional) - -**Symptoms:** Systematic divergence (> 80% logits mismatch) across ALL -test cases, but the model structurally runs correctly. - -**Root cause:** Some models apply LoRA adapters conditionally based on -input modality. For example, Phi4MM applies: -- `input_mode=0` (text): no adapters -- `input_mode=1` (vision): vision LoRA only -- `input_mode=2` (speech): speech LoRA only -- `input_mode=3` (combined): both adapters - -If the ONNX model unconditionally applies all adapters (both vision and -speech LoRA always active), it diverges from HF when HF uses a different -input mode. - -**Quick fix for integration tests:** Run the HF reference with the mode -that matches the ONNX model's behavior (e.g., `input_mode=3` to match -unconditional application of both adapters). - -**Proper fix:** Add an `input_mode` input to the decoder model and use -conditional logic to selectively apply adapters. - -**Detection:** If text-only inference diverges but the model generates -reasonable (not garbage) output, suspect LoRA mode mismatch. Temporarily -zero out all LoRA weights — if base model matches HF perfectly, the -LoRA application mode is the issue. - -### 5. Empty tensor handling (zero-length features) - -**Symptoms:** Crash during text-only inference when no image/audio -features are present. - -**Root cause:** The embedding model's `InputMixer` uses `GatherElements` -to place features at special token positions. With zero features, the -gather indices are empty but the operation may still execute on the -padded dimension, causing shape errors. - -**Fix pattern:** Zero-pad before Gather, then use Where to mask results: -```python -# Pad with one zero row so Gather never accesses out-of-bounds -padded = op.Concat( - op.ConstantOfShape(op.Constant(value_ints=[1, hidden_size])), - features, # may be [0, hidden_size] - axis=0, -) -# After Gather, mask out the padding positions with Where -result = op.Where(feature_mask, gathered, text_embeddings) -``` - -### 6. HD transform image format (5D vs 4D) - -**Symptoms:** Vision model crashes or produces wrong output with multi-crop -HD images. - -**Root cause:** HD-capable vision models expect images in different formats: -- Some expect `[batch, channels, height, width]` (4D, single crop per batch) -- Others expect `[num_images, num_crops, channels, height, width]` (5D) - -The HF processor output format must match the ONNX model's input format. -If using the HF processor for test input preparation, verify it produces -the expected format. - -**Fix:** Check the HF model's preprocessing code for the expected format, -and ensure the ONNX model's input signature matches. For tests, either: -- Use the HF processor: `processor(images=image, return_tensors="np")` -- Or manually construct the correct format for simple test cases - -### 7. Causal mask construction (inputs_embeds vs input_ids) - -**Symptoms:** Attention mask has wrong length, causing decoder crash or -wrong output. - -**Root cause:** When the decoder receives `inputs_embeds` instead of -`input_ids`, the sequence length must be derived from the embeds tensor -shape, not from input_ids. If the mask is built from input_ids length but -the actual sequence includes fused image/audio tokens, the lengths diverge. - -**Fix:** Always derive `seq_len` from `inputs_embeds.shape[1]` when the -model uses inputs_embeds as input. - -### 8. ClippableLinear not used in encoder (Gemma4) - -**Symptoms:** Encoder output has large numerical divergence (max diff > 1.0) -from HuggingFace, despite all weights loading correctly. - -**Root cause:** Some HuggingFace models (e.g. Gemma4) use -`ClippableLinear` — a linear layer with learned finite input/output -activation clipping — for ALL linear layers in their vision and audio -encoders (attention q/k/v/o projections AND MLP gate/up/down projections). -Using plain `Linear` misses the clamping and causes divergence. - -**Detection:** Check HF source for `ClippableLinear`: -```bash -grep -n "ClippableLinear" transformers/models//modeling_.py -``` - -**Fix:** Use `ClippableLinear` from `mobius.components`: -- For attention: use `ClippableLinear` for q/k/v/o_proj -- For MLP: pass `linear_class=ClippableLinear` parameter - -**Impact (Gemma4):** -- Audio: max diff 52.68 → 0.0003 -- Vision: max diff 3.92 → 0.00007 - -### 9. Missing audio/image boundary tokens - -**Symptoms:** Audio transcription or vision description is garbled, but -encoder output matches HuggingFace. - -**Root cause:** HuggingFace wraps modality placeholder tokens with -boundary markers: -- Image: `<|image>` (open) + N × `<|image|>` (pad) + `` (close) -- Audio: `<|audio>` (open) + N × `<|audio|>` (pad) + `` (close) - -Missing boundary markers prevent the model from correctly identifying -modality regions in the input sequence. - -**Fix:** Ensure `build_input_ids` wraps placeholder tokens with the -correct open/close marker token IDs. - -## Step-by-step debugging process - -### Phase 1: Text-only baseline - -1. **Build ONNX model** with tiny config (for fast iteration) or full - weights (for accuracy). -2. **Run text-only** through embedding → decoder (skip encoders). -3. **Compare embedding output** against `hf_model.model.embed_tokens(ids)`. - If this diverges, the issue is in weight loading or embedding model. -4. **Compare decoder logits** against HF full forward. - If embedding matches but logits diverge, issue is in decoder. - -### Phase 2: Isolate decoder issues - -5. **Check weight count** — verify all expected weights are loaded: - ```python - pkg = build(model_id, load_weights=True) - # apply_weights prints statistics: applied, skipped, unmatched - ``` -6. **Disable LoRA** — if the model uses LoRA, zero out adapter weights and - compare base model output against HF with adapters disabled. -7. **Layer-by-layer** — add intermediate outputs to the ONNX graph (see - `debugging-vl-pipeline` skill) to find which decoder layer first diverges. - -### Phase 3: Add modalities - -8. **Vision only** — run vision encoder, feed features to embedding, compare. -9. **Audio only** — run speech encoder, feed features to embedding, compare. -10. **Combined** — all modalities together. - -At each step, if a newly added component causes divergence, isolate that -component's output against HF. - -### Phase 4: LoRA verification - -11. **Match input modes** — ensure HF reference uses the same adapter - activation mode as ONNX (e.g., `input_mode=3` for both adapters). -12. **Compare with LoRA** — verify LoRA scaling factor: `alpha / rank`. -13. **Check adapter routing** — for multi-adapter models, verify the correct - adapter set is active for each modality combination. - -## Integration test patterns - -### Test configuration - -```python -# Always use for HF reference: -hf_model = AutoModelForCausalLM.from_pretrained( - model_id, - trust_remote_code=True, - attn_implementation="eager", # No flash_attn dependency - torch_dtype=torch.float32, # Match ONNX precision -) - -# For models with conditional LoRA: -hf_model.input_mode = 3 # Match ONNX unconditional LoRA -# Or: pass input_mode=3 to forward() if supported -``` - -### Text-only test - -```python -def test_text_only_prefill_logits_match(self): - input_ids = tokenizer.encode("Hello world", return_tensors="np") - empty_image = np.zeros((0, hidden_size), dtype=np.float32) - empty_audio = np.zeros((0, hidden_size), dtype=np.float32) - - embeds = embedding_session.run({ - "input_ids": input_ids, - "image_features": empty_image, - "audio_features": empty_audio, - })["inputs_embeds"] - - onnx_logits = decoder_session.run({ - "inputs_embeds": embeds, - "attention_mask": np.ones((1, seq_len), dtype=np.int64), - "position_ids": np.arange(seq_len).reshape(1, -1), - # ... zero KV cache - })["logits"] - - hf_logits = hf_model(input_ids=..., input_mode=3).logits.numpy() - assert_logits_close(onnx_logits, hf_logits) -``` - -### Audio test - -```python -def test_audio_prefill_logits_match(self): - # Prepare mel spectrogram input - mel = load_audio_as_mel(audio_path) # [1, n_mel, time] - - speech_out = speech_session.run({ - "audio_embeds": mel, - "audio_sizes": np.array([[mel.shape[-1]]], dtype=np.int64), - "audio_projection_mode": np.array(0, dtype=np.int64), - }) - speech_features = speech_out["audio_features"] - - # Build input_ids with audio placeholder tokens - input_ids = build_audio_input_ids(prompt, num_audio_tokens) - - embeds = embedding_session.run({ - "input_ids": input_ids, - "image_features": np.zeros((0, hidden_size), dtype=np.float32), - "audio_features": speech_features, - })["inputs_embeds"] - - onnx_logits = decoder_session.run(...)["logits"] - hf_logits = hf_forward_with_audio(...) - assert_logits_close(onnx_logits, hf_logits) -``` - -### Tolerance guidelines - -| Precision | atol | rtol | Notes | -|-----------|------|------|-------| -| float32 | 1e-4 | 2e-2 | Standard for single-forward-pass | -| float32 (deep model, 32+ layers) | 1e-3 | 5e-2 | Error compounds over layers | -| float16 / bfloat16 | 0.01 | 0.05 | Wider tolerance for mixed precision | -| Cosine similarity (last token) | > 0.98 | — | Primary correctness metric | -| Argmax match (first prediction) | exact | — | Should always match | - -### Weight loading verification - -After `apply_weights`, check the statistics: -```python -# Expected output: -# Applied: 485/485 weights -# Skipped: 0 (weights in state_dict but not in graph) -# Unmatched: 0 (weights in state_dict with no graph match) - -# If unmatched > 0, dump the names to find alignment issues: -pkg = build(model_id) -state_dict = download_weights(model_id) -state_dict = module.preprocess_weights(state_dict) -# Compare state_dict.keys() vs graph initializer names -``` - -## Vision-specific verification (HD multi-crop) - -For models with HD dynamic resolution (Phi4MM, Phi3-Vision): - -### Input preparation - -```python -from transformers import AutoProcessor - -processor = AutoProcessor.from_pretrained( - model_id, trust_remote_code=True -) -inputs = processor( - images=image, - text=prompt, - return_tensors="pt", -) -pixel_values = inputs["pixel_values"] # [num_crops, C, H, W] or 5D -image_sizes = inputs["image_sizes"] # [num_images, 2] -``` - -### HD transform verification - -The HD transform typically: -1. Splits image into base (global) + sub-image crops -2. Encodes each crop through vision encoder → `[num_patches, hidden_size]` -3. Applies spatial merge (e.g., AvgPool2d + reshape) → compressed tokens -4. Adds learned separators (glb_GN between global/sub, sub_GN between subs) -5. Projects to text dimension via MLP - -```python -# Verify token count matches expected: -# global: (image_size/patch_size)^2 / merge^2 tokens -# per sub-image: same count -# separators: 1 glb_GN + (num_subs - 1) sub_GN rows -total_expected = global_tokens + num_subs * sub_tokens + separator_count -assert image_features.shape[0] == total_expected -``` - -### Testing without HD (simpler) - -For initial validation, use base resolution (single crop, no HD): -```python -# Single image at base resolution — bypasses HD transform -pixel_values = np.random.randn(1, 3, 384, 384).astype(np.float32) -image_sizes = np.array([[384, 384]], dtype=np.int64) -``` - -## Reference files - -- **Integration tests:** `tests/phi4mm_integration_test.py`, - `tests/integration_test.py` -- **VL debugging skill:** `.github/skills/debugging-vl-pipeline/SKILL.md` -- **Model implementation:** `src/mobius/models/phi.py`, - `src/mobius/models/gemma4.py` -- **Audio components:** `src/mobius/components/_audio.py`, - `src/mobius/components/_gemma4_audio.py` -- **Vision components:** `src/mobius/components/_vision.py` -- **ClippableLinear:** `src/mobius/components/_gemma4_audio.py` -- **LoRA component:** `src/mobius/components/_lora.py` -- **Weight loading:** `src/mobius/_weight_loading.py` -- **Feature flags:** `src/mobius/_flags.py` -- **ORT GenAI config skill:** `.github/skills/ort-genai-config/SKILL.md` -- **Weight name alignment skill:** - `.github/skills/weight-name-alignment/SKILL.md` -- **Gemma4 example:** `examples/gemma4_multimodal.py` diff --git a/.github/skills/quality-checklist/SKILL.md b/.github/skills/quality-checklist/SKILL.md deleted file mode 100644 index 3aad778d..00000000 --- a/.github/skills/quality-checklist/SKILL.md +++ /dev/null @@ -1,243 +0,0 @@ ---- -name: quality-checklist -description: > - Definition-of-Done checklist for adding a new model to mobius. Covers - all five test confidence levels (L1–L5), ORT GenAI runtime validation, - Foundry Local smoke-test, Olive quantization compatibility, multi-dtype - and multi-EP correctness, documentation, and code review. Use this skill - when verifying that a new model is truly "done" and ready to merge. ---- - -# Skill: Quality Checklist - -## When to use - -Use this checklist before marking a new model addition as **done**. -Every item must be checked — or explicitly waived with a written reason — -before the PR is merged. - ---- - -## The Checklist - -### 1. Code quality - -- [ ] Model file is in `src/mobius/models/` with the Microsoft MIT copyright - header (`# Copyright (c) Microsoft Corporation. / # Licensed under the - MIT License.`) -- [ ] Class has a descriptive one-paragraph docstring (first paragraph is - used in generated docs) -- [ ] Class has `default_task` and `category` class-level attributes if - the model is not a standard text-generation model -- [ ] `preprocess_weights()` correctly maps every HuggingFace state-dict key - to the ONNX initializer name (verified by the weight-alignment test) -- [ ] All new components use `from mobius.components import ...` (public API), - not private submodule paths -- [ ] No explicit protobuf operations anywhere in new code - (`onnx.helper`, `onnx.TensorProto`, etc. are forbidden) -- [ ] Tensor shapes are annotated in comments after non-trivial operations -- [ ] Automated code review (Copilot/PR review) has been run and all - findings are resolved or explicitly dismissed with a reason -- [ ] `lintrunner -a` is clean — zero lint errors before merging - (`lintrunner f --output oneline --all-files` to auto-fix, then re-run to confirm) - -### 2. L1 — Graph builds - -- [ ] Entry exists in `tests/_test_configs.py` (or a dedicated test method - for VLM / audio models) -- [ ] `is_representative=True` if the model has unique behaviour (custom - class, special attention, MoE, hybrid layers, etc.) -- [ ] `python -m pytest tests/build_graph_test.py -k ""` passes -- [ ] Weight-alignment test passes: - `python -m pytest tests/weight_alignment_test.py -k ""` - -### 3. L2 — Config compatible - -- [ ] YAML test case created at `testdata/cases//.yaml` -- [ ] `test_model_id` field set to a real HuggingFace model ID -- [ ] Schema validates: `python -m pytest tests/yaml_schema_test.py` - -### 4. L3 — Synthetic parity - -- [ ] Model type is covered in `tests/synthetic_parity_test.py` (driven by - `_test_configs.py`; added automatically for text-generation models) -- [ ] `python -m pytest tests/synthetic_parity_test.py -k ""` passes - with `atol=1e-3` / `rtol=1e-3` (or `1e-2` for multimodal) -- [ ] Real-weight parity also checked via `tests/integration_test.py` (add - model to `_TEXT_MODELS` or equivalent if a small checkpoint is available) - -### 5. L4 — Golden match - -- [ ] YAML test case has `level: "L4"` or `"L4+L5"` -- [ ] `inputs.prompts: ["Here is my poem:"]` (standard default prompt unless - audio/image model) -- [ ] Golden file generated: - `python scripts/generate_golden.py --level L4 --filter '*'` -- [ ] Golden file committed to `testdata/golden//.json` -- [ ] `python -m pytest tests/e2e_golden_test.py -m golden -k ""` passes - -### 6. L5 — Generation verified - -- [ ] YAML test case has `level: "L5"` or `"L4+L5"` -- [ ] `generation.max_new_tokens` set (≥ 20 recommended) -- [ ] `generation.do_sample: false` (deterministic greedy decode) -- [ ] Generation golden file generated: - `python scripts/generate_golden.py --level L5 --filter '*'` -- [ ] Generation golden file committed to - `testdata/golden//_generation.json` -- [ ] `python -m pytest tests/e2e_golden_test.py -m generation -k ""` passes - -> **Speech-language models:** The golden generation script supports the -> `speech-language` task type for models that process audio inputs (e.g., -> Gemma4). The `_generate_speech_language()` function in -> `scripts/generate_golden.py` handles audio feature extraction and input -> construction for these models. Ensure the correct feature extractor -> (e.g. `Gemma4AudioFeatureExtractor`, not `WhisperFeatureExtractor`) is -> auto-detected via `AutoFeatureExtractor.from_pretrained()`. - -> **Why L4/L5 matter:** Graph-build tests (L1) only verify ONNX graph -> construction; they never execute the graph with real data. A MatMul shape -> mismatch that crashes at runtime, a wrong normalisation type, or a missing -> scaling multiplier can all pass L1 while producing completely wrong output. -> L4/L5 are the only tests that catch these classes of bugs. - -### 7. Multi-dtype correctness - -- [ ] fp32 tests pass (target: exact token match in greedy generation) -- [ ] fp16 tests pass (target: logit parity `atol=1e-2`; token match for - first N tokens) -- [ ] bf16 tests pass (target: logit parity `atol=1e-2`) - -Use the example `--compare-hf --dtype f16/bf16` flag if a comparison script -exists: - -```bash -python examples/_text_generation.py --compare-hf --dtype f16 -python examples/_text_generation.py --compare-hf --dtype bf16 -``` - -### 8. CLI build - -- [ ] `mobius build --model /tmp/out` completes without error -- [ ] Output directory contains the expected ONNX files and `genai_config.json` - -### 8a. Multi-EP correctness (CUDA) - -- [ ] Model runs correctly with `--ep cuda` (or `--device cuda`) -- [ ] CUDA results match CPU results (compare generation output) -- [ ] No crashes from large tensor operations (see ORT Gather int32 - overflow: microsoft/onnxruntime#28107) -- [ ] `ort_lower_opset_for_ep` flag handles opset 24→23 lowering for - CUDA EP (enabled by default in `src/mobius/_flags.py`) - -### 9. ORT GenAI runtime - -- [ ] Model can be loaded with `ort_genai.Model(output_dir)` without error -- [ ] Greedy text generation produces non-empty, coherent output -- [ ] ORT GenAI test added to `tests/ort_genai_test.py` - (or confirmed covered by an existing parametrized test) - -Run the ORT GenAI integration test: - -```bash -python -m pytest tests/ort_genai_test.py -m integration_slow -k "" -sv -``` - -### 10. Foundry Local smoke test - -- [ ] Model exported package can be loaded and run in Foundry Local -- [ ] At minimum, verify that the `genai_config.json` and all ONNX files - are present and the model responds to a short prompt - -> If Foundry Local is not available in the current environment, document the -> skip with a `# TODO: verify with Foundry Local` comment in the PR. - -### 11. Olive quantization compatibility - -- [ ] Model can be loaded from the exported ONNX package by Olive -- [ ] INT4 / INT8 quantization runs to completion without errors -- [ ] Quantized model produces non-degenerate output (coherent text) -- [ ] If quantization changes the graph structure (e.g. MatMulNBits), verify - the `genai_config.json` still loads correctly in ORT GenAI - -Run the quantization integration test suite to confirm existing patterns -are not broken: - -```bash -python -m pytest tests/quantization_integration_test.py -v -``` - -For new architectures, add a quantized variant test if the architecture has -novel weight layouts (e.g. fused QKV, non-standard expert routing). - -### 12. Documentation - -- [ ] Class docstring (first paragraph) clearly describes the model family - and the HuggingFace class it replicates -- [ ] `default_task` and `category` are set correctly (auto-generates the - model page and index entry) -- [ ] If the model has a notable architectural difference from the base class, - a comment in the source file or the skill notes explains it -- [ ] README model table updated if this is a significant new addition - ---- - -## Waiver policy - -Any item that cannot be completed must be waived explicitly in the PR -description: - -``` -**Waivers:** -- L5 golden: Model is 70B — generating golden data exceeds CI resources. - skip_reason added to YAML. -- Foundry Local: Not available in this environment. Tracked in issue #NNN. -``` - -Unchecked items without a waiver are grounds to request changes before merge. - ---- - -## Quick reference commands - -```bash -# Lint (auto-fix then verify clean) -lintrunner f --output oneline --all-files -lintrunner -a - -# L1 – graph build -python -m pytest tests/build_graph_test.py -k "" - -# L1 – weight alignment -python -m pytest tests/weight_alignment_test.py -k "" - -# L2 – YAML schema -python -m pytest tests/yaml_schema_test.py - -# L3 – synthetic parity -python -m pytest tests/synthetic_parity_test.py -k "" -sv - -# L3 – real-weight integration (if small checkpoint available) -python -m pytest tests/integration_test.py -m integration -k "" -sv - -# L4 – generate golden -python scripts/generate_golden.py --level L4 --filter '*' - -# L4 – run golden test -python -m pytest tests/e2e_golden_test.py -m golden -k "" -v - -# L5 – generate generation golden -python scripts/generate_golden.py --level L5 --filter '*' - -# L5 – run generation golden test -python -m pytest tests/e2e_golden_test.py -m generation -k "" -v - -# ORT GenAI runtime -python -m pytest tests/ort_genai_test.py -m integration_slow -k "" -sv - -# Quantization -python -m pytest tests/quantization_integration_test.py -v - -# CLI build -mobius build --model /tmp/out -``` diff --git a/.github/skills/reusable-components/SKILL.md b/.github/skills/reusable-components/SKILL.md deleted file mode 100644 index 17a4ac6b..00000000 --- a/.github/skills/reusable-components/SKILL.md +++ /dev/null @@ -1,698 +0,0 @@ ---- -name: reusable-components -description: > - Guide to the mobius component library and how to create or extend - reusable building blocks (Attention, MLP, RMSNorm, RoPE, etc.). Covers - parameter naming, ONNX op patterns, design principles (subclass over flags, - model-agnostic components). Use this skill when creating new components or - understanding the component architecture. ---- - -# Skill: Reusable Components - -## When to use - -Use this skill when creating or extending the building blocks that models are -composed from — attention layers, MLPs, normalisations, embeddings, RoPE -variants, and activations. - -## Component library overview - -All components live in `src/mobius/components/` and inherit from -`onnxscript.nn.Module`. Each component's `forward(op, ...)` method builds -ONNX nodes via the `OpBuilder`. - -``` -components/ -├── _activations.py # get_activation(), SiLU module -├── _attention.py # Multi-head / GQA attention with KV cache; Qwen35Attention (gated GQA) -├── _audio.py # ConformerEncoder (NeMo subsampling, T5 bias, Conformer layers) -├── _common.py # Embedding, Linear, LayerNorm, LayerNormNoAffine, GroupNorm, create_attention_bias -├── _gemma4_audio.py # Gemma4 audio encoder: ClippableLinear, ConvSubsampling, SlidingWindowAttention -├── _conv.py # Conv2d (2D convolution with bias and groups) -├── _decoder.py # DecoderLayer (pre-norm residual block) -├── _encoder.py # BertEmbeddings, EncoderAttention, EncoderLayer -├── _lora.py # LoRALinear (base + per-adapter A/B/scale) -├── _mlp.py # Gate-up-down MLP -├── _moe.py # MoELayer, TopKGate, SparseMixerGate -├── _multimodal.py # Projectors + InputMixer -├── _qwen3_vl_vision.py # Qwen3-VL block-diagonal vision encoder -├── _gated_deltanet.py # GatedDeltaNet (recurrent linear attention for Qwen3.5 hybrid) -├── _rms_norm.py # RMSNorm, OffsetRMSNorm (1+weight), GatedRMSNorm (norm * SiLU gate) -├── _rotary_embedding.py # RoPE variants (Default, Linear, Dynamic, Llama3, InterleavedMRope, ChunkedMRope) -├── _vision.py # PatchEmbedding, VisionEncoder, VisionModel -└── _whisper.py # Conv1d, WhisperAttention, WhisperDecoderLayer, WhisperEncoderLayer -``` - -Model files import shared primitives from `components/` and alias them with -an underscore prefix for local use: - -```python -from mobius.components import Conv2d as _Conv2d, SiLU as _SiLU -``` - -Model-specific compound blocks (e.g. `_TimestepEmbedding`, `_DiTBlock`, -`_ResNetBlock2D`) remain in the model files they belong to. - -## How to create a new component - -### 1. Define the class - -```python -from onnxscript import nn -from onnxscript._internal import builder - - -class MyComponent(nn.Module): - def __init__(self, hidden_size: int): - super().__init__() - # Name is automatically "weight" from the attribute name - self.weight = nn.Parameter([hidden_size]) - - def forward(self, op: builder.OpBuilder, hidden_states): - # Build ONNX ops - return op.Mul(hidden_states, self.weight) -``` - -### 2. Parameter naming - -Parameter names are **automatically set** from the attribute name by -`nn.Module.__setattr__`. You do **not** need to pass `name=` when the -attribute name matches the desired ONNX name: - -```python -# GOOD — name is automatically "weight" -self.weight = nn.Parameter([hidden_size]) - -# Only use name= when the attribute name differs from the desired ONNX name -self.patch_embedding = nn.Parameter( - [out_ch, in_ch, kH, kW], name="patch_embedding.weight" -) -``` - -When the component is nested in a module tree, names are automatically -prefixed by parent attribute names: - -```python -# In model: self.layer = MyComponent(...) -# Resulting ONNX name: "layer.weight" -``` - -**Critical:** Parameter names must be unique within a component. If two -parameters share the same attribute name at different levels, one will -silently overwrite the other. - -To create a parameter with precomputed data (e.g. frozen positional embeddings), -use the `data=` argument: - -```python -import onnx_ir as ir -self.embed_positions = nn.Parameter( - [max_positions, d_model], - name="embed_positions.weight", - data=ir.tensor(numpy_array), -) -``` - -Do **not** assign `_const_value` directly. - -### 3. Export from `__init__.py` - -Add to `src/mobius/components/__init__.py`: - -```python -__all__ = [..., "MyComponent"] -from mobius.components._my_component import MyComponent -``` - -### 4. Write unit tests - -Create `_my_component_test.py` alongside the source: - -```python -from mobius._testing import create_test_builder, create_test_input - -class TestMyComponent: - def test_forward(self): - comp = MyComponent(hidden_size=64) - b, op, graph = create_test_builder() - x = create_test_input(b, "x", [1, 10, 64]) - result = comp(op, x) - b._adapt_outputs([result]) - assert graph.num_nodes() > 0 - - def test_parameter_names(self): - comp = MyComponent(hidden_size=64) - names = [n for n, _ in comp.named_parameters()] - assert "weight" in names -``` - -## Component reference - -### Attention - -```python -Attention(config) -Attention(config, scale=0.015625) # Override default 1/sqrt(head_dim) scale -# Inputs: hidden_states, attention_bias, position_embeddings, past_key_value -# Outputs: attn_output, (key_cache, value_cache) -``` - -Handles MHA, GQA, and MQA via `num_key_value_heads`. Supports optional QK -norm (`attn_qk_norm=True`) and bias on Q/K/V/O projections. - -The optional `scale` parameter overrides the default `head_dim**-0.5` attention -scale. Use this when a model specifies a custom attention multiplier (e.g. -Granite's `attention_multiplier`). When `None` (default), uses `1/sqrt(head_dim)`. - -The ONNX `Attention` op (opset 23) has an `is_causal` attribute. For -decoder self-attention in encoder-decoder models (e.g., Whisper), set -`is_causal=1` instead of building an explicit causal mask with -`create_attention_bias`. - -Some models (Whisper) require **Q pre-scaling** for numerical parity with -HuggingFace: multiply Q by `head_dim**-0.5` before passing to `op.Attention` -and set `scale=1.0`. This matches HF's order of operations and avoids -floating-point divergence in softmax. - -**Qwen35Attention** (`_attention.py`): Gated GQA variant for Qwen3.5. Doubles -the Q projection to produce both Q and a gate signal, applies per-head -`OffsetRMSNorm` to Q and K, supports partial RoPE, and gates the output with -`attn_output * sigmoid(gate)`. - -### MLP - -```python -MLP(config) -# Uses gate_proj + up_proj + down_proj with configurable activation -``` - -The activation function comes from `config.hidden_act` and is resolved by -`get_activation()`. - -### DecoderLayer - -```python -DecoderLayer(config) -# Pre-norm residual: LayerNorm → Attention → Add → LayerNorm → MLP → Add -``` - -To customise, subclass and override the components: - -```python -class MyDecoderLayer(DecoderLayer): - def __init__(self, config): - super().__init__(config) - # Replace norm with custom variant - self.input_layernorm = MyRMSNorm(config.hidden_size, eps=config.rms_norm_eps) -``` - -### GatedDeltaNet (Linear Attention) - -```python -GatedDeltaNet(config) -# Inputs: hidden_states, position_embeddings (unused), past_key_value (unused) -# Outputs: output, (conv_state, recurrent_state) -``` - -Recurrent linear attention mechanism from the Qwen3.5 hybrid architecture -(`_gated_deltanet.py`). Key operations: fused QKV projection, causal -depthwise Conv1D, L2-normalised Q/K, exponential decay gates, delta rule -recurrence, and gated output via `GatedRMSNorm`. Supports GQA-like key -head grouping (`num_k_heads` → repeat to `num_v_heads`). State is -`conv_state` + `recurrent_state` (currently zero-initialised for stateless -export). - -### RoPE variants - -Created via the factory function `initialize_rope(config)`: - -| `config.rope_type` | Class | Use case | -|--------------------|-------|----------| -| `"default"` | `DefaultRope` | Standard RoPE | -| `"linear"` | `LinearRope` | Linear scaling (factor in `rope_scaling`) | -| `"dynamic"` | `DynamicNTKRope` | Dynamic NTK scaling | -| `"llama3"` | `Llama3Rope` | LLaMA-3 piecewise scaling | - -**MRope (Multimodal RoPE):** Two variants share a `_MRopeBase` base class -that splits frequencies into temporal (T), height (H), and width (W) sections. -`ChunkedMRope` uses a chunked layout `[TTT...HHH...WWW]` (Qwen2-VL). -`InterleavedMRope` uses an interleaved layout `[T,H,W,T,H,W,...]` and -supports `partial_rotary_factor` for partial RoPE (Qwen3-VL, Qwen3.5). - -RoPE embeddings are precomputed as `cos_cache` / `sin_cache` initializers -and looked up at runtime via `Gather` on `position_ids`. - -### RMSNorm - -```python -RMSNorm(hidden_size, eps=1e-6) -``` - -Uses the ONNX `RMSNormalization` op from opset 23. The `eps` is a float -attribute (not a Parameter). - -For Gemma's `weight + 1` variant, subclass: - -```python -class GemmaRMSNorm(RMSNorm): - def forward(self, op, hidden_states): - weight_plus_one = op.Add(self.weight, 1.0) - return apply_rms_norm(op, hidden_states, weight_plus_one, self.variance_epsilon) -``` - -**OffsetRMSNorm** (`_rms_norm.py`): `output * (1 + weight)` variant where -HuggingFace stores weights initialised to 0, so the effective multiplier is -`1 + weight`. Used by Qwen3.5 for per-head Q/K normalisation. - -**GatedRMSNorm** (`_rms_norm.py`): `RMSNorm(x) * SiLU(gate)` — applies -RMS normalisation then element-wise gates the result with a SiLU activation -on a separate gate input. Used by GatedDeltaNet output projection. - -### LayerNorm - -```python -LayerNorm(hidden_size, eps=1e-6) -``` - -Uses the ONNX `LayerNormalization` op. **Always check the model's HF -config for the correct epsilon** — the default `1e-6` does not match all -models. For example, Whisper uses `1e-5`. A wrong epsilon causes large -numerical drift that amplifies through the network. - -### LayerNormNoAffine - -```python -LayerNormNoAffine(dim, eps=1e-5) -``` - -Layer normalization **without learnable parameters** (`elementwise_affine=False` -in PyTorch). Used in AdaLayerNorm blocks where scale/shift come from a -separate modulation projection. Calls `op.LayerNormalization` with no -`Scale` or `Bias` inputs. - -For weight-free LayerNorm that still needs frozen ones/zeros (e.g. OLMo-1B), -create constant parameters with `data=ir.tensor(...)` instead. - -**Key:** RMSNorm vs LayerNorm is NOT interchangeable. LayerNorm subtracts -the mean; RMSNorm does not. Using the wrong type causes max abs diff > 1.0 -that grows through layers. - -### GroupNorm - -```python -GroupNorm(num_groups, num_channels, eps=1e-5) -``` - -Group normalization with learnable `weight` and `bias`. Uses the ONNX -`GroupNormalization` op. Commonly used in diffusion models (UNet, VAE). - -### Conv2d - -```python -Conv2d(in_channels, out_channels, kernel_size=3, stride=1, padding=0, groups=1) -``` - -2D convolution with bias, matching `torch.nn.Conv2d(bias=True)`. Used in -diffusion models (VAE, UNet, ControlNet) and vision patch embeddings. -Parameters: `weight` (`[out, in/groups, kH, kW]`) and `bias` (`[out]`). - -### SiLU - -```python -SiLU() -# SiLU (Swish) activation as a module: x * sigmoid(x) -``` - -Useful in `nn.Sequential` containers where an activation needs to be a -module with a `forward()` method. For functional use, call -`get_activation("silu")` instead. - -### Linear - -```python -Linear(in_features, out_features, bias=False) -# Uses MatMul (+ optional Add for bias) -``` - -### ClippableLinear - -```python -ClippableLinear(in_features, out_features, bias=False) -# Linear with learned input/output activation clamping -# Matches HuggingFace Gemma4ClippableLinear -``` - -Wraps a standard `Linear` with 4 learned scalar parameters: -`input_min`, `input_max`, `output_min`, `output_max`. Inputs are clamped -before the linear projection and outputs are clamped after: - -```python -x = Clip(x, input_min, input_max) -x = MatMul(x, weight.T) [+ bias] -x = Clip(x, output_min, output_max) -``` - -**Critical:** HuggingFace `Gemma4ClippableLinear` stores *finite* learned -bounds (not ±inf). Using plain `Linear` instead of `ClippableLinear` causes -large numerical divergence in Gemma4 vision and audio encoders: - -- Audio encoder: max diff 52.68 → 0.0003 after fix -- Vision encoder: max diff 3.92 → 0.00007 after fix - -The component is exported from the public API: `from mobius.components -import ClippableLinear`. - -**Weight mapping:** HF stores `.linear.weight` for the actual -weight (the `.linear.` segment is stripped by `preprocess_weights`), and -`.input_min`, `.input_max`, `.output_min`, -`.output_max` as direct scalar buffers. - -The MLP component accepts a `linear_class` parameter, so you can pass -`linear_class=ClippableLinear` to use it for all projections in the MLP. - -### Embedding - -```python -Embedding(num_embeddings, embedding_dim, padding_idx=0) -# Uses Gather on weight matrix -``` - -## Design principles - -1. **Favour subclasses over flags.** When a model family has a unique variant - (e.g. Gemma's `weight + 1` norm), create a subclass rather than adding a - boolean flag to the base class. - -2. **Keep components model-agnostic.** A component should work for any model - that has the right config fields. Model-specific wiring belongs in the - model module. - -3. **One file per concern.** Attention in `_attention.py`, RoPE in - `_rotary_embedding.py`, etc. Tests co-located as `_*_test.py`. - -4. **Reuse across model families.** The same `Attention` component is used by - LLaMA, Mistral, Qwen, Phi, and others. Only override when the - architecture genuinely differs. - -5. **Multiple reusable variants, not one-size-fits-all.** When models need - different behaviour (e.g. MoE gates, projector types), create separate - classes rather than cramming everything into one class with many branches. - -6. **Comment generously with architecture context.** Annotate tensor shapes - after ops (e.g. `# (N, num_heads, head_dim)`), explain multi-step - computations (window reordering, RoPE, spatial merge), and document how - the ONNX graph maps to the HuggingFace reference implementation. - -7. **Match HuggingFace's precision behaviour.** Components must work with any - compute dtype (float32, float16, bfloat16). For numerically sensitive ops - (`exp`, `softplus`, RMSNorm variance), upcast to float32 with - `op.Cast(to=ir.DataType.FLOAT)`, compute, then cast back with `op.CastLike(result, input)`. - For dtype-adaptive parameters, use `op.CastLike(param, reference)`. - See "Precision-sensitive ops" below. - -## Common ONNX op patterns - -### Scalar constants - -Many ONNX ops require tensor inputs, not Python scalars: - -```python -# K for TopK must be a 1-D tensor -k = op.Constant(value_ints=[2]) -values, indices = op.TopK(logits, k, axis=-1) - -# Integer constants -one = op.Constant(value_int=1) - -# Float constants -eps = op.Constant(value_float=1e-6) -``` - -### Dtype-agnostic casting with `CastLike` - -When a parameter or constant needs to match an activation tensor's dtype -without knowing what it is at graph-build time, use `op.CastLike`: - -```python -# GOOD — adapts to whatever dtype hidden_states has -scale = op.CastLike(op.Constant(value_float=1e-6), hidden_states) -``` - -**When `op.Cast(to=...)` IS appropriate:** converting between fundamentally -different types (e.g. int64 position_ids to float for arithmetic, or float -timesteps to the model's compute type), and for the fp32 upcast pattern -below. - -### Precision-sensitive ops: fp32 upcast pattern - -Some operations are numerically unstable in float16/bfloat16 and must run -in float32 to match HuggingFace's behaviour. The pattern is: -**upcast → compute → cast back**. - -```python -# Upcast inputs to fp32 for numerically sensitive exp/softplus -dt_f32 = op.Cast(dt, to=ir.DataType.FLOAT) -dt_f32 = op.Softplus(dt_f32) -a_neg = op.Neg(op.Exp(op.Cast(self.A_log, to=ir.DataType.FLOAT))) -... -# Cast output back to input dtype -y = op.CastLike(y_f32, x) -``` - -**Operations that need fp32 (based on HuggingFace source):** - -| Op | Why | HF pattern | -|----|-----|-----------| -| `Exp` on A_log/decay | Overflow/underflow in fp16 range | `self.A_log.float()` | -| `Softplus` (dt) | Uses exp internally | `softplus(dt + dt_bias)` stays in fp32 context | -| `Exp(dt * A)` (discretisation) | Exponential of product | `A.to(dtype=torch.float32)` | -| SSM state update | Accumulates over many steps | `hidden_states.float()`, `B.float()`, `C.float()` | -| RMSNorm variance | Small values squared then averaged | `hidden_states.to(torch.float32)` | -| GatedRMSNorm (SiLU + norm) | Both gate and variance need fp32 | `gate.to(torch.float32)` | - -**When fp32 upcast is NOT needed:** - -- Linear projections (`MatMul`) — handled by the runtime -- SiLU activation on conv output — stays in model dtype in HF -- Standard attention — ONNX `Attention` op handles precision internally -- `RMSNormalization` op — has `stash_type=1` (default) which auto-upcasts - the variance computation to fp32 - -**Rule of thumb:** Check the HuggingFace source for `.float()` or -`.to(torch.float32)` calls. Every such call indicates an fp32 upcast -region that the ONNX component must replicate with explicit -`op.Cast(to=ir.DataType.FLOAT)` ... `op.CastLike(result, input)` bracketing. - -### Shape manipulation - -Use `op.Shape` with `start` and `end` attributes to extract specific -dimensions directly — do **not** use `Gather(Shape(x), index)`: - -```python -# GOOD — single Shape node with start/end -batch_size = op.Shape(x, start=0, end=1) # 1-D [1]-element tensor -seq_len = op.Shape(x, start=1, end=2) -hidden_dim = op.Shape(x, start=2, end=3) - -# BAD — unnecessary Gather -batch_size = op.Gather(op.Shape(x), [0], axis=0) -``` - -Building dynamic shapes for Reshape/Concat: - -```python -new_shape = op.Concat(batch_size, hidden_dim, op.Constant(value_ints=[-1]), axis=0) -reshaped = op.Reshape(x, new_shape) -``` - -Since `Shape(start, end)` returns a 1-D tensor, it can be passed directly -to ops expecting 1-D shape inputs (e.g. `Slice` starts/ends, `Reshape`, -`Concat` for shape building) without intermediate `Reshape` calls. - -### Module lists and sequential containers - -Use `nn.ModuleList` to register a list of child modules. It automatically -registers children with numeric keys (`"0"`, `"1"`, ...) and supports -iteration, indexing, and `len()`: - -```python -# GOOD — nn.ModuleList -self.layers = nn.ModuleList( - [DecoderLayer(config) for _ in range(config.num_hidden_layers)] -) - -# BAD — manual setattr loop -self.layers = [DecoderLayer(config) for _ in range(config.num_hidden_layers)] -for i, layer in enumerate(self.layers): - setattr(self, f"layers.{i}", layer) -``` - -For sequential containers where children should be called in order (e.g. -matching HF `nn.Sequential`), use `nn.Sequential`. It subclasses -`nn.ModuleList` and adds automatic forward chaining: - -```python -from mobius.components import Linear, SiLU - -# nn.Sequential chains forward calls: SiLU → Linear -self.img_mod = nn.Sequential(SiLU(), Linear(dim, 6 * dim)) - -# Clean call — output chains through each child -result = self.img_mod(op, temb) # equivalent to Linear(SiLU(temb)) -``` - -`nn.Sequential` produces the same parameter names as `nn.ModuleList` -(`img_mod.0.weight`, `img_mod.1.weight`). The key implementation detail: -it overrides `_set_name` to keep children with simple "0", "1" names -(not fully-qualified), because `__call__` already pushes the parent name -onto the scope stack. - -**When to use which:** -- `nn.Sequential` — children are called in a fixed chain (e.g. `to_out`, - modulation layers, FFN with activation gaps) -- `nn.ModuleList` — children need custom iteration logic (e.g. decoder - layers with residual connections, down/up blocks with skip connections) - -For non-consecutive indices (e.g. matching HF `nn.Sequential` with -activation/dropout layers at skipped positions), include parameter-free -placeholder modules to fill the gaps: - -```python -class _NoOpModule(nn.Module): - """Placeholder for HF Dropout (no params, identity at inference).""" - def forward(self, op, x): - return x - -# Matches HF net.0.proj.weight, net.2.weight (Dropout at index 1) -self.net = nn.Sequential( - _GELUGate(dim, inner_dim * 2), # index 0 - _NoOpModule(), # index 1 (Dropout placeholder) - Linear(inner_dim, dim), # index 2 -) -result = self.net(op, x) # chains: GELUGate → NoOp → Linear -``` - -If `nn.Sequential` is not available, fall back to `nn.ModuleList` with -explicit indexing: - -```python -self.img_mod = nn.ModuleList([SiLU(), Linear(dim, 6 * dim)]) -# Manual chaining: -result = self.img_mod[1](op, self.img_mod[0](op, temb)) -``` - -### Conditional operations - -```python -mask = op.Equal(input_ids, op.Constant(value_int=token_id)) -result = op.Where(mask, true_value, false_value) -``` - -### Exposing parameters as graph outputs - -Sometimes a generation loop needs access to model weights for external -computation (e.g. embedding lookups in numpy). Use `op.Identity()` to -expose a parameter as a graph output without affecting the initializer -name used for weight loading: - -```python -class MyModel(nn.Module): - def __init__(self, config): - super().__init__() - # Stacked weight exposed for external lookup - self.stacked_embedding = nn.Parameter([num_groups, vocab, hidden]) - - def forward(self, op, ...): - # Use Identity to create a separate output value. - # This prevents the optimizer from renaming the initializer - # when the task sets a custom output name. - embeddings_out = op.Identity(self.stacked_embedding) - return logits, present_key_values, embeddings_out -``` - -In the task, you can safely rename the Identity output: - -```python -# Safe — Identity separates the output name from the initializer name -embeddings_out.name = "codec_embeddings" -graph.outputs.append(embeddings_out) -``` - -**Important:** Without the Identity node, the optimizer may fold the -reference and setting `output.name = "..."` would rename the -initializer itself, breaking `preprocess_weights` name mapping. - -The generation loop extracts the weights once via a dummy inference: - -```python -weights = session.run(dummy_inputs)["codec_embeddings"] # (N, vocab, H) -# Use as numpy lookup: embed = weights[step, code_id, :] -``` - -## Shared weights with per-layer adapters - -Some architectures reuse the same transformer block across multiple layers, -with per-layer low-rank adapters that differentiate each usage (e.g. Zamba2, -which shares one transformer across 6 hybrid layers). - -### The scope challenge - -`onnxscript.nn` determines ONNX initializer names from the module call stack. -A single module instance called multiple times produces the **same** initializer -names each time — which is exactly what we want for shared weights. But -per-layer adapters need **different** names for each layer. - -### Pattern: split shared + per-instance modules - -```python -class _TextModel(nn.Module): - def __init__(self, config): - super().__init__() - # Shared weights: ONE instance → single set of ONNX initializers - self.shared_transformer = SharedAttentionLayer(config) - - # Per-layer adapters: ModuleList → "adapters.0.*", "adapters.1.*" - self.adapters = nn.ModuleList([ - AdapterModule(config) for _ in range(num_layers) - ]) - - # Shared MLP at model scope if adapter output must mix with MLP - self.gate_proj = Linear(hidden, intermediate) -``` - -### Critical: use `__call__` for per-index scope - -When iterating over adapter ModuleList elements, you **must** call the -element (triggering `__call__`) rather than accessing its sub-attributes: - -```python -# ❌ Broken: adapter_out gets "q_adapter.weight" scope (same for all idx!) -adapter_out = self.adapters[idx].q_adapter(op, x) - -# ✅ Correct: adapter_out gets "adapters.{idx}.q_adapter.weight" scope -adapter_out = self.adapters[idx](op, x) -``` - -### Handling the MLP circular dependency - -When per-layer adapter output must be combined with shared MLP weights, and -the adapter input comes from inside the shared module, split the shared -module into phases: - -1. **Phase 1 — Shared attention:** `shared_transformer(op, x)` → returns - `mlp_input` (pre-MLP hidden states) + KV cache -2. **Phase 2 — Per-layer adapter:** `self.adapters[idx](op, mlp_input)` → - per-layer contribution (correct `adapters.{idx}` scope) -3. **Phase 3 — Shared MLP at caller scope:** Apply gate/up/down projections - registered on the caller module, combining with adapter output - -```python -# In _TextModel.forward(): -mlp_input, kv = self.shared_transformer(op, x) # shared scope -adapter_out = self.mlp_adapters[idx](op, mlp_input) # per-layer scope -gate = op.Add(self.gate_proj(op, mlp_input), adapter_out) # model scope -``` - -**Reference implementation:** `models/zamba2.py` — Zamba2 hybrid Mamba2 + -shared attention with Q/K/V/MLP low-rank adapters. diff --git a/.github/skills/scan-and-multi-image/SKILL.md b/.github/skills/scan-and-multi-image/SKILL.md deleted file mode 100644 index 53392c9b..00000000 --- a/.github/skills/scan-and-multi-image/SKILL.md +++ /dev/null @@ -1,392 +0,0 @@ ---- -name: scan-and-multi-image -description: > - How to use the ONNX Scan op and the Scan + Padding + Compaction pattern in - mobius. Covers building Scan body subgraphs, implicit inputs, - carry states, the rename-to-avoid-SSA-violations workaround, and the - compact_scan_output helper. Primary use case: multi-image vision models - where per-image computations have variable output sizes. - Use this skill when building Scan/Loop subgraphs or adding multi-image - support to a vision model. ---- - -# Skill: ONNX Scan Op & Multi-Image Vision - -## When to use - -Use this skill when: - -- Adding multi-image support to a vision encoder (iterating over - `image_grid_thw` rows). -- Building any ONNX subgraph that iterates over a variable-length sequence - with per-element computation that has dynamic output sizes. -- Working with the `Scan` or `Loop` ONNX ops through the `onnxscript` - builder API. - -## Background: the multi-image problem - -ORT GenAI calls the vision model **once** with all images packed together: - -- `pixel_values`: `(total_patches, pixel_dim)` — all images concatenated. -- `image_grid_thw`: `(num_images, 3)` INT64 — one `[T, H, W]` per image. - -HuggingFace iterates with Python `for t, h, w in grid_thw.tolist()` loops, -accumulating per-image results. ONNX has no Python loops, so we use the -**Scan** op. - -### Why not vectorize? - -Per-image computations like rotary position IDs and windowed attention -indices involve `arange(H)`, `Reshape(…, H_m, ms, W_m, ms)`, etc., where -H and W differ across images. These cannot be batched into a single tensor -operation — we need an explicit loop. - -## The Scan + Padding + Compaction pattern - -Since ONNX Scan requires **fixed-size outputs per iteration** but each -image produces variable-size results, we: - -1. **Pre-compute** `max_size = ReduceMax(per_image_sizes)` in the main - graph. -2. **Pad** each iteration's output to `max_size` inside the Scan body. -3. **Compact** the concatenated Scan output by removing padding using a - boolean mask + `Compress`. - -``` -Main graph: - max_patches = ReduceMax(T * H * W for each image) - │ -Scan body (per image): - pos_ids = compute(T_i, H_i, W_i) → (T_i*H_i*W_i, 2) - padded = Pad(pos_ids, max_patches) → (max_patches, 2) - │ -Scan output: → (num_images, max_patches, 2) - │ -Compact: - mask[i,j] = (j < patches_per_image[i]) - result = Compress(flatten(scan_output), flatten(mask)) - → (total_patches, 2) -``` - -## Helpers: `_scan_utils.py` - -Location: `src/mobius/components/_scan_utils.py` - -### `create_body_graph(state_inputs, scan_inputs, name)` - -Creates a Scan body `ir.Graph` and its `GraphBuilder`. - -```python -from mobius.components._scan_utils import ( - compact_scan_output, - create_body_graph, - rename_subgraph_values, -) - -# No carry state, one scan input per iteration -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 -``` - -### `rename_subgraph_values(graph, prefix)` - -**Critical step.** ONNX Scan body graphs share a value namespace with the -parent graph in ORT. Without renaming, node outputs like `v_Constant_11` -in the body collide with identically named values in the main graph, -causing an "SSA form violation" error. - -Call this **after** building all body graph ops and **before** calling -`op.Scan(...)` on the main graph: - -```python -rename_subgraph_values(body_graph, "rotary_body_") -``` - -The prefix must be unique per Scan in the model (e.g. `"rotary_body_"`, -`"win_body_"`, `"cu_body_"`). Graph input/output names are NOT renamed -— they define the Scan interface. - -### `compact_scan_output(op, scan_result, lengths_per_iter)` - -Removes padding from a `(num_iters, max_len, ...)` Scan output using a -boolean mask built from actual per-iteration lengths. - -```python -# After Scan -result = compact_scan_output(op, scan_result, patches_per_image) -# → (total_patches, ...) -``` - -## Step-by-step: building a Scan - -### 1. Compute per-image sizes in the main graph - -Extract column vectors from `grid_thw` using Slice + Squeeze. - -**Important:** Always specify the squeeze axis to avoid collapsing the -batch dimension when `num_images == 1`: - -```python -# GOOD — squeeze only the column axis -T_col = op.Squeeze(op.Slice(grid_thw, [0], [1], [1], [1]), [1]) - -# BAD — squeezes ALL size-1 dims; scalar when num_images=1 -T_col = op.Squeeze(op.Slice(grid_thw, [0], [1], [1], [1])) -``` - -Compute derived values: - -```python -H_col = op.Squeeze(op.Slice(grid_thw, [1], [2], [1], [1]), [1]) -W_col = op.Squeeze(op.Slice(grid_thw, [2], [3], [1], [1]), [1]) -patches_per_image = op.Mul(T_col, op.Mul(H_col, W_col)) # (N,) -max_patches = op.ReduceMax(patches_per_image, keepdims=False) # scalar -``` - -### 2. Create the body graph - -```python -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 -``` - -### 3. Build per-image computation in the body - -Extract T, H, W from the scan input and compute: - -```python -bT = body_op.Squeeze(body_op.Gather(body_thw, body_op.Constant(value_int=0))) -bH = body_op.Squeeze(body_op.Gather(body_thw, body_op.Constant(value_int=1))) -bW = body_op.Squeeze(body_op.Gather(body_thw, body_op.Constant(value_int=2))) - -result = some_per_image_computation(body_op, bT, bH, bW) -``` - -**Helper function pattern:** Extract the per-image logic into a -standalone function that takes `op` as a parameter. This function works -with either the main graph's `op` or the Scan body's `body_op`: - -```python -def _compute_one_image(op, T, H, W, ms): - """Works with any OpBuilder.""" - H_m = op.Div(H, op.Constant(value_int=ms)) - # ... computation ... - return result - -# In main graph (single-image fast path): -result = _compute_one_image(op, T, H, W, ms) - -# In Scan body: -result = _compute_one_image(body_op, bT, bH, bW, ms) -``` - -### 4. Pad the output to max_size - -```python -num_p = body_op.Mul(bT, body_op.Mul(bH, bW)) -pad_len = body_op.Reshape(body_op.Sub(max_patches, num_p), [1]) - -# For 2D output (patches, D): pads = [0, 0, pad_len, 0] -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(result, pads, body_op.Constant(value_int=-1)) -``` - -`max_patches` is an **implicit input** from the main graph — the Scan -body references it directly. The ONNX Scan spec allows body graphs to -reference outer-scope values. This is handled correctly by -`onnxscript`'s builder and `onnx_ir`'s serializer. - -### 5. Set body outputs and rename - -```python -padded.name = "padded_output" -body_graph.outputs.append(padded) - -rename_subgraph_values(body_graph, "my_scan_body_") -``` - -### 6. Call Scan on the main graph - -```python -scan_result = op.Scan( - grid_thw, # scan input (iterated over axis 0) - body=body_graph, # body subgraph - num_scan_inputs=1, # number of scan inputs - _outputs=1, # number of outputs -) -# scan_result: (num_images, max_patches, D) -``` - -### 7. Compact the result - -```python -result = compact_scan_output(op, scan_result, patches_per_image) -# → (total_patches, D) -``` - -## Advanced: carry states - -Carry states (also called "state variables") persist across Scan -iterations. Use them for accumulating offsets. - -**Body graph inputs:** `[state_1, state_2, ..., scan_input_1, ...]` -**Body graph outputs:** `[new_state_1, new_state_2, ..., scan_out_1, ...]` - -```python -# State inputs -body_offset = ir.Value( - name="offset", shape=ir.Shape([]), - type=ir.TensorType(ir.DataType.INT64), -) -body_thw = ir.Value( - name="body_thw", shape=ir.Shape([3]), - type=ir.TensorType(ir.DataType.INT64), -) - -# 1 carry state + 1 scan input -body_graph, body_builder = create_body_graph( - [body_offset], [body_thw], name="window_body", -) -body_op = body_builder.op - -# ... compute per-image result ... -win_idx = body_op.Add(local_indices, body_offset) # add offset -new_offset = body_op.Add(body_offset, total_merged) # update carry - -# Outputs: carry states FIRST, then scan outputs -new_offset.name = "new_offset" -padded_idx.name = "padded_window_index" -body_graph.outputs.extend([new_offset, padded_idx]) - -rename_subgraph_values(body_graph, "win_body_") - -# Call with initial state values -init_offset = op.Constant(value_int=0) -final_offset, scan_idx = op.Scan( - init_offset, # initial carry state - grid_thw, # scan input - body=body_graph, - num_scan_inputs=1, - _outputs=2, # 1 carry output + 1 scan output -) -``` - -### Carry state use cases - -| Use case | Carry state | Updated as | -|----------|-------------|------------| -| Window index global offsets | `merged_offset` (INT64 scalar) | `+= T * llm_h * llm_w` per image | -| cu_window_seqlens offsets | `cu_offset` (INT64 scalar) | `= last cu_window value` per image | -| Running patch count | `patch_offset` (INT64 scalar) | `+= T * H * W` per image | - -## Pitfalls and gotchas - -### 1. SSA violations from name collisions - -**Problem:** ORT rejects models where a body graph value name matches a -main graph value name (e.g. both have `v_Constant_11`). - -**Fix:** Always call `rename_subgraph_values(body_graph, "unique_prefix_")` -before `op.Scan(...)`. - -### 2. Squeeze removes batch dim when N=1 - -**Problem:** `op.Squeeze(tensor)` removes ALL size-1 dims. When -`num_images == 1`, a `(1,)` tensor becomes a scalar, causing downstream -`Unsqueeze` or `Reshape` failures. - -**Fix:** Always specify the axis: `op.Squeeze(tensor, [1])`. - -### 3. Pad format for multi-dim outputs - -ONNX `Pad` pads format for an N-dim tensor: -`[d0_begin, d1_begin, ..., dN_begin, d0_end, d1_end, ..., dN_end]` - -For a 2D `(rows, cols)` tensor padded only on rows: -`pads = [0, 0, pad_rows, 0]` - -For a 1D `(len,)` tensor: -`pads = [0, pad_len]` - -### 4. Implicit inputs from parent graph - -Scan body graphs can reference values from the parent graph (implicit -inputs). This is supported by the ONNX spec, `onnxscript`, and ORT. -Use this for `max_patches`, `max_merged`, learned parameters like -`self.pos_embed`, etc. - -**No special syntax needed:** just use the main-graph `ir.Value` directly -in `body_op` operations. - -### 5. Body graph needs opset imports - -`create_body_graph()` already handles this (sets `opset_imports={"": 23}`). -If building manually, ensure the body graph has opset imports. - -## Reference files - -| File | Content | -|------|---------| -| `src/mobius/components/_scan_utils.py` | `create_body_graph`, `rename_subgraph_values`, `compact_scan_output` | -| `src/mobius/components/_qwen25_vl_vision.py` | Qwen2.5-VL multi-image: rotary, window index, cu_seqlens via Scan | -| `src/mobius/components/_qwen3_vl_vision.py` | Qwen3-VL multi-image: rotary, cu_seqlens, pos embed interpolation via Scan | - -### Qwen2.5-VL examples (3 Scans) - -1. **`_compute_rotary_pos_ids`** — No carry state. Pads `(T*H*W, 2)` to - `(max_patches, 2)`. -2. **`_compute_window_index`** — Two carry states (`merged_offset`, - `cu_offset`). Two scan outputs (window index, cu_window). -3. **`_compute_cu_seqlens`** — No carry state. Pads `(T,)` of hw values - to `(max_T,)`. Post-Scan: compact → CumSum → Pad with leading 0. - -### Qwen3-VL examples (3 Scans) - -1. **`_compute_rotary_pos_ids`** — Same pattern as Qwen2.5-VL but with - block-row/col indexing. -2. **`_compute_cu_seqlens`** — Same as Qwen2.5-VL. -3. **`_interpolate_pos_embed`** — No carry state. Bilinear interpolation - of learned embeddings per image. References `self.pos_embed` as - implicit input. Pads `(T*H*W, hidden_size)` to `(max_patches, D)`. - -## Testing Scan-based code - -Unit tests (`build_graph_test.py`) verify graph construction only. To -verify Scan correctness at runtime, build the vision model, fill -initializers with random weights, and run with ORT: - -```python -import numpy as np -import onnx_ir as ir -import onnxruntime as ort - -# Build and fill weights... -sess = ort.InferenceSession(model_path) - -# Single image -r1 = sess.run(None, { - "pixel_values": np.random.randn(4, pd).astype(np.float32), - "image_grid_thw": np.array([[1, 2, 2]], dtype=np.int64), -}) -assert r1[0].shape[0] == 1 # 4 patches / smu(4) = 1 merged - -# Two different-size images -r2 = sess.run(None, { - "pixel_values": np.random.randn(12, pd).astype(np.float32), - "image_grid_thw": np.array([[1, 2, 2], [1, 2, 4]], dtype=np.int64), -}) -assert r2[0].shape[0] == 3 # (4+8) / 4 = 3 merged -``` diff --git a/.github/skills/weight-name-alignment/SKILL.md b/.github/skills/weight-name-alignment/SKILL.md deleted file mode 100644 index 81f0539d..00000000 --- a/.github/skills/weight-name-alignment/SKILL.md +++ /dev/null @@ -1,416 +0,0 @@ ---- -name: weight-name-alignment -description: > - How to align ONNX parameter names with HuggingFace weight names to simplify - or eliminate preprocess_weights renames. Covers nn.ModuleList for Sequential - patterns, wrapper modules for nesting, placeholder modules, non-consecutive - indices, and which rename categories cannot be eliminated. Use this skill - when adding or modifying a model's preprocess_weights method. ---- - -# Skill: Weight Name Alignment - -## When to use - -Use this skill when: -- Adding a new model and designing `preprocess_weights` -- Simplifying an existing model's `preprocess_weights` method -- Debugging weight loading failures (mismatched parameter names) -- Deciding whether to restructure model construction vs. rename in - `preprocess_weights` - -## Core principle - -**The best `preprocess_weights` is a no-op.** Most renames exist because the -ONNX module hierarchy doesn't match HuggingFace's. By restructuring -`nn.Module` construction to produce parameter names that match HF directly, -you can eliminate renames entirely. - -## How parameter names are formed - -In `onnxscript.nn`, parameter names are built from the Python attribute chain: - -```python -class MyModel(nn.Module): - def __init__(self): - self.layers = nn.ModuleList([MyLayer()]) - # layers[0].weight → "layers.0.weight" - -class MyLayer(nn.Module): - def __init__(self): - self.linear = _Linear(4, 4) - # linear.weight → "linear.weight" -``` - -The full name is `"layers.0.linear.weight"`. - -## Categories of renames - -### ✅ Can be eliminated (restructure model construction) - -#### 1. Sequential index patterns (nn.Sequential / nn.ModuleList) - -**HF pattern:** `nn.Sequential(SiLU(), Linear(...))` → weights at `mod.1.weight` - -**Problem:** Using a plain `_Linear(...)` produces `mod.weight` (no index). - -**Preferred solution — `nn.Sequential`:** - -`nn.Sequential` (from `onnxscript.nn`) registers children with numeric keys -like PyTorch's `nn.Sequential`, AND chains `forward()` calls automatically. -This gives both correct naming and clean call sites: - -```python -from mobius.components import Linear, SiLU - -# Produces "img_mod.1.weight" — matching HF -self.img_mod = nn.Sequential(SiLU(), Linear(dim, 6 * dim)) - -# Forward: output chains through each child automatically -result = self.img_mod(op, temb) -``` - -`nn.Sequential` subclasses `nn.ModuleList`. Key implementation detail: it -overrides `_set_name` to keep children with simple "0", "1" names (not -fully-qualified), because `__call__` already pushes the parent name onto the -scope stack. Without this override, children would be double-prefixed. - -**Fallback — `nn.ModuleList` with manual indexing:** - -If `nn.Sequential` is not yet available, use `nn.ModuleList` with explicit -`[i]` indexing: - -```python -self.img_mod = nn.ModuleList([SiLU(), Linear(dim, 6 * dim)]) - -# Forward: manual chaining -result = self.img_mod[1](op, self.img_mod[0](op, temb)) -``` - -This produces the same parameter names but requires manual forward logic. - -#### 2. Non-consecutive indices with placeholder modules - -**HF pattern:** `nn.Sequential(Linear, GELU, Linear)` → weights at `0.weight` -and `2.weight` (GELU at index 1 has no params). - -**Problem:** `nn.ModuleList([linear1, linear2])` produces indices 0, 1. - -**Solution:** Include activation modules to fill gaps: - -```python -class _NoOpModule(nn.Module): - """Placeholder for HF Dropout (no params, identity at inference).""" - def forward(self, op, x): - return x - -class _GELUGate(nn.Module): - """Matches HF GEGLU wrapper with .proj sub-attribute.""" - def __init__(self, in_features, out_features): - super().__init__() - self.proj = _Linear(in_features, out_features) - -# Matches HF: net.0.proj.weight, net.2.weight -self.net = nn.ModuleList([ - _GELUGate(dim, inner_dim * 2), # index 0 - _NoOpModule(), # index 1 (Dropout placeholder) - _Linear(inner_dim, dim), # index 2 -]) -``` - -#### 3. Wrapper modules for extra nesting - -**HF pattern:** `time_text_embed.timestep_embedder.linear_1.weight` - -**Problem:** Flat structure produces `linear_1.weight` (missing prefix). - -**Solution:** Create wrapper module matching HF nesting: - -```python -class _TimestepMLP(nn.Module): - def __init__(self, in_channels, time_embed_dim): - super().__init__() - self.linear_1 = _Linear(in_channels, time_embed_dim) - self.linear_2 = _Linear(time_embed_dim, time_embed_dim) - -class _TimestepEmbedding(nn.Module): - def __init__(self, in_channels, time_embed_dim): - super().__init__() - self.timestep_embedder = _TimestepMLP(in_channels, time_embed_dim) -``` - -#### 4. Bare Parameter → Module wrapper - -**HF pattern:** `txt_norm.weight` (from `RMSNorm` module) - -**Problem:** Using `nn.Parameter` produces `txt_norm` (no `.weight` suffix). - -**Solution:** Use a proper module: - -```python -# BAD — produces "txt_norm" as a bare parameter name -self.txt_norm = nn.Parameter((dim,)) - -# GOOD — produces "txt_norm.weight" -class _RMSNorm(nn.Module): - def __init__(self, dim, eps=1e-6): - super().__init__() - self.weight = nn.Parameter((dim,)) - self._eps = eps - def forward(self, op, x): - return op.RMSNormalization(x, self.weight, epsilon=self._eps) - -self.txt_norm = _RMSNorm(dim) -``` - -#### 5. Inner model wrapper for prefix nesting - -**HF pattern:** `model.layers.0.self_attn.q_proj.weight` - -**Problem:** Without a `model` wrapper, you get `layers.0.self_attn...`. - -**Solution:** Create inner model class: - -```python -class _TextModel(nn.Module): - def __init__(self, config): - super().__init__() - self.layers = nn.ModuleList([...]) - self.norm = _RMSNorm(config.hidden_size) - -class MyCausalLMModel(nn.Module): - def __init__(self, config): - super().__init__() - self.model = _TextModel(config) # Creates "model." prefix - self.lm_head = _Linear(config.hidden_size, config.vocab_size) -``` - -### ❌ Cannot be eliminated (must stay in preprocess_weights) - -#### 1. QKV splitting - -HuggingFace fuses Q, K, V into a single tensor (`query_key_value`, -`c_attn`, `qkv_proj`), but ONNX uses separate `q_proj`, `k_proj`, `v_proj`. - -```python -def preprocess_weights(self, state_dict): - new_state = {} - for key, tensor in state_dict.items(): - if "query_key_value" in key: - q, k, v = self._split_qkv(tensor, self.config) - new_state[key.replace("query_key_value", "q_proj")] = q - new_state[key.replace("query_key_value", "k_proj")] = k - new_state[key.replace("query_key_value", "v_proj")] = v - else: - new_state[key] = tensor - return new_state -``` - -**Models affected:** GPT-2, Falcon, InternLM2, ChatGLM, Phi3/Phi3Small - -#### 2. Conv1D → Linear transpose - -GPT-2 uses Conv1D `[in, out]` layout; ONNX Linear needs `[out, in]`. - -```python -if key.endswith(".weight") and tensor.ndim == 2: - tensor = tensor.t() -``` - -**Models affected:** GPT-2 - -#### 3. Deep structural naming differences - -BERT, T5, BART have deeply different naming conventions that would require -rewriting fundamental component classes to match. - -```python -# BERT: "encoder.layer.0.attention.self.query.weight" -# Ours: "encoder.layer.0.self_attn.q_proj.weight" -``` - -Changing this would require BERT-specific Attention, MLP components — not -worth the complexity for a simple rename. - -**Models affected:** BERT, DistilBERT, RoBERTa, ALBERT, T5, BART, mBART, -Marian, CLIP, SigLIP - -#### 4. MoE expert weight remapping - -MoE models have mixed naming across architectures (Mixtral: `w1/w2/w3`, -Qwen2-MoE: `gate_proj/up_proj/down_proj`). The current `_rename_moe_expert_weights` -handles both conventions optimally. - -#### 5. Weight tying - -Always needed when `tie_word_embeddings=True`: - -```python -if self.config.tie_word_embeddings: - if "lm_head.weight" in state_dict: - state_dict["model.embed_tokens.weight"] = state_dict["lm_head.weight"] - elif "model.embed_tokens.weight" in state_dict: - state_dict["lm_head.weight"] = state_dict["model.embed_tokens.weight"] -``` - -#### 6. Weight deletion - -Some HF weights are not needed (e.g. `rotary_emb.inv_freq` — RoPE -frequencies are computed at runtime). - -## ⚠️ Scope mechanism pitfalls - -Understanding how `onnxscript.nn` builds parameter names is critical. Names -come from the **call-stack of `__call__` invocations**, not from the Python -attribute chain used to *reach* a module. - -### Rule: only `__call__` pushes scope - -When you call `module(op, x)`, the base `nn.Module.__call__` method: -1. Pushes `module._name` onto the scope stack -2. Calls `module.forward(op, x)` -3. Pops the scope - -**Accessing a child and calling its method directly bypasses the parent scope:** - -```python -# ❌ BAD — self.shared.child(op, x) only pushes "child", NOT "shared" -result = self.shared.child(op, x) -# Parameter name: "child.weight" (missing "shared." prefix!) - -# ✅ GOOD — call shared's forward, which internally calls child -result = self.shared(op, x) -# Parameter name: "shared.child.weight" ✓ -``` - -### ModuleList indexing requires `__call__` - -The same rule applies to `nn.ModuleList`. `self.items[i]` returns the -module object — its `_name` is `"items.{i}"` — but that scope is only -pushed when you call it: - -```python -# ❌ BAD — sub_module's scope is pushed, but items.0 scope is NOT -result = self.items[0].sub_module(op, x) -# Parameter name: "sub_module.weight" (missing "items.0." prefix!) - -# ✅ GOOD — items[0].__call__ pushes "items.0", then forward calls sub_module -result = self.items[0](op, x) -# Parameter name: "items.0.sub_module.weight" ✓ -``` - -**Key takeaway:** If you need per-index distinct parameter names (e.g. -per-layer adapters), you **must** call the ModuleList element via -`self.items[idx](op, ...)`, not reach into its sub-attributes. - -### Shared weights via single module instance - -When multiple layers reuse the same weights (e.g. Zamba2's shared -transformer), register ONE module instance. Calling it multiple times -produces the same initializer names — ONNX uses a single initializer: - -```python -class _TextModel(nn.Module): - def __init__(self, config): - super().__init__() - # ONE shared module → ONE set of initializers - self.shared_transformer = SharedLayer(config) - - def forward(self, op, x): - for i in range(num_uses): - # Same scope "shared_transformer.*" each time → same initializer - x = self.shared_transformer(op, x) -``` - -### Circular dependency: shared weights + per-instance data - -When shared weights and per-instance data (e.g. adapters) must interact in -the same computation, the scope model creates a tension: - -- **Shared weights** must be inside a shared module (for correct scope) -- **Per-instance data** must be outside (different scope per use) - -**Solution: split the computation.** Have the shared module return an -intermediate value. The caller computes per-instance contributions at its -scope, then continues the computation with shared weights at its level: - -```python -class _TextModel(nn.Module): - def __init__(self, config): - super().__init__() - self.shared_attn = SharedAttention(config) # shared weights - self.adapters = nn.ModuleList([...]) # per-layer - self.gate_proj = Linear(...) # shared MLP at model scope - - def forward(self, op, hidden): - for idx in range(num_layers): - # Phase 1: shared attention (inside shared module scope) - intermediate = self.shared_attn(op, hidden) - # Phase 2: per-layer adapter (at "adapters.{idx}" scope) - adapter_out = self.adapters[idx](op, intermediate) - # Phase 3: shared MLP at model scope (gate_proj is model attr) - hidden = self.gate_proj(op, op.Add(intermediate, adapter_out)) -``` - -**Reference implementation:** `models/zamba2.py` — Zamba2 hybrid model with -shared transformer + per-layer Q/K/V/MLP low-rank adapters. - -## How to analyze a model's preprocess_weights - -1. **Compare HF names to ONNX names:** - ```python - # Print ONNX parameter names - module = MyModel(config) - for name, _ in module.named_parameters(): - print(name) - - # Print HF weight names (from safetensors) - from safetensors import safe_open - with safe_open("model.safetensors", framework="pt") as f: - for key in f.keys(): - print(key) - ``` - -2. **Categorize each rename** as one of the types above. - -3. **For eliminable renames**, restructure the module constructor. - -4. **For non-eliminable renames**, keep them in `preprocess_weights`. - -## Non-consecutive index patterns (setattr fallback) - -When HF uses `nn.Sequential` with non-consecutive parameter indices AND the -gap modules have no natural implementation: - -```python -# HF: Sequential(Conv2d, SiLU, Conv2d, SiLU, Conv2d, SiLU, Conv2d) -# Weights at indices: 0, 2, 4, 6 (SiLU at 1, 3, 5 has no params) -# But we process differently in forward, so ModuleList doesn't work - -from mobius.components import Conv2d - -# Fallback: manual setattr -class _SequentialConv2d(nn.Module): - def __init__(self, in_channels, out_channels, **kwargs): - super().__init__() - conv = Conv2d(in_channels, out_channels, **kwargs) - setattr(self, "1", conv) # Matches HF Sequential index -``` - -Use this only when `nn.ModuleList` with activation placeholders doesn't -work (e.g., different forward logic, or HF Sequential wraps padding + conv). - -## Reference implementations - -| Pattern | Model | File | -|---------|-------|------| -| No-op (fully aligned) | QwenImage transformer | `models/qwen_image.py` | -| Weight tying only | CausalLMModel (base) | `models/base.py` | -| Sequential index (ModuleList) | UNet, DiT, VAE | `models/unet.py`, `models/dit.py`, `models/vae.py` | -| Wrapper + placeholder modules | QwenImage (all patterns) | `models/qwen_image.py` | -| QKV splitting | Falcon, GPT-2 | `models/falcon.py`, `models/gpt2.py` | -| Conv1D transpose | GPT-2 | `models/gpt2.py` | -| MoE expert remapping | MoE models | `models/moe.py` | -| Deep structural renames | BERT, T5 | `models/bert.py`, `models/t5.py` | -| Shared weights + per-layer adapters | Zamba2 | `models/zamba2.py` | -| Scope-aware ModuleList adapters | Zamba2 | `models/zamba2.py` | diff --git a/.github/skills/writing-rewrite-rules/SKILL.md b/.github/skills/writing-rewrite-rules/SKILL.md deleted file mode 100644 index 8fa6cd6b..00000000 --- a/.github/skills/writing-rewrite-rules/SKILL.md +++ /dev/null @@ -1,301 +0,0 @@ ---- -name: writing-rewrite-rules -description: > - How to write ONNX rewrite rules using onnxscript.rewriter for the - mobius package. Covers the RewriteRuleClassBase API, pattern - matching, check/rewrite methods, file organization, and testing conventions. - Use this skill when creating rules that transform parts of an ONNX model - (e.g., replacing standard ops with custom/fused ops). ---- - -# Skill: Writing Rewrite Rules - -## When to use - -Use this skill when creating rules that transform parts of an ONNX model graph -— for example, replacing standard ops with custom or fused ops for better -runtime performance. Rewrite rules live in `src/mobius/rewrite_rules/` -and are applied **after** model export. - -## API overview - -Rewrite rules use `RewriteRuleClassBase` from `onnxscript.rewriter`: - -```python -from onnxscript.rewriter._basics import MatchResult -from onnxscript.rewriter._rewrite_rule import RewriteRuleClassBase, RewriteRuleSet -from onnxscript.rewriter import rewrite -``` - -1. **Subclass** `RewriteRuleClassBase` with `pattern()`, `check()`, and - `rewrite()` methods. -2. Call `.rule()` on an instance to create a `RewriteRule`. -3. Wrap one or more rules in a `RewriteRuleSet`. -4. Apply via `rewrite(model, pattern_rewrite_rules=rule_set)`. - -> **Important:** The keyword argument is `pattern_rewrite_rules`, **not** `rules`. - -```python -class MyRule(RewriteRuleClassBase): - def pattern(self, op, ...): ... - def check(self, context, ...): ... - def rewrite(self, op, ...): ... - -def my_rules() -> RewriteRuleSet: - return RewriteRuleSet([MyRule().rule()]) - -# Apply -rewrite(model, pattern_rewrite_rules=my_rules()) -``` - -## Pattern function - -`pattern(self, op, ...)` defines the ONNX subgraph to match. The positional -parameters after `op` become the matched inputs. - -```python -def pattern(self, op, q, k, v, attn_bias_2d): - q_4d = op.Unsqueeze(q, [0]) - k_4d = op.Unsqueeze(k, [0]) - v_4d = op.Unsqueeze(v, [0]) - attn_bias_4d = op.Unsqueeze(attn_bias_2d, [0, 1]) - attn_out = op.Attention( - q_4d, k_4d, v_4d, attn_bias_4d, - _allow_other_attributes=True, - _outputs=["attn_out"], - ) - return op.Squeeze(attn_out, [0]) -``` - -Key options: - -- **`_outputs=["name"]`** — Capture an intermediate value so it can be - referenced in `check()` and `rewrite()` by that keyword name. -- **`_allow_other_attributes=True`** — Match an op even if it has extra - attributes not listed in the pattern (e.g. `scale`, `num_heads`). - -## Check function - -`check(self, context, **kwargs)` validates structural requirements that the -pattern alone cannot express. Matched inputs and captured outputs arrive as -keyword arguments. - -Return `MatchResult()` (no arguments) for success. Call `.fail("reason")` to -reject the match: - -```python -def check(self, context, attn_bias_2d, attn_out, **_): - result = MatchResult() - - # Walk the producer chain to verify structure - where = attn_bias_2d.producer() - if where is None or where.op_type != "Where": - return result.fail("Expected Where producing attention bias") - - # Access attributes on matched nodes - attn = attn_out.producer() - if attn.attributes.get_float("scale", None) is None: - return result.fail("Missing scale attribute on Attention") - if attn.attributes.get_int("q_num_heads", None) is None: - return result.fail("Missing q_num_heads attribute on Attention") - - return result -``` - -Common attribute accessors: - -- `node.attributes.get_float("name")` / `node.attributes.get_float("name", default)` -- `node.attributes.get_int("name")` / `node.attributes.get_int("name", default)` - -## Rewrite function - -`rewrite(self, op, **kwargs)` builds the replacement subgraph. The `op` -parameter is an **IR tape builder** (not the same as `onnxscript`'s -`OpBuilder`). Matched inputs and captured outputs arrive as keyword arguments. - -```python -def rewrite(self, op, q, k, v, attn_bias_2d, attn_out, **_): - attn = attn_out.producer() - scale = attn.attributes.get_float("scale") - num_heads = attn.attributes.get_int("q_num_heads") - - cu_seqlens = self._trace_cu_seqlens(attn_bias_2d) - cu_seqlens_i32 = op.Cast(cu_seqlens, to=6) - - return op.op( - "PackedMultiHeadAttention", - inputs=[q, k, v, None, token_offset, cu_seqlens_i32], - domain="com.microsoft", - attributes={"scale": scale, "num_heads": num_heads}, - ) -``` - -### Critical: constant tensors in rewrite - -Raw Python lists **cannot** be used as inputs in the rewrite function. -Always create constants explicitly: - -```python -# GOOD — explicit Constant node -axes_0 = op.Constant(value_ints=[0]) -neg_one = op.Constant(value_ints=[-1]) -result = op.Squeeze(x, axes_0) - -# BAD — raw list (will fail) -result = op.Squeeze(x, [0]) -``` - -### Custom / domain-specific ops - -Use `op.op(...)` to emit single-output ops from non-default domains: - -```python -op.op( - "PackedMultiHeadAttention", - inputs=[q, k, v, None, token_offset, cu_seqlens], - domain="com.microsoft", - attributes={"scale": scale, "num_heads": num_heads}, -) -``` - -Pass `None` in the inputs list for optional inputs that should be left empty. - -### Multi-output custom ops - -Use `op.op_multi_out(...)` for ops with multiple outputs. -**`op.op()` returns a single `ir.Value`; `op.op_multi_out()` returns -`Sequence[ir.Value]`.** - -```python -outputs = op.op_multi_out( - "GroupQueryAttention", - inputs=[q, k, v, past_key, past_value, seqlens_k, total_seq_len], - domain="com.microsoft", - attributes={"num_heads": num_heads, "kv_num_heads": kv_num_heads}, - num_outputs=3, -) -attn_out, present_key, present_value = outputs[0], outputs[1], outputs[2] -``` - -### Matching patterns with shared intermediate values - -The rewriter will **not** match a pattern if an intermediate node's output -has consumers outside the matched subgraph. For example, matching -`Add → RMSNorm` will fail if the `Add` output is also used by a downstream -residual connection. - -**Workaround:** Match only the end node, then trace back in `check()`: - -```python -def pattern(self, op, add_out, weight): - # Only match RMSNorm — don't include Add in the pattern - return op.RMSNormalization(add_out, weight, _allow_other_attributes=True) - -def check(self, context, add_out, **_): - result = MatchResult() - producer = add_out.producer() - if producer is None or producer.op_type != "Add": - return result.fail("Input is not from Add") - if len(list(add_out.uses())) < 2: - return result.fail("Add has only 1 consumer") - return result - -def rewrite(self, op, add_out, weight, **_): - add_node = add_out.producer() - input_a, input_b = add_node.inputs[0], add_node.inputs[1] - # Create fused op and reroute the shared output - outputs = op.op_multi_out("FusedOp", inputs=[input_a, input_b, weight], ...) - add_out.replace_all_uses_with(outputs[1]) # reroute skip connection - return outputs[0] # return the primary output -``` - -## File organization - -``` -src/mobius/rewrite_rules/ -├── __init__.py # Public exports -├── _packed_attention.py # Rule implementation (private module) -├── _packed_attention_test.py # Unit tests (next to source) -├── _group_query_attention.py # Attention → GQA rule -├── _group_query_attention_test.py # GQA rule tests -├── _skip_norm.py # Add+RMSNorm → SkipNorm rule -└── _skip_norm_test.py # SkipNorm rule tests -``` - -### Conventions - -- Rule files are **private modules**: `_rule_name.py`. -- Unit tests go **next to the source file**: `_rule_name_test.py`. -- Export the public factory function from `__init__.py`: - -```python -# __init__.py -__all__ = ["packed_attention_rules"] -from mobius.rewrite_rules._packed_attention import packed_attention_rules -``` - -- Each rule module should provide a factory function (e.g. - `packed_attention_rules()`) that returns a `RewriteRuleSet`. - -## Testing - -Write unit tests that: - -1. Build a model containing the target pattern (either a tiny model from - the model library or a synthetic graph). -2. Count ops before applying the rule. -3. Apply the rule set via `rewrite(model, pattern_rewrite_rules=rules)`. -4. Count ops after and assert the expected replacements occurred. -5. Verify that non-matching subgraphs are **not** affected. - -```python -from collections import Counter -from onnxscript.rewriter import rewrite -from mobius.rewrite_rules import packed_attention_rules - - -def _count_ops(model) -> Counter: - return Counter(node.op_type for node in model.graph) - - -class TestPackedAttentionRules: - def test_rule_replaces_vision_attention(self): - model = build_model_with_pattern(...) - counts_before = _count_ops(model) - assert counts_before["Attention"] == 4 - - rewrite(model, pattern_rewrite_rules=packed_attention_rules()) - - counts_after = _count_ops(model) - assert counts_after["PackedMultiHeadAttention"] == 2 - assert counts_after["Attention"] == 2 # text decoder untouched - - def test_rule_preserves_non_matching_model(self): - """Models without the pattern are not affected.""" - model = build_text_only_model(...) - counts_before = _count_ops(model) - - rewrite(model, pattern_rewrite_rules=packed_attention_rules()) - - counts_after = _count_ops(model) - assert counts_after["Attention"] == counts_before["Attention"] - assert counts_after.get("PackedMultiHeadAttention", 0) == 0 - - def test_rules_returns_rule_set(self): - from onnxscript.rewriter._rewrite_rule import RewriteRuleSet - rules = packed_attention_rules() - assert isinstance(rules, RewriteRuleSet) -``` - -## Reference files - -- **Full rule implementations:** - - `src/mobius/rewrite_rules/_packed_attention.py` — Block-diagonal → PackedMHA - - `src/mobius/rewrite_rules/_group_query_attention.py` — Attention → GQA - - `src/mobius/rewrite_rules/_skip_norm.py` — Add+RMSNorm → SkipNorm -- **Test examples:** - - `src/mobius/rewrite_rules/_packed_attention_test.py` - - `src/mobius/rewrite_rules/_group_query_attention_test.py` - - `src/mobius/rewrite_rules/_skip_norm_test.py` -- **Exports:** - `src/mobius/rewrite_rules/__init__.py` diff --git a/.github/skills/writing-tests/SKILL.md b/.github/skills/writing-tests/SKILL.md deleted file mode 100644 index e45e92ec..00000000 --- a/.github/skills/writing-tests/SKILL.md +++ /dev/null @@ -1,699 +0,0 @@ ---- -name: writing-tests -description: > - Patterns and conventions for writing tests in mobius. Covers unit - tests (graph construction), integration tests (numerical parity with - HuggingFace), generation tests, testing utilities, and tolerance guidelines. - Use this skill when adding tests for new or modified models and components. ---- - -# Skill: Writing Tests - -## When to use - -Use this skill whenever you add a new model, component, or modify existing -behaviour. The project has a five-level confidence system (L1–L5) plus -shared configuration infrastructure. - -## Confidence levels (L1–L5) - -Each level is detected and counted **independently**. A model can pass L3 -(integration test) without passing L2 (no YAML test case defined), or have -L4 golden data without passing L3. Levels are not hierarchical in detection, -though logically higher levels usually imply lower ones. - -| Level | Name | What it verifies | Data source | -|-------|------|-----------------|-------------| -| **L1** | Graph builds | ONNX graph builds from a tiny synthetic config | `_MODEL_CONFIGS` / `_SPECIALIZED_TEST_MODEL_TYPES` in `tests/build_graph_test.py` | -| **L2** | Config compatible | Full-size HuggingFace config produces a valid graph | `test_model_id` field in YAML test case (`testdata/cases/`) | -| **L3** | Synthetic parity | Random-weight forward pass matches HuggingFace numerically | `tests/integration_test.py` parametrized tests | -| **L4** | Golden match | Real-weight prefill logits match pre-computed golden reference | `*.json` files in `testdata/golden/` | -| **L5** | Generation verified | Full multi-token generation matches golden output | `*_generation.json` files in `testdata/golden/` | - -### How counts work on the dashboard - -The dashboard shows **per-flag counts**: L1 count = models with `l1_graph_build=True`, -L2 count = models with `l2_arch_validation=True`, etc. A model is counted at -every level it passes — not just the highest one. Because all registered model -types have at least one graph build test, L1 equals the total number of -registered models. - -## Test architecture overview - -``` -tests/ -├── build_graph_test.py # L1: graph construction (no weights) -├── _test_configs.py # shared model configs for all tests -├── integration_test.py # L3: real-weight numerical parity -├── e2e_golden_test.py # L4 + L5: golden file comparison -├── yaml_schema_test.py # YAML test case schema validation -└── arch_validation_test.py # L2: full HF config graph build - -testdata/ -├── cases/ # YAML test case definitions (L2, L4, L5) -│ ├── causal-lm/ -│ ├── vision-language/ -│ ├── audio/ -│ └── ... -└── golden/ # Pre-computed reference outputs - ├── causal-lm/ - │ ├── gpt2.json # L4 prefill logits - │ └── gpt2_generation.json # L5 generation tokens - └── ... -``` - -### Running tests - -```bash -# All non-integration tests (fast, no downloads) -python -m pytest tests/build_graph_test.py tests/cli_test.py src/ -q \ - -k "not phi4mm and not apply_weights_unknown" --tb=short - -# Representative models only (~5 seconds) -python -m pytest tests/build_graph_test.py --fast - -# Single model type -python -m pytest tests/build_graph_test.py -k "phi4mm" - -# Integration tests (slow, downloads models) -python -m pytest tests/integration_test.py -m integration -k "qwen2.5-0.5b" - -# L4/L5 golden tests -python -m pytest tests/e2e_golden_test.py -m golden --level L4 -v -python -m pytest tests/e2e_golden_test.py -m golden --level L5 -v -``` - ---- - -## Shared test configuration (`tests/_test_configs.py`) - -All model configs for parametrized tests live in `tests/_test_configs.py`, -organized by category: - -| List | Test class | Task type | -|------|-----------|-----------| -| `CAUSAL_LM_CONFIGS` | `TestBuildGraph` | text-generation | -| `ENCODER_CONFIGS` | `TestBuildEncoderGraph` | feature-extraction | -| `SEQ2SEQ_CONFIGS` | `TestBuildSeq2SeqGraph` | seq2seq | -| `VISION_CONFIGS` | `TestBuildVisionGraph` | image-classification | -| `DETECTION_CONFIGS` | `TestBuildDetectionGraph` | object-detection | - -Each entry is a 3-tuple: `(model_type, config_overrides, is_representative)`. - -### The `is_representative` flag - -Set `True` for models with **unique behaviour** — custom model class, -softcapping, parallel attention, ALiBi, MoE routing, partial rotary, -non-standard activation, etc. Set `False` for models that are simple -aliases of a base class with no special config. - -Representative models are always tested. Non-representative models are -skipped when running with `--fast`. - -### The `--fast` flag - -`pytest --fast` skips non-representative parametrized tests, reducing -run time to ~5 seconds. Non-parametrized tests (VLM, Whisper, TTS, etc.) -always run regardless of this flag. - -The flag is implemented in `tests/conftest.py` via -`pytest_collection_modifyitems`. - -### Auto-generation from registry - -Model types registered with `text-generation` or `hybrid-text-generation` -tasks that have **no explicit entry** in `_test_configs.py` get an -auto-generated `(model_type, {}, False)` entry. This ensures new -registrations get basic graph-build coverage without editing test files. - -Auto-generation only covers text-generation models because other tasks -(vision-language, speech, diffusion) require specialised config overrides -that cannot be guessed. - -### Adding a config for a new model - -```python -# In tests/_test_configs.py, add to the appropriate list: -CAUSAL_LM_CONFIGS: list[tuple[str, dict, bool]] = [ - # ... - ("my_model", {"hidden_act": "gelu", "attn_qkv_bias": True}, True), -] -``` - -Then run: -```bash -python -m pytest tests/build_graph_test.py -k "my_model" -``` - ---- - -## L1: Graph build tests - -Located in `tests/build_graph_test.py`. Uses tiny synthetic -`ArchitectureConfig` objects (64 hidden, 2 layers, 256 vocab) to test -every model type without downloading weights or requiring network access. -Configs are defined in `tests/_test_configs.py` (see above). - -VLM, audio, and other specialised models have dedicated test methods -(not parametrized via `_test_configs.py`) and are tracked in -`_SPECIALIZED_TEST_MODEL_TYPES`. The dashboard L1 scanner covers both. - -### Rewrite rule unit tests - -Rewrite rule unit tests should be placed **next to** the source file they -test, not in the `tests/` directory. For example: - -- Source: `src/mobius/rewrite_rules/_packed_attention.py` -- Test: `src/mobius/rewrite_rules/_packed_attention_test.py` - -This keeps rewrite rule tests co-located with their implementation since -they test internal graph transformation patterns rather than public API -behavior. - -### Pattern: adding a new model type (L1) - -Add an entry to the appropriate list in `tests/_test_configs.py` with the -model type, config overrides, and `is_representative` flag: - -```python -# In tests/_test_configs.py: -CAUSAL_LM_CONFIGS: list[tuple[str, dict, bool]] = [ - ("llama", {}, True), - ("my_model", {"attn_qkv_bias": True, "hidden_act": "gelu"}, True), - # ... -] -``` - -The test framework automatically creates a tiny config, builds the graph, -and checks: -- Graph has inputs (`input_ids`, `attention_mask`, `position_ids`) -- Graph has outputs (`logits`, `present.{i}.key`, `present.{i}.value`) -- Graph has initializers (embedding, attention, MLP/expert parameters) - -### Pattern: model-specific structure tests - -For models with unique structure (e.g. LoRA), add a dedicated test class: - -```python -class TestBuildGraphLoRA: - def test_lora_initializers_present(self): - config = _base_config( - vision_lora={"r": 4, "lora_alpha": 8}, - speech_lora={"r": 8, "lora_alpha": 16}, - ) - model_cls = registry.get("phi4mm") - module = model_cls(config) - task = CausalLMTask() - model = task.build_graph(module, config, opset_version=23) - - init_names = list(model.graph.initializers) - lora_names = [n for n in init_names if "lora" in n] - assert len(lora_names) > 0 -``` - ---- - -## L1 structure tests (weight alignment) - -Located in `tests/weight_alignment_test.py`. Verifies that -`preprocess_weights()` maps HuggingFace state dict keys to ONNX initializer -names correctly — no dropped or mangled weight names. - -For each model type, the test: -1. Builds the ONNX graph with a tiny config -2. Collects all ONNX initializer names -3. Creates a synthetic state dict matching those names -4. Runs `preprocess_weights()` on the dict -5. Asserts all initializers are still covered - -This catches bugs like Falcon's `h.` prefix replacement corrupting names, -MoE fused weight names being dropped, or OPT bias names being mangled. - ---- - -## L2: Config compatibility - -L2 is detected from the `test_model_id` field in the YAML test case -(see YAML format below). If a model has a `test_model_id`, the dashboard -counts it as L2. A model without a YAML case (or without `test_model_id`) -is not counted at L2. - -To add L2 coverage: create a YAML test case with `test_model_id` set to a -real HuggingFace model ID (see YAML format section below). - ---- - -## L3: Integration tests (synthetic parity) - -Located in `tests/integration_test.py`. Require real HuggingFace checkpoints. -Each model is parametrized with `(model_id, trust_remote_code)`. - -### Adding a new model to integration tests - -Add a `pytest.param` to `_TEXT_MODELS`: - -```python -_TEXT_MODELS = [ - pytest.param("Qwen/Qwen2.5-0.5B", False, id="qwen2.5-0.5b"), - pytest.param("my-org/my-small-model", False, id="my-model"), - # ... -] -``` - -Guidelines for choosing models: -- Prefer models ≤ 1B parameters for CI speed -- Models must be publicly accessible (no gated/private repos) -- One representative model per distinct model class - -### Pattern: prefill + decode numerical comparison - -```python -@pytest.mark.integration -@pytest.mark.parametrize("model_id,trust_remote_code", _TEXT_MODELS) -class TestForwardNumerical: - def test_prefill_logits_match(self, model_id, trust_remote_code): - onnx_model = build(model_id, load_weights=True) - torch_model, tokenizer = load_torch_model(model_id) - config = _get_config(model_id, trust_remote_code) - - # Tokenize, run both models, compare - feeds = _make_prefill_feeds(config, input_ids, attention_mask, position_ids) - onnx_outputs = session.run(feeds) - assert_logits_close(onnx_outputs["logits"], torch_logits, rtol=1e-3, atol=1e-3) - - def test_decode_step_logits_match(self, model_id, trust_remote_code): - # Prefill first, then feed next token + KV cache - decode_feeds = _make_decode_feeds(config, ...) - onnx_out_2 = session.run(decode_feeds) - assert_logits_close(onnx_out_2["logits"], torch_logits_2, rtol=1e-3, atol=1e-3) -``` - -### Pattern: greedy generation - -```python -@pytest.mark.integration -class TestGreedyGeneration: - def test_generate_tokens_match(self, model_id, trust_remote_code): - session = OnnxModelSession(onnx_model) - generator = OnnxGenerator(session, config) - onnx_ids = generator.generate(input_ids, max_new_tokens=10, eos_token_id=...) - - torch_ids = torch_generate_greedy(torch_model, input_ids, max_new_tokens=10, eos_token_id=...) - assert_generation_match(onnx_ids[0].tolist(), torch_ids[0].tolist()) -``` - ---- - -## L4 + L5: Golden tests - -L4 and L5 tests compare ONNX model outputs against HuggingFace reference -outputs using pre-computed "golden" files stored in `testdata/golden/`. - -| Level | Golden file | Contents | -|-------|-------------|----------| -| L4 | `testdata/golden//.json` | Prefill top-1/top-2 token IDs + logit summary | -| L5 | `testdata/golden//_generation.json` | Prompt + generated token IDs + generated text | - -### YAML test case format - -**Location:** `testdata/cases//.yaml` - -Categories match task types: `causal-lm`, `encoder`, `seq2seq`, `audio`, -`vision`, `vision-language`, `diffusion`. - -**Required fields:** - -```yaml -model_id: "Qwen/Qwen2.5-1.5B-Instruct" # HuggingFace model ID -revision: "main" # Git revision / commit SHA -task_type: "text-generation" # Task type string -dtype: "float32" # "float32", "float16", or "bfloat16" -level: "L4+L5" # "L4", "L5", or "L4+L5" - -inputs: - prompts: - - "Here is my poem:" # Text prompt(s); use this default -``` - -For image models, use `images:` instead of (or alongside) `prompts:`: - -```yaml -inputs: - images: - - "pipeline-cat-chonk.jpeg" # Path relative to testdata/ -``` - -For audio models: - -```yaml -inputs: - audio: - - "652-129742-0006.flac" -``` - -**Optional fields:** - -```yaml -# Identifier for the test model used in L2 config compatibility check. -# If set, the dashboard counts this model as L2 (full HF config valid). -test_model_id: "Qwen/Qwen2.5-1.5B-Instruct" - -# Skip this test case entirely (model too large, gated repo, etc.). -# Dashboard shows the model as 'skipped' rather than counting it toward coverage. -skip_reason: "Model too large (47B MoE) for CPU golden generation." - -# Pass trust_remote_code=True when loading HuggingFace model (default: false). -trust_remote_code: true - -# Minimum fraction of generated tokens that must match the golden reference. -# Use for VL/audio pipelines where floating-point variance causes later tokens -# to diverge. A value of 0.25 means at least 25% of tokens must match exactly. -# Green (≥0.9) / Yellow (0.5–0.9) / Red (<0.5) on dashboard. -min_token_match_ratio: 0.25 - -# Human-readable notes about this model. -notes: "GPT-2 124M. Absolute positional embeddings, no RoPE." - -generation: - max_new_tokens: 20 # Override token generation limit - do_sample: false -``` - -**`skip_reason` vs `_SKIP_REASONS` dict:** Always use the YAML `skip_reason` -field for new cases. The legacy `_SKIP_REASONS` dict in `e2e_golden_test.py` -has been removed — YAML is the canonical location. - -### Golden file format - -**L4 golden file** (`testdata/golden//.json`): -Generated automatically by `generate_golden.py`. Contains `top1_id`, -`top2_id`, `top10_ids`, `top10_logits`, and `logits_summary` from the last -token position of the prefill pass. - -**L5 generation file** (`testdata/golden//_generation.json`): -Contains `model_id`, `prompt`, `generated_tokens` (list of token IDs), and -`generated_text`. This is the authoritative source for L5 tests — the main -golden JSON does **not** contain generation data. - -### Generating golden data - -```bash -# Generate for all test cases at a given level -python scripts/generate_golden.py --level L4 - -# Generate for a specific task type -python scripts/generate_golden.py --level L4 --task-type causal-lm - -# Generate for a specific model (glob filter on model name) -python scripts/generate_golden.py --level L4 --filter 'llama*' -``` - -Golden files must be committed alongside new test case YAML files. - -### Running L4/L5 tests - -```bash -# L4 tests (single forward pass parity) -python -m pytest tests/e2e_golden_test.py -m golden --level L4 -v - -# L5 tests (multi-token generation parity) -python -m pytest tests/e2e_golden_test.py -m golden --level L5 -v -``` - -### Adding coverage for a new model (step by step) - -**L1 — Graph builds:** -1. Add `("my_model", {config_overrides}, True)` to the appropriate list in - `tests/_test_configs.py` (or add a dedicated method if the model is a VLM/audio). -2. Run `python -m pytest tests/build_graph_test.py -k "my_model"`. - -**L2 — Config compatible:** -1. Create `testdata/cases//my-model.yaml`. -2. Set `test_model_id: "org/my-model-id"`. -3. Run schema validation: `python -m pytest tests/yaml_schema_test.py`. - -**L3 — Synthetic parity:** -1. Add `pytest.param("org/my-model", False, id="my-model")` to the - appropriate parametrized list in `tests/integration_test.py`. -2. Run `python -m pytest tests/integration_test.py -m integration -k "my-model"`. - -**L4 — Golden match:** -1. Create/update `testdata/cases//my-model.yaml` with `level: "L4"`. -2. Set `inputs.prompts: ["Here is my poem:"]` (standard default prompt). -3. Run `python scripts/generate_golden.py --level L4 --filter 'my-model*'`. -4. Commit the generated `testdata/golden//my-model.json`. -5. Run `python -m pytest tests/e2e_golden_test.py -m golden --level L4 -k "my-model"`. - -**L5 — Generation verified:** -1. Update YAML to `level: "L5"` or `"L4+L5"`. -2. Add a `generation:` block with `max_new_tokens` and `do_sample: false`. -3. Optionally set `min_token_match_ratio` if you expect partial divergence - (VL pipelines, long generation sequences). -4. Run `python scripts/generate_golden.py --level L5 --filter 'my-model*'`. -5. Commit `testdata/golden//my-model_generation.json`. -6. Run `python -m pytest tests/e2e_golden_test.py -m golden --level L5 -k "my-model"`. - -### Dashboard coverage - -The dashboard shows L4/L5 coverage per model. A model shows as 'skipped' -(not counted toward coverage) when its YAML test case has a `skip_reason` -field. Models without a YAML case show no L4/L5 coverage. - ---- - -## Tolerances - -| Test type | Recommended rtol/atol | -|-----------|----------------------| -| Standard text models | `1e-3` / `1e-3` | -| Encoder-only (BERT) | `1e-3` / `1e-3` | -| Encoder-decoder (Whisper, BART, T5) | `1e-3` / `1e-3` | -| Multimodal models | `1e-2` / `1e-2` | -| Diffusion models (UNet, DiT, VAE) | `1e-3` / `1e-3` | -| Audio encoder models | `1e-3` / `1e-3` | -| Generation (token IDs) | Exact match | - -Multimodal models use looser tolerances because the vision pipeline -introduces additional floating-point variance. - -`assert_logits_close` uses `strict=True` in `np.testing.assert_allclose`, -which also checks shape and dtype match. If tolerances fail, verify: - -1. **Norm epsilon** — LayerNorm/RMSNorm eps must match HF config exactly - (e.g., Whisper uses `1e-5`, not the default `1e-6`) -2. **Norm type** — Check if the model uses RMSNorm or LayerNorm. OLMo-1B - uses weight-free LayerNorm (not RMSNorm). Using the wrong type causes - max abs diff > 1.0. -3. **Q scaling order** — some models (Whisper) pre-scale Q before attention - and pass `scale=1.0` to the op, which is numerically different from - passing `scale=head_dim**-0.5` -4. **Attention scale** — some models (Granite) replace `1/sqrt(head_dim)` with - a custom `attention_multiplier` from the config -5. **Scaling multipliers** — check HF config for `embedding_multiplier`, - `logits_scaling`, `residual_multiplier` that aren't in standard Llama -6. **Residual pattern** — verify `residual + output * scale` vs - `residual * scale + output` by reading HF source -7. **Weight loading** — compare ONNX initializers against HF state_dict to - rule out name mapping bugs -8. **Float64 contamination** — numpy arrays created from config values default - to float64; always use `dtype=np.float32` - -### Debugging large logit differences - -When max abs diff is large (> 0.5), run this diagnostic: - -```python -import numpy as np -diff = np.abs(onnx_logits[0, -1] - hf_logits) -print(f"Max abs diff: {diff.max():.4f}") -print(f"Mean abs diff: {diff.mean():.4f}") -# If max > 0.5, it's likely a norm or scaling bug, not just floating-point -# If max > 10, weights are probably loaded to wrong parameters -``` - -Check the HF norm class directly: -```python -import inspect -from transformers.models.olmo.modeling_olmo import OlmoLayerNorm -print(inspect.getsource(OlmoLayerNorm)) -``` - -Check for unextracted config fields: -```python -config = AutoConfig.from_pretrained("model-id") -for k, v in config.to_dict().items(): - if any(s in k for s in ("multiplier", "scaling", "factor", "epsilon")): - print(f"{k}: {v}") -``` - -## Testing utilities reference - -| Utility | Import path | Purpose | -|---------|-------------|---------| -| `OnnxModelSession(model)` | `_testing.ort_inference` | Save + load + run ONNX model | -| `OnnxGenerator(session, config)` | `_testing.generation` | Multi-step greedy decoding | -| `load_torch_model(id)` | `_testing.torch_reference` | Load HF model + tokenizer | -| `torch_forward(model, ...)` | `_testing.torch_reference` | Single forward pass | -| `torch_generate_greedy(...)` | `_testing.generation` | Multi-token generation | -| `assert_logits_close(a, b)` | `_testing.comparison` | Logit comparison with diagnostics | -| `assert_generation_match(a, b)` | `_testing.comparison` | Token-ID exact match | - -## Debugging multi-model pipelines (TTS, VLM) - -When a multi-model pipeline produces wrong output but individual -model prefill logits look correct, isolate each model boundary: - -1. **Compare each model's output against HF at the boundary** — e.g. - `last_hidden_state` from the talker, `codec_sum` from embeddings, - `inputs_embeds` constructed for the code predictor. - -2. **Check pre-norm vs post-norm** — `outputs.last_hidden_state` in HF - is typically post-norm. If your ONNX model returns pre-norm hidden - states, downstream models receive wrong values. - -3. **Verify external construction matches HF** — for models where the - generation loop constructs inputs externally (e.g. concatenating - hidden states with embeddings), write a comparison script that - checks the constructed input matches HF token-by-token: - ```python - # Compare inputs_embeds at each generation step - for step in range(num_steps): - onnx_input = construct_inputs_embeds(step, ...) - hf_input = hf_model.get_inputs_embeds(step, ...) - diff = np.abs(onnx_input - hf_input).max() - print(f"Step {step}: max diff = {diff:.6f}") - ``` - -4. **Embedding weight vs lookup mismatch** — if embedding weights are - identical but lookups differ, the issue is usually which code index - or embedding table is being used (off-by-one errors). - -## QA pitfalls: lessons from Qwen3.5 / hybrid-attention models - -These lessons come from debugging a hybrid DeltaNet + full-attention model -(Qwen3.5). They apply to any model that uses ONNX custom functions, Scan -ops, or non-standard dtypes. - -### Build graph tests are necessary but not sufficient - -Build graph tests (L1) verify graph construction — I/O shapes, -initializer existence, no obvious op errors. They do **not** execute the -graph with real data. A Scan body MatMul shape mismatch that would crash at -runtime can still pass all L1 tests. **Always write an integration test -alongside any new custom function or Scan op.** - -### Integration tests must exercise all code paths - -- **Text-only first** — verify generation and logit parity before adding - other modalities -- **Vision with real pixel values** — passing empty features (zeros) doesn't - exercise the vision encoder; use `processor(images=image)` outputs -- **All dtypes**: f32, f16, bf16. Each can expose different bugs: - - f16/bf16 have different overflow points and precision characteristics - - Kernel dispatch is dtype-specific (some kernels only exist for f16) -- **GPU when available** — different kernels activate on CUDA; a model - correct on CPU may diverge on GPU due to different reduction order - -### Test feed creation: symbolic dimensions need real values - -ONNX models export symbolic batch/sequence dimensions. When feeding the -model for ORT inference: - -- **Recurrent state batch dim must match input batch dim** — unlike KV - cache (which initialises to zeros and grows), recurrent state tensors - have a fixed `(B, ...)` shape. Feeding batch=0 produces a zero-sized - carry state that collapses the Scan output. -- **Scan carry state is not KV cache** — do not copy the KV cache - zero-initialisation pattern for recurrent state; the batch dimension - must be the actual inference batch size. - -```python -# WRONG — batch=0 zeros out Scan carry -past_state = np.zeros((0, num_heads, d_k, d_v), dtype=np.float32) - -# CORRECT — must match actual batch size -batch_size = input_ids.shape[0] -past_state = np.zeros((batch_size, num_heads, d_k, d_v), dtype=np.float32) -``` - -### ONNX function registration - -When renaming a custom function's `op_type` (e.g. `CausalConvNdWithState` -→ `CausalConvWithState`), the function must be re-registered under the new -name in ORT's function decomposition list. ORT needs the function embedded -in `model.functions` to decompose the custom op before execution. - -**Checklist when renaming a custom function:** -1. Rename the Python factory function and the `ir.Function.name` -2. Update all call sites that reference the old op_type string -3. Update any integration tests that check the op_type name -4. Verify the function appears in `onnx_model.functions` after build - -### Dtype-specific bugs to watch for - -**fp16 Exp overflow:** `exp(x)` overflows to `inf` for `x > ~11.09` in -fp16. The Softplus activation (`log(1 + exp(x))`) and decay computation -`exp(-softplus(x))` are common overflow sites. Always upcast to float32 -for Exp/Softplus in fp16 models. bf16 has the same 8-bit exponent range as fp32 and does NOT need the upcast: -```python -x_f32 = op.Cast(x, to=ir.DataType.FLOAT) -result = op.Exp(x_f32) -result = op.Cast(result, to=x.dtype) # cast back -``` - -**bf16 vs fp16:** bf16 has the same exponent range as fp32 (no overflow -at 11.09) but much less precision (7-bit mantissa vs 10-bit). If a -computation works in bf16 but not fp16, check for Exp overflow first. - -### Examples as QA tools - -The `--compare-hf` flag in example scripts is the gold-standard correctness -check for a model. Run it as part of every significant change: - -```bash -# Primary correctness check -python examples/qwen35_text_generation.py --compare-hf - -# Test all supported dtypes -python examples/qwen35_text_generation.py --compare-hf --dtype f16 -python examples/qwen35_text_generation.py --compare-hf --dtype bf16 - -# Test on GPU (if available) -python examples/qwen35_text_generation.py --compare-hf --device cuda -``` - -Target: **100% token match** in fp32 greedy generation. fp16/bf16 may -diverge after the first few tokens due to floating-point accumulation, which -is acceptable if logit parity holds at `atol=rtol=1e-2`. - -### Parity testing methodology - -Compare full logit tensors, not just generated tokens. Generated tokens -hide logit divergence (two very-different logit vectors can agree on the -top-1 token): - -```python -# Always compare full logits at every position -assert_logits_close(onnx_logits, hf_logits, atol=1e-3, rtol=1e-3) # fp32 -assert_logits_close(onnx_logits, hf_logits, atol=1e-2, rtol=1e-2) # fp16/bf16 - -# Also check last-position argmax matches (quick sanity check) -assert onnx_logits[0, -1].argmax() == hf_logits[0, -1].argmax() -``` - -If argmax matches but full logit tolerance fails, the model is numerically -correct but some intermediate accumulation differs — this is usually -acceptable for fp16/bf16 and worth a brief comment in the test. - -### Automated code review catches real bugs - -Enable Copilot/automated review on every PR that modifies model or component -code. During Qwen3.5 development, automated review found: - -- The fp16 Exp overflow risk in decay computation (a real runtime bug on fp16 - hardware that would not be caught by fp32 tests) -- Missing input validation (`kernel_size < 1`, `channels <= 0`) that would - cause ZeroDivisionError or silent wrong output - -These were not caught by the 514 build graph tests. Code review + integration -tests together cover the gaps that unit tests cannot. diff --git a/CHANGELOG.md b/CHANGELOG.md index bebd4c57..4fa211dd 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -458,6 +458,6 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 SkipNorm fusion. - 10 examples covering text generation, multimodal, ASR, TTS, and ORT-GenAI integration. -- 11 contribution skills (`.github/skills/`) for AI-agent-assisted +- 11 contribution skills (`.agents/skills/`) for AI-agent-assisted development: adding models, writing tests, debugging VL pipelines, rewrite rules, and more. diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index 094ad9b1..a77bca2c 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -173,7 +173,7 @@ components in `components/_multimodal.py`. Before a new model PR is considered **done**, all items in the quality checklist must be satisfied (or explicitly waived with a written reason in the PR description). The full checklist with explanations lives in the -[quality-checklist skill](.github/skills/quality-checklist/SKILL.md). +[quality-checklist skill](.agents/skills/quality-checklist/SKILL.md). ### Summary checklist diff --git a/README.md b/README.md index c34c5ae1..d5b27ad9 100644 --- a/README.md +++ b/README.md @@ -170,7 +170,7 @@ lintrunner f --all-files ### Adding a new model See the [AI-assisted model support strategy](https://onnxruntime.github.io/mobius/ai-model-support-strategy.html) -and the developer skills in `.github/skills/`: +and the developer skills in `.agents/skills/`: | Skill | Use when | |-------|----------| diff --git a/docs/ai-model-support-strategy.md b/docs/ai-model-support-strategy.md index 5a96c974..d8f6e974 100644 --- a/docs/ai-model-support-strategy.md +++ b/docs/ai-model-support-strategy.md @@ -211,7 +211,7 @@ Add the field to the `ArchitectureConfig` dataclass. #### A.3. Variant (subclass) Create a new model file. The agent should follow the **adding-a-new-model** -skill (`.github/skills/adding-a-new-model/SKILL.md`) which provides: +skill (`.agents/skills/adding-a-new-model/SKILL.md`) which provides: 1. A complete model file template 2. Common variation patterns with solutions @@ -362,7 +362,7 @@ asked to add a new model. ### Inputs - **Model identifier**: HuggingFace model ID (e.g. `meta-llama/Llama-4-Scout-17B-16E`) -- **Skills**: The agent should read the relevant skill files from `.github/skills/` +- **Skills**: The agent should read the relevant skill files from `.agents/skills/` ### Workflow diff --git a/docs/getting-started.md b/docs/getting-started.md index 94679dfa..605b255a 100644 --- a/docs/getting-started.md +++ b/docs/getting-started.md @@ -229,7 +229,7 @@ mobius info meta-llama/Llama-3.2-1B ### Adding a new model -See the [adding-a-new-model skill](../.github/skills/adding-a-new-model/SKILL.md) +See the [adding-a-new-model skill](../.agents/skills/adding-a-new-model/SKILL.md) for copy-paste examples and step-by-step instructions. ## Output Format diff --git a/tests/model_coverage_test.py b/tests/model_coverage_test.py index a389baac..f479fd4a 100644 --- a/tests/model_coverage_test.py +++ b/tests/model_coverage_test.py @@ -29,7 +29,7 @@ 3. Add YAML in ``testdata/cases/`` → L4 4. Generate golden JSON (or add ``skip_reason`` to YAML) → L5 -See ``.github/skills/writing-tests/SKILL.md`` for the full guide. +See ``.agents/skills/writing-tests/SKILL.md`` for the full guide. """ from __future__ import annotations From 440c35b0974c1c3b84155dbadd4d3a481b9f41ab Mon Sep 17 00:00:00 2001 From: Justin Chu Date: Tue, 21 Apr 2026 16:34:55 +0000 Subject: [PATCH 7/8] Address review: use strict=True in assert_logits_close, fix docs MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Update assert_logits_close to use np.testing.assert_allclose with strict=True, which enforces shape and dtype equality - Remove manual shape assertion (strict=True handles it) - Remove protobuf-based extraction method from debugging skill (violates repo zero-protobuf convention) - Fix CastLike usage in tolerance guidelines (not Cast with x.dtype) - Fix skill count in CHANGELOG (11 → 12) - Update skill docs to accurately describe assert_logits_close behavior Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Justin Chu --- .../references/extraction-methods.md | 36 ++----------------- .agents/skills/writing-tests/SKILL.md | 2 +- .../references/test-utilities.md | 2 +- .../references/tolerance-guidelines.md | 6 ++-- CHANGELOG.md | 2 +- src/mobius/_testing/comparison.py | 4 +-- 6 files changed, 10 insertions(+), 42 deletions(-) diff --git a/.agents/skills/debugging-multimodal/references/extraction-methods.md b/.agents/skills/debugging-multimodal/references/extraction-methods.md index 26dbe6fa..e1f8f315 100644 --- a/.agents/skills/debugging-multimodal/references/extraction-methods.md +++ b/.agents/skills/debugging-multimodal/references/extraction-methods.md @@ -1,42 +1,12 @@ # Extracting Intermediate ONNX Values -Three methods for extracting intermediate values from ONNX models to +Two methods for extracting intermediate values from ONNX models to compare block-by-block against HuggingFace. Used when a pipeline stage (e.g., vision encoder) diverges and you need to narrow down the root cause. --- -## Method 1: Add intermediate outputs to the ONNX graph - -The most reliable approach — expose any internal node's output as a graph -output so ORT returns it alongside normal outputs. - -```python -import onnx - -model = onnx.load("vision.onnx") -graph = model.graph - -# Find the node whose output you want to inspect -for node in graph.node: - if node.op_type == "RMSNormalization" and "block_0" in node.output[0]: - # Name the output (if unnamed, give it a name) - target_output = node.output[0] - break - -# Add as a graph output -graph.output.append( - onnx.helper.make_tensor_value_info(target_output, onnx.TensorProto.FLOAT, None) -) -onnx.save(model, "vision_debug.onnx") - -# Now ORT will return this value alongside image_features -session = ort.InferenceSession("vision_debug.onnx") -results = session.run(None, feeds) -# results[-1] is the intermediate value -``` - -## Method 2: Use `ir.Model` graph manipulation (preferred for mobius) +## Method 1: Add intermediate outputs to the `ir.Model` graph When working with `ir.Model` objects from the build pipeline, manipulate the graph directly without saving/loading: @@ -65,7 +35,7 @@ block_0_out = out["block_0_output"] session.close() ``` -## Method 3: Hook HuggingFace model for reference values +## Method 2: Hook HuggingFace model for reference values Use PyTorch hooks to extract intermediate values from HuggingFace at the same points: diff --git a/.agents/skills/writing-tests/SKILL.md b/.agents/skills/writing-tests/SKILL.md index 742bbe2c..79da765c 100644 --- a/.agents/skills/writing-tests/SKILL.md +++ b/.agents/skills/writing-tests/SKILL.md @@ -216,7 +216,7 @@ Compare ONNX outputs against pre-computed golden files in `testdata/golden/`. | fp16/bf16 logits | `1e-2` / `1e-2` | Key rules: -- `assert_logits_close` checks shape + dtype match (`strict=True`) +- `assert_logits_close` checks shape + dtype match via `np.testing.assert_allclose(..., strict=True)` - If max abs diff > 0.5 → likely a norm or scaling bug - If max abs diff > 10 → weights loaded to wrong parameters diff --git a/.agents/skills/writing-tests/references/test-utilities.md b/.agents/skills/writing-tests/references/test-utilities.md index 9dd18db4..dfa2804e 100644 --- a/.agents/skills/writing-tests/references/test-utilities.md +++ b/.agents/skills/writing-tests/references/test-utilities.md @@ -43,7 +43,7 @@ token_ids = generator.generate(input_ids, max_new_tokens=10, eos_token_id=...) ### `assert_logits_close(actual, expected, rtol, atol)` Uses `np.testing.assert_allclose` with `strict=True` (checks shape + dtype). -On failure, prints diagnostic info including max/mean abs diff. +On failure, numpy prints diagnostic info including max abs diff. ### `assert_generation_match(actual_ids, expected_ids)` diff --git a/.agents/skills/writing-tests/references/tolerance-guidelines.md b/.agents/skills/writing-tests/references/tolerance-guidelines.md index 336508ae..af0c7854 100644 --- a/.agents/skills/writing-tests/references/tolerance-guidelines.md +++ b/.agents/skills/writing-tests/references/tolerance-guidelines.md @@ -21,8 +21,8 @@ introduces additional floating-point variance. ## `assert_logits_close` behavior -`assert_logits_close` uses `strict=True` in `np.testing.assert_allclose`, -which also checks shape and dtype match. +`assert_logits_close` uses `np.testing.assert_allclose(..., strict=True)`, +which checks shape, dtype, and value equality within tolerance. ## Tolerance failure checklist @@ -96,7 +96,7 @@ for Exp/Softplus in fp16 models: ```python x_f32 = op.Cast(x, to=ir.DataType.FLOAT) result = op.Exp(x_f32) -result = op.Cast(result, to=x.dtype) # cast back +result = op.CastLike(result, x) # cast back to x's dtype ``` ### bf16 diff --git a/CHANGELOG.md b/CHANGELOG.md index 4fa211dd..023ab65f 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -458,6 +458,6 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 SkipNorm fusion. - 10 examples covering text generation, multimodal, ASR, TTS, and ORT-GenAI integration. -- 11 contribution skills (`.agents/skills/`) for AI-agent-assisted +- 12 contribution skills (`.agents/skills/`) for AI-agent-assisted development: adding models, writing tests, debugging VL pipelines, rewrite rules, and more. diff --git a/src/mobius/_testing/comparison.py b/src/mobius/_testing/comparison.py index d1a0d6dd..df83cde2 100644 --- a/src/mobius/_testing/comparison.py +++ b/src/mobius/_testing/comparison.py @@ -25,14 +25,12 @@ def assert_logits_close( Raises: AssertionError: If shapes differ or values are not close. """ - assert actual.shape == expected.shape, ( - f"Shape mismatch: ONNX {actual.shape} vs reference {expected.shape}" - ) np.testing.assert_allclose( actual, expected, rtol=rtol, atol=atol, + strict=True, err_msg="Logits differ between ONNX and reference model", ) From cbf7d330d8d492d3dd3981cccfd237e92f8c8136 Mon Sep 17 00:00:00 2001 From: Justin Chu Date: Tue, 21 Apr 2026 16:36:10 +0000 Subject: [PATCH 8/8] Remove hardcoded skill count from CHANGELOG Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Justin Chu --- CHANGELOG.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 023ab65f..bb545faa 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -458,6 +458,6 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 SkipNorm fusion. - 10 examples covering text generation, multimodal, ASR, TTS, and ORT-GenAI integration. -- 12 contribution skills (`.agents/skills/`) for AI-agent-assisted +- Contribution skills (`.agents/skills/`) for AI-agent-assisted development: adding models, writing tests, debugging VL pipelines, rewrite rules, and more.