Preserve tied weights across component builds

[Xiaoyu Z (xiaoyu-work)]

Summary

• preserve one Hugging Face component per build while deferring aliases only when the canonical table is planned for compatible quantization; reject incompatible layouts before running builds
• automatically select scoped LM-head, embedding, and vision targets for KQuant/RTN while retaining explicit opt-outs and whole-model defaults
• fail if required shared-weight assembly is skipped, keep deferred aliases through follow-up passes, recover missing float heads for one-sided quantization, and validate shared-weight ownership and metadata
• recreate packed secondary embedding placeholders on Hugging Face reload when their sidecars are present

Validation

• targeted quantizer, assembly, secondary-embedding and cross-repository suites (187 passing); Ruff and lintrunner PYLINT
• network-free tiny Gemma4: actual Olive two-build assembly -> Mobius paired embedding/decoder export -> ONNX Runtime CPU inference; full logits agree with reloaded Hugging Face model (maximum absolute difference ~4.2e-7)

Depends on onnxruntime/mobius#747 for the shared-weight contract and generic export materialization.

[Copilot]

Copilot review overview

🟡 Changes recommended

Unresolved compatibility and tied-weight validation issues remain.

Get a fresh assessment by requesting another Copilot review.

Review effort: Lite
Findings: 1

Open (1)

• Preserve ComponentInfo positional constructor compatibility · New

What changed in this PR

Preserves tied weights across independently built Hugging Face component builds.

Changes:

• Propagates shared-weight metadata.
• Defers alias quantization and restores canonical storage.
• Adds validation, tests, and documentation.

File	Summary
test/​workflows/​test_hf_component_assembly.py	Tests tied-weight assembly and validation.
test/​passes/​pytorch/​test_kquant.py	Tests deferred quantization.
test/​model/​test_composite_model.py	Tests metadata preservation.
test/​common/​test_mobius_utils.py	Tests metadata coercion.
olive/​workflows/​run/​hf_component_assembly.py	Validates and assembles tied quantized weights.
olive/​passes/​pytorch/​quant_utils.py	Defers non-canonical tied-weight quantization.
olive/​model/​config/​model_config.py	Preserves metadata during component selection.
olive/​common/​mobius_utils.py	Defines shared-weight metadata types.
docs/​source/​how-to/​configure-workflows/​build-workflow.md	Documents tied component builds.

---

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

[Copilot]

Copilot review overview

🔵 Needs a closer look

Assembly currently mishandles output-named heads and unsupported shared-weight kinds.

Review effort: Balanced
Findings: None

Resolved since last review (1)

• Preserve ComponentInfo positional constructor compatibility

Previously missed (1)

In code that hasn't changed since last review

Include output aliases to preserve quantized output heads on reload

olive/​workflows/​run/​hf_component_assembly.py:73

Include output here to match both the LM-head fallback accepted by prepare_model (quant_utils.py:923-930) and the linked Mobius shared-weight contract. Otherwise an alias such as output.weight is deferred and omitted from the assembled checkpoint, but shared_output_heads remains false, so reload does not create/tie its quantized placeholder and the output head is left missing or float.

This issue also appears on line 267 of the same file.

[Ti-Tai Wang (titaiwangms)]

Cross-repository review with onnxruntime/mobius#747 (reviewed head 4112c85). The standard Gemma4 embedding/decoder build with compatible layouts is well protected by identity, sidecar, and layout checks. I would address the following before merging:

1. Major — deferred head can remain float after a successful workflow. olive/workflows/run/builds.py:210-230 authorizes deferral when a canonical embedding build plans to quantize; olive/passes/pytorch/quant_utils.py:460-474 then disables the decoder head. But hf_component_assembly.py:190-238 returns without assembly if any selected build produces a non-HF result (or the build set is otherwise ineligible), and olive/workflows/run/run.py:207-210 still returns success. Example: add an ONNX-ending build to the decoder/embedding KQuant workflow. The decoder artifact now has a float head despite its requested quantization, and nothing fulfills the deferred request. Only defer with an enforceable assembly contract, or fail if assembly is skipped with unresolved deferrals.
2. Major — a later quantization pass can select the deferred head again. quant_utils.py:960-986 reapplies auto lm_head selection on each pass but only records/handles deferral when existing_qcfg is None; quant_utils.py:1042-1059 then ORs the new selection into the existing config. A decoder pipeline with KQuant followed by RTN can therefore quantize the head during its second pass while retaining a prior deferral request, leading to conflicting packed tensors or layouts at assembly. Preserve and honor existing deferred endpoints across passes; add a two-pass regression test. The concrete KQuant->RTN outcome is not runtime-verified here.
3. Major — the planner can disagree with actual embedding targeting. _build_quantizes_shared_table() (builds.py:141-155) only checks pass options, while quant_utils.py:958-973,1207-1223 also honors mixed_precision_info.default. With an embedding component pass omitting embeds and mixed-precision default embeds: false, the planner marks the canonical table as planned and the decoder defers, but the embedding pass creates no packed table. Assembly then fails at hf_component_assembly.py:274-284, after both builds. Resolve effective targets once (including mixed-precision defaults and skips), or avoid deferral when canonical production is uncertain.
4. Major / behavior decision — incompatible layouts now fail late rather than remain independent. builds.py:141-155 does not compare canonical and alias qargs before deferral; _resolve_shared_weights() rejects unequal qargs at hf_component_assembly.py:318-320. For an intentionally 8-bit embedding and 4-bit LM head, the builds complete before failing. Previously, independently quantizing those sides could break the tie and retain two tables. Decide whether that mode remains supported; if yes, do not defer the head when qargs differ, and if not, reject early with a targeted message/test.

