| 40 | |
| 41 | template <typename T, typename FoundInfFlagT> |
| 42 | inline HOSTDEVICE void Update(const FoundInfFlagT found_inf_data, |
| 43 | const T* pre_loss_scaling_data, |
| 44 | const int* good_in_data, |
| 45 | const int* bad_in_data, |
| 46 | const int incr_every_n_steps, |
| 47 | const int decr_every_n_nan_or_inf, |
| 48 | const float incr_ratio, |
| 49 | const float decr_ratio, |
| 50 | T* updated_loss_scaling_data, |
| 51 | int* good_out_data, |
| 52 | int* bad_out_data) { |
| 53 | if (IsFoundNanInf(found_inf_data)) { |
| 54 | *good_out_data = 0; |
| 55 | *bad_out_data = *bad_in_data + 1; |
| 56 | if (*bad_out_data == decr_every_n_nan_or_inf) { |
| 57 | T new_loss_scaling = *pre_loss_scaling_data * decr_ratio; |
| 58 | *updated_loss_scaling_data = new_loss_scaling < static_cast<T>(1) |
| 59 | ? static_cast<T>(1) |
| 60 | : new_loss_scaling; |
| 61 | *bad_out_data = 0; |
| 62 | } |
| 63 | } else { |
| 64 | *bad_out_data = 0; |
| 65 | *good_out_data = *good_in_data + 1; |
| 66 | if (*good_out_data == incr_every_n_steps) { |
| 67 | T new_loss_scaling = *pre_loss_scaling_data * incr_ratio; |
| 68 | *updated_loss_scaling_data = CheckFinite(new_loss_scaling) |
| 69 | ? new_loss_scaling |
| 70 | : *pre_loss_scaling_data; |
| 71 | *good_out_data = 0; |
| 72 | } |
| 73 | } |
| 74 | } |
| 75 | |
| 76 | template <typename Context, typename T> |
| 77 | class LazyZeros { |
no test coverage detected