Skip to content

Add multi build support - #2480

Merged
Xiaoyu (xiaoyu-work) merged 58 commits into
mainfrom
xiaoyu/builds-schema
Aug 21, 2026
Merged

Add multi build support#2480
Xiaoyu (xiaoyu-work) merged 58 commits into
mainfrom
xiaoyu/builds-schema

Conversation

@xiaoyu-work

@xiaoyu-work Xiaoyu (xiaoyu-work) commented May 29, 2026

Copy link
Copy Markdown
Collaborator

This pull request introduces a new, flexible "builds" workflow to the Olive engine, enabling multiple independent build pipelines within a single run configuration. It adds support for per-build defaults, validation, and selective execution on model components, primarily targeting composite models. The changes also include robust schema validation and improved modularity for configuring builds.

Key changes include:

Builds Workflow and Configuration:

  • Added a new builds field to RunConfig, allowing users to define multiple named build pipelines, each with its own pipeline, component selection, and system/evaluator overrides. A special _default key enables partial defaults to be merged into sibling builds. (olive/engine/config.py, olive/workflows/run/config.py, olive/workflows/run/run.py) [1] [2] [3]

  • Introduced BuildConfigPartial and BuildConfig schemas for partial and full build configurations, and a merge_build_default function for merging defaults into builds. (olive/engine/config.py)

Validation and Reference Resolution:

  • Added pre- and post-validation to RunConfig to ensure build defaults are correctly merged and that all build references (passes, systems, evaluators) resolve to known entries. (olive/workflows/run/config.py) [1] [2]

Component Selection for Composite Models:

  • Implemented select_components methods in both ModelConfig and CompositeModelHandler to allow builds to operate on specific named components of a composite model, returning either a single component or a sub-composite as needed. (olive/model/config/model_config.py, olive/model/handler/composite.py) [1] [2]

Workflow Execution Logic:

  • Refactored the main run logic to dispatch to a new _run_builds function when builds are present, running each build as an independent workflow with its own engine, pipeline, and input model slice. Includes helper functions for validation, engine config construction, and reference resolution. (olive/workflows/run/run.py) [1] [2]

These improvements make the Olive engine significantly more flexible and modular, supporting advanced workflows for multi-component and multi-pipeline builds

Xiaoyu (xiaoyu-work) and others added 2 commits May 28, 2026 15:27
Introduce a top-level �uilds section on RunConfig that lets users declare
multiple independent execution units (pipelines x devices x components) in one
workflow config.

* Add BuildConfigPartial / BuildConfig and a merge_build_default helper in
  olive/engine/config.py. _default lives inside �uilds as a sentinel key
  whose partial fields are merged into every sibling build with full-replace
  semantics (lists are not deep-merged).
* Add �uilds: dict[str, BuildConfig] to RunConfig with an
  �xpand_build_defaults before-validator that pops _default and merges it
  into siblings, plus a �alidate_builds_references after-validator that
  checks pipeline/host/target/evaluator string refs resolve to known entries.
* Schema-only change: the engine runner does not yet act on �uilds. Existing
  workflows without �uilds keep their current behavior.
* Add 8 unit tests in test/workflows/test_run_config_builds.py covering the
  merge, override, full-replace, missing-field, invalid-ref, absent-builds and
  empty-default cases.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Execute the �uilds schema added in Phase 1.

* Add CompositeModelHandler.select_components(names) that returns the
  unwrapped child handler when one name is given and a sliced
  CompositeModelHandler otherwise. Unknown names raise a clear error.
* Add ModelConfig.select_components(names) so the runner can slice a
  composite input config without materializing the full handler.
* Add a builds-aware execution branch in olive/workflows/run/run.py. When
  �uilds is non-empty, the runner: validates components against the
  composite input model, then loops over builds. For each build it builds a
  per-build engine config (host/target/evaluator/search_strategy overrides
  resolved against systems/evaluators), a per-build pipeline subset from
  passes in the order declared by pipeline, the per-build accelerator
  spec, and calls engine.run with build.output_dir. Returns
  dict[build_name -> WorkflowOutput]. The no-builds path is unchanged and
  still returns a single WorkflowOutput.
* Tests:
  - 7 new composite handler / ModelConfig select_components cases in
    test/model/test_composite_model.py.
  - 7 new runner smoke tests in test/workflows/test_run_builds.py with
    mocked Engine.run covering: no-builds backward compat, multi-build
    dispatch, pipeline-subset ordering, per-build output_dir, host/target
    override, non-composite + components error, unknown component error.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
@xiaoyu-work Xiaoyu (xiaoyu-work) changed the title Xiaoyu/builds schema Add multi build support May 29, 2026
Xiaoyu (xiaoyu-work) and others added 14 commits June 2, 2026 12:56
…us HfModel

