onecomp 1.3.2__tar.gz → 1.3.4__tar.gz

{onecomp-1.3.2/onecomp.egg-info → onecomp-1.3.4}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: onecomp
-Version: 1.3.2
+Version: 1.3.4
 Summary: Python package for LLM compression
 Author: Keiji Kimura
 License: MIT License
@@ -473,7 +473,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
 

{onecomp-1.3.2 → onecomp-1.3.4}/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
 

{onecomp-1.3.2 → onecomp-1.3.4}/onecomp/__version__.py

@@ -6,4 +6,4 @@ Author: Keiji Kimura
 
 """
 
-__version__ = "1.3.2"
+__version__ = "1.3.4"

{onecomp-1.3.2 → onecomp-1.3.4}/onecomp/lpcd/_lpcd_config.py

@@ -25,7 +25,8 @@ class LPCDConfig:
         gd_steps: Number of gradient-descent epochs per sub-problem.
         gd_batch_size: Effective batch size for gradient accumulation.
         gd_base_lr: Base learning rate for gradient-descent solver.
-        device: Device to perform LPCD optimisation on.
+        device: Device to perform LPCD optimisation on. Default is None.
+            When None, Runner synchronises it with ModelConfig.device.
 
     Examples:
         Minimal (residual correction only, fast)::
@@ -53,4 +54,4 @@ class LPCDConfig:
     gd_steps: int = 20
     gd_batch_size: int = 16
     gd_base_lr: float = 1e-4
-    device: str = "cuda:0"
+    device: str = None

{onecomp-1.3.2 → onecomp-1.3.4}/onecomp/model_config.py

@@ -11,7 +11,7 @@ from logging import getLogger
 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 @@ class ModelConfig:
                 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 @@ class ModelConfig:
                 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-1.3.2 → onecomp-1.3.4}/onecomp/qep/_qep_config.py

@@ -25,7 +25,7 @@ class QEPConfig:
             Default is 0.5.
         device (str): Device to use for QEP computations
             (e.g., "cuda", "mps", "cpu").
-            Default is "cuda:0".
+            Default is None. When None, Runner synchronises it with ModelConfig.device.
         exclude_layer_keywords (list[str]): List of keywords to identify
             layers excluded from error propagation. Layers whose names
             contain any of these keywords will be excluded.
@@ -52,6 +52,6 @@ class QEPConfig:
     general: bool = False
     percdamp: float = 0.01
     perccorr: float = 0.5
-    device: str = "cuda:0"
+    device: str = None
     exclude_layer_keywords: list[str] = field(default_factory=lambda: ["mlp.down_proj"])
     # TODO: exclude_layer_keywords depends on the architecture and needs to be fixed

{onecomp-1.3.2 → onecomp-1.3.4}/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,
     )
 

{onecomp-1.3.2 → onecomp-1.3.4}/onecomp/quantized_model_loader.py

@@ -123,20 +123,16 @@ class QuantizedModelLoader:
         elif unfuse_moe_experts(model, logger):
             logger.info("Unfused MoE expert tensors for quantized model load")
 
-        # Align checkpoint key prefixes with the empty model built from config.
-        # Gemma3 VLMs are a common case: weights saved from from_pretrained
-        # use model.language_model.model.layers. (language_model is a
-        # ForCausalLM wrapper) while from_config exposes
-        # model.language_model.layers.* directly.
-        state_dict = cls._remap_state_dict_keys(state_dict, model)
-
-        # Replace quantized layers with empty modules and align quantized
-        # tensor keys with the actual module names in the model built from
-        # config.  This is required when the saved checkpoint and the
-        # from_config model use different wrapper prefixes, e.g.
-        # model.language_model.layers.* vs model.layers.*.
+        # Replace quantized layers first. This resolves every quantizer's
+        # tensor state by layer prefix while the model still exposes the
+        # original empty Linear modules.
         state_dict = cls._replace_quantized_layers(model, state_dict, quant_config)
 
+        # Align remaining checkpoint keys with the model built from config.
+        # Quantized keys now already use the actual module prefix, so this
+        # generic remap only handles non-quantized parameters and buffers.
+        state_dict = cls._remap_state_dict_keys(state_dict, model)
+
         # Load all weights (quantized + non-quantized) in one go.  strict=False
         # is intentional because some wrapper-only components may be absent, but
         # critical language-model and quantized-buffer mismatches must fail fast.
@@ -593,11 +589,10 @@ class QuantizedModelLoader:
         silently skips mismatched keys and leaves layers at their empty-model
         initial values (often all zeros for quantized buffers).
 
-        Remapping runs before _replace_quantized_layers, so the empty
-        model still exposes nn.Linear.weight rather than GPTQ buffers
-        (``qweight``, ``scales``, …).  Known prefix rewrites are therefore
-        applied from checkpoint key patterns alone; they must not require the
-        destination key to already exist in model.named_parameters().
+        Quantized layers are resolved by _replace_quantized_layers before
+        this method runs. Therefore this method only remaps ordinary model
+        parameters and buffers; quantizer-specific tensor fields are already
+        materialized under their actual module prefixes.
 
         Args:
             state_dict: Tensors loaded from *.safetensors.

{onecomp-1.3.2 → onecomp-1.3.4}/onecomp/quantizer/_quantizer.py

@@ -400,12 +400,32 @@ class Quantizer(metaclass=ABCMeta):
         weight[:, dead] = 0
 
         # QEP correction
-        damp = percdamp * torch.mean(torch.diag(hessian))
-        diag = torch.arange(hessian.shape[0], device=hessian.device)
-        hessian[diag, diag] += damp
-        rhs = weight @ delta_hatX
-        delta_weight = _safe_cholesky_and_solve(hessian, rhs).t()
-        weight = weight + (perccorr * delta_weight)
+        damp_scale = 1.0
+        max_retries = 5
+        for attempt in range(max_retries):
+            try:
+                damp = percdamp * torch.mean(torch.diag(hessian))
+                diag = torch.arange(hessian.shape[0], device=hessian.device)
+                hessian[diag, diag] += damp
+                rhs = weight @ delta_hatX
+                delta_weight = _safe_cholesky_and_solve(hessian, rhs).t()
+                weight = weight + (perccorr * delta_weight)
+                break
+            except torch._C._LinAlgError:
+                damp_scale *= 10.0
+                extra = damp_scale * damp
+                hessian[diag, diag] += extra
+                self.logger.warning(
+                    "Cholesky failed (attempt %d/%d); adding extra damping %.2e",
+                    attempt + 1,
+                    max_retries,
+                    extra,
+                )
+        else:
+            raise RuntimeError(
+                "Cholesky decomposition failed after %d damping attempts. "
+                "The Hessian may be severely ill-conditioned." % max_retries
+            )
 
         if isinstance(module, Conv1d):
             weight = weight.t()

{onecomp-1.3.2 → onecomp-1.3.4}/onecomp/quantizer/gptq/gptq_layer.py

@@ -414,12 +414,6 @@ class GPTQLinear(nn.Module):
         if self._gemlite_layer is not None:
             self.using_gemlite = True
 
-        # Permutation order
-        if perm is not None and actorder:
-            self.register_buffer("perm", perm.to(device))
-        else:
-            self.perm = None
-
         # Bias
         if bias is not None:
             self.register_buffer("bias", bias.to(torch.float16).to(device))

{onecomp-1.3.2 → onecomp-1.3.4}/onecomp/runner.py

@@ -261,9 +261,13 @@ class Runner:
         self.qep_config = None
         if qep:
             self.qep_config = qep_config if qep_config is not None else QEPConfig()
+            if self.qep_config.device is None and self.model_config is not None:
+                self.qep_config.device = str(self.model_config.get_device())
         self.lpcd_config = None
         if lpcd:
             self.lpcd_config = lpcd_config if lpcd_config is not None else LPCDConfig()
+            if self.lpcd_config.device is None and self.model_config is not None:
+                self.lpcd_config.device = str(self.model_config.get_device())
         self.report_progress = report_progress
 
     def check(self):
@@ -375,6 +379,11 @@ class Runner:
         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"

{onecomp-1.3.2 → onecomp-1.3.4}/onecomp/utils/blockwise.py

@@ -6,6 +6,7 @@ Author: Yudai Fujimoto, Akihiro Yoshida, Yuma Ichikawa
 
 """
 
