When I launch the following script with the torch.distributed.launch utilility on a 2 GPUs machine, I get a much slower (10x) training than when I launch it on a single GPU.
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.
I profiled the code with Nvidia Nsight Systems and saw that there is a call to ncclAllReduceRingLLKernel_sum_f32 that takes approximately 500 ms each iteration.
Is this expected behaviour with this kind of network? Or am I doing something wrong?
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 > 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() > 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()
I did follow the tutorial. I use the