refresh fp32 params from fp16 model weights if refreshing from state dict fails

Author: sdtblck

deepspeed/runtime/fp16/fused_optimizer.py

@@ -1,26 +1,32 @@
-'''
+"""
 Copyright 2019 The Microsoft DeepSpeed Team
 
 Copyright NVIDIA/apex
 This file is adapted from FP16_Optimizer in NVIDIA/apex
-'''
+"""
 
 import torch
 import math
 from torch._utils import _flatten_dense_tensors, _unflatten_dense_tensors
 
 from deepspeed.runtime.utils import get_grad_norm, CheckOverflow, get_weight_norm
-from deepspeed.runtime.fp16.loss_scaler import INITIAL_LOSS_SCALE, SCALE_WINDOW, MIN_LOSS_SCALE
+from deepspeed.runtime.fp16.loss_scaler import (
+    INITIAL_LOSS_SCALE,
+    SCALE_WINDOW,
+    MIN_LOSS_SCALE,
+)
 from deepspeed.utils import logger, log_dist
 
 from ...ops.adam import FusedAdam
+
 FP16_FUSED_SUPPORTED_OPTIMIZERS = [
     FusedAdam,
 ]
 
 # Add apex FusedAdam to supported list if apex is installed
 try:
     import apex
+
     FP16_FUSED_SUPPORTED_OPTIMIZERS.append(apex.optimizers.FusedAdam)
 except ImportError:
     pass
@@ -34,34 +40,35 @@ def is_fp16_fused_supported_optimizer(optimizer):
         bool: True if ``optimizer`` is compatible with ``FP16_Optimizer``.
     """
     from deepspeed.runtime.config import ONEBIT_ADAM_OPTIMIZER
+
     if isinstance(optimizer, tuple(FP16_FUSED_SUPPORTED_OPTIMIZERS)):
         return True
     if optimizer.__class__.__name__.lower() == ONEBIT_ADAM_OPTIMIZER.lower():
         return True
     return False
 
 
-
-
 class FP16_Optimizer(object):
     """
    FP16 Optimizer for training fp16 models. Handles loss scaling.
 
    For usage example please see, TODO:  DeepSpeed V2 Tutorial
     """
 
-    def __init__(self,
-                 init_optimizer,
-                 deepspeed=None,
-                 static_loss_scale=1.0,
-                 dynamic_loss_scale=False,
-                 initial_dynamic_scale=2 ** 32,
-                 dynamic_loss_args=None,
-                 verbose=True,
-                 mpu=None,
-                 clip_grad=0.0,
-                 fused_adam_legacy=False,
-                 timers=None):
+    def __init__(
+        self,
+        init_optimizer,
+        deepspeed=None,
+        static_loss_scale=1.0,
+        dynamic_loss_scale=False,
+        initial_dynamic_scale=2 ** 32,
+        dynamic_loss_args=None,
+        verbose=True,
+        mpu=None,
+        clip_grad=0.0,
+        fused_adam_legacy=False,
+        timers=None,
+    ):
 
         self.fused_adam_legacy = fused_adam_legacy
         self.timers = timers
@@ -78,23 +85,28 @@ def __init__(self,
         # loop to deal with groups
         for i, param_group in enumerate(self.optimizer.param_groups):
             # push this group to list before modify
-            self.fp16_groups.append(param_group['params'])
+            self.fp16_groups.append(param_group["params"])
             # init fp16 weight buffer, flattened
             self.fp16_groups_flat.append(
-                _flatten_dense_tensors([p.clone().detach()
-                                        for p in self.fp16_groups[i]]))
+                _flatten_dense_tensors(
+                    [p.clone().detach() for p in self.fp16_groups[i]]
+                )
+            )
             # set model fp16 weight to slices of flattened buffer
