Strange behavior of HSDP

yenchenlin: <p>I followed <a href="https://pytorch.org/tutorials/recipes/distributed_device_mesh.html" rel="noopener nofollow ugc">the tutorial on DeviceMesh</a> to test the following HSDP code. When I printed the weights, I expected to see the same values 4 times. However, instead, I see two sets of repeated values. Does this mean HSDP did not sync the weights along the replicate_group axis?</p>
<pre><code class="lang-auto">import torch
import torch.nn as nn

from torch.distributed.device_mesh import init_device_mesh
from torch.distributed.fsdp import FullyShardedDataParallel as FSDP, ShardingStrategy


class ToyModel(nn.Module):
    def __init__(self):
        super(ToyModel, self).__init__()
        self.net1 = nn.Linear(10, 10)
        self.relu = nn.ReLU()
        self.net2 = nn.Linear(10, 5)

    def forward(self, x):
        return self.net2(self.relu(self.net1(x)))


# HSDP: MeshShape(2, 2)
mesh_2d = init_device_mesh("cuda", (2, 2), mesh_dim_names=("replicate", "shard"))
model = FSDP(
    ToyModel().cuda(), 
    device_mesh=mesh_2d, 
    sharding_strategy=ShardingStrategy.HYBRID_SHARD,
    use_orig_params=True
)

with FSDP.summon_full_params(model, with_grads=True):
    print(model.net1.weight.data.sum())
</code></pre>

irisz: <p>From the example, we have the mesh topology of [[0, 1], [2, 3]]:</p>
<ul>
<li>[0, 1] and [2, 3] are the replicate group</li>
<li>[0, 2] and [1, 3] are the shard group</li>
</ul>
<p>So we should see two sets of the same value. Below is the printout from a sample run:</p>
<pre><code class="lang-auto">local_rank=0, model.net1.weight.data.sum()=tensor(-0.4711, device='cuda:0')
local_rank=1, model.net1.weight.data.sum()=tensor(-0.4711, device='cuda:1')
local_rank=2, model.net1.weight.data.sum()=tensor(0.2237, device='cuda:2')
local_rank=3, model.net1.weight.data.sum()=tensor(0.2237, device='cuda:3')
</code></pre>

eli0: <p>My team has been experimenting with something similar, and facing some inconsistencies. Our toy setup includes 2 nodes, with 4 GPUs each and 2 replicas total. That makes for 4 shards per replica (i.e. 1 replica per node). In essence, we used the same code as <a class="mention" href="/u/yenchenlin">@yenchenlin</a>, except we had <code>init_device_mesh("cuda", (2, 4))</code> instead of (2,2).</p>
<p>We’re noticing the same thing as <a class="mention" href="/u/yenchenlin">@yenchenlin</a> when we initialize the models differently before wrapping with FSDP, then wrap with FSDP and print the <code>net1</code> sum, we see 2 unique values being printed 4 times. This makes us wonder if HSDP is really syncing weights, or not. However, if we fix our seed to be consistent across all ranks before initializing our toy model, then the printout of <code>model.net1.weight.data.sum()</code> within the <code>FSDP.summon_full_params()</code> context is the same across all ranks.</p>
<p>It was also our understanding that <code>FSDP.summon_full_params()</code> enabled each rank to reach the respective parameters of its entire replica. And thus, the ranks that report the same sum would actually be part of the same replica, but this understanding doesn’t seem to align with your printouts <a class="mention" href="/u/irisz">@irisz</a> , unless my understanding of “replicate” and “shard” groups is incorrect (which was that a shard group makes up complete replica of the model).</p>
<p>Can you provide us with some info (and sources if available) on whether or not PyTorch HSDP syncs weights across replicas or not?</p>