Using cuda tensor in DataParallel model

username: cybaj
cooked: <p>there are model wrapped by <code>nn.DataParallel</code></p>
<pre><code class="lang-auto">self.model = Bert(6, 12, 513, 384*4, 64, 64, 2, 384, self.base_task.max_vocab_indexes['input_ids'])
self.model = nn.DataParallel(self.model).cuda()
</code></pre>
<p>and in model there are constant tensor, which name is <code>pos</code>.</p>
<pre><code class="lang-auto">class Embedding(nn.Module):
    def __init__(self, maxlen, d_model, n_segments, vocab_size, device='cuda'):
        super(Embedding, self).__init__()
        self.device = device
        self.tok_embed = nn.Embedding(vocab_size, d_model)  # token embedding
        self.pos_embed = nn.Embedding(maxlen, d_model)  # position embedding
        self.seg_embed = nn.Embedding(n_segments, d_model)  # segment(token type) embedding
        self.norm = nn.LayerNorm(d_model)

    def forward(self, x, seg):
        seq_len = x.size(1)
        pos = torch.arange(seq_len, dtype=torch.long, device=self.device)
        pos = pos.unsqueeze(0).expand_as(x)  # (seq_len,) -&gt; (batch_size, seq_len)
        embedding = self.tok_embed(x) + self.pos_embed(pos) + self.seg_embed(seg)
        return self.norm(embedding)
</code></pre>
<p>I used this code with just one GPU, and this works.<br>
But this time I need to use more gpus, and I have to modify that tensor which was manually located <code>device=self.device</code>, to something “which is dynamically located each gpu at DataParallel”.<br>
… But it is hard to me, code below doesn’t works. The tensor just located in cpu, even with <code>.cuda()</code></p>
<p>How can I solve this issue? I am almost struggling with this issue all day…</p>
<pre><code class="lang-auto">class Embedding(nn.Module):
    def __init__(self, maxlen, d_model, n_segments, vocab_size, device='cuda'):
        super(Embedding, self).__init__()
        self.device = torch.device('cuda')
        self.tok_embed = nn.Embedding(vocab_size, d_model)  # token embedding
        self.pos_embed = nn.Embedding(maxlen, d_model).cpu()  # position embedding
        self.seg_embed = nn.Embedding(n_segments, d_model)  # segment(token type) embedding
        self.norm = nn.LayerNorm(d_model)
        self.pos = torch.arange(513, dtype=torch.long, requires_grad=False).unsqueeze(0)

    def forward(self, x, seg):
        # seq_len = x.size(1)
        print(f'pos device: {self.pos.device}') # printed by "cpu"
        pos = self.pos.expand_as(x)  # (seq_len,) -&gt; (batch_size, seq_len)
        cuda_pos = self.pos_embed(pos).cuda()
        print(f'pos device: {cuda_pos.device}, x device : {x.device}, seg device: {seg.device}') # printed by "cpu" for pos
        embedding = self.tok_embed(x) + cuda_pos + self.seg_embed(seg)
        return self.norm(embedding)
</code></pre>

username: cybaj
cooked: <p>I also tried add that tensor to model’s input. But with this case, when I use 2 gpus, the result is just about half of one batch, <code>(64 x ...)</code> tensor returned, just half of one batch (128).</p>
<pre><code class="lang-auto">0/5120 [00:00&lt;? ?it/s] Traceback (most recent call last):
  File "main.py", line 234, in &lt;module&gt;
    trainer.train()
  File "main.py", line 116, in train
    loss_lm = self.criterion(logits_lm.transpose(1, 2), batch.masked_tokens.transpose(0,1)) # for masked LM
  File "/opt/conda/lib/python3.7/site-packages/torch/nn/modules/module.py", line 550, in __call__
    result = self.forward(*input, **kwargs)
  File "/opt/conda/lib/python3.7/site-packages/torch/nn/modules/loss.py", line 932, in forward
    ignore_index=self.ignore_index, reduction=self.reduction)
  File "/opt/conda/lib/python3.7/site-packages/torch/nn/functional.py", line 2317, in cross_entropy
    return nll_loss(log_softmax(input, 1), target, weight, None, ignore_index, None, reduction)
  File "/opt/conda/lib/python3.7/site-packages/torch/nn/functional.py", line 2113, in nll_loss
    .format(input.size(0), target.size(0)))
ValueError: Expected input batch_size (64) to match target batch_size (128).
</code></pre>
<p>and below is training loop.</p>
<pre><code class="lang-auto">        pos = torch.arange(513, dtype=torch.long, requires_grad=False).unsqueeze(0).to(device=torch.device('cuda'))
        for epoch in range(max_epoch):
            loss_sum, acc_sum, len_batch_sum = 0., 0., 0.
            ds_iter.init_epoch()
            tr_total = math.ceil(total_len / self.batch_size)
            tq_iter = tqdm(enumerate(ds_iter), total=tr_total, miniters=min_iters, unit_scale=self.batch_size,
                           bar_format='{n_fmt}/{total_fmt} [{elapsed}&lt;{remaining} {rate_fmt}] {desc}')

            self.model.train()
            print('epoch starts')
            for i, batch in tq_iter:
                self.model.zero_grad()
                print('batch starts')
                device = torch.device('cuda')
                print(device)
                print(f'batch.input_ids device : {batch.input_ids.device}, batch.segment_ids : {batch.segment_ids.device}, batch.masekd_pos : {batch.masked_pos.device}')
                print(f'batch.input_ids shape : {batch.input_ids.shape}, batch.segment_ids : {batch.segment_ids.shape}, batch.masekd_pos : {batch.masked_pos.shape}, pos : {pos.shape}')
                logits_lm, logits_clsf = self.model(batch.input_ids.transpose(0,1).to(device=device), batch.segment_ids.transpose(0,1).to(device=device), batch.masked_pos.transpose(0,1).to(device=device), pos.to(device=device))
                print(f'logits_lm, logits_clsf shape : {logits_lm.shape}, {logits_clsf.shape}')
</code></pre>
<p>and these are logs.</p>
<pre><code class="lang-auto">epoch starts
batch starts
cuda
batch.input_ids device : cuda:0, batch.segment_ids : cuda:0, batch.masekd_pos : cuda:0
batch.input_ids shape : torch.Size([513, 128]), batch.segment_ids : torch.Size([513, 128]), batch.masekd_pos : torch.Size([5, 128]), pos : torch.Size([1, 513])
pos device: cuda:0 # log from Embedding layer in the model
pos device: cuda:0, x device : cuda:0, seg device: cuda:0 # log from Embedding layer in the model
logits_lm, logits_clsf shape : torch.Size([64, 5, 6015]), torch.Size([64, 2])

logits_lm device: cuda:0, batch target device: cuda:0
0/5120 [00:00&lt;? ?it/s] Traceback (most recent call last):
  File "main.py", line 236, in &lt;module&gt;
    trainer.train()
  File "main.py", line 118, in train
    loss_lm = self.criterion(logits_lm.transpose(1, 2), batch.masked_tokens.transpose(0,1)) # for masked LM
  File "/opt/conda/lib/python3.7/site-packages/torch/nn/modules/module.py", line 550, in __call__
    result = self.forward(*input, **kwargs)
  File "/opt/conda/lib/python3.7/site-packages/torch/nn/modules/loss.py", line 932, in forward
    ignore_index=self.ignore_index, reduction=self.reduction)
  File "/opt/conda/lib/python3.7/site-packages/torch/nn/functional.py", line 2317, in cross_entropy
    return nll_loss(log_softmax(input, 1), target, weight, None, ignore_index, None, reduction)
  File "/opt/conda/lib/python3.7/site-packages/torch/nn/functional.py", line 2113, in nll_loss
    .format(input.size(0), target.size(0)))
ValueError: Expected input batch_size (64) to match target batch_size (128).

</code></pre>

