Hi, I come across a problem when I try to train my model in a distributed way.
In summary, the problem is how to solve the memory leak in rank 0.
When I run the code in single gpu, it works well and occupies 10.3G GPU memory. My GPU is 2080ti 11GB. But when I run the code in the distributed way, OOM occurred.
I build a class named Trainer, then initiate dataset and model inside. The rough code is showed below. The process contains several backward operations.
class Trainer():
    def __init__(self):
        self.data_initial()
        self.model_initial()
        self.train()

   def data_initial(self):
        source_data = DataSet(.....)
        source_sampler = data.distributed.DistributedSampler(source_data, seed=1234)
        source_dataloader = data.DataLoader(source_data, batch_size=..., num_workers=4,
            pin_memory=False, sampler=source_sampler)
        self.source_data = enumerate(source_dataloader)

        target_data = DataSet(.....)
        target_sampler = data.distributed.DistributedSampler(target_data, seed=1234)
        target_dataloader = data.DataLoader(source_data, batch_size=..., num_workers=4,
            pin_memory=False, sampler=target_sampler)
        self.target_data = enumerate(target_dataloader)

    def model_initial(self):
        # build backbone
        rank = dist.get_rank()
        self.backbone = ResNet().cuda(rank)
        self.backbone = DDP(self.backbone, device_ids=[rank])
        
        # restore_part
        the parameter succeeds restore. I have checked it.
        self.backbone.train()

        # classifier
        self.classifier = Classifier(...).cuda(rank)
        self.classifier = DDP(self.classifier, device_ids=[rank])
        self.classifier.train()

        # optimizer
        self.backbone_optimizer = optim.SGD(self.backbone.parameters(), ...)
        self.backbone_optimizer.zero_grad()

        self.classifier_optimizer = optim.Adam(self.classifier.parameters(), ...)
        self.classifier_optimizer.zero_grad()

   def train(self):
        # pytorch prework
        rank = dist.get_rank()
        self.criterion = torch.nn.BCEWithLogitsLoss().cuda(rank)
        
        for i in range(1, self.config.num_steps):
            # get data
            c, batch = self.source_data.__next__()
            image, label, _, _ = batch
            _, batch = self.target_data.__next__()
            image_t, label_t, _, _ = batch

            self.step(i, image, label, image_t, label_t, loss_dic)

            gc.collect()

def step(self, i, image, label, image_t, label_t, loss_dic):
        rank = dist.get_rank()

        self.backbone_optimizer.zero_grad()
        self.classifier_optimizer.zero_grad()

        # supervised learning for source
        image = Variable(image).cuda(rank)
        x = self.backbone(image)
        y1, _ = self.classifier(x)
        loss = self.criterion(y, label.long().cuda(rank))
        loss.backward()
      

        for para in self.classifier.parameters():
            para.requires_grad = False
        image_t = Variable(image_t).cuda(rank)
        x = self.backbone(image_t)
        _, y2 = self.classifier(x)

        label2 = Variable(...).cuda(rank)
        loss = self.criterion(y2, label2)
        loss.backward()

        ###
        # optimize the parameter
        self.backbone_optimizer.step()
        self.classifier_optimizer.step()

        ###
        # recycle variable
        #delete some variable in intermedia process
        del x .....
        torch.cuda.empty_cache()

My main function to call distributed training is showed below.
def main():
    world_size = torch.cuda.device_count()
    mp.spawn(sub_process, args=(world_size), nprocs=world_size, join=True)

def sub_process(rank, world_size):
    set_up(rank, world_size)
    trainer = Trainer()
    cleanup()

def set_up(rank, world_size):
    os.environ['MASTER_ADDR'] = '127.0.0.113'
    os.environ['MASTER_PORT'] = '12355'
    dist.init_process_group("nccl", rank=rank, world_size=world_size) 

def cleanup():
    dist.destroy_process_group()

So in process 0(rank 0), the first step training is ok, but in the second step. OOM occurred.
-- Process 0 terminated with the following error:
Traceback (most recent call last):
  File "/home/xx/anaconda3/lib/python3.6/site-packages/torch/multiprocessing/spawn.py", line 20, in _wrap
    fn(i, *args)
.......
File "/home/xx/code/xxx.py", line 225, in step
    loss_aux = self.seg_criterion(out_aux, label.long().cuda(rank))
  File "/home/xx/anaconda3/lib/python3.6/site-packages/torch/nn/modules/module.py", line 722, in _call_impl
    result = self.forward(*input, **kwargs)
  File "/home/xx/code/multitask/utils/utils.py", line 44, in forward
    predict = predict[target_mask.view(n, h, w, 1).repeat(1, 1, 1, c)].view(-1, c)
RuntimeError: CUDA out of memory. Tried to allocate 536.00 MiB (GPU 0; 10.73 GiB total capacity; 8.30 GiB already allocated; 198.56 MiB free; 9.01 GiB reserved in total by PyTorch)

I know that maybe the memory leak occurred. But I have no experience about distributed one. How should I solve this problem