Gradients of some part (RNN in my case) of the model become zero after the model is wrapped with nn.DataParallel

username: grtzsohalf
cooked: <p>When I use a single GPU, all parameters in the model get non-zero gradients;<br>
but after wrapping my model with nn.DataParallel, the gradients of some part (RNN in my case) of the model become zero.</p>
<h2>To Reproduce</h2>
<p>Steps to reproduce the behavior:</p>
<p>test.py:</p>
<pre><code class="lang-python">import sys

import torch
from torch import nn
import torch.nn.functional as F
from torchexp.stat import RunningAvgDict
from torch.nn.utils.rnn import pack_padded_sequence
from torch.nn.utils.rnn import pad_packed_sequence

import numpy as np
import math


gpus = list(map(int, sys.argv[1].split(',')))


def to_device(m, x):
    """Send tensor into the device of the module.

    Args:
        m (torch.nn.Module): Torch module.
        x (Tensor): Torch tensor.

    Returns:
        Tensor: Torch tensor located in the same place as torch module.

    """
    assert isinstance(m, torch.nn.Module)
    device = next(m.parameters()).device
    return x.to(device)


def reset_backward_rnn_state(states):
    """Sets backward BRNN states to zeroes - useful in processing of sliding windows over the inputs"""
    if isinstance(states, (list, tuple)):
        for state in states:
            state[1::2] = 0.
    else:
        states[1::2] = 0.
    return states


class RNNP(nn.Module):
    """RNN with projection layer module

    :param int idim: dimension of inputs, default: odim of CNN
    :param int nlayers: Number of layers
    :param int enc_dim: hidden dim of BLSTMP
    :param int proj_dim: proj. dim of BLSTMP
    :param int odim: ouput of encoder dimension (default: proj_dim)
    """

    def __init__(self, idim, nlayers, enc_dim, proj_dim, odim):
        super(RNNP, self).__init__()
        self.enc_dim = enc_dim
        self.proj_dim = proj_dim
        self.odim = odim
        self.nlayers = nlayers

        self.rnn0 = nn.LSTM(idim, enc_dim, num_layers=1, bidirectional=True, batch_first=True)
        self.bt0 = nn.Linear(2 * enc_dim, proj_dim)


        for i in range(1,self.nlayers):
            rnn = nn.LSTM(proj_dim, enc_dim, num_layers=1, bidirectional=True, batch_first=True)
            setattr(self, f"rnn{i}", rnn)

            if i == self.nlayers - 1:
                setattr(self, f"bt{i}", nn.Linear(2 * enc_dim, odim))
            else:
                setattr(self, f"bt{i}", nn.Linear(2 * enc_dim, proj_dim))


    def forward(self, xs_pad, enc_lens, prev_state=None):
        """RNNP forward

        :param torch.Tensor xs_pad: batch of padded input sequences (B, Tmax, cnn_o_dim)
        :param list of int enc_lens
        :param torch.Tensor prev_state: batch of previous RNN states
        :return: batch of hidden state sequences (B, Tmax, odim)
        :rtype: torch.Tensor
        """
        hid_states = []

        for i in range(self.nlayers):
            total_length = xs_pad.size(1)
            xs_pack = pack_padded_sequence(xs_pad, enc_lens, batch_first = True)
            rnn = getattr(self, f"rnn{i}")
            rnn.flatten_parameters()
            if prev_state is not None:
                prev_state = reset_backward_rnn_state(prev_state)
            ys, states = rnn(xs_pack, hx=None if prev_state is None else prev_state[i])
            hid_states.append(states)

            ys_pad, enc_lens= pad_packed_sequence(ys, batch_first=True, total_length=total_length)
            # ys_pad: (B, T, enc_dim)

            projected = getattr(self, f"bt{i}")(ys_pad.contiguous().view(-1, ys_pad.size(2))) #(B*T, proj_dim)
            xs_pad = torch.tanh(projected.view(ys_pad.size(0), ys_pad.size(1), -1))

        return xs_pad, to_device(self,enc_lens), hid_states


