format

Signed-off-by: gdeng <[email protected]>
NVIDIA · Jan 17, 2025 · 18d9b0e · 18d9b0e
1 parent 6537c10
commit 18d9b0e
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 2 deletions.
diff --git a/nemo/lightning/megatron_init.py b/nemo/lightning/megatron_init.py
@@ -127,7 +127,7 @@ def initialize_model_parallel_for_nemo(
     app_state.encoder_pipeline_model_parallel_size = encoder_pipeline_model_parallel_size
     app_state.use_fp8 = use_fp8
     app_state.init_mpi_proc_group = init_mpi_proc_group
-    app_state.expert_tensor_parallel_size=expert_tensor_parallel_size
+    app_state.expert_tensor_parallel_size = expert_tensor_parallel_size
     (
         app_state.tensor_model_parallel_rank,
         app_state.pipeline_model_parallel_rank,

diff --git a/nemo/lightning/pytorch/strategies/megatron_strategy.py b/nemo/lightning/pytorch/strategies/megatron_strategy.py
@@ -100,7 +100,7 @@ class ParallelismConfig:
     encoder_tensor_model_parallel_size: int = 0
     encoder_pipeline_model_parallel_size: int = 0
     use_te_rng_tracker: bool = False
-    expert_tensor_parallel_size: Optional[int] = None 
+    expert_tensor_parallel_size: Optional[int] = None
 
 
 class MegatronStrategy(DDPStrategy, io.IOMixin):