Load DDP model trained with 8 gpus on only 2 gpus?

kazem: <p>I have many <strong>Distributed Data Parallel models (NOT Data Parallel!)</strong> trained with 8 gpus on a cluster. I have no problem correctly restoring them with same number of gpus (8). But wait time to get 8 is too long. So I want to restore them with only two.</p>
<p>I was wondering if it is even possible? if so what is the correct way to do it?</p>
<p>The script below (test.py) works fine with 8 gpus but produces erroneous results with 2 gpus (in the latter case, the results are the same as a model just initialized with random weights). I use “python -m torch.distributed.launch --nproc_per_node=num_gpus test.py” to run it from terminal.</p>
<pre><code class="lang-auto">import argparse
from torchvision.models import resnet18
from torch.nn.parallel import DistributedDataParallel as DDP
import torch.distributed as dist

def cleanup():
    dist.destroy_process_group()

def main():
    torch.distributed.init_process_group(
backend='nccl', init_method='env://')
    torch.cuda.set_device(args.local_rank)
    model = resnet18()
    model = model.to([args.local_rank][0])
    model = DDP(model, device_ids=[args.local_rank], 
output_device=[args.local_rank][0])

    # load the model
    checkpoint = torch.load(load_path)
    state_dict = checkpoint['model_state_dict']
    model.load_state_dict(state_dict)
    dist.barrier()

    cleanup()

if __name__ == '__main__':
    parser = argparse.ArgumentParser(description="blah")
    parser.add_argument("--local_rank", type=int)
    args, _ = parser.parse_known_args()
    main()
</code></pre>

mrshenli: <p>This should be possible, there is a <code>map_location</code> argument in <code>torch.load</code>. Checkout <a href="https://pytorch.org/tutorials/beginner/saving_loading_models.html" rel="nofollow noopener">this</a>.</p>
<p>The <code>map_location</code> can be a device, a function, a map etc. [<a href="https://pytorch.org/docs/stable/torch.html#torch.load" rel="nofollow noopener">API</a>]</p>

kazem: <p>Thank you for your answer. The documentation does not include a working example for DDP. I have already tried many ways using map function None of which have worked so far. If you could show me a simple working example with mnist dataset to map 8 gpus to 1 or 2 or 4 gpus, or cpu with DistributedData parallel I would greatly appreciate it.</p>

mrshenli: <p>There are a few things to clarify.</p>
<ol>
<li>As you are using the <code>resnet18</code> from torchvision, the model only lives on a single GPU.</li>
<li>The launcher script you use starts <code>num_gpus</code> processes, and each process has its own DDP instance, dataloader, and the model replica.</li>
<li>With 1 and 2, your training scripts only need put the model to one GPU (you can use the rank as the device id), load the data into one GPU, and the DDP instance will handle the comm for you, and make sure that all model replicas are synchronized properly.</li>
<li>With the above 3, the question then would be “how do I load a model to a specific GPU device?”. And the answer is use <code>map_local=torch.device(rank)</code>.</li>
</ol>
<p>The following code works for me with the launching cmd</p>
<pre><code class="lang-auto">python -m torch.distributed.launch --nproc_per_node=2 test.py
</code></pre>
<pre><code class="lang-python">import argparse
from torchvision.models import resnet18
from torch.nn.parallel import DistributedDataParallel as DDP
import torch.distributed as dist
import torch

def cleanup():
    dist.destroy_process_group()

def main(args):
    torch.distributed.init_process_group(backend='nccl', init_method='tcp://localhost:23456', rank=args.local_rank, world_size=2)
    torch.cuda.set_device(args.local_rank)
    model = resnet18()

    path = "save_model.pt"
    if args.local_rank == 0:
        # save CPU model
        torch.save(model, path)

    dist.barrier()
    # local model to GPU
    loaded_model = torch.load(path, map_location=torch.device(args.local_rank))

    model = DDP(loaded_model, device_ids=[args.local_rank])
    print(f"Rank {args.local_rank} traning on device {list(model.parameters())[0].device}")

    # create a dedicated data loader for each process

    cleanup()

if __name__ == '__main__':
    parser = argparse.ArgumentParser(description="blah")
    parser.add_argument("--local_rank", type=int)
    args, _ = parser.parse_known_args()
    main(args)
</code></pre>

kazem: <p>@ mrshenli thanks for your reply. I tried your method after a few minor correction but it still gives me the same erroneous result. I use this <a href="https://github.com/pytorch/vision/blob/master/torchvision/models/resnet.py" rel="nofollow noopener">resnet</a> script to call the model.</p>
<p>I trained it on a large dataset and decided to save it periodically during training. Due to testing slowing down the training I decided to test it later using the saved models. When I train DDP with 8 gpus and test DDP with 8 gpus later, there is no issue. However, when I train DDP with 8 gpus  and test DDP with 2 gpus later the problem occurs.</p>
<p>Also I only want to save and load the state_dict and not the entire model since it takes a lot of space.</p>
<p>I will create a working example for mnist shortly.</p>

mrshenli: <blockquote>
<p>I tried your method after a few minor correction but it still gives me the same erroneous result. I use this <a href="https://github.com/pytorch/vision/blob/master/torchvision/models/resnet.py" rel="noopener nofollow ugc">resnet</a> script to call the model.</p>
</blockquote>
<p>You mean you saw error by running the script as is? What error did you see and what fix did you applied?</p>
<blockquote>
<p>Also I only want to save and load the state_dict and not the entire model since it takes a lot of space.</p>
</blockquote>
<p>It should be doable by just modifying two lines (save and load).</p>
<blockquote>
<p>When I train DDP with 8 gpus and test DDP with 8 gpus later, there is no issue. However, when I train DDP with 8 gpus and test DDP with 2 gpus later the problem occurs.</p>
</blockquote>
<p>The resnet link you posted points to torchvision resnet, so the model only lives on a single device. How did you go from training on 8 gpus to testing on 2 gpus? Did you do the following?</p>
<ol>
<li>After training, use <strong>only</strong> rank 0 to save <code>ddp.module</code> to file.</li>
<li>For testing, as you no longer need comm across models, you don’t need DDP. You can spawn two processes, each load the saved module from file to its dedicated device by setting <code>map_reduce</code>. And use sth like <code>all_gather</code> to collect loss/accuracy data to rank 0?</li>
</ol>

kazem: <p>@ mrshenli thanks again. I will try to answer all your inquiries with more detail in a bit today.</p>
<p>Unfortunately, I could not use your script as is because my already saved DDP (without “.module”) was already saved using a state_dict method.<br>
So as for the minor changes, I did the following. :</p>
<pre><code class="lang-auto">def main(args):
    torch.distributed.init_process_group(backend='nccl', init_method='env://')
    test_loader = DataLoader(
        test_dataset,
        batch_size=args.test_batch_size,
        shuffle=False,
        num_workers=args.num_workers,
        pin_memory=True)

    model = get_model()
#############################################################
   # My changes
    torch.cuda.set_device(args.local_rank)
    model = model.to([args.local_rank][0])
    model = DDP(model, device_ids=[args.local_rank], 
output_device=[args.local_rank][0])
    checkpoint = torch.load(args.load_path)  # , map_location=map_location)
    state_dict = checkpoint['model_state_dict']
    model.load_state_dict(state_dict)
##############################################################
    dist.barrier()
    test_function(model, test_loader, args.local_rank,args.load_path.with_suffix('.csv'))
</code></pre>
<p>I trained resnet18 from scratch. I just copied and used the <a href="https://github.com/pytorch/vision/blob/master/torchvision/models/resnet.py" rel="nofollow noopener">resnet</a> script locally.</p>
<p>As for your last two comments I did use just rank 0 to save the ddp, but I saved the state_dict() for ddp itself (without .module). That is why when I used your script I also had to remove the .module similar to this:<br>
<a href="https://discuss.pytorch.org/t/solved-keyerror-unexpected-key-module-encoder-embedding-weight-in-state-dict/1686">[solved] KeyError: ‘unexpected key “module.encoder.embedding.weight” in state_dict’</a><br>
Is it correctly to do so?</p>

mrshenli: <aside class="quote no-group quote-modified" data-username="kazem" data-post="7" data-topic="73455">
<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/kazem/48/20632_2.png" class="avatar"> kazem:</div>
<blockquote>
<p>As for your last two comments I did use just rank 0 to save the ddp, but I saved the state_dict() for ddp itself (without .module). That is why when I used your script I also had to remove the. Is it correctly to do so?</p>
</blockquote>
</aside>
<p>Yes, that is correct. The saved and loaded model type need to match.</p>

mrshenli: <aside class="quote no-group" data-username="kazem" data-post="7" data-topic="73455">
<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/kazem/48/20632_2.png" class="avatar"> kazem:</div>
<blockquote>
<p>checkpoint = torch.load(args.load_path) # , map_location=map_location)</p>
</blockquote>
</aside>
<p>This line might cause a problem if the model was saved from a device that is not available on the machine that loads the model. But it should be OK in your case, as the model was saved from rank 0 (i.e., “cuda:0”), whose device is available in both envs. However, without <code>map_location</code>, it means the two DDP processes in testing are operating on the same GPU? That could also cause problems.</p>

kazem: <p><a href="https://discuss.pytorch.org/u/mrshenli">mrshenli</a> Sorry for the late reply. Say I want to train the DDP model on 4 gpus and restore it as DDP  on 2. I created an mnist example to illustrate my case while following your example. This whole script is borrowed from <a href="https://github.com/pytorch/examples/blob/master/mnist/main.py" rel="nofollow noopener">mnist</a>, modified and split into three scripts:</p>
<ol start="0">
<li>mnist_common.py</li>
</ol>
<pre><code class="lang-auto">import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.distributed as dist
import argparse
from torchvision import datasets, transforms
from torch.utils.data.distributed import DistributedSampler
from torch.utils.data import DataLoader


def cleanup():
    dist.destroy_process_group()


class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.conv1 = nn.Conv2d(1, 32, 3, 1)
        self.conv2 = nn.Conv2d(32, 64, 3, 1)
        self.dropout1 = nn.Dropout2d(0.25)
        self.dropout2 = nn.Dropout2d(0.5)
        self.fc1 = nn.Linear(9216, 128)
        self.fc2 = nn.Linear(128, 10)

    def forward(self, x):
        x = self.conv1(x)
        x = F.relu(x)
        x = self.conv2(x)
        x = F.relu(x)
        x = F.max_pool2d(x, 2)
        x = self.dropout1(x)
        x = torch.flatten(x, 1)
        x = self.fc1(x)
        x = F.relu(x)
        x = self.dropout2(x)
        x = self.fc2(x)
        output = F.log_softmax(x, dim=1)
        return output


def train(args, model, device, train_loader, optimizer, epoch):
    model.train()
    for batch_idx, (data, target) in enumerate(train_loader):
        data, target = data.to(device, non_blocking=True), \
                       target.to(device, non_blocking=True)
        optimizer.zero_grad()
        output = model(data)
        loss = F.nll_loss(output, target)
        loss.backward()
        optimizer.step()


def test(args, model, device, test_loader):
    model.eval()
    test_loss = 0
    correct = 0
    with torch.no_grad():
        for data, target in test_loader:
            data, target = data.to(device, non_blocking=True), \
                           target.to(device, non_blocking=True)
            output = model(data)
            test_loss += F.nll_loss(
                output,
                target,
                reduction='sum').item()
            pred = output.argmax(dim=1, keepdim=True)
            correct += pred.eq(target.view_as(pred)).sum().item()

    test_loss /= len(test_loader.dataset)

    if args.local_rank == 0:
        print('Test set: Average loss: {:.4f},'
              ' Accuracy: {}/{} ({:.2f}%)'.format(
            test_loss, correct, len(test_loader.dataset),
            100. * correct / len(test_loader.dataset)))


# Training settings
parser = argparse.ArgumentParser(description='PyTorch MNIST Example')
parser.add_argument('--batch-size',
                    type=int,
                    default=64,
                    metavar='N',
                    help='input batch size for training')
parser.add_argument('--test-batch-size',
                    type=int,
                    default=1000,
                    metavar='N',
                    help='input batch size for testing')
parser.add_argument('--epochs', type=int, default=14, metavar='N',
                    help='number of epochs to train (default: 14)')
parser.add_argument('--lr', type=float, default=1.0, metavar='LR',
                    help='learning rate (default: 1.0)')
parser.add_argument('--gamma', type=float, default=0.7, metavar='M',
                    help='Learning rate step gamma (default: 0.7)')
parser.add_argument('--seed', type=int, default=1, metavar='S',
                    help='random seed (default: 1)')
parser.add_argument('--local_rank', type=int)
args = parser.parse_args()

train_dataset = datasets.MNIST(
    '../data',
    train=True,
    download=False,
    transform=transforms.Compose([
       transforms.ToTensor(),
       transforms.Normalize((0.1307,), (0.3081,))
    ]))
train_sampler = DistributedSampler(
    train_dataset,
    num_replicas=torch.cuda.device_count(),
    rank=args.local_rank)
train_loader = DataLoader(train_dataset,
                          batch_size=args.batch_size,
                          shuffle=(train_sampler is None),
                          num_workers=0,
                          pin_memory=True,
                          sampler=train_sampler)
test_loader = DataLoader(
    datasets.MNIST(
        '../data',
        train=False,
        transform=transforms.Compose([
            transforms.ToTensor(),
            transforms.Normalize((0.1307,), (0.3081,))
        ])),
    batch_size=args.test_batch_size,
    shuffle=True,
    num_workers=0,
    pin_memory=True,)
</code></pre>
<ol>
<li>mnist_train.py</li>
</ol>
<pre><code class="lang-auto">from __future__ import print_function
import torch
import torch.optim as optim
import torch.distributed as dist
import torch.backends.cudnn as cudnn
from torch.optim.lr_scheduler import StepLR
from torch.nn.parallel import DistributedDataParallel as DDP
from mnist_common import args, Net, train_loader, train_sampler,\
    test_loader, train, test, cleanup


def main(args):
    dist.init_process_group(backend='nccl',
                            init_method='tcp://localhost:23456',
                            rank=args.local_rank,
                            world_size=torch.cuda.device_count())
    torch.manual_seed(args.seed)
    torch.cuda.set_device(args.local_rank)
    cudnn.benchmark = True

    model = Net()
    model = model.to([args.local_rank][0])  # distribute the model
    # Should we set the output_device value in DPP?
    model = DDP(model, device_ids=[args.local_rank])
    # , output_device=[args.local_rank][0])

    optimizer = optim.Adadelta(model.parameters(), lr=args.lr)
    scheduler = StepLR(optimizer, step_size=1, gamma=args.gamma)
    for epoch in range(1, args.epochs + 1):
        train_sampler.set_epoch(epoch)
        train(args, model, args.local_rank,
              train_loader, optimizer, epoch)
        test(args, model, args.local_rank, test_loader)
        scheduler.step(epoch)

    # I intend to save the model
    # AFTER some training not, not before
    if args.local_rank == 0:
        torch.save(model, "mnist_cnn.pt")
    dist.barrier()
    cleanup()


if __name__ == '__main__':
    main(args)
</code></pre>
<p>Also I intend to test the model, only after training (sometimes up to a few days) has finished, by restoring the saved weights (or model).<br>
2) mnist_test.py</p>
<pre><code class="lang-auto">from __future__ import print_function
import torch
import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP
from mnist_common import args, Net, test_loader, test, cleanup


def main(args):
    dist.init_process_group(backend='nccl',
                            init_method='tcp://localhost:23456',
                            rank=args.local_rank,
                            world_size=2)
    torch.manual_seed(args.seed)
    torch.cuda.set_device(args.local_rank)

    model = torch.load("mnist_cnn.pt",
                       map_location=torch.device(args.local_rank))
    model = DDP(model, device_ids=[args.local_rank])
    print(f"Rank {args.local_rank} "
          f"test on device {list(model.parameters())[0].device}")

    test(args, model, args.local_rank, test_loader)
    cleanup()


if __name__ == '__main__':
    main(args)
</code></pre>
<p>The mnist_train.py runs sucessfully using<br>
<code>python -m torch.distributed.launch nproc_per_node=4 (or 2) mnist_train.py</code>.<br>
but when i run the test script using<br>
<code>python -m torch.distributed.launch nproc_per_node=2 mnist_test.py</code>.<br>
I get the following:</p>
<pre><code class="lang-auto">Rank 0 test on device cuda:0
Rank 1 test on device cuda:1
Test set: Average loss: 0.0274, Accuracy: 9913/10000 (99.13%)

RuntimeError: Expected tensor for argument #1 input to have 
the same device as tensor for argument #2 weight; 
but device 0 does not equal 1 
(while checking arguments for cudnn_convolution)
</code></pre>

mrshenli: <blockquote>
<p>Rank 0 test on device cuda:0<br>
Rank 1 test on device cuda:1<br>
Test set: Average loss: 0.0274, Accuracy: 9913/10000 (99.13%)</p>
<p>RuntimeError: Expected tensor for argument <span class="hashtag-raw">#1</span> input to have<br>
the same device as tensor for argument <span class="hashtag-raw">#2</span> weight;<br>
but device 0 does not equal 1<br>
(while checking arguments for cudnn_convolution)</p>
</blockquote>
<p>This means the first parameter of both models are placed onto the correct device. Can you do the same check for all parameters? i.e., making sure that all parameters are placed to the correct device.</p>
<blockquote>
<p><code>output = model(data)</code></p>
</blockquote>
<p>Before the line above in <code>test(...)</code>, can you print the device ids of the data as well? Looks like the model and data device does not match on rank 1.</p>

kazem: <p>I see. But it should not be the case since both are moved to <code>args.local_rank</code>. Anyways, I did what you suggested and also changed the <code>test-batch-size</code> to <code>1024</code>. Here’s the outcome:</p>
<pre><code class="lang-auto">Rank 0 test on device cuda:0
Rank 1 test on device cuda:1
after data=data.to(device,), before output=model(data) in test function,  batch_idx: 0 device: 1
after data=data.to(device,), before output=model(data) in test function,  batch_idx: 0 device: 0
after data=data.to(device,), before output=model(data) in test function,  batch_idx: 1 device: 0
after data=data.to(device,), before output=model(data) in test function,  batch_idx: 2 device: 0
after data=data.to(device,), before output=model(data) in test function,  batch_idx: 3 device: 0
after data=data.to(device,), before output=model(data) in test function,  batch_idx: 4 device: 0
after data=data.to(device,), before output=model(data) in test function,  batch_idx: 5 device: 0
after data=data.to(device,), before output=model(data) in test function,  batch_idx: 6 device: 0
after data=data.to(device,), before output=model(data) in test function,  batch_idx: 7 device: 0
after data=data.to(device,), before output=model(data) in test function,  batch_idx: 8 device: 0
after data=data.to(device,), before output=model(data) in test function,  batch_idx: 9 device: 0
Test set: Average loss: 0.0275, Accuracy: 9913/10000 (99.13%)

RuntimeError: Expected tensor for argument #1 'input' to have the same device as tensor for argument #2 'weight'; but device 0 does not equal 1 (while checking arguments for cudnn_convolution)
</code></pre>

mrshenli: <blockquote>
<p>Rank 1 test on device cuda:1<br>
after data=data.to(device,), before output=model(data) in test function,  batch_idx: 0 device: 1</p>
</blockquote>
<p>This is weird. This means all model parameters are on <code>cuda:1</code> and the input batch is also on <code>cuda:1</code>, but somehow one of the conv layers still throws device mismatch? I am not sure what happened here, but as the error suggests the mismatch occurs in <code>cudnn_convolution</code>, I would check if the input (<code>x</code>) of and the parameters the two conv layer (<code>self.conv1</code> and <code>self.conv2</code>) match in the <code>forward()</code> function during testing.</p>
<p>BTW, two more comments on the script:</p>
<ol>
<li>As you are only doing forward during testing, it is not necessary to use DDP there, as all comm in DDP occurs during backward.</li>
<li>I noticed you saving a DDP module and then load that DDP module and wrap it with another DDP module. Is this intentional? Shouldn’t <code>mnist_train.py</code> save <code>model.module</code> instead? (or use model.module to initialize DDP instances in testing)</li>
</ol>