[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.

[QingtaoLi1]

@ys950902 This check only makes the training continues without error but the grad_norm and loss are not descending any more. Would it be better to turn this into an error or provide a way to solve the overflow problem by code/by warning?

[tjruwase]

Closing in place of #6976