Skip to content

Commit 1e76021

Browse files
authored
[AMD] Update QuarkQuantization pass for Quark 0.12 (microsoft#2532)
## Describe your changes Updates `QuarkQuantization` to work with AMD Quark 0.12, which introduced two breaking API changes: **ONNX path — calibration reader fix:** Quark 0.12 rejects calibration inputs that are not model inputs. Olive was forwarding the full dataset batch (including label columns like `class`) to the reader, causing `INVALID_ARGUMENT: Invalid input name` during calibration — the workflow would complete with exit code 0 but produce no output model. Fixed by passing `model_path`/`io_config` to `create_calibration_dataloader()` so non-input columns are filtered out, matching the pattern already used by `OnnxQuantization`. **Torch path — API update:** - `revert_model_patching` was removed in Quark 0.12 (patching is now reverted internally during export). Removed the import and call. - `prepare_for_moe_quant` is deprecated and raises on `transformers>=5`. Replaced with `preprocess_for_quantization`. **Version gating:** Both ONNX and Torch paths now gate on `amd-quark>=0.12.0` with a clear `ValueError` so users get an actionable error message instead of a cryptic `ImportError`. ## Checklist before requesting a review - [x] Add unit tests for this change. - [x] Make sure all tests can pass. - [x] Update documents if necessary. - [x] Lint and apply fixes to your code by running `lintrunner -a` - [x] Is this a user-facing change? If yes, give a description of this change to be included in the release notes. The `QuarkQuantization` pass now requires `amd-quark>=0.12.0`. Users on 0.11.x will receive a clear `ValueError` with an upgrade prompt. The ONNX calibration path and Torch path are updated to work correctly with the 0.12 API. ## (Optional) Issue link
1 parent 1d5b84b commit 1e76021

4 files changed

Lines changed: 93 additions & 18 deletions

File tree

olive/olive_config.json

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -605,7 +605,7 @@
605605
"supported_algorithms": [ "awq" ],
606606
"supported_quantization_encodings": [ ],
607607
"run_on_target": true,
608-
"extra_dependencies": [ "amd-quark" ]
608+
"extra_dependencies": [ "amd-quark>=0.12.0" ]
609609
},
610610
"QuarkQuantizationVitisAI": {
611611
"module_path": "olive.passes.quark_vitisai.quark_quantization_vitisai.QuarkQuantizationVitisAI",
@@ -615,7 +615,7 @@
615615
"supported_algorithms": [ "awq" ],
616616
"supported_quantization_encodings": [ ],
617617
"run_on_target": true,
618-
"extra_dependencies": [ "amd-quark" ]
618+
"extra_dependencies": [ "amd-quark>=0.12.0" ]
619619
},
620620
"KQuant": {
621621
"module_path": "olive.passes.pytorch.kquant.KQuant",

olive/passes/quark_quantizer/quark_quantization.py

Lines changed: 12 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -30,7 +30,7 @@ class QuarkQuantization(Pass):
3030
3131
Routes to the appropriate backend based on input model type:
3232
- ONNXModelHandler -> Quark-ONNX quantization (quark.onnx)
33-
- HfModelHandler -> Quark-Torch quantization (quark.torch, Quark 0.11 API)
33+
- HfModelHandler -> Quark-Torch quantization (quark.torch, Quark 0.12 API)
3434
"""
3535

3636
@classmethod
@@ -188,7 +188,7 @@ def _run_for_config(
188188
logger.info("[INFO] Running QuarkQuantization using Quark-ONNX API")
189189
return self._run_quark_onnx(model, config, output_model_path)
190190
else:
191-
logger.info("[INFO] Running QuarkQuantization using Quark-Torch 0.11 API")
191+
logger.info("[INFO] Running QuarkQuantization using Quark-Torch 0.12 API")
192192
return self._run_quark_torch(model, config, output_model_path)
193193

194194
# ── ONNX path ───────────────────────────────────────────
@@ -201,8 +201,8 @@ def _run_quark_onnx(
201201
) -> ONNXModelHandler:
202202
from quark import __version__ as QuarkVersion
203203

204-
if version.parse(QuarkVersion) < version.parse("0.11.0"):
205-
raise ValueError("Quark ONNX Quantization is only supported for amd-quark>=0.11.0")
204+
if version.parse(QuarkVersion) < version.parse("0.12.0"):
205+
raise ValueError("Quark ONNX Quantization is only supported for amd-quark>=0.12.0")
206206

207207
from olive.passes.quark_quantizer.onnx.quantize_quark import run_quark_quantization
208208

@@ -217,7 +217,9 @@ def _run_quark_onnx(
217217
data_reader = None
218218
if config.data_config:
219219
data_config = validate_config(config.data_config, DataConfig)
220-
data_reader = data_config.to_data_container().create_calibration_dataloader()
220+
data_reader = data_config.to_data_container().create_calibration_dataloader(
221+
model_path=model.model_path, io_config=model.io_config
222+
)
221223

222224
run_config = config.model_dump()
223225
to_delete = [
@@ -253,6 +255,11 @@ def _run_quark_torch(
253255
config: BasePassConfig,
254256
output_model_path: str,
255257
) -> HfModelHandler:
258+
from quark import __version__ as QuarkVersion
259+
260+
if version.parse(QuarkVersion) < version.parse("0.12.0"):
261+
raise ValueError("Quark Torch Quantization is only supported for amd-quark>=0.12.0")
262+
256263
from olive.passes.quark_quantizer.torch.quark_torch_quantization import run_quark_torch_quantization
257264

258265
return run_quark_torch_quantization(model, config, output_model_path)

olive/passes/quark_quantizer/torch/quark_torch_quantization.py

Lines changed: 6 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -3,7 +3,7 @@
33
# SPDX-License-Identifier: MIT
44
#
55

6-
"""Quark 0.11 Torch quantization for LLMs.
6+
"""Quark 0.12 Torch quantization for LLMs.
77
88
Uses LLMTemplate + ModelQuantizer from the Quark public API.
99
"""
@@ -27,7 +27,7 @@ def run_quark_torch_quantization(
2727
config: BasePassConfig,
2828
output_model_path: str,
2929
) -> HfModelHandler:
30-
"""Run Quark 0.11 torch quantization on a HuggingFace model.
30+
"""Run Quark 0.12 torch quantization on a HuggingFace model.
3131
3232
Args:
3333
model: Olive HfModelHandler pointing to the source model.
@@ -53,8 +53,7 @@ def run_quark_torch_quantization(
5353
get_calib_dataloader,
5454
get_model,
5555
get_tokenizer,
56-
prepare_for_moe_quant,
57-
revert_model_patching,
56+
preprocess_for_quantization,
5857
)
5958

6059
output_dir = Path(output_model_path)
@@ -75,7 +74,7 @@ def run_quark_torch_quantization(
7574
trust_remote_code=config.trust_remote_code,
7675
)
7776

78-
prepare_for_moe_quant(torch_model)
77+
preprocess_for_quantization(torch_model)
7978

8079
model_type = (
8180
torch_model.config.model_type
@@ -145,15 +144,11 @@ def run_quark_torch_quantization(
145144
logger.info("[INFO] Freezing quantized model")
146145
torch_model = quantizer.freeze(torch_model)
147146

148-
# 6. Revert model patching
149-
logger.info("[INFO] Reverting model patching")
150-
revert_model_patching(torch_model)
151-
152-
# 7. Validate export configuration
147+
# 6. Validate export configuration
153148
if config.custom_mode != "quark" and config.export_weight_format == "fake_quantized":
154149
raise ValueError("'fake_quantized' export is only supported with custom_mode='quark'")
155150

156-
# 8. Export model
151+
# 7. Export model
157152
logger.info("[INFO] Exporting quantized model to: %s", output_dir)
158153

159154
export_formats = config.model_export

test/passes/quark_quantizer/test_quark_onnx_quantization.py

Lines changed: 73 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,9 @@
33
# Licensed under the MIT License.
44
# --------------------------------------------------------------------------
55

6+
from unittest.mock import patch
7+
8+
import numpy as np
69
import pytest
710
from onnxruntime.quantization.calibrate import CalibrationDataReader
811

@@ -194,3 +197,73 @@ def test_static_qdq_u8s8_with_layerwise_mixed_precision_quantization(tmp_path):
194197
p = create_pass_from_dict(QuarkQuantization, config, disable_search=True)
195198
out = p.run(input_model, tmp_path)
196199
assert out is not None
200+
201+
202+
@Registry.register_dataloader()
203+
def _test_quant_dataloader_with_label(dataset, batch_size, **kwargs):
204+
"""Yield model input and a label column as a real dataset pipeline would.
205+
206+
The label must be filtered out before being passed to the ONNX calibration reader.
207+
"""
208+
209+
class _ReaderWithLabel(CalibrationDataReader):
210+
# pylint: disable=W0223
211+
def __init__(self):
212+
super().__init__()
213+
self.samples = [{"input": np.random.randn(1, 1).astype(np.float32), "label": 0} for _ in range(4)]
214+
self._iter = iter(self.samples)
215+
216+
def get_next(self):
217+
return next(self._iter, None)
218+
219+
return _ReaderWithLabel()
220+
221+
222+
def test_calibration_dataloader_filters_label_columns(tmp_path):
223+
"""Regression: Quark 0.12 rejects calibration inputs that are not model inputs.
224+
225+
The pass must pass model_path/io_config to create_calibration_dataloader() so
226+
non-input columns (e.g. a label column) are stripped before reaching Quark.
227+
"""
228+
input_model = get_onnx_model()
229+
config = {
230+
"quant_mode": "static",
231+
"quant_format": "QDQ",
232+
"global_config": {
233+
"activation": {"symmetric": False, "calibration_method": "MinMax", "data_type": "UInt8"},
234+
"weight": {"symmetric": True, "calibration_method": "MinMax", "data_type": "Int8"},
235+
},
236+
"data_config": DataConfig(
237+
name="test_quant_dc_config_with_label",
238+
load_dataset_config=DataComponentConfig(type="simple_dataset"),
239+
dataloader_config=DataComponentConfig(type="_test_quant_dataloader_with_label"),
240+
),
241+
}
242+
p = create_pass_from_dict(QuarkQuantization, config, disable_search=True)
243+
# Should complete without "Invalid input name: label" from Quark
244+
out = p.run(input_model, tmp_path)
245+
assert out is not None
246+
247+
248+
def test_onnx_path_raises_clear_error_for_old_quark(tmp_path):
249+
"""ONNX path must raise a clear ValueError, not an ImportError, for amd-quark < 0.12.0."""
250+
input_model = get_onnx_model()
251+
config = {"quant_mode": "static", "quant_format": "QDQ"}
252+
p = create_pass_from_dict(QuarkQuantization, config, disable_search=True)
253+
with patch("quark.__version__", "0.11.2"), pytest.raises(ValueError, match=r"amd-quark>=0\.12\.0"):
254+
p.run(input_model, tmp_path)
255+
256+
257+
def test_torch_path_raises_clear_error_for_old_quark(tmp_path):
258+
"""Torch path must raise a clear ValueError, not an ImportError, for amd-quark < 0.12.0."""
259+
from unittest.mock import MagicMock
260+
261+
from olive.model import HfModelHandler
262+
263+
# Use a MagicMock so handler construction/validation is bypassed.
264+
# The version gate fires inside _run_quark_torch before any model loading.
265+
input_model = MagicMock(spec=HfModelHandler)
266+
config = {"quant_scheme": "uint4_wo_128"}
267+
p = create_pass_from_dict(QuarkQuantization, config, disable_search=True)
268+
with patch("quark.__version__", "0.11.2"), pytest.raises(ValueError, match=r"amd-quark>=0\.12\.0"):
269+
p.run(input_model, tmp_path)

0 commit comments

Comments
 (0)