Fix fp16 (#6543)

Signed-off-by: MaximumEntropy <sandeep.subramanian.1@umontreal.ca>
NVIDIA · May 3, 2023 · b67297d · b67297d
1 parent d3dee47
commit b67297d
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/examples/nlp/language_modeling/megatron_gpt_eval.py b/examples/nlp/language_modeling/megatron_gpt_eval.py
@@ -196,6 +196,8 @@ def main(cfg) -> None:
             pretrained_cfg.activations_checkpoint_granularity = None
             pretrained_cfg.activations_checkpoint_method = None
             pretrained_cfg.precision = trainer.precision
+            if trainer.precision == "16":
+                pretrained_cfg.megatron_amp_O2 = False
         model = MegatronGPTModel.restore_from(
             restore_path=cfg.gpt_model_file,
             trainer=trainer,