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,6 +449,7 @@ class DistributedOptimizer(MixedPrecisionOptimizer):
 
         # 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:

megatron/optimizer/optimizer.py

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