Fix the bug where DataLoaderDispatcher gets stuck in an infinite wait…

[yuxinyuan]

Fix the bug where DataLoaderDispatcher gets stuck in an infinite wait when the dataset is an IterDataPipe during multi-process training.

In newer version of pytorch, the iterator of DataLoader will try to broadcast a shared seed across all distributed processes (see here). In the current implementation of DataLoaderDispatcher, the iterator is only created in the main process. This causes the training to hang when dataset is an IterDataPipe.

One can try out the script below to see the effect.

import accelerate
from torch.utils.data.dataloader import DataLoader
from torch.utils.data.datapipes.iter import IterableWrapper


def print_rank_by_rank(*args, **kwargs):
    for rank in range(accelerator.num_processes):
        if rank == accelerator.process_index:
            print(*args, **kwargs)
        accelerator.wait_for_everyone()


if __name__ == "__main__":
    accelerator = accelerate.Accelerator()
    accelerator.print(accelerate.__version__)

    dp = IterableWrapper(range(21)).shuffle().map(lambda x: x + 1)

    loader = DataLoader(dp, batch_size=4, shuffle=True)
    loader = accelerator.prepare(loader)

    accelerator.print(type(loader))

    print_rank_by_rank(list(loader))
    print_rank_by_rank(list(loader))

[HuggingFaceDocBuilderDev]

The documentation is not available anymore as the PR was closed or merged.

[muellerzr]

Is there a reason this has been moved into __iter__ and not __init__?

[yuxinyuan]

I didn't go through the rest of this repo to figure out how remainder is used. However, since self.reset() will always set it to -1, it just makes more sense (looking at data_loader.py) to follow DataLoaderShard and set the remainder in __iter__. Otherwise, it should be safe to remove this code completely.

[muellerzr]

It's used for Gradient accumulation. Here it makes sense to have it in __init__ as there's no need to calculate it twice

[sgugger]

Agreed

[yuxinyuan]

Well, I might be missing something, but, in __iter__, self.reset() will set self.remainder to -1. So, self.remainder won't be useful once we start iterating through the dataloader/dataset. Is that correct?

[sgugger]

Looking more, it is also updated at the end of the iter loop. @muellerzr do we actually need those lines?

But while we investigate, I agree that it's safer to just copy this after the reset.

[muellerzr]

Hmmm... will look into this today.

[muellerzr]

This looks great, thanks! Left one comment on moving a chunk of code.

[sgugger]

Thanks for your PR! Let's just keep the remainder logic in the init as mentioned by @muellerzr and we should be good to go.

[sgugger]

Agreed