Repository: huggingface/transformers
Pull request: Remove unconditional train_batch_size assignment

# What does this PR do?

Removes the unconditional `self.state.train_batch_size = self._train_batch_size` assignment that was causing issues when resuming from checkpoint with different batch configurations.

The `train_batch_size` should only be saved to `TrainerState` when `auto_find_batch_size` is enabled, which is already handled in the `auto_find_batch_size` block. The unconditional assignment was redundant and caused the bug where `max_steps` was incorrectly calculated when resuming training with a different batch size configuration (even when keeping the same global batch size).

Fixes #43708

## Before submitting
- [ ] This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
- [x] Did you read the [contributor guideline](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md#create-a-pull-request),
      Pull Request section?
- [x] Was this discussed/approved via a Github issue or the [forum](https://discuss.huggingface.co/)? Please add a link
      to it if that's the case. Yes, discussed in #43708
- [ ] Did you make sure to update the documentation with your changes? Here are the
      [documentation guidelines](https://github.com/huggingface/transformers/tree/main/docs), and
      [here are tips on formatting docstrings](https://github.com/huggingface/transformers/tree/main/docs#writing-source-documentation).
- [ ] Did you write any new necessary tests?


## Who can review?

@SunMarc

Changed files:
filename: src/transformers/trainer.py
patch: @@ -2101,7 +2101,10 @@ def train(
                 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:
+            # Only restore the checkpoint's train_batch_size when using auto_find_batch_size,
+            # as that feature needs to resume with the automatically-found batch size.
+            # Otherwise, use the current args batch size to allow users to change batch configuration.
+            if state.train_batch_size is not None and args.auto_find_batch_size:
                 self._train_batch_size = state.train_batch_size
 
         # If model was re-initialized, put it on the right device and update self.model_wrapped

filename: tests/trainer/test_trainer.py
patch: @@ -3473,6 +3473,49 @@ def test_auto_batch_size_with_resume_from_checkpoint(self):
         # We should be back to 14 again, picking up based upon the last ran Trainer
         self.assertEqual(trainer._train_batch_size, previous_batch_size)
 
+    def test_resume_training_with_different_batch_size(self):
+        # Regression test for https://github.com/huggingface/transformers/issues/43708
+        # When resuming from checkpoint without auto_find_batch_size, user's new batch size should be used
+        train_dataset = RegressionDataset(length=64)
+
+        config = RegressionModelConfig(a=0, b=2)
+        model = RegressionRandomPreTrainedModel(config)
+
+        tmp_dir = self.get_auto_remove_tmp_dir()
+
+        # First training run with batch_size=2
+        args = RegressionTrainingArguments(
+            tmp_dir,
+            do_train=True,
+            max_steps=2,
+            save_steps=1,
+            per_device_train_batch_size=2,
+            auto_find_batch_size=False,
+        )
+        trainer = Trainer(model, args, train_dataset=train_dataset)
+        trainer.train()
+
+        # Verify the checkpoint saved with batch_size=2
+        checkpoint = os.path.join(tmp_dir, "checkpoint-1")
+        state = TrainerState.load_from_json(os.path.join(checkpoint, "trainer_state.json"))
+        self.assertEqual(state.train_batch_size, 2)
+
+        # Resume with a different batch_size=4 (without auto_find_batch_size)
+        # The trainer should use the new batch_size, not the checkpoint's
+        args2 = RegressionTrainingArguments(
+            tmp_dir,
+            do_train=True,
+            max_steps=4,
+            save_steps=1,
+            per_device_train_batch_size=4,
+            auto_find_batch_size=False,
+        )
+        trainer2 = Trainer(model, args2, train_dataset=train_dataset)
+        trainer2.train(resume_from_checkpoint=checkpoint)
+
+        # The trainer should be using the new batch size (4), not the checkpoint's (2)
+        self.assertEqual(trainer2._train_batch_size, 4 * max(trainer2.args.n_gpu, 1))
+
     # regression for this issue: https://github.com/huggingface/transformers/issues/12970
     def test_training_with_resume_from_checkpoint_false(self):
         train_dataset = RegressionDataset(length=128)