Describe the bug
DeepSpeedDataSampler draws from its own generator, self.np_rng = np.random.default_rng(seed), but state_dict() stores np.random.get_state() and load_state_dict() restores it. That is the global legacy RandomState, not the generator the sampler uses, so the saved value is the same whether the sampler has drawn nothing or a thousand batches.
The effect shows up on resume. __init__ seeds self.np_rng fresh, load_state_dict never touches it, and the sampler goes back to the top of its own stream: same cluster mix per step, same shuffles out of get_new_cluster and reshuffle_clusters. Restoring the global state also moves whatever else in the process is drawing from np.random.
To Reproduce
from deepspeed.runtime.data_pipeline.config import get_data_efficiency_config
from deepspeed.runtime.data_pipeline.data_sampling.data_sampler import DeepSpeedDataSampler
metric = {"index_to_sample_path": "dummy", "index_to_metric_path": "dummy", "difficulty_type": "value",
"clustering_type": "single_cluster", "min_difficulty": 8, "max_difficulty": 80,
"schedule_type": "fixed_linear",
"schedule_config": {"total_curriculum_step": 100, "difficulty_step": 8}}
config = get_data_efficiency_config({"data_efficiency": {"enabled": True, "seed": 1234, "data_sampling": {
"enabled": True, "curriculum_learning": {"enabled": True, "data_cluster_path": "/tmp/clusters",
"curriculum_metrics": {"dummy": metric}}}}})
def sampler():
s = DeepSpeedDataSampler(config, 100, 8, 0, 1, None, 1, global_rank=0)
s.data_clusters, s.data_cluster_sizes = [None] * 4, [10, 20, 30, 40]
return s
saved = sampler()
before = saved.state_dict()["np_rng_state"]
first = [saved.sample_from_clusters().tolist() for _ in range(3)]
print("saved rng state changed by 3 draws:", repr(saved.state_dict()["np_rng_state"]) != repr(before))
resumed = sampler()
resumed.load_state_dict(saved.state_dict())
print("resumed draw 1:", resumed.sample_from_clusters().tolist(), " original draw 1:", first[0])On master at 493dafa, python 3.12.14, torch 2.14.0+cpu:
saved rng state changed by 3 draws: False
resumed draw 1: [0, 3, 3, 2] original draw 1: [0, 3, 3, 2]
Expected behavior
A resumed run continues the sampling stream instead of replaying it, and loading a sampler state leaves the global numpy RNG alone.
Additional context
Both lines go back to #2585, which added the library. I found this reading the code, not from a training run. PR on the way.
Describe the bug
DeepSpeedDataSamplerdraws from its own generator,self.np_rng = np.random.default_rng(seed), butstate_dict()storesnp.random.get_state()andload_state_dict()restores it. That is the global legacyRandomState, not the generator the sampler uses, so the saved value is the same whether the sampler has drawn nothing or a thousand batches.The effect shows up on resume.
__init__seedsself.np_rngfresh,load_state_dictnever touches it, and the sampler goes back to the top of its own stream: same cluster mix per step, same shuffles out ofget_new_clusterandreshuffle_clusters. Restoring the global state also moves whatever else in the process is drawing fromnp.random.To Reproduce
On master at
493dafa, python 3.12.14, torch 2.14.0+cpu:Expected behavior
A resumed run continues the sampling stream instead of replaying it, and loading a sampler state leaves the global numpy RNG alone.
Additional context
Both lines go back to #2585, which added the library. I found this reading the code, not from a training run. PR on the way.