mkdir unified_checkpoint directory

Author: DesmonDay

paddlenlp/trainer/trainer.py

@@ -598,7 +598,6 @@ def _load_from_checkpoint(self, resume_from_checkpoint=None):
                 if use_unified_checkpoint:
                     self.unified_checkpoint_handler.load_unified_checkpoint(
                         self.model,
-                        self.optimizer,
                         resume_from_checkpoint,
                     )
                     logger.info(f"Loading model from {resume_from_checkpoint} using unified checkpoint.")
@@ -1241,7 +1240,6 @@ def fused_allreduce_gradients_no_sync(paramlist, hcg):
                 if self.args.unified_checkpoint:
                     self.unified_checkpoint_handler.load_unified_checkpoint(
                         self.model,
-                        self.optimizer,
                         self.state.best_model_checkpoint,
                     )
                     if self.args.sharding_parallel_degree > 1 or self.args.data_parallel_degree > 1:
@@ -1289,7 +1287,6 @@ def _load_best_model_from_peft_checkpoint(self):
         if self.args.unified_checkpoint:
             self.unified_checkpoint_handler.load_unified_checkpoint(
                 self.model,
-                self.optimizer,
                 self.state.best_model_checkpoint,
             )
             if self.args.sharding_parallel_degree > 1 or self.args.data_parallel_degree > 1:

paddlenlp/trainer/plugins/shared_memory_utils.py renamed to paddlenlp/trainer/unified_checkpoint/shared_memory_utils.py

@@ -33,8 +33,7 @@
 
 from paddlenlp.peft import LoRAModel, PrefixModelForCausalLM
 from paddlenlp.trainer.argparser import strtobool
-from paddlenlp.trainer.trainer_utils import ShardingOption
-from paddlenlp.trainer.utils.helper import distributed_file, distributed_isfile
+from paddlenlp.trainer.utils.helper import distributed_isfile
 from paddlenlp.transformers.model_utils import (
     PretrainedModel,
     _add_variant,
@@ -67,7 +66,7 @@
     SAFE_WEIGHTS_NAME,
 )
 from paddlenlp.utils.log import logger
-from paddlenlp.utils.nested import flatten_list, nested_copy
+from paddlenlp.utils.nested import nested_copy
 
 if is_safetensors_available():
     from safetensors.numpy import save_file as safe_save_file
@@ -77,6 +76,7 @@
     else:
         from paddlenlp.utils.safetensors import fast_load_file as load_file
 
+from .check_unified_checkpoint import check_unified_checkpoint, check_unified_optimizer
 from .shared_memory_utils import (
     _read_state_dict_from_shm,
     _traverse_copy_to_shm,
@@ -108,13 +108,13 @@
     get_sharded_file_name,
     get_sharded_index,
     is_need_master_weight,
+    is_sharding_split_param_mode,
     mapping_optimizer_tp_actions,
     merge_tensor_parallel_for_optimizer,
     merge_tensor_parallel_with_shard,
     reduce_master_weights_status,
     rename_shard_file,
-    save_config,
-    save_prefix_past_key_value,
+    save_model_config,
     select_model_weight_index,
     update_master_weight_status,
 )
@@ -361,25 +361,8 @@ def save_unified_checkpoint(self, model, optimizer, output_dir, signal_dir=None)
                         json.dump(sharded_index, f, indent=4)
 
         if self.args.should_save:
-            # Save prefix model past_key_values
-            if isinstance(model_to_save, PrefixModelForCausalLM):
-                save_prefix_past_key_value(model_to_save, save_directory)
-                model_to_save.prefix_config.save_pretrained(save_directory)
-            if isinstance(model_to_save, LoRAModel):
-                model_to_save.lora_config.save_pretrained(save_directory)
-
-        # save the config
-        config_to_save = save_config(model_to_save)
-        # Attach architecture to the config
-        if isinstance(model_to_save, LoRAModel) or isinstance(model_to_save, PrefixModelForCausalLM):
-            config_to_save.architectures = [model_to_save.model.__class__.__name__]
-        else:
-            config_to_save.architectures = [model_to_save.__class__.__name__]
-        if self.args.should_save:
-            config_to_save.save_pretrained(save_directory)
-            # save generation config
-            if model_to_save.can_generate():
-                model_to_save.generation_config.save_pretrained(save_directory)
+            save_model_config(model_to_save, save_directory)
+
         paddle.device.cuda.empty_cache()
 
         if strtobool(os.getenv("FLAG_LLM_PDC", "False")) and self.args.should_save:
@@ -391,7 +374,7 @@ def save_unified_checkpoint(self, model, optimizer, output_dir, signal_dir=None)
             }
             paddle.save(save_info, os.path.join(save_directory, ".saving_info"))
 
-    def load_unified_checkpoint(self, model, optimizer, resume_from_checkpoint: str):
+    def load_unified_checkpoint(self, model, resume_from_checkpoint: str):
         """Load potential model checkpoint
 
         Args:
@@ -539,11 +522,7 @@ def save_unified_optimizer(self, model, optimizer, output_dir, signal_dir):
             save_single_card_optimizer(model, optimizer, output_dir)  # no need to save signal
             return
 
-        if (
-            self.args.sharding_parallel_degree > 1
-            and ShardingOption.SHARD_OP in self.args.sharding
-            and "split_param" in self.args.sharding_parallel_config
-        ):
+        if is_sharding_split_param_mode(self.args):
             optim_state_dict, master_weights = gather_splited_param_for_optimizer(optimizer)
         else:
             optim_state_dict = nested_copy(optimizer.state_dict())
@@ -867,11 +846,7 @@ def unified_checkpoint_into_shards(
 
 def load_unified_optimizer_locally(args, model, optimizer, resume_from_checkpoint, safe_serialization=False):
     # Special process with split param.
-    if (
-        args.sharding_parallel_degree > 1
-        and ShardingOption.SHARD_OP in args.sharding
-        and "split_param" in args.sharding_parallel_config
-    ):
+    if is_sharding_split_param_mode(args):
         returned_optim_state_dict = load_unified_optimizer_split_param(model, optimizer, resume_from_checkpoint)
         return returned_optim_state_dict
 
@@ -1118,211 +1093,3 @@ def unified_optimizer_into_shards(
             (optim_state_dict, shard_optimizer_file, sharded_optim_index),
             (master_weights, shard_master_weight_file, sharded_master_weight_index),
         ]
-
-
-def check_unified_checkpoint(args, model, resume_from_checkpoint, safe_serialization=False):
-    index_filename = select_model_weight_index(model, resume_from_checkpoint, safe_serialization, local=False)
-    index_filename = os.path.join(resume_from_checkpoint, index_filename)
-    # Find index json file and distribute this file in global group.
-    if distributed_isfile(index_filename):
-        distributed_file(index_filename)
-    else:
-        raise Exception(
-            f"Sorry, we can not find {index_filename}. This file should be appear at least on one machine."
-        )
-
-    with open(index_filename, "r") as f:
-        index = json.loads(f.read())
-    all_weight_filenames = sorted(set(index["weight_map"].values()))
-
-    # Get existed weight file list on current machine.
-    existed_filelist = []
-    existed_files = []
-    for filename in os.listdir(resume_from_checkpoint):
-        if filename in all_weight_filenames:
-            existed_files.append(filename)
-
-    # Gather all the existed files in global group.
-    dist.all_gather_object(existed_filelist, existed_files)
-    flatten_existed_filelist = flatten_list(existed_filelist)
-    diff_filelist = list(set(all_weight_filenames).difference(set(flatten_existed_filelist)))
-    if len(diff_filelist) != 0:
-        raise Exception(f"Sorry, the weight file list on the machines is not complete!, missing {diff_filelist}")
-
-    # To decide whether to load the checkpoint locally, or need to dynamically send tensors across machines.
-    local_resume = True
-    if args.dataset_rank == 0 or args.use_expert_parallel:
-        hcg = fleet.get_hybrid_communicate_group()
-        tp_group = hcg.get_model_parallel_group()
-        pp_group = hcg.get_pipe_parallel_group()
-        dp_group = hcg.get_data_parallel_group()
-        dp_rank = dp_group.rank if dp_group.nranks > 1 else 0
-
-        need_files = set()
-        state_dict = get_expected_state_dict(model)
-        for key in state_dict.keys():
-            filename = index["weight_map"][key]
-            # When using expert parallel, there's no need to check tensors with `no_sync=False` when dp_rank > 0.
-            if args.use_expert_parallel and dp_rank > 0 and not getattr(state_dict[key], "no_sync", False):
-                continue
-            need_files.add(filename)
-        diff_filelist = list(need_files.difference(set(existed_files)))
-        num_diff = paddle.to_tensor([len(diff_filelist)])
-        if tp_group.nranks > 1:
-            dist.all_reduce(num_diff, op=dist.ReduceOp.MAX, group=tp_group)
-        if pp_group.nranks > 1:
-            dist.all_reduce(num_diff, op=dist.ReduceOp.MAX, group=pp_group)
-        if args.use_expert_parallel and dp_group.nranks > 1:
-            dist.all_reduce(num_diff, op=dist.ReduceOp.MAX, group=dp_group)
-        if num_diff.item() == 0:
-            local_resume = True
-        else:
-            local_resume = False
-    local_resume = paddle.to_tensor([local_resume])
-    dist.all_reduce(local_resume, op=dist.ReduceOp.PROD)
-    local_resume = local_resume.item()
-    return local_resume
-
-
-def check_unified_optimizer(args, model, optimizer, resume_from_checkpoint, safe_serialization=False):
-    if not safe_serialization:
-        index_filename, index_filename_master_weights = PADDLE_OPTIMIZER_INDEX_NAME, PADDLE_MASTER_WEIGHTS_INDEX_NAME
-    else:
-        index_filename, index_filename_master_weights = SAFE_OPTIMIZER_INDEX_NAME, SAFE_MASTER_WEIGHTS_INDEX_NAME
-    index_filename = os.path.join(resume_from_checkpoint, index_filename)
-    index_filename_master_weights = os.path.join(resume_from_checkpoint, index_filename_master_weights)
-
-    # Find index json file and distribute the file in global group.
-    if distributed_isfile(index_filename):
-        distributed_file(index_filename)
-    else:
-        raise Exception(
-            f"Sorry, we can not find {index_filename}. This file should be appear at least on one machine."
-        )
-
-    with open(index_filename, "r") as f:
-        index = json.loads(f.read())
-    all_optimizer_filenames = sorted(set(index["weight_map"].values()))
-
-    has_master_weights = index["master_weights"]
-    # update has_master_weights and index_filename_master_weights
-    # 1. if the master weight exists, only has_master_weights is set True and loaded when needed
-    # 2. if master weight does not exist, convert model weight to master weight when needed
-    has_master_weights, index_filename_master_weights = update_master_weight_status(
-        args, optimizer, has_master_weights, safe_serialization
-    )
-    if has_master_weights:
-        index_filename_master_weights = os.path.join(resume_from_checkpoint, index_filename_master_weights)
-        if distributed_isfile(index_filename_master_weights):
-            distributed_file(index_filename_master_weights)
-        else:
-            raise Exception(
-                f"Sorry, we can not find {index_filename_master_weights}. This file should be appear at least on one machine."
-            )
-        with open(index_filename_master_weights, "r") as f:
-            index_mw = json.loads(f.read())
-        all_mw_filenames = sorted(set(index_mw["weight_map"].values()))
-
-    hcg = fleet.get_hybrid_communicate_group()
-    tp_group = hcg.get_model_parallel_group()
-    pp_group = hcg.get_pipe_parallel_group()
-    dp_group = hcg.get_data_parallel_group()
-    sharding_group = hcg.get_sharding_parallel_group()
-    sharding_rank = sharding_group.rank
-    dp_rank = dp_group.rank if dp_group.nranks > 1 else 0
-    struct2static_name_mappings = {k: v.name for k, v in model.state_dict().items()}
-
-    if (
-        args.sharding_parallel_degree > 1
-        and ShardingOption.SHARD_OP in args.sharding
-        and "split_param" in args.sharding_parallel_config
-    ):
-        # We do not check optimizer files completion for split_param, since it is very complicated. Directly support local resume.
-        logger.warning("We only support local resume for split_param mode, do not support dynamically loading.")
-        return True
-
-    if sharding_group.nranks > 1:
-        param2rank = optimizer._param2rank
-
-    def check_complete(all_filenames):
-        # Check whether the checkpoint files on machines are complete. If not complete, raise Exception.
-        existed_filelist = []
-        existed_files = []
-        for filename in os.listdir(resume_from_checkpoint):
-            if filename in all_filenames:
-                existed_files.append(filename)
-
-        dist.all_gather_object(existed_filelist, existed_files)
-        flatten_existed_filelist = flatten_list(existed_filelist)
-        diff_filelist = list(set(all_filenames).difference(set(flatten_existed_filelist)))
-        if len(diff_filelist) != 0:
-            raise Exception(
-                f"Sorry, the optimizer file list on `data_parallel_rank==0` machines is not complete!, missing {diff_filelist}"
-            )
-        return existed_files
-
-    def check_dynamic_load(args, weight_map, existed_files, is_master_weights=False, typename_set=None):
-        # To decide whether to load the checkpoint locally, or need to dynamically distribute the checkpoint.
-        local_resume = True
-        if args.data_parallel_rank == 0 or args.use_expert_parallel:
-            need_files = set()
-            state_dict = get_expected_state_dict(model)
-
-            for key in state_dict.keys():
-                if sharding_group.nranks > 1:
-                    static_name = struct2static_name_mappings.get(key, None)
-                    param_rank = param2rank.get(static_name, None)
-                    if param_rank != sharding_rank:
-                        continue
-
-                # When using expert parallel, there's no need to check tensors with `no_sync=False` when dp_rank > 0.
-                if args.use_expert_parallel and dp_rank > 0 and not getattr(state_dict[key], "no_sync", False):
-                    continue
-
-                if is_master_weights and state_dict[key].dtype == core.VarDesc.VarType.FP32:
-                    continue
-
-                if not is_master_weights:
-                    for type_name in typename_set:
-                        type_key = key + "/" + type_name
-                        filename = weight_map[type_key]
-                        need_files.add(filename)
-                else:
-                    filename = weight_map[key]
-                    need_files.add(filename)
-
-            diff_filelist = list(need_files.difference(set(existed_files)))
-            num_diff = paddle.to_tensor([len(diff_filelist)])
-            if tp_group.nranks > 1:
-                dist.all_reduce(num_diff, op=dist.ReduceOp.MAX, group=tp_group)
-            if pp_group.nranks > 1:
-                dist.all_reduce(num_diff, op=dist.ReduceOp.MAX, group=pp_group)
-            if sharding_group.nranks > 1:
-                dist.all_reduce(num_diff, op=dist.ReduceOp.MAX, group=sharding_group)
-            if args.use_expert_parallel and dp_group.nranks > 1:
-                dist.all_reduce(num_diff, op=dist.ReduceOp.MAX, group=dp_group)
-
-            if num_diff.item() == 0:
-                local_resume = True
-            else:
-                local_resume = False
-        local_resume = paddle.to_tensor([local_resume])
-        dist.all_reduce(local_resume, op=dist.ReduceOp.PROD)
-        return local_resume.item()
-
-    # check whether the optimizer checkpoint files are complete.
-    existed_files = check_complete(all_optimizer_filenames)
-    if has_master_weights:
-        existed_files_mw = check_complete(all_mw_filenames)
-    # get optimizer's param type name, like moment1_0.
-    typename_set = set()
-    for key in index["weight_map"].keys():
-        _, typename = key.split("/")
-        typename_set.add(typename)
-    local_resume = check_dynamic_load(
-        args, index["weight_map"], existed_files, is_master_weights=False, typename_set=typename_set
-    )
-    local_resume_rw = True
-    if has_master_weights:
-        local_resume_rw = check_dynamic_load(args, index_mw["weight_map"], existed_files_mw, is_master_weights=True)
-    return local_resume & local_resume_rw

paddlenlp/trainer/plugins/unified_checkpoint.py renamed to paddlenlp/trainer/unified_checkpoint/unified_checkpoint.py

@@ -38,9 +38,10 @@
 )
 
 
-def distributed_send_recv_splited_param(
+def merge_splited_param(
     state_dict, partial_tensor_list, param_shape_info, send_table, recv_table, is_master_weights=False
 ):
+    """Merge the splited param in sharding group."""
     global_rank = dist.get_rank()
     for key in list(state_dict.keys()):
         if state_dict[key].numel().item() == 1:  # for example: beta1, beta2
@@ -144,13 +145,9 @@ def gather_splited_param_for_optimizer(optimizer):
         recv_table[key] = sharding_ranklist[0][0]  # which sharding_rank to recv the splited tensor
         send_table[key] = [(rank, begin, end) for rank, begin, end in sharding_ranklist]
 
-    distributed_send_recv_splited_param(
-        optim_state_dict, partial_tensor_list, param_shape_info, send_table, recv_table, False
-    )
+    merge_splited_param(optim_state_dict, partial_tensor_list, param_shape_info, send_table, recv_table, False)
     if master_weights is not None:
-        distributed_send_recv_splited_param(
-            master_weights, partial_tensor_list, param_shape_info, send_table, recv_table, True
-        )
+        merge_splited_param(master_weights, partial_tensor_list, param_shape_info, send_table, recv_table, True)
     return optim_state_dict, master_weights