Support Qwen3.5/3.6-MoE VL checkpoints in PyTorch-side quantization

[Ti-Tai Wang (titaiwangms)]

Summary

Enables Olive's PyTorch-side quantization passes (Rtn/GPTQ/KQuant, via ModelWrapper/iter_quant_targets) to work on a Qwen3.5/3.6-MoE checkpoint loaded as its vision-language class (Qwen3_5MoeForConditionalGeneration, task="image-text-to-text"), instead of only the text-only decoder class (Qwen3_5MoeForCausalLM) that was previously supported.

Root cause

Two structural gaps, both purely data/mapping issues (no PyTorch quantization logic needed to change):

1. Composite config resolution: the VL checkpoint's top-level config has no flat hidden_size/num_hidden_layers/num_attention_heads/num_key_value_heads/head_dim — they live under config.text_config. ModelWrapper.__init__ resolved these to None and crashed with a TypeError.
2. Decoder module path: the VL decoder is nested at model.language_model.* (alongside model.visual, the vision tower), not model.* directly, so LAYERS/EMBEDDINGS/PRE_HEAD_LAYERNORM/ROTARY_EMBEDDING couldn't resolve it.

Fix

• olive/assets/io_configs/defaults.yaml: add text_config.* fallbacks for num_layers, hidden_size, num_attention_heads, num_kv_heads (mirrors the existing num_experts nested-fallback pattern), plus a new head_dim alias.
• olive/common/hf/wrapper.py: add a "qwen3_5_moe" entry to LAYERS/EMBEDDINGS/PRE_HEAD_LAYERNORM/ROTARY_EMBEDDING pointing at model.language_model.*. Scoped to the VL model_type only — the text-only checkpoint carries model_type == "qwen3_5_moe_text" and its flat config already matches "default", so it is unaffected (see test_hf_wrapper_text_only_config_unaffected_by_vl_aliases).

Two pre-existing (non-VL-specific) quantization gaps, found while validating the above

Both reproduce identically on the text-only checkpoint too, but affect Qwen3.5/3.6-MoE quantization quality generally:

• MAMBA mapping didn't recognize this architecture's linear_attn (GatedDeltaNet) attribute name, so its projections (in_proj_qkv/in_proj_a/in_proj_b/in_proj_z/out_proj) were swept into the generic 2D quantization walk instead of staying full precision like other Mamba/SSM blocks.
• Added a SHARED_EXPERT_GATE mapping + LayerWrapper.get_shared_expert_gate() accessor (also applicable to qwen2_moe/qwen3_next/qwen3_omni_moe, which use the same attribute name) and wired it into iter_quant_targets's exclusion set — the single-row shared-expert sigmoid gate was being quantized like an ordinary Linear, despite being a routing-like signal (same reasoning as excluding the main router).

Verification

• Confirmed against the real Qwen/Qwen3.6-35B-A3B checkpoint's config: alias/mapping resolution (hidden=2048 heads=16 kv=2 head_dim=256 layers=40, decoder path model.language_model.layers) now matches expectations exactly.
• Built a shape-faithful synthetic VL checkpoint (206M params, real vision tower + 4-layer MoE decoder) and ran the actual Rtn pass end-to-end (bits=4, group_size=32, moe=True, modules_to_not_convert=["visual"]):
  • mlp.experts.* (fused 3D MoE weights) — quantized ✅
  • mlp.gate (router), mlp.shared_expert_gate, linear_attn.* — full precision, untouched ✅
  • model.visual.* (vision tower) — 0 quant artifacts ✅
  • Reload of the quantized checkpoint via HfModelHandler succeeds.
• Added 5 new unit tests covering composite-config resolution, text-only-path non-regression, and linear_attn/shared_expert_gate exclusion.
• Full relevant suite: test/common/, test/passes/pytorch/test_rtn.py, test/model/ — 622+ passed, 0 regressions (pre-existing unrelated azure module errors only). lintrunner clean (only the repo-wide CPY001 false positive).

Files

• olive/assets/io_configs/defaults.yaml
• olive/common/hf/wrapper.py
• olive/common/quant/selection.py
• test/common/test_hf_wrapper.py
• test/common/quant/test_selection.py

---

Recreated from #2628 (same branch/commit b7c0cd3, pushed directly to microsoft/Olive instead of a fork) so CI has access to the hf_token secret, which Azure DevOps withholds from fork PR builds.