class BlstmEncoder(nn.Module):
    """
    :param int idim: dim of input
    :param int enc_dim: hidden dim for BLSTM
    :param int proj_dim: projection dim for BLSTMP
    :param int odim: encoder output dimension (usually set as proj_dim)
    """

    def __init__(self, idim, enc_dim, proj_dim, odim):
        super(BlstmEncoder, self).__init__()
        self.idim = idim
        self.enc_dim = enc_dim
        self.proj_dim = proj_dim
        self.odim = odim

        self.cnn_model = nn.Sequential(
            nn.Conv2d(1, 128, kernel_size=3, stride=1, padding=1),
            nn.ReLU(),
            nn.Conv2d(128, 128, kernel_size=3, stride=1, padding=1),
            nn.ReLU(),
            nn.MaxPool2d(2, stride=2, ceil_mode=True),
        )
        cnn_o_dim = np.ceil(np.array(idim, dtype=np.float32) / 2)
        cnn_o_dim = int(cnn_o_dim) * 128

        self.nlayers = 1
        self.blstm = RNNP(cnn_o_dim, self.nlayers, enc_dim, proj_dim, odim)


    def forward(self, xs_pad, ilens, prev_state=None):
        xs_pad = xs_pad.view(xs_pad.size(0), xs_pad.size(1), 1, xs_pad.size(2)).transpose(1,2)
        xs_pad = self.cnn_model(xs_pad)

        if torch.is_tensor(ilens):
            ilens = ilens.cpu().numpy()
        else:
            ilens = np.array(ilens, dtype=np.float32)

        enc_lens = np.array(np.ceil(ilens/2), dtype=np.int64)
        enc_lens = np.array(np.ceil(np.array(enc_lens, dtype=np.float32)/2), dtype=np.int64).tolist()

        xs_pad = xs_pad.transpose(1,2)
        xs_pad = xs_pad.contiguous().view(xs_pad.size(0), xs_pad.size(1), xs_pad.size(2) * xs_pad.size(3))

        out, enc_lens, cur_state = self.blstm(xs_pad, enc_lens, prev_state)

        return out, enc_lens, cur_state


class MultiHeadBLSTM(nn.Module):

    def __init__(self):
        super(MultiHeadBLSTM, self).__init__()

        self.idim = 20
        self.odim = 10

        enc_o_dim = 128

        self.encoder = BlstmEncoder(self.idim, enc_o_dim, enc_o_dim, enc_o_dim)
        self.head = nn.Linear(enc_o_dim, self.odim)


    @property
    def device(self):
        return next(self.parameters()).device


    def forward(self, xs_pad, ilens):
        assert xs_pad.size(0) == ilens.size(0), "Batch size mismatch"
        batch_size = xs_pad.size(0)

        # Put data to device
        xs_pad = to_device(self,xs_pad)
        ilens = to_device(self,ilens)

        enc, enc_lens, _ = self.encoder(xs_pad, ilens)
        out = self.head(enc)

        return out, enc_lens


def run_batch(asr_model, ctc_loss, x, ilens, y_true, olens):

    pred, enc_lens = asr_model(x, ilens)
    if len(gpus) == 1:
        olens = to_device(asr_model, olens)
    else:
        olens = to_device(asr_model.module, olens)
    pred = F.log_softmax(pred, dim=-1)

    loss = ctc_loss(pred.transpose(0,1).contiguous(),
                         y_true.cuda().to(dtype=torch.long),
                         enc_lens.cpu().to(dtype=torch.long),
                         olens.cpu().to(dtype=torch.long))
    print('ctc loss: ', loss)

    info = { 'loss': loss.item() }
    loss.backward()

    return info


def run_task(asr_model, ctc_loss, x, ilens, ys, olens):
    batch_size = len(ys)
    info = run_batch(asr_model, ctc_loss, x, ilens, ys, olens)

    # torch.set_printoptions(profile="full")
    if len(gpus) == 1:
        count = 0
        for n, p in asr_model.named_parameters():
            print('grad of param {}, {}: '.format(count, n), p.grad)#[0, 0, 0])
            count += 1
        print('num of params with grad: ', count)
    else:
        count = 0
        for n, p in asr_model.module.named_parameters():
            print('grad of param {}, {}: '.format(count, n), p.grad)#[0, 0, 0])
            count += 1
        print('num of params with grad: ', count)


if __name__ == '__main__':
    asr_model = MultiHeadBLSTM().cuda()
    if len(gpus) &gt; 1:
        asr_model = nn.DataParallel(asr_model, device_ids=gpus)

    ctc_loss = nn.CTCLoss(blank=0, reduction='mean', zero_infinity=True)

    x = torch.rand(4, 20, 20)#.cuda() 
    ilens = torch.tensor([20, 20, 20, 20])#.cuda()
    ys = [torch.tensor([1, 1, 1]),
          torch.tensor([1, 1, 1, 1]),
          torch.tensor([1, 1, 1, 1, 1]),
          torch.tensor([1, 1, 1, 1, 1, 1])]
    olens = torch.tensor([3, 4, 5, 6])#.cuda()

    y_true = torch.cat(ys)#.cuda()

    run_task(asr_model, ctc_loss, x, ilens, y_true, olens)
