DALI does not work on multi-gpu systems
State: closed | Pull request: no
Labels: name: bug

**Describe the bug**

When setting `use_dali: true` in the dataset config, the following error occurs:

<details>
<summary>Error</summary>
<br>

  ```sh
  Traceback (most recent call last):                                                                                                                                                                                        
    File "/audio-data/tao/nemo_stt_ft.py", line 67, in main                                                                                                                                                                 
      trainer.fit(asr_model)                                                                                                                                                                                                
    File "/opt/conda/lib/python3.8/site-packages/pytorch_lightning/trainer/trainer.py", line 552, in fit                                                                                                                    
      self._run(model)                                                                                                                                                                                                      
    File "/opt/conda/lib/python3.8/site-packages/pytorch_lightning/trainer/trainer.py", line 922, in _run                                                                                                                   
      self._dispatch()                                                                                                                                                                                                      
    File "/opt/conda/lib/python3.8/site-packages/pytorch_lightning/trainer/trainer.py", line 990, in _dispatch                                                                                                              
      self.accelerator.start_training(self)                                                                                                                                                                                 
    File "/opt/conda/lib/python3.8/site-packages/pytorch_lightning/accelerators/accelerator.py", line 92, in start_training                                                                                                 
      self.training_type_plugin.start_training(trainer)                                                                                                                                                                     
    File "/opt/conda/lib/python3.8/site-packages/pytorch_lightning/plugins/training_type/training_type_plugin.py", line 161, in start_training                                                                              
      self._results = trainer.run_stage()                                                                                                                                                                                   
    File "/opt/conda/lib/python3.8/site-packages/pytorch_lightning/trainer/trainer.py", line 1000, in run_stage                                                                                                             
      return self._run_train()                                                                                                                                                                                              
    File "/opt/conda/lib/python3.8/site-packages/pytorch_lightning/trainer/trainer.py", line 1035, in _run_train                                                                                                            
      self._run_sanity_check(self.lightning_module)                                                                                                                                                                         
    File "/opt/conda/lib/python3.8/site-packages/pytorch_lightning/trainer/trainer.py", line 1122, in _run_sanity_check     
      self._evaluation_loop.run()                                                                                                                                                                                  [81/1557]
    File "/opt/conda/lib/python3.8/site-packages/pytorch_lightning/loops/base.py", line 111, in run                                                                                                                         
      self.advance(*args, **kwargs)                                                                                                                                                                                         
    File "/opt/conda/lib/python3.8/site-packages/pytorch_lightning/loops/dataloader/evaluation_loop.py", line 110, in advance                                                                                               
      dl_outputs = self.epoch_loop.run(                                                                                                                                                                                     
    File "/opt/conda/lib/python3.8/site-packages/pytorch_lightning/loops/base.py", line 111, in run                                                                                                                         
      self.advance(*args, **kwargs)                                                                                                                                                                                         
    File "/opt/conda/lib/python3.8/site-packages/pytorch_lightning/loops/epoch/evaluation_epoch_loop.py", line 111, in advance                                                                                              
      output = self.evaluation_step(batch, batch_idx, dataloader_idx)                                                                                                                                                       
    File "/opt/conda/lib/python3.8/site-packages/pytorch_lightning/loops/epoch/evaluation_epoch_loop.py", line 158, in evaluation_step                                                                                      
      output = self.trainer.accelerator.validation_step(step_kwargs)                                                                                                                                                        
    File "/opt/conda/lib/python3.8/site-packages/pytorch_lightning/accelerators/accelerator.py", line 211, in validation_step                                                                                               
      return self.training_type_plugin.validation_step(*step_kwargs.values())                                                                                                                                               
    File "/opt/conda/lib/python3.8/site-packages/pytorch_lightning/plugins/training_type/ddp.py", line 392, in validation_step                                                                                              
      return self.model(*args, **kwargs)                                                                                                                                                                                    
    File "/opt/conda/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1051, in _call_impl                                                                                                                      
      return forward_call(*input, **kwargs)                                                                                                                                                                                 
    File "/opt/conda/lib/python3.8/site-packages/torch/nn/parallel/distributed.py", line 799, in forward                                                                                                                    
      output = self.module(*inputs[0], **kwargs[0])                                                                                                                                                                         
    File "/opt/conda/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1051, in _call_impl                                                                                                                      
      return forward_call(*input, **kwargs)                                                                                                                                                                                 
    File "/opt/conda/lib/python3.8/site-packages/pytorch_lightning/overrides/base.py", line 93, in forward                                                                                                                  
      output = self.module.validation_step(*inputs, **kwargs)                                                                                                                                                               
    File "/opt/conda/lib/python3.8/site-packages/nemo/collections/asr/models/ctc_models.py", line 588, in validation_step                                                                                                   
      log_probs, encoded_len, predictions = self.forward(                                                                                                                                                                   
    File "/opt/conda/lib/python3.8/site-packages/nemo/core/classes/common.py", line 798, in __call__                                                                                                                        
      outputs = wrapped(*args, **kwargs)                                                                                                                                                                                    
    File "/opt/conda/lib/python3.8/site-packages/nemo/collections/asr/models/ctc_models.py", line 545, in forward                                                                                                           
      encoded, encoded_len = self.encoder(audio_signal=processed_signal, length=processed_signal_length)                                                                                                                    
    File "/opt/conda/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1051, in _call_impl                                                                                                                      
      return forward_call(*input, **kwargs)                                                                                                                                                                                 
    File "/opt/conda/lib/python3.8/site-packages/nemo/core/classes/common.py", line 798, in __call__                                                                                                                        
      outputs = wrapped(*args, **kwargs)                                                                                                                                                                                    
    File "/opt/conda/lib/python3.8/site-packages/nemo/collections/asr/modules/conv_asr.py", line 222, in forward                                                                                                            
      s_input, length = self.encoder(([audio_signal], length))                                                                                                                                                              
    File "/opt/conda/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1051, in _call_impl                                                                                                                      
      return forward_call(*input, **kwargs)                                                                                                                                                                                 
    File "/opt/conda/lib/python3.8/site-packages/torch/nn/modules/container.py", line 139, in forward                                                                                                                       
      input = module(input)                                                                                                                                                                                                 
    File "/opt/conda/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1051, in _call_impl                                                                                                                      
      return forward_call(*input, **kwargs)                                                                                                                                                                                 
    File "/opt/conda/lib/python3.8/site-packages/nemo/collections/asr/parts/submodules/jasper.py", line 906, in forward                                                                                                     
      out, lens = l(out, lens)                                                                                                                                                                                              
    File "/opt/conda/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1051, in _call_impl                                                                                                                      
      return forward_call(*input, **kwargs)                                                                                                                                                                                 
    File "/opt/conda/lib/python3.8/site-packages/nemo/collections/asr/parts/submodules/jasper.py", line 298, in forward                                                                                                     
      out = self.conv(x)                                                                                                                                                                                                    
    File "/opt/conda/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1051, in _call_impl                                                                                                                      
      return forward_call(*input, **kwargs)                                                                                                                                                                                 
    File "/opt/conda/lib/python3.8/site-packages/torch/nn/modules/conv.py", line 298, in forward                                                                                                                            
      return self._conv_forward(input, self.weight, self.bias)                                                 
    File "/opt/conda/lib/python3.8/site-packages/torch/nn/modules/conv.py", line 294, in _conv_forward                                                                                                                      
      return F.conv1d(input, weight, bias, self.stride,                                                        
  RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cuda:6! (when checking arugment for argument weight in method wrapper_thnn_conv_depthwise2d_forward) 
  ```

</details>

The error repeats multiple times based on the number of GPUs

**Steps/Code to reproduce bug**

Run any script with `trainer.fit()` and make sure to set `use_dali: true` in any of your `train/validation/test_ds` config.

**Expected behavior**

Training should start as expected.

**Environment overview (please complete the following information)**

 - Environment location: Docker
 - Method of NeMo install: `pip install nemo_toolkit[asr]==1.3.0`

**Environment details**

Docker:
- OS version Ubuntu 20.04
- PyTorch version 1.9.0
- Python version 3.8
- Cuda 11.2, cuDNN 8

**Additional context**

PyTorch compiled from source, but happens with pip installed PyTorch too.