Fix batch_size sanity check logic for split_batches

[izhx]

What does this PR do?

Fixes #2285 and PR/2310 comment: Unfortunately, this broke Megatron DL

Some libraries may not follow PyTorch's design of BatchSampler.
If the batch_size cannot be found, a warning is left to notify the user.

Do anyone have suggestions for such situations?

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.
• 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?

Anyone in the community is free to review the PR once the tests have passed. Feel free to tag
members/contributors who may be interested in your PR.
@muellerzr @BenjaminBossan

[muellerzr]

I think this solution makes sense, but I know this originates from Megatron so cc @pacman100 for guidance.

[pacman100]

Thank you @izhx for the fixes in cases when the Batch Sampler doesn't conform to PyTorch. Makes sense.

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

[muellerzr]

If Sourab is happy, I'm happy. But would definitely be happier if @stas00 could also give this an all-clear? :)

[stas00]

1. OK, so the key fix of this PR is it adding if split_batches: as in my trainer I have never set it to True. So this solves my problem and doesn't require Megatron Sampler to comply with certain requirements. Thank you!

2. Now let's zoom out to the general case where split_batches==True has been set by the user. While this PR makes the dealing with an issue somewhat user-friendlier - this approach swipes the problem under the carpet. If anything it should be an assert and not a warning. i.e. if the user wants split_batches this code path must not proceed w/o having a way to get to sampler.batch_size. In other words split_batches functionality must require a sampler class that provides self.batch_size.

3. While I was investigating this I think you have this problem here as well:

accelerate/src/accelerate/data_loader.py

Lines 487 to 492 in c0a9472

	def total_batch_size(self):
	batch_sampler = self.sampler if isinstance(self.sampler, BatchSampler) else self.batch_sampler
	return (
	batch_sampler.batch_size
	if getattr(batch_sampler, "split_batches", False)
	else (batch_sampler.batch_size * getattr(batch_sampler, "num_processes", 1))

it always calls sampler.batch_size. so that code also probably should check hasattr and assert if it is False. It's probably not in the scope of this PR though, so I'm just flagging to maintainers.

4. In general avoid using warnings as much as possible. Nobody sees those in a sea of thousands of logs lines. It's either a problem and then you assert, or it's not really a problem and then you do nothing. With some rare exceptions there is no in-between state.

[stas00]

let's land this one, @muellerzr?

[muellerzr]

Thanks for iterating! (Good to go once tests pass)