diff --git a/CHANGELOG.md b/CHANGELOG.md index 4ce22a2..41faec5 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -5,6 +5,8 @@ ### 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. +- 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 520eabb..2e1061c 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,14 @@ 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 + if effective_device == "auto": + target_device = get_default_device() + else: + 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), - device_map=effective_device, + device_map=load_device, ) config = self.load_config() @@ -135,6 +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(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/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_model_config.py b/tests/onecomp/test_model_config.py new file mode 100644 index 0000000..8cf5271 --- /dev/null +++ b/tests/onecomp/test_model_config.py @@ -0,0 +1,114 @@ +"""Unit tests for ModelConfig MPS model-loading behavior. + +Copyright 2025-2026 Fujitsu Ltd. + +Author: Yuhki Yano +""" + +from types import SimpleNamespace +from unittest.mock import MagicMock + +import pytest +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() + + +@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) + + 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=device) + model_config.load_model() + + load_model.assert_called_once_with( + "test/model", + dtype=torch.float16, + device_map=device, + ) + model.to.assert_not_called() + model.eval.assert_called_once_with() + + +@pytest.mark.parametrize( + "resolved_device, expected_device_map", + [ + (torch.device("cpu"), "auto"), + (torch.device("mps"), "cpu"), + ], +) +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) + 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=expected_device_map, + ) + 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() 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()