Memory leak of rank 0 in distributed training

JIE_LIU: <p>Hi, I come across a problem when I try to train my model in a distributed way.</p>
<p>In summary, the problem is how to solve the memory leak in rank 0.</p>
<p>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.<br>
I build a class named <code>Trainer</code>, then initiate dataset and model inside. The rough code is showed below. The process contains several backward operations.</p>
<pre><code class="lang-auto">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()
</code></pre>
<p>My main function to call distributed training is showed below.</p>
<pre><code class="lang-auto">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()
</code></pre>
<p>So in process 0(rank 0), the first step training is ok, but in the second step. OOM occurred.</p>
<pre><code class="lang-auto">-- 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)
</code></pre>
<p>I know that maybe the memory leak occurred. But I have no experience about distributed one. How should I solve this problem</p>

JIE_LIU: <p>Does anyone know how to solve it? It’s a little emergent. In summary, how to solve the memory leak in rank 0.</p>

mrshenli: <p>Hey <a class="mention" href="/u/jie_liu">@JIE_LIU</a>, which version of PyTorch are you using? If it is v1.6, I suspect it is  due to the DDP comm bucket reconstruction algorithm temporarily boost memory consumption, and then hit the OOM problem.</p>
<p>cc <a class="mention" href="/u/yanli_zhao">@Yanli_Zhao</a></p>

JIE_LIU: <p>sure, it’s 1.6.0. How can I avoid this situation?</p>

mrshenli: <p>I don’t think there is a decent way to get rid of it in v1.6. One hacky solution might be introducing a tiny unused parameter in the model (e.g., <code>self.unused = nn.Linear(1, 1)</code>), and then set <code>find_unused_parameters=True</code> in DDP ctor, which would disable bucket rebuilt.</p>
<aside class="onebox githubblob">
  <header class="source">
      <a href="https://github.com/pytorch/pytorch/blob/v1.6.0/torch/csrc/distributed/c10d/reducer.cpp#L381" target="_blank" rel="nofollow noopener">github.com</a>
  </header>
  <article class="onebox-body">
    <h4><a href="https://github.com/pytorch/pytorch/blob/v1.6.0/torch/csrc/distributed/c10d/reducer.cpp#L381" target="_blank" rel="nofollow noopener">pytorch/pytorch/blob/v1.6.0/torch/csrc/distributed/c10d/reducer.cpp#L381</a></h4>
<pre class="onebox"><code class="lang-cpp"><ol class="start lines" start="371" style="counter-reset: li-counter 370 ;">
<li>
</li><li>// Rebuild bucket only if 1) it is the first time to rebuild bucket 2)</li>
<li>// unused_parameters_ is empty, currently it does not support when there are</li>
<li>// unused parameters 3) this backward pass needs to run allreduce. Here, we</li>
<li>// just dump tensors and their parameter indices into rebuilt_params_ and</li>
<li>// rebuilt_param_indices_ based on gradient arriving order, and then at the</li>
<li>// end of finalize_backward(), buckets will be rebuilt based on</li>
<li>// rebuilt_params_ and rebuilt_param_indices_, and then will be broadcasted</li>
<li>// and intialized. Also we only need to dump tensors and parameter indcies of</li>
<li>// one replica.</li>
<li class="selected">if (!has_rebuilt_bucket_ &amp;&amp; unused_parameters_.empty() &amp;&amp;</li>
<li>    index.replica_index == 0) {</li>
<li>  rebuilt_params_.push_back(</li>
<li>      replicas_[index.replica_index][index.variable_index]);</li>
<li>  rebuilt_param_indices_.push_back(index.variable_index);</li>
<li>}</li>
<li>
</li><li>// If there are model parameters that went unused when computing the model</li>
<li>// output, they won't be part of the autograd graph, and won't receive</li>
<li>// gradients. These parameters are discovered in the `prepare_for_backward`</li>
<li>// function and their indexes stored in the `unused_parameters_` vector.</li>
</ol></code></pre>


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

<p>This looks like a regression to me. <a class="mention" href="/u/yanli_zhao">@Yanli_Zhao</a> has a some recent work to reduce DDP memory footprint and hopefully that can help.</p>
<p>BTW, if possible, can you try the same script with PyTorch v1.5, it will help to confirm if bucket reconstruction is indeed the culprit.</p>

JIE_LIU: <p>Thanks for your solution. I’ll try this later and give you feedback</p>