Fix issues with PL 1.8

Jenkinsfile

@@ -225,9 +225,11 @@ pipeline {
         stage('Speech to Text EMA') {
           steps {
             sh 'python examples/asr/asr_ctc/speech_to_text_ctc.py \
+            --config-path="../conf/" --config-name="config" \
             model.train_ds.manifest_filepath=/home/TestData/an4_dataset/an4_train.json \
             model.validation_ds.manifest_filepath=/home/TestData/an4_dataset/an4_val.json \
             trainer.devices=2 \
+            trainer.strategy="ddp" \
             trainer.accelerator="gpu" \
             +trainer.fast_dev_run=True \
             +exp_manager.ema.enable=True \

nemo/collections/nlp/parts/nlp_overrides.py

@@ -22,16 +22,16 @@
 
 import pytorch_lightning as pl
 import torch
+from lightning_lite.plugins import ClusterEnvironment
+from lightning_lite.utilities.types import _PATH
 from omegaconf import OmegaConf
 from pytorch_lightning.overrides import LightningDistributedModule
-from pytorch_lightning.plugins.environments.cluster_environment import ClusterEnvironment
 from pytorch_lightning.plugins.io.checkpoint_plugin import CheckpointIO
 from pytorch_lightning.plugins.precision.native_amp import NativeMixedPrecisionPlugin
 from pytorch_lightning.strategies.ddp import DDPStrategy
 from pytorch_lightning.trainer.trainer import Trainer
 from pytorch_lightning.utilities.exceptions import MisconfigurationException
 from pytorch_lightning.utilities.fetching import DataFetcher
-from pytorch_lightning.utilities.types import _PATH
 from torch.distributed.algorithms.ddp_comm_hooks.debugging_hooks import noop_hook
 from torch.nn.parallel import DistributedDataParallel
 

nemo/collections/tts/models/fastpitch.py

@@ -19,7 +19,7 @@
 from hydra.utils import instantiate
 from omegaconf import DictConfig, OmegaConf, open_dict
 from pytorch_lightning import Trainer
-from pytorch_lightning.loggers import LoggerCollection, TensorBoardLogger
+from pytorch_lightning.loggers import TensorBoardLogger
 
 from nemo.collections.common.parts.preprocessing import parsers
 from nemo.collections.tts.helpers.helpers import plot_alignment_to_numpy, plot_spectrogram_to_numpy, process_batch
@@ -228,11 +228,10 @@ def tb_logger(self):
             if self.logger is None and self.logger.experiment is None:
                 return None
             tb_logger = self.logger.experiment
-            if isinstance(self.logger, LoggerCollection):
-                for logger in self.logger:
-                    if isinstance(logger, TensorBoardLogger):
-                        tb_logger = logger.experiment
-                        break
+            for logger in self.trainer.loggers:
+                if isinstance(logger, TensorBoardLogger):
+                    tb_logger = logger.experiment
+                    break
             self._tb_logger = tb_logger
         return self._tb_logger
 

nemo/collections/tts/models/radtts.py

@@ -21,7 +21,7 @@
 from hydra.utils import instantiate
 from omegaconf import DictConfig, OmegaConf
 from pytorch_lightning import Trainer
-from pytorch_lightning.loggers import LoggerCollection, TensorBoardLogger
+from pytorch_lightning.loggers import TensorBoardLogger
 
 from nemo.collections.common.tokenizers.text_to_speech.tts_tokenizers import BaseTokenizer
 from nemo.collections.tts.helpers.helpers import plot_alignment_to_numpy
@@ -389,11 +389,10 @@ def tb_logger(self):
             if self.logger is None and self.logger.experiment is None:
                 return None
             tb_logger = self.logger.experiment
-            if isinstance(self.logger, LoggerCollection):
-                for logger in self.logger:
-                    if isinstance(logger, TensorBoardLogger):
-                        tb_logger = logger.experiment
-                        break
+            for logger in self.trainer.loggers:
+                if isinstance(logger, TensorBoardLogger):
+                    tb_logger = logger.experiment
+                    break
             self._tb_logger = tb_logger
         return self._tb_logger
 

nemo/collections/tts/models/tacotron2.py

@@ -20,7 +20,7 @@
 from hydra.utils import instantiate
 from omegaconf import MISSING, DictConfig, OmegaConf, open_dict
 from omegaconf.errors import ConfigAttributeError
-from pytorch_lightning.loggers import LoggerCollection, TensorBoardLogger, WandbLogger
+from pytorch_lightning.loggers import TensorBoardLogger, WandbLogger
 from torch import nn
 
 from nemo.collections.common.parts.preprocessing import parsers
@@ -284,11 +284,10 @@ def validation_step(self, batch, batch_idx):
     def validation_epoch_end(self, outputs):
         if self.logger is not None and self.logger.experiment is not None:
             logger = self.logger.experiment
-            if isinstance(self.logger, LoggerCollection):
-                for logger in self.logger:
-                    if isinstance(logger, TensorBoardLogger):
-                        logger = logger.experiment
-                        break
+            for logger in self.trainer.loggers:
+                if isinstance(logger, TensorBoardLogger):
+                    logger = logger.experiment
+                    break
             if isinstance(logger, TensorBoardLogger):
                 tacotron2_log_to_tb_func(
                     logger, outputs[0].values(), self.global_step, tag="val", log_images=True, add_audio=False,

nemo/collections/tts/models/waveglow.py

@@ -16,7 +16,7 @@
 import torch
 from hydra.utils import instantiate
 from omegaconf import DictConfig, open_dict
-from pytorch_lightning.loggers import LoggerCollection, TensorBoardLogger
+from pytorch_lightning.loggers import TensorBoardLogger
 
 from nemo.collections.tts.helpers.helpers import OperationMode, waveglow_log_to_tb_func
 from nemo.collections.tts.losses.waveglowloss import WaveGlowLoss
@@ -124,11 +124,10 @@ def validation_step(self, batch, batch_idx):
     def validation_epoch_end(self, outputs):
         if self.logger is not None and self.logger.experiment is not None:
             tb_logger = self.logger.experiment
-            if isinstance(self.logger, LoggerCollection):
-                for logger in self.logger:
-                    if isinstance(logger, TensorBoardLogger):
-                        tb_logger = logger.experiment
-                        break
+            for logger in self.trainer.loggers:
+                if isinstance(logger, TensorBoardLogger):
+                    tb_logger = logger.experiment
+                    break
             waveglow_log_to_tb_func(
                 tb_logger,
                 outputs[0].values(),

nemo/core/classes/exportable.py

@@ -16,6 +16,7 @@
 from typing import List, Union
 
 import torch
+from pytorch_lightning.core.module import _jit_is_scripting
 from torch.onnx import TrainingMode
 
 from nemo.core.classes import typecheck
@@ -128,7 +129,7 @@ def _export(
             # Set module mode
             with torch.onnx.select_model_mode_for_export(
                 self, training
-            ), torch.inference_mode(), torch.jit.optimized_execution(True):
+            ), torch.inference_mode(), torch.jit.optimized_execution(True), _jit_is_scripting():
 
                 if input_example is None:
                     input_example = self.input_module.input_example()

nemo/core/config/hydra_runner.py

@@ -67,6 +67,9 @@ def wrapper(cfg_passthrough: Optional[DictConfig] = None) -> Any:
                 # Set run.dir ONLY for ExpManager "compatibility" - to be removed.
                 overrides.append("hydra.run.dir=.")
 
+                # Set working directory to the job's output directory
+                overrides.append("hydra.job.chdir=True")
+
                 # Check if user set the schema.
                 if schema is not None:
                     # Create config store.

nemo/core/config/pytorch_lightning.py

@@ -84,6 +84,7 @@ class TrainerConfig:
     strategy: Any = None
     enable_checkpointing: bool = False
     enable_model_summary: bool = True
+    inference_mode: bool = True
 
 
 # Register the trainer config.

nemo/utils/exp_manager.py

@@ -32,7 +32,6 @@
 from omegaconf import DictConfig, OmegaConf, open_dict
 from pytorch_lightning.callbacks import Callback, ModelCheckpoint
 from pytorch_lightning.callbacks.timer import Interval, Timer
-from pytorch_lightning.loggers import LoggerCollection as _LoggerCollection
 from pytorch_lightning.loggers import TensorBoardLogger, WandbLogger
 from pytorch_lightning.loops import TrainingEpochLoop
 from pytorch_lightning.strategies.ddp import DDPStrategy
@@ -658,24 +657,6 @@ def get_git_diff():
         return "{}\n".format(err.output.decode("utf-8"))
 
 
-class LoggerList(_LoggerCollection):
-    """ A thin wrapper on Lightning's LoggerCollection such that name and version are better aligned with exp_manager
-    """
-
-    def __init__(self, _logger_iterable, nemo_name=None, nemo_version=""):
-        super().__init__(_logger_iterable)
-        self._nemo_name = nemo_name
-        self._nemo_version = nemo_version
-
-    @property
-    def name(self) -> str:
-        return self._nemo_name
-
-    @property
-    def version(self) -> str:
-        return self._nemo_version
-
-
 def configure_loggers(
     trainer: 'pytorch_lightning.Trainer',
     exp_dir: [Path, str],
@@ -718,9 +699,6 @@ def configure_loggers(
         logger_list.append(wandb_logger)
         logging.info("WandBLogger has been set up")
 
-    logger_list = (
-        LoggerList(logger_list, nemo_name=name, nemo_version=version) if len(logger_list) > 1 else logger_list[0]
-    )
     trainer._logger_connector.configure_logger(logger_list)
 
 

requirements/requirements_lightning.txt

@@ -1,6 +1,6 @@
 hydra-core>=1.2.0,<1.3
 omegaconf>=2.2,<2.3
-pytorch-lightning>=1.7.0,<1.8
+pytorch-lightning>=1.8.0
 pyyaml<6  # Pinned until omegaconf works with pyyaml>=6
 torchmetrics>=0.4.1rc0
 transformers>=4.0.1,<=4.21.2

tests/core/test_config_utils.py

@@ -116,11 +116,8 @@ class DummyDataClass:
     @pytest.mark.unit
     def test_ptl_config(self):
         PTL_DEPRECATED = [
-            'distributed_backend',
-            'automatic_optimization',
             'gpus',
             'num_processes',
-            'weights_save_path',
         ]
 
         result = config_utils.assert_dataclass_signature_match(ptl.Trainer, TrainerConfig, ignore_args=PTL_DEPRECATED)

tests/core/test_optimizers_schedulers.py

@@ -20,6 +20,7 @@
 import pytorch_lightning as pl
 import torch
 import torch.optim
+from pytorch_lightning.utilities import rank_zero_only
 
 from nemo.core import config, optim
 from nemo.core.optim.lr_scheduler import AVAILABLE_SCHEDULERS
@@ -85,7 +86,7 @@ def configure_optimizers(self):
 
 
 class Callback(pl.callbacks.Callback):
-    @pl.utilities.distributed.rank_zero_only
+    @rank_zero_only
     def on_train_end(self, trainer, module):
         count = module.my_opt.param_groups[0]['count']
         if trainer.global_step != count or trainer.global_step != module.max_steps:
@@ -110,13 +111,13 @@ class SchedulerNoOpCallback(Callback):
     def on_train_batch_end(self, trainer: pl.Trainer, pl_module, outputs, batch, batch_idx):
         # pl_module.max_steps is "original" max steps without trainer extra steps.
         if (trainer.global_step + 1) % 3 == 0 and (trainer.global_step + 1) < pl_module.max_steps:
-            schedulers = trainer.lr_schedulers
+            schedulers = trainer.lr_scheduler_configs
 
             for scheduler in schedulers:
                 # Decrement the counter by 2, then perform a scheduler.step() to perform a no-up
                 # as well as update the optimizer lr in all param groups
-                scheduler['scheduler'].last_epoch -= 2
-                scheduler['scheduler'].step()
+                scheduler.scheduler.last_epoch -= 2
+                scheduler.scheduler.step()
 
             # Increase the max step count by 1
             trainer.fit_loop.max_steps = trainer.fit_loop.max_steps + 1

tests/core_ptl/check_for_ranks.py

@@ -18,7 +18,7 @@
 import torch
 from omegaconf import OmegaConf
 from pytorch_lightning import Trainer
-from pytorch_lightning.utilities.distributed import rank_zero_only
+from pytorch_lightning.utilities import rank_zero_only
 
 from nemo.core import ModelPT
 from nemo.utils import logging

tests/core_ptl/test_ptl_stateless_timer.py

@@ -19,7 +19,6 @@
 import torch
 from omegaconf import OmegaConf
 from pytorch_lightning import Trainer
-from pytorch_lightning.utilities.distributed import rank_zero_only
 
 from nemo.core import ModelPT
 from nemo.utils import logging