Shuffle buffer for controlling sample correlation

[Aceticia]

🚀 Feature

Add a sample level shuffle buffer like webdataset to StreamingDataloader, to make shuffling more random when datasets are built with some correlation.

Motivation

I have some data that correlates across samples when building, and I observe that when using a smaller number of workers, data correlation. This is because litdata only shuffles within chunk, so when the chunk is large and many samples in the chunk are correlated, the randomly drawn samples are still correlated with each other. Exact shuffle is, in general, not tractable - but a good heuristic can be to maintain a shuffle buffer explicitly like webdataset. As long as the buffer size is larger than a chunk, data correlation might be improved.

My case

Here's an example:

The logged value is the sparsity of input. You can see that different clusters of sparsity persist for quite a while, likely because they are shuffling from the same chunk.

[deependujha]

hi @Aceticia, is the proposal something like:

ds = StreamingDataset(shuffle_buffer = 1000)

and internally, it creates an array of the next 1000 yielded samples, and if each chunk has ~300 items, the buffer will basically have predownloaded and loaded the next ~4 chunks, and then randomly shuffle the buffer and yield that?

An easy alternative, imo, is just to make larger chunks (512 MB or something).