[BUG] Curriculum data sampler doesn't restore its RNG on resume, so batches change after loading a checkpoint #8787

Describe the bug

If you train with curriculum learning and resume from a checkpoint, the data sampler doesn't pick up where it left off. The batches after the resume are different from the ones an uninterrupted run would have seen.

DeepSpeedDataSampler makes all of its random choices with its own generator:

self.np_rng = np.random.default_rng(self.data_efficiency_config[DATA_EFFICIENCY_SEED])

It uses it to shuffle a new cluster, to pick how many samples come from each cluster, to shuffle each global batch, and to reshuffle a cluster once it runs out.

But state_dict() saves numpy's global RNG, which the sampler never draws from:

CURRICULUM_LEARNING_NP_RNG_STATE: np.random.get_state()

and load_state_dict() puts that back into the global RNG:

np.random.set_state(state_dict[CURRICULUM_LEARNING_NP_RNG_STATE])

So after a resume, self.np_rng starts again from the seed, and the sampler repeats the random choices it made at the start of training instead of continuing them. Loading also rolls numpy's global RNG back to whatever it was at save time, which can surprise other code in the training script.

Nothing errors or warns. You only notice when a resumed run doesn't line up with one that never stopped.

To Reproduce

This runs on CPU in a single process. With one rank there's nothing to broadcast, so the two collectives are stubbed out.

import tempfile
import numpy as np
import deepspeed.comm as dist
from deepspeed.runtime.data_pipeline.data_sampling.data_sampler import DeepSpeedDataSampler

dist.barrier = lambda *a, **k: None
dist.broadcast = lambda *a, **k: None


class Group:

    def size(self):
        return 1


def sampler(path):
    config = {"seed": 1234, "data_sampling": {"num_epochs": 3, "curriculum_learning": {
        "enabled": True, "data_cluster_path": path, "curriculum_metrics": {"m": {
            "index_to_sample_path": "x", "index_to_metric_path": "x",
            "difficulty_type": "value", "clustering_type": "single_cluster",
            "min_difficulty": 2, "max_difficulty": 10, "schedule_type": "fixed_root",
            "schedule_config": {"total_curriculum_step": 8, "difficulty_step": 2, "root_degree": 1}}}}}}
    return DeepSpeedDataSampler(config, one_epoch_total_samples=40, micro_batch_size=4,
                                data_parallel_rank=0, data_parallel_size=1, data_parallel_group=Group(),
                                gradient_accumulation_steps=1, global_rank=0)


with tempfile.TemporaryDirectory() as d:
    it = iter(sampler(d))
    straight = [next(it) for _ in range(30)]

with tempfile.TemporaryDirectory() as d:
    s = sampler(d)
    it = iter(s)
    resumed = [next(it) for _ in range(15)]
    state = s.state_dict()
    s2 = sampler(d)
    s2.load_state_dict(state)
    it2 = iter(s2)
    resumed += [next(it2) for _ in range(15)]

print([i for i in range(30) if straight[i] != resumed[i]])

On master this prints [15, 16, 17, 19, 20, ..., 29], so 14 of the 15 batches after the resume differ. At first only the order inside each batch changes. With more than one data parallel rank, that already changes which rank trains on which sample. From batch 20 on, after the cluster gets reshuffled, the batches hold different samples altogether. With several difficulty clusters, the number of samples taken from each cluster drifts as well.

Expected behavior

A resumed run should get the same batches as a run that never stopped. That's what saving the sampler state is for.

A fix is to save self.np_rng.bit_generator.state and restore it in load_state_dict(). Older checkpoints only have the global state, so loading those can keep the current behavior and log a warning. I'll open a PR with this and a test.

ds_report output

torch version .................... 2.14.0
deepspeed info ................... 0.19.6, master at b7f51df

System info (please complete the following information):

• OS: macOS 26.6
• GPU count and types: none, this reproduces on CPU
• Python version: 3.11

Launcher context

No launcher. The snippet above runs with plain python.

Docker context

No.

Additional context

Only people who use curriculum learning in data_efficiency and resume from checkpoints are affected.