compressed-tensors 0.3.3__py3-none-any.whl → 0.4.0__py3-none-any.whl

compressed_tensors/compressors/model_compressor.py

@@ -0,0 +1,315 @@
+# 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 json
+import logging
+import operator
+import os
+from copy import deepcopy
+from typing import Any, Dict, Optional, Union
+
+from compressed_tensors.base import (
+    COMPRESSION_CONFIG_NAME,
+    QUANTIZATION_CONFIG_NAME,
+    SPARSITY_CONFIG_NAME,
+)
+from compressed_tensors.compressors import Compressor
+from compressed_tensors.config import SparsityCompressionConfig
+from compressed_tensors.quantization import (
+    QuantizationConfig,
+    QuantizationStatus,
+    apply_quantization_config,
+    load_pretrained_quantization,
+)
+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.helpers import fix_fsdp_module_name
+from torch import Tensor
+from torch.nn import Module, Parameter
+from tqdm import tqdm
+from transformers import AutoConfig
+from transformers.file_utils import CONFIG_NAME
+
+
+__all__ = ["ModelCompressor", "map_modules_to_quant_args"]
+
+_LOGGER: logging.Logger = logging.getLogger(__name__)
+
+
+class ModelCompressor:
+    """
+    Handles compression and decompression of a model with a sparsity config and/or
+    quantization config.
+
+    Compression LifeCycle
+        - compressor = ModelCompressor.from_pretrained_model(model)
+        - compressed_state_dict = compressor.compress(model, state_dict)
+            - compressor.quantization_compressor.compress(model, state_dict)
+            - compressor.sparsity_compressor.compress(model, state_dict)
+        - model.save_pretrained(output_dir, state_dict=compressed_state_dict)
+        - compressor.update_config(output_dir)
+
+    Decompression LifeCycle
+        - compressor = ModelCompressor.from_pretrained(comp_model_path)
+        - model = AutoModel.from_pretrained(comp_model_path)
+        - compressor.decompress(comp_model_path, model)
+            - compressor.sparsity_compressor.decompress(comp_model_path, model)
+            - compressor.quantization_compressor.decompress(comp_model_path, model)
+
+    :param sparsity_config: config specifying sparsity compression parameters
+    :param quantization_config: config specifying quantization compression parameters
+    """
+
+    @classmethod
+    def from_pretrained(
+        cls,
+        pretrained_model_name_or_path: str,
+    ) -> Optional["ModelCompressor"]:
+        """
+        Given a path to a model config, extract the sparsity and/or quantization
+        configs and load a ModelCompressor
+
+        :param pretrained_model_name_or_path: path to model config on disk or HF hub
+        :return: compressor for the extracted configs
+        """
+        config = AutoConfig.from_pretrained(pretrained_model_name_or_path)
+        compression_config = getattr(config, COMPRESSION_CONFIG_NAME, None)
+        return cls.from_compression_config(compression_config)
+
+    @classmethod
+    def from_compression_config(cls, compression_config: Dict[str, Any]):
+        """
+        :param compression_config: compression/quantization config dictionary
+            found under key "quantization_config" in HF model config
+        :return: compressor for the extracted configs
+        """
+        if compression_config is None:
+            return None
+
+        try:
+            from transformers.utils.quantization_config import CompressedTensorsConfig
+
+            if isinstance(compression_config, CompressedTensorsConfig):
+                compression_config = compression_config.to_dict()
+        except ImportError:
+            pass
+
+        sparsity_config = cls.parse_sparsity_config(compression_config)
+        quantization_config = cls.parse_quantization_config(compression_config)
+        if sparsity_config is None and quantization_config is None:
+            return None
+
+        if sparsity_config is not None and not isinstance(
+            sparsity_config, SparsityCompressionConfig
+        ):
+            format = sparsity_config.get("format")
+            sparsity_config = SparsityCompressionConfig.load_from_registry(
+                format, **sparsity_config
+            )
+        if quantization_config is not None and not isinstance(
+            quantization_config, QuantizationConfig
+        ):
+            quantization_config = QuantizationConfig.parse_obj(quantization_config)
+
+        return cls(
+            sparsity_config=sparsity_config, quantization_config=quantization_config
+        )
+
+    @classmethod
+    def from_pretrained_model(
+        cls,
+        model: Module,
+        sparsity_config: Union[SparsityCompressionConfig, str, None] = None,
+        quantization_format: Optional[str] = None,
+    ) -> Optional["ModelCompressor"]:
+        """
+        Given a pytorch model and optional sparsity and/or quantization configs,
+        load the appropriate compressors
+
+        :param model: pytorch model to target for compression
+        :param sparsity_config: a filled in sparsity config or string corresponding
+            to a sparsity compression algorithm
+        :param quantization_format: string corresponding to a quantization compression
+            algorithm
+        :return: compressor for the extracted configs
+        """
+        quantization_config = QuantizationConfig.from_pretrained(
+            model, format=quantization_format
+        )
+
+        if isinstance(sparsity_config, str):  # we passed in a sparsity format
+            sparsity_config = SparsityCompressionConfig.load_from_registry(
+                sparsity_config
+            )
+
+        if sparsity_config is None and quantization_config is None:
+            return None
+
+        return cls(
+            sparsity_config=sparsity_config, quantization_config=quantization_config
+        )
+
+    @staticmethod
+    def parse_sparsity_config(compression_config: Dict) -> Union[Dict, None]:
+        if compression_config is None:
+            return None
+        if SPARSITY_CONFIG_NAME not in compression_config:
+            return None
+        if hasattr(compression_config, SPARSITY_CONFIG_NAME):
+            # for loaded HFQuantizer config
+            return getattr(compression_config, SPARSITY_CONFIG_NAME)
+
+        # SparseAutoModel format
+        return compression_config.get(SPARSITY_CONFIG_NAME, None)
+
+    @staticmethod
+    def parse_quantization_config(compression_config: Dict) -> Union[Dict, None]:
+        if compression_config is None:
+            return None
+
+        if hasattr(compression_config, QUANTIZATION_CONFIG_NAME):
+            # for loaded HFQuantizer config
+            return getattr(compression_config, QUANTIZATION_CONFIG_NAME)
+
+        # SparseAutoModel format
+        quantization_config = deepcopy(compression_config)
+        quantization_config.pop(SPARSITY_CONFIG_NAME, None)
+        if len(quantization_config) == 0:
+            quantization_config = None
+        return quantization_config
+
+    def __init__(
+        self,
+        sparsity_config: Optional[SparsityCompressionConfig] = None,
+        quantization_config: Optional[QuantizationConfig] = None,
+    ):
+        self.sparsity_config = sparsity_config
+        self.quantization_config = quantization_config
+        self.sparsity_compressor = None
+        self.quantization_compressor = None
+
+        if sparsity_config is not None:
+            self.sparsity_compressor = Compressor.load_from_registry(
+                sparsity_config.format, config=sparsity_config
+            )
+        if quantization_config is not None:
+            self.quantization_compressor = Compressor.load_from_registry(
+                quantization_config.format, config=quantization_config
+            )
+
+    def compress(
+        self, model: Module, state_dict: Optional[Dict[str, Tensor]] = None
+    ) -> Dict[str, Tensor]:
+        """
+        Compresses a dense state dict or model with sparsity and/or quantization
+
+        :param model: uncompressed model to compress
+        :param model_state: optional uncompressed state_dict to insert into model
+        :return: compressed state dict
+        """
+        if state_dict is None:
+            state_dict = model.state_dict()
+
+        compressed_state_dict = state_dict
+        quantized_modules_to_args = map_modules_to_quant_args(model)
+        if self.quantization_compressor is not None:
+            compressed_state_dict = self.quantization_compressor.compress(
+                state_dict, model_quant_args=quantized_modules_to_args
+            )
+
+        if self.sparsity_compressor is not None:
+            compressed_state_dict = self.sparsity_compressor.compress(
+                compressed_state_dict
+            )
+
+        return compressed_state_dict
+
+    def decompress(self, model_path: str, model: Module):
+        """
+        Overwrites the weights in model with weights decompressed from model_path
+
+        :param model_path: path to compressed weights
+        :param model: pytorch model to load decompressed weights into
+        """
+        model_path = get_safetensors_folder(model_path)
+        if self.sparsity_compressor is not None:
+            dense_gen = self.sparsity_compressor.decompress(model_path)
+            self._replace_weights(dense_gen, model)
+            setattr(model, SPARSITY_CONFIG_NAME, self.sparsity_compressor.config)
+
+        if self.quantization_compressor is not None:
+            apply_quantization_config(model, self.quantization_config)
+            load_pretrained_quantization(model, model_path)
+            dense_gen = self.quantization_compressor.decompress(model_path)
+            self._replace_weights(dense_gen, model)
+
+            def update_status(module):
+                module.quantization_status = QuantizationStatus.FROZEN
+
+            model.apply(update_status)
+            setattr(model, QUANTIZATION_CONFIG_NAME, self.quantization_config)
+
+    def update_config(self, save_directory: str):
+        """
+        Update the model config located at save_directory with compression configs
+        for sparsity and/or quantization
+
+        :param save_directory: path to a folder containing a HF model config
+        """
+        config_file_path = os.path.join(save_directory, CONFIG_NAME)
+        if not os.path.exists(config_file_path):
+            _LOGGER.warning(
+                f"Could not find a valid model config file in "
+                f"{save_directory}. Compression config will not be saved."
+            )
+            return
+
+        with open(config_file_path, "r") as config_file:
+            config_data = json.load(config_file)
+
+        config_data[COMPRESSION_CONFIG_NAME] = {}
+        if self.quantization_config is not None:
+            quant_config_data = self.quantization_config.model_dump()
+            config_data[COMPRESSION_CONFIG_NAME] = quant_config_data
+        if self.sparsity_config is not None:
+            sparsity_config_data = self.sparsity_config.model_dump()
+            config_data[COMPRESSION_CONFIG_NAME][
+                SPARSITY_CONFIG_NAME
+            ] = sparsity_config_data
+
+        with open(config_file_path, "w") as config_file:
+            json.dump(config_data, config_file, indent=2, sort_keys=True)
+
+    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
+
+
+def map_modules_to_quant_args(model: Module) -> Dict:
+    quantized_modules_to_args = {}
+    for name, submodule in iter_named_leaf_modules(model):
+        if is_module_quantized(submodule):
+            if submodule.quantization_scheme.weights is not None:
+                name = fix_fsdp_module_name(name)
+                quantized_modules_to_args[name] = submodule.quantization_scheme.weights
+
+    return quantized_modules_to_args

