Dataloader is restarting on the second resume

[francesco-bertolotti]

Bug description

Dataloader restarts from scratch after resuming twice (upstream datasets bug)

This is not a bug in torchtitan. It originates in IterableDataset from Hugging Facedatasets, but it affects torchtitan because torchtitan relies on IterableDataset for streaming data.

The bug

After a resume, IterableDataset stops tracking its own state. Concretely, load_state_dict() calls _init_state_dict(), which rebinds the example-iterable's _state_dict to a fresh object. The outer IterableDataset._state_dict["examples_iterable"] still points at the original, now stale, zero-initialized dict. During iteration the example-iterable mutates the live dict, while state_dict() keeps deep-copying the stale one, so the reported position stays frozen at zero.

The practical consequence shows up on the second consecutive resume of a training run:

1. The original run tracks state correctly and writes a healthy checkpoint.
2. Resuming from it reads from the right place, but every checkpoint it writes contains a stale, zero-initialized dataloader state_dict.
3. Resuming from that checkpoint restarts the data stream from the very beginning.

Full description: huggingface/datasets#8308

Affected versions

I believe that the regression was introduced in datasets==5.0.0. The rebinding _init_state_dict() call was added to the base load_state_dict in that release. Every release <= 4.8.5 (including 3.6.0) should be unaffected — those versions mutate the state dict in place and never break the reference.

The datasets fix

I authored a fix that re-links the live state after the resume load, at both IterableDataset resume sites. It has been merged on datasets main: huggingface/datasets#8295

Unfortunately, the fix is not yet in any published datasets release, so pip install -U datasets does not resolve the issue today.

Proposed workarounds

Until a datasets release ships the fix, any of the following work. Listed in the order I'd recommend for torchtitan:

1. Pin datasets < 5.0.0. Cleanest option if torchtitan's other requirements tolerate it, no patch to maintain and no dependency on a moving branch. (Worth confirming against torchtitan's current datasets floor; if something already requires >= 5.0.0, this is off the table.)

2. Pin the merged upstream commit. Prefer the merge commit SHA over @main for reproducible builds, since main drifts and pip/uv cache VCS installs aggressively:

   # pyproject.toml
   [project]
   dependencies = [
       "datasets @ git+https://github.com/huggingface/datasets.git@<merge-commit-sha>",
   ]

3. Temporary monkey-patch of IterableDataset. Good stopgap that keeps the released datasets pin intact. I have a tested patch and am happy to open it as an interim guard to be removed once a fixed release lands. Let me know if you'd like it.

---

Happy to open a PR for either.

Hope this helps.

Versions

torch == 2.14.0.dev20260612+cu126
datasets  == 5.0.0

[tianyu-l]

Pin datasets < 5.0.0. Cleanest option if torchtitan's other requirements tolerate it, no patch to maintain and no dependency on a moving branch. (Worth confirming against torchtitan's current datasets floor; if something already requires >= 5.0.0, this is off the table.)

https://github.com/pytorch/torchtitan/blob/main/.ci/docker/requirements.txt#L2
It seems we are doing it already? I vaguely remember we had this because they introduced some bug in newer versions, not sure if it's the one you are reporting. Checking torchtitan GH issue history may find it.

[francesco-bertolotti]

It seems we are doing it already?
Yep. I was taking as reference the one in pyproject.toml.

torchtitan/pyproject.toml

Line 17 in 51c197c

"datasets>=3.6.0",

[tianyu-l]

oh I see. Let's update the pyproject.toml then.