diff --git a/deepspeed/runtime/data_pipeline/data_sampling/data_sampler.py b/deepspeed/runtime/data_pipeline/data_sampling/data_sampler.py index 100bef3f7946..6fd8321a6d93 100644 --- a/deepspeed/runtime/data_pipeline/data_sampling/data_sampler.py +++ b/deepspeed/runtime/data_pipeline/data_sampling/data_sampler.py @@ -321,7 +321,7 @@ def state_dict(self): CURRICULUM_LEARNING_CURRENT_DIFFICULTIES: self.current_difficulties, CURRICULUM_LEARNING_DATA_CLUSTER_PATHS: self.data_cluster_paths, CURRICULUM_LEARNING_DATA_CLUSTER_CURRENT_POSITION: self.data_cluster_current_position, - CURRICULUM_LEARNING_NP_RNG_STATE: np.random.get_state() + CURRICULUM_LEARNING_NP_RNG_STATE: self.np_rng.bit_generator.state } def load_state_dict(self, state_dict): @@ -331,7 +331,13 @@ def load_state_dict(self, state_dict): self.current_difficulties = state_dict[CURRICULUM_LEARNING_CURRENT_DIFFICULTIES] self.data_cluster_paths = state_dict[CURRICULUM_LEARNING_DATA_CLUSTER_PATHS] self.data_cluster_current_position = state_dict[CURRICULUM_LEARNING_DATA_CLUSTER_CURRENT_POSITION] - np.random.set_state(state_dict[CURRICULUM_LEARNING_NP_RNG_STATE]) + np_rng_state = state_dict[CURRICULUM_LEARNING_NP_RNG_STATE] + if isinstance(np_rng_state, dict): + self.np_rng.bit_generator.state = np_rng_state + else: + # checkpoints written before this field held self.np_rng carry the tuple from + # np.random.get_state(), so restore that the way those checkpoints expect + np.random.set_state(np_rng_state) cluster_root_path = self.data_efficiency_config[DATA_SAMPLING][CURRICULUM_LEARNING][ CURRICULUM_LEARNING_CLUSTER_PATH] # Backward compatibility: previously data_cluster_paths were stored as diff --git a/tests/unit/runtime/test_data_efficiency.py b/tests/unit/runtime/test_data_efficiency.py index 50f75f2caddc..62aa8f07de2e 100644 --- a/tests/unit/runtime/test_data_efficiency.py +++ b/tests/unit/runtime/test_data_efficiency.py @@ -5,12 +5,16 @@ import torch import os +import numpy as np import deepspeed from deepspeed.accelerator import get_accelerator import pytest from unit.common import DistributedTest from unit.simple_model import Curriculum_SimpleModel, SimpleModel, random_dataloader, random_dataset +from deepspeed.runtime.data_pipeline.config import get_data_efficiency_config +from deepspeed.runtime.data_pipeline.constants import CURRICULUM_LEARNING_NP_RNG_STATE from deepspeed.runtime.data_pipeline.curriculum_scheduler import CurriculumScheduler +from deepspeed.runtime.data_pipeline.data_sampling.data_sampler import DeepSpeedDataSampler class MPU(): @@ -99,6 +103,74 @@ def test_curriculum_tops_out_at_the_configured_max(): assert scheduler.get_difficulty(200) == 1000 +def _curriculum_data_sampler(cluster_path): + param_dict = { + "data_efficiency": { + "enabled": True, + "seed": 1234, + "data_sampling": { + "enabled": True, + "curriculum_learning": { + "enabled": True, + "data_cluster_path": str(cluster_path), + "curriculum_metrics": { + "dummy": { + "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 + }, + } + }, + }, + }, + } + } + sampler = DeepSpeedDataSampler(data_efficiency_config=get_data_efficiency_config(param_dict), + one_epoch_total_samples=100, + micro_batch_size=8, + data_parallel_rank=0, + data_parallel_size=1, + data_parallel_group=None, + gradient_accumulation_steps=1, + global_rank=0) + # clusters of known size, so sample_from_clusters draws without any cluster file + sampler.data_clusters = [None] * 4 + sampler.data_cluster_sizes = [10, 20, 30, 40] + return sampler + + +def test_curriculum_sampler_checkpoints_its_own_rng(tmp_path): + 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 + + +def test_curriculum_sampler_resumes_its_rng_stream(tmp_path): + saved = _curriculum_data_sampler(tmp_path) + first_draws = [saved.sample_from_clusters().tolist() for _ in range(3)] + + 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)] + assert next_draws != first_draws + + +def test_curriculum_sampler_loads_legacy_rng_state(tmp_path): + sampler = _curriculum_data_sampler(tmp_path) + state_dict = sampler.state_dict() + state_dict[CURRICULUM_LEARNING_NP_RNG_STATE] = np.random.get_state() + sampler.load_state_dict(state_dict) + + @pytest.mark.parametrize('dtype', [torch.bfloat16, torch.float16]) class TestDataEfficiency(DistributedTest): world_size = 2