Fix tied embeddings - #2237
Conversation
There was a problem hiding this comment.
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_embeddingsflags 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_int4behavior). - Updates the Python test runner entrypoint to run the
builder/andmodels/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. |
tianleiwu
left a comment
There was a problem hiding this comment.
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.
### 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.
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.