feat(train): make timeout of wenet_join configurable (#2123)

wenet/bin/train.py

@@ -126,7 +126,7 @@ def main():
 
         # NOTE(xcsong): Why we need a new group? see `train_utils.py::wenet_join`
         group_join = dist.new_group(backend="gloo",
-                                    timeout=datetime.timedelta(seconds=30))
+                                    timeout=datetime.timedelta(seconds=args.timeout))
 
         dist.barrier()  # NOTE(xcsong): Ensure all ranks start Train at the same time.
         executor.train(model, optimizer, scheduler, train_data_loader,

wenet/utils/train_utils.py

@@ -131,6 +131,9 @@ def add_ddp_args(parser):
 
 
 def add_deepspeed_args(parser):
+    parser.add_argument('--timeout', default=30, type=int,
+                        help='timeout (in seconds) of wenet_join. ' +
+                             '30s for aishell & 300s for wenetspeech')
     parser.add_argument('--local_rank', type=int, default=-1,
                         help='local rank passed from distributed launcher')
     parser.add_argument('--deepspeed.save_states',