Remove unconditional train_batch_size assignment

[lordaarush]

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).
• Did you read the contributor guideline,
  Pull Request section?
• Was this discussed/approved via a Github issue or the forum? Please add a link
  to it if that's the case. Yes, discussed in Trainer resume_from_checkpoint incorrectly calculates max_steps when changing per_device_train_batch_size with same global batch size #43708
• Did you make sure to update the documentation with your changes? Here are the
  documentation guidelines, and
  here are tips on formatting docstrings.
• Did you write any new necessary tests?

Who can review?

@SunMarc

[SunMarc]

Nice, can you add a simple test to test that it is indeed not saved when auto_find_batch_size is False ?

[lordaarush]

Sure, I'll do just that.

[lordaarush]

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

[SunMarc]

Thanks !

[lordaarush]

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?

[SunMarc]

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

[lordaarush]

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

[SunMarc]

Thanks a lot ! merging

[HuggingFaceDocBuilderDev]

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.