[FSDP] Can not resume the optimizer with params in mulitple groups

## ❓ Questions and Help
I tried to follow the tutorial to change my codes to use FSDP; however, I do not know how to resume the training properly.
Every time I resume, it seems to restart from scratch.

In order to resume, there are three state_dicts,
1. model weights
2. optimizer
3. learning rate scheduler

I checked that 1 and 3 are correctly after I loaded them in the beginning of the resuming; however, it seems that something wrong with the load_state_dict for the optimizer.

I simply use `optimizer.state_dict()` to get the state_dict and save it to disk and `optimizer.load_state_dict(state_dict)` to recover the state dict, is there anything wrong to resume the state dict of optimizer?

I am not sure is it related to #538 and #539 ?

Thanks for your excellent work and help.

Note:
I am using `AdamW` from the pytorch with multiple parameter groups since each group requires different weight decay.