| |
| |
| |
| |
|
|
| import logging |
|
|
| import torch |
| from fairseq import optim |
| from omegaconf import DictConfig |
|
|
| logger = logging.getLogger(__name__) |
|
|
|
|
| class AMPOptimizer(optim.FairseqOptimizer): |
| """ |
| Wrap an *optimizer* to support AMP (automatic mixed precision) training. |
| """ |
|
|
| def __init__(self, cfg: DictConfig, params, fp32_optimizer, **kwargs): |
| super().__init__(cfg.optimizer) |
| self.fp32_optimizer = fp32_optimizer |
| amp_kwargs = {"init_scale": cfg.common.fp16_init_scale} |
| if getattr(cfg.common, "amp_scale_window", None) is not None: |
| amp_kwargs["growth_interval"] = cfg.common.amp_init_scale |
| self._grad_scaler = torch.cuda.amp.GradScaler(**amp_kwargs) |
| self.min_loss_scale = cfg.common.min_loss_scale |
|
|
| @classmethod |
| def build_optimizer(cls, cfg: DictConfig, params, **kwargs): |
| """ |
| Args: |
| cfg (omegaconf.DictConfig): fairseq args |
| params (iterable): iterable of parameters to optimize |
| """ |
| fp32_optimizer = optim.build_optimizer(cfg.optimizer, params) |
| return cls(cfg, params, fp32_optimizer, **kwargs) |
|
|
| def backward(self, loss): |
| """Computes the sum of gradients of the given tensor w.r.t. graph leaves. |
| |
| Compared to :func:`fairseq.optim.FairseqOptimizer.backward`, this |
| function additionally dynamically scales the loss to avoid gradient |
| underflow. |
| """ |
| self._grad_scaler.scale(loss).backward() |
|
|
| def step(self): |
| self.scaler.step(self.fp32_optimizer) |
| self.scaler.update() |
|
|
| def clip_grad_norm(self, max_norm, aggregate_norm_fn=None): |
| """Clips gradient norm.""" |
| self.scaler.unscale_(self.optimizer) |
| grad_norm = self.fp32_optimizer.clip_grad_norm(max_norm, aggregate_norm_fn) |
| if not torch.isfinite(grad_norm).all(): |
| new_loss_scale = self.next_loss_scale |
| if new_loss_scale <= self.min_loss_scale: |
| raise FloatingPointError( |
| ( |
| "AMP: Minimum loss scale reached ({}). Your loss is probably exploding. " |
| "Try restarting training or use fp32. {}" |
| ).format(self.min_loss_scale, new_loss_scale) |
| ) |
| else: |
| logger.info( |
| "AMP: overflow detected, setting scale to " f"to {new_loss_scale}" |
| ) |
| return grad_norm |
|
|
| @property |
| def scaler(self): |
| return self._grad_scaler |
|
|
| @property |
| def next_loss_scale(self): |
| return self.scaler.get_scale() * self.scaler.get_backoff_factor() |
|
|
| @property |
| def optimizer(self): |
| return self.fp32_optimizer.optimizer |
|
|
| @optimizer.setter |
| def optimizer(self, optimizer): |
| self.fp32_optimizer.optimizer = optimizer |
|
|
| @property |
| def lr_scheduler(self): |
| return getattr(self.fp32_optimizer, "lr_scheduler", None) |
|
|
| @property |
| def optimizer_config(self): |
| return self.fp32_optimizer.optimizer_config |
|
|
| def get_lr(self): |
| return self.fp32_optimizer.get_lr() |
|
|
| def set_lr(self, lr): |
| self.fp32_optimizer.set_lr(lr) |
|
|
| def all_reduce_grads(self, module): |
| self.fp32_optimizer.all_reduce_grads(module) |
|
|
| @property |
| def supports_flat_params(self): |
| return self.fp32_optimizer.supports_flat_params |
|
|
