"Grad scaler with dynamic scale that gets adjusted during training.
(
self,
initial_scale,
min_scale,
growth_factor,
backoff_factor,
growth_interval,
hysteresis,
)
| 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 |
nothing calls this directly
no outgoing calls
no test coverage detected