Repository: pytorch
Issue: Model loaded from full model state checkpoint does not run on multiple GPUs (device mismatch)
Labels: triage review,oncall: distributed,module: nn,module: cuda,module: serialization

## 🐛 Bug

When a full model is saved as a checkpoint (torch.save(model)), restoring from this checkpoint onto multiple GPUs does not work: inputs are not forwarded to the correct device. Restoring the same model when it's saved as model.state_dict works.

## To Reproduce

1. CUDA_VISIBLE_DEVICES=0,1 python test.py                    # train a model and save both full model and model.state_dict
1. CUDA_VISIBLE_DEVICES=0,1 python test.py --checkpoint model.pt           # restore the model from state_dict: works
1. CUDA_VISIBLE_DEVICES=0,1 python test.py --checkpoint model_full.pt    # restore the model from full state: device mismatch error




## Code to reproduce:

```python
import argparse
import os
os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"

import torch
import torch.nn as nn
import torch.nn.parallel
import torch.distributed as dist
import torch.multiprocessing as mp
import torchvision
import torchvision.transforms as transforms


class ConvNet(nn.Module):
    def __init__(self):
        super(ConvNet, self).__init__()
        self.conv1 = nn.Conv2d(3, 50, kernel_size=5)
        self.conv2 = nn.Conv2d(50, 100, kernel_size=5)
        self.fc1 = nn.Linear(100 * 5 * 5, 300)
        self.fc2 = nn.Linear(300, 10)
        self.maxpool = nn.MaxPool2d(2, stride=2)
        self.act_fn = nn.ReLU(inplace=True)

    def forward(self, x, i=0, gpu=0, debug=False):
        if debug and i == 0:
            print(f'\n{i}  inside forward:  GPU {gpu}  input {x.device}  model {self.conv1.weight.device}\n')

        x = self.conv1(x)
        x = self.maxpool(x)
        x = self.act_fn(x)
        x = self.conv2(x)
        x = self.maxpool(x)
        x = self.act_fn(x)
        x = x.reshape(x.shape[0], -1)
        x = self.fc1(x)
        x = self.act_fn(x)
        x = self.fc2(x)
        return x


def run_inference(model, test_loader, args, debug=False):
    model.eval()
    test_correct = 0
    test_total = 0

    with torch.no_grad():
        for i, (inputs, labels) in enumerate(test_loader):
            inputs = inputs.cuda(args.gpu, non_blocking=True)
            labels = labels.cuda(args.gpu, non_blocking=True)
            if debug and i == 0:
                print(f'\n{i} calling forward:  GPU {args.gpu}  inputs {inputs.device} model {model.device}\n')
            outputs = model(inputs, i=i, gpu=args.gpu, debug=debug)
            torch.distributed.barrier()

            _, predicted = torch.max(outputs.data, 1)
            test_total += labels.size(0)
            test_correct += (predicted == labels).sum()

        test_acc = 100. * test_correct / test_total

    return test_acc


def main(gpu_id, args, num_gpus):
    args.gpu = gpu_id

    model = ConvNet()

    print(f"Using GPU {args.gpu}")
    dist.init_process_group(backend='nccl', init_method='tcp://127.0.0.1:12345', world_size=num_gpus, rank=args.gpu)
    torch.cuda.set_device(args.gpu)
    model.cuda(args.gpu)
    model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[args.gpu])

    optimizer = torch.optim.SGD(model.parameters(), lr=0.1, weight_decay=0.0001)

    transform = transforms.Compose([transforms.ToTensor()])
    train_dataset = torchvision.datasets.CIFAR10(root='.', train=True, download=True, transform=transform)
    test_dataset = torchvision.datasets.CIFAR10(root='.', train=False, download=True, transform=transform)
    train_sampler = torch.utils.data.distributed.DistributedSampler(train_dataset)
    train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=500, shuffle=(train_sampler is None),
                                               num_workers=4, pin_memory=True, drop_last=True, sampler=train_sampler)
    test_loader = torch.utils.data.DataLoader(test_dataset, batch_size=500, shuffle=False, num_workers=4)

    if args.checkpoint is not None:
        if 'full' in args.checkpoint:
            model = torch.load(args.checkpoint, map_location=f'cuda:{args.gpu}')
            torch.cuda.set_device(args.gpu)
            model.cuda(args.gpu)
            model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[args.gpu])
        else:
            checkpoint = torch.load(args.checkpoint, map_location=f'cuda:{args.gpu}')
            model.load_state_dict(checkpoint['state_dict'])
            optimizer.load_state_dict(checkpoint['optimizer'])

        test_acc = run_inference(model, test_loader, args, debug=True)
        print(f'\nTest Accuracy {test_acc:.2f}\n')
        return

    criterion = nn.CrossEntropyLoss()

    for epoch in range(2):
        model.train()
        train_sampler.set_epoch(epoch)
        train_correct = 0
        train_total = 0
        for i, (inputs, labels) in enumerate(train_loader):
            inputs = inputs.cuda(args.gpu)
            labels = labels.cuda(args.gpu)

            outputs = model(inputs, i=i, gpu=args.gpu)
            loss = criterion(outputs, labels)

            loss.backward()
            optimizer.step()
            optimizer.zero_grad()

            _, predicted = torch.max(outputs.data, 1)
            train_total += labels.size(0)
            train_correct += (predicted == labels).sum()

        train_acc = 100. * train_correct / train_total
        test_acc = run_inference(model, test_loader, args)

        if args.gpu == 0:
            print(f"Epoch {epoch:>3d} train {train_acc:.2f} test {test_acc:.2f}")
            torch.save(model, 'model_full.pt')
            torch.save({'state_dict': model.state_dict(), 'optimizer': optimizer.state_dict()}, 'model.pt')


if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='')
    parser.add_argument('--checkpoint', type=str, default=None, metavar='', help='path to a model checkpoint')
    args = parser.parse_args()
    num_gpus = torch.cuda.device_count()
    mp.spawn(main, nprocs=num_gpus, args=(args, num_gpus))
```


## Error message:

```
> Traceback (most recent call last):
>   File "test.py", line 138, in <module>
>     mp.spawn(main, nprocs=num_gpus, args=(args, num_gpus))
>   File "/home/mklachko/miniconda3/envs/pt/lib/python3.7/site-packages/torch/multiprocessing/spawn.py", line 230, in spawn
>     return start_processes(fn, args, nprocs, join, daemon, start_method='spawn')
>   File "/home/mklachko/miniconda3/envs/pt/lib/python3.7/site-packages/torch/multiprocessing/spawn.py", line 188, in start_processes
>     while not context.join():
>   File "/home/mklachko/miniconda3/envs/pt/lib/python3.7/site-packages/torch/multiprocessing/spawn.py", line 150, in join
>     raise ProcessRaisedException(msg, error_index, failed_process.pid)
> torch.multiprocessing.spawn.ProcessRaisedException:
> 
> -- Process 1 terminated with the following error:
> Traceback (most recent call last):
>   File "/home/mklachko/miniconda3/envs/pt/lib/python3.7/site-packages/torch/multiprocessing/spawn.py", line 59, in _wrap
>     fn(i, *args)
>   File "/home/mklachko/gfs/MythicNet/test.py", line 96, in main
>     test_acc = run_inference(model, test_loader, args, debug=True)
>   File "/home/mklachko/gfs/MythicNet/test.py", line 52, in run_inference
>     outputs = model(inputs, i=i, gpu=args.gpu, debug=debug)
>   File "/home/mklachko/miniconda3/envs/pt/lib/python3.7/site-packages/torch/nn/modules/module.py", line 889, in _call_impl
>     result = self.forward(*input, **kwargs)
>   File "/home/mklachko/miniconda3/envs/pt/lib/python3.7/site-packages/torch/nn/parallel/distributed.py", line 705, in forward
>     output = self.module(*inputs[0], **kwargs[0])
>   File "/home/mklachko/miniconda3/envs/pt/lib/python3.7/site-packages/torch/nn/modules/module.py", line 889, in _call_impl
>     result = self.forward(*input, **kwargs)
>   File "/home/mklachko/miniconda3/envs/pt/lib/python3.7/site-packages/torch/nn/parallel/distributed.py", line 705, in forward
>     output = self.module(*inputs[0], **kwargs[0])
>   File "/home/mklachko/miniconda3/envs/pt/lib/python3.7/site-packages/torch/nn/modules/module.py", line 889, in _call_impl
>     result = self.forward(*input, **kwargs)
>   File "/home/mklachko/gfs/MythicNet/test.py", line 28, in forward
>     x = self.conv1(x)
>   File "/home/mklachko/miniconda3/envs/pt/lib/python3.7/site-packages/torch/nn/modules/module.py", line 889, in _call_impl
>     result = self.forward(*input, **kwargs)
>   File "/home/mklachko/miniconda3/envs/pt/lib/python3.7/site-packages/torch/nn/modules/conv.py", line 399, in forward
>     return self._conv_forward(input, self.weight, self.bias)
>   File "/home/mklachko/miniconda3/envs/pt/lib/python3.7/site-packages/torch/nn/modules/conv.py", line 396, in _conv_forward
>     self.padding, self.dilation, self.groups)
> 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)
```