[Copilot]

Pull request overview

Extends Olive’s Hugging Face model wrapping and PyTorch-side quantization target selection to correctly handle Qwen3.5/3.6 MoE vision-language checkpoints (composite configs + nested decoder layout), and improves metadata saving for multimodal models.

Changes:

• Add text_config.* fallbacks (and head_dim) to alias resolution so composite VL configs resolve core decoder attributes correctly.
• Introduce Qwen3.5 MoE VL-specific module-path mappings and a model-type normalization to distinguish flat text-only vs composite VL layouts.
• Improve quantization target selection by excluding shared-expert gates, Mamba/linear-attn blocks, and (for composite VL models) the vision tower; add coverage tests.

Reviewed changes

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

Show a summary per file

File	Description
olive/assets/io_configs/defaults.yaml	Adds alias fallbacks for composite configs (text_config.*) and head_dim support.
olive/common/hf/io_config/io_resolver.py	Makes alias resolution robust to heterogeneous-config exceptions and ensures canonical attrs win over fallbacks.
olive/common/hf/wrapper.py	Adds Qwen3.5 MoE VL module-path mappings, shared-expert gate + linear-attn (Mamba) handling, and text-only normalization.
olive/common/quant/selection.py	Excludes shared-expert gates, Mamba/linear-attn modules, and VL vision towers from PyTorch-side quantization targets.
olive/model/handler/mixin/hf.py	Saves processor/preprocessor files without overwriting existing tokenizer artifacts; surfaces failures via warning logs.
test/common/hf/io_config/test_task_config.py	Adds regression tests for alias precedence and ambiguous per-layer attribute behavior.
test/common/test_hf_wrapper.py	Adds synthetic config/model tests for composite VL vs flat text-only Qwen3.5 MoE wrapper behavior.
test/common/quant/test_selection.py	Adds tests for shared-expert gate exclusion, linear-attn exclusion, and vision tower exclusion.
test/model/test_hf_model.py	Adds tests ensuring save_metadata handles processor saving, skip behavior, and warning visibility correctly.

---

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

[Copilot]

Pull request overview

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

Suppressed comments (2)

olive/model/handler/mixin/hf.py:129

• Treating preprocessor_config.json as a completeness sentinel can leave a multimodal checkpoint incomplete. Processors may also emit processor_config.json, chat templates, and modality-specific files; if an earlier step saved only the image-processor config, none of those missing files are copied. Call _save_processor even when this one file exists—the staging/copy logic already preserves every existing destination file.

        if not (output_dir / "preprocessor_config.json").exists():
            saved_filepaths.extend(self._save_processor(output_dir, exclude_load_keys=exclude_load_keys, **kwargs))

olive/model/handler/mixin/hf.py:127

• The PR description characterizes this as mapping-only quantization support and lists five files, but this adds repository-wide processor loading/copy behavior to save_metadata plus a separate test suite. Please either split this unrelated behavior into its own PR or update the description, file list, and verification so reviewers can assess the added metadata/network and compatibility impact.

        # save processor / image processor, skip if one already exists
        # this writes preprocessor_config.json (and any image processor files) so downstream
        # tools that load from this output_dir (e.g. mobius's AutoProcessor.from_pretrained)
        # get the model's real preprocessing config instead of silently falling back to
        # defaults. Only applicable to multimodal models (e.g. VL checkpoints); text-only
        # models have no processor and are already covered by the tokenizer save above.

[Ti-Tai Wang (titaiwangms)]

Re: the two suppressed review comments on olive/model/handler/mixin/hf.py:

• hf.py:129 (preprocessor_config.json completeness sentinel) -- confirmed and fixed in 3b89343. _save_processor's own _copy_missing_files already skips existing files one by one, so gating the whole call on preprocessor_config.json's existence only hid genuinely-missing files (e.g. chat_template.json) whenever some other processor file happened to already exist. Removed the outer guard; added a regression test (test_save_metadata_processor_fills_in_missing_files_without_overwriting_existing).
• hf.py:127 (PR scope/description) -- this is a valid process point (the processor-saving behavior + its own test suite is broader than the PR description's "mapping only" framing), but not a code bug. Leaving the description as-is since this behavior is a real dependency of the VL export path being added here, but flagging it here for visibility.