Model takes twice the memory footprint with distributed data parallel

jastern33: <p>Cross-posting <a href="https://stackoverflow.com/questions/68949954/model-takes-twice-the-memory-footprint-with-distributed-data-parallel" rel="noopener nofollow ugc">from stackoverflow</a>, because it wasn’t getting much attention there. There’s an open bounty, and if anyone answers over there, I’m happy to award it to you.</p>
<p>The question is: when I use distributed data parallel, I see double the memory usage (almost exactly) on distributed data parallel compared to single-GPU training - it looks like two copies of the model are being stored on every GPU. Why does the model take up twice the space in DDP? Is it intended behavior? Is there a way to avoid this extra memory usage?</p>
<p>Here is a minimal working example.</p>
<pre><code class="lang-auto">import os
from torch.nn.parallel import DistributedDataParallel as DDP
import torch.distributed as dist
import torch.multiprocessing as mp
import torch

def train(rank, gpu_list, train_distributed):
    
    device_id = gpu_list[rank]

    model = torch.nn.Linear(1000, 1000)
    print(device_id, torch.cuda.memory_allocated(device_id))
    model.to(device_id)
    print(device_id, torch.cuda.memory_allocated(device_id))

    print(device_id, torch.cuda.memory_allocated(device_id))
    if train_distributed:
        # convert model to DDP
        dist.init_process_group("gloo", rank=rank, world_size=len(gpu_list))
        model = DDP(model, device_ids=[device_id], find_unused_parameters=False)
    print(device_id, torch.cuda.memory_allocated(device_id))

def train_distributed():
    gpu_list = [torch.device(i) for i in [5, 6]]
    os.environ['MASTER_ADDR'] = '127.0.01'
    os.environ['MASTER_PORT'] = '7676'
    mp.spawn(train, args=(gpu_list, True), nprocs=len(gpu_list), join=True)

if __name__ == '__main__':
    # First test one GPU
    print("Single GPU")
    train(0, [torch.device(5)], False)
    print("Multi GPU")
    # Then test multiple GPUs
    train_distributed()
</code></pre>
<p>Output:</p>
<pre><code class="lang-auto">Single GPU
cuda:5 0
cuda:5 4004352
cuda:5 4004352
cuda:5 4004352
Multi GPU
cuda:5 0
cuda:6 0
cuda:5 4004352
cuda:5 4004352
cuda:6 4004352
cuda:6 4004352
cuda:5 8008704
cuda:6 8008704
</code></pre>
<p>I tried rewriting this snippet using the command-line version of DDP, using torch.distributed.launch, and saw the same issue.</p>

zetyquickly: <p>That’s really interesting question. I’d also would like to know why.</p>
<p>As I can see not all models take twice as much memory on the card but <code>Linear(1000, 1000)</code> definitely does. Also it’s seen that with 3 cards available it takes the same amount of GPU memory, so it’s kind of invariant of the number of GPUs greater than 1</p>

jastern33: <p>Interesting, do you have an example model that I could try which doesn’t have the duplicated memory issue?</p>

zetyquickly: <p>Try to change your <code>Linear(1000, 1000)</code> to this model from the tutorials:</p>
<pre><code class="lang-auto">class ToyModel(nn.Module):
    def __init__(self):
        super(ToyModel, self).__init__()
        self.net1 = nn.Linear(10, 10)
        self.relu = nn.ReLU()
        self.net2 = nn.Linear(10, 5)

    def forward(self, x):
        return self.net2(self.relu(self.net1(x)))

</code></pre>
<p>Logs will tell</p>
<pre><code class="lang-auto">...
cuda:1 2048
cuda:1 3072
...
</code></pre>
<p>Looks like it took reserves blocks of <code>1024</code></p>

zetyquickly: <p>Hello <a class="mention" href="/u/ptrblck">@ptrblck</a>, may we attract your attention to it?</p>

ptrblck: <p>You are most likely seeing the same effect described <a href="https://discuss.pytorch.org/t/memory-consumption-for-the-model-get-doubled-after-wrapped-with-ddp/130837/5">here</a>.</p>

jastern33: <p>Not sure how I missed this thread – this is exactly what I am seeing. Thanks for the help, I will now account for an extra copy of the gradients in memory when training with DDP.</p>