+from collections import UserDict
 from logging import getLogger
 
 import torch
@@ -95,6 +96,7 @@ class Catcher(nn.Module):
 _PER_LAYER_INPUTS_KEY = "_per_layer_inputs"
 _POS_EMB_MAP_KEY = "_position_embeddings_map"
 _ATTN_MASK_MAP_KEY = "_attention_mask_map"
+_SHARED_KV_STATES_KEY = "shared_kv_states"
 
 
 def _find_blocks_parent(model, blocks):
@@ -243,9 +245,15 @@ def get_blocks_and_inputs(
 
     # Now capture block inputs for all calibration samples.
     block_inps = []
-    for inp in inp_ids.split(batch_size):
+    for first in range(0, inp_ids.shape[0], batch_size):
+        last = min(first + batch_size, inp_ids.shape[0])
+        inp = inp_ids[first:last]
+        batch_model_kwargs = {
+            k: v[first:last] if isinstance(v, torch.Tensor) and v.dim() >= 1 else v
+            for k, v in model_kwargs.items()
+        }
         try:
-            _ = model(inp, **model_kwargs)
+            _ = model(inp, **batch_model_kwargs)
         except StopForward:
             block_inps.append(blocks[0].inp.cpu())
 
@@ -361,6 +369,10 @@ def move_kwargs_to_device(x, device):
         return [move_kwargs_to_device(v, device) for v in x]
     elif isinstance(x, tuple):
         return tuple(move_kwargs_to_device(v, device) for v in x)
+    elif isinstance(x, UserDict):
+        for k, v in list(x.items()):
+            x[k] = move_kwargs_to_device(v, device)
+        return x
     else:
         return x
 
@@ -396,6 +408,10 @@ def expand_kwargs_batch(kwargs, batch_size):
             return [_expand(t) for t in v]
         elif isinstance(v, dict):
             return {k: _expand(val) for k, val in v.items()}
+        elif isinstance(v, UserDict):
+            for k, val in list(v.items()):
+                v[k] = _expand(val)
+            return v
         return v
 
     return {k: _expand(v) for k, v in kwargs.items()}
@@ -433,9 +449,46 @@ def prepare_block_kwargs(batch_kwargs, block, pli, offset, batch_size, device):
         if layer_type and layer_type in mask_map:
             batch_kwargs["attention_mask"] = mask_map[layer_type]
 
+    # 4) Gemma4 shared KV state. Provider blocks write states into a fresh
+    # batch-local mapping; consumer blocks read the matching calibration slice.
+    shared_kv_states = batch_kwargs.get(_SHARED_KV_STATES_KEY)
+    self_attn = getattr(block, "self_attn", None)
+    if isinstance(shared_kv_states, UserDict) and getattr(
+        self_attn, "store_full_length_kv", False
+    ):
+        batch_kwargs[_SHARED_KV_STATES_KEY] = UserDict()
+    elif isinstance(shared_kv_states, UserDict) and getattr(
+        self_attn, "is_kv_shared_layer", False
+    ):
+        batch_kwargs[_SHARED_KV_STATES_KEY] = _slice_shared_kv_states(
+            shared_kv_states, offset, batch_size, device
+        )
+
     return batch_kwargs
 
 
+def _slice_shared_kv_states(shared_kv_states, offset, batch_size, device):
+    batch_shared_kv_states = UserDict()
+    for key, value in shared_kv_states.items():
+        if isinstance(value, tuple):
+            batch_shared_kv_states[key] = tuple(
+                _slice_shared_kv_tensor(v, offset, batch_size, device) for v in value
+            )
+        else:
+            batch_shared_kv_states[key] = _slice_shared_kv_tensor(
+                value, offset, batch_size, device
+            )
+    return batch_shared_kv_states
+
+
+def _slice_shared_kv_tensor(value, offset, batch_size, device):
+    if isinstance(value, torch.Tensor):
+        if value.dim() >= 1 and value.shape[0] >= offset + batch_size:
+            value = value[offset : offset + batch_size]
+        return value.to(device)
+    return value
+
+
 def _get_block_layer_type(block: nn.Module) -> str | None:
     return (
         getattr(block, "layer_type", None)
@@ -468,6 +521,9 @@ def forward_input(
     pli = kwargs.get(_PER_LAYER_INPUTS_KEY)
     next_inps = []
     offset = 0
+    shared_kv_chunks = {}
+    self_attn = getattr(block, "self_attn", None)
+    stores_shared_kv = getattr(self_attn, "store_full_length_kv", False)
     for inp in inps.split(batch_size):
         bs = inp.shape[0]
         batch_kwargs = expand_kwargs_batch(kwargs, bs)
@@ -475,7 +531,27 @@ def forward_input(
         out = block(inp.to(device), **batch_kwargs)
         out = out[0] if isinstance(out, tuple) else out
         next_inps.append(out.cpu())
+        if stores_shared_kv:
+            for key, value in batch_kwargs[_SHARED_KV_STATES_KEY].items():
+                shared_kv_chunks.setdefault(key, []).append(
+                    tuple(v.detach().cpu() for v in value)
+                    if isinstance(value, tuple)
+                    else value.detach().cpu()
+                )
         offset += bs
+
+    if shared_kv_chunks:
+        shared_kv_states = kwargs.get(_SHARED_KV_STATES_KEY)
+        if not isinstance(shared_kv_states, UserDict):
+            shared_kv_states = UserDict()
+            kwargs[_SHARED_KV_STATES_KEY] = shared_kv_states
+        for key, chunks in shared_kv_chunks.items():
+            if isinstance(chunks[0], tuple):
+                shared_kv_states[key] = tuple(
+                    torch.cat([chunk[i] for chunk in chunks], dim=0) for i in range(len(chunks[0]))
+                )
+            else:
+                shared_kv_states[key] = torch.cat(chunks, dim=0)
     return torch.cat(next_inps)
 
 

{onecomp-1.3.2 → onecomp-1.3.4}/onecomp/utils/unfuse_moe.py

@@ -55,6 +55,9 @@ class _UnfusedExperts(nn.Module):
     def __getitem__(self, idx):
         return getattr(self, str(int(idx)))
 
+    def __iter__(self):
+        return iter(self._modules.values())
+
     def forward(
         self,
         hidden_states: torch.Tensor,
@@ -584,17 +587,19 @@ def _fuse_one(
         expert0 = unfused[0]
         inter = expert0.gate_proj.out_features
         hidden = expert0.gate_proj.in_features
+        up_w0, _ = _dequantized_weight_bias(expert0.up_proj)
+        down_w0, _ = _dequantized_weight_bias(expert0.down_proj)
         gate_up_3d = torch.empty(
             num_experts,
             2 * inter,
             hidden,
-            dtype=expert0.gate_proj.weight.dtype,
+            dtype=up_w0.dtype,
         )
         down_3d = torch.empty(
             num_experts,
             hidden,
             inter,
-            dtype=expert0.down_proj.weight.dtype,
+            dtype=down_w0.dtype,
         )
         for i in range(num_experts):
             expert = unfused[i]

{onecomp-1.3.2 → onecomp-1.3.4/onecomp.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: onecomp
-Version: 1.3.2
+Version: 1.3.4
 Summary: Python package for LLM compression
 Author: Keiji Kimura
 License: MIT License
@@ -473,7 +473,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
 

{onecomp-1.3.2 → onecomp-1.3.4}/onecomp.egg-info/SOURCES.txt

@@ -296,6 +296,7 @@ onecomp/utils/unfuse_moe.py
 onecomp/utils/vram_estimator.py
 scripts/check_copyright_header.py
 scripts/check_no_japanese.py
+scripts/prepare_calibration_cache.py
 vllm_plugins/__init__.py
 vllm_plugins/dbf/__init__.py
 vllm_plugins/dbf/vllm_plugin.py