compressed-tensors 0.5.0__py3-none-any.whl → 0.7.0__py3-none-any.whl

compressed_tensors/compressors/{sparse_bitmask.py → sparse_compressors/sparse_bitmask.py}

@@ -12,17 +12,15 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import logging
-from typing import Dict, Generator, List, Tuple, Union
+from typing import Dict, List, Tuple, Union
 
 import numpy
 import torch
-from compressed_tensors.compressors import Compressor
+from compressed_tensors.compressors.base import BaseCompressor
+from compressed_tensors.compressors.sparse_compressors.base import BaseSparseCompressor
 from compressed_tensors.config import CompressionFormat
-from compressed_tensors.utils import get_nested_weight_mappings, merge_names
-from safetensors import safe_open
+from compressed_tensors.utils import merge_names
 from torch import Tensor
-from tqdm import tqdm
 
 
 __all__ = [
@@ -34,11 +32,9 @@ __all__ = [
     "unpack_bitmasks",
 ]
 
-_LOGGER: logging.Logger = logging.getLogger(__name__)
 
-
-@Compressor.register(name=CompressionFormat.sparse_bitmask.value)
-class BitmaskCompressor(Compressor):
+@BaseCompressor.register(name=CompressionFormat.sparse_bitmask.value)
+class BitmaskCompressor(BaseSparseCompressor):
     """
     Compression for sparse models using bitmasks. Non-zero weights are stored in a 1d
     values tensor, with their locations stored in a 2d bitmask
@@ -46,56 +42,15 @@ class BitmaskCompressor(Compressor):
 
     COMPRESSION_PARAM_NAMES = ["shape", "compressed", "bitmask", "row_offsets"]
 
-    def compress(self, model_state: Dict[str, Tensor]) -> Dict[str, Tensor]:
-        """
-        Compresses a dense state dict using bitmask compression
+    def compress_weight(self, name, value):
+        bitmask_tensor = BitmaskTensor.from_dense(value)
+        bitmask_dict = bitmask_tensor.dict(name_prefix=name, device="cpu")
+        return bitmask_dict
 
-        :param model_state: state dict of uncompressed model
-        :return: compressed state dict
-        """
-        compressed_dict = {}
-        _LOGGER.debug(
-            f"Compressing model with {len(model_state)} parameterized layers..."
-        )
-        for name, value in tqdm(model_state.items(), desc="Compressing model"):
-            bitmask_tensor = BitmaskTensor.from_dense(value)
-            bitmask_dict = bitmask_tensor.dict(name_prefix=name, device="cpu")
-            for key in bitmask_dict.keys():
-                if key in compressed_dict:
-                    _LOGGER.warn(
-                        f"Expected all compressed state_dict keys to be unique, but "
-                        f"found an existing entry for {key}. The existing entry will "
-                        "be replaced."
-                    )
-            compressed_dict.update(bitmask_dict)
-
-        return compressed_dict
-
-    def decompress(
-        self, path_to_model_or_tensors: str, device: str = "cpu", **kwargs
-    ) -> Generator[Tuple[str, Tensor], None, None]:
-        """
-        Reads a bitmask compressed state dict located
-        at path_to_model_or_tensors and returns a generator
-        for sequentially decompressing back to a dense state dict
-
-        :param model_path: path to compressed safetensors model (directory with
-            one or more safetensors files) or compressed tensors file
-        :param device: device to load decompressed weights onto
-        :return: iterator for generating decompressed weights
-        """
-        weight_mappings = get_nested_weight_mappings(
-            path_to_model_or_tensors, self.COMPRESSION_PARAM_NAMES
-        )
-        for weight_name in weight_mappings.keys():
-            weight_data = {}
-            for param_name, safe_path in weight_mappings[weight_name].items():
-                full_name = merge_names(weight_name, param_name)
-                with safe_open(safe_path, framework="pt", device=device) as f:
-                    weight_data[param_name] = f.get_tensor(full_name)
-            data = BitmaskTensor(**weight_data)
-            decompressed = data.decompress()
-            yield weight_name, decompressed
+    def decompress_weight(self, weight_data):
+        data = BitmaskTensor(**weight_data)
+        decompressed = data.decompress()
+        return decompressed
 
 
 class BitmaskTensor:

compressed_tensors/compressors/sparse_quantized_compressors/__init__.py

@@ -0,0 +1,16 @@
+# Copyright (c) 2021 - present / Neuralmagic, Inc. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# flake8: noqa
+
+from .marlin_24 import Marlin24Compressor

compressed_tensors/compressors/{marlin_24.py → sparse_quantized_compressors/marlin_24.py}

@@ -17,7 +17,7 @@ from typing import Dict, Generator, Tuple
 
 import numpy as np
 import torch
-from compressed_tensors.compressors import Compressor
+from compressed_tensors.compressors.base import BaseCompressor
 from compressed_tensors.config import CompressionFormat
 from compressed_tensors.quantization import QuantizationArgs, QuantizationStrategy
 from compressed_tensors.quantization.lifecycle.forward import quantize
@@ -35,8 +35,8 @@ from tqdm import tqdm
 _LOGGER: logging.Logger = logging.getLogger(__name__)
 
 
-@Compressor.register(name=CompressionFormat.marlin_24.value)
-class Marlin24Compressor(Compressor):
+@BaseCompressor.register(name=CompressionFormat.marlin_24.value)
+class Marlin24Compressor(BaseCompressor):
     """
     Compresses a quantized model with 2:4 sparsity structure for inference with the
     Marlin24 kernel. Decompression is not implemented for this compressor.

compressed_tensors/config/base.py

@@ -13,7 +13,7 @@
 # limitations under the License.
 
 from enum import Enum
-from typing import Optional
+from typing import List, Optional
 
 from compressed_tensors.registry import RegistryMixin
 from pydantic import BaseModel
@@ -37,11 +37,16 @@ class SparsityCompressionConfig(RegistryMixin, BaseModel):
     Base data class for storing sparsity compression parameters
 
     :param format: name of compression format
+    :param targets: List of layer names or layer types that aren't sparse and should
+        be ignored during compression. By default, assume all layers are targeted
+    :param ignore: List of layer names (unique) to ignore from targets. Defaults to None
     :param global_sparsity: average sparsity of the entire model
     :param sparsity_structure: structure of the sparsity, such as
     "unstructured", "2:4", "8:16" etc
     """
 
     format: str
+    targets: Optional[List[str]] = None
+    ignore: Optional[List[str]] = None
     global_sparsity: Optional[float] = 0.0
     sparsity_structure: Optional[str] = "unstructured"

compressed_tensors/linear/__init__.py

@@ -0,0 +1,13 @@
+# Copyright (c) 2021 - present / Neuralmagic, Inc. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.

compressed_tensors/linear/compressed_linear.py

@@ -0,0 +1,87 @@
+# Copyright (c) 2021 - present / Neuralmagic, Inc. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import torch
+from compressed_tensors.compressors.base import BaseCompressor
+from compressed_tensors.quantization import (
+    QuantizationScheme,
+    QuantizationStatus,
+    initialize_module_for_quantization,
+)
+from torch import Tensor
+from torch.nn import Parameter
+from torch.nn.functional import linear
+from torch.nn.modules import Linear
+
+
+class CompressedLinear(Linear):
+    """
+    Wrapper module for running a compressed forward pass of a quantized Linear module.
+    The wrapped layer will decompressed on each forward call.
+
+    :param module: dense linear module to replace
+    :param quantization_scheme: quantization config for the module to wrap
+    :param quantization_format: compression format module is stored as
+    """
+
+    @classmethod
+    @torch.no_grad()
+    def from_linear(
+        cls,
+        module: Linear,
+        quantization_scheme: QuantizationScheme,
+        quantization_format: str,
+    ):
+        module.__class__ = CompressedLinear
+        module.compressor = BaseCompressor.load_from_registry(quantization_format)
+        device = next(module.parameters()).device
+
+        # this will initialize all the scales and zero points
+        initialize_module_for_quantization(
+            module, quantization_scheme, force_zero_point=False
+        )
+
+        # get the shape and dtype of compressed parameters
+        compression_params = module.compressor.compression_param_info(
+            module.weight.shape, quantization_scheme.weights
+        )
+
+        # no need for this once quantization is initialized, will be replaced
+        # with the compressed parameter
+        delattr(module, "weight")
+
+        # populate compressed weights and quantization parameters
+        for name, (shape, dtype) in compression_params.items():
+            param = Parameter(
+                torch.empty(shape, device=device, dtype=dtype), requires_grad=False
+            )
+            module.register_parameter(name, param)
+
+        # mark module as compressed
+        module.quantization_status = QuantizationStatus.COMPRESSED
+
+        # handles case where forward is wrapped in new_forward by accelerate hooks
+        if hasattr(module, "_old_forward"):
+            module._old_forward = CompressedLinear.forward.__get__(
+                module, CompressedLinear
+            )
+
+        return module
+
+    def forward(self, input: Tensor) -> Tensor:
+        """
+        Decompresses the weight, then runs the wrapped forward pass
+        """
+        uncompressed_weight = self.compressor.decompress_module(self)
+        return linear(input, uncompressed_weight, self.bias)

compressed_tensors/quantization/__init__.py

@@ -19,3 +19,4 @@ from .quant_args import *
 from .quant_config import *
 from .quant_scheme import *
 from .lifecycle import *
+from .cache import QuantizedKVParameterCache

compressed_tensors/quantization/cache.py

@@ -0,0 +1,201 @@
+# Copyright (c) 2021 - present / Neuralmagic, Inc. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+from enum import Enum
+from typing import Any, Dict, List, Optional, Tuple
+
+from compressed_tensors.quantization.observers import Observer
+from compressed_tensors.quantization.quant_args import QuantizationArgs
+from torch import Tensor
+from transformers import DynamicCache as HFDyanmicCache
+
+
+class KVCacheScaleType(Enum):
+    KEY = "k_scale"
+    VALUE = "v_scale"
+
+
+class QuantizedKVParameterCache(HFDyanmicCache):
+
+    """
+    Quantized KV cache used in the forward call based on HF's dynamic cache.
+    Quantization strategy (tensor, group, channel) set from Quantization arg's strategy
+    Singleton, so that the same cache gets reused in all forward call of self_attn.
+    Each time forward is called, .update() is called, and ._quantize(), ._dequantize()
+     gets called appropriately.
+    The size of tensor is
+     `[batch_size, num_heads, seq_len - residual_length, head_dim]`.
+
+
+    Triggered by adding kv_cache_scheme in the recipe.
+
+    Example:
+
+    ```python3
+    recipe = '''
+    quant_stage:
+        quant_modifiers:
+            QuantizationModifier:
+                kv_cache_scheme:
+                    num_bits: 8
+                    type: float
+                    strategy: tensor
+                    dynamic: false
+                    symmetric: true
+    '''
+
+    """
+
+    _instance = None
+    _initialized = False
+
+    def __new__(cls, *args, **kwargs):
+        """Singleton"""
+        if cls._instance is None:
+            cls._instance = super(QuantizedKVParameterCache, cls).__new__(cls)
+        return cls._instance
+
+    def __init__(self, quantization_args: QuantizationArgs):
+        if not self._initialized:
+            super().__init__()
+
+            self.quantization_args = quantization_args
+
+            self.k_observers: List[Observer] = []
+            self.v_observers: List[Observer] = []
+
+            # each index corresponds to layer_idx of the attention layer
+            self.k_scales: List[Tensor] = []
+            self.v_scales: List[Tensor] = []
+
+            self.k_zps: List[Tensor] = []
+            self.v_zps: List[Tensor] = []
+
+            self._initialized = True
+
+    def update(
+        self,
+        key_states: Tensor,
+        value_states: Tensor,
+        layer_idx: int,
+        cache_kwargs: Optional[Dict[str, Any]] = None,
+    ) -> Tuple[Tensor, Tensor]:
+        """
+        Get the k_scale and v_scale and output the
+         fakequant-ed key_states and value_states
+        """
+
+        if len(self.k_observers) <= layer_idx:
+            k_observer = self.quantization_args.get_observer()
+            v_observer = self.quantization_args.get_observer()
+
+            self.k_observers.append(k_observer)
+            self.v_observers.append(v_observer)
+
+        q_key_states = self._quantize(
+            key_states.contiguous(), KVCacheScaleType.KEY, layer_idx
+        )
+        q_value_states = self._quantize(
+            value_states.contiguous(), KVCacheScaleType.VALUE, layer_idx
+        )
+
+        qdq_key_states = self._dequantize(q_key_states, KVCacheScaleType.KEY, layer_idx)
+        qdq_value_states = self._dequantize(
+            q_value_states, KVCacheScaleType.VALUE, layer_idx
+        )
+
+        keys_to_return, values_to_return = qdq_key_states, qdq_value_states
+
+        return keys_to_return, values_to_return
+
+    def get_seq_length(self, layer_idx: Optional[int] = 0) -> int:
+        """
+        Returns the sequence length of the cached states.
+        A layer index can be optionally passed.
+        """
+        if len(self.key_cache) <= layer_idx:
+            return 0
+        # since we cannot get the seq_length of each layer directly and
+        # rely on `_seen_tokens` which is updated every "layer_idx" == 0,
+        # this is a hack to get the actual seq_length for the given layer_idx
+        # this part of code otherwise fails when used to
+        # verify attn_weight shape in some models
+        return self._seen_tokens if layer_idx == 0 else self._seen_tokens - 1
+
+    def reset_states(self):
+        """reset the kv states (used in calibration)"""
+        self.key_cache: List[Tensor] = []
+        self.value_cache: List[Tensor] = []
+        # Used in `generate` to keep tally of how many tokens the cache has seen
+        self._seen_tokens = 0
+        self._quantized_key_cache: List[Tensor] = []
+        self._quantized_value_cache: List[Tensor] = []
+
+    def reset(self):
+        """
+        Reset the instantiation, create new instance on init
+        """
+        QuantizedKVParameterCache._instance = None
+        QuantizedKVParameterCache._initialized = False
+
+    def _quantize(self, tensor, kv_type, layer_idx):
+        """Quantizes a key/value using a defined quantization method."""
+        from compressed_tensors.quantization.lifecycle.forward import quantize
+
+        if kv_type == KVCacheScaleType.KEY:  # key type
+            observer = self.k_observers[layer_idx]
+            scales = self.k_scales
+            zps = self.k_zps
+        else:
+            assert kv_type == KVCacheScaleType.VALUE
+            observer = self.v_observers[layer_idx]
+            scales = self.v_scales
+            zps = self.v_zps
+
+        scale, zp = observer(tensor)
+        if len(scales) <= layer_idx:
+            scales.append(scale)
+            zps.append(zp)
+        else:
+            scales[layer_idx] = scale
+            zps[layer_idx] = scale
+
+        q_tensor = quantize(
+            x=tensor,
+            scale=scale,
+            zero_point=zp,
+            args=self.quantization_args,
+        )
+        return q_tensor
+
+    def _dequantize(self, qtensor, kv_type, layer_idx):
+        """Dequantizes back the tensor that was quantized by `self._quantize()`"""
+        from compressed_tensors.quantization.lifecycle.forward import dequantize
+
+        if kv_type == KVCacheScaleType.KEY:
+            scale = self.k_scales[layer_idx]
+            zp = self.k_zps[layer_idx]
+        else:
+            assert kv_type == KVCacheScaleType.VALUE
+            scale = self.v_scales[layer_idx]
+            zp = self.v_zps[layer_idx]
+
+        qdq_tensor = dequantize(
+            x_q=qtensor,
+            scale=scale,
+            zero_point=zp,
+            args=self.quantization_args,
+        )
+        return qdq_tensor

compressed_tensors/quantization/lifecycle/apply.py

@@ -14,12 +14,14 @@
 
 import logging
 import re
-from collections import OrderedDict
+from collections import OrderedDict, defaultdict
+from copy import deepcopy
 from typing import Dict, Iterable, List, Optional
 from typing import OrderedDict as OrderedDictType
 from typing import Union
 
 import torch
+from compressed_tensors.config import CompressionFormat
 from compressed_tensors.quantization.lifecycle.calibration import (
     set_module_for_calibration,
 )
@@ -41,8 +43,9 @@ from compressed_tensors.quantization.utils import (
     infer_quantization_status,
     is_kv_cache_quant_scheme,
     iter_named_leaf_modules,
+    iter_named_quantizable_modules,
 )
-from compressed_tensors.utils.helpers import fix_fsdp_module_name
+from compressed_tensors.utils.helpers import fix_fsdp_module_name, replace_module
 from compressed_tensors.utils.offload import update_parameter_data
 from compressed_tensors.utils.safetensors_load import get_safetensors_folder
 from torch.nn import Module
@@ -103,13 +106,25 @@ def load_pretrained_quantization(model: Module, model_name_or_path: str):
             )
 
 
-def apply_quantization_config(model: Module, config: QuantizationConfig) -> Dict:
+def apply_quantization_config(
+    model: Module, config: Union[QuantizationConfig, None], run_compressed: bool = False
+) -> OrderedDict:
     """
     Initializes the model for quantization in-place based on the given config
 
     :param model: model to apply quantization config to
     :param config: quantization config
+    :param run_compressed: Whether the model will be run in compressed mode or
+        decompressed fully on load
     """
+    # Workaround for when HF Quantizer passes None, see PR #180
+    if config is None:
+        return OrderedDict()
+
+    # remove reference to the original `config`
+    # argument. This function can mutate it, and we'd
+    # like to keep the original `config` as it is.
+    config = deepcopy(config)
     # build mapping of targets to schemes for easier matching
     # use ordered dict to preserve target ordering in config
     target_to_scheme = OrderedDict()
@@ -119,21 +134,47 @@ def apply_quantization_config(model: Module, config: QuantizationConfig) -> Dict
         for target in scheme.targets:
             target_to_scheme[target] = scheme
 
+    if run_compressed:
+        from compressed_tensors.linear.compressed_linear import CompressedLinear
+
     # list of submodules to ignore
-    ignored_submodules = []
+    ignored_submodules = defaultdict(list)
     # mark appropriate layers for quantization by setting their quantization schemes
-    for name, submodule in iter_named_leaf_modules(model):
+    for name, submodule in iter_named_quantizable_modules(
+        model,
+        include_children=True,
+        include_attn=True,
+    ):  # child modules and attention modules
         # potentially fix module name to remove FSDP wrapper prefix
         name = fix_fsdp_module_name(name)
-        if find_name_or_class_matches(name, submodule, config.ignore):
-            ignored_submodules.append(name)
+        if matches := find_name_or_class_matches(name, submodule, config.ignore):
+            for match in matches:
+                ignored_submodules[match].append(name)
             continue  # layer matches ignore list, continue
+
         targets = find_name_or_class_matches(name, submodule, target_to_scheme)
+
         if targets:
+            # mark modules to be quantized by adding
+            # quant scheme to the matching layers
+            scheme = _scheme_from_targets(target_to_scheme, targets, name)
+            if run_compressed:
+                format = config.format
+                if format != CompressionFormat.dense.value:
+                    if isinstance(submodule, torch.nn.Linear):
+                        # TODO: expand to more module types
+                        compressed_linear = CompressedLinear.from_linear(
+                            submodule,
+                            quantization_scheme=scheme,
+                            quantization_format=format,
+                        )
+                        replace_module(model, name, compressed_linear)
+
             # target matched - add layer and scheme to target list
             submodule.quantization_scheme = _scheme_from_targets(
                 target_to_scheme, targets, name
             )
+
             names_to_scheme[name] = submodule.quantization_scheme.weights
 
     if config.ignore is not None and ignored_submodules is not None:
@@ -143,8 +184,8 @@ def apply_quantization_config(model: Module, config: QuantizationConfig) -> Dict
                 "not found in the model: "
                 f"{set(config.ignore) - set(ignored_submodules)}"
             )
-    # apply current quantization status across all targeted layers
 
+    # apply current quantization status across all targeted layers
     apply_quantization_status(model, config.quantization_status)
     return names_to_scheme
 
@@ -172,6 +213,9 @@ def process_kv_cache_config(
     :param config: the QuantizationConfig
     :return: the QuantizationConfig with additional "kv_cache" group
     """
+    if targets == KV_CACHE_TARGETS:
+        _LOGGER.info(f"KV cache targets set to default value of: {KV_CACHE_TARGETS}")
+
     kv_cache_dict = config.kv_cache_scheme.model_dump()
     kv_cache_scheme = QuantizationScheme(
         output_activations=QuantizationArgs(**kv_cache_dict),
@@ -192,7 +236,12 @@ def apply_quantization_status(model: Module, status: QuantizationStatus):
     current_status = infer_quantization_status(model)
 
     if status >= QuantizationStatus.INITIALIZED > current_status:
-        model.apply(initialize_module_for_quantization)
+        force_zero_point_init = status != QuantizationStatus.COMPRESSED
+        model.apply(
+            lambda module: initialize_module_for_quantization(
+                module, force_zero_point=force_zero_point_init
+            )
+        )
 
     if current_status < status >= QuantizationStatus.CALIBRATION > current_status:
         # only quantize weights up front when our end goal state is calibration,
@@ -273,9 +322,11 @@ def _load_quant_args_from_state_dict(
     """
     scale_name = f"{base_name}_scale"
     zp_name = f"{base_name}_zero_point"
+    g_idx_name = f"{base_name}_g_idx"
 
     state_dict_scale = state_dict.get(f"{module_name}.{scale_name}", None)
     state_dict_zp = state_dict.get(f"{module_name}.{zp_name}", None)
+    state_dict_g_idx = state_dict.get(f"{module_name}.{g_idx_name}", None)
 
     if state_dict_scale is not None:
         # module is quantized
@@ -285,6 +336,9 @@ def _load_quant_args_from_state_dict(
             state_dict_zp = torch.zeros_like(state_dict_scale, device="cpu")
         update_parameter_data(module, state_dict_zp, zp_name)
 
+    if state_dict_g_idx is not None:
+        update_parameter_data(module, state_dict_g_idx, g_idx_name)
+
 
 def _scheme_from_targets(
     target_to_scheme: OrderedDictType[str, QuantizationScheme],

compressed_tensors/quantization/lifecycle/calibration.py

@@ -36,15 +36,15 @@ def set_module_for_calibration(module: Module, quantize_weights_upfront: bool =
     apply to full model with `model.apply(set_module_for_calibration)`
 
     :param module: module to set for calibration
-    :param quantize_weights_upfront: whether to automatically run weight quantization at the
-    start of calibration
+    :param quantize_weights_upfront: whether to automatically
+       run weight quantization at the start of calibration
     """
     if not getattr(module, "quantization_scheme", None):
         # no quantization scheme nothing to do
         return
     status = getattr(module, "quantization_status", None)
     if not status or status != QuantizationStatus.INITIALIZED:
-        raise _LOGGER.warning(
+        _LOGGER.warning(
             f"Attempting set module with status {status} to calibration mode. "
             f"but status is not {QuantizationStatus.INITIALIZED} - you may "
             "be calibrating an uninitialized module which may fail or attempting "
@@ -54,13 +54,13 @@ def set_module_for_calibration(module: Module, quantize_weights_upfront: bool =
     if quantize_weights_upfront and module.quantization_scheme.weights is not None:
         # set weight scale and zero_point up front, calibration data doesn't affect it
         observer = module.weight_observer
+        g_idx = getattr(module, "weight_g_idx", None)
 
-        offloaded = False
-        if is_module_offloaded(module):
+        offloaded = is_module_offloaded(module)
+        if offloaded:
             module._hf_hook.pre_forward(module)
-            offloaded = True
 
-        scale, zero_point = observer(module.weight)
+        scale, zero_point = observer(module.weight, g_idx=g_idx)
         update_parameter_data(module, scale, "weight_scale")
         update_parameter_data(module, zero_point, "weight_zero_point")