Skip to content

Preserve pretraining epoch across resume - #2313

Open
Alphaxiaoteng wants to merge 1 commit into
Lightning-AI:mainfrom
Alphaxiaoteng:fix-1712-preserve-resume-epoch
Open

Preserve pretraining epoch across resume#2313
Alphaxiaoteng wants to merge 1 commit into
Lightning-AI:mainfrom
Alphaxiaoteng:fix-1712-preserve-resume-epoch

Conversation

@Alphaxiaoteng

Copy link
Copy Markdown

What changed

  • make CycleIterator save and restore its epoch through Fabric checkpoints
  • create the pretraining iterator before checkpoint loading and reuse it in fit
  • keep old checkpoints compatible by loading the existing state first, then restoring the optional iterator state from Fabric's remainder
  • cover mid-epoch and exact epoch-boundary resume with a real LitData StreamingDataLoader

Why

The current litdata==0.2.59 pin already restores the dataloader cursor correctly: after a mid-epoch checkpoint it yields the remaining tail once and then starts a complete new epoch. However, pretrain still constructs CycleIterator after loading and resets its epoch counter to zero.

This addresses the remaining epoch-tracking problem described in #1712 without changing dataloader cursor semantics. The boundary handling is needed because a restored StreamingDataLoader can begin the next epoch itself before CycleIterator observes StopIteration.

Validation

  • unmodified baseline: the real Fabric checkpoint regression fails because CycleIterator contains an unpicklable generator, and the pretraining resume path has no iterator state
  • pytest tests/test_pretrain.py tests/test_utils.py -q: 59 passed, 9 skipped
  • all applicable changed-file pre-commit hooks passed, including Ruff check and format
  • git diff --check passed

@Alphaxiaoteng

Copy link
Copy Markdown
Author

CI note: the tests.yaml / Lit Job (n/a) check failed, but GitHub only exposes a generic “Job Failed” result with no annotations, and the Lightning job log requires separate authentication. The new regression tests in this PR are CPU tests and are not collected by this lane because it runs with RUN_ONLY_CUDA_TESTS=1. The same lane has also failed at a similar runtime on unrelated recent diffs (for example #2309 and #2312), so I do not want to change unrelated code without a traceback.

Could a maintainer please rerun the n/a lane once and approve the fork CPU workflow so the new regression tests run in the standard CI? Locally, pytest tests/test_pretrain.py tests/test_utils.py -q passed (59 passed, 9 skipped), and all applicable pre-commit hooks passed. I will inspect any concrete failing node/traceback before making further code changes.

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.

1 participant