ZERO optimizer.consolidate_state_dict() will hang

FreedWu: <p>I try to use ZERO, however it hang when <code>optimizer.consolidate_state_dict</code>.</p>
<p>Why this code will hang? Is there something silly I did? TIA!</p>
<pre><code class="lang-python">r"""Run ``torchrun --nproc_per_node=2 test.py``
"""
import torch
from torch import distributed as dist
from torchvision.models import resnet18
from torch import nn
from torch.optim import Adam
from torch.distributed.optim.zero_redundancy_optimizer import (
    ZeroRedundancyOptimizer,
)


def main() -&gt; None:
    dist.init_process_group("nccl")
    device = torch.device(dist.get_rank())
    model = nn.parallel.DistributedDataParallel(resnet18().to(device))
    loss_func = nn.CrossEntropyLoss()
    optimizer = ZeroRedundancyOptimizer(model.parameters(), Adam)
    samples = torch.randn(1, 3, 64, 64).to(device)
    targets = torch.randint(1000, (1,)).to(device)
    outputs = model(samples)
    loss = loss_func(outputs, targets)
    loss.backward()
    optimizer.step()
    optimizer.zero_grad(True)
    optimizer.consolidate_state_dict(dist.get_rank())  # will hang
    torch.save({"optimizer": optimizer.state_dict(), }, "test.pth")

if __name__ == "__main__":
    main()
</code></pre>

wanchaol: <p>cc <a class="mention" href="/u/rvarm1">@rvarm1</a> <a class="mention" href="/u/agu">@agu</a> about ZERO optimizer</p>

agu: <aside class="quote no-group" data-username="FreedWu" data-post="1" data-topic="171908">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/freedwu/48/36010_2.png" class="avatar"> FreedWu:</div>
<blockquote>
<p><code>optimizer.consolidate_state_dict(dist.get_rank())  # will hang</code></p>
</blockquote>
</aside>
<p>The argument to <code>consolidate_state_dict()</code> is the target rank (which is 0 by default). If you want to save the unsharded state dict on all ranks, then you need to loop over all ranks as the target:</p>
<pre><code class="lang-auto">for rank in dist.get_world_size():
    optimizer.consolidate_state_dict(rank)
torch.save({"optimizer": optimizer.state_dict(), }, "test.pth")
</code></pre>
<p>However, that will redundantly save the same state dict from all ranks. If that is not desired, you can do something like:</p>
<pre><code class="lang-auto">optimizer.consolidate_state_dict()  # default to rank 0
if dist.get_rank() == 0:
    torch.save({"optimizer": optimizer.state_dict(), }, "test.pth")
</code></pre>

FreedWu: <p>It still hang:</p>
<pre><code class="lang-python">r"""Run ``torchrun --nproc_per_node=2 test.py``
"""
import torch
from torch import distributed as dist
from torchvision.models import resnet18
from torch import nn
from torch.optim import Adam
from torch.distributed.optim.zero_redundancy_optimizer import (
    ZeroRedundancyOptimizer,
)


def main() -&gt; None:
    dist.init_process_group("nccl")
    device = torch.device(dist.get_rank())
    model = nn.parallel.DistributedDataParallel(resnet18().to(device))
    loss_func = nn.CrossEntropyLoss()
    optimizer = ZeroRedundancyOptimizer(model.parameters(), Adam)
    samples = torch.randn(1, 3, 64, 64).to(device)
    targets = torch.randint(1000, (1,)).to(device)
    outputs = model(samples)
    loss = loss_func(outputs, targets)
    loss.backward()
    optimizer.step()
    optimizer.zero_grad(True)
    for rank in range(dist.get_world_size()):
        optimizer.consolidate_state_dict(rank)
    if dist.get_rank() == 0:
        torch.save({"optimizer": optimizer.state_dict()}, "test.pth")

if __name__ == "__main__":
    main()
</code></pre>

agu: <p>My bad. Every time you call <code>optimizer.consolidate_state_dict(rank)</code>, it clears the rank’s previously consolidated state dicts, so the loop <code>for rank in range(dist.get_world_size())</code> was actually clearing rank 0’s consolidated state dict.</p>
<p>This should work (I checked this one locally):</p>
<pre><code class="lang-auto">optimizer.consolidate_state_dict(0)
if dist.get_rank() == 0:
    torch.save(...)
</code></pre>
<p>Or, if you want to save on all ranks, then you should call <code>torch.save()</code> in the loop (I did not check this one):</p>
<pre><code class="lang-auto">for rank in range(dist.get_world_size()):
    optimizer.consolidate_state_dict(rank)
    if dist.get_rank() == rank:
        torch.save(...)
</code></pre>