-            updated_params = _unflatten_dense_tensors(self.fp16_groups_flat[i],
-                                                      self.fp16_groups[i])
+            updated_params = _unflatten_dense_tensors(
+                self.fp16_groups_flat[i], self.fp16_groups[i]
+            )
             for p, q in zip(self.fp16_groups[i], updated_params):
                 p.data = q.data
             # init master weight, flattened
             self.fp32_groups_flat.append(
-                self.fp16_groups_flat[i].clone().float().detach())
+                self.fp16_groups_flat[i].clone().float().detach()
+            )
             # modify optimizer of have flat master weight
             self.fp32_groups_flat[
-                i].requires_grad = True  # keep this in case internal optimizer uses it
-            param_group['params'] = [self.fp32_groups_flat[i]]
+                i
+            ].requires_grad = True  # keep this in case internal optimizer uses it
+            param_group["params"] = [self.fp32_groups_flat[i]]
 
         # we may have a way of fusing dynamic scale. Do not support for now
         if dynamic_loss_scale:
@@ -120,8 +132,8 @@ def __init__(self,
         self.clip_grad = clip_grad
         self.norm_type = 2
 
-        TORCH_MAJOR = int(torch.__version__.split('.')[0])
-        TORCH_MINOR = int(torch.__version__.split('.')[1])
+        TORCH_MAJOR = int(torch.__version__.split(".")[0])
+        TORCH_MINOR = int(torch.__version__.split(".")[1])
         if TORCH_MAJOR == 0 and TORCH_MINOR <= 4:
             self.clip_grad_norm = torch.nn.utils.clip_grad_norm
         else:
@@ -131,16 +143,16 @@ def __init__(self,
         self.mpu = mpu
 
         self.overflow = False
-        self.overflow_checker = CheckOverflow(self.fp16_groups,
-                                              mpu=self.mpu,
-                                              deepspeed=deepspeed)
+        self.overflow_checker = CheckOverflow(
+            self.fp16_groups, mpu=self.mpu, deepspeed=deepspeed
+        )
         self.initialize_optimizer_states()
 
     def initialize_optimizer_states(self):
         for i, group in enumerate(self.fp16_groups):
             self.fp32_groups_flat[i].grad = torch.zeros(
-                self.fp32_groups_flat[i].size(),
-                device=self.fp32_groups_flat[i].device)
+                self.fp32_groups_flat[i].size(), device=self.fp32_groups_flat[i].device
+            )
 
         self.optimizer.step()
 
@@ -172,12 +184,15 @@ def step_fused_adam(self, closure=None):
         norm_groups = []
         for i, group in enumerate(self.fp16_groups):
             grads_groups_flat.append(
-                _flatten_dense_tensors([
-                    torch.zeros(p.size(),
-                                dtype=p.dtype,
-                                device=p.device) if p.grad is None else p.grad
-                    for p in group
-                ]))
+                _flatten_dense_tensors(
+                    [
+                        torch.zeros(p.size(), dtype=p.dtype, device=p.device)
+                        if p.grad is None
+                        else p.grad
+                        for p in group
+                    ]
+                )
+            )
             norm_groups.append(get_weight_norm(grads_groups_flat[i], mpu=self.mpu))
 
         self.overflow = self.overflow_checker.check_using_norm(norm_groups)
@@ -186,23 +201,26 @@ def step_fused_adam(self, closure=None):
 
         if self.overflow:
             if self.verbose:
-                logger.info("[deepspeed] OVERFLOW! Skipping step. Attempted loss "
-                            "scale: {}, reducing to {}".format(
-                    prev_scale,
-                    self.cur_scale))
+                logger.info(
+                    "[deepspeed] OVERFLOW! Skipping step. Attempted loss "
+                    "scale: {}, reducing to {}".format(prev_scale, self.cur_scale)
+                )
             return self.overflow
-        combined_scale = self.unscale_and_clip_grads(grads_groups_flat,
-                                                     norm_groups,
-                                                     apply_scale=False)
+        combined_scale = self.unscale_and_clip_grads(
+            grads_groups_flat, norm_groups, apply_scale=False
+        )
         # 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],
