Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Improve skip take shuffling and distributed #6965

Merged
merged 2 commits into from
Jun 24, 2024
Merged

Improve skip take shuffling and distributed #6965

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
f84a2fa
improve skip take shuffling
lhoestq Jun 12, 2024
768cb35
improve skip take for distributed
lhoestq Jun 12, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
103 changes: 88 additions & 15 deletions src/datasets/iterable_dataset.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1306,10 +1306,18 @@ def n_shards(self) -> int:


class SkipExamplesIterable(_BaseExamplesIterable):
def __init__(self, ex_iterable: _BaseExamplesIterable, n: int):
def __init__(
self,
ex_iterable: _BaseExamplesIterable,
n: int,
block_sources_order_when_shuffling: bool = True,
split_when_sharding: bool = True,
):
super().__init__()
self.ex_iterable = ex_iterable
self.n = n
self.block_sources_order_when_shuffling = block_sources_order_when_shuffling
self.split_when_sharding = split_when_sharding
# TODO(QL): implement iter_arrow

def _init_state_dict(self) -> dict:
Expand All @@ -1322,20 +1330,57 @@ def __iter__(self):
self._state_dict["skipped"] = True
yield from islice(self.ex_iterable, ex_iterable_idx_start, None)

@staticmethod
def split_number(num, n):
quotient = num // n
remainder = num % n
result = [quotient] * n
for i in range(remainder):
result[i] += 1
return result

def shuffle_data_sources(self, generator: np.random.Generator) -> "SkipExamplesIterable":
"""Doesn't shuffle the wrapped examples iterable since it would skip examples from other shards instead."""
return self
"""May not shuffle the wrapped examples iterable since it would skip examples from other shards instead."""
if self.block_sources_order_when_shuffling:
return self
else:
return SkipExamplesIterable(
self.ex_iterable.shuffle_data_sources(generator),
n=self.n,
block_sources_order_when_shuffling=self.block_sources_order_when_shuffling,
split_when_sharding=self.split_when_sharding,
)

def shard_data_sources(self, worker_id: int, num_workers: int) -> "SkipExamplesIterable":
"""Keep only the requested shard."""
if self.split_when_sharding:
return SkipExamplesIterable(
self.ex_iterable.shard_data_sources(worker_id, num_workers),
n=self.split_number(self.n, num_workers)[worker_id],
block_sources_order_when_shuffling=self.block_sources_order_when_shuffling,
split_when_sharding=self.split_when_sharding,
)
else:
return self

@property
def n_shards(self) -> int:
return self.ex_iterable.n_shards


class TakeExamplesIterable(_BaseExamplesIterable):
def __init__(self, ex_iterable: _BaseExamplesIterable, n: int):
def __init__(
self,
ex_iterable: _BaseExamplesIterable,
n: int,
block_sources_order_when_shuffling: bool = True,
split_when_sharding: bool = True,
):
super().__init__()
self.ex_iterable = ex_iterable
self.n = n
self.block_sources_order_when_shuffling = block_sources_order_when_shuffling
self.split_when_sharding = split_when_sharding
# TODO(QL): implement iter_arrow

def _init_state_dict(self) -> dict:
Expand All @@ -1349,10 +1394,6 @@ def __iter__(self):
self._state_dict["num_taken"] += 1
yield key_example

def shuffle_data_sources(self, generator: np.random.Generator) -> "TakeExamplesIterable":
"""Doesn't shuffle the wrapped examples iterable since it would take examples from other shards instead."""
return self

@staticmethod
def split_number(num, n):
quotient = num // n
Expand All @@ -1362,12 +1403,34 @@ def split_number(num, n):
result[i] += 1
return result

def shuffle_data_sources(self, generator: np.random.Generator) -> "TakeExamplesIterable":
"""May not shuffle the wrapped examples iterable since it would take examples from other shards instead."""
if self.block_sources_order_when_shuffling:
return self
else:
return TakeExamplesIterable(
self.ex_iterable.shuffle_data_sources(generator),
n=self.n,
block_sources_order_when_shuffling=self.block_sources_order_when_shuffling,
split_when_sharding=self.split_when_sharding,
)

