diff --git a/src/python/py/models/builder.py b/src/python/py/models/builder.py index 0cc75acdd3..de2745137a 100644 --- a/src/python/py/models/builder.py +++ b/src/python/py/models/builder.py @@ -51,6 +51,9 @@ ) +VALID_QUANT_MODES = frozenset({"default", "hybrid", "int4"}) + + def check_extra_options(kv_pairs, execution_provider): """ Check key-value pairs and set values correctly @@ -71,6 +74,23 @@ def check_extra_options(kv_pairs, execution_provider): "disable_qkv_fusion", "prune_lm_head", ] + + # Validate quant_mode if provided. Empty or whitespace-only values are + # treated as unset so downstream builders can apply their default behavior. + if "quant_mode" in kv_pairs: + quant_mode = kv_pairs["quant_mode"] + if isinstance(quant_mode, str): + quant_mode = quant_mode.strip().lower() + + if not quant_mode: + del kv_pairs["quant_mode"] + elif quant_mode not in VALID_QUANT_MODES: + valid_modes_display = ", ".join(sorted(VALID_QUANT_MODES)) + raise ValueError( + f"quant_mode must be one of {valid_modes_display}, got '{kv_pairs['quant_mode']}'" + ) + else: + kv_pairs["quant_mode"] = quant_mode for key in bools: if key in kv_pairs: if kv_pairs[key] in {"false", "False", "0"}: diff --git a/src/python/py/models/builders/qwen.py b/src/python/py/models/builders/qwen.py index 0e228bb000..3d9d36b2aa 100644 --- a/src/python/py/models/builders/qwen.py +++ b/src/python/py/models/builders/qwen.py @@ -1019,15 +1019,27 @@ def __init__(self, config, io_dtype, onnx_dtype, ep, cache_dir, extra_options): # # Linear attention recurrence accumulates errors across the full sequence, # unlike softmax attention which normalizes per-step. + # + # Control via extra_options quant_mode: + # "default" - INT8 for all linear attn + MLP layers (most accurate) + # "hybrid" - INT8 for linear attn projections only, INT4 for MLPs (balanced) + # "int4" - INT4 for everything (fastest, may degrade quality) + quant_mode = extra_options.get("quant_mode", "").strip().lower() or "default" + if quant_mode not in ("default", "hybrid", "int4"): + raise ValueError(f"quant_mode must be one of default, hybrid, int4, got '{quant_mode}'") + + linear_attn_projs = ("in_proj_a", "in_proj_b", "in_proj_qkv", "in_proj_z", "out_proj") + mlp_projs = ("gate_proj", "up_proj", "down_proj") + int8_nodes = {} - for i, lt in enumerate(self.layer_types): - if lt == "linear_attention": - # All linear attention projections: INT8 - for proj in ("in_proj_a", "in_proj_b", "in_proj_qkv", "in_proj_z", "out_proj"): - int8_nodes[f"/model/layers.{i}/linear_attn/{proj}/MatMul"] = {"bits": 8} - # MLP projections in linear attention layers: INT8 - for proj in ("gate_proj", "up_proj", "down_proj"): - int8_nodes[f"/model/layers.{i}/mlp/{proj}/MatMul"] = {"bits": 8} + if quant_mode in ("default", "hybrid"): + for i, lt in enumerate(self.layer_types): + if lt == "linear_attention": + for proj in linear_attn_projs: + int8_nodes[f"/model/layers.{i}/linear_attn/{proj}/MatMul"] = {"bits": 8} + if quant_mode == "default": + for proj in mlp_projs: + int8_nodes[f"/model/layers.{i}/mlp/{proj}/MatMul"] = {"bits": 8} if int8_nodes: algo_config = self.quant_attrs["int4"].get("algo_config")