interleave_datasets seed with multiple workers

[jonathanasdf]

Describe the bug

Using interleave_datasets with multiple dataloader workers and a seed set causes the same dataset sampling order across all workers.

Should the seed be modulated with the worker id?

Steps to reproduce the bug

See above

Expected behavior

See above

Environment info

• datasets version: 3.5.1
• Platform: macOS-15.4.1-arm64-arm-64bit
• Python version: 3.12.9
• huggingface_hub version: 0.30.2
• PyArrow version: 19.0.1
• Pandas version: 2.2.3
• fsspec version: 2024.12.0

[lhoestq]

Hi ! It's already the case IIRC: the effective seed looks like seed + worker_id. Do you have a reproducible example ?

[jonathanasdf]

here is an example with shuffle

import itertools
import datasets
import multiprocessing
import torch.utils.data


def gen(shard):
  worker_info = torch.utils.data.get_worker_info()
  for i in range(10):
    yield {'value': i, 'worker_id': worker_info.id}


def main():
  ds = datasets.IterableDataset.from_generator(gen, gen_kwargs={'shard': list(range(8))})
  ds = ds.shuffle(buffer_size=100, seed=1234)
  dataloader = torch.utils.data.DataLoader(ds, batch_size=None, num_workers=8)
  for i, ex in enumerate(itertools.islice(dataloader, 50)):
    print(i, ex)


if __name__ == '__main__':
  multiprocessing.set_start_method('spawn')
  main()

python test.py
0 {'value': 8, 'worker_id': 0}
1 {'value': 8, 'worker_id': 1}
2 {'value': 8, 'worker_id': 2}
3 {'value': 8, 'worker_id': 3}
4 {'value': 8, 'worker_id': 4}
5 {'value': 8, 'worker_id': 5}
6 {'value': 8, 'worker_id': 6}
7 {'value': 8, 'worker_id': 7}
8 {'value': 9, 'worker_id': 0}
9 {'value': 9, 'worker_id': 1}
10 {'value': 9, 'worker_id': 2}
11 {'value': 9, 'worker_id': 3}
12 {'value': 9, 'worker_id': 4}
13 {'value': 9, 'worker_id': 5}
14 {'value': 9, 'worker_id': 6}
15 {'value': 9, 'worker_id': 7}
16 {'value': 5, 'worker_id': 0}
17 {'value': 5, 'worker_id': 1}
18 {'value': 5, 'worker_id': 2}
19 {'value': 5, 'worker_id': 3}

[jonathanasdf]

With interleave_datasets

import itertools
import datasets
import multiprocessing
import torch.utils.data


def gen(shard, value):
  while True:
    yield {'value': value}


def main():
  ds = [
    datasets.IterableDataset.from_generator(gen, gen_kwargs={'shard': list(range(8)), 'value': i})
    for i in range(10)
  ]
  ds = datasets.interleave_datasets(ds, probabilities=[1 / len(ds)] * len(ds), seed=1234)
  dataloader = torch.utils.data.DataLoader(ds, batch_size=None, num_workers=8)
  for i, ex in enumerate(itertools.islice(dataloader, 50)):
    print(i, ex)


if __name__ == '__main__':
  multiprocessing.set_start_method('spawn')
  main()

python test.py
0 {'value': 9}
1 {'value': 9}
2 {'value': 9}
3 {'value': 9}
4 {'value': 9}
5 {'value': 9}
6 {'value': 9}
7 {'value': 9}
8 {'value': 3}
9 {'value': 3}
10 {'value': 3}
11 {'value': 3}
12 {'value': 3}
13 {'value': 3}
14 {'value': 3}
15 {'value': 3}
16 {'value': 9}
17 {'value': 9}
18 {'value': 9}
19 {'value': 9}
20 {'value': 9}
21 {'value': 9}
22 {'value': 9}
23 {'value': 9}

[jonathanasdf]

Same results after updating to datasets 3.6.0.

[lhoestq]

Ah my bad, shuffle() uses a global effective seed which is something like seed + epoch, which is used to do the same shards shuffle in each worker so that each worker have a non-overlapping set of shards:

datasets/src/datasets/iterable_dataset.py

Lines 2102 to 2111 in b9efdc6

	def _effective_generator(self):
	if self._shuffling and self.epoch == 0:
	return self._shuffling.generator
	elif self._shuffling:
	# Create effective seed using self.epoch (we subtract in order to avoir overflow in long_scalars)
	effective_seed = deepcopy(self._shuffling.generator).integers(0, 1 << 63) - self.epoch
	effective_seed = (1 << 63) + effective_seed if effective_seed < 0 else effective_seed
	return np.random.default_rng(effective_seed)
	else:
	raise ValueError("This dataset is not shuffled")

I think we should take into account the worker_id in a local seed for the buffer right after this line:

datasets/src/datasets/iterable_dataset.py

Lines 2151 to 2153 in b9efdc6

	ex_iterable = ex_iterable.shard_data_sources(
	num_shards=worker_info.num_workers, index=worker_info.id, contiguous=False
	)

like adding a new step that would propagate in the examples iterables or something like that:

ex_iterable = ex_iterable.shift_rngs(value=worker_id)

is this something you'd like to explore ? contributions on this subject are very welcome

[jonathanasdf]

Potentially, but busy. If anyone wants to take this up please feel free to, otherwise I may or may not revisit when I have free time.

For what it's worth I got around this with

class SeedGeneratorWithWorkerIterable(iterable_dataset._BaseExamplesIterable):
  """ExamplesIterable that seeds the rng with worker id."""

  def __init__(
    self,
    ex_iterable: iterable_dataset._BaseExamplesIterable,
    generator: np.random.Generator,
    rank: int = 0,
  ):
    """Constructor."""
    super().__init__()
    self.ex_iterable = ex_iterable
    self.generator = generator
    self.rank = rank

  def _init_state_dict(self) -> dict:
    self._state_dict = self.ex_iterable._init_state_dict()
    return self._state_dict

  def __iter__(self):
    """Data iterator."""
    effective_seed = copy.deepcopy(self.generator).integers(0, 1 << 63) - self.rank
    effective_seed = (1 << 63) + effective_seed if effective_seed < 0 else effective_seed
    generator = np.random.default_rng(effective_seed)
    self.ex_iterable = self.ex_iterable.shuffle_data_sources(generator)
    if self._state_dict:
      self._state_dict = self.ex_iterable._init_state_dict()
    yield from iter(self.ex_iterable)

  def shuffle_data_sources(self, generator):
    """Shuffle data sources."""
    ex_iterable = self.ex_iterable.shuffle_data_sources(generator)
    return SeedGeneratorWithWorkerIterable(ex_iterable, generator=generator, rank=self.rank)

  def shard_data_sources(self, num_shards: int, index: int, contiguous=True):  # noqa: FBT002
    """Shard data sources."""
    ex_iterable = self.ex_iterable.shard_data_sources(num_shards, index, contiguous=contiguous)
    return SeedGeneratorWithWorkerIterable(ex_iterable, generator=self.generator, rank=index)

  @property
  def is_typed(self):
    return self.ex_iterable.is_typed

  @property
  def features(self):
    return self.ex_iterable.features

  @property
  def num_shards(self) -> int:
    """Number of shards."""
    return self.ex_iterable.num_shards