LSTM forward error when running AMP and dataParallel

Hi, I met the following error for LSTM forward when running AMP and dataParallel. The same code works for dataParallel without using AMP. The opt_level is 'O1'.

I used 2 GPUs and I checked the cuda device for the tensors in new_synthesize_flattened_rnn_weights. It turns out most tensors are from CUDA0 while only a few are from CUDA1. I am not sure if this device mismatch is the cause (as they should be balanced). If so, how can I resolve this?

Any help would be appreciated. Thanks!

  File "/home/anaconda3/envs/LSP/lib/python3.6/site-packages/torch/nn/modules/module.py", line 532, in __call__
    result = self.forward(*input, **kwargs)
  File "/home/anaconda3/envs/LSP/lib/python3.6/site-packages/torch/nn/parallel/data_parallel.py", line 152, in forward
    outputs = self.parallel_apply(replicas, inputs, kwargs)
  File "/home/naconda3/envs/LSP/lib/python3.6/site-packages/torch/nn/parallel/data_parallel.py", line 162, in parallel_apply
    return parallel_apply(replicas, inputs, kwargs, self.device_ids[:len(replicas)])
  File "/home/anaconda3/envs/LSP/lib/python3.6/site-packages/torch/nn/parallel/parallel_apply.py", line 89, in parallel_apply
    output.reraise()
  File "/home/anaconda3/envs/LSP/lib/python3.6/site-packages/torch/_utils.py", line 394, in reraise
    raise self.exc_type(msg)
RuntimeError: Caught RuntimeError in replica 1 on device 1.
Original Traceback (most recent call last):
  File "/home/anaconda3/envs/LSP/lib/python3.6/site-packages/torch/nn/parallel/parallel_apply.py", line 64, in _worker
    output = module(*input, **kwargs)
  File "/home/anaconda3/envs/LSP/lib/python3.6/site-packages/torch/nn/modules/module.py", line 532, in __call__
    result = self.forward(*input, **kwargs)
  File "/mnt/data2/esearch/medical/report-generation/pytorch-generation/utils/models.py", line 786, in forward
    sentence_states)
  File "/mnt/data2/research/medical/report-generation/pytorch-generation/utils/models.py", line 263, in forward
    return self._forward(ctx, prev_hidden_state, states)
  File "/mnt/data2/research/medical/report-generation/pytorch-generation/utils/models.py", line 269, in _forward
    hidden_state, states = self.lstm(ctx, states)
  File "/home/anaconda3/envs/LSP/lib/python3.6/site-packages/torch/nn/modules/module.py", line 532, in __call__
    result = self.forward(*input, **kwargs)
  File "/home/aconda3/envs/LSP/lib/python3.6/site-packages/torch/nn/modules/rnn.py", line 559, in forward
    self.dropout, self.training, self.bidirectional, self.batch_first)
  File "/home/aconda3/envs/LSP/lib/python3.6/site-packages/apex/amp/wrap.py", line 257, in wrapper
    arg, fp16_weight_buf, fn, verbose)
  File "/home/anaconda3/envs/LSP/lib/python3.6/site-packages/apex/amp/utils.py", line 212, in new_synthesize_flattened_rnn_weights
    w_fp32.shape)
RuntimeError: Tensor: invalid storage offset at /opt/conda/conda-bld/pytorch_1579022034529/work/aten/src/THC/THCTensor.cpp:167

Start by examining apex/amp/utils.py at new_synthesize_flattened_rnn_weights and the PyTorch torch/nn/modules/rnn.py path shown in the traceback. Reproduce the LSTM forward case with AMP and DataParallel, then determine whether the reported device placement leads to the invalid storage offset; done means a confirmed cause and a documented fix or regression coverage.