Support the two component-discovery paths from the multi-component design:

- Flow A Option 2 (two steps): load a Mobius export directory as a
  CompositeModel, using per-component subfolder names as component names.
  Adds discover_onnx_components() and directory auto-discovery in
  CompositeModelHandler and ModelConfig.get_components/select_components.
- Flow B (optimize then export): resolve an HfModel's components by querying
  Mobius (olive/common/mobius_utils.inspect_components, lazy import).
  HfModel.get_components returns Mobius component names; select_components
  tags the chosen component's submodule path in model_attributes for
  PyTorch-stage per-component passes.

No 'input' build dependency is used. Build component validation updated for
both sources.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
…exports

When mobius exports a multi-component model, log each component's name and
its ONNX file path so the export layout is visible in the run log.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
mobius.inspect_components returns frozen ComponentInfo dataclasses, but Olive's coercion only handled its own ComponentInfo or a plain dict and crashed calling .get() on a mobius object. Broaden ComponentInfo.coerce (renamed from from_dict) to also accept duck-typed objects exposing name/kind/source_path. Add test/common/test_mobius_utils.py covering the object, dict, passthrough, and missing-mobius paths that the existing inspect_components mocks never exercised.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
The autouse _stub_mobius_module fixture guarded on whether mobius was already in sys.modules, which is False when real mobius is installed but not yet imported at fixture setup. It then injected a non-package stub that shadowed real mobius, breaking test_write_genai_config_requires_real_mobius (imports mobius.integrations). This was masked in CI where mobius is absent and the test is skipped. Guard on _HAS_REAL_MOBIUS instead so the fixture is a true no-op when mobius is installed. Also drop the log-capture assertion from test_multi_component_returns_composite_handler (logging stays in the pass).

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
builds.components could only target CompositeModel (directory) and HfModel (via mobius); a DiffusersModel input fell through to get_components()/select_components() returning None, so the design's per-component diffusion example failed validation. Add an optional components filter to DiffusersModelHandler (restricts get_exportable_components to a subset in canonical variant order) and resolve/select diffusion components in ModelConfig. select_components scopes the handler so each build's conversion emits just that component's ONNX, with later passes auto-mapping over the single-component composite. Variant detection only reads config files, so resolution stays cheap at validation time.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Two runnable Flow A recipes plus a README documenting the two-step flow (CLI Mobius export, then an Olive config whose builds optimize each exported component): SD3 (transformer + VAE encoder/decoder; Mobius skips the CLIP/T5 text encoders) and Qwen3-VL-2B-Instruct (decoder/vision_encoder/embedding, the exact names Mobius produces). Both validate against RunConfig.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
@xiaoyu-work

Copy link
Copy Markdown
Collaborator Author

Copilot resolve the merge conflicts in this pull request

…i/base.py and passes/onnx/mobius_model_builder.py

Copilot AI commented Jun 24, 2026

Copy link
Copy Markdown
Contributor

Copilot resolve the merge conflicts in this pull request

Merge conflicts resolved in commit 97da0323. Two files had conflicts:

  • olive/cli/base.py: Combined the builds-dict output handling (from this branch) with the save_discrepancy_check_results call (from main). The discrepancy check is now skipped when workflow_output is a dict (i.e., a builds run), since that function only handles a single WorkflowOutput.

  • olive/passes/onnx/mobius_model_builder.py: Kept the component_paths.append logging added in this branch, and adopted main's improvements: the component_additional_files variable rename with its descriptive comment, and the shared "additional_files" key on the composite handler to carry genai sidecar files to the package root.

Now that mobius exports text encoders (CLIP-L, CLIP-G, T5-XXL),
update the inference script to use ONNX Runtime for all components
instead of PyTorch text encoders and VAE.

Only tokenizers and the scheduler (pure logic, no neural network)
remain as Python dependencies.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

@github-advanced-security github-advanced-security AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

lintrunner found more than 20 potential problems in the proposed changes. Check the Files changed tab for more details.

Comment thread test/passes/pytorch/test_quant_utils.py Fixed
Comment thread test/passes/pytorch/test_quant_utils.py Fixed
Comment thread test/passes/pytorch/test_quant_utils.py Fixed
Comment thread test/passes/pytorch/test_quant_utils.py Fixed
Comment thread test/passes/pytorch/test_quant_utils.py Fixed
Comment thread test/common/test_hf_wrapper.py Fixed
Comment thread test/common/test_hf_wrapper.py Fixed
Comment thread multi_comp_recipe/sd3_inference.py Fixed
Comment thread multi_comp_recipe/sd3_inference.py Fixed
Comment thread multi_comp_recipe/sd3_inference.py Fixed
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

Copilot-Session: 86248e38-d706-4b53-875d-8eb4c7c69b49
@xiaoyu-work

Copy link
Copy Markdown
Collaborator Author

