Fix issue in IPEX TP (#3409)

Author: blzheng

intel_extension_for_pytorch/transformers/optimize.py

@@ -53,6 +53,7 @@ def lowering_class_cpu(m, target_m, new_class, config, tpp=False, woq=False):
 
 
 distributed = False
+is_deepspeed = False
 
 
 def is_distributed(m, ds_layers):
@@ -62,6 +63,8 @@ def is_distributed(m, ds_layers):
             ds_layers,
         ):
             global distributed
+            global is_deepspeed
+            is_deepspeed = True
             distributed = True
             return
         is_distributed(sub_m, ds_layers)
@@ -465,7 +468,8 @@ def model_convert_reference(_model):
         rank = ipex_comm.get_rank() if ipex_comm.has_ccl else 0
         if world_size > 1:
             global distributed
-            if distributed:
+            global is_deepspeed
+            if is_deepspeed:
                 need_ipex_tp = False
             else:
                 need_ipex_tp = True

intel_extension_for_pytorch/transformers/tensor_parallel.py

@@ -561,7 +561,7 @@ def shard_mha_weights(
                         rank,
                         world_size,
                         shard_by_head=True,
-                        value_with_share_qk=True,
+                        value_with_share_qk=value_with_share_qk,
                     )
                     # del sub_m.__dict__["_modules"][l_name]
                     setattr(sub_m, l_name, TPLinear)

tests/cpu/test_ipex_tensor_parallel.py

@@ -119,7 +119,7 @@ def tensor_parallel_with_optimize_transformers(self, model):
         input_dict["position_ids"] = position_ids.unsqueeze(0)
         ref_m = copy.deepcopy(model)
         for dtype in [torch.float32, torch.bfloat16]:
-            ipex_model = ipex.optimize_transformers(model, dtype=dtype)
+            ipex_model = ipex.llm.optimize(model, dtype=dtype)
             with torch.no_grad(), torch.cpu.amp.autocast(
                 enabled=True if dtype is torch.bfloat16 else False
             ):