TTS: Training FastPitch on multi-GPU results in timeout error
State: closed | Pull request: no
Labels: name: bug

name: TTS

name: stale

**Describe the bug**

Training FastPitch on a 25k samples dataset using 8 A100 GPUs results in the following error:
```
Traceback (most recent call last):
  File "/workspace/scripts/train/fastpitch.py", line 21, in <module>
    run(main)
  File "/miniconda/lib/python3.9/site-packages/typer/main.py", line 1056, in run
    app()
  File "/miniconda/lib/python3.9/site-packages/typer/main.py", line 328, in __call__
    raise e
  File "/miniconda/lib/python3.9/site-packages/typer/main.py", line 311, in __call__
    return get_command(self)(*args, **kwargs)
  File "/miniconda/lib/python3.9/site-packages/click/core.py", line 1130, in __call__
    return self.main(*args, **kwargs)
  File "/miniconda/lib/python3.9/site-packages/typer/core.py", line 716, in main
    return _main(
  File "/miniconda/lib/python3.9/site-packages/typer/core.py", line 216, in _main
    rv = self.invoke(ctx)
  File "/miniconda/lib/python3.9/site-packages/click/core.py", line 1404, in invoke
    return ctx.invoke(self.callback, **ctx.params)
  File "/miniconda/lib/python3.9/site-packages/click/core.py", line 760, in invoke
    return __callback(*args, **kwargs)
  File "/miniconda/lib/python3.9/site-packages/typer/main.py", line 683, in wrapper
    return callback(**use_params)  # type: ignore
  File "/workspace/scripts/train/fastpitch.py", line 17, in main
    train_fastpitch(gpu_ids=gpu_ids, full_config_path=full_config_path)
  File "/workspace/src/tts/train/fastpitch.py", line 66, in train_fastpitch
    trainer.fit(model)
  File "/miniconda/lib/python3.9/site-packages/pytorch_lightning/trainer/trainer.py", line 608, in fit
    call._call_and_handle_interrupt(
  File "/miniconda/lib/python3.9/site-packages/pytorch_lightning/trainer/call.py", line 36, in _call_and_handle_interrupt
    return trainer.strategy.launcher.launch(trainer_fn, *args, trainer=trainer, **kwargs)
  File "/miniconda/lib/python3.9/site-packages/pytorch_lightning/strategies/launchers/subprocess_script.py", line 88, in launch
    return function(*args, **kwargs)
  File "/miniconda/lib/python3.9/site-packages/pytorch_lightning/trainer/trainer.py", line 650, in _fit_impl
    self._run(model, ckpt_path=self.ckpt_path)
  File "/miniconda/lib/python3.9/site-packages/pytorch_lightning/trainer/trainer.py", line 1048, in _run
    self.strategy.setup_environment()
  File "/miniconda/lib/python3.9/site-packages/pytorch_lightning/strategies/ddp.py", line 152, in setup_environment
    self.setup_distributed()
  File "/miniconda/lib/python3.9/site-packages/pytorch_lightning/strategies/ddp.py", line 203, in setup_distributed
    _init_dist_connection(self.cluster_environment, self._process_group_backend, timeout=self._timeout)
  File "/miniconda/lib/python3.9/site-packages/lightning_fabric/utilities/distributed.py", line 245, in _init_dist_connection
    torch.distributed.init_process_group(torch_distributed_backend, rank=global_rank, world_size=world_size, **kwargs)
  File "/miniconda/lib/python3.9/site-packages/torch/distributed/distributed_c10d.py", line 932, in init_process_group
    _store_based_barrier(rank, store, timeout)
  File "/miniconda/lib/python3.9/site-packages/torch/distributed/distributed_c10d.py", line 469, in _store_based_barrier
    raise RuntimeError(
RuntimeError: Timed out initializing process group in store based barrier on rank: 0, for key: store_based_barrier_key:1 (world_size=8, worker_count=1, timeout=0:30:00)
```

Based on my understanding, this is caused by the fact that PyTorch Lightning's DDP invokes the training script N times, where N is the number of GPUs used for training (see [documentation](https://lightning.ai/docs/pytorch/stable/accelerators/gpu_intermediate.html#distributed-data-parallel)). As a consequence, because the dataset preparation logic (loading/normalization/tokenization) is contained within the training script, it is repeated N times. In particular, what I observed is the following (cf. my comment on a related issue, namely #4418):

1. Initially, 1 of the 8 processes begins the dataset preparation
2. Once it's done, all the other processes are unleashed. Said processes perform the dataset preparation again from scratch (on the _whole_ dataset, not a subset) and do so _concurrently_. You can tell this by the fact that their logs get mixed up.
3. Because the dataset is very large, at some point each of these processes hits a timeout (30 minutes) and fails

**Steps/Code to reproduce bug**
Find a dataset with at least 25k samples, then run the following command on a machine with 8 A100 GPUs:

```python
python fastpitch.py \
    --config-path . \
    --config-name fastpitch_align_ipa.yaml \
    train_dataset=/path/to/train_manifest.json \
    validation_datasets=/path/to/val_manifest.json \
    sup_data_path=/path/to/sup_data \
    exp_manager.exp_dir=path/to/exp_dir \
    model.pitch_mean=<some_appropriate_value> \
    model.pitch_std=<some_appropriate_value> \
    trainer.devices=8 \
    trainer.strategy=ddp
```

Where the config file looks like the following: 
[fastpitch_align_ipa.yaml](https://github.com/NVIDIA/NeMo/files/11581619/fastpitch_align_ipa.yaml.txt)


**Expected behavior**

1. I can train FastPitch on multi-GPU without incurring in any error
2. Operations that aren't meant to be repeated, like dataset preparation, are executed only once. Alternatively, each process performs the operation on a different subset of the dataset

**Environment overview (please complete the following information)**

 - Environment location: Docker
 - Method of NeMo install: Docker
   - `docker pull nvcr.io/nvidia/nemo:23.02`
   - `docker run --rm -it --gpus all --ipc=host --ulimit memlock=-1 --ulimit stack=67108864 nvcr.io/nvidia/nemo:23.02`