PyTorch DataLoader and IterableDataset: Worker-Level Data Shuffling

Ensuring Unique Data Distribution

To prevent multiple workers from yielding the same data shards, you must explicitly partition the dataset within the __iter__ method using torch.utils.data.get_worker_info(). Because IterableDataset does not support the shuffle=True argument in DataLoader, the dataset itself is responsible for determining which subset of data each worker process should stream.

Implementation Steps for Partitioning

1. Retrieve Worker Metadata: Call get_worker_info() inside the __iter__ method to access the id and num_workers of the current process.
2. Calculate Shard Offset: Use the worker ID to slice the data source. For example, if you have a list of file paths, each worker should only iterate over paths where index % num_workers == worker_id.
3. Handle Single-Worker Cases: Ensure the logic defaults to processing the full dataset if num_workers is set to 0 (main process).

Implementing Global Shuffling

True global shuffling is impossible with a pure stream without loading the entire index into memory. The recommended mechanism to approximate this is a Shuffle Buffer. This approach maintains a fixed-size window of samples and yields them randomly, providing local stochasticity that scales with the buffer size.

Recommended Shuffling Workflow

• Buffer Initialization: Create a list (the buffer) and fill it with the first N samples from the partitioned stream.
• Random Sampling: For every new sample requested, randomly select one element from the buffer to yield, then immediately replace that element with the next sample from the stream.
• Worker Seeding: Use a worker_init_fn to ensure each worker's random number generator is seeded differently, preventing synchronized shuffling patterns across processes.

# Conceptual implementation within __iter__
worker_info = torch.utils.data.get_worker_info()
if worker_info is None:  # single-process loading
    iter_start = 0
    iter_end = len(self.data)
else:
    # Partition data
    per_worker = int(math.ceil(len(self.data) / float(worker_info.num_workers)))
    worker_id = worker_info.id
    iter_start = worker_id * per_worker
    iter_end = min(iter_start + per_worker, len(self.data))

# Use a buffer for local shuffling
buffer = []
for i in range(iter_start, iter_end):
    item = self.data[i]
    if len(buffer) < buffer_size:
        buffer.append(item)
    else:
        idx = random.randint(0, buffer_size - 1)
        yield buffer[idx]
        buffer[idx] = item
# Drain remaining buffer
random.shuffle(buffer)
for item in buffer: yield item

Verification and Assumptions

This implementation assumes PyTorch 1.7+ where get_worker_info() is stable. To verify the setup, print the worker_id and the first few sample IDs in each worker; they should be distinct and non-overlapping. Warning: Be cautious with buffer_size; since each worker maintains its own buffer, the total memory consumption is buffer_size * num_workers * sample_size.

Diagnostic Detail Needed: Are you streaming from a single large file (like a TFRecord or WebDataset) or a directory of many small files? The partitioning logic differs significantly between these two sources.