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 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