Add MPS (Apple Silicon) device support (#7)

* detect current device and be applicable for mps (apple silicon)

* calculate NLL with cpu even user specify for using mps

* send tensors to cpu at the begging of the run_gtpq

* debug for aligning v1.1.0 release

* raise error if quantizer is not gptq

* fix settings and readme

* bugfix: inference

* remove unused import

* fix: comment No3, 4, and 5

* fix: Comment No.7,8, and, 12

* add new test file for comment No.12

* fix: Comment No6 and 10

* fix Comment No2

* fix: comment No1

* fix No.9; add docs

* update lock file

* fix conflict

* fix conflict in CHANGELOG

* add docs for mps

Author: aki916

CHANGELOG.md

@@ -1,5 +1,13 @@
 # Change log
 
+## [v1.2.0+mps] 2026-06-05
+
+### Apple Silicon / macOS support
+
+- **MPS quantization**: GPTQ (and AutoBit with GPTQ-only candidates) on `device="mps"`; cross-platform `empty_cache()` via new `onecomp/utils/device.py` (`runner.py`, `quantizer/gptq/_gptq.py`, `quantizer/_quantizer.py`)
+- **MPS device placement (GPTQ on CPU, QEP correction on MPS)**: With `device="mps"`, `run_gptq` moves the Hessian and weights to **CPU** for the full column-wise GPTQ loop (including inverse-Hessian Cholesky). The main reason is not absent Cholesky kernels on MPS (recent PyTorch supports them); if the GPTQ loop stayed on MPS, `maxq.item()` inside `quantize()` would run once per column—each call waits for pending MPS work to finish and read back a single scalar to the host (per-column host sync), not a full matrix copy per column—and that overhead is often several times slower than CPU on Apple Silicon (~4× in internal benchmarks with PyTorch 2.12). When QEP weight correction runs (`adjust_weight`, typically under `qep=True`), per-layer work stays on **MPS** (e.g. `weight @ delta_hatX`); only the Cholesky solve uses CPU via `_safe_cholesky_and_solve` (one solve per layer). A full CPU fallback for QEP does not materially improve speed. Calibration forwards may still use MPS. Details: README (macOS / MPS).
+- **MPS inference**: load saved quantized models on Mac with `QuantizedModelLoader` + Transformers `generate()` (GemLite/vLLM remain Linux + CUDA)
+- **macOS `uv sync`**: added `darwin` to `tool.uv.environments`, `--extra mps` for MPS-enabled PyTorch from PyPI; `--extra cpu` is Linux-only (pytorch-cpu index); Linux-only markers on CUDA extras (`cu118`–`cu130`)
 
 ## [v1.2.0] 2026-06-04 (WIP)
 

README.md

@@ -100,6 +100,30 @@ import torch
 print(torch.cuda.is_available())
 ```
 
+#### ✅ macOS (MPS)
+
+On macOS, install PyTorch from PyPI (default wheels include MPS support). You do **not** need the CUDA index URLs above.
+
+```bash
+pip install torch torchvision torchaudio
+```
+
+Verify MPS:
+```python
+import torch
+print(torch.backends.mps.is_available())
+```
+
+Then install OneComp from PyPI (see step 2 below). GPTQ quantization and Hugging Face `generate()` inference on MPS are supported; vLLM serving requires Linux with an NVIDIA GPU. An editable install from a git clone is **not** required for MPS use — see [for developers (pip)](#for-developers-pip) only if you are contributing to OneComp.
+
+> **MPS device placement (GPTQ vs QEP)**  
+> With `device="mps"`, calibration and model forward passes can run on the GPU. The bottleneck is usually not missing Cholesky ops on MPS (recent PyTorch builds implement them); the implementation splits work as follows:
+>
+> - GPTQ (`run_gptq`): Hessian and weights are moved to CPU for the full column-wise loop (including inverse-Hessian Cholesky). If that loop stayed on MPS, `quantize()` would call `maxq.item()` once per column; each call triggers **per-column host sync** (wait for pending MPS ops, then read one scalar—not a full Hessian/weight copy every column)—often several times slower than CPU on Apple Silicon (e.g. ~4× in internal benchmarks with PyTorch 2.12). Keeping GPTQ on CPU avoids that overhead. With `mse=True`, `find_params` also calls `quantize()` in a grid loop and benefits from the same CPU placement.
+> - QEP weight correction (`adjust_weight`, when QEP correction runs—typically `qep=True` with error propagation enabled): Per-layer work stays on MPS (e.g. `weight @ delta_hatX`, diagonal damping). Only the Cholesky solve uses CPU via `_safe_cholesky_and_solve` (one solve per layer, not per column); moving all of QEP to CPU does not materially improve speed. The subsequent GPTQ step still uses the CPU path above.
+>
+> DBF-based AutoBit fallback and multi-GPU quantization are not supported on MPS.
+
 #### 2. Install `onecomp`
 
 Once PyTorch is installed, you can install `onecomp`:
@@ -122,26 +146,41 @@ curl -LsSf https://astral.sh/uv/install.sh | sh
 
 git clone https://github.com/FujitsuResearch/OneCompression.git
 cd OneCompression
-uv sync --extra cu128 --extra dev --extra visualize
 ```
 
 The `uv sync` command creates a Python virtual environment and installs all dependent libraries.
 
+#### Linux (CUDA quantization / vLLM)
+
+```bash
+uv sync --extra cu128 --extra dev --extra visualize
+```
+
 The `--extra cu128` option installs the CUDA-enabled version of PyTorch (along with `torchvision` from the same CUDA index).
 Replace `cu128` with the appropriate variant for your environment: `cpu`, `cu118`, `cu121`, `cu124`, `cu126`, `cu128`, or `cu130`.
 PyTorch will be automatically downloaded by `uv`, so you do not need to install it beforehand.
 
+#### macOS (development / MPS inference)
+
+```bash
+uv sync --extra mps --extra dev --extra visualize
+```
+
+On macOS, use `--extra mps` only. CUDA extras (`cu118`–`cu130`), `--extra cpu` (Linux-only), and `--extra vllm` are not supported on macOS.
+After `uv sync`, you can run GPTQ quantization and Hugging Face `generate()` inference on MPS; vLLM serving still requires Linux with an NVIDIA GPU.
+See the **MPS device placement (GPTQ vs QEP)** note under [macOS (MPS)](#macos-mps) above for why GPTQ runs on CPU while QEP correction uses MPS.
+
 Adding `--extra dev` installs development tools (black, pre-commit, pytest, pylint).
 Adding `--extra visualize` installs matplotlib for visualization features.
 Adding `--extra hydra` installs `hydra-core` for the example scripts and `model_validation/` runners that use Hydra-based configuration.
 
-To use vLLM for serving quantized models, add `--extra vllm` together with `--extra cu130`:
+To use vLLM for serving quantized models on Linux, add `--extra vllm` together with `--extra cu130`:
 
 ```bash
 uv sync --extra cu130 --extra dev --extra visualize --extra vllm
 ```
 
-> **Note:** `--extra vllm` is only compatible with `--extra cu130`. Recent vLLM releases require `torch>=2.10`, whose wheels are only published for the `cu130` index. Combining `--extra vllm` with `cpu` / `cu118` / `cu121` / `cu124` / `cu126` / `cu128` is rejected by `uv` at lock time.
+> **Note:** `--extra vllm` is only compatible with `--extra cu130`. Recent vLLM releases require `torch>=2.10`, whose wheels are only published for the `cu130` index. Combining `--extra vllm` with `cpu` / `mps` / `cu118` / `cu121` / `cu124` / `cu126` / `cu128` is rejected by `uv` at lock time.
 
 > **Note:** `--extra vllm` may take a long time on the first run if a pre-built `xformers` wheel is not available for your Python/CUDA combination (e.g. Python 3.13). Using Python 3.12 typically avoids this.
 
@@ -168,6 +207,8 @@ black --check onecomp/
 
 ### for developers (pip)
 
+> **Note:** The editable install below is for developing OneComp from a local clone. **macOS users who only want MPS inference or quantization should use the [for users (pip)](#for-users-pip) flow** (`pip install torch` then `pip install onecomp` from PyPI); `pip install -e` is not needed for MPS.
+
 ```bash
 git clone <git repository URL>
 cd OneCompression
@@ -217,8 +258,10 @@ uv run pre-commit run --all-files
 
 ### Building Documentation Locally
 
+`--extra docs` alone is sufficient (no PyTorch `mps` / `cu*` extra required):
+
 ```bash
-uv sync --extra cu128 --extra dev --extra docs
+uv sync --extra docs
 uv run mkdocs serve
 ```
 

docs/api/quantized_model_loader.md

@@ -2,6 +2,11 @@
 
 Loader for quantized models saved by OneComp.
 
+On macOS, `load_quantized_model()` places the model on MPS when available
+(CUDA > MPS > CPU via `get_default_device()`). Use Transformers `generate()` for
+inference; vLLM requires Linux with an NVIDIA GPU. See the
+[macOS / MPS guide](../user-guide/mps.md#inference-with-transformers).
+
 ::: onecomp.quantized_model_loader.QuantizedModelLoader
     options:
       show_source: false

docs/getting-started/installation.md

@@ -5,7 +5,7 @@ This page describes how to install Fujitsu One Compression (OneComp).
 ## Requirements
 
 - Python 3.12 or later (< 3.14)
-- PyTorch (CPU or CUDA)
+- PyTorch (CPU, CUDA, or MPS on macOS)
 
 ## For Users (pip)
 
@@ -55,15 +55,40 @@ Install the appropriate version of PyTorch for your system.
     pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu130
     ```
 
-Check your CUDA version:
+=== "macOS (MPS)"
+
+    On macOS, install PyTorch from PyPI (default wheels include MPS support).
+    You do **not** need the CUDA index URLs above.
+
+    ```bash
+    pip install torch torchvision torchaudio
+    ```
+
+    Verify MPS:
+
+    ```python
+    import torch
+    print(torch.backends.mps.is_available())
+    ```
+
+    Then install OneComp (step 2 below). GPTQ quantization and Hugging Face
+    `generate()` inference on MPS are supported; vLLM serving requires Linux with
+    an NVIDIA GPU. An editable install from a git clone is **not** required for
+    MPS use — see [For Developers (pip)](#for-developers-pip) only if you are
+    contributing to OneComp.
+
+    For usage (`device="mps"`, VRAM budget, limitations), see the
+    [macOS / MPS guide](../user-guide/mps.md).
+
+Check your CUDA version (Linux / Windows with NVIDIA GPU):
 
 ```bash
 nvcc --version
 # or
 nvidia-smi
 ```
 
-Verify PyTorch GPU support:
+Verify PyTorch GPU support (CUDA):
 
 ```python
 import torch
@@ -91,26 +116,45 @@ It provides deterministic, reproducible environments via its lockfile.
 # Install uv (macOS or Linux)
 curl -LsSf https://astral.sh/uv/install.sh | sh
 
-# Clone and set up
 git clone https://github.com/FujitsuResearch/OneCompression.git
 cd OneCompression
+```
+
+The `uv sync` command creates a virtual environment and installs all dependencies.
+
+### Linux (CUDA quantization / vLLM)
+
+```bash
 uv sync --extra cu128 --extra dev --extra visualize
 ```
 
-The `uv sync` command creates a virtual environment and installs all dependencies (including `torchvision` from the same CUDA index as PyTorch).
-Replace `cu128` with the appropriate CUDA variant for your system: `cpu`, `cu118`, `cu121`, `cu124`, `cu126`, `cu128`, or `cu130`.
+The `--extra cu128` option installs the CUDA-enabled version of PyTorch (along with `torchvision` from the same CUDA index).
+Replace `cu128` with the appropriate variant for your environment: `cpu`, `cu118`, `cu121`, `cu124`, `cu126`, `cu128`, or `cu130`.
+PyTorch will be automatically downloaded by `uv`, so you do not need to install it beforehand.
+
+### macOS (development / MPS inference)
+
+```bash
+uv sync --extra mps --extra dev --extra visualize
+```
+
+On macOS, use `--extra mps` only. CUDA extras (`cu118`–`cu130`), `--extra cpu` (Linux-only),
+and `--extra vllm` are not supported on macOS.
+After `uv sync`, you can run GPTQ quantization and Hugging Face `generate()` inference on MPS;
+vLLM serving still requires Linux with an NVIDIA GPU.
+See the [macOS / MPS guide](../user-guide/mps.md) for device placement and usage details.
 
 Adding `--extra dev` installs development tools (black, pytest, pylint).
 Adding `--extra visualize` installs matplotlib for visualization features.
 
-To use vLLM for serving quantized models, add `--extra vllm` together with `--extra cu130`:
+To use vLLM for serving quantized models on Linux, add `--extra vllm` together with `--extra cu130`:
 
 ```bash
 uv sync --extra cu130 --extra dev --extra visualize --extra vllm
 ```
 
 !!! note "vLLM requires the `cu130` extra"
-    Recent vLLM releases depend on `torch>=2.10`, whose wheels are only published for the `cu130` index. The `--extra vllm` declaration in `pyproject.toml` therefore conflicts with `cpu`, `cu118`, `cu121`, `cu124`, `cu126`, and `cu128`; combining any of these with `--extra vllm` is rejected by `uv` at lock time.
+    Recent vLLM releases depend on `torch>=2.10`, whose wheels are only published for the `cu130` index. The `--extra vllm` declaration in `pyproject.toml` therefore conflicts with `cpu`, `mps`, `cu118`, `cu121`, `cu124`, `cu126`, and `cu128`; combining any of these with `--extra vllm` is rejected by `uv` at lock time.
 
 !!! warning
     Do **not** install vLLM with `uv pip install vllm` after `uv sync`. Packages installed via `uv pip` are not tracked by the lockfile and will be removed or overwritten by subsequent `uv sync` or `uv run` commands. Always use `--extra vllm` instead.
@@ -138,19 +182,30 @@ uv sync --extra cu130 --extra dev --extra visualize --extra vllm
 
 ## For Developers (pip)
 
+!!! note
+    The editable install below is for developing OneComp from a local clone.
+    **macOS users who only want MPS inference or quantization should use the
+    [For Users (pip)](#for-users-pip) flow** (`pip install torch` then
+    `pip install onecomp` from PyPI); `pip install -e` is not needed for MPS.
+
 ```bash
 git clone https://github.com/FujitsuResearch/OneCompression.git
 cd OneCompression
 
-# Install PyTorch with CUDA support
+# First, install PyTorch for your environment
 pip install torch --index-url https://download.pytorch.org/whl/cu128
-
-# Install onecomp with development dependencies
+# Then install onecomp with development dependencies
 pip install -e ".[dev]"
 ```
 
+Replace `cu128` with the appropriate variant for your environment: `cpu`, `cu118`, `cu121`, `cu124`, `cu126`, `cu128`, or `cu130`.
+On macOS, install PyTorch from PyPI instead (see [macOS (MPS)](#step-1-install-pytorch) above).
+
 ## Building Documentation Locally
 
+`--extra docs` alone is enough. PyTorch extras (`mps`, `cu*`, `cpu`) are not required
+to build or serve the documentation.
+
 ```bash
 uv sync --extra docs
 uv run mkdocs serve

docs/getting-started/quickstart.md

@@ -31,9 +31,9 @@ quantized model is saved to `TinyLlama-1.1B-...-autobit-<X>bit/` by default.
 |------------------------|------------|----------------------------------------------------------|
 | `model_id`             | (required) | Hugging Face model ID or local path                      |
 | `wbits`                | `None`     | Target bitwidth. When `None`, estimated from VRAM        |
-| `total_vram_gb`        | `None`     | VRAM budget in GB. When `None`, detected from GPU        |
+| `total_vram_gb`        | `None`     | VRAM budget in GB. When `None`, detected from CUDA GPU   |
 | `groupsize`            | `128`      | GPTQ group size (`-1` to disable)                        |
-| `device`               | `"cuda:0"` | Device for computation                                   |
+| `device`               | `"cuda:0"` | Device for computation (`"mps"` on macOS — see below)    |
 | `qep`                  | `True`     | Enable QEP (Quantization Error Propagation)              |
 | `evaluate`             | `True`     | Calculate perplexity and zero-shot accuracy              |
 | `eval_original_model`  | `False`    | Also evaluate the original (unquantized) model           |
@@ -66,6 +66,21 @@ Runner.auto_run(
 )
 ```
 
+### macOS (Apple Silicon)
+
+On Mac, set `device="mps"` and pass `total_vram_gb` (VRAM auto-detection uses CUDA only):
+
+
+```python
+Runner.auto_run(
+    model_id="TinyLlama/TinyLlama-1.1B-intermediate-step-1431k-3T",
+    device="mps",
+    total_vram_gb=16,
+)
+```
+See the [macOS / MPS guide](../user-guide/mps.md) for supported features, device
+placement (GPTQ vs QEP), and inference with Transformers.
+
 ---
 
 ## Step-by-step Workflow
@@ -177,3 +192,4 @@ model, tokenizer = load_quantized_model("./output/quantized_model")
 - [Configuration](../user-guide/configuration.md) -- detailed explanation of `ModelConfig`, `QEPConfig`, `LPCDConfig`, and `Runner` parameters
 - [Examples](../user-guide/examples.md) -- more usage patterns including multi-GPU and chunked calibration
 - [Algorithms](../algorithms/overview.md) -- learn about the quantization algorithms available in OneComp
+- [macOS / MPS](../user-guide/mps.md) -- Apple Silicon setup, limitations, and inference

docs/user-guide/basic-usage.md

@@ -60,10 +60,15 @@ model_config = ModelConfig(
 | `model_id` | Hugging Face Hub model ID          | —            |
 | `path`     | Local path to a saved model         | —            |
 | `dtype`    | Data type (`"float16"`, `"float32"`)| `"float16"`  |
-| `device`   | Device (`"cpu"`, `"cuda"`, `"auto"`)| `"auto"`     |
+| `device`   | Device (`"cpu"`, `"cuda"`, `"mps"`, `"auto"`)| `"auto"`     |
 
 You must provide either `model_id` or `path`.
 
+!!! tip "macOS (Apple Silicon)"
+    Use `device="mps"` for quantization on Mac. `Runner.auto_run` defaults to
+    `cuda:0`; pass `device="mps"` and `total_vram_gb` explicitly. See the
+    [macOS / MPS guide](mps.md).
+
 ## Step 2: Choose a Quantizer
 
 ```python

docs/user-guide/cli.md

@@ -43,9 +43,9 @@ onecomp [-h] [--wbits WBITS] [--total-vram-gb GB] [--groupsize GROUPSIZE]
 | Option                    | Default      | Description                                              |
 |---------------------------|--------------|----------------------------------------------------------|
 | `--wbits WBITS`           | `None` (auto)| Target bitwidth. When omitted, estimated from VRAM       |
-| `--total-vram-gb GB`      | `None` (auto)| VRAM budget in GB for bitwidth estimation. When omitted, detected from GPU |
+| `--total-vram-gb GB`      | `None` (auto)| VRAM budget in GB for bitwidth estimation. When omitted, detected from CUDA GPU. **Required on MPS** when `--wbits` is omitted |
 | `--groupsize GROUPSIZE`   | `128`        | GPTQ group size (`-1` to disable grouping)               |
-| `--device DEVICE`         | `cuda:0`     | Device to place the model on                             |
+| `--device DEVICE`         | `cuda:0`     | Device to place the model on (`mps` on macOS)            |
 | `--no-qep`                |              | Disable QEP (enabled by default)                         |
 | `--no-eval`               |              | Skip perplexity and accuracy evaluation                  |
 | `--eval-original`         |              | Also evaluate the original (unquantized) model           |
@@ -122,6 +122,8 @@ onecomp meta-llama/Llama-2-7b-hf --eval-original
 onecomp meta-llama/Llama-2-7b-hf --device cuda:1
 ```
 
+See the [macOS / MPS guide](mps.md) for supported quantizers and limitations.
+
 ## Default Behavior
 
 When run with no options, the `onecomp` command:

docs/user-guide/configuration.md

@@ -21,11 +21,16 @@ model_config = ModelConfig(
 | `model_id` | `str` | Hugging Face Hub model ID              | `None`       |
 | `path`     | `str` | Local path to model directory           | `None`       |
 | `dtype`    | `str` | Model precision (`"float16"`, `"float32"`)| `"float16"` |
-| `device`   | `str` | Device placement (`"cpu"`, `"cuda"`, `"auto"`)| `"auto"` |
+| `device`   | `str` | Device placement (`"cpu"`, `"cuda"`, `"mps"`, `"auto"`)| `"auto"` |
 
 !!! note
     Provide exactly one of `model_id` or `path`. A `ValueError` is raised if neither is specified.
 
+!!! note "macOS (MPS)"
+    On Apple Silicon, set `device="mps"` for GPTQ / AutoBit (GPTQ-only) quantization.
+    Only GPTQ quantizers are supported on MPS; DBF fallback and multi-GPU are not.
+    See the [macOS / MPS guide](mps.md) for details.
+
 ## Runner
 
 `Runner` is the main entry point for quantization. It manages the full pipeline: loading the model, preparing calibration data, executing quantization, and providing evaluation utilities.
@@ -148,7 +153,7 @@ qep_config = QEPConfig(
 | `general`                 | `bool`      | Use generic (architecture-independent) QEP       | `False`              |
 | `percdamp`                | `float`     | Damping percentage for Hessian regularization     | `0.01`               |
 | `perccorr`                | `float`     | Correction percentage for error propagation       | `0.5`                |
-| `device`                  | `str`       | GPU device for QEP computations                   | `"cuda:0"`           |
+| `device`                  | `str`       | Device for QEP computations (`"cuda"`, `"mps"`, `"cpu"`) | `"cuda:0"`           |
 | `exclude_layer_keywords`  | `list[str]` | Layer keywords excluded from error propagation    | `["mlp.down_proj"]`  |
 
 !!! tip

docs/user-guide/examples.md

@@ -370,6 +370,9 @@ outputs = model.generate(**inputs, max_new_tokens=50)
 print(tokenizer.decode(outputs[0], skip_special_tokens=True))
 ```
 
+On macOS, the model is placed on MPS automatically when available. For vLLM serving,
+use a Linux machine with an NVIDIA GPU. See the [macOS / MPS guide](mps.md).
+
 ## Block-wise PTQ
 
 Apply block-wise post-training quantization to improve accuracy after GPTQ quantization: