Hey @Hongkai_Zheng
I tried the code below, but couldn’t reproduce the error either. Which version of PyTorch are you using? I am on 1.10.0a0+git9730d91.
BTW, I also checked that there is no sparse gradients after the backward pass.
import os

import torch
import torch.nn as nn
import torch.multiprocessing as mp
import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP
from torch.optim import Adam

import torchvision.utils as vutil
import argparse


# Number of channels in the training images. For color images this is 3
nc = 3

# Size of z latent vector (i.e. size of generator input)
nz = 100

# Size of feature maps in generator
ngf = 64

# Size of feature maps in discriminator
ndf = 64

# Number of training epochs
num_epochs = 5

# Learning rate for optimizers
lr = 0.0002

# Beta1 hyperparam for Adam optimizers
beta1 = 0.5

# Number of GPUs available. Use 0 for CPU mode.
ngpu = 1


class Generator(nn.Module):
    def __init__(self, ngpu):
        super(Generator, self).__init__()
        self.ngpu = ngpu
        self.main = nn.Sequential(
            # input is Z, going into a convolution
            nn.ConvTranspose2d( nz, ngf * 8, 4, 1, 0, bias=False),
            nn.BatchNorm2d(ngf * 8),
            nn.ReLU(True),
            # state size. (ngf*8) x 4 x 4
            nn.ConvTranspose2d(ngf * 8, ngf * 4, 4, 2, 1, bias=False),
            nn.BatchNorm2d(ngf * 4),
            nn.ReLU(True),
            # state size. (ngf*4) x 8 x 8
            nn.ConvTranspose2d( ngf * 4, ngf * 2, 4, 2, 1, bias=False),
            nn.BatchNorm2d(ngf * 2),
            nn.ReLU(True),
            # state size. (ngf*2) x 16 x 16
            nn.ConvTranspose2d( ngf * 2, ngf, 4, 2, 1, bias=False),
            nn.BatchNorm2d(ngf),
            nn.ReLU(True),
            # state size. (ngf) x 32 x 32
            nn.ConvTranspose2d( ngf, nc, 4, 2, 1, bias=False),
            nn.Tanh()
            # state size. (nc) x 64 x 64
        )

    def forward(self, input):
        return self.main(input)


class Discriminator(nn.Module):
    def __init__(self, ngpu):
        super(Discriminator, self).__init__()
        self.ngpu = ngpu
        self.main = nn.Sequential(
            # input is (nc) x 64 x 64
            nn.Conv2d(nc, ndf, 4, 2, 1, bias=False),
            nn.LeakyReLU(0.2, inplace=True),
            # state size. (ndf) x 32 x 32
            nn.Conv2d(ndf, ndf * 2, 4, 2, 1, bias=False),
            nn.BatchNorm2d(ndf * 2),
            nn.LeakyReLU(0.2, inplace=True),
            # state size. (ndf*2) x 16 x 16
            nn.Conv2d(ndf * 2, ndf * 4, 4, 2, 1, bias=False),
            nn.BatchNorm2d(ndf * 4),
            nn.LeakyReLU(0.2, inplace=True),
            # state size. (ndf*4) x 8 x 8
            nn.Conv2d(ndf * 4, ndf * 8, 4, 2, 1, bias=False),
            nn.BatchNorm2d(ndf * 8),
            nn.LeakyReLU(0.2, inplace=True),
            # state size. (ndf*8) x 4 x 4
            nn.Conv2d(ndf * 8, 1, 4, 1, 0, bias=False),
            nn.Sigmoid()
        )

    def forward(self, input):
        return self.main(input)

# =======================================


def setup(rank, world_size):
    os.environ['MASTER_ADDR'] = 'localhost'
    os.environ['MASTER_PORT'] = '7777'
    # initialize the process group
    dist.init_process_group("nccl", rank=rank, world_size=world_size)


def cleanup():
    dist.destroy_process_group()


def subprocess_fn(rank, args):
    setup(rank, args.num_gpus)
    print(f'running on rank {rank}')
    generator = Generator(args.num_gpus).to(rank)
    discriminator = Discriminator(args.num_gpus).to(rank)

    if args.distributed:
        generator = DDP(generator, device_ids=[rank], broadcast_buffers=False)
        discriminator = DDP(discriminator, device_ids=[rank], broadcast_buffers=False)
    d_optim = Adam(discriminator.parameters(), lr=2e-4)
    g_optim = Adam(generator.parameters(), lr=2e-4)

    discriminator.train()
    generator.train()
    if rank == 0:
        fixed_z = torch.randn(64, 100, 1, 1).to(rank)

    pbar = range(args.iter)
    for e in pbar:
        real_data = torch.randn((args.batchsize, 3, 64, 64)).to(rank)
        real_pred = discriminator(real_data)

        latent = torch.randn((args.batchsize, 100, 1, 1)).to(rank)
        fake_data = generator(latent)
        fake_pred = discriminator(fake_data)

        d_loss = real_pred.mean() - fake_pred.mean()

        d_optim.zero_grad()
        d_loss.backward()
        d_optim.step()

        if rank == 0 and e % 2 == 0:
            print(f'Epoch D loss:{d_loss.item()};')
            with torch.no_grad():
                imgs = generator(fixed_z)

    cleanup()
    print(f'Process {rank} exits...')


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('--num_gpus', type=int, help='Number of GPUs', default=2)
    parser.add_argument('--batchsize', type=int, help='Batch size for each GPU', default=2)
    parser.add_argument('--iter', type=int, help='Total training iterations', default=10)
    args = parser.parse_args()
    args.distributed = args.num_gpus > 1

    if args.distributed:
        mp.spawn(subprocess_fn, args=(args, ), nprocs=args.num_gpus)
    else:
        subprocess_fn(0, args)

    print('Done!')