Add lion algorithm

Author: tfaod

reference_algorithms/paper_baselines/lion/__init__.py

@@ -0,0 +1,284 @@
+from __future__ import annotations
+import collections
+from typing import Tuple, Callable, Any, Dict, Iterator, List, Optional
+
+from absl import logging
+import torch
+from torch.optim.optimizer import Optimizer
+
+import torch.distributed.nn as dist_nn
+from torch.optim.lr_scheduler import CosineAnnealingLR
+from torch.optim.lr_scheduler import LinearLR
+from torch.optim.lr_scheduler import SequentialLR
+
+from algoperf import spec
+from algoperf.pytorch_utils import pytorch_setup
+
+USE_PYTORCH_DDP = pytorch_setup()[0]
+
+# default Lion parameters
+HPARAMS = {
+    "dropout_rate": 0.1,
+    "learning_rate": 2e-4,
+    "one_minus_beta1": 0.05,
+    "beta2": 0.98,
+    "weight_decay": 0.5,
+    "warmup_factor": 0.02
+}
+HPARAMS = collections.namedtuple('Hyperparameters', HPARAMS.keys())(**HPARAMS)
+
+# Modified from https://github.com/google/automl/blob/master/lion/lion_pytorch.py.
+class Lion(Optimizer):
+  def __init__(
+    self,
+    params,
+    lr: float = 1e-4,
+    betas: Tuple[float, float] = (0.9, 0.99),
+    weight_decay: float = 0.0,
+  ):
+    if not 0.0 <= lr:
+      raise ValueError('Invalid learning rate: {}'.format(lr))
+    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]))
+    defaults = dict(lr=lr, betas=betas, weight_decay=weight_decay)
+    super().__init__(params, defaults)
+
+  @torch.no_grad()
+  def step(self, closure=None):
+    """Performs a single optimization step.
+
+    Args:
+      closure (callable, optional): A closure that reevaluates the model
+        and returns the loss.
+
+    Returns:
+      the loss.
+    """
+    loss = None
+    if closure is not None:
+      with torch.enable_grad():
+        loss = closure()
+
+    for group in self.param_groups:
+      for p in group['params']:
+        if p.grad is None:
+          continue
+
+        # Perform stepweight decay
+        p.data.mul_(1 - group['lr'] * group['weight_decay'])
+
+        grad = p.grad
+        state = self.state[p]
+        # State initialization
+        if len(state) == 0:
+          # Exponential moving average of gradient values
+          state['exp_avg'] = torch.zeros_like(p)
+
+        exp_avg = state['exp_avg']
+        beta1, beta2 = group['betas']
+
+        # Weight update
+        update = exp_avg * beta1 + grad * (1 - beta1)
+
+        p.add_(update.sign_(), alpha=-group['lr'])
+
+        # Decay the momentum running average coefficient
+        exp_avg.mul_(beta2).add_(grad, alpha=1 - beta2)
+
+    return loss
+
+
+def init_optimizer_state(workload: spec.Workload,
+                         model_params: spec.ParameterContainer,
+                         model_state: spec.ModelAuxiliaryState,
+                         hyperparameters: spec.Hyperparameters,
+                         rng: spec.RandomState) -> spec.OptimizerState:
+  """Creates a Lion optimizer and a learning rate schedule."""
+  del model_state
+  del rng
+  del hyperparameters
+
+  hyperparameters = HPARAMS
+
+  optimizer_state = {
+      'optimizer':
+          Lion(
+              model_params.parameters(),
+              lr=HPARAMS.learning_rate,
+              betas=(1.0 - HPARAMS.one_minus_beta1,
+                     HPARAMS.beta2),
+              weight_decay=HPARAMS.weight_decay)
+  }
+
+  def pytorch_cosine_warmup(step_hint: int, hyperparameters, optimizer):
+    warmup_steps = int(hyperparameters.warmup_factor * step_hint)
+    warmup = LinearLR(
+        optimizer, start_factor=1e-10, end_factor=1., total_iters=warmup_steps)
+    cosine_steps = max(step_hint - warmup_steps, 1)
+    cosine_decay = CosineAnnealingLR(optimizer, T_max=cosine_steps)
+    return SequentialLR(
+        optimizer, schedulers=[warmup, cosine_decay], milestones=[warmup_steps])
+
+  optimizer_state['scheduler'] = pytorch_cosine_warmup(
+      workload.step_hint, HPARAMS, optimizer_state['optimizer'])
+  optimizer_state['hyperparameters'] = hyperparameters
+
+  return optimizer_state
+
+
+def update_params(
+    workload: spec.Workload,
+    current_param_container: spec.ParameterContainer,
+    current_params_types: spec.ParameterTypeTree,
+    model_state: spec.ModelAuxiliaryState,
+    hyperparameters: spec.Hyperparameters,
+    batch: Dict[str, spec.Tensor],
+    loss_type: spec.LossType,
+    optimizer_state: spec.OptimizerState,
+    eval_results: List[Tuple[int, float]],
+    global_step: int,
+    rng: spec.RandomState,
+    train_state: Optional[Dict[str, Any]] = None) -> spec.UpdateReturn:
+  """Return (updated_optimizer_state, updated_params, updated_model_state)."""
+  del current_params_types
+  del loss_type
+  del train_state
+  del eval_results
+  del hyperparameters
+
+  hyperparameters = HPARAMS
+
+  current_model = current_param_container
+  current_model.train()
+  optimizer_state['optimizer'].zero_grad()
+
+  logits_batch, new_model_state = workload.model_fn(
+      params=current_model,
+      augmented_and_preprocessed_input_batch=batch,
+      model_state=model_state,
+      mode=spec.ForwardPassMode.TRAIN,
+      rng=rng,
+      update_batch_norm=True)
+
+  label_smoothing = (
+      hyperparameters.label_smoothing if hasattr(HPARAMS,
+                                                 'label_smoothing') else 0.0)
+  if hasattr(hyperparameters, 'grad_clip'):
+    grad_clip = hyperparameters.grad_clip
+  else:
+    grad_clip = None
+
+  loss_dict = workload.loss_fn(
+      label_batch=batch['targets'],
+      logits_batch=logits_batch,
+      mask_batch=batch.get('weights'),
+      label_smoothing=label_smoothing)
+  summed_loss = loss_dict['summed']
+  n_valid_examples = loss_dict['n_valid_examples']
+  if USE_PYTORCH_DDP:
+    # Use dist_nn.all_reduce to ensure correct loss and gradient scaling.
+    summed_loss = dist_nn.all_reduce(summed_loss)
+    n_valid_examples = dist_nn.all_reduce(n_valid_examples)
+  loss = summed_loss / n_valid_examples
+
+  loss.backward()
+
+  if grad_clip is not None:
+    torch.nn.utils.clip_grad_norm_(
+        current_model.parameters(), max_norm=grad_clip)
+  optimizer_state['optimizer'].step()
+  optimizer_state['scheduler'].step()
+
+  # Log training metrics - loss, grad_norm, batch_size.
+  if global_step <= 100 or global_step % 500 == 0:
+    with torch.no_grad():
+      parameters = [p for p in current_model.parameters() if p.grad is not None]
+      grad_norm = torch.norm(
+          torch.stack([torch.norm(p.grad.detach(), 2) for p in parameters]), 2)
+    if workload.metrics_logger is not None:
+      workload.metrics_logger.append_scalar_metrics(
+          {
+              'loss': loss.item(),
+              'grad_norm': grad_norm.item(),
+          }, global_step)
+    logging.info('%d) loss = %0.3f, grad_norm = %0.3f',
+                 global_step,
+                 loss.item(),
+                 grad_norm.item())
+
+  return (optimizer_state, current_param_container, new_model_state)
+
+
+def prepare_for_eval(workload: spec.Workload,
+                     current_param_container: spec.ParameterContainer,
+                     current_params_types: spec.ParameterTypeTree,
+                     model_state: spec.ModelAuxiliaryState,
+                     hyperparameters: spec.Hyperparameters,
+                     loss_type: spec.LossType,
+                     optimizer_state: spec.OptimizerState,
+                     eval_results: List[Tuple[int, float]],
+                     global_step: int,
+                     rng: spec.RandomState) -> spec.UpdateReturn:
+  """Return (updated_optimizer_state, updated_params)."""
+  del workload
+  del hyperparameters
+  del current_params_types
+  del loss_type
+  del eval_results
+  del global_step
+  del rng
+  return (optimizer_state, current_param_container, model_state)
+
+
+def get_batch_size(workload_name):
+  # Return the global batch size.
+  if hasattr(HPARAMS, "batch_size"):
+     return HPARAMS.batch_size
+  if workload_name == 'criteo1tb':
+    return 262_144
+  elif workload_name == 'fastmri':
+    return 32
+  elif workload_name == 'imagenet_resnet':
+    return 1024
+  elif workload_name == 'imagenet_resnet_silu':
+    return 512
+  elif workload_name == 'imagenet_resnet_gelu':
+    return 512
+  elif workload_name == 'imagenet_vit':
+    return 1024
+  elif workload_name == 'librispeech_conformer':
+    return 256
+  elif workload_name == 'librispeech_deepspeech':
+    return 256
+  elif workload_name == 'ogbg':
+    return 512
+  elif workload_name == 'wmt':
+    return 128
+  elif workload_name == 'mnist':
+    return 16
+  else:
+    raise ValueError(f'Unsupported workload name: {workload_name}.')
+
+
+def data_selection(workload: spec.Workload,
+                   input_queue: Iterator[Dict[str, spec.Tensor]],
+                   optimizer_state: spec.OptimizerState,
+                   current_param_container: spec.ParameterContainer,
+                   model_state: spec.ModelAuxiliaryState,
+                   hyperparameters: spec.Hyperparameters,
+                   global_step: int,
+                   rng: spec.RandomState) -> Dict[str, spec.Tensor]:
+  """Select data from the infinitely repeating, pre-shuffled input queue.
+  Each element of the queue is a batch of training examples and labels.
+  """
+  del workload
+  del optimizer_state
+  del current_param_container
+  del model_state
+  del hyperparameters
+  del global_step
+  del rng
+  batch = next(input_queue)
+  return batch