bf16+pipeline parallelism

.gitignore

@@ -4,6 +4,7 @@
 *.swp
 *.log
 deepspeed/git_version_info_installed.py
+__pycache__
 
 # Build + installation data
 build/

deepspeed/checkpoint/constants.py

@@ -11,9 +11,11 @@
 
 BASE_OPTIMIZER_STATE = 'base_optimizer_state'
 SINGLE_PARTITION_OF_FP32_GROUPS = "single_partition_of_fp32_groups"
+GROUPS_PADDING = 'groups_padding'
 
 PARTITION_COUNT = 'partition_count'
 ZERO_STAGE = 'zero_stage'
+CLIP_GRAD = 'clip_grad'
 
 #########################################
 # Module checkpoint keys

deepspeed/runtime/config.py

@@ -901,7 +901,6 @@ def _initialize_params(self, param_dict):
         self.fp16_enabled = get_fp16_enabled(param_dict)
         self.bfloat16_enabled = get_bfloat16_enabled(param_dict)
         assert not (self.fp16_enabled and self.bfloat16_enabled), 'bfloat16 and fp16 modes cannot be simultaneously enabled'
-        assert not (self.bfloat16_enabled and (self.zero_optimization_stage not in {1, 2, 3})), f'bfloat16 mode is only enabled for Zero 1,2,3 currently. got {self.zero_optimization_stage}'
         self.fp16_master_weights_and_gradients = get_fp16_master_weights_and_grads_enabled(
             param_dict)
         self.amp_enabled = get_amp_enabled(param_dict)

deepspeed/runtime/constants.py

