From 9ec50bc92d03a854a972c261e905c32da92b250d Mon Sep 17 00:00:00 2001 From: "shengke.sun" Date: Wed, 23 Jul 2025 17:22:58 +0800 Subject: [PATCH] fix: check TrainerState file exists before loading during resume When resuming training from a checkpoint, the Trainer attempts to load \`trainer_state.json\` to recover the train batch size. However, if the file does not exist, a \`FileNotFoundError\` is raised, causing resume to fail. This patch adds a check using \`os.path.isfile\` before loading the state, consistent with other parts of the codebase. If the file is missing, a warning is logged and batch size recovery is skipped, allowing training to continue. Fixes a potential crash and improves robustness when resuming from incomplete or custom checkpoint directories. --- src/transformers/trainer.py | 13 +++++++++---- 1 file changed, 9 insertions(+), 4 deletions(-) diff --git a/src/transformers/trainer.py b/src/transformers/trainer.py index 6dc2a1143605..16acbea5245a 100755 --- a/src/transformers/trainer.py +++ b/src/transformers/trainer.py @@ -2208,10 +2208,15 @@ def train( if not is_sagemaker_mp_enabled() and not self.is_deepspeed_enabled and not self.is_fsdp_enabled: self._load_from_checkpoint(resume_from_checkpoint) # In case of repeating the find_executable_batch_size, set `self._train_batch_size` properly - state = TrainerState.load_from_json(os.path.join(resume_from_checkpoint, TRAINER_STATE_NAME)) - if state.train_batch_size is not None: - self._train_batch_size = state.train_batch_size - + if os.path.isfile(os.path.join(resume_from_checkpoint, TRAINER_STATE_NAME)): + state = TrainerState.load_from_json(os.path.join(resume_from_checkpoint, TRAINER_STATE_NAME)) + if state.train_batch_size is not None: + self._train_batch_size = state.train_batch_size + else: + logger.warning( + f"TrainerState file not found at {os.path.join(resume_from_checkpoint, TRAINER_STATE_NAME)}. " + "Training will continue, but batch size recovery is skipped. " + ) # If model was re-initialized, put it on the right device and update self.model_wrapped if model_reloaded: if self.place_model_on_device: