Implement LAMB optimizer

Summary:
Implement the LAMB Optimizer with support (https://arxiv.org/abs/1904.00962).

The reference implementation can be found here: https://github.com/cybertronai/pytorch-lamb

Reviewed By: AkshatSh

Differential Revision: D18725798

fbshipit-source-id: 91b7de6ced02dac2a85f7d9ef6c8d875a7779cd1

Author: ArmenAg

pytext/optimizer/__init__.py

@@ -6,6 +6,7 @@
     FP16OptimizerApex,
     FP16OptimizerFairseq,
 )
+from pytext.optimizer.lamb import Lamb  # noqa
 from pytext.optimizer.optimizers import (  # noqa
     SGD,
     Adagrad,

pytext/optimizer/lamb.py

@@ -0,0 +1,123 @@
+#!/usr/bin/env python3
+# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved
+
+
+import torch
+from pytext.optimizer.optimizers import Optimizer
+from torch.optim import Optimizer as PT_Optimizer
+
+
+class Lamb(Optimizer, PT_Optimizer):
+    r"""Implements Lamb algorithm.
+    THIS WAS DIRECTLY COPIED OVER FROM pytorch/contrib
+            https://github.com/cybertronai/pytorch-lamb
+    It has been proposed in `Large Batch Optimization for Deep Learning: Training BERT in 76 minutes`_.
+
+    Arguments:
+        params (iterable): iterable of parameters to optimize or dicts defining
+            parameter groups
+        lr (float, optional): learning rate (default: 1e-3)
+        betas (Tuple[float, float], optional): coefficients used for computing
+            running averages of gradient and its square (default: (0.9, 0.999))
+        eps (float, optional): term added to the denominator to improve
+            numerical stability (default: 1e-8)
+        weight_decay (float, optional): weight decay (L2 penalty) (default: 0)
+    .. _Large Batch Optimization for Deep Learning: Training BERT in 76 minutes:
+        https://arxiv.org/abs/1904.00962
+    """
+
+    class Config(Optimizer.Config):
+        lr: float = 0.001
+        weight_decay: float = 0.00001
+        eps: float = 1e-8
+
+    @classmethod
+    def from_config(cls, config: Config, model: torch.nn.Module):
+        return cls(
+            model.parameters(),
+            lr=config.lr,
+            weight_decay=config.weight_decay,
+            eps=config.eps,
+        )
+
+    def __init__(self, params, lr=1e-3, betas=(0.9, 0.999), eps=1e-6, weight_decay=0):
+        if not 0.0 <= lr:
+            raise ValueError("Invalid learning rate: {}".format(lr))
+        if not 0.0 <= eps:
+            raise ValueError("Invalid epsilon value: {}".format(eps))
+        if not 0.0 <= betas[0] < 1.0:
+            raise ValueError("Invalid beta parameter at index 0: {}".format(betas[0]))
+        if not 0.0 <= betas[1] < 1.0:
+            raise ValueError("Invalid beta parameter at index 1: {}".format(betas[1]))
+        PT_Optimizer.__init__(
+            self,
+            params,
+            {"lr": lr, "betas": betas, "eps": eps, "weight_decay": weight_decay},
+        )
+
+    def step(self, closure=None):
+        """Performs a single optimization step.
+
+        Arguments:
+            closure (callable, optional): A closure that reevaluates the model
+                and returns the loss.
+        """
+        loss = None
+        if closure is not None:
+            loss = closure()
+
+        for group in self.param_groups:
+            for p in group["params"]:
+                if p.grad is None:
+                    continue
+                grad = p.grad.data
+                if grad.is_sparse:
+                    raise RuntimeError(
+                        "Lamb does not support sparse gradients, consider SparseAdam instad."
+                    )
+
+                state = self.state[p]
+
+                # State initialization
+                if len(state) == 0:
+                    state["step"] = 0
+                    # Exponential moving average of gradient values
+                    state["exp_avg"] = torch.zeros_like(p.data)
+                    # Exponential moving average of squared gradient values
+                    state["exp_avg_sq"] = torch.zeros_like(p.data)
+
+                exp_avg, exp_avg_sq = state["exp_avg"], state["exp_avg_sq"]
+                beta1, beta2 = group["betas"]
+
+                state["step"] += 1
+
+                # Decay the first and second moment running average coefficient
+                # m_t
+                exp_avg.mul_(beta1).add_(1 - beta1, grad)
+                # v_t
+                exp_avg_sq.mul_(beta2).addcmul_(1 - beta2, grad, grad)
+
+                # Paper v3 does not use debiasing.
+                # bias_correction1 = 1 - beta1 ** state['step']
+                # bias_correction2 = 1 - beta2 ** state['step']
+                # Apply bias to lr to avoid broadcast.
+                step_size = group["lr"]
+                # * math.sqrt(bias_correction2) / bias_correction1
+
+                weight_norm = p.data.pow(2).sum().sqrt().clamp(0, 10)
+
+                adam_step = exp_avg / exp_avg_sq.sqrt().add(group["eps"])
+                if group["weight_decay"] != 0:
+                    adam_step.add_(group["weight_decay"], p.data)
+
+                adam_norm = adam_step.pow(2).sum().sqrt()
+                if weight_norm == 0 or adam_norm == 0:
+                    trust_ratio = 1
+                else:
+                    trust_ratio = weight_norm / adam_norm
+                state["weight_norm"] = weight_norm
+                state["adam_norm"] = adam_norm
+                state["trust_ratio"] = trust_ratio
+                p.data.add_(-step_size * trust_ratio, adam_step)
+
+        return loss

pytext/optimizer/swa.py

@@ -7,6 +7,7 @@
 
 import torch
 from pytext.config.component import create_optimizer
+from pytext.optimizer.lamb import Lamb
 from pytext.optimizer.optimizers import SGD, Adagrad, Adam, AdamW, Optimizer
 from pytext.optimizer.radam import RAdam
 from torch.optim import Optimizer as PT_Optimizer
@@ -15,7 +16,12 @@
 class StochasticWeightAveraging(Optimizer, PT_Optimizer):
     class Config(Optimizer.Config):
         optimizer: Union[
-            SGD.Config, Adam.Config, AdamW.Config, Adagrad.Config, RAdam.Config
+            SGD.Config,
+            Adam.Config,
+            AdamW.Config,
+            Adagrad.Config,
+            RAdam.Config,
+            Lamb.Config,
         ] = SGD.Config()
         start: int = 10
         frequency: int = 5

pytext/optimizer/tests/test_swa.py

@@ -10,7 +10,7 @@
 from unittest import TestCase
 
 import torch
-from pytext.optimizer import StochasticWeightAveraging
+from pytext.optimizer import Lamb, StochasticWeightAveraging
 from torch import nn, optim, sparse
 from torch.autograd import Variable
 from torch.utils import data
@@ -675,3 +675,11 @@ def forward(self, x):
         test(CNN, (objects, channels, height, width), objects, "cpu")
         if torch.cuda.is_available():
             test(CNN, (objects, channels, height, width), objects, "cuda")
+
+    def test_lamb(self):
+        def lamb_constructor(params):
+            return StochasticWeightAveraging(
+                Lamb(params, weight_decay=0.01), swa_start=1000, swa_freq=1, swa_lr=1e-2
+            )
+
+        self._test_rosenbrock(lamb_constructor, automode=False)