[FSDP] Adam Gives Different Results Where Only Difference Is Flattening
Labels: oncall: distributed, module: fsdp

Consider the following unit test (that relies on some imports from `common_fsdp.py`):
```
def test(self):
    local_model = TransformerWithSharedParams.init(
        self.process_group,
        FSDPInitMode.NO_FSDP,
        CUDAInitMode.CUDA_BEFORE,
        deterministic=True,
    )
    fsdp_model = FSDP(
        copy.deepcopy(local_model),
        sharding_strategy=ShardingStrategy.NO_SHARD,
    )
    ddp_model = DDP(local_model, device_ids=[self.rank])
    ddp_optim = torch.optim.Adam(ddp_model.parameters(), lr=1e-2)
    fsdp_optim = torch.optim.Adam(fsdp_model.parameters(), lr=1e-2)
    max_norm = 1
    norm_type = 1
    device = torch.device("cuda")
    for i in range(10):
        ddp_optim.zero_grad(set_to_none=True)
        fsdp_optim.zero_grad(set_to_none=True)
        inp = ddp_model.module.get_input(device)
        for model in (ddp_model, fsdp_model):
            out = model(*inp)
            loss = nn.functional.cross_entropy(
                out.view(-1, out.size(-1)), inp[1].view(-1), reduction="sum"
            )
            loss.backward()
        ddp_total_norm = torch.nn.utils.clip_grad_norm_(
            ddp_model.parameters(),
            max_norm=max_norm,
            norm_type=norm_type,
        )
        fsdp_total_norm = torch.nn.utils.clip_grad_norm_(
            fsdp_model.parameters(),
            max_norm=max_norm,
            norm_type=norm_type,
        )
        self.assertEqual(ddp_total_norm, fsdp_total_norm)
        ddp_flat_grad = torch.cat(tuple(p.grad.flatten() for p in ddp_model.parameters()))
        fsdp_flat_grad = torch.cat(tuple(p.grad.flatten() for p in fsdp_model.parameters()))
        self.assertEqual(ddp_flat_grad, fsdp_flat_grad)
        ddp_flat_param = torch.cat(tuple(p.flatten() for p in ddp_model.parameters()))
        fsdp_flat_param = torch.cat(tuple(p.flatten() for p in fsdp_model.parameters()))
        self.assertEqual(ddp_flat_param, fsdp_flat_param)
        ddp_optim.step()
        fsdp_optim.step()
        ddp_flat_param = torch.cat(tuple(p.flatten() for p in ddp_model.parameters()))
        fsdp_flat_param = torch.cat(tuple(p.flatten() for p in fsdp_model.parameters()))
        self.assertEqual(ddp_flat_param, fsdp_flat_param)
```

On the `i == 3` iteration, the assertion `self.assertEqual(ddp_flat_param, fsdp_flat_param)` *after* the optimizer steps fails.
```
Mismatched elements: 2 / 8427 (0.0%)
Greatest absolute difference: 1.0077477327286033e-05 at index (6610,) (up to 1e-05 allowed)
Greatest relative difference: 8.842818419533154 at index (6610,) (up to 1.3e-06 allowed)
```

The unit test initializes a model (`TransformerWithSharedParams`) and constructs `DDP` and `FSDP` (`NO_SHARD`) instances, which should be semantically equivalent. The _only_ relevant difference should be that FSDP has flattened all parameters into one `FlatParameter`.

We run a training loop that includes `torch.nn.utils.clip_grad_norm_(max_norm=1, norm_type=1)` and uses Adam optimizer. We have 3 checks: (1) gradient elements match after backward and clipping, (2) parameter elements match immediately before optimizer step, and (3) parameter elements match immediately after optimizer step.

Since (1) and (2) pass but (3) does not (on the `i == 3` iteration), this suggests that the optimizer step is not producing the same results. As discussed above, the only difference is that the `fsdp_model` parameters are "bucketed" into a `FlatParameter` (1D containing all the same elements), while the `ddp_model` parameters preserve the original shapes.

A couple of notes:
- [!!] The mismatch does not happen if we pass `use_orig_params=True` to the FSDP constructor. This is a key observation. For `use_orig_params=True`, the optimizer operates on the parameters with their original shapes, just like DDP. This suggests that operating on the flattened parameter is indeed the cause for the difference.
- The mismatch does not happen when using `SGD` instead of `Adam`.
- The mismatch does not happen if we remove the `torch.nn.utils.clip_grad_norm_()`. However, since we have check (1), this should rule out that `clip_grad_norm_()` is producing mismatching results. Rather, we may be relying on `clip_grad_norm_()` to have the gradients be at a sufficiently small magnitude.
- The mismatch also does happen when using `loss = out.sum()` instead of the `cross_entropy` computation.

It requires some nontrivial effort to simplify this repro to be equivalent but not rely on DDP, FSDP, or the FSDP utils from `common_fsdp.py`. I will hold off on that for now.

cc @mrshenli @pritamdamania87 @zhaojuanmao @satgera @rohan-varma @gqchen @aazzolini @osalpekar @jiayisuse @H-Huang @kwen2501