feat: add components_to_skip to OnnxBlockWiseRtnQuantization

[titaiwangms]

Summary

Add optional components_to_skip parameter (list of component names) to the OnnxBlockWiseRtnQuantization pass. When a composite model component's name appears in this list, its model files are copied to the output path unchanged instead of being quantized. Non-listed components are quantized as usual. Default None quantizes all components (fully backward compatible).

Depends on #2456 (MobiusBuilder components_to_export). This PR's diff is scoped to RTN changes only — MobiusBuilder changes are in the base branch.

Motivation

When quantizing large multi-component VLMs (e.g. Mistral-3 with decoder + vision_encoder + embedding), users may want to skip RTN quantization for specific components such as the embedding model — which may need to stay in higher precision (e.g. due to a GatherBlockQuantized bug with certain inputs).

Changes

• olive/passes/onnx/rtn_quantization.py:

  • from typing import Optional (removed unused List)
  • Add components_to_skip: list PassConfigParam (default None)
  • Override run() to intercept CompositeModelHandler processing: copy skipped components via shutil.copytree (handles file vs directory model_path), quantize the rest via recursive self.run()
  • onnx_file_name fallback: use getattr(component_model, "onnx_file_name", None) or "model.onnx" to handle handlers constructed without an explicit file name
  • Mirror base class _initialize() guard and post-processing (model_attributes propagation + _carry_forward_additional_files per component)
  • Non-composite (single) models are unaffected

• test/passes/onnx/test_rtn_quantization.py:

  • New TestRTNQuantizationComponentsToSkip class with 5 tests
  • Skipped component passes through unchanged (MatMul present, no MatMulNBits)
  • Non-skipped components are quantized
  • components_to_skip=None quantizes all (backward compat)
  • Single-model (non-composite) unaffected
  • Unknown component name logs a warning (does not raise)
  • Removed redundant AcceleratorSpec reimport; wrapped long lines (120-char limit); split compound assertions

Testing

All 16 tests pass:

python3 -m pytest test/passes/onnx/test_rtn_quantization.py -v

[Copilot]

Pull request overview

This PR extends Olive’s ONNX tooling to better support multi-component (composite) models by allowing selective processing of components during export and RTN quantization.

Changes:

• Add components_to_skip to OnnxBlockWiseRtnQuantization to copy selected composite components through unmodified while quantizing the rest.
• Add tests validating components_to_skip behavior on composite and non-composite models, including a warning on unknown component names.
• Add components_to_export to MobiusBuilder to export only a subset of components from multi-component mobius packages, with validation + new unit tests.

Reviewed changes

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

File	Description
olive/passes/onnx/rtn_quantization.py	Adds components_to_skip config and custom composite-handling run() logic to skip quantization for selected components.
test/passes/onnx/test_rtn_quantization.py	Adds unit tests covering components_to_skip behavior and logging.
olive/passes/onnx/mobius_model_builder.py	Adds components_to_export config and filters which mobius package components are saved/returned.
test/passes/onnx/test_mobius_model_builder.py	Adds unit tests for components_to_export filtering/validation and updates the fake package save behavior.

[Copilot]

Pull request overview

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

[titaiwangms]

Review (synthesized from 4 reviewers: readability / code / critical / deep)

Convergent findings flagged with [N reviewers]. Severity follows worst-case among reviewers.

Major

1. copytree(src.parent, ...) can copy unrelated siblings [code]

In the skip branch of run(): when component_model.model_path points to an .onnx file (not a directory), the code does src_dir = src.parent and copies the whole parent. If multiple components share a parent directory (flat composite layout), the skip silently bloats/corrupts the output by dragging in siblings. Fix: copy the ONNX file + its .data sidecar + listed additional_files explicitly, not the entire parent dir.

2. Path comparison without .resolve() risks deleting the source [critical]

if src_dir != component_output_path:
    if component_output_path.exists():
        shutil.rmtree(str(component_output_path))
    shutil.copytree(str(src_dir), str(component_output_path))

Lexical Path compare. If absolute and relative forms point to the same dir (very plausible in stacked passes, symlinked workspaces, or case-insensitive FS), the inequality check passes, then rmtree(component_output_path) deletes the actual source before copytree reads from it. Fix: compare .resolve()'d paths, and reject overlapping src/dst.

3. Non-atomic rmtree+copytree on large external-data dirs [critical]

A crash, signal, or disk-full between rmtree and the end of copytree leaves the output unusable. Fix: stage to a sibling dir, validate, then atomic-rename.

4. SKIP branch silently lies about handler type [deep]

output_component = ONNXModelHandler(
    model_path=str(component_output_path),
    onnx_file_name=onnx_file_name,
    model_attributes=component_model.model_attributes,
)

It unconditionally constructs ONNXModelHandler, but CompositeModelHandler permits any OliveModelHandler (nested composite, distributed ONNX, PyTorch). Latent today since Mobius produces only ONNX components, but worth at least an isinstance(component_model, ONNXModelHandler) assert with a clear error, plus a docstring constraint.

5. onnx_file_name fallback hardcodes "model.onnx" [code]

If src.is_file() and the input handler has no explicit onnx_file_name, the output handler points to model.onnx even if the source filename was different (e.g. decoder.onnx). Fix:

if src.is_file():
    onnx_file_name = getattr(component_model, "onnx_file_name", None) or src.name
else:
    onnx_file_name = getattr(component_model, "onnx_file_name", None) or "model.onnx"

6. _initialized guard duplicates a private base-class protocol [readability]

The override copies the _initialized / _initialize() dance from Pass.run() (olive_pass.py:209-211). The comment is honest about it, but the coupling is fragile — if Pass.run() ever adds a second guard or renames the attribute, this silently diverges with no test catching the drift. Long-term: add _ensure_initialized() to the base class so overrides can call it without copying. Short-term: at least add a class-level comment that this assumes _initialize() is a no-op or idempotent.

Minor

• model.get_model_components() is called twice (once for the unknown-name warning, once for the loop). The property has side-effects merging parent attrs into children — idempotent today, but stylistically off and the base only iterates once. Materialize once: pairs = list(model.get_model_components()).
• Inconsistent strictness with feat: add components_to_export filter to MobiusBuilder pass #2456: unknown component name → raise in feat: add components_to_export filter to MobiusBuilder pass #2456, warn in feat: add components_to_skip to OnnxBlockWiseRtnQuantization #2457. Pick one and document.
• Path(output_model_path).with_suffix("") strips dotted-dir names like model.v1 (inherited footgun from base, not a regression).
• Docstring should explain why _accepts_composite_model = True wasn't used instead — short answer: _run_for_config doesn't receive the component name, so a per-component skip check has no name to key on. Worth saying explicitly.
• Add regression tests for: (a) non-model.onnx filenames, (b) file-based components in shared parent dirs.

Praise

• Recursing via self.run(component_model, ...) for the non-skipped branch correctly inherits all base-class per-component behavior — clean design that avoids re-implementing _run_for_config invocation.
• Deep-trace confirmed that model_attributes propagation and _carry_forward_additional_files are correctly preserved with no double-carry on additional files (the helper checks if not output_filepath.exists() before copying).
• _initialize() idempotency: verified safe — base _initialize() is a no-op and not overridden on this pass; the recursive super().run() call will see _initialized = True and skip.

Note

This PR is stacked on #2456 — please rebase / re-target to main once #2456 lands.

Posted by GitHub Copilot CLI on behalf of @titaiwangms.