How to use torch.distributed.optim.ZeroRedundancyOptimizer with overlap_with_ddp=True?

t4rf9: <p>The optional parameter <code>overlap_with_ddp</code> is not well illustrated in the document.<br>
I don’t understand the requirement in the document:</p>
<blockquote>
<p>(2) registering a DDP communication hook constructed from one of the functions in <code>ddp_zero_hook.py</code></p>
</blockquote>
<p>So I just ignored it.</p>
<p>When I was trying to test it, a warning came:</p>
<blockquote>
<p>WARNING:root:<code>step()</code> should not be included in the training loop when <code>overlap_with_ddp=True</code></p>
</blockquote>
<p>Then after I removed <code>optimizer.step()</code>, the warning disappeared, but the parameters are not optimized, or fixed.</p>
<p>I am confused about how to make use of this parameter properly.</p>
<p>Is it possible to add an example to explain this?</p>

Yanli_Zhao: <p>cc <a class="mention" href="/u/awgu">@awgu</a> for using Zero optimizer</p>

agu: <p>Apologies for the delay and lack of documentation. The feature was added experimentally and not revisited.</p>
<p>I think the reason you are not seeing the parameters being optimized is because you skipped that step labeled as (2). You have to register one of the two DDP communication hooks in <code>ddp_zero_hook.py</code>: namely, either <code>hook_with_zero_step()</code> or <code>hook_with_zero_step_interleaved()</code>.</p>
<p>Let me provide some pseudocode for how you might use this:</p>
<pre><code class="lang-auto">from torch.distributed.algorithms.ddp_comm_hooks.ddp_zero_hook import (
    hook_with_zero_step,
    hook_with_zero_step_interleaved,
)
from torch.distributed.algorithms.ddp_comm_hooks.default_hooks import (
    allreduce_hook,
)
from torch.distributed.optim import ZeroRedundancyOptimizer
from torch.nn.parallel import DistributedDataParallel as DDP

model: nn.Module = ...  # define a model
ddp_model = DDP(model, device_ids=[rank])
zero_optim = ZeroRedundancyOptimizer(ddp_model.parameters(), optimizer_class=torch.optim.SGD, overlap_with_ddp=True, lr=1e-3, momentum=0.9)
hook = hook_with_zero_step(allreduce_hook, ddp_model, zero_optim)
ddp_model.register_comm_hook(state=None, hook=hook)
# The first two iterations are warmup and do nothing
for i in range(2):
    output = ddp_model(dummy_input)
    loss = output.sum()  # compute arbitrary loss
    loss.backward()  # compute arbitrary gradients but no optimizer update
# Now we may train overlapping DDP with ZeRO
for (input, label) in dataloader:
    zero_optim.zero_grad()  # need to clear the arbitrary gradients from above
    output = ddp_model(input)
    loss = loss_fn(output, label)  # compute actual loss
    loss.backward()  # compute actual gradients followed by an optimizer update
</code></pre>
<p>I hope this helps, and let me know if you have followup questions.</p>

vgoklani: <p>Hey there <a class="mention" href="/u/agu">@agu</a></p>
<p>We followed the example above, but it failed because we used two parameter groups (one for <code>decay</code> and one for <code>no-decay</code>).</p>
<p>The exact error was:</p>
<pre><code>RuntimeError: Specifying `params_per_rank` only supports a single parameter group
</code></pre>
<p>How do we modify the <code>pseudocode</code> to support multiple parameter groups?</p>
<p>The <code>ZeroRedundancyOptimizer</code> worked correctly with multiple parameter groups when we disabled <code>overlap_with_ddp</code>, but we need <em>ludicrous</em> speed <img src="https://discuss.pytorch.org/images/emoji/apple/wink.png?v=12" title=":wink:" class="emoji" alt=":wink:" loading="lazy" width="20" height="20"> thanks!</p>
<p>Thanks!</p>

agu: <p>Unfortunately, we are not actively developing <code>ZeroRedundancyOptimizer</code> anymore, so supporting multiple parameter groups is unlikely.</p>
<p>You can consider switching your DDP setup to PyTorch’s <code>FullyShardedDataParallel</code> with <code>ShardingStrategy.SHARD_GRAD_OP</code>, which supports multiple parameter groups with <code>use_orig_params=True</code>.</p>