.nemo conversion bug fix (#10598)

* bug fix

Signed-off-by: dimapihtar <dpihtar@gmail.com>

* bug fix

Signed-off-by: dimapihtar <dpihtar@gmail.com>

* Apply isort and black reformatting

Signed-off-by: dimapihtar <dimapihtar@users.noreply.github.com>

---------

Signed-off-by: dimapihtar <dpihtar@gmail.com>
Signed-off-by: dimapihtar <dimapihtar@users.noreply.github.com>
Co-authored-by: dimapihtar <dimapihtar@users.noreply.github.com>

nemo/collections/nlp/parts/nlp_overrides.py

@@ -1068,16 +1068,19 @@ def dummy():
                 torch.distributed.barrier()
 
             # create nemo file from folder with all mp_ranks checkpoints
-            if (
-                app_state.pipeline_model_parallel_rank == 0
-                and app_state.tensor_model_parallel_rank == 0
-                and app_state.data_parallel_rank == 0
-            ):
-                with tempfile.TemporaryDirectory() as tmpdir:
+            if dist_ckpt:
+                should_move_data = is_global_rank_zero()
+            else:
+                should_move_data = (
+                    app_state.pipeline_model_parallel_rank == 0
+                    and app_state.tensor_model_parallel_rank == 0
+                    and app_state.data_parallel_rank == 0
+                )
 
+            if should_move_data:
+                with tempfile.TemporaryDirectory() as tmpdir:
                     if dist_ckpt:
                         shutil.move(str(dist_ckpt_dir), tmpdir)
-
                     elif app_state.pipeline_model_parallel_size == 1:
                         # move weights to the tmpdir
                         for tp_rank in range(app_state.tensor_model_parallel_size):
@@ -1123,6 +1126,9 @@ def dummy():
                         for file in os.listdir(tmpdir):
                             shutil.move(os.path.join(tmpdir, file), folder_path)
 
+            if torch.distributed.is_initialized():
+                torch.distributed.barrier()
+
         else:
             return super().save_to(model, save_path)