compressed_tensors/compressors/pack_quantized.py

@@ -0,0 +1,212 @@
+# 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
+import math
+from typing import Dict, Generator, Tuple
+
+import numpy as np
+import torch
+from compressed_tensors.compressors import Compressor
+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 compressed_tensors.utils import get_nested_weight_mappings, merge_names
+from safetensors import safe_open
+from torch import Tensor
+from tqdm import tqdm
+
+
+__all__ = ["PackedQuantizationCompressor", "pack_4bit_ints", "unpack_4bit_ints"]
+
+_LOGGER: logging.Logger = logging.getLogger(__name__)
+
+
+@Compressor.register(name=CompressionFormat.pack_quantized.value)
+class PackedQuantizationCompressor(Compressor):
+    """
+    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_shape",
+    ]
+
+    def compress(
+        self,
+        model_state: Dict[str, Tensor],
+        model_quant_args: Dict[str, QuantizationArgs],
+        **kwargs,
+    ) -> Dict[str, Tensor]:
+        """
+        Compresses a dense state dict
+
+        :param model_state: state dict of uncompressed model
+        :param model_quant_args: quantization args for each quantized weight, needed for
+        quantize function to calculate bit depth
+        :return: compressed state dict
+        """
+        compressed_dict = {}
+        weight_suffix = ".weight"
+        _LOGGER.debug(
+            f"Compressing model with {len(model_state)} parameterized layers..."
+        )
+
+        for name, value in tqdm(model_state.items(), desc="Compressing model"):
+            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)
+                shape = torch.tensor(value.shape)
+                if scale is not None and zp is not None:
+                    # weight is quantized, compress it
+                    quant_args = model_quant_args[prefix]
+                    if can_quantize(value, quant_args):
+                        # convert weight to an int if not already compressed
+                        value = quantize(
+                            x=value,
+                            scale=scale,
+                            zero_point=zp,
+                            args=quant_args,
+                            dtype=torch.int8,
+                        )
+                    value = pack_4bit_ints(value.cpu())
+                    compressed_dict[merge_names(prefix, "weight_shape")] = shape
+                    compressed_dict[merge_names(prefix, "weight_packed")] = value
+                    continue
+
+            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")
+
+        return compressed_dict
+
+    def decompress(
+        self, path_to_model_or_tensors: str, 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 device: optional device to load intermediate weights into
+        :return: compressed state dict
+        """
+        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)
+
+            if "weight_scale" in weight_data:
+                zero_point = weight_data.get("weight_zero_point", None)
+                scale = weight_data["weight_scale"]
+                if zero_point is None:
+                    # zero_point assumed to be 0 if not included in state_dict
+                    zero_point = torch.zeros_like(scale)
+
+                weight = weight_data["weight_packed"]
+                original_shape = torch.Size(weight_data["weight_shape"])
+                unpacked = unpack_4bit_ints(weight, original_shape)
+                decompressed = dequantize(
+                    x_q=unpacked,
+                    scale=scale,
+                    zero_point=zero_point,
+                )
+                yield merge_names(weight_name, "weight"), decompressed
+
+
+def pack_4bit_ints(value: torch.Tensor) -> torch.Tensor:
+    """
+    Packs a tensor of int4 weights stored in int8 into int32s with padding
+
+    :param value: tensor to pack
+    :returns: packed int32 tensor
+    """
+    if value.dtype is not torch.int8:
+        raise ValueError("Tensor must be quantized to torch.int8 before packing")
+
+    # need to convert to unsigned 8bit to use numpy's pack/unpack
+    temp = (value - 8).to(torch.uint8)
+    bits = np.unpackbits(temp.numpy(), axis=-1, bitorder="little")
+    ranges = np.array([range(x, x + 4) for x in range(0, bits.shape[1], 8)]).flatten()
+    only_4_bits = bits[:, ranges]  # top 4 bits are 0 because we're really uint4
+
+    # pad each row to fill a full 32bit int
+    pack_depth = 32
+    padding = (
+        math.ceil(only_4_bits.shape[1] / pack_depth) * pack_depth - only_4_bits.shape[1]
+    )
+    padded_bits = np.pad(
+        only_4_bits, pad_width=[(0, 0), (0, padding)], constant_values=0
+    )
+
+    # after packbits each uint8 is two packed uint4s
+    # then we keep the bit pattern the same but convert to int32
+    compressed = np.packbits(padded_bits, axis=-1, bitorder="little")
+    compressed = np.ascontiguousarray(compressed).view(np.int32)
+
+    return torch.from_numpy(compressed)
+
+
+def unpack_4bit_ints(value: torch.Tensor, shape: torch.Size) -> torch.Tensor:
+    """
+    Unpacks a tensor packed int4 weights into individual int8s, maintaining the
+    original their int4 range
+
+    :param value: tensor to upack
+    :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."
+        )
+
+    # unpack bits and undo padding to nearest int32 bits
+    individual_depth = 4
+    as_uint8 = value.numpy().view(np.uint8)
+    bits = np.unpackbits(as_uint8, axis=-1, bitorder="little")
+    original_row_size = int(shape[1] * individual_depth)
+    bits = bits[:, :original_row_size]
+
+    # reformat each packed uint4 to a uint8 by filling to top 4 bits with zeros
+    # (uint8 format is required by np.packbits)
+    shape_8bit = (bits.shape[0], bits.shape[1] * 2)
+    bits_as_8bit = np.zeros(shape_8bit, dtype=np.uint8)
+    ranges = np.array([range(x, x + 4) for x in range(0, shape_8bit[1], 8)]).flatten()
+    bits_as_8bit[:, ranges] = bits
+
+    # repack the bits to uint8
+    repacked = np.packbits(bits_as_8bit, axis=-1, bitorder="little")
+
+    # bits are packed in unsigned format, reformat to signed
+    # update the value range from uint4 to int4
+    final = repacked.astype(np.int8) - 8
+
+    return torch.from_numpy(final)

compressed_tensors/compressors/sparse_bitmask.py

@@ -17,7 +17,7 @@ from typing import Dict, Generator, List, Tuple, Union
 
 import numpy
 import torch
-from compressed_tensors.compressors import ModelCompressor
+from compressed_tensors.compressors import Compressor
 from compressed_tensors.config import CompressionFormat
 from compressed_tensors.utils import get_nested_weight_mappings, merge_names
 from safetensors import safe_open
@@ -37,8 +37,8 @@ __all__ = [
 _LOGGER: logging.Logger = logging.getLogger(__name__)
 
 
-@ModelCompressor.register(name=CompressionFormat.sparse_bitmask.value)
-class BitmaskCompressor(ModelCompressor):
+@Compressor.register(name=CompressionFormat.sparse_bitmask.value)
+class BitmaskCompressor(Compressor):
     """
     Compression for sparse models using bitmasks. Non-zero weights are stored in a 1d
     values tensor, with their locations stored in a 2d bitmask

compressed_tensors/compressors/utils/__init__.py

@@ -0,0 +1,19 @@
+# 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 .helpers import *
+from .permutations_24 import *
+from .semi_structured_conversions import *

compressed_tensors/compressors/utils/helpers.py

@@ -0,0 +1,43 @@
+# 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
+
+
+__all__ = ["tensor_follows_mask_structure"]
+
+
+def tensor_follows_mask_structure(tensor, mask: str = "2:4") -> bool:
+    """
+    :param tensor: tensor to check
+    :param mask: mask structure to check for, in the format "n:m"
+    :return: True if the tensor follows the mask structure, False otherwise.
+        Note, some weights can incidentally be zero, so we check for
+        atleast n zeros in each chunk of size m
+    """
+
+    n, m = tuple(map(int, mask.split(":")))
+    # Reshape the tensor into chunks of size m
+    tensor = tensor.view(-1, m)
+
+    # Count the number of zeros in each chunk
+    zero_counts = (tensor == 0).sum(dim=1)
+
+    # Check if the number of zeros in each chunk atleast n
+    # Greater than sign is needed as some weights can incidentally
+    # be zero
+    if not torch.all(zero_counts >= n).item():
+        raise ValueError()
+
+    return True

compressed_tensors/compressors/utils/permutations_24.py

