Skip to content

Samplers

samudra.utils.samplers

EquivalenceGroupBatchSampler(groups, batch_size, shuffle=True, drop_last=False, *, seed)

Bases: Sampler[Batch]

Groups indices into equivalence classes, batches within groups, and optionally shuffles.

This sampler partitions dataset indices into groups. It creates batches within each group, then chains them together. When shuffle=True, batches are globally shuffled each epoch to avoid sequential group processing.

Parameters:

Name Type Description Default
groups list[list[int]]

List of index lists, where each inner list contains indices belonging to the same equivalence group.

 required
batch_size int

Number of samples per batch

 required
shuffle bool

Whether to shuffle indices within groups and shuffle batches globally

 True
drop_last bool

Whether to drop incomplete batches at the end of each group

 False
seed int

Random seed for deterministic shuffling

 required
Source code in src/samudra/utils/samplers.py 
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
def __init__(
    self,
    groups: list[list[int]],
    batch_size: int,
    shuffle: bool = True,
    drop_last: bool = False,
    *,
    seed: int,
):
    super().__init__()
    self.groups = groups
    self.batch_size = batch_size
    self.shuffle = shuffle
    self.drop_last = drop_last
    self.seed = seed
    self.epoch = 0

    self._samplers = [
        BatchSampler(
            _SimpleSubsetSampler(group),
            batch_size=self.batch_size,
            drop_last=self.drop_last,
        )
        for group in self.groups
    ]

set_epoch(epoch)

Set the epoch for deterministic shuffling.

Source code in src/samudra/utils/samplers.py 
87
88
89
def set_epoch(self, epoch: int) -> None:
    """Set the epoch for deterministic shuffling."""
    self.epoch = epoch

from_dataset_sizes(dataset_sizes, batch_size, shuffle=True, drop_last=False, *, seed) classmethod

Create sampler from dataset sizes, treating each as a contiguous group.

Parameters:

Name Type Description Default
dataset_sizes list[int]

List of individual dataset sizes. Groups are created based on cumulative boundaries, where each dataset forms its own equivalence group.

 required
batch_size int

Number of samples per batch

 required
shuffle bool

Whether to shuffle indices within groups and shuffle batches globally

 True
drop_last bool

Whether to drop incomplete batches at the end of each group

 False
seed int

Random seed for deterministic shuffling

 required
Source code in src/samudra/utils/samplers.py 
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
@classmethod
def from_dataset_sizes(
    cls,
    dataset_sizes: list[int],
    batch_size: int,
    shuffle: bool = True,
    drop_last: bool = False,
    *,
    seed: int,
) -> Self:
    """Create sampler from dataset sizes, treating each as a contiguous group.

    Args:
        dataset_sizes: List of individual dataset sizes. Groups are created based on
            cumulative boundaries, where each dataset forms its own equivalence group.
        batch_size: Number of samples per batch
        shuffle: Whether to shuffle indices within groups and shuffle batches globally
        drop_last: Whether to drop incomplete batches at the end of each group
        seed: Random seed for deterministic shuffling
    """
    cumsum = 0
    groups = []
    for size in dataset_sizes:
        groups.append(list(range(cumsum, cumsum + size)))
        cumsum += size
    return cls(groups, batch_size, shuffle, drop_last, seed=seed)

from_datasets(datasets, batch_size, shuffle, drop_last, *, seed) classmethod

Batch datasets with equal public compatibility keys.

Groups follow first-seen dataset order, which must be the same on every DDP rank. Keys need only support equality and hashing, not ordering.

