| 60 | |
| 61 | |
| 62 | class DynamicGradScaler(MegatronGradScaler): |
| 63 | def __init__( |
| 64 | self, |
| 65 | initial_scale, |
| 66 | min_scale, |
| 67 | growth_factor, |
| 68 | backoff_factor, |
| 69 | growth_interval, |
| 70 | hysteresis, |
| 71 | ): |
| 72 | """ "Grad scaler with dynamic scale that gets adjusted |
| 73 | during training.""" |
| 74 | super(DynamicGradScaler, self).__init__(initial_scale) |
| 75 | |
| 76 | # Lower bound on the scale. |
| 77 | assert min_scale > 0.0 |
| 78 | assert min_scale <= initial_scale |
| 79 | self.min_scale = torch.cuda.FloatTensor([min_scale]) |
| 80 | # Growth and backoff factors for the scale. |
| 81 | assert growth_factor > 1.0 |
| 82 | self.growth_factor = torch.cuda.FloatTensor([growth_factor]) |
| 83 | assert backoff_factor < 1.0 |
| 84 | assert backoff_factor > 0.0 |
| 85 | self.backoff_factor = torch.cuda.FloatTensor([backoff_factor]) |
| 86 | # Interval over which if we don't see any inf/nan, |
| 87 | # we will scale the grad scale by the growth factor. |
| 88 | assert growth_interval > 0 |
| 89 | self.growth_interval = growth_interval |
| 90 | # Number of inf/nans we should see before scaling down |
| 91 | # the grad scale by the backoff factor. |
| 92 | assert hysteresis > 0 |
| 93 | self.hysteresis = hysteresis |
| 94 | |
| 95 | # Trackers. |
| 96 | self._growth_tracker = 0 |
| 97 | self._hysteresis_tracker = self.hysteresis |
| 98 | |
| 99 | def update(self, found_inf): |
| 100 | |
| 101 | # If we have an inf/nan, growth tracker is set to 0 |
| 102 | # and hysterisis tracker is reduced by 1. |
| 103 | if found_inf: |
| 104 | self._growth_tracker = 0 |
| 105 | self._hysteresis_tracker -= 1 |
| 106 | # Now if we are out of hysteresis count, scale down the loss. |
| 107 | if self._hysteresis_tracker <= 0: |
| 108 | self._scale = torch.max( |
| 109 | self._scale * self.backoff_factor, self.min_scale |
| 110 | ) |
| 111 | else: |
| 112 | # If there is no nan/inf, increment the growth tracker. |
| 113 | self._growth_tracker += 1 |
| 114 | # If we have had enough consequitive intervals with no nan/inf: |
| 115 | if self._growth_tracker == self.growth_interval: |
| 116 | # Reset the tracker and hysteresis trackers, |
| 117 | self._growth_tracker = 0 |
| 118 | self._hysteresis_tracker = self.hysteresis |
| 119 | # and scale up the loss scale. |
no outgoing calls
no test coverage detected