Support BF16 for FSDP

[yuvalkirstain]

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!

[anj-s]

Thank you for this issue! We are currently working on adding support for bf16 and hope to have it done very soon :)

Assuming that you meant support bf16 with FSDP? Or were you thinking of another API?

[yuvalkirstain]

Exactly, bf16 with FSDP!

[yuvalkirstain]

@anj-s please let me know if there is anything we can do to help, having support for bf16 with FSDP in Fairseq will really really help us! :)

[yuvalkirstain]

Hi, has there been any progress with resolving this issue? @anj-s
Thank you so much

[anj-s]

Hi, has there been any progress with resolving this issue? @anj-s Thank you so much

Hi @yuvalkirstain, I think this should work without any issues. Can you try using bfloat16 by passing the right compute_dtype argument when using FSDP? Unfortunately i haven't had a chance to add a unit test but perhaps someone else on the team has looked into this. cc @anupambhatnagar @min-xu-ai

[wangleiofficial]

bfloat16 support with pytorch lighting will be better, do you have this consideration?

[toriving]

Is there currently any progress on this issue?
Or I'm just wondering if it works if I just apply the above branch.

[anupambhatnagar]

There has been no progress on this so far.