Modify trainer.precision check and other small edits

Signed-off-by: Abhishree <abhishreetm@gmail.com>
NVIDIA · May 1, 2023 · 154dead · 154dead
1 parent 3901c89
commit 154dead
Show file tree

Hide file tree

Showing 3 changed files with 5 additions and 6 deletions.
diff --git a/nemo/collections/nlp/models/language_modeling/megatron_gpt_model.py b/nemo/collections/nlp/models/language_modeling/megatron_gpt_model.py
@@ -147,11 +147,11 @@ def __init__(self, cfg: DictConfig, trainer: Trainer):
             else:
                 self.model = Float16Module(module=self.model, precision=cfg.precision)
 
-        if self.trainer.precision == 'bf16':
+        if self.trainer.precision[0:4] == 'bf16':
             self.autocast_dtype = torch.bfloat16
-        elif int(self.trainer.precision) == 32:
+        elif self.trainer.precision[0:2] == '32':
             self.autocast_dtype = torch.float
-        elif int(self.trainer.precision) == 16:
+        elif self.trainer.precision[0:2] == '16':
             self.autocast_dtype = torch.half
         else:
             raise ValueError('precision must be in [32, 16, "bf16"]')

diff --git a/tests/collections/nlp/test_nlp_exportables.py b/tests/collections/nlp/test_nlp_exportables.py
@@ -82,10 +82,9 @@ def setup_method(self):
                     "accumulate_grad_batches": 1,
                     "precision": 32,
                     "accelerator": "gpu",
-                    "strategy": None,
+                    "strategy": 'auto',
                     "log_every_n_steps": 1,
                     "val_check_interval": 1,
-                    "resume_from_checkpoint": None,
                     "enable_checkpointing": False,
                     "logger": False,
                 },

diff --git a/tests/collections/nlp/test_rampup_batch_size.py b/tests/collections/nlp/test_rampup_batch_size.py
@@ -124,7 +124,7 @@ def trainer_cfg():
         'precision': 16,
         'logger': False,
         'enable_checkpointing': False,
-        'replace_sampler_ddp': False,
+        'use_distributed_sampler': False,
         'max_epochs': 1,
         'max_steps': 150,
         'log_every_n_steps': 10,