Support BF16 for FSDP

## Feature Request
Please support BF16 mixed-precision

## Additional context
Training with BF16 is usually more stable than fp16, which is very important when we want to train large models. Additionally, many models (e.g. T5) are trained with BF16 and if we want to continue training them with mixed-precision, using fp16 will result in NaNs.

Thank you!