Checkpoint the RNG the curriculum sampler actually draws from

[alanhuangyoo]

Fixes #8405, reported by @ebarkhordar.

Problem

Every draw in DeepSpeedDataSampler goes through self.np_rng, a np.random.default_rng seeded in __init__:

self.np_rng = np.random.default_rng(self.data_efficiency_config[DATA_EFFICIENCY_SEED])   # :64
...
self.np_rng.shuffle(new_cluster)                                                          # :217  get_new_cluster
self.np_rng.choice(num_clusters, self.global_batch_size, replace=True, p=weights)         # :236  sample_from_clusters
self.np_rng.shuffle(cluster)                                                              # :246  reshuffle_clusters
self.np_rng.shuffle(batch)                                                                # :289

but the checkpoint stored and restored something else entirely:

CURRICULUM_LEARNING_NP_RNG_STATE: np.random.get_state()      # :324  state_dict
np.random.set_state(state_dict[CURRICULUM_LEARNING_NP_RNG_STATE])   # :334  load_state_dict

That is the process-wide legacy RandomState, which this sampler never touches. So the saved value reads the same whether the sampler has drawn nothing or a thousand batches, and on resume __init__ re-seeds self.np_rng while load_state_dict leaves it alone — the sampler restarts its own stream, reproducing the same cluster mix per step and the same shuffles it produced the first time. Restoring the global state also moved whatever else in the process draws from np.random.

Solution

Save and restore self.np_rng.bit_generator.state.

A checkpoint written before this holds the legacy tuple, and there is no sampler stream recorded in it to restore. _load_np_rng_state recognises the shape, warns once on rank 0, and leaves the sampler on its configured seed — which is where those checkpoints already resumed — rather than refusing to load.

Verification

tests/unit/runtime/data_pipeline/test_curriculum_sampler_rng.py, four cases:

test	on master	here
the saved state moves as the sampler draws	❌ identical before and after 3 draws	✅
resume continues the stream rather than replaying it	❌ [[0, 3, 3, 2], ...] != [[2, 1, 3, 2], ...]	✅
loading leaves the global numpy RNG alone	✅	✅
a legacy checkpoint still loads	✅	✅

The resume test goes through load_state_dict rather than poking the generator, so it fails on master for the behaviour rather than on a type error, and it asserts that a sampler on a fresh seed does not already agree — otherwise it would prove nothing.

1×H20:  4 passed   (this branch)
        2 failed, 2 passed   (master, same test file)

pre-commit run passes on both files.