ds patch for 2.7 (#3638)

* Revert "Handle new linear modules in DeepSpeed v0.16.5 (#3622) (#3631)"

This reverts commit d74bd04.

* update

* fix comment

Signed-off-by: Liu, Mingzhi <[email protected]>

---------

Signed-off-by: Liu, Mingzhi <[email protected]>
Co-authored-by: Chunyuan WU <[email protected]>

Author: inkcherry

intel_extension_for_pytorch/nn/utils/_parameter_wrapper.py

@@ -46,6 +46,15 @@ def IPEX_WEIGHT_PREPACK_MODULE_CPU():
             deepspeed_modules_mapping.update(
                 {LmHeadLinearAllreduce: _IPEXLmHeadLinearAllreduce}
             )
+        if len(deepspeed_modules) > 3:
+            for module in deepspeed_modules[3:]:
+                if module not in deepspeed_modules_mapping:
+                    if issubclass(module, LinearAllreduce):
+                        deepspeed_modules_mapping[module] = _IPEXLinearAllreduce
+                    elif issubclass(module, LinearLayer):
+                        deepspeed_modules_mapping[module] = _IPEXLinear
+                    else:
+                        raise ValueError(f"Unrecognized module type: {module}")
         torch_modules.update(deepspeed_modules_mapping)
 
     return torch_modules
@@ -190,7 +199,9 @@ def get_shared_parameter_status(module, shared_p):
     if deepspeed_modules is not None:
         LinearAllreduce, LinearLayer = deepspeed_modules[:2]
 
-        if isinstance(module, (LinearLayer, LinearAllreduce)):
+        if isinstance(module, (LinearLayer, LinearAllreduce)) or issubclass(
+            type(module), (LinearLayer, LinearAllreduce)
+        ):
             module.weight = torch.nn.Parameter(module.weight, requires_grad=False)
             if module.bias is not None:
                 module.bias = torch.nn.Parameter(module.bias, requires_grad=False)

intel_extension_for_pytorch/nn/utils/_weight_prepack.py

@@ -101,22 +101,22 @@ def may_import_deepspeed_modules():
     try:
         # import deepspeed in a global space will raise circular import error
         # intel-extension-for-deepspeed imports both IPEX and deepspeed
-        from deepspeed.module_inject.layers import (
-            LinearAllreduce,
-            LinearLayer,
-            LmHeadLinearAllreduce,
-            fused_LinearLayer,
-            GateUpPack_LinearLayer,
-        )
+        import deepspeed.module_inject.layers as dslayers
+        from deepspeed.module_inject.layers import LinearAllreduce, LinearLayer
+
+        ds_layers = [LinearAllreduce, LinearLayer]
+
+        from deepspeed.module_inject.layers import LmHeadLinearAllreduce
 
-        ds_layers = [
-            LinearAllreduce,
-            LinearLayer,
-            LmHeadLinearAllreduce,
-            fused_LinearLayer,
-            GateUpPack_LinearLayer,
-        ]
+        ds_layers.append(LmHeadLinearAllreduce)
+        ds_layers += list(
+            cls
+            for cls in dslayers.LinearAllreduce.__subclasses__()
+            if cls is not LmHeadLinearAllreduce
+        )
+        ds_layers += list(cls for cls in dslayers.LinearLayer.__subclasses__())
         return ds_layers
+
     except ImportError:
         return None
 

intel_extension_for_pytorch/quantization/_quantize.py

@@ -137,6 +137,12 @@ def _may_insert_deepspeed_modules(
                     LmHeadLinearAllreduce: q_lm_head_linear_all_reduce_module,
                 }
             )
+        if len(deepspeed_modules) > 3:
+            for module in deepspeed_modules[3:]:
+                if issubclass(module, LinearLayer):
+                    deepspeed_modules_dict[module] = q_linear_layer_module
+                elif issubclass(module, LinearAllreduce):
+                    deepspeed_modules_dict[module] = q_linear_all_reduce_module
         torch_modules.update(deepspeed_modules_dict)
     return torch_modules
 
@@ -231,6 +237,12 @@ def _float_module(cls):
         if deepspeed_modules is not None:
             LinearLayer = deepspeed_modules[1]
             _FLOAT_MODULE.extend([LinearLayer])
+
+        if len(deepspeed_modules) > 3:
+            for module in deepspeed_modules[3:]:
+                if issubclass(module, LinearLayer):
+                    _FLOAT_MODULE.extend([module])
+
         return _FLOAT_MODULE
 
     def __repr__(self):
@@ -260,6 +272,10 @@ def _float_module(cls):
         ), "DynamicQuantizedLinearAllreduce requires deepspeed to be installed"
         LinearAllreduce = deepspeed_modules[0]
         _FLOAT_MODULE = [LinearAllreduce]
+        if len(deepspeed_modules) > 3:
+            for module in deepspeed_modules[3:]:
+                if issubclass(module, LinearAllreduce):
+                    _FLOAT_MODULE.extend([module])
         return _FLOAT_MODULE
 
     def __init__(
@@ -361,6 +377,7 @@ def may_quantize_deepspeed_modules(
     IPEX_QUANTIZATION_MODULE, q_config, module_mappings, qconfig_spec
 ):
     deepspeed_modules = may_import_deepspeed_modules()
+
     if deepspeed_modules is not None:
         LinearAllreduce, LinearLayer = deepspeed_modules[:2]
         module_mappings.update(IPEX_QUANTIZATION_MODULE)
@@ -375,6 +392,13 @@ def may_quantize_deepspeed_modules(
                     LmHeadLinearAllreduce: q_config,
                 }
             )
+        if len(deepspeed_modules) > 3:
+            for module in deepspeed_modules[3:]:
+                deepspeed_qconfig_spec.update(
+                    {
+                        module: q_config,
+                    }
+                )
 
         qconfig_spec.update(deepspeed_qconfig_spec)
     return module_mappings, qconfig_spec

intel_extension_for_pytorch/utils/weight_only_quantization.py

@@ -292,22 +292,23 @@ def _convert_woq_with_low_precision_checkpoint(
 
     deepspeed_modules = may_import_deepspeed_modules()
     if deepspeed_modules is not None:
-        (
-            LinearAllreduce,
-            LinearLayer,
-            LmHeadLinearAllreduce,
-            fused_LinearLayer,
-            GateUpPack_LinearLayer,
-        ) = deepspeed_modules
+        LinearAllreduce, LinearLayer, LmHeadLinearAllreduce, *extra_linear_modules = (
+            deepspeed_modules
+        )
+
         q_op_map.update(
             {
                 LinearAllreduce: IpexWoqLinearAllreduce,
                 LinearLayer: WeightOnlyQuantizedLinear,
-                fused_LinearLayer: WeightOnlyQuantizedLinear,
-                GateUpPack_LinearLayer: WeightOnlyQuantizedLinear,
             }
         )
 
+        if extra_linear_modules:
+            for module in extra_linear_modules:
+                if issubclass(module, LinearAllreduce):
+                    q_op_map[module] = IpexWoqLinearAllreduce
+                elif issubclass(module, LinearLayer):
+                    q_op_map[module] = WeightOnlyQuantizedLinear
     linear_modules = tuple(q_op_map.keys())
 
     def _convert(mod, attr_name):