fix: share tied embeddings with the quantized LM head in ModelBuilder

[Yuri Khrustalev (ykhrustalev)]

Problem
patched_make_embedding writes its own table for every embedding Olive didn't quantize, bypassing the builder's shared_embeddings path (on by default for tied models). Tied models built at INT4 therefore store a second embedding table next to the quantized LM head: full precision under k_quant, which leaves Gather unquantized, and INT4 otherwise.

Solution

• Delegates to the builder's own make_embedding for embeddings Olive didn't quantize, as patched_make_packed_matmul_int4 already does
• Keeps a separate table when the hidden size isn't a multiple of the INT4 block size, since the builder's tied lookup drops the LM head's block padding
• Keeps the builder's method intact when maybe_patch_quant runs more than once
• Makes tied embeddings follow the LM head precision; shared_embeddings: false in extra_options keeps the separate table

Testing

• test_model_builder_int4_embeddings covers tied, untied and padded-block tiny Llama; the tied case fails without the fix
• LFM2.5 INT4 recipes (230M–8B-A1B; CPU, CUDA, WebGPU) rebuild 10–54% smaller, e.g. 1.2B-Instruct CPU 1.36 → 0.86 GiB, with KL divergence to FP32 at most 0.001 higher than before
• Default INT4 Qwen2.5-0.5B, Qwen3-0.6B and Qwen3-1.7B rebuild 15–21% smaller, with KL at most 0.003 higher
• Qwen2.5-0.5B-Instruct with block_size: 256 (hidden size 896) keeps its table and matches the previous build exactly

Refs microsoft/onnxruntime-genai#2628

[azure-pipelines]

Azure Pipelines:
There may be pipelines that require an authorized user to comment /azp run to run.

[Copilot]

Copilot review overview

🟡 Changes recommended

The embedding fallback can be invoked when the native method is unavailable, causing non-quantized builds to fail.

Review effort: Lite
Findings: 1

Open (1)

• Guard missing make_embedding before invoking patched method · New

What changed in this PR

Fixes tied INT4 embedding export by reusing the native builder’s shared-embedding behavior.

Changes:

• Delegates unquantized embeddings to the native builder.
• Preserves the original embedding method across repeated patching.
• Adds tied and untied INT4 embedding tests.

File	Summary
test/​passes/​onnx/​test_model_builder.py	Adds regression coverage for tied embeddings and repeated patching.
olive/​passes/​onnx/​model_builder.py	Updates embedding delegation; requires a callable fallback when the native method is unavailable.

---

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.