Skip to content

Restore memory-mapped datasets when unpickling - #8517

Open
vineethsaivs wants to merge 1 commit into
deepspeedai:masterfrom
vineethsaivs:fix-mmap-dataset-unpickle-20260914
Open

vineethsaivs wants to merge 1 commit into
deepspeedai:masterfrom
vineethsaivs:fix-mmap-dataset-unpickle-20260914

Conversation

@vineethsaivs

Copy link
Copy Markdown
Contributor

Unpickling MMapIndexedDataset raises TypeError, so a DataLoader with spawned workers cannot read it.

Cause: setstate passes only the saved dataset path to _do_init, whose skip_warmup argument has no default.

Fix: Give _do_init the same default warmup behavior as the public constructor.

Test: DS_ACCELERATOR=cpu python -m pytest tests/unit/runtime/test_data.py -k 'mmap or repeating_loader' -q. Three regressions fail before the fix. Four targeted tests pass after, covering signed and unsigned dataset round-trips, a real spawned DataLoader, and the existing repeating-loader control. Validated on Apple Silicon CPU with PyTorch 2.10.0. No GPU or distributed training run. Changed-file pre-commit passes.

Give _do_init the same default warmup behavior as the public constructor.

Signed-off-by: Vineeth Sai <vineethsai4444@gmail.com>
self._do_init(state)

def _do_init(self, path, skip_warmup):
def _do_init(self, path, skip_warmup=False):

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I ran this against 344e22bb10e7f97c339040e0c4e6729d427d4c2c in a clean python:3.11-slim container (torch 2.8.0+cpu, requirements/requirements.txt). The round trip works and the new tests pin what they say.

The default fixes the crash but drops the flag. __getstate__ on line 488 returns self._path alone, so a dataset opened with skip_warmup=True comes back from the round trip with warmup on. I traced the real _warmup_mmap_file (the wrapper calls through, it does not replace it) on a 4.9 MB corpus:

construct with skip_warmup=True  -> warmup calls: []
pickle round trip of that object -> warmup calls: [('corpus.idx', 240050), ('corpus.bin', 5120000)]
__getstate__ returns: '/tmp/ds/corpus'
round trip correct?  len 20000 == 20000, item equal True

That matters here because every construction site in the tree passes skip_warmup=True: data_sampler.py:95,100,221,250,348, data_analyzer.py:238,252,354,403,421,422, variable_batch_size_and_lr.py:398. So on the spawn path your second test covers, each worker sweeps the whole .bin sequentially in 100 MB chunks, which is the cost those call sites set the flag to avoid. It is quiet apart from two warming up ... prints per worker.

Carrying the flag is about the same size as the default:

def __getstate__(self):
    return self._path, self._skip_warmup

def __setstate__(self, state):
    self._do_init(*state)

plus self._skip_warmup = skip_warmup in _do_init. Nothing else in the tree pickles this class, so the state shape stays contained to this file and your new test.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants