Resuming training from an interrupted checkpoint fails to save the final checkpoint.

[rangehow]

System Info

• transformers version: 4.53.0.dev0
• Platform: Linux-4.18.0-147.mt20200626.413.el8_1.x86_64-x86_64-with-glibc2.17
• Python version: 3.12.3
• Huggingface_hub version: 0.30.1
• Safetensors version: 0.5.3
• Accelerate version: 1.6.0
• Accelerate config: not found
• DeepSpeed version: 0.16.7
• PyTorch version (accelerator?): 2.6.0+cu124 (NA)
• Tensorflow version (GPU?): not installed (NA)
• Flax version (CPU?/GPU?/TPU?): not installed (NA)
• Jax version: not installed
• JaxLib version: not installed
• Using distributed or parallel set-up in script?: yes

Information

• The official example scripts
• My own modified scripts

Tasks

• An officially supported task in the examples folder (such as GLUE/SQuAD, ...)
• My own task or dataset (give details below)

Reproduction

Normally, a trainer that saves checkpoints based on steps or epochs is guaranteed to save a final checkpoint upon completion of the training. This is ensured by the following sequence of logic:

First, when the training reaches the final step, do_sync_step is set to True based on the second part of this conditional statement.

transformers/src/transformers/trainer.py

Line 2508 in 0725cd6

do_sync_step = (step + 1) % args.gradient_accumulation_steps == 0 or (step + 1) == steps_in_epoch

This flag indicates that a gradient update is about to be performed, and the global_step is updated precisely within that gradient update branch. (see L2623)

transformers/src/transformers/trainer.py

Lines 2569 to 2635 in 0725cd6

	if do_sync_step:
	# Since we perform prefetching, we need to manually set sync_gradients to True
	self.accelerator.gradient_state._set_sync_gradients(True)
	# Gradient clipping
	if args.max_grad_norm is not None and args.max_grad_norm > 0:
	if is_sagemaker_mp_enabled() and args.fp16:
	_grad_norm = self.optimizer.clip_master_grads(args.max_grad_norm)
	elif self.use_apex:
	from apex import amp
	# Revert to normal clipping otherwise, handling Apex or full precision
	_grad_norm = nn.utils.clip_grad_norm_(
	amp.master_params(self.optimizer),
	args.max_grad_norm,
	)
	else:
	grad_norm_context = contextlib.nullcontext
	if self.is_tp_enabled:
	from torch.distributed._tensor.experimental import implicit_replication
	grad_norm_context = implicit_replication
	with grad_norm_context():
	_grad_norm = self.accelerator.clip_grad_norm_(
	model.parameters(),
	args.max_grad_norm,
	)
	if (
	is_accelerate_available()
	and self.accelerator.distributed_type == DistributedType.DEEPSPEED
	):
	grad_norm = model.get_global_grad_norm()
	# In some cases the grad norm may not return a float
	if hasattr(grad_norm, "item"):
	grad_norm = grad_norm.item()
	else:
	grad_norm = _grad_norm
	self.control = self.callback_handler.on_pre_optimizer_step(args, self.state, self.control)
	self.optimizer.step()
	self.control = self.callback_handler.on_optimizer_step(args, self.state, self.control)
	# get leaning rate before update
	learning_rate = self._get_learning_rate()
	if not self.accelerator.optimizer_step_was_skipped:
	# Delay optimizer scheduling until metrics are generated
	if not isinstance(self.lr_scheduler, torch.optim.lr_scheduler.ReduceLROnPlateau):
	self.lr_scheduler.step()
	model.zero_grad()
	self.state.global_step += 1
	self.state.epoch = epoch + (step + 1 + steps_skipped) / steps_in_epoch
	self.control = self.callback_handler.on_step_end(args, self.state, self.control)
	self._maybe_log_save_evaluate(
	tr_loss,
	grad_norm,
	model,
	trial,
	epoch,
	ignore_keys_for_eval,
	start_time,
	learning_rate=learning_rate,
	)

This coherent logic would work without any issues in an error-free training run. However, when training is resumed, the step counter is not restored; it is still initialized from -1.

transformers/src/transformers/trainer.py

Line 2492 in 0725cd6

step = -1

This prevents a gradient update for the final set of batches in an epoch if the total dataset size is not perfectly divisible by the product of gradient_accumulation_steps and the data-parallel world size.

[rangehow]

I'd be happy to open a PR to fix this issue.
One solution I've considered is to correctly initialize the step variable with a simple pre-calculation at this location:

transformers/src/transformers/trainer.py

