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

compressed_tensors/compressors/{naive_quantized.py → quantized_compressors/base.py}

@@ -16,36 +16,51 @@ import logging
 from typing import Dict, Generator, Tuple
 
 import torch
-from compressed_tensors.compressors import Compressor
-from compressed_tensors.config import CompressionFormat
+from compressed_tensors.compressors.base import BaseCompressor
 from compressed_tensors.quantization import QuantizationArgs
-from compressed_tensors.quantization.lifecycle.forward import dequantize, quantize
-from compressed_tensors.quantization.utils import can_quantize
 from compressed_tensors.utils import get_nested_weight_mappings, merge_names
 from safetensors import safe_open
 from torch import Tensor
 from tqdm import tqdm
 
 
-__all__ = [
-    "QuantizationCompressor",
-    "IntQuantizationCompressor",
-    "FloatQuantizationCompressor",
-]
-
 _LOGGER: logging.Logger = logging.getLogger(__name__)
 
+__all__ = ["BaseQuantizationCompressor"]
+
 
-@Compressor.register(name=CompressionFormat.naive_quantized.value)
-class QuantizationCompressor(Compressor):
+class BaseQuantizationCompressor(BaseCompressor):
     """
-    Implements naive compression for quantized models. Weight of each
-    quantized layer is converted from its original float type to the closest Pytorch
-    type to the type specified by the layer's QuantizationArgs.
+    Base class representing a quant compression algorithm. Each child class should
+    implement compression_param_info, compress_weight and decompress_weight.
+
+    Compressors support compressing/decompressing a full module state dict or a single
+    quantized PyTorch leaf module.
+
+    Model Load Lifecycle (run_compressed=False):
+        - ModelCompressor.decompress()
+            - apply_quantization_config()
+            - BaseQuantizationCompressor.decompress()
+                - BaseQuantizationCompressor.decompress_weight()
+
+    Model Save Lifecycle:
+        - ModelCompressor.compress()
+            - BaseQuantizationCompressor.compress()
+                - BaseQuantizationCompressor.compress_weight()
+
+    Module Lifecycle (run_compressed=True):
+        - apply_quantization_config()
+        - compressed_module = CompressedLinear(module)
+            - initialize_module_for_quantization()
+            - BaseQuantizationCompressor.compression_param_info()
+            - register_parameters()
+        - compressed_module.forward()
+            - compressed_module.decompress()
+
+
+    :param config: config specifying compression parameters
     """
 
