From 0715e9db8c3d95bf2c17b72ba5f42e78f041fd96 Mon Sep 17 00:00:00 2001 From: Yuhki Yano Date: Mon, 7 Sep 2026 20:10:59 +0900 Subject: [PATCH 01/10] Fix Runner behavio on MPS --- CHANGELOG.md | 1 + onecomp/runner.py | 7 ++++++- 2 files changed, 7 insertions(+), 1 deletion(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 4ce22a2..4b1cb86 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -5,6 +5,7 @@ ### Bug Fix - Fix scale layout handling for `groupsize=-1` in RTN fallback. This fallback is used when an MoE expert receives no routed calibration tokens, because GPTQ cannot compute activation-based statistics for that expert. The fix keeps the fallback result compatible with GPTQ's per-channel dequantization path. +- Avoid MPS-specific Runner failure during quantization ## [v1.3.3] 2026-09-03 diff --git a/onecomp/runner.py b/onecomp/runner.py index 6076714..3ef8503 100644 --- a/onecomp/runner.py +++ b/onecomp/runner.py @@ -674,7 +674,12 @@ def quantize(self): def quantize_with_calibration(self): """Quantize the model with calibration""" - model = self.model_config.load_model() + if is_mps_device(self.model_config.get_device()): + # device_map="mps" is unstable for large sharded checkpoints; load on CPU then move. + model = self.model_config.load_model(device_map="cpu") + model = model.to("mps") + else: + model = self.model_config.load_model() logger = self.logger input_device = next(model.parameters()).device inputs = self.prepare_calibration_dataset(input_device, model=model) From 9666baf85edb479d52d613b1e2deea7226cca9f4 Mon Sep 17 00:00:00 2001 From: Yuhki Yano Date: Tue, 8 Sep 2026 15:21:53 +0900 Subject: [PATCH 02/10] Reflect review comments: Add test --- .../onecomp/runner/test_mps_model_loading.py | 81 +++++++++++++++++++ 1 file changed, 81 insertions(+) create mode 100644 tests/onecomp/runner/test_mps_model_loading.py diff --git a/tests/onecomp/runner/test_mps_model_loading.py b/tests/onecomp/runner/test_mps_model_loading.py new file mode 100644 index 0000000..a96bed9 --- /dev/null +++ b/tests/onecomp/runner/test_mps_model_loading.py @@ -0,0 +1,81 @@ +"""Unit tests for Runner MPS model-loading behavior. + +Copyright 2025-2026 Fujitsu Ltd. + +Author: Yuhki Yano +""" + +from unittest.mock import MagicMock + +import torch + +from onecomp.runner import Runner + + +def _mock_model(): + model = MagicMock(name="model") + param = torch.nn.Parameter(torch.empty(1)) + model.parameters.side_effect = lambda: iter([param]) + model.to.return_value = model + return model, param + + +def test_quantize_with_calibration_loads_mps_model_on_cpu_then_moves_to_mps(): + """MPS avoids device_map='mps' for large sharded checkpoints.""" + model, param = _mock_model() + + model_config = MagicMock(name="model_config") + model_config.get_device.return_value = torch.device("mps") + model_config.load_model.return_value = model + + quantizer = MagicMock(name="quantizer") + quantizer.name = "GPTQ_4bit" + quantizer.module_to_name = {} + + inputs = {"input_ids": torch.ones((1, 1), dtype=torch.long)} + + runner = Runner( + model_config=model_config, + quantizer=quantizer, + report_progress=False, + ) + runner.prepare_calibration_dataset = MagicMock(return_value=inputs) + + runner.quantize_with_calibration() + + model_config.load_model.assert_called_once_with(device_map="cpu") + model.to.assert_called_once_with("mps") + runner.prepare_calibration_dataset.assert_called_once_with(param.device, model=model) + quantizer.setup.assert_called_once_with(model) + model.assert_called_once_with(**inputs) + quantizer.execute_post_processing.assert_called_once_with() + + +def test_quantize_with_calibration_non_mps_uses_default_load_model(): + model, param = _mock_model() + + model_config = MagicMock(name="model_config") + model_config.get_device.return_value = torch.device("cpu") + model_config.load_model.return_value = model + + quantizer = MagicMock(name="quantizer") + quantizer.name = "GPTQ_4bit" + quantizer.module_to_name = {} + + inputs = {"input_ids": torch.ones((1, 1), dtype=torch.long)} + + runner = Runner( + model_config=model_config, + quantizer=quantizer, + report_progress=False, + ) + runner.prepare_calibration_dataset = MagicMock(return_value=inputs) + + runner.quantize_with_calibration() + + model_config.load_model.assert_called_once_with() + model.to.assert_not_called() + runner.prepare_calibration_dataset.assert_called_once_with(param.device, model=model) + quantizer.setup.assert_called_once_with(model) + model.assert_called_once_with(**inputs) + quantizer.execute_post_processing.assert_called_once_with() From dece6dc7de3b4c679127f51ca8ae8698847647fd Mon Sep 17 00:00:00 2001 From: Yuhki Yano Date: Wed, 9 Sep 2026 00:34:30 +0900 Subject: [PATCH 03/10] Reflect reviw comments --- .../runner_methods/chunked_quantization.py | 9 +++-- .../onecomp/runner/test_mps_model_loading.py | 36 +++++++++++++------ 2 files changed, 33 insertions(+), 12 deletions(-) diff --git a/onecomp/runner_methods/chunked_quantization.py b/onecomp/runner_methods/chunked_quantization.py index 401ba63..195b5e0 100644 --- a/onecomp/runner_methods/chunked_quantization.py +++ b/onecomp/runner_methods/chunked_quantization.py @@ -34,7 +34,7 @@ from onecomp.calibration import CalibrationConfig, prepare_calibration_dataset from onecomp.model_config import ModelConfig from onecomp.quantizer._quantizer import QuantizationResult, Quantizer -from onecomp.utils.device import empty_cache +from onecomp.utils.device import empty_cache, is_mps_device from onecomp.utils.quantization_progress import QuantizationProgressTracker logger = getLogger(__name__) @@ -74,7 +74,12 @@ def run_chunked_quantization( num_layers_per_group = calibration_config.num_layers_per_group # Load model - model = model_config.load_model() + if is_mps_device(model_config.get_device()): + # device_map="mps" is unstable for large sharded checkpoints; load on CPU then move. + model = model_config.load_model(device_map="cpu") + model = model.to("mps") + else: + model = model_config.load_model() tokenizer = model_config.load_tokenizer() input_device = next(model.parameters()).device diff --git a/tests/onecomp/runner/test_mps_model_loading.py b/tests/onecomp/runner/test_mps_model_loading.py index a96bed9..1d64a60 100644 --- a/tests/onecomp/runner/test_mps_model_loading.py +++ b/tests/onecomp/runner/test_mps_model_loading.py @@ -14,15 +14,26 @@ def _mock_model(): model = MagicMock(name="model") - param = torch.nn.Parameter(torch.empty(1)) - model.parameters.side_effect = lambda: iter([param]) - model.to.return_value = model - return model, param + current_device = torch.device("cpu") + + def mock_to(device): + nonlocal current_device + current_device = torch.device(device) + return model + + def mock_parameters(): + param = MagicMock(name="parameter") + param.device = current_device + return iter([param]) + + model.to.side_effect = mock_to + model.parameters.side_effect = mock_parameters + return model def test_quantize_with_calibration_loads_mps_model_on_cpu_then_moves_to_mps(): """MPS avoids device_map='mps' for large sharded checkpoints.""" - model, param = _mock_model() + model = _mock_model() model_config = MagicMock(name="model_config") model_config.get_device.return_value = torch.device("mps") @@ -45,14 +56,16 @@ def test_quantize_with_calibration_loads_mps_model_on_cpu_then_moves_to_mps(): model_config.load_model.assert_called_once_with(device_map="cpu") model.to.assert_called_once_with("mps") - runner.prepare_calibration_dataset.assert_called_once_with(param.device, model=model) - quantizer.setup.assert_called_once_with(model) - model.assert_called_once_with(**inputs) + runner.prepare_calibration_dataset.assert_called_once_with( + torch.device("mps"), + model=model, + ) quantizer.execute_post_processing.assert_called_once_with() def test_quantize_with_calibration_non_mps_uses_default_load_model(): - model, param = _mock_model() + """Non-MPS devices keep the default model-loading path.""" + model = _mock_model() model_config = MagicMock(name="model_config") model_config.get_device.return_value = torch.device("cpu") @@ -75,7 +88,10 @@ def test_quantize_with_calibration_non_mps_uses_default_load_model(): model_config.load_model.assert_called_once_with() model.to.assert_not_called() - runner.prepare_calibration_dataset.assert_called_once_with(param.device, model=model) + runner.prepare_calibration_dataset.assert_called_once_with( + torch.device("cpu"), + model=model, + ) quantizer.setup.assert_called_once_with(model) model.assert_called_once_with(**inputs) quantizer.execute_post_processing.assert_called_once_with() From 3ba897889e460f78bf9a3f7a4d10d116ae3598ad Mon Sep 17 00:00:00 2001 From: Yuhki Yano Date: Wed, 9 Sep 2026 07:15:21 +0900 Subject: [PATCH 04/10] Reflect review comments --- onecomp/model_config.py | 7 +- onecomp/runner.py | 7 +- .../onecomp/runner/test_mps_model_loading.py | 97 ------------------- tests/onecomp/test_model_config.py | 74 ++++++++++++++ 4 files changed, 80 insertions(+), 105 deletions(-) delete mode 100644 tests/onecomp/runner/test_mps_model_loading.py create mode 100644 tests/onecomp/test_model_config.py diff --git a/onecomp/model_config.py b/onecomp/model_config.py index 520eabb..2371b85 100644 --- a/onecomp/model_config.py +++ b/onecomp/model_config.py @@ -11,7 +11,7 @@ import torch from transformers import AutoConfig, AutoModelForCausalLM, AutoTokenizer -from .utils.device import get_default_device +from .utils.device import get_default_device, is_mps_device from .utils.dtype import needs_bfloat16 try: @@ -95,9 +95,10 @@ def load_model(self, device_map=None): If ``None`` (default), ``self.device`` is used. """ effective_device = device_map if device_map is not None else self.device + load_device = "cpu" if is_mps_device(effective_device) else effective_device kwargs = dict( dtype=self.dtype if self.dtype == "auto" else getattr(torch, self.dtype), - device_map=effective_device, + device_map=load_device, ) config = self.load_config() @@ -135,6 +136,8 @@ def load_model(self, device_map=None): raise self.logger.info("AutoModelForCausalLM failed; trying AutoModelForImageTextToText.") model = _AutoVLM.from_pretrained(self.get_model_id_or_path(), **kwargs) + if is_mps_device(effective_device): + model = model.to(effective_device) model.eval() self.logger.info("Model loaded with dtype=%s", next(model.parameters()).dtype) return model diff --git a/onecomp/runner.py b/onecomp/runner.py index 3ef8503..6076714 100644 --- a/onecomp/runner.py +++ b/onecomp/runner.py @@ -674,12 +674,7 @@ def quantize(self): def quantize_with_calibration(self): """Quantize the model with calibration""" - if is_mps_device(self.model_config.get_device()): - # device_map="mps" is unstable for large sharded checkpoints; load on CPU then move. - model = self.model_config.load_model(device_map="cpu") - model = model.to("mps") - else: - model = self.model_config.load_model() + model = self.model_config.load_model() logger = self.logger input_device = next(model.parameters()).device inputs = self.prepare_calibration_dataset(input_device, model=model) diff --git a/tests/onecomp/runner/test_mps_model_loading.py b/tests/onecomp/runner/test_mps_model_loading.py deleted file mode 100644 index 1d64a60..0000000 --- a/tests/onecomp/runner/test_mps_model_loading.py +++ /dev/null @@ -1,97 +0,0 @@ -"""Unit tests for Runner MPS model-loading behavior. - -Copyright 2025-2026 Fujitsu Ltd. - -Author: Yuhki Yano -""" - -from unittest.mock import MagicMock - -import torch - -from onecomp.runner import Runner - - -def _mock_model(): - model = MagicMock(name="model") - current_device = torch.device("cpu") - - def mock_to(device): - nonlocal current_device - current_device = torch.device(device) - return model - - def mock_parameters(): - param = MagicMock(name="parameter") - param.device = current_device - return iter([param]) - - model.to.side_effect = mock_to - model.parameters.side_effect = mock_parameters - return model - - -def test_quantize_with_calibration_loads_mps_model_on_cpu_then_moves_to_mps(): - """MPS avoids device_map='mps' for large sharded checkpoints.""" - model = _mock_model() - - model_config = MagicMock(name="model_config") - model_config.get_device.return_value = torch.device("mps") - model_config.load_model.return_value = model - - quantizer = MagicMock(name="quantizer") - quantizer.name = "GPTQ_4bit" - quantizer.module_to_name = {} - - inputs = {"input_ids": torch.ones((1, 1), dtype=torch.long)} - - runner = Runner( - model_config=model_config, - quantizer=quantizer, - report_progress=False, - ) - runner.prepare_calibration_dataset = MagicMock(return_value=inputs) - - runner.quantize_with_calibration() - - model_config.load_model.assert_called_once_with(device_map="cpu") - model.to.assert_called_once_with("mps") - runner.prepare_calibration_dataset.assert_called_once_with( - torch.device("mps"), - model=model, - ) - quantizer.execute_post_processing.assert_called_once_with() - - -def test_quantize_with_calibration_non_mps_uses_default_load_model(): - """Non-MPS devices keep the default model-loading path.""" - model = _mock_model() - - model_config = MagicMock(name="model_config") - model_config.get_device.return_value = torch.device("cpu") - model_config.load_model.return_value = model - - quantizer = MagicMock(name="quantizer") - quantizer.name = "GPTQ_4bit" - quantizer.module_to_name = {} - - inputs = {"input_ids": torch.ones((1, 1), dtype=torch.long)} - - runner = Runner( - model_config=model_config, - quantizer=quantizer, - report_progress=False, - ) - runner.prepare_calibration_dataset = MagicMock(return_value=inputs) - - runner.quantize_with_calibration() - - model_config.load_model.assert_called_once_with() - model.to.assert_not_called() - runner.prepare_calibration_dataset.assert_called_once_with( - torch.device("cpu"), - model=model, - ) - quantizer.setup.assert_called_once_with(model) - model.assert_called_once_with(**inputs) - quantizer.execute_post_processing.assert_called_once_with() diff --git a/tests/onecomp/test_model_config.py b/tests/onecomp/test_model_config.py new file mode 100644 index 0000000..029668f --- /dev/null +++ b/tests/onecomp/test_model_config.py @@ -0,0 +1,74 @@ +"""Unit tests for Runner MPS model-loading behavior. + +Copyright 2025-2026 Fujitsu Ltd. + +Author: Yuhki Yano +""" + +from types import SimpleNamespace +from unittest.mock import MagicMock + +import torch + +from onecomp.model_config import ModelConfig + + +def _mock_loaded_model(): + model = MagicMock(name="model") + model.parameters.side_effect = lambda: iter([torch.empty(1, dtype=torch.float16)]) + model.to.return_value = model + return model + + +def test_load_model_loads_mps_model_on_cpu_then_moves_to_mps(monkeypatch): + """MPS loads sharded checkpoints on CPU before moving the model to MPS.""" + model = _mock_loaded_model() + config = SimpleNamespace(quantization_config=None) + + load_config = MagicMock(return_value=config) + load_model = MagicMock(return_value=model) + + monkeypatch.setattr("onecomp.model_config.AutoConfig.from_pretrained", load_config) + monkeypatch.setattr( + "onecomp.model_config.AutoModelForCausalLM.from_pretrained", + load_model, + ) + + model_config = ModelConfig(model_id="test/model", device="mps") + loaded = model_config.load_model() + + assert loaded is model + load_model.assert_called_once_with( + "test/model", + dtype=torch.float16, + device_map="cpu", + ) + model.to.assert_called_once_with("mps") + model.eval.assert_called_once_with() + + +def test_load_model_keeps_non_mps_device_placement(monkeypatch): + """Non-MPS devices retain their requested device_map.""" + model = _mock_loaded_model() + config = SimpleNamespace(quantization_config=None) + + monkeypatch.setattr( + "onecomp.model_config.AutoConfig.from_pretrained", + MagicMock(return_value=config), + ) + load_model = MagicMock(return_value=model) + monkeypatch.setattr( + "onecomp.model_config.AutoModelForCausalLM.from_pretrained", + load_model, + ) + + model_config = ModelConfig(model_id="test/model", device="cpu") + model_config.load_model() + + load_model.assert_called_once_with( + "test/model", + dtype=torch.float16, + device_map="cpu", + ) + model.to.assert_not_called() + model.eval.assert_called_once_with() From 848fcfacdea39f942712ea28e646a3868e985d37 Mon Sep 17 00:00:00 2001 From: Yuhki Yano Date: Wed, 9 Sep 2026 07:18:33 +0900 Subject: [PATCH 05/10] Modify chunked quantization --- onecomp/runner_methods/chunked_quantization.py | 9 ++------- 1 file changed, 2 insertions(+), 7 deletions(-) diff --git a/onecomp/runner_methods/chunked_quantization.py b/onecomp/runner_methods/chunked_quantization.py index 195b5e0..401ba63 100644 --- a/onecomp/runner_methods/chunked_quantization.py +++ b/onecomp/runner_methods/chunked_quantization.py @@ -34,7 +34,7 @@ from onecomp.calibration import CalibrationConfig, prepare_calibration_dataset from onecomp.model_config import ModelConfig from onecomp.quantizer._quantizer import QuantizationResult, Quantizer -from onecomp.utils.device import empty_cache, is_mps_device +from onecomp.utils.device import empty_cache from onecomp.utils.quantization_progress import QuantizationProgressTracker logger = getLogger(__name__) @@ -74,12 +74,7 @@ def run_chunked_quantization( num_layers_per_group = calibration_config.num_layers_per_group # Load model - if is_mps_device(model_config.get_device()): - # device_map="mps" is unstable for large sharded checkpoints; load on CPU then move. - model = model_config.load_model(device_map="cpu") - model = model.to("mps") - else: - model = model_config.load_model() + model = model_config.load_model() tokenizer = model_config.load_tokenizer() input_device = next(model.parameters()).device From 65253cf45c428c0e72e8e70dfef27cda2a51c97c Mon Sep 17 00:00:00 2001 From: Yuhki Yano Date: Wed, 9 Sep 2026 07:23:08 +0900 Subject: [PATCH 06/10] Modify CHANGELOG --- CHANGELOG.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 4b1cb86..14923e7 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -5,7 +5,7 @@ ### Bug Fix - Fix scale layout handling for `groupsize=-1` in RTN fallback. This fallback is used when an MoE expert receives no routed calibration tokens, because GPTQ cannot compute activation-based statistics for that expert. The fix keeps the fallback result compatible with GPTQ's per-channel dequantization path. -- Avoid MPS-specific Runner failure during quantization +- Fix MPS loading of large sharded checkpoints by loading weights on CPU before moving the model to MPS. ## [v1.3.3] 2026-09-03 From f91ab46055531c3bac6bb66c5fc03f9590bef695 Mon Sep 17 00:00:00 2001 From: Yuhki Yano Date: Wed, 9 Sep 2026 15:50:33 +0900 Subject: [PATCH 07/10] Reflect review comments --- onecomp/model_config.py | 10 ++++++--- tests/onecomp/test_model_config.py | 34 ++++++++++++++++++++++++++++-- 2 files changed, 39 insertions(+), 5 deletions(-) diff --git a/onecomp/model_config.py b/onecomp/model_config.py index 2371b85..bba0b3d 100644 --- a/onecomp/model_config.py +++ b/onecomp/model_config.py @@ -95,7 +95,11 @@ def load_model(self, device_map=None): If ``None`` (default), ``self.device`` is used. """ effective_device = device_map if device_map is not None else self.device - load_device = "cpu" if is_mps_device(effective_device) else effective_device + if effective_device == "auto": + target_device = get_default_device() + else: + target_device = torch.device(effective_device) + load_device = "cpu" if is_mps_device(target_device) else effective_device kwargs = dict( dtype=self.dtype if self.dtype == "auto" else getattr(torch, self.dtype), device_map=load_device, @@ -136,8 +140,8 @@ def load_model(self, device_map=None): raise self.logger.info("AutoModelForCausalLM failed; trying AutoModelForImageTextToText.") model = _AutoVLM.from_pretrained(self.get_model_id_or_path(), **kwargs) - if is_mps_device(effective_device): - model = model.to(effective_device) + if is_mps_device(target_device): + model = model.to(target_device) model.eval() self.logger.info("Model loaded with dtype=%s", next(model.parameters()).dtype) return model diff --git a/tests/onecomp/test_model_config.py b/tests/onecomp/test_model_config.py index 029668f..65052c8 100644 --- a/tests/onecomp/test_model_config.py +++ b/tests/onecomp/test_model_config.py @@ -1,4 +1,4 @@ -"""Unit tests for Runner MPS model-loading behavior. +"""Unit tests for ModelConfig MPS model-loading behavior. Copyright 2025-2026 Fujitsu Ltd. @@ -43,7 +43,7 @@ def test_load_model_loads_mps_model_on_cpu_then_moves_to_mps(monkeypatch): dtype=torch.float16, device_map="cpu", ) - model.to.assert_called_once_with("mps") + model.to.assert_called_once_with(torch.device("mps")) model.eval.assert_called_once_with() @@ -72,3 +72,33 @@ def test_load_model_keeps_non_mps_device_placement(monkeypatch): ) model.to.assert_not_called() model.eval.assert_called_once_with() + + +def test_load_model_auto_uses_mps_workaround_when_mps_is_default(monkeypatch): + """device='auto' uses the MPS workaround when MPS is selected.""" + model = _mock_loaded_model() + config = SimpleNamespace(quantization_config=None) + load_model = MagicMock(return_value=model) + + monkeypatch.setattr( + "onecomp.model_config.AutoConfig.from_pretrained", + MagicMock(return_value=config), + ) + monkeypatch.setattr( + "onecomp.model_config.AutoModelForCausalLM.from_pretrained", + load_model, + ) + monkeypatch.setattr( + "onecomp.model_config.get_default_device", + MagicMock(return_value=torch.device("mps")), + ) + + ModelConfig(model_id="test/model", device="auto").load_model() + + load_model.assert_called_once_with( + "test/model", + dtype=torch.float16, + device_map="cpu", + ) + model.to.assert_called_once_with(torch.device("mps")) + model.eval.assert_called_once_with() From b4b46d9107aa6d519a5dd72d8bc3183f16ebaae5 Mon Sep 17 00:00:00 2001 From: Yuhki Yano Date: Wed, 9 Sep 2026 16:14:57 +0900 Subject: [PATCH 08/10] Reflect additional review comments --- onecomp/runner.py | 5 +++++ tests/onecomp/test_runner_check.py | 13 +++++++++++++ 2 files changed, 18 insertions(+) diff --git a/onecomp/runner.py b/onecomp/runner.py index 6076714..75b32bc 100644 --- a/onecomp/runner.py +++ b/onecomp/runner.py @@ -379,6 +379,11 @@ def check(self): if is_mps_device(device): if self.multi_gpu: raise ValueError("multi_gpu is not supported on MPS device.") + if batch_size is not None: + raise ValueError( + "MPS quantization does not support calibration_config.batch_size. " + "Remove batch_size from CalibrationConfig and run without chunked calibration." + ) all_quantizers = self.quantizers if self.quantizers is not None else [self.quantizer] for i, q in enumerate(all_quantizers): label = f"quantizers[{i}]" if self.quantizers else "quantizer" diff --git a/tests/onecomp/test_runner_check.py b/tests/onecomp/test_runner_check.py index 2bf4e48..96319f3 100644 --- a/tests/onecomp/test_runner_check.py +++ b/tests/onecomp/test_runner_check.py @@ -139,3 +139,16 @@ def test_autobit_low_target_on_mps_with_auto_dbf_disabled_passes(self): calibration_config=CalibrationConfig(max_length=128, num_calibration_samples=8), ) runner.check() + + def test_batch_size_is_not_supported_on_mps(self): + """MPS rejects chunked calibration requested via batch_size.""" + runner = Runner( + model_config=self._mps_model_config(), + quantizer=GPTQ(wbits=4, groupsize=128), + calibration_config=CalibrationConfig(batch_size=2), + ) + + with pytest.raises( + ValueError, match=r"MPS quantization does not support calibration_config\.batch_size" + ): + runner.check() From 9c34d5e66120149fd4159ba3a2038a2054e93ad0 Mon Sep 17 00:00:00 2001 From: Yuhki Yano Date: Fri, 11 Sep 2026 12:29:10 +0900 Subject: [PATCH 09/10] Reflect additional review comments --- CHANGELOG.md | 1 + docs/user-guide/mps.md | 3 ++ onecomp/model_config.py | 2 +- tests/onecomp/test_model_config.py | 47 ++++++++++++++++++++++++++---- 4 files changed, 47 insertions(+), 6 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 14923e7..41faec5 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -6,6 +6,7 @@ - Fix scale layout handling for `groupsize=-1` in RTN fallback. This fallback is used when an MoE expert receives no routed calibration tokens, because GPTQ cannot compute activation-based statistics for that expert. The fix keeps the fallback result compatible with GPTQ's per-channel dequantization path. - Fix MPS loading of large sharded checkpoints by loading weights on CPU before moving the model to MPS. +- Reject chunked calibration (`CalibrationConfig(batch_size=...)`) on MPS, where it is not supported. ## [v1.3.3] 2026-09-03 diff --git a/docs/user-guide/mps.md b/docs/user-guide/mps.md index 2282e5c..e9bd96e 100644 --- a/docs/user-guide/mps.md +++ b/docs/user-guide/mps.md @@ -23,6 +23,7 @@ quantization steps intentionally run on CPU for performance. See | DBF, RTN, JointQ, and other quantizers | No | — | | AutoBit DBF fallback | No | — | | Multi-GPU quantization | No | — | +| Chunked calibration | No | - | ## Installation @@ -127,6 +128,8 @@ with an NVIDIA GPU. See the [vLLM Inference guide](vllm-inference.md). - Only **GPTQ** quantizers are allowed (or **AutoBitQuantizer** whose candidates are all GPTQ). - **AutoBit DBF fallback** is rejected when the target bitwidth would require DBF-only assignment. - **`multi_gpu=True`** is not supported. +- **Chunked calibration** (`CalibrationConfig(batch_size=...)`) is not supported. + Leave `batch_size` unset when quantizing on MPS. To avoid DBF fallback on MPS, either set an explicit `wbits` within the GPTQ candidate range or ensure VRAM estimation yields a bitwidth that does not trigger DBF-only paths. diff --git a/onecomp/model_config.py b/onecomp/model_config.py index bba0b3d..2e1061c 100644 --- a/onecomp/model_config.py +++ b/onecomp/model_config.py @@ -98,7 +98,7 @@ def load_model(self, device_map=None): if effective_device == "auto": target_device = get_default_device() else: - target_device = torch.device(effective_device) + target_device = effective_device load_device = "cpu" if is_mps_device(target_device) else effective_device kwargs = dict( dtype=self.dtype if self.dtype == "auto" else getattr(torch, self.dtype), diff --git a/tests/onecomp/test_model_config.py b/tests/onecomp/test_model_config.py index 65052c8..251c6f1 100644 --- a/tests/onecomp/test_model_config.py +++ b/tests/onecomp/test_model_config.py @@ -8,6 +8,7 @@ from types import SimpleNamespace from unittest.mock import MagicMock +import pytest import torch from onecomp.model_config import ModelConfig @@ -43,12 +44,13 @@ def test_load_model_loads_mps_model_on_cpu_then_moves_to_mps(monkeypatch): dtype=torch.float16, device_map="cpu", ) - model.to.assert_called_once_with(torch.device("mps")) + model.to.assert_called_once_with("mps") model.eval.assert_called_once_with() -def test_load_model_keeps_non_mps_device_placement(monkeypatch): - """Non-MPS devices retain their requested device_map.""" +@pytest.mark.parametrize("device", ["cpu", "cuda:0"]) +def test_load_model_preserves_explicit_no_mps_device_map(monkeypatch, device): + """Explicit non-MPS device maps are passed through unchanged.""" model = _mock_loaded_model() config = SimpleNamespace(quantization_config=None) @@ -62,13 +64,13 @@ def test_load_model_keeps_non_mps_device_placement(monkeypatch): load_model, ) - model_config = ModelConfig(model_id="test/model", device="cpu") + model_config = ModelConfig(model_id="test/model", device=device) model_config.load_model() load_model.assert_called_once_with( "test/model", dtype=torch.float16, - device_map="cpu", + device_map=device, ) model.to.assert_not_called() model.eval.assert_called_once_with() @@ -102,3 +104,38 @@ def test_load_model_auto_uses_mps_workaround_when_mps_is_default(monkeypatch): ) model.to.assert_called_once_with(torch.device("mps")) model.eval.assert_called_once_with() + + +@pytest.mark.parametrize( + "resolved_device", + [torch.device("cpu"), torch.device("cuda")], +) +def test_load_model_auto_preserves_auto_for_non_mps_default( + monkeypatch, + resolved_device, +): + """device='auto' keeps Transformers auto placement unless MPS is selected.""" + model = _mock_loaded_model() + config = SimpleNamespace(quantization_config=None) + load_model = MagicMock(return_value=model) + monkeypatch.setattr( + "onecomp.model_config.AutoConfig.from_pretrained", + MagicMock(return_value=config), + ) + monkeypatch.setattr( + "onecomp.model_config.AutoModelForCausalLM.from_pretrained", + load_model, + ) + monkeypatch.setattr( + "onecomp.model_config.get_default_device", + MagicMock(return_value=resolved_device), + ) + ModelConfig(model_id="test/model", device="auto").load_model() + + load_model.assert_called_once_with( + "test/model", + dtype=torch.float16, + device_map="auto", + ) + model.to.assert_not_called() + model.eval.assert_called_once_with() From 957f9ea35bb657edf7364c1af23503a8c1128d7a Mon Sep 17 00:00:00 2001 From: Yuhki Yano Date: Sat, 12 Sep 2026 14:55:38 +0900 Subject: [PATCH 10/10] Modify test model config --- tests/onecomp/test_model_config.py | 49 +++++++----------------------- 1 file changed, 11 insertions(+), 38 deletions(-) diff --git a/tests/onecomp/test_model_config.py b/tests/onecomp/test_model_config.py index 251c6f1..8cf5271 100644 --- a/tests/onecomp/test_model_config.py +++ b/tests/onecomp/test_model_config.py @@ -76,44 +76,14 @@ def test_load_model_preserves_explicit_no_mps_device_map(monkeypatch, device): model.eval.assert_called_once_with() -def test_load_model_auto_uses_mps_workaround_when_mps_is_default(monkeypatch): - """device='auto' uses the MPS workaround when MPS is selected.""" - model = _mock_loaded_model() - config = SimpleNamespace(quantization_config=None) - load_model = MagicMock(return_value=model) - - monkeypatch.setattr( - "onecomp.model_config.AutoConfig.from_pretrained", - MagicMock(return_value=config), - ) - monkeypatch.setattr( - "onecomp.model_config.AutoModelForCausalLM.from_pretrained", - load_model, - ) - monkeypatch.setattr( - "onecomp.model_config.get_default_device", - MagicMock(return_value=torch.device("mps")), - ) - - ModelConfig(model_id="test/model", device="auto").load_model() - - load_model.assert_called_once_with( - "test/model", - dtype=torch.float16, - device_map="cpu", - ) - model.to.assert_called_once_with(torch.device("mps")) - model.eval.assert_called_once_with() - - @pytest.mark.parametrize( - "resolved_device", - [torch.device("cpu"), torch.device("cuda")], + "resolved_device, expected_device_map", + [ + (torch.device("cpu"), "auto"), + (torch.device("mps"), "cpu"), + ], ) -def test_load_model_auto_preserves_auto_for_non_mps_default( - monkeypatch, - resolved_device, -): +def test_load_model_auto_preserves_device_map(monkeypatch, resolved_device, expected_device_map): """device='auto' keeps Transformers auto placement unless MPS is selected.""" model = _mock_loaded_model() config = SimpleNamespace(quantization_config=None) @@ -135,7 +105,10 @@ def test_load_model_auto_preserves_auto_for_non_mps_default( load_model.assert_called_once_with( "test/model", dtype=torch.float16, - device_map="auto", + device_map=expected_device_map, ) - model.to.assert_not_called() + if resolved_device.type == "mps": + model.to.assert_called_once_with(resolved_device) + else: + model.to.assert_not_called() model.eval.assert_called_once_with()