Fix YaRN RoPE bugs in model builder and add parity tests by titaiwangms · Pull Request #2076 · microsoft/onnxruntime-genai

Ti-Tai Wang (titaiwangms) · 2026-04-08T23:15:10Z

Summary

Fixes four bugs in the ModelBuilder's YaRN RoPE configuration resolution that caused completely wrong cos/sin caches, producing garbage output for ALL YaRN-based models (Ministral-3-3B, GPT-OSS-20B, and any model using beta_fast/beta_slow in rope_scaling).

Also adds comprehensive parity tests for 3 YaRN configurations, consolidates mscale logic, and excludes src/python/py/models/ from lintrunner (legacy formatting conflicts).

Bugs Fixed

Bug 1: `hasattr` on dict for `original_max_position_embeddings`

hasattr(config.rope_scaling, 'original_max_position_embeddings') always returns False for dicts. Fixed to use isinstance(Mapping) + in operator.

Bug 2: `rope_theta` fallback

Models that store rope_theta only in rope_scaling dict (not as a top-level config attribute) fell through to default theta=10000. Added fallback chain: config.rope_theta → config.rope_embedding_base → config.rope_scaling["rope_theta"] → 10000.

Bug 3: YaRN `mscale` override

Always computed mscale from factor via make_mscale(), ignoring explicit mscale=1.0 in config. Ministral-3-3B sets mscale=1.0 but was getting mscale≈1.277. Now respects the config value when > 0.

Bug 4: `inv_freq` double-inversion

make_inv_freq_rescaled_with_ntk computed 1/(factor * inv_freq) which double-inverts since inv_freq is already 1/pos_freqs, giving pos_freqs/factor (wrong). Fixed to inv_freq / factor.

Review Feedback Addressed

Changed isinstance(rope_scaling, dict) → isinstance(rope_scaling, Mapping) in 2 locations
Consolidated inline mscale override logic into make_mscale(mscale, config_mscale=0) — config override is now handled centrally with backward-compatible default
Removed make_mscale default parameter values (callers pass explicitly)
Added mscale_all_dim documentation comment
Tests drive real builder code (make_rope_init) instead of reimplementing logic inline
Added GPT-OSS-20B test config per reviewer request

Tests Added

Added test/python/test_yarn_rope_parity.py with 16 tests covering 3 YaRN configurations:

Ministral-3-3B (factor=16, theta=1M, explicit mscale=1.0):

test_ministral_3b_cos_sin_match — end-to-end parity with HF reference
test_bug_a_hasattr_on_dict — dict key access + wrong fallback guard
test_bug_b_rope_theta_fallback — theta from rope_scaling + default guard
test_bug_c_mscale_override — explicit mscale via make_rope_init()
test_mscale_fallback_when_absent — .get("mscale", 0) fallback
test_mscale_all_dim_formula — mscale_all_dim computation
test_bug_d_inv_freq_no_double_inversion — NTK scaling vs HF reference
test_full_cache_length — parity at 2048 positions
test_mapping_isinstance_with_frozen_dict — MappingProxyType support

Synthetic YaRN (factor=40, theta=10K, no explicit mscale):
10. test_yarn_no_mscale_cos_sin_match — parity with HF reference
11. test_yarn_no_mscale_uses_computed_value — computed mscale from factor
12. test_different_yarn_configs_produce_different_caches — configs produce distinct outputs

GPT-OSS-20B (factor=4, theta=500K, top-level rope_theta, no explicit mscale):
13. test_gptoss_20b_cos_sin_match — end-to-end parity with HF reference
14. test_gptoss_20b_top_level_rope_theta — verifies top-level theta is used
15. test_gptoss_20b_computed_mscale — computed mscale from factor=4
16. test_gptoss_20b_full_cache_length — parity at 2048 positions

Impact

Fixes ALL models using YaRN RoPE (beta_fast/beta_slow in rope_scaling)
Verified: cos/sin caches match HF transformers reference (atol=1e-5)
Verified: Ministral-3-3B generates correct output ("The capital of France is Paris.")

