Bug in multi-gpu training

[Pixie8888]

Dear Author,

Thank you for sharing source code! But I find that when using CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 torchrun --master_port 10000 --nproc_per_node 8 train_tiktok.py , every gpu reads the same data in every iteration. I think you should add sampler for DDP in your code.

[eastchun]

@Pixie8888

An iterable custom dataset (i.e., ImageTextControlDataset in this repo) does not support direct indexing like dataset[i]). Accordingly, DistributedSampler won’t work directly, because it relies on indexing and len.
To remedy this situation, you should create a custom dataset wrapper that converts the iterable dataset without indexing into the one with the capability into an indexable one. One way is to use IterableDataset with Manual Sharding instead of using DistributedSampler.

If you are interested in digging further, I can provide you with the code fix (only changing train_tiktok.py, not modifying tiktok_video_arnold_copy.py) that works in DDP environmet without using DistributedSampler. Unfortunately, I only have a single GPU laptop and I will not be able to test this modified code in multi-GPU environment.

Please let me know.