How to combine model parallel with data parallel?

kstarxin: <p>I have designed a big model, its architecture looks like this:</p>
<pre><code class="lang-auto">class BigModel(nn.Module):
    def __init__(self, encoder: nn.Module, component1: nn.Module, component2: nn.Module, component3: nn.Module):
        super(BigModel, self).__init__()
        self.encoder = nn.DataParallel(encoder, device_ids=["cuda:0", "cuda:1","cuda:2", "cuda:3"])
       self.component1 = component1
       self.component2 = component2
       self.component3 = component3
    def deploy(self):
       self.component1 = self.component1.to("cuda:4")
       self.component2 = self.component2.to("cuda:4")
       self.component3 = self.component3.to("cuda:5")
</code></pre>
<p>I have a single machine with 6 Tesla V100 GPU (32 GB)<br>
The encoder is very big (BERT-like model) so I want to use 4 GPUs to perform encoding process and then use the rest GPUs for other works of this model.</p>
<p>It can work but some times the so of the outputs of the encoder are  lost.</p>
<p>For example, I input a batch of sequences of size (16, 256) to the encoder, data parallel should split it into 4 tensor of size (4, 256) and encode them in parallel. Then gather those output and merge into a tensor of size (16, 256, 1024) .</p>
<p>But, sometimes I only got the output of size  (12, 256, 1024),  that means one split of data is lost. I can not figure out the reason of this problem…</p>
<p>So is there anyone can explain this problem or give a way to combine model parallel and data parallel ?</p>

mrshenli: <aside class="quote no-group" data-username="kstarxin" data-post="1" data-topic="115129">
<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/k/c37758/48.png" class="avatar"> kstarxin:</div>
<blockquote>
<p>For example, I input a batch of sequences of size (16, 256) to the encoder, data parallel should split it into 4 tensor of size (4, 256) and encode them in parallel. Then gather those output and merge into a tensor of size (16, 256, 1024) .</p>
</blockquote>
</aside>
<p>This is very weird. Are you using <code>DataParallel</code> instead of <code>DistributedDataParallel</code>? In your <code>encoder</code>’s forward function, could you please add some prints to check if <code>parallel_apply</code> indeed spawned two threads and each thread is getting <code>(4, 256)</code> input and emitting <code>(4, 256, 1024)</code> output?</p>

kstarxin: <p>I do use the nn.Dataparallel instead of nn.DistributedDataParallel.</p>
<p>I have checked the forward function of encoder, and I found that when the program went wrong, there are only 3 encode calls instead of 4. I dont know what is going on…</p>

mrshenli: <aside class="quote no-group" data-username="kstarxin" data-post="3" data-topic="115129">
<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/k/c37758/48.png" class="avatar"> kstarxin:</div>
<blockquote>
<p>when the program went wrong</p>
</blockquote>
</aside>
<p>you mean there is no error message at all. Could you please share a minimum repro?</p>

kstarxin: <p>I think I have found out the cause of the problem, here is my code, it is a runable demo and you need to install:</p>
<p>torch==‘1.7.1+cu110’<br>
transformers==4.3.2</p>
<pre><code class="lang-auto">import transformers
import torch
from transformers import ElectraModel
from transformers import AdamW
from tqdm import tqdm

class Encoder(torch.nn.Module):
    def __init__(self):
        super(Encoder, self).__init__()
        self.encoder = ElectraModel.from_pretrained("google/electra-large-discriminator")
    def forward(self, x1, x2):
        return [self.encoder(x1), self.encoder(x2)]

