Handle new linear modules in DeepSpeed v0.16.5 (#3622)

* Handle fused_LinearLayer in DeepSpeed v0.16.5

* Handle GateUpPack_LinearLayer in DeepSpeed v0.16.5

Author: Xia-Weiwen

intel_extension_for_pytorch/nn/utils/_weight_prepack.py

@@ -101,18 +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
-
-        ds_layers = [LinearAllreduce, LinearLayer]
-
-        # TODO: remove this logic once deepspeed LmHeadLinearAllreduce change has been upstream-ed.
-        try:
-            from deepspeed.module_inject.layers import LmHeadLinearAllreduce
+        from deepspeed.module_inject.layers import (
+            LinearAllreduce,
+            LinearLayer,
+            LmHeadLinearAllreduce,
+            fused_LinearLayer,
+            GateUpPack_LinearLayer,
+        )
 
-            ds_layers.append(LmHeadLinearAllreduce)
-            return ds_layers
-        except ImportError:
-            return ds_layers
+        ds_layers = [
+            LinearAllreduce,
+            LinearLayer,
+            LmHeadLinearAllreduce,
+            fused_LinearLayer,
+            GateUpPack_LinearLayer,
+        ]
+        return ds_layers
     except ImportError:
         return None
 

intel_extension_for_pytorch/utils/weight_only_quantization.py

@@ -292,11 +292,19 @@ def _convert_woq_with_low_precision_checkpoint(
 
     deepspeed_modules = may_import_deepspeed_modules()
     if deepspeed_modules is not None:
-        LinearAllreduce, LinearLayer, LmHeadLinearAllreduce = deepspeed_modules[:]
+        (
+            LinearAllreduce,
+            LinearLayer,
+            LmHeadLinearAllreduce,
+            fused_LinearLayer,
+            GateUpPack_LinearLayer,
+        ) = deepspeed_modules
         q_op_map.update(
             {
                 LinearAllreduce: IpexWoqLinearAllreduce,
                 LinearLayer: WeightOnlyQuantizedLinear,
+                fused_LinearLayer: WeightOnlyQuantizedLinear,
+                GateUpPack_LinearLayer: WeightOnlyQuantizedLinear,
             }
         )
 

tests/cpu/test_deepspeed.py

@@ -197,7 +197,7 @@ def _get_ds_model(self, m_linear):
     def test_ipex_optimize(self):
         deepspeed_modules = may_import_deepspeed_modules()
         if deepspeed_modules is not None:
-            LinearAllreduce, LinearLayer, LmHeadLinearAllreduce = deepspeed_modules
+            LinearAllreduce, LinearLayer, LmHeadLinearAllreduce = deepspeed_modules[:3]
 
             x = torch.randn(2, 3, 64)
             m_linear = DeepSpeedTestM(MyLmHeadModel).eval()
@@ -241,7 +241,7 @@ def _test_quantization(
     ):
         deepspeed_modules = may_import_deepspeed_modules()
         if deepspeed_modules is not None:
-            LinearAllreduce, LinearLayer, LmHeadLinearAllreduce = deepspeed_modules
+            LinearAllreduce, LinearLayer, LmHeadLinearAllreduce = deepspeed_modules[:3]
 
             x = torch.randn(2, 3, 64)
             m_linear = DeepSpeedTestM(MyLmHeadModel).eval()