[NaN check] Add NaN check to support bfloat16.

[ys950902]

No description provided.

[tjruwase]

@ys950902, thanks for helping with this. The problem is a bit more involved and there was a previous attempt that was abandoned. Can you please take a look at #5252

Do you think you can incorporate the learnings into your PR?

[ys950902]

@ys950902, thanks for helping with this. The problem is a bit more involved and there was a previous attempt that was abandoned. Can you please take a look at #5252

Do you think you can incorporate the learnings into your PR?

Sorry for later response, I am doing some other work these days, my understanding is nan check is not for float16 but also bfloat16, and won't add some extra log info on deepspeed side, the uses can only use API was_step_applied() to check whether update successfully.