CL for big science

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']

deepspeed/runtime/engine.py

@@ -1293,13 +1293,14 @@ 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":
+            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

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)