fix: resolve more pytorch lightning import bugs (kd/reinforce) (#442)

Signed-off-by: Terry Kong <[email protected]>
NVIDIA · Dec 10, 2024 · 2bfb3a5 · 2bfb3a5
1 parent 3791aad
commit 2bfb3a5
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 2 deletions.
diff --git a/nemo_aligner/models/nlp/gpt/megatron_gpt_knowledge_distillation.py b/nemo_aligner/models/nlp/gpt/megatron_gpt_knowledge_distillation.py
@@ -16,12 +16,12 @@
 
 import hydra
 import torch
+from lightning.pytorch.trainer.trainer import Trainer
 from megatron.core import parallel_state, tensor_parallel
 from megatron.core.num_microbatches_calculator import get_micro_batch_size, get_num_microbatches
 from megatron.core.parallel_state import get_tensor_model_parallel_group
 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 (

diff --git a/nemo_aligner/models/nlp/gpt/megatron_gpt_reinforce_actor.py b/nemo_aligner/models/nlp/gpt/megatron_gpt_reinforce_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 (