[WIP] Refactor pytorch quant modules, Support MoE

[jambayk]

Describe your changes

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

[github-advanced-security]

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

[Copilot]

Pull request overview

This PR refactors Olive’s PyTorch weight quantization implementation from dedicated quantized nn.Module wrappers (QuantLinear/QuantEmbedding) to a QuantTensor parameter representation, and extends the quantization pipeline to better handle MoE models (including fused 3D expert weights and configurable expert skipping).

Changes:

• Introduces QuantTensor (a torch.Tensor wrapper subclass) + state-dict helpers to store quantized weights as aliased buffers while keeping eager execution working via F.linear/F.embedding dispatch.
• Updates quantization utilities to (a) support N-D quantization along the last dimension and (b) add MoE-aware selection/skip logic plus fused-3D expert parameter quantization.
• Adds regex-aware pattern matching helpers for overrides and modules_to_not_convert, updates ONNX export compatibility and ModelBuilder loading behavior, and refreshes/expands unit tests accordingly.

Reviewed changes

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

Show a summary per file

File	Description
test/passes/pytorch/test_rtn.py	Updates RTN tests to validate quantization via QuantTensor weights rather than quantized module classes.
test/passes/pytorch/test_gptq.py	Updates GPTQ tests to validate QuantTensor-backed weights.
test/passes/pytorch/test_quant_utils.py	Adds offline unit tests for install_quant_tensor_param (incl. fused 3D expert parameters).
test/passes/onnx/test_model_builder.py	Adds coverage ensuring Olive-quantized MoE checkpoints error out in ModelBuilder.
test/common/quant/test_utils.py	Updates assertions and adds tests for N-D quantization/packing helpers.
test/common/quant/test_tensor.py	Adds unit tests for QuantTensor semantics, dispatch behavior, dtype/device propagation, slicing, and ONNX export guards.
test/common/quant/test_patterns.py	Adds tests for regex/literal override + skip matching helpers.
test/common/quant/test_hf_utils.py	Updates HF quantizer tests for QuantTensor and adds MoE-specific behavior tests.
test/common/quant/test_nn.py	Removes tests for deleted QuantLinear/QuantEmbedding module wrappers.
olive/passes/pytorch/rtn.py	Enables MoE option exposure in RTN pass default config.
olive/passes/pytorch/quant_utils.py	Refactors prepare/finalize to install QuantTensor params, adds MoE expert handling and skip patterns.
olive/passes/onnx/model_builder.py	Maps Olive buffer suffixes to expected QuantizedTensorModule attrs and rejects MoE Olive checkpoints explicitly.
olive/common/quant/utils.py	Extends WeightQuantizer + pack/unpack to work on N-D tensors along the last dimension.
olive/common/quant/tensor.py	Adds QuantTensor implementation with eager dispatch + export guards + tensor transform propagation.
olive/common/quant/state_dict.py	Adds buffer naming + state-dict hook + install/refresh helpers for QuantTensor parameters.
olive/common/quant/patterns.py	Adds centralized literal/regex matching helpers for overrides and skip patterns.
olive/common/quant/hf_utils.py	Refactors HF quantizer to use QuantTensor placeholders and MoE-aware skipping; updates tie logic.
olive/common/hf/wrapper.py	Adds LayerWrapper.get_experts / get_router helpers for MoE structure discovery.
olive/common/hf/quant.py	Adds conversion from Olive QuantTensor weights to ONNX-exportable QuantLinearNbit/QuantEmbeddingNbit wrappers.
olive/common/quant/nn.py	Removes legacy quantized module wrapper implementations.

Comments suppressed due to low confidence (1)

olive/common/quant/hf_utils.py:429

• tie_quant_word_embeddings will currently tie the output embedding to the input embedding as long as the input has a QuantTensor weight. This can unintentionally quantize the output embedding even when lm_head=False (or leave dst without Olive’s state_dict hook), and it no longer validates shape/dtype compatibility before aliasing buffers. Please guard so tying only happens when both src and dst already have compatible QuantTensor weights/buffers (and keep the previous shape/dtype assertions) to avoid silently changing quantization behavior or breaking save/load.

def tie_quant_word_embeddings(model: PreTrainedModel) -> None:
    """Tie the input and output embeddings when they share a quantized weight.

    Both modules' ``weight`` ``nn.Parameter`` is set to the **same**
    ``nn.Parameter(QuantTensor)`` object, and the underlying
    ``weight_qweight`` / ``weight_scales`` / ``weight_qzeros`` buffers
    are tied (aliased to the input embedding's buffers). This preserves
    the standard HF tied-weights semantics for the quantized layout.
    """
    src = model.get_input_embeddings()
    dst = model.get_output_embeddings()
    if src is None or dst is None:
        return

    # The input embedding owns the canonical QuantTensor + buffers.
    src_param = src._parameters.get("weight")
    if src_param is None or not isinstance(src_param.data, QuantTensor):
        return

    qname, sname, zname = buffer_names("weight")
    # tie buffers
    for n in (qname, sname, zname):
        src_buf = src._buffers.get(n)
        if src_buf is None:
            continue
        dst._buffers[n] = src_buf

    # tie the QuantTensor parameter itself (same Python Parameter object,
    # so both modules see the same .data and the same inner tensors).
    dst._parameters["weight"] = src_param

[Copilot]

Pull request overview

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

Comments suppressed due to low confidence (1)

olive/common/hf/quant.py:175

• In the dynamo ONNX-export path, the placeholder inserted when qzeros is None and g_idx is not None is created as torch.zeros(0, dtype=torch.uint8) with no device specified. If x/qweight/scales are on CUDA, this can trigger device-mismatch errors during export graph construction. Create the placeholder on x.device (and ideally match qweight's device) to keep tensor args consistent.

                    tensor_args.append(g_idx)
                attrs = {
                    "K": in_features,
                    "N": out_features,
                    "bits": bits,
                    "block_size": group_size,
                    "accuracy_level": 4,

[Copilot]

Pull request overview

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

[Copilot]

Pull request overview

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