add save/load_checkpoint_mode flag

Author: xingmingyyj

paddlenlp/trainer/trainer.py

@@ -223,6 +223,10 @@ def in_auto_parallel_align_mode():
 
 __all__ = ["Trainer"]
 
+FLEX_CKPT_MODEL_STATE_DIR_NAME = "model_state"
+FLEX_CKPT_OPT_STATE_DIR_NAME = "optimizer_states"
+FLEC_CKPT_MASTER_WEIGHTS_INDEX_NAME = "master_weights"
+
 
 class Trainer:
     """
@@ -929,13 +933,13 @@ def train(
         self._memory_tracker.start()
 
         if not self.args.enable_auto_parallel:
-            if not self.args.should_load_sharding_stage1_model and not self.args.using_flex_checkpoint:
+            if not self.args.should_load_sharding_stage1_model and not self.args.load_flex_checkpoint:
                 self._load_from_checkpoint(resume_from_checkpoint)
 
             if self.args.should_load_sharding_stage1_model:
                 model = self._wrap_model_and_load_sharded_checkpoint(resume_from_checkpoint)
 
-            elif self.args.should_save_sharding_stage1_model:
+            elif self.args.should_save_sharding_stage1_model and not self.args.load_flex_checkpoint:
                 # In the non-sharded mode, should invoke _load_from_checkpoint before _wrap_model.
                 # In this mode, the rank0 load all params and the _wrap_model implicitly broadcast params from rank0 to the other ranks.
                 model = self._wrap_model(self.model_wrapped)
@@ -949,36 +953,44 @@ def train(
                 if delay_optimizer_creation:
                     self.create_optimizer_and_scheduler(num_training_steps=max_steps)
                 self._load_optimizer_and_scheduler(resume_from_checkpoint)
-            elif not self.args.using_flex_checkpoint:
+
+            elif self.args.load_flex_checkpoint:
                 model = self._wrap_model(self.model_wrapped)
-                # for the rest of this function `model` is the outside model, whether it was wrapped or not
                 if model is not self.model:
                     self.model_wrapped = model
+
                 if delay_optimizer_creation:
                     self.create_optimizer_and_scheduler(num_training_steps=max_steps)
-                self._load_optimizer_and_scheduler(resume_from_checkpoint)
-            else:
-                assert self.args.using_flex_checkpoint, "default using flex_checkpoint!"
 
+                if resume_from_checkpoint is not None:
+                    if not self.args.ignore_load_lr_and_optim:
+                        model_sharded_state_dict = self.model.sharded_state_dict()
+                        accessible_files = os.listdir(resume_from_checkpoint)
+                        metadata_files = [file for file in accessible_files if file.endswith(".metadata")]
+                        assert len(metadata_files) == 1, "Only support one metadata file now."
+                        metadata = paddle.load(os.path.join(resume_from_checkpoint, metadata_files[0]))
+                        state_dict_metadata = metadata.state_dict_metadata
+                        init_optimizer(self.optimizer, model_sharded_state_dict, state_dict_metadata)
+                        optimizer_sharded_state_dict = self.optimizer.sharded_state_dict(model_sharded_state_dict)
+                        sharded_state_dict = {**model_sharded_state_dict, **optimizer_sharded_state_dict}
+                        dist.load_state_dict(
+                            sharded_state_dict, resume_from_checkpoint, aoa_config=self.args.aoa_config, offload=False
+                        )
+                        self._load_scheduler(resume_from_checkpoint)
+                    else:
+                        model_sharded_state_dict = self.model.sharded_state_dict()
+                        sharded_state_dict = model_sharded_state_dict
+                        dist.load_state_dict(
+                            sharded_state_dict, resume_from_checkpoint, aoa_config=self.args.aoa_config
+                        )
+            else:
                 model = self._wrap_model(self.model_wrapped)
+                # for the rest of this function `model` is the outside model, whether it was wrapped or not
                 if model is not self.model:
                     self.model_wrapped = model
-
                 if delay_optimizer_creation:
                     self.create_optimizer_and_scheduler(num_training_steps=max_steps)
-
-                if resume_from_checkpoint is not None:
-                    model_sharded_state_dict = self.model.sharded_state_dict()
-                    accessible_files = os.listdir(resume_from_checkpoint)
-                    metadata_files = [file for file in accessible_files if file.endswith(".metadata")]
-                    assert len(metadata_files) == 1, "Only support one metadata file now."
-                    metadata = paddle.load(os.path.join(resume_from_checkpoint, metadata_files[0]))
-                    state_dict_metadata = metadata.state_dict_metadata
-                    init_optimizer(self.optimizer, model_sharded_state_dict, state_dict_metadata)
-                    optimizer_sharded_state_dict = self.optimizer.sharded_state_dict(model_sharded_state_dict)
-                    sharded_state_dict = {**model_sharded_state_dict, **optimizer_sharded_state_dict}
-                    dist.load_state_dict(sharded_state_dict, resume_from_checkpoint, aoa_config=self.args.aoa_config)
-                    self._load_scheduler(resume_from_checkpoint)
+                self._load_optimizer_and_scheduler(resume_from_checkpoint)
         else:
             model = self.model_wrapped
             if delay_optimizer_creation:
@@ -2738,7 +2750,7 @@ def _save_checkpoint(self, model, metrics=None):
         else:
             self.save_model(output_dir)
 
-        if self.args.using_flex_checkpoint:
+        if self.args.save_flex_checkpoint:
             model_sharded_state_dict = self.model.sharded_state_dict()
             os.makedirs(output_dir, exist_ok=True)
 
@@ -2801,7 +2813,18 @@ def _save_checkpoint(self, model, metrics=None):
                             signal_dir,
                         )
                     else:
-                        if not self.args.using_flex_checkpoint:
+                        if self.args.save_flex_checkpoint:
+                            optimizer_sharded_state_dict = self.optimizer.sharded_state_dict(model_sharded_state_dict)
+                            dist.save_state_dict(
+                                {**model_sharded_state_dict, **optimizer_sharded_state_dict},
+                                output_dir,
+                            )
+                            if self.args.should_save:
+                                if self.tokenizer is not None and self.args.save_tokenizer:
+                                    self.tokenizer.save_pretrained(output_dir)
+                                # Good practice: save your training arguments together with the trained model
+                                paddle.save(self.args, os.path.join(output_dir, TRAINING_ARGS_NAME))
+                        else:
                             if self.dp_group.rank > 0:  # this should only work for MoE saving
                                 self._save_ckpt_func(
                                     self._filter_moe_no_sync_optimizer_params(),
@@ -2821,12 +2844,7 @@ def _save_checkpoint(self, model, metrics=None):
                                     )
                                 else:
                                     self._save_ckpt_func(state_dict, save_path, saved_signal_path)
-                        else:
-                            optimizer_sharded_state_dict = self.optimizer.sharded_state_dict(model_sharded_state_dict)
-                            dist.save_state_dict(
-                                {**model_sharded_state_dict, **optimizer_sharded_state_dict},
-                                output_dir,
-                            )
+
                 else:
                     if self.args.unified_checkpoint and "async_save" in self.args.unified_checkpoint_config:
                         global_rank = paddle.distributed.get_rank() if paddle.distributed.get_world_size() > 1 else -1
@@ -2852,7 +2870,18 @@ def _save_checkpoint(self, model, metrics=None):
                             output_dir,
                             signal_dir,
                         )
-                    elif not self.args.using_flex_checkpoint:
+                    elif self.args.save_flex_checkpoint:
+                        optimizer_sharded_state_dict = self.optimizer.sharded_state_dict(model_sharded_state_dict)
+                        dist.save_state_dict(
+                            {**model_sharded_state_dict, **optimizer_sharded_state_dict},
+                            output_dir,
+                        )
+                        if self.args.should_save:
+                            if self.tokenizer is not None and self.args.save_tokenizer:
+                                self.tokenizer.save_pretrained(output_dir)
+                            # Good practice: save your training arguments together with the trained model
+                            paddle.save(self.args, os.path.join(output_dir, TRAINING_ARGS_NAME))
+                    else:
                         if self.args.data_parallel_rank > 0 and self.args.use_expert_parallel:
                             self._save_ckpt_func(
                                 self._filter_moe_no_sync_optimizer_params(),
@@ -2866,13 +2895,6 @@ def _save_checkpoint(self, model, metrics=None):
                                 saved_signal_path,
                             )
 
-                    else:
-                        optimizer_sharded_state_dict = self.optimizer.sharded_state_dict(model_sharded_state_dict)
-                        dist.save_state_dict(
-                            {**model_sharded_state_dict, **optimizer_sharded_state_dict},
-                            output_dir,
-                        )
-
                 # FIXME: maybe only save one copy
                 paddle.save(self.lr_scheduler.state_dict(), os.path.join(output_dir, SCHEDULER_NAME))
 
@@ -2893,6 +2915,18 @@ def _save_checkpoint(self, model, metrics=None):
             if self.args.unified_checkpoint and (self.args.offload_optim or self.args.tensorwise_offload_optimizer):
                 self._offload_optimizer()
 
+        else:
+            if self.args.save_flex_checkpoint:
+                dist.save_state_dict(
+                    model_sharded_state_dict,
+                    output_dir,
+                )
+                if self.args.should_save:
+                    if self.tokenizer is not None and self.args.save_tokenizer:
+                        self.tokenizer.save_pretrained(output_dir)
+                    # Good practice: save your training arguments together with the trained model
+                    paddle.save(self.args, os.path.join(output_dir, TRAINING_ARGS_NAME))
+
         self.runtime_timer.stop()
 
         # Maybe delete some older checkpoints.
@@ -3107,6 +3141,7 @@ def _save(
         else:
             if isinstance(self.model, PretrainedModel) and self.args.should_save_sharding_stage1_model:
                 config_to_save = None
+                self.sharding_io.set_optimizer(self.optimizer)
                 state_dict, config_to_save, weight_name_suffix = self.sharding_io.manipulate_state_dict_and_config(
                     self.model, merge_tensor_parallel=merge_tensor_parallel
                 )

paddlenlp/trainer/training_args.py

@@ -407,10 +407,12 @@ class TrainingArguments:
             Whether to release gradients during training. Default is `False`.
         ckpt_quant_stage (`str`, *optional*):
             Whether activate checkpoint quantization. O0: deactivate, O1: Int8 compression, O2: Int4 compression. (default: O0).
-        using_flex_checkpoint(`bool`, *optional*):
-            Whether to use FlexCheckpoint for save and load. Default is False.
         aoa_config (`Optional[dict[str, list[str]]]`, *optional*):
             The AoA configuration of FlexCheckpoint, used to describe the mapping between model weights and the checkpoint content. Default is None.
+        save_checkpoint_mode (`str`, *optional*):
+            Specifies the method for saving checkpoints. Options are: None, 'sharding_io', 'unified_checkpoint', 'flex_checkpoint', and 'safetensor'. (default: None). This setting is ignored if the corresponding switch is configured.
+        load_checkpoint_mode (`str`, *optional*):
+            Specifies the method for loading checkpoints. Options are: None, 'sharding_io', 'unified_checkpoint', 'flex_checkpoint', and 'safetensor'. (default: None). This setting is ignored if the corresponding switch is configured.
     """
 
     output_dir: str = field(
@@ -935,10 +937,6 @@ class TrainingArguments:
         default=False,
         metadata={"help": "Whether to use async_save instead of paddle.save."},
     )
