Repository: NVIDIA/NeMo-Aligner (language: python)

Issue:
chore: main to dev
sync-branches: New code has just landed in main, so let's bring dev up to speed!


Discussion / hints:

Fix patch:
diff --git a/Dockerfile b/Dockerfile
index 4d22c6a21..44a9f8651 100644
--- a/Dockerfile
+++ b/Dockerfile
@@ -115,8 +115,6 @@ RUN pip uninstall -y megatron-core && \
     fi && \
     pip install -e .
 
-RUN pip install --no-cache-dir lightning # can remove this when NEMO_TAG is bumped to include lightning install
-
 COPY --from=aligner-bump /opt/NeMo-Aligner /opt/NeMo-Aligner
 RUN cd /opt/NeMo-Aligner && \
     pip install --no-deps -e .
diff --git a/examples/nlp/gpt/serve_reward_model.py b/examples/nlp/gpt/serve_reward_model.py
index 5f5dfc609..a425f26e7 100644
--- a/examples/nlp/gpt/serve_reward_model.py
+++ b/examples/nlp/gpt/serve_reward_model.py
@@ -13,7 +13,7 @@
 # limitations under the License.
 
 import torch
-from lightning.pytorch.trainer.trainer import Trainer
+from pytorch_lightning.trainer.trainer import Trainer
 
 from nemo.collections.nlp.parts.nlp_overrides import NLPDDPStrategy
 from nemo.core.config import hydra_runner
diff --git a/nemo_aligner/models/nlp/gpt/gpt_sft_model.py b/nemo_aligner/models/nlp/gpt/gpt_sft_model.py
index e8a47763f..d3a615500 100644
--- a/nemo_aligner/models/nlp/gpt/gpt_sft_model.py
+++ b/nemo_aligner/models/nlp/gpt/gpt_sft_model.py
@@ -16,10 +16,10 @@
 
 import hydra
 import torch
-from lightning.pytorch.trainer.trainer import Trainer
 from megatron.core.num_microbatches_calculator import get_micro_batch_size, get_num_microbatches
 from megatron.core.pipeline_parallel.schedules import get_forward_backward_func
 from omegaconf.dictconfig import DictConfig
+from pytorch_lightning.trainer.trainer import Trainer
 
 from nemo.collections.nlp.models.language_modeling.megatron_gpt_model import MegatronGPTModel
 from nemo.collections.nlp.modules.common.megatron.utils import get_iterator_k_split
diff --git a/nemo_aligner/models/nlp/gpt/megatron_gpt_critic.py b/nemo_aligner/models/nlp/gpt/megatron_gpt_critic.py
index a800a6950..a61cd2b8d 100644
--- a/nemo_aligner/models/nlp/gpt/megatron_gpt_critic.py
+++ b/nemo_aligner/models/nlp/gpt/megatron_gpt_critic.py
@@ -15,11 +15,11 @@
 from enum import Enum
 
 import torch
-from lightning.pytorch.trainer.trainer import Trainer
 from megatron.core.num_microbatches_calculator import get_num_microbatches, reconfigure_num_microbatches_calculator
 from megatron.core.pipeline_parallel.schedules import get_forward_backward_func
 from megatron.core.transformer.module import Float16Module
 from omegaconf.dictconfig import DictConfig
