Streaming IterableDataset state stops advancing after first resume — second resume restarts from the beginning

[francesco-bertolotti]

Description

For a streaming IterableDataset, calling load_state_dict() and then continuing to iterate leaves state_dict() frozen at the initial position (shard_idx=0, shard_example_idx=0). Reading resumes correctly the first time, but the tracked state never advances again.

This causes a second consecutive resume to restart the entire data stream from scratch, silently corrupting long training runs that checkpoint and resume more than once.

Root cause

The resume call sites in IterableDataset (_iter_pytorch for num_workers >= 1 and _prepare_ex_iterable_for_iteration for num_workers == 0) set self._state_dict["examples_iterable"] to reference the example-iterable's own state dict, then load the resume state into that example-iterable.

_BaseExamplesIterable.load_state_dict calls self._init_state_dict(), which rebinds self._state_dict to a brand-new object. After the load, dataset._state_dict["examples_iterable"] still points at the stale, zero-initialized state dict. During iteration the example-iterable mutates its new state dict, while dataset.state_dict() keeps deep-copying the old one — so the reported position never moves past zero.

On a fresh run (no resume), the load_state_dict branch is skipped, the reference stays linked, and tracking works correctly — which is why this only surfaces after a resume.

Impact

We hit this training on FineWeb with torchtitan: a resumed-twice run produced a different training loss than a resume-once run, because the second resume silently replayed the dataset from the start.

Fix

Already proposed in #8295 re-points self._state_dict["examples_iterable"] at the live example-iterable state immediately after the resume load, at both resume sites. Includes a regression test (test_resume_dataloader_twice) that fails on main and passes with the fix.

Steps to reproduce the bug

import datasets
def gen():
    for i in range(10): yield {"id": i}
ds = datasets.IterableDataset.from_generator(gen)
full = [x["id"] for x in ds]
it = iter(ds); [next(it) for _ in range(3)]; sd1 = ds.state_dict()
ds.load_state_dict(sd1); it = iter(ds); seen = [next(it)["id"] for _ in range(3)]; sd2 = ds.state_dict()
ds.load_state_dict(sd2); after = [x["id"] for x in ds]
print("reference       :", full)
print("after 1st resume:", seen, "(reads correctly)")
print("after 2nd resume:", after)
print(">>> BUG: 2nd resume restarted from the beginning" if after == full else ">>> OK: continued")

Expected behavior

The dataloader should save the state dict properly so that the next run can be resumed accordingly

Environment info

datasets==5.0.0
python==3.14.0b2