chore: import upstream snapshot with attribution

deepspeed/inference/quantization/__init__.py

@@ -0,0 +1,4 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team

deepspeed/inference/quantization/layers.py

@@ -0,0 +1,114 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+import torch
+
+from torch import nn
+from torch import Tensor
+from torch.nn import functional as F
+from .utils import Quantizer, DeQuantizer, concat_to_compat_param
+from typing import Tuple, Callable, Dict
+from deepspeed.runtime.zero import register_external_parameter
+
+quantized_weight_registry = {}
+is_zero3_enabled = False
+
+
+# deal with weight sharing
+def get_quantized_weight_wrapper(model, pre_quant_weight: nn.Parameter, quantize_weight_fn: Callable) -> nn.Parameter:
+    if id(pre_quant_weight) in quantized_weight_registry:
+        compat_tensor = quantized_weight_registry[id(pre_quant_weight)]
+        if is_zero3_enabled:
+            register_external_parameter(model, compat_tensor)
+
+        return quantized_weight_registry[id(pre_quant_weight)]
+    else:
+        quantized_weights, quant_scale, quant_min = quantize_weight_fn()
+        quantized_weight_registry[id(pre_quant_weight)] = concat_to_compat_param(quantized_weights, quant_scale,
+                                                                                 quant_min)
+        return quantized_weight_registry[id(pre_quant_weight)]
+
+
+def get_quantize_weight_fn(quantizer: Quantizer, pre_quant_weight: nn.Parameter) -> Callable:
+
+    def func() -> Tuple[nn.Parameter, Tensor, Tensor]:
+        quantized_weights, quant_scale, quant_min = quantizer.quantize(pre_quant_weight.data)
+        # A temporary hack as zero Zero3 assume all model weights has the same type. in all_gather_coalesced.get_only_unique_item
+        quantized_weights = quantized_weights.view(pre_quant_weight.dtype)
+        quant_scale = quant_scale.type(pre_quant_weight.dtype)
+        quant_min = quant_min.type(pre_quant_weight.dtype)
+        return quantized_weights, quant_scale, quant_min
+
+    return func
+
+
+class QuantizedLinear(nn.Linear):
+
+    def __init__(self, config: Dict, pre_quant_layer: nn.Linear) -> None:
+        super(QuantizedLinear, self).__init__(in_features=pre_quant_layer.in_features,
+                                              out_features=pre_quant_layer.out_features,
+                                              bias=pre_quant_layer.bias is not None,
+                                              device=pre_quant_layer.weight.device,
+                                              dtype=pre_quant_layer.weight.dtype)
+        self.config = config
+
+        self.quantizer = Quantizer(config=config)
+        self.bias = pre_quant_layer.bias
+        self.weight = get_quantized_weight_wrapper(self, pre_quant_layer.weight,
+                                                   get_quantize_weight_fn(self.quantizer, pre_quant_layer.weight))
+
+        self.weight.dequantizer = DeQuantizer(config, pre_quant_layer.weight.dtype)
+
+    def forward(self, input: Tensor) -> Tensor:
+        quantized_weight, quant_scale, quant_min = self.weight.deconcat(self.weight)
+        temp_dequantized_weight = self.weight.dequantizer.dequantize(quantized_weight.view(torch.uint8), quant_scale,
+                                                                     quant_min)
+
+        # !!! Do not use torch.functional.linear(input, temp_dequantized_weight, self.bias) here as in zero3 torch.functional.linear is
+        # replaced by LinearFunctionForZeroStage3. Which assume weight is non-temporary.
+        # If weight is temp buffer there will be memory leak.
+        return torch._C._nn.linear(input, temp_dequantized_weight, self.bias)
+
+
+class QuantizedEmbedding(nn.Embedding):
+
+    def __init__(self, config: Dict, pre_quant_layer: nn.Embedding) -> None:
+        super(QuantizedEmbedding, self).__init__(num_embeddings=pre_quant_layer.num_embeddings,
+                                                 embedding_dim=pre_quant_layer.embedding_dim,
+                                                 padding_idx=pre_quant_layer.padding_idx,
+                                                 max_norm=pre_quant_layer.max_norm,
+                                                 norm_type=pre_quant_layer.norm_type,
+                                                 scale_grad_by_freq=pre_quant_layer.scale_grad_by_freq,
+                                                 sparse=pre_quant_layer.sparse,
+                                                 _weight=pre_quant_layer.weight,
+                                                 device=pre_quant_layer.weight.device,
+                                                 dtype=pre_quant_layer.weight.dtype)
+
+        assert pre_quant_layer.max_norm is None, 'Not supported'
+        assert pre_quant_layer.norm_type == 2, 'Not supported'
+        assert pre_quant_layer.scale_grad_by_freq == False, 'Not supported'
+        assert pre_quant_layer.sparse == False, 'Not supported'
+
+        self.config = config
+        quantizer = Quantizer(config=config)
+
+        self.weight = get_quantized_weight_wrapper(self, pre_quant_layer.weight,
+                                                   get_quantize_weight_fn(quantizer, pre_quant_layer.weight))
+
+        self.weight.dequantizer = DeQuantizer(config, pre_quant_layer.weight.dtype)
+
+    def forward(self, input: Tensor) -> Tensor:
+        quantized_weight, quant_scale, quant_min = self.weight.deconcat(self.weight)
+        temp_dequantized_weight = self.weight.dequantizer.dequantize(quantized_weight.view(torch.uint8), quant_scale,
+                                                                     quant_min)
+
+        return F.embedding(input, temp_dequantized_weight, self.padding_idx, self.max_norm, self.norm_type,
+                           self.scale_grad_by_freq, self.sparse)
+
+
+QUANTIZATION_LAYER_MAPPINGS = {
+    nn.Linear: QuantizedLinear,
+    nn.Embedding: QuantizedEmbedding,
+}

deepspeed/inference/quantization/quantization.py

@@ -0,0 +1,111 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+import torch
+from torch import nn
+from typing import Dict
+import gc
+from deepspeed.inference.quantization import layers
+from .layers import QUANTIZATION_LAYER_MAPPINGS
+from .utils import get_AsyncPartitionedParameterSwapper, recursive_setattr
+from deepspeed.utils.logging import logger
+from collections import deque
+from transformers.utils.generic import ContextManagers
+from .quantization_context import QuantizationContext
+import contextlib
+
+
+def _init_group_wise_weight_quantization(model: nn.Module, ds_config: Dict) -> nn.Module:
+    """[Experimental] Apply group-wise weight quantization to model. Replace layers module according to config_list
+
+    Args:
+        model (nn.Module): A nn.Module
+        ds_config (Dict, optional): The ds_config dictionary. use None for non-deepspeed managed model.
+
+    Returns:
+        nn.Module: Quantized nn.Module
+    """
+
+    # global quantized_weight_registry
+
+    matched_module_list_by_key = {}
+    matched_module_count = 0
+
+    assert 'weight_quantization' in ds_config, 'Please provide quantization config in ds_config'
+    quantization_config = ds_config['weight_quantization']['post_init_quant']
+
+    # Return nvme swapper if exists, else return None.
+    # For nvme offloading we must use the same swapper here as model initialized.
+    nvme_swapper = get_AsyncPartitionedParameterSwapper(model)
+    is_zero3_enabled = 'zero_optimization' in ds_config and \
+            'stage' in ds_config['zero_optimization'] and \
+            ds_config['zero_optimization']['stage'] == 3
+    is_offloading_enabled = 'zero_optimization' in ds_config and \
+                            'offload_param' in ds_config['zero_optimization']
+
+    layers.is_zero3_enabled = is_zero3_enabled
+
+    context_mgr = ContextManagers([QuantizationContext(config_dict_or_path=ds_config, param_swapper=nvme_swapper)]) \
+                    if is_zero3_enabled else contextlib.suppress()
+    with context_mgr:
+        module_list = list(
+            filter(lambda named_module: type(named_module[1]) in QUANTIZATION_LAYER_MAPPINGS, model.named_modules()))
+
+        # Quantize small weight first then large.
+        if not is_offloading_enabled:
+            module_list.sort(key=lambda named_module: named_module[1].weight.ds_tensor.numel()
+                             if is_zero3_enabled else named_module[1].weight.numel())
+        module_list = deque(module_list)
+
+        while len(module_list) > 0:
+            # Use popleft to timely release module's memory of replaced module after each loop iteration
+            module_name, module = module_list.popleft()
+
+            matched_key = None
+            matched_quantization_config = None
+
+            for key, config in quantization_config.items():
+                if key in module_name:
+                    assert matched_key is None, f'{module_name} matched multiple quantization key word {matched_key} and {key}'
+                    matched_key = key
+                    matched_quantization_config = config
+
+            if matched_key is None:
+                continue
+
+            if is_zero3_enabled:
+                module.weight.all_gather()
+
+            assert module.weight.dtype == torch.float16, 'Model weight is expected in half.'
+
+            new_module = QUANTIZATION_LAYER_MAPPINGS[type(module)](matched_quantization_config, module)
+
+            if is_zero3_enabled:
+                module.weight.partition()
+
+            recursive_setattr(model, module_name, new_module)
+
+            if matched_key not in matched_module_list_by_key:
+                matched_module_list_by_key[matched_key] = []
+            matched_module_list_by_key[matched_key].append(module_name)
+            matched_module_count += 1
+
+            # Timely recycle memory to prevent OOM on large models
+            gc.collect()
+
+    # Clear registry after model construction.
+    layers.quantized_weight_registry.clear()
+
+    logger.info(
+        f'Group-wise weight quantization summary: convert {matched_module_count} node(s) to quantized implementation')
+    summary_str = '\n'
+
+    for key, module_list in matched_module_list_by_key.items():
+        summary_str += f'Key: {key}, matched modules:\n'
+        for module_name in module_list:
+            summary_str += f'\t{module_name}\n'
+    logger.info(summary_str)
+
+    return model

