Iterating a streaming dataset correctly in random order

[unrealwill]

Describe the bug

For training to work properly it is considered good practice in machine learning to sample the dataset randomly (uniformly). I am not sure how the user is expected to do it correctly when the creator of the dataset did a bad job.

https://huggingface.co/docs/hub/datasets-webdataset

 **Shuffle**

Generally, datasets in WebDataset formats are already shuffled and ready to feed to a DataLoader. But you can still reshuffle the data with WebDataset’s approximate shuffling.

In addition to shuffling the list of shards, WebDataset uses a buffer to shuffle a dataset without any cost to speed:

Let me expose the general problem with this specific dataset.
https://huggingface.co/datasets/jackyhate/text-to-image-2M

I'll only consider the single node, single worker case to highlight the problem, which is somewhat hidden and attenuated when training using multiple machines.

Here is the recommended way to load it on the main page

# copy pasted from https://huggingface.co/datasets/jackyhate/text-to-image-2M
num_shards = 46  # Number of webdataset tar files
urls = [base_url.format(i=i) for i in range(num_shards)]
dataset = load_dataset("webdataset", data_files={"train": urls}, split="train", streaming=True)
# Example of iterating through the dataset
for image in dataset:
    print(image)  # single image in row with associated columns
    break

The dataset is composed of 46 shards from multiple sources : Some shards have data generated with flux, and other have data generated from dall-e. There are usually 50000 images per file except for data_000034.tar and data_000046.tar which have less files because they are the end files of some version of the dataset.

