Multi-GPU training capability for the Pytorch Transformer LM training script - https://github.com/kaldi-asr/kaldi/blob/master/egs/wsj/s5/local/pytorchnn/run_nnlm.sh

[saikirandingu]

I used the script ### https://github.com/kaldi-asr/kaldi/blob/master/egs/wsj/s5/local/pytorchnn/run_nnlm.sh, but I could not figure out how we could distribute the training of Transformer based LM on multiple GPUs in order to speed-up the Pytorch training. Please suggest if there is any way to do so.

Thanks!

[danpovey]

It would be necessary to add some wrapping code in that script, and some extra options, to enable DDP
training.
I don't think I have time to do this personally right now.

[danpovey]

If anyones has time to do this I could point you in the right direction.
E.g. look for DDP usage in icefall scripts
https://github.com/k2-fsa/icefall/blob/master/egs/librispeech/ASR/transducer_stateless/train.py
(this is rando, it's just something I'm familiar with). It's basically a few wrappers
including options like --world-size, --rank, the "run" wrapper, setup_dist(). fix_random_seed().
(I'm just pointing out some DDP-related aspects of the script). Also the dataloader has to know about DDP,
i.e. in pytorchnn/train.py, get_batch() needs to figure out the rank and world_size (I think the DDP module has
ways to do this) in order to return disjoint subsets.

In the icefall example you activate DDP by using the --world-size option.

[desh2608]

Related issue: #4468