Follow-up on Issues Regarding Training State Restoration from Interruptions

[rangehow]

Hi team,

I would like to follow up on the status of the following issues. Both of these issues involve erroneous behavior that occurs when resuming from an interruption . One issue is that regardless of when training is interrupted at any given timestep, in most cases, a certain amount of data will be un-trained (#38939). The other issue is that the random state cannot be guaranteed to be consistent when resuming from an interruption, which may affect random operations in the random sampler or collator, thus breaking consistency with a full training run (#39215).

I have provided minimal reproducible code, a detailed description of the problem, and a possible set of fixes in the issue descriptions. However, I have not received any further response.

If you believe this direction for a fix is correct, I would be very happy to create PRs to contribute these fixes.

I hope to get some feedback on whether this solution is feasible. Thank you for your time and excellent work on this project

[Rocketknight1]

cc @SunMarc

[rangehow]

For your reference, the combined code solution for these two issues is as follows. The modifications are very lightweight. Based on a diff against the current trainer code in the main branch and my own tests, everything is working correctly.

model.zero_grad()
        grad_norm: Optional[float] = None
        learning_rate = None
        self.control = self.callback_handler.on_train_begin(args, self.state, self.control)

        if args.eval_on_start:
            self._evaluate(trial, ignore_keys_for_eval, skip_scheduler=True)

        for epoch in range(epochs_trained, num_train_epochs):
            epoch_dataloader = train_dataloader
            if hasattr(epoch_dataloader, "set_epoch"):
                epoch_dataloader.set_epoch(epoch)

            # Reset the past mems state at the beginning of each epoch if necessary.
            if args.past_index >= 0:
                self._past = None

            steps_in_epoch = (
                len(epoch_dataloader)
                if len_dataloader is not None
                else args.max_steps * args.gradient_accumulation_steps
            )
            self.control = self.callback_handler.on_epoch_begin(args, self.state, self.control)

            # ================= START OF THE FIX =================

            # This logic block is carefully designed to handle counter and dataloader state
            # for both new epochs and resumed epochs.

            step = -1
            update_step = -1

            
            if epoch == epochs_trained and resume_from_checkpoint is not None:
                # We are resuming training in the middle of an epoch
                if not args.ignore_data_skip:
                    if steps_trained_in_current_epoch > 0:
                        # 1. Skip the dataloader to the correct batch
                        epoch_dataloader = skip_first_batches(epoch_dataloader, steps_trained_in_current_epoch)
                        
                        # 2. Set the number of batches skipped for correct epoch tracking

                        
                        # 3. Initialize step counters to continue from the interruption point
                        # This is the key to ensuring the last batch of the epoch is handled correctly
                        step = steps_trained_in_current_epoch - 1
                        update_step = self.state.global_step - 1
                
                # 4. Load RNG state to ensure reproducibility
                self._load_rng_state(resume_from_checkpoint)

            # For any new epoch (not the first one after resuming), step, update_step,
            # and steps_skipped will remain at their default initialized values (-1, -1, 0),
            # which is the correct behavior.
            
            # The 'steps_trained_in_current_epoch' is a one-time variable for the first resumed epoch.
            # We no longer need the inner `if steps_trained_in_current_epoch > 0:` loop.
            
            # =================  END OF THE FIX  =================


            epoch_iterator = iter(epoch_dataloader)
            # We chunkify the epoch iterator into gradient accumulation steps `n` batches
            remainder = steps_in_epoch % args.gradient_accumulation_steps
            if remainder == 0:
                remainder = args.gradient_accumulation_steps

            total_updates = steps_in_epoch // args.gradient_accumulation_steps + int(
                remainder < args.gradient_accumulation_steps
            )
            for _ in range(total_updates):
                update_step += 1
                num_batches = args.gradient_accumulation_steps if update_step != (total_updates - 1) else remainder
                batch_samples, num_items_in_batch = self.get_batch_samples(epoch_iterator, num_batches, args.device)
                
                for i, inputs in enumerate(batch_samples):
                    step += 1
                    do_sync_step = (step + 1) % args.gradient_accumulation_steps == 0 or (step + 1) == steps_in_epoch
                    # Since we perform prefetching, we need to manually set sync_gradients
                    self.accelerator.gradient_state._set_sync_gradients(do_sync_step)

                    if self.args.include_num_input_tokens_seen:
                        main_input_name = getattr(self.model, "main_input_name", "input_ids")
                        if main_input_name not in inputs:
                            logger.warning(
                                "Tried to track the number of tokens seen, however the current model is "
                                "not configured properly to know what item is the input. To fix this, add "
                                "a `main_input_name` attribute to the model class you are using."
                            )
                        else:
                            input_tokens = inputs[main_input_name].numel()
                            input_tokens = torch.tensor(input_tokens, device=self.args.device, dtype=torch.int64)
                            self.state.num_input_tokens_seen += self.accelerator.gather(input_tokens).sum().item()
                    
                    # The original `rng_to_sync` and inner skipping logic is now replaced by the block above
                    # and is no longer needed here.

                    if step % args.gradient_accumulation_steps == 0:
                        self.control = self.callback_handler.on_step_begin(args, self.state, self.control)

                    # We explicitly want to avoid relying on `accelerator.accumulate` for generation training
                    context = (
                        functools.partial(self.accelerator.no_sync, model=model)
                        if i != len(batch_samples) - 1
                        and self.accelerator.distributed_type != DistributedType.DEEPSPEED
                        else contextlib.nullcontext
                    )
                    with context():
                        tr_loss_step = self.training_step(model, inputs, num_items_in_batch)

                    if (
                        args.logging_nan_inf_filter
                        and not is_torch_xla_available()
                        and (torch.isnan(tr_loss_step) or torch.isinf(tr_loss_step))
                    ):
                        # if loss is nan or inf simply add the average of previous logged losses
                        tr_loss = tr_loss + tr_loss / (1 + self.state.global_step - self._globalstep_last_logged)
                    else:
                        if tr_loss.device != tr_loss_step.device:
                            raise ValueError(
                                f"Calculated loss must be on the original device: {tr_loss.device} but device in use is {tr_loss_step.device}"
                            )
                        tr_loss = tr_loss + tr_loss_step

                    self.current_flos += float(self.floating_point_ops(inputs))

                    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_in_epoch

[SunMarc]

Really sorry for the wait @rangehow, this month was quite busy and I was away for some days ;)

Let's start with this issue #39215 :

• Indeed it looks like we need to change when we update the rng states since we fetch all the batches in advance compared to before. The only thing that troubles me is the following comment in this PR Skip batches fast with accelerate #21390 when skip_first_batches was added. Not sure what will be the impact if we update the rng states before, mind checking quickly ? If there is truly an issue, we can always update it in get_batch_samples wdty ?

Note that the RNG seeds can't be reloaded before we have started the iteration of the data loader because the random sampler used by the data_loader needs to use the seed at this stage, so there is a small path to load it after the iteration has started

BTW, as you said, the inner if steps_trained_in_current_epoch > 0: loop is not required anymore.

As for the second issue #38939, the tests you added looks nice and I think the fix is the right one

Thanks again for such high quality report, this is the kind of report that makes our work easier. Feel free to open a PR with the fixes you shared here and I will review the pr shortly.

[rangehow]

For a clearer discussion, let's simplify the logic to confirm what the correct process should be.

First, to guarantee complete reproducibility, the process must follow a very specific sequence.

Correct Logic perfectly follows this "narrow path":

1. set_seed(SEED)
2. iterator = iter(dataloader) <-- Iteration starts, using the initial seed.
3. for _ in range(INTERRUPT_STEP): next(iterator) <-- Skip the old batches.
4. load_rng_state(CKPT_DIR) <-- Load the state at the correct point on the "narrow path."
5. inputs, _ = next(iterator) <-- Fetch the new batch. At this point, both data loading randomness (e.g., RandomTransformDataset) and subsequent model computation randomness (e.g., Dropout) are under the correct RNG state.