Initially I thought you were using the webdataset library which has at least tried to think about the problem of shuffling the dataset (although poorly documented and with bad default behavior : See https://github.com/webdataset/webdataset/blob/e0953f9bba17b416d5792d5a263b171c266e78be/src/webdataset/mix.py )

But it seems that you are just iterating the shards in order and then iterating the examples of each shard in order.

datasets/src/datasets/packaged_modules/webdataset/webdataset.py

Lines 111 to 130 in fe7353a

	def _generate_examples(self, tar_paths, tar_iterators):
	image_field_names = [
	field_name for field_name, feature in self.info.features.items() if isinstance(feature, datasets.Image)
	]
	audio_field_names = [
	field_name for field_name, feature in self.info.features.items() if isinstance(feature, datasets.Audio)
	]
	all_field_names = list(self.info.features.keys())
	for tar_idx, (tar_path, tar_iterator) in enumerate(zip(tar_paths, tar_iterators)):
	for example_idx, example in enumerate(self._get_pipeline_from_tar(tar_path, tar_iterator)):
	for field_name in all_field_names:
	if field_name not in example:
	example[field_name] = None
	for field_name in image_field_names + audio_field_names:
	if example[field_name] is not None:
	example[field_name] = {
	"path": example["__key__"] + "." + field_name,
	"bytes": example[field_name],
	}
	yield Key(tar_idx, example_idx), example

If I understand correctly the recommended way to shuffle the data as described in https://huggingface.co/docs/datasets/loading

ds = ds.shuffle(seed=42, buffer_size=10_000) # shuffles the shards order + uses a shuffle buffer

or shuffle with a buffer in the dataloader.

This is problematic because the dataset has not been initially shuffled by the creator of the dataset, which means that samples are highly correlated, for examples samples coming from the flux shard have a lot more details than those coming from the dall-e shard, and they are not mixed in the shuffle buffer, which result in the loss function having a periodic pattern by epoch typical of bad shuffling behavior.

Of course we could just reshuffle the datasets beforehand but these are huge and we would need to do it for every versions of the dataset.

The random mixing and round robin mixing which can be used in the webdataset library, can also be problematic when the shards are of various sizes : When we stopped at the expiration of the first shard, we never visit all the examples from the later portions of other shards. Also if some shards provenance dataset are not uniformly balanced, we have some period of time where the shard is not sampled, for example when the smaller shard has been exhausted, samples from its class won't be seen until the next epoch.

The proper way being sampling from shards so that they are all exhausted roughly at the same time, aka using a smaller sampling probability for smaller shards, but that necessitate to know the sizes of the shard beforehand, and if you also do per shard shuffling like webdataset library do, it needs to open all shards simultaneously and therefore require a lot of memory to load nshards*buffer_size.

Hopefully I have been clear enough, and the problem which occurs with webdataset and probably also occurs with every other streamable format that you handle can be made to work in a standardized, robust and bug-free way, as the types of bugs these generate are usually of the silent kind where training works better but it generalizes less well because the training used some correlation from the sequence of the dataset.

Steps to reproduce the bug

Observe periodic behavior of loss function when training a simple neural network on the dataset

Expected behavior

A behavior somewhat similar to a standard uniform shuffle

Environment info

not relevant

[lhoestq]

Hi, reshuffling is actually pretty easy:

""" shuffle.py """
from datasets import load_dataset

num_proc = 8  # tweak it here depending on your machine
seed = 42
ds = load_dataset("jackyhate/text-to-image-2M", num_proc=num_proc)
ds = ds.shuffle(seed=seed)
ds.push_to_hub("my-username/text-to-image-2M-shuffled", num_proc=num_proc)

If you need a machine with a lot of disk to download the full dataset and do the shuffling, feel free to use HF Jobs

E.g. you can run the above script with

hf jobs uv run --flavor cpu-performance -s HF_TOKEN shuffle.py

PS: the HF_TOKEN is needed for the Job to push the shuffled dataset to HF
PS2: see the other flavors at https://huggingface.co/docs/hub/jobs-configuration#hardware-flavor

[cynricfu]

I guess a simple example to demonstrate the idea is as follows.

Suppose there is a streaming dataset that comes with 3 shards, where each shard contains 4 samples: (1, 2, 3, 4), (11, 22, 33, 44), (111, 222, 333, 444). Note that different shards may originate from different sources and therefore differ in data distributions.

If we start a single-node single-card training with batch_size=2, then even with shuffling (let's say the shuffle buffer size is also 4), the current behavior may yield mini-batches similar to these: [33, 11], [11, 44], [2, 4], [3, 1], ...

This could be problematic as mini-batches are not "random" enough. @unrealwill is this your point?

[lhoestq]

I guess one way to improve this is to interleave the shards together to improve the shuffling quality.

You can already do this using

from datasets import interleave_datasets

seed = 42
num_shards_to_interleave = ds.num_shards
interleaved_ds = interleave_datasets(
    [ds.shard(index=index, num_shards=num_shards_to_interleave) for index in range(num_shards_to_interleave)]
    probabilities=[1/num_shards_to_interleave] * num_shards_to_interleave,
    seed=seed
)
better_shuffled_ds = interleaved_ds.shuffle(seed=seed)

Note that it means that it opens one connection per shard which can be an issue if you have thousands of shards. This can also can memory pressure because it may buffer every shard.

In practice to avoid these issues your can set something like this

num_shards_to_interleave = min(ds.num_shards, 16)

(even better if it ends up with num_shards_to_interleave being a factor of ds.num_shards so the dataset can be more evenly distributed in a DataLoader or in a distributed setup)

if you find this useful we could add a new parameter in shuffle to improve shuffling with shards interleaving

[unrealwill]

@cynricfu yes it's about the quality of the random shuffle. When individual shards contains data with distinct statistics and we iterate them one shard after the next, this cause learning to adapt to the current shard statistics, but then forget the older shards. That's why it's recommended that all shards are statistically similar which can be practically achieved by pre-shuffling.

@lhoestq I hadn't realize it was so easy to reshuffle the dataset, but I find this solution inelegant as it cause massive amount of data duplication. Also most people won't even register the problem. And everyone would have to shuffle the dataset themselves because they can't trust that the dataset has been shuffled correctly. And it breaks replication because everyone use their own shuffled version of the dataset.

Interleaving the dataset is a solution provided that the probability used is number of example in the shard / number of example in total (numbers that we don't have in a streaming mode unless they are given as additional info), rather than what you have written. Because if some shard is smaller, it becomes exhausted earlier, and data from it may not be sampled from it until the next epoch.

The memory pressure cannot be so easily alleviated with your trick, because if some shards become exhausted earlier because it was from one of the first num_shards_to_interleave, it won't be sampled from until the next epoch where its statistics would have been long catastrophically forgotten.

I have not look at the specific code of interleave_dataset, but one tempting solution in webdataset of early iteration termination when the first shard is exhausted, results in later examples of bigger shards never be visited, which is an even bigger problem which is easy to fall into.