[AutoParallel] Refine auto_trainer save load (#8767)

Author: zhangbo9674

paddlenlp/trainer/auto_trainer.py

@@ -122,12 +122,10 @@ def _wrap_for_auto(self, model, train_dataloader):
         if self.args.to_static:
             unified_strategy = dist.Strategy()
             unified_strategy._from_legacy_strategy(self.args.strategy)
-            return (
-                dist.to_static(model, dist_loader, self.criterion, self.optimizer, strategy=unified_strategy),
-                dist_loader,
-            )
-        else:
-            return model, dist_loader
+            model = dist.to_static(model, dist_loader, self.criterion, self.optimizer, strategy=unified_strategy)
+
+        self.model_wrapped = model
+        return model, dist_loader
 
     def _wrap_amp_model(self, args, model):
         logger.info("Using half precision")
@@ -216,7 +214,6 @@ def _inner_training_loop(
             epochs_trained = self.state.global_step // num_update_steps_per_epoch
             if not args.ignore_data_skip:
                 steps_trained_in_current_epoch = self.state.global_step % (num_update_steps_per_epoch)
-                steps_trained_in_current_epoch *= args.gradient_accumulation_steps
             else:
                 steps_trained_in_current_epoch = 0
 
@@ -269,6 +266,9 @@ def _inner_training_loop(
         model, dist_loader = self._wrap_for_auto(model, train_dataloader)
         train_dataloader = dist_loader()
 
+        if resume_from_checkpoint is not None:
+            self._load_from_checkpoint(resume_from_checkpoint)
+
         self.timers and self.timers("read-data").start()
 
         for epoch in range(epochs_trained, num_train_epochs):
@@ -542,14 +542,26 @@ def _save_checkpoint(self, model, metrics=None):
             logger.info(f"Saving checkpoinit files into {output_dir}")
 
             if self.args.should_save_model_state:
-
-                optim_state_dict = self.optimizer.state_dict()
-                optim_state_dict.pop("LR_Scheduler", None)
-
-                state_dict = {
-                    MODEL_NAME: self.model.state_dict(),
-                    OPTIMIZER_NAME: optim_state_dict,
-                }
+                if self.args.to_static:
+                    state_dict = model.state_dict()
+                else:
+                    optim_state_dict = self.optimizer.state_dict()
+                    optim_state_dict.pop("LR_Scheduler", None)
+                    opt_state_keys = ["_moment1_0", "_moment2_0", "_beta1_pow_acc_0", "_beta2_pow_acc_0"]
+                    for p_name, p in model.state_dict().items():
+                        if paddle.distributed.get_rank() not in p.process_mesh.process_ids:
+                            var_name = p.name
+                            for key in opt_state_keys:
+                                if (
+                                    var_name + key in optim_state_dict
+                                    and not optim_state_dict[var_name + key].is_dist()
+                                ):
+                                    optim_state_dict.pop(var_name + key)
+
+                    state_dict = {
+                        MODEL_NAME: model.state_dict(),
+                        OPTIMIZER_NAME: optim_state_dict,
+                    }
 
                 self._save_ckpt_func(state_dict, os.path.join(output_dir, DIST_CKPT_PATH))
                 logger.info(f"Model weights and optimizer states saved in {output_dir}/{DIST_CKPT_PATH}")
@@ -584,13 +596,9 @@ def _save_checkpoint(self, model, metrics=None):
         rng_states = {
             "python": random.getstate(),
             "numpy": np.random.get_state(),
-            "cuda": [k.current_seed() for k in paddle.get_rng_state()],
-            "cpu": paddle.framework.core.default_cpu_generator().get_state().current_seed(),
+            "cuda": paddle.get_rng_state(),
+            "cpu": paddle.framework.core.default_cpu_generator().get_state(),
         }
-        # if self.args.use_hybrid_parallel:
-        #     rng_states[
-        #         "hybrid_parallel_rng_state_tracker"
-        #     ] = fleet.meta_parallel.get_rng_state_tracker().get_states_tracker()
 
         if self.args.world_size > 1:
             rng_states_list = []
@@ -660,15 +668,23 @@ def _load_from_checkpoint(self, resume_from_checkpoint=None):
             if not os.path.isdir(ckpt_path):
                 raise ValueError(f"Can't find a valid checkpoint at {resume_from_checkpoint}")
 
-            optim_state_dict = self.optimizer.state_dict()
-            optim_state_dict.pop("LR_Scheduler", None)
-
-            state_dict = {
-                MODEL_NAME: self.model.state_dict(),
-                OPTIMIZER_NAME: optim_state_dict,
-            }
+            if self.args.to_static:
+                state_dict = self.model_wrapped.state_dict()
+            else:
+                model_state_dict = self.model_wrapped.state_dict()
+                optim_state_dict = self.optimizer.state_dict()
+                optim_state_dict.pop("LR_Scheduler", None)
+                if len(optim_state_dict) == 0:
+                    self.optimizer._create_accumulators(
+                        paddle.base.framework.default_main_program().global_block(), self.optimizer._parameter_list
+                    )
+                    optim_state_dict = self.optimizer.state_dict()
+                    optim_state_dict.pop("LR_Scheduler", None)
 
-            print("state_dict :", state_dict)
+                state_dict = {
+                    MODEL_NAME: model_state_dict,
+                    OPTIMIZER_NAME: optim_state_dict,
+                }
 
             self._load_ckpt_func(state_dict, ckpt_path)
 

paddlenlp/trainer/trainer.py

@@ -748,11 +748,6 @@ def train(
                     os.makedirs(resume_from_checkpoint, exist_ok=True)
                     logger.info(f"Reset resume_from_checkpoint to temp directory : {resume_from_checkpoint}")
 
-        # memory metrics - must set up as early as possible
-        self._memory_tracker.start()
-        if not self.args.should_load_sharding_stage1_model:
-            self._load_from_checkpoint(resume_from_checkpoint)
-
         train_dataloader = self.get_train_dataloader()
 
         total_train_batch_size = args.train_batch_size * args.gradient_accumulation_steps * args.dataset_world_size
@@ -803,34 +798,43 @@ def train(
 
         self.state = TrainerState()
 
-        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:
-            # 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)
-            if self.sharding_io is not None:
-                assert delay_optimizer_creation is False, "delay_optimizer_creation should be False"
-                # the self.optimizer should be wrapped and it is done in _wrap_model
-                self.sharding_io.set_optimizer(self.optimizer)
-            # 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)
+        # memory metrics - must set up as early as possible
+        self._memory_tracker.start()
+
+        if not self.args.enable_auto_parallel:
+            if not self.args.should_load_sharding_stage1_model:
+                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:
+                # 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)
+                if self.sharding_io is not None:
+                    assert delay_optimizer_creation is False, "delay_optimizer_creation should be False"
+                    # the self.optimizer should be wrapped and it is done in _wrap_model
+                    self.sharding_io.set_optimizer(self.optimizer)
+                # 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:
+                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:
-            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)
+            model = self.model_wrapped
+            self.create_optimizer_and_scheduler(num_training_steps=max_steps)
 
         logger.info(f"{self.runtime_timer.log()}")
-
         logger.info("***** Running training *****")
         logger.info(f"  Num examples = {num_examples:,}")
         logger.info(f"  Num Epochs = {num_train_epochs}")