diff --git a/streaming/base/partition/__init__.py b/streaming/base/partition/__init__.py index 28e908cb1..65271d8e2 100644 --- a/streaming/base/partition/__init__.py +++ b/streaming/base/partition/__init__.py @@ -3,6 +3,7 @@ """Apportion shards/samples to nodes/ranks/workers for elastically deterministic sample order.""" +import logging from typing import Optional import numpy as np @@ -11,6 +12,8 @@ from streaming.base.partition.orig import get_partitions_orig from streaming.base.partition.relaxed import get_partitions_relaxed +logger = logging.getLogger(__name__) + algos = { 'orig': get_partitions_orig, 'relaxed': get_partitions_relaxed, @@ -51,6 +54,17 @@ def get_partitions(algo: str, NDArray[np.int64]: Partitions of shape (physical nodes, ranks per node, workers per rank, batches per worker, batch size). """ + world_size = ranks_per_node * num_physical_nodes + num_repeated_samples = world_size - (num_samples % world_size) + if num_samples + num_repeated_samples < drop_first: + raise ValueError(f'Resuming further into the dataset ({drop_first}) than it has samples ' + + f'({num_samples})') + + if num_repeated_samples > 0: + logger.debug(f'Using {num_repeated_samples} repeated samples to ensure that the epoch ' + + f'size is divisible by the number of total devices. This ensures that each ' + + f'device contributes the same number of samples per global batch. ') + get = algos[algo] return get(num_samples, num_canonical_nodes, num_physical_nodes, ranks_per_node, workers_per_rank, batch_size, drop_first, initial_physical_nodes) diff --git a/streaming/base/partition/orig.py b/streaming/base/partition/orig.py index ce8832cf5..cda16ac1d 100644 --- a/streaming/base/partition/orig.py +++ b/streaming/base/partition/orig.py @@ -46,10 +46,6 @@ def get_partitions_orig(num_samples: int, NDArray[np.int64]: Partitions of shape (physical nodes, ranks per node, workers per rank, batches per worker, batch size). """ - if num_samples < drop_first: - raise ValueError(f'Resuming further into the dataset ({drop_first}) than it has samples ' + - f'({num_samples})') - if num_canonical_nodes < num_physical_nodes: if num_physical_nodes % num_canonical_nodes: raise ValueError('Either canonical or physical nodes must be evenly divisible by ' + @@ -81,7 +77,7 @@ def get_partitions_orig(num_samples: int, # For samples to be properly split across canonical nodes, there must be more samples than nodes. # The edge case is when the number of samples is equal to the number of canonical nodes, but this only works when - # there is an equal or greater number of canonical nodes than physical nodes. + # there is an equal or greater number of canonical nodes than physical nodes. # If these conditions are not met, an alternative sampling approach is used that leads to many repeats. if num_samples > num_canonical_nodes or (num_samples == num_canonical_nodes and num_canonical_nodes >= num_physical_nodes): @@ -141,8 +137,7 @@ def get_partitions_orig(num_samples: int, ids = ids.reshape(-1, num_physical_nodes) ids = ids.transpose() - # Interleave the node sample ranges over each node's ranks, padding by repeating the last - # sample. + # Interleave the node sample ranges over each node's ranks, padding with -1 for reshaping. # # ids: (physical nodes, samples per rank, ranks per node). overflow = ids.shape[1] % ranks_per_node diff --git a/streaming/base/partition/relaxed.py b/streaming/base/partition/relaxed.py index 1812b977a..6baa0a48c 100644 --- a/streaming/base/partition/relaxed.py +++ b/streaming/base/partition/relaxed.py @@ -49,10 +49,6 @@ def get_partitions_relaxed(num_samples: int, NDArray[np.int64]: Partitions of shape (physical nodes, ranks per node, workers per rank, batches per worker, batch size). """ - if num_samples < drop_first: - raise ValueError(f'Resuming further into the dataset ({drop_first}) than it has samples ' + - f'({num_samples})') - if initial_physical_nodes is None or (num_physical_nodes <= num_canonical_nodes and num_canonical_nodes % num_physical_nodes == 0) or \ (num_physical_nodes > num_canonical_nodes and diff --git a/tests/test_partition.py b/tests/test_partition.py index aa26a63d1..42cfaa1f6 100644 --- a/tests/test_partition.py +++ b/tests/test_partition.py @@ -38,16 +38,22 @@ def test_partition_walk(partition_algo: str): assert x.shape == (22, 8, 8, 1, 10) -@pytest.mark.parametrize('num_samples', [400, 1000]) -@pytest.mark.parametrize('num_canonical_nodes', [1, 4]) -@pytest.mark.parametrize('num_physical_nodes', [1, 4]) +@pytest.mark.parametrize('num_samples', [405, 812, 1111]) +@pytest.mark.parametrize('num_canonical_nodes', [1, 2]) +@pytest.mark.parametrize('num_physical_nodes', [2, 8]) @pytest.mark.parametrize('ranks_per_node', [1, 8]) @pytest.mark.parametrize('workers_per_rank', [1, 8]) @pytest.mark.parametrize('batch_size', [4]) @pytest.mark.parametrize('partition_algo', ['orig', 'relaxed']) -def test_partition_drop_all(num_samples: int, num_canonical_nodes: int, num_physical_nodes: int, - ranks_per_node: int, workers_per_rank: int, batch_size: int, - partition_algo: str): +def test_partition_drop_all( + num_samples: int, + num_canonical_nodes: int, + num_physical_nodes: int, + ranks_per_node: int, + workers_per_rank: int, + batch_size: int, + partition_algo: str, +): initial_physical_nodes = None if partition_algo == 'relaxed' and num_canonical_nodes == 4 and ranks_per_node == 8: num_canonical_nodes = 3 @@ -55,7 +61,11 @@ def test_partition_drop_all(num_samples: int, num_canonical_nodes: int, num_phys batch_size = batch_size * 3 num_samples = 3 * num_samples - drop_first = num_samples + # Partitioning should repeat samples so that the epoch size is divisible by the world size. + # To drop all samples, we need to drop all repeated samples as well. + world_size = num_physical_nodes * ranks_per_node + num_repeated_samples = world_size - (num_samples % world_size) + drop_first = num_samples + num_repeated_samples x = get_partitions(partition_algo, num_samples, num_canonical_nodes, num_physical_nodes, ranks_per_node, workers_per_rank, batch_size, drop_first, @@ -77,7 +87,14 @@ def test_partition_invalid_drop_first(num_samples: int, drop_additional: int, num_canonical_nodes: int, num_physical_nodes: int, ranks_per_node: int, workers_per_rank: int, batch_size: int, partition_algo: str): - drop_first = num_samples + drop_additional + + # Partitioning should repeat samples so that the epoch size is divisible by the world size. + # For `drop_first` to be invalid, we need to exceed the number of unique samples plus the + # number of repeated samples. + world_size = num_physical_nodes * ranks_per_node + num_repeated_samples = world_size - (num_samples % world_size) + drop_first = num_samples + num_repeated_samples + drop_additional + with pytest.raises(ValueError, match=f'Resuming further into the dataset*'): _ = get_partitions(partition_algo, num_samples, num_canonical_nodes, num_physical_nodes, ranks_per_node, workers_per_rank, batch_size, drop_first)