bf16 gradient clipping fix

bf16 checkpoint save/load

Author: tjruwase

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_gradient'
 
 #########################################
 # Module checkpoint keys

deepspeed/runtime/bf16_optimizer.py

@@ -2,7 +2,9 @@
 import torch.distributed as dist
 from deepspeed.runtime.constants import PIPE_REPLICATED
 from deepspeed.ops.op_builder import UtilsBuilder
+from packaging import version as pkg_version
 
+from deepspeed.git_version_info import version
 from deepspeed.runtime.utils import (get_global_norm_of_tensors,
                                      clip_tensors_by_global_norm,
                                      get_grad_norm,
@@ -13,6 +15,13 @@
                                      is_model_parallel_parameter,
                                      see_memory_usage)
 
+from deepspeed.checkpoint.constants import (DS_VERSION,
+                                            PARTITION_COUNT,
+                                            BASE_OPTIMIZER_STATE,
+                                            SINGLE_PARTITION_OF_FP32_GROUPS,
+                                            CLIP_GRAD,
+                                            GROUPS_PADDING)
+
 
 class BF16_Optimizer:
     def __init__(self,
@@ -36,6 +45,10 @@ def __init__(self,
         self.real_dp_process_group = [
             dp_process_group for i in range(len(self.optimizer.param_groups))
         ]
+        dp_world_size = dist.get_world_size(group=self.dp_process_group)
+        self.partition_count = [
+            dp_world_size for i in range(len(self.optimizer.param_groups))
+        ]
 
         # Load pre-built or JIT compile (un)flatten ops
         util_ops = UtilsBuilder().load()
@@ -58,9 +71,9 @@ def __init__(self,
         self.fp32_groups_actual_gradients_flat = []
         self.fp32_groups_gradient_flat_partition = []
         self.fp32_groups_has_gradients = []
-        self.step_count = 0
 
-        dp_world_size = dist.get_world_size(group=self.dp_process_group)
+        self.step_count = 0
+        self.groups_padding = []
 
         for i, param_group in enumerate(self.optimizer.param_groups):
             see_memory_usage(f'before initializing group {i}', force=True)
@@ -127,6 +140,15 @@ def __init__(self,
             # track fp32 gradient updates
             self.fp32_groups_has_gradients.append([False] * len(self.bf16_groups[i]))
 
+            # Record padding required for alignment
+            if partition_id == dist.get_world_size(
+                    group=self.real_dp_process_group[i]) - 1:
+                padding = self.bf16_groups_flat[i].numel() - length_without_padding
+            else:
+                padding = 0
+
+            self.groups_padding.append(padding)
+
             # update optimizer param groups to reference fp32 params partition
             param_group['params'] = [self.fp32_groups_flat_partition[i]]
 
@@ -186,8 +208,8 @@ def step(self, closure=None):
         if self.clip_grad > 0.:
             clip_tensors_by_global_norm(input_tensors=self.get_grads_for_norm(),
                                         max_norm=self.clip_grad,
-                                        mpu=self.mpu,
-                                        global_grad_norm=all_groups_norm)
+                                        global_norm=all_groups_norm,
+                                        mpu=self.mpu)
 
         self.optimizer.step()
 
@@ -278,18 +300,47 @@ def clear_lp_grads(self):
                 param.grad = None
 
     def state_dict(self):
-        # TODO capture all training state for checkpointing
         state_dict = {}
-        state_dict['optimizer_state_dict'] = self.optimizer.state_dict()
-        state_dict['clip_grad'] = self.clip_grad
+        state_dict[CLIP_GRAD] = self.clip_grad
+        state_dict[BASE_OPTIMIZER_STATE] = self.optimizer.state_dict()
+        state_dict[SINGLE_PARTITION_OF_FP32_GROUPS] = self.fp32_groups_flat_partition
+        state_dict[GROUPS_PADDING] = self.groups_padding
+        state_dict[PARTITION_COUNT] = self.partition_count
+        state_dict[DS_VERSION] = version
+
         return state_dict
 
-    def load_state_dict(self, state_dict, load_optimizer_states=True):
+    def load_state_dict(self,
+                        state_dict_list,
+                        load_optimizer_states=True,
+                        load_from_fp32_weights=False):
+        dp_rank = dist.get_rank(group=self.dp_process_group)
+        current_rank_sd = state_dict_list[dp_rank]
+
+        ckpt_version = current_rank_sd.get(DS_VERSION, False)
+        assert ckpt_version, f"Empty ds_version in checkpoint, not clear how to proceed"
+        ckpt_version = pkg_version.parse(ckpt_version)
+
+        self.clip_grad = current_rank_sd[CLIP_GRAD]
+
         if load_optimizer_states:
-            self.optimizer.load_state_dict(state_dict['optimizer_state_dict'])
-        self.clip_grad = state_dict['clip_grad']
+            self.optimizer.load_state_dict(current_rank_sd[BASE_OPTIMIZER_STATE])
+
+        if load_from_fp32_weights:
+            for current, saved in zip(self.fp32_groups_flat_partition, current_rank_sd[SINGLE_PARTITION_OF_FP32_GROUPS]):
+                src_tensor = _get_padded_tensor(saved, current.numel())
+                current.data.copy_(src_tensor.data)
 
     @property
     def param_groups(self):
         """Forward the wrapped optimizer's parameters."""
         return self.optimizer.param_groups
+
+
+def _get_padded_tensor(src_tensor, size):
+    if src_tensor.numel() >= size:
+        return src_tensor
+    padded_tensor = torch.zeros(size, dtype=src_tensor.dtype, device=src_tensor.device)
+    slice_tensor = torch.narrow(padded_tensor, 0, 0, src_tensor.numel())
+    slice_tensor.data.copy_(src_tensor.data)
+    return padded_tensor

deepspeed/runtime/engine.py

@@ -807,7 +807,7 @@ def _configure_checkpointing(self, dist_init_required):
         self.save_non_zero_checkpoint = (
             dp_rank == 0) or self.zero_optimization_partition_weights()
 
-        if self.zero_optimization():
+        if self.zero_optimization() or self.bfloat16_enabled():
             param_rank = torch.distributed.get_rank(
                 group=self.optimizer.dp_process_group)
 
@@ -2370,7 +2370,8 @@ def load_module_state_dict(self, state_dict, strict=True):
         self.module.load_state_dict(state_dict, strict=strict)
 
     def _get_rank_zero_ckpt_name(self, checkpoints_path, tag, mp_rank, dp_rank):
-        filename = "zero_pp_rank_{}".format(dp_rank)
+        filename = "bf16_zero_pp_rank_{}".format(
+            dp_rank) if self.bfloat16_enabled() else "zero_pp_rank_{}".format(dp_rank)
         zero_ckpt_name = os.path.join(
             checkpoints_path,
             str(tag),
@@ -2495,7 +2496,8 @@ def load_checkpoint(self,
                                                          load_lr_scheduler_states=load_lr_scheduler_states,
                                                          load_module_only=load_module_only)
 
-        if self.zero_optimization() and load_path is not None:
+        load_zero_checkpoint = self.zero_optimization() or self.bfloat16_enabled()
+        if load_zero_checkpoint and load_path is not None:
             success = self._load_zero_checkpoint(
                 load_dir,
                 tag,
@@ -2567,8 +2569,9 @@ def _load_checkpoint(self,
             else:
                 optim_checkpoint = checkpoint
 
-            if load_optimizer_states and self.optimizer is not None and not self.zero_optimization(
-            ):
+            has_zero_optimizer_state = self.zero_optimization() or self.bfloat16_enabled(
+            )
+            if load_optimizer_states and self.optimizer is not None and not has_zero_optimizer_state:
                 if self.fp16_enabled():
                     self.optimizer.load_state_dict(
                         optim_checkpoint['optimizer'],
@@ -2964,13 +2967,13 @@ def _save_checkpoint(self, save_dir, tag, client_state={}):
         # module_state_dict() and uses this path to save the model. module_state_dict()
         # then instead just returns None.
         self._curr_ckpt_path = os.path.join(save_dir, tag)
-
+        zero_optimizer_state = self.zero_optimization() or self.bfloat16_enabled()
         state = dict(module=self.module_state_dict(),
                      buffer_names=self._get_buffer_names(),
                      optimizer=self.optimizer.state_dict()
-                     if self.optimizer and not self.zero_optimization() else None,
+                     if self.optimizer and not zero_optimizer_state else None,
                      param_shapes=self._get_zero_param_shapes()
-                     if self.optimizer and self.zero_optimization() else None,
+                     if self.optimizer and zero_optimizer_state else None,
                      lr_scheduler=self.lr_scheduler.state_dict()
                      if self.lr_scheduler is not None else None,
                      sparse_tensor_module_names=self.sparse_tensor_module_names,
@@ -3028,6 +3031,8 @@ def _get_zero_param_shapes(self):
         # if we don't use it, we get parameters ordered incorrectly
         if hasattr(self.optimizer, "round_robin_bit16_groups"):
             bit16_groups = self.optimizer.round_robin_bit16_groups
+        elif self.bfloat16_enabled() and not self.zero_optimization():
+            bit16_groups = self.optimizer.bf16_groups
         else:
             bit16_groups = self.optimizer.bit16_groups if self.zero_optimization_stage(
             ) == 2 else self.optimizer.fp16_groups
@@ -3068,7 +3073,8 @@ def _save_zero_checkpoint(self, save_path, tag):
         torch.save(zero_sd, zero_checkpoint_name)
         if self.global_rank == 0:
             self._copy_recovery_script(save_path)
-        logger.info('zero checkpoint saved {}'.format(zero_checkpoint_name))
+        ckpt_type = 'zero' if self.zero_optimization() else 'bfl6_zero'
+        logger.info(f'{ckpt_type} checkpoint saved {zero_checkpoint_name}')
 
     def _zero3_consolidated_16bit_state_dict(self):
         """

deepspeed/runtime/utils.py

@@ -944,7 +944,7 @@ def clip_tensors_by_global_norm(input_tensors,
     """Clip list of tensors by global norm.
     Args:
         input_tensors: List of tensors to be clipped
-        global_grad_norm (float, optional): Precomputed norm. Defaults to None.
+        global_norm (float, optional): Precomputed norm. Defaults to None.
         mpu (optional): model parallelism unit. Defaults to None.
         eps (float, optional): epsilon value added to grad norm. Defaults to 1e-6
     Returns:
@@ -953,7 +953,7 @@ def clip_tensors_by_global_norm(input_tensors,
     if global_norm is None:
         global_norm = get_global_norm_of_tensors(input_tensors, mpu=mpu)
 
-    clip_coef = max_norm / (global_grad_norm + eps)
+    clip_coef = max_norm / (global_norm + eps)
 
     if clip_coef < 1:
         for t in input_tensors:

deepspeed/runtime/zero/stage_1_and_2.py

@@ -23,7 +23,7 @@
 from deepspeed.utils import logger
 from deepspeed.moe.utils import is_moe_param
 from deepspeed.git_version_info import version
-
+from deepspeed.runtime.constants import PIPE_REPLICATED
 from deepspeed.checkpoint.constants import (DS_VERSION,
                                             PARTITION_COUNT,
                                             SINGLE_PARTITION_OF_FP32_GROUPS,
@@ -1747,41 +1747,6 @@ def step(self, closure=None):
             start_alignment_factor=self.nccl_start_alignment_factor,
             allgather_bucket_size=self.allgather_bucket_size)
 
-        #        for group_id, partitioned_params in enumerate(self.parallel_partitioned_bit16_groups):
-        #
-        #            # Sequential AllGather Best of both worlds
-        #            dp_world_size = dist.get_world_size(
-        #                group=self.real_dp_process_group[group_id])
-        #            num_shards = max(
-        #                1,
-        #                partitioned_params[partition_id].numel() * dp_world_size //
-        #                self.allgather_bucket_size)
-        #
-        #            shard_size = partitioned_params[partition_id].numel() // num_shards
-        #
-        #            # Enforce nccl/rccl alignment of start location of each shard
-        #            shard_size = shard_size - (shard_size % self.nccl_start_alignment_factor)
-        #
-        #            num_elements = shard_size
-        #
-        #            assert shard_size * num_shards <= partitioned_params[partition_id].numel()
-        #
-        #            for shard_id in range(num_shards):
-        #
-        #                if shard_id == (num_shards - 1):
-        #                    num_elements = partitioned_params[partition_id].numel(
-        #                    ) - shard_id * shard_size
-        #
-        #                shard_list = []
-        #                for dp_id in range(dp_world_size):
-        #                    curr_shard = partitioned_params[dp_id].narrow(
-        #                        0,
-        #                        shard_id * shard_size,
-        #                        num_elements).detach()
-        #                    shard_list.append(curr_shard)
-        #                dist.all_gather(shard_list,
-        #                                shard_list[partition_id],
-        #                                group=self.real_dp_process_group[group_id])
         self.stop_timers([OPTIMIZER_ALLGATHER])
 
         # TODO: we probably don't need this? just to be safe