Fix non-QEP Runner quantization for large sharded MPS checkpoints

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
 

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.

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

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"

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()

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()