Fix SmolLM3 inverted no_rope_layers logic and add L3/L5 tests - #233
Merged
Conversation
SmolLM3's no_rope_layers config was interpreted backwards: the HF convention is no_rope_layers[i] == 1 means USE RoPE (despite the misleading name), but mobius was treating it as skip RoPE. This caused 27/36 layers to have wrong RoPE gating, resulting in max logit diffs of ~6.9 vs HF (vs 0.00005 after fix). Changes: - Fix inverted RoPE gating in SmolLM3TextModel.forward() - Remove L5 xfail now that generation matches HF exactly - Add L3 integration test entry for SmolLM3-3B - Mark smollm3 as representative in L1 test configs - Regenerate L4 golden data with updated HF reference Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Justin Chu <justinchu@microsoft.com>
Exercises per-layer RoPE gating in the unit test so the inverted polarity bug is directly covered at L1 level (layer 0 = use RoPE, layer 1 = skip RoPE). Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Justin Chu <justinchu@microsoft.com>
🏗️ Architecture Diff
No architecture changes detected. ✅ Legend: ⚪ No change · 🔵 Minor (attrs/inits) · 🟡 Moderate (nodes added/removed) · 🔴 Major (interface changed) |
Performance Comparison
|
Contributor
There was a problem hiding this comment.
Pull request overview
Fixes SmolLM3’s per-layer RoPE gating to match HuggingFace’s no_rope_layers convention (where 1 means use RoPE), and updates test coverage/goldens accordingly.
Changes:
- Correct SmolLM3 RoPE gating logic (
no_rope_layers[i] == 1now enables RoPE). - Expand test coverage by adding SmolLM3-3B to integration coverage and un-xfailing L5 golden generation.
- Add
no_rope_layersto the tiny SmolLM3 config and update golden outputs.
Reviewed changes
Copilot reviewed 5 out of 5 changed files in this pull request and generated no comments.
Show a summary per file
| File | Description |
|---|---|
src/mobius/models/smollm.py |
Fixes SmolLM3 per-layer RoPE enable/disable logic and clarifies naming/comments. |
tests/_test_configs.py |
Adds no_rope_layers to tiny SmolLM3 config to exercise gating in fast graph tests. |
tests/integration_test.py |
Adds SmolLM3-3B to the integration model catalogue. |
tests/e2e_golden_test.py |
Removes SmolLM3-3B from L5-only expected-fail list so generation is asserted. |
testdata/golden/causal-lm/smollm3-3b.json |
Updates L4 golden logits to reflect corrected RoPE behavior. |
Codecov Report❌ Patch coverage is
📢 Thoughts on this report? Let us know! |
The ONNX RotaryEmbedding op requires all tensor inputs (x, cos, sin) to share the same type T. The RoPE cos/sin cache is computed in FP32 (numpy limitation), so when the model uses BF16 or FP16, the RotaryEmbedding op received mixed types: BF16 query/key + FP32 cos/sin. This cascaded through Attention -> o_proj -> residual Add, causing ORT to throw: Type Error: Type parameter (T) of Optype (Add) bound to different types (tensor(bfloat16) and tensor(float)) Fix: Add a dtype parameter to BaseRope that all RoPE subclasses thread from config.dtype. BaseRope._cast_embeddings() inserts Cast ops after gathering cos/sin when the model dtype differs from FP32. For FP32 models, no Cast is added (zero overhead). Covers all RoPE variants: DefaultRope, ProportionalRope, LinearRope, DynamicNTKRope, Llama3Rope, LongRope, YarnRope, and MRope (chunked/interleaved). Also fixes YarnRope's attn_scale dtype derivation to use the stored model dtype instead of inspecting cos_cache.dtype (which stays FP32). Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Justin Chu <justinchu@microsoft.com>
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Justin Chu <justinchu@microsoft.com>
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Justin Chu <justinchu@microsoft.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Fix inverted no_rope_layers interpretation in SmolLM3, fix BF16 type mismatch in RoPE, and add missing L3/L5 tests.
Bug Fix: Inverted no_rope_layers
HuggingFace convention: no_rope_layers[i] == 1 means USE RoPE (despite the misleading name). Mobius had it backwards, causing 27/36 layers to have wrong RoPE gating.
Bug Fix: BF16 RoPE Type Mismatch
The ONNX RotaryEmbedding op requires all tensor inputs (x, cos, sin) to share the same type T. The RoPE cos/sin cache is always computed in FP32 (numpy limitation), so BF16/FP16 models had a type constraint violation: BF16 query/key + FP32 cos/sin. This cascaded through Attention -> o_proj -> residual Add causing ORT type errors.
Fix: Added a dtype parameter to BaseRope that all RoPE subclasses thread from config.dtype. BaseRope._cast_embeddings() inserts Cast ops after gathering cos/sin when the model dtype differs from FP32. For FP32 models, no Cast is added (zero overhead). Covers all RoPE variants.
Also fixes a latent bug in YarnRope where attn_scale dtype was derived from cos_cache.dtype (always FP32), fixing Ministral3/Mistral4 BF16 support.
Results
Tests