|
import torch |
|
from torch import nn |
|
|
|
|
|
class DDPGradientStatsHook: |
|
def __init__(self, ddp_module): |
|
try: |
|
ddp_module.register_comm_hook(self, self._hook_fn) |
|
except AttributeError: |
|
raise ValueError('DDPGradientStatsHook does not support non-DDP wrapped modules') |
|
self._clear_state() |
|
|
|
def _clear_state(self): |
|
self.bucket_sq_norms_small_batch = [] |
|
self.bucket_sq_norms_large_batch = [] |
|
|
|
@staticmethod |
|
def _hook_fn(self, bucket): |
|
buf = bucket.buffer() |
|
self.bucket_sq_norms_small_batch.append(buf.pow(2).sum()) |
|
fut = torch.distributed.all_reduce(buf, op=torch.distributed.ReduceOp.AVG, async_op=True).get_future() |
|
def callback(fut): |
|
buf = fut.value()[0] |
|
self.bucket_sq_norms_large_batch.append(buf.pow(2).sum()) |
|
return buf |
|
return fut.then(callback) |
|
|
|
def get_stats(self): |
|
sq_norm_small_batch = sum(self.bucket_sq_norms_small_batch) |
|
sq_norm_large_batch = sum(self.bucket_sq_norms_large_batch) |
|
self._clear_state() |
|
stats = torch.stack([sq_norm_small_batch, sq_norm_large_batch]) |
|
torch.distributed.all_reduce(stats, op=torch.distributed.ReduceOp.AVG) |
|
return stats[0].item(), stats[1].item() |
|
|
|
|
|
class GradientNoiseScale: |
|
"""Calculates the gradient noise scale (1 / SNR), or critical batch size, |
|
from _An Empirical Model of Large-Batch Training_, |
|
https://arxiv.org/abs/1812.06162). |
|
|
|
Args: |
|
beta (float): The decay factor for the exponential moving averages used to |
|
calculate the gradient noise scale. |
|
Default: 0.9998 |
|
eps (float): Added for numerical stability. |
|
Default: 1e-8 |
|
""" |
|
|
|
def __init__(self, beta=0.9998, eps=1e-8): |
|
self.beta = beta |
|
self.eps = eps |
|
self.ema_sq_norm = 0. |
|
self.ema_var = 0. |
|
self.beta_cumprod = 1. |
|
self.gradient_noise_scale = float('nan') |
|
|
|
def state_dict(self): |
|
"""Returns the state of the object as a :class:`dict`.""" |
|
return dict(self.__dict__.items()) |
|
|
|
def load_state_dict(self, state_dict): |
|
"""Loads the object's state. |
|
Args: |
|
state_dict (dict): object state. Should be an object returned |
|
from a call to :meth:`state_dict`. |
|
""" |
|
self.__dict__.update(state_dict) |
|
|
|
def update(self, sq_norm_small_batch, sq_norm_large_batch, n_small_batch, n_large_batch): |
|
"""Updates the state with a new batch's gradient statistics, and returns the |
|
current gradient noise scale. |
|
|
|
Args: |
|
sq_norm_small_batch (float): The mean of the squared 2-norms of microbatch or |
|
per sample gradients. |
|
sq_norm_large_batch (float): The squared 2-norm of the mean of the microbatch or |
|
per sample gradients. |
|
n_small_batch (int): The batch size of the individual microbatch or per sample |
|
gradients (1 if per sample). |
|
n_large_batch (int): The total batch size of the mean of the microbatch or |
|
per sample gradients. |
|
""" |
|
est_sq_norm = (n_large_batch * sq_norm_large_batch - n_small_batch * sq_norm_small_batch) / (n_large_batch - n_small_batch) |
|
est_var = (sq_norm_small_batch - sq_norm_large_batch) / (1 / n_small_batch - 1 / n_large_batch) |
|
self.ema_sq_norm = self.beta * self.ema_sq_norm + (1 - self.beta) * est_sq_norm |
|
self.ema_var = self.beta * self.ema_var + (1 - self.beta) * est_var |
|
self.beta_cumprod *= self.beta |
|
self.gradient_noise_scale = max(self.ema_var, self.eps) / max(self.ema_sq_norm, self.eps) |
|
return self.gradient_noise_scale |
|
|
|
def get_gns(self): |
|
"""Returns the current gradient noise scale.""" |
|
return self.gradient_noise_scale |
|
|
|
def get_stats(self): |
|
"""Returns the current (debiased) estimates of the squared mean gradient |
|
and gradient variance.""" |
|
return self.ema_sq_norm / (1 - self.beta_cumprod), self.ema_var / (1 - self.beta_cumprod) |
|
|
