Skip to content

Commit 80e13e8

Browse files
authored
[DLTP-85730] optimize save cost for zcc (#10939)
1 parent 816ba4e commit 80e13e8

File tree

1 file changed

+1
-1
lines changed

1 file changed

+1
-1
lines changed

paddlenlp/trainer/utils/zero_cost_checkpoint.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -414,7 +414,7 @@ def on_step_end(self, args, state, control, model, lr_scheduler, optimizer, **kw
414414
self.maybe_update_zcc_worker(args, model, optimizer, state.global_step)
415415
checkpoint_folder = f"{PREFIX_CHECKPOINT_DIR}-{state.global_step}"
416416
save_infos = self._get_save_infos_based_on_steps(state, args, checkpoint_folder)
417-
non_cached_objects = (lr_scheduler.state_dict(), copy.deepcopy(state), self.get_rng_states(args))
417+
non_cached_objects = (lr_scheduler.state_dict(), state, self.get_rng_states(args))
418418
self.manager.get_idle_worker_for_saving((save_infos, non_cached_objects))
419419
self.runtime_timer.stop()
420420
if not isinstance(model, PipelineLayer):

0 commit comments

Comments
 (0)