class BigModel(torch.nn.Module):
    def __init__(self):
        super(BigModel, self).__init__()
        self.encoder = torch.nn.DataParallel(Encoder(),device_ids=["cuda:0", "cuda:1", "cuda:2", "cuda:3"] )
        self.l1 = torch.nn.Linear(1024, 1024)
        self.l2 = torch.nn.Linear(1024, 1024)
        self.l3 = torch.nn.Linear(1024, 2)
        self.loss = torch.nn.CrossEntropyLoss()

    def deploy(self):
        self.l1 = self.l1.to("cuda:4")
        self.l2 = self.l2.to("cuda:4")
        self.l3 = self.l3.to("cuda:5")
        self.loss = self.loss.to("cuda:5")

    def forward(self, x: torch.LongTensor, x1: torch.LongTensor, y: torch.LongTensor):
        res = self.encoder(x, x1)
        x2 = self.l1(res[0][0].to("cuda:4")).to("cuda:5")
        x3 = self.l2(res[1][0].to("cuda:4")).to("cuda:5")
        print(x2.shape)
        print(x3.shape)
        x4 = self.l3((x2.unsqueeze(1) + x3.unsqueeze(0)).mean(1))
        return self.loss(x4.reshape((-1, 2)), y.reshape((-1,))).mean()


def main():
    model = BigModel().cuda()
    model.deploy()
    optimizer = AdamW(model.parameters(), lr=3e-5)
    # assume we train a token binary classificaiton task for 10 epochs, and the data set has 100 batch
    for epoch in range(10):
        for d in tqdm(range(100)):
            ids = torch.randint(0, 30000, (16, 256))
            ids1 = torch.randint(0, 30000, (6, 256)) # strange size that lead to the problem
            y = torch.randint(0, 2, (16, 256))
            model.zero_grad()
            ids = ids.cuda()
            ids1 = ids1.cuda()
            y = y.to("cuda:5")
            loss = model(ids, ids1, y)
            loss.backward()
            optimizer.step()


main()
</code></pre>
<p>And after I start this python program, and get the error output</p>
<pre><code class="lang-auto">  0%|                                                                                                                           | 0/100 [00:00&lt;?, ?it/s]
torch.Size([12, 256, 1024])
torch.Size([6, 256, 1024])
  0%|                                                                                                                           | 0/100 [00:09&lt;?, ?it/s]
Traceback (most recent call last):
  File "demo.py", line 58, in &lt;module&gt;
    main()
  File "demo.py", line 53, in main
    loss = model(ids, ids1, y)
  File "/usr/local/lib/python3.7/dist-packages/torch/nn/modules/module.py", line 727, in _call_impl
    result = self.forward(*input, **kwargs)
  File "demo.py", line 36, in forward
    return self.loss(x4.reshape((-1, 2)), y.reshape((-1,))).mean()
  File "/usr/local/lib/python3.7/dist-packages/torch/nn/modules/module.py", line 727, in _call_impl
    result = self.forward(*input, **kwargs)
  File "/usr/local/lib/python3.7/dist-packages/torch/nn/modules/loss.py", line 962, in forward
    ignore_index=self.ignore_index, reduction=self.reduction)
  File "/usr/local/lib/python3.7/dist-packages/torch/nn/functional.py", line 2468, in cross_entropy
    return nll_loss(log_softmax(input, 1), target, weight, None, ignore_index, None, reduction)
  File "/usr/local/lib/python3.7/dist-packages/torch/nn/functional.py", line 2262, in nll_loss
    .format(input.size(0), target.size(0)))
ValueError: Expected input batch_size (3072) to match target batch_size (4096).
</code></pre>
<p>This error happens in loss function and torch not check the the input batch size of all input tensors. And at the same time, a split of one encode result is lost.</p>

kstarxin: <p><a class="mention" href="/u/mrshenli">@mrshenli</a>  Here is my demo program, is there a potential bug ?</p>

mrshenli: <p>I can reproduce this locally, and here is what happened:</p>
<p>DP will try to scatter the inputs to the given devices:</p>
<aside class="onebox githubblob">
  <header class="source">
      <a href="https://github.com/pytorch/pytorch/blob/679f07a017d7387dce76936ced0eca55881d99f5/torch/nn/parallel/data_parallel.py#L157" target="_blank" rel="noopener nofollow ugc">github.com</a>
  </header>
  <article class="onebox-body">
    <h4><a href="https://github.com/pytorch/pytorch/blob/679f07a017d7387dce76936ced0eca55881d99f5/torch/nn/parallel/data_parallel.py#L157" target="_blank" rel="noopener nofollow ugc">pytorch/pytorch/blob/679f07a017d7387dce76936ced0eca55881d99f5/torch/nn/parallel/data_parallel.py#L157</a></h4>