Copilot resolve the merge conflicts in this pull request

Co-authored-by: xiaoyu-work <85524621+xiaoyu-work@users.noreply.github.com>

Copilot AI commented Aug 18, 2026

Copy link
Copy Markdown
Contributor

Copilot resolve the merge conflicts in this pull request

Resolved in commit 81aaa00e by merging origin/main into this branch and fixing the conflicts in olive/passes/pytorch/quant_utils.py and test/passes/pytorch/test_quant_utils.py.

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Pull request overview

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

Suppressed comments (4)

olive/passes/pytorch/quant_utils.py:330

  • The same subtree-saving problem affects decoder/embedding fallbacks. For example, selecting T5's shared embedding makes wrapper.model an nn.Embedding; finalize then accesses .config and .save_pretrained, so the newly added shared-embedding test cannot complete. Wrap the root model and let the component-path filter select only the requested weights.
    if component_role in {"decoder", "embedding"}:
        component_model = get_attr(root_model, slice_path) if slice_path else root_model
        config = getattr(component_model, "config", root_model.config)
        return _GenericComponentWrapper(component_model, config), (f"{slice_path}." if slice_path else "")

olive/passes/pytorch/quant_utils.py:606

  • skip_patterns contains root-relative modules_to_not_convert names, while iter_quant_targets sees names relative to a sliced component. Thus an exclusion such as vision.proj does not match local proj, and the explicitly excluded weight is quantized. Apply the pattern after converting full_name to root_name (while retaining identity-based skips).
            skip_patterns=skip_patterns,
            extra_skip_modules=excluded_attn_inputs,
        ):
            root_name = _root_module_name(full_name, name_prefix)
            # When the slice spans more than the component (multi-path components slice

olive/cli/base.py:346

  • Changing this elif to if makes --test mode fall through to the normal model message. In test mode the model is actually written under <output_path>/model (lines 315-320), but this branch prints Model is saved at <output_path>, which is misleading in addition to the test-report message. Preserve the mutually exclusive branch.
            if not workflow_output.has_output_model():

olive/workflows/run/builds.py:56

  • The new format is only recognized through this helper; RunConfig itself still has no builds field. Existing schema/validation consumers (docs/source/dump_schema.py and skills/olive/scripts/validate_config.py) continue to use RunConfig, so they omit or reject every valid multi-build configuration even though the runtime accepts it. Expose the multi-build schema through the public run-config model or update all public schema/validation entry points to use a common top-level parser/model.
def parse_run_config(
    run_config: Union[str, Path, dict],
) -> Union[RunConfig, OrderedDict[str, RunConfig]]:

Comment thread olive/passes/pytorch/quant_utils.py
Comment thread olive/workflows/run/run.py Outdated
Comment thread olive/passes/openvino/optimum_intel.py
Comment thread olive/common/quant/hf_utils.py Fixed
Comment thread olive/common/quant/hf_utils.py Fixed
Comment thread olive/common/quant/hf_utils.py Fixed
Comment thread olive/common/quant/hf_utils.py Fixed
Comment thread olive/common/quant/hf_utils.py Fixed
Comment thread olive/common/quant/hf_utils.py Fixed
Comment thread olive/common/quant/hf_utils.py Fixed
Comment thread olive/common/quant/hf_utils.py Fixed
@titaiwangms

Copy link
Copy Markdown
Contributor

Full-team review: PR #2480 "Add multi build support"

Ran a 5-reviewer pass (readability / correctness / adversarial / spec-deep / cross-module integration) against the current diff (81aaa00e on xiaoyu/builds-schema), with tests executed and several claims verified by live repro (probes discarded, not committed). Summary below, ordered by severity. Happy to share repro scripts on request.