Source code in src/samudra/utils/samplers.py 
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
@classmethod
def from_datasets(
    cls,
    datasets: list[DatasetT],
    batch_size: int,
    shuffle: bool,
    drop_last: bool,
    *,
    seed: int,
) -> Self:
    """Batch datasets with equal public compatibility keys.

    Groups follow first-seen dataset order, which must be the same on every
    DDP rank. Keys need only support equality and hashing, not ordering.
    """
    # Preserve first-seen group order. Compatibility keys define equality,
    # but need not be comparable or identical objects across processes.
    groups: dict[Hashable, list[int]] = {}

    cumsum = 0
    for ds in datasets:
        key = ds.batch_compatibility_key
        assert isinstance(key, Hashable), (
            "`batch_compatibility_key` must be hashable."
        )
        groups.setdefault(key, []).extend(range(cumsum, cumsum + len(ds)))
        cumsum += len(ds)

    return cls(list(groups.values()), batch_size, shuffle, drop_last, seed=seed)

__len__()

Calculate total number of batches across all groups.

Source code in src/samudra/utils/samplers.py 
166
167
168
169
170
171
172
def __len__(self):
    """Calculate total number of batches across all groups."""
    total_batches = 0
    for sampler in self._samplers:
        total_batches += len(sampler)

    return total_batches

DistributedEquivalenceGroupBatchSampler(datasets, batch_size, num_replicas, rank, shuffle=True, drop_last=False, seed=0)

Bases: Sampler[Batch]

Distributed version of EquivalenceGroupBatchSampler for multi-GPU training.

Uses composition to delegate batching logic to EquivalenceGroupBatchSampler, handling only the distribution and epoch-based shuffling.

Ensures uniform batch counts across all ranks to prevent DDP hangs at collective sync points. Each equivalence group is chunked into logical DDP steps of num_replicas batches, so ranks process the same group at the same step. For an incomplete per-group step: - drop_last=True: drops that incomplete group step - drop_last=False: pads it by duplicating batches from the same group

Note: Compared to the non-distributed sampler: this one won't shuffle within batches, only between batches, when shuffle=True.

Parameters:

Name Type Description Default
datasets list[DatasetT]

List of TorchTrainDataset instances to group

 required
batch_size int

Number of samples per batch

 required
num_replicas int

Number of distributed workers (world size)

 required
rank int

Index of current worker (0 to num_replicas-1)

 required
shuffle bool

Whether to shuffle batches

 True
drop_last bool

Whether to drop incomplete batches within each group, and whether to trim (vs pad) when distributing batches across ranks

 False
seed int

Random seed for shuffling (default: 0)

 0
Source code in src/samudra/utils/samplers.py 
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
def __init__(
    self,
    datasets: list[DatasetT],
    batch_size: int,
    num_replicas: int,
    rank: int,
    shuffle: bool = True,
    drop_last: bool = False,
    seed: int = 0,
):
    super().__init__()
    if num_replicas <= 0:
        raise ValueError(f"num_replicas must be positive, got {num_replicas}.")
    if rank >= num_replicas or rank < 0:
        raise ValueError(
            f"Invalid rank {rank}, must be in range [0, {num_replicas})"
        )

    self.num_replicas = num_replicas
    self.rank = rank
    self.shuffle = shuffle
    self.drop_last = drop_last
    self.seed = seed
    self.epoch = 0

    # Delegate batching logic to inner sampler (without shuffle for determinism)
    self._inner = EquivalenceGroupBatchSampler.from_datasets(
        datasets=datasets,
        batch_size=batch_size,
        shuffle=False,  # We handle shuffling with seeded RNG
        drop_last=drop_last,
        seed=seed,
    )

set_epoch(epoch)

Set the epoch for deterministic shuffling across workers.

Source code in src/samudra/utils/samplers.py 
236
237
238
def set_epoch(self, epoch: int) -> None:
    """Set the epoch for deterministic shuffling across workers."""
    self.epoch = epoch

__len__()

Number of batches for this worker (same for all ranks).

Source code in src/samudra/utils/samplers.py 
312
313
314
315
316
317
318
319
320
321
322
323
324
325
def __len__(self):
    """Number of batches for this worker (same for all ranks)."""
    n = self.num_replicas
    total = 0
    for sampler in self._inner._samplers:
        group_len = len(sampler)
        if self.drop_last:
            total += (group_len // n) * n
        else:
            total += math.ceil(group_len / n) * n
    if self.drop_last:
        return total // n
    else:
        return math.ceil(total / n)