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-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/debugging-multimodal/references/debugging-cookbook.md b/.agents/skills/debugging-multimodal/references/debugging-cookbook.md new file mode 100644 index 00000000..33e0ab19 --- /dev/null +++ b/.agents/skills/debugging-multimodal/references/debugging-cookbook.md @@ -0,0 +1,183 @@ +# 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 + +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 + `references/extraction-methods.md`) 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) +``` + +## 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/debugging-multimodal/references/extraction-methods.md b/.agents/skills/debugging-multimodal/references/extraction-methods.md new file mode 100644 index 00000000..e1f8f315 --- /dev/null +++ b/.agents/skills/debugging-multimodal/references/extraction-methods.md @@ -0,0 +1,114 @@ +# Extracting Intermediate ONNX Values + +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 `ir.Model` graph + +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 2: 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-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/.github/skills/diffusion-models/SKILL.md b/.agents/skills/diffusion-models/SKILL.md similarity index 96% rename from .github/skills/diffusion-models/SKILL.md rename to .agents/skills/diffusion-models/SKILL.md index 90cf20b3..6a03d10c 100644 --- a/.github/skills/diffusion-models/SKILL.md +++ b/.agents/skills/diffusion-models/SKILL.md @@ -1,11 +1,12 @@ --- 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. + 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 diff --git a/.github/skills/moe-models/SKILL.md b/.agents/skills/moe-models/SKILL.md similarity index 97% rename from .github/skills/moe-models/SKILL.md rename to .agents/skills/moe-models/SKILL.md index 8a03fb46..662cb582 100644 --- a/.github/skills/moe-models/SKILL.md +++ b/.agents/skills/moe-models/SKILL.md @@ -1,10 +1,11 @@ --- 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. + 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 diff --git a/.github/skills/multi-agent-coordination/SKILL.md b/.agents/skills/multi-agent-coordination/SKILL.md similarity index 96% rename from .github/skills/multi-agent-coordination/SKILL.md rename to .agents/skills/multi-agent-coordination/SKILL.md index 5d31d210..eb13d385 100644 --- a/.github/skills/multi-agent-coordination/SKILL.md +++ b/.agents/skills/multi-agent-coordination/SKILL.md @@ -1,6 +1,12 @@ --- 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. +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 diff --git a/.agents/skills/multimodal-models/SKILL.md b/.agents/skills/multimodal-models/SKILL.md new file mode 100644 index 00000000..f3241135 --- /dev/null +++ b/.agents/skills/multimodal-models/SKILL.md @@ -0,0 +1,179 @@ +--- +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. + +> 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): +- `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-multimodal/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 + +- **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` +- **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..18e623fe --- /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 `references/scan-pattern.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/.github/skills/scan-and-multi-image/SKILL.md b/.agents/skills/multimodal-models/references/scan-pattern.md similarity index 95% rename from .github/skills/scan-and-multi-image/SKILL.md rename to .agents/skills/multimodal-models/references/scan-pattern.md index 53392c9b..52e83335 100644 --- a/.github/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: > - 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 +# 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/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/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) +``` diff --git a/.github/skills/quality-checklist/SKILL.md b/.agents/skills/quality-checklist/SKILL.md similarity index 95% rename from .github/skills/quality-checklist/SKILL.md rename to .agents/skills/quality-checklist/SKILL.md index 3aad778d..284ce6ff 100644 --- a/.github/skills/quality-checklist/SKILL.md +++ b/.agents/skills/quality-checklist/SKILL.md @@ -1,11 +1,12 @@ --- 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. + 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 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, :] +``` diff --git a/.github/skills/weight-name-alignment/SKILL.md b/.agents/skills/weight-name-alignment/SKILL.md similarity index 96% rename from .github/skills/weight-name-alignment/SKILL.md rename to .agents/skills/weight-name-alignment/SKILL.md index 81f0539d..422e74ad 100644 --- a/.github/skills/weight-name-alignment/SKILL.md +++ b/.agents/skills/weight-name-alignment/SKILL.md @@ -1,11 +1,12 @@ --- 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. + 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 diff --git a/.github/skills/writing-rewrite-rules/SKILL.md b/.agents/skills/writing-rewrite-rules/SKILL.md similarity index 96% rename from .github/skills/writing-rewrite-rules/SKILL.md rename to .agents/skills/writing-rewrite-rules/SKILL.md index 8fa6cd6b..14f3b96a 100644 --- a/.github/skills/writing-rewrite-rules/SKILL.md +++ b/.agents/skills/writing-rewrite-rules/SKILL.md @@ -1,11 +1,12 @@ --- 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). + 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 diff --git a/.agents/skills/writing-tests/SKILL.md b/.agents/skills/writing-tests/SKILL.md new file mode 100644 index 00000000..79da765c --- /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 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 + +> 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..dfa2804e --- /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, numpy prints diagnostic info including max 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..af0c7854 --- /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 `np.testing.assert_allclose(..., strict=True)`, +which checks shape, dtype, and value equality within tolerance. + +## 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.CastLike(result, x) # cast back to x's dtype +``` + +### 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. 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/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/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/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..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. -- 11 contribution skills (`.github/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. 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/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", ) 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