Guard RTN MoE quantization by expert layout

[Ti-Tai Wang (titaiwangms)]

Problem

PR #2584 merged RTN MoE quantization (moe=True) support without a layout check on fused-expert weights. RTN's WeightQuantizer groups unconditionally along a tensor's last dimension, which is only correct when the fused expert weight is stored (num_experts, out_features, in_features) -- K last. Architectures such as gpt-oss store the transposed (num_experts, in, out) layout instead, and are silently mis-quantized (wrong axis grouped, no error) today on main when run through moe=True.

Fix

Adds olive/passes/pytorch/moe_support.py with check_moe_layout_support, gated into RTN's _run_for_config after prepare_model() and before finalize(). The check trusts transformers's own is_transposed attribute (set by the use_experts_implementation decorator, not derived from config/checkpoint data) directly:

• Accept only experts modules that report is_transposed is False.
• Reject anything where is_transposed is missing or not a bool (covers older transformers releases, undecorated architectures such as llama4/aria, and unrecognized implementations).
• Reject is_transposed=True (e.g. gpt-oss).
• Exempt classic per-expert nn.ModuleList experts (e.g. Mixtral/PhiMoE on older transformers) that carry no direct 3D parameter.

A trust_remote_code custom experts implementation that misreports its own is_transposed is out of scope: that is treated as user-introduced misuse of an explicitly opted-in trust boundary, not a layout Olive can independently verify. No architecture allow-list is used.

Only affects the moe=True path -- gated behind if qcfg.moe: and only runs when experts modules are actually detected, so non-MoE quantization is unaffected.

Testing

28/28 relevant tests pass (test/passes/pytorch/test_moe_support.py, test/passes/pytorch/test_rtn.py); lintrunner clean.

Note: PR #2610 (GPTQ MoE) will stack on top of this branch to reuse check_moe_layout_support and avoid duplicating the layout logic.

[Copilot]

Pull request overview

This PR adds a fail-closed safety gate for PyTorch RTN MoE quantization to prevent silent mis-quantization of fused expert weights when the expert layout is transposed (K not last). It introduces a shared check_moe_layout_support helper and updates RTN to invoke it only for MoE quantization flows, with accompanying unit tests and user-facing documentation updates.

Changes:

• Added olive/passes/pytorch/moe_support.py with MoeSupportError and check_moe_layout_support to validate fused-experts layout safety via is_transposed.
• Wired the layout check into olive/passes/pytorch/rtn.py for moe=True runs (before finalize/save).
• Added/updated tests and docs to cover/reflection the new fail-closed behavior.

Reviewed changes

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

Show a summary per file

File	Description
olive/passes/pytorch/rtn.py	Adds the MoE layout validation gate before finalizing RTN quantization.
olive/passes/pytorch/moe_support.py	Implements shared layout-safety checks and a dedicated exception type for MoE quantization support.
test/passes/pytorch/test_rtn.py	Adds end-to-end RTN tests ensuring layout rejection happens before finalize and that moe=False bypasses the gate.
test/passes/pytorch/test_moe_support.py	Adds unit tests for check_moe_layout_support behavior (missing/invalid is_transposed, transposed layouts, ModuleList exemptions).
docs/source/features/quantization.md	Documents the new MoE support constraints for fused experts and the fail-closed behavior.

---

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

copilot's comment looks reasonable

[Ti-Tai Wang (titaiwangms)]

Jambay Kinley (@jambayk) Done! This is ready!