Refactor: move int4 code to deepspeed/inference (#528)

* Move int 4 code to deepspeed/inference

* fix

* fix

* fix

deepspeed/inference/__init__.py

@@ -2,5 +2,3 @@
 # SPDX-License-Identifier: Apache-2.0
 
 # DeepSpeed Team
-
-from .engine import InferenceEngine

deepspeed/inference/config.py

@@ -101,6 +101,8 @@ class BaseQuantConfig(DeepSpeedConfigModel):
 
 class WeightQuantConfig(BaseQuantConfig):
     enabled = True
+    quantized_initialization: Dict = {}
+    post_init_quant: Dict = {}
 
 
 class ActivationQuantConfig(BaseQuantConfig):

deepspeed/compression/inference/quantization.py → deepspeed/inference/quantization/quantization.py

@@ -7,10 +7,9 @@
 from torch import nn
 from typing import Dict
 import gc
-from deepspeed.compression.inference import layers
+from deepspeed.inference.quantization import layers
 from .layers import QUANTIZATION_LAYER_MAPPINGS
-from .utils import get_AsyncPartitionedParameterSwapper
-from ..helper import recursive_setattr
+from .utils import get_AsyncPartitionedParameterSwapper, recursive_setattr
 from deepspeed.utils.logging import logger
 from collections import deque
 from transformers.utils.generic import ContextManagers
@@ -35,7 +34,7 @@ def _init_group_wise_weight_quantization(model: nn.Module, ds_config: Dict) -> n
     matched_module_count = 0
 
     assert 'weight_quantization' in ds_config, 'Please provide quantization config in ds_config'
-    quantization_config = ds_config['weight_quantization']
+    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.

deepspeed/compression/inference/utils.py → deepspeed/inference/quantization/utils.py

@@ -192,6 +192,24 @@ def get_AsyncPartitionedParameterSwapper(model: nn.Module):
     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,

deepspeed/runtime/config.py

@@ -29,7 +29,7 @@
 from .activation_checkpointing.config import DeepSpeedActivationCheckpointingConfig
 from ..comm.config import DeepSpeedCommsConfig
 from ..monitor.config import get_monitor_config
-from ..compression.inference.config import WeightQuantizationConfig
+from ..inference.config import WeightQuantConfig
 
 from deepspeed import comm as dist
 from deepspeed.runtime.config_utils import DeepSpeedConfigModel
@@ -871,7 +871,8 @@ def _initialize_params(self, param_dict):
 
         self.nebula_config = DeepSpeedNebulaConfig(param_dict)
 
-        self.weight_quantization_config = WeightQuantizationConfig(param_dict)
+        self.weight_quantization_config = WeightQuantConfig(
+            **param_dict['weight_quantization']) if 'weight_quantization' in param_dict else None
 
     def _batch_assertion(self):
 

deepspeed/runtime/zero/partition_parameters.py

@@ -31,7 +31,7 @@
                                    debug_param2name_id, debug_param2name_id_shape_status)
 from deepspeed.accelerator import get_accelerator
 from ..swap_tensor.partitioned_param_swapper import AsyncPartitionedParameterSwapper, PartitionedParamStatus
-from deepspeed.compression.inference.utils import _quantize_param, WEIGHT_QUANTIZATION_LAYERS, wrap_quantized_functional, wrap_load_from_state_dict
+from deepspeed.inference.quantization.utils import _quantize_param, WEIGHT_QUANTIZATION_LAYERS, wrap_quantized_functional, wrap_load_from_state_dict
 
 param_count = 0
 partitioned_param_data_shape = [0]
@@ -295,7 +295,7 @@ def __init__(self, enabled=True, mem_efficient_linear=True, ds_config=None, dtyp
         self.wrapped_cls = set()
 
         self.quantized_initialization = None
-        if ds_config is not None and ds_config.weight_quantization_config.quantized_initialization:
+        if ds_config is not None and ds_config.weight_quantization_config and ds_config.weight_quantization_config.quantized_initialization:
             self.quantized_initialization = ds_config.weight_quantization_config.quantized_initialization
 
     def __enter__(self):

tests/unit/compression/inference/test_int4_quantization.py → tests/unit/inference/quantization/test_int4_quantization.py

@@ -7,9 +7,9 @@
 import torch
 import torch.nn as nn
 from deepspeed.accelerator import get_accelerator
-from deepspeed.compression.inference.quantization import _init_group_wise_weight_quantization
-from deepspeed.compression.inference.utils import Quantizer, DeQuantizer
-from deepspeed.compression.inference.layers import QuantizedLinear
+from deepspeed.inference.quantization.quantization import _init_group_wise_weight_quantization
+from deepspeed.inference.quantization.utils import Quantizer, DeQuantizer
+from deepspeed.inference.quantization.layers import QuantizedLinear
 from transformers.models.opt.modeling_opt import OPTDecoderLayer
 from transformers import AutoConfig, OPTConfig, AutoModel
 import pytest
@@ -258,35 +258,37 @@ def test_model_quantization():
 
         ds_config = {
             'weight_quantization': {
-                'fc': {
-                    'num_bits': bits,
-                    'group_size': 64,
-                    'group_dim': 0,
-                    'symmetric': False
-                },
-                'self_attn.q_proj': {
-                    'num_bits': bits,
-                    'group_size': 64,
-                    'group_dim': 0,
-                    'symmetric': False
-                },
-                'self_attn.k_proj': {
-                    'num_bits': bits,
-                    'group_size': 64,
-                    'group_dim': 0,
-                    'symmetric': False
-                },
-                'self_attn.v_proj': {
-                    'num_bits': bits,
-                    'group_size': 64,
-                    'group_dim': 0,
-                    'symmetric': False
-                },
-                'self_attn.out_proj': {
-                    'num_bits': bits,
-                    'group_size': 64,
-                    'group_dim': 0,
-                    'symmetric': False
+                'post_init_quant': {
+                    'fc': {
+                        'num_bits': bits,
+                        'group_size': 64,
+                        'group_dim': 0,
+                        'symmetric': False
+                    },
+                    'self_attn.q_proj': {
+                        'num_bits': bits,
+                        'group_size': 64,
+                        'group_dim': 0,
+                        'symmetric': False
+                    },
+                    'self_attn.k_proj': {
+                        'num_bits': bits,
+                        'group_size': 64,
+                        'group_dim': 0,
+                        'symmetric': False
+                    },
+                    'self_attn.v_proj': {
+                        'num_bits': bits,
+                        'group_size': 64,
+                        'group_dim': 0,
+                        'symmetric': False
+                    },
+                    'self_attn.out_proj': {
+                        'num_bits': bits,
+                        'group_size': 64,
+                        'group_dim': 0,
+                        'symmetric': False
+                    }
                 }
             }
         }
@@ -321,11 +323,13 @@ def test_quantized_linear():
 
         ds_config = {
             'weight_quantization': {
-                'layer': {
-                    'num_bits': 4,
-                    'group_size': 64,
-                    'group_dim': 0,
-                    'symmetric': False
+                'post_init_quant': {
+                    'layer': {
+                        'num_bits': 4,
+                        'group_size': 64,
+                        'group_dim': 0,
+                        'symmetric': False
+                    }
                 }
             }
         }