@@ -446,3 +446,8 @@ class ValidationMode:
 '''
 DATALOADER_DROP_LAST = "dataloader_drop_last"
 DATALOADER_DROP_LAST_DEFAULT = False
+
+#########################################
+# PIPELINE PARALLELISM
+#########################################
+PIPE_REPLICATED = 'ds_pipe_replicated'

deepspeed/runtime/fp16/fused_optimizer.py

@@ -11,7 +11,7 @@
 from deepspeed.runtime.utils import get_global_norm, get_grad_norm, CheckOverflow, get_weight_norm
 from deepspeed.runtime.fp16.loss_scaler import INITIAL_LOSS_SCALE, SCALE_WINDOW, MIN_LOSS_SCALE
 from deepspeed.utils import groups, logger, log_dist
-from deepspeed.checkpoint.constants import OPTIMIZER_STATE_DICT
+from deepspeed.checkpoint.constants import OPTIMIZER_STATE_DICT, CLIP_GRAD
 import torch.distributed as dist
 
 
@@ -32,11 +32,13 @@ def __init__(self,
                  mpu=None,
                  clip_grad=0.0,
                  fused_adam_legacy=False,
+                 has_moe_layers=False,
                  timers=None):
 
         self.fused_adam_legacy = fused_adam_legacy
         self.timers = timers
         self.deepspeed = deepspeed
+        self.has_moe_layers = has_moe_layers
         self.using_pipeline = self.deepspeed.pipeline_parallelism
         if not torch.cuda.is_available:
             raise SystemError("Cannot use fp16 without CUDA.")
@@ -93,6 +95,7 @@ def __init__(self,
 
         self.clip_grad = clip_grad
         self.norm_type = 2
+        self.step_count = 0
 
         TORCH_MAJOR = int(torch.__version__.split('.')[0])
         TORCH_MINOR = int(torch.__version__.split('.')[1])
@@ -167,11 +170,15 @@ def step_fused_adam(self, closure=None):
                                                        self.cur_scale))
             return self.overflow
 
-        self._global_grad_norm = get_global_norm(norm_list=norm_groups)
+        scaled_grad_norm = get_global_norm(norm_list=norm_groups)
 
         combined_scale = self.unscale_and_clip_grads(grads_groups_flat,
-                                                     self._global_grad_norm,
+                                                     scaled_grad_norm,
                                                      apply_scale=False)
+
+        # Stash unscaled gradient norm
+        self._global_grad_norm = scaled_global_grad_norm / self.cur_scale
+
         # norm is in fact norm*cur_scale
         self.optimizer.step(grads=[[g] for g in grads_groups_flat],
                             output_params=[[p] for p in self.fp16_groups_flat],
@@ -258,26 +265,19 @@ def step(self, closure=None):
         self.start_timers([COMPUTE_NORM])
 
         all_groups_norm = get_grad_norm(self.fp32_groups_flat, mpu=self.mpu)
-        #all_groups_norm_old = all_groups_norm
-        # Need to allreduce (avg) the norms across different ranks because moe params will not be synced during allreduce
-        if self.using_pipeline:
-            pg = self.deepspeed.mpu.get_data_parallel_group()
-        else:
-            pg = groups._get_data_parallel_group()
-        scaled_norm = all_groups_norm * 1.0 / float(dist.get_world_size(group=pg))
-        scaled_norm_tensor = torch.tensor(scaled_norm,
-                                          device=self.fp32_groups_flat[i].device,
-                                          dtype=torch.float)
-        dist.all_reduce(scaled_norm_tensor, group=pg)
-        all_groups_norm = scaled_norm_tensor.item()
-        #print(f"old = {all_groups_norm_old} and new = {all_groups_norm} at rank: {torch.distributed.get_rank()}")
 
         self.stop_timers([COMPUTE_NORM])
 
-        self._global_grad_norm = get_global_norm(norm_list=[all_groups_norm])
+        if self.has_moe_layers:
+            all_groups_norm = self._get_norm_with_moe_layers(all_groups_norm)
+
+        scaled_global_grad_norm = get_global_norm(norm_list=[all_groups_norm])
+
+        # Stash unscaled gradient norm
+        self._global_grad_norm = scaled_global_grad_norm / self.cur_scale
 
         self.start_timers([UNSCALE_AND_CLIP])
-        self.unscale_and_clip_grads(grads_groups_flat, self._global_grad_norm)
+        self.unscale_and_clip_grads(grads_groups_flat, scaled_global_grad_norm)
         self.stop_timers([UNSCALE_AND_CLIP])
 
         self.start_timers([BASIC_STEP])
@@ -300,8 +300,26 @@ def step(self, closure=None):
 
         self.log_timers(STEP_TIMERS)
 
+        self.step_count += 1
+
         return self.overflow
 
+    def _get_norm_with_moe_layers(self, all_groups_norm):
+        #all_groups_norm_old = all_groups_norm
+        # Need to allreduce (avg) the norms across different ranks because moe params will not be synced during allreduce
+        if self.using_pipeline:
+            pg = self.deepspeed.mpu.get_data_parallel_group()
+        else:
+            pg = groups._get_data_parallel_group()
+        scaled_norm = all_groups_norm * 1.0 / float(dist.get_world_size(group=pg))
+        scaled_norm_tensor = torch.tensor(scaled_norm,
+                                          device=self.fp32_groups_flat[0].device,
+                                          dtype=torch.float)
+        dist.all_reduce(scaled_norm_tensor, group=pg)
+        all_groups_norm = scaled_norm_tensor.item()
+        #print(f"old = {all_groups_norm_old} and new = {all_groups_norm} at rank: {torch.distributed.get_rank()}")
+        return all_groups_norm
+
     def unscale_and_clip_grads(self, grad_groups_flat, total_norm, apply_scale=True):
         # compute combined scale factor for this group
         combined_scale = self.cur_scale
@@ -325,8 +343,8 @@ def backward(self, loss, create_graph=False, retain_graph=False):
         2. scaled_loss = fp32_loss*loss_scale
         3. scaled_loss.backward(), which accumulates scaled gradients into the ``.grad`` attributes of the model's fp16 leaves
         """
-        scaled_loss = (loss.float()) * self.cur_scale
 
+        scaled_loss = (loss.float()) * self.cur_scale
         scaled_loss.backward(create_graph=create_graph, retain_graph=retain_graph)
 
     def _update_scale(self, skip):
@@ -399,7 +417,7 @@ def state_dict(self):
             state_dict['scale_window'] = self.scale_window
         state_dict[OPTIMIZER_STATE_DICT] = self.optimizer.state_dict()
         state_dict['fp32_groups_flat'] = self.fp32_groups_flat
-        state_dict['clip_grad'] = self.clip_grad
+        state_dict[CLIP_GRAD] = self.clip_grad
         return state_dict
 
     # Refresh fp32 master params from fp16 copies
@@ -433,7 +451,7 @@ def load_state_dict(self, state_dict, load_optimizer_states=True):
             self.scale_window = state_dict['scale_window']
         if load_optimizer_states:
             self.optimizer.load_state_dict(state_dict[OPTIMIZER_STATE_DICT])
-        self.clip_grad = state_dict['clip_grad']
+        self.clip_grad = state_dict[CLIP_GRAD]
         # At this point, the optimizer's references to the model's fp32 parameters are up to date.
         # The optimizer's hyperparameters and internal buffers are also up to date.
         # However, the fp32 master copies of the model's fp16 params stored by the optimizer are still

deepspeed/runtime/pipe/engine.py

@@ -241,14 +241,33 @@ def _exec_reduce_tied_grads(self):
         # (see https://github.com/EleutherAI/gpt-neox/issues/62#issuecomment-761471944)
         if self.zero_optimization_partition_gradients():
             self.optimizer.overlapping_partition_gradients_reduce_epilogue()
-        self.module.allreduce_tied_weight_gradients()
+
+        if self.bfloat16_enabled():
+            weight_group_list = self.module.get_tied_weights_and_groups()
+            for weight, group in weight_group_list:
+                dist.all_reduce(weight._hp_grad, group=group)
+        else:
+            self.module.allreduce_tied_weight_gradients()
 
     def _exec_reduce_grads(self):
         self._force_grad_boundary = True
         if self.pipeline_enable_backward_allreduce:
-            self.allreduce_gradients(bucket_size=MEMORY_OPT_ALLREDUCE_SIZE)
+            if self.bfloat16_enabled():
+                if self.zero_optimization_stage() == 0:
+                    self._bf16_reduce_grads()
+                else:
+                    assert self.zero_optimization_stage() == 1, "only bf16 + z1 are supported"
+                    raise NotImplementedError()
+            else:
+                self.allreduce_gradients(bucket_size=MEMORY_OPT_ALLREDUCE_SIZE)
         self._force_grad_boundary = False
 
+    def _bf16_reduce_grads(self):
+        # Make our own list of gradients from the optimizer's FP32 grads
+        grads = []
+        self.buffered_allreduce_fallback(grads=self.optimizer.get_grads_for_reduction(),
+                                         elements_per_buffer=MEMORY_OPT_ALLREDUCE_SIZE)
+
     def _reserve_pipe_buffers(self, num_buffers):
         """Ensure that each pipeline buffer has at least ``num_buffers`` slots.
 
@@ -726,6 +745,10 @@ def _exec_backward_pass(self, buffer_id):
             part_grad = None
             #print(f'RANK={self.global_rank} BEFORE-BWD restored grad={self.grad_layer[0].size()} {self.grad_layer[1].size()}')
 
+        if self.bfloat16_enabled() and not self.is_last_stage():
+            # manually call because we don't call optimizer.backward()
+            self.optimizer.clear_lp_grads()
+
         # This handles either a single tensor or tuple of tensors.
         if isinstance(outputs, tuple):
             out_tensors = [t for t in outputs if t.is_floating_point()]
@@ -734,6 +757,10 @@ def _exec_backward_pass(self, buffer_id):
         else:
             torch.autograd.backward(tensors=(outputs, ), grad_tensors=(grad_tensors, ))
 
+        if self.bfloat16_enabled() and not self.is_last_stage():
+            # manually call because we don't call optimizer.backward()
+            self.optimizer.update_hp_grads(clear_lp_grads=False)
+
         # Free up the memory from the output of forward()
         self.pipe_buffers['output_tensors'][buffer_id] = None
         self.pipe_buffers['outputs'][buffer_id] = None

deepspeed/runtime/pipe/module.py

@@ -419,6 +419,13 @@ def allreduce_tied_weight_gradients(self):
             weight = getattr(self.tied_modules[key], comm['weight_attr'])
             dist.all_reduce(weight.grad, group=comm['group'])
 
+    def get_tied_weights_and_groups(self):
+        weight_group_list = []
+        for key, comm in self.tied_comms.items():
+            weight = getattr(self.tied_modules[key], comm['weight_attr'])
+            weight_group_list.append((weight, comm['group']))
+        return weight_group_list
+
     def _synchronize_tied_weights(self):
         for key, comm in self.tied_comms.items():
             dist.broadcast(