xpu use allgather

PaddlePaddle · Jul 2, 2024 · 1e81eb9 · 1e81eb9
1 parent a53477c
commit 1e81eb9
Showing 1 changed file with 6 additions and 2 deletions.
diff --git a/paddlenlp/transformers/conversion_utils.py b/paddlenlp/transformers/conversion_utils.py
@@ -37,7 +37,8 @@
 from paddle import Tensor
 from paddle.nn import Layer
 
-from paddlenlp.utils.distributed import distributed_gather
+from paddlenlp.utils.distributed import distributed_gather, distributed_allgather
+from paddlenlp.utils.tools import get_env_device
 from paddlenlp.utils.env import CONFIG_NAME, PADDLE_WEIGHTS_NAME, PYTORCH_WEIGHTS_NAME
 from paddlenlp.utils.import_utils import (
     is_package_available,
@@ -1269,7 +1270,10 @@ def merge_tensor_parallel(cls, state_dict, config) -> None:
         for key in state_dict.keys():
             tensor = state_dict[key]
             if key in name_action_mappings:
-                ret = distributed_gather(tensor, group=mp_group, offload=True)
+                if get_env_device == 'xpu':
+                    ret = distributed_allgather(tensor, group=mp_group, offload=True)
+                else:
+                    ret = distributed_gather(tensor, group=mp_group, offload=True)
                 action = name_action_mappings.pop(key)
                 tensor = action(ret) if is_dst else None
             else: