Skip to content

Fix tied embeddings - #2237

Merged
kunal-vaishnavi merged 10 commits into
mainfrom
kvaishnavi/fix-tied-embeds
Jul 2, 2026
Merged

Fix tied embeddings#2237
kunal-vaishnavi merged 10 commits into
mainfrom
kvaishnavi/fix-tied-embeds

Conversation

@kunal-vaishnavi

Copy link
Copy Markdown
Contributor

Description

This PR rewrites how tied embeddings are determined in the model builder. It also adds a unit test file to ensure that the different cases for tied embeddings are tested in the CIs.

Motivation and Context

There are still some small edge cases that haven't been validated for tied embeddings.

Copilot AI review requested due to automatic review settings June 23, 2026 01:10
@kunal-vaishnavi
kunal-vaishnavi requested a review from a team as a code owner June 23, 2026 01:10

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Pull request overview

This PR updates the Python model builder’s tied-embedding initialization logic and adds CI coverage to exercise tied-embedding edge cases via pytest.

Changes:

  • Refactors tied-embedding determination in the model builder to compute explicit tied_quantized_embeddings / tied_unquantized_embeddings flags and uses them in embedding construction.
  • Adds a new pytest module to cover shared-embedding configuration permutations (plus a small unit test for make_matmul_int4 behavior).
  • Updates the Python test runner entrypoint to run the builder/ and models/ pytest suites as part of the standard pipeline run.

Reviewed changes

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

File Description
src/python/py/models/builders/base.py Reworks tied-embedding / quantization eligibility flags and switches embedding logic to use the new tied flags.
test/python/builder/test_tied_embeddings.py Adds unit tests for tied-embedding flag behavior and INT4 MatMul fallback/emit behavior.
test/python/test_onnxruntime_genai.py Expands the invoked pytest targets to include builder/ and models/ test suites.

Comment thread src/python/py/models/builders/base.py
Comment thread test/python/builder/test_tied_embeddings.py Outdated
Comment thread src/python/py/models/quantized_model.py Fixed

@tianleiwu tianleiwu left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks for tightening up the tied-embeddings logic — splitting the decision into explicit tied_quantized_embeddings / tied_unquantized_embeddings flags (and gating on exclude_embeds/exclude_lm_head/prune_lm_head up front) is a clear improvement, and the new make_tied_quantized_embedding_input_names helper plus test_tied_embeddings.py give good coverage. This also fixes a latent mismatch where make_algo_config forced /lm_head/MatMul to 8 bits for plain k_quant while make_embedding used a 4-bit initializer name.

The two earlier bot threads ("MatMul quantized but Gather not" and exclude_lm_head) look addressed by the new guards on the current head. A couple of small items remain — see inline comments. Additionally: test_onnxruntime_genai.py drops the explicit run_gemma4_vision_tests / run_qwen_fara_vision_tests calls in favor of pytest ... builder models; since those helpers are run_* (not test_*) functions, please confirm they are still exercised somewhere so vision coverage isn't silently lost.

Non-blocking overall.

Comment thread src/python/py/models/quantized_model.py
Comment thread src/python/py/models/builder.py
@tianleiwu
tianleiwu enabled auto-merge (squash) July 1, 2026 22:32
@kunal-vaishnavi
kunal-vaishnavi disabled auto-merge July 2, 2026 00:54
@kunal-vaishnavi
kunal-vaishnavi merged commit 9722726 into main Jul 2, 2026
62 of 65 checks passed
@kunal-vaishnavi
kunal-vaishnavi deleted the kvaishnavi/fix-tied-embeds branch July 2, 2026 00:55
tianleiwu pushed a commit that referenced this pull request Jul 11, 2026
### Description

This PR rewrites how tied embeddings are determined in the model
builder. It also adds a unit test file to ensure that the different
cases for tied embeddings are tested in the CIs.

### Motivation and Context

There are still some small edge cases that haven't been validated for
tied embeddings.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

OGA export generating an invalid model when a lm_head unquantized model is exported

4 participants