Remove unconditional train_batch_size assignment
Author: lordaarush
Pull request: yes

# 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

Comments:
Sure, I'll do just that.

Done @SunMarc ! Added a test that verifies `train_batch_size` is `None` in the saved checkpoint when `auto_find_batch_size=False`.

Hi @SunMarc, looks like simply removing that line broke some tests:
FAILED test_resume_training_with_frozen_params - TypeError: unsupported operand type(s) for //: 'NoneType' and 'int'
FAILED test_resume_training_with_checkpoint - TypeError: unsupported operand type(s) for //: 'NoneType' and 'int'
FAILED test_auto_batch_size_with_resume_from_checkpoint - TypeError: '>=' not supported between instances of 'NoneType' and 'int'
... (9 total trainer resume tests failed)

The issue is that other resume code expects `state.train_batch_size` to have a value for calculations.

I think the correct fix is to:
1. **Keep saving** `train_batch_size` to the checkpoint (as before)
2. **Only restore** it to `_train_batch_size` when `auto_find_batch_size=True`

This way:
- `train_batch_size` is still available in state for other resume calculations
- But it only overwrites the user's batch config when `auto_find_batch_size` is enabled (which needs it)

Should I update the PR with this approach?

Okay let's do that for now ! Thanks for checking 

Updated @SunMarc ! I've implemented the approach we discussed. Also added `test_resume_training_with_different_batch_size`to verify users can change batch size when resuming. All tests passing now.

Thanks a lot ! merging 

The docs for this PR live [here](https://moon-ci-docs.huggingface.co/docs/transformers/pr_43770). All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.