Flawed Trainer Logic deviates from this path:

1. set_seed(SEED)
2. iterator = iter(dataloader) <-- Iteration starts.
3. for _ in range(INTERRUPT_STEP): next(iterator) <-- Skip the old batches.
4. inputs, _ = next(iterator) <-- Point of failure: Fetches the new batch too early. The randomness in data loading (e.g., RandomTransformDataset) uses the wrong RNG state at this moment.
5. load_rng_state(CKPT_DIR) <-- Loading the state is too late; the data is already incorrect.

---

Note that although the trainer's code might seem to call a function like skip_first_batches before creating the iterator, based on my observation of the source code:
https://github.com/huggingface/accelerate/blob/7c25f696b87e0f82cdb915cdba3945cce3ded0bc/src/accelerate/data_loader.py#L1375-L1451

The skip_first_batches function does not perform iteration itself. It is a configuration function. It returns a new, specially configured DataLoader that is set to skip a specified number of batches when iteration begins (i.e., when next() is called).

After carefully considering your description and sgugger's NOTE, I am confident I now have a more comprehensive understanding of the issue with RNG states and the transformers repository. The code modifications I proposed earlier still have some issues. I now believe the problem lies in the fact that skipping old batches and fetching new ones are not separate, atomic operations within the trainer. Instead, they are a tightly coupled process handled by get_batch_samples. This prevents us from finding a "narrow path" where we can insert the load_rng_state call.

I believe there are two potential solutions:

Solution 1: Load RNG State Inside get_batch_samples

Based on my understanding of your description, I will first propose a flawed solution to illustrate the problem. To implement this, we need to let get_batch_samples know that this is the "first call after resuming." We can achieve this by passing a flag.

1. Modify _inner_training_loop

In the main training loop, we need to modify the call to get_batch_samples:

# ... inside _inner_training_loop ...
            for _ in range(total_updates):
                update_step += 1
                num_batches = args.gradient_accumulation_steps if update_step != (total_updates - 1) else remainder
                
                # We pass the rng_to_sync flag into the function
                batch_samples, num_items_in_batch = self.get_batch_samples(
                    epoch_iterator, num_batches, args.device, load_rng_if_needed=rng_to_sync
                )
                # After the call, we must reset the flag
                if rng_to_sync:
                    rng_to_sync = False

                # ... the rest of the loop ...

2. Modify get_batch_samples

Now, the function needs to handle this new flag.

def get_batch_samples(
    self, epoch_iterator: Iterator, num_batches: int, device: torch.device, load_rng_if_needed: bool = False
) -> tuple[list, Optional[torch.Tensor]]:
    
    # *** Key Modification ***
    # If the flag is set, load the RNG state BEFORE any `next()` call.
    if load_rng_if_needed:
        # Assuming `resume_from_checkpoint` is accessible, e.g., self.state.resume_from_checkpoint
        self._load_rng_state(self.state.resume_from_checkpoint)

    batch_samples = []
    num_items_in_batch = None

    # The loop proceeds as normal. The first `next()` call will now happen
    # under the correct, restored RNG state.
    for _ in range(num_batches):
        try:
            batch_samples.append(next(epoch_iterator))
        except StopIteration:
            break
    
    # ... rest of the function is unchanged ...
    return batch_samples, num_items_in_batch

However, this solution is still fundamentally flawed because it loads the RNG state before the next() calls that are meant to skip old batches, changing the correct sequence (1->2->3->4->5) into an incorrect one (1->2->4->3->5). I might have misunderstood your idea and welcome any corrections.

---

Solution 2

Fundamentally decouple "skipping" from "fetching."

This should be done at the _inner_training_loop level, rather than pushing the problem down into the get_batch_samples function.

# ... inside _inner_training_loop, before the main step loop ...

