Merge branch 'fix_grad_scalar_warning' into 'main'

fixed grad scalar warning for bf16

See merge request ADLR/megatron-lm!442

megatron/optimizer/distrib_optimizer.py

@@ -449,8 +449,9 @@ class DistributedOptimizer(MixedPrecisionOptimizer):
 
         # Grad scaler.
         if 'grad_scaler' not in state_dict:
-            print_rank_0('***WARNING*** found an old checkpoint, will not '
-                         'load grad scaler ...')
+            if self.fp16:
+                print_rank_0('***WARNING*** found an old checkpoint, will not '
+                             'load grad scaler ...')
         else:
             if self.grad_scaler:
                 self.grad_scaler.load_state_dict(state_dict['grad_scaler'])

megatron/optimizer/optimizer.py

@@ -679,9 +679,10 @@ class Float16OptimizerWithFloat16Params(MixedPrecisionOptimizer):
         self.optimizer.load_state_dict(state_dict[optimizer_key])
 
         # Grad scaler.
-        if self.fp16 and 'grad_scaler' not in state_dict:
-            print_rank_0('***WARNING*** found an old checkpoint, will not '
-                         'load grad scaler ...')
+        if 'grad_scaler' not in state_dict:
+            if self.fp16:
+                print_rank_0('***WARNING*** found an old checkpoint, will not '
+                             'load grad scaler ...')
         else:
             if self.grad_scaler:
                 self.grad_scaler.load_state_dict(state_dict['grad_scaler'])