DistributedDataParallel slows down training with big fully connected layers

username: Martin_Castin
cooked: <p>When I launch the following script with the <code>torch.distributed.launch</code> utilility on a 2 GPUs machine, I get a much slower (10x) training than when I launch it on a single GPU.</p>
<p>I realized that it seems to come from the big fully connected layer at the end of the network (130000x1024), and I suppose it is because the gradients that need to be synchronized at each iteration represent a big amount of memory.<br>
I profiled the code with Nvidia Nsight Systems and saw that there is a call to <code>ncclAllReduceRingLLKernel_sum_f32</code> that takes approximately 500 ms each iteration.</p>
<p>Is this expected behaviour with this kind of network? Or am I doing something wrong?</p>
<pre><code class="lang-python">import torch
import torch.nn as nn
import argparse
from torch.nn.parallel import DistributedDataParallel as DPP
import torch.nn.functional as F
from tqdm import tqdm


class Discriminator(nn.Module):
    def __init__(self):
        super(Discriminator, self).__init__()

        in_channels = 3
        out_channels = 64
        depth = 7

        m_features = [
            nn.Conv2d(in_channels, out_channels, 3, padding=1),
        ]
        for i in range(depth):
            in_channels = out_channels
            if i % 2 == 1:
                stride = 1
                out_channels *= 2
            else:
                stride = 2
            m_features.append(nn.Conv2d(
                in_channels, out_channels, 3, padding=1, stride=stride,
            ))

        self.features = nn.Sequential(*m_features)

        patch_size = 256 // (2 ** ((depth + 1) // 2))
        m_classifier = [
            nn.Linear(out_channels * patch_size ** 2, 1024),
            nn.LeakyReLU(negative_slope=0.2, inplace=False),
            nn.Linear(1024, 1)
        ]
        self.classifier = nn.Sequential(*m_classifier)

    def forward(self, f0):
        features = self.features(f0)
        output = self.classifier(features.view(features.size(0), -1))

        return output

torch.backends.cudnn.enabled = True  # make sure to use cudnn for computational performance

parser = argparse.ArgumentParser()
parser.add_argument("--local_rank", type=int, default=0)

args = parser.parse_args()

def train(rank, world_size):

    if world_size &gt; 1:
        torch.distributed.init_process_group("nccl", init_method="env://", rank=rank, world_size=world_size)
        torch.cuda.set_device(rank)

    discriminator = Discriminator()
    discriminator.to(rank)
    optimizer = torch.optim.Adam(
        discriminator.parameters(),
        lr=1e-5
    )

    # -- Initialize model for distributed training --
    if torch.cuda.device_count() &gt; 1:
        discriminator = DPP(discriminator, device_ids=[rank])

    frame = torch.rand((1, 3, 256, 256), device=f"cuda:{rank}")
    d_01 = discriminator(frame)

    label_01 = torch.zeros_like(d_01)

    for i in tqdm(range(30)):

        # - Compute loss -
        d_01 = discriminator(frame)
        loss = F.binary_cross_entropy_with_logits(d_01, label_01)

        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

def main():
    world_size = torch.cuda.device_count()

    with torch.autograd.profiler.emit_nvtx():
        train(args.local_rank, world_size)

if __name__ == '__main__':
    main()


</code></pre>

username: Martin_Castin
cooked: <p>Does anyone have an insight on this?</p>

username: wanchaol
cooked: <p>so looking at your code, it looks like you didn’t create the process groups in different processes, and ended up just using one process, or are you just running the scripts on multi hosts? If not, did you try following the <a href="https://pytorch.org/tutorials/intermediate/ddp_tutorial.html" rel="noopener nofollow ugc">DDP tutorial</a>, launch it in multiple processes and see if improves the performance?</p>

username: Martin_Castin
cooked: <p>Thank you for your answer <img src="https://discuss.pytorch.org/images/emoji/apple/slight_smile.png?v=9" title=":slight_smile:" class="emoji" alt=":slight_smile:"> I did follow the tutorial. I use the <a href="https://pytorch.org/docs/stable/distributed.html#launch-utility" rel="noopener nofollow ugc"><code>torch.distributed.launch</code></a> utility that takes care of creating one process per GPU. I am using it on one machine with two GPUs.</p>
<p>It is true that I put my two models in the same (default) process group, but I also tried to use different process groups for both models and it did not change anything.</p>

username: wayi
cooked: <aside class="quote no-group" data-username="Martin_Castin" data-post="1" data-topic="121633">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/martin_castin/48/31078_2.png" class="avatar"> Martin_Castin:</div>
<blockquote>
<p>ncclAllReduceRingLLKernel_sum_f32</p>
</blockquote>
</aside>
<p>If the cost is dominated by allreduce communication, can you try <a href="https://pytorch.org/docs/stable/generated/torch.nn.parallel.DistributedDataParallel.html#torch.nn.parallel.DistributedDataParallel.no_sync" rel="noopener nofollow ugc">no_sync context manager</a> to reduce the sync frequency?</p>
<p>Moreover, you probably can try to register a <a href="https://pytorch.org/docs/master/ddp_comm_hooks.html#torch.distributed.algorithms.ddp_comm_hooks.default_hooks.fp16_compress_hook" rel="noopener nofollow ugc">FP16 compression communication hook</a> to compress the gradients before allreduce. It’s one-line code change. You can try even more advanced gradient compression if interested.</p>
<p>I am really surprised at such a high allreduce cost.</p>

username: Martin_Castin
cooked: <p>Thank you for your answer, I am sure your suggestion would have helped!</p>
<p>I finally understood the source of the problem. This network contains this layer: <code>Linear(in_features=131072, out_features=1024, bias=True)</code></p>
<p>Which requires a huge number of gradients to be synchronized (~400 MB per iteration in fp32). As I don’t have any NVLink it makes sense that this synchronisation takes 500 ms I think.</p>
<p>So I ended up redesigning the network to avoid having such a huge fully connected layer.</p>