fix: Update PTL import (#410)

Signed-off-by: Maanu Grover <[email protected]>
Signed-off-by: NeMo-Aligner CI <[email protected]>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Terry Kong <[email protected]>

Dockerfile

@@ -115,6 +115,8 @@ 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 .

examples/nlp/gpt/serve_reward_model.py

@@ -13,7 +13,7 @@
 # limitations under the License.
 
 import torch
-from pytorch_lightning.trainer.trainer import Trainer
+from lightning.pytorch.trainer.trainer import Trainer
 
 from nemo.collections.nlp.parts.nlp_overrides import NLPDDPStrategy
 from nemo.core.config import hydra_runner

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

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,

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 (

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 (

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 (

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,

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 (

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 (

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 (

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

setup/requirements.txt

@@ -1,4 +1,5 @@
 jsonlines
+lightning>2.2.1
 megatron_core>=0.8
 nemo_toolkit[nlp]
 nvidia-pytriton

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