chore: import upstream snapshot with attribution

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