<pre class="onebox"><code class="lang-py"><ol class="start lines" start="147" style="counter-reset: li-counter 146 ;">
<li>def forward(self, *inputs, **kwargs):</li>
<li>    if not self.device_ids:</li>
<li>        return self.module(*inputs, **kwargs)</li>
<li>
<li>    for t in chain(self.module.parameters(), self.module.buffers()):</li>
<li>        if t.device != self.src_device_obj:</li>
<li>            raise RuntimeError("module must have its parameters and buffers "</li>
<li>                               "on device {} (device_ids[0]) but found one of "</li>
<li>                               "them on device: {}".format(self.src_device_obj, t.device))</li>
<li>
<li class="selected">    inputs, kwargs = self.scatter(inputs, kwargs, self.device_ids)</li>
<li>    # for forward function without any inputs, empty list and dict will be created</li>
<li>    # so the module can be executed on one device which is the first one in device_ids</li>
<li>    if not inputs and not kwargs:</li>
<li>        inputs = ((),)</li>
<li>        kwargs = ({},)</li>
<li>
<li>    if len(self.device_ids) == 1:</li>
<li>        return self.module(*inputs[0], **kwargs[0])</li>
<li>    replicas = self.replicate(self.module, self.device_ids[:len(inputs)])</li>
<li>    outputs = self.parallel_apply(replicas, inputs, kwargs)</li>
</ol></code></pre>


  </article>
  <div class="onebox-metadata">
    
    
  </div>
  <div style="clear: both"></div>
</aside>

<p>In this case, the given devices are [0, 1, 2, 3], and the inputs are 16X* and 6X* tensors. The scatter function traverses the list and try to scatter every tensor in the list.</p>
<aside class="onebox githubblob">
  <header class="source">
      <a href="https://github.com/pytorch/pytorch/blob/454dd7ba8647ac11735c6563ac8e2b60313789ad/torch/nn/parallel/scatter_gather.py#L11-L39" target="_blank" rel="noopener nofollow ugc">github.com</a>
  </header>
  <article class="onebox-body">
    <h4><a href="https://github.com/pytorch/pytorch/blob/454dd7ba8647ac11735c6563ac8e2b60313789ad/torch/nn/parallel/scatter_gather.py#L11-L39" target="_blank" rel="noopener nofollow ugc">pytorch/pytorch/blob/454dd7ba8647ac11735c6563ac8e2b60313789ad/torch/nn/parallel/scatter_gather.py#L11-L39</a></h4>
<pre class="onebox"><code class="lang-py"><ol class="start lines" start="11" style="counter-reset: li-counter 10 ;">
<li>def scatter(inputs, target_gpus, dim=0):</li>
<li>    r"""</li>
<li>    Slices tensors into approximately equal chunks and</li>
<li>    distributes them across given GPUs. Duplicates</li>
<li>    references to objects that are not tensors.</li>
<li>    """</li>
<li>    def scatter_map(obj):</li>
<li>        if isinstance(obj, torch.Tensor):</li>
<li>            return Scatter.apply(target_gpus, None, dim, obj)</li>
<li>        if is_namedtuple(obj):</li>
<li>            return [type(obj)(*args) for args in zip(*map(scatter_map, obj))]</li>
<li>        if isinstance(obj, tuple) and len(obj) &gt; 0:</li>
<li>            return list(zip(*map(scatter_map, obj)))</li>
<li>        if isinstance(obj, list) and len(obj) &gt; 0:</li>
<li>            return [list(i) for i in zip(*map(scatter_map, obj))]</li>
<li>        if isinstance(obj, dict) and len(obj) &gt; 0:</li>
<li>            return [type(obj)(i) for i in zip(*map(scatter_map, obj.items()))]</li>
<li>        return [obj for targets in target_gpus]</li>
<li>
<li>    # After scatter_map is called, a scatter_map cell will exist. This cell</li>
</ol></code></pre>

  This file has been truncated. <a href="https://github.com/pytorch/pytorch/blob/454dd7ba8647ac11735c6563ac8e2b60313789ad/torch/nn/parallel/scatter_gather.py#L11-L39" target="_blank" rel="noopener nofollow ugc">show original</a>

  </article>
  <div class="onebox-metadata">
    
    
  </div>
  <div style="clear: both"></div>
