Skip to content

Fix stateful dataloader checkpoint cache after loading - #4244

Open
MrCapricornLiu wants to merge 1 commit into
huggingface:mainfrom
MrCapricornLiu:MrCapricornLiu/fix/stateful-loader-recheckpoint
Open

Fix stateful dataloader checkpoint cache after loading#4244
MrCapricornLiu wants to merge 1 commit into
huggingface:mainfrom
MrCapricornLiu:MrCapricornLiu/fix/stateful-loader-recheckpoint

Conversation

@MrCapricornLiu

Copy link
Copy Markdown

Loading a stateful dataloader checkpoint updates the underlying loader but leaves the adapter's cached state unchanged. Saving again before reading another batch therefore writes the previous cursor. A fresh adapter can replay data; an adapter that had advanced further can skip data after restoration.

Refresh the cache after load_state_dict succeeds. The regression covers shard, dispatcher and skip loaders with zero or two workers, both fresh and reused adapters, and checks the actual remaining samples after another save/load.

Validation: the 12 new cases fail before the fix, and all 49 data-loader tests pass afterward. An Accelerator.save_state/load_state training check passes on CPU and eight H800 GPUs, matching data order and SGD momentum updates against uninterrupted training. Repository Ruff lint/format and applicable pre-commit checks pass.

Signed-off-by: Chenghao Liu <chliu@stu.pku.edu.cn>
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