deepspeed/inference/quantization/quantization_context.py

@@ -0,0 +1,13 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+from deepspeed.runtime.zero import partition_parameters
+from deepspeed.runtime.swap_tensor.partitioned_param_swapper import AsyncPartitionedParameterSwapper
+
+
+class QuantizationContext(partition_parameters.Init):
+
+    def __init__(self, config_dict_or_path, param_swapper: AsyncPartitionedParameterSwapper = None) -> None:
+        super().__init__(config_dict_or_path=config_dict_or_path, param_swapper=param_swapper)

deepspeed/inference/quantization/utils.py

@@ -0,0 +1,288 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+import torch
+import deepspeed
+from torch import Tensor
+from typing import Tuple
+import torch.nn as nn
+from typing import Dict, Callable, Union
+from deepspeed.accelerator import get_accelerator
+import functools
+
+device = get_accelerator().device_name() if get_accelerator().is_available() else 'cpu'
+
+quantizer_module = None
+
+
+def get_quantizer_module():
+    global quantizer_module
+    if quantizer_module is None:
+        quantizer_module = deepspeed.ops.op_builder.QuantizerBuilder().load()
+    return quantizer_module
+
+
+def tensor_clamp(tensor: Tensor, min, max) -> Tensor:
+    if tensor.device.type == 'cpu' and tensor.dtype == torch.float16:
+        # CPU does not support FP16 clamp
+        return tensor.to(dtype=torch.float32).clamp_(min, max).to(dtype=torch.float16)
+    else:
+        return tensor.clamp_(min, max)
+
+
+def tensor_round(tensor: Tensor) -> Tensor:
+    if tensor.device.type == 'cpu' and tensor.dtype == torch.float16:
+        # CPU does not support FP16 round
+        return tensor.to(dtype=torch.float32).round_().to(dtype=torch.float16)
+    else:
+        return tensor.round_()
+
+
+class Quantizer:
+
+    def __init__(self, config: Dict) -> None:
+        self.config = config
+        assert self.config['num_bits'] == 4 or self.config[
+            'num_bits'] == 8, 'Only INT4 and INT8 quantization is supported.'
+        assert self.config['symmetric'] == False, 'Only asymmetric quantization is supported at this moment.'
+
+    def quantize(self, tensor: Tensor) -> Tuple[Tensor, Tensor, Tensor]:
+        assert tensor.shape[self.config['group_dim']] % self.config['group_size'] == 0 \
+            , f'Tensor shape: {tensor.shape} quantization config {self.config}'
+
+        tensor = torch.clone(tensor)
+
+        shape = tensor.shape
+        num_groups = shape[self.config['group_dim']] // self.config['group_size']
+        new_shape = (shape[:self.config['group_dim']] + (num_groups, self.config['group_size']) +
+                     shape[self.config['group_dim'] + 1:])
+        tensor = tensor.view(new_shape)
+
+        quantized_tensor, scale, min_value = self._quantize_int8(tensor)
+        quantized_tensor = quantized_tensor.view(shape)
+
+        if self.config['num_bits'] == 4:
+            return self._compress_uint8_to_uint4(quantized_tensor), scale, min_value
+        if self.config['num_bits'] == 8:
+            return quantized_tensor, scale, min_value
+
+        assert False, 'Unsupported quantization bits {}'.format(self.config['num_bits'])
+
+    def _quantize_int8(self, tensor: Tensor) -> Tuple[Tensor, Tensor, Tensor]:
+        q_range = 2**self.config['num_bits'] - 1
+        min_value = tensor.amin(dim=self.config['group_dim'] + 1, keepdim=True)
+        max_value = tensor.amax(dim=self.config['group_dim'] + 1, keepdim=True)
+
+        scale = q_range / (max_value - min_value)
+
+        tensor = tensor.sub_(min_value).mul_(scale)
+        tensor = tensor_round(tensor_clamp(tensor, 0, q_range)).to(torch.uint8)
+        return tensor, scale, min_value
+
+    def _compress_uint8_to_uint4(self, tensor: Tensor) -> Tensor:
+        assert tensor.shape[-1] % 2 == 0
+
+        new_data_shape = list(tensor.shape)
+        new_data_shape[-1] = new_data_shape[-1] // 2
+
+        data = torch.empty(new_data_shape, dtype=torch.uint8, device=tensor.device)
+        data = torch.bitwise_or(tensor[..., 0::2].bitwise_left_shift(4), tensor[..., 1::2])
+
+        return data
+
+
+class DeQuantizer:
+
+    def __init__(self, config: Dict, dtype: torch.dtype) -> None:
+        self.config = config
+        self.dtype = dtype
+        assert self.config['num_bits'] == 4 or self.config[
+            'num_bits'] == 8, 'Only INT4 and INT8 quantization is supported.'
+        assert self.config['symmetric'] == False, 'Only asymmetric quantization is supported at this moment.'
+
+    def dequantize(self, tensor: Tensor, quant_scale: Tensor, quant_min: Tensor) -> Tensor:
+        # Use customized CUDA quantization kernel if possible.
+        if self.config['group_size'] % 8 == 0 and \
+                (self.config['num_bits'] == 4 or self.config['num_bits'] == 8) and \
+                self.config['group_dim'] == len(tensor.shape) - 1 and \
+                    self.dtype == torch.float16 and device == get_accelerator().device_name():
+
+            last_dimension_size = self.config['group_size']
+            if self.config['num_bits'] == 4:
+                last_dimension_size = last_dimension_size // 2
+                quantized_tensor = get_quantizer_module().dequantize_int4_to_half_experimental(
+                    tensor.reshape(-1, last_dimension_size), quant_scale, quant_min,
+                    tensor.numel() // last_dimension_size, self.config['group_size'])
+                shape = list(tensor.shape)
+                shape[-1] = shape[-1] * 2
+            elif self.config['num_bits'] == 8:
+                # last_dimension_size = last_dimension_size // 2
+                quantized_tensor = get_quantizer_module().dequantize_int8_to_half_experimental(
+                    tensor.reshape(-1, last_dimension_size), quant_scale, quant_min,
+                    tensor.numel() // last_dimension_size, self.config['group_size'])
+                shape = list(tensor.shape)
+
+            return quantized_tensor.reshape(shape)
+
+        if self.config['num_bits'] == 4:
+            tensor = self._decompress_uint4_to_uint8(tensor)
+        elif self.config['num_bits'] != 8:
+            assert False, 'Unsupported quantization bits {}'.format(self.config['num_bits'])
+
+        shape = tensor.shape
+        num_groups = shape[self.config['group_dim']] // self.config['group_size']
+        new_shape = (shape[:self.config['group_dim']] + (num_groups, self.config['group_size']) +
+                     shape[self.config['group_dim'] + 1:])
+        tensor = tensor.view(new_shape)
+
+        dequantized_tensor = self._dequantize_int8(tensor, quant_scale, quant_min).view(shape)
+        return dequantized_tensor
+
+    def _dequantize_int8(self, tensor: Tensor, quant_scale: Tensor, quant_min: Tensor) -> Tensor:
+        assert tensor.dtype == torch.uint8
+        data = torch.zeros_like(tensor, dtype=self.dtype, device=tensor.device)
+        data = data.copy_(tensor)
+        data = data.div_(quant_scale).add_(quant_min)
+
+        return data
+
+    def _decompress_uint4_to_uint8(self, tensor: Tensor) -> Tensor:
+        new_data_shape = list(tensor.shape)
+        new_data_shape[-1] = new_data_shape[-1] * 2
+        data = torch.empty(new_data_shape, dtype=torch.uint8, device=tensor.device)
+        data[..., 0::2] = tensor.bitwise_right_shift(4)
+        data[..., 1::2] = tensor.bitwise_and(0xF)
+
+        return data
+
+
+def get_AsyncPartitionedParameterSwapper(model: nn.Module):
+    for param_name, param in model.named_parameters():
+        if hasattr(param, 'nvme_swapper') and param.nvme_swapper is not None:
+            return param.nvme_swapper
+    return None
+
+
+def recursive_setattr(model, module_name, module):
+    """
+    Recursively set the attribute of a module.
+    Args:
+        model (`torch.nn.Module`)
+            The model to set the attribute in.
+        module_name (`str`)
+            The name of the module to set the attribute in.
+        module (`torch.nn.Module`)
+            The module to set the attribute to.
+    """
+    split_list = module_name.split('.')
+    output = model
+    for name in split_list[:-1]:
+        output = getattr(output, name)
+    output.__setattr__(split_list[-1], module)
+
+
+def concat_to_compat_param(quantized_weight: Tensor,
+                           quant_scale: Tensor,
+                           quant_min: Tensor,
+                           return_param: bool = True) -> Union[nn.Parameter, Tensor]:
+    shape_wieght = quantized_weight.shape
+    shape_scale = quant_scale.shape
+    shape_min = quant_min.shape
+
+    quantized_weight = torch.flatten(quantized_weight)
+    quant_scale = torch.flatten(quant_scale)
+    quant_min = torch.flatten(quant_min)
+
+    def deconcat_individual_tensors(shape_wieght: torch.Size, shape_scale: torch.Size,
+                                    shape_min: torch.Size) -> Callable:
+
+        def fn(compat_tensor: nn.Parameter) -> Tuple[Tensor, Tensor, Tensor]:
+            weight = torch.narrow(compat_tensor, 0, 0, shape_wieght.numel()).view(shape_wieght)
+            scale = torch.narrow(compat_tensor, 0, shape_wieght.numel(), shape_scale.numel()).view(shape_scale)
+            min_val = torch.narrow(compat_tensor, 0,
+                                   shape_wieght.numel() + shape_scale.numel(), shape_min.numel()).view(shape_min)
+
+            return weight, scale, min_val
+
+        return fn
+
+    compat_tensor = torch.concat([quantized_weight, quant_scale, quant_min])
+    if return_param:
+        compat_tensor = nn.Parameter(compat_tensor, requires_grad=False)
+    compat_tensor.deconcat = deconcat_individual_tensors(shape_wieght, shape_scale, shape_min)
+
+    return compat_tensor
+
+
+def _quantize_param(param: nn.Parameter, quant_config: Dict):
+    assert not hasattr(param, 'weight_quantized'), 'Parameter has already been quantized.'
+    quantizer = Quantizer(quant_config)
+    dequantizer = DeQuantizer(quant_config, param.dtype)
+
+    quantized_weight, quant_scale, quant_min = quantizer.quantize(param.data)
+
+    quantized_weight = quantized_weight.view(param.dtype)
+    quant_scale = quant_scale.view(param.dtype)
+    quant_min = quant_min.view(param.dtype)
+
+    quantized_compat_tensor = concat_to_compat_param(quantized_weight, quant_scale, quant_min)
+    param.data = quantized_compat_tensor
+    param.deconcat = quantized_compat_tensor.deconcat
+
+    param.quantizer = quantizer
+    param.dequantizer = dequantizer
+    setattr(param, 'weight_quantized', True)
+
+
+def wrap_quantized_functional(f):
+
+    @functools.wraps(f)
+    def wrapper(input: Tensor, weight: nn.Parameter, *args, **kwargs) -> Tensor:
+        if hasattr(weight, 'weight_quantized') and getattr(weight, 'weight_quantized'):
+            quantized_weight, quant_scale, quant_min = weight.deconcat(weight)
+            temp_dequantized_weight = weight.dequantizer.dequantize(quantized_weight.view(torch.uint8), quant_scale,
+                                                                    quant_min)
+            return f(input, temp_dequantized_weight, *args, **kwargs)
+        else:
+            return f(input, weight, *args, **kwargs)
+
+    return wrapper
+
+
+def wrap_load_from_state_dict(f):
+
+    @functools.wraps(f)
+    def wrapper(model, state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs):
+        replaced_old_value = None
+        key = None
+        # We may have nested wrappers if we launch multiple initialization context.
+        # Use state_dict_quantized flag to quantize state_dict only once
+        if hasattr(model.weight, 'weight_quantized') and getattr(
+                model.weight, 'weight_quantized') and not hasattr(model.weight, 'state_dict_quantized'):
+            setattr(model.weight, 'state_dict_quantized', True)
+            key = prefix + 'weight'
+            if key in state_dict:
+                quantized_weight, quant_scale, quant_min = model.weight.quantizer.quantize(state_dict[key])
+                quantized_weight = quantized_weight.view(model.weight.dtype)
+                quant_scale = quant_scale.view(model.weight.dtype)
+                quant_min = quant_min.view(model.weight.dtype)
+
+                replaced_old_value = state_dict[key]
+
+                state_dict[key] = concat_to_compat_param(quantized_weight, quant_scale, quant_min)
+
+        f(model, state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs)
+
+        if replaced_old_value is not None:
+            state_dict[key] = replaced_old_value
+            delattr(model.weight, 'state_dict_quantized')
+
+    return wrapper
+
+
+WEIGHT_QUANTIZATION_LAYERS = (
+    nn.Linear,
+    nn.Embedding,
+)