Repository: pytorch
Issue: get_model_state_dict failed after FSDP_model.to(dtype)
Labels: oncall: distributed,triaged,module: fsdp

### 🐛 Describe the bug

```python
# fsdp_model with mixed precision (fp32 parameters)
fsdp_model.to(torch.bfloat16)
save_policy = FullStateDictConfig(offload_to_cpu=True, rank0_only=True)
with FSDP.state_dict_type(model, StateDictType.FULL_STATE_DICT, save_policy):
    return fsdp_model.state_dict()
```

The above code fails with error:

```
[rank0]: Traceback (most recent call last):
......
[rank0]:   File "xxx", line 118, in xxx
[rank0]:     return model.state_dict()
[rank0]:   File "/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py", line 2199, in state_dict
[rank0]:     hook(self, prefix, keep_vars)
[rank0]:   File "/usr/local/lib/python3.10/dist-packages/torch/utils/_contextlib.py", line 116, in decorate_context
[rank0]:     return func(*args, **kwargs)
[rank0]:   File "/usr/local/lib/python3.10/dist-packages/torch/distributed/fsdp/_state_dict_utils.py", line 782, in _pre_state_dict_hook
[rank0]:     _pre_state_dict_hook_fn[fsdp_state._state_dict_type](
[rank0]:   File "/usr/local/lib/python3.10/dist-packages/torch/distributed/fsdp/_state_dict_utils.py", line 303, in _full_pre_state_dict_hook
[rank0]:     _common_unshard_pre_state_dict_hook(
[rank0]:   File "/usr/local/lib/python3.10/dist-packages/torch/distributed/fsdp/_state_dict_utils.py", line 170, in _common_unshard_pre_state_dict_hook
[rank0]:     _enter_unshard_params_ctx(
[rank0]:   File "/usr/local/lib/python3.10/dist-packages/torch/distributed/fsdp/_state_dict_utils.py", line 134, in _enter_unshard_params_ctx
[rank0]:     fsdp_state._unshard_params_ctx[module].__enter__()
[rank0]:   File "/usr/lib/python3.10/contextlib.py", line 135, in __enter__
[rank0]:     return next(self.gen)
[rank0]:   File "/usr/local/lib/python3.10/dist-packages/torch/distributed/fsdp/_unshard_param_utils.py", line 198, in _unshard_fsdp_state_params
[rank0]:     _unshard(state, handle, computation_stream, computation_stream)
[rank0]:   File "/usr/local/lib/python3.10/dist-packages/torch/distributed/fsdp/_runtime_utils.py", line 301, in _unshard
[rank0]:     handle.unshard()
[rank0]:   File "/usr/local/lib/python3.10/dist-packages/torch/distributed/fsdp/_flat_param.py", line 1312, in unshard
[rank0]:     padded_unsharded_flat_param = self._all_gather_flat_param(unsharded_flat_param)
[rank0]:   File "/usr/local/lib/python3.10/dist-packages/torch/distributed/fsdp/_flat_param.py", line 1403, in _all_gather_flat_param
[rank0]:     dist.all_gather_into_tensor(
[rank0]:   File "/usr/local/lib/python3.10/dist-packages/torch/distributed/c10d_logger.py", line 83, in wrapper
[rank0]:     return func(*args, **kwargs)
[rank0]:   File "/usr/local/lib/python3.10/dist-packages/torch/distributed/distributed_c10d.py", line 3410, in all_gather_into_tensor
[rank0]:     work = group._allgather_base(output_tensor, input_tensor, opts)
[rank0]: TypeError: output tensor must have the same type as input tensor
```

### Versions

torch 2.4+

cc @H-Huang @awgu @kwen2501 @wanchaol @fegin @fduwjj @wz337 @wconstab @d4l3k @c-p-i-o @zhaojuanmao @mrshenli @rohan-varma @chauhang