re-introduce: stage3: efficient compute of scaled_global_grad_norm (#…

…5493)

reverting previous revert of this feature:

nelyahu@bc48371
in addition,
bug fix for offload mode.

deepspeed/runtime/zero/stage3.py

@@ -1409,7 +1409,7 @@ def complete_grad_norm_calculation_for_cpu_offload(self, params):
         norm_is_nan = total_norm.isnan()
         inf_or_nan = norm_is_nan.logical_or(norm_is_inf)
 
-        err = torch.tensor(-1.0, device=self.device, dtype=torch.float)
+        err = torch.tensor(-1.0, device=inf_or_nan.device, dtype=torch.float)
         total_norm = inf_or_nan * err + inf_or_nan.logical_not() * total_norm
 
         return total_norm