## Expected behavior
Both checkpoints should work the same in the above example

## Environment
Collecting environment information...
PyTorch version: 1.8.0
Is debug build: False
CUDA used to build PyTorch: 10.1
ROCM used to build PyTorch: N/A

OS: Ubuntu 18.04.2 LTS (x86_64)
GCC version: (Ubuntu 7.5.0-3ubuntu1~18.04) 7.5.0
Clang version: Could not collect
CMake version: version 3.13.4

Python version: 3.7 (64-bit runtime)
Is CUDA available: True
CUDA runtime version: Could not collect
GPU models and configuration:
GPU 0: TITAN Xp
GPU 1: TITAN Xp
GPU 2: TITAN Xp
GPU 3: TITAN RTX
GPU 4: TITAN Xp
GPU 5: TITAN RTX
GPU 6: TITAN RTX
GPU 7: TITAN Xp

Nvidia driver version: 418.87.00
cuDNN version: /usr/lib/x86_64-linux-gnu/libcudnn.so.7.6.5
HIP runtime version: N/A
MIOpen runtime version: N/A

Versions of relevant libraries:
[pip3] numpy==1.19.2
[pip3] torch==1.8.0
[pip3] torchvision==0.9.0
[pip3] torchviz==0.0.1
[conda] blas                      1.0                         mkl
[conda] cudatoolkit               10.1.243             h6bb024c_0
[conda] ffmpeg                    4.3                  hf484d3e_0    pytorch
[conda] mkl                       2020.2                      256
[conda] mkl-service               2.3.0            py37he8ac12f_0
[conda] mkl_fft                   1.2.1            py37h54f3939_0
[conda] mkl_random                1.1.1            py37h0573a6f_0
[conda] numpy                     1.19.2           py37h54aff64_0
[conda] numpy-base                1.19.2           py37hfa32c7d_0
[conda] pytorch                   1.8.0           py3.7_cuda10.1_cudnn7.6.3_0    pytorch
[conda] torchvision               0.9.0                py37_cu101    pytorch
[conda] torchviz                  0.0.1                    pypi_0    pypi


## Additional context
Same problem exists in Pytorch 1.7.1


cc @ezyang @gchanan @zou3519 @bdhirsh @jbschlosser @anjali411 @ngimel @aocsa @albanD @mruberry @pietern @mrshenli @pritamdamania87 @zhaojuanmao @satgera @rohan-varma @gqchen @aazzolini @osalpekar @jiayisuse @agolynski @SciPioneer @H-Huang @mrzzd @cbalioglu