Update trainer.py

* change variable name

* support unified checkpoint

legacy/model_zoo/bert/static/run_glue.py

@@ -159,7 +159,7 @@ def reset_program_state_dict(args, model, state_dict, pretrained_state_dict):
             reset_parameter_names.append(n)
         else:
             dtype_str = "float32"
-            if str(p.dtype) == "VarType.FP64":
+            if p.dtype == paddle.float64:
                 dtype_str = "float64"
             reset_state_dict[p.name] = np.random.normal(loc=0.0, scale=scale, size=p.shape).astype(dtype_str)
     logger.info("the following parameter had reset, please check. {}".format(reset_parameter_names))

legacy/model_zoo/bert/static/run_glue_with_sparaity.py

@@ -159,7 +159,7 @@ def reset_program_state_dict(args, model, state_dict, pretrained_state_dict):
             reset_parameter_names.append(n)
         else:
             dtype_str = "float32"
-            if str(p.dtype) == "VarType.FP64":
+            if p.dtype == paddle.float64:
                 dtype_str = "float64"
             reset_state_dict[p.name] = np.random.normal(loc=0.0, scale=scale, size=p.shape).astype(dtype_str)
     logger.info("the following parameter had reset, please check. {}".format(reset_parameter_names))

legacy/model_zoo/bert/static/run_pretrain.py

@@ -151,7 +151,7 @@ def reset_program_state_dict(model, state_dict):
     for n, p in state_dict.items():
         if "layer_norm" not in p.name:
             dtype_str = "float32"
-            if str(p.dtype) == "VarType.FP64":
+            if p.dtype == paddle.float64:
                 dtype_str = "float64"
             new_state_dict[p.name] = np.random.normal(loc=0.0, scale=scale, size=p.shape).astype(dtype_str)
     return new_state_dict

legacy/model_zoo/gpt-3/ppfleetx/optims/optimizer.py

@@ -98,7 +98,7 @@ class FusedOffloadAdamW(paddle.optimizer.AdamW):
     def _add_moments_pows(self, p):
         acc_dtype = p.dtype
         if self._is_dtype_fp16_or_bf16(acc_dtype):
-            acc_dtype = core.VarDesc.VarType.FP32
+            acc_dtype = paddle.float32
         self._add_accumulator(self._moment1_acc_str, p, dtype=acc_dtype, device="cpu")
         self._add_accumulator(self._moment2_acc_str, p, dtype=acc_dtype, device="cpu")
         self._add_accumulator(

legacy/model_zoo/moe/dygraph/framework/group_sharded.py

@@ -37,7 +37,6 @@ from paddle.distributed.fleet.meta_parallel.sharding.group_sharded_optimizer_sta
 from paddle.distributed.fleet.meta_parallel.sharding.group_sharded_stage2 import (
     GroupShardedStage2,
 )
-from paddle.framework import core
 from paddle.incubate.distributed.models.moe.grad_clip import ClipGradForMOEByGlobalNorm
 from paddle.optimizer import Optimizer
 
@@ -99,7 +98,7 @@ class ClipGradForShardedMOEByGlobalNorm(ClipGradForMOEByGlobalNorm):
                 params_and_grads.append((p, g))
                 continue
             # TODO(wangxi): use inplace elementwise_mul
-            clip_input = clip_var.astype("float16") if g.dtype == core.VarDesc.VarType.FP16 else clip_var
+            clip_input = clip_var.astype("float16") if g.dtype == paddle.float16 else clip_var
             new_grad = paddle.multiply(x=g, y=clip_input)
             params_and_grads.append((p, new_grad))
         return params_and_grads

legacy/model_zoo/moe/dygraph/run_moe_pretrain.py

@@ -37,7 +37,6 @@ from paddle.distributed.fleet.meta_parallel import get_rng_state_tracker
 from paddle.distributed.fleet.meta_parallel.sharding.group_sharded_utils import (
     GroupShardedScaler,
 )
-from paddle.framework import core
 from paddle.incubate.distributed.models import moe
 from utils import get_timers, set_timers
 from visualdl import LogWriter
@@ -158,12 +157,8 @@ def initialize_mp_dp_parameters(model, hcg):
 def unscale_method(self, optimizer):
     if not self._enable:
         return
-    if paddle.framework.use_pir_api():
-        type_float16 = core.DataType.FLOAT16
-        type_float32 = core.DataType.FLOAT32
-    else:
-        type_float16 = core.VarDesc.VarType.FP16
-        type_float32 = core.VarDesc.VarType.FP32
+    type_float16 = paddle.float16
+    type_float32 = paddle.float32
 
     if getattr(optimizer, "_param_groups", None) and isinstance(optimizer._param_groups[0], dict):
         param_grads_fp16 = []

llm/alignment/dpo/dpo_argument.py

@@ -51,6 +51,15 @@ class DPOTrainingArguments(TrainingArguments):
         default=False,
         metadata={"help": "Whether to run benchmark by autotuner. True for from_scratch."},
     )
+    continue_training: bool = field(
+        default=True,
+        metadata={
+            "help": (
+                "Whether to train from existing paddlenlp model weights.\n"
+                "If set True, the model_name_or_path argument must exist in the paddlenlp models."
+            )
+        },
+    )
 
     def __post_init__(self):
         super().__post_init__()

llm/alignment/dpo/run_dpo.py

@@ -117,7 +117,9 @@ def main():
         model_class = AutoModelForCausalLMPipe
     else:
         model_class = AutoModelForCausalLM
-    if not training_args.autotuner_benchmark or model_args.weight_quantize_algo is not None:
+    if training_args.continue_training and (
+        not training_args.autotuner_benchmark or model_args.weight_quantize_algo is not None
+    ):
         model = model_class.from_pretrained(**model_kwargs)
         # for DPO save
         model.config.dpo_config = None
@@ -131,6 +133,9 @@ def main():
         config = AutoConfig.from_pretrained(**model_kwargs)
         model = model_class.from_config(config, dtype=dtype)
         if not dpo_config.reference_free and not dpo_config.lora:
+            # NOTE(gongenlei): architectures will be None after from_config,
+            # so, we need from_pretrained to create new config.
+            config = AutoConfig.from_pretrained(**model_kwargs)
             ref_model = model_class.from_config(config, dtype=dtype)
         else:
             ref_model = None

llm/run_finetune.py

@@ -41,7 +41,7 @@ from paddlenlp.peft import (
     VeRAConfig,
     VeRAModel,
 )
-from paddlenlp.trainer import PdArgumentParser, get_last_checkpoint
+from paddlenlp.trainer import PdArgumentParser, get_last_checkpoint, set_seed
 from paddlenlp.trainer.trainer_callback import TrainerState
 from paddlenlp.transformers import (
     AutoConfig,
@@ -91,6 +91,7 @@ def main():
 
     # Setup GPU & distributed training
     paddle.set_device(training_args.device)
+    set_seed(seed=training_args.seed)
     logger.warning(
         f"Process rank: {training_args.local_rank}, device: {training_args.device}, world_size: {training_args.world_size}, "
         + f"distributed training: {bool(training_args.local_rank != -1)}, 16-bits training: {training_args.fp16 or training_args.bf16}"

llm/tools/preprocess/README.md

@@ -126,12 +126,9 @@ head -1 baike_sample.jsonl
 ```
 optional arguments:
   -h, --help            show this help message and exit
-  --model_name MODEL_NAME
+  --model_name_or_path MODEL_NAME_OR_PATH
                         What model to use.
                         必须设置，如：idea-ccnl/ziya-llama-13b-v1, 可以参考已有的模型名称 https://github.com/PaddlePaddle/PaddleNLP/blob/develop/llm
-  --tokenizer_name {LlamaTokenizer}
-                        What type of tokenizer to use.
-                        模型对应的tokenizer, Llama模型需使用LlamaTokenizer
 data input/output:
   --input_path INPUT_PATH
                         Path to input JSON files.
@@ -183,8 +180,7 @@ common config:
 * 针对 llama 模型
 ```shell
 python -u  create_pretraining_data.py \
-    --model_name "idea-ccnl/ziya-llama-13b-v1" \
-    --tokenizer_name "LlamaTokenizer" \
+    --model_name_or_path "idea-ccnl/ziya-llama-13b-v1" \
     --input_path "baike_sample.jsonl" \
     --output_prefix "baike_sample"  \
     --data_format "JSON" \
@@ -199,8 +195,7 @@ python -u  create_pretraining_data.py \
 * 针对 ernie 模型
 ```shell
 python -u  create_pretraining_data.py \
-    --model_name "ernie-3.0-base-zh" \
-    --tokenizer_name "ErnieTokenizer" \
+    --model_name_or_path "ernie-3.0-base-zh" \
     --input_path "baike_sample.jsonl" \
     --output_prefix "baike_sample"  \
     --data_format "JSON" \

llm/tools/preprocess/create_pretraining_data.py

@@ -23,8 +23,8 @@ import time
 import numpy as np
 from tqdm import tqdm
 
-import paddlenlp.transformers as tfs
 from paddlenlp.data import indexed_dataset
+from paddlenlp.transformers import AutoTokenizer
 from paddlenlp.utils.log import logger
 
 try:
@@ -44,23 +44,7 @@ def print_datetime(string):
 
 def get_args():
     parser = argparse.ArgumentParser()
-    parser.add_argument("--model_name", type=str, required=True, help="What model to use.")
-    parser.add_argument(
-        "--tokenizer_name",
-        type=str,
-        required=True,
-        choices=[
-            "ErnieTokenizer",
-            "BertTokenizer",
-            "GPTTokenizer",
-            "GPTChineseTokenizer",
-            "LlamaTokenizer",
-            "ElectraTokenizer",
-            "T5Tokenizer",
-            "Qwen2Tokenizer"
-        ],
-        help="What type of tokenizer to use.",
-    )
+    parser.add_argument("--model_name_or_path", type=str, required=True, help="What model to use.")
     group = parser.add_argument_group(title="data input/output")
     group.add_argument("--input_path", type=str, required=True, help="Path to input JSON files.")
     group.add_argument("--output_prefix", type=str, required=True, help="Output prefix to store output file.")
@@ -227,7 +211,7 @@ class Converter(object):
         self.args = args
 
     def initializer(self):
-        Converter.tokenizer = getattr(tfs, self.args.tokenizer_name).from_pretrained(self.args.model_name)
+        Converter.tokenizer = AutoTokenizer.from_pretrained(self.args.model_name_or_path)
         if self.args.cn_whole_word_segment:
             # Extend chinese char vocab for ErnieTokinzer
             Converter.tokenizer.extend_chinese_char()
@@ -333,7 +317,7 @@ def main():
     convert = Converter(args)
 
     # Try tokenizer is availiable
-    sample_tokenizer = getattr(tfs, args.tokenizer_name).from_pretrained(args.model_name)
+    sample_tokenizer = AutoTokenizer.from_pretrained(args.model_name_or_path)
     if sample_tokenizer.vocab_size < 2**16 - 1:
         save_dtype = np.uint16
     else:

paddlenlp/__init__.py

@@ -20,7 +20,7 @@ PADDLENLP_STABLE_VERSION = "PADDLENLP_STABLE_VERSION"
 
 # this version is used for develop and test.
 # release version will be added fixed version by setup.py.
-__version__ = "3.0.0b1.post"
+__version__ = "3.0.0b2.post"
 if os.getenv(PADDLENLP_STABLE_VERSION):
     __version__ = __version__.replace(".post", "")
 else:
@@ -48,6 +48,7 @@ from . import (
     experimental,
     layers,
     losses,
+    mergekit,
     metrics,
     ops,
     peft,

paddlenlp/data/dist_dataloader.py

@@ -66,6 +66,7 @@ class DistDataLoader(paddle.io.DataLoader):
 
         eval = kwargs.pop("eval", False)
         is_iterable_dataset = kwargs.pop("is_iterable_dataset", False)
+        self._pp_data_group = kwargs.pop("pp_data_group", None)
 
         if dataset is None:
             dataset = DummyDataset() if not is_iterable_dataset else IterableDummyDataset()
@@ -78,10 +79,8 @@ class DistDataLoader(paddle.io.DataLoader):
 
         # Init pp data comm group.
         if self._hcg.get_pipe_parallel_world_size() > 1:
-            self._pp_data_group = self._init_dataloader_comm_group()
             self._pp_group = self._hcg.get_pipe_parallel_group()
         else:
-            self._pp_data_group = None
             self._pp_group = None
 
         self.mp_group = self._hcg.get_model_parallel_group()
@@ -132,18 +131,6 @@ class DistDataLoader(paddle.io.DataLoader):
         else:
             raise ValueError("raise error for `paddlenlp.trainer.trainer_utils.has_length`")
 
-    def _init_dataloader_comm_group(self):
-        topo = self._hcg._topo
-        parallel_comm_group = None
-        parallel_groups = topo.get_comm_list("pipe")
-
-        for group in parallel_groups:
-            ranks = [group[0], group[-1]]
-            comm_group = paddle.distributed.new_group(ranks=ranks)
-            if paddle.distributed.get_rank() in ranks:
-                parallel_comm_group = comm_group
-        return parallel_comm_group
-
     def __iter__(self):
         return self
 
@@ -212,3 +199,17 @@ class DistDataLoader(paddle.io.DataLoader):
                 logger.debug(e)
         data = self._broadcast_data(data)
         return data
+
+
+def init_dataloader_comm_group():
+    hcg = fleet.get_hybrid_communicate_group()
+    topo = hcg._topo
+    parallel_groups = topo.get_comm_list("pipe")
+    parallel_comm_group = None
+
+    for group in parallel_groups:
+        ranks = [group[0], group[-1]]
+        comm_group = paddle.distributed.new_group(ranks=ranks)
+        if paddle.distributed.get_rank() in ranks:
+            parallel_comm_group = comm_group
+    return parallel_comm_group

paddlenlp/mergekit/__init__.py

@@ -0,0 +1,19 @@
+# Copyright (c) 2024 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you smay 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 .merge_config import *
+from .merge_method import *
+from .merge_model import *
+from .merge_utils import *
+from .sparsify_method import *

paddlenlp/mergekit/merge_config.py

@@ -0,0 +1,168 @@
+# Copyright (c) 2024 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 json
+import os
+from dataclasses import asdict, dataclass, field
+from typing import List, Optional
+
+import paddle
+
+from paddlenlp.utils.env import MERGE_CONFIG_NAME
+from paddlenlp.utils.log import logger
+
+
+@dataclass
+class MergeConfig:
+    """
+    This is the configuration class to store the configuration of a [`MergeKit`].
+    """
+
+    # Common parameters
+    device: str = field(default="cpu", metadata={"help": "Device to use for the merge.ex cpu、 gpu、low_gpu_mem"})
+    tensor_type: str = field(
+        default="np", metadata={"help": "Tensor type to use for the merge. Choose np(CPU Only) or pd (CPU/GPU)"}
+    )
+    n_process: int = field(default=1, metadata={"help": "Number of processes to use for the merge."})
+    merge_preifx: str = field(default="model", metadata={"help": "Prefix name: model or master_weights"})
+    merge_method: str = field(default="linear", metadata={"help": "The merge strategy."})
+    merge_type: str = field(default="linear", metadata={"help": "The type of merge process."})
+    sparsify_type: str = field(default=None, metadata={"help": "The type of sparsify process."})
+
+    # Model parameters
+    model_path_list: Optional[List[str]] = field(default=None, metadata={"help": "Merge model name or path list"})
+    model_path_str: Optional[str] = field(
+        default=None, metadata={"help": "Merge model name or path string.(split by ',')"}
+    )
+    base_model_path: str = field(default=None, metadata={"help": "Base model name or path."})
+    output_path: str = field(default=None, metadata={"help": "Base model name or path."})
+    # merge parameters
+    weight_list: Optional[List[float]] = field(
+        default=None, metadata={"help": "Relative (or absolute if normalize=False) weighting of a given tensor"}
+    )
+    normalize: bool = field(default=True, metadata={"help": "Whether to normalize the weighting."})
+    slerp_alpha: float = field(default=0.5, metadata={"help": "Slerp alpha."})
+    slerp_normalize_eps: float = field(default=1e-8, metadata={"help": "Slerp normalization epsilon value"})
+    slerp_dot_threshold: float = field(
+        default=0.9995,
+        metadata={
+            "help": "Slerp dot threshold. If dot value exceeds this threshold, then we consider them as colinear, so use linear instead."
+        },
+    )
+    ties_elect_type: str = field(default="sum", metadata={"help": "The type of ties mask. 'sum' or 'count'"})
+
+    # Sparsify parameters
+    rescale: bool = field(default=True, metadata={"help": "Rescale the weights after sparsifying."})
+    reserve_p: float = field(default=0.7, metadata={"help": "Random reserve probability for the sparsify model."})
+    epsilon: float = field(default=0.14, metadata={"help": "Epsilon value for magprune."})
+
+    def __post_init__(self):
+        self.config_check()
+
+    def config_check(self):
+        if self.output_path is not None:
+            os.makedirs(self.output_path, exist_ok=True)
+        if self.tensor_type not in ["np"]:
+            raise ValueError(f"Unsupported tensor type: {self.tensor_type}. Support 'np' only.")
+        if self.device != "cpu":
+            logger.warning(f"Currently only support cpu device, but got {self.device}. Setting `device` to `cpu`.")
+            self.device = "cpu"
+            self.tensor_type = "np"
+
+        elif self.merge_method not in ["linear", "ties", "slerp", "della_linear", "della", "dare_linear", "dare_ties"]:
+            raise ValueError(
+                f"Unsupported merge strategy: {self.merge_method}. Please choose one from ['linear', 'slerp']."
+            )
+        if self.model_path_str is not None:
+            self.model_path_list = self.model_path_str.split(",")
+        if self.model_path_list is not None:
+            if not isinstance(self.model_path_list, list) or len(self.model_path_list) < 2:
+                raise ValueError(f"Please specify the model_path_list at least two. But got {self.model_path_list}")
+            if self.weight_list is None:
+                self.weight_list = [1.0] * len(self.model_path_list)
+                self.normalize = True
+            if len(self.model_path_list) != len(self.weight_list):
+                raise ValueError("The length of model_path_list and weight_list must be the same.")
+        if self.reserve_p < 0 or self.reserve_p > 1:
+            raise ValueError("reserve_p must be between 0 and 1.")
+        if "della" in self.merge_method or self.sparsify_type == "magprune":
+            if self.reserve_p <= self.epsilon / 2 or self.reserve_p >= (1 - self.epsilon):
+                raise ValueError(
+                    f"Error: reserve_p +- epsilon/2 must be in the range (0, 1). reserve_p + epsilon/2 = {self.reserve_p + self.epsilon / 2 }, reserve_p - epsilon/2 = {self.reserve_p - self.epsilon / 2 }"
+                )
+        paddle.set_device(self.device)
+
+    @property
+    def __dict__(self):
+        return asdict(self)
+
+    def to_dict(self):
+        return self.__dict__
+
+    def save_pretrained(self, save_directory):
+        r"""
+        This method saves the configuration of your adapter model in a directory.
+        Args:
+            save_directory (`str`):
+                The directory where the configuration will be saved.
+        """
+        if os.path.isfile(save_directory):
+            raise AssertionError(f"Provided path ({save_directory}) should be a directory, not a file")
+
+        os.makedirs(save_directory, exist_ok=True)
+
+        output_dict = self.__dict__
+        output_path = os.path.join(save_directory, MERGE_CONFIG_NAME)
+
+        # save it
+        with open(output_path, "w") as writer:
+            writer.write(json.dumps(output_dict, indent=2, sort_keys=True))
+
+    @classmethod
+    def from_pretrained(cls, pretrained_model_path, **kwargs):
+        r"""
+        This method loads the configuration of your adapter model from a directory.
+        Args:
+            pretrained_model_path (`str`):
+                The directory or the hub-id where the configuration is saved.
+            **kwargs:
+                Additional keyword arguments passed along to the child class initialization.
+        """
+        if os.path.isfile(os.path.join(pretrained_model_path, MERGE_CONFIG_NAME)):
+            config_file = os.path.join(pretrained_model_path, MERGE_CONFIG_NAME)
+        else:
+            raise ValueError(f"Can't find merge_config.json at '{pretrained_model_path}'")
+
+        loaded_attributes = cls.from_json_file(config_file)
+
+        config = cls(**kwargs)
+
+        for key, value in loaded_attributes.items():
+            if hasattr(config, key):
+                setattr(config, key, value)
+
+        return config
+
+    @classmethod
+    def from_json_file(cls, path_json_file):
+        r"""
+        Loads a configuration file from a json file.
+        Args:
+            path_json_file (`str`):
+                The path to the json file.
+        """
+        with open(path_json_file, "r") as file:
+            json_object = json.load(file)
+
+        return json_object

paddlenlp/mergekit/merge_method.py

@@ -0,0 +1,130 @@
+# Copyright (c) 2024 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
+
+
+class MergeMethod:
+    def __init__(self, merge_config, sparsify_method=None):
+        self.merge_config = merge_config
+        self.sparsify_method = sparsify_method
+
+    def merge(self, tensor_list):
+        if self.sparsify_method is not None:
+            tensor_list = [self.sparsify_method.sparsify(tensor) for tensor in tensor_list]
+        if self.merge_config.merge_type == "linear":
+            return self.linear(tensor_list)
+        elif self.merge_config.merge_type == "slerp":
+            return self.slerp(tensor_list)
+        elif self.merge_config.merge_type == "ties":
+            return self.ties(tensor_list)
+        else:
+            raise NotImplementedError(f"{self.merge_config.merge_type} is not supported yet.")
+
+    def linear(self, tensor_list):
+        """
+        Linear interpolation between multiple values.
+        """
+        # init weight
+        weight_list = self.merge_config.weight_list
+        if self.merge_config.normalize:
+            weight_sum = sum(weight_list)
+            weight_list = [weight / weight_sum for weight in weight_list]
+
+        # merge
+        if self.merge_config.tensor_type == "np":
+            tensor_output = sum(weight * tensor for weight, tensor in zip(weight_list, tensor_list))
+            return tensor_output
+        else:
+            raise NotImplementedError("Paddle Tensor is not supported yet.")
+
+    def slerp(self, tensor_list):
+        """
+        Spherical linear interpolation
+        """
+        # check tensor_list length
+        if len(tensor_list) != 2:
+            raise ValueError("Slerp only support two tensors merge.")
+
+        if self.merge_config.tensor_type == "np":
+            t0, t1 = tensor_list
+            # Copy the vectors to reuse them later
+            t0_copy = np.copy(t0)
+            t1_copy = np.copy(t1)
+
+            # Normalize the vectors to get the directions and angles
+            t0 = self.normalize(t0)
+            t1 = self.normalize(t1)
+
+            # Dot product with the normalized vectors (can't use np.dot in W)
+            dot = np.sum(t0 * t1)
+            # If absolute value of dot product is almost 1, vectors are ~colinear, so use lerp
+            if np.abs(dot) > self.merge_config.slerp_dot_threshold:
+                return (1 - self.merge_config.slerp_alpha) * t0_copy + self.merge_config.slerp_alpha * t1_copy
+
+            # Calculate initial angle between t0 and t1
+            theta_0 = np.arccos(dot)
+            sin_theta_0 = np.sin(theta_0)
+
+            # Angle at timestep t
+            theta_t = theta_0 * self.merge_config.slerp_alpha
+            sin_theta_t = np.sin(theta_t)
+
+            # Finish the slerp algorithm
+            s0 = np.sin(theta_0 - theta_t) / sin_theta_0
+            s1 = sin_theta_t / sin_theta_0
+
+            return s0 * t0_copy + s1 * t1_copy
+        else:
+            raise NotImplementedError("Paddle Tensor is not supported yet.")
+
+    def ties(self, tensor_list):
+        if self.merge_config.tensor_type == "np":
+            # Get weight tensor
+            mask_dtype = tensor_list[0].dtype
+            weight_list = self.merge_config.weight_list
+            tensor_list = [weight * tensor for (weight, tensor) in zip(weight_list, tensor_list)]
+
+            # Elect majority sign
+            sign_tensor_list = [np.sign(tensor).astype(mask_dtype) for tensor in tensor_list]
+            if self.merge_config.ties_elect_type == "sum":
+                majority_sign = (np.sum(tensor_list, axis=0) >= 0).astype(mask_dtype) * 2 - 1
+            elif self.merge_config.ties_elect_type == "count":
+                majority_sign = (np.sum(sign_tensor_list, axis=0) >= 0).astype(mask_dtype) * 2 - 1
+            else:
+                raise NotImplementedError(f"ties_elect_type: {self.merge_config.ties_elect_type} is unknown.")
+
+            # Merge
+            mask_list = [sign_tensor == majority_sign for sign_tensor in sign_tensor_list]
+            tensor_list = [mask * tensor for mask, tensor in zip(mask_list, tensor_list)]
+            merge_tensor = np.sum(tensor_list, axis=0)
+
+            # Normalize
+            if self.merge_config.normalize:
+                weight_mask = [mask * weight for mask, weight in zip(mask_list, weight_list)]
+                divisor = np.sum(weight_mask, axis=0)
+                divisor[np.abs(divisor) < 1e-8] = 1
+                merge_tensor /= divisor
+            return merge_tensor
+        else:
+            raise NotImplementedError("Paddle Tensor is not supported yet.")
+
+    def normalize(self, t):
+        """
+        Normalize a vector by its L2 norm.
+        """
+        norm_t = np.linalg.norm(t)
+        if norm_t > self.merge_config.slerp_normalize_eps:
+            t = t / norm_t
+        return t

paddlenlp/mergekit/merge_model.py

@@ -0,0 +1,337 @@
+# Copyright (c) 2024 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 gc
+import json
+import os
+from multiprocessing import Process
+
+import numpy as np
+import paddle
+from safetensors import safe_open
+from safetensors.numpy import save_file
+
+from paddlenlp.utils.env import (
+    PADDLE_MASTER_WEIGHTS_NAME,
+    PADDLE_WEIGHTS_NAME,
+    SAFE_MASTER_WEIGHTS_INDEX_NAME,
+    SAFE_MASTER_WEIGHTS_NAME,
+    SAFE_WEIGHTS_INDEX_NAME,
+    SAFE_WEIGHTS_NAME,
+)
+from paddlenlp.utils.safetensors import fast_safe_open
+
+from .merge_method import MergeMethod
+from .merge_utils import divide_positions
+from .sparsify_method import SparsifyMethod
+
+SPARSIFY_MERGE_MAPPING = {
+    "linear": (None, "linear"),
+    "ties": ("trim", "ties"),
+    "slerp": (None, "slerp"),
+    "della_linear": ("magprune", "linear"),
+    "della": ("magprune", "ties"),
+    "dare_linear": ("dare", "linear"),
+    "dare_ties": ("dare", "ties"),
+}
+
+
+class MergeModel:
+    def __init__(self, merge_config):
+        self.reset_merge_model(merge_config=merge_config)
+        self.numpy_dtype_map = {"float32": 4, "float16": 2, "uint16": 2}
+
+    def reset_merge_model(self, merge_config=None, merge_param_dict=None):
+        if merge_config is not None:
+            self.merge_config = merge_config
+        elif merge_param_dict is not None:
+            for k, v in merge_param_dict.items():
+                setattr(self.merge_config, k, v)
+            self.merge_config.config_check()
+
+        # map sparsify & merge type
+        if self.merge_config.merge_method is not None:
+            self.merge_config.sparsify_type, self.merge_config.merge_type = SPARSIFY_MERGE_MAPPING[
+                self.merge_config.merge_method
+            ]
+        if self.merge_config.merge_type is None:
+            raise ValueError("Either merge_type or merge_method must be specified.")
+        # init merge method
+        sparsify_method = SparsifyMethod(self.merge_config)
+        self.merge_method = MergeMethod(merge_config, sparsify_method)
+
+    def merge_model(self):
+        file_type_list = []
+        for model_path in self.merge_config.model_path_list:
+            file_type_list.append(self.check_model_path(model_path))
+        if self.merge_config.base_model_path is not None:
+            file_type_list.append(self.check_model_path(self.merge_config.base_model_path))
+        if all(file_type == "safetensors" or file_type == "safetensors_without_index" for file_type in file_type_list):
+            self.merge_safetensor_model(file_type_list)
+        else:
+            self.merge_mix_model(file_type_list)
+
+    def merge_mix_model(self, file_type_list):
+        state_dict_list = []
+        for i, model_path in enumerate(self.merge_config.model_path_list):
+            state_dict_list.append(self.get_model_state_dict(model_path, file_type_list[i]))
+        if self.merge_config.base_model_path is not None:
+            state_dict_list.append(self.get_model_state_dict(self.merge_config.base_model_path, file_type_list[-1]))
+        if not all(state_dict_list[0].keys() == state_dict.keys() for state_dict in state_dict_list):
+            raise ValueError("State dict keys mismatch. Please make sure you load the correct weight file")
+        if self.merge_config.base_model_path is not None:
+            base_state_dict = state_dict_list.pop()
+            base_file_type = file_type_list.pop()
+        merge_state_dict = {}
+        total_size = 0
+        weight_map = {}
+        for key in state_dict_list[0].keys():
+            is_bf16 = False
+            tensor_list = []
+            for state_dict, file_type in zip(state_dict_list, file_type_list):
+                if file_type == "pdparams":
+                    if str(state_dict[key].dtype) == "paddle.bfloat16":
+                        is_bf16 = True
+                        state_dict[key] = state_dict[key].astype("float32").numpy()
+                    else:
+                        state_dict[key] = state_dict[key].numpy()
+                elif str(state_dict[key].dtype) == "uint16":
+                    is_bf16 = True
+                    state_dict[key] = paddle.to_tensor(state_dict[key], dtype="bfloat16").astype("float32").numpy()
+                tensor_list.append(state_dict[key])
+            if self.merge_config.base_model_path is not None:
+                if base_file_type == "pdparams":
+                    if str(base_state_dict[key].dtype) == "paddle.bfloat16":
+                        base_state_dict[key] = base_state_dict[key].astype("float32").numpy()
+                    else:
+                        base_state_dict[key] = base_state_dict[key].numpy()
+                elif str(base_state_dict[key].dtype) == "uint16":
+                    base_state_dict[key] = (
+                        paddle.to_tensor(base_state_dict[key], dtype="bfloat16").astype("float32").numpy()
+                    )
+                tensor_list = [tensor - base_state_dict[key] for tensor in tensor_list]
+            merge_state_dict[key] = self.merge_method.merge(tensor_list)
+            if self.merge_config.base_model_path is not None:
+                merge_state_dict[key] += base_state_dict[key]
+            # dtype==bfloat16: numpy(float32) -> paddle(float32) -> paddle(bfloat16) -> numpy(uint16)
+            if is_bf16:
+                merge_state_dict[key] = (
+                    paddle.to_tensor(merge_state_dict[key], dtype="float32").astype("bfloat16").numpy()
+                )
+            total_size += np.prod(merge_state_dict[key].shape) * self.numpy_dtype_map[str(merge_state_dict[key].dtype)]
+            weight_map[key] = f"{self.merge_config.merge_preifx}-00001-of-00001.safetensors"
+        # save safetensor file
+        save_file(
+            merge_state_dict,
+            os.path.join(
+                self.merge_config.output_path, f"{self.merge_config.merge_preifx}-00001-of-00001.safetensors"
+            ),
+            metadata={"format": "np"},
+        )
+        # save safe index file
+        index = {"metadata": {"total_size": int(total_size)}, "weight_map": weight_map}
+        save_index_file = os.path.join(self.merge_config.output_path, self.safe_index_name())
+        with open(save_index_file, "w", encoding="utf-8") as f:
+            content = json.dumps(index, indent=2) + "\n"
+            f.write(content)
+        # save merge config file
+        self.merge_config.save_pretrained(self.merge_config.output_path)
+        del state_dict_list
+        del merge_state_dict
+        if self.merge_config.base_model_path is not None:
+            del base_state_dict
+        gc.collect()
+
+    def get_model_state_dict(self, model_path, file_type):
+        if file_type == "safetensors":
+            state_dict = {}
+            with open(os.path.join(model_path, self.safe_index_name()), "r", encoding="utf-8") as f:
+                index = json.load(f)
+            for key in index["weight_map"].keys():
+                with fast_safe_open(
+                    os.path.join(model_path, index["weight_map"][key]),
+                    framework="np",
+                ) as f:
+                    state_dict[key] = f.get_tensor(key)
+        elif file_type == "safetensors_without_index":
+            state_dict = {}
+            with fast_safe_open(os.path.join(model_path, self.safe_weight_name()), framework="numpy") as f:
+                for k in f.keys():
+                    state_dict[k] = f.get_tensor(k)
+        elif file_type == "pdparams":
+            state_dict = paddle.load(os.path.join(model_path, self.weight_name()))
+        else:
+            raise ValueError(f"Unsupported file_type: {file_type}")
+        return state_dict
+
+    def create_safetensor_index(self, model_path):
+        weight_map = {}
+        total_size = 0
+
+        with safe_open(os.path.join(model_path, self.safe_weight_name()), framework="numpy") as f:
+            for key in f.keys():
+                tensor = f.get_tensor(key)
+                total_size += np.prod(tensor.shape) * self.numpy_dtype_map[str(tensor.dtype)]
+                weight_map[key] = self.safe_weight_name()
+        index = {"metadata": {"total_size": total_size}, "weight_map": weight_map}
+        return index
+
+    def merge_safetensor_model(self, file_type_list):
+        # load index
+        index_list = []
+        model_path_list = self.merge_config.model_path_list
+        if self.merge_config.base_model_path is not None:
+            model_path_list += [self.merge_config.base_model_path]
+
+        for model_path, file_type in zip(model_path_list, file_type_list):
+            if file_type == "safetensors":
+                with open(os.path.join(model_path, self.safe_index_name()), "r", encoding="utf-8") as f:
+                    index_list.append(json.load(f))
+            else:
+                index = self.create_safetensor_index(model_path)
+                index_list.append(index)
+        # check index
+        if not all(index_list[0]["metadata"]["total_size"] == index["metadata"]["total_size"] for index in index_list):
+            raise ValueError("Weights total_size mismatch. Please make sure you load the correct weight file")
+        if not all(index_list[0]["weight_map"].keys() == index["weight_map"].keys() for index in index_list):
+            raise ValueError("Weights weight_map mismatch. Please make sure you load the correct weight file")
+        # init new index
+        index = {}
+        index["metadata"] = index_list[0]["metadata"]
+        index["metadata"]["total_size"] = int(index["metadata"]["total_size"])
+        index["weight_map"] = {}
+
+        # Multi-process update
+        key_list = list(index_list[0]["weight_map"].keys())
+        positions = divide_positions(len(key_list), self.merge_config.n_process)
+        threads = []
+        if self.merge_config.tensor_type == "np":
+            target = self.shard_merge_np
+        else:
+            target = self.shard_merge_pd
+        for i in range(len(positions) - 1):
+            shard_file = f"{self.merge_config.merge_preifx}-{i+1:05d}-of-{self.merge_config.n_process:05d}.safetensors"
+            t = Process(
+                target=target,
+                args=(
+                    key_list[positions[i] : positions[i + 1]],  # key_list
+                    index_list,  # index_list
+                    shard_file,  # shard_file name
+                ),
+            )
+            threads.append(t)
+            for k in key_list[positions[i] : positions[i + 1]]:
+                index["weight_map"][k] = shard_file
+        for t in threads:
+            t.start()
+        for t in threads:
+            t.join()
+
+        # save safe index file
+        save_index_file = os.path.join(self.merge_config.output_path, self.safe_index_name())
+        with open(save_index_file, "w", encoding="utf-8") as f:
+            content = json.dumps(index, indent=2) + "\n"
+            f.write(content)
+        self.merge_config.save_pretrained(self.merge_config.output_path)
+
+    def shard_merge_np(
+        self,
+        key_list,
+        index_list,
+        shard_file,
+    ):
+        merge_state_dict = {}
+        for k in key_list:
+            tensor_list = []
+
+            for i, model_path in enumerate(self.merge_config.model_path_list):
+                with fast_safe_open(os.path.join(model_path, index_list[i]["weight_map"][k]), framework="np") as w:
+                    tensor = w.get_tensor(k)
+                    dtype = tensor.dtype
+                    # dtype==bfloat16: numpy(uint16) -> paddle(bfloat16) -> paddle(float32) -> numpy(float32)
+                    if tensor.dtype == np.uint16:
+                        tensor = paddle.to_tensor(tensor, dtype="bfloat16").astype("float32").numpy()
+                    tensor_list.append(tensor)
+            if self.merge_config.base_model_path is not None:
+                with fast_safe_open(
+                    os.path.join(self.merge_config.base_model_path, index_list[-1]["weight_map"][k]),
+                    framework="np",
+                ) as w:
+                    base_tensor = w.get_tensor(k)
+                    if base_tensor.dtype == np.uint16:
+                        base_tensor = paddle.to_tensor(base_tensor, dtype="bfloat16").astype("float32").numpy()
+                tensor_list = [tensor - base_tensor for tensor in tensor_list]
+            merge_state_dict[k] = self.merge_method.merge(tensor_list)
+            if self.merge_config.base_model_path is not None:
+                merge_state_dict[k] += base_tensor
+            # dtype==bfloat16: numpy(float32) -> paddle(float32) -> paddle(bfloat16) -> numpy(uint16)
+            if dtype == np.uint16:
+                merge_state_dict[k] = paddle.to_tensor(merge_state_dict[k], dtype="float32").astype("bfloat16").numpy()
+            del tensor_list
+            if self.merge_config.base_model_path is not None:
+
+                del base_tensor
+        save_file(
+            merge_state_dict,
+            os.path.join(self.merge_config.output_path, shard_file),
+            metadata={"format": "np"},
+        )
+        del merge_state_dict
+        gc.collect()
+
+    def shard_merge_pd(
+        self,
+        key_list,
+        index_list,
+        shard_file,
+    ):
+        raise NotImplementedError("Not support paddle tensors.")
+
+    def check_model_path(self, model_path):
+        if os.path.exists(os.path.join(model_path, self.safe_index_name())):
+            with open(os.path.join(model_path, self.safe_index_name()), "r", encoding="utf-8") as f:
+                index = json.load(f)
+                safe_file_list = list(set(index["weight_map"][k] for k in index["weight_map"]))
+                for i in range(len(safe_file_list)):
+                    if os.path.exists(os.path.join(model_path, safe_file_list[i])):
+                        continue
+                    else:
+                        ValueError(f"Not found {os.path.join(model_path, safe_file_list[i])}.")
+            file_type = "safetensors"
+        elif os.path.exists(os.path.join(model_path, self.safe_weight_name())):
+            file_type = "safetensors_without_index"
+        elif os.path.exists(os.path.join(model_path, self.weight_name())):
+            file_type = "pdparams"
+        else:
+            raise ValueError(
+                f"Please check path {model_path} is correct. Support safetensors and pdparams in complete parameter format (not TP or PP format) only."
+            )
+        return file_type
+
+    def weight_name(self):
+        if self.merge_config.merge_preifx == "model":
+            return PADDLE_WEIGHTS_NAME
+        else:
+            return PADDLE_MASTER_WEIGHTS_NAME
+
+    def safe_weight_name(self):
+        if self.merge_config.merge_preifx == "model":
+            return SAFE_WEIGHTS_NAME
+        else:
+            return SAFE_MASTER_WEIGHTS_NAME
+
+    def safe_index_name(self):
+        if self.merge_config.merge_preifx == "model":
+            return SAFE_WEIGHTS_INDEX_NAME
+        else:
+            return SAFE_MASTER_WEIGHTS_INDEX_NAME

paddlenlp/mergekit/merge_utils.py

@@ -0,0 +1,31 @@
+# Copyright (c) 2024 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.
+
+
+def divide_positions(m, n):
+    if n == 0:
+        raise ValueError("n should be greater than zero")
+    if m < n:
+        raise ValueError("tensor number should be greater than or equal to processor number")
+    base_value = m // n
+    remainder = m % n
+    positions = [0]
+    for i in range(1, n):
+        if remainder > 0:
+            positions.append(positions[-1] + base_value + 1)
+            remainder -= 1
+        else:
+            positions.append(positions[-1] + base_value)
+    positions.append(m)
+    return positions

paddlenlp/mergekit/sparsify_method.py

@@ -0,0 +1,87 @@
+# Copyright (c) 2024 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
+
+
+class SparsifyMethod:
+    def __init__(self, merge_config):
+        self.merge_config = merge_config
+
+    def sparsify(self, tensor):
+        if self.merge_config.sparsify_type is None:
+            return tensor
+        if self.merge_config.sparsify_type == "dare":
+            return self.dare(tensor)
+        elif self.merge_config.sparsify_type == "magprune":
+            return self.magprune(tensor)
+        elif self.merge_config.sparsify_type == "trim":
+            return self.trim(tensor)
+        else:
+            raise ValueError(f"Unknown sparsify method: {self.merge_config.sparsify_type}")
+
+    def dare(self, tensor):
+        if self.merge_config.tensor_type == "np":
+            mask = np.random.binomial(1, self.merge_config.reserve_p, size=tensor.shape).astype(tensor.dtype)
+            tensor *= mask
+            if self.merge_config.rescale:
+                tensor /= self.merge_config.reserve_p
+            return tensor
+        else:
+            raise NotImplementedError("Paddle Tensor is not supported yet.")
+
+    def magprune(self, tensor):
+        if self.merge_config.tensor_type == "np":
+            if np.all(tensor == 0):
+                return tensor
+            drop_p = 1 - self.merge_config.reserve_p
+            # 1: ranking(descending)
+            abs_tensor = np.abs(tensor)
+            sorted_indices = np.argsort(-abs_tensor.flatten())
+
+            # 2: caclculate drop rate p_i
+            probs = np.empty_like(sorted_indices)
+            probs[sorted_indices] = np.arange(tensor.size).astype(tensor.dtype)
+            probs = probs.reshape(tensor.shape)  # r_i ∈ {0，1，... ,n}
+            probs = probs * self.merge_config.epsilon / tensor.size  # Δ_i =  ε/n * r_i
+            p_min = drop_p - self.merge_config.epsilon / 2  # minimal drop rate
+            probs += p_min  # p_i for each parameter
+
+            # 3: drop parameters according to their probabilities
+            mask = np.random.binomial(1, probs)
+            tensor *= (1 - mask).astype(tensor.dtype)
+            if self.merge_config.rescale:
+                tensor /= 1 - probs
+            return tensor
+        else:
+            raise NotImplementedError("Paddle Tensor is not supported yet.")
+
+    def trim(self, tensor):
+        if self.merge_config.tensor_type == "np":
+            shape = tensor.shape
+            org_sum = np.sum(np.abs(tensor))
+            tensor = tensor.flatten()
+            abs_tensor = np.abs(tensor)
+            threshold = np.quantile(abs_tensor, 1 - self.merge_config.reserve_p)
+
+            if self.merge_config.rescale:
+                org_sum = np.sum(np.abs(tensor))
+                tensor[abs_tensor < threshold] = 0
+                new_sum = np.sum(np.abs(tensor))
+                if org_sum >= 1e-8 and new_sum >= 1e-8:
+                    tensor *= org_sum / new_sum
+            else:
+                tensor[abs_tensor < threshold] = 0
+            return tensor.reshape(shape)
+        else:
+            raise NotImplementedError("Paddle Tensor is not supported yet.")

paddlenlp/peft/lora/lora_config.py

@@ -76,6 +76,7 @@ class LoRAConfig:
     do_qat: bool = field(default=False, metadata={"help": "Whether the lora model would do quant-aware training"})
     rslora: bool = field(default=False, metadata={"help": "Whether to use RsLoRA"})
     pissa: bool = field(default=False, metadata={"help": "Whether to use Pissa: https://arxiv.org/pdf/2404.02948.pdf"})
+    loraga: bool = field(default=False, metadata={"help": "Whether to LoRA-GA"})
     lora_plus_scale: float = field(default=1.0, metadata={"help": "Lora B scale in LoRA+"})
     base_model_name_or_path: Optional[str] = field(
         default=None, metadata={"help": "The name of the base model to use."}

paddlenlp/peft/lora/lora_layers.py

@@ -48,6 +48,7 @@ from ...transformers.mc2_parallel_linear import (
     MC2RowSeqParallelCoreLinear,
 )
 from .lora_quick_layers import quick_lora
+from .utils import rng_ctx
 
 
 class LoRALinear(nn.Linear):
@@ -198,14 +199,15 @@ class RowParallelLoRALinear(RowParallelLinear):
         self.name = self._name
 
         # Actual trainable parameters
-        self.lora_A = self.create_parameter(
-            shape=[self.input_size_per_partition, r],
-            dtype=self._dtype,
-            is_bias=False,
-            attr=paddle.ParamAttr(
-                initializer=nn.initializer.KaimingUniform(negative_slope=math.sqrt(5), nonlinearity="leaky_relu")
-            ),
-        )
+        with rng_ctx(self.is_mp, paddle.in_dynamic_mode()):
+            self.lora_A = self.create_parameter(
+                shape=[self.input_size_per_partition, r],
+                dtype=self._dtype,
+                is_bias=False,
+                attr=paddle.ParamAttr(
+                    initializer=nn.initializer.KaimingUniform(negative_slope=math.sqrt(5), nonlinearity="leaky_relu")
+                ),
+            )
         self.lora_B = self.create_parameter(
             shape=[r, self.out_features],
             dtype=self._dtype,
@@ -345,14 +347,15 @@ class RowSequenceParallelLoRALinear(RowSequenceParallelLinear):
         self.name = self._name
 
         # Actual trainable parameters
-        self.lora_A = self.create_parameter(
-            shape=[self.input_size_per_partition, r],
-            dtype=self._dtype,
-            is_bias=False,
-            attr=paddle.ParamAttr(
-                initializer=nn.initializer.KaimingUniform(negative_slope=math.sqrt(5), nonlinearity="leaky_relu")
-            ),
-        )
+        with rng_ctx(self.is_mp, paddle.in_dynamic_mode()):
+            self.lora_A = self.create_parameter(
+                shape=[self.input_size_per_partition, r],
+                dtype=self._dtype,
+                is_bias=False,
+                attr=paddle.ParamAttr(
+                    initializer=nn.initializer.KaimingUniform(negative_slope=math.sqrt(5), nonlinearity="leaky_relu")
+                ),
+            )
         self.lora_B = self.create_parameter(
             shape=[r, self.out_features],
             dtype=self._dtype,
@@ -468,15 +471,16 @@ class ColumnParallelLoRALinear(ColumnParallelLinear):
             attr=lora_A_weight_attr,
         )
         self.lora_A.is_distributed = False
-        self.lora_B = self.create_parameter(
-            shape=[r, self.output_size_per_partition],
-            dtype=self._dtype,
-            is_bias=False,
-            attr=paddle.ParamAttr(
-                initializer=paddle.nn.initializer.Constant(value=0.0),
-                learning_rate=lora_plus_scale,
-            ),
-        )
+        with rng_ctx(self.is_mp, paddle.in_dynamic_mode()):
+            self.lora_B = self.create_parameter(
+                shape=[r, self.output_size_per_partition],
+                dtype=self._dtype,
+                is_bias=False,
+                attr=paddle.ParamAttr(
+                    initializer=paddle.nn.initializer.Constant(value=0.0),
+                    learning_rate=lora_plus_scale,
+                ),
+            )
 
         self.lora_B.is_distributed = True
         self.lora_B.split_axis = 1
@@ -599,15 +603,16 @@ class ColumnSequenceParallelLoRALinear(ColumnSequenceParallelLinear):
         self.lora_A.is_distributed = False
         mark_as_sequence_parallel_parameter(self.lora_A)
 
-        self.lora_B = self.create_parameter(
-            shape=[r, self.output_size_per_partition],
-            dtype=self._dtype,
-            is_bias=False,
-            attr=paddle.ParamAttr(
-                initializer=paddle.nn.initializer.Constant(value=0.0),
-                learning_rate=lora_plus_scale,
-            ),
-        )
+        with rng_ctx(self.is_mp, paddle.in_dynamic_mode()):
+            self.lora_B = self.create_parameter(
+                shape=[r, self.output_size_per_partition],
+                dtype=self._dtype,
+                is_bias=False,
+                attr=paddle.ParamAttr(
+                    initializer=paddle.nn.initializer.Constant(value=0.0),
+                    learning_rate=lora_plus_scale,
+                ),
+            )
 
         self.lora_B.is_distributed = True
         self.lora_B.split_axis = 1

paddlenlp/peft/lora/lora_model.py

@@ -160,6 +160,9 @@ class LoRAModel(nn.Layer):
                 f"Reset tensor_parallel_degree of lora_config to {self.model.config.tensor_parallel_degree}."
             )
         self.forward = self.model.forward
+        if lora_config.loraga:
+            self.loraga_init_dict = {}
+            self.reinit_base_model = False
 
         logger.info("Mark only lora and trainable_module as trainable.")
         self.mark_only_lora_as_trainable()
@@ -262,9 +265,11 @@ class LoRAModel(nn.Layer):
                     pre_tensor_parallel_split = True
                     tp_actions = lora_model._get_tensor_parallel_convert_actions(loaded_keys, is_split=True)
                 state_dict = load_state_dict(
-                    shard_file, tp_actions if pre_tensor_parallel_split else None, expected_keys
+                    shard_file,
+                    tp_actions if pre_tensor_parallel_split else None,
+                    expected_keys,
                 )
-                error_msgs += _load_state_dict_into_model(lora_model.model, state_dict, "")
+                error_msgs += _load_state_dict_into_model(lora_model, state_dict, "")
                 del state_dict
                 gc.collect()
 
@@ -314,6 +319,41 @@ class LoRAModel(nn.Layer):
         warnings.filterwarnings(
             action="ignore", message=".*Skip loading for.*", category=Warning, lineno=0, append=False
         )
+
+        model_state_dict = self.model.state_dict()
+        if self.lora_config.loraga:
+
+            def process_split_and_assign(name, concat_tensor, axis, init_dict, state_dict):
+                if isinstance(concat_tensor, np.ndarray):
+                    final_lora, init_lora = np.split(concat_tensor, 2, axis=axis)
+                    init_lora = paddle.to_tensor(init_lora)
+                else:
+                    final_lora, init_lora = paddle.split(concat_tensor, 2, axis=axis)
+                init_dict[name] = init_lora
+                state_dict[name] = final_lora
+                return init_lora
+
+            for name in state_dict.keys():
+                if "lora_A" in name:
+                    concat_lora_A = state_dict[name]
+                    init_loraA = process_split_and_assign(
+                        name, concat_lora_A, axis=1, init_dict=self.loraga_init_dict, state_dict=state_dict
+                    )
+
+                    loraB_name = name.replace("lora_A", "lora_B")
+                    concat_lora_B = state_dict[loraB_name]
+                    init_loraB = process_split_and_assign(
+                        loraB_name, concat_lora_B, axis=0, init_dict=self.loraga_init_dict, state_dict=state_dict
+                    )
+
+                    base_name = name.replace("lora_A", "weight")
+                    if not self.reinit_base_model:
+                        # Reinit base model
+                        offset = init_loraA.cuda() @ init_loraB.cuda()
+                        ori_weight = model_state_dict[base_name]
+                        model_state_dict[base_name].set_value(ori_weight - self.lora_config.scaling * offset)
+        del model_state_dict
+        gc.collect()
         self.model.set_state_dict(state_dict)
         logger.info("Load lora weight successfully")
 
@@ -395,8 +435,9 @@ class LoRAModel(nn.Layer):
 
         lora_config_to_save = LoRAConfig(**self.lora_config.to_dict())
 
+        trainable_state_dict = self.get_trainable_state_dict(concat_init_lora=lora_config_to_save.loraga)
+
         if merge_tensor_parallel and lora_config_to_save.tensor_parallel_degree > 1:
-            trainable_state_dict = self.get_trainable_state_dict()
             trainable_state_dict = self._merge_trainable_tensor_parallel(trainable_state_dict)
             if not is_main_process:
                 logger.info("Saving with merge_tensor_parallel, tensor_parallel_rank > 0 don't need save")
@@ -405,7 +446,6 @@ class LoRAModel(nn.Layer):
                 variant = "_".join([x for x in variant.split("_") if "tp" not in x])
             lora_config_to_save.tensor_parallel_degree = -1
         else:
-            trainable_state_dict = self.get_trainable_state_dict()
             if lora_config_to_save.tensor_parallel_degree > 1:
                 if variant is None:
                     variant = weight_name_suffix()
@@ -635,12 +675,18 @@ class LoRAModel(nn.Layer):
             original_module.bias = module.bias
         setattr(parent_module, attribute_chain[-1], original_module)
 
-    def get_trainable_state_dict(self):
+    def get_trainable_state_dict(self, concat_init_lora=False):
         trainable_state_dict = OrderedDict()
         for name, weight in self.model.state_dict().items():
             # get lora parameter & QAT scale parameter
             if not weight.stop_gradient or "activation_quanter" in name or "weight_quanter" in name:
-                trainable_state_dict[name] = weight
+                if concat_init_lora:
+                    if "lora_A" in name:
+                        trainable_state_dict[name] = paddle.concat([weight, self.loraga_init_dict[name]], axis=1)
+                    else:
+                        trainable_state_dict[name] = paddle.concat([weight, self.loraga_init_dict[name]], axis=0)
+                else:
+                    trainable_state_dict[name] = weight
         return trainable_state_dict
 
     def print_trainable_parameters(self) -> None:

paddlenlp/peft/lora/lora_quantization_layers.py

@@ -25,6 +25,7 @@ from ...quantization.quantization_linear import (
     QuantizationLinear,
     RowParallelQuantizationLinear,
 )
+from .utils import rng_ctx
 
 
 class QuantizationLoRALinear(QuantizationLinear):
@@ -241,12 +242,13 @@ class ColumnParallelQuantizationLoRALinear(ColumnParallelQuantizationLinear):
             attr=lora_A_weight_attr,
         )
         self.lora_A.is_distributed = False
-        self.lora_B = self.create_parameter(
-            shape=[r, self.output_size_per_partition],
-            dtype=self._dtype,
-            is_bias=False,
-            default_initializer=nn.initializer.Constant(value=0.0),
-        )
+        with rng_ctx(self.is_mp, paddle.in_dynamic_mode()):
+            self.lora_B = self.create_parameter(
+                shape=[r, self.output_size_per_partition],
+                dtype=self._dtype,
+                is_bias=False,
+                default_initializer=nn.initializer.Constant(value=0.0),
+            )
         self.lora_B.is_distributed = True
         self.lora_B.split_axis = 1
         self.scaling = self.lora_alpha / self.r
@@ -386,14 +388,15 @@ class RowParallelQuantizationLoRALinear(RowParallelQuantizationLinear):
             self.lora_dropout = lambda x: x
 
         # Actual trainable parameters
-        self.lora_A = self.create_parameter(
-            shape=[self.input_size_per_partition, r],
-            dtype=self._dtype,
-            is_bias=False,
-            attr=paddle.ParamAttr(
-                initializer=nn.initializer.KaimingUniform(negative_slope=math.sqrt(5), nonlinearity="leaky_relu")
-            ),
-        )
+        with rng_ctx(self.is_mp, paddle.in_dynamic_mode()):
+            self.lora_A = self.create_parameter(
+                shape=[self.input_size_per_partition, r],
+                dtype=self._dtype,
+                is_bias=False,
+                attr=paddle.ParamAttr(
+                    initializer=nn.initializer.KaimingUniform(negative_slope=math.sqrt(5), nonlinearity="leaky_relu")
+                ),
+            )
         self.lora_B = self.create_parameter(
             shape=[r, self.out_features],
             dtype=self._dtype,

paddlenlp/peft/lora/loraga_utils.py

@@ -0,0 +1,381 @@
+# Copyright (c) 2024 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 paddle
+import paddle.distributed as dist
+from paddle.distributed import fleet
+
+from paddlenlp.peft import LoRAModel
+from paddlenlp.peft.lora.lora_layers import (
+    ColumnParallelLoRALinear,
+    ColumnSequenceParallelLoRALinear,
+    LoRALinear,
+    RowParallelLoRALinear,
+    RowSequenceParallelLoRALinear,
+)
+from paddlenlp.trainer import Trainer, TrainingArguments
+from paddlenlp.trainer.trainer_utils import ShardingOption
+from paddlenlp.transformers.model_utils import PretrainedModel, unwrap_model
+from paddlenlp.utils.log import logger
+
+
+class LoRAGATrainer(Trainer):
+    """A Trainer class for Lora-GA gradient estimation."""
+
+    def __init__(self, loraga_init_iters: int, gradient_offload: bool, **kwargs):
+        """
+        Initialize the Trainer class for Lora-GA gradient estimation.
+
+        Args:
+        loraga_init_iters (int): The number of forward and backward process in initializing Lora-GA.
+        gradient_offload (bool): Whether to offload gradients to CPU memory.
+
+        """
+        super().__init__(**kwargs)
+        logger.info(f"Initialization iterations for LoraGA: {loraga_init_iters}")
+        self.loraga_init_iters = loraga_init_iters
+        self.gradient_offload = gradient_offload
+
+    def estimate_gradient(self, model: PretrainedModel):
+        """
+        Estimate the gradient of the model on the given dataset
+        Args:
+            model (PretrainedModel): The base model to be trained.
+
+        Returns:
+            dict: A dictionary containing the estimated gradients for each named layer.
+                  Note: In tensor parallel mode, the gradients in the dict are not gathered.
+        """
+        gradient_dict = {}
+        logger.info("Estimating gradient for LoraGA.")
+
+        model = self._wrap_model(model)
+        model.train()
+        dataloader = self.get_train_dataloader()
+        iters = 0
+
+        with GradientOffloadHookContext(
+            model=model,
+            gradient_dict=gradient_dict,
+            local_rank=self.args.local_rank,
+            loraga_init_iters=self.loraga_init_iters,
+            gradient_offload=self.gradient_offload,
+        ):
+            for batch in dataloader:
+                iters += 1
+                # Pipeline parallel not supported currently
+                with paddle.amp.auto_cast(enable=True, custom_black_list=self.args.amp_custom_black_list):
+                    loss, _ = model(**batch)
+                loss.backward()
+
+                if iters == self.loraga_init_iters:
+                    break
+        return gradient_dict
+
+    def _wrap_model(self, model):
+        """Wrap Model without optimizer, support dp, tp and sharding"""
+
+        in_pipeline_parallel_mode = self.args.pipeline_parallel_degree > 1
+        in_sharding_parallel_mode = self.sharding is not None
+        in_tensor_parallel_mode = self.args.tensor_parallel_degree > 1
+        in_sep_parallel_mode = self.args.sep_parallel_degree > 1
+        in_cp_parallel_mode = self.args.context_parallel_degree > 1
+
+        if in_pipeline_parallel_mode:
+            raise ValueError("LoRA-GA do not supported pipeline parallel currently.")
+
+        # Multi-gpu training
+        if self.args.world_size > 1 and (not self.args.use_hybrid_parallel):
+            # MOE use DDP to broadcaset parameters.
+            ddp_kwargs = {}
+            if self.args.ddp_find_unused_parameters is not None:
+                ddp_kwargs["find_unused_parameters"] = self.args.ddp_find_unused_parameters
+            elif isinstance(model, PretrainedModel):
+                # find_unused_parameters breaks checkpointing as per
+                # https://github.com/huggingface/transformers/pull/4659#issuecomment-643356021
+                ddp_kwargs["find_unused_parameters"] = not any(
+                    hasattr(m, "enable_recompute") and m.enable_recompute for m in model.sublayers(include_self=True)
+                )
+            else:
+                ddp_kwargs["find_unused_parameters"] = True
+            model = paddle.DataParallel(model, **ddp_kwargs)
+
+        # sharding
+        if in_sharding_parallel_mode:
+            # Sharded DDP!
+            if self.args.tensor_parallel_degree > 1:
+                hcg = fleet.get_hybrid_communicate_group()
+                assert (
+                    ShardingOption.SHARD_GRAD_OP in self.args.sharding or ShardingOption.SHARD_OP in self.args.sharding
+                ), "Only support tensor parallel + sharding stage1/stage2 hybrid parallel now."
+                model = paddle.distributed.fleet.meta_parallel.TensorParallel(model, hcg, strategy=None)
+            if ShardingOption.SHARD_OP in self.args.sharding:
+                model = fleet.distributed_model(model)
+
+        if not in_sharding_parallel_mode and (in_tensor_parallel_mode or in_sep_parallel_mode or in_cp_parallel_mode):
+            model = fleet.distributed_model(model)
+
+        return model
+
+
+def get_module_gradient(
+    grad_name,
+    base_model_prefix,
+    gradient_dict,
+    base_model_split_mappings,
+    tp_degree,
+    sharding_degree,
+    dp_degree,
+    local_rank,
+):
+    """
+    Gather modules gradient in tensor parallel mode.
+    Average module gradient in data parallel mode and sharding parallel mode.
+
+    Args:
+        grad_name (str): The name of the gradient parameter.
+        base_model_prefix (str): The prefix of the base model's parameter names.
+        gradient_dict (dict): A dictionary containing the estimated gradients for each named layer.
+        base_model_split_mappings (dict): A mapping of model keys to merge functions.
+        sharding_degree (int): The sharding parallel degree.
+        dp_degree (int): The data parallel degree.
+        local_rank (int): The local rank of the current process.
+
+    Returns:
+        Tensor: The processed gradient tensor.
+    """
+
+    rank_suffix = "_" + str(local_rank)
+    local_grad_name = ".".join(grad_name.split(".")[1:]) + ".weight" + rank_suffix
+    gradient = gradient_dict.pop(local_grad_name).cuda()
+
+    is_fleet_init = True
+    try:
+        hcg = fleet.get_hybrid_communicate_group()
+        model_parallel_group = hcg.get_model_parallel_group()
+        sharding_parallel_group = hcg.get_sharding_parallel_group()
+        data_parallel_group = hcg.get_data_parallel_group()
+    except:
+        is_fleet_init = False
+
+    if tp_degree > 1:
+        # remove prefix and suffix in name
+        model_split_key = local_grad_name.split(base_model_prefix)[-1].rsplit(rank_suffix, 1)[0]
+        if model_split_key in base_model_split_mappings:
+            merge_func = base_model_split_mappings[model_split_key]
+            output_tensors = []
+            dist.all_gather(output_tensors, gradient, group=model_parallel_group)
+
+            output_tensors = [t if len(t.shape) > 0 else t.reshape_([-1]) for t in output_tensors]
+            gradient = merge_func(output_tensors).cuda()
+
+    # sharding
+    if sharding_degree > 1:
+        if sharding_parallel_group.nranks > 1:
+            dist.all_reduce(gradient, op=dist.ReduceOp.SUM, group=sharding_parallel_group)
+            gradient /= sharding_parallel_group.nranks
+
+    # dp
+    if dp_degree > 1:
+        if data_parallel_group.nranks > 1:
+            if is_fleet_init:
+                dist.all_reduce(gradient, op=dist.ReduceOp.SUM, group=data_parallel_group)
+            else:
+                dist.all_reduce(gradient, op=dist.ReduceOp.SUM)
+            gradient /= data_parallel_group.nranks
+    return gradient
+
+
+def loraga_svd_reinit(
+    model: LoRAModel, gradient_dict: dict, stable_gamma: int, training_args: TrainingArguments, **kwargs
+) -> None:
+    """
+    Perform SVD to gradients and reinitialize base model weight and lora adapter weight.
+
+    Args:
+        model (LoRAModel): The LoRAModel containing LoRA layers.
+        gradient_dict (dict): A dictionary containing the estimated gradients for each named layer.
+        stable_gamma (int): A scaling factor for LoRA-GA initialization.
+        training_args (TrainingArguments): Training arguments.
+
+    Returns:
+        None: Updates the model's weights and LoRA adapter weights in place.
+    """
+    tensor_parallel_degree = training_args.tensor_parallel_degree
+    in_tensor_parallel_mode = tensor_parallel_degree > 1
+    lora_split_mapping = None
+    base_model_split_mappings = None
+    if in_tensor_parallel_mode:
+        base_model_split_mappings = model.model._get_tensor_parallel_mappings(config=model.config, is_split=False)
+
+    base_model_prefix = unwrap_model(model).base_model_prefix + "."
+    if in_tensor_parallel_mode:
+        lora_split_mapping = model._get_tensor_parallel_mappings(model.config)
+    loraga_init_dict = {}
+    for name, module in model.named_sublayers():
+        if isinstance(
+            module,
+            (
+                LoRALinear,
+                RowSequenceParallelLoRALinear,
+                ColumnSequenceParallelLoRALinear,
+                RowParallelLoRALinear,
+                ColumnParallelLoRALinear,
+            ),
+        ):
+            # gather gradient if in tensor parallel mode, average gradient if in data parallel mode
+            module_gradient = get_module_gradient(
+                name,
+                base_model_prefix,
+                gradient_dict,
+                base_model_split_mappings,
+                training_args.tensor_parallel_degree,
+                training_args.sharding_parallel_degree,
+                training_args.data_parallel_degree,
+                training_args.local_rank,
+            )
+            # perform SVD to reinit base model weight and lora adapter weight
+            loraga_svd_module(
+                name,
+                module,
+                module_gradient,
+                stable_gamma,
+                loraga_init_dict,
+                in_tensor_parallel_mode,
+                lora_split_mapping,
+                **kwargs,
+            )
+    model.reinit_base_model = True
+    model.loraga_init_dict = loraga_init_dict
+
+
+def loraga_svd_module(
+    name,
+    module,
+    grads,
+    stable_gamma,
+    loraga_init_dict,
+    in_tensor_parallel_mode=False,
+    lora_split_mapping=None,
+    **kwargs
+):
+    with paddle.no_grad():
+        lora_r = module.r
+
+        loraA_name = ".".join(name.split(".")[1:]) + ".lora_A"
+        loraB_name = ".".join(name.split(".")[1:]) + ".lora_B"
+        # Perform SVD to gradients
+        U, S, V = paddle.linalg.svd_lowrank(grads.astype("float32"), q=4 * lora_r, niter=4)
+
+        V = V.T
+        # get new low-rank adapter after SVD
+        A = U[:, lora_r : 2 * lora_r]
+        B = V[:lora_r, :]
+
+        m, n = grads.shape
+        # If stable_gamma is not -1, scale the matrices A and B by the square root of the stable_gamma
+        if stable_gamma != -1:
+            A = A * m**0.25 / stable_gamma**0.5
+            B = B * m**0.25 / stable_gamma**0.5
+        else:
+            A = A / module.scaling
+            B = B / module.scaling
+
+        if in_tensor_parallel_mode:
+            # split lora adapter weight if in tensor parallel mode
+            if module.lora_A.is_distributed and lora_split_mapping is not None:
+                split_function = lora_split_mapping[loraA_name]
+                A = paddle.to_tensor(split_function(A))
+            if module.lora_B.is_distributed and lora_split_mapping is not None:
+                split_function = lora_split_mapping[loraB_name]
+                B = paddle.to_tensor(split_function(B))
+        A = A.astype(module.lora_A.dtype)
+        B = B.astype(module.lora_B.dtype)
+        loraga_init_dict[loraA_name] = A
+        loraga_init_dict[loraB_name] = B
+        # reinit lora adapter weight
+        module.lora_A.set_value(A)
+        module.lora_B.set_value(B)
+
+        offset = module.lora_A @ module.lora_B
+        # reinit base model weight
+        module.weight.data -= module.scaling * offset
+
+
+def set_hook_enable(value=False):
+    global ENABLE_HOOK
+    ENABLE_HOOK = value
+
+
+def get_hook_enable():
+    global ENABLE_HOOK
+    return ENABLE_HOOK
+
+
+class GradientOffloadHookContext:
+    """Context manager for offloading gradient memory to CPU."""
+
+    def __init__(
+        self,
+        model,
+        gradient_dict: dict,
+        local_rank: int = 0,
+        loraga_init_iters: int = 4,
+        gradient_offload: bool = False,
+        *args,
+        **kwargs,
+    ):
+        self.model = model
+        self.gradient_dict = gradient_dict
+        self.local_rank = local_rank
+        self.loraga_init_iters = loraga_init_iters
+        self.gradient_offload = gradient_offload
+
+    def __enter__(self):
+        set_hook_enable(True)
+        self.register_gradient_hook()
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        set_hook_enable(False)
+
+    def register_gradient_hook(self):
+        """Register gradient hooks for all model parameters."""
+        for grad_name, param in self.model.named_parameters():
+            param._register_backward_hook(
+                self.get_record_gradient_hook(self.model, self.gradient_dict, grad_name, param)
+            )
+
+    def get_record_gradient_hook(self, model, gradient_dict, grad_name, param):
+        """Create a gradient recording hook for a parameter."""
+
+        def record_gradient_hook(*_):
+            if get_hook_enable():
+                grad = param.grad
+                local_grad_name = grad_name.split("_layers.")[-1] + "_" + str(self.local_rank)
+                if not param.stop_gradient and grad is not None:
+                    if local_grad_name not in gradient_dict:
+                        if self.gradient_offload:
+                            gradient_dict[local_grad_name] = (grad / self.loraga_init_iters).cpu()
+                        else:
+                            gradient_dict[local_grad_name] = grad.clone() / self.loraga_init_iters
+                    else:
+                        if self.gradient_offload:
+                            new_grad = gradient_dict[local_grad_name].cuda() + grad / self.loraga_init_iters
+                            gradient_dict[local_grad_name] = new_grad.cpu()
+                        else:
+                            gradient_dict[local_grad_name] += grad / self.loraga_init_iters
+                param.clear_gradient(False)  # release gradient memory
+
+        return record_gradient_hook

paddlenlp/peft/lora/utils.py

@@ -0,0 +1,21 @@
+# 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 contextlib import nullcontext
+
+from paddle.distributed.fleet.meta_parallel import get_rng_state_tracker
+
+
+def rng_ctx(is_mp: bool, in_dynamic_mode: bool):
+    return get_rng_state_tracker().rng_state() if (is_mp and in_dynamic_mode) else nullcontext()

paddlenlp/peft/prefix/prefix_model.py

@@ -333,7 +333,9 @@ class PrefixModelForCausalLM(paddle.nn.Layer):
                     pre_tensor_parallel_split = True
                     tp_actions = prefix_model._get_tensor_parallel_convert_actions(is_split=True)
                 state_dict = load_state_dict(
-                    shard_file, tp_actions if pre_tensor_parallel_split else None, expected_keys
+                    shard_file,
+                    tp_actions if pre_tensor_parallel_split else None,
+                    expected_keys,
                 )
                 error_msgs += _load_state_dict_into_model(prefix_model.prefix_encoder, state_dict, "")
                 del state_dict

paddlenlp/quantization/checkpoint_quantization_utils.py

@@ -0,0 +1,364 @@
+# Copyright (c) 2024 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
+
+
+def cal_ratio(m, v, eps=1e-8):
+    """
+    cal part adam update ratio.
+    Args:
+        m (`paddle.Tensor`):
+            moment in Adam optimizer.
+        v (`paddle.Tensor`):
+            variance in Adam optimizer.
+        eps (`int`):
+            epsilon in Adam optimizer.
+    """
+    return 1 / (np.sqrt(v) + eps)
+
+
+def group_wise_quant_dequant(
+    inputs,
+    mins=None,
+    maxs=None,
+    quant_bits=4,
+    group_size=32,
+    quant=True,
+    tp_rank=-1,
+    tp_degree=1,
+    use_pd=False,
+    symmetry=False,
+):
+    """
+    group-wise quantization (support symmetry, asymmetry).
+    Args:
+        inputs (`paddle.Tensor`):
+            The tensor to quantize.
+        mins (`paddle.Tensor`):
+            Min scales tensor in asymmetry quantization.
+        maxs (`paddle.Tensor`):
+            Max scales tensor in asymmetry quantization, or Abs max tensor in symmetry quantization.
+        quant_bits (`int`):
+            Quantization bits.
+        group_size (`int`):
+            Group size of group-wise quantization.
+        quant (`bool`):
+            True when quantization, False in dequantization.
+        tp_rank (`int`):
+            Tensor parallel rank.
+        tp_degree (`int`):
+            Tensor parallel world size.
+        use_pd (`bool`):
+            Whether to use paddle caculation. If False will use numpy.
+        symmetry (`bool`):
+            Whether to use symmetry quantization.
+    """
+
+    qmax = (1 << (quant_bits)) - 1
+    qmin = 0
+    shape = inputs.shape
+
+    if quant:
+        inputs_processed = inputs.reshape([shape[0] // group_size, group_size, shape[1]])
+        if symmetry:
+            bnt = (1 << (quant_bits - 1)) - 1
+            scales = np.max(np.abs(inputs_processed), axis=1)
+            new_scales = np.repeat(scales, repeats=group_size, axis=0)
+            quant_tensor = np.clip(np.round(inputs / new_scales * bnt), -bnt - 1, bnt)
+            return quant_tensor.astype("int8"), scales
+
+        # scales: [shape[0] // group_size, shape[1]]
+        maxs = np.max(inputs_processed, axis=1)
+        mins = np.min(inputs_processed, axis=1)
+        scales = maxs - mins
+        # new_scales: [shape[0], shape[1]]
+        new_scales = np.repeat(scales, repeats=group_size, axis=0)
+        new_mins = np.repeat(mins, repeats=group_size, axis=0)
+        # add eps to avoid devide zero
+        quant_tensor = np.clip(np.round((inputs - new_mins) / (new_scales) * qmax), qmin, qmax)
+        quant_tensor = np.nan_to_num(quant_tensor)
+        return quant_tensor.astype("uint8"), mins, maxs
+    else:
+        if symmetry:
+            scales = mins
+            bnt = (1 << (quant_bits - 1)) - 1
+            if use_pd:
+                new_scales = paddle.repeat_interleave(scales, group_size, 0)
+            else:
+                new_scales = np.repeat(scales, repeats=group_size, axis=0)
+
+            if tp_rank == -1:
+                dequant_tensor = inputs.astype("float32") * new_scales / bnt
+            elif len(new_scales.shape) == 0 or inputs.shape[-1] == new_scales.shape[-1]:
+                # input tensor was row parallel in tp.
+                dequant_tensor = (
+                    inputs.astype("float32")
+                    * new_scales[
+                        tp_rank * new_scales.shape[0] // tp_degree : (tp_rank + 1) * new_scales.shape[0] // tp_degree
+                    ]
+                    / bnt
+                )
+            else:
+                # input tensor was column parallel in tp.
+                dequant_tensor = (
+                    inputs.astype("float32")
+                    * new_scales[
+                        :,
+                        tp_rank
+                        * new_scales.shape[-1]
+                        // tp_degree : (tp_rank + 1)
+                        * new_scales.shape[-1]
+                        // tp_degree,
+                    ]
+                    / bnt
+                )
+            return dequant_tensor
+
+        scales = maxs - mins
+        if use_pd:
+            new_scales = paddle.repeat_interleave(scales, group_size, 0)
+            new_mins = paddle.repeat_interleave(mins, group_size, 0)
+        else:
+            new_scales = np.repeat(scales, repeats=group_size, axis=0)
+            new_mins = np.repeat(mins, repeats=group_size, axis=0)
+
+        if tp_rank == -1:
+            dequant_tensor = (inputs.astype("float32") / qmax * new_scales) + new_mins
+        elif len(new_scales.shape) == 0 or inputs.shape[-1] == new_scales.shape[-1]:
+            # input tensor was row parallel in tp.
+            dequant_tensor = (
+                inputs.astype("float32")
+                / qmax
+                * new_scales[
+                    tp_rank * new_scales.shape[0] // tp_degree : (tp_rank + 1) * new_scales.shape[0] // tp_degree
+                ]
+            ) + new_mins[tp_rank * new_mins.shape[0] // tp_degree : (tp_rank + 1) * new_mins.shape[0] // tp_degree]
+        else:
+            # input tensor was column parallel in tp.
+            dequant_tensor = (
+                inputs.astype("float32")
+                / qmax
+                * new_scales[
+                    :, tp_rank * new_scales.shape[-1] // tp_degree : (tp_rank + 1) * new_scales.shape[-1] // tp_degree
+                ]
+            ) + new_mins[
+                :, tp_rank * new_mins.shape[-1] // tp_degree : (tp_rank + 1) * new_mins.shape[-1] // tp_degree
+            ]
+        return dequant_tensor
+
+
+def merge_int4(x, y):
+    """
+    merge 2 signed int4 to 1 int8
+    Args:
+        x (`numpy.array`):
+            4bits signed int x.
+        y (`numpy.array`):
+            4bits signed int y.
+    """
+    int4_high = x << 4
+    int4_low = y & 0x0F
+    final = int4_high | int4_low
+    return final.astype("int8")
+
+
+def split_int8(final):
+    """
+    split an int8 to 2 int4 elems
+    Args:
+        final (`numpy.array`):
+            8bits signed int.
+    """
+    int4_high = final >> 4
+    int4_low = final & 0x0F
+
+    int4_high = np.where(int4_high > 8, int4_high - 16, int4_high)
+
+    high_tensor = paddle.Tensor(int4_high)
+    low_tensor = paddle.Tensor(int4_low)
+
+    return high_tensor, low_tensor
+
+
+def cal_abs_min_max_channel(inputs, quant_axis=1):
+    """
+    channel-wise min max scales calculation
+    Args:
+        inputs (`numpy.array`):
+            input tensor for quantization.
+        quant_axis (`int`):
+            dimension where calulating inputs' abs min and max scales on.
+    """
+    eps = 1e-8
+    reduce_axis = tuple([i for i in range(len(inputs.shape)) if i != quant_axis])
+    abs_max_values = np.max(inputs, axis=reduce_axis)
+    abs_min_values = np.min(inputs, axis=reduce_axis)
+    abs_max_values = np.where(
+        abs_max_values == np.array(0, dtype=inputs.dtype), np.array(eps, dtype=inputs.dtype), abs_max_values
+    )
+    abs_min_values = np.where(
+        abs_min_values == np.array(0, dtype=inputs.dtype), np.array(eps, dtype=inputs.dtype), abs_min_values
+    )
+    return abs_max_values, abs_min_values
+
+
+def asymmetry_qdq_weight(
+    x, quant_bit=8, quant_axis=-1, mins=None, maxs=None, dequant=False, tp_rank=-1, tp_degree=1, use_pd=False
+):
+    """
+    channel-wise asymmetry quantization
+    Args:
+        x (`paddle.Tensor`):
+            The tensor to quantize.
+        quant_bits (`int`):
+            Quantization bits.
+        quant_axis (`int`):
+            Scales caculation axis.
+        mins (`paddle.Tensor`):
+            Min scales tensor in asymmetry quantization.
+        maxs (`paddle.Tensor`):
+            Max scales tensor in asymmetry quantization.
+        dequant (`bool`):
+            True when dequantization, False in quantization.
+        tp_rank (`int`):
+            Model parallel rank.
+        tp_degree (`int`):
+            Model parallel world size.
+        use_pd (`bool`):
+            Whether to use paddle caculation. If False will use numpy.
+    """
+
+    if mins is None:
+        maxs, mins = cal_abs_min_max_channel(x)
+    bnt = (1 << (quant_bit)) - 1
+    scales = maxs - mins
+    if not dequant:
+        # quant
+        quant_x = np.clip(np.round((x - mins) / scales * bnt), 0, bnt)
+        return quant_x.astype(np.uint8), mins, maxs
+    else:
+        quant_x = x
+        # dequant
+        if not use_pd:
+            if len(scales.shape) == 0 or quant_x.shape[-1] == scales.shape[-1]:
+                # input tensor was row parallel in tp.
+                qdq_x = (quant_x / bnt * scales) + mins
+            else:
+                # input tensor was column parallel in tp.
+                qdq_x = (
+                    quant_x
+                    / bnt
+                    * scales[tp_rank * scales.shape[0] // tp_degree : (tp_rank + 1) * scales.shape[0] // tp_degree]
+                ) + mins[tp_rank * mins.shape[0] // tp_degree : (tp_rank + 1) * mins.shape[0] // tp_degree]
+            return qdq_x.astype(np.float32), scales
+        else:
+            if len(scales.shape) == 0 or quant_x.shape[-1] == scales.shape[-1]:
+                # input tensor was row parallel in tp.
+                qdq_x = (quant_x / bnt * scales.unsqueeze(0).expand(quant_x.shape)) + mins
+            else:
+                # input tensor was column parallel in tp.
+                qdq_x = (
+                    quant_x
+                    / bnt
+                    * scales[tp_rank * scales.shape[0] // tp_degree : (tp_rank + 1) * scales.shape[0] // tp_degree]
+                    .unsqueeze(0)
+                    .expand(quant_x.shape)
+                ) + mins[tp_rank * mins.shape[0] // tp_degree : (tp_rank + 1) * mins.shape[0] // tp_degree]
+            return qdq_x.astype(paddle.float32), scales
+
+
+def cal_abs_max_channel(inputs, quant_axis=1):
+    """
+    channel-wise abs max calculation
+    Args:
+        inputs (`numpy.array`):
+            input tensor for quantization.
+        quant_axis (`int`):
+            dimension where calulating inputs' abs max scales on.
+    """
+    epsilon = 1e-8
+    reduce_axis = tuple([i for i in range(len(inputs.shape)) if i != quant_axis])
+    abs_max_values = np.max(np.abs(inputs), axis=reduce_axis)
+    # maybe all elements are zero in one group,
+    # so set the scales from those group to an actual number
+    # from divide 0.
+    abs_max_values = np.where(
+        abs_max_values == np.array(0, dtype=inputs.dtype), np.array(epsilon, dtype=inputs.dtype), abs_max_values
+    )
+    return abs_max_values
+
+
+def qdq_weight(x, quant_bit=8, quant_axis=-1, scales=None, dequant=False, tp_rank=-1, tp_degree=1, use_pd=False):
+    """
+    channel-wise symmetry quantization
+    Args:
+        x (`paddle.Tensor`):
+            The tensor to quantize.
+        quant_bits (`int`):
+            Quantization bits.
+        quant_axis (`int`):
+            Scales caculation axis.
+        scales (`paddle.Tensor`):
+            Abs max scales tensor in symmetry quantization.
+        dequant (`bool`):
+            True when dequantization, False in quantization.
+        tp_rank (`int`):
+            Model parallel rank.
+        tp_degree (`int`):
+            Model parallel world size.
+        use_pd (`bool`):
+            Whether to use paddle caculation. If False will use numpy.
+    """
+
+    if scales is None:
+        scales = cal_abs_max_channel(x)
+    bnt = (1 << (quant_bit - 1)) - 1
+    if not dequant:
+        # quant
+        quant_x = np.clip(np.round(x / scales * bnt), -bnt - 1, bnt)
+        return quant_x.astype(np.int8), scales
+    else:
+        quant_x = x
+        # dequant
+        if not use_pd:
+            if len(scales.shape) == 0 or quant_x.shape[-1] == scales.shape[-1]:
+                # input tensor was row parallel in tp.
+                qdq_x = quant_x / bnt * scales
+            else:
+                # input tensor was column parallel in tp.
+                qdq_x = (
+                    quant_x
+                    / bnt
+                    * scales[tp_rank * scales.shape[0] // tp_degree : (tp_rank + 1) * scales.shape[0] // tp_degree]
+                )
+            # fp32 , int8, int, fp32 or fp64
+            return qdq_x.astype(np.float32), scales
+        else:
+            if len(scales.shape) == 0 or quant_x.shape[-1] == scales.shape[-1]:
+                # input tensor was row parallel in tp.
+                qdq_x = quant_x / bnt * scales.unsqueeze(0).expand(quant_x.shape)
+            else:
+                # input tensor was column parallel in tp.
+                qdq_x = (
+                    quant_x
+                    / bnt
+                    * scales[tp_rank * scales.shape[0] // tp_degree : (tp_rank + 1) * scales.shape[0] // tp_degree]
+                    .unsqueeze(0)
+                    .expand(quant_x.shape)
+                )
+            # fp32 , int8, int, fp32 or fp64
+            return qdq_x.astype(paddle.float32), scales

paddlenlp/quantization/unified_checkpoint_quantization.py

@@ -0,0 +1,221 @@
+# Copyright (c) 2024 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 import fleet
+
+from paddlenlp.quantization.checkpoint_quantization_utils import (
+    asymmetry_qdq_weight,
+    cal_ratio,
+    group_wise_quant_dequant,
+    merge_int4,
+    qdq_weight,
+    split_int8,
+)
+from paddlenlp.utils.env import (
+    ASYMMETRY_QUANT_SCALE_MAX,
+    ASYMMETRY_QUANT_SCALE_MIN,
+    MOMENT1_KEYNAME,
+    MOMENT2_KEYNAME,
+    SYMMETRY_QUANT_SCALE,
+)
+from paddlenlp.utils.log import logger
+
+
+def dequant_unified_optimizer(state_dict, ckpt_quant_stage, scale_dict, use_pd=False):
+    """
+    dequantize unified optimizer state dict.
+    Args:
+        state_dict (`dict`):
+            unified checkpoint optimizer state dict.
+        ckpt_quant_stage (`str`):
+            checkpoint quantization stage, chosen in ["O0", "O1", "O2"].
+        scale_dict (`int`):
+            compression checkpoint scale dict.
+    """
+    logger.info(f"Start unified checkpoint dequantization, stage {ckpt_quant_stage}.")
+    tp_rank, tp_degree = -1, 1
+    if paddle.distributed.get_world_size() > 1:
+        hcg = fleet.get_hybrid_communicate_group()
+        tp_group = hcg.get_model_parallel_group()
+        tp_rank, tp_degree = tp_group.rank, tp_group.nranks
+
+    if ckpt_quant_stage == "O1":
+        # set eps
+        eps = 1e-8
+        for quant_key in state_dict.keys():
+            is_moment1 = MOMENT1_KEYNAME in quant_key
+            is_moment2 = MOMENT2_KEYNAME in quant_key
+            if is_moment1:
+                # dequant m1
+                scale_key = quant_key + SYMMETRY_QUANT_SCALE
+                weight = state_dict[quant_key]
+                scales = scale_dict[scale_key]
+                weight, _ = qdq_weight(
+                    weight,
+                    scales=scales,
+                    quant_bit=8,
+                    dequant=True,
+                    tp_rank=tp_rank,
+                    tp_degree=tp_degree,
+                    use_pd=use_pd,
+                )
+                state_dict[quant_key] = weight
+            elif is_moment2:
+                # dequant ratio
+                weight = state_dict[quant_key]
+                min_scale_key = quant_key + ASYMMETRY_QUANT_SCALE_MIN
+                max_scale_key = quant_key + ASYMMETRY_QUANT_SCALE_MAX
+                mins, maxs = scale_dict[min_scale_key], scale_dict[max_scale_key]
+                weight, _ = asymmetry_qdq_weight(
+                    weight,
+                    mins=mins,
+                    maxs=maxs,
+                    quant_bit=8,
+                    dequant=True,
+                    tp_rank=tp_rank,
+                    tp_degree=tp_degree,
+                    use_pd=use_pd,
+                )
+                # cal m2
+                if use_pd:
+                    weight = paddle.square(1.0 / weight - eps)
+                else:
+                    weight = np.square(1.0 / weight - eps)
+                state_dict[quant_key] = weight
+    elif ckpt_quant_stage == "O2":
+        # set eps
+        eps = 1e-8
+        m1_state_dict = {}
+        for quant_key in state_dict.keys():
+            # not all optimizer weights in O2 stage were quantized to int8,
+            # the norm-like weights were still remain in float32.
+            if state_dict[quant_key].dtype != paddle.int8:
+                logger.info(f"{quant_key} skip.")
+                continue
+            # split int8
+            weight = state_dict[quant_key]
+            m1_quant, ratio_quant = split_int8(weight.numpy())
+            # dequant ratio
+            ratio_min_scale_key = quant_key + ASYMMETRY_QUANT_SCALE_MIN
+            ratio_max_scale_key = quant_key + ASYMMETRY_QUANT_SCALE_MAX
+            m1_scale_key = quant_key[: -len(MOMENT2_KEYNAME)] + MOMENT1_KEYNAME + SYMMETRY_QUANT_SCALE
+            m1_scales = scale_dict[m1_scale_key]
+            ratio_mins, ratio_maxs = scale_dict[ratio_min_scale_key], scale_dict[ratio_max_scale_key]
+            m1_weight = group_wise_quant_dequant(
+                m1_quant,
+                mins=m1_scales,
+                maxs=None,
+                quant_bits=4,
+                quant=False,
+                tp_rank=tp_rank,
+                tp_degree=tp_degree,
+                use_pd=use_pd,
+                symmetry=True,
+            )
+            ratio_weight = group_wise_quant_dequant(
+                ratio_quant,
+                mins=ratio_mins,
+                maxs=ratio_maxs,
+                quant_bits=4,
+                quant=False,
+                tp_rank=tp_rank,
+                tp_degree=tp_degree,
+                use_pd=use_pd,
+            )
+
+            if use_pd:
+                ratio_weight = paddle.square(1.0 / ratio_weight - eps)
+            else:
+                ratio_weight = np.square(1.0 / ratio_weight - eps)
+            state_dict[quant_key] = ratio_weight
+            m1_state_dict[quant_key[: -len(MOMENT2_KEYNAME)] + MOMENT1_KEYNAME] = m1_weight
+            state_dict.update(m1_state_dict)
+
+    logger.info(f"Unified checkpoint dequantization done, stage {ckpt_quant_stage}.")
+
+    return state_dict
+
+
+def quant_unified_optimizer(state_dict, state_dict_type, ckpt_quant_stage, async_save=False):
+    """
+    quantize unified optimizer state dict.
+    Args:
+        state_dict (`dict`):
+            unified checkpoint optimizer state dict.
+        state_dict_type (`str`):
+            state_dict type, chosen in ["model_weight", "master_weight", "optimizer_weight"].
+        ckpt_quant_stage (`str`):
+            checkpoint quantization stage, chosen in ["O0", "O1", "O2"].
+        async_save (`bool`):
+            whether use async_save.
+    """
+    logger.info(f"Start unified checkpoint quantization, stage {ckpt_quant_stage}.")
+
+    quant = False
+    if ckpt_quant_stage != "O0":
+        quant = True
+    del_key = []
+    if quant and state_dict_type == "optimizer_weight":
+        scales_dict = {}
+        for k in state_dict.keys():
+            momentum1 = k.endswith(MOMENT1_KEYNAME)
+            momentum2 = k.endswith(MOMENT2_KEYNAME)
+
+            quant_weight = None
+
+            if ckpt_quant_stage == "O1":
+                # m1: wint8, 1/(sqrt(m2)+eps): wint8
+                if momentum2:
+                    # m1: m1_quant_weight, m2: ratio
+                    m1_key = k.split("/")[0] + "/" + MOMENT1_KEYNAME
+                    ratio = cal_ratio(state_dict[m1_key], state_dict[k])
+                    m1_quant, scales = qdq_weight(state_dict[m1_key], quant_bit=8)
+                    quant_weight, mins, maxs = asymmetry_qdq_weight(ratio, quant_bit=8)
+                    state_dict[m1_key] = m1_quant
+                    scales_dict[m1_key + SYMMETRY_QUANT_SCALE] = scales
+                    scales_dict[k + ASYMMETRY_QUANT_SCALE_MIN] = mins
+                    scales_dict[k + ASYMMETRY_QUANT_SCALE_MAX] = maxs
+                elif not momentum1:
+                    quant_weight = state_dict[k]
+            elif ckpt_quant_stage == "O2":
+                # m1: bw-wint4, 1/(sqrt(m2)+eps): bw-wint4
+                if momentum2:
+                    # skip norm-like parameters
+                    if len(state_dict[k].shape) < 2:
+                        continue
+                    # m1: m1_quant_weight, m2: ratio
+                    m1_key = k.split("/")[0] + "/" + MOMENT1_KEYNAME
+                    ratio = cal_ratio(state_dict[m1_key], state_dict[k])
+                    m1_quant, m1_scales = group_wise_quant_dequant(state_dict[m1_key], quant_bits=4, symmetry=True)
+                    quant_weight, r_mins, r_maxs = group_wise_quant_dequant(ratio, quant_bits=4)
+                    quant_weight = merge_int4(m1_quant, quant_weight)
+                    scales_dict[m1_key + SYMMETRY_QUANT_SCALE] = m1_scales
+                    scales_dict[k + ASYMMETRY_QUANT_SCALE_MIN] = r_mins
+                    scales_dict[k + ASYMMETRY_QUANT_SCALE_MAX] = r_maxs
+                    del_key.append(m1_key)
+                elif not momentum1:
+                    quant_weight = state_dict[k]
+
+            if quant_weight is not None:
+                state_dict[k] = quant_weight
+
+        for k in del_key:
+            state_dict.pop(k, None)
+
+        state_dict.update(scales_dict)
+    logger.info(f"Unified checkpoint quantization done, stage {ckpt_quant_stage}.")
+
+    return state_dict

paddlenlp/trainer/trainer.py

@@ -79,6 +79,7 @@ from ..data import (
     DataCollatorWithPadding,
     DistDataLoader,
     default_data_collator,
+    init_dataloader_comm_group,
 )
 from ..peft import LoRAModel, PrefixModelForCausalLM, VeRAModel
 
@@ -110,10 +111,10 @@ from ..utils.env import (
 )
 from ..utils.import_utils import is_datasets_available, is_paddle_cuda_available
 from ..utils.log import logger
+from ..utils.tools import get_env_device
 from .argparser import strtobool
 from .integrations import get_reporting_integration_callbacks
 from .plugins.timer import RuntimeTimer, get_timers, set_timers
-from .plugins.unified_checkpoint import UnifiedCheckpointHandler
 from .trainer_callback import (
     CallbackHandler,
     DefaultFlowCallback,
@@ -142,8 +143,10 @@ from .trainer_utils import (  # set_hyrbid_parallel_seed,
     set_seed,
     should_skip_data,
     speed_metrics,
+    split_parallel_config,
 )
 from .training_args import TrainingArguments
+from .unified_checkpoint import UnifiedCheckpointHandler
 from .utils import reshard as reshard_util
 from .utils.async_save import AsyncSaver
 from .utils.helper import (  # nested_truncate,
@@ -421,6 +424,10 @@ class Trainer:
 
             model.apply(fn)
 
+        self._pp_data_group = None
+        if self.args.pipeline_parallel_degree > 1 and self.args.distributed_dataloader:
+            self._pp_data_group = init_dataloader_comm_group()
+
         default_label_names = (
             ["start_positions", "end_positions"]
             if "QusetionAnswering" in type(self.model).__name__ or "UIE" in type(self.model).__name__
@@ -570,7 +577,9 @@ class Trainer:
         # Load potential model checkpoint
         if isinstance(resume_from_checkpoint, bool) and resume_from_checkpoint:
             uc_async_save = self.args.unified_checkpoint and "async_save" in self.args.unified_checkpoint_config
-            resume_from_checkpoint = get_last_checkpoint(self.args.output_dir, uc_async_save)
+            resume_from_checkpoint = get_last_checkpoint(
+                self.args.output_dir, signal_folder=self.args.output_signal_dir, uc_async_save=uc_async_save
+            )
             if resume_from_checkpoint is None:
                 raise ValueError(f"No valid checkpoint found in output directory ({self.args.output_dir})")
 
@@ -585,9 +594,10 @@ class Trainer:
                 if use_unified_checkpoint:
                     self.unified_checkpoint_handler.load_unified_checkpoint(
                         self.model,
-                        self.optimizer,
                         resume_from_checkpoint,
                     )
+                    if isinstance(self.model, LoRAModel) and self.model.lora_config.loraga:
+                        self.model.reinit_base_model = True
                     logger.info(f"Loading model from {resume_from_checkpoint} using unified checkpoint.")
                     self.runtime_timer.stop()
                     return
@@ -598,6 +608,8 @@ class Trainer:
             or isinstance(self.model, VeRAModel)
         ):
             self._load_from_peft_checkpoint(resume_from_checkpoint)
+            if isinstance(self.model, LoRAModel) and self.model.lora_config.loraga:
+                self.model.reinit_base_model = True
             self.runtime_timer.stop()
             return
 
@@ -930,6 +942,7 @@ class Trainer:
         self.state.num_train_epochs = num_train_epochs
         self.state.is_local_process_zero = self.is_local_process_zero()
         self.state.is_world_process_zero = self.is_world_process_zero()
+        self.state.consumed_samples = 0
 
         self.control = self.callback_handler.on_train_begin(args, self.state, self.control)
 
@@ -1005,6 +1018,12 @@ class Trainer:
                         self._skip_steps_since_last_logged += 1
 
                         self.state.epoch = epoch + (step + 1) / steps_in_epoch
+                        self.state.consumed_samples = (
+                            self.state.global_step
+                            * args.per_device_train_batch_size
+                            * args.gradient_accumulation_steps
+                            * args.dataset_world_size
+                        )
 
                         if self.state.global_step == 1 and self.args.logging_first_step:
                             self.control.should_log = True
@@ -1054,12 +1073,12 @@ class Trainer:
                 if dp_master_grad:
                     is_no_sync = True
 
-                if is_no_sync:
-                    # Avoid unnecessary DDP synchronization since there will be no backward pass on this example.
-                    with model.no_sync():
+                sync_context = model.no_sync() if is_no_sync else contextlib.nullcontext()
+                with sync_context:
+                    if "step_control" in inspect.signature(self.training_step).parameters:
+                        tr_loss_step = self.training_step(model, inputs, step_control=step_control)
+                    else:
                         tr_loss_step = self.training_step(model, inputs)
-                else:
-                    tr_loss_step = self.training_step(model, inputs)
 
                 tr_loss += tr_loss_step
 
@@ -1134,6 +1153,10 @@ class Trainer:
                         args, self.state, self.control, scaler=self.scaler if self.do_grad_scaling else None
                     )
                     optimizer_was_run = True
+
+                    if self.args.offload_optim:
+                        self._reload_optimizer()
+
                     if self.do_grad_scaling:
                         if args.pipeline_parallel_degree > 1:
                             assert not self.args.use_expert_parallel, "pipeline moe not work under fp16"
@@ -1157,6 +1180,9 @@ class Trainer:
                     else:
                         self.optimizer.step()
 
+                    if self.args.offload_optim:
+                        self._offload_optimizer()
+
                     self.timers and self.timers("optimizer-step").stop()
 
                     if optimizer_was_run:
@@ -1177,6 +1203,12 @@ class Trainer:
 
                     self.state.global_step += 1
                     self.state.epoch = epoch + (step + 1) / steps_in_epoch
+                    self.state.consumed_samples = (
+                        self.state.global_step
+                        * args.per_device_train_batch_size
+                        * args.gradient_accumulation_steps
+                        * args.dataset_world_size
+                    )
                     self.control = self.callback_handler.on_step_end(args, self.state, self.control)
                     self._maybe_log_save_evaluate(tr_loss, model, epoch, ignore_keys_for_eval, inputs=inputs)
                     self._print_timer()
@@ -1228,7 +1260,6 @@ class Trainer:
                 if self.args.unified_checkpoint:
                     self.unified_checkpoint_handler.load_unified_checkpoint(
                         self.model,
-                        self.optimizer,
                         self.state.best_model_checkpoint,
                     )
                     if self.args.sharding_parallel_degree > 1 or self.args.data_parallel_degree > 1:
@@ -1276,7 +1307,6 @@ class Trainer:
         if self.args.unified_checkpoint:
             self.unified_checkpoint_handler.load_unified_checkpoint(
                 self.model,
-                self.optimizer,
                 self.state.best_model_checkpoint,
             )
             if self.args.sharding_parallel_degree > 1 or self.args.data_parallel_degree > 1:
@@ -1505,6 +1535,7 @@ class Trainer:
             train_dataset = self._remove_unused_columns(train_dataset, description="training")
         _DataLoader = DistDataLoader if self.args.distributed_dataloader else DataLoader
 
+        additional_configs = {}
         if is_iterable_dataset:  # For iterable dataset
             if self.args.dataset_world_size > 1 and train_dataset is not None:
                 train_dataset = IterableDatasetShard(
@@ -1517,9 +1548,7 @@ class Trainer:
 
             if self.args.distributed_dataloader:
                 logger.info("Training using DistDataLoader.")
-                additional_configs = {"is_iterable_dataset": True}
-            else:
-                additional_configs = {}
+                additional_configs = {"is_iterable_dataset": True, "pp_data_group": self._pp_data_group}
             return _DataLoader(
                 train_dataset,
                 batch_size=self.args.per_device_train_batch_size,
@@ -1531,17 +1560,18 @@ class Trainer:
             train_sampler = self._get_train_sampler()
             if self.args.distributed_dataloader:
                 logger.info("Training using DistDataLoader.")
+                additional_configs = {"pp_data_group": self._pp_data_group}
             return _DataLoader(
                 train_dataset,
                 batch_sampler=train_sampler,
                 collate_fn=self.data_collator,
                 num_workers=self.args.dataloader_num_workers,
+                **additional_configs,
             )
 
     def _get_eval_sampler(self, eval_dataset: Dataset):
         if eval_dataset is None or not has_length(eval_dataset):
             return None
-
         if self.args.world_size <= 1:
             return paddle.io.BatchSampler(
                 eval_dataset,
@@ -1550,21 +1580,28 @@ class Trainer:
                 drop_last=False,
             )
         else:
-            drop_last = False
             if self.args.pipeline_parallel_degree > 1:
-                drop_last = True
-                logger.warning(
-                    "In parallel mode, the batch_size is strictly checked. set DistributedBatchSampler drop_last=True."
-                )
+                # In pipeline parallelism, batch size will be strictly checked
+                # Use LastBatchPaddingSampler to pad the last batch with the first batch
+                from .trainer_utils import LastBatchPaddingSampler
 
-            return DistributedBatchSampler(
-                eval_dataset,
-                num_replicas=self.args.dataset_world_size,
-                rank=self.args.dataset_rank,
-                batch_size=self.args.per_device_eval_batch_size,
-                shuffle=False,
-                drop_last=drop_last,
-            )
+                return LastBatchPaddingSampler(
+                    eval_dataset,
+                    num_replicas=self.args.dataset_world_size,
+                    rank=self.args.dataset_rank,
+                    batch_size=self.args.per_device_eval_batch_size,
+                    shuffle=False,
+                    drop_last=False,
+                )
+            else:
+                return DistributedBatchSampler(
+                    eval_dataset,
+                    num_replicas=self.args.dataset_world_size,
+                    rank=self.args.dataset_rank,
+                    batch_size=self.args.per_device_eval_batch_size,
+                    shuffle=False,
+                    drop_last=False,
+                )
 
     def get_eval_dataloader(self, eval_dataset: Optional[Dataset] = None) -> DataLoader:
         """
@@ -1591,8 +1628,11 @@ class Trainer:
             eval_dataset = self._remove_unused_columns(eval_dataset, description="evaluation")
         _DataLoader = DistDataLoader if self.args.distributed_dataloader else DataLoader
 
+        additional_configs = {}
         if is_iterable_dataset:
-            if self.args.dataset_world_size > 1 and eval_dataset is not None:
+            if (
+                self.args.dataset_world_size > 1 or self.args.pipeline_parallel_degree > 1
+            ) and eval_dataset is not None:
                 eval_dataset = IterableDatasetShard(
                     eval_dataset,
                     batch_size=self.args.per_device_eval_batch_size,
@@ -1600,11 +1640,10 @@ class Trainer:
                     num_processes=self.args.dataset_world_size,
                     process_index=self.args.dataset_rank,
                 )
+
             if self.args.distributed_dataloader:
                 logger.info("Eval using DistDataLoader.")
-                additional_configs = {"eval": True, "is_iterable_dataset": True}
-            else:
-                additional_configs = {}
+                additional_configs = {"eval": True, "is_iterable_dataset": True, "pp_data_group": self._pp_data_group}
             return _DataLoader(
                 eval_dataset,
                 batch_size=self.args.per_device_eval_batch_size,
@@ -1616,14 +1655,12 @@ class Trainer:
             eval_sampler = self._get_eval_sampler(eval_dataset)
             if self.args.distributed_dataloader:
                 logger.info("Eval using DistDataLoader.")
-                additional_configs = {"eval": True}
-            else:
-                additional_configs = {}
+                additional_configs = {"eval": True, "pp_data_group": self._pp_data_group}
             return _DataLoader(
                 eval_dataset,
                 batch_sampler=eval_sampler,
                 collate_fn=self.data_collator,
-                num_workers=self.args.dataloader_num_workers,
+                num_workers=0,
                 **additional_configs,
             )
 
@@ -1651,6 +1688,7 @@ class Trainer:
             test_dataset = self._remove_unused_columns(test_dataset, description="test")
         _DataLoader = DistDataLoader if self.args.distributed_dataloader else DataLoader
 
+        additional_config = {}
         if is_iterable_dataset:
             if self.args.dataset_world_size > 1 and test_dataset is not None:
                 test_dataset = IterableDatasetShard(
@@ -1663,23 +1701,19 @@ class Trainer:
 
             if self.args.distributed_dataloader:
                 logger.info("Test using DistDataLoader.")
-                additional_config = {"eval": True, "is_iterable_dataset": True}
-            else:
-                additional_config = {}
+                additional_config = {"eval": True, "is_iterable_dataset": True, "pp_data_group": self._pp_data_group}
             return _DataLoader(
                 test_dataset,
                 batch_size=self.args.per_device_eval_batch_size * self.world_size,
                 collate_fn=self.data_collator,
-                num_workers=self.args.dataloader_num_workers,
+                num_workers=0,
                 **additional_config,
             )
         else:
             test_sampler = self._get_eval_sampler(test_dataset)
             if self.args.distributed_dataloader:
                 logger.info("Test using DistDataLoader.")
-                additional_config = {"eval": True}
-            else:
-                additional_config = {}
+                additional_config = {"eval": True, "pp_data_group": self._pp_data_group}
             # We use the same batch_size as for eval.
             return _DataLoader(
                 test_dataset,
@@ -1735,6 +1769,38 @@ class Trainer:
 
         return self.optimizer
 
+    def _apply_to_optimizer(self, action):
+        attributes = [
+            ("_accumulators", "_moment1_acc_str"),
+            ("_accumulators", "_moment2_acc_str"),
+            ("_master_weights",),
+            ("_accumulators_holder",),
+        ]
+
+        for attr in attributes:
+            if all(hasattr(self.optimizer, a) for a in attr):
+                target_attr = getattr(self.optimizer, attr[0])
+                if len(attr) == 2:
+                    target_attr = target_attr[getattr(self.optimizer, attr[1])]
+
+                for key, value in target_attr.items():
+                    if get_env_device() == "gpu":
+                        target_attr[key] = getattr(value, action)()
+                    else:
+                        target_attr[key] = getattr(value, "to")(action)
+
+    def _offload_optimizer(self):
+        if get_env_device() == "gpu":
+            self._apply_to_optimizer("pin_memory")
+        else:
+            self._apply_to_optimizer("cpu")
+
+    def _reload_optimizer(self):
+        if get_env_device() == "gpu":
+            self._apply_to_optimizer("cuda")
+        else:
+            self._apply_to_optimizer(get_env_device())
+
     def _load_rng_state(self, checkpoint):
         # Load RNG states from `checkpoint`
         if checkpoint is None:
@@ -1989,6 +2055,14 @@ class Trainer:
                 self.optimizer = mix_precision_utils.MixPrecisionOptimizer(self.optimizer)
             self.optimizer = fleet.distributed_optimizer(self.optimizer)
 
+            if (
+                hasattr(self.args, "enable_sharding_comm_overlap")
+                and self.args.enable_sharding_comm_overlap
+                and self.args.unified_checkpoint
+                and "split_param" in split_parallel_config(self.args.sharding_parallel_config)
+            ):
+                model.register_sharding_comm_overlap_hook(self.optimizer)
+
         # No pipeline mode, sharding only
         if not in_pipeline_parallel_mode and in_sharding_parallel_mode:
             # Sharded DDP!
@@ -2191,7 +2265,9 @@ class Trainer:
         else:
             return False
 
-    def training_step(self, model: nn.Layer, inputs: Dict[str, Union[paddle.Tensor, Any]]) -> paddle.Tensor:
+    def training_step(
+        self, model: nn.Layer, inputs: Dict[str, Union[paddle.Tensor, Any]], step_control=0
+    ) -> paddle.Tensor:
         """
         Perform a training step on a batch of inputs.
 
@@ -2283,7 +2359,11 @@ class Trainer:
 
         return loss.detach()
 
-    def save_model(self, output_dir: Optional[str] = None, merge_tensor_parallel: Optional[bool] = False):
+    def save_model(
+        self,
+        output_dir: Optional[str] = None,
+        merge_tensor_parallel: Optional[bool] = False,
+    ):
         """
         Will save the model, so you can reload it using `from_pretrained()`.
 
@@ -2293,26 +2373,22 @@ class Trainer:
         if output_dir is None:
             output_dir = self.args.output_dir
 
+        if PREFIX_CHECKPOINT_DIR in os.path.split(output_dir)[-1]:
+            signal_dir = os.path.join(self.args.output_signal_dir, os.path.split(output_dir)[-1])
+        else:
+            signal_dir = self.args.output_signal_dir
+
         if ShardingOption.FULL_SHARD in self.args.sharding:
             self.model_wrapped.get_all_parameters(convert2cpu=True)
 
         if self.args.should_save_model_state:
-            unified_checkpoint_config_backup = self.args.unified_checkpoint_config
-            # backup and remove unified_checkpoint_config for not trine stage
-            if not self.is_in_train:
-                self.args.unified_checkpoint_config = []
-
             self._save(output_dir=output_dir, merge_tensor_parallel=merge_tensor_parallel)
-
-            # recover unified_checkpoint_config for not trine stage
-            if not self.is_in_train:
-                self.args.unified_checkpoint_config = unified_checkpoint_config_backup
         else:
             if self.args.unified_checkpoint and "async_save" in self.args.unified_checkpoint_config:
-                os.makedirs(output_dir, exist_ok=True)
+                os.makedirs(signal_dir, exist_ok=True)
                 if self.is_in_train:
                     global_rank = paddle.distributed.get_rank() if paddle.distributed.get_world_size() > 1 else -1
-                    paddle.save(global_rank, os.path.join(output_dir, f".model_weight.done.{global_rank}"))
+                    paddle.save(global_rank, os.path.join(signal_dir, f".model_weight.done.{global_rank}"))
 
         if strtobool(os.getenv("FLAG_LLM_PDC", "False")):
             # save model_done file to ensure model is complete
@@ -2328,9 +2404,9 @@ class Trainer:
             and "async_save" in self.args.unified_checkpoint_config
             and not self.is_in_train
         ):
-            os.makedirs(output_dir, exist_ok=True)
+            os.makedirs(signal_dir, exist_ok=True)
             global_rank = paddle.distributed.get_rank() if paddle.distributed.get_world_size() > 1 else -1
-            paddle.save(self.state.global_step, os.path.join(output_dir, f".model_weight.done.{global_rank}"))
+            paddle.save(self.state.global_step, os.path.join(signal_dir, f".model_weight.done.{global_rank}"))
 
     def _filter_moe_no_sync_optimizer_params(self):
         """
@@ -2341,7 +2417,7 @@ class Trainer:
         filter_optimzier_state_dict = OrderedDict()
         param_names_in_master_weights = list(optimzier_state_dict["master_weights"].keys()) if self.args.bf16 else []
         filter_optimzier_state_dict["master_weights"] = OrderedDict()
-        for k, v in state_dict.items():
+        for _, v in state_dict.items():
             if getattr(v, "no_sync", False):
                 if v.name in param_names_in_master_weights:
                     filter_optimzier_state_dict["master_weights"][v.name] = optimzier_state_dict["master_weights"][
@@ -2360,8 +2436,10 @@ class Trainer:
         checkpoint_folder = f"{PREFIX_CHECKPOINT_DIR}-{self.state.global_step}"
 
         run_dir = self.args.output_dir
+        run_signal_dir = self.args.output_signal_dir
 
         output_dir = os.path.join(run_dir, checkpoint_folder)
+        signal_dir = os.path.join(run_signal_dir, checkpoint_folder)
 
         if isinstance(self.model, LoRAModel) and (self.model.quantized or self.args.pipeline_parallel_degree > 1):
             self.save_model(output_dir)
@@ -2384,6 +2462,7 @@ class Trainer:
                             self.model,
                             self.optimizer,
                             output_dir,
+                            signal_dir,
                         )
                     else:
                         if self.dp_group.rank > 0:  # this should only work for MoE saving
@@ -2406,10 +2485,13 @@ class Trainer:
                 else:
                     if self.args.unified_checkpoint and "async_save" in self.args.unified_checkpoint_config:
                         global_rank = paddle.distributed.get_rank() if paddle.distributed.get_world_size() > 1 else -1
-                        os.makedirs(output_dir, exist_ok=True)
-                        paddle.save(global_rank, os.path.join(output_dir, f".optimizer_weight.done.{global_rank}"))
-                        if "skip_save_model_weight" not in self.args.unified_checkpoint_config:
-                            paddle.save(global_rank, os.path.join(output_dir, f".master_weight.done.{global_rank}"))
+                        os.makedirs(signal_dir, exist_ok=True)
+                        paddle.save(global_rank, os.path.join(signal_dir, f".optimizer_weight.done.{global_rank}"))
+                        if (
+                            "skip_save_model_weight" not in self.args.unified_checkpoint_config
+                            or "remove_master_weight" not in self.args.unified_checkpoint_config
+                        ):
+                            paddle.save(global_rank, os.path.join(signal_dir, f".master_weight.done.{global_rank}"))
             if self.args.should_save or self.args.use_expert_parallel:
                 if not self.args.use_hybrid_parallel:
                     logger.info("Saving optimizer files.")
@@ -2418,6 +2500,7 @@ class Trainer:
                             self.model,
                             self.optimizer,
                             output_dir,
+                            signal_dir,
                         )
                     else:
                         if self.args.data_parallel_rank > 0 and self.args.use_expert_parallel:
@@ -2442,10 +2525,13 @@ class Trainer:
                 if self.args.unified_checkpoint and not self.args.use_hybrid_parallel:
                     if "async_save" in self.args.unified_checkpoint_config:
                         global_rank = paddle.distributed.get_rank() if paddle.distributed.get_world_size() > 1 else -1
-                        os.makedirs(output_dir, exist_ok=True)
-                        paddle.save(global_rank, os.path.join(output_dir, f".optimizer_weight.done.{global_rank}"))
-                        if "skip_save_model_weight" not in self.args.unified_checkpoint_config:
-                            paddle.save(global_rank, os.path.join(output_dir, f".master_weight.done.{global_rank}"))
+                        os.makedirs(signal_dir, exist_ok=True)
+                        paddle.save(global_rank, os.path.join(signal_dir, f".optimizer_weight.done.{global_rank}"))
+                        if (
+                            "skip_save_model_weight" not in self.args.unified_checkpoint_config
+                            or "remove_master_weight" not in self.args.unified_checkpoint_config
+                        ):
+                            paddle.save(global_rank, os.path.join(signal_dir, f".master_weight.done.{global_rank}"))
 
         self.runtime_timer.stop()
         # Determine the new best metric / best model checkpoint
@@ -2494,7 +2580,7 @@ class Trainer:
         # For hybrid parallel training, the checkpoint files maybe on different node.
         need_to_rotate_checkpoints = False
         if self.args.use_hybrid_parallel:
-            if self.dp_group.rank <= 0:
+            if self.dp_group.rank <= 0 or self.args.use_expert_parallel:
                 need_to_rotate_checkpoints = True
         else:
             need_to_rotate_checkpoints = self.args.should_save_model_state
@@ -2503,6 +2589,7 @@ class Trainer:
         need_to_rotate_checkpoints = need_to_rotate_checkpoints and self.args.local_rank == 0
         if need_to_rotate_checkpoints:
             self._rotate_checkpoints(use_mtime=True, output_dir=run_dir)
+            self._rotate_checkpoints(use_mtime=True, output_dir=run_signal_dir)
 
         if strtobool(os.getenv("FLAG_LLM_PDC", "False")) and not ("async_save" in self.args.unified_checkpoint_config):
             # save checkpoint_done file to ensure checkpoint is complete
@@ -2577,30 +2664,46 @@ class Trainer:
             # ignore_errors for shared disks between train nodes.
             shutil.rmtree(checkpoint, ignore_errors=True)
 
-    def _save(self, output_dir: Optional[str] = None, state_dict=None, merge_tensor_parallel=False):
+    def _save(
+        self,
+        output_dir: Optional[str] = None,
+        state_dict=None,
+        merge_tensor_parallel=False,
+    ):
         output_dir = output_dir if output_dir is not None else self.args.output_dir
         os.makedirs(output_dir, exist_ok=True)
         logger.info(f"Saving model checkpoint to {output_dir}")
+
+        # signal_dir is used for asynchronous saving situations.
+        signal_dir = self.args.output_signal_dir
+        if self.args.unified_checkpoint and "async_save" in self.args.unified_checkpoint_config:
+            if PREFIX_CHECKPOINT_DIR in os.path.split(output_dir)[-1]:
+                signal_dir = os.path.join(signal_dir, os.path.split(output_dir)[-1])
+            os.makedirs(signal_dir, exist_ok=True)
+            logger.info(f"Saving model checkpoint finish signal to {signal_dir}")
+
         # Save a trained model and configuration using `save_pretrained()`.
         # They can then be reloaded using `from_pretrained()`
 
-        local_rank = int(os.getenv("PADDLE_RANK_IN_NODE", 0))
         if (
             strtobool(os.getenv("FLAG_LLM_PDC", "False"))
-            and local_rank == 0
+            and paddle.distributed.get_rank() == 0
             and self.args.unified_checkpoint
             and "async_save" in self.args.unified_checkpoint_config
         ):
-            os.makedirs(self.args.logging_dir, exist_ok=True)
             world_size = paddle.distributed.get_world_size()
             save_info = {
                 "world_size": world_size,
                 "ignore_save_lr_and_optim": self.args.ignore_save_lr_and_optim,
                 "skip_save_model_weight": "skip_save_model_weight" in self.args.unified_checkpoint_config,
+                "remove_master_weight": "remove_master_weight" in self.args.unified_checkpoint_config,
             }
-            if not os.path.exists(os.path.join(self.args.logging_dir, "async_save_info.json")):
-                with open(os.path.join(self.args.logging_dir, "async_save_info.json"), "w") as f:
-                    json.dump(save_info, f)
+            if os.path.exists(
+                os.path.join(self.args.output_signal_dir, "async_save_info.json")
+            ):  # afs cannot overwrite
+                os.remove(os.path.join(self.args.output_signal_dir, "async_save_info.json"))
+            with open(os.path.join(self.args.output_signal_dir, "async_save_info.json"), "w") as f:
+                json.dump(save_info, f)
 
         if self.args.should_save:
             if self.tokenizer is not None:
@@ -2609,7 +2712,17 @@ class Trainer:
             paddle.save(self.args, os.path.join(output_dir, TRAINING_ARGS_NAME))
 
         if self.args.unified_checkpoint:
-            self.unified_checkpoint_handler.save_unified_checkpoint(self.model, self.optimizer, output_dir)
+            unified_checkpoint_config_backup = self.args.unified_checkpoint_config
+            # backup and remove unified_checkpoint_config for not trine stage
+            if not self.is_in_train:
+                self.args.unified_checkpoint_config = []
+
+            self.unified_checkpoint_handler.save_unified_checkpoint(self.model, self.optimizer, output_dir, signal_dir)
+
+            # recover unified_checkpoint_config for not trine stage
+            if not self.is_in_train:
+                self.args.unified_checkpoint_config = unified_checkpoint_config_backup
+
             return
 
         merge_tensor_parallel = merge_tensor_parallel and self.args.use_hybrid_parallel
@@ -2732,9 +2845,15 @@ class Trainer:
                 else:
                     opt_state_dict = None
             else:
+                model = self.model
+                if (
+                    hasattr(self.args, "enable_sharding_comm_overlap")
+                    and self.args.enable_sharding_comm_overlap
+                    and "split_param" in split_parallel_config(self.args.sharding_parallel_config)
+                ):
+                    model = self.model_wrapped
                 opt_state_dict = self.unified_checkpoint_handler.load_unified_optimizer(
-                    args=self.args,
-                    model=self.model,
+                    model=model,
                     optimizer=self.optimizer,
                     resume_from_checkpoint=checkpoint,
                 )
@@ -2776,6 +2895,11 @@ class Trainer:
                 self.scaler.load_state_dict(
                     paddle.load(distributed_file(os.path.join(checkpoint, SCALER_NAME)), return_numpy=True)
                 )
+
+        if self.args.offload_optim:
+            logger.info("Offloading optimizer state...")
+            self._offload_optimizer()
+
         self.runtime_timer.stop()
 
     def log(self, logs: Dict[str, float], **kwargs) -> None:
@@ -2891,8 +3015,10 @@ class Trainer:
         prediction_loss_only = prediction_loss_only if prediction_loss_only is not None else args.prediction_loss_only
 
         if self.args.pipeline_parallel_degree > 1:
+            from paddle.distributed.fleet.meta_parallel import PipelineLayer
+
             # Only accept wrapped model for pipeline_parallel mode
-            if self.model is self.model_wrapped:
+            if self.model is self.model_wrapped and isinstance(self.model_wrapped, PipelineLayer):
                 # NOTE(gongenlei): when do_train=False, do_eval=True, we need to wrap model for pipeline
                 self.model_wrapped = fleet.distributed_model(self.model_wrapped)
             model = self.model_wrapped
@@ -3058,7 +3184,9 @@ class Trainer:
 
         # Metrics!
         if self.compute_metrics is not None and all_preds is not None and all_labels is not None:
-            metrics = self.compute_metrics(EvalPrediction(predictions=all_preds, label_ids=all_labels))
+            # all_labels maybe is a tuple when prediction_steps output label_mask
+            batch_labels = all_labels[0] if isinstance(all_labels, (list, tuple)) else all_labels
+            metrics = self.compute_metrics(EvalPrediction(predictions=all_preds, label_ids=batch_labels))
         else:
             metrics = {}
 
@@ -3153,6 +3281,8 @@ class Trainer:
             else:
                 labels = None
             inputs = inputs.pop("input_ids")
+        # train & eval share the same p2p_helper, so clear it before and after each step
+        model._p2p_helper.clear_meta_cache()
 
         with paddle.no_grad():
             if has_labels:
@@ -3162,6 +3292,8 @@ class Trainer:
                 loss = loss.mean().detach()
             else:
                 raise ValueError("pipeline mode eval need label!")
+        # train & eval share the same p2p_helper, so clear it before and after each step
+        model._p2p_helper.clear_meta_cache()
 
         return (loss, None, labels)
 

paddlenlp/trainer/trainer_utils.py

@@ -240,13 +240,15 @@ PREFIX_CHECKPOINT_DIR = "checkpoint"
 _re_checkpoint = re.compile(r"^" + PREFIX_CHECKPOINT_DIR + r"\-(\d+)$")
 
 
-def _check_checkpoint_files(folder_path, world_size, ignore_save_lr_and_optim, skip_save_model_weight):
+def _check_checkpoint_files(
+    folder_path, world_size, ignore_save_lr_and_optim, skip_save_model_weight, remove_master_weight
+):
     files = os.listdir(folder_path)
     model_weight_files = [f for f in files if f.startswith(".model_weight")]
     a = len(model_weight_files) == world_size
     if not ignore_save_lr_and_optim:
         b = True
-        if not skip_save_model_weight:
+        if not skip_save_model_weight or not remove_master_weight:
             master_weight_file = [f for f in files if f.startswith(".master_weight")]
             b = len(master_weight_file) == world_size
         optimizer_file = [f for f in files if f.startswith(".optimizer_weight")]
@@ -256,7 +258,7 @@ def _check_checkpoint_files(folder_path, world_size, ignore_save_lr_and_optim, s
         return a
 
 
-def get_last_checkpoint(folder, uc_async_save=False):
+def get_last_checkpoint(folder, signal_folder=None, uc_async_save=False):
     content = os.listdir(folder)
     checkpoints = [
         path
@@ -266,6 +268,9 @@ def get_last_checkpoint(folder, uc_async_save=False):
     if len(checkpoints) == 0:
         return
 
+    if uc_async_save:
+        assert signal_folder is not None
+
     if strtobool(os.getenv("FLAG_LLM_PDC", "False")):
         for i in sorted(checkpoints, key=lambda x: int(_re_checkpoint.search(x).groups()[0]), reverse=True):
             current_path = os.path.join(folder, i)
@@ -275,11 +280,17 @@ def get_last_checkpoint(folder, uc_async_save=False):
                     return current_path
             else:
                 saving_info = paddle.load(distributed_file(os.path.join(current_path, ".saving_info")))
+                current_signal_path = os.path.join(signal_folder, i)
                 pre_world_size = saving_info.get("world_size", 1)
                 ignore_save_lr_and_optim = saving_info.get("ignore_save_lr_and_optim", False)
                 skip_save_model_weight = saving_info.get("skip_save_model_weight", False)
+                remove_master_weight = saving_info.get("remove_master_weight", False)
                 if _check_checkpoint_files(
-                    current_path, pre_world_size, ignore_save_lr_and_optim, skip_save_model_weight
+                    current_signal_path,
+                    pre_world_size,
+                    ignore_save_lr_and_optim,
+                    skip_save_model_weight,
+                    remove_master_weight,
                 ):
                     return current_path
         return
@@ -1029,6 +1040,66 @@ class IterableDatasetShard(IterableDataset):
             return math.ceil(len(self.dataset) / (self.batch_size * self.num_processes)) * self.batch_size
 
 
+class LastBatchPaddingSampler(paddle.io.DistributedBatchSampler):
+    """The sampler which pads the first batch to the last batch"""
+
+    def __iter__(self):
+        local_batch_size = self.batch_size * self._acc_steps
+        num_samples = len(self.dataset)
+        indices = np.arange(num_samples).tolist()
+        global_eval_batch_size = self.batch_size * self.nranks
+        last_batch_size = num_samples % global_eval_batch_size
+
+        # Padding the first batch if the last batch is not full
+        if last_batch_size > 0:
+            padding_size = global_eval_batch_size - last_batch_size
+            # Select the first batch of indices for padding
+            if global_eval_batch_size <= len(indices):
+                first_batch_idx = indices[:global_eval_batch_size]
+            else:
+                first_batch_idx = indices.copy()
+            while padding_size > 0:
+                # Repeatedly pad the indices until the padding size is fulfilled
+                if padding_size > len(first_batch_idx):
+                    indices += first_batch_idx
+                    padding_size -= len(first_batch_idx)
+                else:
+                    indices += first_batch_idx[:padding_size]
+                    padding_size = 0
+
+        # Update the total number of indices
+        self.total_size = len(indices)
+        if self.shuffle:
+            np.random.RandomState(self.epoch).shuffle(indices)
+            self.epoch += 1
+
+        # subsample
+        def _get_indices_by_batch_size(indices):
+            subsampled_indices = []
+            # Iterate over the indices and extract batches that belong to the current device
+            for i in range(
+                self.local_rank * self.batch_size,
+                len(indices),
+                self.batch_size * self.nranks,
+            ):
+                subsampled_indices.extend(indices[i : i + self.batch_size])
+
+            return subsampled_indices
+
+        if self.nranks > 1:
+            indices = _get_indices_by_batch_size(indices)
+
+        _sample_iter = iter(indices)
+        batch_indices = []
+        for idx in _sample_iter:
+            batch_indices.append(idx)
+            if len(batch_indices) == local_batch_size:
+                yield batch_indices
+                batch_indices = []
+        # Ensure that there are no leftover indices after batching
+        assert len(batch_indices) == 0
+
+
 def find_batch_size(tensors):
     """
     Find the first dimension of a tensor in a nested list/tuple/dict of tensors.
@@ -1122,3 +1193,11 @@ def should_skip_data(global_step, skip_data_intervals):
             skip_flag = True
             break
     return skip_flag
+
+
+def split_parallel_config(parallel_config):
+    if "," in parallel_config:
+        parallel_config = set(parallel_config.split(","))
+    else:
+        parallel_config = set(parallel_config.split(" "))
+    return parallel_config

paddlenlp/trainer/training_args.py

@@ -37,6 +37,7 @@ from .trainer_utils import (
     OptimizerNames,
     SchedulerType,
     ShardingOption,
+    split_parallel_config,
 )
 
 __all__ = [
@@ -436,6 +437,7 @@ class TrainingArguments:
         },
     )
     logging_dir: Optional[str] = field(default=None, metadata={"help": "VisualDL log dir."})
+    output_signal_dir: Optional[str] = field(default=None, metadata={"help": "Asynchronous saving signal dir."})
     logging_strategy: IntervalStrategy = field(
         default="steps",
         metadata={"help": "The logging strategy to use."},
@@ -839,11 +841,16 @@ class TrainingArguments:
                 "- skip_save_model_weight: do not save model weights when the masters weight exist\n"
                 "- master_weight_compatible: 1. if the master weights exist, only load when needed\n"
                 "                            2. if master weights does not exist, convert model weights to master weights when needed\n"
+                "- remove_master_weight: same with `master_weight_compatible`, use in checkpoint quantization.\n"
                 "- async_save: enable asynchronous saving checkpoints to disk\n"
                 "- enable_all_options: enable all optimization configurations\n"
             )
         },
     )
+    ckpt_quant_stage: str = field(
+        default="O0",
+        metadata={"help": "checkpoint quantization stage."},
+    )
     ignore_load_lr_and_optim: Optional[bool] = field(
         default=False,
         metadata={"help": "whether to ignore load optimizer and scheduler."},
@@ -871,6 +878,10 @@ class TrainingArguments:
         default=None,
         metadata={"help": "The intervals to skip, pass start global step and end global step at each interval"},
     )
+    offload_optim: Optional[bool] = field(
+        default=False,
+        metadata={"help": "Offload optimizer after optimizer.step()"},
+    )
 
     def __post_init__(self):
         env_local_rank = int(os.environ.get("PADDLE_RANK_IN_NODE", -1))
@@ -897,6 +908,10 @@ class TrainingArguments:
             self.logging_dir = os.path.join(self.output_dir, default_logdir())
         if self.logging_dir is not None:
             self.logging_dir = os.path.expanduser(self.logging_dir)
+        if self.output_signal_dir is None and self.output_dir is not None:
+            self.output_signal_dir = self.output_dir
+        if self.output_signal_dir is not None:
+            self.output_signal_dir = os.path.expanduser(self.output_signal_dir)
 
         if self.disable_tqdm is None:
             self.disable_tqdm = False  # logger.getEffectiveLevel() > logging.WARN
@@ -1064,13 +1079,6 @@ class TrainingArguments:
                 logger.warning("set amp_master_grad to false since amp is disabled.")
                 self.amp_master_grad = False
 
-        def split_parallel_config(parallel_config):
-            if "," in parallel_config:
-                parallel_config = set(parallel_config.split(","))
-            else:
-                parallel_config = set(parallel_config.split(" "))
-            return parallel_config
-
         # use_hybrid_parallel
         if self.use_hybrid_parallel:
 
@@ -1122,12 +1130,12 @@ class TrainingArguments:
                         or "enable_dp_comm_overlap" in pipeline_parallel_config
                     )
                     enable_dp_comm_overlap = using_comm_overlap and self.data_parallel_degree > 1
-                    enable_sharding_comm_overlap = using_comm_overlap and self.sharding_parallel_degree > 1
+                    self.enable_sharding_comm_overlap = using_comm_overlap and self.sharding_parallel_degree > 1
                     assert not (
-                        enable_dp_comm_overlap and enable_sharding_comm_overlap
+                        enable_dp_comm_overlap and self.enable_sharding_comm_overlap
                     ), "dp_comm_overlap and sharding_comm_overlap cannot be enabled at the same time"
 
-                    if enable_sharding_comm_overlap and not self.amp_master_grad:
+                    if self.enable_sharding_comm_overlap and not self.amp_master_grad:
                         raise ValueError(
                             "If `enable_sharding_comm_overlap` in pipeline_parallel_configs, `amp_master_grad` must be True."
                         )
@@ -1135,7 +1143,7 @@ class TrainingArguments:
                     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,
-                        "sharding_comm_overlap": enable_sharding_comm_overlap,
+                        "sharding_comm_overlap": self.enable_sharding_comm_overlap,
                         "enable_timer": "enable_timer" in pipeline_parallel_config,
                         "release_gradients": "enable_release_grads" in pipeline_parallel_config or self.release_grads,
                         "overlap_p2p_comm": "enable_overlap_p2p_comm" in pipeline_parallel_config,
@@ -1378,6 +1386,13 @@ class TrainingArguments:
                             f"but got logging_steps={self.logging_steps}."
                         )
 
+                    if "split_param" in sharding_parallel_config:
+                        if ShardingOption.SHARD_OP not in self.sharding:
+                            logger.warning("Only sharding stage1 support split_param.")
+                        assert (
+                            self.amp_master_grad
+                        ), "If `split_param` in sharding_parallel_config, `amp_master_grad` must be True."
+
                 fleet.init(is_collective=True, strategy=strategy)
                 logger.info(strategy)
 
@@ -1623,6 +1638,7 @@ class TrainingArguments:
                     if x not in [
                         "skip_save_model_weight",
                         "master_weight_compatible",
+                        "remove_master_weight",
                         "async_save",
                         "enable_all_options",
                         "ignore_merge_optimizer",

paddlenlp/trainer/unified_checkpoint/__init__.py

@@ -0,0 +1,15 @@
+# Copyright (c) 2024 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 .unified_checkpoint import UnifiedCheckpointHandler

paddlenlp/trainer/unified_checkpoint/async_handler.py

@@ -0,0 +1,268 @@
+# Copyright (c) 2024 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.
+"""Asynchronous unified checkpoint handler."""
+
+import multiprocessing
+import os
+import time
+from multiprocessing import shared_memory
+
+import paddle
+import paddle.distributed as dist
+
+from paddlenlp.transformers.utils import is_safetensors_available
+from paddlenlp.utils.log import logger
+
+if is_safetensors_available():
+    from safetensors.numpy import save_file as safe_save_file
+
+from paddlenlp.quantization.unified_checkpoint_quantization import (
+    quant_unified_optimizer,
+)
+
+from .shared_memory_utils import (
+    _read_state_dict_from_shm,
+    _traverse_copy_to_shm,
+    create_meta_dict,
+)
+
+__all__ = ["AsyncCheckpointHandler"]
+
+
+class AsyncCheckpointHandler:
+    def __init__(self, args):
+        # Mainly for asynchronous saving.
+        self.args = args
+        self.global_rank = paddle.distributed.get_rank() if paddle.distributed.get_world_size() > 1 else -1
+
+        self._shm_model_weight = None
+        self._shm_master_weight = None
+        self._shm_optimizer_weight = None
+        self._meta_dict_model = None
+        self._meta_dict_master_weight = None
+        self._meta_dict_optim = None
+        self._process_model_weight = None
+        self._process_master_weight = None
+        self._process_optimizer_weight = None
+        self._lock = None
+        self._shared_save_model_flag = None
+        self._shared_save_master_weight_flag = None
+        self._shared_save_optimizer_flag = None
+
+        if "async_save" in self.args.unified_checkpoint_config:
+            self._lock = multiprocessing.Lock()
+            self._shared_save_model_path = multiprocessing.Array("c", 100000)
+            self._shared_save_model_signal_path = multiprocessing.Array("c", 100000)
+            self._shared_save_master_weight_path = multiprocessing.Array("c", 100000)
+            self._shared_save_master_weight_signal_path = multiprocessing.Array("c", 100000)
+            self._shared_save_optimizer_path = multiprocessing.Array("c", 100000)
+            self._shared_save_optimizer_signal_path = multiprocessing.Array("c", 100000)
+            self._shared_save_model_flag = multiprocessing.Array("i", 1)
+            self._shared_save_master_weight_flag = multiprocessing.Array("i", 1)
+            self._shared_save_optimizer_flag = multiprocessing.Array("i", 1)
+
+    def _file_save_async_or_sync(
+        self, state_dict, path, signal_path=None, is_sync=True, state_dict_type="model_weight", ckpt_quant_stage="O0"
+    ):
+        if is_sync:
+            for k in list(state_dict.keys()):
+                if isinstance(state_dict[k], paddle.Tensor):
+                    state_dict[k] = state_dict.pop(k).cpu().numpy()
+
+            if state_dict_type == "optimizer_weight" and ckpt_quant_stage != "O0":
+                state_dict = quant_unified_optimizer(state_dict, state_dict_type, ckpt_quant_stage)
+            safe_save_file(state_dict, path, metadata={"format": "np"})
+        else:
+            if len(state_dict.keys()) == 0:
+                saved_signal_path = os.path.join(signal_path, f".{state_dict_type}.done.{self.global_rank}")
+                paddle.save(self.global_rank, saved_signal_path)
+                return
+
+            if state_dict_type == "model_weight":
+                if self._shm_model_weight is None:
+                    self._meta_dict_model, buffer_size = create_meta_dict(state_dict)
+                    self._shm_model_weight = shared_memory.SharedMemory(create=True, size=buffer_size)
+                shm_state_dict = self._shm_model_weight
+                meta_dict = self._meta_dict_model
+                shared_save_flag = self._shared_save_model_flag
+                shared_save_path = self._shared_save_model_path
+                shared_save_signal_path = self._shared_save_model_signal_path
+                if self._process_model_weight is None:
+                    self._process_model_weight = multiprocessing.Process(
+                        target=self._save_file_async_in_process,
+                        args=(
+                            meta_dict,
+                            self._shm_model_weight.name,
+                            self._shared_save_model_flag,
+                            self._shared_save_model_path,
+                            self._shared_save_model_signal_path,
+                            self._lock,
+                            state_dict_type,
+                            self.global_rank,
+                        ),
+                    )
+                    self._process_model_weight.start()
+                process = self._process_model_weight
+            elif state_dict_type == "master_weight":
+                if self._shm_master_weight is None:
+                    self._meta_dict_master_weight, buffer_size = create_meta_dict(state_dict)
+                    self._shm_master_weight = shared_memory.SharedMemory(create=True, size=buffer_size)
+                shm_state_dict = self._shm_master_weight
+                meta_dict = self._meta_dict_master_weight
+                shared_save_flag = self._shared_save_master_weight_flag
+                shared_save_path = self._shared_save_master_weight_path
+                shared_save_signal_path = self._shared_save_master_weight_signal_path
+                if self._process_master_weight is None:
+                    self._process_master_weight = multiprocessing.Process(
+                        target=self._save_file_async_in_process,
+                        args=(
+                            meta_dict,
+                            self._shm_master_weight.name,
+                            self._shared_save_master_weight_flag,
+                            self._shared_save_master_weight_path,
+                            self._shared_save_master_weight_signal_path,
+                            self._lock,
+                            "model_weight"
+                            if "skip_save_model_weight" in self.args.unified_checkpoint_config
+                            else state_dict_type,
+                            self.global_rank,
+                        ),
+                    )
+                    self._process_master_weight.start()
+                process = self._process_master_weight
+            elif state_dict_type == "optimizer_weight":
+                if self._shm_optimizer_weight is None:
+                    self._meta_dict_optim, buffer_size = create_meta_dict(state_dict)
+                    self._shm_optimizer_weight = shared_memory.SharedMemory(create=True, size=buffer_size)
+                shm_state_dict = self._shm_optimizer_weight
+                meta_dict = self._meta_dict_optim
+                shared_save_flag = self._shared_save_optimizer_flag
+                shared_save_path = self._shared_save_optimizer_path
+                shared_save_signal_path = self._shared_save_optimizer_signal_path
+                if self._process_optimizer_weight is None:
+                    self._process_optimizer_weight = multiprocessing.Process(
+                        target=self._save_file_async_in_process,
+                        args=(
+                            meta_dict,
+                            self._shm_optimizer_weight.name,
+                            self._shared_save_optimizer_flag,
+                            self._shared_save_optimizer_path,
+                            self._shared_save_optimizer_signal_path,
+                            self._lock,
+                            state_dict_type,
+                            self.global_rank,
+                            ckpt_quant_stage,
+                        ),
+                    )
+                    self._process_optimizer_weight.start()
+                process = self._process_optimizer_weight
+
+            while True:  # wait until no process is saving.
+                flag_value = shared_save_flag[0]
+                if flag_value == 0:
+                    break
+                if not process.is_alive():
+                    raise RuntimeError(f"The process that saves {state_dict_type} has been killed unexpectedly.")
+                time.sleep(0.5)
+                logger.info(f"Wait for the previous save process to finish saving {state_dict_type}")
+            # only save model weight or save master weight, we enter this loop.
+            self._reset_and_update(shared_save_path, path)
+            self._reset_and_update(shared_save_signal_path, signal_path)
+            _traverse_copy_to_shm(state_dict, meta_dict, shm_state_dict.buf)
+            with self._lock:
+                shared_save_flag[0] = 1
+
+    def _save_file_async_in_process(
+        self,
+        meta_dict,
+        shm_name,
+        shared_save_flag,
+        shared_save_path,
+        shared_save_signal_path,
+        lock,
+        state_dict_type,
+        global_rank,
+        ckpt_quant_stage="O0",
+    ):
+        shm = shared_memory.SharedMemory(name=shm_name)
+        while True:
+            flag_value = shared_save_flag[0]  # if process uses `spawn`, cannot read this value.
+            if flag_value == -1:  # stop process
+                break
+            if flag_value == 0:  # nothing to save
+                continue
+            if flag_value == 1:  # need to save
+                path = shared_save_path[:].decode("utf-8").rstrip("\x00")
+                signal_path = shared_save_signal_path[:].decode("utf-8").rstrip("\x00")
+                logger.info(f"Start to async save {path}")
+                state_dict = _read_state_dict_from_shm(meta_dict, shm)  # numpy array
+                if state_dict_type == "optimizer_weight" and ckpt_quant_stage != "O0":
+                    state_dict = quant_unified_optimizer(
+                        state_dict, state_dict_type, ckpt_quant_stage, async_save=True
+                    )  # ckpt quantization
+                safe_save_file(state_dict, path, {"format": "np"})
+                del state_dict
+                saved_signal_path = os.path.join(signal_path, f".{state_dict_type}.done.{global_rank}")
+                paddle.save(global_rank, saved_signal_path)
+                with lock:
+                    shared_save_flag[0] = 0
+            time.sleep(0.5)
+        shm.close()
+
+    def _reset_and_update(self, shared_array, new_value):
+        # clear array
+        for i in range(len(shared_array)):
+            shared_array[i] = b"\0"
+        # update array
+        encoded_value = new_value.encode("utf-8")
+        shared_array[: len(encoded_value)] = encoded_value
+
+    def unlink_shared_memory(self):
+        if not ("async_save" in self.args.unified_checkpoint_config):
+            return
+
+        if self._shared_save_model_flag is not None:
+            while self._shared_save_model_flag[0] > 0:  # async process is saving
+                if not self._process_model_weight.is_alive():
+                    raise RuntimeError("The process that saves model_weight has been killed unexpectedly.")
+                time.sleep(0.5)
+            self._shared_save_model_flag[0] = -1
+        if self._shared_save_master_weight_flag is not None:
+            while self._shared_save_master_weight_flag[0] > 0:
+                if not self._process_master_weight.is_alive():
+                    raise RuntimeError("The process that saves master_weight has been killed unexpectedly.")
+                time.sleep(0.5)
+            self._shared_save_master_weight_flag[0] = -1
+        if self._shared_save_optimizer_flag is not None:
+            while self._shared_save_optimizer_flag[0] > 0:
+                if not self._process_optimizer_weight.is_alive():
+                    raise RuntimeError("The process that saves optimizer_weight has been killed unexpectedly.")
+                time.sleep(0.5)
+            self._shared_save_optimizer_flag[0] = -1
+
+        if self._shm_model_weight is not None:
+            self._shm_model_weight.close()
+            self._shm_model_weight.unlink()
+            self._shm_model_weight = None
+        if self._shm_master_weight is not None:
+            self._shm_master_weight.close()
+            self._shm_master_weight.unlink()
+            self._shm_master_weight = None
+        if self._shm_optimizer_weight is not None:
+            self._shm_optimizer_weight.close()
+            self._shm_optimizer_weight.unlink()
+            self._shm_optimizer_weight = None
+
+        if paddle.distributed.get_world_size() > 1:
+            dist.barrier()

paddlenlp/trainer/unified_checkpoint/check_completion.py

@@ -0,0 +1,247 @@
+# Copyright (c) 2024 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.
+"""Unfied checkpoint check functions."""
+
+import json
+import os
+
+import paddle
+import paddle.distributed as dist
+from paddle.distributed import fleet
+
+from paddlenlp.trainer.utils.helper import distributed_file, distributed_isfile
+from paddlenlp.utils.env import (
+    PADDLE_MASTER_WEIGHTS_INDEX_NAME,
+    PADDLE_OPTIMIZER_INDEX_NAME,
+    SAFE_MASTER_WEIGHTS_INDEX_NAME,
+    SAFE_OPTIMIZER_INDEX_NAME,
+)
+from paddlenlp.utils.log import logger
+from paddlenlp.utils.nested import flatten_list
+
+from .utils import (
+    get_expected_state_dict,
+    is_sharding_split_param_mode,
+    select_model_weight_index,
+    update_master_weight_status,
+)
+
+__all__ = ["check_unified_checkpoint", "check_unified_optimizer"]
+
+
+def check_unified_checkpoint(args, model, resume_from_checkpoint, safe_serialization=False):
+    index_filename = select_model_weight_index(model, resume_from_checkpoint, safe_serialization, local=False)
+    index_filename = os.path.join(resume_from_checkpoint, index_filename)
+    # Find index json file and distribute this file in global group.
+    if distributed_isfile(index_filename):
+        distributed_file(index_filename)
+    else:
+        raise Exception(
+            f"Sorry, we can not find {index_filename}. This file should be appear at least on one machine."
+        )
+
+    with open(index_filename, "r") as f:
+        index = json.loads(f.read())
+    all_weight_filenames = sorted(set(index["weight_map"].values()))
+
+    # Get existed weight file list on current machine.
+    existed_filelist = []
+    existed_files = []
+    for filename in os.listdir(resume_from_checkpoint):
+        if filename in all_weight_filenames:
+            existed_files.append(filename)
+
+    # Gather all the existed files in global group.
+    dist.all_gather_object(existed_filelist, existed_files)
+    flatten_existed_filelist = flatten_list(existed_filelist)
+    diff_filelist = list(set(all_weight_filenames).difference(set(flatten_existed_filelist)))
+    if len(diff_filelist) != 0:
+        raise Exception(f"Sorry, the weight file list on the machines is not complete!, missing {diff_filelist}")
+
+    # To decide whether to load the checkpoint locally, or need to dynamically send tensors across machines.
+    local_resume = True
+    if args.dataset_rank == 0 or args.use_expert_parallel:
+        hcg = fleet.get_hybrid_communicate_group()
+        tp_group = hcg.get_model_parallel_group()
+        pp_group = hcg.get_pipe_parallel_group()
+        dp_group = hcg.get_data_parallel_group()
+        dp_rank = dp_group.rank if dp_group.nranks > 1 else 0
+
+        need_files = set()
+        state_dict = get_expected_state_dict(model)
+        for key in state_dict.keys():
+            filename = index["weight_map"][key]
+            # When using expert parallel, there's no need to check tensors with `no_sync=False` when dp_rank > 0.
+            if args.use_expert_parallel and dp_rank > 0 and not getattr(state_dict[key], "no_sync", False):
+                continue
+            need_files.add(filename)
+        diff_filelist = list(need_files.difference(set(existed_files)))
+        num_diff = paddle.to_tensor([len(diff_filelist)])
+        if tp_group.nranks > 1:
+            dist.all_reduce(num_diff, op=dist.ReduceOp.MAX, group=tp_group)
+        if pp_group.nranks > 1:
+            dist.all_reduce(num_diff, op=dist.ReduceOp.MAX, group=pp_group)
+        if args.use_expert_parallel and dp_group.nranks > 1:
+            dist.all_reduce(num_diff, op=dist.ReduceOp.MAX, group=dp_group)
+        if num_diff.item() == 0:
+            local_resume = True
+        else:
+            local_resume = False
+    local_resume = paddle.to_tensor([local_resume])
+    dist.all_reduce(local_resume, op=dist.ReduceOp.MIN)
+    local_resume = local_resume.item()
+    return local_resume
+
+
+def check_unified_optimizer(args, model, optimizer, resume_from_checkpoint, safe_serialization=False):
+    if not safe_serialization:
+        index_filename, index_filename_master_weights = PADDLE_OPTIMIZER_INDEX_NAME, PADDLE_MASTER_WEIGHTS_INDEX_NAME
+    else:
+        index_filename, index_filename_master_weights = SAFE_OPTIMIZER_INDEX_NAME, SAFE_MASTER_WEIGHTS_INDEX_NAME
+    index_filename = os.path.join(resume_from_checkpoint, index_filename)
+    index_filename_master_weights = os.path.join(resume_from_checkpoint, index_filename_master_weights)
+
+    # Find index json file and distribute the file in global group.
+    if distributed_isfile(index_filename):
+        distributed_file(index_filename)
+    else:
+        raise Exception(
+            f"Sorry, we can not find {index_filename}. This file should be appear at least on one machine."
+        )
+
+    with open(index_filename, "r") as f:
+        index = json.loads(f.read())
+    all_optimizer_filenames = sorted(set(index["weight_map"].values()))
+
+    has_master_weights = index["master_weights"]
+    # update has_master_weights and index_filename_master_weights
+    # 1. if the master weight exists, only has_master_weights is set True and loaded when needed
+    # 2. if master weight does not exist, convert model weight to master weight when needed
+    has_master_weights, index_filename_master_weights = update_master_weight_status(
+        args, optimizer, has_master_weights, safe_serialization
+    )
+    if has_master_weights:
+        index_filename_master_weights = os.path.join(resume_from_checkpoint, index_filename_master_weights)
+        if distributed_isfile(index_filename_master_weights):
+            distributed_file(index_filename_master_weights)
+        else:
+            raise Exception(
+                f"Sorry, we can not find {index_filename_master_weights}. This file should be appear at least on one machine."
+            )
+        with open(index_filename_master_weights, "r") as f:
+            index_mw = json.loads(f.read())
+        all_mw_filenames = sorted(set(index_mw["weight_map"].values()))
+
+    hcg = fleet.get_hybrid_communicate_group()
+    tp_group = hcg.get_model_parallel_group()
+    pp_group = hcg.get_pipe_parallel_group()
+    dp_group = hcg.get_data_parallel_group()
+    sharding_group = hcg.get_sharding_parallel_group()
+    sharding_rank = sharding_group.rank
+    dp_rank = dp_group.rank if dp_group.nranks > 1 else 0
+    model_state_dict = get_expected_state_dict(model)
+    struct2static_name_mappings = {k: v.name for k, v in model_state_dict.items()}
+
+    if is_sharding_split_param_mode(args):
+        # We do not check optimizer files completion for split_param, since it is very complicated. Directly support local resume.
+        logger.warning("We only support local resume for split_param mode, do not support dynamically loading.")
+        return True
+
+    if sharding_group.nranks > 1:
+        param2rank = optimizer._param2rank
+
+    def check_complete(all_filenames):
+        # Check whether the checkpoint files on machines are complete. If not complete, raise Exception.
+        existed_filelist = []
+        existed_files = []
+        for filename in os.listdir(resume_from_checkpoint):
+            if filename in all_filenames:
+                existed_files.append(filename)
+
+        dist.all_gather_object(existed_filelist, existed_files)
+        flatten_existed_filelist = flatten_list(existed_filelist)
+        diff_filelist = list(set(all_filenames).difference(set(flatten_existed_filelist)))
+        if len(diff_filelist) != 0:
+            raise Exception(
+                f"Sorry, the optimizer file list on `data_parallel_rank==0` machines is not complete!, missing {diff_filelist}"
+            )
+        return existed_files
+
+    def check_dynamic_load(args, weight_map, existed_files, is_master_weights=False, typename_set=None):
+        # To decide whether to load the checkpoint locally, or need to dynamically distribute the checkpoint.
+        local_resume = True
+        if args.data_parallel_rank == 0 or args.use_expert_parallel:
+            need_files = set()
+            state_dict = get_expected_state_dict(model)
+
+            for key in state_dict.keys():
+                if model._keys_to_ignore_on_load_missing is not None and key in model._keys_to_ignore_on_load_missing:
+                    continue
+                if sharding_group.nranks > 1:
+                    static_name = struct2static_name_mappings.get(key, None)
+                    param_rank = param2rank.get(static_name, None)
+                    if param_rank != sharding_rank:
+                        continue
+
+                # When using expert parallel, there's no need to check tensors with `no_sync=False` when dp_rank > 0.
+                if args.use_expert_parallel and dp_rank > 0 and not getattr(state_dict[key], "no_sync", False):
+                    continue
+
+                if is_master_weights and state_dict[key].dtype == paddle.float32:
+                    continue
+
+                if not is_master_weights:
+                    for type_name in typename_set:
+                        type_key = key + "/" + type_name
+                        filename = weight_map[type_key]
+                        need_files.add(filename)
+                else:
+                    filename = weight_map[key]
+                    need_files.add(filename)
+
+            diff_filelist = list(need_files.difference(set(existed_files)))
+            num_diff = paddle.to_tensor([len(diff_filelist)])
+            if tp_group.nranks > 1:
+                dist.all_reduce(num_diff, op=dist.ReduceOp.MAX, group=tp_group)
+            if pp_group.nranks > 1:
+                dist.all_reduce(num_diff, op=dist.ReduceOp.MAX, group=pp_group)
+            if sharding_group.nranks > 1:
+                dist.all_reduce(num_diff, op=dist.ReduceOp.MAX, group=sharding_group)
+            if args.use_expert_parallel and dp_group.nranks > 1:
+                dist.all_reduce(num_diff, op=dist.ReduceOp.MAX, group=dp_group)
+
+            if num_diff.item() == 0:
+                local_resume = True
+            else:
+                local_resume = False
+        local_resume = paddle.to_tensor([local_resume])
+        dist.all_reduce(local_resume, op=dist.ReduceOp.MIN)
+        return local_resume.item()
+
+    # check whether the optimizer checkpoint files are complete.
+    existed_files = check_complete(all_optimizer_filenames)
+    if has_master_weights:
+        existed_files_mw = check_complete(all_mw_filenames)
+    # get optimizer's param type name, like moment1_0.
+    typename_set = set()
+    for key in index["weight_map"].keys():
+        _, typename = key.split("/")
+        typename_set.add(typename)
+    local_resume = check_dynamic_load(
+        args, index["weight_map"], existed_files, is_master_weights=False, typename_set=typename_set
+    )
+    local_resume_rw = True
+    if has_master_weights:
+        local_resume_rw = check_dynamic_load(args, index_mw["weight_map"], existed_files_mw, is_master_weights=True)
+    return local_resume & local_resume_rw

paddlenlp/trainer/unified_checkpoint/load_dynamic.py

@@ -0,0 +1,490 @@
+# Copyright (c) 2024 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.
+"""Unified Checkpoint Dynamic Loading Functions"""
+
+import copy
+import json
+import os
+import sys
+
+import paddle
+import paddle.distributed as dist
+from paddle.distributed import fleet
+
+from paddlenlp.peft import LoRAModel, PrefixModelForCausalLM
+from paddlenlp.transformers.model_utils import _load_state_dict_into_model
+from paddlenlp.transformers.utils import device_guard, is_safetensors_available
+from paddlenlp.utils.env import (
+    PADDLE_MASTER_WEIGHTS_INDEX_NAME,
+    PADDLE_OPTIMIZER_INDEX_NAME,
+    SAFE_MASTER_WEIGHTS_INDEX_NAME,
+    SAFE_OPTIMIZER_INDEX_NAME,
+)
+from paddlenlp.utils.log import logger
+from paddlenlp.utils.nested import nested_copy
+
+if is_safetensors_available():
+    if sys.platform.startswith("win"):
+        from safetensors import safe_open
+    else:
+        from paddlenlp.utils.safetensors import fast_safe_open as safe_open
+
+from .utils import (
+    FP32_MASTER,
+    get_expected_state_dict,
+    mapping_optimizer_tp_actions,
+    optimizer_non_scaler_name,
+    optimizer_scalar_name,
+    select_model_weight_index,
+    update_master_weight_status,
+)
+
+__all__ = ["load_unified_checkpoint_dynamically", "load_unified_optimizer_dynamically"]
+
+
+def create_send_table(file_keyname_mappings, file_machine_mappings):
+    send_table = {}
+    global_rank = dist.get_rank()
+    local_rank = int(os.getenv("PADDLE_RANK_IN_NODE", 0))
+    local_device_count = int(os.getenv("PADDLE_LOCAL_SIZE"))
+    for filename, keys in file_keyname_mappings.items():
+        machine = file_machine_mappings[filename][0]
+        is_src = (global_rank // local_device_count) == machine
+        for i, key in enumerate(keys):
+            if is_src and local_rank == i % local_device_count:
+                send_table[key] = global_rank
+    dispatch_list = []
+    dist.all_gather_object(dispatch_list, send_table)
+    send_table = {}
+    for dl in dispatch_list:
+        send_table.update(dl)
+    return send_table
+
+
+def create_dispatch_table(args, model, file_keyname_mappings, file_machine_mappings):
+    """Create dispatch table for dynamically loading state dict.
+
+    Args:
+        args
+    """
+
+    hcg = fleet.get_hybrid_communicate_group()
+    tp_group = hcg.get_model_parallel_group()
+    tp_rank = tp_group.rank
+
+    # Create tensor receive table, contains {"key0": [global_rank, tp_rank], "key1": [global_rank, tp_rank]}
+    dispatch_list = []
+    recv_table = {}
+    if args.dataset_rank == 0:
+        state_dict = get_expected_state_dict(model)
+        for (k, v) in state_dict.items():
+            if hasattr(v, "is_distributed") and v.is_distributed:
+                recv_table[k] = [(dist.get_rank(), tp_rank)]
+            else:
+                recv_table[k] = [(dist.get_rank(), -1)]
+
+    # Gather receive table in global group.
+    dist.all_gather_object(dispatch_list, recv_table)
+    recv_table = {}
+    for dl in dispatch_list:
+        for key, value in dl.items():
+            if key not in recv_table:
+                recv_table[key] = value
+            else:
+                recv_table[key] += value
+
+    # Create send table, to decide which worker to send the key. Contains {"key0:" global_rank, "key1": global_rank, ...}
+    send_table = create_send_table(file_keyname_mappings, file_machine_mappings)
+
+    return send_table, recv_table
+
+
+def create_optimizer_dispatch_table(
+    args,
+    model,
+    optimizer,
+    file_keyname_mappings,
+    file_machine_mappings,
+    struct2static_name_mappings,
+    is_master_weights=False,
+    typename_set=None,
+):
+    hcg = fleet.get_hybrid_communicate_group()
+    tp_group = hcg.get_model_parallel_group()
+    sharding_group = hcg.get_sharding_parallel_group()
+    sharding_rank = sharding_group.rank
+    if sharding_group.nranks > 1:
+        param2rank = optimizer._param2rank
+    tp_rank = tp_group.rank
+
+    # Create receive table, contains {"param_key0": [global_rank, tp_rank], "param_key1": [global_rank, tp_rank]}
+    dispatch_list = []
+    recv_table = {}
+    if args.data_parallel_rank == 0:
+        state_dict = get_expected_state_dict(model)
+        for (k, v) in state_dict.items():
+            if sharding_group.nranks > 1:
+                static_name = struct2static_name_mappings[k]
+                param_rank = param2rank.get(static_name, None)
+                if param_rank != sharding_rank:
+                    continue
+            if is_master_weights:
+                if hasattr(v, "is_distributed") and v.is_distributed:
+                    recv_table[k] = [(dist.get_rank(), tp_rank)]
+                else:
+                    recv_table[k] = [(dist.get_rank(), -1)]
+            else:
+                for typename in typename_set:
+                    type_key = k + "/" + typename
+                    if typename in optimizer_non_scaler_name:
+                        if hasattr(v, "is_distributed") and v.is_distributed:
+                            recv_table[type_key] = [(dist.get_rank(), tp_rank)]
+                        else:
+                            recv_table[type_key] = [(dist.get_rank(), -1)]
+                    else:
+                        recv_table[type_key] = [(dist.get_rank(), -1)]
+
+    dist.all_gather_object(dispatch_list, recv_table)
+    recv_table = {}
+    for dl in dispatch_list:
+        for k, v in dl.items():
+            if k not in recv_table:
+                recv_table[k] = v
+            else:
+                recv_table[k] += v
+
+    # Create send table, to decide which worker to send the key. Contains {"param_key0:" 0, "param_key1": 1, ...}
+    send_table = create_send_table(file_keyname_mappings, file_machine_mappings)
+    return send_table, recv_table
+
+
+def get_file_mappings(index, resume_from_checkpoint):
+    file_keyname_mappings = {}
+    for k, v in index["weight_map"].items():
+        if v not in file_keyname_mappings:
+            file_keyname_mappings[v] = []
+        file_keyname_mappings[v].append(k)
+    for k in file_keyname_mappings.keys():
+        file_keyname_mappings[k] = sorted(file_keyname_mappings[k])
+
+    local_device_count = int(os.getenv("PADDLE_LOCAL_SIZE"))
+    local_rank = int(os.getenv("PADDLE_RANK_IN_NODE", 0))
+    global_rank = dist.get_rank()
+    file_machine_mappings = {}
+    for filename in file_keyname_mappings.keys():
+        if local_rank == 0 and os.path.exists(os.path.join(resume_from_checkpoint, filename)):
+            file_machine_mappings[filename] = [global_rank // local_device_count]
+    file_machine_list = []
+    dist.all_gather_object(file_machine_list, file_machine_mappings)
+    file_machine_mappings = {}
+    for mappings in file_machine_list:
+        for k, v in mappings.items():
+            if k not in file_machine_mappings:
+                file_machine_mappings[k] = v
+            else:
+                file_machine_mappings[k] += v
+    return file_keyname_mappings, file_machine_mappings
+
+
+def distributed_send_recv(
+    state_dict,
+    tp_actions,
+    send_table,
+    recv_table,
+    resume_from_checkpoint,
+    file_keyname_mappings,
+    file_machine_mappings,
+):
+
+    local_device_count = int(os.getenv("PADDLE_LOCAL_SIZE"))
+    global_rank = dist.get_rank()
+    for filename in file_keyname_mappings.keys():
+        machine = file_machine_mappings[filename][0]
+        is_src = global_rank // local_device_count == machine
+        if is_src:
+            f = safe_open(os.path.join(resume_from_checkpoint, filename), framework="np")
+
+        for key in file_keyname_mappings[filename]:
+            recv_info = recv_table[key]
+            recv_ranklist = [a for (a, _) in recv_info]
+            if is_src and global_rank == send_table[key]:
+                py_safe_slice_ = f.get_slice(key)
+                # send
+                if key in tp_actions:
+                    weight = tp_actions[key](py_safe_slice_)
+                    # copy weight to GPU
+                    for j in range(len(weight)):
+                        with device_guard():
+                            weight[j] = paddle.Tensor(weight[j], zero_copy=True)
+                        weight[j] = weight[j]._copy_to(paddle.framework._current_expected_place(), False)
+
+                    for recv_rank, split_index in recv_info:
+                        if recv_rank == global_rank:
+                            state_dict[key] = weight[split_index]
+                        else:
+                            dist.stream.send(weight[split_index], dst=recv_rank)
+                else:
+                    # no need to tp split
+                    weight = py_safe_slice_[:]
+                    with device_guard():
+                        weight = paddle.Tensor(weight, zero_copy=True)
+                    weight = weight._copy_to(paddle.framework._current_expected_place(), False)
+                    for recv_rank, _ in recv_info:
+                        if recv_rank == global_rank:
+                            state_dict[key] = weight
+                        else:
+                            dist.stream.send(weight, dst=recv_rank)
+
+            if global_rank != send_table[key] and global_rank in recv_ranklist:
+                dist.stream.recv(state_dict[key], src=send_table[key])
+
+        if is_src:
+            f.__exit__(None, None, None)
+
+    return state_dict
+
+
+def load_unified_checkpoint_dynamically(args, model, resume_from_checkpoint, safe_serialization=False):
+    index_filename = select_model_weight_index(model, resume_from_checkpoint, safe_serialization, local=False)
+    index_filename = os.path.join(resume_from_checkpoint, index_filename)
+
+    with open(index_filename, "r") as f:
+        index = json.loads(f.read())
+
+    # `file_keyname_mappings` indicates which keys each file contains. For example, {"model-00001-of-00002.safetensors": ["llama.embed_tokens.weight", "llama.layers.0.self_attn.q_proj.weight", ...]}
+    # `file_machine_mappings` indicates the machine where the files appear. For example, {"model-00001-of-00002.safetensors": [machine_0, machine_1], "model-00002-of-00002.safetensors": [machine_0]}
+    file_keyname_mappings, file_machine_mappings = get_file_mappings(index, resume_from_checkpoint)
+
+    logger.debug("Creating dispatch table for unified checkpoint load ...")
+    # Get send_table and recv_table. The send table indicates which workers are responsible for sending tensors, and the recv table indicates which workers should receive the tensors.
+    send_table, recv_table = create_dispatch_table(
+        args,
+        model,
+        file_keyname_mappings,
+        file_machine_mappings,
+    )
+
+    # Get all the keys that are splited by tensor parallelism.
+    all_tp_keys = set()
+    for k, v in recv_table.items():
+        if v[0][1] != -1:
+            all_tp_keys.add(k)
+
+    config_revise = copy.deepcopy(model.config)
+    config_revise.tensor_parallel_rank = None
+    if len(all_tp_keys) == 0:
+        tp_actions = {}
+    else:
+        # Get corresponding tensor parallel actions.
+        if isinstance(model, LoRAModel) or isinstance(model, PrefixModelForCausalLM):
+            tp_actions = model._get_tensor_parallel_convert_actions(
+                set(all_tp_keys), is_split=True, ignore_error=True, config=config_revise
+            )
+        else:
+            tp_actions = model.get_tensor_parallel_convert_actions(config_revise, all_tp_keys, ignore_error=True)
+
+    logger.debug("Distributed send recv for state dict load ...")
+    # Distribute the checkpoint tensor dynamically, using the `send_table` and `recv_table` we create before.
+    state_dict = distributed_send_recv(
+        get_expected_state_dict(model),
+        tp_actions,
+        send_table,
+        recv_table,
+        resume_from_checkpoint,
+        file_keyname_mappings,
+        file_machine_mappings,
+    )
+    dist.barrier()
+    logger.debug("Setting state dict into model ...")
+    error_msgs = _load_state_dict_into_model(model, state_dict, "")
+    if len(error_msgs) > 0:
+        error_msg = "\n\t".join(error_msgs)
+        raise RuntimeError(f"Error(s) in loading dynamic state_dict for {model.__class__.__name__}:\n\t{error_msg}")
+
+
+def load_unified_optimizer_dynamically(args, model, optimizer, resume_from_checkpoint, safe_serialization=False):
+    optim_state_dict = nested_copy(optimizer.state_dict())
+    if "master_weights" in optim_state_dict.keys():
+        optim_state_dict.pop("master_weights")
+
+    if safe_serialization:
+        index_filename, index_filename_mw = SAFE_OPTIMIZER_INDEX_NAME, SAFE_MASTER_WEIGHTS_INDEX_NAME
+    else:
+        index_filename, index_filename_mw = PADDLE_OPTIMIZER_INDEX_NAME, PADDLE_MASTER_WEIGHTS_INDEX_NAME
+
+    with open(os.path.join(resume_from_checkpoint, index_filename), "r") as f:
+        index = json.loads(f.read())
+
+    # `file_keyname_mappings` indicates which keys each file contains. For example, {"optimizer-00001-of-00002.safetensors": ["llama.embed_tokens.weight/moment1_0", "llama.layers.1.mlp.gate_proj.weight/moment1_0", ...]}
+    # `file_machine_mappings` indicates the machine where the files appear. For example, {"optimizer-00001-of-00002.safetensors": [machine_0, machine_1], "optimizer-00002-of-00002.safetensors": [machine_0]}
+    file_keyname_mappings, file_machine_mappings = get_file_mappings(index, resume_from_checkpoint)
+
+    has_master_weights = index["master_weights"]
+    # update has_master_weights and index_filename_master_weights
+    # 1. if the master weights exists, only has_master_weights is set True and load master weights when needed
+    # 2. if master weights does not exist, convert model weights to master weights when needed
+    has_master_weights, index_filename_mw = update_master_weight_status(
+        args, optimizer, has_master_weights, safe_serialization
+    )
+
+    if has_master_weights:
+        with open(os.path.join(resume_from_checkpoint, index_filename_mw), "r") as f:
+            index_mw = json.loads(f.read())
+        file_keyname_mappings_mw, file_machine_mappings_mw = get_file_mappings(index_mw, resume_from_checkpoint)
+
+    # Get optimizer param type name, like moment1_0, moment2_0, beta1_pow_acc_0.
+    typename_set = set()
+    for key in index["weight_map"].keys():
+        _, typename = key.split("/")
+        typename_set.add(typename)
+
+    model_state_dict = get_expected_state_dict(model)
+    struct2static_name_mappings = {k: v.name for k, v in model_state_dict.items()}
+    static2struct_name_mappings = {v.name: k for k, v in model_state_dict.items()}
+    # Get send_table and recv_table. The send table indicates which workers are responsible for sending tensors, and the recv table indicates which workers should receive the tensors.
+    send_table, recv_table = create_optimizer_dispatch_table(
+        args,
+        model,
+        optimizer,
+        file_keyname_mappings,
+        file_machine_mappings,
+        struct2static_name_mappings,
+        is_master_weights=False,
+        typename_set=typename_set,
+    )
+    if has_master_weights:
+        send_table_mw, recv_table_mw = create_optimizer_dispatch_table(
+            args,
+            model,
+            optimizer,
+            file_keyname_mappings_mw,
+            file_machine_mappings_mw,
+            struct2static_name_mappings,
+            is_master_weights=True,
+        )
+
+    # Initialize optimizer state dict.
+    hcg = fleet.get_hybrid_communicate_group()
+    sharding_group = hcg.get_sharding_parallel_group()
+    if sharding_group.nranks > 1:
+        param2rank = optimizer._param2rank
+    optim_state_dict_mw = {}
+
+    def check_optimizer_param(parameter):
+        if sharding_group.nranks > 1:
+            param_rank = param2rank.get(parameter.name, None)
+            if param_rank != sharding_group.rank:
+                return False
+        if parameter.stop_gradient:
+            return False
+        return True
+
+    optimizer_keys_with_shape = []
+    if isinstance(optimizer._parameter_list[0], dict):
+        for param_group in optimizer._parameter_list:
+            # If parameter groups are set, there must be `params` key. This is guaranteed by the optimizer's initialization code.
+            for parameter in param_group["params"]:
+                if check_optimizer_param(parameter):
+                    optimizer_keys_with_shape.append((parameter.name, parameter.shape))
+    else:
+        for parameter in optimizer._parameter_list:
+            if check_optimizer_param(parameter):
+                optimizer_keys_with_shape.append((parameter.name, parameter.shape))
+
+    # see how to change
+    for static_name, shape in optimizer_keys_with_shape:
+        k = static2struct_name_mappings[static_name]
+        for typename in typename_set:
+            new_k = k + "/" + typename
+            if typename in optimizer_scalar_name:
+                optim_state_dict[new_k] = paddle.empty([1], dtype="float32")
+            else:
+                optim_state_dict[new_k] = paddle.empty(shape, dtype="float32")
+        if has_master_weights:
+            optim_state_dict_mw[k] = paddle.empty(shape, dtype="float32")
+
+    # Get all the keys that are splited by tensor parallelism.
+    all_tp_keys = set()
+    for k, v in recv_table.items():
+        structure_name, typename = k.split("/")
+        if typename in optimizer_non_scaler_name:
+            if v[0][1] != -1:
+                all_tp_keys.add(structure_name)
+
+    # Get corresponding tensor parallel actions.
+    config_revise = copy.deepcopy(model.config)
+    config_revise.tensor_parallel_rank = None
+    if len(all_tp_keys) == 0:
+        tp_actions = {}
+    else:
+        if isinstance(model, LoRAModel) or isinstance(model, PrefixModelForCausalLM):
+            tp_actions = model._get_tensor_parallel_convert_actions(
+                set(all_tp_keys), is_split=True, ignore_error=True, config=config_revise
+            )
+        else:
+            tp_actions = model.get_tensor_parallel_convert_actions(config_revise, all_tp_keys, ignore_error=True)
+    optimizer_keys = list(index["weight_map"].keys())
+    optimizer_tp_actions = mapping_optimizer_tp_actions(tp_actions, optimizer_keys)
+    if has_master_weights:
+        optimizer_tp_actions.update(tp_actions)
+
+    # Distribute the optimizer checkpoint dynamically, using the `send_table` and `recv_table` we create before.
+    optim_state_dict = distributed_send_recv(
+        optim_state_dict,
+        optimizer_tp_actions,
+        send_table,
+        recv_table,
+        resume_from_checkpoint,
+        file_keyname_mappings,
+        file_machine_mappings,
+    )
+    dist.barrier()
+    if has_master_weights:
+        optim_state_dict_mw = distributed_send_recv(
+            optim_state_dict_mw,
+            optimizer_tp_actions,
+            send_table_mw,
+            recv_table_mw,
+            resume_from_checkpoint,
+            file_keyname_mappings_mw,
+            file_machine_mappings_mw,
+        )
+        dist.barrier()
+
+    # Rename optimizer state dict.
+    for key in list(optim_state_dict.keys()):
+        if key == "LR_Scheduler":
+            continue
+        key_name = key.split("/")
+        static_name = struct2static_name_mappings[key_name[0]]
+        if has_master_weights:
+            if model_state_dict[key_name[0]].dtype != paddle.float32:
+                key_name = "_".join([static_name, FP32_MASTER, key_name[1]])
+            else:
+                key_name = "_".join([static_name, key_name[1]])
+        else:
+            key_name = "_".join([static_name, key_name[1]])
+        optim_state_dict[key_name] = optim_state_dict.pop(key)
+        optim_state_dict[key_name].name = key_name
+
+    if has_master_weights:
+        optim_state_dict["master_weights"] = {}
+        for key in list(optim_state_dict_mw.keys()):
+            static_name = struct2static_name_mappings[key]
+            optim_state_dict["master_weights"][static_name] = optim_state_dict_mw.pop(key)
+            optim_state_dict["master_weights"][static_name].name = "_".join([static_name, FP32_MASTER])
+
+    if args.data_parallel_rank == 0:
+        return optim_state_dict
+    return None

paddlenlp/trainer/unified_checkpoint/load_local.py

@@ -0,0 +1,297 @@
+# Copyright (c) 2024 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.
+"""Unfied checkpoint locally loading functions."""
+
+import gc
+import json
+import os
+
+import paddle
+from tqdm.auto import tqdm
+
+from paddlenlp.peft import LoRAModel, PrefixModelForCausalLM
+from paddlenlp.transformers.model_utils import (
+    _load_state_dict_into_model,
+    faster_set_state_dict,
+    load_state_dict,
+)
+from paddlenlp.transformers.utils import get_checkpoint_shard_files
+from paddlenlp.utils.env import (
+    PADDLE_MASTER_WEIGHTS_INDEX_NAME,
+    PADDLE_OPTIMIZER_INDEX_NAME,
+    SAFE_MASTER_WEIGHTS_INDEX_NAME,
+    SAFE_OPTIMIZER_INDEX_NAME,
+)
+from paddlenlp.utils.log import logger
+from paddlenlp.utils.nested import nested_copy
+
+from .sharding_split_param_utils import load_unified_optimizer_split_param
+from .utils import (
+    FP32_MASTER,
+    get_expected_keys,
+    get_expected_state_dict,
+    get_optimizer_shard_files,
+    is_sharding_split_param_mode,
+    mapping_optimizer_tp_actions,
+    select_model_weight_index,
+    update_master_weight_status,
+)
+
+__all__ = ["load_unified_checkpoint_locally", "load_unified_optimizer_locally"]
+
+
+def load_unified_checkpoint_locally(args, model, resume_from_checkpoint: str, safe_serialization=False):
+    """
+    Only dataset_rank == 0 or using expert parallel can enter this function.
+    """
+    index_filename = select_model_weight_index(model, resume_from_checkpoint, safe_serialization, local=True)
+
+    resolved_archive_file, sharded_metadata = get_checkpoint_shard_files(
+        pretrained_model_name_or_path=resume_from_checkpoint,
+        index_filename=os.path.join(resume_from_checkpoint, index_filename),
+    )
+    loaded_keys = sharded_metadata["all_checkpoint_keys"]
+
+    model_state_dict = get_expected_state_dict(model)
+    # If using expert parallel, when dp_rank > 0, need to modify the expected_keys here.
+    if not args.use_expert_parallel or (args.use_expert_parallel and args.data_parallel_rank == 0):
+        expected_keys = set(list(model_state_dict.keys()))
+    else:
+        expected_keys = set()
+        for key in model_state_dict.keys():
+            if getattr(model_state_dict[key], "no_sync", False):
+                expected_keys.add(key)
+    missing_keys = expected_keys - set(loaded_keys)
+
+    use_fast_set = True
+    if isinstance(model, LoRAModel) or isinstance(model, PrefixModelForCausalLM):
+        use_fast_set = False
+
+    if len(missing_keys) > 0:
+        raise ValueError(f"missing_keys: {missing_keys}")
+
+    def _remove_unused_keys(
+        state_dict,
+        model_state_dict,
+    ):
+        unused_keys = set(state_dict.keys()) - set(model_state_dict.keys())
+        for unused_key in unused_keys:
+            del state_dict[unused_key]
+        return unused_keys
+
+    # This should always be a list but, just to be sure.
+    if not isinstance(resolved_archive_file, list):
+        resolved_archive_file = [resolved_archive_file]
+
+    error_msgs = []
+
+    if len(resolved_archive_file) > 1:
+        resolved_archive_file = tqdm(resolved_archive_file, desc="Loading checkpoint shards")
+
+    for shard_file in resolved_archive_file:
+        # TODO: check if  no expected_keys in shard_file, then don't load it
+        if expected_keys.isdisjoint(sharded_metadata["file_map"][os.path.split(shard_file)[-1]]):
+            continue
+
+        pre_tensor_parallel_split = False
+        if shard_file.endswith(".safetensors") and model.config.tensor_parallel_degree > 1:
+            pre_tensor_parallel_split = True
+            assert loaded_keys is not None, "loaded_keys is not None."
+            if isinstance(model, LoRAModel) or isinstance(model, PrefixModelForCausalLM):
+                tp_actions = model._get_tensor_parallel_convert_actions(
+                    set(loaded_keys), is_split=True, ignore_error=True
+                )
+            else:
+                tp_actions = model.get_tensor_parallel_convert_actions(model.config, loaded_keys, ignore_error=True)
+        # Here we use expected_keys to optimize weights loading for pipeline model. Only works for safetensors
+        state_dict = load_state_dict(
+            shard_file, tp_actions if pre_tensor_parallel_split else None, expected_keys, device="expected"
+        )
+
+        if not pre_tensor_parallel_split:
+            # Since we load all keys but we only need one of pipeline stages
+            _ = _remove_unused_keys(state_dict, model_state_dict)
+
+        if model.config.tensor_parallel_degree > 1 and not pre_tensor_parallel_split:
+            logger.info("Converting state_dict to Tensor Parallel Format")
+            # ignore error for multi shard, since only parts of data
+            state_dict = model.convert_tensor_parallel(
+                None, model.config, state_dict=state_dict, ignore_error=len(resolved_archive_file) > 1
+            )
+
+        if use_fast_set:
+            error_msgs += faster_set_state_dict(model, state_dict, strict_dtype=False)
+        else:
+            error_msgs += _load_state_dict_into_model(model, state_dict, "")
+
+        # force memory release
+        del state_dict
+        # gc.collect()
+
+    if len(error_msgs) > 0:
+        error_msg = "\n\t".join(error_msgs)
+        if " but the expected shape is" in error_msg:
+            error_msg += (
+                "\n\tYou may consider adding `ignore_mismatched_sizes=True` in the model `from_pretrained` method."
+            )
+        raise RuntimeError(f"Error(s) in loading state_dict for {model.__class__.__name__}:\n\t{error_msg}")
+
+
+def load_unified_optimizer_locally(args, model, optimizer, resume_from_checkpoint, safe_serialization=False):
+    # Special process with split param.
+    if is_sharding_split_param_mode(args):
+        returned_optim_state_dict = load_unified_optimizer_split_param(args, model, optimizer, resume_from_checkpoint)
+        return returned_optim_state_dict
+
+    # init and get optimizer LR_Scheduler
+    returned_optim_state_dict = nested_copy(optimizer.state_dict())
+
+    if not safe_serialization:
+        index_filename, index_filename_master_weights = (
+            PADDLE_OPTIMIZER_INDEX_NAME,
+            PADDLE_MASTER_WEIGHTS_INDEX_NAME,
+        )
+    else:
+        index_filename, index_filename_master_weights = SAFE_OPTIMIZER_INDEX_NAME, SAFE_MASTER_WEIGHTS_INDEX_NAME
+
+    resolved_archive_file, sharded_metadata = get_optimizer_shard_files(
+        optimizer_path=resume_from_checkpoint,
+        index_filename=os.path.join(resume_from_checkpoint, index_filename),
+    )
+    has_master_weights = True if sharded_metadata["master_weights"] else False
+
+    model_state_dict = get_expected_state_dict(model)
+    model_keys = list(model_state_dict.keys())
+    struct2static_name_mappings = {k: v.name for k, v in model_state_dict.items()}  # get optimizer param mappings
+
+    expected_keys = get_expected_keys(args, sharded_metadata, model, optimizer)
+
+    # This should always be a list but, just to be sure.
+    if not isinstance(resolved_archive_file, list):
+        resolved_archive_file = [resolved_archive_file]
+
+    if len(resolved_archive_file) > 1:
+        resolved_archive_file = tqdm(resolved_archive_file, desc="Loading optimizer shards")
+
+    with open(os.path.join(resume_from_checkpoint, index_filename), "r") as f:
+        index = json.loads(f.read())
+
+    ckpt_quant_stage = "O0"
+    if "ckpt_quant_stage" in index:
+        ckpt_quant_stage = index["ckpt_quant_stage"]
+
+    # update has_master_weights and index_filename_master_weights
+    # 1. if the master weight exists, only has_master_weights is set True and loaded when needed
+    # 2. if master weight does not exist, convert model weight to master weight when needed
+    has_master_weights, index_filename_master_weights = update_master_weight_status(
+        args, optimizer, has_master_weights, safe_serialization
+    )
+
+    if has_master_weights:
+        returned_optim_state_dict["master_weights"] = {}
+
+        resolved_archive_file_mw, sharded_metadata_mw = get_optimizer_shard_files(
+            optimizer_path=resume_from_checkpoint,
+            index_filename=os.path.join(resume_from_checkpoint, index_filename_master_weights),
+        )
+
+        expected_keys_mw = get_expected_keys(args, sharded_metadata_mw, model, optimizer, is_master_weights=True)
+        if not isinstance(resolved_archive_file_mw, list):
+            resolved_archive_file_mw = [resolved_archive_file_mw]
+        if len(resolved_archive_file_mw) > 1:
+            resolved_archive_file_mw = tqdm(resolved_archive_file_mw, desc="Loading master weights shards")
+
+    def load_resolved_archive_file(
+        resolved_archive_file, sharded_metadata, expected_keys, is_master_weights=False, ckpt_quant_stage="O0"
+    ):
+        returned_state_dict = {}
+        # load optimizer
+        for shard_file in resolved_archive_file:
+            # TODO: check if no expected_keys in shard_file, then don't load it
+            if expected_keys.isdisjoint(sharded_metadata["file_map"][os.path.split(shard_file)[-1]]):
+                continue
+
+            if shard_file.endswith(".safetensors"):
+                # assert model_keys is not None, "model_keys is None." TODO: correct the assert
+                if model.config.tensor_parallel_degree > 1:
+                    if isinstance(model, LoRAModel) or isinstance(model, PrefixModelForCausalLM):
+                        tp_actions = model._get_tensor_parallel_convert_actions(
+                            model_keys, is_split=True, ignore_error=True
+                        )
+                    else:
+                        tp_actions = model.get_tensor_parallel_convert_actions(
+                            model.config, model_keys, ignore_error=True
+                        )
+                    if not is_master_weights:
+                        tp_actions = mapping_optimizer_tp_actions(tp_actions, expected_keys)
+
+                    # Here we use expected_keys to optimize weights loading for pipeline model. Only works for safetensors
+                    state_dict = load_state_dict(
+                        shard_file,
+                        tp_actions,
+                        expected_keys,
+                        device="expected",
+                        ckpt_quant_stage=ckpt_quant_stage,
+                    )
+                else:
+                    # for pipeline model, we don't need to use tp_actions
+                    state_dict = load_state_dict(
+                        shard_file,
+                        None,
+                        expected_keys,
+                        device="expected",
+                        ckpt_quant_stage=ckpt_quant_stage,
+                    )
+
+            returned_state_dict.update(state_dict)
+            # force memory release
+            del state_dict
+            gc.collect()
+        return returned_state_dict
+
+    state_dict_optim = load_resolved_archive_file(
+        resolved_archive_file, sharded_metadata, expected_keys, ckpt_quant_stage=ckpt_quant_stage
+    )
+    if has_master_weights:
+        state_dict_master_weight = load_resolved_archive_file(
+            resolved_archive_file_mw, sharded_metadata_mw, expected_keys_mw, is_master_weights=True
+        )
+    # rename optimizer param
+    for key in list(state_dict_optim.keys()):
+        key_name = key.split("/")
+        model_weight_key = key_name[0]
+        static_name = struct2static_name_mappings[model_weight_key]
+        if has_master_weights:
+            if model_state_dict[model_weight_key].dtype != paddle.float32:
+                key_name = "_".join([static_name, FP32_MASTER, key_name[1]])
+            else:
+                key_name = "_".join([static_name, key_name[1]])
+        else:
+            key_name = "_".join([static_name, key_name[1]])
+        returned_optim_state_dict[key_name] = state_dict_optim.pop(key)
+        returned_optim_state_dict[key_name].name = key_name
+
+        # master weight cast (only in remove_master_weight)
+        if has_master_weights and state_dict_master_weight[model_weight_key].dtype != paddle.float32:
+            state_dict_master_weight[model_weight_key] = paddle.cast(
+                state_dict_master_weight[model_weight_key], dtype=paddle.float32
+            )
+
+    if has_master_weights:
+        for key in list(state_dict_master_weight.keys()):
+            static_name = struct2static_name_mappings[key]
+            returned_optim_state_dict["master_weights"][static_name] = state_dict_master_weight.pop(key)
+            returned_optim_state_dict["master_weights"][static_name].name = "_".join([static_name, FP32_MASTER])
+
+    return returned_optim_state_dict

paddlenlp/trainer/unified_checkpoint/load_save_single_card.py

@@ -0,0 +1,237 @@
+# Copyright (c) 2024 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.
+"""Save and load single card checkpoint for Unified Checkpoint"""
+
+import gc
+import json
+import os
+
+import paddle
+
+from paddlenlp.peft import LoRAModel, PrefixModelForCausalLM
+from paddlenlp.transformers.model_utils import (
+    _load_state_dict_into_model,
+    load_state_dict,
+)
+from paddlenlp.transformers.utils import (
+    dtype_byte_size,
+    get_checkpoint_shard_files,
+    is_safetensors_available,
+)
+from paddlenlp.utils.env import (
+    SAFE_MASTER_WEIGHTS_INDEX_NAME,
+    SAFE_OPTIMIZER_INDEX_NAME,
+    SAFE_PEFT_WEIGHTS_INDEX_NAME,
+    SAFE_WEIGHTS_INDEX_NAME,
+)
+from paddlenlp.utils.log import logger
+from paddlenlp.utils.nested import nested_copy
+
+if is_safetensors_available():
+    from safetensors.numpy import save_file as safe_save_file
+
+from .utils import (
+    FP32_MASTER,
+    generate_base_static_name,
+    get_expected_state_dict,
+    get_optimizer_shard_files,
+    save_model_config,
+)
+
+__all__ = [
+    "load_single_card_checkpoint",
+    "load_single_card_optimizer",
+    "save_single_card_checkpoint",
+    "save_single_card_optimizer",
+]
+
+
+def save_file_sync(state_dict, path):
+    for k in list(state_dict.keys()):
+        if isinstance(state_dict[k], paddle.Tensor):
+            state_dict[k] = state_dict.pop(k).cpu().numpy()
+    safe_save_file(state_dict, path, metadata={"format": "np"})
+
+
+def save_single_card_checkpoint(model_to_save, output_dir):
+    """Save checkpoint for non-distributed environment."""
+
+    state_dict = get_expected_state_dict(model_to_save, concat_additional_adapter=True)
+    if isinstance(model_to_save, LoRAModel) or isinstance(model_to_save, PrefixModelForCausalLM):
+        weight_filename = "peft_model-00001-of-00001.safetensors"
+        index_filename = SAFE_PEFT_WEIGHTS_INDEX_NAME
+    else:
+        weight_filename = "model-00001-of-00001.safetensors"
+        index_filename = SAFE_WEIGHTS_INDEX_NAME
+    # get index json
+    index_weight_file = {}
+    total_size = 0
+    for key, weight in state_dict.items():
+        index_weight_file[key] = weight_filename
+        total_size += weight.numel().item() * dtype_byte_size(weight.dtype)
+    sharded_index_json = {}
+    sharded_index_json["metadata"] = {"total_size": total_size}
+    sharded_index_json["weight_map"] = index_weight_file
+    if isinstance(model_to_save, LoRAModel):
+        sharded_index_json["type"] = "lora"
+    elif isinstance(model_to_save, PrefixModelForCausalLM):
+        sharded_index_json["type"] = "ptuning"
+
+    os.makedirs(output_dir, exist_ok=True)
+    path = os.path.join(output_dir, index_filename)
+    with open(path, "w") as f:
+        json.dump(sharded_index_json, f, indent=4)
+
+    # save checkpoint, do no support asynchronous save for single card currently.
+    logger.warning("Asynchronous saving is not supported for single card environment currently.")
+    save_file_sync(state_dict, path=os.path.join(output_dir, weight_filename))
+
+    save_model_config(model_to_save, output_dir)
+
+
+def save_single_card_optimizer(model, optimizer, output_dir):
+    """ "Save optimizer for non-distributed environment."""
+    # Split into optimizer params and master weights.
+    optim_state_dict = nested_copy(optimizer.state_dict())
+    master_weights = None
+    if "master_weights" in optim_state_dict.keys():
+        master_weights = optim_state_dict.pop("master_weights")
+    if "LR_Scheduler" in optim_state_dict.keys():
+        optim_state_dict.pop("LR_Scheduler")
+
+    static2struct_name_mappings = {}
+    state_dict = get_expected_state_dict(model)
+    fp32_weight = {}
+    for k, v in state_dict.items():
+        static2struct_name_mappings[v.name] = k
+        if master_weights is not None and v.dtype == paddle.float32:
+            fp32_weight[k] = v
+
+    # rename optimizer param
+    for key in list(optim_state_dict.keys()):
+        static_name, type_name = generate_base_static_name(key)
+        new_name = static2struct_name_mappings[static_name] + "/" + type_name
+        optim_state_dict[new_name] = optim_state_dict.pop(key)
+    if master_weights is not None:
+        for key in list(master_weights.keys()):
+            master_weights[static2struct_name_mappings[key]] = master_weights.pop(key)
+        master_weights.update(fp32_weight)
+
+    # save index json
+    index_optimizer_file, index_master_weight_file = {}, {}
+    total_optim_size, total_master_weight_size = 0, 0
+    for key, weight in optim_state_dict.items():
+        index_optimizer_file[key] = "optimizer-00001-of-00001.safetensors"
+        total_optim_size += weight.numel().item() * dtype_byte_size(weight.dtype)
+    if master_weights is not None:
+        for key, weight in master_weights.items():
+            index_master_weight_file[key] = "master_weights-00001-of-00001.safetensors"
+            total_master_weight_size += weight.numel().item() * dtype_byte_size(weight.dtype)
+    path = os.path.join(output_dir, SAFE_OPTIMIZER_INDEX_NAME)
+    master_path = os.path.join(output_dir, SAFE_MASTER_WEIGHTS_INDEX_NAME)
+    with open(path, "w") as f:
+        has_master_weights = master_weights is not None
+        json.dump(
+            {
+                "metadata": {"total_size": total_optim_size},
+                "weight_map": index_optimizer_file,
+                "master_weights": has_master_weights,
+            },
+            f,
+            indent=4,
+        )
+    if master_weights is not None:
+        with open(master_path, "w") as f:
+            json.dump(
+                {"metadata": {"total_size": total_master_weight_size}, "weight_map": index_master_weight_file},
+                f,
+                indent=4,
+            )
+
+    # save optimizer state dict
+    save_file_sync(optim_state_dict, path=os.path.join(output_dir, "optimizer-00001-of-00001.safetensors"))
+    if master_weights is not None:
+        save_file_sync(master_weights, path=os.path.join(output_dir, "master_weights-00001-of-00001.safetensors"))
+
+
+def load_single_card_checkpoint(model, resume_from_checkpoint: str):
+    if isinstance(model, LoRAModel) or isinstance(model, PrefixModelForCausalLM):
+        index_filename = SAFE_PEFT_WEIGHTS_INDEX_NAME
+    else:
+        index_filename = SAFE_WEIGHTS_INDEX_NAME
+    resolved_archive_file, sharded_metadata = get_checkpoint_shard_files(
+        pretrained_model_name_or_path=resume_from_checkpoint,
+        index_filename=os.path.join(resume_from_checkpoint, index_filename),
+    )
+
+    loaded_keys = sharded_metadata["all_checkpoint_keys"]
+    model_state_dict = get_expected_state_dict(model)
+    expected_keys = set(list(model_state_dict.keys()))
+    missing_keys = expected_keys - set(loaded_keys)
+
+    if len(missing_keys) > 0:
+        raise ValueError(f"Missing keys: {missing_keys}")
+
+    state_dict = load_state_dict(resolved_archive_file[0], None, expected_keys)
+    error_msgs = _load_state_dict_into_model(model, state_dict, "")
+    del state_dict
+    gc.collect()
+
+    if error_msgs:
+        raise RuntimeError(f"Error(s) in loading state dict for {model.__class__.__name__}:\n\t{error_msgs}")
+
+
+def load_single_card_optimizer(model, optimizer, resume_from_checkpoint: str):
+    returned_optim_state_dict = nested_copy(optimizer.state_dict())
+
+    resolved_archive_file, sharded_metadata = get_optimizer_shard_files(
+        optimizer_path=resume_from_checkpoint,
+        index_filename=os.path.join(resume_from_checkpoint, SAFE_OPTIMIZER_INDEX_NAME),
+    )
+    has_master_weights = True if sharded_metadata["master_weights"] else False
+
+    model_state_dict = get_expected_state_dict(model)
+    struct2static_name_mappings = {k: v.name for k, v in model_state_dict.items()}
+    expected_keys = sharded_metadata["all_optimizer_keys"]
+
+    if has_master_weights:
+        returned_optim_state_dict["master_weights"] = {}
+        resolved_archive_file_mw, sharded_metadata_mw = get_optimizer_shard_files(
+            optimizer_path=resume_from_checkpoint,
+            index_filename=os.path.join(resume_from_checkpoint, SAFE_MASTER_WEIGHTS_INDEX_NAME),
+        )
+        expected_keys_mw = sharded_metadata_mw["all_optimizer_keys"]
+
+    state_dict_optim = load_state_dict(resolved_archive_file[0], None, expected_keys)
+    if has_master_weights:
+        state_dict_optim_mw = load_state_dict(resolved_archive_file_mw[0], None, expected_keys_mw)
+
+    for key in list(state_dict_optim.keys()):
+        key_name = key.split("/")
+        static_name = struct2static_name_mappings[key_name[0]]
+        if has_master_weights:
+            if model_state_dict[key_name[0]].dtype != paddle.float32:
+                key_name = "_".join([static_name, FP32_MASTER, key_name[1]])
+            else:
+                key_name = "_".join([static_name, key_name[1]])
+        else:
+            key_name = "_".join([static_name, key_name[1]])
+        returned_optim_state_dict[key_name] = state_dict_optim.pop(key)
+        returned_optim_state_dict[key_name].name = key_name
+    if has_master_weights:
+        for key in list(state_dict_optim_mw.keys()):
+            static_name = struct2static_name_mappings[key]
+            returned_optim_state_dict["master_weights"][static_name] = state_dict_optim_mw.pop(key)
+            returned_optim_state_dict["master_weights"][static_name].name = "_".join([static_name, FP32_MASTER])
+    return returned_optim_state_dict

paddlenlp/trainer/unified_checkpoint/sharding_split_param_utils.py

@@ -0,0 +1,308 @@
+# Copyright (c) 2024 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.
+"""Support Sharding Stage1 V2(split param) for Unified Checkpoint"""
+
+import gc
+import os
+from itertools import chain
+
+import paddle
+import paddle.distributed as dist
+from paddle.distributed import fleet
+from tqdm.auto import tqdm
+
+from paddlenlp.peft import LoRAModel, PrefixModelForCausalLM
+from paddlenlp.transformers.model_utils import load_state_dict, unwrap_model
+from paddlenlp.utils.env import (
+    SAFE_MASTER_WEIGHTS_INDEX_NAME,
+    SAFE_OPTIMIZER_INDEX_NAME,
+)
+from paddlenlp.utils.nested import nested_copy
+
+from .utils import (
+    FP32_MASTER,
+    generate_base_static_name,
+    get_expected_state_dict,
+    get_optimizer_shard_files,
+    mapping_optimizer_tp_actions,
+)
+
+__all__ = ["gather_splited_param_for_optimizer", "load_unified_optimizer_split_param"]
+
+
+def merge_splited_param(
+    state_dict, partial_tensor_list, param_shape_info, send_table, recv_table, is_master_weights=False
+):
+    """Merge the splited param in sharding group."""
+    global_rank = dist.get_rank()
+    for key in list(state_dict.keys()):
+        if state_dict[key].numel().item() == 1:  # for example: beta1, beta2
+            continue
+
+        static_name = key if is_master_weights else generate_base_static_name(key)[0]
+        shape, numel, index, padded_size = param_shape_info[static_name]
+        if static_name not in partial_tensor_list:
+            state_dict[key] = state_dict[key].reshape(shape)
+            continue
+
+        recv_rank = recv_table[static_name]
+        send_info = send_table[static_name]
+
+        base_padding_start = index + numel
+        base_padding_end = index + padded_size
+
+        if global_rank == recv_rank:
+            tmp_tensor_list = []
+            for send_rank, begin, end in send_info:
+                padding_start = max(begin, base_padding_start)
+                padding_end = min(end, base_padding_end)
+
+                if send_rank == recv_rank:
+                    tensor = (
+                        state_dict[key] if padding_start >= padding_end else state_dict[key][: padding_start - begin]
+                    )
+                    tmp_tensor_list.append(tensor)
+                else:
+                    length = end - begin if padding_start >= padding_end else padding_start - begin
+                    tmp_tensor = paddle.empty(shape=[length], dtype=state_dict[key].dtype)
+                    dist.stream.recv(tmp_tensor, src=send_rank)
+                    tmp_tensor_list.append(tmp_tensor)
+            state_dict[key] = paddle.concat(tmp_tensor_list, axis=0).reshape(shape)
+        else:
+            for send_rank, begin, end in send_info:
+                padding_start = max(begin, base_padding_start)
+                padding_end = min(end, base_padding_end)
+                if global_rank == send_rank:
+                    tensor = (
+                        state_dict[key] if padding_start >= padding_end else state_dict[key][: padding_start - begin]
+                    )
+                    dist.stream.send(tensor, dst=recv_rank)
+                    state_dict.pop(key)
+    return state_dict
+
+
+def gather_splited_param_for_optimizer(optimizer):
+    hcg = fleet.get_hybrid_communicate_group()
+    sharding_group = hcg.get_sharding_parallel_group()
+    global_rank = dist.get_rank()
+    param_slice_info = {}
+    param_shape_info = {}
+
+    for buffer in optimizer._inner_opt._comm_buffer_list:
+        for key in buffer._sharding_param_grad_view.keys():
+            param_slice_info[key] = (
+                buffer._sharding_param_grad_view[key]._param_begin,
+                buffer._sharding_param_grad_view[key]._param_end,
+            )
+            param_shape_info[key] = (
+                buffer._sharding_param_grad_view[key]._param.shape,
+                buffer._sharding_param_grad_view[key]._param.numel().item(),
+                buffer._sharding_param_grad_view[key]._index,
+                buffer._sharding_param_grad_view[key]._padded_size,
+            )
+    param_slice_info["global_rank"] = global_rank
+    param_slice_info_list = []
+    dist.all_gather_object(param_slice_info_list, param_slice_info, group=sharding_group)
+
+    optim_state_dict = nested_copy(optimizer.state_dict())
+    master_weights = None
+    if "master_weights" in optim_state_dict.keys():
+        master_weights = optim_state_dict.pop("master_weights")
+    if "LR_Scheduler" in optim_state_dict.keys():
+        optim_state_dict.pop("LR_Scheduler")
+
+    # deal with optimizer param
+    partial_tensor_list = []
+    for key in list(optim_state_dict.keys()):
+        static_name, _ = generate_base_static_name(key)
+        if static_name in param_slice_info.keys():
+            if optim_state_dict[key].numel().item() == 1:  # for example: beta1, beta2
+                continue
+            begin, end = param_slice_info[static_name]
+            shape, numel, _, _ = param_shape_info[static_name]
+            if end - begin == numel:  # full tensor
+                optim_state_dict[key] = optim_state_dict[key].reshape(shape)
+            elif end <= begin:  # empty tensor
+                continue
+            else:  # partial tensor, end > begin but end - begin < numel
+                partial_tensor_list.append(static_name)
+
+    send_table = {}
+    recv_table = {}
+    for key in partial_tensor_list:
+        sharding_ranklist = []
+        for slice_info in param_slice_info_list:
+            begin, end = slice_info[key]
+            if end > begin:
+                sharding_ranklist.append((slice_info["global_rank"], begin, end))
+        recv_table[key] = sharding_ranklist[0][0]  # which sharding_rank to recv the splited tensor
+        send_table[key] = [(rank, begin, end) for rank, begin, end in sharding_ranklist]
+
+    merge_splited_param(optim_state_dict, partial_tensor_list, param_shape_info, send_table, recv_table, False)
+    if master_weights is not None:
+        merge_splited_param(master_weights, partial_tensor_list, param_shape_info, send_table, recv_table, True)
+    return optim_state_dict, master_weights
+
+
+def load_unified_optimizer_split_param(args, model, optimizer, resume_from_checkpoint):
+    returned_optim_state_dict = nested_copy(optimizer.state_dict())
+
+    index_filename, index_filename_master_weights = SAFE_OPTIMIZER_INDEX_NAME, SAFE_MASTER_WEIGHTS_INDEX_NAME
+
+    resolved_archive_file, sharded_metadata = get_optimizer_shard_files(
+        optimizer_path=resume_from_checkpoint,
+        index_filename=os.path.join(resume_from_checkpoint, index_filename),
+    )
+    has_master_weights = True if sharded_metadata["master_weights"] else False
+
+    typename_set = set()
+    for key in sharded_metadata["weight_map"].keys():
+        _, typename = key.split("/")
+        typename_set.add(typename)
+
+    model_state_dict = get_expected_state_dict(model)
+    model_keys = list(model_state_dict.keys())
+    static2struct_name_mappings = {v.name: k for k, v in model_state_dict.items()}  # get optimizer param mappings
+    struct2static_name_mappings = {k: v.name for k, v in model_state_dict.items()}
+
+    expected_keys = []
+    param_slice_info = {}
+    param_shape_info = {}
+
+    comm_buffer_list = optimizer._inner_opt._comm_buffer_list
+    if hasattr(args, "enable_sharding_comm_overlap") and args.enable_sharding_comm_overlap:
+        comm_buffer_list = list(chain(*model._chunk_2_comm_buffers.values()))
+        model = unwrap_model(model)
+
+    for buffer in comm_buffer_list:
+        for key in buffer._sharding_param_grad_view.keys():
+            begin = buffer._sharding_param_grad_view[key]._param_begin
+            end = buffer._sharding_param_grad_view[key]._param_end
+            if end > begin:
+                expected_keys.append(key)
+                shape = buffer._sharding_param_grad_view[key]._param.shape
+                numel = buffer._sharding_param_grad_view[key]._param.numel().item()
+                index = buffer._sharding_param_grad_view[key]._index
+                padded_size = buffer._sharding_param_grad_view[key]._padded_size
+                param_slice_info[key] = (begin, end)
+                param_shape_info[key] = (shape, numel, index, padded_size)
+
+    expected_keys = set([static2struct_name_mappings.get(name, None) for name in expected_keys])
+    expected_keys_optim = []
+    for key in expected_keys:
+        for typename in typename_set:
+            expected_keys_optim.append(f"{key}/{typename}")
+    expected_keys_optim = set(expected_keys_optim)
+
+    if len(resolved_archive_file) > 1:
+        resolved_archive_file = tqdm(resolved_archive_file, desc="Loading optimizer shards")
+
+    if has_master_weights:
+        returned_optim_state_dict["master_weights"] = {}
+        resolved_archive_file_mw, sharded_metadata_mw = get_optimizer_shard_files(
+            optimizer_path=resume_from_checkpoint,
+            index_filename=os.path.join(resume_from_checkpoint, index_filename_master_weights),
+        )
+        if len(resolved_archive_file_mw) > 1:
+            resolved_archive_file_mw = tqdm(resolved_archive_file_mw, desc="Loading master weights shards")
+
+    def load_resolved_archive_file(resolved_archive_file, sharded_metadata, expected_keys, is_master_weights=False):
+        returned_state_dict = {}
+
+        if model.config.tensor_parallel_degree > 1:
+            if isinstance(model, LoRAModel) or isinstance(model, PrefixModelForCausalLM):
+                tp_actions = model._get_tensor_parallel_convert_actions(model_keys, is_split=True, ignore_error=True)
+            else:
+                tp_actions = model.get_tensor_parallel_convert_actions(model.config, model_keys, ignore_error=True)
+            if not is_master_weights:
+                tp_actions = mapping_optimizer_tp_actions(tp_actions, expected_keys)
+
+        for shard_file in resolved_archive_file:
+            if expected_keys.isdisjoint(sharded_metadata["file_map"][os.path.split(shard_file)[-1]]):
+                continue
+            if model.config.tensor_parallel_degree > 1:
+                state_dict = load_state_dict(shard_file, tp_actions, expected_keys, device="cpu")
+            else:
+                state_dict = load_state_dict(shard_file, None, expected_keys, device="cpu")
+            returned_state_dict.update(state_dict)
+            del state_dict
+            gc.collect()
+
+        return returned_state_dict
+
+    # get tp params
+    state_dict_optim = load_resolved_archive_file(resolved_archive_file, sharded_metadata, expected_keys_optim)
+
+    # need to split param for different sharding rank, maybe need to deal with oom issue.
+    for key in list(state_dict_optim.keys()):
+        key_name = key.split("/")
+        static_name = struct2static_name_mappings.get(key_name[0], None)
+
+        if state_dict_optim[key].numel().item() > 1:
+            begin, end = param_slice_info[static_name]
+            shape, numel, index, padded_size = param_shape_info[static_name]
+            state_dict_optim[key] = state_dict_optim[key].reshape([-1])
+            state_dict_optim[key] = state_dict_optim[key][begin - index : end - index]
+
+            padding_start = max(begin, index + numel)
+            padding_end = min(end, index + padded_size)
+            if padding_start < padding_end:
+                state_dict_optim[key] = paddle.concat(
+                    (
+                        state_dict_optim[key],
+                        paddle.zeros([padding_end - padding_start], dtype=state_dict_optim[key].dtype),
+                    )
+                )
+        if has_master_weights:
+            key_name = "_".join([static_name, FP32_MASTER, key_name[1]])
+        else:
+            key_name = "_".join([static_name, key_name[1]])
+
+        state_dict_optim[key] = state_dict_optim[key]._copy_to(paddle.framework._current_expected_place(), False)
+
+        returned_optim_state_dict[key_name] = state_dict_optim.pop(key)
+        returned_optim_state_dict[key_name].name = key_name
+
+    if has_master_weights:
+        state_dict_master_weight = load_resolved_archive_file(
+            resolved_archive_file_mw,
+            sharded_metadata_mw,
+            expected_keys,
+            is_master_weights=True,
+        )
+
+        for key in list(state_dict_master_weight.keys()):
+            static_name = struct2static_name_mappings.get(key, None)
+            if state_dict_master_weight[key].numel().item() > 1:
+                begin, end = param_slice_info[static_name]
+                shape, numel, index, padded_size = param_shape_info[static_name]
+                state_dict_master_weight[key] = state_dict_master_weight[key].reshape([-1])
+                state_dict_master_weight[key] = state_dict_master_weight[key][begin - index : end - index]
+
+                padding_start = max(begin, index + numel)
+                padding_end = min(end, index + padded_size)
+                if padding_start < padding_end:
+                    state_dict_master_weight[key] = paddle.concat(
+                        (
+                            state_dict_master_weight[key],
+                            paddle.zeros([padding_end - padding_start], dtype=state_dict_master_weight[key].dtype),
+                        )
+                    )
+            state_dict_master_weight[key] = state_dict_master_weight[key]._copy_to(
+                paddle.framework._current_expected_place(), False
+            )
+            returned_optim_state_dict["master_weights"][static_name] = state_dict_master_weight.pop(key)
+            returned_optim_state_dict["master_weights"][static_name].name = "_".join([static_name, FP32_MASTER])
+
+    return returned_optim_state_dict

paddlenlp/trainer/unified_checkpoint/unified_checkpoint.py

@@ -0,0 +1,702 @@
+# 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 copy
+import json
+import os
+
+import paddle
+from paddle.distributed import fleet
+
+from paddlenlp.peft import LoRAModel, PrefixModelForCausalLM
+from paddlenlp.trainer.argparser import strtobool
+from paddlenlp.trainer.utils.helper import distributed_isfile
+from paddlenlp.transformers.model_utils import (
+    PretrainedModel,
+    _add_variant,
+    load_state_dict,
+    unwrap_model,
+)
+from paddlenlp.transformers.utils import dtype_byte_size
+from paddlenlp.utils import infohub
+from paddlenlp.utils.env import (
+    LORA_WEIGHTS_NAME,
+    MAX_QUANTIZATION_TIMES,
+    PADDLE_MASTER_WEIGHTS_NAME,
+    PADDLE_OPTIMIZER_NAME,
+    PADDLE_WEIGHTS_NAME,
+    PREFIX_WEIGHTS_NAME,
+    SAFE_MASTER_WEIGHTS_INDEX_NAME,
+    SAFE_MASTER_WEIGHTS_NAME,
+    SAFE_OPTIMIZER_INDEX_NAME,
+    SAFE_OPTIMIZER_NAME,
+    SAFE_PEFT_WEIGHTS_INDEX_NAME,
+    SAFE_PEFT_WEIGHTS_NAME,
+    SAFE_WEIGHTS_INDEX_NAME,
+    SAFE_WEIGHTS_NAME,
+)
+from paddlenlp.utils.log import logger
+from paddlenlp.utils.nested import nested_copy
+
+from .async_handler import AsyncCheckpointHandler
+from .check_completion import check_unified_checkpoint, check_unified_optimizer
+from .load_dynamic import (
+    load_unified_checkpoint_dynamically,
+    load_unified_optimizer_dynamically,
+)
+from .load_local import load_unified_checkpoint_locally, load_unified_optimizer_locally
+from .load_save_single_card import (
+    load_single_card_checkpoint,
+    load_single_card_optimizer,
+    save_single_card_checkpoint,
+    save_single_card_optimizer,
+)
+from .sharding_split_param_utils import gather_splited_param_for_optimizer
+from .utils import (
+    FP32_MASTER,
+    UnifiedCheckpointOption,
+    filter_params,
+    gather_sharded_object,
+    generate_base_static_name,
+    get_expected_state_dict,
+    get_sharded_file_name,
+    get_sharded_index,
+    is_need_master_weight,
+    is_sharding_split_param_mode,
+    merge_tensor_parallel_for_optimizer,
+    merge_tensor_parallel_with_shard,
+    reduce_master_weights_status,
+    rename_shard_file,
+    save_model_config,
+)
+
+__all__ = ["UnifiedCheckpointHandler"]
+
+
+class UnifiedCheckpointHandler:
+    def __init__(self, args):
+        self.args = args
+        self.async_handler = AsyncCheckpointHandler(args)
+
+    def save_unified_checkpoint(self, model, optimizer, output_dir, signal_dir=None):
+        """save unified checkpoint
+
+        Args:
+            model (PretrainedModel): model to save
+            output_dir (str): save dir
+            safe_serialization (bool, optional): use safetensors. Defaults to False.
+
+        Raises:
+            ValueError: if model is not an instance of `PretrainedModel` and the model cannot be saved
+        """
+        if isinstance(model, PretrainedModel):
+            model_to_save = model
+        elif isinstance(unwrap_model(model), PretrainedModel):
+            model_to_save = unwrap_model(model)
+        elif isinstance(model, PrefixModelForCausalLM) or isinstance(model, LoRAModel):
+            model_to_save = model
+        else:
+            raise ValueError("Unified checkpoint only supports PretrainedModel, LoRAModel and PrefixModelForCausalLM!")
+
+        # Under non distributed environment.
+        if paddle.distributed.get_world_size() <= 1:
+            save_single_card_checkpoint(model_to_save, output_dir)
+            return
+
+        skip_save_model_weight = False
+        if UnifiedCheckpointOption.SKIP_SAVE_MODEL_WEIGHT.value in self.args.unified_checkpoint_config:
+            if is_need_master_weight(optimizer, is_fp16_or_bp16=(self.args.fp16 or self.args.bf16)):
+                logger.info(
+                    f"With {UnifiedCheckpointOption.SKIP_SAVE_MODEL_WEIGHT.value}, skip the model checkpoint save."
+                    " The master weight will be loaded as model weights for next resumption."
+                )
+                # not save model weight, load from master weight
+                skip_save_model_weight = True
+
+        save_directory = output_dir
+        os.makedirs(save_directory, exist_ok=True)
+        if signal_dir is not None:
+            os.makedirs(signal_dir, exist_ok=True)  # only for async save
+
+        # save model weights
+        if not skip_save_model_weight:
+            state_dict, shard_file, sharded_index = unified_checkpoint_into_shards(
+                self.args, model_to_save, safe_serialization=True
+            )
+            is_sync_save = True
+            if "async_save" in self.args.unified_checkpoint_config:
+                is_sync_save = False
+            self.async_handler._file_save_async_or_sync(
+                state_dict,
+                path=os.path.join(save_directory, shard_file),
+                signal_path=signal_dir,
+                is_sync=is_sync_save,
+                state_dict_type="model_weight",
+            )
+            if sharded_index is not None:
+                if isinstance(model_to_save, LoRAModel) or isinstance(model_to_save, PrefixModelForCausalLM):
+                    index_name = SAFE_PEFT_WEIGHTS_INDEX_NAME
+                else:
+                    index_name = SAFE_WEIGHTS_INDEX_NAME
+                path = os.path.join(output_dir, index_name)
+
+                if self.args.should_save:
+                    with open(path, "w") as f:
+                        json.dump(sharded_index, f, indent=4)
+
+        if self.args.should_save:
+            save_model_config(model_to_save, save_directory)
+
+        paddle.device.cuda.empty_cache()
+
+        if strtobool(os.getenv("FLAG_LLM_PDC", "False")) and self.args.should_save:
+            world_size = paddle.distributed.get_world_size()
+            save_info = {
+                "world_size": world_size,
+                "ignore_save_lr_and_optim": self.args.ignore_save_lr_and_optim,
+                "skip_save_model_weight": "skip_save_model_weight" in self.args.unified_checkpoint_config,
+                "remove_master_weight": "remove_master_weight" in self.args.unified_checkpoint_config,
+            }
+            paddle.save(save_info, os.path.join(save_directory, ".saving_info"))
+
+    def load_unified_checkpoint(self, model, resume_from_checkpoint: str):
+        """Load potential model checkpoint
+
+        Args:
+            model (PretrainedModel): Your model to load
+            resume_from_checkpoint (str): path of the checkpoint to load
+
+        Returns:
+            None
+        """
+        if paddle.distributed.get_world_size() <= 1:
+            load_single_card_checkpoint(model, resume_from_checkpoint)
+            return
+
+        local_resume = check_unified_checkpoint(self.args, model, resume_from_checkpoint, safe_serialization=True)
+
+        if not local_resume:
+            logger.info("Begin to dynamically load unified checkpoint!")
+            load_unified_checkpoint_dynamically(self.args, model, resume_from_checkpoint, safe_serialization=True)
+            return
+
+        if self.args.dataset_rank == 0 or self.args.use_expert_parallel:
+            load_unified_checkpoint_locally(self.args, model, resume_from_checkpoint, safe_serialization=True)
+
+    def save_non_merge_optimizer(self, model, optim_state_dict, master_weights, output_dir, signal_dir):
+        paddle.device.cuda.empty_cache()
+
+        # gather global master_weights status.
+        global_master_weights = reduce_master_weights_status(master_weights is not None)
+        if master_weights is None and global_master_weights:
+            master_weights = {}
+
+        # get optimizer param mappings
+        static2struct_name_mappings = {}
+        state_dict = get_expected_state_dict(model)
+        for k, v in state_dict.items():
+            static2struct_name_mappings[v.name] = k
+
+        # rename optimizer param name
+        for key in list(optim_state_dict.keys()):
+            static_name, type_name = generate_base_static_name(key)
+            new_name = static2struct_name_mappings[static_name] + "/" + type_name
+            optim_state_dict[new_name] = optim_state_dict.pop(key)
+
+        if master_weights is not None:
+            for key in list(master_weights.keys()):
+                master_weights[static2struct_name_mappings[key]] = master_weights.pop(key)
+
+        no_sync_kname = []
+        model_state_dict = get_expected_state_dict(model)
+        for k, v in model_state_dict.items():
+            if getattr(v, "no_sync", False):
+                no_sync_kname.append(k)
+
+        hcg = fleet.get_hybrid_communicate_group()
+        dp_group = hcg.get_data_parallel_group()
+        dp_rank = dp_group.rank if dp_group.nranks > 1 else 0
+        if self.args.use_expert_parallel:
+            for k in list(optim_state_dict.keys()):
+                model_k = k.split("/")[0]
+                if dp_rank > 0 and model_k not in no_sync_kname:
+                    optim_state_dict.pop(k)
+            if master_weights is not None:
+                for k in list(master_weights.keys()):
+                    model_k = k.split("/")[0]
+                    if dp_rank > 0 and model_k not in no_sync_kname:
+                        master_weights.pop(k)
+
+        optimizer_name = _add_variant(SAFE_OPTIMIZER_NAME, self.args.optimizer_name_suffix)
+        master_weights_name = _add_variant(SAFE_MASTER_WEIGHTS_NAME, self.args.optimizer_name_suffix)
+
+        sharded_optim_index = {}
+        # save opt index json if checkpoint quantization is on.
+        if self.args.ckpt_quant_stage != "O0" and "quant_reach_limit" not in infohub:
+            sharded_optim_index["ckpt_quant_stage"] = self.args.ckpt_quant_stage
+
+        sharded_optim_index["quant_ckpt_resume_times"] = (
+            infohub["quant_ckpt_resume_times"] if "quant_ckpt_resume_times" in infohub else 0
+        )
+
+        if len(sharded_optim_index) > 0:
+            optimizer_index_name = SAFE_OPTIMIZER_INDEX_NAME
+            path = os.path.join(output_dir, optimizer_index_name)
+            if self.args.should_save:
+                with open(path, "w") as f:
+                    json.dump(sharded_optim_index, f, indent=4)
+
+        is_sync_save = True
+        if "async_save" in self.args.unified_checkpoint_config:
+            is_sync_save = False
+        self.async_handler._file_save_async_or_sync(
+            optim_state_dict,
+            path=os.path.join(output_dir, optimizer_name),
+            signal_path=signal_dir,
+            is_sync=is_sync_save,
+            state_dict_type="optimizer_weight",
+            ckpt_quant_stage=self.args.ckpt_quant_stage if "quant_reach_limit" not in infohub else "O0",
+        )
+        if master_weights is not None:
+            self.async_handler._file_save_async_or_sync(
+                master_weights,
+                path=os.path.join(output_dir, master_weights_name),
+                signal_path=signal_dir,
+                is_sync=is_sync_save,
+                state_dict_type="master_weight",
+            )
+
+    def load_non_merge_optimizer(self, model, optimizer, resume_from_checkpoint, ckpt_quant_stage="O0"):
+        # init and get optimizer LR_Scheduler
+        returned_optim_state_dict = nested_copy(optimizer.state_dict())
+
+        optimizer_name = _add_variant(SAFE_OPTIMIZER_NAME, self.args.optimizer_name_suffix)
+        master_weights_name = _add_variant(SAFE_MASTER_WEIGHTS_NAME, self.args.optimizer_name_suffix)
+        optimizer_path = os.path.join(resume_from_checkpoint, optimizer_name)
+        master_weights_path = os.path.join(resume_from_checkpoint, master_weights_name)
+        # no quantization & no master weight represent O1 AMP strategy.
+        is_amp_o1 = self.args.fp16_opt_level == "O1"
+
+        model_state_dict = get_expected_state_dict(model)
+        struct2static_name_mappings = {k: v.name for k, v in model_state_dict.items()}  # get optimizer param mappings
+        optimizer_state_dict = load_state_dict(
+            optimizer_path, None, None, device="expected", ckpt_quant_stage=ckpt_quant_stage
+        )
+        master_weights = {}
+        # normal AMP O2
+        if not is_amp_o1 and os.path.isfile(master_weights_path):
+            master_weights = load_state_dict(master_weights_path, None, None, device="expected")
+
+        # rename and move to paddle.Tensor
+        for key in list(optimizer_state_dict.keys()):
+            key_name = key.split("/")
+            model_weight_key = key_name[0]
+            static_name = struct2static_name_mappings[key_name[0]]
+            if not is_amp_o1:
+                if model_state_dict[key_name[0]].dtype != paddle.float32:
+                    key_name = "_".join([static_name, FP32_MASTER, key_name[1]])
+                else:
+                    key_name = "_".join([static_name, key_name[1]])
+            else:
+                key_name = "_".join([static_name, key_name[1]])
+            returned_optim_state_dict[key_name] = optimizer_state_dict.pop(key)
+            returned_optim_state_dict[key_name].name = key_name
+
+            # master weight cast (only in AMP O2 + remove_master_weight)
+            if not is_amp_o1 and not os.path.isfile(master_weights_path):
+                master_weights[model_weight_key] = paddle.cast(
+                    model_state_dict[model_weight_key], dtype=paddle.float32
+                )
+
+        if not is_amp_o1:
+            returned_optim_state_dict["master_weights"] = {}
+            for key in list(master_weights.keys()):
+                static_name = struct2static_name_mappings[key]
+                returned_optim_state_dict["master_weights"][static_name] = master_weights.pop(key)
+                returned_optim_state_dict["master_weights"][static_name].name = "_".join([static_name, FP32_MASTER])
+
+        return returned_optim_state_dict
+
+    def save_unified_optimizer(self, model, optimizer, output_dir, signal_dir):
+        """save unified optimizer
+
+        Args:
+            model (PretrainedModel): model used to get key mapping.
+            optimizer (Optimizer): optimizer to save
+            output_dir (str): Save directory.
+            signal_dir (str): Asynchronous saving signal directory.
+
+        """
+
+        if paddle.distributed.get_world_size() <= 1:
+            save_single_card_optimizer(model, optimizer, output_dir)  # no need to save signal
+            return
+
+        if is_sharding_split_param_mode(self.args):
+            optim_state_dict, master_weights = gather_splited_param_for_optimizer(optimizer)
+        else:
+            optim_state_dict = nested_copy(optimizer.state_dict())
+            master_weights = None
+            if "master_weights" in optim_state_dict.keys():
+                master_weights = optim_state_dict["master_weights"]
+                optim_state_dict.pop("master_weights")
+            if "LR_Scheduler" in optim_state_dict.keys():
+                optim_state_dict.pop("LR_Scheduler")
+
+        if UnifiedCheckpointOption.REMOVE_MASTER_WEIGHT.value in self.args.unified_checkpoint_config:
+            logger.info("Skip master weight saving.")
+            master_weights = None
+
+        if "ignore_merge_optimizer" in self.args.unified_checkpoint_config:
+            self.save_non_merge_optimizer(model, optim_state_dict, master_weights, output_dir, signal_dir)
+            return
+
+        # Split into naive optimizer params and master weights.
+        results = unified_optimizer_into_shards(
+            self.args, model, optim_state_dict, master_weights, safe_serialization=True
+        )
+        master_weight_state_dict = None
+        if len(results) == 1:
+            optim_state_dict, shard_optim_file, sharded_optim_index = results[0]
+        else:
+            optim_state_dict, shard_optim_file, sharded_optim_index = results[0]
+            master_weight_state_dict, shard_master_weight_file, sharded_master_weight_index = results[1]
+
+        paddle.device.cuda.empty_cache()
+        save_directory = output_dir
+        os.makedirs(save_directory, exist_ok=True)
+        if signal_dir is not None:
+            os.makedirs(signal_dir, exist_ok=True)
+
+        is_sync_save = True
+        if "async_save" in self.args.unified_checkpoint_config:
+            is_sync_save = False
+        self.async_handler._file_save_async_or_sync(
+            optim_state_dict,
+            path=os.path.join(save_directory, shard_optim_file),
+            signal_path=signal_dir,
+            is_sync=is_sync_save,
+            state_dict_type="optimizer_weight",
+            ckpt_quant_stage=self.args.ckpt_quant_stage if "quant_reach_limit" not in infohub else "O0",
+        )
+        if master_weight_state_dict is not None:
+            self.async_handler._file_save_async_or_sync(
+                master_weight_state_dict,
+                path=os.path.join(save_directory, shard_master_weight_file),
+                signal_path=signal_dir,
+                is_sync=is_sync_save,
+                state_dict_type="master_weight",
+            )
+
+        if sharded_optim_index is not None:
+            optimizer_index_name = SAFE_OPTIMIZER_INDEX_NAME
+            path = os.path.join(output_dir, optimizer_index_name)
+            if self.args.should_save:
+                with open(path, "w") as f:
+                    json.dump(sharded_optim_index, f, indent=4)
+
+            master_weights_name = SAFE_MASTER_WEIGHTS_INDEX_NAME
+            if UnifiedCheckpointOption.SKIP_SAVE_MODEL_WEIGHT.value in self.args.unified_checkpoint_config:
+                master_weights_name = SAFE_WEIGHTS_INDEX_NAME
+            master_path = os.path.join(output_dir, master_weights_name)
+            if master_weight_state_dict is not None:
+                if self.args.should_save:
+                    with open(master_path, "w") as f:
+                        json.dump(sharded_master_weight_index, f, indent=4)
+
+    def load_unified_optimizer(self, model, optimizer, resume_from_checkpoint):
+        """Load potential model checkpoint
+
+        Args:
+            model (PretrainedModel): Your model to load
+            resume_from_checkpoint (str): path of the checkpoint to load
+
+        Returns:
+            None
+        """
+
+        if paddle.distributed.get_world_size() <= 1:
+            optim_state_dict = load_single_card_optimizer(model, optimizer, resume_from_checkpoint)
+            return optim_state_dict
+
+        index = {}
+        has_merge_optimizer_safetensors = distributed_isfile(
+            os.path.join(resume_from_checkpoint, SAFE_OPTIMIZER_INDEX_NAME)
+        )
+        if has_merge_optimizer_safetensors:
+            with open(os.path.join(resume_from_checkpoint, SAFE_OPTIMIZER_INDEX_NAME), "r") as f:
+                index = json.loads(f.read())
+
+        # get quant ckpt info `ckpt_quant_stage` and `quant_ckpt_resume_times`
+        ckpt_quant_stage = "O0"
+        if "ckpt_quant_stage" in index:
+            ckpt_quant_stage = index["ckpt_quant_stage"]
+
+        quant_ckpt_resume_times = 0
+        if "quant_ckpt_resume_times" in index:
+            quant_ckpt_resume_times = index["quant_ckpt_resume_times"]
+        # increment and save resume times in infohub
+        if ckpt_quant_stage != "O0":
+            quant_ckpt_resume_times += 1
+        infohub["quant_ckpt_resume_times"] = quant_ckpt_resume_times
+
+        # Quantization times exceeds the limit. Turn off the quantization strategy.
+        if quant_ckpt_resume_times >= MAX_QUANTIZATION_TIMES:
+            infohub["quant_reach_limit"] = True
+            logger.info("Checkpoint quantization time reach limit and will be closed.")
+
+        # If not having merge optimizer, then load non-merge optimizer.
+        if "weight_map" not in index:
+            if self.args.data_parallel_rank == 0 or self.args.use_expert_parallel:
+                returned_optim_state_dict = self.load_non_merge_optimizer(
+                    model,
+                    optimizer,
+                    resume_from_checkpoint,
+                    ckpt_quant_stage=ckpt_quant_stage,
+                )
+                return returned_optim_state_dict
+            else:
+                return None
+
+        local_resume = check_unified_optimizer(
+            self.args, model, optimizer, resume_from_checkpoint, safe_serialization=True
+        )
+        if not local_resume:
+            logger.info("Begin to dynamically load unified optimizer!")
+            returned_optim_state_dict = load_unified_optimizer_dynamically(
+                self.args, model, optimizer, resume_from_checkpoint, safe_serialization=True
+            )
+            return returned_optim_state_dict
+
+        if self.args.data_parallel_rank == 0 or self.args.use_expert_parallel:
+            returned_optim_state_dict = load_unified_optimizer_locally(
+                self.args, model, optimizer, resume_from_checkpoint, safe_serialization=True
+            )
+            return returned_optim_state_dict
+        return None
+
+    def unlink_shared_memory(self):
+        return self.async_handler.unlink_shared_memory()
+
+
+def unified_checkpoint_into_shards(
+    args,
+    model_to_save,
+    safe_serialization=False,
+):
+    """Get state_dict and config to save
+
+    Args:
+        model_to_save (nn.Layer): model to, save
+        safe_serialization (bool, optional): safe serialization using safetensors. Defaults to False.
+
+    Returns:
+        tuple: state_dict, config, shard_file: file name, sharded_index: map for weight to file name.
+    """
+    paddle.device.cuda.empty_cache()
+    assert hasattr(model_to_save, "config")
+
+    state_dict = get_expected_state_dict(model_to_save, concat_additional_adapter=True)
+    all_filter_keys = filter_params(model_to_save, state_dict, args)
+
+    config_to_save = copy.deepcopy(model_to_save.config)
+
+    if config_to_save.tensor_parallel_degree > 1:
+        if isinstance(model_to_save, LoRAModel) or isinstance(model_to_save, PrefixModelForCausalLM):
+            tp_actions = model_to_save._get_tensor_parallel_convert_actions(
+                all_filter_keys, is_split=False, ignore_error=True
+            )
+        else:
+            tp_actions = model_to_save.get_tensor_parallel_convert_actions(
+                model_to_save.config, state_dict.keys(), is_split=False, ignore_error=True
+            )
+        logger.info("Unified model tensor parallel weights in shards")
+        state_dict = merge_tensor_parallel_with_shard(state_dict, tp_actions, all_filter_keys)
+
+    # build index json file
+    index_weight_file = {}
+    total_size = 0
+    if isinstance(model_to_save, LoRAModel):
+        weights_name = SAFE_PEFT_WEIGHTS_NAME if safe_serialization else LORA_WEIGHTS_NAME
+    elif isinstance(model_to_save, PrefixModelForCausalLM):
+        weights_name = SAFE_PEFT_WEIGHTS_NAME if safe_serialization else PREFIX_WEIGHTS_NAME
+    else:
+        weights_name = SAFE_WEIGHTS_NAME if safe_serialization else PADDLE_WEIGHTS_NAME
+
+    shard_file = get_sharded_file_name(args, weights_name)
+    # renumerize shard_file name for expert_parallel.
+    if args.use_expert_parallel:
+        shard_file = rename_shard_file(args, shard_file, weights_name)
+
+    for key, weight in state_dict.items():
+        index_weight_file[key] = shard_file
+        total_size += weight.numel().item() * dtype_byte_size(weight.dtype)
+
+    index_file_list, total_size_list = gather_sharded_object(
+        index_weight_file, total_size, use_expert_parallel=args.use_expert_parallel
+    )
+    sharded_index = get_sharded_index(
+        index_file_list,
+        total_size_list,
+    )
+    if sharded_index is not None:
+        if isinstance(model_to_save, LoRAModel):
+            sharded_index["type"] = "lora"
+        elif isinstance(model_to_save, PrefixModelForCausalLM):
+            sharded_index["type"] = "ptuning"
+
+    paddle.device.cuda.empty_cache()
+
+    return state_dict, shard_file, sharded_index
+
+
+def unified_optimizer_into_shards(
+    args,
+    model,
+    optim_state_dict,
+    master_weights,
+    safe_serialization=False,
+):
+    """Get optimizer state dict and master weight state dict.
+
+    Args:
+        optimizer (Optimizer): optimizer to save.
+        safe_serialization (bool, optional): safe serialization using safetensors. Defaults to False.
+    """
+    paddle.device.cuda.empty_cache()
+
+    # gather global master_weights status.
+    global_master_weights = reduce_master_weights_status(master_weights is not None)
+    if master_weights is None and global_master_weights:
+        master_weights = {}
+
+    # get optimizer param mappings
+    static2struct_name_mappings = {}
+    state_dict = get_expected_state_dict(model)
+    fp32_weight = {}
+    for k, v in state_dict.items():
+        static2struct_name_mappings[v.name] = k
+        if master_weights is not None and v.dtype == paddle.float32:
+            if args.dataset_rank > 0:  # deal with different dataset rank.
+                continue
+            fp32_weight[k] = v
+
+    # rename optimizer param
+    for key in list(optim_state_dict.keys()):
+        static_name, type_name = generate_base_static_name(key)
+        new_name = static2struct_name_mappings[static_name] + "/" + type_name
+        optim_state_dict[new_name] = optim_state_dict.pop(key)
+
+    if master_weights is not None:
+        for key in list(master_weights.keys()):
+            master_weights[static2struct_name_mappings[key]] = master_weights.pop(key)
+        master_weights.update(fp32_weight)
+
+    # filter optimizer param
+    if master_weights is not None:
+        filter_master_keys = filter_params(model, master_weights, args, is_optimizer=True)
+    filter_optim_keys = filter_params(model, optim_state_dict, args, is_optimizer=True)
+
+    tp_group = fleet.get_hybrid_communicate_group().get_model_parallel_group()
+    tp_size = tp_group.nranks
+
+    if tp_size > 1:
+        # get tp_actions
+        model_keys = []
+        for key in optim_state_dict.keys():
+            base_model_key = key.split("/")[0]
+            if base_model_key not in model_keys:
+                model_keys.append(base_model_key)
+        if isinstance(model, LoRAModel) or isinstance(model, PrefixModelForCausalLM):
+            tp_actions = model._get_tensor_parallel_convert_actions(model_keys, is_split=False, ignore_error=True)
+        else:
+            tp_actions = model.get_tensor_parallel_convert_actions(
+                model.config, model_keys, is_split=False, ignore_error=True
+            )
+        logger.info("Unified optimizer tensor parallel in shards")
+        optim_state_dict = merge_tensor_parallel_for_optimizer(
+            optim_state_dict,
+            tp_actions,
+            filter_optim_keys,
+            state_dict if args.use_expert_parallel else None,
+        )
+        paddle.device.cuda.empty_cache()
+
+        if master_weights is not None:
+            logger.info("Unified master weight tensor parallel in shards")
+            master_weights = merge_tensor_parallel_for_optimizer(
+                master_weights,
+                tp_actions,
+                filter_master_keys,
+                state_dict if args.use_expert_parallel else None,
+            )
+            paddle.device.cuda.empty_cache()
+
+    # build index json file
+    index_optimizer_file, index_master_weight_file = {}, {}
+    total_optim_size, total_master_weight_size = 0, 0
+    optimizer_name = SAFE_OPTIMIZER_NAME if safe_serialization else PADDLE_OPTIMIZER_NAME
+    master_weights_name = SAFE_MASTER_WEIGHTS_NAME if safe_serialization else PADDLE_MASTER_WEIGHTS_NAME
+    if UnifiedCheckpointOption.SKIP_SAVE_MODEL_WEIGHT.value in args.unified_checkpoint_config:
+        master_weights_name = SAFE_WEIGHTS_NAME if safe_serialization else PADDLE_WEIGHTS_NAME
+    shard_optimizer_file = get_sharded_file_name(args, optimizer_name, is_optimizer=True)
+    shard_master_weight_file = get_sharded_file_name(args, master_weights_name, is_optimizer=True)
+
+    for key, weight in optim_state_dict.items():
+        index_optimizer_file[key] = shard_optimizer_file
+        total_optim_size += weight.numel().item() * dtype_byte_size(weight.dtype)
+
+    if master_weights is not None:
+        for key, weight in master_weights.items():
+            index_master_weight_file[key] = shard_master_weight_file
+            total_master_weight_size += weight.numel().item() * dtype_byte_size(weight.dtype)
+
+    index_optimizer_filelist, total_optim_size_list = gather_sharded_object(
+        index_optimizer_file,
+        total_optim_size,
+        is_optimizer=True,
+        use_expert_parallel=args.use_expert_parallel,
+    )
+    sharded_optim_index = get_sharded_index(index_optimizer_filelist, total_optim_size_list)
+
+    if args.should_save:
+        if args.ckpt_quant_stage in ["O1", "O2"] and "quant_reach_limit" not in infohub:
+            sharded_optim_index["ckpt_quant_stage"] = args.ckpt_quant_stage
+        sharded_optim_index["quant_ckpt_resume_times"] = (
+            infohub["quant_ckpt_resume_times"] if "quant_ckpt_resume_times" in infohub else 0
+        )
+
+    if master_weights is not None:
+        index_master_weight_filelist, total_master_weight_size_list = gather_sharded_object(
+            index_master_weight_file,
+            total_master_weight_size,
+            is_optimizer=True,
+            use_expert_parallel=args.use_expert_parallel,
+        )
+        sharded_master_weight_index = get_sharded_index(index_master_weight_filelist, total_master_weight_size_list)
+
+    if sharded_optim_index is not None:
+        if master_weights is not None:
+            sharded_optim_index["master_weights"] = True
+        else:
+            sharded_optim_index["master_weights"] = False
+
+    paddle.device.cuda.empty_cache()
+    if master_weights is None:
+        return [(optim_state_dict, shard_optimizer_file, sharded_optim_index)]
+    else:
+        return [
+            (optim_state_dict, shard_optimizer_file, sharded_optim_index),
+            (master_weights, shard_master_weight_file, sharded_master_weight_index),
+        ]

paddlenlp/trainer/unified_checkpoint/utils.py

@@ -0,0 +1,781 @@
+# Copyright (c) 2024 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.
+"""Unified Checkpoint Utility Functions"""
+
+import copy
+import os
+
+import numpy as np
+import paddle
+import paddle.distributed as dist
+from paddle.distributed import fleet
+
+from paddlenlp.peft import LoRAModel, PrefixModelForCausalLM
+from paddlenlp.trainer.trainer_utils import ExplicitEnum, ShardingOption
+from paddlenlp.trainer.utils.helper import distributed_isfile
+from paddlenlp.transformers.model_utils import (
+    PretrainedModel,
+    get_parameter_dtype,
+    unwrap_model,
+)
+from paddlenlp.transformers.utils import dtype_byte_size
+from paddlenlp.utils.distributed import distributed_allgather, distributed_gather
+from paddlenlp.utils.env import (
+    BETA1_KEYNAME,
+    BETA2_KEYNAME,
+    MOMENT1_KEYNAME,
+    MOMENT2_KEYNAME,
+    PADDLE_MASTER_WEIGHTS_INDEX_NAME,
+    PADDLE_PEFT_WEIGHTS_INDEX_NAME,
+    PADDLE_WEIGHTS_INDEX_NAME,
+    PAST_KEY_VALUES_FILE_NAME,
+    SAFE_MASTER_WEIGHTS_INDEX_NAME,
+    SAFE_PEFT_WEIGHTS_INDEX_NAME,
+    SAFE_WEIGHTS_INDEX_NAME,
+)
+from paddlenlp.utils.log import logger
+from paddlenlp.utils.nested import flatten_list
+from paddlenlp.utils.tools import get_env_device
+
+FP32_MASTER = "fp32_master_0"
+optimizer_scalar_name = [
+    "beta1_pow_acc_0",
+    "beta2_pow_acc_0",
+]
+optimizer_non_scaler_name = [
+    "moment1_0",
+    "moment2_0",
+    "velocity_0",
+]  # to be added
+
+
+DEST_PLACE = paddle.CPUPlace()
+if paddle.device.is_compiled_with_cuda():
+    DEST_PLACE = paddle.CUDAPinnedPlace()
+
+
+class UnifiedCheckpointOption(ExplicitEnum):
+    """
+    "- skip_save_model_weight: do not save model weights when the masters weight exist\n"
+    "- master_weight_compatible: 1. if the master weights exist, only load when needed\n"
+    "                            2. if master weights does not exist, convert model weights to master weights when needed\n"
+    "- async_save: enable asynchronous saving checkpoints to disk\n"
+    "- enable_all_options: enable all optimization configurations\n"
+    """
+
+    SKIP_SAVE_MODEL_WEIGHT = "skip_save_model_weight"
+    MASTER_WEIGHT_COMPATIBLE = "master_weight_compatible"
+    REMOVE_MASTER_WEIGHT = "remove_master_weight"
+    ASYNC_SAVE = "async_save"
+    IGNORE_MERGE_OPTIMIZER = "ignore_merge_optimizer"
+
+
+def unwrap_optimizer(optimizer):
+    while hasattr(optimizer, "_inner_opt") or hasattr(optimizer, "_optim"):
+        if hasattr(optimizer, "_inner_opt"):
+            optimizer = optimizer._inner_opt
+        if hasattr(optimizer, "_optim"):
+            optimizer = optimizer._optim
+    return optimizer
+
+
+def is_need_master_weight(optimizer, is_fp16_or_bp16):
+    optimizer = unwrap_optimizer(optimizer)
+    if hasattr(optimizer, "_multi_precision"):
+        return optimizer._multi_precision and is_fp16_or_bp16
+    else:
+        return False
+
+
+def update_master_weight_status(args, optimizer, has_master_weight, safe_serialization):
+    if is_need_master_weight(optimizer, is_fp16_or_bp16=(args.fp16 or args.bf16)):
+        if not has_master_weight:
+            if (
+                UnifiedCheckpointOption.REMOVE_MASTER_WEIGHT.value in args.unified_checkpoint_config
+                or UnifiedCheckpointOption.MASTER_WEIGHT_COMPATIBLE.value in args.unified_checkpoint_config
+            ):
+                index_filename_master_weights = (
+                    PADDLE_WEIGHTS_INDEX_NAME if not safe_serialization else SAFE_WEIGHTS_INDEX_NAME
+                )
+                has_master_weight = True
+                logger.warning(
+                    "The unified checkpoint does not contain master weight, "
+                    "the model weight will be loaded as master weight."
+                )
+            else:
+                raise ValueError(
+                    "Can't find a valid unified master weight checkpoint,"
+                    f"add '{UnifiedCheckpointOption.MASTER_WEIGHT_COMPATIBLE.value}'"
+                    f" or '{UnifiedCheckpointOption.REMOVE_MASTER_WEIGHT.value}' into 'unified_checkpoint_config' to "
+                    "load model checkpoint as master weight"
+                )
+        else:
+            has_master_weight = True
+            index_filename_master_weights = (
+                PADDLE_MASTER_WEIGHTS_INDEX_NAME if not safe_serialization else SAFE_MASTER_WEIGHTS_INDEX_NAME
+            )
+            if UnifiedCheckpointOption.SKIP_SAVE_MODEL_WEIGHT.value in args.unified_checkpoint_config:
+                index_filename_master_weights = (
+                    PADDLE_WEIGHTS_INDEX_NAME if not safe_serialization else SAFE_WEIGHTS_INDEX_NAME
+                )
+    else:
+        has_master_weight = False
+        index_filename_master_weights = None
+
+    return has_master_weight, index_filename_master_weights
+
+
+def reduce_master_weights_status(has_master_weights=False):
+    """
+    Get master_weight status througn tp, pp and sharding group.
+    """
+    data = paddle.to_tensor([has_master_weights], dtype="int32")
+
+    hcg = fleet.get_hybrid_communicate_group()
+    tp_group = hcg.get_model_parallel_group()
+    pp_group = hcg.get_pipe_parallel_group()
+    sharding_group = hcg.get_sharding_parallel_group()
+
+    if tp_group.nranks > 1:
+        dist.all_reduce(data, op=dist.ReduceOp.SUM, group=tp_group)
+    if pp_group.nranks > 1:
+        dist.all_reduce(data, op=dist.ReduceOp.SUM, group=pp_group)
+    if sharding_group.nranks > 1:
+        dist.all_reduce(data, op=dist.ReduceOp.SUM, group=sharding_group)
+
+    return data.item() > 0
+
+
+def select_model_weight_index(model, resume_from_checkpoint, safe_serialization, local=True):
+    """
+    try select model weight index from model weight or master weight index.
+    """
+
+    # find model weight index file
+    if isinstance(model, LoRAModel) or isinstance(model, PrefixModelForCausalLM):
+        index_filename = SAFE_PEFT_WEIGHTS_INDEX_NAME if safe_serialization else PADDLE_PEFT_WEIGHTS_INDEX_NAME
+    else:
+        index_filename = SAFE_WEIGHTS_INDEX_NAME if safe_serialization else PADDLE_WEIGHTS_INDEX_NAME
+
+    index_filename_path = os.path.join(resume_from_checkpoint, index_filename)
+    identify_func = os.path.isfile if local else distributed_isfile
+
+    if identify_func(index_filename_path):
+        return index_filename
+    else:
+        index_filename = PADDLE_MASTER_WEIGHTS_INDEX_NAME if not safe_serialization else SAFE_MASTER_WEIGHTS_INDEX_NAME
+        index_filename_path = os.path.join(resume_from_checkpoint, index_filename)
+
+        if identify_func(index_filename_path):
+            return index_filename
+        else:
+            raise ValueError("Can't find a valid unified model or master weight checkpoint to load.")
+
+
+def mapping_optimizer_tp_actions(tp_actions, optimizer_loaded_keys):
+    """# convert param.name to
+    param.key/moment1_0
+    or param.key/beta1_XXX
+    or param.key/beta2_XXX
+    Args:
+        tp_actions (dict): dictionay of tensor parallel actions {key: action}
+        optimizer_loaded_keys (list or set): [param.key1/moment1_0, param.key2/beta1_XXX, param.key3/beta2_XXX]
+    Returns:
+        dict: new dictionay of tensor parallel actions {key: action}
+    """
+    new_actions = {}
+    for key in optimizer_loaded_keys:
+        key_base, typename = key.split("/")
+        if typename in optimizer_non_scaler_name and key_base in tp_actions:
+            new_actions[key] = tp_actions[key_base]
+    return new_actions
+
+
+def get_expected_state_dict(model_to_save, **kwargs):
+    """
+    Get trainable state_dict of model_to_save.
+    """
+    model_to_save = unwrap_model(model_to_save)
+
+    if isinstance(model_to_save, PretrainedModel):
+        state_dict = model_to_save.state_dict()
+        if (
+            hasattr(model_to_save.config, "tie_word_embeddings")
+            and model_to_save.config.tie_word_embeddings
+            and hasattr(model_to_save, "_tied_weights_keys")
+            and model_to_save._tied_weights_keys is not None
+        ):
+            for key in model_to_save._tied_weights_keys:
+                if key in state_dict:
+                    state_dict.pop(key)
+    elif isinstance(model_to_save, LoRAModel):
+        concat_additional_adapter = kwargs.get("concat_additional_adapter", False)
+        concat_init_lora = model_to_save.lora_config.loraga and concat_additional_adapter
+        state_dict = model_to_save.get_trainable_state_dict(concat_init_lora=concat_init_lora)
+    elif isinstance(model_to_save, PrefixModelForCausalLM):
+        state_dict = model_to_save.prefix_encoder.state_dict()
+
+    return state_dict
+
+
+def get_expected_keys(args, sharded_metadata, model, optimizer, is_master_weights=False):
+    hcg = fleet.get_hybrid_communicate_group()
+    sharding_group = hcg.get_sharding_parallel_group()
+    sharding_rank = sharding_group.rank
+    in_sharding_parallel_model = sharding_group.nranks > 1
+    if in_sharding_parallel_model:
+        params2rank = optimizer._param2rank
+
+    model_state_dict = get_expected_state_dict(model)
+    struct2static_name_mappings = {k: v.name for k, v in model_state_dict.items()}
+
+    expected_keys = []
+    for key in list(sharded_metadata["all_optimizer_keys"]):
+        key_name = key.split("/")[0]
+        if is_master_weights and key_name in model_state_dict and model_state_dict[key_name].dtype == paddle.float32:
+            continue
+
+        if args.use_expert_parallel and args.data_parallel_rank > 0:
+            if key_name in model_state_dict and not getattr(model_state_dict[key_name], "no_sync", False):
+                continue
+
+        static_name = struct2static_name_mappings.get(key_name, None)
+
+        if in_sharding_parallel_model:
+            params_rank = params2rank.get(static_name, None)
+            if params_rank == sharding_rank:
+                expected_keys.append(key)
+        else:
+            if static_name is not None:
+                expected_keys.append(key)
+    expected_keys = set(expected_keys)
+
+    loaded_keys = sharded_metadata["all_optimizer_keys"]
+    missing_keys = expected_keys - set(loaded_keys)
+    if len(missing_keys) > 0:
+        raise ValueError(f"optimizer missing weights keys: {missing_keys}")
+
+    return expected_keys
+
+
+def get_optimizer_shard_files(optimizer_path, index_filename):
+    """
+    For a given model:
+    - download and cache all the shards of a sharded checkpoint if `pretrained_model_name_or_path` is a model ID on the
+      Hub
+    - returns the list of paths to all the shards, as well as some metadata.
+    For the description of each arg, see [`PretrainedModel.from_pretrained`]. `index_filename` is the full path to the
+    index (downloaded and cached if `pretrained_model_name_or_path` is a model ID on the Hub).
+    """
+
+    import json
+
+    if not os.path.isfile(index_filename):
+        raise ValueError(f"Can't find a optimizer index ({index_filename}) in {optimizer_path}.")
+
+    with open(index_filename, "r") as f:
+        index = json.loads(f.read())
+
+    shard_filenames = sorted(set(index["weight_map"].values()))
+    sharded_metadata = index["metadata"]
+    sharded_metadata["all_optimizer_keys"] = list(index["weight_map"].keys())
+    sharded_metadata["weight_map"] = index["weight_map"].copy()
+    sharded_metadata["master_weights"] = index.get("master_weights", False)
+
+    file_map = {file: set() for file in shard_filenames}
+    for weight, file in index["weight_map"].items():
+        file_map[file].add(weight)
+
+    sharded_metadata["file_map"] = file_map
+
+    # First, let's deal with local folder.
+    # TODO: if optimizer_path is a folder, we should check if the optimizer is already cached or not.
+    if os.path.isdir(optimizer_path):
+        shard_filenames = [os.path.join(optimizer_path, f) for f in shard_filenames]
+        return shard_filenames, sharded_metadata
+
+
+def generate_base_static_name(vname):
+    """
+    Return base static name and specific type name, like [embedding_0.w_0, moment1_0]
+    """
+    if FP32_MASTER in vname:
+        vname = vname.split("_" + FP32_MASTER + "_")
+        return vname[0], vname[1]
+    else:
+        # Directly deal with type names, for example: moe_gate_1_moment1_0.
+        type_names = optimizer_scalar_name + optimizer_non_scaler_name
+        for name in type_names:
+            if name in vname:
+                a = vname.split(name)[0][:-1]
+                b = name
+                return a, b
+
+
+def merge_large_tensor_parallel(tensor, tp_group, tp_action, dst_rank, is_dst):
+    """
+    Move large tensor merge process to CPU, in order to avoid OOM.
+    """
+    num_rows = tensor.shape[0]
+    num_splits = 4
+    parts = np.array_split(np.arange(num_rows), num_splits)
+    splits = [len(part) for part in parts]
+    split_parts = np.insert(np.cumsum(splits), 0, 0)
+    split_tensors = []
+    for i in range(num_splits):
+        if get_env_device() == "xpu":
+            ret = distributed_allgather(tensor[split_parts[i] : split_parts[i + 1], :], group=tp_group, offload=False)
+        else:
+            ret = distributed_gather(
+                tensor[split_parts[i] : split_parts[i + 1], :], dst=dst_rank, group=tp_group, offload=False
+            )
+        # Copy to CPUPlace temporarily, may lower speed.
+        if ret is not None:
+            ret = [t.cpu() for t in ret]
+        split_tensors.append(ret)
+    concat_tensors = []
+    if is_dst:
+        for i in range(tp_group.nranks):
+            tmp = []
+            for j in range(num_splits):
+                tmp.append(split_tensors[j][i])
+            concat_tensors.append(paddle.concat(tmp))
+        tensor = tp_action(concat_tensors)
+    else:
+        tensor = None
+    return tensor
+
+
+def merge_tensor_parallel_with_shard(state_dict, tp_actions, all_filter_keys):
+    """
+    Merge tensor parallel according to tp_actions, used for model weight.
+    """
+    hcg = fleet.get_hybrid_communicate_group()
+    tp_group = hcg.get_model_parallel_group()
+    dp_group = hcg.get_data_parallel_group()
+    tp_rank = tp_group.rank
+    dp_rank = dp_group.rank if dp_group.nranks > 1 else 0
+
+    # filter actions for pipeline mode
+    if hcg.get_pipe_parallel_group().nranks > 1:
+        filter_keys = set([y for x in all_filter_keys for y in x])
+        for key in list(tp_actions.keys()):
+            if key not in filter_keys:
+                tp_actions.pop(key)
+
+    state_dict_to_save = {}
+    max_key_len = max([len(_) for _ in all_filter_keys])
+    for i in range(max_key_len):
+        for j, filter_keys in enumerate(all_filter_keys):
+            is_dst = tp_rank == j
+            if i > len(filter_keys) - 1:
+                continue
+            key = filter_keys[i]
+            tensor = state_dict[key]
+            # When using expert parallel, there's no need to save tensors with `no_sync=False` when dp_rank > 0.
+            if dp_rank > 0 and not getattr(tensor, "no_sync", False):
+                continue
+            if key in tp_actions:
+                # Get tensor size
+                tensor_bytes = tensor.numel().item() * dtype_byte_size(tensor.dtype) * tp_group.nranks
+                if tensor_bytes >= 5 * 1024 * 1024 * 1024:  # temporarily set 5GB as threshold
+                    tensor = merge_large_tensor_parallel(tensor, tp_group, tp_actions[key], j, is_dst)
+                else:
+                    if get_env_device() == "xpu":
+                        ret = distributed_allgather(tensor, group=tp_group, offload=False)
+                    else:
+                        ret = distributed_gather(tensor, dst=j, group=tp_group, offload=False)
+                    action = tp_actions.pop(key)
+                    tensor = action(ret) if is_dst else None
+            else:
+                if is_dst:
+                    tensor = tensor._copy_to(DEST_PLACE, False) if tensor.place.is_cpu_place() else tensor
+                else:
+                    tensor = None
+
+            if is_dst:
+                state_dict_to_save[key] = tensor
+
+    if len(tp_actions) > 0:
+        for x in tp_actions.keys():
+            logger.debug(f"key <{x}> need to merge tensor parallel but we can't find in model state.")
+
+    return state_dict_to_save
+
+
+def merge_tensor_parallel_for_optimizer(state_dict, tp_actions, all_filter_keys, model_state_dict=None):
+    """
+    Merge tensor parallel according to tp_actions, used for master_weight and optimizer weight.
+    """
+    hcg = fleet.get_hybrid_communicate_group()
+    tp_group = hcg.get_model_parallel_group()
+    dp_group = hcg.get_data_parallel_group()
+    tp_rank = tp_group.rank
+    dp_rank = dp_group.rank if dp_group.nranks > 1 else 0
+
+    no_sync_kname = []
+    if model_state_dict is not None:
+        for k, v in model_state_dict.items():
+            if getattr(v, "no_sync", False):
+                no_sync_kname.append(k)
+
+    state_dict_to_save = {}
+    max_key_len = max([len(_) for _ in all_filter_keys])
+    for i in range(max_key_len):
+        for j, filter_keys in enumerate(all_filter_keys):
+            is_dst = tp_rank == j
+            if i > len(filter_keys) - 1:
+                continue
+            # get base model key
+            model_key = filter_keys[i].split("/")[0]
+            tensor = state_dict[filter_keys[i]]
+            # When using expert parallel, there's no need to save tensors with `no_sync=False` when dp_rank > 0.
+            if dp_rank > 0 and model_key not in no_sync_kname:
+                continue
+            if model_key in tp_actions:
+                # for example: beta1, beta2
+                if tensor.numel().item() == 1:
+                    if is_dst:
+                        tensor = tensor._copy_to(DEST_PLACE, False) if not tensor.place.is_cpu_place() else tensor
+                    else:
+                        tensor = None
+                else:
+                    # Get tensor size
+                    tensor_bytes = tensor.numel().item() * dtype_byte_size(tensor.dtype) * tp_group.nranks
+                    if tensor_bytes >= 5 * 1024 * 1024 * 1024:  # temporarily set 5GB as threshold
+                        tensor = merge_large_tensor_parallel(tensor, tp_group, tp_actions[model_key], j, is_dst)
+                    else:
+                        if get_env_device() == "xpu":
+                            ret = distributed_allgather(tensor, group=tp_group, offload=False)
+                        else:
+                            ret = distributed_gather(tensor, dst=j, group=tp_group, offload=False)
+                        action = tp_actions[model_key]
+                        tensor = action(ret) if is_dst else None
+            else:
+                if is_dst:
+                    tensor = tensor._copy_to(DEST_PLACE, False) if not tensor.place.is_cpu_place() else tensor
+                else:
+                    tensor = None
+
+            if is_dst:
+                state_dict_to_save[filter_keys[i]] = tensor
+
+    return state_dict_to_save
+
+
+def filter_params(model_to_save, state_dict, args, is_optimizer=False):
+    """
+    Group according to the size of the tensor, aiming to make the weight size
+    stored on each device as equal as possible.
+    """
+    hcg = fleet.get_hybrid_communicate_group()
+    tp_group = hcg.get_model_parallel_group()
+
+    tp_size = tp_group.nranks
+    tp_rank = tp_group.rank
+
+    # for pure sharding or pure pp
+    if tp_size <= 1:
+        return [list(state_dict.keys())]
+
+    filter_tensor_list = [[] for _ in range(tp_size)]
+    is_master_weights = False
+
+    if tp_rank == 0:
+        quant = False
+        if args.ckpt_quant_stage != "O0":
+            quant = True
+        tensor_bytes_dict = {}
+        model_state_dict = get_expected_state_dict(model_to_save)
+        for (k, v) in state_dict.items():
+            # master weight has same key as model weight
+            if not is_master_weights and k in model_state_dict:
+                is_master_weights = True
+
+            weight_key = k.split("/")[0]
+            model_v = model_state_dict[weight_key] if is_optimizer else v
+            if not quant or not is_optimizer:
+                if hasattr(model_v, "is_distributed") and model_v.is_distributed:
+                    tensor_bytes_dict[k] = v.numel().item() * tp_size * dtype_byte_size(v.dtype)
+                else:
+                    tensor_bytes_dict[k] = v.numel().item() * dtype_byte_size(v.dtype)
+            else:
+                if weight_key not in tensor_bytes_dict:
+                    tensor_bytes_dict[weight_key] = 0
+
+                if hasattr(model_v, "is_distributed") and model_v.is_distributed:
+                    tensor_bytes_dict[weight_key] += v.numel().item() * tp_size * dtype_byte_size(v.dtype)
+                else:
+                    tensor_bytes_dict[weight_key] += v.numel().item() * dtype_byte_size(v.dtype)
+
+        filter_tensor_list = []
+        current_block = []
+        current_block_size = 0
+        total_size = 0
+
+        max_shard_size = (sum(tensor_bytes_dict.values()) + tp_size - 1) // tp_size
+
+        for index, (key, weight_size) in enumerate(tensor_bytes_dict.items()):
+            # If this weight is going to tip up over the maximal size, we split.
+            # if current_block_size + weight_size > max_shard_size:
+            if total_size + weight_size > max_shard_size * (len(filter_tensor_list) + 1) or (
+                len(tensor_bytes_dict) - index < (tp_size - len(filter_tensor_list))
+            ):
+                # fix if the first param is large than max_shard_size
+                if len(current_block) > 0:
+                    filter_tensor_list.append(current_block)
+                current_block = []
+                current_block_size = 0
+
+            if not quant or not is_optimizer or is_master_weights:
+                current_block.append(key)
+            else:
+                current_block.append(key + "/" + MOMENT1_KEYNAME)
+                current_block.append(key + "/" + MOMENT2_KEYNAME)
+                current_block.append(key + "/" + BETA1_KEYNAME)
+                current_block.append(key + "/" + BETA2_KEYNAME)
+
+            current_block_size += weight_size
+            total_size += weight_size
+
+        filter_tensor_list.append(current_block)
+        if len(filter_tensor_list) < tp_size:
+            filter_tensor_list.extend([[] for i in range(tp_size - len(filter_tensor_list))])
+
+    dist.broadcast_object_list(
+        filter_tensor_list,
+        src=hcg.get_model_parallel_group_src_rank(),
+        group=tp_group,
+    )
+
+    return filter_tensor_list
+
+
+def get_sharded_file_name(args, file_name, is_optimizer=False):
+    """
+    Get safetensors file name for saving.
+    """
+    if not is_optimizer:
+        sd_degree = args.sharding_parallel_degree if args.sharding_parallel_degree > 1 else 1
+        size = sd_degree if args.use_expert_parallel else args.dataset_world_size
+        shard_file = file_name.replace(
+            ".pdparams",
+            f"-{args.logical_process_index + 1:05d}-of-{args.world_size//size:05d}.pdparams",
+        )
+        shard_file = shard_file.replace(
+            ".safetensors",
+            f"-{args.logical_process_index + 1:05d}-of-{args.world_size//size:05d}.safetensors",
+        )
+    else:
+        hcg = fleet.get_hybrid_communicate_group()
+        dp_group = hcg.get_data_parallel_group()
+        size = dp_group.nranks if not args.use_expert_parallel else 1
+        shard_file = file_name.replace(
+            ".pdparams", f"-{args.logical_process_index + 1:05d}-of-{args.world_size//size:05d}.pdparams"
+        )
+        shard_file = shard_file.replace(
+            ".safetensors",
+            f"-{args.logical_process_index + 1:05d}-of-{args.world_size//size:05d}.safetensors",
+        )
+        shard_file = shard_file.replace(
+            ".pdopt", f"-{args.logical_process_index + 1:05d}-of-{args.world_size//size:05d}.pdopt"
+        )
+    return shard_file
+
+
+def get_sharded_index(
+    index_file_list,
+    total_size_list,
+):
+    """
+    Save safetensors index json file, including metadata and weight_map.
+    """
+    local_rank = int(os.getenv("PADDLE_RANK_IN_NODE", 0))
+    if local_rank == 0:
+        sharded_index_json = {}
+
+        sharded_index_json["metadata"] = {"total_size": sum(total_size_list)}
+
+        weight_map = {}
+        for i, _ in enumerate(index_file_list):
+            weight_map.update(index_file_list[i])
+
+        sharded_index_json["weight_map"] = weight_map
+        return sharded_index_json
+
+    return None
+
+
+def gather_sharded_object(index_file, total_size, is_optimizer=False, use_expert_parallel=False):
+    """
+    All gather sharded files list across different groups.
+    """
+    index_file_list, total_size_list = [], []
+
+    hcg = fleet.get_hybrid_communicate_group()
+    tp_group = hcg.get_model_parallel_group()
+    pp_group = hcg.get_pipe_parallel_group()
+
+    logger.info(
+        f"Unified checkpoint: generating sharded_index json files for {'optimizer or master weight' if is_optimizer else 'model weight'}."
+    )
+
+    if tp_group.nranks > 1:
+        dist.all_gather_object(index_file_list, index_file, tp_group)
+        dist.all_gather_object(total_size_list, total_size, tp_group)
+    if pp_group.nranks > 1:
+        pp_index_file_list = []
+        pp_total_size_list = []
+        dist.all_gather_object(
+            pp_index_file_list, index_file_list if len(index_file_list) > 0 else index_file, pp_group
+        )
+        dist.all_gather_object(
+            pp_total_size_list, total_size_list if len(total_size_list) > 0 else total_size, pp_group
+        )
+        index_file_list = pp_index_file_list
+        total_size_list = pp_total_size_list
+
+    index_file_list = flatten_list(index_file_list)
+    total_size_list = flatten_list(total_size_list)
+
+    # for pure sharding
+    if len(index_file_list) == 0 and len(total_size_list) == 0:
+        index_file_list = [index_file]
+        total_size_list = [total_size]
+
+    if use_expert_parallel:
+        data_group = hcg.get_data_parallel_group()
+        if data_group.nranks > 1:
+            data_index_file_list = []
+            data_total_size_list = []
+            dist.all_gather_object(data_index_file_list, index_file_list, data_group)
+            dist.all_gather_object(data_total_size_list, total_size_list, data_group)
+            index_file_list = flatten_list(data_index_file_list)
+            total_size_list = flatten_list(data_total_size_list)
+
+    if is_optimizer:
+        sharding_group = hcg.get_sharding_parallel_group()
+        if sharding_group.nranks > 1:
+            sharding_index_file_list = []
+            sharding_total_size_list = []
+            dist.all_gather_object(sharding_index_file_list, index_file_list, sharding_group)
+            dist.all_gather_object(sharding_total_size_list, total_size_list, sharding_group)
+            index_file_list = flatten_list(sharding_index_file_list)
+            total_size_list = flatten_list(sharding_total_size_list)
+
+    return index_file_list, total_size_list
+
+
+def rename_shard_file(args, shard_file, file_name):
+    """
+    Rename shard file when using expert_parallel.
+    """
+    assert args.use_expert_parallel, "only expert_parallel need to use this function"
+
+    shard_file_list = []
+
+    hcg = fleet.get_hybrid_communicate_group()
+    tp_group = hcg.get_model_parallel_group()
+    pp_group = hcg.get_pipe_parallel_group()
+    data_group = hcg.get_data_parallel_group()
+
+    if tp_group.nranks > 1:
+        dist.all_gather_object(shard_file_list, shard_file, tp_group)
+    if pp_group.nranks > 1:
+        pp_shard_file_list = []
+        dist.all_gather_object(
+            pp_shard_file_list, shard_file_list if len(shard_file_list) > 0 else shard_file, pp_group
+        )
+        shard_file_list = flatten_list(pp_shard_file_list)
+    if data_group.nranks > 1:
+        data_shard_file_list = []
+        dist.all_gather_object(
+            data_shard_file_list, shard_file_list if len(shard_file_list) > 0 else shard_file, data_group
+        )
+        shard_file_list = flatten_list(data_shard_file_list)
+
+    new_index = shard_file_list.index(shard_file)
+    sd_degree = args.sharding_parallel_degree if args.sharding_parallel_degree > 1 else 1
+    shard_file = file_name.replace(
+        ".pdparams",
+        f"-{new_index + 1:05d}-of-{args.world_size//sd_degree:05d}.pdparams",
+    )
+    shard_file = shard_file.replace(
+        ".safetensors",
+        f"-{new_index + 1:05d}-of-{args.world_size//sd_degree:05d}.safetensors",
+    )
+    return shard_file
+
+
+def save_prefix_past_key_value(model_to_save, save_directory):
+    """
+    Used only for PrefixModelForCausalLM.
+    """
+    past_key_value = model_to_save.prefix_encoder(model_to_save.prefix_tokens.unsqueeze(0).expand([1, -1]))
+    past_key_value = past_key_value.reshape(
+        [
+            model_to_save.prefix_config.num_prefix_tokens,
+            2,
+            model_to_save.prefix_config.num_hidden_layers,
+            model_to_save.num_heads,
+            model_to_save.head_dim,
+        ]
+    )
+    past_key_value = paddle.transpose(past_key_value, perm=[2, 1, 3, 0, 4]).cpu().numpy()
+    np.save(os.path.join(save_directory, PAST_KEY_VALUES_FILE_NAME), past_key_value)
+
+
+def is_sharding_split_param_mode(args):
+    return (
+        args.sharding_parallel_degree > 1
+        and ShardingOption.SHARD_OP in args.sharding
+        and "split_param" in args.sharding_parallel_config
+    )
+
+
+def save_model_config(model_to_save, save_directory):
+    """
+    Save model config.
+    """
+
+    def save_config(model_to_save):
+        dtype = get_parameter_dtype(model_to_save)
+        model_to_save.config.dtype = str(dtype).split(".")[1]
+        config_to_save = copy.deepcopy(model_to_save.config)
+
+        if config_to_save.tensor_parallel_degree > 1:
+            # do we need to change?
+            config_to_save.tensor_parallel_degree = 1
+
+        return config_to_save
+
+    # Save prefix model past_key_values
+    if isinstance(model_to_save, PrefixModelForCausalLM):
+        save_prefix_past_key_value(model_to_save, save_directory)
+        model_to_save.prefix_config.save_pretrained(save_directory)
+    if isinstance(model_to_save, LoRAModel):
+        model_to_save.lora_config.save_pretrained(save_directory)
+
+    # save the config
+    config_to_save = save_config(model_to_save)
+    # Attach architecture to the config
+    if isinstance(model_to_save, LoRAModel) or isinstance(model_to_save, PrefixModelForCausalLM):
+        config_to_save.architectures = [model_to_save.model.__class__.__name__]
+    else:
+        config_to_save.architectures = [model_to_save.__class__.__name__]
+
+    config_to_save.save_pretrained(save_directory)
+    # save generation config
+    if model_to_save.can_generate():
+        model_to_save.generation_config.save_pretrained(save_directory)

paddlenlp/transformers/contrastive_loss.py

@@ -0,0 +1,65 @@
+# Copyright (c) 2024 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 typing import List, Optional
+
+import paddle
+import paddle.nn as nn
+
+
+class SimpleContrastiveLoss(nn.Layer):
+    def __init__(self, embedding_temperature: float = 0.02):
+        super().__init__()
+        self.embedding_temperature = embedding_temperature
+        self.cross_entropy = nn.CrossEntropyLoss(reduction="mean")
+
+    def forward(self, q_reps, p_reps):
+        scores = paddle.matmul(q_reps, p_reps.transpose([1, 0]))
+        scores = scores / self.embedding_temperature
+
+        group_size = p_reps.shape[0] // q_reps.shape[0]
+        batch_size = q_reps.shape[0]
+
+        target = paddle.arange(batch_size, dtype="int64")
+        target = target * group_size
+
+        loss = self.cross_entropy(scores, target)
+        return loss
+
+
+class MatryoshkaContrastiveLoss(nn.Layer):
+    def __init__(self, embedding_temperature: float = 0.02, embedding_matryoshka_dims: Optional[List[int]] = None):
+        super().__init__()
+        self.embedding_temperature = embedding_temperature
+        if embedding_matryoshka_dims is None:
+            self.embedding_matryoshka_dims = []
+        else:
+            self.embedding_matryoshka_dims = embedding_matryoshka_dims
+        self.loss_fn = SimpleContrastiveLoss(embedding_temperature)
+
+    def forward(self, q_reps, p_reps):
+        if len(self.embedding_matryoshka_dims) > 0:
+            loss = 0.0
+            for dim in self.embedding_matryoshka_dims:
+                reduced_q_reps = q_reps[:, :dim]
+                reduced_q_reps = nn.functional.normalize(reduced_q_reps, axis=-1)
+
+                reduced_p_reps = p_reps[:, :dim]
+                reduced_p_reps = nn.functional.normalize(reduced_p_reps, axis=-1)
+
+                dim_loss = self.loss_fn(reduced_q_reps, reduced_p_reps)
+                loss += dim_loss
+        else:
+            loss = self.loss_fn(q_reps, p_reps)
+        return loss

paddlenlp/transformers/embedding_utils.py

@@ -0,0 +1,51 @@
+# Copyright (c) 2024 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 paddle
+from paddle.distributed import fleet
+
+
+def dist_gather_tensor_with_gradient(tensor):
+    if tensor is None:
+        return None
+
+    if paddle.distributed.get_world_size() <= 1:
+        return tensor
+
+    hcg = fleet.get_hybrid_communicate_group()
+    sharding_group = hcg.get_sharding_parallel_group()
+    sharding_rank = sharding_group.rank
+    data_group = hcg.get_data_parallel_group()
+    data_rank = data_group.rank
+
+    if sharding_group.nranks == 1 and data_group.nranks == 1:
+        return tensor
+
+    if sharding_group.nranks > 1:
+        all_tensors = []
+        paddle.distributed.all_gather(all_tensors, tensor.contiguous(), group=sharding_group)
+        all_tensors[sharding_rank] = tensor
+        all_tensors = paddle.concat(all_tensors, axis=0)
+    else:
+        all_tensors = tensor
+
+    if data_group.nranks > 1:
+        final_tensors = []
+        paddle.distributed.all_gather(final_tensors, all_tensors.contiguous(), group=data_group)
+        final_tensors[data_rank] = all_tensors
+        final_tensors = paddle.concat(final_tensors, axis=0)
+    else:
+        final_tensors = all_tensors
+
+    return final_tensors

paddlenlp/transformers/llama/fusion_ops.py

@@ -219,13 +219,23 @@ def fusion_flash_attention(
                     assert alibi is None, "flash_attention_with_sparse_mask not support alibi"
                     if len(attn_mask_startend_row_indices.shape) == 2:
                         attn_mask_startend_row_indices = paddle.unsqueeze(attn_mask_startend_row_indices, axis=1)
-                    attn_output = F.flash_attention_with_sparse_mask(
-                        query_states,
-                        key_states,
-                        value_states,
-                        attn_mask_start_row_indices=attn_mask_startend_row_indices,
-                        is_causal=True,
-                    )
+                    # NOTE: supports new paddle which has no flash_attention_with_sparse_mask api.
+                    if hasattr(F, "flashmask_attention"):
+                        attn_output = F.flashmask_attention(
+                            query_states,
+                            key_states,
+                            value_states,
+                            startend_row_indices=attn_mask_startend_row_indices.unsqueeze(-1),
+                            causal=True,
+                        )
+                    else:
+                        attn_output = F.flash_attention_with_sparse_mask(
+                            query_states,
+                            key_states,
+                            value_states,
+                            attn_mask_start_row_indices=attn_mask_startend_row_indices,
+                            is_causal=True,
+                        )
                 else:
                     attn_output = F.scaled_dot_product_attention(
                         query_states,

paddlenlp/transformers/model_utils.py

@@ -13,6 +13,7 @@
 # limitations under the License.
 from __future__ import annotations
 
+import concurrent.futures
 import contextlib
 import copy
 import gc
@@ -53,6 +54,8 @@ from paddle.utils.download import is_url as is_remote_url
 from tqdm.auto import tqdm
 
 from paddlenlp.utils.env import (
+    ASYMMETRY_QUANT_SCALE_MAX,
+    ASYMMETRY_QUANT_SCALE_MIN,
     CONFIG_NAME,
     LEGACY_CONFIG_NAME,
     PADDLE_WEIGHTS_INDEX_NAME,
@@ -63,10 +66,12 @@ from paddlenlp.utils.env import (
     SAFE_PEFT_WEIGHTS_INDEX_NAME,
     SAFE_WEIGHTS_INDEX_NAME,
     SAFE_WEIGHTS_NAME,
+    SYMMETRY_QUANT_SCALE,
 )
 from paddlenlp.utils.log import logger
 
 from ..generation import GenerationConfig, GenerationMixin
+from ..quantization.unified_checkpoint_quantization import dequant_unified_optimizer
 from ..utils import device_guard
 from ..utils.download import resolve_file_path
 from .configuration_utils import PretrainedConfig
@@ -319,12 +324,98 @@ def get_parameter_dtype(parameter: nn.Layer) -> paddle.dtype:
     return last_dtype
 
 
+def _split_keys_evenly(keys: list, n: int) -> list:
+    """Split a list into n lists with an equal number of elements.
+
+    Args:
+        keys (list): the list to be split
+        n (int): number of splits
+
+    Returns:
+        result: list of lists
+    """
+
+    total_len = len(keys)
+    base_size = total_len // n
+    extra = total_len % n
+
+    result = []
+    index = 0
+    for _ in range(n):
+        part_size = base_size + 1 if extra > 0 else base_size
+        extra -= 1
+        result.append(keys[index : index + part_size])
+        index += part_size
+
+    return result
+
+
+def _load_part_state_dict(
+    keys, checkpoint_file: Union[str, os.PathLike], tensor_parallel_split_mapping, fliter_dict_keys, device
+):
+    """load part state dict from checkpoint file.
+
+    Args:
+        keys (list): the keys of part state dict
+        checkpoint_file (str): the path of checkpoint file
+        tensor_parallel_split_mapping (dict): mapping from key to function
+        fliter_dict_keys (list): filter keys in state dict
+
+    Returns:
+        part_state_dict (dict): the part state dict
+
+    """
+    part_state_dict = {}
+    scale_dict = {}
+    with safe_open(checkpoint_file, framework="np") as f:
+        for key in keys:
+            # 1. non-merge ckpt loading dont have filter key.
+            # 2. merge ckpt will skip quant scale by `fliter_dict_keys`
+            if (
+                key.endswith(SYMMETRY_QUANT_SCALE)
+                or key.endswith(ASYMMETRY_QUANT_SCALE_MIN)
+                or key.endswith(ASYMMETRY_QUANT_SCALE_MAX)
+            ):
+                continue
+
+            if fliter_dict_keys is not None and key not in fliter_dict_keys:
+                continue
+
+            py_safe_slice_ = f.get_slice(key)
+            if key in tensor_parallel_split_mapping:
+                weight = tensor_parallel_split_mapping[key](py_safe_slice_)
+            else:
+                weight = py_safe_slice_[:]
+            if device == "expected":
+                with device_guard():
+                    weight = paddle.Tensor(weight, zero_copy=True)
+                weight = weight._copy_to(paddle.framework._current_expected_place(), False)
+            part_state_dict[key] = weight
+        for key in keys:
+            if (
+                key.endswith(SYMMETRY_QUANT_SCALE)
+                or key.endswith(ASYMMETRY_QUANT_SCALE_MIN)
+                or key.endswith(ASYMMETRY_QUANT_SCALE_MAX)
+            ):
+                scale = f.get_tensor(key)
+                with device_guard():
+                    scale = paddle.Tensor(scale, zero_copy=True)
+                scale = scale._copy_to(paddle.framework._current_expected_place(), False)
+                scale_dict[key] = scale
+    return part_state_dict, scale_dict
+
+
 def load_state_dict(
-    checkpoint_file: Union[str, os.PathLike], tensor_parallel_split_mapping=None, fliter_dict_keys=None, device="cpu"
+    checkpoint_file: Union[str, os.PathLike],
+    tensor_parallel_split_mapping=None,
+    fliter_dict_keys=None,
+    device="cpu",
+    ckpt_quant_stage="O0",
 ):
     """
     Reads a PaddlePaddle checkpoint file, returning properly formatted errors if they arise.
     """
+
     if tensor_parallel_split_mapping is None:
         tensor_parallel_split_mapping = {}
 
@@ -343,27 +434,50 @@ def load_state_dict(
         if metadata.get("format", "np") == "pd":
             raise ValueError("Currently unsupport paddle weights file, use numpy instead.")
         if metadata.get("format", "np") == "np":
-            state_dict = {}
-            with safe_open(checkpoint_file, framework="np") as f:
-                for key in f.keys():
-                    if fliter_dict_keys is not None and key not in fliter_dict_keys:
-                        continue
-                    py_safe_slice_ = f.get_slice(key)
-                    if key in tensor_parallel_split_mapping:
-                        weight = tensor_parallel_split_mapping[key](py_safe_slice_)
-                    else:
-                        weight = py_safe_slice_[:]
-                    if device == "expected":
-                        with device_guard():
-                            weight = paddle.Tensor(weight, zero_copy=True)
-                        weight = weight._copy_to(paddle.framework._current_expected_place(), False)
-                    state_dict[key] = weight
+            thread_num = int(os.environ.get("LOAD_STATE_DICT_THREAD_NUM", "1"))
+            if thread_num > 1:
+                logger.info(f"Set loading state_dict thread num to {thread_num}")
+            state_dict, scale_dict = {}, {}
+            if thread_num <= 1:
+                with safe_open(checkpoint_file, framework="np") as f:
+                    state_dict, scale_dict = _load_part_state_dict(
+                        list(f.keys()),
+                        checkpoint_file,
+                        tensor_parallel_split_mapping,
+                        fliter_dict_keys,
+                        device,
+                    )
+            else:
+                # Load state dict in multi-thread to speed up loading
+                with safe_open(checkpoint_file, framework="np") as f:
+                    keys_groups = _split_keys_evenly(list(f.keys()), thread_num)
+                with concurrent.futures.ThreadPoolExecutor(max_workers=thread_num) as executor:
+                    future_to_key = {
+                        executor.submit(
+                            _load_part_state_dict,
+                            keys,
+                            checkpoint_file,
+                            tensor_parallel_split_mapping,
+                            fliter_dict_keys,
+                            device,
+                        ): keys
+                        for keys in keys_groups
+                    }
+                    for future in concurrent.futures.as_completed(future_to_key):
+                        res_state_dict, res_scale_dict = future.result()
+                        state_dict.update(res_state_dict)
+                        scale_dict.update(res_scale_dict)
 
             if device == "cpu":
                 for k in list(state_dict.keys()):
                     with device_guard():
                         state_dict[k] = paddle.Tensor(state_dict.pop(k), zero_copy=True)
 
+            if len(scale_dict) != 0:
+                if ckpt_quant_stage == "O0":
+                    raise ValueError('optimizer weight has quantization scales but `ckpt_quant_stage` is set to "O0"')
+                state_dict = dequant_unified_optimizer(state_dict, ckpt_quant_stage, scale_dict, use_pd=True)
+
             return state_dict
 
     state_dict = paddlenlp_load(checkpoint_file, map_location="cpu")
@@ -1963,9 +2077,10 @@ class PretrainedModel(Layer, GenerationMixin, ConversionMixin):
 
                 if config.quantization_config.is_weight_quantize():
                     filter_dict_keys = None
-
                 state_dict = load_state_dict(
-                    shard_file, tp_actions if pre_tensor_parallel_split else None, filter_dict_keys
+                    shard_file,
+                    tp_actions if pre_tensor_parallel_split else None,
+                    filter_dict_keys,
                 )
 
                 # convert for fusing or splitting weights
@@ -2279,7 +2394,6 @@ class PretrainedModel(Layer, GenerationMixin, ConversionMixin):
             else:
                 raise ValueError(f"Unexpected file: {resolved_archive_file} for weight conversion.")
             # load pt weights early so that we know which dtype to init the model under
-
         if not is_sharded and state_dict is None:
             # 4. loading non-sharded ckpt from the state dict
             if config.tensor_parallel_degree > 1 and resolved_archive_file.endswith("model_state.pdparams"):

paddlenlp/transformers/tensor_parallel_utils.py

@@ -244,22 +244,32 @@ class FusedHeadAndCrossEntropy(PyLayer):
                 grad_hidden_states = None
 
         # initialize outputs
-        token_loss = paddle.empty((n_tokens,), dtype=hidden_states.dtype)
+        token_loss = paddle.empty((n_tokens,), dtype=paddle.float32)
 
         # blockwise calculations
         for i in range(0, n_tokens, loop_chunk_size):
             token_start_idx = i
             token_end_idx = min(i + loop_chunk_size, n_tokens)
-            hidden_states_chunk = hidden_states[token_start_idx:token_end_idx]
-            labels_chunk = labels[token_start_idx:token_end_idx]
+            cur_chunk_range = paddle.arange(token_start_idx, token_end_idx)
+            hidden_states_chunk = paddle.gather(hidden_states, cur_chunk_range, axis=0)
+            labels_chunk = paddle.gather(labels, cur_chunk_range, axis=0)
+            loss_mask_chunk = paddle.gather(loss_mask, cur_chunk_range, axis=0)
 
             # logits calculations
-            logits_chunk_cast = paddle.matmul(hidden_states_chunk, lm_head_weight_cast, transpose_y=transpose_y)
+            logits_chunk_cast = paddle.matmul(
+                hidden_states_chunk,
+                lm_head_weight_cast,
+                transpose_y=transpose_y,
+            )
             if lm_head_bias is not None:
                 logits_chunk_cast += lm_head_bias_cast
             if tensor_parallel_degree > 1 and not tensor_parallel_output:
                 logits_chunk_cast_lst = []
-                dist.all_gather(logits_chunk_cast_lst, logits_chunk_cast, group=model_parallel_group)
+                dist.all_gather(
+                    logits_chunk_cast_lst,
+                    logits_chunk_cast,
+                    group=model_parallel_group,
+                )
                 logits_chunk_cast = paddle.concat(logits_chunk_cast_lst, axis=-1)
             logits_chunk = logits_chunk_cast.astype("float32")
 
@@ -271,22 +281,34 @@ class FusedHeadAndCrossEntropy(PyLayer):
             exp_logits = paddle.exp(normalized_logits)
             sum_exp_logits = paddle.sum(exp_logits, axis=-1, keepdim=True)
             if tensor_parallel_degree > 1 and tensor_parallel_output:
-                dist.all_reduce(sum_exp_logits, op=dist.ReduceOp.SUM, group=model_parallel_group)
+                dist.all_reduce(
+                    sum_exp_logits,
+                    op=dist.ReduceOp.SUM,
+                    group=model_parallel_group,
+                )
             log_sum_exp_logits = paddle.log(sum_exp_logits)
 
             # cross entropy
             labels_one_hot = labels_chunk.unsqueeze(1) == indices
             label_logits = paddle.sum(
-                paddle.where(labels_one_hot, normalized_logits, paddle.zeros_like(normalized_logits)),
+                paddle.where(
+                    labels_one_hot,
+                    normalized_logits,
+                    paddle.zeros_like(normalized_logits),
+                ),
                 axis=-1,
                 keepdim=True,
             )
             if tensor_parallel_degree > 1 and tensor_parallel_output:
-                dist.all_reduce(label_logits, op=dist.ReduceOp.SUM, group=model_parallel_group)
+                dist.all_reduce(
+                    label_logits,
+                    op=dist.ReduceOp.SUM,
+                    group=model_parallel_group,
+                )
             token_loss_chunk = (log_sum_exp_logits - label_logits).squeeze(1) / divisor
-            cond = loss_mask[token_start_idx:token_end_idx].astype("bool")
+            cond = loss_mask_chunk.astype("bool")
             token_loss_chunk = paddle.where(cond, token_loss_chunk, paddle.zeros_like(token_loss_chunk))
-            token_loss[token_start_idx:token_end_idx] = token_loss_chunk * loss_mask[token_start_idx:token_end_idx]
+            paddle.scatter_(token_loss, cur_chunk_range, token_loss_chunk, overwrite=True)
 
             # gradients calculations
             if not return_token_loss:
@@ -298,18 +320,31 @@ class FusedHeadAndCrossEntropy(PyLayer):
                 grad_logits_chunk = (exp_logits / sum_exp_logits - labels_one_hot.astype("float32")) / divisor
                 grad_logits_chunk = grad_logits_chunk.astype(dtype)
                 grad_logits_chunk = paddle.where(
-                    cond.unsqueeze(1), grad_logits_chunk, paddle.zeros_like(grad_logits_chunk)
+                    cond.unsqueeze(1),
+                    grad_logits_chunk,
+                    paddle.zeros_like(grad_logits_chunk),
                 )
 
                 if grad_hidden_states is not None:
-                    grad_hidden_states[token_start_idx:token_end_idx] = paddle.matmul(
-                        grad_logits_chunk, lm_head_weight_cast, transpose_y=not transpose_y
+                    paddle.scatter_(
+                        grad_hidden_states,
+                        cur_chunk_range,
+                        paddle.matmul(grad_logits_chunk, lm_head_weight_cast, transpose_y=not transpose_y),
+                        overwrite=True,
                     )
                 if grad_lm_head_weight is not None:
                     if transpose_y:
-                        grad_lm_head_weight += paddle.matmul(grad_logits_chunk, hidden_states_chunk, transpose_x=True)
+                        grad_lm_head_weight += paddle.matmul(
+                            grad_logits_chunk,
+                            hidden_states_chunk,
+                            transpose_x=True,
+                        )
                     else:
-                        grad_lm_head_weight += paddle.matmul(hidden_states_chunk, grad_logits_chunk, transpose_x=True)
+                        grad_lm_head_weight += paddle.matmul(
+                            hidden_states_chunk,
+                            grad_logits_chunk,
+                            transpose_x=True,
+                        )
                 if grad_lm_head_bias is not None:
                     grad_lm_head_bias += grad_logits_chunk.astype("float32").sum(axis=0).astype(dtype)
 
@@ -340,7 +375,11 @@ class FusedHeadAndCrossEntropy(PyLayer):
             grad_args = []
             if ctx.hidden_states_has_grad:
                 if tensor_parallel_degree > 1:
-                    dist.all_reduce(grad_hidden_states, op=dist.ReduceOp.SUM, group=model_parallel_group)
+                    dist.all_reduce(
+                        grad_hidden_states,
+                        op=dist.ReduceOp.SUM,
+                        group=model_parallel_group,
+                    )
                 grad_args.append(grad_hidden_states.reshape(original_shape))
             if ctx.lm_head_weight_has_grad:
                 grad_args.append(grad_lm_head_weight)
@@ -376,9 +415,20 @@ class FusedHeadAndCrossEntropy(PyLayer):
                 grad_lm_head_bias = None
 
             if ctx.aux_num == 1:
-                return grad_hidden_states, grad_lm_head_weight, grad_lm_head_bias, None
+                return (
+                    grad_hidden_states,
+                    grad_lm_head_weight,
+                    grad_lm_head_bias,
+                    None,
+                )
             else:
-                return grad_hidden_states, grad_lm_head_weight, grad_lm_head_bias, None, None
+                return (
+                    grad_hidden_states,
+                    grad_lm_head_weight,
+                    grad_lm_head_bias,
+                    None,
+                    None,
+                )
 
         # return_token_loss = True
         grad_token_loss = grad_output.reshape([-1])
@@ -440,16 +490,26 @@ class FusedHeadAndCrossEntropy(PyLayer):
         for i in range(0, n_tokens, loop_chunk_size):
             token_start_idx = i
             token_end_idx = min(i + loop_chunk_size, n_tokens)
-            hidden_states_chunk = hidden_states[token_start_idx:token_end_idx]
-            labels_chunk = labels[token_start_idx:token_end_idx]
+            cur_chunk_range = paddle.arange(token_start_idx, token_end_idx)
+            hidden_states_chunk = paddle.gather(hidden_states, cur_chunk_range, axis=0)
+            labels_chunk = paddle.gather(labels, cur_chunk_range, axis=0)
+            loss_mask_chunk = paddle.gather(loss_mask, cur_chunk_range, axis=0)
 
             # logits calculations
-            logits_chunk_cast = paddle.matmul(hidden_states_chunk, lm_head_weight_cast, transpose_y=transpose_y)
+            logits_chunk_cast = paddle.matmul(
+                hidden_states_chunk,
+                lm_head_weight_cast,
+                transpose_y=transpose_y,
+            )
             if lm_head_bias is not None:
                 logits_chunk_cast += lm_head_bias_cast
             if tensor_parallel_degree > 1 and not tensor_parallel_output:
                 logits_chunk_cast_lst = []
-                dist.all_gather(logits_chunk_cast_lst, logits_chunk_cast, group=model_parallel_group)
+                dist.all_gather(
+                    logits_chunk_cast_lst,
+                    logits_chunk_cast,
+                    group=model_parallel_group,
+                )
                 logits_chunk_cast = paddle.concat(logits_chunk_cast_lst, axis=-1)
             logits_chunk = logits_chunk_cast.astype("float32")
 
@@ -461,7 +521,11 @@ class FusedHeadAndCrossEntropy(PyLayer):
             exp_logits = paddle.exp(normalized_logits)
             sum_exp_logits = paddle.sum(exp_logits, axis=-1, keepdim=True)
             if tensor_parallel_degree > 1 and tensor_parallel_output:
-                dist.all_reduce(sum_exp_logits, op=dist.ReduceOp.SUM, group=model_parallel_group)
+                dist.all_reduce(
+                    sum_exp_logits,
+                    op=dist.ReduceOp.SUM,
+                    group=model_parallel_group,
+                )
 
             labels_one_hot = labels_chunk.unsqueeze(1) == indices
             if tensor_parallel_degree > 1 and not tensor_parallel_output:
@@ -469,16 +533,21 @@ class FusedHeadAndCrossEntropy(PyLayer):
                 labels_one_hot = labels_one_hot.split(model_parallel_group.nranks, axis=-1)[model_parallel_group.rank]
             grad_logits_chunk = exp_logits / sum_exp_logits - labels_one_hot.astype("float32")
             # NOTE(hehuang): scaling grad_logits_chunk by grad_token_loss
-            grad_logits_chunk *= grad_token_loss[token_start_idx:token_end_idx].unsqueeze(1)
+            grad_logits_chunk *= paddle.gather(grad_token_loss, cur_chunk_range, axis=0).unsqueeze(1)
             grad_logits_chunk = grad_logits_chunk.astype(dtype)
-            cond = loss_mask[token_start_idx:token_end_idx].astype("bool")
+            cond = loss_mask_chunk.astype("bool")
             grad_logits_chunk = paddle.where(
-                cond.unsqueeze(1), grad_logits_chunk, paddle.zeros_like(grad_logits_chunk)
+                cond.unsqueeze(1),
+                grad_logits_chunk,
+                paddle.zeros_like(grad_logits_chunk),
             )
 
             if grad_hidden_states is not None:
-                grad_hidden_states[token_start_idx:token_end_idx] = paddle.matmul(
-                    grad_logits_chunk, lm_head_weight_cast, transpose_y=not transpose_y
+                paddle.scatter_(
+                    grad_hidden_states,
+                    cur_chunk_range,
+                    paddle.matmul(grad_logits_chunk, lm_head_weight_cast, transpose_y=not transpose_y),
+                    overwrite=True,
                 )
             if grad_lm_head_weight is not None:
                 if transpose_y:
@@ -490,10 +559,25 @@ class FusedHeadAndCrossEntropy(PyLayer):
 
         if grad_hidden_states is not None:
             if tensor_parallel_degree > 1:
-                dist.all_reduce(grad_hidden_states, op=dist.ReduceOp.SUM, group=model_parallel_group)
+                dist.all_reduce(
+                    grad_hidden_states,
+                    op=dist.ReduceOp.SUM,
+                    group=model_parallel_group,
+                )
             grad_hidden_states = grad_hidden_states.reshape(ctx.original_shape)
 
         if ctx.aux_num == 1:
-            return grad_hidden_states, grad_lm_head_weight, grad_lm_head_bias, None
+            return (
+                grad_hidden_states,
+                grad_lm_head_weight,
+                grad_lm_head_bias,
+                None,
+            )
         else:
-            return grad_hidden_states, grad_lm_head_weight, grad_lm_head_bias, None, None
+            return (
+                grad_hidden_states,
+                grad_lm_head_weight,
+                grad_lm_head_bias,
+                None,
+                None,
+            )

paddlenlp/trl/__init__.py

@@ -14,5 +14,8 @@
 
 from .dpo_criterion import DPOCriterion
 from .dpo_trainer import DPOTrainer
+from .embedding_trainer import EmbeddingTrainer
+from .kto_criterion import KTOCriterion
+from .kto_trainer import KTOTrainer
 from .trl_data import *
 from .trl_utils import *

paddlenlp/trl/dpo_criterion.py

@@ -287,10 +287,10 @@ class DPOCriterion(nn.Layer):
         )
         loss = dpo_loss + sft_loss
         if self.use_infohub:
-            infohub.policy_chosen_logps.append(policy_chosen_logps)
-            infohub.policy_rejected_logps.append(policy_rejected_logps)
-            infohub.sft_loss.append(sft_loss)
-            infohub.dpo_loss.append(dpo_loss)
+            infohub.policy_chosen_logps.append(policy_chosen_logps.detach())
+            infohub.policy_rejected_logps.append(policy_rejected_logps.detach())
+            infohub.sft_loss.append(sft_loss.detach())
+            infohub.dpo_loss.append(dpo_loss.detach())
             return loss
         else:
             return policy_chosen_logps, policy_rejected_logps, sft_loss, dpo_loss, loss

paddlenlp/trl/dpo_trainer.py

@@ -279,6 +279,7 @@ class DPOTrainer(Trainer):
         """
         prediction_step function for pipeline parallel mode.
         """
+        model._p2p_helper.clear_meta_cache()
         concatenated_inputs = {}
         # consider no drop last
         per_device_train_batch_size = self.args.per_device_train_batch_size
@@ -349,6 +350,7 @@ class DPOTrainer(Trainer):
         )
         self.log_metric(**metric_inputs)
         self.reset_dpo_infohub()
+        model._p2p_helper.clear_meta_cache()
         return (loss, None, None)
 
     def log_metric(

paddlenlp/trl/embedding_trainer.py

@@ -0,0 +1,184 @@
+# Copyright (c) 2024 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 contextlib import nullcontext
+
+import paddle
+try:
+    from paddle.base import core
+except:
+    core = None
+from paddle.distributed import fleet
+
+from paddlenlp.trainer import Trainer
+from paddlenlp.transformers.contrastive_loss import (
+    MatryoshkaContrastiveLoss,
+    SimpleContrastiveLoss,
+)
+from paddlenlp.transformers.embedding_utils import dist_gather_tensor_with_gradient
+
+__all__ = ["EmbeddingTrainer"]
+
+
+class EmbeddingTrainer(Trainer):
+    def __init__(self, model_args, **kwargs):
+        super().__init__(**kwargs)
+
+        self.model_args = model_args
+        self.embedding_negatives_cross_device = model_args.embedding_negatives_cross_device
+        self.accum_data = []
+        self.accum_freq = 0
+        self.accum_q_features = []
+        self.accum_p_features = []
+        self.accum_rng_states = {}
+        self.accum_rng_states["cpu"] = []
+        self.accum_rng_states["cuda"] = []
+        self.accum_rng_states["hybrid"] = []
+
+        if model_args.embedding_matryoshka_dims is not None and len(model_args.embedding_matryoshka_dims) > 0:
+            self.loss_fn = MatryoshkaContrastiveLoss(
+                model_args.embedding_temperature, model_args.embedding_matryoshka_dims
+            )
+        else:
+            self.loss_fn = SimpleContrastiveLoss(model_args.embedding_temperature)
+
+    def clear_memory(self):
+        self.accum_q_features.clear()
+        self.accum_p_features.clear()
+        paddle.device.cuda.empty_cache()
+
+    def clear_state(self):
+        self.accum_data.clear()
+        self.accum_rng_states["cpu"].clear()
+        self.accum_rng_states["cuda"].clear()
+        self.accum_rng_states["hybrid"].clear()
+        self.accum_freq = 0
+
+    @paddle.no_grad()
+    def forward_no_grad(self, model, inputs):
+        # Step1: graph-less forward
+        self.accum_data.append(inputs)
+        inputs = self._prepare_inputs(inputs)
+        with self.autocast_smart_context_manager():
+            # collect rand states
+            self.accum_rng_states["cpu"].append(paddle.framework.core.default_cpu_generator().get_state())
+            self.accum_rng_states["cuda"].append(paddle.get_rng_state())
+            if self.args.use_hybrid_parallel:
+                self.accum_rng_states["hybrid"].append(
+                    fleet.meta_parallel.get_rng_state_tracker().get_states_tracker()
+                )
+
+            query_reps, passage_reps = model(**inputs, return_encode=True)
+
+            if self.embedding_negatives_cross_device:
+                query_reps = dist_gather_tensor_with_gradient(query_reps)
+                passage_reps = dist_gather_tensor_with_gradient(passage_reps)
+
+            self.accum_q_features.append(query_reps)
+            self.accum_p_features.append(passage_reps)
+
+        self.accum_freq += 1
+
+    def get_current_rng_state(self):
+        return {
+            "cpu": [paddle.framework.core.default_cpu_generator().get_state()],
+            "cuda": [paddle.get_rng_state()],
+            "hybrid": [fleet.meta_parallel.get_rng_state_tracker().get_states_tracker()]
+            if self.args.use_hybrid_parallel
+            else [],
+        }
+
+    def reset_rng_state(self, states, index=0):
+        # set random states
+        if len(states) != 3:
+            raise ValueError("The length of state should be 3")
+        cpu_state = states["cpu"][index]
+        cuda_state = states["cuda"][index]
+        paddle.framework.core.default_cpu_generator().set_state(cpu_state)
+        # TODO(daisiming): support xpu and other custom devices.
+        if core.is_compiled_with_cuda():
+            for j in range(core.get_cuda_device_count()):
+                core.default_cuda_generator(j).set_state(cuda_state[j])
+        if self.args.use_hybrid_parallel:
+            hybrid_state = states["hybrid"][index]
+            fleet.meta_parallel.get_rng_state_tracker().set_states_tracker(hybrid_state)
+
+    def accum_forward_backward(self, model):
+        # Step2: representation gradient computation and caching
+        for i in range(len(self.accum_q_features)):
+            self.accum_q_features[i].stop_gradient = False
+        q_reps = paddle.concat(self.accum_q_features, axis=0)
+        for i in range(len(self.accum_p_features)):
+            self.accum_p_features[i].stop_gradient = False
+        p_reps = paddle.concat(self.accum_p_features, axis=0)
+
+        loss = self.loss_fn(q_reps, p_reps)
+        if self.do_grad_scaling:
+            self.scaler.scale(loss).backward()
+        else:
+            loss.backward()
+        # get represetation gradient cache
+        accum_q_grads = [q.grad for q in self.accum_q_features]
+        accum_p_grads = [p.grad for p in self.accum_p_features]
+        del q_reps, p_reps
+
+        # clear trash memory
+        self.clear_memory()
+
+        current_rng_state = self.get_current_rng_state()
+        # Step3: sub-batch gradient accumulation
+        for i in range(self.accum_freq):
+            inputs = self.accum_data[i]
+            inputs = self._prepare_inputs(inputs)
+
+            sync_context = model.no_sync() if i != self.accum_freq - 1 and hasattr(model, "no_sync") else nullcontext()
+            with sync_context:
+                self.reset_rng_state(self.accum_rng_states, index=i)
+
+                with self.autocast_smart_context_manager():
+                    query_reps, passage_reps = model(**inputs, return_encode=True)
+
+                if self.embedding_negatives_cross_device:
+                    query_reps = dist_gather_tensor_with_gradient(query_reps)
+                    passage_reps = dist_gather_tensor_with_gradient(passage_reps)
+
+                _loss = paddle.dot(query_reps.flatten(), accum_q_grads[i].flatten()) + paddle.dot(
+                    passage_reps.flatten(), accum_p_grads[i].flatten()
+                )
+                _loss.backward()
+
+        self.reset_rng_state(current_rng_state)
+        self.clear_state()
+        return loss.detach()
+
+    def training_step(
+        self,
+        model,
+        inputs,
+        step_control=0,
+    ):
+        if self.args.pipeline_parallel_degree > 1:
+            raise NotImplementedError("Cannot support pipeline parallel for Embedding training now.")
+
+        if self.args.gradient_accumulation_steps == 1:
+            return super().training_step(model, inputs)
+        else:
+            self.forward_no_grad(model, inputs)
+
+            # if (step_control + 1) % self.args.gradient_accumulation_steps is not zero, move on to next batch.
+            if (step_control + 1) % self.args.gradient_accumulation_steps != 0:
+                return 0.0
+
+            loss = self.accum_forward_backward(model)
+        return loss

paddlenlp/trl/kto_criterion.py

@@ -0,0 +1,262 @@
+# Copyright (c) 2024 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 copy
+import os
+
+import paddle
+import paddle.nn as nn
+import paddle.nn.functional as F
+from paddle.distributed import fleet
+from paddle.distributed.fleet.meta_parallel import ParallelCrossEntropy
+
+from paddlenlp.transformers import (
+    AllGatherVarlenOp,
+    fused_head_and_loss_fn,
+    parallel_linear,
+    parallel_matmul,
+    sequence_parallel_sparse_mask_labels,
+)
+from paddlenlp.utils import infohub
+
+
+class KTOCriterion(nn.Layer):
+    """KTO Criterion"""
+
+    def __init__(self, config, kto_config=None, ignore_label=0, use_infohub=False):
+        super(KTOCriterion, self).__init__()
+        self.config = config
+        if kto_config is None:
+            if getattr(self.config, "kto_config", None) is None:
+                raise ValueError("KTO Criterion requires model_config.kto_config.")
+            self.kto_config = copy.deepcopy(config.kto_config)
+        else:
+            self.kto_config = kto_config
+        if self.config.tensor_parallel_output and self.config.tensor_parallel_degree > 1:
+            self.logprobs = ParallelCrossEntropy()
+        else:
+            self.logprobs = nn.CrossEntropyLoss(reduction="none")
+        self.use_infohub = use_infohub
+        self.ignore_label = ignore_label
+        # allgather kl in criterion
+        topo = fleet.get_hybrid_communicate_group()._topo
+        parallel_groups = topo.get_comm_list("pipe")
+        ranks = []
+        for group in parallel_groups:
+            ranks.append(group[-1])
+        self.comm_group = paddle.distributed.new_group(ranks=ranks)
+
+    def _nested_gather(self, tensors):
+        """
+        Gather value of `tensors` (tensor or list/tuple of nested tensors) and convert them to numpy before
+        concatenating them to `gathered`
+        """
+        local_rank = -1
+        env_local_rank = int(os.environ.get("PADDLE_RANK_IN_NODE", -1))
+        if env_local_rank != -1 and env_local_rank != local_rank and paddle.distributed.get_world_size() > 1:
+            local_rank = env_local_rank
+        if tensors is None:
+            return
+        if local_rank != -1:
+            output_tensors = []
+            paddle.distributed.all_gather(
+                output_tensors, paddle.tile(tensors, repeat_times=[1, 1]), group=self.comm_group
+            )
+            tensors = paddle.concat(output_tensors, axis=0)
+        return tensors
+
+    def kto_logps(self, logits, response_labels, response_kl_labels, response_indexs):
+        """KTO logprobs"""
+        labels = response_labels + response_kl_labels
+        if self.config.use_fused_head_and_loss_fn:
+            hidden_states, weight, bias, transpose_y = logits
+        elif self.config.use_sparse_head_and_loss_fn:
+            hidden_states, weight, bias = logits
+        if self.config.use_sparse_head_and_loss_fn:
+            if self.config.tensor_parallel_degree > 1 and self.config.sequence_parallel:
+                labels, sparse_tgt_idx = sequence_parallel_sparse_mask_labels(labels, self.ignore_label)
+
+                hidden_states = paddle.take_along_axis(hidden_states, sparse_tgt_idx, axis=0)
+                hidden_states = AllGatherVarlenOp.apply(hidden_states)
+            else:
+                labels = labels.flatten()
+                sparse_tgt_idx = paddle.nonzero(labels != self.ignore_label).flatten()
+                labels = paddle.take_along_axis(labels, sparse_tgt_idx, axis=0)
+
+                hidden_states = hidden_states.reshape([-1, hidden_states.shape[-1]])
+                hidden_states = paddle.take_along_axis(hidden_states, sparse_tgt_idx.unsqueeze(-1), axis=0)
+        if self.config.use_fused_head_and_loss_fn:
+            per_token_logps = -fused_head_and_loss_fn(
+                hidden_states,
+                weight,
+                bias,
+                labels,
+                None,
+                transpose_y,
+                self.config.vocab_size,
+                self.config.tensor_parallel_degree,
+                self.config.tensor_parallel_output,
+                self.config.fused_linear,
+                getattr(self.config, "chunk_size", 1024),
+                return_token_loss=True,
+                ignore_index=self.ignore_label,
+            )
+        elif self.config.use_sparse_head_and_loss_fn:
+            if bias is None:
+                logits = parallel_matmul(hidden_states, weight, self.config.tensor_parallel_output)
+            else:
+                logits = parallel_linear(
+                    hidden_states,
+                    weight,
+                    bias,
+                    self.config.tensor_parallel_output,
+                )
+            logits = logits.astype("float32")
+            per_token_logps = -self.logprobs(logits, labels)
+        else:
+            logits = logits.astype("float32")
+            if logits.shape[:-1] != labels.shape:
+                raise ValueError("Logits (batch and sequence length dim) and labels must have the same shape.")
+            # bs, seq
+            per_token_logps = -self.logprobs(logits, labels.unsqueeze(2)).squeeze(2)
+
+        if len(response_indexs.shape) == 3:
+            response_indexs = response_indexs[0]
+        if self.config.use_sparse_head_and_loss_fn:
+            chosen_logps_list = [
+                (per_token_logps[response_index[1] : response_index[2]]).sum()
+                for response_index in response_indexs
+                if response_index[4] == 1
+            ]
+            rejected_logps_list = [
+                (per_token_logps[response_index[1] : response_index[2]]).sum()
+                for response_index in response_indexs
+                if response_index[4] == 0
+            ]
+            kl_logps_list = [
+                (per_token_logps[response_index[2] : response_index[3]]).sum() for response_index in response_indexs
+            ]
+        else:
+            chosen_logps_list = [
+                (per_token_logps[response_index[0]][response_index[1] : response_index[2]]).sum()
+                for response_index in response_indexs
+                if response_index[4] == 1
+            ]
+            rejected_logps_list = [
+                (per_token_logps[response_index[0]][response_index[1] : response_index[2]]).sum()
+                for response_index in response_indexs
+                if response_index[4] == 0
+            ]
+            kl_logps_list = [
+                (per_token_logps[response_index[0]][response_index[2] : response_index[3]]).sum()
+                for response_index in response_indexs
+            ]
+        if len(chosen_logps_list) == 0:
+            chosen_logps = paddle.zeros([0], dtype="float32")
+        else:
+            chosen_logps = paddle.stack(chosen_logps_list, axis=0)
+        if len(rejected_logps_list) == 0:
+            rejected_logps = paddle.zeros([0], dtype="float32")
+        else:
+            rejected_logps = paddle.stack(rejected_logps_list, axis=0)
+        kl_logps = paddle.stack(kl_logps_list, axis=0)
+        return chosen_logps, rejected_logps, kl_logps
+
+    def kto_loss(
+        self,
+        policy_chosen_logps,
+        policy_rejected_logps,
+        policy_kl_logps,
+        reference_chosen_logps,
+        reference_rejected_logps,
+        reference_kl_logps,
+    ):
+        """KTO Loss"""
+        kl = (policy_kl_logps - reference_kl_logps).mean().detach()
+        kl = self._nested_gather(paddle.tile(kl, repeat_times=[1, 1])).mean().clip(min=0)
+        if policy_chosen_logps.shape[0] == 0 or reference_chosen_logps.shape[0] == 0:
+            chosen_losses = paddle.zeros([0])
+        else:
+            chosen_logratios = policy_chosen_logps - reference_chosen_logps
+            chosen_losses = 1 - F.sigmoid(self.kto_config.beta * (chosen_logratios - kl))
+        if policy_rejected_logps.shape[0] == 0 or reference_rejected_logps.shape[0] == 0:
+            rejected_losses = paddle.zeros([0])
+        else:
+            rejected_logratios = policy_rejected_logps - reference_rejected_logps
+            rejected_losses = 1 - F.sigmoid(self.kto_config.beta * (kl - rejected_logratios))
+        losses = paddle.concat(
+            (
+                self.kto_config.desirable_weight * chosen_losses,
+                self.kto_config.undesirable_weight * rejected_losses,
+            ),
+            0,
+        )
+        return losses.mean(), kl
+
+    def forward(
+        self,
+        logits,
+        labels,
+    ):
+        """Forward"""
+        (
+            response_labels,
+            response_kl_labels,
+            response_indexs,
+            reference_chosen_logps,
+            reference_rejected_logps,
+            reference_kl_logps,
+        ) = labels
+        if reference_chosen_logps is None or reference_rejected_logps is None or reference_kl_logps is None:
+            (
+                reference_chosen_logps,
+                reference_rejected_logps,
+                reference_kl_logps,
+            ) = self.kto_logps(logits, response_labels, response_kl_labels, response_indexs)
+            if self.use_infohub:
+                infohub.reference_chosen_logps.append(reference_chosen_logps)
+                infohub.reference_rejected_logps.append(reference_rejected_logps)
+                infohub.reference_kl_logps.append(reference_kl_logps)
+                # pipeline mode requires return loss when self._compute_loss is True
+                return paddle.zeros([1])
+            else:
+                return (
+                    reference_chosen_logps,
+                    reference_rejected_logps,
+                    reference_kl_logps,
+                )
+        policy_chosen_logps, policy_rejected_logps, policy_kl_logps = self.kto_logps(
+            logits, response_labels, response_kl_labels, response_indexs
+        )
+        loss, kl = self.kto_loss(
+            policy_chosen_logps,
+            policy_rejected_logps,
+            policy_kl_logps,
+            reference_chosen_logps,
+            reference_rejected_logps,
+            reference_kl_logps,
+        )
+        if self.use_infohub:
+            infohub.policy_chosen_logps.append(policy_chosen_logps.detach())
+            infohub.policy_rejected_logps.append(policy_rejected_logps.detach())
+            infohub.policy_kl_logps.append(policy_kl_logps.detach())
+            infohub.kl.append(kl.detach())
+            return loss
+        else:
+            return (
+                policy_chosen_logps,
+                policy_rejected_logps,
+                policy_kl_logps,
+                loss,
+                kl,
+            )

paddlenlp/trl/kto_trainer.py

@@ -0,0 +1,555 @@
+# Copyright (c) 2024 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.
+
+""" KTO Trainer """
+from collections import OrderedDict, defaultdict
+
+import paddle
+from paddle.distributed import fleet
+
+from paddlenlp.trainer import Trainer
+from paddlenlp.transformers.model_utils import unwrap_model
+from paddlenlp.trl import KTOCriterion
+from paddlenlp.utils import infohub
+
+
+def disable_dropout_in_model(model: paddle.nn.Layer) -> None:
+    """ "disable dropout"""
+    for module in model.children():
+        if isinstance(module, paddle.nn.Dropout):
+            module.p = 0
+
+
+try:
+    from paddlenlp.peft.lora.lora_model import AVAILABLE_LAYERS
+except:
+    from paddlenlp.peft.lora.lora_model import AVALIABLE_LAYERS
+
+    AVAILABLE_LAYERS = AVALIABLE_LAYERS
+
+KTO_INFO_KEYS = [
+    "reference_chosen_logps",
+    "reference_rejected_logps",
+    "reference_kl_logps",
+    "policy_chosen_logps",
+    "policy_rejected_logps",
+    "policy_kl_logps",
+    "kl",
+]
+
+
+class KTOTrainer(Trainer):
+    """
+    Initialize KTOTrainer.
+    """
+
+    def __init__(
+        self,
+        model,
+        data_collator,
+        ref_model=None,
+        kto_config=None,
+        disable_dropout: bool = True,
+        padding_value: int = 0,
+        kto_criterion=None,
+        ignore_label: int = 0,
+        **kwargs,
+    ):
+        super().__init__(model, data_collator=data_collator, **kwargs)
+        if kto_config is None:
+            raise ValueError("kto_config is None")
+        else:
+            self.kto_config = kto_config
+        if ref_model:
+            self.ref_model = ref_model
+            self.ref_model_wrapped = self._wrap_ref_model(self.ref_model)
+            self.ref_model_wrapped.eval()
+        elif self.kto_config.lora:
+            self.ref_model = None
+            self.ref_model_wrapped = None
+        else:
+            raise ValueError("ref_model is None! KTO requires a reference model")
+        if not self.args.pipeline_parallel_degree > 1:
+            if kto_criterion is None:
+                self.kto_criterion = KTOCriterion(self.model.config, kto_config=kto_config, ignore_label=ignore_label)
+            elif isinstance(kto_criterion, KTOCriterion):
+                self.kto_criterion = kto_criterion
+            else:
+                raise ValueError(f"kto_criterion should be None or KTOCriterion. Got {type(kto_criterion)}")
+        if disable_dropout:
+            disable_dropout_in_model(model)
+            if self.ref_model is not None:
+                disable_dropout_in_model(self.ref_model)
+
+        self.padding_value = padding_value
+        self._stored_metrics = defaultdict(lambda: defaultdict(list))
+        if self.model.config.tensor_parallel_output and self.model.config.tensor_parallel_degree > 1:
+            self.logprobs = paddle.distributed.fleet.meta_parallel.ParallelCrossEntropy()
+        else:
+            self.logprobs = paddle.nn.CrossEntropyLoss(reduction="none")
+        self.reset_dpo_infohub()
+
+    def get_batch_metrics(self, ref_model, model, batch, train_eval="train"):
+        """Compute the KTO loss and other metrics for the given batch of inputs for train or test."""
+        inputs = {
+            "input_ids": batch["input_ids"],
+            "position_ids": batch["position_ids"],
+        }
+        if "attention_mask" in batch:
+            inputs["attention_mask"] = batch["attention_mask"]
+        elif "attn_mask_start_row_indices" in batch:
+            inputs["attn_mask_start_row_indices"] = batch["attn_mask_start_row_indices"]
+        else:
+            raise ValueError("No attention mask found in batch.")
+        labels = (
+            batch["response_labels"],
+            batch["response_kl_labels"],
+            batch["response_indexs"],
+            None,
+            None,
+            None,
+        )
+        with paddle.no_grad():
+            if self.kto_config.lora:
+                self.disable_lora(model)
+                model.eval()
+                logits = model(**inputs)
+                self.enable_lora(model)
+                model.train()
+            else:
+                logits = ref_model(**inputs)
+            (
+                reference_chosen_logps,
+                reference_rejected_logps,
+                reference_kl_logps,
+            ) = self.kto_criterion(logits, labels)
+        labels = labels[:3] + (
+            reference_chosen_logps,
+            reference_rejected_logps,
+            reference_kl_logps,
+        )
+        logits = model(**inputs)
+        (
+            policy_chosen_logps,
+            policy_rejected_logps,
+            policy_kl_logps,
+            loss,
+            kl,
+        ) = self.kto_criterion(logits, labels)
+
+        # metrics
+        metric_inputs = dict(
+            policy_chosen_logps=policy_chosen_logps,
+            policy_rejected_logps=policy_rejected_logps,
+            reference_chosen_logps=reference_chosen_logps,
+            reference_rejected_logps=reference_rejected_logps,
+            kl=kl,
+            train_eval=train_eval,
+        )
+        self.log_metric(**metric_inputs)
+        return loss
+
+    def log_metric(
+        self,
+        policy_chosen_logps,
+        policy_rejected_logps,
+        reference_chosen_logps,
+        reference_rejected_logps,
+        kl,
+        train_eval,
+    ):
+        metrics = {}
+        chosen_rewards = self.kto_config.beta * (policy_chosen_logps - reference_chosen_logps).detach()
+        rejected_rewards = self.kto_config.beta * (policy_rejected_logps - reference_rejected_logps).detach()
+
+        prefix = "eval_" if train_eval == "eval" else ""
+        metrics[f"{prefix}count/chosen"] = paddle.to_tensor(chosen_rewards.shape[0])
+        metrics[f"{prefix}count/rejected"] = paddle.to_tensor(rejected_rewards.shape[0])
+
+        if policy_chosen_logps.shape[0] == 0 or len(reference_chosen_logps.shape) == 0:
+            metrics[f"{prefix}rewards/chosen"] = paddle.zeros([])
+            metrics[f"{prefix}logps/chosen"] = paddle.zeros([])
+        else:
+            metrics[f"{prefix}rewards/chosen"] = chosen_rewards.mean()
+            metrics[f"{prefix}logps/chosen"] = policy_chosen_logps.mean()
+        if policy_rejected_logps.shape[0] == 0 or reference_rejected_logps.shape[0] == 0:
+            metrics[f"{prefix}rewards/rejected"] = paddle.zeros([])
+            metrics[f"{prefix}logps/rejected"] = paddle.zeros([])
+        else:
+            metrics[f"{prefix}rewards/rejected"] = rejected_rewards.mean()
+            metrics[f"{prefix}logps/rejected"] = policy_rejected_logps.mean()
+
+        for key in metrics:
+            if "count" in key:
+                metrics[key] = self._nested_gather(paddle.tile(metrics[key], repeat_times=[1, 1])).sum().cpu()
+                metrics[key] /= max(self.args.tensor_parallel_degree, 1)
+            else:
+                metrics[key] = self._nested_gather(paddle.tile(metrics[key], repeat_times=[1, 1])).mean().cpu()
+        metrics[f"{prefix}kl"] = kl
+        metrics[f"{prefix}rewards/margins"] = metrics[f"{prefix}rewards/chosen"] - metrics[f"{prefix}rewards/rejected"]
+        if self.args.should_save:
+            self.store_metrics(metrics, train_eval=train_eval)
+
+    def compute_loss(self, model, inputs):
+        """Compute the KTO loss for the given batch of inputs."""
+        loss = self.get_batch_metrics(self.ref_model_wrapped, model, inputs, train_eval="train")
+        return loss
+
+    def _wrap_ref_model(self, model):
+        """Wrap reference model."""
+        if unwrap_model(model) is not model:
+            return model
+        self.amp_dtype = "float16" if self.args.fp16 else "bfloat16"
+        model = paddle.amp.decorate(
+            models=model,
+            level=self.args.fp16_opt_level,
+            dtype=self.amp_dtype,
+        )
+        model = fleet.distributed_model(model)
+        if self.args.pipeline_parallel_degree > 1:
+            model._prepare_pipeline_inputs_func = prepare_pipeline_dpo_inputs_func
+        return model
+
+    def _wrap_model(self, model, training=True):
+        """Wrap model."""
+        model = super()._wrap_model(model, training)
+        if self.args.pipeline_parallel_degree > 1:
+            model._prepare_pipeline_inputs_func = prepare_pipeline_dpo_inputs_func
+        return model
+
+    def evaluate(self, eval_dataset=None, ignore_keys=None, metric_key_prefix="eval"):
+        """evaluate"""
+        self.model_wrapped = self._wrap_ref_model(self.model_wrapped)
+        return super().evaluate(eval_dataset, ignore_keys, metric_key_prefix)
+
+    def prediction_step(self, model, inputs, prediction_loss_only=False, ignore_keys=None):
+        """prediction_step"""
+        if self.args.pipeline_parallel_degree > 1:
+            # hack for pipeline mode
+            inputs = self._prepare_inputs(inputs)
+            return self.prediction_pipeline_step(self.ref_model_wrapped, model, inputs)
+        if ignore_keys is None:
+            if hasattr(model, "config"):
+                ignore_keys = getattr(model.config, "keys_to_ignore_at_inference", [])
+            else:
+                ignore_keys = []
+
+        with paddle.no_grad():
+            with self.autocast_smart_context_manager():
+                loss = self.get_batch_metrics(self.ref_model_wrapped, model, inputs, train_eval="eval")
+
+        if prediction_loss_only:
+            return (loss.detach(), None, None)
+        else:
+            raise NotImplementedError("KTOTrainer only supports prediction_loss_only=True for now.")
+
+    def store_metrics(self, metrics, train_eval="train"):
+        """store_metrics"""
+        for key, value in metrics.items():
+            self._stored_metrics[train_eval][key].append(value)
+
+    def log(self, logs, **kwargs):
+        """
+        Log `logs` on the various objects watching training, including stored metrics.
+
+        Args:
+            logs (`Dict[str, float]`):
+                The values to log.
+        """
+        # logs either has 'loss' or 'eval_loss'
+        train_eval = "train" if "loss" in logs else "eval"
+        # Add averaged stored metrics to logs
+        for key, metrics in self._stored_metrics[train_eval].items():
+            if "count" in key:
+                logs[key] = paddle.to_tensor(metrics).sum().item()
+            else:
+                logs[key] = paddle.to_tensor(metrics).mean().item()
+        del self._stored_metrics[train_eval]
+        if self.state.epoch is not None and train_eval == "train":
+            self.state.epoch *= self.args.num_train_epochs
+        return super().log(logs, **kwargs)
+
+    def disable_lora(self, model):
+        """Disable LORA layers."""
+        for _, layer in model.named_sublayers():
+            if any(isinstance(layer, lora_layer) for lora_layer in AVAILABLE_LAYERS):
+                layer.disable_lora = True
+
+    def enable_lora(self, model):
+        """Enable LORA layers."""
+        for _, layer in model.named_sublayers():
+            if any(isinstance(layer, lora_layer) for lora_layer in AVAILABLE_LAYERS):
+                layer.disable_lora = False
+
+    def training_pipeline_step(self, model, inputs):
+        """
+        Perform a training step on a batch of inputs.
+        """
+        # accumulation data
+        if not hasattr(self, "_pp_data_buffer"):
+            self._pp_data_buffer = []
+        self._pp_data_buffer.append(inputs)
+        if len(self._pp_data_buffer) != self.args.gradient_accumulation_steps:
+            return paddle.zeros([])
+
+        concatenated_inputs = {}
+        for key in self._pp_data_buffer[0].keys():
+            concatenated_inputs[key] = [
+                self._pp_data_buffer[i][key] for i in range(self.args.gradient_accumulation_steps)
+            ]
+        concatenated_inputs["reference_chosen_logps"] = None
+        concatenated_inputs["reference_rejected_logps"] = None
+        concatenated_inputs["reference_kl_logps"] = None
+        self._pp_data_buffer = []
+        inputs, labels = model._prepare_pipeline_inputs_func(concatenated_inputs)
+        model_config_backup = model.micro_batch_size, model.accumulate_steps
+        model.micro_batch_size = self.args.per_device_train_batch_size
+        model.accumulate_steps = self.args.gradient_accumulation_steps
+
+        if self.kto_config.lora:
+            self.disable_lora(model)
+            model.eval()
+            with paddle.no_grad():
+                with self.autocast_smart_context_manager():
+                    model.eval_batch(data=[inputs, labels], compute_loss=True)
+            self.enable_lora(model)
+            model._p2p_helper.clear_meta_cache()
+            model.train()
+        else:
+            ref_model = self.ref_model_wrapped
+            ref_model_config_backup = (
+                ref_model.micro_batch_size,
+                ref_model.accumulate_steps,
+            )
+            ref_model.accumulate_steps = model.accumulate_steps
+            ref_model.micro_batch_size = model.micro_batch_size
+            with paddle.no_grad():
+                with self.autocast_smart_context_manager():
+                    ref_model.eval_batch(data=[inputs, labels], compute_loss=True)
+            ref_model.micro_batch_size, ref_model.accumulate_steps = ref_model_config_backup
+        reference_chosen_logps = infohub.reference_chosen_logps
+        reference_rejected_logps = infohub.reference_rejected_logps
+        reference_kl_logps = infohub.reference_kl_logps
+
+        if model.is_pipeline_last_stage(ignore_virtual=model._layers._num_virtual_pipeline_stages > 1):
+            labels = labels[:3] + (
+                reference_chosen_logps,
+                reference_rejected_logps,
+                reference_kl_logps,
+            )
+        train_inputs = [inputs, labels]
+        train_inputs = model._prepare_training(train_inputs, self.optimizer, self.lr_scheduler)
+        model.optimizer = None  # we do not use `PipelineParallel` to handler optimizer step
+        model.lr_scheduler = None
+        with self.autocast_smart_context_manager():
+            loss = model.forward_backward_pipeline(train_inputs, self.scaler if self.do_grad_scaling else None)
+        model.micro_batch_size, model.accumulate_steps = model_config_backup
+
+        # broadcast KTO_INFO_KEYS
+        self.broadcast_last_stage_infohub_tensor()
+
+        # metrics
+        metric_inputs = dict(
+            policy_chosen_logps=infohub.policy_chosen_logps,
+            policy_rejected_logps=infohub.policy_rejected_logps,
+            reference_chosen_logps=infohub.reference_chosen_logps,
+            reference_rejected_logps=infohub.reference_rejected_logps,
+            kl=infohub.kl,
+            train_eval="train",
+        )
+        self.log_metric(**metric_inputs)
+        self.reset_dpo_infohub()
+        return loss.detach()
+
+    def prediction_pipeline_step(
+        self,
+        ref_model,
+        model,
+        batch,
+    ):
+        """
+        prediction_step function for pipeline parallel mode.
+        """
+        concatenated_inputs = {}
+        # consider no drop last
+        per_device_train_batch_size = self.args.per_device_train_batch_size
+        gradient_accumulation_steps = self.args.gradient_accumulation_steps
+        # preprocess inputs: tuple(List[Tensor])
+        for key in batch.keys():
+            if key not in "response_indexs":
+                concatenated_inputs[key] = [
+                    batch[key][i * per_device_train_batch_size : (i + 1) * per_device_train_batch_size]
+                    for i in range(gradient_accumulation_steps)
+                ]
+            else:
+                concatenated_inputs["response_indexs"] = [[] for _ in range(gradient_accumulation_steps)]
+                for i in range(gradient_accumulation_steps):
+                    for response_index in batch[key]:
+                        if response_index[0] in list(
+                            range(
+                                i * per_device_train_batch_size,
+                                (i + 1) * per_device_train_batch_size,
+                            )
+                        ):
+                            response_index[0] -= i * per_device_train_batch_size
+                            concatenated_inputs["response_indexs"][i].append(response_index)
+                    concatenated_inputs["response_indexs"][i] = paddle.stack(concatenated_inputs["response_indexs"][i])
+                    if model._layers.config.use_sparse_head_and_loss_fn:
+                        last_batch_response_length = concatenated_inputs["response_indexs"][i][0, 1]
+                        concatenated_inputs["response_indexs"][i][:, 1:] -= last_batch_response_length
+
+        concatenated_inputs["reference_chosen_logps"] = None
+        concatenated_inputs["reference_rejected_logps"] = None
+        concatenated_inputs["reference_kl_logps"] = None
+
+        self._pp_data_buffer = []
+        inputs, labels = model._prepare_pipeline_inputs_func(concatenated_inputs)
+
+        if self.kto_config.lora:
+            self.disable_lora(model)
+            model.eval()
+            with paddle.no_grad():
+                with self.autocast_smart_context_manager():
+                    model.eval_batch(data=[inputs, labels], compute_loss=True)
+            self.enable_lora(model)
+            model._p2p_helper.clear_meta_cache()
+            model.train()
+        else:
+            ref_model = self.ref_model_wrapped
+            with paddle.no_grad():
+                with self.autocast_smart_context_manager():
+                    ref_model.eval_batch(data=[inputs, labels], compute_loss=True)
+        reference_chosen_logps = infohub.reference_chosen_logps
+        reference_rejected_logps = infohub.reference_rejected_logps
+        reference_kl_logps = infohub.reference_kl_logps
+
+        if model.is_pipeline_last_stage(ignore_virtual=model._layers._num_virtual_pipeline_stages > 1):
+            labels = labels[:3] + (
+                reference_chosen_logps,
+                reference_rejected_logps,
+                reference_kl_logps,
+            )
+        with paddle.no_grad():
+            with self.autocast_smart_context_manager():
+                loss = model.eval_batch(data=[inputs, labels], compute_loss=True)
+
+        # broadcast KTO_INFO_KEYS
+        self.broadcast_last_stage_infohub_tensor()
+        # metrics
+        metric_inputs = dict(
+            policy_chosen_logps=infohub.policy_chosen_logps,
+            policy_rejected_logps=infohub.policy_rejected_logps,
+            reference_chosen_logps=infohub.reference_chosen_logps,
+            reference_rejected_logps=infohub.reference_rejected_logps,
+            kl=infohub.kl,
+            train_eval="eval",
+        )
+        self.log_metric(**metric_inputs)
+        self.reset_dpo_infohub()
+        return (loss, None, None)
+
+    def reset_dpo_infohub(self):
+        """Initialize infohub"""
+        for key in KTO_INFO_KEYS:
+            setattr(infohub, key, [])
+
+    def broadcast_last_stage_infohub_tensor(self):
+        for key in KTO_INFO_KEYS:
+            if self.model_wrapped.is_pipeline_last_stage(
+                ignore_virtual=self.model_wrapped._layers._num_virtual_pipeline_stages > 1
+            ):
+                if key == "kl":
+                    tensor = paddle.stack(getattr(infohub, key)).mean().detach()
+                elif "logps" in key:
+                    logps_list = getattr(infohub, key)
+                    if all(logps.shape == [0] for logps in logps_list):
+                        tensor = paddle.zeros([1])
+                    else:
+                        tensor = paddle.concat(getattr(infohub, key), axis=0).detach()
+                    tensor_shape = paddle.to_tensor(tensor.shape, dtype="int64")
+                    paddle.distributed.broadcast(
+                        tensor_shape,
+                        src=self.model_wrapped.global_rank,
+                        group=self.model_wrapped.pp_group,
+                    )
+                else:
+                    raise ValueError(f"Invalid key: {key}")
+                paddle.distributed.broadcast(
+                    tensor,
+                    src=self.model_wrapped.global_rank,
+                    group=self.model_wrapped.pp_group,
+                )
+            else:
+                if key == "kl":
+                    tensor = paddle.zeros([], "float32")
+                elif "logps" in key:
+                    tensor_shape = paddle.empty([1], dtype="int64")
+                    paddle.distributed.broadcast(
+                        tensor_shape,
+                        src=self.model_wrapped._hcg.get_rank_from_stage(self.model_wrapped.num_stages - 1),
+                        group=self.model_wrapped.pp_group,
+                    )
+                    tensor = paddle.zeros(tensor_shape, "float32")
+                else:
+                    raise ValueError(f"Invalid key: {key}")
+                paddle.distributed.broadcast(
+                    tensor,
+                    src=self.model_wrapped._hcg.get_rank_from_stage(self.model_wrapped.num_stages - 1),
+                    group=self.model_wrapped.pp_group,
+                )
+            setattr(infohub, key, tensor)
+
+
+def prepare_pipeline_dpo_inputs_func(inputs):
+    """Prepare pipeline inputs"""
+    if "attention_mask" in inputs:
+        first_stage_keys = [
+            "input_ids",
+            "attention_mask",
+            "position_ids",
+        ]
+    else:
+        first_stage_keys = [
+            "input_ids",
+            "attn_mask_start_row_indices",
+            "position_ids",
+        ]
+
+    last_stage_keys = [
+        "response_labels",
+        "response_kl_labels",
+        "response_indexs",
+        "reference_chosen_logps",
+        "reference_rejected_logps",
+        "reference_kl_logps",
+    ]
+
+    def get_expected_keys(inputs, keys):
+        ret = tuple([inputs.pop(k) for k in keys if k in inputs])
+        if len(ret) == 1:
+            ret = ret[0]
+        return ret
+
+    if type(inputs) is dict or type(inputs) is OrderedDict:
+        return [
+            get_expected_keys(inputs, first_stage_keys),
+            get_expected_keys(inputs, last_stage_keys),
+        ]
+
+    keys = list(inputs[0].keys())
+    inputs_batch = {key: [data.pop(key) for data in inputs] for key in keys}
+    return [
+        get_expected_keys(inputs_batch, first_stage_keys),
+        get_expected_keys(inputs_batch, last_stage_keys),
+    ]

paddlenlp/utils/env.py

@@ -75,7 +75,7 @@ GENERATION_CONFIG_NAME = "generation_config.json"
 # Fast tokenizers (provided by HuggingFace tokenizer's library) can be saved in a single file
 FULL_TOKENIZER_NAME = "tokenizer.json"
 
-
+MERGE_CONFIG_NAME = "merge_config.json"
 LORA_CONFIG_NAME = "lora_config.json"
 LORA_WEIGHTS_NAME = "lora_model_state.pdparams"
 
@@ -111,3 +111,19 @@ SAFE_MASTER_WEIGHTS_INDEX_NAME = "master_weights.safetensors.index.json"
 
 SAFE_PEFT_WEIGHTS_NAME = "peft_model.safetensors"
 SAFE_PEFT_WEIGHTS_INDEX_NAME = "peft_model.safetensors.index.json"
+
+# Checkpoint quantization
+MOMENT1_KEYNAME = "moment1_0"
+MOMENT2_KEYNAME = "moment2_0"
+BETA1_KEYNAME = "beta1_pow_acc_0"
+BETA2_KEYNAME = "beta2_pow_acc_0"
+SYMMETRY_QUANT_SCALE = "@scales"
+ASYMMETRY_QUANT_SCALE_MIN = "@min_scales"
+ASYMMETRY_QUANT_SCALE_MAX = "@max_scales"
+MAX_QUANTIZATION_TIMES = 1
+
+# LLM Inference related environment variables
+# Note(@Wanglongzhi2001): MAX_BSZ, SPECULATE_MAX_BSZ, MAX_DRAFT_TOKENS must be the same as definition in get_output / save_output
+MAX_BSZ = 512
+SPECULATE_MAX_BSZ = 256
+MAX_DRAFT_TOKENS = 6

paddlenlp/utils/nested.py

@@ -116,3 +116,13 @@ def nested_copy_place(inputs, place=None, blocking=False):
     if isinstance(inputs, paddle.Tensor):
         inputs = inputs if inputs.place == place else inputs._copy_to(place, blocking)
     return inputs
+
+
+def flatten_list(nested_list):
+    flattened_list = []
+    for item in nested_list:
+        if isinstance(item, list):
+            flattened_list.extend(flatten_list(item))
+        else:
+            flattened_list.append(item)
+    return flattened_list

requirements.txt

@@ -26,4 +26,4 @@ jinja2
 regex
 numpy<=1.26.4
 tiktoken
-tokenizers
+tokenizers<0.21

setup.py

@@ -112,7 +112,7 @@ def show():
 
 
 # only use this file to contral the version
-__version__ = "3.0.0b1.post"
+__version__ = "3.0.0b2.post"
 if os.getenv(PADDLENLP_STABLE_VERSION):
     __version__ = __version__.replace(".post", "")
 else:

tests/fixtures/llm/finetune.yaml

@@ -63,4 +63,53 @@ inference-infer:
     dtype: float16
     batch_size: 2
     decode_strategy: greedy_search
-    max_length: 20
+    max_length: 20
+
+ckpt_quant:
+  base:
+    dataset_name_or_path: "./data"
+    per_device_train_batch_size: 4
+    gradient_accumulation_steps: 4
+    per_device_eval_batch_size: 8
+    eval_accumulation_steps: 16
+    num_train_epochs: 3
+    learning_rate: 3e-05
+    warmup_steps: 30
+    logging_steps: 1
+    evaluation_strategy: "steps"
+    save_strategy: "steps"
+    save_steps: 1
+    max_steps: 1
+    src_length: 1024
+    max_length: 2048
+    fp16: true
+    fp16_opt_level: "O2"
+    ckpt_quant_stage: "O1"
+    do_train: true
+    do_eval: true
+    use_flash_attention: true
+    unified_checkpoint: true
+    unified_checkpoint_config: "async_save remove_master_weight"
+    disable_tqdm: true
+    load_best_model_at_end: true
+    eval_with_do_generation: false
+    metric_for_best_model: "accuracy"
+    recompute: true
+    tensor_parallel_degree: 2
+    pipeline_parallel_degree: 1
+
+  default:
+    llama:
+      model_name_or_path: __internal_testing__/tiny-random-llama
+    chatglm:
+      model_name_or_path: __internal_testing__/tiny-fused-chatglm
+    chatglm2:
+      model_name_or_path: __internal_testing__/tiny-fused-chatglm2
+    bloom:
+      model_name_or_path: __internal_testing__/tiny-fused-bloom
+    qwen:
+      model_name_or_path: __internal_testing__/tiny-fused-qwen
+    baichuan:
+      model_name_or_path: __internal_testing__/tiny-fused-baichuan
+    qwen2:
+      model_name_or_path: __internal_testing__/tiny-random-qwen2

tests/llm/test_finetune.py

@@ -13,11 +13,15 @@
 # limitations under the License.
 from __future__ import annotations
 
+import json
+import os
 import sys
 import unittest
 
 from parameterized import parameterized_class
 
+from paddlenlp.utils.env import SAFE_OPTIMIZER_INDEX_NAME
+from tests.parallel_launch import TestMultipleGpus
 from tests.testing_utils import argv_context_guard, load_test_config
 
 from .testing_utils import LLMTest
@@ -55,3 +59,46 @@ class FinetuneTest(LLMTest, unittest.TestCase):
             self.run_predictor({"inference_model": True})
 
         self.run_predictor({"inference_model": False})
+
+
+@parameterized_class(
+    ["model_dir"],
+    [
+        ["llama"],
+    ],
+)
+class CkptQuantTest(LLMTest, TestMultipleGpus):
+    config_path: str = "./tests/fixtures/llm/finetune.yaml"
+    model_dir: str = None
+
+    def setUp(self) -> None:
+        LLMTest.setUp(self)
+
+        sys.path.insert(0, self.model_dir)
+        self.run_sft = "llm/run_finetune.py"
+
+    def tearDown(self) -> None:
+        LLMTest.tearDown(self)
+
+    def test_ckpt_quant(self):
+        finetune_config = load_test_config(self.config_path, "ckpt_quant", self.model_dir)
+
+        finetune_config["dataset_name_or_path"] = self.data_dir
+        finetune_config["output_dir"] = self.output_dir
+
+        self.runfirst(finetune_config)
+
+        # get `quant_ckpt_resume_times`
+        with open(os.path.join(self.output_dir, "checkpoint-1", SAFE_OPTIMIZER_INDEX_NAME), "r") as r:
+            index = json.loads(r.read())
+        quant_ckpt_resume_times = index["quant_ckpt_resume_times"]
+
+        self.rerun(finetune_config)
+
+        self.assertEqual(quant_ckpt_resume_times, 0)
+
+    def runfirst(self, train_args):
+        self.run_n1c2(self.run_sft, **train_args)
+
+    def rerun(self, train_args):
+        self.run_n1c2(self.run_sft, **train_args)

tests/trainer/test_unified_checkpoint.py

@@ -18,7 +18,7 @@ import shutil
 import numpy as np
 import pytest
 
-from paddlenlp.trainer.plugins.unified_checkpoint import UnifiedCheckpointOption
+from paddlenlp.trainer.unified_checkpoint.utils import UnifiedCheckpointOption
 from tests.parallel_launch import TestMultipleGpus
 from tests.testing_utils import (
     require_paddle_at_least_2_gpu,