compressed-tensors-nightly 0.4.0.20240721__py3-none-any.whl → 0.4.0.20240731__py3-none-any.whl

compressed_tensors/compressors/model_compressor.py

@@ -39,10 +39,10 @@ from compressed_tensors.quantization.utils import (
     is_module_quantized,
     iter_named_leaf_modules,
 )
-from compressed_tensors.utils import get_safetensors_folder
+from compressed_tensors.utils import get_safetensors_folder, update_parameter_data
 from compressed_tensors.utils.helpers import fix_fsdp_module_name
 from torch import Tensor
-from torch.nn import Module, Parameter
+from torch.nn import Module
 from tqdm import tqdm
 from transformers import AutoConfig
 from transformers.file_utils import CONFIG_NAME
@@ -307,12 +307,10 @@ class ModelCompressor:
 
     def _replace_weights(self, dense_weight_generator, model):
         for name, data in tqdm(dense_weight_generator, desc="Decompressing model"):
-            # loading the decompressed weights into the model
-            model_device = operator.attrgetter(name)(model).device
-            data_old = operator.attrgetter(name)(model)
-            data_dtype = data_old.dtype
-            data_new = Parameter(data.to(model_device).to(data_dtype))
-            data_old.data = data_new.data
+            split_name = name.split(".")
+            prefix, param_name = ".".join(split_name[:-1]), split_name[-1]
+            module = operator.attrgetter(prefix)(model)
+            update_parameter_data(module, data, param_name)
 
 
 def map_modules_to_quant_args(model: Module) -> Dict:

compressed_tensors/quantization/lifecycle/apply.py

@@ -43,6 +43,7 @@ from compressed_tensors.quantization.utils import (
     iter_named_leaf_modules,
 )
 from compressed_tensors.utils.helpers import fix_fsdp_module_name
+from compressed_tensors.utils.offload import update_parameter_data
 from compressed_tensors.utils.safetensors_load import get_safetensors_folder
 from torch.nn import Module
 
@@ -265,19 +266,17 @@ def _load_quant_args_from_state_dict(
     """
     scale_name = f"{base_name}_scale"
     zp_name = f"{base_name}_zero_point"
-    device = next(module.parameters()).device
-
-    scale = getattr(module, scale_name, None)
-    zp = getattr(module, zp_name, None)
-    if scale is not None:
-        state_dict_scale = state_dict[f"{module_name}.{scale_name}"]
-        scale.data = state_dict_scale.to(device).to(scale.dtype)
-    if zp is not None:
-        zp_from_state = state_dict.get(f"{module_name}.{zp_name}", None)
-        if zp_from_state is not None:  # load the non-zero zero points
-            zp.data = zp_from_state.to(device).to(zp.dtype)
-        else:  # fill with zeros matching scale shape
-            zp.data = torch.zeros_like(scale, dtype=zp.dtype).to(device)
+
+    state_dict_scale = state_dict.get(f"{module_name}.{scale_name}", None)
+    state_dict_zp = state_dict.get(f"{module_name}.{zp_name}", None)
+
+    if state_dict_scale is not None:
+        # module is quantized
+        update_parameter_data(module, state_dict_scale, scale_name)
+        if state_dict_zp is None:
+            # fill in zero point for symmetric quantization
+            state_dict_zp = torch.zeros_like(state_dict_scale, device="cpu")
+        update_parameter_data(module, state_dict_zp, zp_name)
 
 
 def _scheme_from_targets(

compressed_tensors/quantization/lifecycle/calibration.py

@@ -16,6 +16,7 @@
 import logging
 
 from compressed_tensors.quantization.quant_config import QuantizationStatus
+from compressed_tensors.utils import is_module_offloaded, update_parameter_data
 from torch.nn import Module
 
 
@@ -48,4 +49,20 @@ def set_module_for_calibration(module: Module):
             "to re-calibrate a frozen module"
         )
 
+    if 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
+
+        offloaded = False
+        if is_module_offloaded(module):
+            module._hf_hook.pre_forward(module)
+            offloaded = True
+
+        scale, zero_point = observer(module.weight)
+        update_parameter_data(module, scale, "weight_scale")
+        update_parameter_data(module, zero_point, "weight_zero_point")
+
+        if offloaded:
+            module._hf_hook.post_forward(module, None)
+
     module.quantization_status = QuantizationStatus.CALIBRATION

compressed_tensors/quantization/lifecycle/forward.py

@@ -25,6 +25,7 @@ from compressed_tensors.quantization.quant_args import (
 )
 from compressed_tensors.quantization.quant_config import QuantizationStatus
 from compressed_tensors.quantization.quant_scheme import QuantizationScheme
+from compressed_tensors.utils import update_parameter_data
 from torch.nn import Module
 
 
@@ -312,16 +313,19 @@ def maybe_calibrate_or_quantize(
         scale = getattr(module, f"{base_name}_scale")
         zero_point = getattr(module, f"{base_name}_zero_point")
 
-        if module.quantization_status == QuantizationStatus.CALIBRATION:
+        if (
+            module.quantization_status == QuantizationStatus.CALIBRATION
+            and base_name != "weight"
+        ):
             # calibration mode - get new quant params from observer
             observer = getattr(module, f"{base_name}_observer")
 
             updated_scale, updated_zero_point = observer(value)
 
             # update scale and zero point
-            device = next(module.parameters()).device
-            scale.data = updated_scale.to(device)
-            zero_point.data = updated_zero_point.to(device)
+            update_parameter_data(module, updated_scale, f"{base_name}_scale")
+            update_parameter_data(module, updated_zero_point, f"{base_name}_zero_point")
+
     return fake_quantize(value, scale, zero_point, args)
 
 

compressed_tensors/quantization/lifecycle/initialize.py

@@ -17,6 +17,8 @@ import logging
 from typing import Optional
 
 import torch
+from accelerate.hooks import add_hook_to_module, remove_hook_from_module
+from accelerate.utils import PrefixedDataset
 from compressed_tensors.quantization.lifecycle.forward import (
     wrap_module_forward_quantized,
 )
@@ -26,6 +28,7 @@ from compressed_tensors.quantization.quant_args import (
 )
 from compressed_tensors.quantization.quant_config import QuantizationStatus
 from compressed_tensors.quantization.quant_scheme import QuantizationScheme
+from compressed_tensors.utils import get_execution_device, is_module_offloaded
 from torch.nn import Module, Parameter
 
 
@@ -81,9 +84,32 @@ def initialize_module_for_quantization(
     module.quantization_scheme = scheme
     module.quantization_status = QuantizationStatus.INITIALIZED
 
+    offloaded = False
+    if is_module_offloaded(module):
+        offloaded = True
+        hook = module._hf_hook
+        prefix_dict = module._hf_hook.weights_map
+        new_prefix = {}
+
+        # recreate the prefix dict (since it is immutable)
+        # and add quantization parameters
+        for key, data in module.named_parameters():
+            if key not in prefix_dict:
+                new_prefix[f"{prefix_dict.prefix}{key}"] = data
+            else:
+                new_prefix[f"{prefix_dict.prefix}{key}"] = prefix_dict[key]
+        new_prefix_dict = PrefixedDataset(new_prefix, prefix_dict.prefix)
+        remove_hook_from_module(module)
+
     # wrap forward call of module to perform quantized actions based on calltime status
     wrap_module_forward_quantized(module, scheme)
 
+    if offloaded:
+        # we need to re-add the hook for offloading now that we've wrapped forward
+        add_hook_to_module(module, hook)
+        if prefix_dict is not None:
+            module._hf_hook.weights_map = new_prefix_dict
+
 
 def _initialize_scale_zero_point_observer(
     module: Module,
@@ -99,6 +125,8 @@ def _initialize_scale_zero_point_observer(
         return  # no need to register a scale and zero point for a dynamic observer
 
     device = next(module.parameters()).device
+    if is_module_offloaded(module):
+        device = get_execution_device(module)
 
     # infer expected scale/zero point shape
     expected_shape = 1  # per tensor

compressed_tensors/quantization/quant_config.py

@@ -239,3 +239,14 @@ class QuantizationConfig(BaseModel):
             format=format,
             ignore=consolidated_ignore,
         )
+
+    def requires_calibration_data(self):
+        for _, scheme in self.config_groups.items():
+            if scheme.input_activations is not None:
+                if not scheme.input_activations.dynamic:
+                    return True
+            if scheme.output_activations is not None:
+                if not scheme.output_activations.dynamic:
+                    return True
+
+        return False

compressed_tensors/quantization/quant_scheme.py

@@ -165,7 +165,7 @@ W4A8 = dict(
     input_activations=QuantizationArgs(
         num_bits=8,
         type=QuantizationType.INT,
-        strategy=QuantizationStrategy.TENSOR,
+        strategy=QuantizationStrategy.TOKEN,
         symmetric=True,
         dynamic=True,
     ),
@@ -189,6 +189,24 @@ FP8 = dict(
     ),
 )
 
+# FP8 weights and FP8 dynamic activations quantization
+FP8_DYNAMIC = dict(
+    weights=QuantizationArgs(
+        num_bits=8,
+        type=QuantizationType.FLOAT,
+        strategy=QuantizationStrategy.CHANNEL,
+        symmetric=True,
+        dynamic=False,
+    ),
+    input_activations=QuantizationArgs(
+        num_bits=8,
+        type=QuantizationType.FLOAT,
+        strategy=QuantizationStrategy.TOKEN,
+        symmetric=True,
+        dynamic=True,
+    ),
+)
+
 PRESET_SCHEMES = {
     # Integer weight only schemes
     "W8A16": W8A16,
@@ -198,4 +216,5 @@ PRESET_SCHEMES = {
     "W4A8": W4A8,
     # Float weight and activation schemes
     "FP8": FP8,
+    "FP8_DYNAMIC": FP8_DYNAMIC,
 }

compressed_tensors/utils/__init__.py

@@ -14,6 +14,7 @@
 # flake8: noqa
 
 from .helpers import *
+from .offload import *
 from .permutations_24 import *
 from .safetensors_load import *
 from .semi_structured_conversions import *

compressed_tensors/utils/offload.py

@@ -0,0 +1,104 @@
+# 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 torch.nn import Module
+
+
+__all__ = [
+    "is_module_offloaded",
+    "get_execution_device",
+    "get_offloaded_device",
+    "update_prefix_dict",
+    "update_parameter_data",
+]
+
+
+def is_module_offloaded(module: Module) -> bool:
+    """
+    :param module: layer to check
+    :return: True if layer is offloaded from GPU, False otherwise
+    """
+    return hasattr(module, "_hf_hook") and module._hf_hook.offload
+
+
+def get_execution_device(module: Module) -> torch.device:
+    """
+    :param module: layer to check
+    :return: device layer is loaded onto during forward pass
+    """
+    if is_module_offloaded(module):
+        return module._hf_hook.execution_device
+    return next(module.parameters()).device
+
+
+def get_offloaded_device(module: Module) -> torch.device:
+    """
+    :param module: layer to check
+    :return: device layer is offloaded to onto after forward pass
+    """
+    if is_module_offloaded(module):
+        first_key = list(module._hf_hook.weights_map.keys())[0]
+        prefix_dataset = module._hf_hook.weights_map.dataset
+        return prefix_dataset[first_key].device
+    return next(module.parameters()).device
+
+
+def update_prefix_dict(module: Module, key: str, data: torch.Tensor):
+    """
+    Updates the offloaded state dict for a given module. Parameter named key is replaced
+    by data. This is neccesary because parameter updates for offloaded modules do not
+    persist automatically between loads. This function only affects the offloaded
+    state dict and not the current state of the loaded module.
+
+    :param module: layer containing the parameter to update
+    :param key: name of parameter to update
+    :param data: tensor to update parameter with in the offloaded state dict
+    """
+    if not is_module_offloaded(module):
+        raise ValueError("Prefix dict is only applicable to offloaded modules")
+    prefix_dict = module._hf_hook.weights_map
+    prefix_dict.dataset[f"{prefix_dict.prefix}{key}"] = data
+
+
+def update_parameter_data(
+    module: Module, new_param_data: torch.Tensor, param_name: str
+):
+    """
+    Updates the paramter value named param_name for a given module. This function
+    updates both the current loaded module state and the offloaded state dict if
+    the module is offloaded. This is neccesary because parameter updates for offloaded
+    modules do not persist automatically between loads.
+
+    :param module: layer containing the parameter to update
+    :param new_param_data: tensor to update parameter with
+    :param param_name:
+    """
+    device = next(module.parameters()).device
+
+    offloaded = False
+    if is_module_offloaded(module):
+        offload_device = get_offloaded_device(module)
+        offloaded = True
+
+    parameter = getattr(module, param_name, None)
+    dtype = parameter.dtype
+    parameter.data = new_param_data.to(device).to(dtype)
+
+    if offloaded:
+        prefix_dict = module._hf_hook.weights_map.dataset
+        prefix = module._hf_hook.weights_map.prefix
+        prefix_dict[f"{prefix}{param_name}"] = new_param_data.to(offload_device).to(
+            dtype
+        )

{compressed_tensors_nightly-0.4.0.20240721.dist-info → compressed_tensors_nightly-0.4.0.20240731.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: compressed-tensors-nightly
-Version: 0.4.0.20240721
+Version: 0.4.0.20240731
 Summary: Library for utilization of compressed safetensors of neural network models
 Home-page: https://github.com/neuralmagic/compressed-tensors
 Author: Neuralmagic, Inc.
@@ -10,6 +10,7 @@ Description-Content-Type: text/markdown
 License-File: LICENSE
 Requires-Dist: torch >=1.7.0
 Requires-Dist: transformers
+Requires-Dist: accelerate
 Requires-Dist: pydantic >=2.0
 Provides-Extra: dev
 Requires-Dist: black ==22.12.0 ; extra == 'dev'

{compressed_tensors_nightly-0.4.0.20240721.dist-info → compressed_tensors_nightly-0.4.0.20240731.dist-info}/RECORD

@@ -6,7 +6,7 @@ compressed_tensors/compressors/base.py,sha256=-rqT2h9G2iwDkwrVj0d0jxxn9h0dccJA1m
 compressed_tensors/compressors/dense.py,sha256=xcWECjcRY4INN6jC7vHx5wvUX3NmnKlxA9SVE1A6m2Q,1267
 compressed_tensors/compressors/helpers.py,sha256=k9avlkmeYj6vkOAvl-MgcixtP7ib24SCfhzZ-RusXfw,5403
 compressed_tensors/compressors/marlin_24.py,sha256=e7fGUyZbjUpA5VUMCPxqcYPGNiwoDKupHJaXWCoVKRw,9410
-compressed_tensors/compressors/model_compressor.py,sha256=9dyM2mvAgO7QeFTBWXBzT29JtmRMKQWWU7xh8StaFyI,13446
+compressed_tensors/compressors/model_compressor.py,sha256=b7jPE4czwP9uulIZML5qUQAvQaQzElwzUGwat7jlpgI,13352
 compressed_tensors/compressors/naive_quantized.py,sha256=6_1wuTF96-lw-UzzrsiEX_ipciKiQQJoZ8uotVwtbyQ,5569
 compressed_tensors/compressors/pack_quantized.py,sha256=tnhqvkko6fIaTywI2JNvh5lE2xXWKJ_hYShv_s6C9Vk,8506
 compressed_tensors/compressors/sparse_bitmask.py,sha256=kiDwBlFV0sJGLcIdDYxIiuF64ccgwDfqq1hWRQThYDc,8647
@@ -16,16 +16,16 @@ compressed_tensors/config/dense.py,sha256=NgSxnFCnckU9-iunxEaqiFwqgdO7YYxlWKR74j
 compressed_tensors/config/sparse_bitmask.py,sha256=pZUboRNZTu6NajGOQEFExoPknak5ynVAUeiiYpS1Gt8,1308
 compressed_tensors/quantization/__init__.py,sha256=83J5bPB7PavN2TfCoW7_vEDhfYpm4TDrqYO9vdSQ5bk,760
 compressed_tensors/quantization/quant_args.py,sha256=Vc_tWSTcbZZsMJlACpLq4JEPvGx87izc8VEx-mcXjoM,5621
-compressed_tensors/quantization/quant_config.py,sha256=PU3BchHm09ks6_yAderrHoIZI07zBlU9ejC87v3A-54,9568
-compressed_tensors/quantization/quant_scheme.py,sha256=IKTtMfusSe7x31t7hipBfptTbCwGd9eGMtrWC1sPM9o,5522
+compressed_tensors/quantization/quant_config.py,sha256=NpVu8YJ4Xw2pIQW_PGaNaml8kx1bUnxkvb0jBYWbKdE,9971
+compressed_tensors/quantization/quant_scheme.py,sha256=_RKOFJI0T5xJVBLX63UeYkSY4EFAecsBnqzUIVBjeU0,6014
 compressed_tensors/quantization/lifecycle/__init__.py,sha256=MXE2E7GfIfRRfhrdGy2Og3AZOz5N59B0ZGFcsD89y6c,821
-compressed_tensors/quantization/lifecycle/apply.py,sha256=fyv5ujZC0__oG1ESOTmMyMsKK7DGAxG7uQI7_sxT7Mw,13308
-compressed_tensors/quantization/lifecycle/calibration.py,sha256=mLns4jlaWmBwOW8Jtlm5bMX-JET1AiZYUBO7qa-XuxI,1776
+compressed_tensors/quantization/lifecycle/apply.py,sha256=aamouy1IWCSGl5_lfC7rZrUNbMEfhyHQrZFGEm2VH4w,13242
+compressed_tensors/quantization/lifecycle/calibration.py,sha256=n-m4xwa9Ds2xrltp6r0rhuVLJhE8bQ1LnifrcrbA-ig,2448
 compressed_tensors/quantization/lifecycle/compressed.py,sha256=VreB10xPwgSLQQlTu20UCrFpRS--cA7-lx5s7nrPPrg,2247
-compressed_tensors/quantization/lifecycle/forward.py,sha256=yz9l93EXQtOR-DkGbrEA62rdSleNnmwZfraD0whH758,12271
+compressed_tensors/quantization/lifecycle/forward.py,sha256=6PSXYcf-R1dOY8zsuIWnBaoyARNymYc3-qvV6-L7SlI,12397
 compressed_tensors/quantization/lifecycle/frozen.py,sha256=h1XYt89MouBTf3jTYLG_6OdFxIu5q2N8tPjsy6J4E6Y,1726
 compressed_tensors/quantization/lifecycle/helpers.py,sha256=xDkM3yVpGVnwAdg2aUOmrlDPaOksi-bavSQ5mMeOQlk,1651
-compressed_tensors/quantization/lifecycle/initialize.py,sha256=kIEx6a7UyqAIG_ZPNBhijrDiAHnp2wR7K_GC3envz4M,4631
+compressed_tensors/quantization/lifecycle/initialize.py,sha256=oCD8pgmHT3lW5J7zdsSN3YzEQIhTfE7M01R5Wb0wpck,5801
 compressed_tensors/quantization/observers/__init__.py,sha256=DNH31NQYrIBBcmHsMyFA6whh4pbRsLwuNa6L8AeXaGc,745
 compressed_tensors/quantization/observers/base.py,sha256=2WO7N2eyXf1r1gxVidos1bUS5o7pcrpug4gQgHIazrQ,6794
 compressed_tensors/quantization/observers/helpers.py,sha256=s_A23Qa_BLfOdHJCN5bm-qPWkhjjj_RIVrhSp1Y9Dtk,4211
@@ -35,13 +35,14 @@ compressed_tensors/quantization/utils/__init__.py,sha256=VdtEmP0bvuND_IGQnyqUPc5
 compressed_tensors/quantization/utils/helpers.py,sha256=YjXABJQUnelof-z7qcwck6fnrFLh4uMSrOmPiqNp_RY,8591
 compressed_tensors/registry/__init__.py,sha256=FwLSNYqfIrb5JD_6OK_MT4_svvKTN_nEhpgQlQvGbjI,658
 compressed_tensors/registry/registry.py,sha256=fxjOjh2wklCvJhQxwofdy-zV8q7MkQ85SLG77nml2iA,11890
-compressed_tensors/utils/__init__.py,sha256=dvAatm3p0He4J7u5ZmZYVa8Iwpwq3ZSm6S9ZJleta5M,762
+compressed_tensors/utils/__init__.py,sha256=rvbIJlvdKYn4iX7r3KP6peCbU5uyMzgxwhsQstLoMxQ,785
 compressed_tensors/utils/helpers.py,sha256=d3yP9ViQ8R3GzMHfohxNlaokzyrRuj2PyjxWAJZmSws,3156
+compressed_tensors/utils/offload.py,sha256=BL7_cNAHTKbSta179R5R4ASk6oXuZhTJDY4D_8Lv2OE,3717
 compressed_tensors/utils/permutations_24.py,sha256=kx6fsfDHebx94zsSzhXGyCyuC9sVyah6BUUir_StT28,2530
 compressed_tensors/utils/safetensors_load.py,sha256=0MheXwx1jeY12PeISppiSIZHs6rmN2YddwPpFb9V67I,8527
 compressed_tensors/utils/semi_structured_conversions.py,sha256=g1EZHzdv-ko7ufPX430dp7wE33o6FWJXuSP4zZydCu0,13488
-compressed_tensors_nightly-0.4.0.20240721.dist-info/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
-compressed_tensors_nightly-0.4.0.20240721.dist-info/METADATA,sha256=0rv0MZclUora99WyqSEnHjoJu-2n-YzYNpoJ85eWGNU,5668
-compressed_tensors_nightly-0.4.0.20240721.dist-info/WHEEL,sha256=GJ7t_kWBFywbagK5eo9IoUwLW6oyOeTKmQ-9iHFVNxQ,92
-compressed_tensors_nightly-0.4.0.20240721.dist-info/top_level.txt,sha256=w2i-GyPs2s1UwVxvutSvN_lM22SXC2hQFBmoMcPnV7Y,19
-compressed_tensors_nightly-0.4.0.20240721.dist-info/RECORD,,
+compressed_tensors_nightly-0.4.0.20240731.dist-info/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
+compressed_tensors_nightly-0.4.0.20240731.dist-info/METADATA,sha256=2uwF7iMEXcX3ZvX-mCcklcr-qdqhlap0Ld5WDxL41pQ,5694
+compressed_tensors_nightly-0.4.0.20240731.dist-info/WHEEL,sha256=GJ7t_kWBFywbagK5eo9IoUwLW6oyOeTKmQ-9iHFVNxQ,92
+compressed_tensors_nightly-0.4.0.20240731.dist-info/top_level.txt,sha256=w2i-GyPs2s1UwVxvutSvN_lM22SXC2hQFBmoMcPnV7Y,19
+compressed_tensors_nightly-0.4.0.20240731.dist-info/RECORD,,