-    COMPRESSION_PARAM_NAMES = ["weight", "weight_scale", "weight_zero_point"]
-
     def compress(
         self,
         model_state: Dict[str, Tensor],
@@ -57,7 +72,7 @@ class QuantizationCompressor(Compressor):
 
         :param model_state: state dict of uncompressed model
         :param names_to_scheme: quantization args for each quantized weight, needed for
-        quantize function to calculate bit depth
+            quantize function to calculate bit depth
         :return: compressed state dict
         """
         compressed_dict = {}
@@ -66,42 +81,50 @@ class QuantizationCompressor(Compressor):
             f"Compressing model with {len(model_state)} parameterized layers..."
         )
 
-        for name, value in tqdm(model_state.items(), desc="Compressing model"):
+        for name, value in tqdm(model_state.items(), desc="Quantized Compression"):
             if name.endswith(weight_suffix):
                 prefix = name[: -(len(weight_suffix))]
                 scale = model_state.get(merge_names(prefix, "weight_scale"), None)
                 zp = model_state.get(merge_names(prefix, "weight_zero_point"), None)
-                if scale is not None and zp is not None:
+                g_idx = model_state.get(merge_names(prefix, "weight_g_idx"), None)
+                if scale is not None:
                     # weight is quantized, compress it
                     quant_args = names_to_scheme[prefix]
-                    if can_quantize(value, quant_args):
-                        # only quantize if not already quantized
-                        value = quantize(
-                            x=value,
-                            scale=scale,
-                            zero_point=zp,
-                            args=quant_args,
-                            dtype=quant_args.pytorch_dtype(),
-                        )
-            elif name.endswith("zero_point"):
-                if torch.all(value == 0):
-                    # all zero_points are 0, no need to include in
-                    # compressed state_dict
-                    continue
-            compressed_dict[name] = value.to("cpu")
+                    compressed_data = self.compress_weight(
+                        weight=value,
+                        scale=scale,
+                        zero_point=zp,
+                        g_idx=g_idx,
+                        quantization_args=quant_args,
+                        device="cpu",
+                    )
+                    for key, value in compressed_data.items():
+                        compressed_dict[merge_names(prefix, key)] = value
+                else:
+                    compressed_dict[name] = value.to("cpu")
+            elif name.endswith("zero_point") and torch.all(value == 0):
+                continue
+            elif name.endswith("g_idx") and torch.any(value <= -1):
+                continue
+            else:
+                compressed_dict[name] = value.to("cpu")
 
         return compressed_dict
 
     def decompress(
-        self, path_to_model_or_tensors: str, device: str = "cpu", **kwargs
+        self,
+        path_to_model_or_tensors: str,
+        names_to_scheme: Dict[str, QuantizationArgs],
+        device: str = "cpu",
     ) -> Generator[Tuple[str, Tensor], None, None]:
         """
         Reads a 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 path_to_model_or_tensors: path to compressed safetensors model (directory
+            with one or more safetensors files) or compressed tensors file
+        :param names_to_scheme: quantization args for each quantized weight
         :param device: optional device to load intermediate weights into
         :return: compressed state dict
         """
@@ -116,29 +139,8 @@ class QuantizationCompressor(Compressor):
                     weight_data[param_name] = f.get_tensor(full_name)
 
             if "weight_scale" in weight_data:
-                zero_point = weight_data.get("weight_zero_point", None)
-                scale = weight_data["weight_scale"]
-                decompressed = dequantize(
-                    x_q=weight_data["weight"],
-                    scale=scale,
-                    zero_point=zero_point,
+                quant_args = names_to_scheme[weight_name]
+                decompressed = self.decompress_weight(
+                    compressed_data=weight_data, quantization_args=quant_args
                 )
                 yield merge_names(weight_name, "weight"), decompressed
-
-
-@Compressor.register(name=CompressionFormat.int_quantized.value)
-class IntQuantizationCompressor(QuantizationCompressor):
-    """
-    Alias for integer quantized models
-    """
-
-    pass
-
-
-@Compressor.register(name=CompressionFormat.float_quantized.value)
-class FloatQuantizationCompressor(QuantizationCompressor):
-    """
-    Alias for fp quantized models
-    """
-
-    pass

compressed_tensors/compressors/quantized_compressors/naive_quantized.py

@@ -0,0 +1,140 @@
+# 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 typing import Dict, Optional, Tuple
+
+import torch
+from compressed_tensors.compressors.base import BaseCompressor
+from compressed_tensors.compressors.quantized_compressors.base import (
+    BaseQuantizationCompressor,
+)
+from compressed_tensors.config import CompressionFormat
+from compressed_tensors.quantization import QuantizationArgs
+from compressed_tensors.quantization.lifecycle.forward import dequantize, quantize
+from compressed_tensors.quantization.utils import can_quantize
+from torch import Tensor
+
+
+__all__ = [
+    "NaiveQuantizationCompressor",
+    "IntQuantizationCompressor",
+    "FloatQuantizationCompressor",
+]
+
+
+@BaseCompressor.register(name=CompressionFormat.naive_quantized.value)
+class NaiveQuantizationCompressor(BaseQuantizationCompressor):
+    """
+    Implements naive compression for quantized models. Weight of each
+    quantized layer is converted from its original float type to the closest Pytorch
+    type to the type specified by the layer's QuantizationArgs.
+    """
+
+    COMPRESSION_PARAM_NAMES = [
+        "weight",
+        "weight_scale",
+        "weight_zero_point",
+        "weight_g_idx",
+    ]
+
+    def compression_param_info(
+        self,
+        weight_shape: torch.Size,
+        quantization_args: Optional[QuantizationArgs] = None,
+    ) -> Dict[str, Tuple[torch.Size, torch.dtype]]:
+        """
+        Creates a dictionary of expected shapes and dtypes for each compression
+            parameter used by the compressor
+
+        :param weight_shape: uncompressed weight shape
+        :param quantization_args: quantization parameters for the weight
+        :return: dictionary mapping compressed parameter names to shape and dtype
+        """
+        dtype = quantization_args.pytorch_dtype()
+        return {"weight": (weight_shape, dtype)}
+
+    def compress_weight(
+        self,
+        weight: Tensor,
+        scale: Tensor,
+        zero_point: Optional[Tensor] = None,
+        g_idx: Optional[torch.Tensor] = None,
+        quantization_args: Optional[QuantizationArgs] = None,
+        device: Optional[torch.device] = None,
+    ) -> Dict[str, torch.Tensor]:
+        """
+        Compresses a single uncompressed weight
+
+        :param weight: uncompressed weight tensor
+        :param scale: quantization scale for weight
+        :param zero_point: quantization zero point for weight
+        :param g_idx: optional mapping from column index to group index
+        :param quantization_args: quantization parameters for weight
+        :param device: optional device to move compressed output to
+        :return: dictionary of compressed weight data
+        """
+        if can_quantize(weight, quantization_args):
+            quantized_weight = quantize(
+                x=weight,
+                scale=scale,
+                zero_point=zero_point,
+                g_idx=g_idx,
+                args=quantization_args,
+                dtype=quantization_args.pytorch_dtype(),
+            )
+
+            if device is not None:
+                quantized_weight = quantized_weight.to(device)
+
+        return {"weight": quantized_weight}
+
+    def decompress_weight(
+        self,
+        compressed_data: Dict[str, Tensor],
+        quantization_args: Optional[QuantizationArgs] = None,
+    ) -> torch.Tensor:
+        """
+        Decompresses a single compressed weight
+
+        :param compressed_data: dictionary of data needed for decompression
+        :param quantization_args: quantization parameters for the weight
+        :return: tensor of the decompressed weight
+        """
+        weight = compressed_data["weight"]
+        scale = compressed_data["weight_scale"]
+        zero_point = compressed_data.get("weight_zero_point", None)
+        g_idx = compressed_data.get("weight_g_idx", None)
+        decompressed_weight = dequantize(
+            x_q=weight, scale=scale, zero_point=zero_point, g_idx=g_idx
+        )
+
+        return decompressed_weight
+
+
+@BaseCompressor.register(name=CompressionFormat.int_quantized.value)
+class IntQuantizationCompressor(NaiveQuantizationCompressor):
+    """
+    Alias for integer quantized models
+    """
+
+    pass
+
+
+@BaseCompressor.register(name=CompressionFormat.float_quantized.value)
+class FloatQuantizationCompressor(NaiveQuantizationCompressor):
+    """
+    Alias for fp quantized models
+    """
+
+    pass

compressed_tensors/compressors/quantized_compressors/pack_quantized.py

@@ -0,0 +1,211 @@
+# 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 math
+from typing import Dict, Optional, Tuple
+
+import numpy as np
+import torch
+from compressed_tensors.compressors.base import BaseCompressor
+from compressed_tensors.compressors.quantized_compressors.base import (
+    BaseQuantizationCompressor,
+)
+from compressed_tensors.config import CompressionFormat
+from compressed_tensors.quantization import QuantizationArgs
+from compressed_tensors.quantization.lifecycle.forward import dequantize, quantize
+from compressed_tensors.quantization.utils import can_quantize
+from torch import Tensor
+
+
+__all__ = ["PackedQuantizationCompressor", "pack_to_int32", "unpack_from_int32"]
+
+
+@BaseCompressor.register(name=CompressionFormat.pack_quantized.value)
+class PackedQuantizationCompressor(BaseQuantizationCompressor):
+    """
+    Compresses a quantized model by packing every eight 4-bit weights into an int32
+    """
+
+    COMPRESSION_PARAM_NAMES = [
+        "weight_packed",
+        "weight_scale",
+        "weight_zero_point",
+        "weight_g_idx",
+        "weight_shape",
+    ]
+
+    def compression_param_info(
+        self,
+        weight_shape: torch.Size,
+        quantization_args: Optional[QuantizationArgs] = None,
+    ) -> Dict[str, Tuple[torch.Size, torch.dtype]]:
+        """
+        Creates a dictionary of expected shapes and dtypes for each compression
+            parameter used by the compressor
+
+        :param weight_shape: uncompressed weight shape
+        :param quantization_args: quantization parameters for the weight
+        :return: dictionary mapping compressed parameter names to shape and dtype
+        """
+        pack_factor = 32 // quantization_args.num_bits
+        packed_size = math.ceil(weight_shape[1] / pack_factor)
+        return {
+            "weight_packed": (torch.Size((weight_shape[0], packed_size)), torch.int32),
+            "weight_shape": (torch.Size((2,)), torch.int32),
+        }
+
+    def compress_weight(
+        self,
+        weight: Tensor,
+        scale: Tensor,
+        zero_point: Optional[Tensor] = None,
+        g_idx: Optional[torch.Tensor] = None,
+        quantization_args: Optional[QuantizationArgs] = None,
+        device: Optional[torch.device] = None,
+    ) -> Dict[str, torch.Tensor]:
+        """
+        Compresses a single uncompressed weight
+
+        :param weight: uncompressed weight tensor
+        :param scale: quantization scale for weight
+        :param zero_point: quantization zero point for weight
+        :param g_idx: optional mapping from column index to group index
+        :param quantization_args: quantization parameters for weight
+        :param device: optional device to move compressed output to
+        :return: dictionary of compressed weight data
+        """
+        compressed_dict = {}
+        if can_quantize(weight, quantization_args):
+            quantized_weight = quantize(
+                x=weight,
+                scale=scale,
+                zero_point=zero_point,
+                g_idx=g_idx,
+                args=quantization_args,
+                dtype=torch.int8,
+            )
+
+        packed_weight = pack_to_int32(quantized_weight, quantization_args.num_bits)
+        weight_shape = torch.tensor(weight.shape)
+        if device is not None:
+            packed_weight = packed_weight.to(device)
+            weight_shape = weight_shape.to(device)
+
+        compressed_dict["weight_shape"] = weight_shape
+        compressed_dict["weight_packed"] = packed_weight
+
+        return compressed_dict
+
+    def decompress_weight(
+        self,
+        compressed_data: Dict[str, Tensor],
+        quantization_args: Optional[QuantizationArgs] = None,
+    ) -> torch.Tensor:
+        """
+        Decompresses a single compressed weight
+
+        :param compressed_data: dictionary of data needed for decompression
+        :param quantization_args: quantization parameters for the weight
+        :return: tensor of the decompressed weight
+        """
+        weight = compressed_data["weight_packed"]
+        scale = compressed_data["weight_scale"]
+        zero_point = compressed_data.get("weight_zero_point", None)
+        g_idx = compressed_data.get("weight_g_idx", None)
+        original_shape = torch.Size(compressed_data["weight_shape"])
+        num_bits = quantization_args.num_bits
+        unpacked = unpack_from_int32(weight, num_bits, original_shape)
+        decompressed_weight = dequantize(
+            x_q=unpacked, scale=scale, zero_point=zero_point, g_idx=g_idx
+        )
+
+        return decompressed_weight
+
+
+def pack_to_int32(value: torch.Tensor, num_bits: int) -> torch.Tensor:
+    """
+    Packs a tensor of quantized weights stored in int8 into int32s with padding
+
+    :param value: tensor to pack
+    :param num_bits: number of bits used to store underlying data
+    :returns: packed int32 tensor
+    """
+    if value.dtype is not torch.int8:
+        raise ValueError("Tensor must be quantized to torch.int8 before packing")
+
+    if num_bits > 8:
+        raise ValueError("Packing is only supported for less than 8 bits")
+
+    # convert to unsigned for packing
+    offset = pow(2, num_bits) // 2
+    value = (value + offset).to(torch.uint8)
+    value = value.cpu().numpy().astype(np.uint32)
+    pack_factor = 32 // num_bits
+
+    # pad input tensor and initialize packed output
+    packed_size = math.ceil(value.shape[1] / pack_factor)
+    packed = np.zeros((value.shape[0], packed_size), dtype=np.uint32)
+    padding = packed.shape[1] * pack_factor - value.shape[1]
+    value = np.pad(value, pad_width=[(0, 0), (0, padding)], constant_values=0)
+
+    # pack values
+    for i in range(pack_factor):
+        packed |= value[:, i::pack_factor] << num_bits * i
+
+    # convert back to signed and torch
+    packed = np.ascontiguousarray(packed).view(np.int32)
+    return torch.from_numpy(packed)
+
+
+def unpack_from_int32(
+    value: torch.Tensor, num_bits: int, shape: torch.Size
+) -> torch.Tensor:
+    """
+    Unpacks a tensor of packed int32 weights into individual int8s, maintaining the
+    original their bit range
+
+    :param value: tensor to upack
+    :param num_bits: number of bits to unpack each data point into
+    :param shape: shape to unpack into, used to remove padding
+    :returns: unpacked int8 tensor
+    """
+    if value.dtype is not torch.int32:
+        raise ValueError(
+            f"Expected {torch.int32} but got {value.dtype}, Aborting unpack."
+        )
+
+    if num_bits > 8:
+        raise ValueError("Unpacking is only supported for less than 8 bits")
+
+    pack_factor = 32 // num_bits
+
+    # unpack
+    mask = pow(2, num_bits) - 1
+    unpacked = torch.zeros(
+        (value.shape[0], value.shape[1] * pack_factor),
+        device=value.device,
+        dtype=torch.int32,
+    )
+    for i in range(pack_factor):
+        unpacked[:, i::pack_factor] = (value >> (num_bits * i)) & mask
+
+    # remove padding
+    original_row_size = int(shape[1])
+    unpacked = unpacked[:, :original_row_size]
+
+    # bits are packed in unsigned format, reformat to signed
+    # update the value range from unsigned to signed
+    offset = pow(2, num_bits) // 2
+    unpacked = (unpacked - offset).to(torch.int8)
+
+    return unpacked

compressed_tensors/compressors/sparse_compressors/__init__.py

@@ -0,0 +1,18 @@
+# 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 .base import *
+from .dense import *
+from .sparse_bitmask import *

compressed_tensors/compressors/sparse_compressors/base.py

@@ -0,0 +1,110 @@
+# 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 logging
+from typing import Dict, Generator, Tuple
+
+from compressed_tensors.compressors.base import BaseCompressor
+from compressed_tensors.utils import get_nested_weight_mappings, merge_names
+from safetensors import safe_open
+from torch import Tensor
+from tqdm import tqdm
+
+
+__all__ = ["BaseSparseCompressor"]
+
+_LOGGER: logging.Logger = logging.getLogger(__name__)
+
+
+class BaseSparseCompressor(BaseCompressor):
+    """
+    Base class representing a sparse compression algorithm. Each child class should
+    implement compression_param_info, compress_weight and decompress_weight.
+
+    Compressors support compressing/decompressing a full module state dict or a single
+    quantized PyTorch leaf module.
+
+    Model Load Lifecycle (run_compressed=False):
+        - ModelCompressor.decompress()
+            - apply_quantization_config()
+            - BaseSparseCompressor.decompress()
+                - BaseSparseCompressor.decompress_weight()
+
+    Model Save Lifecycle:
+        - ModelCompressor.compress()
+            - BaseSparseCompressor.compress()
+                - BaseSparseCompressor.compress_weight()
+
+    Module Lifecycle (run_compressed=True):
+        - apply_quantization_config()
+        - compressed_module = CompressedLinear(module)
+            - initialize_module_for_quantization()
+            - BaseSparseCompressor.compression_param_info()
+            - register_parameters()
+        - compressed_module.forward()
+            - compressed_module.decompress()
+
+
+    :param config: config specifying compression parameters
+    """
+
+    def compress(self, model_state: Dict[str, Tensor]) -> Dict[str, Tensor]:
+        """
+        Compresses a dense state dict using bitmask compression
+
+        :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"):
+            compression_data = self.compress_weight(name, value)
+            for key in compression_data.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(compression_data)
+
+        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)
+            decompressed = self.decompress_weight(weight_data)
+            yield weight_name, decompressed

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

@@ -14,13 +14,13 @@
 
 from typing import Dict, Generator, Tuple
 
-from compressed_tensors.compressors import Compressor
+from compressed_tensors.compressors.base import BaseCompressor
 from compressed_tensors.config import CompressionFormat
 from torch import Tensor
 
 
-@Compressor.register(name=CompressionFormat.dense.value)
-class DenseCompressor(Compressor):
+@BaseCompressor.register(name=CompressionFormat.dense.value)
+class DenseCompressor(BaseCompressor):
     """
     Identity compressor for dense models, returns the original state_dict
     """