Trainer `resume_from_checkpoint` incorrectly calculates `max_steps` when changing `per_device_train_batch_size` with same global batch size
Author: KHPan
Pull request: no
Labels: name: bug

### System Info

- `transformers` version: 4.57.3
- Platform: linux
- Python version: 3.11.14
- PyTorch version: 2.9.1

### Who can help?

@SunMarc

### Information

- [ ] The official example scripts
- [x] My own modified scripts

### Tasks

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

### Reproduction

## Bug Description

When resuming training from a checkpoint with a **different `per_device_train_batch_size`** but the **same global batch size** (by adjusting `gradient_accumulation_steps`), the Trainer incorrectly recalculates `max_steps` based on the new batch size configuration, causing:

1. **Incorrect total training steps**: The resumed training runs more (or fewer) steps than expected
2. **Incorrect learning rate schedule**: The linear scheduler recalculates based on wrong `max_steps`

## Minimal Reproducible Example

### Script 1: Initial Training (will be interrupted)

```python
import torch
from torch import nn
from datasets import Dataset
from transformers import PreTrainedModel, PretrainedConfig, Trainer, TrainingArguments
import time

class TinyConfig(PretrainedConfig):
    model_type = "tiny"
    def __init__(self, **kwargs):
        super().__init__(**kwargs)

data_dim = 32
data_cnt = 32

class TinyModel(PreTrainedModel):
    config_class = TinyConfig
    def __init__(self, config):
        super().__init__(config)
        self.fc = nn.Linear(data_dim, data_dim)
    
    def forward(self, input_ids, attention_mask=None, labels=None, **kwargs):
        time.sleep(1)
        loss = self.fc(input_ids.float()).mean(dim=-1)
        return {"loss": loss}

dataset = Dataset.from_dict({
    "input_ids": torch.zeros(data_cnt, data_dim, dtype=torch.long).tolist(),
    "attention_mask": torch.ones(data_cnt, data_dim, dtype=torch.long).tolist(),
    "labels": torch.zeros(data_cnt, data_dim, dtype=torch.long).tolist(),
})

model = TinyModel(TinyConfig())
trainer_args = TrainingArguments(
    output_dir="./issue_output", 
    num_train_epochs=1, 
    report_to="none", 
    logging_steps=1,
    save_steps=1,
    learning_rate=0.08,
    per_device_train_batch_size=1,
    gradient_accumulation_steps=4,
    lr_scheduler_type="linear",
)
trainer = Trainer(
    model=model,
    args=trainer_args,
    train_dataset=dataset,
)
trainer.train()
```

### Script 2: Resume Training with Different Batch Config (same global batch)

```python
import torch
from torch import nn
from datasets import Dataset
from transformers import PreTrainedModel, PretrainedConfig, Trainer, TrainingArguments
import time

class TinyConfig(PretrainedConfig):
    model_type = "tiny"
    def __init__(self, **kwargs):
        super().__init__(**kwargs)

data_dim = 32
data_cnt = 32

class TinyModel(PreTrainedModel):
    config_class = TinyConfig
    def __init__(self, config):
        super().__init__(config)
        self.fc = nn.Linear(data_dim, data_dim)
    
    def forward(self, input_ids, attention_mask=None, labels=None, **kwargs):
        time.sleep(1)
        loss = self.fc(input_ids.float()).mean(dim=-1)
        return {"loss": loss}

dataset = Dataset.from_dict({
    "input_ids": torch.zeros(data_cnt, data_dim, dtype=torch.long).tolist(),
    "attention_mask": torch.ones(data_cnt, data_dim, dtype=torch.long).tolist(),
    "labels": torch.zeros(data_cnt, data_dim, dtype=torch.long).tolist(),
})

model = TinyModel(TinyConfig())
trainer_args = TrainingArguments(
    output_dir="./issue_output", 
    num_train_epochs=1, 
    report_to="none", 
    logging_steps=1,
    save_steps=1,
    learning_rate=0.08,
    per_device_train_batch_size=4,   # Changed from 1 to 4
    gradient_accumulation_steps=1,    # Changed from 4 to 1
    lr_scheduler_type="linear",
)
trainer = Trainer(
    model=model,
    args=trainer_args,
    train_dataset=dataset,
)
trainer.train(resume_from_checkpoint=True)
```

## Steps to Reproduce

1. Run Script 1 and **interrupt it at 50%** (around step 4/8)
2. Run Script 2 to resume from the checkpoint

## Observed Behavior

### Script 1 Log (interrupted at 50%):
```
{'loss': -0.0153, 'grad_norm': 0.7071067690849304, 'learning_rate': 0.08, 'epoch': 0.12}
{'loss': -0.3353, 'grad_norm': 0.7071067690849304, 'learning_rate': 0.07, 'epoch': 0.25}
{'loss': -0.6153, 'grad_norm': 0.7071067690849304, 'learning_rate': 0.06, 'epoch': 0.38}
{'loss': -0.8553, 'grad_norm': 0.7071067690849304, 'learning_rate': 0.05, 'epoch': 0.5}
 50%|██████████████████████████████████████████▌                                          | 4/8 [00:16<00:16,  4.07s/it]
```

- Total steps: **8** (32 samples / (batch_size=1 × gradient_accumulation=4) = 8)
- Completed: **4 steps** (50%)

