Error in prepared DataLoader with BatchSampler

[etiennebeaulac]

System Info

accelerate: 0.12.0
OS: Linux 5.4.188+ (Colab)
Python: 3.7.13
numpy: 1.21.6
torch: 1.12.1+cu113
config: 1 CPU

Information

• The official example scripts
• My own modified scripts

Tasks

• One of the scripts in the examples/ folder of Accelerate or an officially supported no_trainer script in the examples folder of the transformers repo (such as run_no_trainer_glue.py)
• My own task or dataset (give details below)

Reproduction

MRE : https://colab.research.google.com/drive/17krCJCF_nWtNFSiMBo3oz12l7eX1bBZ6

First of all, thanks for this library and the great docs and examples that comes with it 😄!

I am using a custom torch Dataset that contains a Hugging Face Dataset (pyarrow) instance. Therefore, as indicated in the Datasets docs (https://huggingface.co/docs/datasets/v2.4.0/en/use_with_pytorch#use-a-batchsampler), I tried to use a BatchSampler to reduce the number of queries. However, I have not been able yet to make it work yet with accelerate.

I tried many different possibilities, one of which works one CPU or one GPU, but gets stuck when using distributed training.

Thanks for your help!

[sgugger]

You are passing your batch sampler using the sampler argument, not the batch_sampler argument when creating your DataLoader. That may be the source of the issue.

[etiennebeaulac]

Thanks @sgugger for your answer. This confused me too at first, but it is apparently the way it should be done, as shown in the HuggingFace docs and in this StackOverflow answer.

[sgugger]

I believe this a typo in both cases.

[etiennebeaulac]

Thanks for your suggestion. I tried it and added a cell in the MRE above where I pass the batch sampler to the batch_sampler argument.

Sadly,

sampler = BatchSampler(RandomSampler(dataset), batch_size=4, drop_last=False)
loader = DataLoader(dataset, batch_sampler=sampler)

produces the same results as the "basic" version :

loader = DataLoader(dataset, batch_size=4, shuffle=True, drop_last=False)

which accesses the dataset one index at a time (not in bulk to reduce overhead). I believe this is why the HF docs uses the sampler argument.

[sgugger]

Oh ok. I did not know about this use of a sampler as a batch sampler. It is indeed not supported by Accelerate. Will have a look on how to add support when I have a bit of time.

[sgugger]

Would you like to give a try to the PR linked above? I believe it should solve your issue.

[etiennebeaulac]

I tried it and it works partially. On more than 1 GPU, it no longer crashes or hangs at the end!

However, the RandomSampler inside the BatchSampler yields duplicated indices and gather_for_metrics does not remove the extra elements as it is supposed to.

[pacman100]

Hello @etiennebeaulac , please test the above PR and let us know if that solves the issue.