FSDP/HSDP with `device_mesh` multiple replica intra node

vasqu: <p>It’s somewhat related to <a href="https://github.com/pytorch/pytorch/issues/117470" class="inline-onebox" rel="noopener nofollow ugc">[FSDP] HYBRID_SHARD Apply FULL_SHARD across multiple nodes instead of just intar-node · Issue #117470 · pytorch/pytorch · GitHub</a> but in the opposite direction:</p>
<ul>
<li>Replicas within a node and across node(s)</li>
<li>Sharding within a node but only to a limited number of devices</li>
</ul>
<p>For example, if we had 2 nodes with 8 GPUs each, I’d like to have FSDP/HSDP with 4GPUs for sharding and 4 replicas (2 within a node). I’d assume that setting <code>device_mesh</code> to 4x4 would (maybe?) do this. So, my question is the following I guess:</p>
<ul>
<li>Is it guaranteed to init sharding within a node first if I create a device mesh, i.e. sharding inter nodes is avoided?</li>
</ul>

irisz: <p><a class="mention" href="/u/vasqu">@vasqu</a></p>
<p>When you create a 2D DeviceMesh and pass it to HSDP, it always assumes the outer dimension is the replicate group and the inner dimension is the shard group. Under the hood, it is just using the process group (code pointer here: <a href="https://github.com/pytorch/pytorch/blob/main/torch/distributed/fsdp/_init_utils.py#L162-L167" class="inline-onebox" rel="noopener nofollow ugc">pytorch/torch/distributed/fsdp/_init_utils.py at main · pytorch/pytorch · GitHub</a>). DeviceMesh is just helping initialize the process group easily.</p>

vasqu: <p><a class="mention" href="/u/irisz">@irisz</a></p>
<p>Sorry, just for clarification as I’m not familiar with the device mesh initialization (/abstraction), I’d expand on my example. Assuming I have node 1 with ranks 0 to 7 (8 in total) and node 2 with ranks 8 to 15 (also 8 in total), I’d then expect with a (4,4) device mesh to get the following pgs:</p>
<ul>
<li>[0, 1, 2, 3], [4, 5, 6, 7], [8, 9, 10, 11], [12, 13, 14, 15] (sharding pgs)</li>
<li>[0, 4, 8, 12], [1, 5, 9, 13], [2, 6, 10, 14], [3, 7, 11, 15] (replication pgs)</li>
</ul>
<p>Is that correct? Thanks in advance!</p>
<p>Edit: In essence, I want to achieve this:<br>
<div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/c/3/c3a96d48c18d37f15aec919de0a222634d2440db.png" data-download-href="https://discuss.pytorch.org/uploads/default/c3a96d48c18d37f15aec919de0a222634d2440db" title="Screenshot from 2024-09-06 12-55-29"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/c/3/c3a96d48c18d37f15aec919de0a222634d2440db_2_401x375.png" alt="Screenshot from 2024-09-06 12-55-29" data-base62-sha1="rUU6kLX6ucCunTArtzQLtBabkFR" width="401" height="375" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/c/3/c3a96d48c18d37f15aec919de0a222634d2440db_2_401x375.png, https://discuss.pytorch.org/uploads/default/optimized/3X/c/3/c3a96d48c18d37f15aec919de0a222634d2440db_2_601x562.png 1.5x, https://discuss.pytorch.org/uploads/default/original/3X/c/3/c3a96d48c18d37f15aec919de0a222634d2440db.png 2x" data-dominant-color="F0E9E7"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">Screenshot from 2024-09-06 12-55-29</span><span class="informations">678×633 28.6 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div><br>
Maybe the graphic can help here.</p>

irisz: <p>Yes. This is right. You would get this:</p>
<ul>
<li>[0, 1, 2, 3], [4, 5, 6, 7], [8, 9, 10, 11], [12, 13, 14, 15] (sharding pgs)</li>
<li>[0, 4, 8, 12], [1, 5, 9, 13], [2, 6, 10, 14], [3, 7, 11, 15] (replication pgs)</li>
</ul>
<p>To get this, you would want to initialize a DeviceMesh like this:</p>
<pre><code class="lang-auto">mesh_2d = init_device_mesh("cuda", (4, 4), mesh_dim_names=("replicate", "shard"))
</code></pre>
<p>Here is a bit information: <a href="https://pytorch.org/tutorials/recipes/distributed_device_mesh.html" class="inline-onebox" rel="noopener nofollow ugc">Getting Started with DeviceMesh — PyTorch Tutorials 2.4.0+cu121 documentation</a></p>

vasqu: <p>Thank you! I think I got it now.</p>