I have many Distributed Data Parallel models (NOT Data Parallel!) 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.
I was wondering if it is even possible? if so what is the correct way to do it?
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.
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()