From 84af3df0b90a54547ff25d6609ee4f6a99d95e6b Mon Sep 17 00:00:00 2001 From: kimura-keiji Date: Mon, 7 Sep 2026 17:08:51 +0900 Subject: [PATCH 01/17] Define v1.3.4 --- CHANGELOG.md | 2 ++ onecomp/__version__.py | 2 +- 2 files changed, 3 insertions(+), 1 deletion(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index a913548..cd9c207 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,5 +1,7 @@ # Change log +## [v1.3.4] 2026-09-07 + ## [v1.3.3] 2026-09-03 ### Enhancement diff --git a/onecomp/__version__.py b/onecomp/__version__.py index 022c0c7..17be442 100644 --- a/onecomp/__version__.py +++ b/onecomp/__version__.py @@ -6,4 +6,4 @@ """ -__version__ = "1.3.3" +__version__ = "1.3.4" From 75fb6a0614fba30f173aef2da3cc094527571fbb Mon Sep 17 00:00:00 2001 From: Yuhki Yano Date: Mon, 7 Sep 2026 20:05:27 +0900 Subject: [PATCH 02/17] Fix RTN fallback scale layout for groupsize=-1 --- CHANGELOG.md | 4 ++++ onecomp/qep/_quantize_with_qep_arch.py | 15 +++++++++++---- .../test_qep_expert_recovery_integration.py | 8 +++++--- tests/onecomp/test_rtn_fallback_result.py | 13 +++++++++++-- 4 files changed, 31 insertions(+), 9 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index cd9c207..9c15f88 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -2,6 +2,10 @@ ## [v1.3.4] 2026-09-07 +### Bug Fix + +- Fix RTN fallback scale layout for groupsize=-1 + ## [v1.3.3] 2026-09-03 ### Enhancement diff --git a/onecomp/qep/_quantize_with_qep_arch.py b/onecomp/qep/_quantize_with_qep_arch.py index 1bd3402..811de0c 100644 --- a/onecomp/qep/_quantize_with_qep_arch.py +++ b/onecomp/qep/_quantize_with_qep_arch.py @@ -298,8 +298,15 @@ def _rtn_fallback_result(module: nn.Module, quantizer: Quantizer, name: str) -> result_dict = run_rtn(module, wbits=wbits, groupsize=groupsize, sym=quantizer.sym) - # RTN's raw scale/zero are (out_features, num_groups); GPTQResult - # expects (num_groups, out_features). + scales = result_dict["scale"] + qzeros = result_dict["zero"] + + if groupsize != -1: + # RTN's raw scale/zero are (out_features, num_groups); GPTQResult + # expects (num_groups, out_features). + scales = scales.T + qzeros = qzeros.T + return GPTQResult( dequantized_weight=result_dict["dequantized_weight"], wbits=wbits, @@ -307,8 +314,8 @@ def _rtn_fallback_result(module: nn.Module, quantizer: Quantizer, name: str) -> actorder=False, sym=quantizer.sym, qweight=result_dict["quantized_weight"], - scales=result_dict["scale"].T, - qzeros=result_dict["zero"].T, + scales=scales, + qzeros=qzeros, perm=None, ) diff --git a/tests/onecomp/test_qep_expert_recovery_integration.py b/tests/onecomp/test_qep_expert_recovery_integration.py index fae173c..276f8a5 100644 --- a/tests/onecomp/test_qep_expert_recovery_integration.py +++ b/tests/onecomp/test_qep_expert_recovery_integration.py @@ -12,6 +12,7 @@ import logging +import pytest import torch import torch.nn as nn import torch.nn.functional as F @@ -24,7 +25,7 @@ # in_features of every expert Linear must be divisible by the 4-bit pack factor # (32 // 4 == 8): gate/up consume HIDDEN, down consumes INTERMEDIATE. HIDDEN = 8 -INTERMEDIATE = 8 +INTERMEDIATE = 16 NUM_EXPERTS = 3 SEQ_LEN = 3 @@ -230,7 +231,8 @@ def test_expert_never_selected_falls_back_to_rtn(monkeypatch, caplog): assert all(quantizer.results[n].actorder is False for n in names) -def test_rtn_fallback_weight_is_actually_applied_to_the_module(monkeypatch, caplog): +@pytest.mark.parametrize("groupsize", [-1, 2]) +def test_rtn_fallback_weight_is_actually_applied_to_the_module(monkeypatch, caplog, groupsize): """The RTN-fallback ``GPTQResult`` must actually flow into the module's live weight, not just sit unused in ``quantizer.results``. A @@ -248,7 +250,7 @@ def test_rtn_fallback_weight_is_actually_applied_to_the_module(monkeypatch, capl "onecomp.qep._quantize_with_qep_arch.prepare_calibration_dataset", _fake_prepare_calibration_dataset, ) - quantizer = GPTQ(wbits=4, groupsize=2, sym=True, include_layer_keywords=["experts"]) + quantizer = GPTQ(wbits=4, groupsize=groupsize, sym=True, include_layer_keywords=["experts"]) qep_config = QEPConfig(device="cpu", percdamp=0.01, perccorr=0.5) caplog.set_level(logging.INFO, logger="onecomp.qep._quantize_with_qep_arch") diff --git a/tests/onecomp/test_rtn_fallback_result.py b/tests/onecomp/test_rtn_fallback_result.py index 4142149..2b4277a 100644 --- a/tests/onecomp/test_rtn_fallback_result.py +++ b/tests/onecomp/test_rtn_fallback_result.py @@ -8,6 +8,7 @@ Copyright 2025-2026 Fujitsu Ltd. """ +import pytest import torch import torch.nn as nn @@ -84,18 +85,26 @@ def test_dequantized_weight_matches_shape(self): result = _rtn_fallback_result(module, quantizer, "mlp.experts.0.down_proj") assert result.dequantized_weight.shape == module.weight.data.shape - def test_compute_dequantized_weight_roundtrip(self): + @pytest.mark.parametrize("groupsize", [-1, 32]) + def test_compute_dequantized_weight_roundtrip(self, groupsize): """The packaged qweight/scales/qzeros must reconstruct a weight consistent with the shapes GPTQResult.compute_dequantized_weight expects (this is what create_inference_layer / export relies on downstream). """ module = _linear(in_features=32, out_features=16) - quantizer = GPTQ(wbits=4, groupsize=16, sym=True) + quantizer = GPTQ(wbits=4, groupsize=groupsize, sym=True) result = _rtn_fallback_result(module, quantizer, "mlp.experts.0.down_proj") reconstructed = result.compute_dequantized_weight() + expected = result.dequantized_weight.to(dtype=reconstructed.dtype) assert reconstructed.shape == module.weight.data.shape + torch.testing.assert_close( + reconstructed, + expected, + rtol=0, + atol=2e-4, + ) class TestResolveGptqForRtnFallback: From f6adf4f0907ccaa2948ca45d3148b2badb9d0c21 Mon Sep 17 00:00:00 2001 From: kimura-keiji Date: Mon, 7 Sep 2026 22:54:50 +0900 Subject: [PATCH 03/17] Allow Pages job to run immediately --- .gitlab-ci.yml | 1 + 1 file changed, 1 insertion(+) diff --git a/.gitlab-ci.yml b/.gitlab-ci.yml index 6b75917..9da0291 100644 --- a/.gitlab-ci.yml +++ b/.gitlab-ci.yml @@ -288,6 +288,7 @@ sast-gate: pages: image: python:3.11 stage: deploy + needs: [] script: - pip install mkdocs mkdocs-material mkdocstrings mkdocstrings-python - mkdocs build --site-dir public From 3e1f5501cea81f0b218af9e0004ab755e166f7f5 Mon Sep 17 00:00:00 2001 From: Yuhki Yano Date: Tue, 8 Sep 2026 14:35:15 +0900 Subject: [PATCH 04/17] Reflect review comments --- CHANGELOG.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 9c15f88..4ce22a2 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,7 +4,7 @@ ### Bug Fix -- Fix RTN fallback scale layout for groupsize=-1 +- 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. ## [v1.3.3] 2026-09-03 From 0715e9db8c3d95bf2c17b72ba5f42e78f041fd96 Mon Sep 17 00:00:00 2001 From: Yuhki Yano Date: Mon, 7 Sep 2026 20:10:59 +0900 Subject: [PATCH 05/17] 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 06/17] 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 07/17] 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 700e485d2a60f6b33a1f82fe7caa7fefccb58cc5 Mon Sep 17 00:00:00 2001 From: Yuhki Yano <30323722+y-vectorfield@users.noreply.github.com> Date: Wed, 9 Sep 2026 01:29:51 +0900 Subject: [PATCH 08/17] Modify docs to clarify Qwen3.6 full-wrapper save format (#62) * Modify docs to clarify Qwen3.6 full-wrapper save format * Reflect review comments --- CHANGELOG.md | 4 ++++ docs/user-guide/basic-usage.md | 12 ++++++++++++ docs/user-guide/examples.md | 6 ++++++ docs/user-guide/vllm-inference.md | 21 ++++++--------------- 4 files changed, 28 insertions(+), 15 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 4ce22a2..2570915 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -6,6 +6,10 @@ - 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. +### Documentation + +- Clarify the confirmed Qwen3.6 `save_format="full_wrapper"` workflows, including vLLM serving and the current GGUF export workflow. + ## [v1.3.3] 2026-09-03 ### Enhancement diff --git a/docs/user-guide/basic-usage.md b/docs/user-guide/basic-usage.md index c91e9f8..62aba0c 100644 --- a/docs/user-guide/basic-usage.md +++ b/docs/user-guide/basic-usage.md @@ -155,6 +155,18 @@ runner.save_dequantized_model("./output/dequantized") runner.save_quantized_model("./output/quantized") ``` +!!! note "Qwen3.6 save format" + Qwen3.6 is quantized through its text-model layout. For confirmed + downstream workflows that require the full Hugging Face wrapper layout, + including vLLM serving and the current GGUF export workflow, save Qwen3.6 + checkpoints with `save_format="full_wrapper"`: + + `runner.save_quantized_model("./output/qwen36_quantized", save_format="full_wrapper")` + + The option is specific to Qwen3.6 and raises `RuntimeError` for other + models. Leave `save_format` at its default (`"auto"`) for other + architectures. + !!! note "vLLM serving is method-specific" `save_quantized_model()` produces a model loadable by the OneComp loader for any quantizer that supports saving (see the table below). vLLM serving, however, is only diff --git a/docs/user-guide/examples.md b/docs/user-guide/examples.md index cc7aee8..d4d238f 100644 --- a/docs/user-guide/examples.md +++ b/docs/user-guide/examples.md @@ -361,6 +361,12 @@ runner.save_quantized_model("./output/my_quantized_model") runner.save_dequantized_model("./output/my_dequantized_model") ``` +!!! note "Qwen3.6 save format" + For confirmed downstream workflows that require the full Hugging Face + wrapper layout, including vLLM serving and the current GGUF export workflow, + save Qwen3.6 models with `save_format="full_wrapper"`. See + [Basic Usage](basic-usage.md#step-5-save-the-model). + ### Load a saved quantized model ```python diff --git a/docs/user-guide/vllm-inference.md b/docs/user-guide/vllm-inference.md index 094e6a1..c1e08b8 100644 --- a/docs/user-guide/vllm-inference.md +++ b/docs/user-guide/vllm-inference.md @@ -126,24 +126,15 @@ runner.run() runner.save_quantized_model("./Llama-3.1-8B-Instruct-gptq-4bit") ``` -!!! note "Qwen3.6: use `save_format=\"full_wrapper\"`" - Qwen3.6 quantizes as a text-only checkpoint, whose native Hugging Face - layout (`model.layers.*`) does not match what vLLM's composite - `Qwen3_5ForConditionalGeneration` loader expects (`model.language_model.layers.*`). - Pass `save_format="full_wrapper"` to `save_quantized_model()` to remap the - checkpoint for vLLM serving: - - ```python - runner.save_quantized_model("./Qwen3.6-gptq-4bit-vllm", save_format="full_wrapper") - ``` - - This option is specific to Qwen3.6 and will raise `RuntimeError` for any - other model. Leave `save_format` at its default (`"auto"`) for everything - else, including other VLMs. +!!! note "Qwen3.6 save format" + Qwen3.6 checkpoints used for vLLM serving should be saved with + `save_quantized_model(..., save_format="full_wrapper")`. See + [Basic Usage](basic-usage.md#step-5-save-the-model) for the shared + Qwen3.6 save-format guidance, including other confirmed workflows. For MoE variants (e.g. Qwen3.6-A3B), `full_wrapper` also drops each expert's trivial `g_idx` buffer, since vLLM's GPTQ `FusedMoE` kernel has no `g_idx` - parameter and an unmapped one would crash weight loading — see the `desc_act`/ + parameter and an unmapped one would crash weight loading -- see the `desc_act`/ `actorder` warning above for when this isn't safe to drop. ### 2. Serve with vLLM From 3ba897889e460f78bf9a3f7a4d10d116ae3598ad Mon Sep 17 00:00:00 2001 From: Yuhki Yano Date: Wed, 9 Sep 2026 07:15:21 +0900 Subject: [PATCH 09/17] 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 10/17] 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 11/17] 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 1e1ecc0e639e681887d75b7180b3ad948fcc5e9f Mon Sep 17 00:00:00 2001 From: aki916f Date: Wed, 9 Sep 2026 12:03:18 +0900 Subject: [PATCH 12/17] clarify the license for dependency OSS may change when updated --- CHANGELOG.md | 4 ++++ README.md | 7 ++++++- docs/index.md | 7 ++++++- 3 files changed, 16 insertions(+), 2 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index cd9c207..1f54db7 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -2,6 +2,10 @@ ## [v1.3.4] 2026-09-07 +### Documentation + +- Clarified that OneComp is released under the MIT License and that licenses for dependency OSS may change when dependencies are updated. + ## [v1.3.3] 2026-09-03 ### Enhancement diff --git a/README.md b/README.md index cd58c9d..48d6cf6 100644 --- a/README.md +++ b/README.md @@ -359,7 +359,12 @@ See the [GPT-OSS guide](docs/user-guide/gptoss.md) for HF save/load, patch detai ## 📄 License -See [LICENSE](./LICENSE) for more details. +OneComp is licensed under the [MIT License](./LICENSE). + +The dependencies installed with OneComp are separate open-source software (OSS) +projects and are distributed under their respective licenses. Their licenses +may change when the dependencies are updated, so please check the license +terms of the installed versions as well. ## Citation diff --git a/docs/index.md b/docs/index.md index ad5e8cc..895e571 100644 --- a/docs/index.md +++ b/docs/index.md @@ -184,6 +184,11 @@ MDBF (Multi-Envelope Double Binary Factorization): ## License -Fujitsu One Compression is released under the terms of the [LICENSE](https://github.com/FujitsuResearch/OneCompression/blob/main/LICENSE) file included in the repository. +Fujitsu One Compression is released under the [MIT License](https://github.com/FujitsuResearch/OneCompression/blob/main/LICENSE). + +The dependencies installed with OneComp are separate open-source software (OSS) +projects and are distributed under their respective licenses. Their licenses +may change when the dependencies are updated, so please check the license +terms of the installed versions as well. Copyright 2025-2026 Fujitsu Ltd. From f91ab46055531c3bac6bb66c5fc03f9590bef695 Mon Sep 17 00:00:00 2001 From: Yuhki Yano Date: Wed, 9 Sep 2026 15:50:33 +0900 Subject: [PATCH 13/17] 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 14/17] 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 cc370e3280b5614f28f7984410fae31150a74e45 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Kimura=2C=20Keiji/=E6=9C=A8=E6=9D=91=20=E5=9C=AD=E5=85=90?= Date: Thu, 10 Sep 2026 02:45:56 +0000 Subject: [PATCH 15/17] Use a fixed offline calibration cache in cluster CI --- .gitlab-ci.yml | 3 +- CHANGELOG.md | 13 + scripts/cluster_test.sh | 55 +++- scripts/prepare_calibration_cache.py | 284 ++++++++++++++++++ .../test_prepare_calibration_cache.py | 118 ++++++++ 5 files changed, 457 insertions(+), 16 deletions(-) create mode 100644 scripts/prepare_calibration_cache.py create mode 100644 tests/onecomp/calibration/test_prepare_calibration_cache.py diff --git a/.gitlab-ci.yml b/.gitlab-ci.yml index 9da0291..c2abe40 100644 --- a/.gitlab-ci.yml +++ b/.gitlab-ci.yml @@ -84,7 +84,7 @@ lint:format: before_script: - | for var in \ - CI_BASTION_HOST CI_CLUSTER_HOST CI_CLUSTER_USER ONECOMP_REPO \ + CI_BASTION_HOST CI_CLUSTER_HOST CI_CLUSTER_USER ONECOMP_REPO ONECOMP_CALIB_CACHE \ CI_SLURM_PARTITION CI_SLURM_MEM CI_SLURM_CPUS CI_SLURM_TIME CI_SLURM_GPUS \ CI_UV_VENV CI_TORCH_EXTRA; do eval "test -n \"\${${var}:-}\" || { echo \"Set ${var} in CI/CD Variables\" >&2; exit 1; }" @@ -97,6 +97,7 @@ lint:format: - | { printf 'export ONECOMP_REPO=%q\n' "${ONECOMP_REPO}" + printf 'export ONECOMP_CALIB_CACHE=%q\n' "${ONECOMP_CALIB_CACHE}" printf 'export CI_COMMIT_SHA=%q\n' "${CI_COMMIT_SHA}" printf 'export CI_COMMIT_REF_NAME=%q\n' "${CI_COMMIT_REF_NAME:-}" printf 'export CI_JOB_TOKEN=%q\n' "${CI_JOB_TOKEN}" diff --git a/CHANGELOG.md b/CHANGELOG.md index 1f54db7..efe6740 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -2,6 +2,19 @@ ## [v1.3.4] 2026-09-07 +### Bug Fix + +- Make cluster CI use an explicitly prepared, content-verified C4 calibration + cache from shared storage. This removes its dependency on Hugging Face Hub + connectivity and generated config hashes, which previously caused C4 loading + to fail on offline compute nodes despite unrelated cached configs being + present. Add an intentional cache preparation/verification script and fail + fast when `ONECOMP_CALIB_CACHE/c4` is missing or invalid. +- Fix parallel cluster test jobs racing on the shared repository's + `.git/config.lock`. The cluster test orchestrator now fetches directly from + the authenticated CI URL under the existing lock without temporarily + rewriting the `origin` remote. + ### Documentation - Clarified that OneComp is released under the MIT License and that licenses for dependency OSS may change when dependencies are updated. diff --git a/scripts/cluster_test.sh b/scripts/cluster_test.sh index 3654bf7..21593ef 100755 --- a/scripts/cluster_test.sh +++ b/scripts/cluster_test.sh @@ -5,6 +5,7 @@ set -euo pipefail : "${ONECOMP_REPO:?ONECOMP_REPO is required}" +: "${ONECOMP_CALIB_CACHE:?ONECOMP_CALIB_CACHE is required}" : "${CI_COMMIT_SHA:?CI_COMMIT_SHA is required}" : "${CI_JOB_TOKEN:?CI_JOB_TOKEN is required}" : "${CI_SERVER_HOST:?CI_SERVER_HOST is required}" @@ -43,23 +44,21 @@ echo "target: ${PYTEST_TARGET}" cd "${ONECOMP_REPO}" mkdir -p output error .cache -# Sync repo to the MR commit. Temporarily swap origin to CI_JOB_TOKEN auth; restore on exit. +# Sync repo to the MR commit without modifying the shared repository config. # flock: parallel matrix jobs share this Lustre checkout — serialize fetch/checkout. -ORIGIN_URL="$(git remote get-url origin)" -git remote set-url origin "https://gitlab-ci-token:${CI_JOB_TOKEN}@${CI_SERVER_HOST}/${CI_PROJECT_PATH}.git" -trap 'git remote set-url origin "${ORIGIN_URL}"' EXIT - +AUTH_REPO_URL="https://gitlab-ci-token:${CI_JOB_TOKEN}@${CI_SERVER_HOST}/${CI_PROJECT_PATH}.git" REF="${CI_COMMIT_REF_NAME:-}" -flock "${ONECOMP_REPO}/.cache/git-sync.lock" bash -c ' - set -euo pipefail - cd "'"${ONECOMP_REPO}"'" - if [[ -n "'"${REF}"'" ]]; then - git fetch origin "'"${REF}"'" +( + flock 9 + if [[ -n "${REF}" ]]; then + git fetch "${AUTH_REPO_URL}" "${REF}" else - git fetch origin + # Detached/manual pipelines may not provide a ref; fetch the exact commit + # because fetching the remote default branch need not include it. + git fetch "${AUTH_REPO_URL}" "${CI_COMMIT_SHA}" fi - git checkout "'"${CI_COMMIT_SHA}"'" -' + git checkout "${CI_COMMIT_SHA}" +) 9>"${ONECOMP_REPO}/.cache/git-sync.lock" PYTEST_TARGET_Q="" for target in ${PYTEST_TARGET}; do @@ -134,6 +133,7 @@ uv --version export UV_PROJECT_ENVIRONMENT="${CI_UV_VENV}" export UV_CACHE_DIR="${ONECOMP_REPO}/.cache/uv" +export ONECOMP_CALIB_CACHE=$(printf '%q' "${ONECOMP_CALIB_CACHE}") if [[ "${RUN_UV_SYNC}" -eq 1 ]]; then echo "=== uv sync (${CI_TORCH_EXTRA}, dev, vllm, visualize) -> ${CI_UV_VENV} on \$(uname -m) ===" @@ -143,13 +143,38 @@ if [[ "${RUN_UV_SYNC}" -eq 1 ]]; then uv sync --extra ${CI_TORCH_EXTRA} --extra dev --extra vllm --extra visualize --frozen fi +# Keep every test offline, not only the preflight verification. This makes an +# unexpected calibration cache miss fail locally instead of contacting the Hub. +export HF_DATASETS_OFFLINE=1 +export HF_HUB_OFFLINE=1 + +cache_lock="\${ONECOMP_CALIB_CACHE}/.c4.lock" +if [[ ! -r "\${cache_lock}" ]]; then + echo "ERROR: fixed calibration cache lock is not readable: \${cache_lock}" >&2 + echo "Regenerate the cache with scripts/prepare_calibration_cache.py." >&2 + exit 1 +fi +exec 8<"\${cache_lock}" +# Hold a reader lock through pytest so regeneration cannot replace the cache +# between preflight verification and a later calibration-data load. +flock -s 8 + +if [[ ! -r "\${ONECOMP_CALIB_CACHE}/c4/dataset_dict.json" ]]; then + echo "ERROR: fixed calibration cache is not readable: \${ONECOMP_CALIB_CACHE}/c4" >&2 + echo "Prepare it with scripts/prepare_calibration_cache.py before running cluster CI." >&2 + exit 1 +fi +uv run --no-sync python scripts/prepare_calibration_cache.py --verify + echo "=== job info ===" echo "host: \$(hostname)" echo "arch: \$(uname -m)" echo "job id: \${SLURM_JOB_ID:-N/A}" echo "node list: \${SLURM_JOB_NODELIST:-N/A}" echo "cpus: \${SLURM_CPUS_PER_TASK:-N/A}" -uv run python - <<'PY' +# The setup job is the only writer to the shared venv. Without --no-sync, +# parallel matrix jobs can reinstall torch while another job is importing it. +uv run --no-sync python - <<'PY' import platform, torch print(f"python arch: {platform.machine()}") print(f"torch: {torch.__version__}") @@ -160,7 +185,7 @@ PY echo "================" if [[ "${RUN_PYTEST}" -eq 1 ]]; then - uv run pytest -v --maxfail=0 --durations=20 ${PYTEST_M_ARGS} ${PYTEST_TARGET_Q} + uv run --no-sync pytest -v --maxfail=0 --durations=20 ${PYTEST_M_ARGS} ${PYTEST_TARGET_Q} else echo "=== cluster setup passed ===" fi diff --git a/scripts/prepare_calibration_cache.py b/scripts/prepare_calibration_cache.py new file mode 100644 index 0000000..aa514b8 --- /dev/null +++ b/scripts/prepare_calibration_cache.py @@ -0,0 +1,284 @@ +#!/usr/bin/env python3 +"""Prepare and verify the fixed C4 calibration cache used by cluster CI. + +Copyright 2025-2026 Fujitsu Ltd. +""" + +import argparse +import fcntl +import hashlib +import json +import os +import shutil +import sys +from contextlib import contextmanager +from pathlib import Path + +import datasets +from huggingface_hub import hf_hub_download + +DATASET_ID = "allenai/c4" +DATASET_REVISION = "1588ec454efa1a09f29cd18ddd04fe05fc8653a2" +DATA_FILE = "en/c4-train.00001-of-01024.json.gz" +SOURCE_SHA256 = "b945059cd1a343cabe311881b7840a6f0363f570e745a0eff0e687e266f6b55d" +EXPECTED_TRAIN_ROWS = 356318 +MANIFEST_NAME = "onecomp-calibration-manifest.json" + + +def _sha256(path): + """Return the SHA-256 digest of a file without loading it all into memory.""" + + digest = hashlib.sha256() + with path.open("rb") as file: + for chunk in iter(lambda: file.read(1024 * 1024), b""): + digest.update(chunk) + return digest.hexdigest() + + +def _cache_path(cache_root): + """Return the save-to-disk directory consumed by the C4 loader.""" + + return Path(cache_root).expanduser().resolve() / "c4" + + +def _file_records(root): + """Describe every persisted cache file so content changes are detectable.""" + + records = [] + for path in sorted(item for item in root.rglob("*") if item.is_file()): + if path.name == MANIFEST_NAME: + # The manifest cannot include its own digest without becoming self-referential. + continue + records.append( + { + "path": path.relative_to(root).as_posix(), + "size": path.stat().st_size, + "sha256": _sha256(path), + } + ) + return records + + +@contextmanager +def _cache_lock(cache_root, *, exclusive): + """Coordinate cache readers with the short final installation step. + + Verification and CI use a shared lock for the entire period in which cache + files may be read. Regeneration uses an exclusive lock only while replacing + the verified temporary cache, so building and hashing it does not block CI. + """ + + cache_root = Path(cache_root).expanduser().resolve() + lock_path = cache_root / ".c4.lock" + if exclusive: + cache_root.mkdir(parents=True, exist_ok=True) + lock_path.touch(exist_ok=True) + mode = "r+" + operation = fcntl.LOCK_EX + else: + mode = "r" + operation = fcntl.LOCK_SH + + try: + lock_file = lock_path.open(mode, encoding="utf-8") + except FileNotFoundError as exc: + raise FileNotFoundError( + f"Fixed calibration cache lock is missing: {lock_path}. " + "Regenerate the cache with this script." + ) from exc + with lock_file: + fcntl.flock(lock_file, operation) + yield + + +def prepare(cache_root, *, local_files_only=False, force=False): + """Build, verify, and transactionally install the pinned C4 cache.""" + + destination = _cache_path(cache_root) + if destination.exists() and not force: + raise FileExistsError( + f"Calibration cache already exists: {destination}. " + "Use --force only for an intentional regeneration." + ) + + source = Path( + hf_hub_download( + repo_id=DATASET_ID, + filename=DATA_FILE, + repo_type="dataset", + revision=DATASET_REVISION, + local_files_only=local_files_only, + ) + ) + source_sha256 = _sha256(source) + if source_sha256 != SOURCE_SHA256: + raise ValueError( + f"Unexpected SHA-256 for {DATA_FILE}: {source_sha256}; " f"expected {SOURCE_SHA256}" + ) + + dataset = datasets.load_dataset("json", data_files={"train": str(source)}) + if len(dataset["train"]) != EXPECTED_TRAIN_ROWS: + raise ValueError( + f"Unexpected C4 train row count: {len(dataset['train'])}; " + f"expected {EXPECTED_TRAIN_ROWS}" + ) + + destination.parent.mkdir(parents=True, exist_ok=True) + temporary = destination.with_name(f".{destination.name}.tmp-{os.getpid()}") + if temporary.exists(): + shutil.rmtree(temporary) + + try: + dataset.save_to_disk(temporary) + # Fingerprints can change during save_to_disk(), so record the value that + # consumers will observe after loading the persisted cache. + persisted_dataset = datasets.load_from_disk(temporary) + manifest = { + "dataset_id": DATASET_ID, + "dataset_revision": DATASET_REVISION, + "data_files": {"train": DATA_FILE}, + "source_sha256": source_sha256, + "datasets_version": datasets.__version__, + "files": _file_records(temporary), + "splits": { + "train": { + "num_rows": len(persisted_dataset["train"]), + "fingerprint": persisted_dataset["train"]._fingerprint, + "features": persisted_dataset["train"].features.to_dict(), + } + }, + } + (temporary / MANIFEST_NAME).write_text( + json.dumps(manifest, indent=2, sort_keys=True) + "\n", + encoding="utf-8", + ) + # Reject an incomplete or internally inconsistent cache before it can + # become visible at the shared destination. + _verify_cache(temporary) + + backup = destination.with_name(f".{destination.name}.backup-{os.getpid()}") + with _cache_lock(cache_root, exclusive=True): + if destination.exists() and not force: + raise FileExistsError( + f"Calibration cache already exists: {destination}. " + "Use --force only for an intentional regeneration." + ) + if backup.exists(): + raise FileExistsError(f"Calibration cache backup already exists: {backup}") + + moved_existing = False + try: + # Keep the previous cache available for rollback until the + # already-verified replacement has been installed. + if destination.exists(): + destination.rename(backup) + moved_existing = True + temporary.rename(destination) + except Exception: + if moved_existing and not destination.exists(): + backup.rename(destination) + raise + + if moved_existing: + try: + shutil.rmtree(backup) + except OSError as exc: + print( + f"WARNING: failed to remove calibration cache backup: {exc}", + file=sys.stderr, + ) + except Exception: + shutil.rmtree(temporary, ignore_errors=True) + raise + + print(f"Prepared fixed C4 calibration cache: {destination}") + + +def _verify_cache(destination): + """Validate the pinned source identity and every persisted cache artifact.""" + + manifest_path = destination / MANIFEST_NAME + if not manifest_path.is_file(): + raise FileNotFoundError(f"Fixed calibration cache manifest is missing: {manifest_path}") + + manifest = json.loads(manifest_path.read_text(encoding="utf-8")) + expected = { + "dataset_id": DATASET_ID, + "dataset_revision": DATASET_REVISION, + "data_files": {"train": DATA_FILE}, + "source_sha256": SOURCE_SHA256, + } + for key, expected_value in expected.items(): + if manifest.get(key) != expected_value: + raise ValueError( + f"Invalid fixed calibration cache manifest field {key!r}: " + f"{manifest.get(key)!r}; expected {expected_value!r}" + ) + + files = _file_records(destination) + if manifest.get("files") != files: + raise ValueError("Fixed calibration cache files do not match its manifest") + + dataset = datasets.load_from_disk(destination) + if "train" not in dataset: + raise ValueError("Fixed calibration cache has no train split") + train_manifest = manifest.get("splits", {}).get("train", {}) + if len(dataset["train"]) != EXPECTED_TRAIN_ROWS: + raise ValueError( + f"Invalid fixed calibration cache row count: {len(dataset['train'])}; " + f"expected {EXPECTED_TRAIN_ROWS}" + ) + if train_manifest.get("num_rows") != EXPECTED_TRAIN_ROWS: + raise ValueError("Fixed calibration cache manifest has an invalid row count") + if dataset["train"]._fingerprint != train_manifest.get("fingerprint"): + raise ValueError("Fixed calibration cache fingerprint does not match its manifest") + if dataset["train"].features.to_dict() != train_manifest.get("features"): + raise ValueError("Fixed calibration cache schema does not match its manifest") + + print( + "Verified fixed C4 calibration cache: " + f"{destination} ({len(dataset['train'])} train rows)" + ) + + +def verify(cache_root): + """Verify a cache while preventing concurrent regeneration from replacing it.""" + + with _cache_lock(cache_root, exclusive=False): + _verify_cache(_cache_path(cache_root)) + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument( + "--cache-root", + default=os.environ.get("ONECOMP_CALIB_CACHE"), + help="Cache root containing the c4 directory (default: ONECOMP_CALIB_CACHE)", + ) + parser.add_argument("--verify", action="store_true", help="Verify without writing") + parser.add_argument( + "--local-files-only", + action="store_true", + help="Use an already downloaded Hugging Face source shard", + ) + parser.add_argument( + "--force", + action="store_true", + help="Replace an existing cache intentionally", + ) + args = parser.parse_args() + if not args.cache_root: + parser.error("--cache-root or ONECOMP_CALIB_CACHE is required") + + if args.verify: + verify(args.cache_root) + else: + prepare( + args.cache_root, + local_files_only=args.local_files_only, + force=args.force, + ) + + +if __name__ == "__main__": + main() diff --git a/tests/onecomp/calibration/test_prepare_calibration_cache.py b/tests/onecomp/calibration/test_prepare_calibration_cache.py new file mode 100644 index 0000000..ce3b773 --- /dev/null +++ b/tests/onecomp/calibration/test_prepare_calibration_cache.py @@ -0,0 +1,118 @@ +"""Test fixed-cache integrity and replacement without network or shared storage.""" + +import hashlib +import importlib.util +import json +import shutil +from pathlib import Path +from unittest.mock import patch + +import datasets +import pytest + +_SCRIPT_PATH = Path(__file__).parents[3] / "scripts" / "prepare_calibration_cache.py" +_SPEC = importlib.util.spec_from_file_location("prepare_calibration_cache", _SCRIPT_PATH) +assert _SPEC and _SPEC.loader +_CACHE_SCRIPT = importlib.util.module_from_spec(_SPEC) +_SPEC.loader.exec_module(_CACHE_SCRIPT) + + +def _dataset(values): + """Create a minimal C4-shaped dataset for fast, self-contained tests.""" + + return datasets.DatasetDict( + { + "train": datasets.Dataset.from_dict( + { + "text": values, + "url": [f"https://example.com/{index}" for index in range(len(values))], + } + ) + } + ) + + +def _write_cache(cache_root, values): + """Write a valid synthetic cache and the manifest expected by verification.""" + + destination = cache_root / "c4" + _dataset(values).save_to_disk(destination) + persisted = datasets.load_from_disk(destination) + manifest = { + "dataset_id": _CACHE_SCRIPT.DATASET_ID, + "dataset_revision": _CACHE_SCRIPT.DATASET_REVISION, + "data_files": {"train": _CACHE_SCRIPT.DATA_FILE}, + "source_sha256": _CACHE_SCRIPT.SOURCE_SHA256, + "datasets_version": datasets.__version__, + "files": _CACHE_SCRIPT._file_records(destination), + "splits": { + "train": { + "num_rows": len(persisted["train"]), + "fingerprint": persisted["train"]._fingerprint, + "features": persisted["train"].features.to_dict(), + } + }, + } + (destination / _CACHE_SCRIPT.MANIFEST_NAME).write_text(json.dumps(manifest), encoding="utf-8") + (cache_root / ".c4.lock").touch() + return destination + + +@pytest.fixture(autouse=True) +def _small_expected_dataset(monkeypatch): + """Match production row-count validation to the two-row test datasets.""" + + monkeypatch.setattr(_CACHE_SCRIPT, "EXPECTED_TRAIN_ROWS", 2) + + +def test_verify_accepts_matching_file_checksums(tmp_path): + """An untouched cache whose files match the manifest must verify successfully.""" + + _write_cache(tmp_path, ["first", "second"]) + + _CACHE_SCRIPT.verify(tmp_path) + + +def test_verify_rejects_replaced_arrow_file(tmp_path): + """File checksums must detect Arrow replacement missed by saved metadata.""" + + destination = _write_cache(tmp_path, ["first", "second"]) + replacement = tmp_path / "replacement" + _dataset(["other-a", "other-b"]).save_to_disk(replacement) + source_arrow = next(replacement.rglob("*.arrow")) + destination_arrow = next(destination.rglob("*.arrow")) + shutil.copyfile(source_arrow, destination_arrow) + + with pytest.raises(ValueError, match="files do not match"): + _CACHE_SCRIPT.verify(tmp_path) + + +def test_force_restores_existing_cache_when_install_fails(tmp_path, monkeypatch): + """A failed forced install must restore the previous cache and remove its backup.""" + + destination = tmp_path / "c4" + destination.mkdir() + marker = destination / "existing" + marker.write_text("keep", encoding="utf-8") + source = tmp_path / "source.json" + source.write_text("source", encoding="utf-8") + monkeypatch.setattr(_CACHE_SCRIPT, "SOURCE_SHA256", hashlib.sha256(b"source").hexdigest()) + monkeypatch.setattr(_CACHE_SCRIPT, "hf_hub_download", lambda **_kwargs: str(source)) + monkeypatch.setattr( + _CACHE_SCRIPT.datasets, "load_dataset", lambda *_args, **_kwargs: _dataset(["a", "b"]) + ) + + original_rename = Path.rename + + def fail_install(path, target): + # Let the old cache move to backup, then fail only the replacement move. + if path.name.startswith(".c4.tmp-"): + raise OSError("simulated install failure") + return original_rename(path, target) + + with patch.object(Path, "rename", fail_install): + with pytest.raises(OSError, match="simulated install failure"): + _CACHE_SCRIPT.prepare(tmp_path, force=True) + + assert marker.read_text(encoding="utf-8") == "keep" + assert not list(tmp_path.glob(".c4.backup-*")) From 9c34d5e66120149fd4159ba3a2038a2054e93ad0 Mon Sep 17 00:00:00 2001 From: Yuhki Yano Date: Fri, 11 Sep 2026 12:29:10 +0900 Subject: [PATCH 16/17] 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 17/17] 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()