Alternating Parameters in DDP

ziqipang: <p>Hi Guys,</p>
<p>My question and context might a bit long, but I think it will help a lot of people trying to train <strong>large models</strong> with <strong>limited GPU memory</strong> under a <strong>DDP setup</strong>. After trying a lot of stuff (flash-attention, fp16, etc.), I am trying to leverage the property that <strong>setting</strong> <code>requires_grad=False</code> <strong>saves GPU memory</strong>.</p>
<p>A most intuitive description of what I want is that:</p>
<ul>
<li>I have a neural network with two layers <code>layer1</code> and <code>layer2</code>.</li>
<li>I use two separate optimizers: <code>opt1</code> for <code>layer1</code> and <code>opt2</code> for <code>layer2</code>.</li>
<li>During the training process, I <strong>alternate between the two optimizers and set the layers with</strong> <code>requires_grad=False</code>.</li>
</ul>
<p>Then <strong>the following code works perfectly well under the single-GPU use-case</strong> and squeeze the training right below the GPU memory:</p>
<pre><code class="lang-auto"># initialization
net = ...
opt1, opt2 = ..., ...

# training
for step, data in enumerate(data_loader):
    # alternating the parameter sets
    if step % 2 == 0:
        for param in net.layer1.parameters():
            param.requires_grad = True
            param.grad = torch.zeros_like(param.data)
        for param in net.layer2.parameters():
            param.requires_grad = False
            del param.grad
        opt = opt1
    else:
        for param in net.layer2.parameters():
            param.requires_grad = True
            param.grad = torch.zeros_like(param.data)
        for param in net.layer1.parameters():
            param.requires_grad = False
            del param.grad
        opt = opt2
    
    loss = net(data)
    loss.backward()
    opt.step()
    opt.zero_grad()
</code></pre>
<p>Then I try to switch to multi-GPU training with DDP as below, adding something supporting the DDP. Please note that I wrap DDP after each alternation of parameter set so that DDP can correctly register and reduce the gradients.</p>
<pre><code class="lang-auto"># initialization
net = ...
ddp_net = DDP(model)
opt1, opt2 = ..., ...

for step, data in enumerate(data_loader):
    unwrap_net = unwrap(ddp_net) # remove the ddp wrapper
    # alternating the parameter sets
    if step % 2 == 0:
        for param in unwrap_net.layer1.parameters():
            param.requires_grad = True
            param.grad = torch.zeros_like(param.data)
        for param in unwrap_net.layer2.parameters():
            param.requires_grad = False
            del param.grad
        opt = opt1
    else:
        for param in unwrap_net.layer2.parameters():
            param.requires_grad = True
            param.grad = torch.zeros_like(param.data)
        for param in unwrap_net.layer1.parameters():
            param.requires_grad = False
            del param.grad
        opt = opt2
    
    new_ddp_net = DDP(unwrap_net)
    del ddp_net
    ddp_net = new_ddp_net
    torch.cuda.empty_cache()
    
    loss = ddp_net(data)
    loss.backward()
    opt.step()
    opt.zero_grad()
</code></pre>
<p>Here’s the question: I found that <strong>[alternating the parameter set] occupies more GPU memory than [training only layer1 or layer2 alone with DDP]</strong>. Although it is just 1-2GB, I can no longer squeeze the large model into my GPU anymore ( <img src="https://discuss.pytorch.org/images/emoji/apple/sob.png?v=12" title=":sob:" class="emoji" alt=":sob:" loading="lazy" width="20" height="20"> <img src="https://discuss.pytorch.org/images/emoji/apple/sob.png?v=12" title=":sob:" class="emoji" alt=":sob:" loading="lazy" width="20" height="20"> <img src="https://discuss.pytorch.org/images/emoji/apple/sob.png?v=12" title=":sob:" class="emoji" alt=":sob:" loading="lazy" width="20" height="20">). So I am wondering:</p>
<ul>
<li>Are there any place I am doing wrong that causes this issue?</li>
<li>Are there better approaches to alternating the parameter sets under a DDP setting?</li>
<li>Considering that I am using huggingface’s <code>accelerate</code>, is my observed issue not a pytorch one, actually?</li>
</ul>
<p>If you have read through here, my friend, I am truly grateful for your patience and help! May the force of optimization be with you!</p>
<p>Best,</p>
<p>Ziqi</p>