Distributed model saves weights different from current state dict?

username: Priyank_Pathak
cooked: <p>I’m using torch ‘1.13.1+cu117’ and if I run multi-gpu torch code</p>
<pre><code class="lang-auto">NUM_GPU=2
CUDA_VISIBLE_DEVICES=0,1 python -W ignore -m torch.distributed.launch --nproc_per_node=$NUM_GPU .... 
</code></pre>
<p>and if I do</p>
<pre><code class="lang-auto">import torch.distributed as dist
local_rank = dist.get_rank()
if local_rank ==0 :     
     torch.save(model.state_dict(), "temp.pth")
     for name, param in model.state_dict().items():
            print(name, param.mean().item())
</code></pre>
<p>I get values like</p>
<pre><code class="lang-auto">mlp.fc1_g.weight 2.249192903036601e-07
mlp.fc1_g.bias -0.00013029106776230037
mlp.fc1_x.weight -1.5291587942556362e-06
mlp.fc1_x.bias 0.00014961455599404871
mlp.norm.weight 1.0
mlp.norm.bias 0.0
mlp.fc2.weight -8.485401536972859e-08
mlp.fc2.bias 2.6364039513282478e-05
</code></pre>
<p>Now when I load the same weights on 1 GPU</p>
<pre><code class="lang-auto">checkpoint = torch.load("temp.pth", map_location='cpu')
model.load_state_dict( checkpoint , strict = True )
for name, param in model.state_dict().items():
    print(name, param.mean().item())
</code></pre>
<p>I get</p>
<pre><code class="lang-auto">mlp.fc1_g.weight 2.249192903036601e-07
mlp.fc1_g.bias -0.00013029106776230037
mlp.fc1_x.weight -1.5291587942556362e-06
mlp.fc1_x.bias 0.00014961455599404871
mlp.norm.weight 1.0
mlp.norm.bias 0.0
mlp.fc2.weight -8.485401536972859e-08
mlp.fc2.bias 2.6364039513282478e-05
</code></pre>
<p>There is a difference in decimal places of weights like <code>mlp.fc2.weight</code> , <code>mlp.fc1_x.weight</code>, <code>mlp.fc1_g.weight</code>.</p>
<p>As a result, validation accuracy during training and later testing the same model is different (up to 2-3%).</p>
<p>What am I doing wrong? Note: I have just posted a few weights of the entire transformers, but the weights differ across all layers.<br>
After training the model for 1 epoch, I have compared values across all weights <a href="https://www.diffchecker.com/QxKX87EW/" rel="noopener nofollow ugc">here</a></p>

username: Malik_Arshad
cooked: <p>This usually happens due to model wrapping (e.g., <code>DistributedDataParallel</code>), so use <code>model.module.state dict()</code> when saving or loading weights.</p>

username: Priyank_Pathak
cooked: <p>Turns out I didnt wrap the model to begin with <img src="https://discuss.pytorch.org/images/emoji/apple/sweat_smile.png?v=14" title=":sweat_smile:" class="emoji" alt=":sweat_smile:" loading="lazy" width="20" height="20"></p>
<p>doing</p>
<pre><code class="lang-auto">model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[local_rank],find_unused_parameters=True)
</code></pre>
<p>Fixed the problem!</p>
<p>Apparently <code>torch.cuda.synchronize()</code> doesnt work unless <code>DistributedDataParallel</code> and GPUs had different weights thus which ever rank saved the model (during validation) was evaluating on testing</p>

username: Malik_Arshad
cooked: <p>Yes, in distributed training, the saved weights may differ in format (e.g., wrapped in <code>module.</code>) due to <code>DataParallel</code> or <code>DistributedDataParallel</code>, requiring state dict key adjustment when loading.</p>