-                            scale=combined_scale,
-                            grad_norms=norm_groups)
+        self.optimizer.step(
+            grads=[[g] for g in grads_groups_flat],
+            output_params=[[p] for p in self.fp16_groups_flat],
+            scale=combined_scale,
+            grad_norms=norm_groups,
+        )
         # TODO: we probably don't need this? just to be safe
         for i in range(len(norm_groups)):
-            updated_params = _unflatten_dense_tensors(self.fp16_groups_flat[i],
-                                                      self.fp16_groups[i])
+            updated_params = _unflatten_dense_tensors(
+                self.fp16_groups_flat[i], self.fp16_groups[i]
+            )
             for p, q in zip(self.fp16_groups[i], updated_params):
                 p.data = q.data
         return self.overflow
@@ -230,11 +248,11 @@ def step(self, closure=None):
             return self.step_fused_adam()
 
         COMPUTE_NORM = "compute_norm"
-        OVERFLOW_CHECK = 'overflow_check'
+        OVERFLOW_CHECK = "overflow_check"
         OVERFLOW_TIMERS = [COMPUTE_NORM, OVERFLOW_CHECK]
-        UNSCALE_AND_CLIP = 'unscale_and_clip'
-        BASIC_STEP = 'basic_step'
-        UPDATE_FP16 = 'update_fp16'
+        UNSCALE_AND_CLIP = "unscale_and_clip"
+        BASIC_STEP = "basic_step"
+        UPDATE_FP16 = "update_fp16"
         STEP_TIMERS = OVERFLOW_TIMERS + [UNSCALE_AND_CLIP, BASIC_STEP, UPDATE_FP16]
 
         # First determine if there is overflow.
@@ -252,7 +270,8 @@ def step(self, closure=None):
                 log_dist(
                     "Overflow detected. Skipping step. Attempted loss "
                     f"scale: {prev_scale}, reducing to {self.cur_scale}",
-                    ranks=[0])
+                    ranks=[0],
+                )
             # Clear gradients
             for i, group in enumerate(self.fp16_groups):
                 for p in group:
@@ -266,12 +285,15 @@ def step(self, closure=None):
             data_type = self.fp32_groups_flat[i].dtype
 
             grads_groups_flat.append(
-                _flatten_dense_tensors([
-                    torch.zeros(p.size(),
-                                dtype=data_type,
-                                device=p.device)
-                    if p.grad is None else p.grad.to(data_type) for p in group
-                ]))
+                _flatten_dense_tensors(
+                    [
+                        torch.zeros(p.size(), dtype=data_type, device=p.device)
+                        if p.grad is None
+                        else p.grad.to(data_type)
+                        for p in group
+                    ]
+                )
+            )
 
             for p in group:
                 p.grad = None
@@ -296,8 +318,9 @@ def step(self, closure=None):
 
         self.start_timers([UPDATE_FP16])
         for i in range(len(self.fp16_groups)):
-            updated_params = _unflatten_dense_tensors(self.fp32_groups_flat[i],
-                                                      self.fp16_groups[i])
+            updated_params = _unflatten_dense_tensors(
+                self.fp32_groups_flat[i], self.fp16_groups[i]
+            )
             for p, q in zip(self.fp16_groups[i], updated_params):
                 p.data.copy_(q.data)
         self.stop_timers([UPDATE_FP16])
