fix load ema and model_meta path (#11203)

Author: xingmingyyj

paddlenlp/trainer/trainer.py

@@ -784,27 +784,27 @@ def get_metadata_file_name(path):
                 metadata = paddle.load(metadata_file)
                 state_dict_metadata.update(metadata.state_dict_metadata)
 
-            init_optimizer(self.optimizer, model_sharded_state_dict, state_dict_metadata)
+            if not self.args.sharded_model_from_ema:
+                init_optimizer(self.optimizer, model_sharded_state_dict, state_dict_metadata)
 
-            optimizer_sharded_state_dict = self.optimizer.sharded_state_dict(model_sharded_state_dict)
+                optimizer_sharded_state_dict = self.optimizer.sharded_state_dict(model_sharded_state_dict)
 
-            opt_states = {}
-            master_weights = {}
-            for k, v in optimizer_sharded_state_dict.items():
-                if k.endswith(".w_0"):
-                    master_weights[k] = v
-                else:
-                    opt_states[k] = v
+                opt_states = {}
+                master_weights = {}
+                for k, v in optimizer_sharded_state_dict.items():
+                    if k.endswith(".w_0"):
+                        master_weights[k] = v
+                    else:
+                        opt_states[k] = v
 
-            dist.load_state_dict(
-                opt_states,
-                opt_states_path,
-                aoa_config=self.args.aoa_config,
-                offload=self.args.load_via_cpu,
-                comm_method=self.args.flex_ckpt_comm_method,
-            )
+                dist.load_state_dict(
+                    opt_states,
+                    opt_states_path,
+                    aoa_config=self.args.aoa_config,
+                    offload=self.args.load_via_cpu,
+                    comm_method=self.args.flex_ckpt_comm_method,
+                )
 
-            if not self.args.sharded_model_from_ema:
                 dist.load_state_dict(
                     master_weights,
                     master_weights_path,
@@ -819,12 +819,8 @@ def get_metadata_file_name(path):
             ema_states_path = os.path.join(resume_from_checkpoint, EMA_STATE_DIC, f"{dist.get_rank()}_0.distcp")
             ema_state_dict = paddle.load(ema_states_path)
             ema_master_weights = ema_state_dict.pop("master_weights", None)
-            opt_master_weights = self.optimizer.state_dict()["master_weights"]
-            for k, v in opt_master_weights.items():
-                assert (
-                    k in ema_master_weights
-                ), f"{k} not in ema_master_weights, emas_master_weight keys {ema_master_weights.keys()}"
-                paddle.assign(ema_master_weights[k], opt_master_weights[k])
+            opt_state_dict = {"master_weights": ema_master_weights}
+            self.optimizer.set_state_dict(opt_state_dict)
 
             self.model.set_state_dict(ema_state_dict)
         else:

paddlenlp/trainer/utils/zero_cost_checkpoint.py

@@ -1253,6 +1253,7 @@ def __init__(self, resume_from_checkpoint, args, offload=True, hcg=None, model=N
         self.optimizer = optimizer
         self.dist_info_collector_and_validator = DistInfoCollectorValidator(args, hcg)
 
+        self.device_id = int(os.getenv("FLAGS_selected_gpus"))
         super().__init__(resume_from_checkpoint, args, offload)
 
     def _get_model_meta(self):
@@ -1262,7 +1263,7 @@ def _ema_path(self, base_path):
         return os.path.join(base_path, "ema_state", f"{dist.get_rank()}_0.distcp")
 
     def _check_consistent_dist_strategy(self, resume_from_checkpoint):
-        return self.dist_info_collector_and_validator.check_same_strategy(os.path.dirname(resume_from_checkpoint))
+        return self.dist_info_collector_and_validator.check_same_strategy(resume_from_checkpoint)
 
     def _get_model_state(self):
         assert self.model is not None, "expected model is not None"
@@ -1274,9 +1275,12 @@ def _get_master_weight(self):
 
     def save(self, global_step):
         model_meta_content = self._get_model_meta()
-        model_meta_path = os.path.join(self.args.output_dir, MODEL_META_NAME)
-        with open(model_meta_path, "w") as f:
-            json.dump(model_meta_content, f)
+        base_path = os.path.join(self.args.output_dir, f"{PREFIX_CHECKPOINT_DIR}-{global_step}")
+        os.makedirs(base_path, exist_ok=True)
+        model_meta_path = os.path.join(base_path, MODEL_META_NAME)
+        if self.device_id == 0:
+            with open(model_meta_path, "w") as f:
+                json.dump(model_meta_content, f)
 
         super().save(global_step)