From eb83b8b58bfc126985ac80eb926e2fd918456146 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Stardust=C2=B7=E5=87=8F?= <star_dust_chen@foxmail.com>
Date: Thu, 7 Sep 2023 20:45:51 +0800
Subject: [PATCH] use gloo and update Dataloader

---
 train_ms.py | 9 ++++++---
 1 file changed, 6 insertions(+), 3 deletions(-)

diff --git a/train_ms.py b/train_ms.py
index 37f9087..c43441a 100644
--- a/train_ms.py
+++ b/train_ms.py
@@ -49,7 +49,7 @@ global_step = 0
 
 def run():
     dist.init_process_group(
-        backend="nccl", init_method="env://"
+        backend="gloo", init_method="env://"  # Due to some training problem,we proposed to use gloo instead of nccl.
     )  # Use torchrun instead of mp.spawn
     rank = dist.get_rank()
     n_gpus = dist.get_world_size()
@@ -71,16 +71,19 @@ def run():
         num_replicas=n_gpus,
         rank=rank,
         shuffle=True,
+        persistent_workers=True,
+        prefetch_factor=4
     )
     collate_fn = TextAudioSpeakerCollate()
     train_loader = DataLoader(
         train_dataset,
-        num_workers=0,
+        num_workers=16,
         shuffle=False,
         pin_memory=True,
         collate_fn=collate_fn,
         batch_sampler=train_sampler,
-    )  # 0 worker
+        
+    )  # DataLoader config could be adjusted.
     if rank == 0:
         eval_dataset = TextAudioSpeakerLoader(hps.data.validation_files, hps.data)
         eval_loader = DataLoader(