Skip to content

fix(train): reconstruct MixedDataLoader RNG stream on resume#14

Merged
AuraIis merged 3 commits into
mainfrom
fix/resume-dataloader-state
Jul 10, 2026
Merged

fix(train): reconstruct MixedDataLoader RNG stream on resume#14
AuraIis merged 3 commits into
mainfrom
fix/resume-dataloader-state

Conversation

@AuraIis

@AuraIis AuraIis commented Jul 10, 2026

Copy link
Copy Markdown
Owner

Root cause

MixedDataLoader owns independent NumPy generators for each language plus a shuffle generator and fractional row-credit state. Checkpoints restore global RNGs but not these loader-owned streams, so resumed training restarted sampling at batch zero.

Measured on the pinned clean checkout: LR and eval loss restored, while train loss diverged immediately; an isolated token-stream comparison confirmed different post-resume batches.

Fix

  • count mixed batches consumed by the loader
  • add a data-free deterministic fast-forward that reconstructs:
    • each per-language RNG
    • the batch-row shuffle RNG
    • fractional mix row credit
  • after loading a checkpoint, advance by
    trainer_state.step * gradient_accumulation
  • reconstruction uses each rank's own seed/shard, so DDP ranks recover their own streams
  • no token data is read during fast-forward

This works with the existing step_60000 checkpoint format; no migration or checkpoint rewrite is required.

Regression coverage

  • exact next-batch equality after 160 consumed batches (step 20 × grad accumulation 8)
  • ten subsequent batches remain byte-identical
  • rank-specific DDP-shard reconstruction
  • invalid negative fast-forward rejected

The existing code/config/tokenizer provenance checks remain in force; changing loader configuration across a resume is still a provenance mismatch.

@AuraIis
AuraIis merged commit c5a7edd into main Jul 10, 2026
1 check passed
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