Question about loading the model that was trained using 4GPU with distributed dataparallel to only 1 GPU job

zhengrchan: <p>Hi <img src="https://discuss.pytorch.org/images/emoji/apple/smiley.png?v=9" title=":smiley:" class="emoji" alt=":smiley:">  , I’m having trouble loading the distributed dataparallel model to just 1 GPU. And I want to know how to load the model (trained by 4 GPU with distributed dataparallel) to another job using only 1 GPU.</p>
<p>I have trained a model using 4 GPU and distributed dataparallel, and I saved it as the tutorial:<br>
<a href="https://pytorch.org/tutorials/intermediate/ddp_tutorial.html#save-and-load-checkpoints" class="onebox" target="_blank" rel="nofollow noopener">https://pytorch.org/tutorials/intermediate/ddp_tutorial.html#save-and-load-checkpoints</a><br>
However, I don’t know how to load it using just 1 GPU for some simple job like validation test.</p>
<pre><code class="lang-auto">if rank == 0:
   torch.save(ddp_model.state_dict(), CHECKPOINT_PATH)
dist.barrier()
</code></pre>
<p>I’m now using this method:</p>
<pre><code class="lang-auto">    # initialize
    torch.distributed.init_process_group(backend="nccl")
    local_rank = torch.distributed.get_rank()
    print(local_rank)
    torch.cuda.set_device(local_rank)
    device = torch.device("cuda", local_rank)
    print(device)
    
    # only gpu with rank0 can remain running

    model = resnet50()
    model.to(device)
    model = torch.nn.parallel.DistributedDataParallel(model, 
                                                        device_ids=[local_rank],
                                                        output_device=local_rank)
        
    model.load_state_dict(torch.load(cfg.MODEL.pretrained_model_path))
    model.eval()
    if local_rank == 0:
        acc, acc_std, th = lfw_test(model, cfg.TEST.lfw_root, cfg.TEST.lfw_test_list)
</code></pre>
<p>and the command code:</p>
<pre><code class="lang-auto">CUDA_VISIBLE_DEVICES=0,1,2,3 python -m torch.distributed.launch --nproc_per_node=4 test.py
</code></pre>
<p>This method works, and through <code>nvidia-smi</code> I saw that only GPU0 is working, but when I run another test process using GPU device 1 (when the previous one is still running):</p>
<pre><code class="lang-auto">CUDA_VISIBLE_DEVICES=1,2,3 python -m torch.distributed.launch --nproc_per_node=3 test.py
</code></pre>
<p>The previous process throw a runtime error:</p>
<pre><code class="lang-auto">RuntimeError: NCCL error in: /opt/conda/conda-bld/pytorch_1549633347309/work/torch/lib/c10d/ProcessGroupNCCL.cpp:260, unhandled system error
</code></pre>
<p>There are 2 reasons for me to load the model using 1GPU:</p>
<ol>
<li>Some jobs have file-writing part and distributed parallel may cause wrong order.</li>
<li>Running 4 tiny experiment with 1 GPU per process is more efficient for me to test my idea and finding bugs.</li>
</ol>
<p>So is there a way that I can load the model like the common ways :<code>torch.load_state_dict(torch.load()).to(torch.device("cuda:0))</code>? <img src="https://discuss.pytorch.org/images/emoji/apple/smile.png?v=9" title=":smile:" class="emoji" alt=":smile:"></p>

zhengrchan: <p>Oh… the first method is not work, neither. I find that using:</p>
<pre><code class="lang-auto">if local_rank == 0:
    output = model(input)
</code></pre>
<p>the <code>model(input)</code> will never output. And the code is just blocked there.</p>

ptrblck: <p>I’m not sure to understand the use case.<br>
It seems you would like to load the <code>state_dict</code> to a single GPU machine, but in your code you are wrapping the model again in DDP.</p>
<p>Would creating the model, loading the <code>state_dict</code>, and pushing the model to the single GPU not work?</p>

zhengrchan: <p>It works! Thank you <img src="https://discuss.pytorch.org/images/emoji/apple/smile.png?v=12" title=":smile:" class="emoji" alt=":smile:" loading="lazy" width="20" height="20"></p>
<p>I used:</p>
<pre><code class="lang-auto">model = resnet50()
model.to(device)
model.load_state_dict(torch.load(cfg.MODEL.pretrained_model_path))
</code></pre>
<p>and I got a Runtime error:</p>
<pre><code class="lang-auto">RuntimeError: Error(s) in loading state_dict for Resnet:
    Missing key(s) in state_dict: "conv1.weight", "bn1.weight" ... ...
    Unexpected key(s) in state_dict: "module.conv1.weight", "module.bn1.weight" ... ...
</code></pre>
<p>Seems the Distributed DataParallel save the model in <code>module</code>. Then I find a <strong>solution</strong> there:</p><aside class="quote" data-post="4" data-topic="1686">
  <div class="title">
    <div class="quote-controls"></div>
    <img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/fmassa/48/58_2.png" class="avatar">
    <a href="https://discuss.pytorch.org/t/solved-keyerror-unexpected-key-module-encoder-embedding-weight-in-state-dict/1686/4">[solved] KeyError: 'unexpected key "module.encoder.embedding.weight" in state_dict'</a> 
  </div>
  <blockquote>
    I was thinking about something like the following: 
# original saved file with DataParallel
state_dict = torch.load('myfile.pth.tar')
# create new OrderedDict that does not contain `module.`
from collections import OrderedDict
new_state_dict = OrderedDict()
for k, v in state_dict.items():
    name = k[7:] # remove `module.`
    new_state_dict[name] = v
# load params
model.load_state_dict(new_state_dict)
  </blockquote>
</aside>

<pre><code class="lang-auto">    state_dict = torch.load(cfg.MODEL.pretrained_model_path)
    new_state_dict = OrderedDict()
    for k, v in state_dict.items():
        name = k[7:]
        new_state_dict[name] = v
    model.load_state_dict(new_state_dict)
</code></pre>