Merge branch 'r2.0.0rc1' into r2.0.0rc1

rohitrango · Jul 9, 2024 · 25ce6f2 · 25ce6f2
2 parents 0f97b47 + dffbaf4
commit 25ce6f2
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