Fix gemma4 multi component quantization

[Xiaoyu (xiaoyu-work)]

Describe your changes

This pull request improves quantization robustness for small floating point values and enhances the ModelConfig component selection logic to support multiple components, along with corresponding test updates.

Quantization robustness:

• In olive/common/quant/utils.py, the find_qparams method now clamps the calculated scales to a minimum positive value to prevent underflow to zero for tiny FP16 values, ensuring numerical stability during quantization.
• Added a test (test_find_qparams_tiny_fp16_values_do_not_underflow) to verify that quantizing nonzero FP16 groups does not produce zero scales, preventing underflow issues.

ModelConfig component selection enhancements:

• Removed the restriction that only one HuggingFace model component can be selected at a time in ModelConfig.select_components, allowing aggregation of multiple components for component-scoped PyTorch optimization.
• Updated the logic in _select_hf_component to handle multiple selected components, aggregating their names and source paths into model_attributes, and improved error handling for missing source paths.
• Updated the test for multiple component selection to verify the correct aggregation of component names and source paths, replacing the previous test that expected an error.

Checklist before requesting a review

• Add unit tests for this change.
• Make sure all tests can pass.
• Update documents if necessary.
• Lint and apply fixes to your code by running lintrunner -a
• Is this a user-facing change? If yes, give a description of this change to be included in the release notes.

(Optional) Issue link

[Copilot]

Pull request overview

This PR improves robustness in Olive’s weight quantization for tiny FP16 values and extends ModelConfig.select_components to support selecting multiple Hugging Face model components (aggregating their runtime source paths) for component-scoped PyTorch optimization.

Changes:

• Clamp WeightQuantizer.find_qparams scales to the smallest positive representable value to avoid FP16 underflow to zero.
• Allow HfModel component selection of multiple components and aggregate their component_source_paths into model_attributes.
• Update/add tests covering multi-component selection aggregation and tiny-FP16 scale underflow prevention.

Reviewed changes

Copilot reviewed 4 out of 4 changed files in this pull request and generated 1 comment.

File	Description
olive/common/quant/utils.py	Clamps computed quantization scales to a minimum positive value to prevent underflow to zero.
olive/model/config/model_config.py	Updates HF component selection logic to support multiple selected components and aggregate source paths into model_attributes.
test/common/quant/test_utils.py	Adds a unit test ensuring tiny nonzero FP16 weights do not yield zero scales.
test/model/test_composite_model.py	Updates HF-model component selection tests to validate multi-component aggregation instead of expecting an error.

---

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

[Ti-Tai Wang (titaiwangms)]

Review team verdict: Request changes

The review team found one critical model-correctness issue and several major checkpoint integrity and integration problems.

Severity	Location	Finding
Critical	olive/workflows/run/hf_component_assembly.py:182-224	Quantization metadata is merged by copying the first build and reconciling only selected fields. tie_word_embeddings can therefore depend on build order, causing missing quantized lm_head buffers or incorrectly tied embeddings and silently corrupting the loaded model. Merge every load-affecting field explicitly and reject incompatible configurations.
Major	olive/workflows/run/hf_component_assembly.py:192-206	Generated exact overrides are appended after inherited regex overrides, but Olive uses first-match-wins ordering. A broad earlier override can reconstruct component tensors with the wrong group size or bit width. Resolve effective settings per module or place exact overrides first.
Major	olive/workflows/run/hf_component_assembly.py:208-224	Removing any skip pattern that matches one quantized module also removes protection for unrelated float modules matched by that pattern. Derive exact skip entries from assembled float-weight keys, following the existing logic in quant_utils.py.
Major	olive/workflows/run/hf_component_assembly.py:367-406	Publication is non-atomic: the new index is copied before referenced component shards, then existing outputs and metadata are deleted. ENOSPC, interruption, or concurrent assembly can leave a corrupt checkpoint and destroy valid build outputs. Publish shards first and atomically replace the index last; use unique staging and locking.
Major	olive/workflows/run/hf_component_assembly.py:381-395	Cleanup does not remove a pre-existing top-level model.safetensors or pytorch_model.bin. Transformers prefers that file over the new index, so loading can silently use a stale model. Reject or remove all conflicting checkpoint formats before publication.
Major	olive/workflows/run/hf_component_assembly.py:384-406, olive/cli/run.py:114-131	Assembly deletes each build's model/ and model_config.json, while CLI output and persisted footprints still point to those locations. Chained commands, packaging, and provenance can reference deleted artifacts. Preserve or rewrite those artifacts and print the assembled parent path.
Major	olive/model/config/model_config.py:213-219	Multi-component selection removes component_role, but downstream PyTorch quantization still uses that field for wrapper selection, decoder handling, and target eligibility. Mixed decoder/vision selections can now succeed while only partially applying the requested optimization. Either restore one-component-per-build or fully support component_names downstream.
Major	olive/passes/pytorch/quant_utils.py:1106-1117	The meta-parameter fallback ties output embeddings to input embeddings based only on equal shape. Models with intentionally untied embeddings can be silently changed. Require explicit tying metadata such as config.tie_word_embeddings.
Major	olive/workflows/run/hf_component_assembly.py:344-358	Unoptimized tensors are copied exclusively from the first build without validating other builds' key sets, shapes, or configs. Keys unique to later builds can disappear, and incompatible component outputs can be combined. Validate cross-build checkpoint/config compatibility before assembly.
Major	olive/workflows/run/builds.py:84-96	Automatic assembly is inferred from whether named builds explicitly set output_dir, rather than whether resolved outputs share a parent as documented. This both rejects valid sibling directories and enables destructive assembly without an explicit opt-out. Derive the parent from resolved paths and add an assembly setting.
Major	olive/cli/capture_onnx.py:199-208	The new Diffusers Mobius path uses get_diffusers_input_model, which does not forward trust_remote_code; commands requiring custom model code now invoke Mobius with it disabled. Propagate the shared load arguments.

Non-blocking compatibility questions

• Safetensors shards are written without metadata={"format": "pt"}; compatibility with older Transformers versions needs a targeted run.
• Component shards appear in subdirectories. Transformers supports index-relative paths, but consumers that scan only root-level *.safetensors may load an incomplete model.
• The new ONNX IR round-trip should be checked against the pass's documented support for models larger than 2 GB.

[Xiaoyu (xiaoyu-work)]

Copilot resolve the merge conflicts in this pull request

[Copilot]

Copilot resolve the merge conflicts in this pull request

Merged origin/main into this branch in 2c41035.

There was one conflict, in olive/passes/onnx/mobius_model_builder.py, where main's #2634 added revision from model.get_load_kwargs() while this branch had changed that block to also handle DiffusersModelHandler. Resolved by keeping both: a single load_kwargs lookup that uses get_load_kwargs() for HfModelHandler and .load_kwargs for DiffusersModelHandler (which has no get_load_kwargs()), then reading revision and trust_remote_code from it.

test/passes/onnx/test_mobius_model_builder.py passes (35 passed, 1 skipped), and lintrunner is clean.

[Ti-Tai Wang (titaiwangms)]

Updated review team verdict: Request changes

The latest commits resolve most of the original findings, including override precedence, loss of skip-pattern coverage, stale checkpoint detection, CLI/footprint path updates, explicit rejection of unsupported multi-component PyTorch optimization, guarded embedding retying, resolved sibling-output detection, and Diffusers trust_remote_code forwarding.

The following blockers remain:

Severity	Location	Finding
Critical	olive/workflows/run/builds.py:90-103, olive/workflows/run/hf_component_assembly.py:556-596	When assemble_components is omitted, sibling build outputs still enable assembly automatically. If their shared parent already contains a checkpoint, publication replaces its config/index and removes existing model.safetensors, shard, or PyTorch checkpoint files. The rollback copies are stored inside the temporary assembly directory and deleted after a successful publication, making the data loss permanent. Assembly should be explicit opt-in, or replacement should require an Olive ownership manifest plus an explicit replace option. Cleanup must only remove files owned by a previous Olive assembly.
Major	olive/workflows/run/hf_component_assembly.py:187-210,325-334	Compatibility validation compares model configuration fields that quantization legitimately changes per build, particularly tie_word_embeddings. Decoder and vision builds can therefore both succeed and then fail assembly as incompatible. Conversely, the merged quantization config still computes tie_word_embeddings with any(...), which can disagree with the tensors present in the assembled checkpoint. Exclude pass-mutated fields from equality checks and reconstruct and validate tying semantics explicitly from the final checkpoint.
Major	olive/workflows/run/hf_component_assembly.py:547-619	Publication still lacks a single crash-consistent commit boundary. The config and index are made public before auxiliary files, model configs, footprints, and source cleanup are completed. A failure after index publication can leave a publicly visible but partially finalized assembly, and some source build artifacts may already be deleted. Complete all fallible writes before the commit point; make post-commit cleanup ownership-scoped, idempotent, and non-fatal.
Major	olive/workflows/run/hf_component_assembly.py:471-482,631-634	The assembly lock is acquired before checking whether the workflow is eligible for assembly. Unrelated multi-build workflows sharing an output parent can therefore block or fail each other even when neither would assemble. A killed process also leaves a permanent lock because the recorded owner PID is never checked. Perform the applicability check first and use an OS-backed advisory lock or reliable stale-owner recovery.
Major	olive/cli/run.py:131-143	CLI output infers assembly whenever the actual model path is outside the configured build output. Packaging and cache-backed outputs can satisfy the same condition, causing ordinary builds to be reported as component artifacts with an assembled model in the cache. Detect assembly using model_attributes.assembled_components or another explicit marker rather than path shape.
Major	olive/workflows/run/hf_component_assembly.py:548-554	Shards are copied directly to their final destinations, but a destination is added to the rollback list only after copy2 succeeds. ENOSPC or interruption can therefore leave large partial shards that rollback does not remove. Copy each shard to a unique temporary destination and atomically rename it after the copy completes.