Files Changed

.lintrunner.toml — exclude src/python/py/models/** from lintrunner (legacy formatting)
src/python/py/models/builders/base.py — 4 bug fixes + Mapping import + mscale consolidation
test/python/test_yarn_rope_parity.py — 16 parity tests for 3 YaRN configs (new file)

Copilot

Pull request overview

Fixes YaRN RoPE configuration resolution in the Python ModelBuilder to prevent incorrect rotary cos/sin cache generation (which can produce unusable outputs for YaRN-based models).

Changes:

Fixes original_max_position_embeddings detection when rope_scaling is stored as a dict-like object.
Adds rope_theta fallback to read from rope_scaling when not present as a top-level config field.
Respects explicit YaRN mscale from config and corrects NTK rescaling math to avoid double-inversion.

Copilot

Pull request overview

Copilot reviewed 2 out of 2 changed files in this pull request and generated 5 comments.

Ti-Tai Wang (titaiwangms)

Reverted all cosmetic spacing changes in base.py. Only functional YaRN fixes remain — diff is now 36 changed lines (29+, 7-) vs ~487 before. Thanks for the catch!

Squashed into single clean commit: f24ee36

Tianlei Wu (tianleiwu)

Review Summary

All four YaRN RoPE bug fixes are correct and well-tested. The fixes are minimal, targeted, and produce cos/sin caches matching the HuggingFace transformers reference.

Positives:

Using collections.abc.Mapping instead of isinstance(_, dict) correctly handles MappingProxyType, ordered dicts, and other Mapping implementations.
The rope_theta fallback chain covers all known config layouts.
Bug 4 (inv_freq double-inversion) fix is mathematically clean: inv_freq / factor correctly gives 1/(factor * pos_freqs) matching HF.
Test suite is comprehensive: 11 tests with HF reference parity, guard assertions preventing false passes, and real builder code exercised via make_rope_init.

One suggestion on make_mscale generalization (see inline comment). Two nitpicks on test file (below).

Test nitpicks (non-blocking):

_make_builder_cos_sin uses object.__new__(Model) and manually initializes attributes. If make_rope_init starts depending on new attributes from __init__, tests will fail with opaque AttributeErrors. Consider adding a comment listing assumed attributes.
YARN_NO_MSCALE_CONFIG docstring says "Modeled after DeepSeek-V2 style YaRN scaling" but DeepSeek-V2 uses longrope/su-type RoPE, not yarn. The inline comment ("Synthetic YaRN config") is accurate — consider aligning the docstring.

Ti-Tai Wang (titaiwangms)

Thanks for the review feedback!

Tianlei Wu (@tianleiwu) — Good suggestion on the mscale_all_dim formula. We went ahead and implemented it in this PR since it was straightforward: make_mscale() now accepts config_mscale_all_dim and uses the full HF formula (get_mscale(factor, mscale) / get_mscale(factor, mscale_all_dim)) when both are provided. Added a test covering equal, different, and fallback cases. Backward compatible — existing callers default to config_mscale_all_dim=0.

kunal-vaishnavi — Confirmed: all cosmetic spacing changes have been reverted. Only functional YaRN RoPE fixes remain (36→40 changed lines in base.py). The make_mscale consolidation you suggested is also in place with config_mscale parameter.

Ti-Tai Wang (titaiwangms) · 2026-04-15T22:21:01Z

I constantly found CI failures that are not from this PR but they are required. How did people merge their PR in genai? Or I miss something.

Fix four bugs in the YaRN RoPE implementation in base.py: 1. Use isinstance(Mapping) instead of hasattr() for dict-based rope_scaling — hasattr always returns False for dict keys 2. Resolve rope_theta from rope_scaling when top-level attr is absent 3. Add config_mscale parameter to make_mscale() so explicit mscale values (e.g. Ministral-3-3B's mscale=1.0) override computed values 4. Fix inv_freq double-inversion: use inv_freq/factor instead of 1/(factor*inv_freq) Add 11 parity tests comparing builder cos/sin caches against HuggingFace reference implementation for two YaRN configurations (Ministral-3-3B and a DeepSeek-V2 style config). Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

Add Mistral3TextModel builder for mistralai/Ministral-3-3B-Instruct-2512: - VLM text decoder with FP8 dequantization (weight * scale_inv) - Patches genai_config.json with image_token_id from HF config - Excludes embedding layer (handled by separate embedding model) - Registers Mistral3ForConditionalGeneration in builder.py dispatch Also includes from PR #2076 (base branch): - Fix 4 YaRN RoPE bugs (hasattr, theta fallback, mscale, inv_freq) - isinstance(Mapping) for robust rope_scaling checks - Consolidated make_mscale() with config_mscale parameter - 12 YaRN RoPE parity tests Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

Add src/python/py/models/** to exclude_patterns for both RUFF and RUFF-FORMAT linters in .lintrunner.toml. This prevents the linter from introducing formatting-only changes to model builder files, keeping PR diffs focused on functional changes. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

Convert lambda assignment to a proper function definition to satisfy ruff E731. Add required blank lines around nested def per ruff-format. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

Add Mistral3TextModel builder for mistralai/Ministral-3-3B-Instruct-2512: - VLM text decoder with FP8 dequantization (weight * scale_inv) - Patches genai_config.json with image_token_id from HF config - Excludes embedding layer (handled by separate embedding model) - Registers Mistral3ForConditionalGeneration in builder.py dispatch Also includes from PR #2076 (base branch): - Fix 4 YaRN RoPE bugs (hasattr, theta fallback, mscale, inv_freq) - isinstance(Mapping) for robust rope_scaling checks - Consolidated make_mscale() with config_mscale parameter - 12 YaRN RoPE parity tests Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

## Summary Adds Mistral3/Pixtral VLM support to onnxruntime-genai with multi-image inference. Includes C++ image processor, PixtralVisionState for per-image vision processing, Python export support, and comprehensive tests. ## Changes ### C++ Runtime - **Mistral3 image processor** — `[IMG]`/`[IMG_BREAK]`/`[IMG_END]` token expansion based on image resolution and patch geometry, multi-image support - **PixtralVisionState** — per-image vision processing loop with bounds checks and overflow guard; slices from padded batch tensor using `image_sizes` metadata from ort-extensions `PixtralImageSizes` op - **Virtual `SetExtraInputs`** — proper polymorphic dispatch for vision state subclasses - **`IsPixtralFamily()` model type detection** — enables Pixtral-specific codepath - **`processor_config.json`** — with `PixtralImageSizes` preprocessing step - **`context_length` / `max_length` separation** — `context_length` controls KV cache allocation while `max_length` controls generation stopping, preventing premature EOS with large image token counts - **INT32 `input_ids`** — token IDs above 32767 (Pixtral `[IMG]`=128011) require int32 ### Python Export Support - Mistral3 model classes (`Mistral3Config`, `Mistral3ForConditionalGeneration`) - FP8 dtype promotion for checkpoint loading - `get_user_content()` handler for Mistral3 prompt formatting ## Multi-Image Architecture Pixtral uses dynamic image sizes (28×28 to 1540×1540) so images can't be batched in the vision encoder. `PixtralVisionState` processes each image individually by: 1. Reading `image_sizes` tensor from ort-extensions `PixtralImageSizes` op (provides per-image H×W) 2. Slicing the padded `[N, C, max_H, max_W]` batch tensor to extract each image's actual pixels 3. Running vision encoder on each image separately 4. Concatenating vision embeddings for the decoder ## Dependencies - **onnxruntime-extensions PR #1050** — `PixtralImageSizes` custom op for image size metadata - **PR #2076** — YaRN RoPE parity fixes (merged) ✅ ## Testing - 5 multi-image token expansion tests (various image sizes and counts) - Virtual dispatch verification tests for `SetExtraInputs` - YaRN RoPE parity tests (from merged #2076) - E2E verified: multi-image (fish.jpg + challenge.jpg) correctly describes both images --------- Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

Copilot AI review requested due to automatic review settings April 8, 2026 23:15

Ti-Tai Wang (titaiwangms) mentioned this pull request Apr 8, 2026

[Mistral3] Add VLM support with multi-image inference #2077

Merged

Copilot started reviewing on behalf of Ti-Tai Wang (titaiwangms) April 8, 2026 23:16 View session

Ti-Tai Wang (titaiwangms) marked this pull request as draft April 8, 2026 23:17

Copilot AI reviewed Apr 8, 2026

View reviewed changes

Comment thread src/python/py/models/builders/base.py

Comment thread src/python/py/models/builders/base.py

Ti-Tai Wang (titaiwangms) force-pushed the pr/fix-yarn-config-resolution branch 2 times, most recently from dbb6147 to d786f35 Compare April 9, 2026 20:28

Ti-Tai Wang (titaiwangms) marked this pull request as ready for review April 9, 2026 21:00

Ti-Tai Wang (titaiwangms) changed the title ~~Fix 4 ModelBuilder config resolution bugs for YaRN RoPE models~~ Fix YaRN RoPE bugs in model builder and add parity tests Apr 9, 2026

Ti-Tai Wang (titaiwangms) requested review from Baiju Meswani (baijumeswani), Copilot, kunal-vaishnavi and Xavier Dupré (xadupre) April 9, 2026 21:04

Copilot started reviewing on behalf of Ti-Tai Wang (titaiwangms) April 10, 2026 20:49 View session

Copilot AI reviewed Apr 10, 2026

View reviewed changes

Ti-Tai Wang (titaiwangms) force-pushed the pr/fix-yarn-config-resolution branch from d786f35 to 216066d Compare April 10, 2026 21:10

Ti-Tai Wang (titaiwangms) requested a review from Tianlei Wu (tianleiwu) April 10, 2026 21:52

kunal-vaishnavi requested a review from Akshay Sonawane (apsonawane) April 13, 2026 20:14

kunal-vaishnavi reviewed Apr 13, 2026

View reviewed changes

Comment thread test/python/test_yarn_rope_parity.py

kunal-vaishnavi reviewed Apr 13, 2026

View reviewed changes

Comment thread src/python/py/models/builders/base.py Outdated

Ti-Tai Wang (titaiwangms) force-pushed the pr/fix-yarn-config-resolution branch from 216066d to 0b8c74a Compare April 14, 2026 19:52

Ti-Tai Wang (titaiwangms) requested a review from kunal-vaishnavi April 14, 2026 19:53

Ti-Tai Wang (titaiwangms) force-pushed the pr/fix-yarn-config-resolution branch from 0b8c74a to 81b1c13 Compare April 14, 2026 22:41

kunal-vaishnavi reviewed Apr 15, 2026

View reviewed changes

Comment thread src/python/py/models/builders/base.py Outdated

Ti-Tai Wang (titaiwangms) force-pushed the pr/fix-yarn-config-resolution branch from 81b1c13 to f24ee36 Compare April 15, 2026 19:59

Ti-Tai Wang (titaiwangms) commented Apr 15, 2026

View reviewed changes

Ti-Tai Wang (titaiwangms) requested a review from kunal-vaishnavi April 15, 2026 20:01

Tianlei Wu (tianleiwu) previously approved these changes Apr 15, 2026

View reviewed changes

Comment thread src/python/py/models/builders/base.py

Ti-Tai Wang (titaiwangms) commented Apr 15, 2026

View reviewed changes

Ti-Tai Wang (titaiwangms) dismissed Tianlei Wu (tianleiwu)’s stale review via 71a0f29 April 15, 2026 22:12

Ti-Tai Wang (titaiwangms) force-pushed the pr/fix-yarn-config-resolution branch from f24ee36 to 71a0f29 Compare April 15, 2026 22:12

Ti-Tai Wang (titaiwangms) force-pushed the pr/fix-yarn-config-resolution branch from 71a0f29 to a959221 Compare April 15, 2026 22:31

kunal-vaishnavi reviewed Apr 15, 2026

View reviewed changes

Comment thread src/python/py/models/builders/base.py

kunal-vaishnavi reviewed Apr 16, 2026

View reviewed changes

Comment thread test/python/test_yarn_rope_parity.py

Ti-Tai Wang (titaiwangms) force-pushed the pr/fix-yarn-config-resolution branch from b2c21e0 to a959221 Compare April 16, 2026 19:57

Ti-Tai Wang (titaiwangms) and others added 2 commits April 17, 2026 00:48

Ti-Tai Wang (titaiwangms) requested a review from kunal-vaishnavi April 17, 2026 20:22

kunal-vaishnavi approved these changes Apr 18, 2026

View reviewed changes

kunal-vaishnavi enabled auto-merge (squash) April 18, 2026 00:22

kunal-vaishnavi merged commit b297562 into main Apr 20, 2026
15 of 19 checks passed

kunal-vaishnavi deleted the pr/fix-yarn-config-resolution branch April 20, 2026 19:09

dependabot Bot mentioned this pull request May 30, 2026

Bump Microsoft.ML.OnnxRuntimeGenAI from 0.13.1 to 0.14.0 dvai-global/dvai-bridge#33

Closed

dependabot Bot mentioned this pull request Jun 7, 2026

Bump Microsoft.ML.OnnxRuntimeGenAI.DirectML and Microsoft.ML.OnnxRuntimeGenAI.Managed EricCogen/GauntletCI#226

Merged

This was referenced Aug 1, 2026

Bump Microsoft.ML.OnnxRuntimeGenAI and Microsoft.ML.OnnxRuntimeGenAI.Managed trackdubllc/Trackdub#53

Merged

Bump Microsoft.ML.OnnxRuntime.Gpu and Microsoft.ML.OnnxRuntimeGenAI.Cuda trackdubllc/Trackdub#54

Closed

Navigation Menu

Uh oh!

Fix YaRN RoPE bugs in model builder and add parity tests - #2076

Fix YaRN RoPE bugs in model builder and add parity tests#2076
kunal-vaishnavi merged 3 commits into
mainfrom
pr/fix-yarn-config-resolution

Ti-Tai Wang (titaiwangms) commented Apr 8, 2026 •

edited

Loading

Uh oh!

Copilot AI left a comment

Uh oh!

Uh oh!

Uh oh!

Copilot AI left a comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Ti-Tai Wang (titaiwangms) left a comment

Uh oh!

Tianlei Wu (tianleiwu) left a comment

Uh oh!

Uh oh!

Ti-Tai Wang (titaiwangms) left a comment

Uh oh!

Ti-Tai Wang (titaiwangms) commented Apr 15, 2026

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

Uh oh!

Conversation

Ti-Tai Wang (titaiwangms) commented Apr 8, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary

Bugs Fixed

Bug 1: hasattr on dict for original_max_position_embeddings

Bug 2: rope_theta fallback

Bug 3: YaRN mscale override

Bug 4: inv_freq double-inversion

Review Feedback Addressed

Tests Added

Impact

Files Changed

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Uh oh!

Uh oh!

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Ti-Tai Wang (titaiwangms) left a comment

Choose a reason for hiding this comment

Uh oh!

Tianlei Wu (tianleiwu) left a comment

Choose a reason for hiding this comment

Review Summary

Uh oh!

Uh oh!

Ti-Tai Wang (titaiwangms) left a comment

Choose a reason for hiding this comment

Uh oh!

Ti-Tai Wang (titaiwangms) commented Apr 15, 2026

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

Ti-Tai Wang (titaiwangms) commented Apr 8, 2026 •

edited

Loading

Bug 1: `hasattr` on dict for `original_max_position_embeddings`

Bug 2: `rope_theta` fallback

Bug 3: YaRN `mscale` override

Bug 4: `inv_freq` double-inversion