A confusing problem about DDP

kabu1204: <p>I’m running the code of Mean Teacher where the teacher is updated via EMA.<br>
The Teacher’s state_dicts of 10 continuous EMA updates are saved (marked as Tn-1, Tn-2, …, Tn-10), and I want to check the distance between the output of Tn and Tn-10:<br>
<div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/7/5/75c6ca45316ba7f4bc8daf08f12f36f44141d7bb.png" data-download-href="https://discuss.pytorch.org/uploads/default/75c6ca45316ba7f4bc8daf08f12f36f44141d7bb" title="image"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/7/5/75c6ca45316ba7f4bc8daf08f12f36f44141d7bb_2_690x101.png" alt="image" data-base62-sha1="gNTLJwkKry6Dlp5nYqYPgRmT8cH" width="690" height="101" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/7/5/75c6ca45316ba7f4bc8daf08f12f36f44141d7bb_2_690x101.png, https://discuss.pytorch.org/uploads/default/original/3X/7/5/75c6ca45316ba7f4bc8daf08f12f36f44141d7bb.png 1.5x, https://discuss.pytorch.org/uploads/default/original/3X/7/5/75c6ca45316ba7f4bc8daf08f12f36f44141d7bb.png 2x" data-dominant-color="000000"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">image</span><span class="informations">884×130 10.7 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div></p>
<p>When I use traditional DataParallel, the value is usually 0.3~0.7.<br>
But when I use AMP+DDP on single device but multi GPUs, the value is 0.01 ~ 0.1.<br>
And also, I tried AMP+DataParallel and AMP+DDP on single device and single GPU, the value back to 0.3~0.7.<br>
So it looks like something goes different from usual when DDP is apply on multi GPUs.</p>
<p>My code looks like:</p>
<pre data-code-wrap="python"><code class="lang-python"># Model Initialize
scaler = amp.GradScaler()
student_model = Model().cuda()
teacher_model = Model().cuda()
teacher_dict = student_model.state_dict().copy()
student_model = nn.SyncBatchNorm.convert_sync_batchnorm(student_model)
# student_model = nn.DataParallel(student_model, device_ids = gpu_ids)
student_model = DDP(student_model, device_ids = [local_rank])
teacher_model = nn.SyncBatchNorm.convert_sync_batchnorm(teacher_model)
# teacher_model = nn.DataParallel(teacher_model, device_ids = gpu_ids)
teacher_model = DDP(teacher_model, device_ids = [local_rank])

loss_fn = nn.SmoothL1Loss().cuda()

opt = torch.optim.Adam([{"params":[p for p in list(student_model.parameters()) if p.requires_grad]}], lr=lr, betas=(0.5, 0.999))

# Training
teacher_dict_list = []
teacher_dict_prev = None  # Tn-10, at the beginning is None
for epoch in range(num_epochs):
	for data in trainloader:
		images, labels = data
		images = images.cuda()
		labels = labels.cuda()

		with amp.autocast():
			preds = student_model(images)
			loss = loss_fn(preds, labels)

		scaler.scale(loss).backward()
		scaler.step(opt)
		opt.zero_grad()

		t = teacher_dict.copy()
		for k, v in t.items():
			t[k] = t[k].cpu()	# store on cpu memory, or the cuda memory would overflow
		teacher_dict_list.append(t)	# store Tn-1
		# EMA update teacher_dict
		for k, v in student_model.module.state_dict().items():
			if k.find('num_batches_tracked') == -1:
				teacher_dict[k] = 0.999 * teacher_dict[k] + 0.001 * v
		# we have stored 10 teacher's state_dicts
		if len(teacher_dict_list) == 10:
			teacher_dict_prev = teacher_dict_list[0]	# pick Tn-10
			del(teacher_dict_list[0])

		if teacher_dict_prev is None:
			continue
		else:
			teacher_model.module.load_state_dict(teacher_dict)	# load the newest state_dict
			with torch.no_grad():
				with amp.autocast():
					output_cur = teacher_model(images)

			teacher_model.module.load_state_dict(teacher_dict_prev)	# load the old state_dict
			with torch.no_grad():
				with amp.autocast():
					output_prev = teacher_model(images)
			print(torch.abs(output_cur - output_prev).max())
		scaler.update()
</code></pre>
<p>The code above is a simplified version, the structure is the same. To some reason, I can’t post all my code for you to reproduce the problem.<br>
So, the printed value is 0.3 ~ 0.7 (which i think should be correct) at the setting of:</p>
<ul>
<li>DataParallel w/o AMP on single device and multi GPUs</li>
<li>DataParallel with AMP on single device and multi GPUs</li>
<li>DDP with AMP on single device and <strong>single</strong> GPU</li>
</ul>
<p>but it becomes much smaller (0.01 ~ 0.1) on DDP with AMP on single device and multi GPUs.</p>

blowtorch: <p>Hey did u solve your problem ? If yes, could you maybe explain the solution?</p>