From dba07e656ec4829afe6d47cab92afa8889a6f0b7 Mon Sep 17 00:00:00 2001 From: Maanu Grover <109391026+maanug-nv@users.noreply.github.com> Date: Fri, 22 Nov 2024 10:21:26 -0800 Subject: [PATCH] fix: Update PTL import (#410) Signed-off-by: Maanu Grover Signed-off-by: NeMo-Aligner CI Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: Terry Kong --- Dockerfile | 2 ++ examples/nlp/gpt/serve_reward_model.py | 2 +- nemo_aligner/models/nlp/gpt/gpt_sft_model.py | 2 +- nemo_aligner/models/nlp/gpt/megatron_gpt_critic.py | 2 +- nemo_aligner/models/nlp/gpt/megatron_gpt_dpo_model.py | 2 +- nemo_aligner/models/nlp/gpt/megatron_gpt_kto_model.py | 2 +- nemo_aligner/models/nlp/gpt/megatron_gpt_ppo_actor.py | 2 +- .../models/nlp/gpt/megatron_gpt_regression_reward_model.py | 2 +- nemo_aligner/models/nlp/gpt/megatron_gpt_reward_model.py | 2 +- nemo_aligner/models/nlp/gpt/megatron_gpt_rs_actor.py | 2 +- nemo_aligner/models/nlp/gpt/megatron_gpt_spin_model.py | 2 +- nemo_aligner/utils/train_script_utils.py | 4 ++-- setup/requirements.txt | 1 + tests/conftest.py | 2 +- 14 files changed, 16 insertions(+), 13 deletions(-) diff --git a/Dockerfile b/Dockerfile index 44a9f8651..4d22c6a21 100644 --- a/Dockerfile +++ b/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 . diff --git a/examples/nlp/gpt/serve_reward_model.py b/examples/nlp/gpt/serve_reward_model.py index a425f26e7..5f5dfc609 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 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 diff --git a/nemo_aligner/models/nlp/gpt/gpt_sft_model.py b/nemo_aligner/models/nlp/gpt/gpt_sft_model.py index d3a615500..e8a47763f 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 a61cd2b8d..a800a6950 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 952b4e897..715914be7 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 0be08e808..5a16d18b8 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 275e02e82..43643faab 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 eb9333c12..06b719610 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 596436c67..e1d2fcc24 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 49dd3846c..ebd5adf4b 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 4598eaaac..7277becb6 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 c6f6f8089..bce7bff5a 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 198d2e07a..3ce01d911 100644 --- a/setup/requirements.txt +++ b/setup/requirements.txt @@ -1,4 +1,5 @@ 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 8ac1c2af7..9a70eba31 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