From 34bfe1b8436a3834af3257650f6bbdb45f58274a Mon Sep 17 00:00:00 2001 From: Anna Shors <71393111+ashors1@users.noreply.github.com> Date: Mon, 15 Jul 2024 19:15:27 -0700 Subject: [PATCH] [NeMo-UX] Minor bug fix when TE/Apex not installed (#9749) * minor 2.0 bug fix when TE/Apex not installed Signed-off-by: ashors1 * Apply isort and black reformatting Signed-off-by: ashors1 --------- Signed-off-by: ashors1 Signed-off-by: ashors1 Co-authored-by: ashors1 --- nemo/lightning/megatron_parallel.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/nemo/lightning/megatron_parallel.py b/nemo/lightning/megatron_parallel.py index 0c141e82ab52..74db25af64bf 100644 --- a/nemo/lightning/megatron_parallel.py +++ b/nemo/lightning/megatron_parallel.py @@ -445,8 +445,8 @@ def infer_num_microbatches(self, data: Union[DataT, Iterator[DataT], List[Iterat raise ValueError("Cannot infer `num_microbatches` from data, please specify it manually") def init_model_parallel(self): - from apex.transformer.tensor_parallel.layers import set_defaults_if_not_set_tensor_model_parallel_attributes from megatron.core import parallel_state + from megatron.core.tensor_parallel.layers import set_defaults_if_not_set_tensor_model_parallel_attributes for model_module in self: if not self._cpu: