Can't train in DDP mode after recent update

[wudashuo]

🐛 Bug

When I pull the latest code, I found that DDP training would get stuck in the first few epochs.
I ran some tests to see which commit caused this bug and I found commit a3ecf0fd640465f9a7c009e81bcc5ecabf381004 on Mar 3 worked well.
But when I git checkout commit e931b9da33f45551928059b8d61bddd50e401e48 on Mar 4, the bug appeared.
And the bug still exists in the latest commit.

To Reproduce (REQUIRED)

python3 -m torch.distributed.launch --nproc_per_node 4 train.py

The training process would get stuck forever unless you terminate it manually.
And it still occupied the GPU memory unless killing the process by kill -9 xxxxx

Expected behavior

Roll back to the older code, and get the expected behavior.

$ git checkout a3ecf0fd640465f9a7c009e81bcc5ecabf381004
$ python3 -m torch.distributed.launch --nproc_per_node 4 train.py

Environment

If applicable, add screenshots to help explain your problem.

• OS: Ubuntu 20.04
• GPU: 1080 Ti * 4
• Python: 3.8
• pytorch: 1.7.1
• CUDA: 11.1
• Driver: 455.32

Additional

It seems like the latest commit working fine on 2 * 3090, I'm not sure yet, I will do some further tests on 3090 or other GPU.

[glenn-jocher]

@wudashuo thanks for the bug report! I'm not currently training DDP so can't verify correct operation of master myself. I looked at the compare between the commit you said worked well and the one said produced the error. There is only 1 commit between those, but I don't see anything obvious that would impact DDP training:
a3ecf0f...e931b9d

In general when training DDP I always use the Docker image, for stability and reproducibility but also for speed. If you are having issues in a local or cloud environment I would recommend retrying your trainings in the docker image instead. Once docker is installed it's very easy to start training in an image, and I've found that sometimes its a lot faster also than using the native GCP/AWS deep learning AMIs.

[glenn-jocher]

@wudashuo the docker image is on pytorch 1.8 by the way, and its built 'FROM' the latest Nvidia and Facebook PyTorch image https://ngc.nvidia.com/catalog/containers/nvidia:pytorch, so the authors are incentivized to produce working images.

[wudashuo]

I currently don't have the sudo permission, so I can't test the bug in Docker images for now, but I'll test it later today after I got the root user.

I ran some further DDP training tests on GTX1650 * 2, and I got the same problem as I described before.
Maybe something wrong with commit e931b9da33f45551928059b8d61bddd50e401e48

I don't know if this helps you to debug:
When training on the earlier commit, everything worked well, but there is always a warning that said Reducer buckets have been rebuilt in this iteration. in the first epoch. It comes from torch.nn.parallel.distributed.py, I am not familiar with it.

But there is no such warning in the commit e931b9d and later commit, which are all have stucking problem in DDP training:

If I get the root permission I will test it in the Docker image and update this issue.

[roytseng-tw]

yolov5/train.py

Lines 184 to 218 in e8a2b83

	# DDP mode
	if cuda and rank != -1:
	model = DDP(model, device_ids=[opt.local_rank], output_device=opt.local_rank)
	# Trainloader
	dataloader, dataset = create_dataloader(train_path, imgsz, batch_size, gs, opt,
	hyp=hyp, augment=True, cache=opt.cache_images, rect=opt.rect, rank=rank,
	world_size=opt.world_size, workers=opt.workers,
	image_weights=opt.image_weights, quad=opt.quad, prefix=colorstr('train: '))
	mlc = np.concatenate(dataset.labels, 0)[:, 0].max() # max label class
	nb = len(dataloader) # number of batches
	assert mlc < nc, 'Label class %g exceeds nc=%g in %s. Possible class labels are 0-%g' % (mlc, nc, opt.data, nc - 1)
	# Process 0
	if rank in [-1, 0]:
	testloader = create_dataloader(test_path, imgsz_test, batch_size * 2, gs, opt, # testloader
	hyp=hyp, cache=opt.cache_images and not opt.notest, rect=True, rank=-1,
	world_size=opt.world_size, workers=opt.workers,
	pad=0.5, prefix=colorstr('val: '))[0]
	if not opt.resume:
	labels = np.concatenate(dataset.labels, 0)
	c = torch.tensor(labels[:, 0]) # classes
	# cf = torch.bincount(c.long(), minlength=nc) + 1. # frequency
	# model._initialize_biases(cf.to(device))
	if plots:
	plot_labels(labels, save_dir, loggers)
	if tb_writer:
	tb_writer.add_histogram('classes', c, 0)
	# Anchors
	if not opt.noautoanchor:
	check_anchors(dataset, model=model, thr=hyp['anchor_t'], imgsz=imgsz)
	model.half().float() # pre-reduce anchor precision

line 217: model.half().float() # pre-reduce anchor precision is the reason causing ddp hang.
Currently, the model is wrapped by DDP before line 217 at line 186.
Need to do the wrapping AFTER model.half().float().
Like,

            # Anchors
            if not opt.noautoanchor:
                check_anchors(dataset, model=model, thr=hyp['anchor_t'], imgsz=imgsz)
            model.half().float()  # pre-reduce anchor precision 

    # DDP mode
    if cuda and rank != -1:
        model = DDP(model, device_ids=[opt.local_rank], output_device=opt.local_rank)

BTW, what's the purpose of model.half().float()? Why do we need to reduce anchor precision?

[glenn-jocher]

@roytseng-tw ah interesting. Maybe this is because the if rank in [-1, 0] command means this line is only applying to the rank 0 process?

Anchor precision is preemptively reduced by this line before training starts to allow the model to train on the FP16-precision anchors. Without this the model would train on FP32 anchors, which may be altered slightly when the model is saved as FP16. I introduced this line after observing autoanchor compute an FP32 anchor as 8.249, which then was degraded to 8.250 when the model was cast to model.half().

test.py is casting the EMA to .half() everytime it tests at the end of each epoch, this line does the same for the model before training start. I'll have to think about the best solution (other than removing the line).

[glenn-jocher]

@roytseng-tw yes I think the DDP line can be moved to where you suggested. Have you been able to train this way without issue?

[roytseng-tw]

@glenn-jocher Yes, I'am able to train without hang.

[glenn-jocher]

@roytseng-tw ok, thanks for the feedback. I'll make the change then!