remove synchronize (#8745)

* Support sync params in tensor parallel config

* Update code

---------

Co-authored-by: chenruibiao <chenruibiao@gmail.com>

docs/trainer.md

@@ -510,6 +510,20 @@ Trainer 是一个简单，但功能完整的 Paddle训练和评估模块，并
                         default -1 for not use tensor parallel,  Suggest tensor_parallel_degree<=8 for better proformance.
                         Note, this need model support in source code, currently GPT/BLOOM/LLAMA/BLOOM/CLM/CHATGLM is supported.
 
+  --tensor_parallel_config
+                        对于张量并行,一些选项会影响训练性能,这里将一些选项配置集中管理,以str形式传入配置.
+                        支持如下选项:
+                            enable_delay_scale_loss : 在优化器阶段做梯度累加，将所有梯度除以累加次数，而不是直接对loss除以累加次数。
+                            sync_param : 在优化器阶段使用broadcast同步所有is_distributed=False的参数
+                            sync_grad : 在优化器阶段使用broadcast同步所有is_distributed=False的梯度
+                            sync_moment : 在优化器阶段使用broadcast同步所有is_distributed=False的momentum
+
+                        Some additional config it highly affect the usage of tensor parallel, we provide some option to config it.
+                        following config is support:
+                            enable_delay_scale_loss, accumulate gradients until optimizer step, all gradients div by accumute step. instead of div accumute step on loss directly.
+                            sync_param, in optimizer step, use broadcast to sync parameters those attr 'is_distributed' is False.
+                            sync_grad, in optimizer step, use broadcast to sync gradients those attr 'is_distributed' is False.
+                            sync_moment, in optimizer step, use broadcast to sync momentums those attr 'is_distributed' is False.
 
   --pipeline_parallel_degree
                         流水线并行是Megatron论文针对多层Transformer结构提出的按层划分方法.
@@ -538,7 +552,7 @@ Trainer 是一个简单，但功能完整的 Paddle训练和评估模块，并
                         following config is support:
                           disable_p2p_cache_shape, if you max sequence length is varying, please set disable_p2p_cache_shape.
                           disable_partial_send_recv, optmize send speed for tensor parallel.
-                          enable_delay_scale_loss, accumulate gradients util optimizer step, all gradients div by inner pipeline accumute step. instead of div accumute step on loss directly.
+                          enable_delay_scale_loss, accumulate gradients until optimizer step, all gradients div by inner pipeline accumute step. instead of div accumute step on loss directly.
                           enable_dp_comm_overlap, fuse data parallel gradient communication.
 
 

model_zoo/gpt-3/external_ops/fast_ln/ln_fwd_kernels.h

@@ -116,10 +116,10 @@ __global__ __launch_bounds__(Ktraits::THREADS_PER_CTA) void ln_fwd_kernel(
     for (int it = 0; it < LDGS; it++) {
 #pragma unroll
       for (int jt = 0; jt < NUM_ELTS; jt++) {
-        output_t y_ij = output_t(rs * (xf[it * NUM_ELTS + jt] - mu));
-        output_t g_ij = gamma[it].data.elt[jt];
-        output_t b_ij = beta[it].data.elt[jt];
-        z[it].data.elt[jt] = (g_ij * y_ij + b_ij);
+        compute_t y_ij = rs * (xf[it * NUM_ELTS + jt] - mu);
+        compute_t g_ij = gamma[it].data.elt[jt];
+        compute_t b_ij = beta[it].data.elt[jt];
+        z[it].data.elt[jt] = output_t(g_ij * y_ij + b_ij);
       }
       z[it].store_to(params.y, idx);
       idx += VEC_COLS_PER_LDG;

paddlenlp/taskflow/taskflow.py

@@ -14,7 +14,6 @@
 # limitations under the License.
 
 import threading
-import warnings
 
 import paddle
 
@@ -43,7 +42,10 @@ from .text_summarization import TextSummarizationTask
 from .word_segmentation import SegJiebaTask, SegLACTask, SegWordTagTask
 from .zero_shot_text_classification import ZeroShotTextClassificationTask
 
-warnings.simplefilter(action="ignore", category=Warning, lineno=0, append=False)
+# import warnings
+
+
+# warnings.simplefilter(action="ignore", category=Warning, lineno=0, append=False)
 
 TASKS = {
     "dependency_parsing": {

paddlenlp/trainer/plugins/timer.py

@@ -0,0 +1,124 @@
+# Copyright 2020-present the HuggingFace Inc. team.
+# Copyright (c) 2022 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import time
+
+import paddle
+
+from paddlenlp.utils.log import logger
+
+
+class _Timer:
+    """Profile Timer for recording time taken by forward/ bacward/ reduce/ step."""
+
+    def __init__(self, name):
+        self.name = name
+        self.elapsed_ = 0.0
+        self.started_ = False
+        self.start_time = time.time()
+
+    def start(self):
+        """Start the timer."""
+        assert not self.started_, "timer has already started"
+        paddle.device.cuda.synchronize()
+        self.start_time = time.time()
+        self.started_ = True
+
+    def stop(self):
+        """Stop the timers."""
+        assert self.started_, "timer is not started."
+        paddle.device.cuda.synchronize()
+        self.elapsed_ += time.time() - self.start_time
+        self.started_ = False
+
+    def reset(self):
+        """Reset timer."""
+        self.elapsed_ = 0.0
+        self.started_ = False
+
+    def elapsed(self, reset=True):
+        """Calculate the elapsed time."""
+        started_ = self.started_
+        # If the timing in progress, end it first.
+        if self.started_:
+            self.stop()
+        # Get the elapsed time.
+        elapsed_ = self.elapsed_
+        # Reset the elapsed time
+        if reset:
+            self.reset()
+        # If timing was in progress, set it back.
+        if started_:
+            self.start()
+        return elapsed_
+
+
+class Timers:
+    """Group of timers."""
+
+    def __init__(self):
+        self.timers = {}
+
+    def __call__(self, name):
+        if name not in self.timers:
+            self.timers[name] = _Timer(name)
+        return self.timers[name]
+
+    def write(self, names, writer, iteration, normalizer=1.0, reset=True):
+        """Write timers to a tensorboard writer"""
+        assert normalizer > 0.0
+        for name in names:
+            value = self.timers[name].elapsed(reset=reset) / normalizer
+            writer.add_scalar("timers/" + name, value, iteration)
+
+    def log(self, names, normalizer=1.0, reset=True):
+        """Log a group of timers."""
+        assert normalizer > 0.0
+        # string = "time (ms) / rate"
+        string = "time (ms)"
+
+        time_dict = {}
+        for name in names:
+            time_dict[name] = self.timers[name].elapsed(reset=reset) * 1000.0 / normalizer
+
+        # total_time = sum(list(time_dict.values()))
+        # string += " | total_time : {:.2f} ".format(total_time)
+        time_dict = sorted(time_dict.items(), key=lambda x: x[1], reverse=True)
+
+        for time_tuple in time_dict:
+            name, value = time_tuple
+            # string += " | {} : {:.2f} ({:.2f}%) ".format(name, value, value * 100.0 / total_time)
+            string += " | {} : {:.2f}".format(name, value)
+        return string
+
+
+_GLOBAL_TIMERS = None
+
+
+def get_timers():
+    global _GLOBAL_TIMERS
+    return _GLOBAL_TIMERS
+
+
+def set_timers():
+    global _GLOBAL_TIMERS
+    logger.info("enable PaddleNLP timer")
+    _GLOBAL_TIMERS = Timers()
+
+
+def disable_timers():
+    global _GLOBAL_TIMERS
+    logger.info("disable PaddleNLP timer")
+    _GLOBAL_TIMERS = None

paddlenlp/trainer/trainer_utils.py

@@ -218,10 +218,10 @@ def speed_metrics(split, start_time, num_samples=None, num_steps=None):
     result = {f"{split}_runtime": round(runtime, 4)}
     if num_samples is not None:
         samples_per_second = num_samples / runtime
-        result[f"{split}_samples_per_second"] = round(samples_per_second, 3)
+        result[f"{split}_samples_per_second"] = samples_per_second
     if num_steps is not None:
         steps_per_second = num_steps / runtime
-        result[f"{split}_steps_per_second"] = round(steps_per_second, 3)
+        result[f"{split}_steps_per_second"] = steps_per_second
     return result
 
 

paddlenlp/trainer/training_args.py

@@ -20,6 +20,7 @@ import contextlib
 import json
 import math
 import os
+import time
 import types
 import warnings
 from dataclasses import asdict, dataclass, field
@@ -206,6 +207,13 @@ class TrainingArguments:
             tensor_parallel_degree means split the transformer layer to how many parts.
             default -1 for not use tensor parallel,  Suggest tensor_parallel_degree<=8 for better proformance.
             Note, this need model support in source code, currently GPT/BLOOM/LLAMA/BLOOM/CLM/CHATGLM is supported.
+        tensor_parallel_config (`str`, *optional*) (
+            Some additional config it highly affect the usage of tensor parallel, we provide some option to config it.
+            following config is support:
+                enable_delay_scale_loss, accumulate gradients until optimizer step, all gradients div by accumute step. instead of div accumute step on loss directly.
+                sync_param, in optimizer step, use broadcast to sync parameters those attr 'is_distributed' is False.
+                sync_grad, in optimizer step, use broadcast to sync gradients those attr 'is_distributed' is False.
+                sync_moment, in optimizer step, use broadcast to sync momentums those attr 'is_distributed' is False.
         pipeline_parallel_degree (`int`, *optional*, defaults to `-1`)
             Pipeline parallelism is parallel technique proposed in (https://arxiv.org/pdf/2104.04473.pdf see 2.2 Pipeline Model Parallelism).
             Pipeline parallelism assigns multi-transformer layers to different cards, the micro batch data stream passed between cards like pipelines.
@@ -217,8 +225,15 @@ class TrainingArguments:
             following config is support:
               disable_p2p_cache_shape, if you max sequence length is varying, please set disable_p2p_cache_shape.
               disable_partial_send_recv, optmize send speed for tensor parallel.
-              enable_delay_scale_loss, accumulate gradients util optimizer step, all gradients div by inner pipeline accumute step. instead of div accumute step on loss directly.
+              enable_delay_scale_loss, accumulate gradients until optimizer step, all gradients div by inner pipeline accumute step. instead of div accumute step on loss directly.
               enable_dp_comm_overlap, fuse data parallel gradient communication.
+              enable_overlap_p2p_comm, overlap p2p communication with computation.
+              enable_clear_every_step_cache, clear every step cache for pipeline parallel.
+
+        sharding_parallel_config (`str`, *optional*)(
+            Some additional config it highly affect the useage of sharding parallel, we provide some option to config it.
+            following config is support:
+                split_param, use Sharding stage1 V2
         recompute (`bool`, *optional*, defaults to `False`):
             Recompute the forward pass to calculate gradients. Used for saving memory.
             Only support for networks with transformer blocks.
@@ -293,6 +308,8 @@ class TrainingArguments:
             scripts](https://github.com/PaddlePaddle/PaddleNLP/tree/develop/examples) for more details.
         flatten_param_grads (`bool`, *optional*):
             Whether use flatten_param_grads method in optimizer, only used on NPU devices. Default is `False`.
+        skip_profile_timer (`bool`, *optional*):
+            Whether skip profile timer, timer will record time usage of forward/ backward/ step, etc.
     """
 
     output_dir: str = field(
@@ -463,6 +480,15 @@ class TrainingArguments:
             )
         },
     )
+    save_sharded_model: bool = field(
+        default=False,
+        metadata={"help": ("whether saved sharded model when sharding_parallel_degree > 1")},
+    )
+
+    load_sharded_model: bool = field(
+        default=False,
+        metadata={"help": ("whether load a sharded model when sharding_parallel_degree > 1")},
+    )
     tensor_parallel_degree: int = field(
         default=-1,
         metadata={
@@ -475,6 +501,19 @@ class TrainingArguments:
             )
         },
     )
+    tensor_parallel_config: str = field(
+        default="",
+        metadata={
+            "help": (
+                "Some additional config it highly affect the usage of tensor parallel, we provide some option to config it."
+                "following config is support:\n"
+                "enable_delay_scale_loss, accumulate gradients until optimizer step, all gradients div by accumute step. instead of div accumute step on loss directly.\n"
+                "sync_param, in optimizer step, use broadcast to sync parameters those attr 'is_distributed' is False.\n"
+                "sync_grad, in optimizer step, use broadcast to sync gradients those attr 'is_distributed' is False.\n"
+                "sync_moment, in optimizer step, use broadcast to sync momentums those attr 'is_distributed' is False.\n"
+            )
+        },
+    )
     pipeline_parallel_degree: int = field(
         default=-1,
         metadata={
@@ -495,11 +534,24 @@ class TrainingArguments:
                 "following config is support:\n"
                 "disable_p2p_cache_shape, if you max sequence length is varying, please set disable_p2p_cache_shape. \n"
                 "disable_partial_send_recv, optmize send speed for tensor parallel.\n"
-                "enable_delay_scale_loss, accumulate gradients util optimizer step, all gradients div by inner pipeline accumute step. instead of div accumute step on loss directly.\n"
+                "enable_delay_scale_loss, accumulate gradients until optimizer step, all gradients div by inner pipeline accumute step. instead of div accumute step on loss directly.\n"
                 "enable_dp_comm_overlap, fuse data parallel gradient communication. \n"
+                "enable_overlap_p2p_comm, overlap p2p communication with computation. \n"
+                "enable_clear_every_step_cache, clear cache every step. \n"
             )
         },
     )
+    sharding_parallel_config: str = field(
+        default="",
+        metadata={
+            "help": (
+                "Some additional config it highly affect the useage of sharding parallel, we provide some option to config it."
+                "following config is support: \n"
+                "split_param, sharding split param\n"
+            )
+        },
+    )
+
     recompute: bool = field(
         default=False,
         metadata={
@@ -589,6 +641,23 @@ class TrainingArguments:
         default=True,
         metadata={"help": "Whether use lazy data processing."},
     )
+    use_async_save: Optional[bool] = field(
+        default=False,
+        metadata={"help": "Whether to use paddle.async_save instead of paddle.save."},
+    )
+    skip_profile_timer: Optional[bool] = field(
+        default=True,
+        metadata={"help": "enable framework timer, will output timeline informatoin in logging and visualdl"},
+    )
+    use_moe: Optional[bool] = field(
+        default=False,
+        metadata={"help": "开启moe训练"},
+    )
+    enable_optimizer_timer: Optional[bool] = field(
+        default=False,
+        metadata={"help": "是否开启Optimzier的timer"},
+    )
+    ignore_load_lr_and_optim: Optional[bool] = field(default=False, metadata={"help": "热启时，不加载lr与optimizer"})
 
     def __post_init__(self):
         env_local_rank = int(os.environ.get("PADDLE_RANK_IN_NODE", -1))
@@ -695,14 +764,14 @@ class TrainingArguments:
 
         if self.amp_master_grad:
             if (
-                self.pipeline_parallel_degree <= 1 and self.tensor_parallel_degree <= 1
-            ) or self.fp16_opt_level != "O2":
+                self.pipeline_parallel_degree <= 1
+                and self.tensor_parallel_degree <= 1
+                and not (self.sharding and ShardingOption.SHARD_OP in self.sharding)
+            ):
                 raise ValueError(
-                    "Temporarily amp master grad only suport for tensor/pipeline parallel with fp16_opt_level O2. please set amp_master_grad to False."
+                    "Temporarily amp master grad only support for tensor/pipeline/sharding parallel. "
+                    "Please set amp_master_grad to False."
                 )
-            if not (self.bf16 or self.fp16):
-                logger.warning("set amp_master_grad to false since amp is disabled.")
-                self.amp_master_grad = False
 
         if self.use_hybrid_parallel:
             world_size = paddle.distributed.get_world_size()
@@ -753,6 +822,10 @@ class TrainingArguments:
                                 "disable_partial_send_recv",
                                 "enable_delay_scale_loss",
                                 "enable_dp_comm_overlap",
+                                "enable_timer",
+                                "enable_release_grads",
+                                "enable_overlap_p2p_comm",
+                                "enable_clear_every_step_cache",
                             ]:
                                 raise ValueError(
                                     f"Found unknown pipeline mode config {x}, accpet config is disable_p2p_cache_shape, disable_partial_send_recv."
@@ -766,9 +839,24 @@ class TrainingArguments:
                         # "delay_scale_loss": True, Fix ME
                     }
                     logger.info(f"PP configs:{strategy.pipeline_configs}, use master_grad: {self.amp_master_grad}")
+
+                    debug_hang = int(os.environ.get("PADDLE_DEBUG_HANG", 0))
                     dygraph_pp_configs = {
                         "delay_scale_loss": True if "enable_delay_scale_loss" in pipeline_parallel_config else False,
-                        "dp_comm_overlap": "enable_dp_comm_overlap" in pipeline_parallel_config,
+                        "dp_comm_overlap": "enable_dp_comm_overlap" in pipeline_parallel_config
+                        and not debug_hang
+                        and self.data_parallel_degree > 1,
+                        "sharding_comm_overlap": "enable_dp_comm_overlap" in pipeline_parallel_config
+                        and not debug_hang
+                        and self.sharding_parallel_degree > 1,
+                        "enable_timer": "enable_timer" in pipeline_parallel_config,
+                        "release_gradients": "enable_release_grads" in pipeline_parallel_config,
+                        "overlap_p2p_comm": False
+                        if "enable_overlap_p2p_comm" not in pipeline_parallel_config
+                        else True,
+                        "clear_every_step_cache": False
+                        if "enable_clear_every_step_cache" not in pipeline_parallel_config
+                        else True,
                     }
 
                     if self.do_eval:
@@ -780,23 +868,100 @@ class TrainingArguments:
                             "Please set per_device_eval_batch_size=per_device_train_batch_size * gradient_accumulation_steps."
                         )
 
+                if sharding_parallel_degree > 1:
+                    sharding_parallel_config = (
+                        set(self.sharding_parallel_config.split(" ")) if self.sharding_parallel_config else set()
+                    )
+                    for x in sharding_parallel_config:
+                        if len(x) > 0:
+                            if x not in [
+                                "split_param",
+                                "shardingv1_comm_overlap",  # will remove in next version
+                                "sharding_comm_overlap",
+                            ]:
+                                raise ValueError(
+                                    f"Found unknown pipeline mode config {x}, accpet config `split_param` and `shardingv1_comm_overlap`."
+                                )
+                    sharding_split_param = "split_param" in sharding_parallel_config
+                    sharding_comm_overlap_non_pp = (
+                        True
+                        if "shardingv1_comm_overlap" in sharding_parallel_config
+                        or "sharding_comm_overlap" in sharding_parallel_config
+                        else False
+                    )
+
                 if tensor_parallel_degree > 1:
                     strategy.tensor_parallel_configs = {"tensor_init_seed": self.seed}
+                    tensor_parallel_config = set(self.tensor_parallel_config.split(" "))
+                    for x in tensor_parallel_config:
+                        if len(x) > 0:
+                            if x not in ["enable_delay_scale_loss", "sync_param", "sync_grad", "sync_moment"]:
+                                raise ValueError(
+                                    f"Found unknown tensor parallel mode config {x} , accept config is enable_delay_scale_loss, sync_param, sync_grad, sync_moment."
+                                )
+                else:
+                    logger.warning("The tensor_parallel_config would be ignored when tensor parallel is not enabled.")
+                    self.tensor_parallel_config = ""
+
+                if self.use_moe:
+                    order = ["sharding", "pp", "dp", "mp"]
+                elif tensor_parallel_degree == 1 and sharding_parallel_degree == 1:
+                    order = ["pp", "dp", "sharding", "mp"]
+                else:
+                    order = ["dp", "sharding", "pp", "mp"]
 
                 hybrid_configs = {
                     "dp_degree": self.data_parallel_degree,
                     "mp_degree": tensor_parallel_degree,
                     "pp_degree": pipeline_parallel_degree,
+                    "order": order,
                     "sharding_degree": sharding_parallel_degree,
                 }
+                if self.enable_optimizer_timer:
+                    hybrid_configs["enable_optimizer_timer"] = True
 
                 if pipeline_parallel_degree > 1:
                     hybrid_configs["pp_configs"] = dygraph_pp_configs
-                    logger.info(f"using pipline configs:{dygraph_pp_configs}")
+                    logger.info(f"using pipeline configs:{dygraph_pp_configs}")
+
+                if tensor_parallel_degree > 1:
+                    sync_param = "sync_param" in tensor_parallel_config
+                    sync_grad = "sync_grad" in tensor_parallel_config
+                    sync_moment = "sync_moment" in tensor_parallel_config
+
+                    # sync_param_name = [""] matches any parameter name.
+                    # If sync_param, sync_grad and sync_moment are not set, the default value in Paddle is :
+                    # sync_param = True, sync_grad = False, sync_moment = False, sync_param_name = ["embedding", "layer_norm", ".b_"].
+                    mp_configs = {}
+                    if sync_param:
+                        mp_configs["sync_param"] = True
+                        mp_configs["sync_param_name"] = [""]
+                    if sync_grad:
+                        mp_configs["sync_grad"] = True
+                        mp_configs["sync_param_name"] = [""]
+                    if sync_moment:
+                        mp_configs["sync_moment"] = True
+                        mp_configs["sync_param_name"] = [""]
+
+                    if mp_configs != {}:
+                        hybrid_configs["mp_configs"] = mp_configs
+                        logger.info(f"using mp configs: {mp_configs}")
 
                 # setter once https://github.com/PaddlePaddle/Paddle/blob/b7295120b0e78b293cd7ae29706e21769d06a3cc/python/paddle/distributed/fleet/base/distributed_strategy.py#L1692
                 strategy.hybrid_configs = hybrid_configs
+
+                if sharding_parallel_degree > 1:
+                    if sharding_split_param:
+                        strategy.hybrid_configs["sharding_configs"].split_param = True
+                    if sharding_comm_overlap_non_pp:
+                        strategy.hybrid_configs["sharding_configs"].comm_overlap = True
+                        strategy.hybrid_configs["sharding_configs"].accumulate_steps = self.gradient_accumulation_steps
+
+                start_time = time.time()
                 fleet.init(is_collective=True, strategy=strategy)
+                paddle.device.cuda.synchronize()
+                elapsed = time.time() - start_time
+                logger.info("NCCL-Connection costs {:.2f} ms.".format(elapsed))
 
                 logger.info(strategy)
 
@@ -936,9 +1101,28 @@ class TrainingArguments:
                 name.append(f"pp{self.pipeline_parallel_rank:0>2d}")
             if self.sharding_parallel_degree > 1:
                 name.append(f"shard{self.sharding_parallel_rank:0>2d}")
-
+            if self.use_moe:
+                name.append(f"moe{self.data_parallel_rank:0>2d}")
             return "_".join(name)
         else:
+            if self.use_moe:
+                return f"moe{self.data_parallel_rank:0>2d}"
+            return None
+
+    @property
+    def old_weight_name_suffix(self):
+        if self.use_hybrid_parallel:
+            name = []
+            if self.tensor_parallel_degree > 1:
+                name.append(f"tp{self.tensor_parallel_rank:0>2d}")
+            if self.pipeline_parallel_degree > 1:
+                name.append(f"pp{self.pipeline_parallel_rank:0>2d}")
+            if self.use_moe:
+                name.append(f"moe{self.data_parallel_rank:0>2d}")
+            return "_".join(name)
+        else:
+            if self.use_moe:
+                return f"moe{self.data_parallel_rank:0>2d}"
             return None
 
     @property
@@ -949,8 +1133,14 @@ class TrainingArguments:
                 name.append(f"tp{self.tensor_parallel_rank:0>2d}")
             if self.pipeline_parallel_degree > 1:
                 name.append(f"pp{self.pipeline_parallel_rank:0>2d}")
+            if self.save_sharding_stage1_model:
+                name.append(f"shard{self.sharding_parallel_rank:0>2d}")
+            if self.use_moe:
+                name.append(f"moe{self.data_parallel_rank:0>2d}")
             return "_".join(name)
         else:
+            if self.use_moe:
+                return f"moe{self.data_parallel_rank:0>2d}"
             return None
 
     @property
@@ -1011,7 +1201,9 @@ class TrainingArguments:
         if self.save_on_each_node:
             return self.local_process_index == 0
         else:
-            if self.tensor_parallel_degree > 1:
+            if self.save_sharding_stage1_model:
+                return True
+            elif self.use_hybrid_parallel:
                 # save on dataset rank 0
                 return self.sharding_parallel_rank == 0 and self.data_parallel_rank == 0
             else:
@@ -1024,6 +1216,18 @@ class TrainingArguments:
         """
         return True
 
+    @property
+    def save_sharding_stage1_model(self):
+        return (
+            ShardingOption.SHARD_OP in self.sharding and self.sharding_parallel_degree > 1 and self.save_sharded_model
+        )
+
+    @property
+    def load_sharding_stage1_model(self):
+        return (
+            ShardingOption.SHARD_OP in self.sharding and self.sharding_parallel_degree > 1 and self.load_sharded_model
+        )
+
     @contextlib.contextmanager
     def main_process_first(self, local=True, desc="work"):
         """

paddlenlp/trainer/utils/__init__.py

@@ -12,10 +12,9 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from .helper import *
-
 from .doc import (
     add_end_docstrings,
     add_start_docstrings,
     add_start_docstrings_to_model_forward,
 )
+from .helper import *

paddlenlp/trainer/utils/reshard/__init__.py

@@ -0,0 +1,23 @@
+# Copyright (c) 2023 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from . import sharding_v1, sharding_v2
+from .common import (
+    SHARDING_STRATEGY_V1,
+    SHARDING_STRATEGY_V2,
+    NodeModelState,
+    all_gather_state_dict,
+    get_sharding_strategy,
+    is_sharding_opt,
+)

paddlenlp/trainer/utils/reshard/common.py

@@ -0,0 +1,579 @@
+# Copyright (c) 2023 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from collections import OrderedDict
+
+import numpy as np
+import paddle
+from paddle.distributed.fleet.meta_optimizers.dygraph_optimizer.dygraph_sharding_optimizer import (
+    DygraphShardingOptimizer,
+)
+from paddle.distributed.fleet.utils.log_util import logger
+
+try:
+    from paddle.distributed.fleet.meta_optimizers.dygraph_optimizer.dygraph_sharding_optimizer import (
+        DygraphShardingOptimizerV2,
+    )
+except:
+    DygraphShardingOptimizerV2 = None
+
+
+from ....transformers.model_utils import unwrap_optimizer
+
+SHARDING_STRATEGY_V1 = "ShardingV1"
+SHARDING_STRATEGY_V2 = "ShardingV2"
+
+
+def is_sharding_opt(optimizer):
+    def check(cls):
+        tmp = unwrap_optimizer(optimizer, cls)
+        if tmp is not None:
+            return True
+        return False
+
+    if check(DygraphShardingOptimizer):
+        return True
+
+    if DygraphShardingOptimizerV2 is not None:
+        if check(DygraphShardingOptimizerV2):
+            return True
+
+    return False
+
+
+def get_sharding_strategy(optimizer):
+    if DygraphShardingOptimizerV2 is not None:
+        tmp = unwrap_optimizer(optimizer, DygraphShardingOptimizerV2)
+        if tmp is not None:
+            return SHARDING_STRATEGY_V2
+    return SHARDING_STRATEGY_V1
+
+
+class NodeModelState:
+    def __init__(self, mp_rank=None, sharding_rank=None, pp_rank=None):
+        self._model_weights = OrderedDict()
+        self._opt_state = OrderedDict()
+        self._master_weights = OrderedDict()
+        self._lr_scheduler = None
+        self.set_node_rank(mp_rank, sharding_rank, pp_rank)
+
+    def set_node_rank(self, mp_rank, sharding_rank, pp_rank):
+        self._mp_rank = mp_rank
+        self._sharding_rank = sharding_rank
+        self._pp_rank = pp_rank
+
+    def _add_kv(self, d, k, v):
+        assert k not in d
+        d[k] = v
+
+    @property
+    def model_weights(self):
+        return self._model_weights
+
+    def add_weight(self, k, v):
+        self._add_kv(self._model_weights, k, v)
+
+    def add_weights(self, model_state_dict, rank=None):
+        for (k, v) in model_state_dict.items():
+            if rank is not None:
+                k = (k, rank)
+            self.add_weight(k, v)
+
+    def set_weights(self, model_state_dict):
+        self._master_weights = model_state_dict
+
+    def set_opt_state(self, opt_state_dict):
+        self._opt_state = opt_state_dict
+
+    def set_master_weights(self, master_weights):
+        self._master_weights = master_weights
+
+    @property
+    def opt_state(self):
+        return self._opt_state
+
+    def add_opt(self, k, v):
+        self._add_kv(self._opt_state, k, v)
+
+    def add_opts(self, opts, rank=None):
+        if "master_weights" in opts:
+            s_master = opts["master_weights"]
+            opts.pop("master_weights")
+            self.add_master_weights(s_master, rank)
+
+        if "LR_Scheduler" in opts:
+            lr_scheduler = opts["LR_Scheduler"]
+            opts.pop("LR_Scheduler")
+            self.set_lr_scheduler(lr_scheduler)
+
+        for (k, v) in opts.items():
+            if rank is not None:
+                k = (k, rank)
+            self.add_opt(k, v)
+
+    @property
+    def master_weights(self):
+        return self._master_weights
+
+    def add_master_weight(self, k, v):
+        self._add_kv(self._master_weights, k, v)
+
+    def add_master_weights(self, master, rank=None):
+        for (k, v) in master.items():
+            if rank is not None:
+                k = (k, rank)
+            self.add_master_weight(k, v)
+
+    @property
+    def lr_scheduler(self):
+        return self._lr_scheduler
+
+    def set_lr_scheduler(self, lr_scheduler):
+        if lr_scheduler is not None:
+            self._lr_scheduler = lr_scheduler
+
+    def map_names(self, map_func):
+        """
+        rename param names and change the keys of the dicts(model_weights, opt, master_weights) accordingly
+        """
+        # model weighs
+        model_weights_tmp = OrderedDict()
+        (self._model_weights, model_weights_tmp) = (model_weights_tmp, self._model_weights)
+        for key in list(model_weights_tmp.keys()):
+            structure_name, t_name = key
+            t_name_new = map_func(structure_name, t_name)
+            self._model_weights[(structure_name, t_name_new)] = model_weights_tmp[key]
+            del model_weights_tmp[key]
+
+        # opt
+        opt_tmp = OrderedDict()
+        (self._opt_state, opt_tmp) = (opt_tmp, self._opt_state)
+        for key in list(opt_tmp.keys()):
+            structure_name, t_name, opt_name = key
+            t_name_new = map_func(structure_name, t_name)
+            if self._model_weights:
+                assert (structure_name, t_name_new) in self._model_weights
+            opt_name_new = t_name_new + opt_name[len(t_name) :]
+            self._opt_state[(structure_name, t_name_new, opt_name_new)] = opt_tmp[key]
+            del opt_tmp[key]
+
+        # master weights
+        master_weights_tmp = OrderedDict()
+        (self._master_weights, master_weights_tmp) = (master_weights_tmp, self._master_weights)
+        for key in list(master_weights_tmp.keys()):
+            structure_name, master_weights_name = key
+            master_weights_name_new = map_func(structure_name, master_weights_name)
+            if self._model_weights:
+                assert (structure_name, t_name_new) in self._model_weights
+            self._master_weights[(structure_name, master_weights_name_new)] = master_weights_tmp[key]
+            del master_weights_tmp[key]
+
+        return self
+
+    def drop_rank(self):
+        """
+        drop rank in the keys of the state dict
+        change dict of (key, rank)=>tensor to dict of key =>tensor
+        """
+
+        def drop(state, l=2):
+            tmp_state = OrderedDict()
+            (state, tmp_state) = (tmp_state, state)
+            for key in list(tmp_state.keys()):
+                k, rank = key
+                assert len(key) == 2
+                assert len(k) == l
+                state[k] = tmp_state[key]
+                del tmp_state[key]
+            return state
+
+        self._model_weights = drop(self._model_weights, 2)
+        self._opt_state = drop(self._opt_state, 3)
+        self._master_weights = drop(self._master_weights, 2)
+        return self
+
+    def collapse_key(self):
+        """
+        collapse dict of (key, rank)=>tensor to dict of key=>list[(rank, tensor)]
+        """
+
+        def collapse(state, l):
+            tmp_state = OrderedDict()
+            (state, tmp_state) = (tmp_state, state)
+            state_keys = list(tmp_state.keys())
+            state_keys = sorted(state_keys)
+            pre = None
+            for key in state_keys:
+                assert len(key) == 2
+                k, rank = key
+                assert len(k) == l
+                if k != pre:
+                    pre = k
+                    state[k] = []
+                state[k].append((rank, tmp_state[key]))
+                del tmp_state[key]
+            return state
+
+        self._model_weights = collapse(self._model_weights, 2)
+        self._opt_state = collapse(self._opt_state, 3)
+        self._master_weights = collapse(self._master_weights, 2)
+        return self
+
+    def flatten_key(self):
+        """
+        flatten dict of key=>list[(rank, tensor)], to dict of (key, rank)=>tensor
+        """
+
+        def flatten(state, l):
+            tmp_state = OrderedDict()
+            (state, tmp_state) = (tmp_state, state)
+            state_keys = list(tmp_state.keys())
+            for key in state_keys:
+                assert len(key) == l
+                for (rank, items) in tmp_state[key]:
+                    state[(key, rank)] = items
+                del tmp_state[key]
+            return state
+
+        self._model_weights = flatten(self._model_weights, 2)
+        self._opt_state = flatten(self._opt_state, 3)
+        self._master_weights = flatten(self._master_weights, 2)
+        return self
+
+    def pack_keys(self, structure_name_mapping=None):
+        """
+        change the key of model_weights dict from param_name to (structure_name, param_name);
+        change the key of opt dict from opt_name to (structure_name, param_name, opt_name);
+        chnage the key of master weights dict from param_name to (structure_name, param_name)
+        """
+        # pack key for pp convert
+        def _opt_name_to_tname(tensor_names, opt_names):
+            suffix = [
+                "_fp32_master_0_beta1_pow_acc_0",
+                "_fp32_master_0_beta2_pow_acc_0",
+                "_fp32_master_0_moment1_0",
+                "_fp32_master_0_moment2_0",
+                "_beta1_pow_acc_0",
+                "_beta2_pow_acc_0",
+                "_moment1_0",
+                "_moment2_0",
+            ]
+            opt_to_t = {}
+
+            for t in opt_names:
+                _find = False
+                for s in suffix:
+                    if t.endswith(s):
+                        print(t, t[:-len(s)], t[:-len(s)] in tensor_names)
+                        opt_to_t[t] = t[:-len(s)]
+                        _find = True
+                        break
+                assert _find
+            # tensor_names = set(tensor_names)
+            # all_names = []
+            # all_names.extend(list(tensor_names))
+            # all_names.extend(opt_names)
+            # all_names.sort()
+            # pre_t_name = ""
+            # opt_to_t = {}
+            # for n in all_names:
+            #     if n in tensor_names:
+            #         # we get a param
+            #         pre_t_name = n
+            #     else:
+            #         assert pre_t_name
+            #         opt_to_t[n] = pre_t_name
+            return opt_to_t
+
+        if structure_name_mapping is not None:
+            tname_to_structure_name = {v: k for (k, v) in structure_name_mapping.items()}
+        else:
+            structure_name_mapping = {k: v for (k, v) in self._model_weights.items()}
+            tname_to_structure_name = {v: k for (k, v) in structure_name_mapping.items()}
+
+        tensor_names = list(tname_to_structure_name.keys())
+        opt_names = list(self._opt_state.keys())
+        opt_name_to_tname = _opt_name_to_tname(tensor_names, opt_names)
+
+        # model state
+        model_weights_tmp = OrderedDict()
+        (self._model_weights, model_weights_tmp) = (model_weights_tmp, self._model_weights)
+        for k in list(model_weights_tmp.keys()):
+            t_name = structure_name_mapping[k]
+            self._model_weights[(k, t_name)] = model_weights_tmp[k]
+            del model_weights_tmp[k]
+
+        # opt
+        opt_tmp = OrderedDict()
+        (self._opt_state, opt_tmp) = (opt_tmp, self._opt_state)
+        for opt_name in list(opt_tmp.keys()):
+            assert opt_name in opt_name_to_tname
+            t_name = opt_name_to_tname[opt_name]
+            assert t_name in tname_to_structure_name
+            structure_name = tname_to_structure_name[t_name]
+            self._opt_state[(structure_name, t_name, opt_name)] = opt_tmp[opt_name].cpu()
+            del opt_tmp[opt_name]
+
+        # master weights
+        master_weights_tmp = OrderedDict()
+        (self._master_weights, master_weights_tmp) = (master_weights_tmp, self._master_weights)
+        for master_weights_name in list(master_weights_tmp.keys()):
+            assert master_weights_name in tname_to_structure_name
+            structure_name = tname_to_structure_name[master_weights_name]
+            self._master_weights[(structure_name, master_weights_name)] = master_weights_tmp[master_weights_name].cpu()
+            del master_weights_tmp[master_weights_name]
+
+        return self
+
+    def unpack_keys(self):
+        """
+        the opposite of pack_keys,
+        revert the key of model_weights dict from  (structure_name, param_name) to param_name
+        revert the key of opt dict from  (structure_name, param_name, opt_name) to opt_name
+        revert the key of master weights dict from (structure_name, param_name) to param_name
+        """
+        # model weights
+        model_weights_tmp = OrderedDict()
+        (self._model_weights, model_weights_tmp) = (model_weights_tmp, self._model_weights)
+        for key in list(model_weights_tmp.keys()):
+            structure_name, t_name = key
+            self._model_weights[structure_name] = model_weights_tmp[key]
+            self._model_weights[structure_name].name = t_name
+            del model_weights_tmp[key]
+        # opt
+        opt_tmp = OrderedDict()
+        (self._opt_state, opt_tmp) = (opt_tmp, self._opt_state)
+        for key in list(opt_tmp.keys()):
+            structure_name, t_name, opt_name = key
+            if structure_name in self._model_weights:
+                assert self._model_weights[structure_name].name == t_name
+            self._opt_state[opt_name] = opt_tmp[key]
+            del opt_tmp[key]
+
+        # master weights
+        master_weights_tmp = OrderedDict()
+        (self._master_weights, master_weights_tmp) = (master_weights_tmp, self._master_weights)
+        for key in list(master_weights_tmp.keys()):
+            structure_name, master_weights_name = key
+            if structure_name in self._model_weights:
+                assert self._model_weights[structure_name].name == master_weights_name
+            self._master_weights[master_weights_name] = master_weights_tmp[key]
+            del master_weights_tmp[key]
+        return self
+
+    def split_state(self, split_func):
+        """
+        split this node state to multiple node state according to the passed in split_func
+        """
+        node_model_states = {}
+        for (k, v) in self._model_weights.items():
+            rank = split_func(k)
+            if rank not in node_model_states:
+                node_model_states[rank] = NodeModelState()
+            node_model_states[rank].add_weight(k, v)
+
+        for (k, v) in self._opt_state.items():
+            rank = split_func(k)
+            if rank not in node_model_states:
+                node_model_states[rank] = NodeModelState()
+            node_model_states[rank].add_opt(k, v)
+
+        for (k, v) in self._master_weights.items():
+            rank = split_func(k)
+            if rank not in node_model_states:
+                node_model_states[rank] = NodeModelState()
+            node_model_states[rank].add_master_weight(k, v)
+
+        return node_model_states
+
+    def even_distribute(self, group):
+        """
+        distribute the node state evenly among all workers in group， and make sure
+        in the dicts of (key, rank)=>tensor, items keys of the same key but different rank are distributed to the
+        same worker
+        """
+        # sharding degree == 1
+        if group is None or group.nranks < 2:
+            return self
+
+        def build_router(state_dict):
+            state_keys_list = all_gather_simple_object([(k, v.shape) for (k, v) in state_dict.items()], group)
+
+            key_to_size = {}
+            for l in state_keys_list:
+                for (k, shape) in l:
+                    key, rank = k
+                    if key not in key_to_size:
+                        key_to_size[key] = 0
+                    key_to_size[key] = key_to_size[key] + np.prod(shape)
+
+            key_to_size = sorted(list(key_to_size.items()), key=lambda x: x[1], reverse=True)
+            node_distributed = [0 for _ in range(group.nranks)]
+            key_to_rank = {}
+            for (k, v) in key_to_size:
+                min_val = min(node_distributed)
+                min_index = node_distributed.index(min_val)
+                key_to_rank[k] = min_index
+                node_distributed[min_index] = node_distributed[min_index] + v
+
+            return key_to_rank
+
+        def distribute(state_dict):
+
+            key_to_rank = build_router(state_dict)
+
+            def filter_func(key):
+                assert key[0] in key_to_rank, key
+                dst_rank = key_to_rank[key[0]]
+                return dst_rank == group.rank
+
+            return _all_gather_state_dict(state_dict, filter_func, group)
+
+        self._model_weights = distribute(self._model_weights)
+        self._opt_state = distribute(self._opt_state)
+        self._master_weights = distribute(self._master_weights)
+        return self
+
+    def reshard(self, group, filter_func):
+        """
+        reshard according to the passed in filter_func
+        """
+        self._model_weights = _all_gather_state_dict(self._model_weights, filter_func, group)
+        self._opt_state = _all_gather_state_dict(self._opt_state, filter_func, group)
+        self._master_weights = _all_gather_state_dict(self._master_weights, filter_func, group)
+        lr_schedulers = all_gather_simple_object(self._lr_scheduler, group)
+        self._lr_scheduler = lr_schedulers[0]
+        return self
+
+    def split_items(self, split_func):
+        """
+        split tensor in the dicts of key=tensor, change the dicts to dicts of key=>list[(rank, tensor)]
+        """
+
+        def split(state, l):
+            tmp_state = OrderedDict()
+            (state, tmp_state) = (tmp_state, state)
+            state_keys = list(tmp_state.keys())
+            for key in state_keys:
+                assert len(key) == l
+                v = tmp_state[key]
+                state[key] = split_func(key, v)
+                del tmp_state[key]
+            return state
+
+        self._model_weights = split(self._model_weights, 2)
+        self._opt_state = split(self._opt_state, 3)
+        self._master_weights = split(self._master_weights, 2)
+        return self
+
+    def merge_items(self, merge_func):
+        """
+        merge list in the dicts of key=>list[(rank, tensor)]  a tensor, change the dicts to dicts of key=>tensor
+        """
+
+        def merge(state, l):
+            tmp_state = OrderedDict()
+            (state, tmp_state) = (tmp_state, state)
+            state_keys = list(tmp_state.keys())
+            for key in state_keys:
+                assert len(key) == l
+                v = tmp_state[key]
+                v = sorted(v, key=lambda x: x[0])
+                state[key] = merge_func(key, v)
+                del tmp_state[key]
+            return state
+
+        self._model_weights = merge(self._model_weights, 2)
+        self._opt_state = merge(self._opt_state, 3)
+        self._master_weights = merge(self._master_weights, 2)
+        return self
+
+    def merge_from(self, other, rank=None):
+        self.add_opts(other.opt_state, rank)
+        self.add_master_weights(other.master_weights, rank)
+        self.set_lr_scheduler(other.lr_scheduler)
+
+    def get_opt_state_dict(self):
+        opt_state_dict = OrderedDict()
+        for (k, v) in self.opt_state.items():
+            opt_state_dict[k] = v
+        if self._lr_scheduler is not None:
+            opt_state_dict["LR_Scheduler"] = self._lr_scheduler
+        opt_state_dict["master_weights"] = self._master_weights
+        return opt_state_dict
+
+
+def all_gather_simple_object(obj, group):
+    res = []
+    paddle.distributed.all_gather_object(res, obj, group)
+    return res
+
+
+def all_gather_state_dict(state_dict, filter_func, group):
+    res = OrderedDict()
+
+    def map_func(weight):
+        if isinstance(weight, paddle.Tensor):
+            weight = weight.numpy()
+        return weight
+
+    state_dict = {k: map_func(v) for (k, v) in state_dict.items()}
+
+    meta_dict = {}
+    for (k, v) in state_dict.items():
+        # src rank
+        meta_dict[k] = (v.dtype, v.shape, group.rank)
+
+    meta_dict_list = all_gather_simple_object(meta_dict, group)
+
+    total_meta_dict = {}
+    for meta_dict in meta_dict_list:
+        for (k, v) in meta_dict.items():
+            assert k not in total_meta_dict
+            total_meta_dict[k] = v
+
+    meta_list = list(total_meta_dict.items())
+    meta_list = sorted(meta_list, key=lambda x: x[0])
+    for (k, meta) in meta_list:
+        dtype, shape, rank = meta
+        if rank == group.rank:
+            assert k in state_dict
+            tensor = paddle.to_tensor(state_dict[k])
+            del state_dict[k]
+        else:
+            tensor = paddle.to_tensor(np.empty(shape, dtype))
+        logger.info(f"broadcast {k} from {rank}")
+        # broadcast the tensor
+        paddle.distributed.broadcast(
+            tensor,
+            src=group.ranks[rank],
+            group=group,
+            sync_op=True,
+        )
+        if filter_func(k):
+            res[k] = tensor.cpu()
+        del tensor
+    return res
+
+
+def _all_gather_state_dict(state_dict, filter_func, group):
+    remote_state_dict_keys = [k for k in state_dict.keys() if not filter_func(k)]
+    tmp_state_dict = OrderedDict()
+    for k in remote_state_dict_keys:
+        tmp_state_dict[k] = state_dict[k]
+        state_dict.pop(k)
+    tmp_state_dict = all_gather_state_dict(tmp_state_dict, filter_func, group)
+    for (k, v) in tmp_state_dict.items():
+        state_dict[k] = v
+    return state_dict

paddlenlp/trainer/utils/reshard/sharding_v1.py

@@ -0,0 +1,42 @@
+# Copyright (c) 2023 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from paddle.distributed.fleet.meta_optimizers.dygraph_optimizer.dygraph_sharding_optimizer import (
+    DygraphShardingOptimizer,
+)
+
+from ....transformers.model_utils import unwrap_optimizer
+
+
+def shard(node_model_state, model, optimizer, hcg):
+    group = hcg.get_sharding_parallel_group()
+    cur_rank = group.rank
+    optimizer = unwrap_optimizer(optimizer, DygraphShardingOptimizer)
+    assert optimizer is not None
+    param2rank = optimizer._param2rank
+
+    def filter_func(key):
+        names = key
+        param_name = names[1]
+        assert param_name in param2rank
+        dst_rank = param2rank[param_name]
+        return dst_rank == cur_rank
+
+    node_model_state.reshard(group, filter_func)
+    return node_model_state
+
+
+def restore(node_model_state, model, optimizer, hcg):
+    node_model_state.drop_rank()
+    return node_model_state

paddlenlp/trainer/utils/reshard/sharding_v2.py

@@ -0,0 +1,175 @@
+# Copyright (c) 2023 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import numpy as np
+import paddle
+from paddle.distributed.fleet.meta_optimizers.dygraph_optimizer import (
+    HybridParallelOptimizer,
+)
+from paddle.distributed.fleet.model import PipelineParallel
+
+from ....transformers.model_utils import unwrap_optimizer
+
+try:
+    from paddle.distributed.fleet.meta_optimizers.dygraph_optimizer.dygraph_sharding_optimizer import (
+        DygraphShardingOptimizerV2,
+    )
+except:
+    DygraphShardingOptimizerV2 = None
+
+
+def shard(node_model_state, model, optimizer, hcg):
+    assert DygraphShardingOptimizerV2 is not None
+    group = hcg.get_sharding_parallel_group()
+    cur_rank = group.rank
+    split_infos = collect_split_info(optimizer, model)
+
+    def split_func(k, v):
+        param_name = k[1]
+        opt_name = k[-1]
+        assert param_name in split_infos
+        is_beta = is_bata(opt_name)
+        index, padded_size, buffer_size, has_slice_grad = split_infos[param_name]
+
+        if not is_beta:
+            v = pad_tensor(v, padded_size)
+
+        def get_slice(v, begin, end):
+            if is_beta:
+                return v
+            return slice_tensor(v, begin, end)
+
+        assert buffer_size % group.nranks == 0, f"buffer_size {buffer_size} group.nranks {group.nranks}"
+        buffer_slice = buffer_size // group.nranks
+
+        # has slice grad in cur rank
+        if has_slice_grad:
+            assert index < (cur_rank + 1) * buffer_slice
+            assert index + padded_size > cur_rank * buffer_slice
+
+        offset = buffer_slice - index % buffer_slice
+        tensors = []
+        tensors.append((index // buffer_slice, get_slice(v, 0, min(offset, padded_size))))
+
+        left_size = padded_size - offset
+
+        if left_size > 0:
+            for _ in range((left_size + buffer_slice - 1) // buffer_slice):
+                end = min(offset + buffer_slice, padded_size)
+                assert end <= buffer_size
+                tensors.append(((offset + index) // buffer_slice, get_slice(v, offset, end)))
+                offset = end
+
+        return tensors
+
+    node_model_state.split_items(split_func).flatten_key()
+
+    def filter_func(k):
+        names, rank = k
+        assert rank < group.nranks
+        return rank == cur_rank
+
+    # reshard
+    node_model_state.reshard(group, filter_func)
+    node_model_state.drop_rank()
+    return node_model_state
+
+
+def restore(node_model_state, model, optimizer, hcg):
+    group = hcg.get_sharding_parallel_group()
+    # evenly distribute param
+    node_model_state.even_distribute(group)
+    param_shapes = {k: v.shape for (k, v) in model.state_dict().items()}
+
+    def merge_func(k, v):
+        structure_name = k[0]
+        opt_name = k[-1]
+        assert structure_name in param_shapes, structure_name
+        tensor_list = [e[1] for e in v]
+        # do not merge beta acc
+        if is_bata(opt_name):
+            return tensor_list[0]
+        shape = param_shapes[structure_name]
+        return merge_tensors(tensor_list, shape)
+
+    node_model_state.collapse_key().merge_items(merge_func)
+    return node_model_state
+
+
+def merge_tensors(tensor_list, shape):
+    assert len(tensor_list) > 0
+    if len(tensor_list) == 1:
+        t = tensor_list[0]
+    else:
+        assert len(tensor_list[0].shape) == 1
+        t = paddle.concat(x=tensor_list, axis=0)
+    tensor_size = np.prod(shape)
+    padded_size = t._numel()
+    assert padded_size >= tensor_size
+    t = t._slice(0, tensor_size)
+    t.get_tensor()._set_dims(shape)
+    return t
+
+
+def pad_tensor(tensor, padded_size):
+    tensor_shape = tensor.shape
+    tensor_size = np.prod(tensor_shape)
+    assert tensor_size <= padded_size
+    t = paddle.zeros([padded_size], dtype=tensor.dtype)
+    tensor.flatten_()
+    t[0:tensor_size] = tensor
+    tensor.get_tensor()._set_dims(tensor_shape)
+    return t
+
+
+def slice_tensor(tensor, begin, end):
+    return tensor[begin:end]
+
+
+def collect_split_info(optimizer, model):
+    split_infos = {}
+
+    def gather_infos(comm_buffer):
+        for (k, v) in comm_buffer._sharding_param_grad_view.items():
+            index = v._index
+            padded_size = v._padded_size
+            buffer_size = v._param_buffer._numel()
+            has_slice_grad = v._slice_grad is not None
+            split_infos[k] = (index, padded_size, buffer_size, has_slice_grad)
+
+    if isinstance(model, PipelineParallel) and model._sharding_comm_overlap > 0:
+        optimizer = unwrap_optimizer(optimizer, HybridParallelOptimizer)
+        assert optimizer is not None
+        # dalayed comm_overlap_hook register
+        model.register_sharding_comm_overlap_hook(optimizer)
+        for (k, v) in model._chunk_2_comm_buffers.items():
+            for comm_buffer in v:
+                gather_infos(comm_buffer)
+
+    else:
+        optimizer = unwrap_optimizer(optimizer, DygraphShardingOptimizerV2)
+        assert optimizer is not None
+        for comm_buffer in optimizer._comm_buffer_list:
+            gather_infos(comm_buffer)
+
+    assert len(split_infos)
+    return split_infos
+
+
+def is_bata(name):
+    if "_beta1_pow_acc_" in name:
+        return True
+    if "_beta2_pow_acc_" in name:
+        return True
+    return False

paddlenlp/transformers/conversion_utils.py

@@ -256,7 +256,65 @@ class StateDictKeysChecker:
         return all_diff_keys
 
 
-def merge_tensor_parallel_weight(weight_list, is_column=True):
+def naive_fuse_merge_tp(weight_list, is_column=True, fuse_tensor_parts=2):
+    """
+
+    [A1 B1],[A2 B2]  => [A1, A2, B1, B2]
+
+    Args:
+        weight_list (List[np.ndarray]): The splited tensor parallel weight list.
+        is_column (bool, optional): Is ColumnLinear or RowLinear. Defaults to True.
+
+    Returns:
+        weight (np.ndarray): the merged weight.
+    """
+    if is_column:
+        axis = -1
+    else:
+        axis = 0
+
+    reorder = []
+    for item in weight_list:
+        reorder.extend(np.split(item, fuse_tensor_parts, axis=axis))
+    # 0 1 2 3 -> 0 2 1 3
+    index = (
+        np.transpose(np.arange(len(reorder)).reshape([len(weight_list), fuse_tensor_parts]), [1, 0])
+        .reshape(-1)
+        .tolist()
+    )
+    return np.concatenate([reorder[i] for i in index], axis=axis)
+
+
+def naive_fuse_split_tp(
+    weight, tensor_parallel_degree, tensor_parallel_rank=None, is_column=True, fuse_tensor_parts=2
+):
+    """
+
+    [A1, A2, B1, B2] => [A1 B1],[A2 B2]
+
+    Args:
+        weight (numpy.ndarray): the tensor weight,
+        tensor_parallel_degree (int): tensor_parallel_degree
+        tensor_parallel_rank (int): tensor_parallel_rank
+        is_column (bool, optional): is ColumnLinear . Defaults to True.
+
+    Returns:
+        tensor (numpy.ndarray): splited weight.
+
+    """
+    axis = -1 if is_column else 0
+    splited = np.split(weight, fuse_tensor_parts * tensor_parallel_degree, axis=axis)
+
+    if tensor_parallel_rank is None:
+        ret = []
+        for tensor_parallel_rank in range(tensor_parallel_degree):
+            ret.append(np.concatenate(splited[tensor_parallel_rank::tensor_parallel_degree], axis=axis))
+        return ret
+
+    return np.concatenate(splited[tensor_parallel_rank::tensor_parallel_degree], axis=axis)
+
+
+def normal_fuse_merge_tp(weight_list, is_column=True):
     """
 
     [A1],[A2]  => [A1, A2]
@@ -274,7 +332,7 @@ def merge_tensor_parallel_weight(weight_list, is_column=True):
         return np.concatenate(weight_list, axis=0)
 
 
-def split_tensor_parallel_weight(weight, tensor_parallel_degree, tensor_parallel_rank=None, is_column=True):
+def normal_fuse_split_tp(weight, tensor_parallel_degree, tensor_parallel_rank=None, is_column=True):
     """
 
     [A1, A2]  =>  [A1],[A2]
@@ -288,6 +346,29 @@ def split_tensor_parallel_weight(weight, tensor_parallel_degree, tensor_parallel
     Returns:
         tensor (numpy.ndarray): splited weight.
     """
+    dim = -1 if is_column else 0
+    if "PySafeSlice" in str(type(weight)):
+        size = weight.get_shape()[dim]
+        block_size = size // tensor_parallel_degree
+        start = tensor_parallel_rank * block_size
+        stop = (tensor_parallel_rank + 1) * block_size
+        assert (
+            size % tensor_parallel_degree == 0
+        ), f"The choosen size {size} is not compatible with sharding on {tensor_parallel_degree} shards"
+
+        if dim == 0 or len(weight.get_shape()) == 1:
+            tensor = weight[start:stop]
+        elif dim == -1:
+            tensor = weight[:, start:stop]
+        else:
+            raise NotImplementedError("Let's make that generic when needed")
+        return tensor
+
+    size = weight.shape[dim]
+    assert (
+        size % tensor_parallel_degree == 0
+    ), f"The choosen size {size} is not compatible with sharding on {tensor_parallel_degree} shards. for tensor shape {weight.shape}"
+
     if is_column:
         splited_weights = np.split(weight, tensor_parallel_degree, axis=-1)
     else:
@@ -362,13 +443,17 @@ def splited_qkv_to_tensor_parallel_qkv(weight_list, num_attention_heads):
 
 
 def get_tensor_parallel_merge_func(tensor_parallel_degree, tensor_parallel_rank, num_attention_heads=None):
-    def fn(x, is_column=True, transpose=False, is_old_qkv=False):
+    def fn(x, is_column=True, transpose=False, is_old_qkv=False, is_naive_2fuse=False, is_naive_3fuse=False):
         if x is None:
             return None
-        x = merge_tensor_parallel_weight(
-            x,
-            is_column=is_column,
-        )
+
+        if is_naive_2fuse:
+            return naive_fuse_merge_tp(x, is_column=is_column, fuse_tensor_parts=2)
+        elif is_naive_3fuse:
+            return naive_fuse_merge_tp(x, is_column=is_column, fuse_tensor_parts=3)
+        else:
+            x = normal_fuse_merge_tp(x, is_column=is_column)
+
         if is_old_qkv:
             assert is_column, "QKV tensor should be column parallel linear."
             assert num_attention_heads is not None, "is_old_qkv need num_attention_heads"
@@ -382,7 +467,7 @@ def get_tensor_parallel_merge_func(tensor_parallel_degree, tensor_parallel_rank,
 
 
 def get_tensor_parallel_split_func(tensor_parallel_degree, tensor_parallel_rank, num_attention_heads=None):
-    def fn(x, is_column=True, transpose=False, is_old_qkv=False):
+    def fn(x, is_column=True, transpose=False, is_old_qkv=False, is_naive_2fuse=False, is_naive_3fuse=False):
         if x is None:
             return None
         if transpose:
@@ -391,12 +476,16 @@ def get_tensor_parallel_split_func(tensor_parallel_degree, tensor_parallel_rank,
             assert is_column, "QKV tensor should be column parallel linear."
             assert num_attention_heads is not None, "is_old_qkv need num_attention_heads"
             x = naive_merged_qkv_to_tensor_parallel_qkv(x, num_attention_heads)
-        return split_tensor_parallel_weight(
-            x,
-            tensor_parallel_degree=tensor_parallel_degree,
-            tensor_parallel_rank=tensor_parallel_rank,
-            is_column=is_column,
-        )
+        if is_naive_2fuse:
+            return naive_fuse_split_tp(
+                x, tensor_parallel_degree, tensor_parallel_rank, is_column=is_column, fuse_tensor_parts=2
+            )
+        if is_naive_3fuse:
+            return naive_fuse_split_tp(
+                x, tensor_parallel_degree, tensor_parallel_rank, is_column=is_column, fuse_tensor_parts=3
+            )
+
+        return normal_fuse_split_tp(x, tensor_parallel_degree, tensor_parallel_rank, is_column=is_column)
 
     return fn
 
@@ -925,6 +1014,14 @@ class ConversionMixin:
         """
         raise NotImplementedError
 
+    @classmethod
+    def get_tensor_parallel_convert_actions(cls, config: PretrainedConfig, loaded_state_dict_keys, ignore_error=False):
+        name_action_mappings = cls._get_tensor_parallel_mappings(config)
+        state_keys_map = cls._resolve_prefix_keys(name_action_mappings.keys(), loaded_state_dict_keys, ignore_error)
+        for k, v in state_keys_map.items():
+            name_action_mappings[v] = name_action_mappings.pop(k)
+        return name_action_mappings
+
     @classmethod
     def convert_tensor_parallel(
         cls, weight_file: str, config: PretrainedConfig, state_dict=None, ignore_error=False
@@ -932,13 +1029,14 @@ class ConversionMixin:
         """the entry of converting config and converting model file
 
         Args:
-            input_dir (str | None): the input dir which contains `pytorch_model.bin` and `config.json` file
+            weight_file (str | None): the weight file path of `model_state.pdparams` file
             config (PretrainedConfig): the PretrainedConfig instance of model
         """
         name_action_mappings = cls._get_tensor_parallel_mappings(config)
         if state_dict is None:
             with device_guard("cpu"):
                 state_dict = paddle.load(weight_file, return_numpy=False)
+            logger.info("Starting to convert orignal state_dict to tensor parallel state_dict.")
 
         state_keys_map = cls._resolve_prefix_keys(name_action_mappings.keys(), state_dict.keys(), ignore_error)
 
@@ -948,7 +1046,7 @@ class ConversionMixin:
         for name, action in name_action_mappings.items():
             if name not in state_dict:
                 if not ignore_error:
-                    logger.warning(f"key<{name}> not in the model state weight file.")
+                    logger.warning(f"Key <{name}> not in the model state weight file.")
                 continue
             tensor = state_dict.pop(name)
             new_tensor = action(tensor)
@@ -986,6 +1084,11 @@ class ConversionMixin:
             else:
                 tensor = tensor.numpy() if is_dst else None
 
+            # keep state dict use paddle.tensor
+            if isinstance(tensor, np.ndarray):
+                with device_guard("cpu"):
+                    tensor = paddle.Tensor(tensor, zero_copy=True)
+
             state_dict_to_save[key] = tensor
 
         if len(name_action_mappings) > 0:
@@ -1024,7 +1127,7 @@ class ConversionMixin:
                     break
             if key not in state_keys_map:
                 if not ignore_error:
-                    logger.error(f"could not find name {key} in loaded state dict!")
+                    logger.error(f"tensor parallel conversion: could not find name {key} in loaded state dict!")
             else:
                 state_keys_real.remove(state_keys_map[key])
 

paddlenlp/transformers/generation_utils.py

@@ -1151,6 +1151,7 @@ class GenerationMixin(object):
             else:
                 next_tokens = paddle.multinomial(probs)
 
+            #  paddle.distributed.broadcast(next_tokens, src=0)
             next_scores = paddle.index_sample(origin_probs, next_tokens)
 
             if eos_token_id is not None:

paddlenlp/transformers/model_utils.py

@@ -19,6 +19,7 @@ import os
 import re
 import shutil
 import tempfile
+from collections import OrderedDict
 from contextlib import contextmanager
 from typing import Any, Callable, Dict, List, Optional, Set, Tuple, Type
 
@@ -79,6 +80,87 @@ __all__ = [
 ]
 
 
+def unwrap_optimizer(optimizer, optimizer_instances=()):
+    if optimizer is None:
+        return None
+    while hasattr(optimizer, "_inner_opt") and not isinstance(optimizer, optimizer_instances):
+        optimizer = optimizer._inner_opt
+    if isinstance(optimizer, optimizer_instances):
+        return optimizer
+    return None
+
+
+def filter_sharded_params(state_dict, optimizer, sharding_group):
+    sharding_rank = sharding_group.rank
+    sharding_world_size = sharding_group.nranks
+    from paddlenlp.trainer.utils import reshard as reshard_util
+
+    logger.info(f"filter sharded_params not placed in sharding_rank {sharding_rank} .")
+    if not reshard_util.is_sharding_opt(optimizer):
+        return state_dict
+    filtered_state_dict = OrderedDict()
+    if reshard_util.get_sharding_strategy(optimizer) == reshard_util.SHARDING_STRATEGY_V1:
+        from paddle.distributed.fleet.meta_optimizers.dygraph_optimizer.dygraph_sharding_optimizer import (
+            DygraphShardingOptimizer,
+        )
+
+        optimizer = unwrap_optimizer(optimizer, DygraphShardingOptimizer)
+        for (k, v) in state_dict.items():
+            if v.name in optimizer._param2rank:
+                sharded_rank = optimizer._param2rank[v.name]
+                if sharded_rank != sharding_rank:
+                    continue
+                filtered_state_dict[k] = v
+            else:
+                if sharding_rank == 0:
+                    filtered_state_dict[k] = v
+    else:
+        from paddle.distributed.fleet.meta_optimizers.dygraph_optimizer.dygraph_sharding_optimizer import (
+            DygraphShardingOptimizerV2,
+        )
+
+        optimizer = unwrap_optimizer(optimizer, DygraphShardingOptimizerV2)
+        parameters = optimizer._parameter_list
+        all_parameters = [p.name for p in parameters]
+        filtered_parameters = [p.name for (i, p) in enumerate(parameters) if i % sharding_world_size == sharding_rank]
+        filtered_parameters = set(filtered_parameters)
+
+        for (k, v) in state_dict.items():
+            if v.name in filtered_parameters:
+                filtered_state_dict[k] = v
+            else:
+                if v.name not in all_parameters and sharding_rank == 0:
+                    filtered_state_dict[k] = v
+
+    return filtered_state_dict
+
+
+def exlclude_paramters_in_state_dict(
+    model_state_dict, param_names_in_master_weights, sharding_group, save_sharding_stage1_model=True
+):
+    assert sharding_group is not None
+    assert isinstance(model_state_dict, dict) and isinstance(
+        param_names_in_master_weights, (list, set)
+    ), "param_names_in_master_weights type:{}".format(type(param_names_in_master_weights))
+    state_param_names = [v.name for k, v in model_state_dict.items()]
+    logger.debug(
+        "param_names_in_master_weights:{}, state_param_names:{}".format(
+            param_names_in_master_weights, state_param_names
+        )
+    )
+    # allgather parameter names in sharding group
+    tmp = []
+    paddle.distributed.all_gather_object(tmp, param_names_in_master_weights, group=sharding_group)
+    param_names_in_master_weights = set([v for item in tmp for v in item])
+    logger.info("sharding_group_param_names:{}".format(param_names_in_master_weights))
+    non_parameters_state_dict = copy.copy(model_state_dict)
+    for k, v in model_state_dict.items():
+        if v.name in param_names_in_master_weights:
+            non_parameters_state_dict.pop(k)
+
+    return non_parameters_state_dict
+
+
 def prune_linear_layer(layer: nn.Linear, index: paddle.Tensor, dim: int = 0) -> nn.Linear:
     """
     Prune a linear layer to keep only entries in index.
@@ -281,6 +363,7 @@ def _find_weight_file_path(
     cache_dir: str,
     model_class: Type[PretrainedModel],
     resource_uri: Optional[str] = None,
+    config: Optional[PretrainedConfig] = None,
 ) -> str | None:
     """find the target weight file under the cache dir, because there are some conflicts about weight file names.
 
@@ -305,7 +388,7 @@ def _find_weight_file_path(
     # 3. find the target weight file name for splited tensor parallel
     # fix for load hybrid parallel
     hybrid_parallel_weight_file_path = os.path.join(
-        cache_dir, _add_variant(resource_weight_file_name, weight_name_suffix())
+        cache_dir, _add_variant(resource_weight_file_name, weight_name_suffix(config))
     )
     if os.path.isfile(hybrid_parallel_weight_file_path):
         return hybrid_parallel_weight_file_path
@@ -954,7 +1037,10 @@ class PretrainedModel(Layer, GenerationMixin, ConversionMixin):
 
             # find the weight file with the above two branch: `bert-base-uncased.pdparams`, `model_state.pdparams`
             weight_file_path = _find_weight_file_path(
-                cache_dir=cache_dir, model_class=cls, resource_uri=pretrained_model_name_or_path
+                cache_dir=cache_dir,
+                model_class=cls,
+                resource_uri=pretrained_model_name_or_path,
+                config=config,
             )
 
             return weight_file_path
@@ -964,7 +1050,7 @@ class PretrainedModel(Layer, GenerationMixin, ConversionMixin):
             # in-order to compatible with old style:
             # file name in pretrained_resouce_file_maps is https://path/to/bert-base-uncased.pdparams, but the registered model-state file name in `resouce_file_maps` is `model_state.pdparams`
 
-            return _find_weight_file_path(cache_dir=pretrained_model_name_or_path, model_class=cls)
+            return _find_weight_file_path(cache_dir=pretrained_model_name_or_path, model_class=cls, config=config)
 
         # 4. download from community or hf-hub
         else:
@@ -1187,6 +1273,7 @@ class PretrainedModel(Layer, GenerationMixin, ConversionMixin):
         import paddlenlp.ops.fast_transformer.transformer.decoding as ft_decoding
 
         state_to_load = ft_decoding.get_ft_para_conf().fit_partial_model(model_to_load, state_dict)
+
         if paddle.in_dynamic_mode():
             model_to_load.set_state_dict(state_to_load)
 
@@ -1414,6 +1501,12 @@ class PretrainedModel(Layer, GenerationMixin, ConversionMixin):
         merge_tensor_parallel = kwargs.get("merge_tensor_parallel", False)
         variant = kwargs.get("variant", None)
         is_main_process = kwargs.get("is_main_process", True)
+        is_bf16 = kwargs.get("is_bf16", False)
+        param_names_in_master_weights = list(kwargs.get("param_names_in_master_weights", []))
+        sharding_group = kwargs.get("sharding_group", None)
+        optimizer = kwargs.get("optimizer", None)
+        save_sharding_stage1_model = kwargs.get("save_sharding_stage1_model", False)
+        use_async_save = kwargs.get("use_async_save", False)
 
         # 1. retrieve the model related config
 
@@ -1426,10 +1519,12 @@ class PretrainedModel(Layer, GenerationMixin, ConversionMixin):
         dtype = get_parameter_dtype(model_to_save)
         model_to_save.config.dtype = str(dtype).split(".")[1]
 
-        state_dict_to_save = None
         config_to_save = copy.deepcopy(model_to_save.config)
+        state_dict_to_save = model_to_save.state_dict()
+        if save_sharding_stage1_model:
+            state_dict_to_save = filter_sharded_params(state_dict_to_save, optimizer, sharding_group)
         if merge_tensor_parallel and config_to_save.tensor_parallel_degree > 1:
-            state_dict_to_save = model_to_save.merge_tensor_parallel(model_to_save.state_dict(), config_to_save)
+            state_dict_to_save = model_to_save.merge_tensor_parallel(state_dict_to_save, config_to_save)
             config_to_save.tensor_parallel_degree = 1
             # set variant to None for merge_tensor_parallel, but there should no relationship with variant setting
             variant = None
@@ -1442,7 +1537,15 @@ class PretrainedModel(Layer, GenerationMixin, ConversionMixin):
                     variant = f"tp{config_to_save.tensor_parallel_rank:0>2d}"
                 # WEIGHTS_NAME = _add_variant(WEIGHTS_NAME, variant)
 
-            state_dict_to_save = self.state_dict()
+        if is_bf16 and save_sharding_stage1_model:
+            state_dict_to_save = exlclude_paramters_in_state_dict(
+                state_dict_to_save, param_names_in_master_weights, sharding_group
+            )
+            logger.info(
+                "param_names_in_master_weights len:{}, bf16 state_dict_to_save len:{}, :{}".format(
+                    len(param_names_in_master_weights), len(state_dict_to_save), state_dict_to_save.keys()
+                )
+            )
 
         if is_main_process:
             # Attach architecture to the config
@@ -1452,7 +1555,10 @@ class PretrainedModel(Layer, GenerationMixin, ConversionMixin):
         # Save model
         if paddle.in_dynamic_mode():
             file_name = os.path.join(save_dir, _add_variant(WEIGHTS_NAME, variant))
-            paddle.save(state_dict_to_save, file_name)
+            if use_async_save:
+                paddle.async_save(state_dict_to_save, file_name)
+            else:
+                paddle.save(state_dict_to_save, file_name)
             del model_to_save
         else:
             logger.warning("Save pretrained model only supported dygraph mode for now!")

paddlenlp/transformers/utils.py

@@ -391,7 +391,7 @@ def optimizer_name_suffix():
         return None
 
 
-def weight_name_suffix():
+def weight_name_suffix(config=None):
     hcg = use_hybrid_parallel()
     if hcg is not None:
         name = []
@@ -399,6 +399,12 @@ def weight_name_suffix():
             name.append(f"tp{hcg.get_model_parallel_rank():0>2d}")
         if hcg.get_pipe_parallel_world_size() > 1:
             name.append(f"pp{hcg.get_stage_id():0>2d}")
+        if config and getattr(config, "moe_num_experts", 0):
+            dp_group = hcg.get_data_parallel_group()
+            name.append(f"moe{dp_group.rank:0>2d}")
         return "_".join(name)
     else:
+        if config and getattr(config, "moe_num_experts", 0):
+            rank = paddle.distributed.get_rank()
+            return f"moe{rank:0>2d}"
         return None