### Script 2 Log (resumed training):
```
{'loss': -0.2638, 'grad_norm': 0.1767766922712326, 'learning_rate': 0.04, 'epoch': 0.16}
{'loss': -0.3001, 'grad_norm': 0.1767766922712326, 'learning_rate': 0.0675, 'epoch': 0.19}
{'loss': -0.3568, 'grad_norm': 0.1767766922712326, 'learning_rate': 0.065, 'epoch': 0.22}
... (many more steps)
{'loss': -0.9351, 'grad_norm': 0.1767766922712326, 'learning_rate': 0.0025, 'epoch': 1.0}
100%|███████████████████████████████████████████████████████████████████████████████████| 32/32 [00:28<00:00,  1.11it/s]
```

- **max_steps was recalculated to 32** instead of maintaining the original 8
- The training ran **28 additional steps** (32 - 4 = 28) instead of **4** (8 - 4 = 4)
- The learning rate schedule was recalculated based on wrong max_steps

## Analysis

| Configuration | Script 1 | Script 2 |
|---------------|----------|----------|
| `per_device_train_batch_size` | 1 | 4 |
| `gradient_accumulation_steps` | 4 | 1 |
| **Global batch size** | **4** | **4** |
| Calculated `max_steps` | 8 | 32 (incorrect!) |

Since the global batch size remains the same (4), the total number of optimization steps should remain **8**. However, when resuming, the Trainer recalculates `max_steps` based on the new `per_device_train_batch_size` without considering that the global batch size is unchanged.

The learning rate schedule also gets recalculated based on the new (incorrect) `max_steps`, causing the linear decay to behave incorrectly.

### Expected behavior

When resuming from a checkpoint with the same global batch size (but different `per_device_train_batch_size` and `gradient_accumulation_steps`), the Trainer should:

1. **Maintain the original `max_steps`** from the initial training run, or at least calculate based on the global batch size
2. **Continue the learning rate schedule correctly** based on the original `max_steps`

In this case:
- Original `max_steps` = 8
- Resumed at step 4
- Expected remaining steps = 4
- Expected final step = 8

The resumed training should complete at step 8, not step 32.

Comments:
I've investigated this issue and can confirm the root cause. The problem appears to be in the `set_initial_training_values` method in `trainer.py`.

When resuming from a checkpoint, the Trainer recalculates `num_update_steps_per_epoch` based on the new dataloader length and gradient_accumulation_steps, without considering:

1. The original training's `max_steps` should be preserved
2. The global batch size is unchanged, so total optimization steps should remain the same

The fix should involve checking if we're resuming from a checkpoint and preserving the original `max_steps` from the `TrainerState`, rather than recalculating it.

I'm interested in working on a fix. @SunMarc, could you confirm this is the right approach?

Feel free to work on it @lordaarush ! I think there might be an issue with how max_steps is calculated maybe, but I don't think we should just use the value from TrainerState

Thanks @SunMarc for the guidance! I dug deeper and found the actual root cause.

**The issue wasn't in `set_initial_training_values` or the `max_steps` calculation itself** - that formula is correct. The problem was earlier in the `train()` method at line 2150-2151.

When resuming from a checkpoint, `self._train_batch_size` was unconditionally restored from the checkpoint state, causing the dataloader to be created with the OLD batch size. This led to incorrect `len_dataloader` and ultimately wrong `max_steps`.

**The fix:** Only restore the checkpoint's `train_batch_size` when `auto_find_batch_size` is enabled (since that feature needs the auto-discovered batch size). Otherwise, respect the user's current configuration.

```python
# Before:
if state.train_batch_size is not None:
    self._train_batch_size = state.train_batch_size

# After:
if state.train_batch_size is not None and args.auto_find_batch_size:
    self._train_batch_size = state.train_batch_size

We only save the `train_batch_size` if `auto_find_batch_size` was used I think. So it shouldn't happen for this issue. Can you try to confirm ? 

> We only save the `train_batch_size` if `auto_find_batch_size` was used I think. So it shouldn't happen for this issue. Can you try to confirm ?


I checked this and it looks like `train_batch_size` is actually saved unconditionally, not just when `auto_find_batch_size` is used.

There are two places in `_inner_training_loop` where `state.train_batch_size` is set:

1. **Line 2251** - Inside the `if self.args.auto_find_batch_size:` block
2. **Line 2308** - Unconditionally: `self.state.train_batch_size = self._train_batch_size`

That second assignment runs every training, regardless of the `auto_find_batch_size` setting.

I verified this by running training with `auto_find_batch_size=False` and checking the saved checkpoint:

```python
args = TrainingArguments(
    ...,
    per_device_train_batch_size=2,
    auto_find_batch_size=False,
)
trainer.train()
# Then check checkpoint:
state = TrainerState.load_from_json("checkpoint-2/trainer_state.json")
print(state.train_batch_size)  # Output: 2

```
So the checkpoint contains `train_batch_size = 2` even though `auto_find_batch_size=False`. This causes the bug when resuming with a different batch configuration.

Arf indeed, thanks for finding this @ ! We should remove this then. Would you like to open a PR for that ? 
`Line 2308 - Unconditionally: self.state.train_batch_size = self._train_batch_size`
Would this fix your issue @KHPan ?

Sure @SunMarc, I have opened a new PR (#43770) which simply removes the unconditional `self.state.train_batch_size = self._train_batch_size` assignment.