diff --git a/.github/skills/adding-a-new-model/SKILL.md b/.github/skills/adding-a-new-model/SKILL.md index d48108ab..60cf4d8b 100644 --- a/.github/skills/adding-a-new-model/SKILL.md +++ b/.github/skills/adding-a-new-model/SKILL.md @@ -819,3 +819,150 @@ When adding a new model, use these files as canonical references: | **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/moe-models/SKILL.md b/.github/skills/moe-models/SKILL.md index 2d203ff7..8a03fb46 100644 --- a/.github/skills/moe-models/SKILL.md +++ b/.github/skills/moe-models/SKILL.md @@ -260,3 +260,124 @@ should be small enough for CI (~1-4B params). The test pattern: 3. Optionally test greedy generation (token-ID matching) See `tests/moe_integration_test.py` for the complete pattern. + +## Direct MoE op emission (com.microsoft.MoE) + +OnnxRuntime ships a fused `com.microsoft.MoE` contrib op (CUDA float32/fp16/bf16, +CPU float32/fp16). For new model architectures, **emit it directly** — like +`com.microsoft.GroupQueryAttention` — rather than relying on a rewrite rule. + +### When to use + +Use `com.microsoft.MoE` when: +- The model uses top-k MoE routing (standard softmax gate → TopK) +- All expert weights are the same shape (no dynamic expert counts) +- The EP's `caps.supports_fused_moe` is `True` + +Fall back to the loop-over-experts path when `supports_fused_moe` is `False` +(CPU EP without contrib ops, or EPs that don't support the custom op). + +### Gate output: full pre-topk router_probs + +The op takes the **full** `(num_tokens, num_experts)` probability tensor and +performs top-k selection internally via the `k` attribute. The gate must +produce the full softmax distribution — not already-selected top-k indices. + +```python +# In your gate forward(), return shape [num_tokens, num_experts] +router_probs = op.Softmax(op.MatMul(hidden_states, self.weight), axis=-1) +``` + +### Emission pattern (from Gemma 4 implementation) + +```python +from mobius._build_context import ep_capabilities + +caps = ep_capabilities() +if caps.supports_fused_moe: + moe_out = op.CastLike( + op.MoE( # type: ignore[attr-defined] + normed_hidden, # [num_tokens, hidden_size] + router_probs, # [num_tokens, num_experts] — full pre-topk + self.fc1_experts_weights, # [E, inter_size, hidden_size] + self.fc2_experts_weights, # [E, hidden_size, inter_size] + activation_type="silu", + k=self._top_k, + normalize_routing_weights=1, + _domain="com.microsoft", + ), + normed_hidden, # CastLike: preserve bf16/fp16/fp32 — NOT hardcoded float32 + ) +else: + moe_out = self._dispatch_moe_fallback(op, normed_hidden, router_probs) +``` + +**Critical: use `CastLike` after the MoE op.** The `com.microsoft.MoE` custom +op has `type=None` on its output — ONNX type propagation cannot infer the +output dtype. `op.CastLike(moe_output, target=input)` restores the correct +dtype (bf16/fp16/fp32), which allows downstream ops to share scalar +initializers and avoids hard-coded `Cast` to float32. + +### preprocess_weights: expert weight mapping + +HuggingFace Gemma 4 stores experts as a 3D tensor per projection: +`layers.N.experts.gate_up_proj [E, 2*inter, H]` +`layers.N.experts.down_proj [E, H, inter]` + +Map these to the parameter names used by the ONNX MoE op: + +```python +def preprocess_weights(self, state_dict): + for key in list(state_dict.keys()): + if ".experts.gate_up_proj" in key: + new_key = key.replace(".experts.gate_up_proj", ".fc1_experts_weights") + state_dict[new_key] = state_dict.pop(key) + elif ".experts.down_proj" in key: + new_key = key.replace(".experts.down_proj", ".fc2_experts_weights") + state_dict[new_key] = state_dict.pop(key) + return super().preprocess_weights(state_dict) +``` + +For models that store per-expert weights separately (one matrix per expert), +stack them into 3D tensors in `preprocess_weights`: + +```python +n = config.num_local_experts +gate = torch.stack([state_dict.pop(f"experts.{i}.gate_proj.weight") for i in range(n)]) +down = torch.stack([state_dict.pop(f"experts.{i}.down_proj.weight") for i in range(n)]) +state_dict["fc1_experts_weights"] = gate # [E, inter, hidden] +state_dict["fc2_experts_weights"] = down # [E, hidden, inter] +``` + +### EP capability check (matches GQA pattern) + +```python +# In _execution_providers.py EpCapabilities: +supports_fused_moe: bool = True # set False for EPs without com.microsoft.MoE + +# In model forward(): +from mobius._execution_providers import ep_capabilities +caps = ep_capabilities() +if caps.supports_fused_moe: + # emit com.microsoft.MoE +else: + # fallback loop +``` + +### Fallback: TopKGate + loop dispatch + +When `supports_fused_moe` is False, implement a static unroll: + +```python +def _dispatch_moe_fallback(self, op, hidden, router_probs): + k_tensor = op.Constant(value_ints=[self._top_k]) + top_weights, top_indices = op.TopK(router_probs, k_tensor, axis=-1) + top_weights = op.Softmax(top_weights, axis=-1) # renormalize + output = op.CastLike(op.ConstantOfShape(op.Shape(hidden), value=0.0), hidden) + for e_idx in range(self._num_experts): + w1 = op.Squeeze(op.Gather(self.fc1_experts_weights, [e_idx], axis=0), [0]) + w2 = op.Squeeze(op.Gather(self.fc2_experts_weights, [e_idx], axis=0), [0]) + # expert output, gated by routing weight + ... + return output +``` diff --git a/.github/skills/multimodal-models/SKILL.md b/.github/skills/multimodal-models/SKILL.md index 8247d7de..a7e9a512 100644 --- a/.github/skills/multimodal-models/SKILL.md +++ b/.github/skills/multimodal-models/SKILL.md @@ -533,3 +533,87 @@ 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) ``` + +## Any-to-Any 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). +Each tier uses a different number of exported ONNX models. + +### Tier split + +| Tier | Models | ONNX split | Example | +|------|--------|-----------|---------| +| Small Any-to-Any | E2B, E4B | 4 models: decoder + vision + **audio** + 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` | + +The task class detects which tier to use from the config (e.g. whether +`config.audio is not None` — `ArchitectureConfig.from_transformers` populates +the `audio` field when the HuggingFace config contains an audio sub-config). + +### 4-model task structure + +``` +decoder inputs_embeds [B, S, H] → logits + KV cache +vision_encoder pixel_values [B, 3, H, W] → image_features [num_image_tokens, H] +audio_encoder 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: `Gemma4AnyToAnyTask` in +`src/mobius/tasks/_gemma4.py`. This follows the same 4-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. + +### Audio encoder wiring + +The audio encoder takes raw mel-spectrogram frames and outputs token-level +features at the text hidden size: + +```python +# In Gemma4AnyToAnyTask._build_audio(): +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 `"audio"` key in the +`ModelPackage`. + +### Embedding model fuses all modalities + +The embedding model receives `input_ids`, `image_features`, and +`audio_features` as separate inputs and splices them into the token +embedding sequence at the placeholder positions: + +```python +# In Gemma4AnyToAnyTask._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] +) +# returns inputs_embeds: [B, S, H] +``` + +### Task class tier detection + +```python +class MyAnyToAnyTask(ModelTask): + def build(self, module, config): + models = {} + models["decoder"] = self._build_decoder(module.decoder, config) + models["vision"] = self._build_vision(module.vision_encoder, config) + models["embedding"] = self._build_embedding(module.embedding, config) + # Build audio encoder only when audio config is present + if config.audio is not None: + models["audio"] = self._build_audio(module.audio_encoder, config) + return ModelPackage(models, config=config) +```