Documentation

The checkpoint tree still shows fixed names such as model-unoptimized-00001.safetensors and decoder/model-00001.safetensors, while the implementation now inserts a per-assembly generation identifier. The example should show that identifier and explain why filenames are generation-specific.

[Ti-Tai Wang (titaiwangms)]

Updated review team verdict: Request changes

The latest update resolves four of the six previous blockers:

• Assembly is now explicit opt-in and disabled by default.
• Applicability is checked before taking an OS-backed advisory lock.
• CLI output uses the explicit assembled_components marker.
• Shards are copied through temporary files and atomically renamed.
• Documentation now shows generation-specific shard names.

However, the rewritten tying and ownership/publication logic introduces one reproducible critical regression and leaves several hardening gaps.

Severity	Location	Finding
Critical	olive/workflows/run/hf_component_assembly.py:335-355,707-711	When no build quantizes embeddings or lm_head—the normal default—tying_configs is empty, so the merged tie_word_embeddings becomes false and overwrites the source model config. For an originally tied model whose checkpoint has no separate lm_head.weight, reload reports the head as missing and randomly initializes it. This was reproduced with Transformers 5.15. Preserve the source tying state when assembly did not change embedding/head storage, and validate it against the final checkpoint keys.
Major	olive/workflows/run/hf_component_assembly.py:195-205,707-711	For composite/VLM configurations such as Gemma, the effective tying setting may be text_config.tie_word_embeddings. Normalization recursively ignores that field, while publication updates only the top-level key. Nested tying therefore remains dependent on the first build's config. Reconcile and publish the effective submodel configuration rather than only the root key.
Major	olive/workflows/run/hf_component_assembly.py:559-581	Ownership preflight does not include pytorch_model-*.bin, so an existing sharded PyTorch checkpoint may not be recognized as foreign checkpoint content. Include sharded .bin files and test this layout.
Major	olive/workflows/run/hf_component_assembly.py:559-595	Cleanup trusts metadata.olive_assembly.files as its own proof of ownership. It does not require those paths to equal the old weight_map, match the recorded generation/naming convention, or even be safetensors shards. A crafted or colliding index can authorize deletion of arbitrary descendant files under the parent. Validate the manifest strictly against the old index and Olive generation format, or retain old generations instead of deleting from self-asserted metadata.
Major	olive/workflows/run/hf_component_assembly.py:630-658	Tokenizer files, custom Python files, component manifests, and model_config.json can be overwritten before the config/index commit, but they are absent from both the ownership manifest and rollback. A failed publication can therefore preserve the old checkpoint while irreversibly replacing its ancillary assets. Track every published file in the ownership transaction and back up or reject unowned destinations.
Major	olive/workflows/run/hf_component_assembly.py:648-658	The existing config.json backup still uses direct shutil.copy2. If ENOSPC leaves a partial backup, the exception handler sees that partial file and replaces the valid original config with it. Create the backup atomically and restore it only after recording successful completion.
Major	olive/workflows/run/hf_component_assembly.py:621-674	A process kill between config and index replacement can expose mismatched generations. A kill after publishing generation shards but before publishing or cleaning the index leaves files that the next run treats as foreign, permanently blocking retries until manual cleanup. Add a recoverable generation/journal protocol and reclaim unmistakably incomplete Olive generations while holding the lock.
Major	olive/workflows/run/hf_component_assembly.py:687-689	With explicit assemble_components: true, one ineligible build causes the entire assembly to silently return None; the workflow still succeeds and reports ordinary independent outputs. An explicit assembly request should either fail with the disqualifying reason or emit a clear warning.

