Can't load checkpoint in HSDP, stuck at synchronization in `optim_state_dict_to_load`

MFajcik1: <p>Hi, I am trying to load my checkpoint in HSDP (HYBRID_SHARD) mode. Training runs fine, and checkpoint saving runs fine. However there is a problem with checkpoint loading, it blocks forever on <code>optim_state_dict_to_load</code> (or until some long timeout). I am using torch 2.1.2, and the code works fine with FULL_SHARD FSDP.</p>
<pre><code class="lang-auto">from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
from torch.distributed.fsdp.fully_sharded_data_parallel import FullOptimStateDictConfig

fsdp_state_dict_type = StateDictType.FULL_STATE_DICT
state_dict_config = FullStateDictConfig(offload_to_cpu=True, rank0_only=True)
optim_state_dict_config = FullOptimStateDictConfig(offload_to_cpu=True, rank0_only=True)

with FSDP.state_dict_type(module,
                                  state_dict_type=fsdp_state_dict_type,
                                  state_dict_config=state_dict_config,
                                  optim_state_dict_config=optim_state_dict_config):
  optim_state_dict = FSDP.optim_state_dict_to_load(  # blocks on this line
                              optim_state_dict=optim_state_dict, model=self.model, optim=optimizer)
</code></pre>
<p>I am not the author of the code, this is part of the composer library, <a href="https://github.com/mosaicml/composer/blob/7e0e40a72a6b0c87cca4d408792336e4e5c742da/composer/core/state.py#L1173" rel="noopener nofollow ugc">right here</a> (In this example, I explicitly expanded the code variables, and context managers for easier understanding).</p>
<p>Since this works in FULL_SHARD, I figured out that there is likely a problem with waiting for processes from different node, which the HSDP should not wait for (sharding with HSDP here is done only on 1 node, not across nodes).</p>
<p>To debug this, I used 2 nodes, each with 4 GPUs.<br>
I printed the following information about process groups</p>
<pre><code class="lang-auto">print(f"TOTAL PGs ranks: {str(torch.distributed.distributed_c10d._pg_group_ranks)}")
print(f"TAGS TO PG: {str(torch.distributed.distributed_c10d._tags_to_pg)}")
</code></pre>
<p>outputs on node 0</p>
<pre><code class="lang-auto">TOTAL PGs ranks: {
&lt;torch.distributed.distributed_c10d.ProcessGroup object at 0x2b971ac009f0&gt;: {0: 0, 1: 1, 2: 2, 3: 3, 4: 4, 5: 5, 6: 6, 7: 7}, 
&lt;torch.distributed.distributed_c10d.ProcessGroup object at 0x2b975cc2d5b0&gt;: {0: 0, 1: 1, 2: 2, 3: 3}, 
-100: {3: 0, 7: 1}, 
&lt;torch.distributed.distributed_c10d.ProcessGroup object at 0x2b975cc574f0&gt;: {0: 0, 4: 1}}

TAGS TO PG: {
'': [
&lt;torch.distributed.distributed_c10d.ProcessGroup object at 0x2b971ac009f0&gt;, 
&lt;torch.distributed.distributed_c10d.ProcessGroup object at 0x2b975cc2d5b0&gt;, 
&lt;torch.distributed.distributed_c10d.ProcessGroup object at 0x2b975cc574f0&gt;], 
'ptd:0': [&lt;torch.distributed.distributed_c10d.ProcessGroup object at 0x2b971ac009f0&gt;], 
'ptd:1': [&lt;torch.distributed.distributed_c10d.ProcessGroup object at 0x2b975cc2d5b0&gt;], 
'ptd:3': [&lt;torch.distributed.distributed_c10d.ProcessGroup object at 0x2b975cc574f0&gt;]}
</code></pre>
<p>outputs on node 1</p>
<pre><code class="lang-auto">TOTAL PGs ranks: {
&lt;torch.distributed.distributed_c10d.ProcessGroup object at 0x2b970d06e6b0&gt;: {0: 0, 1: 1, 2: 2, 3: 3, 4: 4, 5: 5, 6: 6, 7: 7},
-100: {3: 0, 7: 1}, 
&lt;torch.distributed.distributed_c10d.ProcessGroup object at 0x2b971bcf5670&gt;: {4: 0, 5: 1, 6: 2, 7: 3}, &lt;torch.distributed.distributed_c10d.ProcessGroup object at 0x2b971bdda5b0&gt;: {0: 0, 4: 1}}

TAGS TO PG: {
'': [
&lt;torch.distributed.distributed_c10d.ProcessGroup object at 0x2b970d06e6b0&gt;, 
&lt;torch.distributed.distributed_c10d.ProcessGroup object at 0x2b971bcf5670&gt;, 
&lt;torch.distributed.distributed_c10d.ProcessGroup object at 0x2b971bdda5b0&gt;], 
'ptd:0': [&lt;torch.distributed.distributed_c10d.ProcessGroup object at 0x2b970d06e6b0&gt;], 
'ptd:2': [&lt;torch.distributed.distributed_c10d.ProcessGroup object at 0x2b971bcf5670&gt;], 
'ptd:3': [&lt;torch.distributed.distributed_c10d.ProcessGroup object at 0x2b971bdda5b0&gt;]}
</code></pre>
<p>I guessed the single node group is <code>ptd:1</code> for rank 0, and <code>ptd:2</code> for rank 1, and passed these to <code>optim_state_dict_to_load</code> via <code>pg</code> parameter (if <code>pg</code> is not passed, default group is considered, which was <code>ptd:0</code>. This one contains all 8 GPUs, so “it made sense” that this could 've caused the block). This worked for node 0 (or at least, it didn’t block), but failed on node 2 raising this</p><aside class="onebox githubblob" data-onebox-src="https://github.com/pytorch/pytorch/blob/64efec9953cebfeef2185eaec626ce0db37c6aa6/torch/distributed/distributed_c10d.py#L790">
  <header class="source">

      <a href="https://github.com/pytorch/pytorch/blob/64efec9953cebfeef2185eaec626ce0db37c6aa6/torch/distributed/distributed_c10d.py#L790" target="_blank" rel="noopener nofollow ugc">github.com</a>
  </header>

  <article class="onebox-body">
    <h4><a href="https://github.com/pytorch/pytorch/blob/64efec9953cebfeef2185eaec626ce0db37c6aa6/torch/distributed/distributed_c10d.py#L790" target="_blank" rel="noopener nofollow ugc">pytorch/pytorch/blob/64efec9953cebfeef2185eaec626ce0db37c6aa6/torch/distributed/distributed_c10d.py#L790</a></h4>



    <pre class="onebox"><code class="lang-py">
      <ol class="start lines" start="780" style="counter-reset: li-counter 779 ;">
          <li>        Group rank of ``global_rank`` relative to ``group``</li>
          <li></li>
          <li>    N.B. calling this function on the default process group returns identity</li>
          <li>    """</li>
          <li>    if group is GroupMember.WORLD:</li>
          <li>        return global_rank</li>
          <li>    if group not in _world.pg_group_ranks:</li>
          <li>        raise ValueError(f"Group {group} is not registered, please create group with torch.distributed.new_group API")</li>
          <li>    group_ranks = _world.pg_group_ranks[group]</li>
          <li>    if global_rank not in group_ranks:</li>
          <li class="selected">        raise ValueError(f"Global rank {global_rank} is not part of group {group}")</li>
          <li></li>
          <li>    return group_ranks[global_rank]</li>
          <li></li>
          <li>def get_global_rank(group: ProcessGroup, group_rank: int) -&gt; int:</li>
          <li>    """</li>
          <li>    Translate a group rank into a global rank.</li>
          <li></li>
          <li>    ``group_rank`` must be part of `group` otherwise this raises RuntimeError.</li>
          <li></li>
          <li>    Args:</li>
      </ol>
    </code></pre>



  </article>

  <div class="onebox-metadata">
    
    
  </div>

  <div style="clear: both"></div>
</aside>
<p>
with a message that global rank 0 is not part of the group (which is not, as node 1 group has global ranks 4-7).</p>
<p>Any ideas of what could be the issue, or what could I try? I spent days on this one <img src="https://discuss.pytorch.org/images/emoji/apple/frowning.png?v=12" title=":frowning:" class="emoji" alt=":frowning:" loading="lazy" width="20" height="20"> .</p>

agu: <p>Sorry that this was a tough bug! Would it be easier to submit an issue to the composer repo? There have been some fixes recently in PyTorch, but I am not sure off the top of my head whether this is related.</p>