Checkpoint the curriculum data sampler's own RNG state - #8406
Conversation
DeepSpeedDataSampler draws from self.np_rng, but state_dict() saved np.random.get_state() and load_state_dict() restored it. That is the global legacy RandomState, so the sampler's generator was never in the checkpoint and a resumed run replayed its sampling stream from the seed. Save and restore self.np_rng.bit_generator.state instead. A state dict written by an older version carries the old tuple, which still goes through np.random.set_state. Signed-off-by: Ehsan Barkhordar <realbarkhordar@gmail.com>
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 8384ae8662
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
| resumed = _curriculum_data_sampler(tmp_path) | ||
| resumed.load_state_dict(saved.state_dict()) | ||
| next_draws = [resumed.sample_from_clusters().tolist() for _ in range(3)] | ||
| assert next_draws == [saved.sample_from_clusters().tolist() for _ in range(3)] |
There was a problem hiding this comment.
The resume test passes an in-memory dictionary directly between samplers, so it cannot catch failures in the actual DeepSpeedEngine.save_checkpoint()/load_checkpoint() path, serialization, dataloader restoration, or distributed execution. Because this change modifies checkpoint-resume behavior observable by a training loop, add an integration test that compares uninterrupted and checkpoint-resumed sampling on actual devices and report the hardware result.
AGENTS.md reference: AGENTS.md:L35-L36
Useful? React with 👍 / 👎.
| sampler = _curriculum_data_sampler(tmp_path) | ||
| for _ in range(3): | ||
| sampler.sample_from_clusters() | ||
| assert sampler.state_dict()[CURRICULUM_LEARNING_NP_RNG_STATE] == sampler.np_rng.bit_generator.state |
There was a problem hiding this comment.
This assertion pins the test to the sampler's private np_rng.bit_generator.state representation rather than the checkpoint-resume contract; an otherwise correct implementation that copies, normalizes, or encodes the generator state differently would fail it. The following test already checks continued sampling behavior, so remove this assertion or replace it with a round-trip assertion over observable samples.
AGENTS.md reference: AGENTS.md:L30-L32
Useful? React with 👍 / 👎.
|
On the RNG assertion: It also pins something the resume test does not, which is which generator gets checkpointed. Putting the sampler back on the process-global On the engine-level resume test: |
|
#8425 opened a couple of hours ago with the same source change as this one, legacy-tuple branch included, so what actually differs is the tests. This PR adds three to I am happy to close this in favour of #8425 if you would rather take that one, though the tests are worth carrying over either way, since nothing else in the suite catches this coming back. Updated 2026-09-08: there are now three PRs on #8405, since #8460 opened this morning. All three make the same one-file change, saving and restoring My view is that whichever one lands should carry tests, and past that I do not think it matters which. Happy to close this in favour of #8460 if you would rather take that one. |
1 participant
DeepSpeedDataSamplerdraws fromself.np_rng, the generator it seeds in__init__, butstate_dict()savednp.random.get_state()andload_state_dict()restored that. The global legacyRandomStatewas never the RNG the sampler used, so the generator stayed out of the checkpoint and a resumed run started its stream again from the seed: same cluster mix per step, same shuffles.Saves and restores
self.np_rng.bit_generator.stateinstead. A state dict written by an older version carries the old tuple, and that shape still goes throughnp.random.set_stateso old checkpoints keep loading the way they did.Added two tests for the resume path and one that loads an old-style state dict. The two behaviour ones fail on master.
Fixes #8405