Fix drop_first checking in partitioning to account for world_size divisibility

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)

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

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

tests/test_partition.py

@@ -38,24 +38,34 @@ 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
         initial_physical_nodes = 3
         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)