diff --git a/deepspeed/runtime/data_pipeline/curriculum_scheduler.py b/deepspeed/runtime/data_pipeline/curriculum_scheduler.py index 5f676ab7905a..24ff1f8d975a 100644 --- a/deepspeed/runtime/data_pipeline/curriculum_scheduler.py +++ b/deepspeed/runtime/data_pipeline/curriculum_scheduler.py @@ -17,6 +17,7 @@ def __init__(self, config): self.state['max_difficulty'] = config['max_difficulty'] self.state['current_difficulty'] = config['min_difficulty'] self.state['schedule_type'] = config['schedule_type'] + self.first_step = True if config['schedule_type'] == 'fixed_discrete': """ The schedule_config is a list of difficulty and a list of max @@ -36,11 +37,7 @@ def __init__(self, config): assert len(config['schedule_config']['difficulty']) > 0 assert len(config['schedule_config']['difficulty']) == len( config['schedule_config']['max_step']) + 1 - self.state['schedule'] = {} - for i in range(len(config['schedule_config']['max_step'])): - self.state['schedule'][config['schedule_config']['difficulty'][i]] = \ - [config['schedule_config']['max_step'][i], - config['schedule_config']['difficulty'][i+1]] + self.state['schedule'] = config['schedule_config'] elif config['schedule_type'] == 'fixed_root': """ The schedule_config includes: @@ -99,13 +96,15 @@ def get_state(self): def set_state(self, state): self.state = state - def __fixed_discrete_update_difficulty(self, global_steps): - s_state = self.state['schedule'][self.state['current_difficulty']] - if global_steps > s_state[0]: - self.state['current_difficulty'] = s_state[1] - return self.state['current_difficulty'] + def __fixed_discrete_get_difficulty(self, global_steps): + s_state = self.state['schedule'] + if global_steps > s_state['max_step'][-1]: + return s_state['difficulty'][-1] + for i in range(len(s_state['max_step'])): + if global_steps <= s_state['max_step'][i]: + return s_state['difficulty'][i] - def __fixed_root_update_difficulty(self, global_steps, root_degree=None): + def __fixed_root_get_difficulty(self, global_steps, root_degree=None): s_state = self.state['schedule'] if root_degree is None: root_degree = s_state['root_degree'] @@ -116,18 +115,20 @@ def __fixed_root_update_difficulty(self, global_steps, root_degree=None): (self.state['max_difficulty'] - self.state['min_difficulty']) + self.state['min_difficulty']) next_difficulty -= (next_difficulty % s_state['difficulty_step']) - self.state['current_difficulty'] = min(next_difficulty, - self.state['max_difficulty']) - return self.state['current_difficulty'] + next_difficulty = min(next_difficulty, self.state['max_difficulty']) + return next_difficulty - def update_difficulty(self, global_steps): - if self.state['current_difficulty'] >= self.state['max_difficulty']: - return self.state['current_difficulty'] + def get_difficulty(self, global_steps): if self.state['schedule_type'] == 'fixed_discrete': - return self.__fixed_discrete_update_difficulty(global_steps) + return self.__fixed_discrete_get_difficulty(global_steps) elif self.state['schedule_type'] == 'fixed_linear': - return self.__fixed_root_update_difficulty(global_steps, 1) + return self.__fixed_root_get_difficulty(global_steps, 1) elif self.state['schedule_type'] == 'fixed_root': - return self.__fixed_root_update_difficulty(global_steps) + return self.__fixed_root_get_difficulty(global_steps) else: raise RuntimeError('Unsupported curriculum schedule type') + + def update_difficulty(self, global_steps): + if self.state['current_difficulty'] < self.state['max_difficulty']: + self.state['current_difficulty'] = self.get_difficulty(global_steps) + return self.state['current_difficulty'] diff --git a/deepspeed/runtime/engine.py b/deepspeed/runtime/engine.py index 48e89457ec3f..dcb0efd6f56b 100755 --- a/deepspeed/runtime/engine.py +++ b/deepspeed/runtime/engine.py @@ -1293,13 +1293,16 @@ def forward(self, *inputs, **kwargs): if self.module.training and self.progressive_layer_drop: kwargs.update(self.progressive_layer_drop.get_state()) - if self.module.training and self.curriculum_enabled(): - self.curriculum_scheduler.update_difficulty(self.global_steps + 1) - if self.curriculum_params()["curriculum_type"] == "seqlen": - kwargs.update({ - "curriculum_seqlen": - self.curriculum_scheduler.get_current_difficulty() - }) + if self.__class__.__name__ != "PipelineEngine": + # TODO: The above if condition is a HACK since for PipelineEngine + # it's difficult to inject argument in forward pass. + if self.module.training and self.curriculum_enabled(): + self.curriculum_scheduler.update_difficulty(self.global_steps + 1) + if self.curriculum_params()["curriculum_type"] == "seqlen": + kwargs.update({ + "curriculum_seqlen": + self.curriculum_scheduler.get_current_difficulty() + }) if self.zero_optimization_partition_weights(): # Enable automated discovery of external parameters by indicating that diff --git a/deepspeed/runtime/pipe/engine.py b/deepspeed/runtime/pipe/engine.py index 5e9297916b86..0905222c281c 100644 --- a/deepspeed/runtime/pipe/engine.py +++ b/deepspeed/runtime/pipe/engine.py @@ -265,6 +265,16 @@ def _reserve_pipe_buffers(self, num_buffers): self.pipe_buffers[key].extend([None] * num_added) self.num_pipe_buffers = num_buffers + def reset_activation_shape(self): + """Reset the buffers when the shape of activation and gradient change. + For example, for curriculum learning that changes the seqlen of each + sample, we need to call this whenever the seqlen is going to change. + """ + self.first_output_send = True + self.pipe_recv_buf = None + self.grad_layer = None + self.meta_buffer = None + def train_batch(self, data_iter=None): """Progress the pipeline to train the next batch of data. The engine will ingest ``self.train_batch_size()`` total samples collectively across all workers. @@ -293,6 +303,17 @@ def train_batch(self, data_iter=None): raise RuntimeError( f'train_batch() requires gradients enabled. Use eval_batch() instead.') + # Curriculum learning could change activation shape + if self.curriculum_enabled(): + new_difficulty = self.curriculum_scheduler.update_difficulty( \ + self.global_steps + 1) + if self.global_steps == 0 or self.curriculum_scheduler.first_step: + self.reset_activation_shape() + self.curriculum_scheduler.first_step = False + elif new_difficulty != self.curriculum_scheduler.get_difficulty( \ + self.global_steps): + self.reset_activation_shape() + if data_iter: self.set_dataiterator(data_iter)