username: mrshenli
cooked: <p>Hey <a class="mention" href="/u/cybaj">@cybaj</a>, the model needs to be moved to GPU before passing it to <code>DataParallel</code> ctor. Have you tried changing the following code:</p>
<aside class="quote no-group" data-username="cybaj" data-post="1" data-topic="90392">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/c/f04885/48.png" class="avatar"> cybaj:</div>
<blockquote>
<p><code>self.model = nn.DataParallel(self.model).cuda()</code></p>
</blockquote>
</aside>
<p>to</p>
<pre><code class="lang-auto">self.model = nn.DataParallel(self.model.to("cuda:0"))
</code></pre>

username: mrshenli
cooked: <p>BTW, here is the <code>DataParallel</code> tutorial: <a href="https://pytorch.org/tutorials/beginner/blitz/data_parallel_tutorial.html" rel="nofollow noopener">https://pytorch.org/tutorials/beginner/blitz/data_parallel_tutorial.html</a></p>

username: cybaj
cooked: <p>Thank you for reply! <a class="mention" href="/u/mrshenli">@mrshenli</a>, I already tried with <code>to("cuda:0")</code> . and as far as I knew <code>cuda()</code> and <code>to(device="cuda")</code> is same.<br>
Below is the issues.</p>
<ol>
<li>In my model, I need to use some constant tensor, and I defined it in model’s <code>forward</code> as you can see top post.</li>
<li>So I give <code>torch.device('cuda')</code> to model, and in model, locate that tensor to the device. more specified,</li>
</ol>
<pre><code class="lang-auto">device = torch.device('cuda')
...
self.model = Bert(6, 12, 513, 384*4, 64, 64, 2, 384, 
						self.base_task.max_vocab_indexes['input_ids'], device=device)
self.model = nn.DataParallel(self.model)
self.model = self.model.cuda()
</code></pre>
<ul>
<li>I ran it with 2 gpus, and it shows up right.</li>
</ul>
<pre><code class="lang-auto">    device_count = torch.cuda.device_count()
    print(f'gpu count: {torch.cuda.device_count()}')

gpu count: 2
</code></pre>
<ul>
<li>dataset iterator is <code>device=torch.device('cuda')</code>
</li>
<li>DataParallel wrapped model with <code>cuda()</code>
</li>
<li>and model get returned, <code>logits_clsf</code> and <code>logits_lm</code>
</li>
<li>loss was calculated.</li>
<li>and <code>loss.backward()</code> phase, I think stuck. <code>print</code> after <code>loss.backward()</code> wasn’t printed.<br>
belows are logs and the model.</li>
</ul>
<pre><code class="lang-auto">batch starts
logits_lm device: cuda:0, batch target device: cuda:0
loss_lm calculated
loss_clsf calculated
</code></pre>
<pre><code class="lang-auto">            print('epoch starts')
            for i, batch in tq_iter:
                self.model.zero_grad()
                print('batch starts')
                logits_lm, logits_clsf = self.model(batch.input_ids.transpose(0,1), batch.segment_ids.transpose(0,1), batch.masked_pos.transpose(0,1))
                print(f'logits_lm device: {logits_lm.device}, batch target device: {batch.masked_tokens.device}')
                loss_lm = self.criterion(logits_lm.transpose(1, 2), batch.masked_tokens.transpose(0,1)) # for masked LM
                print('loss_lm calculated')
                loss_lm = (loss_lm.float()).mean()
                loss_clsf = self.criterion(logits_clsf, batch.is_next) # for sentence classification
                print('loss_clsf calculated')
                loss = loss_lm + loss_clsf
                loss.backward()
                print('loss backwarded')
</code></pre>
<p>without any error message or logs, the log <code>loss backwarded</code> didn’t shows up. I think it’s stuck.<br>
I don’t know what I have to do.<br>
All the post in this thread, are my tries after this happened…</p>
<p>I assumed that,<br>
in backward phase,<br>
‘replica 0’, which is dedicated to <code>cuda:0</code> is worked well(backward completed well), and waiting for ‘replica 1’, which is dedicated to <code>cuda:1</code> to be completed with it’s backward, but in this case, some error occurred to him and failed to backward, so ‘replica 0’ keeps waiting and thus <code>'loss backwarded'</code> log didn’t show up.</p>
<p>Further more assuming that,<br>
even embedding tensor in each replica’s model, looks like splitted well and dedicated well like this,</p>
<pre><code class="lang-auto">class Embedding(nn.Module):
    def __init__ 
        ...
    def forward(self, x, seg):
        seq_len = x.size(1)
        pos = torch.arange(seq_len, dtype=torch.long, device=self.device)
        pos = pos.unsqueeze(0).expand_as(x)  # (seq_len,) -&gt; (batch_size, seq_len)
        pos.requires_grad = False
        print(f'pos tensor device: {pos.device}, shape: {pos.shape}')
        embedding = self.tok_embed(x) + self.pos_embed(pos) + self.seg_embed(seg)
        return self.norm(embedding)


pos tensor device: cuda:1, shape: torch.Size([64, 513])
pos tensor device: cuda:0, shape: torch.Size([64, 513])

embedding tensor device: cuda:1, shape: torch.Size([64, 513, 384])
embedding tensor device: cuda:0, shape: torch.Size([64, 513, 384])
</code></pre>
<p>and about other <code>cuda()</code> defined tensor in the model, too</p>
<pre><code class="lang-auto">context tensor device: cuda:1, shape: torch.Size([64, 513, 768])
multihead output tensor device: cuda:1, shape: torch.Size([64, 513, 384])
context tensor device: cuda:1, shape: torch.Size([64, 513, 768])
multihead output tensor device: cuda:1, shape: torch.Size([64, 513, 384])
context tensor device: cuda:1, shape: torch.Size([64, 513, 768])
multihead output tensor device: cuda:1, shape: torch.Size([64, 513, 384])
context tensor device: cuda:0, shape: torch.Size([64, 513, 768])
multihead output tensor device: cuda:0, shape: torch.Size([64, 513, 384])
context tensor device: cuda:1, shape: torch.Size([64, 513, 768])
multihead output tensor device: cuda:1, shape: torch.Size([64, 513, 384])
context tensor device: cuda:0, shape: torch.Size([64, 513, 768])
multihead output tensor device: cuda:0, shape: torch.Size([64, 513, 384])
context tensor device: cuda:1, shape: torch.Size([64, 513, 768])
multihead output tensor device: cuda:1, shape: torch.Size([64, 513, 384])
context tensor device: cuda:0, shape: torch.Size([64, 513, 768])
multihead output tensor device: cuda:0, shape: torch.Size([64, 513, 384])
context tensor device: cuda:1, shape: torch.Size([64, 513, 768])
multihead output tensor device: cuda:1, shape: torch.Size([64, 513, 384])
context tensor device: cuda:0, shape: torch.Size([64, 513, 768])
multihead output tensor device: cuda:0, shape: torch.Size([64, 513, 384])
context tensor device: cuda:0, shape: torch.Size([64, 513, 768])
multihead output tensor device: cuda:0, shape: torch.Size([64, 513, 384])
context tensor device: cuda:0, shape: torch.Size([64, 513, 768])
multihead output tensor device: cuda:0, shape: torch.Size([64, 513, 384])
</code></pre>
<p>and the last logit gathered well, too,</p>
<pre><code class="lang-auto">logtis_lm shape : torch.Size([128, 5, 6015])
</code></pre>
<p>There are some error could happened at backward phase, but I cannot even imagine about it…</p>
<p>Or maybe the sum about loss, <code>loss = loss_lm + loss_clsf</code> or <code>loss_lm.float().mean()</code> after get gathered outputs, affects to backward phase, I think but … I don’t know what should I do.</p>
<pre><code class="lang-auto">                loss_lm = self.criterion(logits_lm.transpose(1, 2), batch.masked_tokens.transpose(0,1)) # for masked LM
                print('loss_lm calculated')
                print(f'loss_lm tensor device: {loss_lm.device}, shape: {loss_lm.shape}')
                loss_lm = (loss_lm.float()).mean()
                print(f'loss_lm tensor device: {loss_lm.device}, shape: {loss_lm.shape}')
                loss_clsf = self.criterion(logits_clsf, batch.is_next) # for sentence classification
                print('loss_clsf calculated')
                print(f'loss_clsf tensor device: {loss_clsf.device}, shape: {loss_clsf.shape}')
                loss = loss_lm + loss_clsf
                print(f'sumloss tensor device: {loss.device}, shape: {loss.shape}')
                loss.backward()
                self.optimizer.step()
                print('stepped')
</code></pre>

username: mrshenli
cooked: <aside class="quote no-group" data-username="cybaj" data-post="6" data-topic="90392">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/c/f04885/48.png" class="avatar"> cybaj:</div>
<blockquote>
<pre><code class="lang-auto">    def forward(self, x, seg):
        seq_len = x.size(1)
        pos = torch.arange(seq_len, dtype=torch.long, device=self.device)
        pos = pos.unsqueeze(0).expand_as(x)  # (seq_len,) -&gt; (batch_size, seq_len)
        pos.requires_grad = False
        print(f'pos tensor device: {pos.device}, shape: {pos.shape}')
        embedding = self.tok_embed(x) + self.pos_embed(pos) + self.seg_embed(seg)
        return self.norm(embedding)


pos tensor device: cuda:1, shape: torch.Size([64, 513])
pos tensor device: cuda:0, shape: torch.Size([64, 513])
</code></pre>
</blockquote>
</aside>
<p>Hey <a class="mention" href="/u/cybaj">@cybaj</a>, I am a little confused about the above code. IIUC, the <code>self.device</code> attribute on both replicas point to <code>cuda:0</code>, as <code>replicate.py</code> is not smart enough to  change that for you. In that case how did <code>pos</code> successfully located to different devices? Looks like <code>expand_as</code> would not automatically change device either:</p>
<pre data-code-wrap="python"><code class="lang-python">&gt;&gt;&gt; import torch
&gt;&gt;&gt; x = torch.arange(10, device="cuda:0")
&gt;&gt;&gt; y = torch.ones(10, 10).to(1)
&gt;&gt;&gt; z = x.expand_as(y)
&gt;&gt;&gt; z.device
device(type='cuda', index=0)
&gt;&gt;&gt; y.device
device(type='cuda', index=1)
&gt;&gt;&gt; x.device
device(type='cuda', index=0)
</code></pre>
<p>If you would like to get the correct device, can you read that from <code>x.device</code>? (input to forward should be scattered properly.)</p>

username: cybaj
cooked: <p>Thank you, <a class="mention" href="/u/mrshenli">@mrshenli</a> . I change all <code>self.device</code> to something like <code>x.device</code> on forwarding.</p>
<p>… It looks like works but still, stuck after <code>loss.backward()</code> .<br>
All loss calculated and, loss tensor is on <code>cuda:0</code>, which is default output device.</p>
<pre><code class="lang-auto">loss_lm calculated
loss_lm : 72.27577209472656
loss_lm tensor device: cuda:0, shape: torch.Size([])
after mean loss_lm : 72.27577209472656
after mean loss_lm tensor device: cuda:0, shape: torch.Size([])
loss_clsf calculated
loss_clsf : 0.7298979759216309
loss_clsf tensor device: cuda:0, shape: torch.Size([])
sumloss tensor device: cuda:0, shape: torch.Size([])
</code></pre>
<p>Why it is stuck when <code>cuda:0</code> loss tensor starts to backward?<br>
…</p>

username: mrshenli
cooked: <p>Hey <a class="mention" href="/u/cybaj">@cybaj</a>, could you please share a self-contained min repro program? It will be hard to tell with just printed outputs.</p>

username: cybaj
cooked: <p>It is totally my bad, thank you <a class="mention" href="/u/mrshenli">@mrshenli</a>. I learned how to create and use the tensor at module’s forward in DataParallel from your advices.</p>
<p><code>os.environ['CUDA_LAUNCH_BLOCKING'] = '1' </code>  is the fault…</p>
<p>Someone who use this cuda option for checking the logs, because it stops DataParallel process, it is recommend to comment it…</p>