fix pipeline parallel dtype bug (#9637)

Signed-off-by: ashors1 <[email protected]>
NVIDIA · Jul 9, 2024 · 53cbc7d · 53cbc7d
1 parent 8898b76
commit 53cbc7d
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/nemo/lightning/_strategy_lib.py b/nemo/lightning/_strategy_lib.py
@@ -136,6 +136,7 @@ def set_model_parallel_attributes(model, parallelism):
         config.expert_model_parallel_size = parallelism.expert_model_parallel_size
         config.moe_extended_tp = parallelism.moe_extended_tp
         config.sequence_parallel = parallelism.sequence_parallel
+        config.pipeline_dtype = parallelism.pipeline_dtype
 
         return config