@@ -0,0 +1,65 @@
+# 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 numpy
+import torch
+
+
+__all__ = ["get_permutations_24"]
+
+
+# Precompute permutations for Marlin24 weight and scale shuffling
+# Originally implemented in nm-vllm/vllm/model_executor/layers/quantization/utils/marlin_24_perms.py # noqa: E501
+#
+# Marlin works on [16*2,64] tiles. The goal of the permutations is to reorder the weight
+# data so that it is compatible with the tensor-core format that is described here:
+# https://docs.nvidia.com/cuda/parallel-thread-execution/index.html#matrix-fragments-for-mma-m16n8k16-with-floating-point-type # noqa: E501
+#
+# As a result of this reordering, the vector loads inside the kernel will get the data
+# as it is needed for tensor-core (without the need to use ldmatrix instructions)
+def get_permutations_24(num_bits):
+    perm_list = []
+    for i in range(32):
+        perm1 = []
+        col = i // 4
+        col_o = col // 2
+        for block in [0, 1]:
+            for row in [
+                2 * (i % 4),
+                2 * (i % 4) + 1,
+                2 * (i % 4 + 4),
+                2 * (i % 4 + 4) + 1,
+            ]:
+                perm1.append(16 * row + col_o * 256 + 8 * (col % 2) + 4 * block)
+        for j in range(4):
+            perm_list.extend([p + 1 * j for p in perm1])
+    perm = numpy.array(perm_list)
+
+    if num_bits == 4:
+        interleave = numpy.array([0, 2, 4, 6, 1, 3, 5, 7])
+    elif num_bits == 8:
+        interleave = numpy.array([0, 2, 1, 3])
+    else:
+        raise ValueError("num_bits must be 4 or 8, got {}".format(num_bits))
+
+    perm = perm.reshape((-1, len(interleave)))[:, interleave].ravel()
+    perm = torch.from_numpy(perm)
+    scale_perm = []
+    for i in range(8):
+        scale_perm.extend([i * 8 + j for j in [0, 4, 1, 5, 2, 6, 3, 7]])
+    scale_perm_single = []
+    for i in range(8):
+        scale_perm_single.extend([8 * i + j for j in [0, 1, 2, 3, 4, 5, 6, 7]])
+    return perm, scale_perm, scale_perm_single