Change accelerator to 'auto' in nlp_checkpoint_port.py (#7747)

* Change accelerator to auto Signed-off-by: Abhishree <[email protected]> * Pass omegaconf object to trainer in nlp_checkpoint_port.py Signed-off-by: Abhishree <[email protected]> * Pass omegaconf object to trainer in export.py Signed-off-by: Abhishree <[email protected]> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Signed-off-by: Abhishree <[email protected]> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: Eric Harper <[email protected]>
NVIDIA · Oct 19, 2023 · 027ec11 · 027ec11
1 parent e57aff0
commit 027ec11
Show file tree

Hide file tree

Showing 3 changed files with 9 additions and 6 deletions.
diff --git a/nemo/core/config/pytorch_lightning.py b/nemo/core/config/pytorch_lightning.py
@@ -54,7 +54,7 @@ class TrainerConfig:
     limit_test_batches: Any = 1.0
     val_check_interval: Any = 1.0
     log_every_n_steps: int = 50
-    accelerator: Optional[str] = None
+    accelerator: Optional[str] = 'auto'
     sync_batchnorm: bool = False
     precision: Any = 32
     num_sanity_val_steps: int = 2
@@ -68,8 +68,8 @@ class TrainerConfig:
     gradient_clip_algorithm: str = 'norm'
     max_time: Optional[Any] = None  # can be one of Union[str, timedelta, Dict[str, int], None]
     reload_dataloaders_every_n_epochs: int = 0
-    devices: Any = None
-    strategy: Any = None
+    devices: Any = 'auto'
+    strategy: Any = 'auto'
     enable_checkpointing: bool = False
     enable_model_summary: bool = True
     inference_mode: bool = True

diff --git a/scripts/export.py b/scripts/export.py
@@ -30,6 +30,7 @@
 import sys
 
 import torch
+from omegaconf import OmegaConf
 from pytorch_lightning import Trainer
 
 import nemo
@@ -103,7 +104,8 @@ def nemo_export(argv):
         logger=False,
         enable_checkpointing=False,
     )
-    trainer = Trainer(cfg_trainer)
+    cfg_trainer = OmegaConf.to_container(OmegaConf.create(cfg_trainer))
+    trainer = Trainer(**cfg_trainer)
 
     logging.info("Restoring NeMo model from '{}'".format(nemo_in))
     try:

diff --git a/scripts/nemo_legacy_import/nlp_checkpoint_port.py b/scripts/nemo_legacy_import/nlp_checkpoint_port.py
@@ -82,13 +82,14 @@ def nemo_convert(argv):
     # Create a PL trainer object which is required for restoring Megatron models
     cfg_trainer = TrainerConfig(
         devices=1,
-        accelerator="ddp",
+        accelerator='auto',
         num_nodes=1,
         # Need to set the following two to False as ExpManager will take care of them differently.
         logger=False,
         enable_checkpointing=False,
     )
-    trainer = pl.Trainer(cfg_trainer)
+    cfg_trainer = OmegaConf.to_container(OmegaConf.create(cfg_trainer))
+    trainer = pl.Trainer(**cfg_trainer)
 
     logging.info("Restoring NeMo model from '{}'".format(nemo_in))
     try: