t5x/adafactor.py

# Copyright 2022 The T5X Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Adafactor Optimizer.

Specialized Adafactor implementation for T5X with:
 - custom factorization specification rules.
 - support for stacked parameters from scanned layers and parameter fusions.

Why do we need custom factorization?  In the Adafactor paper, scalar, vector and
matrix parameters are considered. This is sufficiently general because higher
dimensional parameters can be reshaped. In practice, there are situations where
higher dimensional parameters are desirable.  For example, consider the
multi-headed attention. It has projection kernels.  This is naturally
represented as 3-dimensional array [d_model, num_head, head_dim]. Keeping the
3-dimensional structure can be beneficial for performance optimization, e.g., by
giving compilers additional degree of freedom to do layout optimization.

The default heuristic behavior for the second-moment estimator can lead to an
unexpected result because it assumes that the parameters are matrices (vectors
and scalars are not factored). The dimensions are sorted and the smaller
dimension is assigned to the row dim and the larger dim to the col dim (unless
the two largest dims have an equal size and then the original ordering of the
dimensions is used). Then `v_row` (i.e., the optimizer state for the row) is
obtained by removing the col dim. In other words, `rank(v_row) = rank(v) - 1`.
If the parameter is higher dimensional, v_row and v_col are higher dimensional.
Therefore, the outer product of v_row and v_col do not necessarily corresponds
to the row rank approximation that minimizes the generalized Kullback-Leibler
divergence (the original Adafactor formulation).

This Adafactor implementation generalized the default behavior such that we
obtain the correct second moment estimator even for higher dimensional
parameters.

"""
import enum
import re
from typing import Any, Mapping, Optional, Sequence, Tuple, Union

from absl import logging
from flax import struct
from flax.core import freeze
from flax.core import FrozenDict
from flax.core import unfreeze
from flax.serialization import from_state_dict
from flax.serialization import to_state_dict
from flax.traverse_util import flatten_dict
from flax.traverse_util import unflatten_dict
import jax
import jax.numpy as jnp
import numpy as np
from t5x import utils
from t5x.optimizers import OptimizerDef
from t5x.optimizers import OptimizerState

Dtype = Any


class FactorDim(enum.Enum):
  # Don't factorize this dimension.
  NONE = None
  # A batch-like dimension that we should not average over.
  BATCH = 1
  ROW = 2
  COLUMN = 3


# Sentinel value signifying the legacy heuristic factorization rule.
class HeuristicRule(enum.Enum):
  token = 1


HEURISTIC_RULE = HeuristicRule.token
FactorRule = Union[HeuristicRule, Tuple[FactorDim]]


def _restore(target, flat):
  state_dict = unflatten_dict({tuple(k.split('/')): v for k, v in flat.items()})
  if isinstance(target, FrozenDict):
    return freeze(state_dict)
  else:
    return state_dict


def _insert(tpl, idx, x):
  tmp = list(tpl)
  tmp.insert(idx, x)
  return tuple(tmp)


def standard_logical_factor_rules():
  return freeze({
      'vocab': FactorDim.COLUMN,
      'embed': FactorDim.ROW,
      'mlp': FactorDim.COLUMN,
      'heads': FactorDim.COLUMN,
      'kv': FactorDim.COLUMN,
      'joined_kv': FactorDim.COLUMN,
      'relpos_buckets': FactorDim.NONE,
      'layers': FactorDim.BATCH,  # used in scanned layers
      'stack': FactorDim.BATCH,  # used in stacked params
      # 'batch', 'length' should not occur in parameters
      'q_wi_fused': FactorDim.COLUMN,
      'o_wo_fused': FactorDim.COLUMN,
      'multiquery_heads': FactorDim.COLUMN,
      'kv_fused': FactorDim.COLUMN,
      'layer_norm_scale': FactorDim.NONE,
      'mlp_activations': FactorDim.COLUMN,
  })


def factor_name_to_factordim(name):
  if not isinstance(name, str):
    return name
  name = name.lower()
  return {
      'row': FactorDim.ROW,
      'col': FactorDim.COLUMN,
      'column': FactorDim.COLUMN,
      'batch': FactorDim.BATCH,
      'none': FactorDim.NONE,
      'unfactorized': FactorDim.NONE
  }[name]


class HParamMap:
  """Maps parameter path names to hparams.

  Names of parameters nested in a PyTree (e.g., an Optimizer) are formed by
  joining the names along the path to the parameter leaf with '/'.
  """

  def __init__(self, rules):
    self._rules = [(re.compile(r), p) for r, p in rules]

  def __getitem__(self, key: str) -> Any:
    for r, p in self._rules:
      if r.search(key):
        return p
    raise KeyError(f'No factor rule found for parameter: {key}')

  def __call__(self, params):
    """Returns a copy of the params with mapped hparams in leaves."""
    flat_state_dict = flatten_dict(to_state_dict(params))
    flat_rules_dict = {k: self['/'.join(k)] for k in flat_state_dict.keys()}
    return from_state_dict(params, unflatten_dict(flat_rules_dict))


@struct.dataclass
class _AdafactorHyperParams:
  """Hparams for Adafactor optimizer."""
  learning_rate: Optional[float]
  factored: bool
  multiply_by_parameter_scale: Union[bool, HParamMap]
  beta1: Optional[float]
  decay_rate: float
  step_offset: int
  clipping_threshold: Optional[float]
  weight_decay_rate: Optional[float]
  min_dim_size_to_factor: int
  epsilon1: float
  epsilon2: float
  factor_map: Optional[HParamMap] = None
  logical_factor_rules: Any = None
  weight_decay_rate_lr_exponent: Optional[float] = None
  global_norm_clip_threshold: Optional[float] = None
  max_parameter_scale: Optional[float] = None
  skip_nan_updates: Optional[bool] = False


@struct.dataclass
class _AdafactorParamState:
  v_row: np.ndarray  # used in normal factored version
  v_col: np.ndarray
  v: np.ndarray  # only used without factoring
  m: np.ndarray  # only used with momentum


class Adafactor(OptimizerDef):
  """Adafactor optimizer.

  Adafactor is described in https://arxiv.org/abs/1804.04235.
  """

  def __init__(self,
               learning_rate: Optional[float] = None,
               factored: bool = True,
               multiply_by_parameter_scale: Union[bool, HParamMap] = True,
               beta1: Optional[float] = None,
               decay_rate: float = 0.8,
               step_offset: int = 0,
               clipping_threshold: Optional[float] = 1.0,
               weight_decay_rate: Optional[float] = None,
               min_dim_size_to_factor: int = 128,
               epsilon1: float = 1e-30,
               epsilon2: float = 1e-3,
               dtype_momentum: Dtype = jnp.float32,
               factor_map: Optional[HParamMap] = None,
               logical_factor_rules: Optional[Mapping[str, FactorDim]] = None,
               weight_decay_rate_lr_exponent: Optional[float] = None,
               global_norm_clip_threshold: Optional[float] = None,
               max_parameter_scale: Optional[float] = None,
               skip_nan_updates: Optional[bool] = False):
    """Constructor for the Adafactor optimizer.


    Args:
      learning_rate: float: learning rate.  NB: the natural scale for adafactor
        LR is markedly different from Adam, one doesn't use the 1/sqrt(hidden)
        correction for this optimizer with attention-based models.
      factored: boolean: whether to use factored second-moment estimator for 2d
        variables.
      multiply_by_parameter_scale: boolean: if True, then scale provided
        learning_rate by parameter norm. if False, provided learning_rate is
        absolute step size.
      beta1: an optional float value between 0 and 1, enables momentum and uses
        extra memory if non-None! None by default.
      decay_rate: float: controls second-moment exponential decay schedule.
      step_offset: for finetuning, one may optionally set this to the starting
        step-number of the finetuning phase to reset the second moment
        accumulators after pretraining. Does not affect the momentum even if it
        was used during pretraining.
      clipping_threshold: an optional float >= 1, if None no update clipping.
      weight_decay_rate: optional rate at which to decay weights.
      min_dim_size_to_factor: only factor accumulator if two array dimensions
        are at least this size.
      epsilon1: Regularization constant for squared gradient.
      epsilon2: Regularization constant for parameter scale.
      dtype_momentum: dtype of momentum buffers.
      factor_map: hparam-map from key path to manual factorization rules.
      logical_factor_rules: factorization rules provided as a set of mappings
        from logical axis name to ROW, COLUMN, BATCH, or NONE.  Supercedes
        factor_map if `set_param_axes` is called.
      weight_decay_rate_lr_exponent: If present, weight decay rate is computed
        as (learning_rate ** weight_decay_rate_lr_exponent).  If
        weight_decay_rate is also present, then multiply by it.
      global_norm_clip_threshold: If set, will clip gradients by global norm
        before Adafactor stats are applied.
      max_parameter_scale: If set, clips the parameter scale to a maximum value,
        which helps prevent parameters from growing without bound.
      skip_nan_updates: If set, any parameter that would have been updated by a
        NaN value after a applying gradients will be kept with the earlier
        value it had.
    """
    if not factored and factor_map is not None:
      raise ValueError('Adafactor factored is False but factorization rules '
                       'have been provided.')
    if not isinstance(multiply_by_parameter_scale, (bool, HParamMap)):
      raise TypeError(
          '`multiply_by_parameter_scale` must be either bool or `HParamMap` '
          f'type. Got {type(multiply_by_parameter_scale)}')

    if not isinstance(factor_map, (type(None), HParamMap)):
      raise TypeError(
          '`factor_map` must be either None or `HParamMap` type. Got '
          f'{type(factor_map)}')

    hyper_params = _AdafactorHyperParams(
        learning_rate, factored, multiply_by_parameter_scale, beta1, decay_rate,
        step_offset, clipping_threshold, weight_decay_rate,
        min_dim_size_to_factor, epsilon1, epsilon2, factor_map,
        logical_factor_rules, weight_decay_rate_lr_exponent,
        global_norm_clip_threshold, max_parameter_scale, skip_nan_updates)
    self.dtype_momentum = jax.dtypes.canonicalize_dtype(dtype_momentum)
    super().__init__(hyper_params)

  @staticmethod
  def _decay_rate_pow(i: int, exponent: float = 0.8) -> float:
    """Default Adafactor second-moment decay schedule."""
    t = jnp.array(i, jnp.float32) + 1.0
    return 1.0 - t**(-exponent)

  @staticmethod
  def _parse_rule(
      rule: Optional[FactorRule],
      shape: Sequence[int],
      path: str,
      fallback_to_heuristics=True
  ) -> Tuple[Tuple[int, ...], Optional[Union[HeuristicRule, Tuple[Tuple[
      int, ...], Tuple[int, ...]]]]]:
    """Parses specification and return factored dims and dims for averaging.

    Adafactor needs to know the two largest dimensions to factorize along.
    Traditionally it used a heuristic, but we want finer control over these
    factorization dimensions.  Additionally, there are situations where
    parameters are batched together for e.g. scanned layers and QKV fusion,
    and we want to ensure that the scale updates and clipping thresholds are
    calculated _within_ each array and not across the entire batched array.

    Args:
      rule: the rule is either None (default to heuristic behavior) or a tuple
        of the same rank as the `param` array containing a FactorDim.ROW or
        FactorDim.COLUMN to mark dimensions to factorize in two row and column
        sets, and optionally dimensions marked FactorDim.BATCH to denote batched
        dimensions that should not be averaged over. e.g. (BATCH, ROW, COLUMN,
        COLUMN)
      shape: shape of the variable
      path: '/' joined parameter path.
      fallback_to_heuristics: whether to fallback to heuristic factorization
        rule. For most cases this should be set to `True`.

    Returns:
      tuple of: tuple of dimensions to average over, 2-tuple of dimensions to
      factorize over.
    """
    param_ndim = len(shape)

    if rule is None:
      # No factorization.
      return tuple(np.arange(param_ndim)), None

    if rule is HEURISTIC_RULE:
      if param_ndim > 2:
        raise ValueError(
            f'A parameter with rank strictly higher than 2 must have an '
            f'explicit factorization rule: {path}, {shape}')
      # Even if no explicit rule is provided for the param, we still want to
      # average over all the dimensions for computing the RMS scale.
      return tuple(np.arange(param_ndim)), HEURISTIC_RULE

    if len(rule) != param_ndim:
      raise ValueError(f'Factorization rule {rule} has incorrect rank '
                       f'for param of rank {param_ndim}: {path}, {shape}')

    row_dims = tuple(idx for idx, d in enumerate(rule) if d == FactorDim.ROW)
    col_dims = tuple(idx for idx, d in enumerate(rule) if d == FactorDim.COLUMN)
    batched_dims = tuple(
        idx for idx, d in enumerate(rule) if d == FactorDim.BATCH)
    averaging_dims = tuple(np.delete(np.arange(param_ndim), batched_dims))
    factor_dims = (row_dims, col_dims)
    if factor_dims == ((), ()):
      factor_dims = None

    if fallback_to_heuristics and param_ndim <= 2 and not batched_dims:
      logging.warning(
          'Since rank of parameter %s %d is less than or equal to 2, the '
          'factorization method falls back to heuristics and the provided '
          'factor rule %s is ignored.', path, param_ndim, rule)
      return tuple(np.arange(param_ndim)), HEURISTIC_RULE

    return averaging_dims, factor_dims

  def _factored_dims(
      self, shape: Sequence[int]) -> Optional[Tuple[Tuple[int], Tuple[int]]]:
    """Whether to use a factored second moment estimator.

    If there are not two dimensions of size >= min_dim_size_to_factor, then we
    do not factor. If we do factor the accumulator, then this function returns a
    tuple of the two largest axes to reduce over.

    Args:
      shape: a Shape

    Returns:
      None or a tuple of ints
    """
    if not self.hyper_params.factored or len(shape) < 2:
      return None
    sorted_dims = np.argsort(shape)
    if shape[sorted_dims[-2]] < self.hyper_params.min_dim_size_to_factor:
      return None
    return (int(sorted_dims[-2]),), (int(sorted_dims[-1]),)

  def init_param_state(self, param, path):
    shape = param.shape
    state = {k: jnp.zeros((1,)) for k in ['v_row', 'v_col', 'v', 'm']}
    if self.hyper_params.factored:
      factor_rule = (
          self.hyper_params.factor_map[path]
          if self.hyper_params.factor_map else HEURISTIC_RULE)
    else:
      factor_rule = None
    _, factored_dims = self._parse_rule(factor_rule, param.shape, path)
    if factored_dims is HEURISTIC_RULE:
      factored_dims = self._factored_dims(shape)
    if factored_dims is not None:
      d1, d0 = factored_dims
      vr_shape = np.delete(shape, d0)
      vc_shape = np.delete(shape, d1)
      state['v_row'] = jnp.zeros(vr_shape, dtype=jnp.float32)
      state['v_col'] = jnp.zeros(vc_shape, dtype=jnp.float32)
    else:
      state['v'] = jnp.zeros(param.shape, dtype=jnp.float32)
    if self.hyper_params.beta1 is not None:
      state['m'] = jnp.zeros(param.shape, dtype=self.dtype_momentum)
    return _AdafactorParamState(**state)

  def init_state(self, params):
    params_flat = utils.flatten_dict_string_keys(params)
    param_states_flat = [
        self.init_param_state(param, path)
        for path, param in params_flat.items()
    ]
    param_states_flat = {
        k: v for k, v in zip(params_flat.keys(), param_states_flat)
    }
    param_states = _restore(params, param_states_flat)
    state = OptimizerState(jnp.asarray(0, dtype=jnp.int32), param_states)
    return state

  def apply_param_gradient(self, step, hyper_params, param, state, grad, path):
    assert hyper_params.learning_rate is not None, 'no learning rate provided.'
    learning_rate = hyper_params.learning_rate
    beta1 = hyper_params.beta1
    decay_rate = hyper_params.decay_rate
    step_offset = hyper_params.step_offset
    multiply_by_parameter_scale = hyper_params.multiply_by_parameter_scale
    max_parameter_scale = hyper_params.max_parameter_scale
    clipping_threshold = hyper_params.clipping_threshold
    weight_decay_rate = hyper_params.weight_decay_rate
    epsilon1 = hyper_params.epsilon1
    epsilon2 = hyper_params.epsilon2
    if hyper_params.weight_decay_rate_lr_exponent:
      weight_decay_rate = (
          (weight_decay_rate or 1.0) *
          learning_rate**hyper_params.weight_decay_rate_lr_exponent)

    if self.hyper_params.factored:
      factor_rule = (
          self.hyper_params.factor_map[path]
          if self.hyper_params.factor_map else HEURISTIC_RULE)
    else:
      factor_rule = None
    averaging_dims, factored_dims = self._parse_rule(factor_rule, param.shape,
                                                     path)

    grad = grad.astype(jnp.float32)

    updates = {k: jnp.zeros((1,)) for k in ['v_row', 'v_col', 'v', 'm']}
    decay_rate = self._decay_rate_pow(step - step_offset, exponent=decay_rate)
    update_scale = learning_rate

    if isinstance(multiply_by_parameter_scale, HParamMap):
      multiply_by_parameter_scale = multiply_by_parameter_scale[path]
    if multiply_by_parameter_scale:
      param_scale = jnp.sqrt(
          jnp.mean(param * param, axis=averaging_dims, keepdims=True))
      # Clip param_scale to a minimum value of epsilon2.
      param_scale = jnp.maximum(param_scale, epsilon2)
      # Clip param_scale to a maximum value, if specified.
      if max_parameter_scale is not None:
        param_scale = jnp.minimum(param_scale, max_parameter_scale)
      update_scale *= param_scale
    mixing_rate = 1.0 - decay_rate

    grad_sqr = grad * grad + epsilon1
    if factored_dims is HEURISTIC_RULE:
      factored_dims = self._factored_dims(param.shape)
    if factored_dims is not None:
      d1, d0 = factored_dims
      new_v_row = (
          decay_rate * state.v_row + mixing_rate * jnp.mean(grad_sqr, axis=d0))
      new_v_col = (
          decay_rate * state.v_col + mixing_rate * jnp.mean(grad_sqr, axis=d1))
      updates['v_row'] = new_v_row
      updates['v_col'] = new_v_col
      reduced_d1 = tuple(d - len([e for e in d0 if e < d]) for d in d1)

      row_col_mean = jnp.mean(new_v_row, axis=reduced_d1, keepdims=True)
      row_factor = (new_v_row / row_col_mean)**-0.5
      col_factor = (new_v_col)**-0.5
      y = (
          grad * jnp.expand_dims(row_factor, axis=d0) *
          jnp.expand_dims(col_factor, axis=d1))
    else:
      new_v = decay_rate * state.v + mixing_rate * grad_sqr
      updates['v'] = new_v
      y = grad * (new_v)**-0.5

    if clipping_threshold is not None:
      clipping_denom = (
          jnp.maximum(
              1.0,
              jnp.sqrt(jnp.mean(y * y, axis=averaging_dims, keepdims=True)) /
              clipping_threshold))
      y /= clipping_denom

    subtrahend = update_scale * y
    if beta1 is not None:
      new_m = beta1 * state.m + (1.0 - beta1) * subtrahend
      subtrahend = new_m
      updates['m'] = new_m.astype(self.dtype_momentum)

    if weight_decay_rate is not None:
      new_param = (1.0 - weight_decay_rate) * param - subtrahend
    else:
      new_param = param - subtrahend

    if hyper_params.skip_nan_updates:
      updates['v_row'] = jnp.where(
          jnp.isnan(updates['v_row']), state.v_row, updates['v_row'])
      updates['v_col'] = jnp.where(
          jnp.isnan(updates['v_col']), state.v_col, updates['v_col'])
      updates['v'] = jnp.where(jnp.isnan(updates['v']), state.v, updates['v'])
      updates['m'] = jnp.where(jnp.isnan(updates['m']), state.m, updates['m'])
      new_param = jnp.where(jnp.isnan(new_param), param, new_param)
    new_state = _AdafactorParamState(**updates)

    return new_param.astype(param.dtype), new_state

  def apply_gradient(self, hyper_params, params, state, grads):
    """Applies a gradient for a set of parameters.

    Args:
      hyper_params: a named tuple of hyper parameters.
      params: the parameters that should be updated.
      state: a named tuple containing the state of the optimizer
      grads: the gradient tensors for the parameters.

    Returns:
      A tuple containing the new parameters and the new optimizer state.
    """
    step = state.step
    # We assume that params, param_states, and grads are all dict-like here.
    params_flat_dict = utils.flatten_dict_string_keys(params)
    params_paths = params_flat_dict.keys()
    params_flat = params_flat_dict.values()
    # extra paranoia to guarantee identical value ordering
    states_flat = utils.flatten_dict_string_keys(state.param_states)
    states_flat = [states_flat[k] for k in params_paths]
    grads_flat = utils.flatten_dict_string_keys(grads)
    grads_flat = [grads_flat[k] for k in params_paths]

    if hyper_params.global_norm_clip_threshold:
      # Paper: http://proceedings.mlr.press/v28/pascanu13.pdf
      # TF: https://www.tensorflow.org/api_docs/python/tf/clip_by_global_norm
      squared_l2_norms = [jnp.sum(jnp.square(g)) for g in grads_flat]
      global_norm = jnp.sqrt(jnp.sum(jnp.array(squared_l2_norms)))
      scale = hyper_params.global_norm_clip_threshold * jnp.minimum(
          1.0 / hyper_params.global_norm_clip_threshold, 1.0 / global_norm)
      grads_flat = [g * scale for g in grads_flat]

    out = [
        self.apply_param_gradient(step, hyper_params, param, state, grad, path)
        for param, state, grad, path in zip(params_flat, states_flat,
                                            grads_flat, params_paths)
    ]

    new_params_flat, new_states_flat = list(zip(*out)) if out else ((), ())
    new_params_flat = {k: v for k, v in zip(params_paths, new_params_flat)}
    new_states_flat = {k: v for k, v in zip(params_paths, new_states_flat)}
    new_params = _restore(params, new_params_flat)
    new_param_states = _restore(params, new_states_flat)
    new_state = OptimizerState(step + 1, new_param_states)

    return new_params, new_state

  def set_param_axes(self, param_logical_axes):
    """Sets Adafactor factorization map from logical axis names tree."""
    logical_factor_rules = self.hyper_params.logical_factor_rules
    if logical_factor_rules is None:
      return

    # pylint:disable=invalid-name
    NONE = FactorDim.NONE
    COLUMN = FactorDim.COLUMN
    ROW = FactorDim.ROW

    # pylint:enable=invalid-name

    def apply_rules(axes):
      # Partially factorized params are marked as unfactorized, preserving
      # only BATCH axis annotations. We also check for incompletely factorized
      # params that have ROW, COLUMN but also accidental NONE dimensions and
      # raise an error in that case.
      axis_rules = tuple(logical_factor_rules[x] for x in axes)
      axis_rules = tuple(factor_name_to_factordim(x) for x in axis_rules)
      if ROW in axis_rules and COLUMN in axis_rules and NONE in axis_rules:
        raise ValueError(f'Incomplete adafactor spec {axis_rules} for {axes}!')
      if ROW not in axis_rules or COLUMN not in axis_rules:
        axis_rules = tuple(
            NONE if x in (ROW, COLUMN) else x for x in axis_rules)
      return axis_rules

    factor_map = jax.tree_map(apply_rules, param_logical_axes)
    factor_map = utils.flatten_dict_string_keys(factor_map)

    self.hyper_params = self.hyper_params.replace(factor_map=factor_map)

  def derive_logical_axes(self, optimizer_state, param_logical_axes):
    """Derives optimizer logical partitioning from model logical partitions."""
    optimizer_logical_axes = jax.tree_map(lambda x: None,
                                          optimizer_state.state_dict())
    optimizer_logical_axes['target'] = param_logical_axes

    def factor_rule(logical_axes, adafactor_leaf):
      return dict(
          v_row=None,
          v_col=None,
          v=logical_axes if adafactor_leaf['v'].shape != (1,) else None,
          m=logical_axes if self.hyper_params.beta1 else None)

    optimizer_logical_axes['state']['param_states'] = jax.tree_map(
        factor_rule, unfreeze(param_logical_axes),
        optimizer_state.state_dict()['state']['param_states'])

    return optimizer_state.restore_state(unfreeze(optimizer_logical_axes))