def shard_data_sources(self, worker_id: int, num_workers: int) -> "TakeExamplesIterable":
"""Keep only the requested shard."""
return TakeExamplesIterable(
self.ex_iterable.shard_data_sources(worker_id, num_workers),
n=self.split_number(self.n, num_workers)[worker_id],
)
if self.split_when_sharding:
return TakeExamplesIterable(
self.ex_iterable.shard_data_sources(worker_id, num_workers),
n=self.split_number(self.n, num_workers)[worker_id],
block_sources_order_when_shuffling=self.block_sources_order_when_shuffling,
split_when_sharding=self.split_when_sharding,
)
else:
return TakeExamplesIterable(
self.ex_iterable.shard_data_sources(worker_id, num_workers),
n=self.n,
block_sources_order_when_shuffling=self.block_sources_order_when_shuffling,
split_when_sharding=self.split_when_sharding,
)

@property
def n_shards(self) -> int:
Expand Down Expand Up @@ -2254,7 +2317,7 @@ def shuffle(
return IterableDataset(
ex_iterable=BufferShuffledExamplesIterable(
self._ex_iterable, buffer_size=buffer_size, generator=generator
).shuffle_data_sources(generator),
),
info=self._info.copy(),
split=self._split,
formatting=self._formatting,
Expand Down Expand Up @@ -2294,7 +2357,12 @@ def skip(self, n: int) -> "IterableDataset":
'text': 'if you sometimes like to go to the movies to have fun , wasabi is a good place to start .'}]
```
"""
ex_iterable = SkipExamplesIterable(self._ex_iterable, n)
ex_iterable = SkipExamplesIterable(
self._ex_iterable,
n,
block_sources_order_when_shuffling=self._shuffling is None,
split_when_sharding=self._distributed is None,
)
return IterableDataset(
ex_iterable=ex_iterable,
info=self._info.copy(),
Expand Down Expand Up @@ -2326,7 +2394,12 @@ def take(self, n: int) -> "IterableDataset":
'text': 'the gorgeously elaborate continuation of " the lord of the rings " trilogy is so huge that a column of words cannot adequately describe co-writer/director peter jackson\'s expanded vision of j . r . r . tolkien\'s middle-earth .'}]
```
"""
ex_iterable = TakeExamplesIterable(self._ex_iterable, n)
ex_iterable = TakeExamplesIterable(
self._ex_iterable,
n,
block_sources_order_when_shuffling=self._shuffling is None,
split_when_sharding=self._distributed is None,
)
return IterableDataset(
ex_iterable=ex_iterable,
info=self._info.copy(),
Expand Down
58 changes: 51 additions & 7 deletions tests/test_iterable_dataset.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -10,6 +10,7 @@

from datasets import Dataset, load_dataset
from datasets.combine import concatenate_datasets, interleave_datasets
from datasets.distributed import split_dataset_by_node
from datasets.features import (
ClassLabel,
Features,
Expand Down Expand Up @@ -1653,17 +1654,60 @@ def test_iterable_dataset_take(dataset: IterableDataset, n):


@pytest.mark.parametrize("method", ["skip", "take"])
def test_iterable_dataset_shuffle_after_skip_or_take(method):
@pytest.mark.parametrize("after_shuffle", [False, True])
@pytest.mark.parametrize("count", [2, 5, 11])
def test_iterable_dataset_skip_or_take_after_shuffle(method, after_shuffle, count):
seed = 42
n, n_shards = 3, 10
count = 7
ex_iterable = ExamplesIterable(generate_examples_fn, {"n": n, "filepaths": [f"{i}.txt" for i in range(n_shards)]})
dataset = IterableDataset(ex_iterable)
dataset = dataset.skip(n) if method == "skip" else dataset.take(count)
shuffled_dataset = dataset.shuffle(seed, buffer_size=DEFAULT_N_EXAMPLES)
# shuffling a skip/take dataset should keep the same examples and don't shuffle the shards
key = lambda x: f"{x['filepath']}_{x['id']}" # noqa: E731
assert sorted(dataset, key=key) == sorted(shuffled_dataset, key=key)
shuffled_dataset = dataset
if after_shuffle:
shuffled_dataset = shuffled_dataset.shuffle(seed, buffer_size=DEFAULT_N_EXAMPLES)
shuffled_dataset = shuffled_dataset.skip(count) if method == "skip" else shuffled_dataset.take(count)
# skip/take a shuffled dataset should not keep the same examples and shuffle the shards
key = lambda x: f"{x['filepath']}_{x['id']}" # noqa: E731
assert (len(list(dataset)) - count if method == "skip" else count) == len(list(shuffled_dataset))
assert sorted(list(dataset)[count:] if method == "skip" else list(dataset)[:count], key=key) != sorted(
shuffled_dataset, key=key
)
else:
shuffled_dataset = shuffled_dataset.skip(count) if method == "skip" else shuffled_dataset.take(count)
shuffled_dataset = shuffled_dataset.shuffle(seed, buffer_size=DEFAULT_N_EXAMPLES)
# shuffling a skip/take dataset should keep the same examples and don't shuffle the shards
key = lambda x: f"{x['filepath']}_{x['id']}" # noqa: E731
assert (len(list(dataset)) - count if method == "skip" else count) == len(list(shuffled_dataset))
assert sorted(list(dataset)[count:] if method == "skip" else list(dataset)[:count], key=key) == sorted(
shuffled_dataset, key=key
)


@pytest.mark.parametrize("method", ["skip", "take"])
@pytest.mark.parametrize("after_split_by_node", [False, True])
@pytest.mark.parametrize("count", [2, 5, 11])
def test_iterable_dataset_skip_or_take_after_split_by_node(method, after_split_by_node, count):
n, n_shards = 3, 10
rank, world_size = 1, 2
ex_iterable = ExamplesIterable(generate_examples_fn, {"n": n, "filepaths": [f"{i}.txt" for i in range(n_shards)]})
dataset = IterableDataset(ex_iterable)
distributed_dataset = dataset
true_distributed_dataset = split_dataset_by_node(dataset, rank=rank, world_size=world_size)
if after_split_by_node:
distributed_dataset = split_dataset_by_node(distributed_dataset, rank=rank, world_size=world_size)
distributed_dataset = distributed_dataset.skip(count) if method == "skip" else distributed_dataset.take(count)
assert (
list(true_distributed_dataset)[count:]
Copy link

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

this looks good to me!

if method == "skip"
else list(true_distributed_dataset)[:count] == list(distributed_dataset)
)
else:
distributed_dataset = distributed_dataset.skip(count) if method == "skip" else distributed_dataset.take(count)
distributed_dataset = split_dataset_by_node(distributed_dataset, rank=rank, world_size=world_size)
assert len(
list(true_distributed_dataset)[count // world_size :]
Copy link

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I don't know the implementation of split_dataset_by_node very well, just trying to raise a concern on this test case and please correct me if I'm wrong.

Assume world_size = 2, count = 1. Assume underlying data is [0, 1, 2, ..., 9].
true_distributed_dataset on rank 0 would be [0, 2, 4, 6, 8]
true_distributed_dataset on rank 1 would be [1, 3, 5, 7, 9]
and after calling skip and split_dataset_by_node
distributed_dataset on rank 0 would be [1, 3, 5, 7, 9]
distributed_dataset on rank 1 would be [2, 4, 6, 8]

Copy link
Member Author

@lhoestq lhoestq Jun 24, 2024 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Yes that's correct. The test doesn't ensure that some examples are not skipped if count is not a factor of world_size at the moment, bu we can improve that later

if method == "skip"
else list(true_distributed_dataset)[: count // world_size]
) == len(list(distributed_dataset))


def test_iterable_dataset_add_column(dataset_with_several_columns: IterableDataset):
Expand Down
Loading