[Python] Issue with running training on multigpu using DDP

I am training a classifier model but since it is taking far too long I want to use multigpu for the training. The current code is

rank = int(os.environ["RANK"])
world_size = int(os.environ["WORLD_SIZE"])
local_rank = int(os.environ["LOCAL_RANK"])

torch.cuda.set_device(local_rank)
device = torch.device("cuda", local_rank)

dist.init_process_group(backend='nccl')

model = DDP(model, device_ids=[local_rank])

for epoch in range(1, epochs + 1):
tr_loss = train_epoch(model, train_loader, epoch, backbone, optimizer, device, logger)
te_loss, te_acc = eval_epoch(model, test_loader, epoch, backbone, device, logger, num_classes, all_classes)
if rank == 0:
if te_loss < best_loss:
ckpt = { ... }
torch.save(ckpt, path)

dist.destroy_process_group()


Running the code using

CUDA_VISIBLE_DEVICES=0,1,2,3 torchrun --nproc_per_node 4 train.py


I am getting an error on the line

model = DDP(model, device_ids=[local_rank])


Error: Duplicate GPU detected: rank 2 and rank 0 both on CUDA device 7000

And the same error for rank 0 and rank 1, rank 3 and rank 0, rank 1 and rank 0.

I'm assuming all the models are being loaded to GPU 0, but I'm not sure why. Also I understand why are saving the checkpoint only on GPU 0 but want to make sure the code is correct for that. Thanks

Continue reading...