+from pytorch_lightning.trainer.trainer import Trainer
 
 from nemo.collections.nlp.modules.common.megatron.utils import (
     average_losses_across_data_parallel_group,
diff --git a/nemo_aligner/models/nlp/gpt/megatron_gpt_dpo_model.py b/nemo_aligner/models/nlp/gpt/megatron_gpt_dpo_model.py
index 715914be7..952b4e897 100644
--- a/nemo_aligner/models/nlp/gpt/megatron_gpt_dpo_model.py
+++ b/nemo_aligner/models/nlp/gpt/megatron_gpt_dpo_model.py
@@ -16,11 +16,11 @@
 from functools import partial
 
 import torch
-from lightning.pytorch.trainer.trainer import Trainer
 from megatron.core.num_microbatches_calculator import get_num_microbatches
 from megatron.core.pipeline_parallel.schedules import get_forward_backward_func
 from megatron.core.utils import divide
 from omegaconf.dictconfig import DictConfig
+from pytorch_lightning.trainer.trainer import Trainer
 
 from nemo.collections.nlp.models.language_modeling.megatron_gpt_model import MegatronGPTModel
 from nemo.collections.nlp.modules.common.megatron.utils import (
diff --git a/nemo_aligner/models/nlp/gpt/megatron_gpt_kto_model.py b/nemo_aligner/models/nlp/gpt/megatron_gpt_kto_model.py
index 5a16d18b8..0be08e808 100644
--- a/nemo_aligner/models/nlp/gpt/megatron_gpt_kto_model.py
+++ b/nemo_aligner/models/nlp/gpt/megatron_gpt_kto_model.py
@@ -16,11 +16,11 @@
 from functools import partial
 
 import torch
-from lightning.pytorch.trainer.trainer import Trainer
 from megatron.core.num_microbatches_calculator import get_num_microbatches
 from megatron.core.pipeline_parallel.schedules import get_forward_backward_func
 from megatron.core.utils import divide
 from omegaconf.dictconfig import DictConfig
+from pytorch_lightning.trainer.trainer import Trainer
 
 from nemo.collections.nlp.models.language_modeling.megatron_gpt_model import MegatronGPTModel
 from nemo.collections.nlp.modules.common.megatron.utils import (
diff --git a/nemo_aligner/models/nlp/gpt/megatron_gpt_ppo_actor.py b/nemo_aligner/models/nlp/gpt/megatron_gpt_ppo_actor.py
index 43643faab..275e02e82 100644
--- a/nemo_aligner/models/nlp/gpt/megatron_gpt_ppo_actor.py
+++ b/nemo_aligner/models/nlp/gpt/megatron_gpt_ppo_actor.py
@@ -16,12 +16,12 @@
 
 import torch
 import torch.distributed
-from lightning.pytorch.trainer.trainer import Trainer
 from megatron.core.num_microbatches_calculator import get_num_microbatches
 from megatron.core.pipeline_parallel.schedules import get_forward_backward_func
 from megatron.core.utils import divide
 from omegaconf import OmegaConf
 from omegaconf.dictconfig import DictConfig
+from pytorch_lightning.trainer.trainer import Trainer
 
 from nemo.collections.nlp.models.language_modeling.megatron_gpt_model import MegatronGPTModel
 from nemo.collections.nlp.modules.common.megatron.utils import (
diff --git a/nemo_aligner/models/nlp/gpt/megatron_gpt_regression_reward_model.py b/nemo_aligner/models/nlp/gpt/megatron_gpt_regression_reward_model.py
index 06b719610..eb9333c12 100644
--- a/nemo_aligner/models/nlp/gpt/megatron_gpt_regression_reward_model.py
+++ b/nemo_aligner/models/nlp/gpt/megatron_gpt_regression_reward_model.py
@@ -14,10 +14,10 @@
 
 
 import torch
-from lightning.pytorch.trainer.trainer import Trainer
 from megatron.core.num_microbatches_calculator import get_num_microbatches
 from megatron.core.pipeline_parallel.schedules import get_forward_backward_func
 from omegaconf.dictconfig import DictConfig
+from pytorch_lightning.trainer.trainer import Trainer
 
 from nemo.collections.nlp.modules.common.megatron.utils import (
     average_losses_across_data_parallel_group,
diff --git a/nemo_aligner/models/nlp/gpt/megatron_gpt_reward_model.py b/nemo_aligner/models/nlp/gpt/megatron_gpt_reward_model.py
index e1d2fcc24..596436c67 100644
--- a/nemo_aligner/models/nlp/gpt/megatron_gpt_reward_model.py
+++ b/nemo_aligner/models/nlp/gpt/megatron_gpt_reward_model.py
@@ -17,11 +17,11 @@
 from typing import List, Tuple, Union
 
 import torch
-from lightning.pytorch.trainer.trainer import Trainer
 from megatron.core.num_microbatches_calculator import get_num_microbatches
 from megatron.core.pipeline_parallel.schedules import get_forward_backward_func
 from megatron.core.utils import divide
 from omegaconf.dictconfig import DictConfig
+from pytorch_lightning.trainer.trainer import Trainer
 
 from nemo.collections.nlp.models.language_modeling.megatron_gpt_model import MegatronGPTModel, get_specs
 from nemo.collections.nlp.modules.common.megatron.utils import (
diff --git a/nemo_aligner/models/nlp/gpt/megatron_gpt_rs_actor.py b/nemo_aligner/models/nlp/gpt/megatron_gpt_rs_actor.py
index ebd5adf4b..49dd3846c 100644
--- a/nemo_aligner/models/nlp/gpt/megatron_gpt_rs_actor.py
+++ b/nemo_aligner/models/nlp/gpt/megatron_gpt_rs_actor.py
@@ -15,13 +15,13 @@
 from contextlib import nullcontext
 
 import torch
-from lightning.pytorch.trainer.trainer import Trainer
 from megatron.core import parallel_state
 from megatron.core.num_microbatches_calculator import get_num_microbatches
 from megatron.core.pipeline_parallel.schedules import get_forward_backward_func
 from megatron.core.utils import divide
 from omegaconf import OmegaConf
 from omegaconf.dictconfig import DictConfig
+from pytorch_lightning.trainer.trainer import Trainer
 
 from nemo.collections.nlp.models.language_modeling.megatron_gpt_model import MegatronGPTModel
 from nemo.collections.nlp.modules.common.megatron.utils import (
diff --git a/nemo_aligner/models/nlp/gpt/megatron_gpt_spin_model.py b/nemo_aligner/models/nlp/gpt/megatron_gpt_spin_model.py
index 7277becb6..4598eaaac 100644
--- a/nemo_aligner/models/nlp/gpt/megatron_gpt_spin_model.py
+++ b/nemo_aligner/models/nlp/gpt/megatron_gpt_spin_model.py
@@ -16,11 +16,11 @@
 from contextlib import nullcontext
 
 import torch
-from lightning.pytorch.trainer.trainer import Trainer
 from megatron.core.num_microbatches_calculator import get_num_microbatches
 from megatron.core.pipeline_parallel.schedules import get_forward_backward_func
 from omegaconf import OmegaConf
 from omegaconf.dictconfig import DictConfig
+from pytorch_lightning.trainer.trainer import Trainer
 
 from nemo.collections.nlp.models.language_modeling.megatron_gpt_model import MegatronGPTModel
 from nemo.collections.nlp.modules.common.megatron.utils import (
diff --git a/nemo_aligner/utils/train_script_utils.py b/nemo_aligner/utils/train_script_utils.py
index 0c36da0d7..eeed1a538 100644
--- a/nemo_aligner/utils/train_script_utils.py
+++ b/nemo_aligner/utils/train_script_utils.py
@@ -17,10 +17,10 @@
 from dataclasses import dataclass
 from functools import partial
 
-from lightning.pytorch.trainer import call
-from lightning.pytorch.trainer.states import TrainerFn
 from omegaconf import open_dict
 from omegaconf.omegaconf import OmegaConf
+from pytorch_lightning.trainer import call
+from pytorch_lightning.trainer.states import TrainerFn
 
 from nemo.collections.nlp.parts.megatron_trainer_builder import MegatronTrainerBuilder
 from nemo.collections.nlp.parts.peft_config import PEFT_CONFIG_MAP
diff --git a/setup/requirements.txt b/setup/requirements.txt
index b65bcd768..d074f3672 100644
--- a/setup/requirements.txt
+++ b/setup/requirements.txt
@@ -1,6 +1,5 @@
 Jinja2~=3.1.4
 jsonlines
-lightning>2.2.1
 megatron_core>=0.8
 nemo_toolkit[nlp]
 nvidia-pytriton


Test patch:
diff --git a/tests/conftest.py b/tests/conftest.py
index 9a70eba31..8ac1c2af7 100644
--- a/tests/conftest.py
+++ b/tests/conftest.py
@@ -15,8 +15,8 @@
 import os
 
 import pytest
-from lightning.pytorch import Trainer
 from omegaconf import DictConfig
+from pytorch_lightning import Trainer
 
 from nemo.collections.common.tokenizers.huggingface.auto_tokenizer import AutoTokenizer
 from nemo.collections.nlp.models.language_modeling.megatron_gpt_model import MegatronGPTModel