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

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>

@SunMarc SunMarc left a comment

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks, just a nit

Comment thread tests/test_data_loader.py


class StatefulDataLoaderTester(AccelerateTestCase):
@parameterized.expand(

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

simplify the test pleas , we don't need that much parametrization.

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.

2 participants