Kimirock/PaddleNLP
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

fix load ema and model_meta path (#11203)

Browse Source
This commit is contained in:
xingmingyyj
2025-12-17 17:19:22 +08:00
committed by GitHub
parent 32b9a2e257
commit 587a5cdf97
2 changed files with 30 additions and 30 deletions
Download Patch File Download Diff File Expand all files Collapse all files

48
paddlenlp/trainer/trainer.py
View File

@@ -784,27 +784,27 @@ class Trainer:
metadata = paddle.load(metadata_file)
state_dict_metadata.update(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)
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,
)
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)
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(
master_weights,
master_weights_path,
@@ -819,12 +819,8 @@ class Trainer:
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:

12
paddlenlp/trainer/utils/zero_cost_checkpoint.py
View File

@@ -1253,6 +1253,7 @@ class EMABufferFcBased(EMABuffer):
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 @@ class EMABufferFcBased(EMABuffer):
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 @@ class EMABufferFcBased(EMABuffer):
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)