Additional correctness guards:

• hf_component_assembly.py:68-73 omits output from _OUTPUT_HEAD_MODULE_NAMES, although Mobius Fix pytorch tensor parallel implementation #747 includes it and Olive's head discovery also accepts it. A tied output.weight would be materialized but its assembled decoder lm_head flag could remain false. Align the name sets and test that endpoint.
• olive/common/mobius_utils.py:63-84 accepts duplicate/canonical-as-alias endpoints even though Mobius rejects them. _shared_weight_exclusions() (hf_component_assembly.py:363-372) would then exclude canonical sidecars too. Validate serialized shared-weight endpoint uniqueness and ownership before deduplicating.

Not established by this review: whether quantizing only the embedding while explicitly leaving the LM head float always retains a loadable float head after assembly; please add a checkpoint reload regression for this boundary. I did not run the PR's tests or an end-to-end ONNX export, so the follow-up-pass and one-sided-save runtime outcomes should be confirmed with targeted tests. The docs example also changes component-target defaults without explaining the changed lm_head/embeds/quantize_vision behavior; a short migration note would help.

[Xiaoyu Z (xiaoyu-work)]

Thanks Ti-Tai Wang (@titaiwangms) — addressed the review in 9c5534e. Deferred heads now survive follow-up quantizers, and workflows fail if required HF assembly is skipped. The planner uses effective mixed-precision settings and rejects incompatible layouts before running builds. Assembly restores a float head when only the canonical embedding was quantized, recognizes output.weight, and rejects malformed or unsupported shared-weight declarations. Targeted real-checkpoint reload tests cover the follow-up and one-sided cases.

[Ti-Tai Wang (titaiwangms)]

Follow-up review at head 9c5534e (paired with onnxruntime/mobius#747): I'm holding approval for now, but the four Major issues in my previous comment have been substantively addressed. Deferred aliases persist through subsequent quantization passes; skipped required assembly fails rather than returning success; planning now accounts for effective mixed-precision options; and incompatible endpoint layouts are rejected before builds run. The new float-head reload, output.weight, and malformed-metadata coverage also address the related follow-ups.

The remaining gate is the cross-repository contract. Mobius #747 still needs to clarify how conflicting parent/nested tie_word_embeddings declarations are resolved, and I have not seen an end-to-end test that starts with an assembled Olive checkpoint and actually exports/loads the paired Mobius ONNX components. Please confirm the supported scope and run or add that integration check before the final approval decision.

One behavioral point to document explicitly: incompatible embedding/head quantization layouts now fail early rather than being retained as two independently quantized, untied tables. If rejecting that configuration is intentional, it is no longer a late-failure bug; it is a compatibility decision users should know about.

This is a follow-up comment, not a formal GitHub changes-requested review.

[Xiaoyu Z (xiaoyu-work)]

Ti-Tai Wang (@titaiwangms) Addressed the cross-repository gate with onnxruntime/mobius#747 (fb6b087c) and be1a942 here. A locally generated tiny tied Gemma4 checkpoint now runs through actual olive.run assembly, mobius.build, and both ONNX Runtime CPU components; full logits match reloaded Hugging Face within ~4.2e-7 maximum absolute difference. This is checked by test/workflows/test_tied_gemma4_export.py (optional when the companion Mobius version is unavailable). The test also exposed secondary packed embeddings not being restored on HF reload; the quantizer now reconstructs those placeholders from checkpoint sidecars. A single sentence in the workflow guide notes that incompatible tied layouts fail before builds run.

[Ti-Tai Wang (titaiwangms)]

Reviewed the latest cross-component quantization and assembly changes. The earlier review issues are addressed and CI is green. The author reports a successful local Olive -> Mobius -> ONNX Runtime parity run; the optional cross-repository test is not asserted to have run in Olive CI. Approved.