[zero] revert PR microsoft#3166, it disabled grad clip for bf16 (micr…

…osoft#3790)

* zero++ tutorial PR (microsoft#3783)

* [Fix] _conv_flops_compute when padding is a str and stride=1 (microsoft#3169)

* fix conv_flops_compute when padding is a str when stride=1

* fix error

* change type of paddings to tuple

* fix padding calculation

* apply formatting check

---------

Co-authored-by: Cheng Li <[email protected]>
Co-authored-by: Olatunji Ruwase <[email protected]>

* fix interpolate flops compute (microsoft#3782)

* use `Flops Profiler` to test `model.generate()` (microsoft#2515)

* Update profiler.py

* pre-commit run --all-files

* Delete .DS_Store

* Delete .DS_Store

* Delete .DS_Store

---------

Co-authored-by: Jeff Rasley <[email protected]>
Co-authored-by: Cheng Li <[email protected]>

* revert PR microsoft#3166, it disabled grad clip for bf16

* ensure no loss scaling for non-fp16 dtypes

* revert PR microsoft#3611 (microsoft#3786)

* bump to 0.9.6

* ZeRO++ chinese blog (microsoft#3793)

* zeropp chinese blog

* try better quality images

* make title larger

* even larger...

* various fix

* center captions

* more fixes

* fix format

* remove staging trigger (microsoft#3792)

* DeepSpeed-Triton for Inference (microsoft#3748)

Co-authored-by: Stephen Youn <[email protected]>
Co-authored-by: Arash Bakhtiari <[email protected]>
Co-authored-by: Cheng Li <[email protected]>
Co-authored-by: Ethan Doe <[email protected]>
Co-authored-by: yidoe <[email protected]>
Co-authored-by: Jeff Rasley <[email protected]>

* ZeRO++ (microsoft#3784)

Co-authored-by: HeyangQin <[email protected]>
Co-authored-by: GuanhuaWang <[email protected]>
Co-authored-by: cmikeh2 <[email protected]>
Co-authored-by: Ammar Ahmad Awan <[email protected]>
Co-authored-by: Jeff Rasley <[email protected]>
Co-authored-by: Michael Wyatt <[email protected]>
Co-authored-by: Olatunji Ruwase <[email protected]>
Co-authored-by: Reza Yazdani <[email protected]>

* adding zero++ to navigation panel of deepspeed.ai (microsoft#3796)

* Add ZeRO++ Japanese blog (microsoft#3797)

* zeropp chinese blog

* try better quality images

* make title larger

* even larger...

* various fix

* center captions

* more fixes

* fix format

* add ZeRO++ Japanese blog

* add links

---------

Co-authored-by: HeyangQin <[email protected]>
Co-authored-by: Conglong Li <[email protected]>

* Bug Fixes for autotuner and flops profiler (microsoft#1880)

* fix autotuner when backward is not called

* fix format

---------

Co-authored-by: Olatunji Ruwase <[email protected]>

* Missing strided copy for gated MLP (microsoft#3788)

Co-authored-by: Ammar Ahmad Awan <[email protected]>
Co-authored-by: Jeff Rasley <[email protected]>
Co-authored-by: Logan Adams <[email protected]>

* Requires grad checking. (microsoft#3789)

Co-authored-by: Jeff Rasley <[email protected]>

* bump to 0.10.0

* Fix Bug in transform.cu (microsoft#3534)

* Bug fix

* Fixed formatting error

---------

Co-authored-by: Logan Adams <[email protected]>

* bug fix: triton importing error (microsoft#3799)

Co-authored-by: Stephen Youn <[email protected]>
Co-authored-by: Jeff Rasley <[email protected]>

---------

Co-authored-by: Heyang Qin <[email protected]>
Co-authored-by: Bill Luo <[email protected]>
Co-authored-by: Cheng Li <[email protected]>
Co-authored-by: Olatunji Ruwase <[email protected]>
Co-authored-by: Guorun <[email protected]>
Co-authored-by: stephen youn <[email protected]>
Co-authored-by: Stephen Youn <[email protected]>
Co-authored-by: Arash Bakhtiari <[email protected]>
Co-authored-by: Ethan Doe <[email protected]>
Co-authored-by: yidoe <[email protected]>
Co-authored-by: GuanhuaWang <[email protected]>
Co-authored-by: cmikeh2 <[email protected]>
Co-authored-by: Ammar Ahmad Awan <[email protected]>
Co-authored-by: Michael Wyatt <[email protected]>
Co-authored-by: Reza Yazdani <[email protected]>
Co-authored-by: Masahiro Tanaka <[email protected]>
Co-authored-by: Conglong Li <[email protected]>
Co-authored-by: Logan Adams <[email protected]>
Co-authored-by: Joe Mayer <[email protected]>
Co-authored-by: Ramya Ramineni <[email protected]>

deepspeed/runtime/zero/stage_1_and_2.py

@@ -481,6 +481,11 @@ def __init__(self,
                                             dynamic_loss_args=dynamic_loss_args)
         self.dynamic_loss_scale = self.loss_scaler.dynamic
 
+        if self.dtype != torch.float16:
+            # Only fp16 should use dynamic loss scaling
+            assert self.loss_scaler.cur_scale == 1.0
+            assert not self.dynamic_loss_scale
+
         see_memory_usage("Before initializing optimizer states", force=True)
         self.initialize_optimizer_states()
         see_memory_usage("After initializing optimizer states", force=True)
@@ -1670,21 +1675,19 @@ def step(self, closure=None):
             self.stop_timers(timer_names)
             return
 
-        # Step 1:- Calculate gradient norm using fp-16 grads
-        if self.dtype == torch.float16:
-            see_memory_usage('Before norm calculation')
-            scaled_global_grad_norm = self.scaled_global_norm()
-            self._global_grad_norm = scaled_global_grad_norm / prev_scale
-            see_memory_usage('After norm before optimizer')
+        # Step 1:- Calculate gradient norm using bit-16 grads
+        see_memory_usage('Before norm calculation')
+        scaled_global_grad_norm = self.scaled_global_norm()
+        self._global_grad_norm = scaled_global_grad_norm / prev_scale
+        see_memory_usage('After norm before optimizer')
 
         # Step 2:- run optimizer and upscaling simultaneously
         for i, group in enumerate(self.bit16_groups):
             self.start_timers([OPTIMIZER_GRADIENTS])
             partition_id = dist.get_rank(group=self.real_dp_process_group[i])
             if self.cpu_offload:
                 single_grad_partition = self.single_partition_of_fp32_groups[i].grad
-                if self.dtype == torch.float16:
-                    self.unscale_and_clip_grads([single_grad_partition], scaled_global_grad_norm)
+                self.unscale_and_clip_grads([single_grad_partition], scaled_global_grad_norm)
 
                 self.stop_timers([OPTIMIZER_GRADIENTS])
                 self.start_timers([OPTIMIZER_STEP])
@@ -1724,8 +1727,7 @@ def step(self, closure=None):
 
                 self.averaged_gradients[i] = None
 
-                if self.dtype == torch.float16:
-                    self.unscale_and_clip_grads([single_grad_partition], scaled_global_grad_norm)
+                self.unscale_and_clip_grads([single_grad_partition], scaled_global_grad_norm)
 
                 self.stop_timers([OPTIMIZER_GRADIENTS])