Calling new_group cost long time when training model with large cluster，such as using 5000+ gpus

username: xuexiao1987
cooked: <p>Calling new_group cost long time when training model with large cluster， using 5000+ gpus， cost 30 minutes.<br>
Having the configs as follows:<br>
1,TP=8，PP=12<br>
2, torch2.0<br>
3, torch.distributed.init_process_group env init<br>
4, Megatron mpu.initialize_model_parallel cost almost 30 minutes, calling plenty of new_group<br>
we see the feature in version torch2.1, wether the feature fix new_group?</p><aside class="onebox githubcommit" data-onebox-src="https://github.com/pytorch/pytorch/commit/ad21890f8fab73a15e758c7b893e129e9db1a81a">
  <header class="source">

      <a href="https://github.com/pytorch/pytorch/commit/ad21890f8fab73a15e758c7b893e129e9db1a81a" target="_blank" rel="noopener nofollow ugc">github.com/pytorch/pytorch</a>
  </header>

  <article class="onebox-body">
    <div class="github-row">
  <div class="github-icon-container" title="Commit">
    <svg width="60" height="60" class="github-icon" viewBox="0 0 14 16" aria-hidden="true"><path fill-rule="evenodd" d="M10.86 7c-.45-1.72-2-3-3.86-3-1.86 0-3.41 1.28-3.86 3H0v2h3.14c.45 1.72 2 3 3.86 3 1.86 0 3.41-1.28 3.86-3H14V7h-3.14zM7 10.2c-1.22 0-2.2-.98-2.2-2.2 0-1.22.98-2.2 2.2-2.2 1.22 0 2.2.98 2.2 2.2 0 1.22-.98 2.2-2.2 2.2z"></path></svg>
  </div>

  <div class="github-info-container">
    <h4>
      <a href="https://github.com/pytorch/pytorch/commit/ad21890f8fab73a15e758c7b893e129e9db1a81a" target="_blank" rel="noopener nofollow ugc">[c10d] Scalable PG initiation. (#99931)</a>
    </h4>

    <div class="github-info">
      <div class="date">
        committed <span class="discourse-local-date" data-format="ll" data-date="2023-04-27" data-time="13:44:02" data-timezone="UTC">01:44PM - 27 Apr 23 UTC</span>
      </div>

      <div class="user">
        <a href="https://github.com/kumpera" target="_blank" rel="noopener nofollow ugc">
          <img alt="kumpera" src="https://avatars.githubusercontent.com/u/5738?v=4" class="onebox-avatar-inline" width="20" height="20">
          kumpera
        </a>
      </div>

      <div class="lines" title="changed 4 files with 225 additions and 40 deletions">
        <a href="https://github.com/pytorch/pytorch/commit/ad21890f8fab73a15e758c7b893e129e9db1a81a" target="_blank" rel="noopener nofollow ugc">
          <span class="added">+225</span>
          <span class="removed">-40</span>
        </a>
      </div>
    </div>
  </div>
</div>

  <div class="github-row">
    <p class="github-body-container">Add use_local_synchronization argument to new_group.

When this argument is True<span class="show-more-container"><a href="https://github.com/pytorch/pytorch/commit/ad21890f8fab73a15e758c7b893e129e9db1a81a" target="_blank" rel="noopener nofollow ugc" class="show-more">…</a></span><span class="excerpt hidden">, is change new_group to do a store_barrier only on the ranks that are park of the group and not the whole cluster.

This addressess both scalability and composability problems associated with new_group.

Fixes #81291.

This is relanding #84224
As part of the original PR I did a quick benchmark of creating 3 PGs per rank using both functions and perf is the following:

new_group use_local_synchronization=False:
| World Size | Time (in secs) |
| --- | ----------- |
| 4 | 0.12 |
| 8 | 0.25 |
| 16 | 0.51 |
| 32 | 0.87 |
| 64 | 1.50 |
| 128 | 2.87 |

new_group use_local_synchronization=True:
| World Size | Time (in secs) |
| --- | ----------- |
| 4 | 0.05 |
| 8 | 0.04 |
| 16 | 0.03 |
| 32 | 0.03 |
| 64 | 0.04 |
| 128 | 0.04 |

Scaling for `use_local_synchronization=False` is sub linear because the number of process groups created as a multiple of world_size decreases as we go up. It's 6 with world_size 4 and 192 with world_size 128.

Scaling for `use_local_synchronization=True` is constant as the number of store barriers executed per rank remains constant at 3.

Setup:

1 AWS host, backend gloo.

Pull Request resolved: https://github.com/pytorch/pytorch/pull/99931
Approved by: https://github.com/xw285cornell</span></p>
  </div>

  </article>

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

  <div style="clear: both"></div>
</aside>


username: ptrblck
cooked: <p>Did you profile what exactly takes this amount of time?</p>

username: xuexiao1987
cooked: <p>Not yet，profiling on small cluster is ongoing，large cluster is in busy use and has no chance to proflie owing to need rebuild torch.we only profile  this function mpu.initialize_model_parallel in megatron, and it cost 30 minutes when restart training.</p>