# 1. Decoupling Step 1: Explicitly and manually skip the old batches.
if steps_trained_in_current_epoch > 0:
    for _ in range(steps_trained_in_current_epoch):
        try:
            next(epoch_iterator)
        except StopIteration:
            # Epoch ended while skipping, which is fine.
            break 
    
    # 2. Insertion Point: Load the RNG state right after skipping and before fetching.
    self._load_rng_state(resume_from_checkpoint)

# 3. Decoupling Step 2: Now, enter a clean main loop that is only responsible for fetching new batches.
for _ in range(total_updates):
    # ...
    # `get_batch_samples` is now only responsible for fetching data and no longer
    # needs to worry about RNG restoration.
    batch_samples, ... = self.get_batch_samples(epoch_iterator, num_batches, args.device)
    # ...

[rangehow]

This is a complex issue. My earlier confidence in my changes to the trainer code was based on the belief that I had correctly integrated the "1 2 3 4 5" logic. However, I never realized that skip_first_batch itself doesn't actually handle skipping the old batch data.

[SunMarc]

I mean in this old pr codebase https://github.com/huggingface/transformers/pull/21390, we are updating the RNG state only after we get the first input. Not sure why it worked in the past.

step = -1
  for step, inputs in enumerate(epoch_iterator):
      if rng_to_sync:
          self._load_rng_state(resume_from_checkpoint)
          rng_to_sync = False

Ideally, we should indeed so that as you said only after we skipped the old batches but I think it will require major changes. Maybe it should be fine if we revert to the old behavior. WDYT ?

    batch_samples = []
    num_items_in_batch = None

    batch_samples.append(next(epoch_iterator))
    if load_rng_if_needed: 
        self._load_rng_state(self.state.resume_from_checkpoint)

    for _ in range(num_batches-1):
        try:
            batch_samples.append(next(epoch_iterator))
        except StopIteration:
            break

[rangehow]

The solution you proposed seems correct. I will try to refine the test code to verify it. I'm not entirely sure why the previous approach was used; it might have been an unintentional bug, as the validation related to randomness could be somewhat obscure.

[rangehow]

I've carefully reviewed the source code you provided and found some issues with the approach you suggested.

transformers/src/transformers/trainer.py

Line 2612 in 8365f70

num_batches = args.gradient_accumulation_steps if update_step != (total_updates - 1) else remainder

The num_batches variable is used for batches within a gradient accumulation step. Therefore, when we split this code into two parts, we did not find the narrow window we expected to insert our logic between skipping the old batches and unpacking the new one.

transformers/src/transformers/trainer.py

Lines 5434 to 5438 in 8365f70

	for _ in range(num_batches):
	try:
	batch_samples.append(next(epoch_iterator))
	except StopIteration:
	break

Your idea of fetching one batch first would still cause the first valid batch to be retrieved with an incorrect RNG state. If this valid batch itself contains stochastic operations, the RNG state we load afterward would not have been advanced by them, causing all subsequent data to be loaded with an incorrect state as well.

Your proposed idea is somewhat analogous to this sequence:
1 2 3 5.1 4 5.(2-infinity) not 1 2 3 4 5

[SunMarc]

Your idea of fetching one batch first would still cause the first valid batch to be retrieved with an incorrect RNG state. If this valid batch itself contains stochastic operations, the RNG state we load afterward would not have been advanced by them, causing all subsequent data to be loaded with an incorrect state as well.

Yeah I know this is a bit flawed but I feel like this is what we did in the past here no ? #21390

Otherwise, if we can't to really have the right implementation, we would have to update skip_first_batches to take into account the new rng states. But maybe do that if we are able to craft a test that indeed fails if we don't load the rng state right after skipping the first batches wdyt ?

[rangehow]

It seems the current point of contention is that adjusting skip_first_batch would involve significant effort. Therefore, we need to confirm whether this subtle deviation from the ideal sequence would actually compromise the reproducibility of resumed training. I will start implementing a test script tomorrow at work to validate if this is necessary.

If this minor difference doesn't affect reproducibility, I will then open a PR to update the existing trainer with the RNG loading timing that you mentioned.

[SunMarc]

Sounds good with me !

[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.