Cherry-pick bugfix PRs (#9089)

* quick fix from pretrained. (#8487)

* quick fix os.path.split (#8508)

* Cp/fix (#8569)

* [Safetensors] Fix fast safe open slice. (#8512)
* [FIX DDP] fix ddp (#8549)

* [BUG] Fix build train valid test datasets (#8823)

* Update causal_dataset.py

* Add twenty redundant data in post pretrain (#8777)

* 给dataset再添加20条数据,防止blend dataset出现错误

* num_samples向下去整,防止数据集的溢出 (#8691)

* update release_grads (#8834)

* update release_grads (#8834)

* [Trainer] Fix release_grads (#9085)

* fix pp release_grads

* add dataloader_drop_last to evaldataloader (#8773)

* bugfix

* Fix eval hang (#9052)

* fix pipeline eval

* fix eval dataloader_num_workers

---------

Co-authored-by: Zhong Hui <zhonghui.net@gmail.com>
Co-authored-by: yujun <50394665+JunnYu@users.noreply.github.com>
Co-authored-by: gongel <ainlp88@qq.com>

llm/finetune_generation.py

@@ -45,6 +45,7 @@ from paddlenlp.transformers import (
     AutoConfig,
     AutoModelForCausalLM,
     AutoTokenizer,
+    Llama3Tokenizer,
     LlamaTokenizer,
 )
 from paddlenlp.utils.log import logger
@@ -232,7 +233,7 @@ def main():
     if tokenizer.chat_template is not None:
         data_args.eval_with_do_generation = False
 
-    if isinstance(tokenizer, LlamaTokenizer):
+    if isinstance(tokenizer, LlamaTokenizer) or isinstance(tokenizer, Llama3Tokenizer):
         tokenizer.pad_token_id = tokenizer.eos_token_id
 
     if data_args.dataset_name_or_path is None:

model_zoo/gpt-3/ppfleetx/models/language_model/gpt/dygraph/hybrid_model.py

@@ -48,13 +48,16 @@ from .processor import (
     MinLengthLogitsProcessor,
     RepetitionPenaltyLogitsProcessor,
 )
-from paddle.distributed.fleet.utils.sequence_parallel_utils import (
-    ColumnSequenceParallelLinear,
-    GatherOp,
-    RowSequenceParallelLinear,
-    ScatterOp,
-    mark_as_sequence_parallel_parameter,
-)
+try:
+    from paddle.distributed.fleet.utils.sequence_parallel_utils import (
+        ColumnSequenceParallelLinear,
+        GatherOp,
+        RowSequenceParallelLinear,
+        ScatterOp,
+        mark_as_sequence_parallel_parameter,
+    )
+except:
+    pass
 
 from paddlenlp.transformers.segment_parallel_utils  import ReshardLayer
 

model_zoo/gpt-3/ppfleetx/models/language_model/language_module.py

@@ -24,9 +24,12 @@ from paddle.static import InputSpec
 from ppfleetx.core.module.basic_module import BasicModule
 from ppfleetx.data.tokenizers import GPTTokenizer
 from ppfleetx.distributed.apis import env
-from paddle.distributed.fleet.utils.sequence_parallel_utils import (
-    register_sequence_parallel_allreduce_hooks,
-)
+try:
+    from paddle.distributed.fleet.utils.sequence_parallel_utils import (
+        register_sequence_parallel_allreduce_hooks,
+    )
+except:
+    pass
 from ppfleetx.utils.log import logger
 
 # TODO(haohongxiang): to solve the problem of cross-reference

paddlenlp/__init__.py

@@ -18,7 +18,7 @@ import sys
 PADDLENLP_STABLE_VERSION = "PADDLENLP_STABLE_VERSION"
 
 
-__version__ = "2.7.1.post"
+__version__ = "2.8.0.post"
 if os.getenv(PADDLENLP_STABLE_VERSION):
     __version__ = __version__.replace(".post", "")
 

paddlenlp/data/causal_dataset.py

@@ -94,10 +94,11 @@ def get_datasets_weights_and_num_samples(data_prefix, train_val_test_num_samples
     # Add 0.5% (the 1.005 factor) so in case the bleding dataset does
     # not uniformly distribute the number of samples, we still have
     # samples left to feed to the network.
+    # (NOTE, yujun06): This is a workaround to avoid issues with indexing in the blending dataset. Therefore, we need to add 20 samples to each dataset.
     datasets_train_valid_test_num_samples = []
     for weight in weights:
         datasets_train_valid_test_num_samples.append(
-            [int(math.ceil(val * weight * 1.005)) for val in train_val_test_num_samples]
+            [int(math.ceil(val * weight * 1.005)) + 20 for val in train_val_test_num_samples]
         )
 
     return prefixes, weights, datasets_train_valid_test_num_samples
@@ -146,7 +147,9 @@ def build_train_valid_test_datasets(
     # Parse the values.
     output = get_datasets_weights_and_num_samples(data_prefix, train_val_test_num_samples)
     prefixes, weights, datasets_train_valid_test_num_samples = output
-    train_num_samples, valid_num_samples, test_num_samples = map(sum, zip(*datasets_train_valid_test_num_samples))
+    # NOTE: megatron/gpt_dataset.py has been updated. When creating BlendableDataset, we will use the raw train_val_test_num_samples instead of the expanded ones.
+    # Please refer to https://github.com/NVIDIA/NeMo/blob/72f630d087d45655b1a069dc72debf01dfdbdb2d/nemo/collections/nlp/data/language_modeling/megatron/gpt_dataset.py#L74-L80 for more information
+    train_num_samples, valid_num_samples, test_num_samples = train_val_test_num_samples
 
     # Build individual datasets.
     train_datasets = []

paddlenlp/data/dist_dataloader.py

@@ -12,13 +12,16 @@
 # 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.utils.log import logger
-
-_MAX_DATA_DIM = 64
+from paddlenlp.utils.nested import (
+    nested_broadcast_tensor,
+    nested_copy_place,
+    nested_empty_tensor,
+    nested_reduce_tensor,
+)
 
 
 class DummyDataset(paddle.io.Dataset):
@@ -53,6 +56,7 @@ class DistDataLoader(paddle.io.DataLoader):
         timeout=0,
         worker_init_fn=None,
         persistent_workers=False,
+        eval=False,
     ):
 
         if dataset is None:
@@ -62,12 +66,15 @@ class DistDataLoader(paddle.io.DataLoader):
         super().__init__(dataset=dataset, batch_sampler=batch_sampler, collate_fn=collate_fn, num_workers=num_workers)
 
         self._hcg = fleet.get_hybrid_communicate_group()
+        self.eval = eval
 
         # 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()
         self.mp_rank = self._hcg.get_model_parallel_rank()
@@ -78,10 +85,6 @@ class DistDataLoader(paddle.io.DataLoader):
         sharding_rank = self._hcg.get_sharding_parallel_rank()
         self._need_data = (self.mp_rank == 0) and (self.pp_rank == 0)
 
-        # When needed other data types, we can modify dtype_list.
-        self.dtype_list = [paddle.int64, paddle.float32, paddle.int32]
-        self._data_keys_list, self._data_keys_size = None, None
-
         if self._need_data:
             self._dataloader = paddle.io.DataLoader(
                 dataset,
@@ -127,7 +130,6 @@ class DistDataLoader(paddle.io.DataLoader):
         parallel_groups = topo.get_comm_list("pipe")
 
         for group in parallel_groups:
-            # only first rank and last rank
             ranks = [group[0], group[-1]]
             comm_group = paddle.distributed.new_group(ranks=ranks)
             if paddle.distributed.get_rank() in ranks:
@@ -137,127 +139,68 @@ class DistDataLoader(paddle.io.DataLoader):
     def __iter__(self):
         return self
 
-    def __next__(self):
-        data_keys_size = [0 for i in range(len(self.dtype_list))]
-        if self._need_data:
-            data = next(self._dataloader_iter)
-            data_keys = list(data.keys())
-
-            for key in data_keys:
-                if data[key].dtype not in self.dtype_list:
-                    raise ValueError(
-                        f"Dist dataloader requires dtype as `int64`, `float32` or `int32` currently, but got: {data[key].dtype}"
+    def _broadcast_data(self, data):
+        process_rank = paddle.distributed.get_rank()
+        if self.mp_group.nranks > 1:
+            if process_rank == self.mp_src_rank:
+                fake_data = [nested_reduce_tensor(data)]
+            else:
+                if data is not None:
+                    logger.warning(
+                        f"Your local rank {paddle.distributed.get_rank()} are forbidden to have a state_dict."
                     )
+                fake_data = [None]
+        if self._pp_group is not None:
+            if process_rank == self._pp_group.ranks[0]:
+                fake_data = [nested_reduce_tensor(data)]
+            else:
+                if data is not None:
+                    logger.warning(
+                        f"Your local rank {paddle.distributed.get_rank()} are forbidden to have a state_dict."
+                    )
+                fake_data = [None]
+        if self.mp_group.nranks > 1 and self.pp_rank == 0:
+            paddle.distributed.broadcast_object_list(
+                fake_data,
+                src=self.mp_src_rank,
+                group=self.mp_group,
+            )
+        if self._pp_group is not None:
+            paddle.distributed.broadcast_object_list(
+                fake_data,
+                src=self._pp_group.ranks[0],
+                group=self._pp_group,
+            )
 
-            data_list, data_keys_list = [], []
-            for i, dtype in enumerate(self.dtype_list):
-                data_list.append([data[key] for key in data_keys if data[key].dtype == dtype])
-                data_keys_list.append([key for key in data_keys if data[key].dtype == dtype])
-            data_keys_size = [len(keys) for keys in data_keys_list]
+        fake_data = fake_data[0]
+        if fake_data is None:
+            raise StopIteration
 
-        # Broadcast data keys size.
-        if self._data_keys_size is None:
-            if self.mp_group.nranks > 1 and self.pp_rank == 0:
-                paddle.distributed.broadcast_object_list(data_keys_size, src=self.mp_src_rank, group=self.mp_group)
-            if self._pp_data_group is not None:
-                paddle.distributed.broadcast_object_list(
-                    data_keys_size, src=self._pp_data_group.ranks[0], group=self._pp_data_group
-                )
-            self._data_keys_size = data_keys_size
-
-        if not self._need_data:
-            data_keys_list = [[None for i in range(keys_size)] for keys_size in self._data_keys_size]
-
-        # Broadcast data keys name.
-        if self._data_keys_list is None:
-            if self.mp_group.nranks > 1 and self.pp_rank == 0:
-                paddle.distributed.broadcast_object_list(data_keys_list, src=self.mp_src_rank, group=self.mp_group)
-            if self._pp_data_group is not None:
-                paddle.distributed.broadcast_object_list(
-                    data_keys_list, src=self._pp_data_group.ranks[0], group=self._pp_data_group
-                )
-            self._data_keys_list = data_keys_list
-
-        # Broadcast data.
-        if not self._need_data:
-            data_list = [[None for i in range(keys_size)] for keys_size in self._data_keys_size]
+        dst_pp_group = self._pp_group if self.eval else self._pp_data_group
+        if self.mp_group.nranks > 1:
+            if process_rank != self.mp_src_rank:
+                data = nested_empty_tensor(fake_data)
+        if dst_pp_group is not None:
+            if process_rank != dst_pp_group.ranks[0]:
+                data = nested_empty_tensor(fake_data)
 
         if self.mp_group.nranks > 1 and self.pp_rank == 0:
-            for i, dtype in enumerate(self.dtype_list):
-                if self._data_keys_size[i] > 0:
-                    data_list[i] = broadcast_data_list(
-                        data_list[i], dtype, self.mp_rank, self.mp_group, self.mp_src_rank
-                    )
+            data = nested_broadcast_tensor(data, src=self.mp_src_rank, group=self.mp_group)
+        if dst_pp_group is not None:
+            data = nested_broadcast_tensor(data, src=dst_pp_group.ranks[0], group=dst_pp_group)
+        # for pp1 - pp_{n-1}, Paddle need to recevie empty dict for pipeline parallel.
+        if data is None:
+            data = {}
 
-        if self._pp_data_group is not None:
-            # Note(daisimng): In last stage of pp, we don't need input_ids.
-            # It will be removed in future.
-            for i, dtype in enumerate(self.dtype_list):
-                if self._data_keys_size[i] > 0:
-                    data_list[i] = broadcast_data_list(
-                        data_list[i],
-                        dtype,
-                        self.pp_rank,
-                        self._pp_data_group,
-                        self._pp_data_group.ranks[0],
-                    )
+        return data
 
-        out_data = {}
-        for keys, datas in zip(self._data_keys_list, data_list):
-            out_data.update([(k, d) for k, d in zip(keys, datas)])
-
-        return out_data
-
-
-def broadcast_data_list(data_list, datatype, comm_rank=0, comm_group=None, src_rank=0):
-    """
-    Broadcast data from src_rank to all ranks in comm_group.
-    """
-    # Move to GPU and broadcast.
-    size_cpu = []
-    if comm_rank == 0:
-        for data in data_list:
-            size_cpu.append(len(data.shape))
-            size_cpu += data.shape
-    size_cpu = size_cpu + [0] * (_MAX_DATA_DIM - len(size_cpu))
-    size_cuda = paddle.to_tensor(size_cpu)
-    paddle.distributed.broadcast(size_cuda, src_rank, group=comm_group).wait()
-
-    size_cpu = size_cuda.tolist()
-    i = 0
-    numel = 0
-    sizes = []
-    while size_cpu[i] > 0:
-        rank = size_cpu[i]
-        this_size = size_cpu[i + 1 : i + 1 + rank]
-        numel += int(np.prod(this_size))
-        sizes.append(this_size)
-        i += rank + 1
-
-    if comm_rank == 0:
-        assert data.dtype == datatype, "input has data type {} which " "is different than {}".format(
-            data.dtype, datatype
-        )
-        if paddle.is_compiled_with_cuda():
-            data_b = paddle.concat([d.cuda().reshape([-1]) for d in data_list], 0)
-        else:
-            data_b = paddle.concat([d.reshape([-1]) for d in data_list], 0)
-
-        assert numel == sum([d.numel().item() for d in data_list]), (numel, [d.numel().item() for d in data_list])
-    else:
-        if paddle.is_compiled_with_cuda():
-            data_b = paddle.empty([numel], dtype=datatype).cuda()
-        else:
-            data_b = paddle.empty([numel], dtype=datatype)
-
-    # Broadcast
-    paddle.distributed.broadcast(data_b, src_rank, group=comm_group).wait()
-
-    ret = []
-    offset = 0
-    for size in sizes:
-        numel = int(np.prod(size))
-        ret.append(data_b[offset : offset + numel].reshape(size))
-        offset += numel
-
-    return ret
+    def __next__(self):
+        data = None
+        if self._need_data:
+            try:
+                data = next(self._dataloader_iter)
+                data = nested_copy_place(data, place=paddle.framework._current_expected_place())
+            except:
+                pass
+        data = self._broadcast_data(data)
+        return data

paddlenlp/peft/lora/lora_config.py

@@ -13,6 +13,7 @@
 # limitations under the License.
 
 import json
+import math
 import os
 from dataclasses import asdict, dataclass, field
 from typing import List, Optional, Union
@@ -94,6 +95,15 @@ class LoRAConfig:
             )
             self.use_quick_lora = False
 
+    @property
+    def scaling(self):
+        if not self.rslora and not self.pissa:
+            return self.lora_alpha / self.r
+        elif self.pissa:
+            return 1.0
+        else:
+            return self.lora_alpha / math.sqrt(self.r)
+
     @property
     def __dict__(self):
         return asdict(self)
@@ -114,6 +124,7 @@ class LoRAConfig:
         os.makedirs(save_directory, exist_ok=True)
 
         output_dict = self.__dict__
+        output_dict["scaling"] = self.scaling
         output_path = os.path.join(save_directory, LORA_CONFIG_NAME)
 
         # save it
@@ -136,7 +147,7 @@ class LoRAConfig:
             raise ValueError(f"Can't find lora_config.json at '{pretrained_model_name_or_path}'")
 
         loaded_attributes = cls.from_json_file(config_file)
-
+        loaded_attributes.pop("scaling", None)
         config = cls(**kwargs)
 
         for key, value in loaded_attributes.items():

paddlenlp/peft/lora/lora_layers.py

@@ -13,7 +13,6 @@
 # limitations under the License.
 
 import math
-import os
 from typing import List, Optional
 
 import paddle
@@ -25,13 +24,34 @@ from paddle.distributed.fleet.meta_parallel import (
     RowParallelLinear,
 )
 
-from .lora_quick_layers import quick_lora
+try:
+    from paddle.distributed.fleet.utils.sequence_parallel_utils import (
+        AllGatherOp,
+        ColumnSequenceParallelLinear,
+        ReduceScatterOp,
+        RowSequenceParallelLinear,
+        mark_as_sequence_parallel_parameter,
+    )
+except:
+    AllGatherOp = None
+    ReduceScatterOp = None
+    mark_as_sequence_parallel_parameter = None
 
-if "npu" in paddle.device.get_all_custom_device_type():
-    from .mc2_lora_npu import MC2LoRaColumnParallelLinear, MC2LoRaRowParallelLinear
-else:
-    MC2LoRaRowParallelLinear = None
-    MC2LoRaColumnParallelLinear = None
+    class ColumnSequenceParallelLinear:
+        pass
+
+    class RowSequenceParallelLinear:
+        pass
+
+
+from paddlenlp.transformers.mc2_parallel_linear import (
+    MC2ColumnParallelCoreLinear,
+    MC2ColumnSeqParallelCoreLinear,
+    MC2RowParallelCoreLinear,
+    MC2RowSeqParallelCoreLinear,
+)
+
+from .lora_quick_layers import quick_lora
 
 
 class LoRALinear(nn.Linear):
@@ -266,9 +286,7 @@ class RowParallelLoRALinear(RowParallelLinear):
             )
         else:
             # x @ W : [bz, in_f / ws] ===> [bz, out_f]
-            if "npu" in paddle.device.get_all_custom_device_type() and int(os.getenv("MC2", "0")):
-                output = MC2LoRaRowParallelLinear.apply(input_mp, self.weight, self.model_parallel_group)
-            else:
+            if MC2RowParallelCoreLinear is None:
                 result_mp = F.linear(x=input_mp, weight=self.weight, name=self.name)
                 output = mp_ops._mp_allreduce(
                     result_mp,
@@ -276,6 +294,8 @@ class RowParallelLoRALinear(RowParallelLinear):
                     use_calc_stream=True,
                     use_model_parallel=True,
                 )
+            else:
+                output = MC2RowParallelCoreLinear.apply(input_mp, self.weight, self.model_parallel_group)
 
             if not self.merged:
                 # x @ A: [bz, in_f/ ws] ===> [bz, r]
@@ -298,6 +318,120 @@ class RowParallelLoRALinear(RowParallelLinear):
         return f"in_features={self.weight.shape[0]}, out_features={self.weight.shape[1]}, rank={self.r}{name}"
 
 
+class RowSequenceParallelLoRALinear(RowSequenceParallelLinear):
+    def __init__(
+        self,
+        in_features: int,
+        out_features: int,
+        r: int = 0,
+        lora_alpha: int = 1,
+        lora_dropout: float = 0.0,
+        rslora: bool = False,
+        lora_plus_scale: float = 1.0,
+        merge_weights: bool = True,
+        use_quick_lora: bool = False,
+        **kwargs
+    ):
+        RowSequenceParallelLinear.__init__(self, in_features, out_features, **kwargs)
+        if not isinstance(r, int) or r <= 0:
+            raise ValueError("Lora rank r should be a positive integer")
+        self.r = r
+        self.lora_alpha = lora_alpha
+        # Optional dropout
+        if lora_dropout > 0.0:
+            self.lora_dropout = nn.Dropout(p=lora_dropout)
+        else:
+            self.lora_dropout = lambda x: x
+        # Mark the weight as unmerged
+        self.merged = False
+        self.merge_weights = merge_weights
+
+        # compatible
+        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")
+            ),
+        )
+        self.lora_B = self.create_parameter(
+            shape=[r, self.out_features],
+            dtype=self._dtype,
+            is_bias=False,
+            attr=paddle.ParamAttr(
+                initializer=paddle.nn.initializer.Constant(value=0.0),
+                learning_rate=lora_plus_scale,
+            ),
+        )
+
+        self.lora_A.is_distributed = True
+        self.lora_A.split_axis = 0
+        self.lora_B.is_distributed = False
+        mark_as_sequence_parallel_parameter(self.lora_B)
+        if not rslora:
+            self.scaling = self.lora_alpha / self.r
+        else:
+            self.scaling = self.lora_alpha / math.sqrt(self.r)
+
+        # Freezing the pre-trained weight matrix
+        self.weight.stop_gradient = True
+        self._use_quick_lora = use_quick_lora and lora_dropout == 0.0
+
+    @property
+    def use_quick_lora(self):
+        # TODO(@gexiao): support qlora
+        return False  # self._use_quick_lora and self.training and not self.merged
+
+    def train(self):
+        super().train()
+        if self.merge_weights and self.merged:
+            # Make sure that the weights are not merged
+            new_weight = self.weight - self.lora_A @ self.lora_B * self.scaling
+            self.weight.set_value(new_weight)
+            self.merged = False
+
+    def eval(self):
+        super().eval()
+        if self.merge_weights and not self.merged:
+            # Merge the weights and mark it
+            new_weight = self.weight + self.lora_A @ self.lora_B * self.scaling
+            self.weight.set_value(new_weight)
+            self.merged = True
+
+    def forward(self, x: paddle.Tensor):
+        if not self.input_is_parallel:
+            input_mp = mp_ops._c_split(x, group=self.model_parallel_group)
+        else:
+            input_mp = x
+
+        if MC2RowSeqParallelCoreLinear is None:
+            output_parallel = self.linear(input_mp, self.weight, name=self._name)
+            output_ = ReduceScatterOp.apply(output_parallel)
+            result_mp = output_ + self.bias if self.bias is not None else output_
+        else:
+            output_ = MC2RowSeqParallelCoreLinear.apply(input_mp, self.weight, self.model_parallel_group)
+            result_mp = output_ + self.bias if self.bias is not None else output_
+
+        if not self.merged:
+            input_mp = self.lora_dropout(input_mp)
+            if MC2RowSeqParallelCoreLinear is None:
+                input_mp = input_mp @ self.lora_A
+                input_mp = ReduceScatterOp.apply(input_mp)
+            else:
+                input_mp = MC2RowSeqParallelCoreLinear.apply(input_mp, self.lora_A, self.model_parallel_group)
+            delta_mp = (input_mp @ self.lora_B) * self.scaling
+            result_mp += delta_mp
+        return result_mp
+
+    def extra_repr(self):
+        name = f", name={self.name}" if self.name else ""
+        return f"in_features={self.weight.shape[0]}, out_features={self.weight.shape[1]}, rank={self.r}{name}"
+
+
 class ColumnParallelLoRALinear(ColumnParallelLinear):
     def __init__(
         self,
@@ -400,21 +534,21 @@ class ColumnParallelLoRALinear(ColumnParallelLinear):
                 world_size=self.world_size,
             )
         else:
-            if "npu" in paddle.device.get_all_custom_device_type() and int(os.getenv("MC2", "0")):
-                res_mp = MC2LoRaColumnParallelLinear.apply(input, self.weight, self.model_parallel_group)
-                result_mp = res_mp + self.bias
-            else:
+            if MC2ColumnParallelCoreLinear is None:
                 input_mp = mp_ops._c_identity(input, group=self.model_parallel_group)
                 result_mp = F.linear(x=input_mp, weight=self.weight, bias=self.bias, name=self.name)
+            else:
+                res_mp = MC2ColumnParallelCoreLinear.apply(input, self.weight, self.model_parallel_group)
+                result_mp = res_mp + self.bias
 
             if not self.merged:
                 input_a = self.lora_dropout(input) @ self.lora_A
-                if "npu" in paddle.device.get_all_custom_device_type() and int(os.getenv("MC2", "0")):
-                    tmp = MC2LoRaColumnParallelLinear.apply(input_a, self.lora_B, self.model_parallel_group)
-                    delta_mp = tmp * self.scaling
-                else:
+                if MC2ColumnParallelCoreLinear is None:
                     input_a_mp = mp_ops._c_identity(input_a, group=self.model_parallel_group)
                     delta_mp = (input_a_mp @ self.lora_B) * self.scaling
+                else:
+                    tmp = MC2ColumnParallelCoreLinear.apply(input_a, self.lora_B, self.model_parallel_group)
+                    delta_mp = tmp * self.scaling
                 result_mp += delta_mp
 
         if self.gather_output and self.is_mp:
@@ -428,6 +562,123 @@ class ColumnParallelLoRALinear(ColumnParallelLinear):
         return f"in_features={self.weight.shape[0]}, out_features={self.weight.shape[1]}, rank={self.r}{name}"
 
 
+class ColumnSequenceParallelLoRALinear(ColumnSequenceParallelLinear):
+    def __init__(
+        self,
+        in_features: int,
+        out_features: int,
+        r: int = 0,
+        lora_alpha: int = 1,
+        lora_dropout: float = 0.0,
+        rslora: bool = False,
+        lora_plus_scale: float = 1.0,
+        merge_weights: bool = True,
+        lora_A_weight_attr: Optional[paddle.ParamAttr] = None,
+        use_quick_lora: bool = False,
+        **kwargs
+    ):
+        ColumnSequenceParallelLinear.__init__(self, in_features, out_features, **kwargs)
+        if not isinstance(r, int) or r <= 0:
+            raise ValueError("Lora rank r should be a positive integer")
+        self.r = r
+        self.lora_alpha = lora_alpha
+        # Optional dropout
+        if lora_dropout > 0.0:
+            self.lora_dropout = nn.Dropout(p=lora_dropout)
+        else:
+            self.lora_dropout = lambda x: x
+        # Mark the weight as unmerged
+        self.merged = False
+        self.merge_weights = merge_weights
+
+        # compatible
+        self.name = self._name
+
+        # Actual trainable parameters
+        self.lora_A = self.create_parameter(
+            shape=[in_features, r],
+            dtype=self._dtype,
+            is_bias=False,
+            attr=lora_A_weight_attr,
+        )
+        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,
+            ),
+        )
+
+        self.lora_B.is_distributed = True
+        self.lora_B.split_axis = 1
+        if not rslora:
+            self.scaling = self.lora_alpha / self.r
+        else:
+            self.scaling = self.lora_alpha / math.sqrt(self.r)
+
+        # Freezing the pre-trained weight matrix
+        self.weight.stop_gradient = True
+        self._use_quick_lora = use_quick_lora and lora_dropout == 0.0
+
+    @property
+    def use_quick_lora(self):
+        # TODO(@gexiao): support qlora
+        return False  # self._use_quick_lora and self.training and not self.merged
+
+    def train(self):
+        super().train()
+        if self.merge_weights and self.merged:
+            # Make sure that the weights are not merged
+            new_weight = self.weight - self.lora_A @ self.lora_B * self.scaling
+            self.weight.set_value(new_weight)
+            self.merged = False
+
+    def eval(self):
+        super().eval()
+        if self.merge_weights and not self.merged:
+            # Merge the weights and mark it
+            new_weight = self.weight + self.lora_A @ self.lora_B * self.scaling
+            self.weight.set_value(new_weight)
+            self.merged = True
+
+    def forward(self, x: paddle.Tensor):
+        if MC2ColumnSeqParallelCoreLinear is None:
+            if self.is_mp:
+                input_parallel = AllGatherOp.apply(x)
+            else:
+                input_parallel = x
+            result_mp = self.linear(input_parallel, self.weight, self.bias, name=self._name)
+        else:
+            result_mp = MC2ColumnSeqParallelCoreLinear.apply(x, self.weight, self.model_parallel_group)
+            if self.bias is not None:
+                result_mp += self.bias
+
+        if not self.merged:
+            input_a = self.lora_dropout(x) @ self.lora_A
+            if MC2ColumnSeqParallelCoreLinear is None:
+                input_a = AllGatherOp.apply(input_a)
+                delta_mp = (input_a @ self.lora_B) * self.scaling
+            else:
+                input_a = MC2ColumnSeqParallelCoreLinear.apply(input_a, self.lora_B, self.model_parallel_group)
+                delta_mp = input_a * self.scaling
+            result_mp += delta_mp
+
+        if self.gather_output and self.is_mp:
+            result = mp_ops._c_concat(result_mp, group=self.model_parallel_group)
+        else:
+            result = result_mp
+        return result
+
+    def extra_repr(self):
+        name = f", name={self.name}" if self.name else ""
+        return f"in_features={self.weight.shape[0]}, out_features={self.weight.shape[1]}, rank={self.r}{name}"
+
+
 class LoRAMergedLinear(nn.Linear):
     # LoRA implemented in a dense layer  with merged linear weights for q, k, v
     def __init__(

paddlenlp/peft/lora/lora_model.py

@@ -45,13 +45,31 @@ from ...utils.distributed import distributed_gather
 from ...utils.env import LORA_WEIGHTS_NAME, SAFE_PEFT_WEIGHTS_INDEX_NAME
 from ...utils.log import logger
 from .lora_config import LoRAConfig
+
+try:
+    from paddle.distributed.fleet.utils.sequence_parallel_utils import (
+        ColumnSequenceParallelLinear,
+        RowSequenceParallelLinear,
+    )
+
+except:
+
+    class ColumnSequenceParallelLinear:
+        pass
+
+    class RowSequenceParallelLinear:
+        pass
+
+
 from .lora_layers import (
     ColumnParallelLoRALinear,
     ColumnParallelLoRAMergedLinear,
+    ColumnSequenceParallelLoRALinear,
     LoRAConv2D,
     LoRALinear,
     LoRAMergedLinear,
     RowParallelLoRALinear,
+    RowSequenceParallelLoRALinear,
 )
 
 try:
@@ -454,6 +472,58 @@ class LoRAModel(nn.Layer):
                 # Lora column parallel will spilt lora A matrix
                 self.add_lora_split_mapping(module_name + ".lora_A", is_column=False)
 
+                # for lora qat
+                if self.lora_config.do_qat:
+                    self.add_lora_split_mapping(module_name + ".weight_quanter._scale", is_column=False)
+                    self.add_lora_split_mapping(module_name + ".activation_quanter._scale", is_column=False)
+                    self.add_lora_split_mapping(module_name + ".activation_quanter.quanter._scale", is_column=False)
+            elif isinstance(module, ColumnSequenceParallelLinear):
+                # recover the original output_features
+                output_features = module.weight.shape[1] * module.world_size
+                lora_module = ColumnSequenceParallelLoRALinear(
+                    in_features=module.weight.shape[0],
+                    out_features=output_features,
+                    gather_output=module.gather_output,
+                    has_bias=module.bias is not None,
+                    r=lora_config.r,
+                    lora_alpha=lora_config.lora_alpha,
+                    lora_dropout=lora_config.lora_dropout,
+                    rslora=lora_config.rslora,
+                    lora_plus_scale=lora_config.lora_plus_scale,
+                    merge_weights=lora_config.merge_weights,
+                    lora_A_weight_attr=paddle.ParamAttr(
+                        initializer=nn.initializer.KaimingUniform(
+                            negative_slope=math.sqrt(5), nonlinearity="leaky_relu"
+                        )
+                    ),
+                    use_quick_lora=lora_config.use_quick_lora,
+                )
+                # Lora column parallel will spilt lora B matrix
+                self.add_lora_split_mapping(module_name + ".lora_B", is_column=True)
+
+                # for lora qat
+                if self.lora_config.do_qat:
+                    self.add_lora_split_mapping(module_name + ".weight_quanter._scale", is_column=True)
+                    self.add_lora_split_mapping(module_name + ".activation_quanter._scale", is_column=False)
+                    self.add_lora_split_mapping(module_name + ".activation_quanter.quanter._scale", is_column=False)
+            elif isinstance(module, RowSequenceParallelLinear):
+                # recover the original output_features
+                lora_module = RowSequenceParallelLoRALinear(
+                    in_features=module.weight.shape[0] * module.world_size,
+                    out_features=module.weight.shape[1],
+                    has_bias=module.bias is not None,
+                    input_is_parallel=module.input_is_parallel,
+                    r=lora_config.r,
+                    lora_alpha=lora_config.lora_alpha,
+                    lora_dropout=lora_config.lora_dropout,
+                    rslora=lora_config.rslora,
+                    lora_plus_scale=lora_config.lora_plus_scale,
+                    merge_weights=lora_config.merge_weights,
+                    use_quick_lora=lora_config.use_quick_lora,
+                )
+                # Lora column parallel will spilt lora A matrix
+                self.add_lora_split_mapping(module_name + ".lora_A", is_column=False)
+
                 # for lora qat
                 if self.lora_config.do_qat:
                     self.add_lora_split_mapping(module_name + ".weight_quanter._scale", is_column=False)
@@ -539,7 +609,7 @@ class LoRAModel(nn.Layer):
                 )
         if lora_module is None:
             raise ValueError(
-                f"LoRA strategy only supports paddle.nn.Linear or paddle.distributed.fleet.meta_parallel.ColumnParallelLinear. {module}({module_name}) is not supported。"
+                f"LoRA strategy only supports paddle.nn.Linear or paddle.distributed.fleet.meta_parallel.ColumnParallelLinear or paddlenlp.transformers.sequence_utils. {module}({module_name} {type(module).__name__}) is not supported。"
             )
         if getattr(lora_module, "quant_weight", None) is not None:
             lora_module.quant_weight = module.quant_weight
@@ -597,6 +667,8 @@ class LoRAModel(nn.Layer):
                 or isinstance(layer, LoRAConv2D)
                 or isinstance(layer, ColumnParallelLoRALinear)
                 or isinstance(layer, RowParallelLoRALinear)
+                or isinstance(layer, ColumnSequenceParallelLoRALinear)
+                or isinstance(layer, RowSequenceParallelLoRALinear)
                 or isinstance(layer, LoRAMergedLinear)
                 or isinstance(layer, ColumnParallelLoRAMergedLinear)
                 or (QuantizationLoRALinear is not None and isinstance(layer, QuantizationLoRALinear))
@@ -684,9 +756,11 @@ class LoRAModel(nn.Layer):
                 self._find_and_restore_module(layer_name)
             elif (
                 isinstance(layer, ColumnParallelLoRALinear)
+                or isinstance(layer, ColumnSequenceParallelLoRALinear)
                 or isinstance(layer, LoRAConv2D)
                 or isinstance(layer, ColumnParallelLoRAMergedLinear)
                 or isinstance(layer, RowParallelLoRALinear)
+                or isinstance(layer, RowSequenceParallelLoRALinear)
                 or (QuantizationLoRALinear is not None and isinstance(layer, QuantizationLoRALinear))
                 or (
                     ColumnParallelQuantizationLoRALinear is not None

paddlenlp/peft/lora/mc2_lora_npu.py

@@ -1,80 +0,0 @@
-# !/usr/bin/env python3
-
-# 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.
-
-""" mc2(tp overlap) """
-
-import paddle
-import paddle_custom_device
-from paddle.autograd import PyLayer
-
-
-class MC2LoRaRowParallelLinear(PyLayer):
-    @staticmethod
-    def forward(ctx, input_, weight, group):
-        ctx.save_for_backward(input_, weight)
-        rank = paddle.distributed.get_rank()
-        hcom_name = group.process_group.get_comm_name(rank)
-        x = input_.reshape([-1, input_.shape[-1]])
-        out = paddle_custom_device.npu.fused_mm_allreduce(
-            x, weight, bias=None, hcom=hcom_name, reduce_op="sum", comm_turn=0
-        )
-        output = out.reshape([input_.shape[0], input_.shape[1], weight.shape[1]])
-        ctx.ring_id = group.id
-        return output
-
-    @staticmethod
-    def backward(ctx, dy):
-        input_, weight = ctx.saved_tensor()
-        out_grad = dy
-        sub_grad = out_grad.reshape([-1, out_grad.shape[-1]])
-        input_grad = paddle.matmul(sub_grad, weight, transpose_y=True)
-        if weight.stop_gradient:
-            return input_grad.reshape(input_.shape)
-        else:
-            input_reshape = input_.reshape([-1, input_.shape[-1]])
-            weight_grad = paddle.matmul(input_reshape, sub_grad, transpose_x=True)
-            return input_grad.reshape(input_.shape), weight_grad
-
-
-class MC2LoRaColumnParallelLinear(PyLayer):
-    @staticmethod
-    def forward(ctx, input_, weight, group):
-        ctx.save_for_backward(input_, weight)
-        ctx.group = group
-        input_mp = input_
-        result_mp = paddle.matmul(input_mp, weight)
-        return result_mp
-
-    @staticmethod
-    def backward(ctx, dy):
-        input_, weight = ctx.saved_tensor()
-        sub_grad = dy.reshape([-1, dy.shape[-1]])
-        rank = paddle.distributed.get_rank()
-        hcom_name = ctx.group.process_group.get_comm_name(rank)
-
-        d_weight = (
-            paddle.matmul(input_.reshape([-1, input_.shape[-1]]), sub_grad, transpose_x=True)
-            if not weight.stop_gradient
-            else None
-        )
-        d_input = paddle_custom_device.npu.fused_mm_allreduce(
-            sub_grad, weight.t(), bias=None, hcom=hcom_name, reduce_op="sum", comm_turn=0
-        )
-
-        if d_weight is not None:
-            return d_input.reshape(input_.shape), d_weight
-        else:
-            return d_input.reshape(input_.shape)

paddlenlp/trainer/plugins/unified_checkpoint.py

@@ -30,6 +30,7 @@ from paddlenlp.trainer.utils.helper import distributed_file, distributed_isfile
 from paddlenlp.transformers.model_utils import (
     PretrainedModel,
     _load_state_dict_into_model,
+    faster_set_state_dict,
     get_parameter_dtype,
     load_state_dict,
     unwrap_model,
@@ -62,11 +63,13 @@ from paddlenlp.utils.env import (
     SAFE_WEIGHTS_NAME,
 )
 from paddlenlp.utils.log import logger
+from paddlenlp.utils.nested import nested_copy, nested_copy_place
 
 if is_safetensors_available():
-    from safetensors import safe_open
+    # from safetensors import safe_open
     from safetensors.numpy import save_file as safe_save_file
 
+    from paddlenlp.utils.safetensors import fast_safe_open as safe_open
 
 FP32_MASTER = "fp32_master_0"
 optimizer_scalar_name = [
@@ -90,6 +93,11 @@ __all__ = [
 async_save_queue = []
 
 
+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"
@@ -195,7 +203,6 @@ def load_unified_checkpoint(args, model, optimizer, resume_from_checkpoint: str,
     Returns:
         None
     """
-
     if paddle.distributed.get_world_size() <= 1:
         load_single_card_checkpoint(args, model, resume_from_checkpoint)
         return
@@ -221,7 +228,6 @@ def load_unified_checkpoint_locally(args, model, resume_from_checkpoint: str, sa
         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)
@@ -265,7 +271,9 @@ def load_unified_checkpoint_locally(args, model, resume_from_checkpoint: str, sa
             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)
+        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
@@ -278,11 +286,12 @@ def load_unified_checkpoint_locally(args, model, resume_from_checkpoint: str, sa
                 None, model.config, state_dict=state_dict, ignore_error=len(resolved_archive_file) > 1
             )
 
-        error_msgs += _load_state_dict_into_model(model, state_dict, "")
+        # error_msgs += _load_state_dict_into_model(model, state_dict, "")
+        error_msgs += faster_set_state_dict(model, state_dict, strict_dtype=False)
 
         # force memory release
         del state_dict
-        gc.collect()
+        # gc.collect()
 
     if len(error_msgs) > 0:
         error_msg = "\n\t".join(error_msgs)
@@ -336,6 +345,7 @@ def unified_checkpoint_into_shards(
             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
@@ -489,6 +499,7 @@ def load_unified_optimizer_locally(args, model, optimizer, resume_from_checkpoin
     # 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")
 
@@ -536,10 +547,10 @@ def load_unified_optimizer_locally(args, model, optimizer, resume_from_checkpoin
                         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)
+                    state_dict = load_state_dict(shard_file, tp_actions, expected_keys, device="expected")
                 else:
                     # for pipeline model, we don't need to use tp_actions
-                    state_dict = load_state_dict(shard_file, None, expected_keys)
+                    state_dict = load_state_dict(shard_file, None, expected_keys, device="expected")
 
             returned_state_dict.update(state_dict)
             # force memory release
@@ -552,7 +563,6 @@ def load_unified_optimizer_locally(args, model, optimizer, resume_from_checkpoin
         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("/")
@@ -561,13 +571,13 @@ def load_unified_optimizer_locally(args, model, optimizer, resume_from_checkpoin
             key_name = "_".join([static_name, FP32_MASTER, key_name[1]])
         else:
             key_name = "_".join([static_name, key_name[1]])
-        returned_optim_state_dict[key_name] = state_dict_optim[key]
+        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_master_weight.keys()):
             static_name = struct2static_name_mappings[key]
-            returned_optim_state_dict["master_weights"][static_name] = state_dict_master_weight[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])
 
     returned_optim_state_dict = nested_copy_place(
@@ -639,6 +649,7 @@ def unified_optimizer_into_shards(
             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,
@@ -647,6 +658,7 @@ def unified_optimizer_into_shards(
         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,
@@ -702,7 +714,6 @@ def unified_optimizer_into_shards(
 def check_unified_checkpoint(args, model, resume_from_checkpoint, safe_serialization=False):
     index_filename = select_model_weight_index(args, 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)
@@ -1604,7 +1615,9 @@ def gather_sharded_object(index_file, total_size, is_optimizer=False):
     tp_group = hcg.get_model_parallel_group()
     pp_group = hcg.get_pipe_parallel_group()
 
-    logger.info("Unified checkpoint generating sharded_index json files.")
+    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)
@@ -1713,8 +1726,6 @@ def filter_params(model_to_save, state_dict, is_optimizer=False):
 
 
 def merge_tensor_parallel_with_shard(state_dict, tp_actions, all_filter_keys):
-    logger.info("Unified checkpoint merge tensor parallel in shards")
-
     hcg = fleet.get_hybrid_communicate_group()
     tp_group = hcg.get_model_parallel_group()
     tp_rank = tp_group.rank
@@ -1740,7 +1751,7 @@ def merge_tensor_parallel_with_shard(state_dict, tp_actions, all_filter_keys):
                 action = tp_actions.pop(key)
                 tensor = action(ret) if is_dst else None
             else:
-                tensor = tensor._copy_to(paddle.CPUPlace(), False) if is_dst else None
+                tensor = tensor._copy_to(DEST_PLACE, False) if is_dst else None
 
             if is_dst:
                 state_dict_to_save[key] = tensor
@@ -1753,8 +1764,7 @@ def merge_tensor_parallel_with_shard(state_dict, tp_actions, all_filter_keys):
 
 
 def merge_tensor_parallel_for_optimizer(state_dict, tp_actions, all_filter_keys):
-    logger.info("Unified optimizer tensor parallel in shards")
-
+    # Core function for UC
     hcg = fleet.get_hybrid_communicate_group()
     tp_group = hcg.get_model_parallel_group()
     tp_rank = tp_group.rank
@@ -1772,15 +1782,13 @@ def merge_tensor_parallel_for_optimizer(state_dict, tp_actions, all_filter_keys)
             if model_key in tp_actions:
                 # for example: beta1, beta2
                 if tensor.numel().item() == 1:
-                    tensor = (
-                        tensor._copy_to(paddle.CPUPlace(), False) if is_dst else None
-                    )  # Need broadcast when loaded
+                    tensor = tensor._copy_to(DEST_PLACE, False) if is_dst else None  # Need broadcast when loaded
                 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:
-                tensor = tensor._copy_to(paddle.CPUPlace(), False) if is_dst else None
+                tensor = tensor._copy_to(DEST_PLACE, False) if is_dst else None
 
             if is_dst:
                 state_dict_to_save[filter_keys[i]] = tensor
@@ -1876,26 +1884,6 @@ def mapping_optimizer_tp_actions(tp_actions, optimizer_loaded_keys):
     return new_actions
 
 
-def nested_copy(inputs):
-    if isinstance(inputs, dict):
-        outputs = {}
-        for key in list(inputs.keys()):
-            outputs[key] = nested_copy(inputs[key])
-        return outputs
-    return inputs
-
-
-def nested_copy_place(inputs, place=None, blocking=False):
-    if isinstance(inputs, dict):
-        outputs = {}
-        for key in list(inputs.keys()):
-            outputs[key] = nested_copy_place(inputs[key], place, blocking)
-        return outputs
-    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:

paddlenlp/trainer/trainer.py

@@ -1011,6 +1011,7 @@ class Trainer:
                     self.timers and self.timers("optimizer-step").start()
 
                     if self.args.gradient_accumulation_steps > 1 and self._enable_delay_scale_loss():
+                        paddle.device.synchronize()
                         for p in model._layers.parameters():
                             with paddle.no_grad():
                                 if hasattr(p, "main_grad") and p.main_grad is not None:
@@ -1050,11 +1051,12 @@ class Trainer:
                     if optimizer_was_run:
                         self.lr_scheduler.step()
 
-                    if enable_release_grads and args.pipeline_parallel_degree > 1:
+                    if args.release_grads or enable_release_grads:
                         self.optimizer.clear_grad(set_to_zero=False)
-                        for _, buffers in model._chunk_2_comm_buffers.items():
-                            for buffer in buffers:
-                                buffer._clear_grad_storage()
+                        if args.pipeline_parallel_degree > 1:
+                            for _, buffers in model._chunk_2_comm_buffers.items():
+                                for buffer in buffers:
+                                    buffer._clear_grad_storage()
                     else:
                         self.optimizer.clear_grad()
 
@@ -1418,8 +1420,6 @@ class Trainer:
         if is_datasets_available() and eval_dataset is not None and isinstance(eval_dataset, datasets.Dataset):
             eval_dataset = self._remove_unused_columns(eval_dataset, description="evaluation")
 
-        _DataLoader = DistDataLoader if self.args.distributed_dataloader else DataLoader
-
         if self._is_iterable_dataset(eval_dataset):
             if self.args.dataset_world_size > 1:
                 eval_dataset = IterableDatasetShard(
@@ -1430,24 +1430,43 @@ class Trainer:
                     process_index=self.args.dataset_rank,
                 )
 
-            return _DataLoader(
-                eval_dataset,
-                batch_size=self.args.per_device_eval_batch_size,
-                collate_fn=self.data_collator,
-                num_workers=self.args.dataloader_num_workers,
-            )
+            if self.args.distributed_dataloader:
+                return DistDataLoader(
+                    eval_dataset,
+                    batch_size=self.args.per_device_eval_batch_size,
+                    collate_fn=self.data_collator,
+                    drop_last=self.args.dataloader_drop_last,
+                    num_workers=0,
+                    eval=True,
+                )
+            else:
+                return DataLoader(
+                    eval_dataset,
+                    batch_size=self.args.per_device_eval_batch_size,
+                    collate_fn=self.data_collator,
+                    drop_last=self.args.dataloader_drop_last,
+                    num_workers=0,
+                )
 
         eval_sampler = self._get_eval_sampler(eval_dataset)
 
         if self.args.distributed_dataloader:
             logger.info("Eval using DistDataLoader.")
 
-        return _DataLoader(
-            eval_dataset,
-            batch_sampler=eval_sampler,
-            collate_fn=self.data_collator,
-            num_workers=self.args.dataloader_num_workers,
-        )
+            return DistDataLoader(
+                eval_dataset,
+                batch_sampler=eval_sampler,
+                collate_fn=self.data_collator,
+                num_workers=0,
+                eval=True,
+            )
+        else:
+            return DataLoader(
+                eval_dataset,
+                batch_sampler=eval_sampler,
+                collate_fn=self.data_collator,
+                num_workers=0,
+            )
 
     def get_test_dataloader(self, test_dataset: Dataset) -> DataLoader:
         """
@@ -1468,8 +1487,6 @@ class Trainer:
         if is_datasets_available() and test_dataset is not None and isinstance(test_dataset, datasets.Dataset):
             test_dataset = self._remove_unused_columns(test_dataset, description="test")
 
-        _DataLoader = DistDataLoader if self.args.distributed_dataloader else DataLoader
-
         if self._is_iterable_dataset(test_dataset):
             if self.args.dataset_world_size > 1:
                 test_dataset = IterableDatasetShard(
@@ -1480,25 +1497,46 @@ class Trainer:
                     process_index=self.args.dataset_rank,
                 )
 
-            return _DataLoader(
-                test_dataset,
-                batch_size=self.args.per_device_eval_batch_size * self.world_size,
-                collate_fn=self.data_collator,  # _get_collator_with_removed_columns
-                num_workers=self.args.dataloader_num_workers,
-            )
+            if self.args.distributed_dataloader:
+                return DistDataLoader(
+                    test_dataset,
+                    batch_size=self.args.per_device_eval_batch_size * self.world_size,
+                    collate_fn=self.data_collator,  # _get_collator_with_removed_columns
+                    num_workers=0,
+                    drop_last=self.args.dataloader_drop_last,
+                    eval=True,
+                )
+            else:
+                return DataLoader(
+                    test_dataset,
+                    batch_size=self.args.per_device_eval_batch_size * self.world_size,
+                    collate_fn=self.data_collator,  # _get_collator_with_removed_columns
+                    drop_last=self.args.dataloader_drop_last,
+                    num_workers=0,
+                )
 
         test_sampler = self._get_eval_sampler(test_dataset)
 
         if self.args.distributed_dataloader:
             logger.info("Test using DistDataLoader.")
 
-        # We use the same batch_size as for eval.
-        return _DataLoader(
-            test_dataset,
-            batch_sampler=test_sampler,
-            collate_fn=self.data_collator,
-            drop_last=self.args.dataloader_drop_last,
-        )
+            # We use the same batch_size as for eval.
+            return DistDataLoader(
+                test_dataset,
+                batch_sampler=test_sampler,
+                collate_fn=self.data_collator,
+                num_workers=0,
+                drop_last=self.args.dataloader_drop_last,
+                eval=True,
+            )
+        else:
+            return DataLoader(
+                test_dataset,
+                batch_sampler=test_sampler,
+                collate_fn=self.data_collator,
+                drop_last=self.args.dataloader_drop_last,
+                num_workers=0,
+            )
 
     def create_optimizer_and_scheduler(self, num_training_steps: int):
         """
@@ -1560,16 +1598,13 @@ class Trainer:
                 if os.path.isfile(rng_file):
                     rng_file_list = paddle.load(rng_file, return_numpy=True)
             paddle.distributed.broadcast_object_list(rng_file_list, src=0)
-            # if rng_file_list still empty, then use old style rng_state
+            # if rng_file_list still empty, not log rng state.
             if rng_file_list[0] is None:
-                rng_file = os.path.join(checkpoint, f"rng_state_{process_index}.pth")
-                if not os.path.isfile(rng_file):
-                    logger.info(
-                        f"Didn't find an RNG file for process {process_index}, if you are resuming a training that "
-                        "wasn't launched in a distributed fashion, reproducibility is not guaranteed."
-                    )
-                    return
-                checkpoint_rng_state = paddle.load(rng_file, return_numpy=True)
+                logger.info(
+                    f"Didn't find an RNG file for process {process_index}, if you are resuming a training that "
+                    "wasn't launched in a distributed fashion, reproducibility is not guaranteed."
+                )
+                return
             else:
                 checkpoint_rng_state = rng_file_list[process_index]
         else:
@@ -1720,16 +1755,8 @@ class Trainer:
         in_sep_parallel_mode = self.args.sep_parallel_degree > 1
 
         # Multi-gpu training
-        if (
-            self.args.world_size > 1
-            and not self.args.use_hybrid_parallel
-            or not (
-                in_pipeline_parallel_mode
-                or in_sharding_parallel_mode
-                or in_tensor_parallel_mode
-                or in_sep_parallel_mode
-            )
-        ):
+        if self.args.world_size > 1 and (not self.args.use_hybrid_parallel):
+            # MOE use DDP to broadcaset parameters.
             model = paddle.DataParallel(model)
             # Distributed training (should be after fp16 initialization)
 
@@ -2096,10 +2123,10 @@ class Trainer:
             if not self.is_in_train:
                 self.args.unified_checkpoint_config = unified_checkpoint_config_backup
         if strtobool(os.getenv("FLAG_LLM_PDC", "False")):
-            # save checkpoint_done file to ensure checkpoint is complete
+            # save model_done file to ensure model is complete
             if self.args.should_save_model_state and self.args.should_save:
                 # For ckpt integrity
-                paddle.save(self.state.global_step, os.path.join(output_dir, ".checkpoint_done"))
+                paddle.save(self.state.global_step, os.path.join(output_dir, ".model_done"))
 
     def _save_checkpoint(self, model, metrics=None):
         # assert unwrap_model(model) is self.model, "internal model should be a reference to self.model"
@@ -2391,6 +2418,7 @@ class Trainer:
             self.runtime_timer.stop()
             return
 
+        logger.info("Loading optimizer and scheduler...")
         if (not self.args.should_load_sharding_stage1_model) and self.args.ignore_load_lr_and_optim:
             self.runtime_timer.stop()
             return
@@ -2720,11 +2748,15 @@ class Trainer:
         # Number of losses has been rounded to a multiple of batch_size and in a distributed training, the number of
         # samplers has been rounded to a multiple of batch_size, so we truncate.
         if all_losses is not None:
-            all_losses = all_losses[:num_samples]
+            all_losses = all_losses[: num_samples * int(self.args.world_size / self.args.dataset_world_size)]
         if all_preds is not None:
-            all_preds = nested_truncate(all_preds, num_samples)
+            all_preds = nested_truncate(
+                all_preds, num_samples * int(self.args.world_size / self.args.dataset_world_size)
+            )
         if all_labels is not None:
-            all_labels = nested_truncate(all_labels, num_samples)
+            all_labels = nested_truncate(
+                all_labels, num_samples * int(self.args.world_size / self.args.dataset_world_size)
+            )
 
         model.train()
 

paddlenlp/trainer/training_args.py

@@ -344,6 +344,8 @@ class TrainingArguments:
             Whether skip profile timer, timer will record time usage of forward/ backward/ step, etc.
         distributed_dataloader (`bool`, *optional*):
             Whether to use distributed dataloader. Default is `False`.
+        release_grads (`bool`, *optional*):
+            Whether to release gradients during training. Default is `False`.
     """
 
     output_dir: str = field(
@@ -787,6 +789,9 @@ class TrainingArguments:
         default=False,
         metadata={"help": "whether to run distributed training in auto parallel mode"},
     )
+    release_grads: Optional[bool] = field(
+        default=False, metadata={"help": "Whether to release gradients during training. Default is `False`."}
+    )
 
     def __post_init__(self):
         env_local_rank = int(os.environ.get("PADDLE_RANK_IN_NODE", -1))
@@ -1020,12 +1025,17 @@ class TrainingArguments:
                         enable_dp_comm_overlap and 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:
+                        raise ValueError(
+                            "If `enable_sharding_comm_overlap` in pipeline_parallel_configs, `amp_master_grad` must be True."
+                        )
+
                     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,
                         "enable_timer": "enable_timer" in pipeline_parallel_config,
-                        "release_gradients": "enable_release_grads" 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,
                         "clear_every_step_cache": "enable_clear_every_step_cache" in pipeline_parallel_config,
                         "use_batch_p2p_comm": "disable_batch_p2p_comm" not in pipeline_parallel_config,
@@ -1395,7 +1405,7 @@ class TrainingArguments:
             if world_size > 1:
                 if not paddle.distributed.parallel.parallel_helper._is_parallel_ctx_initialized():
                     if self.unified_checkpoint:
-                        self.use_hybrid_parallel = True
+                        # DP use hybrid group
                         strategy = fleet.DistributedStrategy()
                         fleet.init(is_collective=True, strategy=strategy)
                     else:

paddlenlp/trainer/utils/helper.py

@@ -16,8 +16,6 @@
 # This file is modified from
 #  https://github.com/huggingface/transformers/blob/main/src/transformers
 
-import collections
-import copy
 import os
 from typing import Any, Optional
 
@@ -27,6 +25,11 @@ import paddle.distributed as dist
 from paddle.distributed import fleet
 
 from paddlenlp.utils.log import logger
+from paddlenlp.utils.nested import (
+    nested_broadcast_tensor,
+    nested_empty_tensor,
+    nested_reduce_tensor,
+)
 
 __all__ = [
     "distributed_concat",
@@ -180,52 +183,6 @@ def distributed_file(filename):
         return filename
 
 
-TensorHolder = collections.namedtuple("TensorHolder", ["shape", "dtype", "name"])
-
-
-def nested_reduce_tensor(tensor):
-    if isinstance(tensor, dict):
-        # copy tensor since it will be inplace modified dict
-        tensor = copy.copy(tensor)
-        for key in list(tensor.keys()):
-            tensor[key] = nested_reduce_tensor(tensor[key])
-    if isinstance(tensor, (tuple, list)):
-        return type(tensor)(nested_reduce_tensor(t) for t in tensor)
-
-    if isinstance(tensor, paddle.Tensor):
-        return TensorHolder(tensor.shape, tensor.dtype, tensor.name)
-
-    return tensor
-
-
-def nested_empty_tensor(tensor):
-    if isinstance(tensor, dict):
-        for key in list(tensor.keys()):
-            tensor[key] = nested_empty_tensor(tensor[key])
-    if isinstance(tensor, list):
-        return type(tensor)(nested_empty_tensor(t) for t in tensor)
-
-    # TensorHolder is tuple
-    if isinstance(tensor, TensorHolder):
-        t = paddle.empty(tensor.shape, dtype=tensor.dtype, name=tensor.name)
-        t.name = tensor.name
-        return t
-
-    return tensor
-
-
-def nested_broadcast_tensor(tensor, src=0, group=None):
-    if isinstance(tensor, dict):
-        for key in list(tensor.keys()):
-            tensor[key] = nested_broadcast_tensor(tensor[key], src=src, group=group)
-    if isinstance(tensor, list):
-        return type(tensor)(nested_broadcast_tensor(t, src=src, group=group) for t in tensor)
-
-    if isinstance(tensor, paddle.Tensor):
-        paddle.distributed.broadcast(tensor, src=src, group=group, sync_op=True)
-    return tensor
-
-
 def broadcast_dp_optimizer(state_dict):
     if paddle.distributed.get_world_size() <= 1:
         return state_dict

paddlenlp/transformers/__init__.py

@@ -29,16 +29,20 @@ from .processing_utils import ProcessorMixin
 from .feature_extraction_utils import BatchFeature, FeatureExtractionMixin
 from .image_processing_utils import ImageProcessingMixin
 from .attention_utils import create_bigbird_rand_mask_idx_list
-from paddle.distributed.fleet.utils.sequence_parallel_utils import (
-    GatherOp,
-    ScatterOp,
-    AllGatherOp,
-    ReduceScatterOp,
-    ColumnSequenceParallelLinear,
-    RowSequenceParallelLinear,
-    mark_as_sequence_parallel_parameter,
-    register_sequence_parallel_allreduce_hooks,
-)
+
+try:
+    from paddle.distributed.fleet.utils.sequence_parallel_utils import (
+        GatherOp,
+        ScatterOp,
+        AllGatherOp,
+        ReduceScatterOp,
+        ColumnSequenceParallelLinear,
+        RowSequenceParallelLinear,
+        mark_as_sequence_parallel_parameter,
+        register_sequence_parallel_allreduce_hooks,
+    )
+except:
+    pass
 from .export import export_model
 
 # isort: split

paddlenlp/transformers/auto/tokenizer.py

@@ -189,13 +189,20 @@ class AutoTokenizer:
             init_class = init_kwargs.pop("tokenizer_class", None)
 
         if init_class:
-            class_name = cls._name_mapping[init_class]
-            import_class = import_module(f"paddlenlp.transformers.{class_name}.tokenizer")
-            tokenizer_class = getattr(import_class, init_class)
-            if use_fast:
-                fast_tokenizer_class = cls._get_fast_tokenizer_class(init_class, class_name)
-                tokenizer_class = fast_tokenizer_class if fast_tokenizer_class else tokenizer_class
-            return tokenizer_class
+            if init_class in cls._name_mapping:
+                class_name = cls._name_mapping[init_class]
+                import_class = import_module(f"paddlenlp.transformers.{class_name}.tokenizer")
+                tokenizer_class = getattr(import_class, init_class)
+                if use_fast:
+                    fast_tokenizer_class = cls._get_fast_tokenizer_class(init_class, class_name)
+                    tokenizer_class = fast_tokenizer_class if fast_tokenizer_class else tokenizer_class
+                return tokenizer_class
+            else:
+                import_class = import_module("paddlenlp.transformers")
+                tokenizer_class = getattr(import_class, init_class, None)
+                assert tokenizer_class is not None, f"Can't find tokenizer {init_class}"
+                return tokenizer_class
+
         # If no `init_class`, we use pattern recognition to recognize the tokenizer class.
         else:
             # TODO: Potential issue https://github.com/PaddlePaddle/PaddleNLP/pull/3786#discussion_r1024689810

paddlenlp/transformers/conversion_utils.py

@@ -285,8 +285,12 @@ def naive_fuse_merge_tp(weight_list, is_column=True, fuse_tensor_parts=2):
 
     if isinstance(weight_list[0], np.ndarray):
         return np.concatenate([reorder[i] for i in index], axis=axis)
+    else:
+        tensor = paddle.concat([reorder[i] for i in index], axis=axis)
 
-    return paddle.concat([reorder[i] for i in index], axis=axis)._copy_to(paddle.CPUPlace(), False)
+        if tensor.place.is_gpu_place():
+            tensor = tensor._copy_to(paddle.CUDAPinnedPlace(), False)
+        return tensor
 
 
 def naive_fuse_split_tp(
@@ -361,12 +365,18 @@ def normal_fuse_merge_tp(weight_list, is_column=True):
         if isinstance(weight_list[0], np.ndarray):
             return np.concatenate(weight_list, axis=-1)
         else:
-            return paddle.concat(weight_list, axis=-1)._copy_to(paddle.CPUPlace(), False)
+            tensor = paddle.concat(weight_list, axis=-1)
+            if tensor.place.is_gpu_place():
+                tensor = tensor._copy_to(paddle.CUDAPinnedPlace(), False)
+            return tensor
     else:
         if isinstance(weight_list[0], np.ndarray):
             return np.concatenate(weight_list, axis=0)
         else:
-            return paddle.concat(weight_list, axis=0)._copy_to(paddle.CPUPlace(), False)
+            tensor = paddle.concat(weight_list, axis=0)
+            if tensor.place.is_gpu_place():
+                tensor = tensor._copy_to(paddle.CUDAPinnedPlace(), False)
+            return tensor
 
 
 def normal_fuse_split_tp(weight, tensor_parallel_degree, tensor_parallel_rank=None, is_column=True):

paddlenlp/transformers/gpt/modeling.py

@@ -29,13 +29,17 @@ import paddle.tensor as tensor
 from paddle.distributed import fleet
 from paddle.distributed.fleet.meta_parallel import get_rng_state_tracker
 from paddle.distributed.fleet.utils import recompute
-from paddle.distributed.fleet.utils.sequence_parallel_utils import (
-    ColumnSequenceParallelLinear,
-    GatherOp,
-    RowSequenceParallelLinear,
-    ScatterOp,
-    mark_as_sequence_parallel_parameter,
-)
+
+try:
+    from paddle.distributed.fleet.utils.sequence_parallel_utils import (
+        ColumnSequenceParallelLinear,
+        GatherOp,
+        RowSequenceParallelLinear,
+        ScatterOp,
+        mark_as_sequence_parallel_parameter,
+    )
+except:
+    pass
 from paddle.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss
 from paddle.utils import try_import
 

paddlenlp/transformers/gpt/modeling_auto.py

@@ -30,10 +30,14 @@ import paddle.tensor as tensor
 from paddle.distributed import fleet
 from paddle.distributed.fleet.meta_parallel import get_rng_state_tracker
 from paddle.distributed.fleet.utils import recompute
-from paddle.distributed.fleet.utils.sequence_parallel_utils import (
-    ScatterOp,
-    mark_as_sequence_parallel_parameter,
-)
+
+try:
+    from paddle.distributed.fleet.utils.sequence_parallel_utils import (
+        ScatterOp,
+        mark_as_sequence_parallel_parameter,
+    )
+except:
+    pass
 
 from ...utils.converter import StateDictNameMapping
 from .. import PretrainedModel, register_base_model

paddlenlp/transformers/gpt/modeling_pp.py

@@ -19,9 +19,13 @@ from paddle.distributed.fleet.meta_parallel import (
     SharedLayerDesc,
 )
 from paddle.distributed.fleet.utils import recompute
-from paddle.distributed.fleet.utils.sequence_parallel_utils import (
-    mark_as_sequence_parallel_parameter,
-)
+
+try:
+    from paddle.distributed.fleet.utils.sequence_parallel_utils import (
+        mark_as_sequence_parallel_parameter,
+    )
+except:
+    pass
 
 from paddlenlp.transformers.model_utils import PipelinePretrainedModel
 

paddlenlp/transformers/llama/configuration.py

@@ -147,6 +147,7 @@ class LlamaConfig(PretrainedConfig):
         num_key_value_heads=None,
         initializer_range=0.02,
         rms_norm_eps=1e-6,
+        rope_theta=10000.0,
         use_cache=True,
         use_recompute=False,
         recompute_granularity="full",
@@ -188,6 +189,7 @@ class LlamaConfig(PretrainedConfig):
 
         self.initializer_range = initializer_range
         self.rms_norm_eps = rms_norm_eps
+        self.rope_theta = rope_theta
 
         self.use_cache = use_cache
         self.use_recompute = use_recompute

paddlenlp/transformers/llama/modeling.py

@@ -45,13 +45,16 @@ except ImportError:
         return F.silu(x) * y
 
 
-from paddle.distributed.fleet.utils.sequence_parallel_utils import (
-    ColumnSequenceParallelLinear,
-    GatherOp,
-    RowSequenceParallelLinear,
-    ScatterOp,
-    mark_as_sequence_parallel_parameter,
-)
+try:
+    from paddle.distributed.fleet.utils.sequence_parallel_utils import (
+        ColumnSequenceParallelLinear,
+        GatherOp,
+        RowSequenceParallelLinear,
+        ScatterOp,
+        mark_as_sequence_parallel_parameter,
+    )
+except:
+    pass
 from paddle.utils import try_import
 
 from paddlenlp.transformers.conversion_utils import (
@@ -59,6 +62,10 @@ from paddlenlp.transformers.conversion_utils import (
     init_name_mappings,
 )
 from paddlenlp.transformers.long_sequence_strategies import LongSequenceStrategies
+from paddlenlp.transformers.mc2_parallel_linear import (
+    MC2ColumnSeqParallelLinear,
+    MC2RowSeqParallelLinear,
+)
 from paddlenlp.transformers.model_outputs import (
     BaseModelOutputWithPastAndCrossAttentions,
     CausalLMOutputWithCrossAttentions,
@@ -93,13 +100,6 @@ __all__ = [
 ]
 
 
-def is_mc2_valid():
-    current_device = get_env_device()
-    if current_device == "npu":
-        return True
-    return False
-
-
 def _get_interleave(n):
     def _get_interleave_power_of_2(n):
         start = 2 ** (-(2 ** -(math.log2(n) - 3)))
@@ -571,12 +571,7 @@ class LlamaMLP(nn.Layer):
         self.fuse_attention_ffn = config.fuse_attention_ffn
 
         if config.sequence_parallel:
-            if is_mc2_valid and int(os.getenv("FLAGS_NPU_MC2", 0)):
-                from paddlenlp.transformers.mc2_seqence_parallel_linear import (
-                    MC2ColumnSeqParallelLinear,
-                    MC2RowSeqParallelLinear,
-                )
-
+            if MC2ColumnSeqParallelLinear is not None and MC2RowSeqParallelLinear is not None:
                 ColumnParallelLinear = MC2ColumnSeqParallelLinear
                 RowParallelLinear = MC2RowSeqParallelLinear
             else:
@@ -694,12 +689,7 @@ class LlamaAttention(nn.Layer):
                 self.use_fused_rope = False
 
         if config.sequence_parallel:
-            if is_mc2_valid and int(os.getenv("FLAGS_NPU_MC2", 0)):
-                from paddlenlp.transformers.mc2_seqence_parallel_linear import (
-                    MC2ColumnSeqParallelLinear,
-                    MC2RowSeqParallelLinear,
-                )
-
+            if MC2ColumnSeqParallelLinear is not None and MC2RowSeqParallelLinear is not None:
                 ColumnParallelLinear = MC2ColumnSeqParallelLinear
                 RowParallelLinear = MC2RowSeqParallelLinear
             else:
@@ -810,24 +800,28 @@ class LlamaAttention(nn.Layer):
             self.rotary_emb = LlamaRotaryEmbedding(
                 self.head_dim,
                 max_position_embeddings=self.max_position_embeddings,
+                base=self.config.rope_theta,
             )
         elif self.config.rope_scaling_type == "linear":
             self.rotary_emb = LlamaLinearScalingRotaryEmbedding(
                 self.head_dim,
                 max_position_embeddings=self.max_position_embeddings,
                 scaling_factor=self.config.rope_scaling_factor,
+                base=self.config.rope_theta,
             )
         elif self.config.rope_scaling_type == "ntk":
             self.rotary_emb = LlamaNTKScalingRotaryEmbedding(
                 self.head_dim,
                 max_position_embeddings=self.max_position_embeddings,
                 scaling_factor=self.config.rope_scaling_factor,
+                base=self.config.rope_theta,
             )
         elif self.config.rope_scaling_type == "dynamic_ntk":
             self.rotary_emb = LlamaDynamicNTKScalingRotaryEmbedding(
                 self.head_dim,
                 max_position_embeddings=self.max_position_embeddings,
                 scaling_factor=self.config.rope_scaling_factor,
+                base=self.config.rope_theta,
             )
         else:
             raise ValueError(f"Unknown RoPE scaling type {self.config.rope_scaling_type}")
@@ -900,6 +894,7 @@ class LlamaAttention(nn.Layer):
             query_states = self.q_proj(hidden_states)
             key_states = self.k_proj(hidden_states)
             value_states = self.v_proj(hidden_states)
+
             if self.reshard_layer is not None:
                 if self.sequence_parallel:
                     assert self.seq_length % self.config.sep_parallel_degree == 0
@@ -1024,7 +1019,6 @@ class LlamaAttention(nn.Layer):
             value_states = paddle.concat([past_key_value[1], value_states], axis=1)
 
         past_key_value = (key_states, value_states) if use_cache else None
-
         if self.kv_indices is not None:
             key_states = paddle.index_select(key_states, self.kv_indices, axis=2)
             value_states = paddle.index_select(value_states, self.kv_indices, axis=2)
@@ -1033,7 +1027,7 @@ class LlamaAttention(nn.Layer):
         # repeat k/v heads if n_kv_heads < n_heads
         # paddle version > 2.6 or develop support flash-attn with gqa/mqa
         paddle_version = float(paddle.__version__[:3])
-        if (paddle_version != 0.0) and (paddle_version <= 2.6):
+        if not self.config.use_flash_attention or ((paddle_version != 0.0) and (paddle_version <= 2.6)):
             key_states = repeat_kv(key_states, self.num_key_value_groups)
             value_states = repeat_kv(value_states, self.num_key_value_groups)
 
@@ -1557,7 +1551,6 @@ class LlamaModel(LlamaPretrainedModel):
             else:
                 attention_mask = attention_mask.astype("bool")
         hidden_states = inputs_embeds
-
         # decoder layers
         all_hidden_states = () if output_hidden_states else None
         all_self_attns = () if output_attentions else None

paddlenlp/transformers/llama/tokenizer.py

@@ -24,7 +24,7 @@ from ...utils.log import logger
 from .. import PretrainedTokenizer
 from ..tokenizer_utils_base import BatchEncoding, EncodedInput, PaddingStrategy
 
-__all__ = ["LlamaTokenizer"]
+__all__ = ["LlamaTokenizer", "Llama3Tokenizer"]
 
 
 class LlamaTokenizer(PretrainedTokenizer):
@@ -199,6 +199,7 @@ class LlamaTokenizer(PretrainedTokenizer):
         """
         Create a mask from the two sequences passed to be used in a sequence-pair classification task. T5 does not make
         use of token type ids, therefore a list of zeros is returned.
+
         Args:
             token_ids_0 (`List[int]`):
                 List of IDs.
@@ -270,3 +271,292 @@ class LlamaTokenizer(PretrainedTokenizer):
                         constant_values=0,
                     )
         return encoded_inputs
+
+
+"""Copied Tokenization classes for QWen."""
+
+import base64
+import unicodedata
+from typing import Collection, Dict, List, Optional, Set, Tuple, Union
+
+from ...utils.import_utils import is_tiktoken_available
+from .. import PretrainedTokenizer
+from ..tokenizer_utils_base import (
+    AddedToken,
+    BatchEncoding,
+    EncodedInput,
+    PaddingStrategy,
+)
+
+VOCAB_FILES_NAMES = {"vocab_file": "tokenizer.model"}
+
+PAT_STR = "(?i:'s|'t|'re|'ve|'m|'ll|'d)|[^\\r\\n\\p{L}\\p{N}]?\\p{L}+|\\p{N}{1,3}| ?[^\\s\\p{L}\\p{N}]+[\\r\\n]*|\\s*[\\r\\n]+|\\s+(?!\\S)|\\s+"
+BEGINOFTEXT = "<|begin_of_text|>"
+ENDOFTEXT = "<|end_of_text|>"
+IMSTART = "<|start_header_id|>"
+IMEND = "<|end_header_id|>"
+EOTID = "<|eot_id|>"
+# as the default behavior is changed to allow special tokens in
+# regular texts, the surface forms of special tokens need to be
+# as different as possible to minimize the impact
+EXTRAS = tuple((f"<|reserved_special_token_{i}|>" for i in range(251)))
+SPECIAL_TOKENS = (BEGINOFTEXT, ENDOFTEXT) + EXTRAS[0:4] + (IMSTART, IMEND, EXTRAS[4], EOTID) + EXTRAS[5:]
+
+tiktoken = None
+
+
+def _load_tiktoken_bpe(tiktoken_bpe_file: str) -> Dict[bytes, int]:
+    with open(tiktoken_bpe_file, "rb") as f:
+        contents = f.read()
+    return {
+        base64.b64decode(token): int(rank) for token, rank in (line.split() for line in contents.splitlines() if line)
+    }
+
+
+class Llama3Tokenizer(PretrainedTokenizer):
+    """QWen tokenizer."""
+
+    model_input_names = ["input_ids", "attention_mask", "position_ids"]
+    resource_files_names = VOCAB_FILES_NAMES
+
+    def __init__(
+        self,
+        vocab_file,
+        errors="replace",
+        padding_side="left",
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        if not is_tiktoken_available():
+            raise ValueError("tiktoken is not installed, please install it use: pip install tiktoken")
+
+        import tiktoken as tk
+
+        tiktoken = tk
+
+        self.errors = errors  # how to handle errors in decoding
+
+        self.mergeable_ranks = _load_tiktoken_bpe(vocab_file)  # type: dict[bytes, int]
+        self.special_tokens = {
+            token: index for index, token in enumerate(SPECIAL_TOKENS, start=len(self.mergeable_ranks))
+        }
+        enc = tiktoken.Encoding(
+            "Llama3",
+            pat_str=PAT_STR,
+            mergeable_ranks=self.mergeable_ranks,
+            special_tokens=self.special_tokens,
+        )
+        assert (
+            len(self.mergeable_ranks) + len(self.special_tokens) == enc.n_vocab
+        ), f"{len(self.mergeable_ranks) + len(self.special_tokens)} != {enc.n_vocab} in encoding"
+
+        self.decoder = {v: k for k, v in self.mergeable_ranks.items()}  # type: dict[int, bytes|str]
+        self.decoder.update({v: k for k, v in self.special_tokens.items()})
+
+        self.tokenizer = enc  # type: tiktoken.Encoding
+
+        self.bod_id = self.special_tokens[BEGINOFTEXT]
+        self.eod_id = self.special_tokens[ENDOFTEXT]
+        self.start_header_id = self.special_tokens[IMSTART]
+        self.end_header_id = self.special_tokens[IMEND]
+        self.eot_id = self.special_tokens[EOTID]
+
+        if "pad_token_id" in kwargs:
+            self.pad_token_id = kwargs["pad_token_id"]
+        if "eos_token_id" in kwargs:
+            self.eos_token_id = kwargs["eos_token_id"]
+
+    def __len__(self) -> int:
+        return self.tokenizer.n_vocab
+
+    def get_vocab(self) -> Dict[bytes, int]:
+        return self.mergeable_ranks
+
+    def convert_tokens_to_ids(self, tokens: Union[bytes, str, List[Union[bytes, str]]]) -> List[int]:
+        ids = []
+        if isinstance(tokens, (str, bytes)):
+            if tokens in self.special_tokens:
+                return self.special_tokens[tokens]
+            else:
+                return self.mergeable_ranks.get(tokens)
+        for token in tokens:
+            if token in self.special_tokens:
+                ids.append(self.special_tokens[token])
+            else:
+                ids.append(self.mergeable_ranks.get(token))
+        return ids
+
+    def _add_tokens(self, new_tokens: Union[List[str], List[AddedToken]], special_tokens: bool = False) -> int:
+        if not special_tokens and new_tokens:
+            raise ValueError("Adding regular tokens is not supported")
+        for token in new_tokens:
+            surface_form = token.content if isinstance(token, AddedToken) else token
+            if surface_form not in SPECIAL_TOKENS:
+                raise ValueError("Adding unknown special tokens is not supported")
+        return 0
+
+    def save_vocabulary(self, save_directory: str, **kwargs) -> Tuple[str]:
+        """
+        Save only the vocabulary of the tokenizer (vocabulary).
+
+        Returns:
+            `Tuple(str)`: Paths to the files saved.
+        """
+        file_path = os.path.join(save_directory, "tokenizer.model")
+        with open(file_path, "w", encoding="utf8") as w:
+            for k, v in self.mergeable_ranks.items():
+                line = base64.b64encode(k).decode("utf8") + " " + str(v) + "\n"
+                w.write(line)
+        return (file_path,)
+
+    def tokenize(
+        self,
+        text: str,
+        allowed_special: Union[Set, str] = "all",
+        disallowed_special: Union[Collection, str] = (),
+        **kwargs,
+    ) -> List[Union[bytes, str]]:
+        """
+        Converts a string in a sequence of tokens.
+
+        Args:
+            text (`str`):
+                The sequence to be encoded.
+            allowed_special (`Literal["all"]` or `set`):
+                The surface forms of the tokens to be encoded as special tokens in regular texts.
+                Default to "all".
+            disallowed_special (`Literal["all"]` or `Collection`):
+                The surface forms of the tokens that should not be in regular texts and trigger errors.
+                Default to an empty tuple.
+
+            kwargs (additional keyword arguments, *optional*):
+                Will be passed to the underlying model specific encode method.
+
+        Returns:
+            `List[bytes|str]`: The list of tokens.
+        """
+        tokens = []
+        text = unicodedata.normalize("NFC", text)
+
+        # this implementation takes a detour: text -> token id -> token surface forms
+        for t in self.tokenizer.encode(text, allowed_special=allowed_special, disallowed_special=disallowed_special):
+            tokens.append(self.decoder[t])
+        return tokens
+
+    def convert_tokens_to_string(self, tokens: List[Union[bytes, str]]) -> str:
+        """
+        Converts a sequence of tokens in a single string.
+        """
+        text = ""
+        temp = b""
+        for t in tokens:
+            if isinstance(t, str):
+                if temp:
+                    text += temp.decode("utf-8", errors=self.errors)
+                    temp = b""
+                text += t
+            elif isinstance(t, bytes):
+                temp += t
+            else:
+                raise TypeError("token should only be of type types or str")
+        if temp:
+            text += temp.decode("utf-8", errors=self.errors)
+        return text
+
+    @property
+    def vocab_size(self):
+        return self.tokenizer.n_vocab
+
+    def _convert_id_to_token(self, index: int) -> Union[bytes, str]:
+        """Converts an id to a token, special tokens included"""
+        if index in self.decoder:
+            return self.decoder[index]
+        raise ValueError("unknown ids")
+
+    def _convert_token_to_id(self, token: Union[bytes, str]) -> int:
+        """Converts a token to an id using the vocab, special tokens included"""
+        if token in self.special_tokens:
+            return self.special_tokens[token]
+        if token in self.mergeable_ranks:
+            return self.mergeable_ranks[token]
+        raise ValueError("unknown token")
+
+    def _tokenize(self, text: str, **kwargs):
+        """
+        Converts a string in a sequence of tokens (string), using the tokenizer. Split in words for word-based
+        vocabulary or sub-words for sub-word-based vocabularies (BPE/SentencePieces/WordPieces).
+
+        Do NOT take care of added tokens.
+        """
+        raise NotImplementedError
+
+    def _decode(
+        self,
+        token_ids: Union[int, List[int]],
+        skip_special_tokens: bool = False,
+        errors: str = None,
+        **kwargs,
+    ) -> str:
+        if isinstance(token_ids, int):
+            token_ids = [token_ids]
+        if skip_special_tokens:
+            token_ids = [i for i in token_ids if i < self.eod_id]
+        return self.tokenizer.decode(token_ids, errors=errors or self.errors)
+
+    def _pad(
+        self,
+        encoded_inputs: Union[Dict[str, EncodedInput], BatchEncoding],
+        max_length: Optional[int] = None,
+        padding_strategy: PaddingStrategy = PaddingStrategy.DO_NOT_PAD,
+        pad_to_multiple_of: Optional[int] = None,
+        return_attention_mask: Optional[bool] = None,
+    ) -> dict:
+        """
+        Pad encoded inputs (on left/right and up to predefined length or max length in the batch)
+
+        Args:
+            encoded_inputs:
+                Dictionary of tokenized inputs (`List[int]`) or batch of tokenized inputs (`List[List[int]]`).
+            max_length: maximum length of the returned list and optionally padding length (see below).
+                Will truncate by taking into account the special tokens.
+            padding_strategy: PaddingStrategy to use for padding.
+
+                - PaddingStrategy.LONGEST Pad to the longest sequence in the batch
+                - PaddingStrategy.MAX_LENGTH: Pad to the max length (default)
+                - PaddingStrategy.DO_NOT_PAD: Do not pad
+                The tokenizer padding sides are defined in self.padding_side:
+
+                    - 'left': pads on the left of the sequences
+                    - 'right': pads on the right of the sequences
+            pad_to_multiple_of: (optional) Integer if set will pad the sequence to a multiple of the provided value.
+                This is especially useful to enable the use of Tensor Core on NVIDIA hardware with compute capability
+                >= 7.5 (Volta).
+            return_attention_mask:
+                (optional) Set to False to avoid returning attention mask (default: set to model specifics)
+        """
+        # Load from model defaults
+
+        # attention_mask shape [1,seq_len,seq_len]
+        if "attention_mask" in encoded_inputs and len(np.shape(encoded_inputs["attention_mask"])) > 2:
+            attention_mask = encoded_inputs["attention_mask"]
+            encoded_inputs.pop("attention_mask")
+        else:
+            attention_mask = None
+
+        required_input = encoded_inputs[self.model_input_names[0]]
+        encoded_inputs = super()._pad(
+            encoded_inputs, max_length, padding_strategy, pad_to_multiple_of, return_attention_mask
+        )
+        if attention_mask is not None and len(np.shape(attention_mask)) > 2:
+            encoded_inputs["attention_mask"] = attention_mask
+            needs_to_be_padded = padding_strategy != PaddingStrategy.DO_NOT_PAD and len(required_input) != max_length
+            if needs_to_be_padded:
+                difference = max_length - len(required_input)
+                if "attention_mask" in encoded_inputs:
+                    encoded_inputs["attention_mask"] = np.pad(
+                        encoded_inputs["attention_mask"],
+                        pad_width=[(0, 0), (difference, 0), (difference, 0)],
+                        mode="constant",
+                        constant_values=0,
+                    )
+        return encoded_inputs

paddlenlp/transformers/mc2_parallel_linear.py

@@ -0,0 +1,230 @@
+# 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 os
+
+import paddle
+
+try:
+    import paddle_custom_device
+except ImportError:
+    pass
+
+from paddle import distributed as dist
+from paddle.autograd import PyLayer
+
+try:
+    from paddle.distributed.fleet.utils.sequence_parallel_utils import (
+        ColumnSequenceParallelLinear,
+        RowSequenceParallelLinear,
+    )
+except:
+    pass
+from paddlenlp.utils.tools import get_env_device
+
+__all_gather_recomputation__ = False
+if int(os.getenv("MC2_Recompute", 0)):
+    __all_gather_recomputation__ = True
+
+
+def is_mc2_valid():
+    current_device = get_env_device()
+    if current_device == "npu":
+        return int(os.getenv("MC2", 0))
+    return 0
+
+
+if is_mc2_valid():
+
+    class MC2ColumnParallelCoreLinear(PyLayer):
+        @staticmethod
+        def forward(ctx, input_, weight, group):
+            ctx.save_for_backward(input_, weight)
+            ctx.group = group
+            input_mp = input_
+            result_mp = paddle.matmul(input_mp, weight)
+            return result_mp
+
+        @staticmethod
+        def backward(ctx, dy):
+            input_, weight = ctx.saved_tensor()
+            sub_grad = dy.reshape([-1, dy.shape[-1]])
+            rank = paddle.distributed.get_rank()
+            hcom_name = ctx.group.process_group.get_comm_name(rank)
+
+            d_weight = (
+                paddle.matmul(input_.reshape([-1, input_.shape[-1]]), sub_grad, transpose_x=True)
+                if not weight.stop_gradient
+                else None
+            )
+            d_input = paddle_custom_device.npu.fused_mm_allreduce(
+                sub_grad, weight.t(), bias=None, hcom=hcom_name, reduce_op="sum", comm_turn=0
+            )
+
+            if d_weight is not None:
+                return d_input.reshape(input_.shape), d_weight
+            else:
+                return d_input.reshape(input_.shape), None
+
+    class MC2RowParallelCoreLinear(PyLayer):
+        @staticmethod
+        def forward(ctx, input_, weight, group):
+            ctx.save_for_backward(input_, weight)
+            rank = paddle.distributed.get_rank()
+            hcom_name = group.process_group.get_comm_name(rank)
+            x = input_.reshape([-1, input_.shape[-1]])
+            out = paddle_custom_device.npu.fused_mm_allreduce(
+                x, weight, bias=None, hcom=hcom_name, reduce_op="sum", comm_turn=0
+            )
+            output = out.reshape([input_.shape[0], input_.shape[1], weight.shape[1]])
+            ctx.ring_id = group.id
+            return output
+
+        @staticmethod
+        def backward(ctx, dy):
+            input_, weight = ctx.saved_tensor()
+            out_grad = dy
+            sub_grad = out_grad.reshape([-1, out_grad.shape[-1]])
+            input_grad = paddle.matmul(sub_grad, weight, transpose_y=True)
+            if weight.stop_gradient:
+                return input_grad.reshape(input_.shape), None
+            else:
+                input_reshape = input_.reshape([-1, input_.shape[-1]])
+                weight_grad = paddle.matmul(input_reshape, sub_grad, transpose_x=True)
+                return input_grad.reshape(input_.shape), weight_grad
+
+    class MC2ColumnSeqParallelCoreLinear(PyLayer):
+        @staticmethod
+        def forward(ctx, input_, weight, group):
+            ctx.weight_stop_gradient = weight.stop_gradient
+            ctx.save_for_backward(input_, weight)
+
+            rank = dist.get_rank()
+            hcomm_info = group.process_group.get_comm_name(rank)
+
+            world_size = group.nranks
+            output, gather_out = paddle_custom_device.npu.fused_allgather_mm(
+                input_,
+                weight,
+                bias=None,
+                hcom=hcomm_info,
+                world_size=world_size,
+                gather_index=0,
+                gather_output=(not __all_gather_recomputation__),
+                comm_turn=0,
+            )
+
+            ctx.all_gather_output = gather_out
+            ctx.world_size = world_size
+            ctx.group = group
+            return output
+
+        @staticmethod
+        def backward(ctx, grad_output):
+            input_, weight = ctx.saved_tensor()
+
+            if __all_gather_recomputation__:
+                dim_size = input_.shape
+                dim_size[0] = dim_size[0] * ctx.world_size
+                all_gather_output = paddle.empty(dim_size, dtype=input_.dtype)
+                all_gather_output.stop_gradient = True
+                all_gather_work = dist.stream.all_gather(all_gather_output, input_, group=ctx.group, sync_op=False)
+            else:
+                all_gather_output = ctx.all_gather_output
+
+            grad_input = paddle.matmul(grad_output, weight, transpose_y=True)
+            sub_grad_input = paddle.empty(input_.shape, dtype=input_.dtype)
+            reduce_scatter_work = dist.stream.reduce_scatter(
+                sub_grad_input, grad_input, group=ctx.group, sync_op=False
+            )
+
+            if __all_gather_recomputation__:
+                all_gather_work.wait()
+
+            grad_weight = (
+                paddle.matmul(all_gather_output, grad_output, transpose_x=True)
+                if not ctx.weight_stop_gradient
+                else None
+            )
+            reduce_scatter_work.wait()
+
+            return sub_grad_input, grad_weight
+
+    class MC2RowSeqParallelCoreLinear(PyLayer):
+        @staticmethod
+        def forward(ctx, input_, weight, group):
+            ctx.weight_stop_gradient = weight.stop_gradient
+            ctx.save_for_backward(input_, weight)
+
+            rank = dist.get_rank()
+            hcomm_info = group.process_group.get_comm_name(rank)
+            world_size = group.nranks
+
+            output = paddle_custom_device.npu.fused_mm_reduce_scatter(
+                input_,
+                weight,
+                bias=None,
+                hcom=hcomm_info,
+                world_size=world_size,
+                reduce_op="sum",
+                comm_turn=0,
+            )
+
+            ctx.hcomm_info = hcomm_info
+            ctx.world_size = world_size
+            return output
+
+        @staticmethod
+        def backward(ctx, grad_output):
+            input_, weight = ctx.saved_tensor()
+            hcomm_info = ctx.hcomm_info
+            world_size = ctx.world_size
+
+            grad_input, all_gather_grad_output = paddle_custom_device.npu.fused_allgather_mm(
+                grad_output,
+                weight.t(),
+                bias=None,
+                hcom=hcomm_info,
+                world_size=world_size,
+                gather_index=0,
+                gather_output=True,
+                comm_turn=0,
+            )
+            grad_weight = (
+                paddle.matmul(input_, all_gather_grad_output, transpose_x=True)
+                if not ctx.weight_stop_gradient
+                else None
+            )
+
+            return grad_input, grad_weight
+
+    class MC2ColumnSeqParallelLinear(ColumnSequenceParallelLinear):
+        def forward(self, x):
+            output = MC2ColumnSeqParallelCoreLinear.apply(x, self.weight, self.model_parallel_group)
+            output = output + self.bias if self.bias is not None else output
+            return output
+
+    class MC2RowSeqParallelLinear(RowSequenceParallelLinear):
+        def forward(self, x):
+            output = MC2RowSeqParallelCoreLinear.apply(x, self.weight, self.model_parallel_group)
+            output = output + self.bias if self.bias is not None else output
+            return output
+
+else:
+    MC2ColumnSeqParallelCoreLinear = None
+    MC2RowSeqParallelCoreLinear = None
+    MC2ColumnSeqParallelLinear = None
+    MC2RowSeqParallelLinear = None
+    MC2ColumnParallelCoreLinear = None
+    MC2RowParallelCoreLinear = None

paddlenlp/transformers/mc2_seqence_parallel_linear.py

@@ -1,142 +0,0 @@
-# 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 os
-
-import paddle
-
-try:
-    import paddle_custom_device
-except ImportError:
-    raise ImportError("Current device does not support MC2!")
-
-from paddle import distributed as dist
-from paddle.autograd import PyLayer
-from paddle.distributed.fleet.utils.sequence_parallel_utils import (
-    ColumnSequenceParallelLinear,
-    RowSequenceParallelLinear,
-)
-
-__all_gather_recomputation__ = False
-if int(os.getenv("MC2_Recompute", 0)):
-    __all_gather_recomputation__ = True
-
-
-class MC2Column(PyLayer):
-    @staticmethod
-    def forward(ctx, input_, weight, group):
-        ctx.save_for_backward(input_, weight)
-
-        rank = dist.get_rank()
-        hcomm_info = group.process_group.get_comm_name(rank)
-
-        world_size = group.nranks
-        output, gather_out = paddle_custom_device.npu.fused_allgather_mm(
-            input_,
-            weight,
-            bias=None,
-            hcom=hcomm_info,
-            world_size=world_size,
-            gather_index=0,
-            gather_output=(not __all_gather_recomputation__),
-            comm_turn=0,
-        )
-
-        ctx.all_gather_output = gather_out
-        ctx.world_size = world_size
-        ctx.group = group
-        return output
-
-    @staticmethod
-    def backward(ctx, grad_output):
-        input_, weight = ctx.saved_tensor()
-
-        if __all_gather_recomputation__:
-            dim_size = input_.shape
-            dim_size[0] = dim_size[0] * ctx.world_size
-            all_gather_output = paddle.empty(dim_size, dtype=input_.dtype)
-            all_gather_output.stop_gradient = True
-            all_gather_work = dist.stream.all_gather(all_gather_output, input_, group=ctx.group, sync_op=False)
-        else:
-            all_gather_output = ctx.all_gather_output
-
-        grad_input = paddle.matmul(grad_output, weight, transpose_y=True)
-        sub_grad_input = paddle.empty(input_.shape, dtype=input_.dtype)
-        reduce_scatter_work = dist.stream.reduce_scatter(sub_grad_input, grad_input, group=ctx.group, sync_op=False)
-
-        if __all_gather_recomputation__:
-            all_gather_work.wait()
-
-        grad_weight = paddle.matmul(all_gather_output, grad_output, transpose_x=True)
-        reduce_scatter_work.wait()
-
-        return sub_grad_input, grad_weight
-
-
-class MC2Row(PyLayer):
-    @staticmethod
-    def forward(ctx, input_, weight, group):
-        ctx.save_for_backward(input_, weight)
-
-        rank = dist.get_rank()
-        hcomm_info = group.process_group.get_comm_name(rank)
-        world_size = group.nranks
-
-        output = paddle_custom_device.npu.fused_mm_reduce_scatter(
-            input_,
-            weight,
-            bias=None,
-            hcom=hcomm_info,
-            world_size=world_size,
-            reduce_op="sum",
-            comm_turn=0,
-        )
-
-        ctx.hcomm_info = hcomm_info
-        ctx.world_size = world_size
-        return output
-
-    @staticmethod
-    def backward(ctx, grad_output):
-        input_, weight = ctx.saved_tensor()
-        hcomm_info = ctx.hcomm_info
-        world_size = ctx.world_size
-
-        grad_input, all_gather_grad_output = paddle_custom_device.npu.fused_allgather_mm(
-            grad_output,
-            weight.t(),
-            bias=None,
-            hcom=hcomm_info,
-            world_size=world_size,
-            gather_index=0,
-            gather_output=True,
-            comm_turn=0,
-        )
-        grad_weight = paddle.matmul(input_, all_gather_grad_output, transpose_x=True)
-
-        return grad_input, grad_weight
-
-
-class MC2ColumnSeqParallelLinear(ColumnSequenceParallelLinear):
-    def forward(self, x):
-        output = MC2Column.apply(x, self.weight, self.model_parallel_group)
-        output = output + self.bias if self.bias is not None else output
-        return output
-
-
-class MC2RowSeqParallelLinear(RowSequenceParallelLinear):
-    def forward(self, x):
-        output = MC2Row.apply(x, self.weight, self.model_parallel_group)
-        output = output + self.bias if self.bias is not None else output
-        return output

paddlenlp/transformers/mixtral/modeling.py

@@ -33,13 +33,16 @@ try:
 except ImportError:
     fused_rotary_position_embedding = None
 
-from paddle.distributed.fleet.utils.sequence_parallel_utils import (
-    ColumnSequenceParallelLinear,
-    GatherOp,
-    RowSequenceParallelLinear,
-    ScatterOp,
-    mark_as_sequence_parallel_parameter,
-)
+try:
+    from paddle.distributed.fleet.utils.sequence_parallel_utils import (
+        ColumnSequenceParallelLinear,
+        GatherOp,
+        RowSequenceParallelLinear,
+        ScatterOp,
+        mark_as_sequence_parallel_parameter,
+    )
+except:
+    pass
 
 from paddlenlp.transformers.conversion_utils import (
     StateDictNameMapping,

paddlenlp/transformers/model_utils.py

@@ -109,10 +109,13 @@ def unwrap_optimizer(optimizer, optimizer_instances=()):
 
 if is_safetensors_available():
 
-    from safetensors import safe_open
-    from safetensors.numpy import load_file as safe_load_file
+    # from safetensors import safe_open
+    # from safetensors.numpy import load_file as safe_load_file
     from safetensors.numpy import save_file as safe_save_file
 
+    from paddlenlp.utils.safetensors import fast_load_file as safe_load_file
+    from paddlenlp.utils.safetensors import fast_safe_open as safe_open
+
 
 def prune_linear_layer(layer: nn.Linear, index: paddle.Tensor, dim: int = 0) -> nn.Linear:
     """
@@ -313,7 +316,7 @@ def get_parameter_dtype(parameter: nn.Layer) -> paddle.dtype:
 
 
 def load_state_dict(
-    checkpoint_file: Union[str, os.PathLike], tensor_parallel_split_mapping=None, fliter_dict_keys=None
+    checkpoint_file: Union[str, os.PathLike], tensor_parallel_split_mapping=None, fliter_dict_keys=None, device="cpu"
 ):
     """
     Reads a PaddlePaddle checkpoint file, returning properly formatted errors if they arise.
@@ -346,11 +349,16 @@ def load_state_dict(
                         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
 
-            for k in list(state_dict.keys()):
-                with device_guard():
-                    state_dict[k] = paddle.Tensor(state_dict.pop(k), zero_copy=True)
+            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)
 
             return state_dict
 
@@ -672,8 +680,10 @@ def load_sharded_checkpoint(model, folder, variant=None, strict=True, prefer_saf
     return missing_keys, unexpected_keys
 
 
-def faster_set_state_dict(model, state_dict):
+def faster_set_state_dict(model, state_dict, strict_dtype=True):
     # the state_dict will be destroied.
+    unused_keys = set(state_dict.keys())
+    unset_keys = set(model.state_dict().keys())
     with paddle.no_grad():
         for k, v in model.state_dict().items():
             if k in state_dict:
@@ -683,8 +693,10 @@ def faster_set_state_dict(model, state_dict):
                         f"faster_set_state_dict need state dict with paddle.Tensor, but got {type(v_new)}"
                     )
                 # 2. cast param / Tensor to dtype
+                #
                 if v.dtype != v_new.dtype:
-                    raise ValueError(f"for key: {k}, expect dtype {v.dtype}, but got {v_new.dtype}")
+                    if strict_dtype or (not v.is_floating_point() or not v_new.is_floating_point()):
+                        raise ValueError(f"for key: {k}, expect dtype {v.dtype}, but got {v_new.dtype}")
                 # check shape
                 if list(v.shape) != list(v_new.shape):
                     raise ValueError(f"for key: {k}, expect shape {v.shape}, but got {v_new.shape}")
@@ -700,9 +712,22 @@ def faster_set_state_dict(model, state_dict):
                 else:
                     new_t = v_new
 
+                if not strict_dtype and v.dtype != new_t.dtype:
+                    new_t = new_t.astype(v.dtype)
+
                 # 4. share Tensor to origin param / Tensor
                 src_tensor = new_t.value().get_tensor()
                 dst_tensor._share_data_with(src_tensor)
+                unset_keys.remove(k)
+                unused_keys.remove(k)
+
+    error_msgs = []
+    # if len(unset_keys) > 0:
+    #    error_msgs.append(f"Those weight of model is not initialized: {list(unset_keys)}")
+    if len(unused_keys) > 0:
+        error_msgs.append(f"Those state dict keys are not using in model: {list(unused_keys)}")
+
+    return error_msgs
 
 
 def _load_state_dict_into_model(model_to_load, state_dict, start_prefix):
@@ -734,9 +759,8 @@ def _convert_state_dict_dtype_and_shape(state_dict, model_to_load):
     def is_0d_or_1d(tensor):
         return len(tensor.shape) == 0 or list(tensor.shape) == [1]
 
-    expected_place = paddle.framework._current_expected_place()
     for key, value in model_to_load.state_dict().items():
-        if key in state_dict:
+        if key in list(state_dict.keys()):
             if isinstance(state_dict[key], np.ndarray):
                 raise ValueError(
                     "convert_state_dict_dtype expected paddle.Tensor not numpy.ndarray, plase convert numpy.ndarray to paddle.Tensor"
@@ -744,12 +768,7 @@ def _convert_state_dict_dtype_and_shape(state_dict, model_to_load):
             # confirm parameter cast is executed on the same device as model
             # TODO: cast(FP32 -> FP16) has diff on different devices, need to fix it
             if state_dict[key].is_floating_point() and state_dict[key].dtype != value.dtype:
-                value_pop = state_dict.pop(key)
-                value_new_place = (
-                    value_pop if value_pop.place == expected_place else value_pop._copy_to(expected_place, False)
-                )
-                state_dict[key] = paddle.cast(value_new_place, value.dtype)._copy_to(value_pop.place, False)
-                del value_new_place
+                state_dict[key] = paddle.cast(state_dict.pop(key), value.dtype)
             # unified 0d and 1d tensor
             if is_0d_or_1d(value) and is_0d_or_1d(state_dict[key]):
                 if list(value.shape) != list(state_dict[key].shape):
@@ -779,7 +798,7 @@ def _load_state_dict_into_meta_model(
 
     dtype = convert_np_dtype_to_dtype_(dtype)
     error_msgs = []
-
+    model_state_dict = model.state_dict()
     for param_name, param in state_dict.items():
         # First part of the test is always true as loaded_state_dict_keys always contains state_dict keys.
         if param_name not in loaded_state_dict_keys or param_name not in expected_keys:
@@ -814,7 +833,7 @@ def _load_state_dict_into_meta_model(
             if old_param is not None:
                 param = param.astype(dtype=old_param.dtype)
         with paddle.no_grad():
-            model.state_dict()[param_name].get_tensor()._share_data_with(param.value().get_tensor())
+            model_state_dict[param_name].get_tensor()._share_data_with(param.value().get_tensor())
             param.value().get_tensor()._clear()
     return error_msgs
 
@@ -1871,7 +1890,7 @@ class PretrainedModel(Layer, GenerationMixin, ConversionMixin):
                 if (
                     shard_file.endswith(".safetensors")
                     and config.tensor_parallel_degree > 1
-                    and "tp" not in shard_file
+                    and "tp" not in os.path.split(shard_file)[-1]
                 ):
                     pre_tensor_parallel_split = True
                     assert loaded_keys is not None, "loaded_keys is not None."

paddlenlp/utils/batch_sampler.py

@@ -14,8 +14,6 @@
 
 from __future__ import division, print_function
 
-import math
-
 import paddle
 
 __all__ = ["DistributedBatchSampler"]
@@ -110,7 +108,7 @@ class DistributedBatchSampler(paddle.io.BatchSampler):
             # In pre-training mode when using distributed dataloader, the input dataset can be None. We should handle this situation.
             self.num_samples = 0
         else:
-            self.num_samples = int(math.ceil(len(self.dataset) * 1.0 / self.nranks))
+            self.num_samples = int(len(self.dataset) * 1.0 / self.nranks)
         self.total_size = self.num_samples * self.nranks
 
     def get_start_end_idx(self):
@@ -125,7 +123,7 @@ class DistributedBatchSampler(paddle.io.BatchSampler):
             self.consumed_samples,
             self.nranks,
         )
-        self.remain_num_samples = int(math.ceil((len(self.dataset) - self.consumed_samples) * 1.0 / self.nranks))
+        self.remain_num_samples = int((len(self.dataset) - self.consumed_samples) * 1.0 / self.nranks)
         self.remain_total_size = self.remain_num_samples * self.nranks
         self.batch_size_times_rank_size = self.batch_size * self.nranks
 

paddlenlp/utils/nested.py

@@ -0,0 +1,83 @@
+# 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 collections
+import copy
+
+import paddle
+
+TensorHolder = collections.namedtuple("TensorHolder", ["shape", "dtype", "name"])
+
+
+def nested_reduce_tensor(tensor):
+    if isinstance(tensor, dict):
+        # copy tensor since it will be inplace modified dict
+        tensor = copy.copy(tensor)
+        for key in list(tensor.keys()):
+            tensor[key] = nested_reduce_tensor(tensor[key])
+    if isinstance(tensor, (tuple, list)):
+        return type(tensor)(nested_reduce_tensor(t) for t in tensor)
+
+    if isinstance(tensor, paddle.Tensor):
+        return TensorHolder(tensor.shape, tensor.dtype, tensor.name)
+
+    return tensor
+
+
+def nested_empty_tensor(tensor):
+    if isinstance(tensor, dict):
+        for key in list(tensor.keys()):
+            tensor[key] = nested_empty_tensor(tensor[key])
+    if isinstance(tensor, list):
+        return type(tensor)(nested_empty_tensor(t) for t in tensor)
+
+    # TensorHolder is tuple
+    if isinstance(tensor, TensorHolder):
+        t = paddle.empty(tensor.shape, dtype=tensor.dtype, name=tensor.name)
+        t.name = tensor.name
+        return t
+
+    return tensor
+
+
+def nested_broadcast_tensor(tensor, src=0, group=None):
+    if isinstance(tensor, dict):
+        for key in list(tensor.keys()):
+            tensor[key] = nested_broadcast_tensor(tensor[key], src=src, group=group)
+    if isinstance(tensor, list):
+        return type(tensor)(nested_broadcast_tensor(t, src=src, group=group) for t in tensor)
+
+    if isinstance(tensor, paddle.Tensor):
+        paddle.distributed.broadcast(tensor, src=src, group=group, sync_op=True)
+    return tensor
+
+
+def nested_copy(inputs):
+    if isinstance(inputs, dict):
+        outputs = {}
+        for key in list(inputs.keys()):
+            outputs[key] = nested_copy(inputs[key])
+        return outputs
+    return inputs
+
+
+def nested_copy_place(inputs, place=None, blocking=False):
+    if isinstance(inputs, dict):
+        outputs = {}
+        for key in list(inputs.keys()):
+            outputs[key] = nested_copy_place(inputs[key], place, blocking)
+        return outputs
+    if isinstance(inputs, paddle.Tensor):
+        inputs = inputs if inputs.place == place else inputs._copy_to(place, blocking)
+    return inputs

paddlenlp/utils/safetensors.py

@@ -0,0 +1,312 @@
+# 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 json
+import mmap
+from collections import OrderedDict
+
+import numpy as np
+
+__all__ = [
+    "fast_safe_open",
+    "fast_load_file",
+]
+
+
+MAX_HEADER_SIZE = 100 * 1000 * 1000
+
+dtype_size = {
+    "BOOL": 1,
+    "U8": 1,
+    "I8": 1,
+    "F8_E5M2": 1,
+    "F8_E4M3": 1,
+    "I16": 2,
+    "U16": 2,
+    "I32": 4,
+    "U32": 4,
+    "I64": 8,
+    "U64": 8,
+    "F16": 2,
+    "BF16": 2,
+    "F32": 4,
+    "F64": 8,
+}
+
+numpy_dtype = {
+    "BOOL": np.bool_,
+    "U8": np.uint8,
+    "I8": np.int8,
+    "F8_E5M2": 1,  # no fp8
+    "F8_E4M3": 1,  # no fp8
+    "I16": np.int16,
+    "U16": np.uint16,
+    "I32": np.int32,
+    "U32": np.uint32,
+    "I64": np.int64,
+    "U64": np.uint64,
+    "F16": np.float16,
+    "BF16": 2,  # no bf16
+    "F32": np.float32,
+    "F64": np.float64,
+}
+
+
+def getSize(fileobject):
+    fileobject.seek(0, 2)  # move the cursor to the end of the file
+    size = fileobject.tell()
+    fileobject.seek(0)  # move the cursor to the start of the file
+    return size
+
+
+def metadata_validate(metadata):
+    start = 0
+    for key, info in metadata.items():
+        s, e = info["data_offsets"]
+        if s != start or e < s:
+            raise ValueError(f"SafeTensorError::InvalidOffset({key})")
+        start = e
+        nelements = np.prod(info["shape"])
+        nbytes = nelements * dtype_size[info["dtype"]]
+        if (e - s) != nbytes:
+            raise ValueError("SafeTensorError::TensorInvalidInfo")
+    return start
+
+
+def read_metadata(buffer):
+    buffer_len = getSize(buffer)
+    if buffer_len < 8:
+        raise ValueError("SafeTensorError::HeaderTooSmall")
+
+    n = np.frombuffer(buffer.read(8), dtype=np.uint64).item()
+    if n > MAX_HEADER_SIZE:
+        raise ValueError("SafeTensorError::HeaderTooLarge")
+
+    stop = n + 8
+    if stop > buffer_len:
+        raise ValueError("SafeTensorError::InvalidHeaderLength")
+
+    tensors = json.loads(buffer.read(n), object_pairs_hook=OrderedDict)
+    metadata = tensors.pop("__metadata__", None)
+    buffer_end = metadata_validate(tensors)
+
+    if buffer_end + 8 + n != buffer_len:
+        raise ValueError("SafeTensorError::MetadataIncompleteBuffer")
+
+    return stop, tensors, metadata
+
+
+def readinto_numpy(meta, buffer, base_ptr):
+    def create_empty(info):
+        return np.empty(shape=info["shape"], dtype=numpy_dtype[info["dtype"]])
+
+    ret = {}
+    for k, v in meta.items():
+        t = create_empty(v)
+        buffer.seek(base_ptr + v["data_offsets"][0])
+        buffer.readinto(memoryview(t))
+        ret[k] = t
+    return ret
+
+
+class PySafeSlice:
+    def __init__(self, info, bufferfile, base_ptr, buffermmap):
+        self.info = info
+        self.bufferfile = bufferfile
+        self.buffermmap = buffermmap
+        self.base_ptr = base_ptr
+
+        self.start = [0 for dim in self.shape]
+        self.stop = [dim for dim in self.shape]
+        self.step = [1 for dim in self.shape]
+
+    @property
+    def ndim(self):
+        return len(self.shape)
+
+    def __getitem__(self, index):
+        # https://github.com/numpy/numpy/blob/4d652465cea38e9504f954ac708d91e4954bd13a/numpy/lib/_arrayterator_impl.py#L96-L126
+        # Fix index, handling ellipsis and incomplete slices.
+        if not isinstance(index, tuple):
+            index = (index,)
+        fixed = []
+        length, dims = len(index), self.ndim
+        for slice_ in index:
+            if slice_ is Ellipsis:
+                fixed.extend([slice(None)] * (dims - length + 1))
+                length = len(fixed)
+            elif isinstance(slice_, int):
+                fixed.append(slice(slice_, slice_ + 1, 1))
+            else:
+                fixed.append(slice_)
+        index = tuple(fixed)
+        if len(index) < dims:
+            index += (slice(None),) * (dims - len(index))
+
+        out_start, out_stop, out_step = copy.deepcopy((self.start, self.stop, self.step))
+        for i, (start, stop, step, slice_) in enumerate(zip(self.start, self.stop, self.step, index)):
+            out_start[i] = slice_.start if slice_.start is not None else 0
+            out_step[i] = slice_.step if slice_.step is not None else 1
+            out_stop[i] = slice_.stop if slice_.stop is not None else stop - start
+            out_stop[i] = min(stop, out_stop[i])
+
+        target_shape = []
+        for x, y, z, sli in zip(out_start, out_stop, out_step, index):
+            assert z == 1, "only support step = 1"
+            if y - x > 1 or sli.step is None:
+                target_shape.append(max(int(y - x), 0))
+
+        if len(target_shape) == 0:
+            if self.shape == [1]:
+                target_shape = self.shape
+
+        # https://github.com/huggingface/safetensors/blob/b947b59079a6197d7930dfb535818ac4896113e8/safetensors/src/slice.rs#L297-L315
+        indices = []
+        span = self.bits
+        for i, (start, stop, step) in enumerate(zip(out_start[::-1], out_stop[::-1], out_step[::-1])):
+            if len(indices) == 0:
+                if start == 0 and stop == self.shape[i]:
+                    pass
+                    #  We haven't started to slice yet, just increase the span
+                else:
+                    offset = start * span
+                    small_span = stop * span - offset
+                    indices.append((offset, offset + small_span))
+
+            else:
+                capacity = (stop - start) * len(indices)
+                newindices = []
+                for n in range(start, stop):
+                    offset = n * span
+                    for (old_start, old_stop) in indices:
+                        newindices.append((old_start + offset, old_stop + offset))
+                indices = newindices
+                assert len(indices) == capacity, f"error {capacity} {len(indices)}"
+            span *= self.shape[-(i + 1)]
+
+        if len(indices) == 0:
+            indices.append((0, self.nbytes))
+
+        merge_indices = []
+        last_end = -1
+        last_start = -1
+        for start, end in indices:
+            if start == last_end:
+                last_end = end
+                continue
+            else:
+                if last_start != -1:
+                    merge_indices.append((last_start, last_end))
+                last_start = start
+                last_end = end
+        if last_start != -1:
+            merge_indices.append((last_start, last_end))
+        tensor = np.empty(shape=[1] if len(target_shape) == 0 else np.prod(target_shape), dtype=self.dtype)
+
+        tensor_view = memoryview(tensor.view(np.uint8).reshape(-1))
+        curr_data_ptr = 0
+        # if to many slice and each slice < 1M
+        if len(merge_indices) > 128 and (merge_indices[0][1] - merge_indices[0][0] < 1024 * 1024):
+            # Use mmap for random access
+            for start, end in merge_indices:
+                data_len = end - start
+                tensor_view[curr_data_ptr : curr_data_ptr + data_len] = self.buffermmap[
+                    self.start_offset + start : self.start_offset + end
+                ]
+                curr_data_ptr += data_len
+        else:
+            # Use file read for sequence access
+            for start, end in merge_indices:
+                data_len = end - start
+                self.bufferfile.seek(self.start_offset + start)
+                view = tensor_view[curr_data_ptr : curr_data_ptr + data_len]
+                self.bufferfile.readinto(view)
+                curr_data_ptr += data_len
+
+        return tensor.reshape(target_shape)
+
+    def get(self, *args, **kwargs):
+        tensor = np.empty(shape=self.shape, dtype=self.dtype)
+        self.bufferfile.seek(self.start_offset)
+        self.bufferfile.readinto(memoryview(tensor))
+        return tensor
+
+    @property
+    def start_offset(self):
+        return self.base_ptr + self.info["data_offsets"][0]
+
+    def get_shape(self):
+        return self.shape
+
+    @property
+    def shape(self):
+        return self.info["shape"]
+
+    @property
+    def dtype(self):
+        return numpy_dtype[self.info["dtype"]]
+
+    @property
+    def nelements(self):
+        return np.prod(self.info["shape"])
+
+    @property
+    def bits(self):
+        return dtype_size[self.info["dtype"]]
+
+    @property
+    def nbytes(self):
+        return self.nelements * dtype_size[self.info["dtype"]]
+
+
+# a simple file writer object
+class fast_safe_open:
+    def __init__(self, filename, framework=None, device="cpu"):
+        self.filename = filename
+        self.framework = framework
+        self.file = open(self.filename, "rb")
+        self.file_mmap = mmap.mmap(self.file.fileno(), 0, flags=mmap.MAP_PRIVATE)
+        self.base, self.tensors_decs, self.__metadata__ = read_metadata(self.file)
+        self.tensors = OrderedDict()
+        for key, info in self.tensors_decs.items():
+            self.tensors[key] = PySafeSlice(info, self.file, self.base, self.file_mmap)
+            self.tensors[key].key = key
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, *args):
+        self.file_mmap.close()
+        self.file.close()
+
+    def metadata(self):
+        return self.__metadata__
+
+    def keys(self):
+        return list(self.tensors.keys())
+
+    def get_tensor(self, name):
+        return self.tensors[name].get()
+
+    def get_slice(self, name):
+        return self.tensors[name]
+
+
+def fast_load_file(filename):
+    result = {}
+    with fast_safe_open(filename, framework="np") as f:
+        for k in f.keys():
+            result[k] = f.get_tensor(k)
+    return result

pyproject.toml

@@ -10,7 +10,7 @@ exclude = ['.flake8']
 
 [tool.pytest.ini_options]
 minversion = "6.0"
-addopts = "-ra -q --ignore model_zoo/gpt-3/"
+addopts = "-ra -q --dist loadgroup"
 pythonpath = ["."]
 testpaths = [
     "tests/data",
@@ -28,7 +28,7 @@ testpaths = [
     "tests/prompt",
     # "tests/taskflow",  TODO (paddle 2.5.1 breaks this test suite, debug later)
     "tests/utils",
-    "model_zoo",
+    # "model_zoo",
 ]
 python_files = [
     "test.py",

requirements.txt

@@ -23,3 +23,4 @@ safetensors
 tool_helpers
 aistudio-sdk>=0.1.3
 jinja2
+regex

setup.py

@@ -109,7 +109,7 @@ def show():
         f.write(content)
 
 
-__version__ = "2.7.1.post"
+__version__ = "2.8.0.post"
 if os.getenv(PADDLENLP_STABLE_VERSION):
     __version__ = __version__.replace(".post", "")
 

tests/trainer/test_lora_unified_checkpoint.py

@@ -149,7 +149,7 @@ class TestUnifiedCheckpointBase(TestMultipleGpus):
 
     def setUp(self):
         """
-        1. update runfrist and rerun to run defined different config
+        1. update runfirst and rerun to run defined different config
         2. update need_allclose to True if you want to check the result
         3. update rtol to the relative value you want to check
         """
@@ -169,7 +169,7 @@ class TestUnifiedCheckpointBase(TestMultipleGpus):
 
         self.run_lora_file = "llm/finetune_generation.py"
 
-    def runfrist(self, train_args):
+    def runfirst(self, train_args):
         self.run_n1c8(self.run_lora_file, **train_args)
 
     def rerun(self, train_args):
@@ -181,7 +181,7 @@ class TestUnifiedCheckpointBase(TestMultipleGpus):
         remove_ckpt(lora_arguments["output_dir"])
 
         train_args = self.configs["TP4PP2"]
-        self.runfrist(train_args)
+        self.runfirst(train_args)
         self.rerun(train_args)
 
         if self.need_allclose:
@@ -196,7 +196,7 @@ class TestUnifiedCheckpointBase(TestMultipleGpus):
         remove_ckpt(lora_arguments["output_dir"])
 
         train_args = self.configs["TP2Sharding4"]
-        self.runfrist(train_args)
+        self.runfirst(train_args)
         self.rerun(train_args)
 
         if self.need_allclose:
@@ -213,7 +213,7 @@ class TestUnifiedCheckpointFull(TestUnifiedCheckpointBase):
         remove_ckpt(lora_arguments["output_dir"])
 
         train_args = self.configs["TP8"]
-        self.runfrist(train_args)
+        self.runfirst(train_args)
         self.rerun(train_args)
 
         if self.need_allclose:
@@ -227,7 +227,7 @@ class TestUnifiedCheckpointFull(TestUnifiedCheckpointBase):
         remove_ckpt(lora_arguments["output_dir"])
 
         train_args = self.configs["TP4DP2"]
-        self.runfrist(train_args)
+        self.runfirst(train_args)
         self.rerun(train_args)
 
         if self.need_allclose:
@@ -242,7 +242,7 @@ class TestUnifiedCheckpointFull(TestUnifiedCheckpointBase):
         remove_ckpt(lora_arguments["output_dir"])
 
         train_args = self.configs["TP4Sharding2"]
-        self.runfrist(train_args)
+        self.runfirst(train_args)
         self.rerun(train_args)
 
         if self.need_allclose:
@@ -257,7 +257,7 @@ class TestUnifiedCheckpointFull(TestUnifiedCheckpointBase):
         remove_ckpt(lora_arguments["output_dir"])
 
         train_args = self.configs["TP2PP4"]
-        self.runfrist(train_args)
+        self.runfirst(train_args)
         self.rerun(train_args)
 
         if self.need_allclose:
@@ -272,7 +272,7 @@ class TestUnifiedCheckpointFull(TestUnifiedCheckpointBase):
         remove_ckpt(lora_arguments["output_dir"])
 
         train_args = self.configs["PP8"]
-        self.runfrist(train_args)
+        self.runfirst(train_args)
         self.rerun(train_args)
 
         if self.need_allclose:
@@ -287,7 +287,7 @@ class TestUnifiedCheckpointFull(TestUnifiedCheckpointBase):
         remove_ckpt(lora_arguments["output_dir"])
 
         train_args = self.configs["PP4DP2"]
-        self.runfrist(train_args)
+        self.runfirst(train_args)
         self.rerun(train_args)
 
         if self.need_allclose:
@@ -302,7 +302,7 @@ class TestUnifiedCheckpointFull(TestUnifiedCheckpointBase):
         remove_ckpt(lora_arguments["output_dir"])
 
         train_args = self.configs["PP4Sharding2"]
-        self.runfrist(train_args)
+        self.runfirst(train_args)
         self.rerun(train_args)
 
         if self.need_allclose:
@@ -317,7 +317,7 @@ class TestUnifiedCheckpointFull(TestUnifiedCheckpointBase):
         remove_ckpt(lora_arguments["output_dir"])
 
         train_args = self.configs["Sharding8S1"]
-        self.runfrist(train_args)
+        self.runfirst(train_args)
         self.rerun(train_args)
 
         if self.need_allclose:
@@ -332,7 +332,7 @@ class TestUnifiedCheckpointFull(TestUnifiedCheckpointBase):
         remove_ckpt(lora_arguments["output_dir"])
 
         train_args = self.configs["Sharding8S2"]
-        self.runfrist(train_args)
+        self.runfirst(train_args)
         self.rerun(train_args)
 
         if self.need_allclose:
@@ -347,7 +347,7 @@ class TestUnifiedCheckpointFull(TestUnifiedCheckpointBase):
         remove_ckpt(lora_arguments["output_dir"])
 
         train_args = self.configs["Sharding4S1DP2"]
-        self.runfrist(train_args)
+        self.runfirst(train_args)
         self.rerun(train_args)
 
         if self.need_allclose:
@@ -362,7 +362,7 @@ class TestUnifiedCheckpointFull(TestUnifiedCheckpointBase):
         remove_ckpt(lora_arguments["output_dir"])
 
         train_args = self.configs["Sharding4S2DP2"]
-        self.runfrist(train_args)
+        self.runfirst(train_args)
         self.rerun(train_args)
 
         if self.need_allclose:
@@ -377,7 +377,7 @@ class TestUnifiedCheckpointFull(TestUnifiedCheckpointBase):
         remove_ckpt(lora_arguments["output_dir"])
 
         train_args = self.configs["Sharding2S1DP4"]
-        self.runfrist(train_args)
+        self.runfirst(train_args)
         self.rerun(train_args)
 
         if self.need_allclose:
@@ -392,7 +392,7 @@ class TestUnifiedCheckpointFull(TestUnifiedCheckpointBase):
         remove_ckpt(lora_arguments["output_dir"])
 
         train_args = self.configs["Sharding2S2DP4"]
-        self.runfrist(train_args)
+        self.runfirst(train_args)
         self.rerun(train_args)
 
         if self.need_allclose:
@@ -407,7 +407,7 @@ class TestUnifiedCheckpointFull(TestUnifiedCheckpointBase):
         remove_ckpt(lora_arguments["output_dir"])
 
         train_args = self.configs["DP8"]
-        self.runfrist(train_args)
+        self.runfirst(train_args)
         self.rerun(train_args)
 
         if self.need_allclose:
@@ -416,19 +416,21 @@ class TestUnifiedCheckpointFull(TestUnifiedCheckpointBase):
             np.testing.assert_allclose(res[0], res[1], self.rtol)
 
 
+@pytest.mark.skipif(True, reason="Skip for None CE")
 class TestUnifiedCheckpointOnN2C4(TestUnifiedCheckpointBase):
     def setUp(self):
         super().setUp()
         self.need_allclose = True
         self.rtol = 1e-7
 
-    def runfrist(self, train_args):
+    def runfirst(self, train_args):
         self.run_n2c4(self.run_lora_file, **train_args)
 
     def rerun(self, train_args):
         self.run_n2c4(self.run_lora_file, **train_args)
 
 
+@pytest.mark.skipif(True, reason="Skip for None CE")
 class TestUnifiedCheckpointOnN1C8CheckpointCompatible(TestUnifiedCheckpointBase):
     def setUp(self):
         super().setUp()
@@ -436,7 +438,7 @@ class TestUnifiedCheckpointOnN1C8CheckpointCompatible(TestUnifiedCheckpointBase)
         self.need_allclose = True
         self.rtol = 1e-7
 
-    def runfrist(self, train_args):
+    def runfirst(self, train_args):
         train_args["unified_checkpoint"] = 0
         self.run_n1c8(self.run_lora_file, **train_args)
 
@@ -445,6 +447,7 @@ class TestUnifiedCheckpointOnN1C8CheckpointCompatible(TestUnifiedCheckpointBase)
         self.run_n1c8(self.run_lora_file, **train_args)
 
 
+@pytest.mark.skipif(True, reason="Skip for None CE")
 class TestPaddleCheckpointOnN1C8Reset(TestUnifiedCheckpointBase):
     def setUp(self):
         super().setUp()
@@ -452,7 +455,7 @@ class TestPaddleCheckpointOnN1C8Reset(TestUnifiedCheckpointBase):
         self.need_allclose = True
         self.rtol = 1e-7
 
-    def runfrist(self, train_args):
+    def runfirst(self, train_args):
         train_args["unified_checkpoint"] = 0
         self.run_n1c8(self.run_lora_file, **train_args)
 
@@ -469,7 +472,7 @@ class TestUnifiedCheckpointOnN2C4CheckpointCompatible(TestUnifiedCheckpointBase)
         self.need_allclose = True
         self.rtol = 1e-7
 
-    def runfrist(self, train_args):
+    def runfirst(self, train_args):
         train_args["unified_checkpoint"] = 0
         self.run_n2c4(self.run_lora_file, **train_args)
 

tests/trainer/test_unified_checkpoint.py

@@ -48,6 +48,7 @@ environment_variables = {
     "Flags_skip_mp_c_identity": "1",
     "FLAGS_shard_norm_align_dp": "0",
     "FLAGS_shard_use_reduce": "1",
+    "FLAGS_eager_communication_connection": "1",  # no lazy init comm group
     "test_ci_no_save_model": "1",
 }
 
@@ -174,7 +175,7 @@ class TestUnifiedCheckpointBase(TestMultipleGpus):
 
     def setUp(self):
         """
-        1. update runfrist and rerun to run defined diffrent config
+        1. update runfirst and rerun to run defined diffrent config
         2. update need_allclose to True if you want to check the result
         3. update rtol to the relative value you want to check
         """
@@ -193,7 +194,7 @@ class TestUnifiedCheckpointBase(TestMultipleGpus):
 
         self.run_pretrain_file = "llm/llama/run_pretrain.py"
 
-    def runfrist(self, train_args):
+    def runfirst(self, train_args):
         self.run_n1c8(self.run_pretrain_file, **train_args)
 
     def rerun(self, train_args):
@@ -205,7 +206,7 @@ class TestUnifiedCheckpointBase(TestMultipleGpus):
         remove_ckpt(pretrain_arguments["output_dir"])
 
         train_args = self.configs["TP4PP2"]
-        self.runfrist(train_args)
+        self.runfirst(train_args)
         self.rerun(train_args)
 
         if self.need_allclose:
@@ -220,7 +221,7 @@ class TestUnifiedCheckpointBase(TestMultipleGpus):
         remove_ckpt(pretrain_arguments["output_dir"])
 
         train_args = self.configs["TP2Sharding4"]
-        self.runfrist(train_args)
+        self.runfirst(train_args)
         self.rerun(train_args)
 
         if self.need_allclose:
@@ -237,7 +238,7 @@ class TestUnifiedCheckpointFull(TestUnifiedCheckpointBase):
         remove_ckpt(pretrain_arguments["output_dir"])
 
         train_args = self.configs["TP8"]
-        self.runfrist(train_args)
+        self.runfirst(train_args)
         self.rerun(train_args)
 
         if self.need_allclose:
@@ -251,7 +252,7 @@ class TestUnifiedCheckpointFull(TestUnifiedCheckpointBase):
         remove_ckpt(pretrain_arguments["output_dir"])
 
         train_args = self.configs["TP4DP2"]
-        self.runfrist(train_args)
+        self.runfirst(train_args)
         self.rerun(train_args)
 
         if self.need_allclose:
@@ -266,7 +267,7 @@ class TestUnifiedCheckpointFull(TestUnifiedCheckpointBase):
         remove_ckpt(pretrain_arguments["output_dir"])
 
         train_args = self.configs["TP4Sharding2"]
-        self.runfrist(train_args)
+        self.runfirst(train_args)
         self.rerun(train_args)
 
         if self.need_allclose:
@@ -281,7 +282,7 @@ class TestUnifiedCheckpointFull(TestUnifiedCheckpointBase):
         remove_ckpt(pretrain_arguments["output_dir"])
 
         train_args = self.configs["TP2PP4"]
-        self.runfrist(train_args)
+        self.runfirst(train_args)
         self.rerun(train_args)
 
         if self.need_allclose:
@@ -296,7 +297,7 @@ class TestUnifiedCheckpointFull(TestUnifiedCheckpointBase):
         remove_ckpt(pretrain_arguments["output_dir"])
 
         train_args = self.configs["PP8"]
-        self.runfrist(train_args)
+        self.runfirst(train_args)
         self.rerun(train_args)
 
         if self.need_allclose:
@@ -311,7 +312,7 @@ class TestUnifiedCheckpointFull(TestUnifiedCheckpointBase):
         remove_ckpt(pretrain_arguments["output_dir"])
 
         train_args = self.configs["PP4DP2"]
-        self.runfrist(train_args)
+        self.runfirst(train_args)
         self.rerun(train_args)
 
         if self.need_allclose:
@@ -326,7 +327,7 @@ class TestUnifiedCheckpointFull(TestUnifiedCheckpointBase):
         remove_ckpt(pretrain_arguments["output_dir"])
 
         train_args = self.configs["PP4Sharding2"]
-        self.runfrist(train_args)
+        self.runfirst(train_args)
         self.rerun(train_args)
 
         if self.need_allclose:
@@ -341,7 +342,7 @@ class TestUnifiedCheckpointFull(TestUnifiedCheckpointBase):
         remove_ckpt(pretrain_arguments["output_dir"])
 
         train_args = self.configs["Sharding8S1"]
-        self.runfrist(train_args)
+        self.runfirst(train_args)
         self.rerun(train_args)
 
         if self.need_allclose:
@@ -356,7 +357,7 @@ class TestUnifiedCheckpointFull(TestUnifiedCheckpointBase):
         remove_ckpt(pretrain_arguments["output_dir"])
 
         train_args = self.configs["Sharding8S2"]
-        self.runfrist(train_args)
+        self.runfirst(train_args)
         self.rerun(train_args)
 
         if self.need_allclose:
@@ -371,7 +372,7 @@ class TestUnifiedCheckpointFull(TestUnifiedCheckpointBase):
         remove_ckpt(pretrain_arguments["output_dir"])
 
         train_args = self.configs["Sharding4S1DP2"]
-        self.runfrist(train_args)
+        self.runfirst(train_args)
         self.rerun(train_args)
 
         if self.need_allclose:
@@ -386,7 +387,7 @@ class TestUnifiedCheckpointFull(TestUnifiedCheckpointBase):
         remove_ckpt(pretrain_arguments["output_dir"])
 
         train_args = self.configs["Sharding4S2DP2"]
-        self.runfrist(train_args)
+        self.runfirst(train_args)
         self.rerun(train_args)
 
         if self.need_allclose:
@@ -401,7 +402,7 @@ class TestUnifiedCheckpointFull(TestUnifiedCheckpointBase):
         remove_ckpt(pretrain_arguments["output_dir"])
 
         train_args = self.configs["Sharding2S1DP4"]
-        self.runfrist(train_args)
+        self.runfirst(train_args)
         self.rerun(train_args)
 
         if self.need_allclose:
@@ -416,7 +417,7 @@ class TestUnifiedCheckpointFull(TestUnifiedCheckpointBase):
         remove_ckpt(pretrain_arguments["output_dir"])
 
         train_args = self.configs["Sharding2S2DP4"]
-        self.runfrist(train_args)
+        self.runfirst(train_args)
         self.rerun(train_args)
 
         if self.need_allclose:
@@ -431,7 +432,7 @@ class TestUnifiedCheckpointFull(TestUnifiedCheckpointBase):
         remove_ckpt(pretrain_arguments["output_dir"])
 
         train_args = self.configs["DP8"]
-        self.runfrist(train_args)
+        self.runfirst(train_args)
         self.rerun(train_args)
 
         if self.need_allclose:
@@ -440,13 +441,14 @@ class TestUnifiedCheckpointFull(TestUnifiedCheckpointBase):
             np.testing.assert_allclose(res[0], res[1], self.rtol)
 
 
+@pytest.mark.skipif(True, reason="Skip for None CE")
 class TestUnifiedCheckpointOnN2C4(TestUnifiedCheckpointBase):
     def setUp(self):
         super().setUp()
         self.need_allclose = True
         self.rtol = 1e-7
 
-    def runfrist(self, train_args):
+    def runfirst(self, train_args):
         self.run_n2c4(self.run_pretrain_file, **train_args)
 
     def rerun(self, train_args):
@@ -462,7 +464,7 @@ class TestUnifiedCheckpointOnN1C8Dynamic(TestUnifiedCheckpointFull):
         self.rtol = 1e-4
         self.k = MAX_CONVERT_CONFIGS  # max: 16, min: 1
 
-    def runfrist(self, train_args):
+    def runfirst(self, train_args):
         self.run_n1c8(self.run_pretrain_file, **train_args)
 
     def rerun(self, train_args):
@@ -484,7 +486,7 @@ class TestUnifiedCheckpointOnN2C4Dynamic(TestUnifiedCheckpointBase):
         self.rtol = 1e-4
         self.k = MAX_CONVERT_CONFIGS  # max: 16, min: 1
 
-    def runfrist(self, train_args):
+    def runfirst(self, train_args):
         self.run_n2c4(self.run_pretrain_file, **train_args)
 
     def rerun(self, train_args):
@@ -506,7 +508,7 @@ class TestUnifiedCheckpointOnN1C8ToN2C4(TestUnifiedCheckpointBase):
         self.rtol = 1e-4
         self.k = MAX_CONVERT_CONFIGS  # max: 16, min: 1
 
-    def runfrist(self, train_args):
+    def runfirst(self, train_args):
         self.run_n1c8(self.run_pretrain_file, **train_args)
         move_checkpoint_N1C8_to_N2C4()
 
@@ -528,7 +530,7 @@ class TestUnifiedCheckpointOnN2C4ToN1C8(TestUnifiedCheckpointBase):
         self.rtol = 1e-4
         self.k = MAX_CONVERT_CONFIGS  # max: 16, min: 1
 
-    def runfrist(self, train_args):
+    def runfirst(self, train_args):
         self.run_n2c4(self.run_pretrain_file, **train_args)
         move_checkpoint_N2C4_to_N1C8()
 
@@ -556,7 +558,7 @@ class TestUnifiedCheckpointOnN1C8SkipSaveModelWeight(TestUnifiedCheckpointBase):
         self.need_allclose = True
         self.rtol = 1e-7
 
-    def runfrist(self, train_args):
+    def runfirst(self, train_args):
         self.run_n1c8(self.run_pretrain_file, **train_args)
 
     def rerun(self, train_args):
@@ -575,7 +577,7 @@ class TestUnifiedCheckpointOnN1C8MasterWeightCompatibleO1ToO2(TestUnifiedCheckpo
 
         self.need_allclose = False
 
-    def runfrist(self, train_args):
+    def runfirst(self, train_args):
         train_args["fp16_opt_level"] = "O1"
         self.run_n1c8(self.run_pretrain_file, **train_args)
 
@@ -584,6 +586,7 @@ class TestUnifiedCheckpointOnN1C8MasterWeightCompatibleO1ToO2(TestUnifiedCheckpo
         self.run_n1c8(self.run_pretrain_file, **train_args)
 
 
+@pytest.mark.skipif(True, reason="Skip for None CE")
 class TestUnifiedCheckpointOnN1C8MasterWeightCompatibleO2ToO1(TestUnifiedCheckpointBase):
     def setUp(self):
         super().setUp()
@@ -595,7 +598,7 @@ class TestUnifiedCheckpointOnN1C8MasterWeightCompatibleO2ToO1(TestUnifiedCheckpo
 
         self.need_allclose = False
 
-    def runfrist(self, train_args):
+    def runfirst(self, train_args):
         train_args["fp16_opt_level"] = "O2"
         self.run_n1c8(self.run_pretrain_file, **train_args)
 
@@ -604,6 +607,7 @@ class TestUnifiedCheckpointOnN1C8MasterWeightCompatibleO2ToO1(TestUnifiedCheckpo
         self.run_n1c8(self.run_pretrain_file, **train_args)
 
 
+@pytest.mark.skipif(True, reason="Skip for None CE")
 class TestUnifiedCheckpointOnN1C8CheckpointCompatible(TestUnifiedCheckpointBase):
     def setUp(self):
         super().setUp()
@@ -611,7 +615,7 @@ class TestUnifiedCheckpointOnN1C8CheckpointCompatible(TestUnifiedCheckpointBase)
         self.need_allclose = True
         self.rtol = 1e-7
 
-    def runfrist(self, train_args):
+    def runfirst(self, train_args):
         train_args["unified_checkpoint"] = 0
         self.run_n1c8(self.run_pretrain_file, **train_args)
 
@@ -620,6 +624,7 @@ class TestUnifiedCheckpointOnN1C8CheckpointCompatible(TestUnifiedCheckpointBase)
         self.run_n1c8(self.run_pretrain_file, **train_args)
 
 
+@pytest.mark.skipif(True, reason="Skip for None CE")
 class TestPaddleCheckpointOnN1C8Reset(TestUnifiedCheckpointBase):
     def setUp(self):
         super().setUp()
@@ -627,7 +632,7 @@ class TestPaddleCheckpointOnN1C8Reset(TestUnifiedCheckpointBase):
         self.need_allclose = True
         self.rtol = 1e-7
 
-    def runfrist(self, train_args):
+    def runfirst(self, train_args):
         train_args["unified_checkpoint"] = 0
         self.run_n1c8(self.run_pretrain_file, **train_args)
 
@@ -636,6 +641,7 @@ class TestPaddleCheckpointOnN1C8Reset(TestUnifiedCheckpointBase):
         self.run_n1c8(self.run_pretrain_file, **train_args)
 
 
+@pytest.mark.skipif(True, reason="Skip for None CE")
 class TestPaddleCheckpointOnN1C2Reset(TestMultipleGpus):
     def setUp(self):
         self.configs = get_pretrain_arguments(pretrain_arguments)
@@ -652,7 +658,7 @@ class TestPaddleCheckpointOnN1C2Reset(TestMultipleGpus):
 
         self.run_pretrain_file = "llm/llama/run_pretrain.py"
 
-    def runfrist(self, train_args):
+    def runfirst(self, train_args):
         train_args["unified_checkpoint"] = 0
         self.run_n1c2(self.run_pretrain_file, **train_args)
 
@@ -668,7 +674,7 @@ class TestPaddleCheckpointOnN1C2Reset(TestMultipleGpus):
 
         train_args = self.configs["TP2"]
 
-        self.runfrist(train_args)
+        self.runfirst(train_args)
         self.rerun(train_args)
 
         if self.need_allclose:
@@ -677,6 +683,7 @@ class TestPaddleCheckpointOnN1C2Reset(TestMultipleGpus):
             np.testing.assert_allclose(res[0], res[1], self.rtol)
 
 
+@pytest.mark.skipif(True, reason="Skip for None CE")
 class TestUnifiedCheckpointOnN1C2Reset(TestMultipleGpus):
     def setUp(self):
         self.configs = get_pretrain_arguments(pretrain_arguments)
@@ -713,7 +720,7 @@ class TestUnifiedCheckpointOnN1C2Reset(TestMultipleGpus):
             "training_args.bin",
         ]
 
-    def runfrist(self, train_args):
+    def runfirst(self, train_args):
         train_args["unified_checkpoint"] = 1
         self.run_n1c2(self.run_pretrain_file, **train_args)
 
@@ -729,7 +736,7 @@ class TestUnifiedCheckpointOnN1C2Reset(TestMultipleGpus):
 
         train_args = self.configs["TP2"]
 
-        self.runfrist(train_args)
+        self.runfirst(train_args)
         self.rerun(train_args)
 
         if self.need_allclose:
@@ -747,7 +754,7 @@ class TestUnifiedCheckpointOnN1C2Reset(TestMultipleGpus):
         base_ckpt_path = os.path.join(pretrain_arguments["output_dir"], "checkpoint-%d" % save_steps)
 
         train_args = self.configs["TP2"]
-        self.runfrist(train_args)
+        self.runfirst(train_args)
         assert sorted(self.filelists) == sorted(os.listdir(base_ckpt_path))
         self.rerun(train_args)
 
@@ -760,7 +767,7 @@ class TestUnifiedCheckpointOnN1C2Reset(TestMultipleGpus):
         remove_logs()
         remove_ckpt(pretrain_arguments["output_dir"])
         train_args["unified_checkpoint_config"] = "skip_save_model_weight"
-        self.runfrist(train_args)
+        self.runfirst(train_args)
         unsave_filelists = [
             "master_weights-00001-of-00002.safetensors",
             "master_weights-00002-of-00002.safetensors",
@@ -787,7 +794,7 @@ class TestUnifiedCheckpointOnN1C8AsyncSaveToDisk(TestUnifiedCheckpointBase):
         self.need_allclose = True
         self.rtol = 1e-7
 
-    def runfrist(self, train_args):
+    def runfirst(self, train_args):
         self.run_n1c8(self.run_pretrain_file, **train_args)
 
     def rerun(self, train_args):
@@ -808,7 +815,7 @@ class TestUnifiedCheckpointOnN2C4SkipSaveModelWeight(TestUnifiedCheckpointBase):
         self.need_allclose = True
         self.rtol = 1e-7
 
-    def runfrist(self, train_args):
+    def runfirst(self, train_args):
         self.run_n2c4(self.run_pretrain_file, **train_args)
 
     def rerun(self, train_args):
@@ -827,7 +834,7 @@ class TestUnifiedCheckpointOnN2C4MasterWeightCompatibleO1ToO2(TestUnifiedCheckpo
 
         self.need_allclose = False
 
-    def runfrist(self, train_args):
+    def runfirst(self, train_args):
         train_args["fp16_opt_level"] = "O1"
         self.run_n2c4(self.run_pretrain_file, **train_args)
 
@@ -848,7 +855,7 @@ class TestUnifiedCheckpointOnN2C4MasterWeightCompatibleO2ToO1(TestUnifiedCheckpo
 
         self.need_allclose = False
 
-    def runfrist(self, train_args):
+    def runfirst(self, train_args):
         train_args["fp16_opt_level"] = "O2"
         self.run_n2c4(self.run_pretrain_file, **train_args)
 
@@ -865,7 +872,7 @@ class TestUnifiedCheckpointOnN2C4CheckpointCompatible(TestUnifiedCheckpointBase)
         self.need_allclose = True
         self.rtol = 1e-7
 
-    def runfrist(self, train_args):
+    def runfirst(self, train_args):
         train_args["unified_checkpoint"] = 0
         self.run_n2c4(self.run_pretrain_file, **train_args)
 
@@ -885,7 +892,7 @@ class TestUnifiedCheckpointOnN2C4AsyncSaveToDisk(TestUnifiedCheckpointBase):
         self.need_allclose = True
         self.rtol = 1e-7
 
-    def runfrist(self, train_args):
+    def runfirst(self, train_args):
         self.run_n2c4(self.run_pretrain_file, **train_args)
 
     def rerun(self, train_args):
@@ -908,7 +915,7 @@ class TestUnifiedCheckpointOnN1C8ToN2C4SkipSaveModelWeight(TestUnifiedCheckpoint
         self.rtol = 1e-4
         self.k = MAX_CONVERT_CONFIGS  # max: 16, min: 1
 
-    def runfrist(self, train_args):
+    def runfirst(self, train_args):
         self.run_n1c8(self.run_pretrain_file, **train_args)
         move_checkpoint_N1C8_to_N2C4()
 
@@ -936,7 +943,7 @@ class TestUnifiedCheckpointOnN1C8ToN2C4MasterWeightCompatibleO1ToO2(TestUnifiedC
         self.rtol = 1e-4
         self.k = MAX_CONVERT_CONFIGS  # max: 16, min: 1
 
-    def runfrist(self, train_args):
+    def runfirst(self, train_args):
         train_args["fp16_opt_level"] = "O1"
         self.run_n1c8(self.run_pretrain_file, **train_args)
         move_checkpoint_N1C8_to_N2C4()
@@ -966,7 +973,7 @@ class TestUnifiedCheckpointOnN1C8ToN2C4MasterWeightCompatibleO2ToO1(TestUnifiedC
         self.rtol = 1e-4
         self.k = MAX_CONVERT_CONFIGS  # max: 16, min: 1
 
-    def runfrist(self, train_args):
+    def runfirst(self, train_args):
         train_args["fp16_opt_level"] = "O2"
         self.run_n1c8(self.run_pretrain_file, **train_args)
         move_checkpoint_N1C8_to_N2C4()
@@ -994,7 +1001,7 @@ class TestUnifiedCheckpointOnN1C8ToN2C4AsyncSaveToDisk(TestUnifiedCheckpointBase
         self.rtol = 1e-4
         self.k = MAX_CONVERT_CONFIGS  # max: 16, min: 1
 
-    def runfrist(self, train_args):
+    def runfirst(self, train_args):
         self.run_n1c8(self.run_pretrain_file, **train_args)
         move_checkpoint_N1C8_to_N2C4()
 
@@ -1022,7 +1029,7 @@ class TestUnifiedCheckpointOnN2C4ToN1C8SkipSaveModelWeight(TestUnifiedCheckpoint
         self.rtol = 1e-4
         self.k = MAX_CONVERT_CONFIGS  # max: 16, min: 1
 
-    def runfrist(self, train_args):
+    def runfirst(self, train_args):
         self.run_n2c4(self.run_pretrain_file, **train_args)
         move_checkpoint_N2C4_to_N1C8()
 
@@ -1050,7 +1057,7 @@ class TestUnifiedCheckpointOnN2C4ToN1C8MasterWeightCompatibleO1ToO2(TestUnifiedC
         self.rtol = 1e-4
         self.k = MAX_CONVERT_CONFIGS  # max: 16, min: 1
 
-    def runfrist(self, train_args):
+    def runfirst(self, train_args):
         train_args["fp16_opt_level"] = "O1"
         self.run_n2c4(self.run_pretrain_file, **train_args)
         move_checkpoint_N2C4_to_N1C8()
@@ -1080,7 +1087,7 @@ class TestUnifiedCheckpointOnN2C4ToN1C8MasterWeightCompatibleO2ToO1(TestUnifiedC
         self.rtol = 1e-4
         self.k = MAX_CONVERT_CONFIGS  # max: 16, min: 1
 
-    def runfrist(self, train_args):
+    def runfirst(self, train_args):
         train_args["fp16_opt_level"] = "O2"
         self.run_n2c4(self.run_pretrain_file, **train_args)
         move_checkpoint_N2C4_to_N1C8()
@@ -1108,7 +1115,7 @@ class TestUnifiedCheckpointOnN2C4ToN1C8AsyncSaveToDisk(TestUnifiedCheckpointBase
         self.rtol = 1e-4
         self.k = MAX_CONVERT_CONFIGS  # max: 16, min: 1
 
-    def runfrist(self, train_args):
+    def runfirst(self, train_args):
         self.run_n2c4(self.run_pretrain_file, **train_args)
         move_checkpoint_N2C4_to_N1C8()
 
@@ -1122,6 +1129,7 @@ class TestUnifiedCheckpointOnN2C4ToN1C8AsyncSaveToDisk(TestUnifiedCheckpointBase
             np.testing.assert_allclose(res[0], res[-1], rtol=self.rtol)
 
 
+@pytest.mark.skipif(True, reason="Skip for None CE")
 class TestUnifiedCheckpointOnN1C8EnableAll(TestUnifiedCheckpointBase):
     def setUp(self):
         super().setUp()
@@ -1132,8 +1140,47 @@ class TestUnifiedCheckpointOnN1C8EnableAll(TestUnifiedCheckpointBase):
         self.need_allclose = True
         self.rtol = 1e-7
 
-    def runfrist(self, train_args):
+    def runfirst(self, train_args):
         self.run_n1c8(self.run_pretrain_file, **train_args)
 
     def rerun(self, train_args):
         self.run_n1c8(self.run_pretrain_file, **train_args)
+
+
+@pytest.mark.skipif(True, reason="Skip for None CE")
+class TestUnifiedCheckpointOnN1C8SaveLoadSpeed(TestUnifiedCheckpointFull):
+    def setUp(self):
+        super().setUp()
+        for config_key in self.configs:
+            self.configs[config_key]["skip_profile_timer"] = 0
+            self.configs[config_key]["unified_checkpoint"] = 1
+            self.configs[config_key]["save_steps"] = 6
+            self.configs[config_key]["unified_checkpoint_config"] = "skip_save_model_weight master_weight_compatible"
+
+        self.need_allclose = False
+        self.rtol = 1e-7
+
+    def runfirst(self, train_args):
+        self.run_n1c8(self.run_pretrain_file, log_dir="log_uc", **train_args)
+
+    def rerun(self, train_args):
+        self.run_n1c8(self.run_pretrain_file, log_dir="log_uc", **train_args)
+
+
+@pytest.mark.skipif(True, reason="Skip for None CE")
+class TestPaddleCheckpointOnN1C8SaveLoadSpeed(TestUnifiedCheckpointFull):
+    def setUp(self):
+        super().setUp()
+        for config_key in self.configs:
+            self.configs[config_key]["skip_profile_timer"] = 0
+            self.configs[config_key]["unified_checkpoint"] = 0
+            self.configs[config_key]["save_steps"] = 6
+
+        self.need_allclose = False
+        self.rtol = 1e-7
+
+    def runfirst(self, train_args):
+        self.run_n1c8(self.run_pretrain_file, log_dir="log_pd", **train_args)
+
+    def rerun(self, train_args):
+        self.run_n1c8(self.run_pretrain_file, log_dir="log_pd", **train_args)

tests/transformers/clap/test_feature_extraction.py

@@ -68,6 +68,7 @@ class ClapFeatureExtractionTester(unittest.TestCase):
         self.feature_size = feature_size
         self.chunk_length = chunk_length
         self.hop_length = hop_length
+        super().__init__()
 
     def prepare_feat_extract_dict(self):
         return {

tests/transformers/ernie_vil/test_image_processing.py

@@ -58,6 +58,7 @@ class ErnieViLImageProcessingTester(unittest.TestCase):
         self.image_mean = image_mean
         self.image_std = image_std
         self.do_convert_rgb = do_convert_rgb
+        super().__init__()
 
     def prepare_image_processor_dict(self):
         return {

tests/transformers/speecht5/test_feature_extraction.py

@@ -81,6 +81,7 @@ class SpeechT5FeatureExtractionTester(unittest.TestCase):
         self.fmax = fmax
         self.mel_floor = mel_floor
         self.return_attention_mask = return_attention_mask
+        super().__init__()
 
     def prepare_feat_extract_dict(self):
         return {

tests/transformers/test_safetensors.py

@@ -0,0 +1,69 @@
+# 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 os
+import tempfile
+import unittest
+
+import numpy as np
+
+# from safetensors import safe_open
+from safetensors.numpy import load_file, save_file
+
+from paddlenlp.utils.safetensors import fast_load_file, fast_safe_open
+
+
+class FastSafetensors(unittest.TestCase):
+    def setUp(self):
+        super().setUp()
+        self.weigth_map = {}
+        tensors = [
+            ([10, 1, 10], "float32"),
+            ([1, 1, 10], "float32"),
+            ([1, 1, 1, 10], "float32"),
+            ([10, 10], "float32"),
+            ([8], "float16"),
+            ([5, 5, 5], "int32"),
+        ]
+        count = 0
+        for shape, dtype in tensors:
+            self.weigth_map[f"weight_{count}"] = (np.random.random(shape) * 100).astype(dtype)
+            count += 1
+        print(self.weigth_map)
+
+    def test_load_file(self):
+        with tempfile.TemporaryDirectory() as tmpdirname:
+            path = os.path.join(tmpdirname, "test.safetensors")
+            save_file(self.weigth_map, path, metadata={"format": "np"})
+            sf_load = load_file(path)
+            fs_sf_load = fast_load_file(path)
+            for k, v in self.weigth_map.items():
+                np.testing.assert_equal(v, sf_load[k])
+                np.testing.assert_equal(v, fs_sf_load[k])
+
+    def test_safe_open(self):
+        with tempfile.TemporaryDirectory() as tmpdirname:
+            path = os.path.join(tmpdirname, "test.safetensors")
+            save_file(self.weigth_map, path, metadata={"format": "np"})
+
+            with fast_safe_open(path, framework="np") as f:
+                for key in f.keys():
+                    safe_slice = f.get_slice(key)
+                    # np.testing.assert_equal(self.weigth_map[key][2:1, ...], safe_slice[2:1, ...])
+                    np.testing.assert_equal(self.weigth_map[key][0, ...], safe_slice[0, ...])
+                    np.testing.assert_equal(self.weigth_map[key][0:1, ...], safe_slice[0:1, ...])
+                    np.testing.assert_equal(self.weigth_map[key][..., 2:], safe_slice[..., 2:])
+                    np.testing.assert_equal(self.weigth_map[key][..., 1], safe_slice[..., 1])
+                    np.testing.assert_equal(self.weigth_map[key][:2, ...], safe_slice[:2, ...])
+                    np.testing.assert_equal(self.weigth_map[key][..., :4], safe_slice[..., :4])