def _test_grad_accum_func(rank, world_size, tempfile_name):
_dist_init(rank, world_size, tempfile_name, backend="gloo") # Covers gloo
model = Linear(4, 2, bias=False)
model.to("cuda")
model = DDP(model, device_ids=[rank])
optim = AdaScale(SGD(model.parameters(), lr=0.1), num_gradients_to_accumulate=2)
with model.no_sync():
# iter 1, input vectors are pointing dim0 and dim1
in_data = Tensor([0.0] * 4)
in_data[rank] = 1.0
in_data = in_data.cuda()
out = model(in_data)
out.sum().backward()
# iter 2, input vectors are pointing dim2 and dim3
in_data = Tensor([0.0] * 4)
in_data[rank + 2] = 1.0
in_data = in_data.cuda()
out = model(in_data)
out.sum().backward()
# since all inputs are orthogonal, the gain should be exactly 4.0.
assert np.allclose(optim.gain(), 4.0), optim.gain()
optim.step()
optim.zero_grad()
dist.destroy_process_group()
