minor updates and added dataclass check

Signed-off-by: arendu <[email protected]>

examples/nlp/language_modeling/conf/megatron_t5_prompt_learning.yaml

@@ -34,6 +34,7 @@ exp_manager:
     filename: "megatron_t5_prompt_tune--{${exp_manager.checkpoint_callback_params.monitor}:.3f}-{step}"
     model_parallel_size: ${model.tensor_model_parallel_size}
     save_best_model: True
+  create_early_stopping_callback: True
   early_stopping_callback_params:
     monitor: "val_loss"
     mode: "min"

nemo/utils/exp_manager.py

@@ -284,7 +284,8 @@ def exp_manager(trainer: 'pytorch_lightning.Trainer', cfg: Optional[Union[DictCo
                 pytorch lightning trainer. The ModelCheckpoint saves the top 3 models with the best "val_loss", the most
                 recent checkpoint under ``*last.ckpt``, and the final checkpoint after training completes under ``*end.ckpt``.
                 Defaults to True.
-            - create_early_stopping_callback (bool): Flag to decide if early stopping should be used to stop training. Default is False. See EarlyStoppingParams dataclass above.
+            - create_early_stopping_callback (bool): Flag to decide if early stopping should be used to stop training. Default is False.
+             See EarlyStoppingParams dataclass above.
             - files_to_copy (list): A list of files to copy to the experiment logging directory. Defaults to None which
                 copies no files.
             - log_local_rank_0_only (bool): Whether to only create log files for local rank 0. Defaults to False.

tests/core/test_config_utils.py

@@ -17,6 +17,8 @@
 
 import pytest
 import pytorch_lightning as ptl
+from pytorch_lightning.callbacks.early_stopping import EarlyStopping
+from nemo.utils.exp_manager import EarlyStoppingParams
 
 from nemo.core.config.pytorch_lightning import TrainerConfig
 from nemo.utils import config_utils
@@ -126,3 +128,12 @@ def test_ptl_config(self):
         assert signatures_match
         assert cls_subset is None
         assert dataclass_subset is None
+
+    @pytest.mark.unit
+    def test_early_stopping_config(self,):
+        result = config_utils.assert_dataclass_signature_match(EarlyStopping, EarlyStoppingParams)
+        signatures_match, cls_subset, dataclass_subset = result
+
+        assert signatures_match
+        assert cls_subset is None
+        assert dataclass_subset is None