https://github.com/commaai/gigashuffle
https://github.com/commaai/gigashuffle
Last synced: about 1 month ago
JSON representation
- Host: GitHub
- URL: https://github.com/commaai/gigashuffle
- Owner: commaai
- Created: 2026-02-18T02:49:29.000Z (6 months ago)
- Default Branch: master
- Last Pushed: 2026-06-18T01:39:06.000Z (about 2 months ago)
- Last Synced: 2026-06-18T03:22:00.538Z (about 2 months ago)
- Language: Python
- Size: 176 KB
- Stars: 3
- Watchers: 0
- Forks: 0
- Open Issues: 1
-
Metadata Files:
- Readme: README.md
Awesome Lists containing this project
README
# gigashuffle
Shuffle-buffer dataloading for PyTorch training.
## Usage
Start a redis server before constructing a loader:
```bash
redis-server
```
Use in your training script:
```python
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 the buffer with `Tensor.share_memory_()` using PyTorch's `file_descriptor` CPU sharing strategy, publishes only an `AF_UNIX` attach-socket path in Redis, and sends the shared tensor objects over that socket; this follows the same fd-transfer mechanism PyTorch uses for CPU tensor IPC. The simpler alternatives do not work well here: a `multiprocessing.Queue` object cannot be shared across independent `torchrun` ranks, Redis cannot transmit process-local file descriptors, explicit `/dev/shm` or `torch.from_file` names can leak after killed workers, and `/proc//fd/` attachment is Linux/container-permission fragile and races owner death.
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.