gigashuffle

Shuffle-buffer dataloading for PyTorch training.

Usage

Use in your training script:

import os
import torch
import torch.distributed as dist
from torch.utils.data import IterableDataset
from gigashuffle import DataloaderConfig, MultiprocessShuffledDataloader

class MyDataset(IterableDataset):
  def __iter__(self):
    while True:
      x = torch.randn(64, 3, 224, 224)
      y = torch.randint(0, 1000, (64,))
      yield [{'x': x, 'y': y}]

config = DataloaderConfig(
  bs=32,
  shuffle_size=1000,
  num_writers=2,
  num_readers=2,
  local_rank=int(os.environ.get('LOCAL_RANK', '0')),
  global_rank=int(os.environ.get('RANK', '0')),
  local_world_size=int(os.environ.get('LOCAL_WORLD_SIZE', '1')),
  global_world_size=int(os.environ.get('WORLD_SIZE', '1')),
  queue_name='my-training-loader', # must be unique across your dataloaders
  fill_once=False, # see below
)

loader = MultiprocessShuffledDataloader(MyDataset(), config=config)
dummy_batch = loader.get_dummy_batch()

for batch in loader:
  x = batch[0]['x']
  y = batch[0]['y']
  break

Datasets yield a Buffer: list[dict[str, Tensor | ndarray]]. Every tensor or array in a sample must share the same first dimension; that is the input batch size. config.bs is the output batch size.

Call loader.get_dummy_batch() when you need a sanity-check batch before the shuffle buffer reaches min_mixing. It repeats the initial sample to config.bs and does not advance the normal shuffled iterator.

Set config.fill_once=True, config.min_mixing=1, and config.num_readers=1 to populate the shuffle buffer once, then yield one ordered pass over it without returning indices to the writers. Reader waits for the full shuffle_size before yielding batches in this mode.

By default (config.evict_on_read=True) a buffer index is returned to the writers as soon as it is read, so each sample is consumed once. Set config.evict_on_read=False to keep read entries in the buffer for re-sampling and have the client manage eviction with loader.evict(indices). Never evicting will freeze the buffer once it fills, and evicting faster than writers can refill will make reads wait. Every batch carries the buffer index of each row in batch[0][INDEX_KEY] (a (bs,) int64 tensor). config.evict_on_read=False is not supported with fill_once=True.

Notes

Each Dataloader owns one shared CPU shuffle buffer. The owner writer (proc_idx=0) allocates it with Tensor.share_memory_() and hosts a local AF_UNIX coordinator socket for attachment and queue state.

The output batches are views into reusable shared-memory reader buffers, not copies. Their contents may be overwritten after the iterator advances or the batch is released, so callers that need to keep a batch must clone/copy it before requesting another batch or dropping the iterator.