Why does ZeroRedundancyOptimizer consume less memory during training than FullyShardedDataParallel?

Renshuang_Zhou: <p>I trained a Transformer model with 1B parameters on servers with 8 A100 GPUs.</p>
<p>As far as I know, <code>ZeroRedundancyOptimizer</code> is based on ZeRO-1 and <code>FullyShardedDataParallel</code> is based on ZeRO-3, and <code>FSDP</code> should reduce more GPU memory consumption.</p>
<p>However, in reality, when using <code>ZeroRedundancyOptimizer</code>, the GPU memory consumption during training is lower.</p>
<p>Is this a normal phenomenon?</p>
<p>The pytorch version I am using is 2.0.1, the following is the code I used when testing:</p>
<pre><code class="lang-python"># FSDP
import torch
from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
from torch.distributed.fsdp.api import ShardingStrategy


device = torch.device('cuda')
sharded_module = FSDP(my_module, sharding_strategy=ShardingStrategy.FULL_SHARD)

optim = torch.optim.Adam(sharded_module.parameters(), lr=0.0001)
dummy_input = torch.randn((1, 3, 224, 224), device=device)

sharded_module.train()
x = sharded_module(dummy_input)
loss = x.sum()
loss.backward()
optim.step()
optim.zero_grad()


# ZeroRedundancyOptimizer
import torch
from torch.distributed.optim.zero_redundancy_optimizer import ZeroRedundancyOptimizer
from torch.nn.parallel import DataParallel, DistributedDataParallel


device = torch.device('cuda')
model = DistributedDataParallel(
    my_module, device_ids=[torch.cuda.current_device()],
            find_unused_parameters=True)

optim = ZeroRedundancyOptimizer(model.parameters(), torch.optim.Adam,
                                lr=0.0001)
dummy_input = torch.randn((1, 3, 224, 224), device=device)

model.train()
x = model(dummy_input)
loss = x.sum()
loss.backward()
optim.step()
optim.zero_grad()
</code></pre>

fduwjj: <p>When you use <code>ZeroRedundancyOptimizer</code> did you use DDP?</p>
<p>cc: <a class="mention" href="/u/agu">@agu</a></p>

Renshuang_Zhou: <p>Yes, I used DDP when I used ZeroRedundancyOptimizer.</p>

agu: <p>How are you using FSDP? Are you using nested FSDP wrapping?</p>

Renshuang_Zhou: <p>My code is already posted in the description, I am not using nested FSDP warpping.</p>

agu: <p>Sorry, I missed the code. If you do not use nested FSDP wrapping, then you will incur <code>1 + 1 / W</code> memory overhead from the parameters for <code>W</code> workers (due to FSDP always maintaining the <code>1 / W</code> shard), which could be why you see higher GPU memory usage with FSDP.</p>

Renshuang_Zhou: <p>Sorry, I did not understand your meaning. Are you saying that if FSDP wants to achieve lower GPU memory consumption, it needs nested warpping? Is there any sample code for reference?</p>

agu: <p>Yes, nested wrapping helps lower the GPU memory usage.</p>
<p>You can try something like:</p>
<pre><code class="lang-auto">from torch.distributed.fsdp.wrap import ModuleWrapPolicy
policy = ModuleWrapPolicy({TransformerBlock})
sharded_module = FSDP(my_module, auto_wrap_policy=policy, ...)
</code></pre>
<p>where you can replace <code>TransformerBlock</code> with whatever your transformer block’s module class is.</p>

Renshuang_Zhou: <p>Thanks, it does work! Is it true that ZeroRedundancyOptimizer is already nested wrapped, while FSDP requires setting a policy to achieve nested wrapping?</p>

agu: <p>The FSDP wrapping effectively determines the max parameter size that is unsharded (all-gathered) at once, so if you only wrap at the top-level, then all parameters contribute to that max size. Therefore, nested wrapping is necessary to decrease the max memory contribution from parameters.</p>
<p>For <code>ZeroRedundancyOptimizer</code>, the implementation is different. The optimizer states are greedily partitioned across ranks to minimize unevenness. It does not shard parameters or gradients.</p>

Renshuang_Zhou: <p>Thanks for your reply, it’s really helpful. But I still have some questions. Is ZeroRedundancyOptimizer an implementation of ZeRO-1? If so, why not continue implementing ZeRO-2 and ZeRO-3 based on ZeroRedundancyOptimizer instead of introducing FSDP separately? Is there a reason for this?</p>

agu: <p>One reason is that <code>ZeroRedundancyOptimizer</code> is a <code>torch.optim.Optimizer</code>, but in order to shard parameters/gradients, we need something that interfaces with <code>nn.Module</code>. Hence, <code>FullyShardedDataParallel</code> is an <code>nn.Module</code> (wrapper).</p>

Renshuang_Zhou: <p>I see, thanks for your patience.</p>

Renshuang_Zhou: <p>Hello, I still have a question that I don’t understand. When FSDP does not do nested warpping, the memory consumption is higher than ZeroRedundancyOptimizer, but FSDP actually includes ZeRO-1. Is this caused by the different implementation between FSDP and ZeroRedundancyOptimizer? Where is the extra memory consumption of FSDP reflected?</p>

agu: <p>It is an implementation detail of FSDP. See the description from above:</p>
<blockquote>
<p>If you do not use nested FSDP wrapping, then you will incur <code>1 + 1 / W</code> memory overhead from the parameters for <code>W</code> workers (due to FSDP always maintaining the <code>1 / W</code> shard), which could be why you see higher GPU memory usage with FSDP.</p>
</blockquote>

Renshuang_Zhou: <p>I understand, thank you!</p>