</code></pre>
<p>If I run <code>python3 test.py 0</code> (gpu devices = [0]), everything is fine;<br>
however, if I run <code>python3 test.py 0,1</code> (gpu devices = [0, 1]), rnn receives no gradient.</p>
<h2>Expected behavior</h2>
<p>All parameters of both single- and multi-gpu models should have non-zero gradients.</p>
<h2>Environment</h2>
<pre><code class="lang-auto">PyTorch version: 1.4.1
Is debug build: No
CUDA used to build PyTorch: 10.2

OS: Arch Linux
GCC version: (Arch Linux 9.3.0-1) 9.3.0
CMake version: version 3.17.0

Python version: 3.8
Is CUDA available: Yes
CUDA runtime version: 10.2.89
GPU models and configuration: 
GPU 0: GeForce GTX TITAN X
GPU 1: GeForce GTX 1060 6GB

Nvidia driver version: 440.64
cuDNN version: /usr/lib/libcudnn.so.7.6.5

Versions of relevant libraries:
[pip3] numpy==1.18.2
[pip3] torch==1.4.1
[pip3] torchexp==0.1.0
[pip3] torchvision==0.5.0a0+c81ac87
[conda] Could not collect
</code></pre>

username: grtzsohalf
cooked: <p>There is a bug in version 1.4<br>
<aside class="onebox githubpullrequest">
  <header class="source">
      <a href="https://github.com/pytorch/pytorch/pull/33907" target="_blank" rel="nofollow noopener">github.com/pytorch/pytorch</a>
  </header>
  <article class="onebox-body">
    <div class="github-row">
  <div class="github-icon-container" title="Pull Request">
    <svg width="60" height="60" class="github-icon" viewbox="0 0 12 16" aria-hidden="true"><path d="M11 11.28V5c-.03-.78-.34-1.47-.94-2.06C9.46 2.35 8.78 2.03 8 2H7V0L4 3l3 3V4h1c.27.02.48.11.69.31.21.2.3.42.31.69v6.28A1.993 1.993 0 0 0 10 15a1.993 1.993 0 0 0 1-3.72zm-1 2.92c-.66 0-1.2-.55-1.2-1.2 0-.65.55-1.2 1.2-1.2.65 0 1.2.55 1.2 1.2 0 .65-.55 1.2-1.2 1.2zM4 3c0-1.11-.89-2-2-2a1.993 1.993 0 0 0-1 3.72v6.56A1.993 1.993 0 0 0 2 15a1.993 1.993 0 0 0 1-3.72V4.72c.59-.34 1-.98 1-1.72zm-.8 10c0 .66-.55 1.2-1.2 1.2-.65 0-1.2-.55-1.2-1.2 0-.65.55-1.2 1.2-1.2.65 0 1.2.55 1.2 1.2zM2 4.2C1.34 4.2.8 3.65.8 3c0-.65.55-1.2 1.2-1.2.65 0 1.2.55 1.2 1.2 0 .65-.55 1.2-1.2 1.2z"></path></svg>
  </div>

  <div class="github-info-container">
    <h4>
      <a href="https://github.com/pytorch/pytorch/pull/33907" target="_blank" rel="nofollow noopener">fix handling of replica parameters in DataParallel</a>
    </h4>

    <div class="branches">
      <code>pytorch:master</code> ← <code>ngimel:data_parallel</code>
    </div>

    <div class="github-info">
      <div class="date">
        opened <span class="discourse-local-date" data-format="ll" data-date="2020-02-27" data-time="22:25:12" data-timezone="UTC">10:25PM - 27 Feb 20 UTC</span>
      </div>

      <div class="user">
        <a href="https://github.com/ngimel" target="_blank" rel="nofollow noopener">
          <img alt="ngimel" src="https://avatars3.githubusercontent.com/u/15841449?v=4" class="onebox-avatar-inline" width="20" height="20">
          ngimel
        </a>
      </div>

      <div class="lines" title="4 commits changed 2 files with 9 additions and 9 deletions">
        <a href="https://github.com/pytorch/pytorch/pull/33907/files" target="_blank" rel="nofollow noopener">
          <span class="added">+9</span>
          <span class="removed">-9</span>
        </a>
      </div>
    </div>

  </div>
</div>
  </article>
  <div class="onebox-metadata">
    
    
  </div>
  <div style="clear: both"></div>
</aside>
<br>
After using nightly version, the issue is solved.</p>