Add batch sampler for StreamDataloader to enable flexiable training strategy.

Nobody has claimed this yet.

Hi, I just encounter a problem that associated with litdata. I reformat my dataset and dataloader with StreamDataset and StreamingDataLoader, but I found that the StreamingDataLoader can't accept the sampler function, which makes my original training strategy fails. I use bucket data training, which will put the images with same width/height ratio into the same batch, so I need to realize a special batch sampler to realize this. I need to deliever this sampler to the dataloader. But when I use StreamingDataLoader, I found it doesn't support this, so I can't realize the bucket training strategy, which will make my training sub-optimal. For example, I have prepare my data seqlen like [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15], when I use a batch-size 4, I can guarantee that 0 1 2 3 and 4 5 6 7 and 8 9 10 11 and 12 13 14 15 will have the same bucket size, so in single GPU training, it runs fine. But when I switch to multi-gpu traning, the StreamingDataLoader will shuffle the seqlen, so I will get images that belong to different buckets, so it will raise error in the following collate function. Please add the sampler function for the StreamingDataLoader, or can you have any way to help me solve this problem?

Start by tracing StreamingDataLoader's sampler and shuffle behavior, then reproduce the multi-GPU case with the provided sequence and batch size. Done means a sampler or equivalent mechanism preserves bucket boundaries across distributed batches so the collate function receives compatible samples.

Written by the indexing model from the issue text.