def test_one_iteration(self):
"""Test FSDP with uneven divide of parameter shards."""
model = Linear(3, 3, bias=False)
input = torch.rand(8, 3)
my_lr = 0.1
ref_forward_output_my_rank, ref_weight_out = self._get_ref_results(
model, input, my_lr)
model.to(self.rank)
model = FSDP(model)
optim = SGD(model.parameters(), lr=my_lr)
self.assertTrue(len(input) >= self.world_size)
in_data = torch.Tensor(input[self.rank]).to(self.rank)
out = model(in_data)
out.float().sum().backward()
optim.step()
optim.zero_grad()
with model.summon_full_params():
torch.cuda.synchronize() # TODO: This is here because it was
# originally part of get_full_params(), debug why it is needed here.
weight_out = model.module.weight.T.clone()
self.assertEqual(ref_forward_output_my_rank, out)
self.assertEqual(ref_weight_out, weight_out)