Priority

The tying regression, nested VLM tying reconciliation, and unsafe partial-config restoration are merge blockers. Ancillary-file ownership and crash recovery should also be addressed before relying on assembly as a safe checkpoint publication mechanism. The explicit no-op behavior and comprehensive .bin preflight could be follow-ups only if their behavior is made explicit and documented.

[Ti-Tai Wang (titaiwangms)]

Updated review team verdict: Request changes

The latest commit fixes the original tied-model regression for flat configurations such as Llama and Qwen. However, the same random-lm_head corruption remains possible for composite Gemma/VLM configurations, so the tying issue is only partially resolved.

Severity	Location	Finding
Critical	olive/workflows/run/hf_component_assembly.py:209-230,684-689	_effective_tie_word_embeddings gives text_config.tie_word_embeddings precedence and then writes that value back to the root config. In Transformers 5.15, Gemma3's top-level model owns lm_head and uses the root tie_word_embeddings value. For root=true, text=false, the helper changes root to false; because the tied checkpoint has no independent lm_head.weight, reload reports it missing and randomly initializes the head. This was reproduced with a tiny Gemma3 configuration. When assembly did not change embedding/head storage, the safest fix is to leave all source tying fields unchanged rather than normalize their values.
Major	olive/workflows/run/hf_component_assembly.py:234-240	The compatibility check uses the same text-first helper, so it can miss root-level tying disagreement. For example, builds with root values true and false but matching text_config=true are accepted, and the published behavior remains dependent on build order. Compare the relevant root and nested fields without collapsing them through the helper.
Major	olive/workflows/run/hf_component_assembly.py:550-556,667	Automatic assembly remains non-reentrant. After the first successful run, the output root contains files and the second identical workflow fails only after all component builds have completed. A process crash that leaves partially published files creates the same unrecoverable state.
Major	olive/workflows/run/builds.py:84-93, olive/workflows/run/hf_component_assembly.py:662-663	The assembly destination accepts only top-level output_dir or engine.output_dir; it does not use the existing _default.output_dir path or the engine's resolved default. Such workflows complete their builds and then fail during automatic assembly. Use one canonical resolved output parent, or skip assembly rather than failing after successful builds.
Major	olive/workflows/run/hf_component_assembly.py:577-632	Clean-root and destination checks remain vulnerable to symlink redirection and check/use races. A component directory symlink can redirect publication outside the workflow root, and multiple individual renames provide no recovery after process termination. Reject symlinked path components and publish through a transaction that can be resumed or rolled back.
Major	olive/workflows/run/hf_component_assembly.py:639-647	Assembly reads from artifact.model_dir, but cleanup recursively deletes the hard-coded artifact.output_dir / "model". If those paths differ, cleanup can remove an unrelated directory while retaining the actual source checkpoint. Avoid deleting source builds, or verify and delete only the exact assembly-owned source path.
Major	olive/workflows/run/hf_component_assembly.py:217-274	Unowned tensors are compared only by shape and dtype. Builds can contain different shared tensor values with identical metadata; assembly silently selects the first build's copy, making the checkpoint depend on build order. Compare tensor contents or stable digests before assembly.

Tying fix guidance

• If no build quantizes lm_head or embeddings, do not rewrite any source-model tying field.
• If storage actually changes, derive and validate tying from the final checkpoint keys and the architecture's real input/output embedding paths rather than applying a generic root-versus-text_config precedence rule.

[Ti-Tai Wang (titaiwangms)]

Approved

The latest update resolves the previously blocking tying issues:

• Assembly no longer rewrites source tying fields when embedding/head storage is unchanged.
• Root and text_config.tie_word_embeddings values are compared independently across builds.
• Cleanup no longer removes output_dir/model unless it is the actual source model directory.

The remaining concerns are non-blocking and can be handled as follow-ups:

• Use the same canonical parent for automatic assembly when a workflow configures only builds._default.output_dir; currently this can fail after successful builds because assembly requires a top-level engine.output_dir.
• Define rerun and crash-recovery behavior for an already populated assembly root.
• Harden component publication against symlink redirection and check/use races.
• Optionally compare unowned tensor contents, rather than only shape and dtype, when passes may mutate tensors outside their selected component.
• Generalize tying-field discovery if models with deeper nested configurations such as thinker_config.text_config need assembly support.
• Validate storage-changing tying metadata against the final input/output embedding tensor keys.

These should be tracked separately so they do not expand the scope of the Gemma4 tying fix further.