</aside>

<p>In your code, it’s same as calling <code>x.chunk(4, 0)</code> and <code>x1.chunk(4, 0)</code>. However, <code>x.chunk(4, 0)</code> will return 4 tensors and 16 can be divided by 4. But <code>x1.chunk(4, 0)</code> only returns 3 tensors, as the chunk algorithm there is, when not divisible, put <code>6 / chunks-1</code> in the first <code>chunks-1</code> splits and the reminder in the last split. But <code>6 / (4-1)</code> is divisible, as a result, the last split has nothing.</p>
<p>See: <a href="https://pytorch.org/docs/stable/generated/torch.chunk.html" class="inline-onebox" rel="noopener nofollow ugc">torch.chunk — PyTorch 1.8.0 documentation</a></p>
<p>Then the scatter tries to zip together splits from <code>x</code> and <code>x1</code>, as a result, <code>x</code>'s last split is dropped.</p>
<aside class="onebox githubblob">
  <header class="source">
      <a href="https://github.com/pytorch/pytorch/blob/454dd7ba8647ac11735c6563ac8e2b60313789ad/torch/nn/parallel/scatter_gather.py#L27" target="_blank" rel="noopener nofollow ugc">github.com</a>
  </header>
  <article class="onebox-body">
    <h4><a href="https://github.com/pytorch/pytorch/blob/454dd7ba8647ac11735c6563ac8e2b60313789ad/torch/nn/parallel/scatter_gather.py#L27" target="_blank" rel="noopener nofollow ugc">pytorch/pytorch/blob/454dd7ba8647ac11735c6563ac8e2b60313789ad/torch/nn/parallel/scatter_gather.py#L27</a></h4>
<pre class="onebox"><code class="lang-py"><ol class="start lines" start="17" style="counter-reset: li-counter 16 ;">
<li>def scatter_map(obj):</li>
<li>    if isinstance(obj, torch.Tensor):</li>
<li>        return Scatter.apply(target_gpus, None, dim, obj)</li>
<li>    if is_namedtuple(obj):</li>
<li>        return [type(obj)(*args) for args in zip(*map(scatter_map, obj))]</li>
<li>    if isinstance(obj, tuple) and len(obj) &gt; 0:</li>
<li>        return list(zip(*map(scatter_map, obj)))</li>
<li>    if isinstance(obj, list) and len(obj) &gt; 0:</li>
<li>        return [list(i) for i in zip(*map(scatter_map, obj))]</li>
<li>    if isinstance(obj, dict) and len(obj) &gt; 0:</li>
<li class="selected">        return [type(obj)(i) for i in zip(*map(scatter_map, obj.items()))]</li>
<li>    return [obj for targets in target_gpus]</li>
<li>
<li># After scatter_map is called, a scatter_map cell will exist. This cell</li>
<li># has a reference to the actual function scatter_map, which has references</li>
<li># to a closure that has a reference to the scatter_map cell (because the</li>
<li># fn is recursive). To avoid this reference cycle, we set the function to</li>
<li># None, clearing the cell</li>
<li>try:</li>
<li>    res = scatter_map(inputs)</li>
<li>finally:</li>
</ol></code></pre>


  </article>
  <div class="onebox-metadata">
    
    
  </div>
  <div style="clear: both"></div>
</aside>

<p>Curious, what is the expected behavior here? is it <code>[4, 2], [4, 2], [4, 2], [4, 0]</code>, or <code>[4, 2], [4, 2], [4, 1], [4, 1]</code>?</p>