Critical

  1. olive/common/quant/hf_utils.py:264-266QuantEmbedding/QuantLinear are never imported → guaranteed NameError on every quantized-model load.

    if isinstance(model.get_input_embeddings(), (QuantEmbedding, QuantLinear)) or isinstance(
        model.get_output_embeddings(), (QuantEmbedding, QuantLinear)
    ):

    Verified by AST (neither name is bound in the module). The source module (olive/common/quant/nn.py) was deleted upstream in 96fb44d9 (Extend PyTorch RTN weight quantization to MoE experts #2584); the surviving classes are QuantLinearNbit/QuantEmbeddingNbit in olive/common/hf/quant.py. Same stale import appears in test/passes/pytorch/test_quant_utils.py:24 (from olive.common.quant.nn import QuantEmbedding, QuantLinear), so the entire new component-quant test file fails at collection — none of the new tests for this feature actually ran on this branch. This needs fixing (and the full new suite re-run) before anything else below can be meaningfully re-verified.

  2. quant_utils.py:234-269 (_GenericComponentWrapper.maybe_untie_word_embeddings) — untying silently rescales T5-family logits by √d_model.
    For T5/UMT5/LongT5/SwitchTransformers/UDOP/Pop2Piano, tie_word_embeddings isn't just a storage hint — the forward pass does sequence_output * (model_dim ** -0.5) gated on that flag (e.g. modeling_umt5.py:1052-1055). Flipping it to False after cloning weights (as this method does for every alias) changes the whole model's output, not just the sliced component. Reproduced on a tiny UMT5 (no quantization involved at all, just the untie): output ratio == √d_model exactly. Plain T5 already migrated off this overload (scale_decoder_outputs), but the other five architectures still use it.

  3. quant_utils.py:643-653 + finalize — persisted modules_to_not_convert/overrides use root-relative names, but a non-decoder component's saved checkpoint is slice-relative.
    When _component_model_wrapper returns a non-empty name_prefix (e.g. an encoder/vision component), finalize saves the sliced sub-module, but the quant config that gets written into it is built from _root_module_name(...) (root-relative). At load time match_skip does substring matching against the slice-relative names in the checkpoint, so root-relative patterns never match — reproduced end-to-end: float modules that should have been skipped get a placeholder QuantTensor installed instead (weight UNEXPECTED / weight_qweight MISSING), and 8-bit overrides silently degrade to the 4-bit default. Related: finalize also assumes wrapper.model is a PreTrainedModel (.config, .save_pretrained) — for a bare leaf-module slice (nn.Linear/nn.Embedding, exactly the PR's own vision/T5-embedding test fixtures) this raises AttributeError once Update Perf Tuning Build  #1 is fixed and the tests can actually run.

Major

  1. Common-ancestor slicing leaks other components' float weights into a component's saved artifact (quant_utils.py slice logic + finalize). Target-selection correctly isolates siblings during quantization (_is_in_component's dotted-prefix check is solid), but when saving, wrapper.model is the common ancestor of a multi-path component, which can also contain sibling components' modules — so a build that selects only "vision" can ship the full text backbone's float weights alongside it. Combined with OLive Web App #3, those extra weights are also misclassified at load.

  2. T5-style "embedding" component quantization is a no-op at inference. After untying, only the wrapper's own module (shared) gets quantized, but each stack's forward reads its own encoder.embed_tokens/decoder.embed_tokens (the fresh clones), never shared — so the "quantized" embedding component is still float at inference despite being reported as quantized. The new test for this (test_finalize_t5_shared_embedding_preserves_float_aliases) currently canonizes this behavior; worth confirming intent.

  3. embeds/lm_head gates resolve against the root model's get_input_embeddings()/get_output_embeddings(), not the selected component. For components whose embedding/head isn't the root's own (exactly the non-standard leaf names this PR added — codec_embedding, text_embedding, tok_embeddings, proj_out, codec_head, output_projection), embeds: true silently quantizes nothing, and lm_head: false silently fails to protect the component's actual head.

  4. all_tied_weights_keys filtering in hf_utils.py:247-262 appears to be dead code for the case it targets. install_quant_tensor_param keeps an nn.Parameter-typed object in place (now wrapping a QuantTensor), so the has_parameter(...) guard this PR adds never actually filters anything — reproduced: get_parameter() succeeds on the quantized param, so nothing is dropped from the mapping.

  5. Thread-based parallel build execution violates several existing process-global invariants. rotate.py/slicegpt.py call torch.manual_seed (process-global RNG), QAIRT passes mutate QAIRT_LOG_LEVEL and never restore it, and file logging (enable_filelog) adds a handler to the single global olive logger that's never removed — with N concurrent builds, all subsequent log records go to all N cache log files, and handlers/fds leak across the run. Two builds with different configs can observe each other's RNG seed or logging state.

  6. Unbounded concurrency + no per-build GPU allocation. max_workers=len(build_configs) lets a run.json spawn arbitrarily many concurrent builds, each of which loads the full root HF model before slicing, with accelerator selection always picking the first configured accelerator — straightforward path to concurrent OOM on CPU/GPU for any run with more than a couple of builds.

  7. Rebase/merge-conflict risk with the already-landed PR Support Qwen3.5/3.6-MoE VL checkpoints in PyTorch-side quantization #2630 (quantize_vision flag threaded through iter_quant_targets/OliveHfQuantizationConfig/get_quantizer_config in this same file). This PR's new _iter_component_quant_targets wrapper calls iter_quant_targets(...) with an explicit argument list that doesn't include quantize_vision — once both land, component-scoped (builds-driven) quantization would silently drop vision-tower exclusion behavior. Will need quantize_vision=getattr(quant_cfg, "quantize_vision", False) threaded through when rebasing.

  8. docs/source/dump_schema.py still dumps the plain RunConfig.schema_json(), which has no builds property — generated schema.json will reject the builds key for anyone relying on IDE/editor schema validation.

  9. select_components() sub-composites reuse the parent's model_path and inherited attributes (including no_flatten) (composite.py:142-147). olive/cache.py's no_flatten packaging path (shutil.copytree(source_path, ...)) then copies the entire parent tree, not just the selected subset — a 2-of-5-component build ships all 5 components' files, and the resulting model_path still points at the (possibly-cleaned-up) parent cache dir.

  10. _path_with_leaf is first-match / order-dependent (quant_utils.py:165-169) — for a component spanning two sub-trees that both happen to have a layers leaf (plausible for VLM vision towers), which one gets declared as the decoder's LAYERS depends on incoming path order, with no error on ambiguity; layerwise GPTQ/KQuant would then silently calibrate the wrong stack.

  11. prepare_model is a ~200-line function doing 6+ distinct jobs (config validation, wrapper resolution, lm_head/embeds detection with multiple fallback heuristics, quant-target iteration, merge-with-existing-state, modules_to_not_convert pruning). Recommend extracting _resolve_component_wrapper, _detect_lm_head_and_embeds, _merge_with_existing_quant_config as named helpers — several of these are already comment-delimited blocks that map cleanly to functions.

  12. Duplicated/slightly-inconsistent "head"/"embedding" leaf-name sets across _root_component_model_wrapper (quant_utils.py:201,214) and prepare_model (:549,559,574) — e.g. "output" present in one head-name set but not the other, with nothing documenting whether the divergence is intentional. Worth extracting to shared module-level constants.

Minor

  • A selected component with zero quantizable targets currently "succeeds" as a silent no-op quantization (only the embeds=True-with-no-embeddings case is guarded).
  • modules_to_not_convert entries are full dotted paths but matched via substring (patterns.py), so short top-level component names can over-skip unrelated modules that merely contain that substring, and the field can bloat config.json with thousands of entries on real models.
  • composite.py:139-141's single-component select_components() branch mutates child.model_attributes in place on a handler still owned by the parent's _model_components list — the model_components property already performs the same merge on read, making this both redundant and a shared-object side effect.
  • olive/workflows/run/run.py's _run_builds_in_parallel collects errors as dict[str, list[Exception]], but each build can only ever raise once — a plain dict[str, Exception] would avoid implying multi-raise semantics that don't exist.
  • ModelWrapper's new olive_root_model/olive_component_path/olive_component_role attributes (wrapper.py:408-410) have no inline comment explaining the olive_ prefix, inconsistent with the rest of the class's unprefixed attributes.
  • Build name regex (BUILD_NAME_PATTERN) correctly blocks path traversal (.., slashes) — verified; flagged only as an open question whether Windows reserved device names / trailing-dot normalization need consideration (not verified, no Windows env available).

Open questions (not verified, need author input)

  • For non-decoder components, is the intended saved artifact the slice (current finalize behavior) or the full model with only that component quantized (decoder's behavior)? OLive Web App #3/Notebook Update #4/Doc Terms Unified; Bug Fixed #7 all stem from this being inconsistent between paths.
  • existing_modules_to_not_convert merge logic drops a user-specified exclusion as soon as any newly-quantized name contains it as a substring (reverse-direction match) — losing the exclusion permanently across re-runs. Intended?
  • Whether _GenericComponentWrapper.maybe_untie_word_embeddings should scope its alias-splitting to aliases that actually intersect the selected component's source_paths, rather than always operating on the whole root model.

Praise

  • olive/workflows/run/builds.py is clean, small, single-purpose functions (_validate_build_host, _get_build_write_dirs, _paths_overlap, _set_engine_value); the pairwise write-dir overlap check is easy to follow.
  • _is_in_component's dotted-prefix check correctly avoids the classic model.language_model vs. model.language_model_2 false-positive.
  • _component_slice_path's greatest-common-ancestor computation is done segment-wise, not by naive string prefix — handles ["model.lang", "model.language"] correctly.
  • _root_component_model_wrapper sets LAYERS/LM_HEAD/EMBEDDINGS as instance attributes rather than mutating the shared class-level dicts.
  • _validate_component_source_paths and prepare_model's "no source paths for a named component" check both fail closed with actionable error messages.
  • Clean Union[WorkflowOutput, dict[str, WorkflowOutput]] adaptation across mcp/src/olive_mcp/worker.py and olive/cli/api.py for the new multi-build return shape.

Net: the workflow-level builds mechanism (builds.py, cache isolation design, write-dir validation) looks solid, but the PyTorch component-scoped quantization slicing in quant_utils.py has several verified correctness bugs (some via live repro) that would need fixing before this is safe to merge, plus the immediate blocker that the new test suite doesn't currently execute (ModuleNotFoundError) due to a stale import from a module deleted in a separate already-merged PR.

@xiaoyu-work

Copy link
Copy Markdown
Collaborator Author

Copilot resolve the merge conflicts in this pull request

…ntLinear)

Co-authored-by: xiaoyu-work <85524621+xiaoyu-work@users.noreply.github.com>
# Conflicts:
#	olive/common/hf/wrapper.py
#	olive/model/handler/mixin/hf.py
#	olive/passes/pytorch/quant_utils.py

Co-authored-by: xiaoyu-work <85524621+xiaoyu-work@users.noreply.github.com>

Copilot AI commented Aug 18, 2026

Copy link
Copy Markdown
Contributor

Copilot resolve the merge conflicts in this pull request

Resolved. Merged origin/main into this branch, fixing conflicts in olive/common/hf/wrapper.py, olive/model/handler/mixin/hf.py, and olive/passes/pytorch/quant_utils.py.

@titaiwangms

Copy link
Copy Markdown
Contributor

Review: no Critical, but 6 Major findings

Multi-model review pass over the full diff. Every claim below was verified against the diff before posting (a couple of reviewer-generated "criticals" about mcp/src/olive_mcp/worker.py were false positives — the diff already handles the dict result and already switched to parse_run_config — so they were dropped).

Major (worth fixing before merge)

1. olive/workflows/run/builds.py — the writable-dir overlap check is buggy and non-deterministic

artifact_dir = output_dir.parent if output_dir.suffix and not output_dir.is_dir() else output_dir

BUILD_NAME_PATTERN explicitly allows ., so a build named llama.q4 produces output/llama.q4, which has a suffix and therefore collapses to output — overlapping with every other build and failing with "overlapping writable directories". Worse, not output_dir.is_dir() makes the result depend on whether the directory already exists, so the same config fails on the first run and passes on the second. Validation should be a pure function of the config. Suggested fix: drop the suffix heuristic entirely — output_dir is documented as a directory.

2. _default.output_dir is a declared field that can never be used

BuildConfigPartial.output_dir is merged shallowly, so the literal value is copied verbatim into every sibling build; expand_builds then assigns the identical engine.output_dir to all of them and finding #1's check immediately rejects the config — with an error that doesn't point at _default. Either resolve it as a parent (<default>/<build_name>) or reject output_dir in BuildConfigPartial. The same shallow-merge logic is also duplicated in olive/cli/run.py::_print_build_outputs, which would print the same directory for every build.

3. olive/model/handler/composite.py::select_components mutates the source composite

For a single selection:

child = selected[0]
child.model_attributes = {**(self.model_attributes or {}), **(child.model_attributes or {})}
return child

child is a reference into self._model_components, so merely selecting a component permanently mutates the original composite's child. Later selections and any other holder of the composite observe attributes that were never in the input model, and behavior becomes order-dependent. Clone (deepcopy / rebuild from config) before merging attributes. Separately, names is not deduplicated — select_components(["decoder", "decoder"]) builds a composite with two duplicate component names, and component names are treated as keys downstream.

4. olive/passes/pytorch/quant_utils.py — modules quantized in this run can land in modules_to_not_convert

unquantized_modules is derived from fresh_names, which comes from _iter_component_quant_targets(fresh_qcfg), but the loop that actually quantizes and populates new_qargs iterates the merged qcfg. When an on-disk config has a category flag True that this run sets False, and the module is not yet a QuantTensor, that module gets both quantized (packed buffers written) and listed in modules_to_not_convert. On reload, OliveHfQuantizer skips it, so it stays a plain nn.Linear while the checkpoint carries packed buffers → key mismatch. The invariant set(qcfg.modules_to_not_convert) & set(new_qargs) == {} should hold; compute the skip set after the new_qargs loop, from new_qargs.keys().

Secondary, same hunk: modules_to_not_convert uses substring matching (olive/common/quant/patterns.py::match_skip). Pre-existing entries are filtered with excluded in fresh_name, but the newly generated unquantized_modules are added unfiltered — so a skipped ...mlp.1 will silently un-convert a quantized ...mlp.10. This assignment is also unconditional, so it changes the saved quantization_config for every existing non-component run (potentially hundreds of entries for a VLM with quantize_vision=False). Worth a test pinning the non-component path; the new tests only cover the component path.

5. olive/cli/base.py:343elifif is an unrelated behavior change

In --test mode the CLI now prints both Test report saved at ... and either No output model produced or Model is saved at <output_path>. The latter is misleading in test mode, where the model is redirected to <output_path>/model. This looks unrelated to multi-build support; if intentional, please add a comment and print the model subdirectory. Note _run_workflow is shared by ~10 other CLI commands that do not have the new builds guards from cli/run.py, so a builds config reaching them would hit a bare AttributeError on dict.has_output_model().

6. New public config surface has no schema and no docs

docs/source/dump_schema.py generates the JSON schema from RunConfig.schema_json(), but builds and max_concurrent_builds are not RunConfig fields — they're handled by dict manipulation in builds.py. So the entire multi-build syntax is missing from IDE autocompletion, editors will flag builds as an invalid key, and docs/source/ has zero mention of builds / _default / max_concurrent_builds. Also skills/olive/scripts/validate_config.py:71 still calls RunConfig.parse_file_or_obj directly, so it reports on the wrong, un-expanded config for multi-build inputs. Consider elevating builds into a real pydantic model and pointing both the schema dump and that script at it.

Minor / worth noting

  • ContextVar cache isolation does not cross thread or process boundaries. While isolated_cache_env is active, set_cache_env stops writing os.environ["OLIVE_CACHE_DIR"]. Per the contextvars contract a new threading.Thread starts with a fresh context and multiprocessing spawn children inherit only os.environ, so anything reaching OliveCache.from_cache_env() from such a child (e.g. the spawn context in olive/passes/onnx/conversion.py) silently falls back to the default cache dir. Mostly latent today since max_concurrent_builds defaults to 1 — consider still setting os.environ in the serial case to preserve the pre-existing contract.
  • Parallel builds share process-global state. enable_filelog adds a handler to the single global Olive logger and never removes it, so with max_concurrent_builds > 1 every build's records land in every build's log file (and handlers leak after the run). Zip packaging via shutil.make_archive also stages in the process CWD under a shared archive name, so concurrent builds with packaging enabled can overwrite each other.
  • Config validation now performs I/O. expand_builds eagerly resolves components, which can inspect/download a HF repo — a meaningful contract change for the MCP validate_config entry point that now routes through parse_run_config.
  • olive/cli/run.py hardcodes the "_default" literal instead of importing BUILD_DEFAULT_KEY, and re-derives the output dir instead of reusing expand_builds' result.
  • _validate_build_write_dirs runs unconditionally and says "Parallel builds ..." even when max_concurrent_builds == 1.
  • Partial results are discarded when any build fails — the successful builds' WorkflowOutputs are thrown away with the RuntimeError.
  • Test coverage quietly reduced: test/common/test_hf_wrapper.py drops the LFM2 get_attention_outputs() assertions with no replacement, and test/passes/onnx/test_mnb_to_qdq.py loosens the 8-bit tolerance from 1e-2 to 2e-2. Both can mask regressions — please justify or restore.

Nits

  • _print_build_outputs prints the raw configured string while the engine resolved it to an absolute path — inconsistent with the single-build message.
  • errors.setdefault(name, []).append(exc) — a Future can only fail once, so a flat dict[str, Exception] is clearer.
  • get_build_output_dir's docstring doesn't state the actual output/<build_name> convention, which is user-visible and duplicated in the --output_path error message.

What's good

  • Pre-validating all expanded builds before running any of them is exactly the right failure mode for a long multi-build run, and it's pinned by a test.
  • get_build_cache_dir mirrors CacheConfig.create_cache's layout so isolated_cache_env seeds the correct value before Engine.initialize() calls set_cache_env — a non-obvious ordering fix.
  • assertValueError in CompositeModelHandler (with a test citing python -O) is the right call for validating external input.
  • Build names are constrained before being used in filesystem paths.
  • discover_onnx_components has an excellent docstring with a concrete directory-layout example.
  • DiffusersModelHandler.get_exportable_components preserving the variant's canonical order rather than the requested order matches the documented contract.

AI-assisted review. Findings were cross-checked against the diff, but please verify before acting — particularly #4, which would benefit from a test with an on-disk quant config whose category flags conflict with the current run.

@titaiwangms

Copy link
Copy Markdown
Contributor

Re-review of the update — all 6 Major findings resolved ✅

Re-ran the review against the updated diff and checked out the branch to verify empirically rather than by reading alone. Every Major from the previous pass is genuinely fixed, and several fixes are better than what I suggested.

Verified fixed

1. Writable-dir overlap heuristic — the output_dir.suffix and not output_dir.is_dir() heuristic is gone; _get_build_write_dirs now just resolves output_dir. Validation is a pure function of the config again. Verified:

builds {"llama.q4": ..., "plain": ...}
  -> {'llama.q4': 'output/llama.q4', 'plain': 'output/plain'}   # previously raised "overlapping writable directories"

2. _default.output_dir — now a parent directory, resolved in _parse_builds via get_build_output_dir(build_name, default_output_dir=...), with BuildConfigPartial.output_dir carrying a description that says so and the docs stating it explicitly. Verified:

{"_default": {"output_dir": "models"}, "a": ..., "b": ...}
  -> {'a': 'models/a', 'b': 'models/b'}

3. select_components mutationselected = [deepcopy(component_map[n]) for n in names] plus an explicit uniqueness check ("select_components requires unique component names."). The source composite is no longer mutated and the duplicate-name case is rejected.

4. quant_utils skip-set — this is now correct and better factored than my suggestion. quantized_names = already_quantized | set(new_qargs) is computed after the quantization loop, and persisted_skip_patterns rewrites any pattern that overlaps a quantized module into exact re:^...$ patterns for the genuinely-unquantized reload targets. That fixes both halves of the finding: the skip/quantize contradiction and the substring-matching hazard. Generated patterns are also now gated on component_source_paths, so existing non-component runs no longer get their quantization_config rewritten. The new tests pin the exact case I was worried about:

assert match_skip("blocks.1", qcfg.modules_to_not_convert)
assert not match_skip("blocks.10", qcfg.modules_to_not_convert)

5. olive/cli/base.py elifif — reverted; the file is no longer in the diff.

6. Schema and docsbuilds and max_concurrent_builds are now real RunConfig fields, so dump_schema.py picks them up automatically (pinned by test_run_config_schema_includes_multi_build_fields). docs/source/how-to/configure-workflows/build-workflow.md gains a complete "Run multiple builds" section covering _default, the parent-directory semantics, the serial default, the OpenVINOOptimumConversion caveat and the local-host restriction. skills/olive/scripts/validate_config.py now routes through parse_run_config and reports per-build.

Verification run

test/workflows/test_run_config_builds.py
test/workflows/test_run_builds.py
test/model/test_composite_model.py          -> 50 passed
test/passes/pytorch/test_quant_utils.py
test/common/quant/test_tensor.py            -> 96 passed
lintrunner --skip PYLINT (changed files)    -> clean (only the usual CPY001 false positives)

test/test_cache.py shows 5 failed / 10 errored, but that is identical on main (missing azure-storage-blob / azure-identity locally) — the PR adds one passing test there.

Remaining, all non-blocking

  • test/common/test_hf_wrapper.py still drops the LFM2 get_attention_outputs() assertions with no corresponding behavior change — get_attention_outputs still returns ([], []) when self.attn is None, and nothing in wrapper.py touches attn / ATTENTION_OUTPUTS. This looks like gratuitous coverage loss; consider restoring the three lines.
  • Per-build file logging is still process-global. Engine.initializeenable_filelog adds a handler to the single Olive logger and never removes it, and _run_single calls set_default_logger_severity globally. With max_concurrent_builds > 1 — which the docs now actively advertise — every build's records land in every build's log file and handlers leak after the run. Worth at least a docs caveat next to the parallelism note, or a build-scoped filtering handler.
  • Zip packaging under parallel builds. _package_zipfile_model stages via shutil.make_archive(output_name, ...) relative to the process CWD under a shared archive name, so concurrent builds inheriting the same top-level packaging_config can clobber each other outside their validated output dirs.
  • olive/cli/run.py:118 still hardcodes the "_default" literal instead of importing BUILD_DEFAULT_KEY, and re-derives the output dir from raw JSON rather than reusing expand_builds' result. It happens to agree with _parse_builds today, except when a build sets "output_dir": null explicitly.
  • Partial results discarded_run_builds_in_parallel raises and throws away the WorkflowOutputs of builds that already succeeded.
  • ContextVar cache isolation still doesn't cross thread/spawn boundaries, but I traced all three from_cache_env call sites (mlflow.py, image_data_container.py, diffusers/lora.py) and they all run in the owning build thread, and the spawn context in conversion.py doesn't touch it. Latent rather than a live bug — noting it only so the constraint is known if a future pass reads the cache from a worker.

Nits

  • RunConfig.builds is typed dict[str, BuildConfigPartial], so the generated schema advertises pipeline as optional for named builds even though BuildConfig requires it. Harmless at runtime (_parse_builds validates with BuildConfig), but the published schema is slightly looser than reality.
  • _validate_build_write_dirs still says "Parallel builds ... have overlapping writable directories" even when max_concurrent_builds == 1.
  • errors.setdefault(name, []).append(exc) — a Future can only fail once, so a flat dict[str, Exception] would read better.

Nice work on #4 in particular — the reload_target_names - quantized_names rewrite handles a case I'd only half-identified. The test_mnb_to_qdq.py tolerance change also reads much better now that it's paired with fork_rng + manual_seed(0) + default_rng(0); the looser bound is on a deterministic fixture rather than a random one.


AI-assisted review. Fixes were verified by checking out the branch, running the listed test targets, and executing expand_builds directly for findings 1 and 2.

@xiaoyu-work
Xiaoyu (xiaoyu-work) merged commit ee6f40f into main Aug 21, 2026
12 checks passed
@xiaoyu-work
Xiaoyu (xiaoyu-work) deleted the xiaoyu/builds-schema branch August 21, 2026 02:24
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

5 participants