How to save a single final model from distributed multi GPU training?

AndreaSottana: <p>Hello,</p>
<p>There is something I seem to struggle to understand regarding how to use the DistributedDataParallel correctly. Below is a reproducible example of my code (I tried to make it as short and general as possible, and removed the evaluation step from the training).</p>
<p>I’m running the code on a machine with two GPUs, and my problem is that the code will save two separate torch models, one for each GPU process I’ve spawned. My assumption is that distributed multiprocessing should eventually reconvene everything under the same model, and so a single model should be saved by the code after the training. Could someone please tell me what I’m doing wrong in the code below?</p>
<p>Thanks in advance</p>
<pre><code class="lang-auto">import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torch.utils.data import DataLoader
from torch.nn.parallel import DistributedDataParallel as DDP
from torch.utils.data.distributed import DistributedSampler
import torch.multiprocessing as mp
from argparse import ArgumentParser
import os

class MyModel(nn.Module):
    def __init__(self, input_dim, inner_layer_1, inner_layer_2, output_dim):
        super().__init__()
        self.fc1 = nn.Linear(input_dim, inner_layer_1)
        self.fc2 = nn.Linear(inner_layer_1, inner_layer_2)
        self.fc3 = nn.Linear(inner_layer_2, output_dim)

    def forward(self, x):
        x = self.fc1(x) 
        x = F.relu(x)
        x = self.fc2(x)
        x = F.relu(x)
        x = self.fc3(x)
        x = F.softmax(x, dim=1)
        return x

def train(gpu_number, n_epochs, model, train_data, optimizer, loss_fn, log_interval=2):
    os.environ['MASTER_ADDR'] = 'localhost'
    os.environ['MASTER_PORT'] = '29500'
    torch.distributed.init_process_group(
        backend='nccl',
        init_method='env://',
        world_size=2,  # total number of gpus
        rank=gpu_number
    )

    sampler = DistributedSampler(train_data, num_replicas=2, rank=gpu_number)
    trainloader = DataLoader(train_data, batch_size=8, sampler=sampler)

    #torch.cuda.set_device(gpu_number)
    model = model.cuda(gpu_number)
    model = DDP(model, device_ids=[gpu_number], output_device=gpu_number)

    for epoch in range(n_epochs):
        for i, batch in enumerate(trainloader):
            inputs, labels = batch[:,:8].cuda(gpu_number), batch[:,-2:].cuda(gpu_number)
            optimizer.zero_grad()
            outputs = model(inputs)

            loss = loss_fn(outputs, labels)
            loss.backward()
            optimizer.step()
    torch.save(model.state_dict(), f"model_{gpu_number}.pt")

if __name__ == "__main__":
    train_data = torch.rand(100, 10)
    n_epochs = 3
    learning_rate = 0.001
    model = MyModel(8, 800, 300, 2)
    loss_fn = nn.MSELoss() # use nn.CrossEntropyLoss() for binary classification
    optimizer = optim.Adam(model.parameters(), lr=learning_rate)
    mp.spawn(train, nprocs=2, args=(n_epochs, model, train_data, optimizer, loss_fn, 2))

</code></pre>

yoshitomo-matsubara: <p>Hi <a class="mention" href="/u/andreasottana">@AndreaSottana</a></p>
<p>When using the distributed training mode, one of the processes should be treated as the main process, and you can save the model only for the main process.</p>
<p>Check one of the <a href="https://github.com/pytorch/vision/blob/master/references/classification/utils.py#L224-L230" rel="noopener nofollow ugc">torchvision’s examples</a>, which will give you a good idea for your problem.</p>
<p>Say you’re using two GPUs in the distributed training mode, and then there will be two processes of your script.<br>
Both the processes will call <code>save_on_master</code> function in the above example at some time, but only one of the processes will save the model (i.e., when <code>is_main_process()</code> returns True)</p>

AndreaSottana: <p>Hi <a class="mention" href="/u/yoshitomo-matsubara">@yoshitomo-matsubara</a><br>
Many thanks for your prompt reply. This is the first time I hear about <code>save_on_master</code> and couldn’t find it in any torch distributed documentation I’ve seen so far, it would be good to make it more visible I think.<br>
Anyway, I have 3 follow up questions.</p>
<ol>
<li>If I understand correctly, the functions you linked would simply save the model only if the rank is zero. Is this arbitrary, i.e. could I have just chosen rank one instead for example?</li>
<li>The code you linked does not contain any explicit action of bringing together the two spawned processes, which I’m therefore assuming has to happen in the background. Does this mean that at any point after calling optimiser.step(), the weights on both GPUs will synchronise automatically? If true, this would imply that the two models that my original code would have saved were completely identical and took into account the training from both GPUs, as opposed to being two different models trained on a different GPU each with a different subset of the data. If my understanding is correct, this would also imply that the <code>save_on_master</code> function you linked has the only advantage of not saving the model multiple times, but even if I left my code as original, it would just save multiple copies of the same model with exactly the same weights, and I could use any of these saved copies, so it wouldn’t make any real difference. Please correct me if I’m wrong.</li>
<li>Given that the two (or more) GPUs are independent, one might run the code slightly faster than the other. If you save the model when the rank is zero and the GPU from rank 1 hasn’t yet finished running, is there a chance that you might save the model before the whole training has actually completed?</li>
</ol>
<p>Thank you very much for your kind help!</p>

yoshitomo-matsubara: <p>Hi <a class="mention" href="/u/andreasottana">@AndreaSottana</a></p>
<ol>
<li>Yes, technically you can choose the rank you like as the main process. Since the rank number starts from zero, I’d suggest using rank 0 as the main process</li>
<li>Using DDP, the two processes synchronize gradient across processes. You can confirm the detail <a href="https://pytorch.org/tutorials/intermediate/ddp_tutorial.html" rel="noopener nofollow ugc">here</a>. If you don’t control the timing of saving your model when allowing both the processes to save a model, it may cause file handler error like one of the processes attempts to overwrite a file whose save operation is not complete.</li>
<li>It can be controlled by <code>torch.distributed.barrier()</code>, which waits for other processes. For instance, I’d put <code>torch.distributed.barrier()</code> right before <code>torch.save</code> in your code above to wait for all the iterations in the epoch done before saving the model. <a href="https://github.com/pytorch/vision/blob/master/references/classification/train.py#L66-L67" rel="noopener nofollow ugc">The same example in torchvision</a> uses it in evaluation to compute global accuracy across the processes. The <code>synchronize_between_processes()</code> uses <code>torch.distributed.barrier()</code> as <a href="https://github.com/pytorch/vision/blob/7947fc8fb38b1d3a2aca03f22a2e6a3caa63f2a0/references/classification/utils.py#L38" rel="noopener nofollow ugc"><code>dist.barrier()</code></a>
</li>
</ol>

AndreaSottana: <p>Hi <a class="mention" href="/u/yoshitomo-matsubara">@yoshitomo-matsubara</a><br>
Many thanks for the very detailed reply.<br>
I will try to study the code you linked and implement your changes, and will let you know if I have further queries down the line.<br>
Thanks again</p>

AndreaSottana: <p>Hi <a class="mention" href="/u/yoshitomo-matsubara">@yoshitomo-matsubara</a><br>
I’d have a small follow up from this.</p>
<p>I’ve noticed that you don’t call <code>dist.barrier()</code> before <code>save_on_master</code> (for example <a href="https://github.com/pytorch/vision/blob/f52ddb0c23e2590ae45d2afa70ff935eb4037a91/references/classification/train.py#L228" rel="noopener nofollow ugc">here</a>) but suggested I do that in my code. Is there a reason why you don’t need to do this in your code?</p>
<p>I’ve also noticed that when you save the model, you save the <code>model_without_ddp.state_dict()</code> instead of simply <code>model.state_dict()</code>. I’ve not seen this done before. Is there any advantage in doing this over simply saving the DDP model’s state dict?</p>
<p>Thanks again</p>

yoshitomo-matsubara: <p>Hi <a class="mention" href="/u/andreasottana">@AndreaSottana</a></p>
<p>I don’t know which one you referred to as “your code”, but I suggested calling <code>dist.barrrier()</code> before <code>torch.save</code> in the code you showed above as follows:</p>
<blockquote>
<p>For instance, I’d put <code>torch.distributed.barrier()</code> right before <code>torch.save</code> in your code above to wait for all the iterations in the epoch done before saving the model.</p>
</blockquote>
<p>The example code actually calls <code>dist.barrier()</code> before <code>save_on_master</code>, and that happens in <a href="https://github.com/pytorch/vision/blob/f52ddb0c23e2590ae45d2afa70ff935eb4037a91/references/classification/train.py#L68" rel="noopener nofollow ugc"><code>evaluate</code> function</a>, which is called before <code>save_on_master</code></p>
<blockquote>
<p>I’ve also noticed that when you save the model, you save the <code>model_without_ddp.state_dict()</code> instead of simply <code>model.state_dict()</code> . I’ve not seen this done before. Is there any advantage in doing this over simply saving the DDP model’s state dict?</p>
</blockquote>
<p>Since DDP and DP will wrap your model, the state_dict saved by <code>model.stage_dict()</code> cannot be directly loaded a model without DDP/DP.<br>
e.g.,<br>
Try this minimal example</p>
<pre><code class="lang-auto">from torchvision import models

model = models.resnet18()
print(model.state_dict().keys())

model = torch.nn.parallel.DistributedDataParallel(model)
# or use DP instead of DDP
# model = torch.nn.DataParallel(model)
print(model.state_dict().keys())
</code></pre>
<p>From the first print statement, you’ll see pure module paths used in ResNet-18. But from the second print statement, you can confirm “module.” as a prefix for all the module paths shown by the first print statement.</p>
<p>It can be confirmed that your model is referred to as <code>module</code> in <a href="https://github.com/pytorch/pytorch/blob/master/torch/nn/parallel/data_parallel.py#L137" rel="noopener nofollow ugc">DP</a> and <a href="https://github.com/pytorch/pytorch/blob/master/torch/nn/parallel/distributed.py#L536" rel="noopener nofollow ugc">DDP</a> implementations</p>

AndreaSottana: <p>Thanks again!<br>
By “your code” I meant the code <a href="https://github.com/pytorch/vision/blob/f52ddb0c23e2590ae45d2afa70ff935eb4037a91/references/classification/train.py#L228" rel="noopener nofollow ugc">here</a> but I didn’t realise that the code actually called the <code>evaluate</code> function just before <code>save_on_master</code>, which was calling <code>dist.barrier</code>, makes sense now. I should have said the torch vision code that you linked.<br>
I will convert my <code>torch.save</code> to the <code>save_on_master</code> function and call <code>dict.barrier()</code> before it in my code as well then.</p>
<p>Thanks for the small example with resnet as well, all clear now. I’m still new to multi GPU training but this helped me a lot.</p>

Kartik_jain: <p>Hi <a class="mention" href="/u/andreasottana">@AndreaSottana</a> <a class="mention" href="/u/yoshitomo-matsubara">@yoshitomo-matsubara</a>, I also had a similar doubt and this discussion really helped. Can you please share your final code which allowed you to save the model on master.</p>
<p>Thanks in advance!</p>