(fix) remove sampler_is_batch_sampler code in prepare_data_loader(..)

[suzyahyah]

What does this PR do?

This PR fixes various confusions wrt to torch.utils.data.DataLoader, torch.utils.data.BatchSampler, and accelerate/data_loader.py:: prepare_data_loader(..).

#3322 Edit, not quite this
#3014
#2091

Motivation

1. accelerate/data_loader.py has various patches allowing for a BatchSampler object to be passed as an argument to sampler.

2. In the code, it allows this behavior using sampler_is_batch_sampler = isinstance(dataloader.sampler, BatchSampler).

3. Allowing this is unintuitive for developers as it directly conflicts with torch.utils.data.DataLoader documentation, when accelerate should only wrap the pytorch dataloader and not change/allow different logic or arguments.

4. It also permits various unintended behavior from other libraries relying on accelerate. For instance, when passing a custom BatchSampler to the dataloader argument in HuggingFace Trainer, the sampler kwargs is currently used, regardless of whether the sampler is BatchSampler or just RandomSampler. https://github.com/huggingface/transformers/blob/3b07ca78bb696825feee3e976795fab58f2b6d0c/src/transformers/trainer.py#L1026 (I'll be making a separate PR in HuggingFace on this)

5. I believe this is due to a misunderstanding stemming from this Issue #Error in prepared DataLoader with BatchSampler #679. As @sgugger initially suspected, this is probably a typo or misunderstanding from HuggingFace.

This should not be allowed behavior, based on the following basic test case which will throw the datasets/formatting/formatting.py: TypeError: int() argument must be a string, a bytes-like object or a real number, not 'list'

from torch.utils.data import BatchSampler, RandomSampler
from datasets import load_dataset

ds = load_dataset("wikitext", "wikitext-103-raw-v1", split='validation')

train_dataloader = DataLoader(
    ds,
    sampler=BatchSampler(RandomSampler(dev_ds), batch_size=32, drop_last=False),
    num_workers=0,  # Adjust based on your setup
    pin_memory=True,
)

for batch in train_dataloader:
  print(batch)

Official Pytorch Documentation for reference:

DataLoader(dataset, batch_size=1, shuffle=False, sampler=None, 
           batch_sampler=None, num_workers=0, collate_fn=None,
           pin_memory=False, drop_last=False, timeout=0,
           worker_init_fn=None, *, prefetch_factor=2,
           persistent_workers=False)

This PR

1. Reverses this historical PR: Fix DataLoader with samplers that are batch samplers #687 which supports the wrong logic.
2. Removes all traces of sampler_is_batch_sampler, simplifying the code
3. Immediately throws an error Edit: Throws a warning if a BatchSampler had been passed as an argument to sampler

Tests

Passes all tests in tests/test_data_loader.py, but does not introduce any new tests. Open to suggestions.

Considerations

1. The PR Immediately throws an error if a BatchSampler had been passed as an argument to Sampler. Technically this error should be thrown earlier or in torch.utils.data.Sampler but we can make it explicit that the problem is not coming from the Accelerate Library, since it had previously been allowed.

2. Upon reading the torch.utils.data source code v2.6.0, I figured that torch.utils.data.DataLoader will attempt to construct a BatchSampler from a Sampler if batch_sampler=None, and will also construct a default sampler if sampler=None.

This means, we should always be able to recover a dataloader.batch_sampler when wrapping accelerate around an already constructed pytorch dataloader, and there is no need to check whether sampler_is_batch_sampler again, if the intent is just to get "new_batch_sampler".

Before submitting

• [ Yes ] Did you read the contributor guideline,
  Pull Request section?
• [ No ] Was this discussed/approved via a Github issue or the forum? Please add a link
  to it if that's the case.
• [NA; no docs mention batch_sampler ] Did you make sure to update the documentation with your changes? Here are the
  documentation guidelines, and
  here are tips on formatting docstrings.
• [No, open to suggestion ] Did you write any new necessary tests?

Who can review?

@SunMarc @BenjaminBossan @zach-huggingface @muellerzr

[muellerzr]

Can you please show the exact paragraph/section of the pytorch docs that state this? You reference many issues, but I dont' see where you found this

[suzyahyah]

Thanks for the quick reply! Yes it’s referenced under motivation #5.

Official PyTorch Documentation for Reference showing two arguments for sampler and batch_sampler,

“””users may use the sampler argument to specify a custom Sampler object that at each time yields the next index/key to fetch.
A custom Sampler that yields a list of batch indices at a time can be passed as the batch_sampler argument.”””

And in the docs they write this is “mutually exclusive with sampler”

batch_sampler (Sampler or Iterable, optional) – like sampler, but returns a batch of indices at a time. Mutually exclusive with batch_size, shuffle, sampler, and drop_last.

[suzyahyah]

If we pass a BatchSampler (yields a list of indices) as an argument to sampler, we will get a list of list of indices, because pytorch will convert whatever was passed in the sampler argument to “batch mode”.

Maybe it’s legal behavior for some very custom parallel processing on batch of batches, but it seems rare and pretty unlikely to me. Should I convert ValueError to Warning

[SunMarc]

can you rebase again ? the diff doesn't look correct

[suzyahyah]

can you rebase again ? the diff doesn't look correct

Yeah I think I fixed it now.

I dont' see where you found this

Re @muellerzr's very valid concern, I changed this to a warning and reference PyTorch documentation.

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

[SunMarc]

Some tests are still failing @suzyahyah

[suzyahyah]

Thanks, @SunMarc

The tests were failing because I misunderstood and had wrongly changed the logic in the function src/accelerate.data_loader.py: set_sampler(). I have reverted the logic in that function back to main branch.

After rebase against main, the tests successfully run (w/o hardware accelerators):

make test_core
make test_prod
pytest tests/test_data_loader.py
make test

Quality checks run:

make style
make quality

[SunMarc]

Thanks for your work ! I feel like it is better if we first do a deprecation message instead of breaking everything right now WDYT @suzyahyah ? This will help users to fix their code. Also, it would be nice to fix the HF docs, otherwise users will still use sampler arg by default

[SunMarc]

let's keep the isinstance(dataloader.sampler, BatchSampler) check

[SunMarc]

The warning is nice, maybe we should add a deprecation message also saying that we won't allow passing BatchSampler to sampler anymore ?

[github-actions]

This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the contributing guidelines are likely to be ignored.

[SunMarc]

LMK if you plan to finish the PR or I will leave it as a feature request / bug to fix @suzyahyah