juancopi81's picture
Add t5x and mt3 models
b100e1c
raw
history blame
24.9 kB
# 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))