Lines 2423 to 2443 in 0725cd6

	if resume_from_checkpoint is not None and os.path.isfile(
	os.path.join(resume_from_checkpoint, TRAINER_STATE_NAME)
	):
	self.state = TrainerState.load_from_json(os.path.join(resume_from_checkpoint, TRAINER_STATE_NAME))
	self.compare_trainer_and_checkpoint_args(self.args, self.state)
	self._load_callback_state()
	epochs_trained = int(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
	logger.info(" Continuing training from checkpoint, will skip to saved global_step")
	logger.info(f" Continuing training from epoch {epochs_trained}")
	logger.info(f" Continuing training from global step {self.state.global_step}")
	if not args.ignore_data_skip:
	logger.info(
	f" Will skip the first {epochs_trained} epochs then the first"
	f" {steps_trained_in_current_epoch} batches in the first epoch."
	)

I'm open to any better ideas or suggestions as well.

modification would be like below:

if resume_from_checkpoint is not None and os.path.isfile(
            os.path.join(resume_from_checkpoint, TRAINER_STATE_NAME)
        ):
            self.state = TrainerState.load_from_json(os.path.join(resume_from_checkpoint, TRAINER_STATE_NAME))
            self.compare_trainer_and_checkpoint_args(self.args, self.state)
            self._load_callback_state()
            epochs_trained = int(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

            logger.info("  Continuing training from checkpoint, will skip to saved global_step")
            logger.info(f"  Continuing training from epoch {epochs_trained}")
            logger.info(f"  Continuing training from global step {self.state.global_step}")
            if not args.ignore_data_skip:
                logger.info(
                    f"  Will skip the first {epochs_trained} epochs then the first"
                    f" {steps_trained_in_current_epoch} batches in the first epoch."
                )
            
            step = steps_trained_in_current_epoch
            update_step = self.state.global_step
        else:
            step = -1
            update_step = -1

---

While this code solves the original problem, it introduces new ones. Do not use.

[Rocketknight1]

cc @SunMarc

[SunMarc]

This sounds reasonable @rangehow ! If you can open a PR with a small test to check that we are indeed resuming with step correctly initialized, that would be perfect

[rangehow]

Hi @SunMarc, I'm not very familiar with contributing test code. Below is a test code snippet. The training should run for a total of 3 steps. I intentionally interrupted the training at step 2 and then resumed it to verify whether the checkpoint for the third step would be saved. Running this on the current main version triggers an assert error, confirming the issue I described.

I have two questions:

1. Should this code be added to the tests folder, and are there any additional steps needed to align it with the repository's testing framework?
2. Are there any oversights or areas for improvement in this test code itself?

I look forward to your guidance.

import os
import shutil
import torch
from torch.utils.data import TensorDataset, Dataset
from transformers import (
    AutoModelForSequenceClassification,
    AutoTokenizer,
    Trainer,
    TrainingArguments,
)
import torch.nn as nn

class DummyModel(nn.Module):
    def __init__(self, input_dim=10, num_labels=2):
        super().__init__()
        self.linear = nn.Linear(input_dim, num_labels)

    def forward(self, input_ids=None, attention_mask=None, labels=None):
        logits = self.linear(input_ids.float())
        loss = None
        if labels is not None:
            loss_fn = nn.CrossEntropyLoss()
            loss = loss_fn(logits, labels)
        return {"loss": loss, "logits": logits}

class DummyDictDataset(Dataset):
    def __init__(self, input_ids, attention_mask, labels):
        self.input_ids = input_ids
        self.attention_mask = attention_mask
        self.labels = labels

    def __len__(self):
        return len(self.input_ids)

    def __getitem__(self, idx):
        return {
            "input_ids": self.input_ids[idx],
            "attention_mask": self.attention_mask[idx],
            "labels": self.labels[idx],
        }

def test_resume_from_checkpoint_saves_final_model():
    # 1. Set up a dummy model
    model = DummyModel(input_dim=10, num_labels=2)

    # 2. Create a dataset with a size that is not perfectly divisible
    # Assume batch_size=2, gradient_accumulation_steps=2, total steps should be 10 / (2*2) = 2.5, which will have a remainder
    # Simulate a non-divisible case
    num_samples = 11
    input_dim = 10
    dummy_input_ids = torch.rand(num_samples, input_dim)
    dummy_attention_mask = torch.ones(num_samples, input_dim)  # attention_mask is actually not used, but kept for interface consistency
    dummy_labels = torch.randint(0, 2, (num_samples,))
    dummy_dataset = DummyDictDataset(
        dummy_input_ids,
        dummy_attention_mask,
        dummy_labels,
    )

    # 3. First training (simulate interruption)
    output_dir_initial = "./test_resume_training_initial"
    training_args_initial = TrainingArguments(
        output_dir=output_dir_initial,
        num_train_epochs=1,
        per_device_train_batch_size=2,
        gradient_accumulation_steps=2,
        save_strategy="steps",
        save_steps=1, # Save at every step
        report_to=[], # Disable wandb/tensorboard and other loggers
        max_steps=2,  # Stop after step 2 to simulate interruption
    )

    trainer_initial = Trainer(
        model=model,
        args=training_args_initial,
        train_dataset=dummy_dataset,
    )
    trainer_initial.train()

    # Make sure we have a checkpoint before interruption
    checkpoint_path = os.path.join(output_dir_initial, "checkpoint-2")
    assert os.path.exists(checkpoint_path)

    # 4. Resume training from checkpoint
    output_dir_resumed = "./test_resume_training_resumed"
    training_args_resumed = TrainingArguments(
        output_dir=output_dir_resumed,
        num_train_epochs=1,
        per_device_train_batch_size=2,
        gradient_accumulation_steps=2,
        save_strategy="steps",
        save_steps=1, # Keep the same save strategy

    )

    trainer_resumed = Trainer(
        model=model,
        args=training_args_resumed,
        train_dataset=dummy_dataset,
    )
    # Resume from the interrupted checkpoint and finish the remaining training
    trainer_resumed.train(resume_from_checkpoint=checkpoint_path)

    # 5. Assertion: Check if the final model has been saved
    # According to the logic, after training is complete, the final model should be saved in output_dir_resumed
    # Before the fix, this assertion should fail
    final_model_path = os.path.join(output_dir_resumed, "pytorch_model.bin")
    assert os.path.exists(final_model_path), "The final model checkpoint was not saved!"

    # Clean up test directories
    shutil.rmtree(output_dir_initial)
    shutil.rmtree(output_dir_resumed)

# Run the test
test_resume_from_checkpoint_saves_final_model()

[github-actions]

This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the contributing guidelines are likely to be ignored.

[rangehow]

It's a shame this bug was closed if it still exists.

[github-actions]

This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the contributing guidelines are likely to be ignored.

[rangehow]

bump

[github-actions]

This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the contributing guidelines are likely to be ignored.