Inconsistent settings for FSDP Precision

[awaelchli]

Bug description

The FSDPPrecision class controls the dtype of the model parameters, autocasting and grad scaler.
It has two possible inputs: precision="16-mixed" or "bf16-mixed":
https://github.com/Lightning-AI/lightning/blob/6464650a3c91a75d4a7c4009d08a68a84037a9a1/src/lightning/fabric/plugins/precision/fsdp.py#L32-L34

In all other precision plugins, "mixed" refers to mixed precision training with model weights in float32 and autocasting inputs and operations to lower precision. However, the FSDPPrecision plugin sets the dtype of the model parameters to float16/bfloat16 regardless. This means we are actually running "16-true".

https://github.com/Lightning-AI/lightning/blob/6464650a3c91a75d4a7c4009d08a68a84037a9a1/src/lightning/fabric/plugins/precision/fsdp.py#L50-L60

Proposal:

For the mixed precision settings, set param_type=torch.float32. For the current case, introduce the precision settings "16-true" and "bf16-true".

cc @awaelchli @carmocca

[leng-yue]

This bug causes training to fail when performing certain fp32 operations, such as bicubic interpolation. I will submit a pull request later tonight to address this issue.