-    using_flex_checkpoint: Optional[bool] = field(
-        default=False,
-        metadata={"help": "Whether use FlexCheckpoint."},
-    )
     ordered_save_group_size: int = field(
         default=0,
         metadata={
@@ -1111,6 +1109,30 @@ class TrainingArguments:
         },
     )
 
+    save_checkpoint_mode: Optional[str] = field(
+        default=None,
+        metadata={
+            "help": (
+                "Specifies the method used to save checkpoints. "
+                "Available options: 'sharding_io', 'unified_checkpoint', "
+                "'flex_checkpoint', 'safetensor'."
+                "This setting is ignored if the corresponding switch is configured."
+            )
+        },
+    )
+
+    load_checkpoint_mode: Optional[str] = field(
+        default=None,
+        metadata={
+            "help": (
+                "Specifies the method used to load checkpoints. "
+                "Available options: 'sharding_io', 'unified_checkpoint', "
+                "'flex_checkpoint', 'safetensor'."
+                "This setting is ignored if the corresponding switch is configured."
+            )
+        },
+    )
+
     def __post_init__(self):
         world_size = paddle.distributed.get_world_size()
         if in_auto_parallel_align_mode():
@@ -1210,6 +1232,8 @@ def __post_init__(self):
             raise ValueError("AdamW Mini currently doesn't support tensor parallelism.")
 
         self._post_init_parallel_degree()
+        self._post_init_save_checkpoint_mode()
+        self._post_init_load_checkpoint_mode()
 
         if self.to_static:
             assert world_size == 1 or self.enable_auto_parallel, (
@@ -1862,7 +1886,7 @@ def is_context_parallel_supported():
                         # DP use hybrid group
                         strategy = fleet.DistributedStrategy()
                         fleet.init(is_collective=True, strategy=strategy)
-                    elif self.using_flex_checkpoint:
+                    elif self.save_flex_checkpoint or self.load_flex_checkpoint:
                         strategy = fleet.DistributedStrategy()
                         fleet.init(is_collective=True, strategy=strategy)
                     else:
@@ -2131,6 +2155,64 @@ def _post_init_parallel_degree(self):
         if self.use_hybrid_parallel and self.enable_auto_parallel:
             self.use_hybrid_parallel = False
 
+    def _post_init_save_checkpoint_mode(self):
+        if not self.save_checkpoint_mode:
+            return
+
+        # Ensure that only one checkpoint mode is set at a time
+        if self.unified_checkpoint or self.save_sharded_model:
+            return
+
+        self.save_flex_checkpoint = False
+
+        valid_modes = ["unified_checkpoint", "sharding_io", "safetensor", "flex_checkpoint"]
+        assert (
+            self.save_checkpoint_mode in valid_modes
+        ), f"Invalid save_checkpoint_mode: {self.save_checkpoint_mode}, Only these modes are allowed: {valid_modes}."
+
+        if self.save_checkpoint_mode == "safetensor":
+            raise NotImplementedError("safetensor checkpoint saving is not implemented yet.")
+        elif self.save_checkpoint_mode == "unified_checkpoint":
+            assert (
+                getattr(self, "load_checkpoint_mode", None) == "unified_checkpoint"
+            ), "When saving in unified_checkpoint mode, load_checkpoint_mode must also be 'unified_checkpoint'."
+            self.unified_checkpoint = True
+        elif self.save_checkpoint_mode == "sharding_io":
+            self.save_sharded_model = True
+        elif self.save_checkpoint_mode == "flex_checkpoint":
+            self.save_flex_checkpoint = True
+        else:
+            raise NotImplementedError(f"Checkpoint mode '{self.save_checkpoint_mode}' is not supported.")
+
+    def _post_init_load_checkpoint_mode(self):
+        if not self.load_checkpoint_mode:
+            return
+
+        self.load_flex_checkpoint = False
+
+        # Ensure that only one checkpoint mode is set at a time
+        if self.unified_checkpoint or self.load_sharded_model:
+            return
+
+        valid_modes = ["unified_checkpoint", "sharding_io", "safetensor", "flex_checkpoint"]
+        assert (
+            self.load_checkpoint_mode in valid_modes
+        ), f"Invalid load_checkpoint_mode: {self.load_checkpoint_mode}, Only these modes are allowed: {valid_modes}."
+
+        if self.load_checkpoint_mode == "safetensor":
+            raise NotImplementedError("safetensor checkpoint loading is not implemented yet.")
+        elif self.load_checkpoint_mode == "unified_checkpoint":
+            assert (
+                getattr(self, "save_checkpoint_mode", None) == "unified_checkpoint"
+            ), "When loading in unified_checkpoint mode, save_checkpoint_mode must also be 'unified_checkpoint'."
+            self.unified_checkpoint = True
+        elif self.load_checkpoint_mode == "sharding_io":
+            self.load_sharded_model = True
+        elif self.load_checkpoint_mode == "flex_checkpoint":
+            self.load_flex_checkpoint = True
+        else:
+            raise NotImplementedError(f"Checkpoint mode '{self.load_checkpoint_mode}' is not supported.")
+
     def add_moe_comm_group(self):
         hybrid_configs = fleet.fleet._user_defined_strategy.hybrid_configs
         hcg = fleet.get_hybrid_communicate_group()
@@ -2459,7 +2541,7 @@ def should_save_model_state(self):
                 return True
             elif self.enable_auto_parallel:
                 return True
-            elif self.using_flex_checkpoint:
+            elif self.save_flex_checkpoint:
                 return False
             elif self.use_hybrid_parallel:
                 # save on dataset rank 0