fixed grad scalar warning for bf16

aaa5715a · Mohammad Shoeybi · Jared Casper · 91384a5a · aaa5715a · aaa5715a
Commit aaa5715a authored Aug 16, 2022 by Mohammad Shoeybi Committed by Jared Casper Aug 16, 2022
Hide whitespace changes
Inline Side-by-side

Showing with 7 additions and 5 deletions

megatron/optimizer/distrib_optimizer.py megatron/optimizer/distrib_optimizer.py +3 -2

megatron/optimizer/optimizer.py megatron/optimizer/optimizer.py +4 -3

No files found.
--- a/megatron/optimizer/distrib_optimizer.py
+++ b/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'])

--- a/megatron/optimizer/optimizer.py
+++ b/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'])