@@ -314,15 +337,15 @@ def unscale_and_clip_grads(self, grad_groups_flat, norm_groups, apply_scale=True
 
         # compute combined scale factor for this group
         combined_scale = self.cur_scale
-        if self.clip_grad > 0.:
+        if self.clip_grad > 0.0:
             # norm is in fact norm*scale
             clip = ((total_norm / self.cur_scale) + 1e-6) / self.clip_grad
             if clip > 1:
                 combined_scale = clip * self.cur_scale
 
         if apply_scale:
             for grad in grad_groups_flat:
-                grad.data.mul_(1. / combined_scale)
+                grad.data.mul_(1.0 / combined_scale)
 
         return combined_scale
 
@@ -341,8 +364,9 @@ def _update_scale(self, skip):
         if self.dynamic_loss_scale:
             prev_scale = self.cur_scale
             if skip:
-                self.cur_scale = max(self.cur_scale / self.scale_factor,
-                                     self.min_loss_scale)
+                self.cur_scale = max(
+                    self.cur_scale / self.scale_factor, self.min_loss_scale
+                )
                 self.last_overflow_iter = self.cur_iter
                 if self.verbose:
                     logger.info(f"\nGrad overflow on iteration {self.cur_iter}")
@@ -356,7 +380,8 @@ def _update_scale(self, skip):
                     self.cur_scale *= self.scale_factor
                     if self.verbose:
                         logger.info(
-                            f"No Grad overflow for {self.scale_window} iterations")
+                            f"No Grad overflow for {self.scale_window} iterations"
+                        )
                         logger.info(
                             f"Increasing dynamic loss scale from {prev_scale} to {self.cur_scale}"
                         )
@@ -398,16 +423,16 @@ def state_dict(self):
             torch.save(checkpoint, "saved.pth")
         """
         state_dict = {}
-        state_dict['dynamic_loss_scale'] = self.dynamic_loss_scale
-        state_dict['cur_scale'] = self.cur_scale
-        state_dict['cur_iter'] = self.cur_iter
-        if state_dict['dynamic_loss_scale']:
-            state_dict['last_overflow_iter'] = self.last_overflow_iter
-            state_dict['scale_factor'] = self.scale_factor
-            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["dynamic_loss_scale"] = self.dynamic_loss_scale
+        state_dict["cur_scale"] = self.cur_scale
+        state_dict["cur_iter"] = self.cur_iter
+        if state_dict["dynamic_loss_scale"]:
+            state_dict["last_overflow_iter"] = self.last_overflow_iter
+            state_dict["scale_factor"] = self.scale_factor
+            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
         return state_dict
 
     # Refresh fp32 master params from fp16 copies
@@ -432,16 +457,16 @@ def load_state_dict(self, state_dict, load_optimizer_states=True):
             optimizer.load_state_dict(checkpoint['optimizer'])
         """
         # I think it should actually be ok to reload the optimizer before the model.
-        self.dynamic_loss_scale = state_dict['dynamic_loss_scale']
-        self.cur_scale = state_dict['cur_scale']
-        self.cur_iter = state_dict['cur_iter']
-        if state_dict['dynamic_loss_scale']:
-            self.last_overflow_iter = state_dict['last_overflow_iter']
-            self.scale_factor = state_dict['scale_factor']
-            self.scale_window = state_dict['scale_window']
+        self.dynamic_loss_scale = state_dict["dynamic_loss_scale"]
+        self.cur_scale = state_dict["cur_scale"]
+        self.cur_iter = state_dict["cur_iter"]
+        if state_dict["dynamic_loss_scale"]:
+            self.last_overflow_iter = state_dict["last_overflow_iter"]
+            self.scale_factor = state_dict["scale_factor"]
+            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.optimizer.load_state_dict(state_dict["optimizer_state_dict"])
+        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
@@ -456,8 +481,17 @@ def load_state_dict(self, state_dict, load_optimizer_states=True):
         # the current optimizer instance.  In our case, as long as the current FP16_Optimizer has been
         # constructed in the same way as the one whose state_dict we are loading, the same master params
         # are guaranteed to exist, so we can just copy_() from the saved master params.
-        for current, saved in zip(self.fp32_groups_flat, state_dict['fp32_groups_flat']):
-            current.data.copy_(saved.data)
+        try:
+            for current, saved in zip(
+                self.fp32_groups_flat, state_dict["fp32_groups_flat"]
+            ):
+                current.data.copy_(saved.data)
+        except RuntimeError as error:
+            print(error)
+            print(
+                "Error in loading fp32 model parameters!\nRefreshing fp32 model params from the model's fp16 params instead. This may incur some precision loss."
+            )
+            self.refresh_fp32_params()
 
     def __repr__(self):
         return repr(self.optimizer)