# 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. """General utility functions for t5x.""" import collections.abc from concurrent.futures import thread import contextlib import dataclasses import functools import importlib import inspect import os import re import time import typing from typing import Any, Callable, Iterable, Mapping, Optional, Sequence, Tuple, Type, Union import warnings from absl import logging import clu.data from flax import traverse_util import flax.core from flax.core import scope as flax_scope from flax.linen import partitioning as flax_partitioning import jax from jax import prng from jax import pxla from jax.experimental import multihost_utils from jax.experimental.global_device_array import GlobalDeviceArray import jax.numpy as jnp import numpy as np import orbax.checkpoint import seqio from t5x import checkpoints from t5x import optimizers from t5x import partitioning from t5x import state_utils from t5x import train_state as train_state_lib import tensorflow as tf from tensorflow.io import gfile import typing_extensions Array = Union[np.ndarray, jnp.ndarray, jax.pxla.ShardedDeviceArray, tf.Tensor] PyTreeDef = type(jax.tree_structure(None)) PartitionSpec = partitioning.PartitionSpec DType = Union[np.dtype, type(jnp.bfloat16)] Shape = Tuple[int, ...] # TODO(adarob): Remove namespace mapping after client gin files are updated. TensorBoardLogger = seqio.TensorBoardLogger # ----------------------------------------------------------------------------- # Configurations # ----------------------------------------------------------------------------- @dataclasses.dataclass class SaveCheckpointConfig: """Configuration for saving model checkpoints.""" # The dtype to save ('float32' or 'bfloat16'). dtype: str = 'float32' # Number of steps between writing checkpoints. period: Optional[int] = None # Number of most recent checkpoints to keep, or None to keep them all. keep: Optional[int] = None # Number of dataset checkpoints to keep, or None to keep them all. # Note: Dataset checkpoints are also affected by `keep`. keep_dataset_checkpoints: Optional[int] = None # Whether to save dataset checkpoints. save_dataset: bool = False # The checkpointer class to use. checkpointer_cls: checkpoints.CheckpointerConstructor = checkpoints.Checkpointer # Transformations to apply, in order, to the state before writing. state_transformation_fns: Sequence[checkpoints.SaveStateTransformationFn] = ( dataclasses.field(default_factory=list)) def __post_init__(self): if self.dtype not in ('float32', 'bfloat16'): raise ValueError( "`SaveCheckpointConfig.dtype` must be one of 'float32' or " f"'bfloat16'. Got {self.dtype}.") @dataclasses.dataclass class RestoreCheckpointConfig: """Configuration for restoring model from checkpoint.""" # Path(s) to checkpoint to restore from or directory (depending on `mode`). path: Union[str, Sequence[str]] # One of 'specific', 'latest', or 'all'. # specific: load the checkpoint specified by `path`. # latest: load most recent checkpoint in the directory specified by `path`. # all: sequentially load all of checkpoints in the directory `path`. mode: str = 'latest' # An optional sequence of (pattern, replacement) regex pairs. The pattern # matches parameters in the model and the replacement matches the checkpoint # (after substitutions). The replacement may be None, in which case the # parameter can be dropped. Use `fallback_to_scratch` to fill them in with # newly initialized values. assignment_map: Optional[Sequence[Tuple[str, Optional[str]]]] = None # Whether to restore all optimizer parameters from the checkpoint. strict: bool = True # Whether to initialize parameters that are in the model being restored but # are missing from the checkpoint (after `assignment_map` is applied). fallback_to_scratch: bool = False # The dtype to restore ('float32' or 'bfloat16'), or None to load as saved. dtype: Optional[str] = None # Whether to restore the dataset checkpoint. Fails if checkpoint not present. restore_dataset: bool = False # The checkpointer class to use. checkpointer_cls: checkpoints.CheckpointerConstructor = checkpoints.Checkpointer # Transformations to apply, in order, to the state after reading. These will # be applied after the `assignment_map` transformations. state_transformation_fns: Sequence[ checkpoints.RestoreStateTransformationFn] = () def __post_init__(self): if self.mode not in ('specific', 'latest', 'all'): raise ValueError( "`RestoreCheckpointConfig.mode` must be one of 'specific', 'latest', " f"or 'all'. Got {self.mode}.") if self.dtype not in (None, 'float32', 'bfloat16'): raise ValueError( "`RestoreCheckpointConfig.dtype` must be one of `None`, 'float32', " f"or 'bfloat16'. Got {self.dtype}.") if self.assignment_map is not None: # Turns `assignment_map` into a transformation function. assignment_map_fn = functools.partial( state_utils.apply_assignment_map, assignment_map=self.assignment_map) # Prepends the `assignment_map` transformation to the front of the list. self.state_transformation_fns = (assignment_map_fn, *self.state_transformation_fns) @dataclasses.dataclass class CheckpointConfig: """Configuration for checkpointing of model and dataset.""" save: Optional[SaveCheckpointConfig] = None restore: Optional[RestoreCheckpointConfig] = None class LegacyCheckpointer(orbax.checkpoint.Checkpointer): """Implementation of Checkpointer interface for T5X. Relies on underlying save_checkpointer and restore_checkpointer, which are t5x.checkpoints.Checkpointer objects. """ def __init__(self, save_checkpointer: checkpoints.Checkpointer, restore_checkpointer: checkpoints.Checkpointer, *, strict: Optional[bool] = False): self._save_checkpointer = save_checkpointer self._restore_checkpointer = restore_checkpointer self._strict = strict async def async_save(self, path: str, item: Any): raise NotImplementedError async def async_restore(self, path: str, item: Optional[Any] = None) -> Any: raise NotImplementedError def save(self, path: str, item: train_state_lib.TrainState, state_transformation_fns: Sequence[ checkpoints.SaveStateTransformationFn] = (), *, concurrent_gb: int = 128): """Performs save operation using save_checkpointer. Args: path: path to save item to. item: a TrainState PyTree to save. state_transformation_fns: Transformations to apply, in order, to the state before writing. concurrent_gb: the approximate number of gigabytes of partitionable parameters to process in parallel. Useful to preserve RAM. """ train_state = item del path # stored in save_checkpointer # dataset_iterator is also saved, but is provided in checkpointer init self._save_checkpointer.save( train_state, state_transformation_fns, concurrent_gb=concurrent_gb) def restore(self, path: str, item: Optional[train_state_lib.TrainState], state_transformation_fns: Sequence[ checkpoints.RestoreStateTransformationFn] = (), fallback_state: Optional[Mapping[str, Any]] = None, lazy_parameters: bool = False) -> train_state_lib.TrainState: """Performs restore operation using restore_checkpointer. Determines whether the indicated path is a Tensorflow checkpoint. Args: path: the string path to restore from. item: a TrainState PyTree to restore. Unused. state_transformation_fns: Transformations to apply, in order, to the state before writing. fallback_state: a state dict of an optimizer to fall back to for loading params that do not exist in the checkpoint (after applying all `state_transformation_fns`), but do exist in `Checkpointer.optimizer`. The union of `fallback_state` and state loaded from the checkpoint must match `Checkpointer.optimizer`. lazy_parameters: whether to load the parameters as LazyArrays to preserve memory. Returns: The restored train state. """ del item # not needed for restore in T5X from_tensorflow = gfile.exists(path + '.index') if from_tensorflow and state_transformation_fns: raise ValueError('Cannot initialize from a TensorFlow checkpoint using ' '`state_transformation_fns`.') if from_tensorflow: logging.info('Initializing parameters from TensorFlow checkpoint %s', path) return self._restore_checkpointer.restore_from_tf_checkpoint( path, strict=self._strict) return self._restore_checkpointer.restore( path=path, state_transformation_fns=state_transformation_fns, fallback_state=fallback_state, lazy_parameters=lazy_parameters) class LegacyCheckpointManager(orbax.checkpoint.CheckpointManager): """Implementation of CheckpointManager interface for T5X. Uses underlying LegacyCheckpointer to handle save/restore for Dataset and TrainState. """ def __init__(self, save_cfg: SaveCheckpointConfig, restore_cfg: RestoreCheckpointConfig, train_state_shape: train_state_lib.TrainState, partitioner: partitioning.BasePartitioner, ds_iter: Optional[tf.data.Iterator] = None, model_dir: Optional[str] = None, use_gda: Optional[bool] = False): if save_cfg.save_dataset: assert ds_iter is not None save_checkpointer = save_cfg.checkpointer_cls( train_state=train_state_shape, partitioner=partitioner, checkpoints_dir=model_dir, dataset_iterator=ds_iter if save_cfg.save_dataset else None, save_dtype=save_cfg.dtype, keep=save_cfg.keep, use_gda=use_gda, keep_dataset_checkpoints=save_cfg.keep_dataset_checkpoints) if restore_cfg: restore_checkpointer = restore_cfg.checkpointer_cls( train_state=train_state_shape, partitioner=partitioner, checkpoints_dir='', # unused for restore dataset_iterator=ds_iter if restore_cfg.restore_dataset else None, restore_dtype=jnp.dtype(restore_cfg.dtype) if restore_cfg.dtype else None) strict = restore_cfg.strict else: restore_checkpointer = None strict = False self._checkpointer = LegacyCheckpointer( save_checkpointer, restore_checkpointer, strict=strict) def save(self, train_state: train_state_lib.TrainState, state_transformation_fns: Sequence[ checkpoints.SaveStateTransformationFn] = ()): """Performs save operation. Args: train_state: a TrainState PyTree to save. state_transformation_fns: Transformations to apply, in order, to the state before writing. """ self._checkpointer.save( path='', # not used item=train_state, state_transformation_fns=state_transformation_fns) def restore( self, paths: Sequence[str], restore_cfg: RestoreCheckpointConfig, fallback_state: Optional[Mapping[str, Any]] = None ) -> Union[train_state_lib.TrainState, Sequence[train_state_lib.TrainState]]: """Performs restore operation using restore_checkpointer. Determines whether the indicated path is a Tensorflow checkpoint. Args: paths: A sequence of paths to restore from. restore_cfg: RestoreCheckpointConfig specifying restoration information. fallback_state: a state dict of an optimizer to fall back to for loading params that do not exist in the checkpoint (after applying all `state_transformation_fns`), but do exist in `Checkpointer.optimizer`. The union of `fallback_state` and state loaded from the checkpoint must match `Checkpointer.optimizer`. Returns: The restored TrainState if only one TrainState can be restored from the given paths, otherwise a sequence of TrainStates. """ if restore_cfg is None or paths is None: return None restored = [] for path in paths: logging.info('Initializing parameters from specific T5X checkpoint %s', path) restored.append( self._checkpointer.restore( path=path, item=None, # not used state_transformation_fns=restore_cfg.state_transformation_fns, fallback_state=fallback_state)) if len(restored) == 1: restored = restored[0] return restored @dataclasses.dataclass class DatasetConfig: """Configuration for loading a dataset from a SeqIO Task or Mixture.""" mixture_or_task_name: str task_feature_lengths: Mapping[str, int] split: str batch_size: int shuffle: bool seed: Optional[int] # Whether to use a precomputed version of the dataset from a cache dir. use_cached: bool = False pack: bool = False # Whether to use tensor2tensor custom ops for more efficient packing. use_custom_packing_ops: bool = False # An optional module to import for registering the referenced Mixture or Task. # DEPRECATED. module: Optional[str] = None # Whether to cache the dataset in memory (only applies to evaluation data). use_memory_cache: bool = True #------------------------------------------------------------------------------ # Fast *nondeterministic* hardware RNG for faster Dropout #------------------------------------------------------------------------------ def _hardware_uniform( rng_key: Array, shape: Shape, dtype: jnp.dtype = np.float32, minval: Array = np.float32(0), maxval: Array = np.float32(1) ) -> Array: """Random uniform method that uses non-deterministic accelerator hardware.""" del rng_key # non-deterministic prng. minval = jax.lax.convert_element_type(minval, dtype) maxval = jax.lax.convert_element_type(maxval, dtype) return jax.lax.rng_uniform(minval, maxval, shape) # For dropout-only hardware rng. def _hardware_bernoulli( rng_key: Array, p: np.ndarray = np.float32(0.5), shape: Shape = ()) -> Array: del rng_key # non-deterministic prng. return jax.lax.rng_uniform(0.0, 1.0, shape) < p def set_hardware_rng_ops(): """Enable JAX Custom PRNG extension.""" jax.config.update('jax_enable_custom_prng', True) # Use only fast TPU hardware PRNG with iterated-hash "split" substitute. # Expected to be deterministic for a fixed partitioning. # Monkey-patch JAX PRNGKey to use unsafe_rbg_prng_impl # TODO(levskaya): replace with jax global config option once we debug it. rbg_prng_key = functools.partial(prng.seed_with_impl, prng.unsafe_rbg_prng_impl) jax.random.PRNGKey = rbg_prng_key jax._src.random.PRNGKey = rbg_prng_key # pylint: disable=protected-access # ----------------------------------------------------------------------------- # Training utility functions. # ----------------------------------------------------------------------------- def get_zeros_batch_like_spec( batch_spec: Mapping[str, jax.ShapeDtypeStruct]) -> Mapping[str, jnp.ndarray]: return {k: jnp.zeros(t.shape, t.dtype) for k, t in batch_spec.items()} def get_zeros_batch_like_dataset(dataset: tf.data.Dataset, batch_size=None) -> Mapping[str, jnp.ndarray]: reshape = lambda s: (batch_size,) + s[1:] if batch_size else tuple(s) batch_spec = { k: jax.ShapeDtypeStruct(reshape(t.shape), t.dtype.as_numpy_dtype) for k, t in dataset.element_spec.items() } return get_zeros_batch_like_spec(batch_spec) class InitFnCallable(typing_extensions.Protocol): """A callable that initializes model variables.""" def __call__( self, rng: Array, input_shapes: Mapping[str, Array], input_types: Optional[Mapping[str, DType]]) -> flax_scope.FrozenVariableDict: ... class LearningRateCallable(typing_extensions.Protocol): def __call__(self, step: jnp.ndarray) -> jnp.ndarray: ... def create_learning_rate_scheduler( factors: str = 'constant * linear_warmup * rsqrt_decay', base_learning_rate: float = 0.5, warmup_steps: int = 1000, decay_factor: float = 0.5, steps_per_decay: int = 20000, steps_per_cycle: int = 100000, step_offset: int = 0, min_learning_rate: float = 1e-8) -> LearningRateCallable: """Creates learning rate schedule. Interprets factors in the factors string which can consist of: * constant: interpreted as the constant value, * linear_warmup: interpreted as linear warmup until warmup_steps, * linear_decay: linear decay from warmup_steps with decay_factor slope. Note this option implies 'constant * linear_warmup', and should not be used in in conjunction with `constant` or `linear_warmup` factors. * rsqrt_decay: divide by square root of max(step, warmup_steps) * rsqrt_normalized_decay: divide by square root of max(step/warmup_steps, 1) * decay_every: Every k steps decay the learning rate by decay_factor. * cosine_decay: Cyclic cosine decay, uses steps_per_cycle parameter. Args: factors: string, factors separated by '*' that defines the schedule. base_learning_rate: float, the starting constant for the lr schedule. warmup_steps: int, how many steps to warm up for in the warmup schedule. decay_factor: float, the amount to decay the learning rate by. steps_per_decay: int, how often to decay the learning rate. steps_per_cycle: int, steps per cycle when using cosine decay. step_offset: int, an offset that the step parameters to this function are relative to. min_learning_rate: float, minimum learning rate to output. Useful for cases when a decay function is (mis)configured to decay to non-positive values. Returns: a function learning_rate(step): float -> {'learning_rate': float}, the step-dependent lr. """ factors = [n.strip() for n in factors.split('*')] def step_fn(step: jnp.ndarray) -> jnp.ndarray: """Step to learning rate function.""" step = jnp.maximum(0, step - step_offset) ret = 1.0 for name in factors: if name == 'constant': ret *= base_learning_rate elif name == 'linear_warmup': ret *= jnp.minimum(1.0, step / warmup_steps) elif name == 'linear_decay': ret *= base_learning_rate * jnp.minimum( step / warmup_steps, 1.0 + decay_factor * (warmup_steps - step)) elif name == 'rsqrt_decay': ret /= jnp.sqrt(jnp.maximum(step, warmup_steps)) elif name == 'rsqrt_normalized_decay': ret *= jnp.sqrt(warmup_steps) ret /= jnp.sqrt(jnp.maximum(step, warmup_steps)) elif name == 'decay_every': ret *= (decay_factor**(step // steps_per_decay)) elif name == 'cosine_decay': progress = jnp.maximum(0.0, (step - warmup_steps) / float(steps_per_cycle)) ret *= jnp.maximum(0.0, 0.5 * (1.0 + jnp.cos(jnp.pi * (progress % 1.0)))) else: raise ValueError('Unknown factor %s.' % name) ret = jnp.maximum(ret, min_learning_rate) return jnp.asarray(ret, dtype=jnp.float32) return step_fn def get_first_valid_restore_config_and_paths( restore_cfgs: Sequence[RestoreCheckpointConfig] ) -> Tuple[Optional[RestoreCheckpointConfig], Sequence[str]]: """Returns first valid restore_cfg and the paths to restore. Args: restore_cfgs: a sequence of RestoreCheckpointConfig objects, which should be filtered to determine the first valid object. Returns: Tuple of valid RestoreCheckpointConfig and a sequence of paths. If the first config encountered has mode 'specfic', it is immediately returned, along with its specified paths. If the mode is 'all' or 'latest', checks to ensure that there are valid checkpoints at each of the provided paths and filters the returned paths accordingly. """ for restore_cfg in restore_cfgs: paths = ([restore_cfg.path] if isinstance(restore_cfg.path, str) else restore_cfg.path) if restore_cfg.mode == 'specific': return restore_cfg, paths elif restore_cfg.mode in ('all', 'latest'): for ckpt_dir in paths: if not gfile.isdir(ckpt_dir): raise ValueError( 'Checkpoint path(s) must be valid directories when using ' "restore mode 'all' or 'latest'.") # Check if this is a TensorFlow checkpoint dir. tf_ckpt_state = tf.train.get_checkpoint_state(ckpt_dir) if tf_ckpt_state: ckpt_paths = tf_ckpt_state.all_model_checkpoint_paths else: ckpt_paths = [ os.path.join(ckpt_dir, f'checkpoint_{step}') for step in checkpoints.all_steps(ckpt_dir) ] if not ckpt_paths: logging.info('No checkpoints found in specified directory: %s', ckpt_dir) continue if restore_cfg.mode == 'latest': logging.info('Using latest T5X checkpoint.') ckpt_paths = ckpt_paths[-1:] return restore_cfg, ckpt_paths else: logging.error('Unsupported checkpoint restore mode: %s', restore_cfg.mode) return None, [] def get_fallback_state(restore_cfg: RestoreCheckpointConfig, init_fn: Callable[[jnp.ndarray], Mapping[str, Any]], init_rng: jnp.ndarray) -> Optional[Mapping[str, Any]]: """Returns the fallback_state that can be used in restore().""" if restore_cfg is None: return if restore_cfg.fallback_to_scratch: if not restore_cfg.state_transformation_fns: raise ValueError('`state_transformation_fns` must be provided with ' '`fallback_to_scratch`') if init_rng is None: raise ValueError('An `init_rng` must be provided with ' '`fallback_to_scratch`') fallback_state = init_fn(init_rng) else: fallback_state = None return fallback_state class TrainStateInitializer: """Helper for initializing partitioned TrainState from checkpoints or scratch. Common use cases: * To restore from a single checkpoint, use `from_checkpoint`. * To iterate over multiple checkpoints without recompiling the model, use `from_checkpoints`. * To initialize from scratch, use `from_scratch`. * To restore from a checkpoint with a fallback to initializing from scratch, use `from_checkpoint_or_scratch`. Attributes: global_train_state_shape: a TrainState containing the global (unpartitioned) shape (in `jax.ShapeDtypeStruct`) of each parameter instead of its value. train_state_axes: a TrainState object containing a PartitionSpec (or None) for each parameter, in place of the parameter itself. """ # TODO(adarob): Replace input_shapes and input_types with sample batch. def __init__(self, optimizer_def: Optional[optimizers.OptimizerDefType], init_fn: InitFnCallable, input_shapes: Mapping[str, Array], partitioner: partitioning.BasePartitioner, input_types: Optional[Mapping[str, DType]] = None): """TrainStateInitializer constructor. Args: optimizer_def: Optimizer def to be initialized, or None to create a `InferenceState` without an optimizer. init_fn: callable that initializes model variables from a PRNGKey and the input shapes. input_shapes: a mapping from key to array shape for each feature in the global (unsharded) input batch. partitioner: the partitioner to use. input_types: a mapping from key to array type for each feature in the global (unshared) input batch. If not provided, the type is assumed to be `jnp.float32`. """ def initialize_train_state(rng: Array): initial_variables = init_fn( rng=rng, input_shapes=input_shapes, input_types=input_types) if optimizer_def: return train_state_lib.FlaxOptimTrainState.create( optimizer_def, initial_variables) return train_state_lib.InferenceState.create(initial_variables) self._partitioner = partitioner self.global_train_state_shape = jax.eval_shape( initialize_train_state, rng=jax.random.PRNGKey(0)) self.train_state_axes = partitioner.get_mesh_axes( self.global_train_state_shape) self._initialize_train_state = initialize_train_state # Currently scanned layers require passing annotations through to the # point of the scan transformation to resolve an XLA SPMD issue. # init_fn is always(?) equal to model.get_initial_variables, fetch the model # instance from the bound method. model = init_fn.__self__ # pytype: disable=attribute-error if (hasattr(model, 'module') and hasattr(model.module, 'scan_layers') and model.module.scan_layers): if hasattr(model.module, 'spmd_annotations'): # update top-level module with spmd annotations. model.module = model.module.clone( parent=None, spmd_annotations=self.train_state_axes.params) def from_scratch(self, init_rng: Array) -> train_state_lib.TrainState: """Initializes the partitioned Optimizer from scratch.""" logging.info('Initializing parameters from scratch.') # If pretraining and no checkpoint imported, we jit the (sharded-) init # function to minimize fragmentation. We use the same partition # setup as the training step/loop to initialize everything "in-place" and # avoid communication or OOM. p_initialize_train_state_fn = self._partitioner.partition( self._initialize_train_state, in_axis_resources=None, out_axis_resources=self.train_state_axes) return p_initialize_train_state_fn(init_rng) # TODO(b/216650048) deprecate this function and use orbax. def from_checkpoints( self, restore_cfgs: Sequence[RestoreCheckpointConfig], ds_iter: Optional[tf.data.Iterator] = None, init_rng: Optional[jnp.ndarray] = None, ) -> Iterable[train_state_lib.TrainState]: """Yields 0 or more restored partitioned Optimizers, and maybe datasets. The manner in which parameters are initialized depends on `restore_cfgs` and `restore_cfgs` is iterated over and the first config that matches one or more existing checkpoints is used to generate restored optimizers from the checkpoint(s). Any remaining configs are ignored. Args: restore_cfgs: ordered sequence of configurations specifying checkpoint(s) to restore from. The first config to match a checkpoint will be used. ds_iter: a tf.data.Iterator for the input data, or None. If provided, the referenced iterator's state may be silently restored (depending on the config's `restore_dataset` value) along with the optimizer. init_rng: for initializing parameters from scratch when they are not available in the checkpoint and `fallback_to_scratch` is True Yields: TrainState with initialized optimizer, with parameters copied to devices. """ def _restore_path(path, cfg): restore_checkpointer = cfg.checkpointer_cls( train_state=self.global_train_state_shape, partitioner=self._partitioner, checkpoints_dir='', # unused for restore dataset_iterator=ds_iter if cfg.restore_dataset else None, restore_dtype=jnp.dtype(cfg.dtype) if cfg.dtype else None) from_tensorflow = gfile.exists(path + '.index') if from_tensorflow and cfg.state_transformation_fns: raise ValueError('Cannot initialize from a TensorFlow checkpoint using ' '`state_transformation_fns`.') if from_tensorflow: logging.info('Initializing parameters from TensorFlow checkpoint %s', path) return restore_checkpointer.restore_from_tf_checkpoint( path, strict=cfg.strict) else: fallback_state = get_fallback_state( cfg, lambda rng: self.from_scratch(rng).state_dict(), init_rng) logging.info('Initializing parameters from specific T5X checkpoint %s', path) return restore_checkpointer.restore( path=path, state_transformation_fns=cfg.state_transformation_fns, fallback_state=fallback_state) restore_cfg, paths = get_first_valid_restore_config_and_paths(restore_cfgs) for path in paths: yield _restore_path(path, restore_cfg) def from_checkpoint( self, ckpt_cfgs: Sequence[RestoreCheckpointConfig], *, ds_iter: Optional[tf.data.Iterator] = None, init_rng: Optional[jnp.ndarray] = None ) -> Optional[train_state_lib.TrainState]: """Restores (at most) 1 checkpoint using `from_checkpoints`, or dies.""" train_states = list( self.from_checkpoints(ckpt_cfgs, ds_iter=ds_iter, init_rng=init_rng)) if len(train_states) > 1: raise ValueError( f'Expected at most 1 checkpoint but got {len(train_states)} for ' f'config(s): {ckpt_cfgs}') return (train_states[0]) if train_states else None def from_checkpoint_or_scratch( self, ckpt_cfgs: Sequence[RestoreCheckpointConfig], *, init_rng: Array, ds_iter: Optional[tf.data.Iterator] = None) -> train_state_lib.TrainState: """Initializes from checkpoint, if found, or from scratch.""" return (self.from_checkpoint(ckpt_cfgs, ds_iter=ds_iter, init_rng=init_rng) or self.from_scratch(init_rng)) # ----------------------------------------------------------------------------- # Logging utility functions # ----------------------------------------------------------------------------- def log_model_info(log_file: Optional[str], full_train_state: train_state_lib.TrainState, partitioner: partitioning.BasePartitioner): """Log the variable shapes information and optionally write it to a file.""" # Only write logs on host 0. if jax.process_index() != 0: return state_dict = full_train_state.state_dict() total_num_params = jax.tree_util.tree_reduce( np.add, jax.tree_map(np.size, state_dict['target'])) logical_axes = partitioner.get_logical_axes(full_train_state).state_dict() mesh_axes = jax.tree_map( lambda x: tuple(x) if x is not None else None, partitioner.get_mesh_axes(full_train_state).state_dict()) def _log_info_and_write_to_file(writer, format_str, *args): logging.info(format_str, *args) if writer is not None: writer.write(format_str % args + '\n') with contextlib.ExitStack() as stack: writer = stack.enter_context(gfile.GFile( log_file, 'w')) if log_file is not None else None # Log params def _log_variable(name: str, arr: Optional[np.ndarray], logical_axes: Optional[partitioning.AxisNames], mesh_axes: Optional[partitioning.PartitionSpec]): # Log nothing on empty dict leaves, which occur with optax EmptyState(). if isinstance(arr, dict) and not arr: return if arr is None: _log_info_and_write_to_file(writer, 'Variable %-80s None', name) return if logical_axes is None or len(logical_axes) != len(arr.shape): shape_str = str(arr.shape) else: shape_str = '({})'.format(', '.join( f'{name}={dimension}' for name, dimension in zip(logical_axes, arr.shape))) _log_info_and_write_to_file( writer, 'Variable %-80s size %-12s shape %-40s partition spec %s', name, arr.size, shape_str, mesh_axes) jax.tree_map( _log_variable, state_utils.get_name_tree(state_dict['target'], keep_empty_nodes=True), state_dict['target'], logical_axes['target'], mesh_axes['target']) _log_info_and_write_to_file(writer, 'Total number of parameters: %d', total_num_params) # Add a blank line between params and states. _log_info_and_write_to_file(writer, '') jax.tree_map( _log_variable, state_utils.get_name_tree(state_dict['state'], keep_empty_nodes=True), state_dict['state'], logical_axes['state'], mesh_axes['state']) # ----------------------------------------------------------------------------- # Utility functions for prediction and evaluation. # ----------------------------------------------------------------------------- class InferStepWithRngCallable(typing_extensions.Protocol): def __call__(self, params: Mapping[str, Any], batch: Mapping[str, jnp.ndarray], rng: jnp.ndarray = None) -> PyTreeDef: """Runs an inference step returning a prediction or score.""" ... class InferStepWithoutRngCallable(typing_extensions.Protocol): def __call__(self, params: Mapping[str, Any], batch: Mapping[str, jnp.ndarray]) -> PyTreeDef: """Runs an inference step returning a prediction or score.""" ... InferStepCallable = Union[InferStepWithRngCallable, InferStepWithoutRngCallable] # NOTE: We're not more prescriptive than PyTreeDef because that's what # InferStepCallable expects. _InferFnResult = Sequence[Tuple[int, PyTreeDef]] _InferFnWithAuxResult = Tuple[_InferFnResult, Mapping[str, Sequence[Any]]] class InferFnCallable(typing_extensions.Protocol): def __call__( self, ds: tf.data.Dataset, train_state: train_state_lib.TrainState, rng: Optional[jnp.ndarray] = None ) -> Union[_InferFnResult, _InferFnWithAuxResult]: """Runs inference on the dataset.""" ... def _remove_padding(all_inferences, all_indices): """Remove padded examples. Args: all_inferences: PyTree[total_examples + padding_count, ...]. all_indices: [total_examples + padding_count]. Returns: all_inferences in shape PyTree[total_examples, ...]. all_indices in shape [total_exmamples]. """ non_pad_idxs = np.where(all_indices >= 0) all_indices = all_indices[non_pad_idxs] all_inferences = jax.tree_map(lambda x: x[non_pad_idxs], all_inferences) return all_inferences, all_indices def get_infer_fn(infer_step: InferStepCallable, batch_size: int, train_state_axes: train_state_lib.TrainState, partitioner: partitioning.BasePartitioner) -> InferFnCallable: """Get prediction function for the SeqIO evaluator. The returned prediction function should take in an enumerated dataset, make predictions and return in an enumerated form with the original indices and examples zipped together. This ensures that the predictions are compared to the targets in a correct order even if the dataset is sharded across multiple hosts and gathered in a nondeterministic way. jax.process_index == 0 is used as a "main host", i.e., it gathers all inference results and returns. Shape notation: Per replica set num replicas: R Per replica set batch size: B Number of replica sets: H Length: L Some transformations have shape transformation annotation, e.g., [B, L] -> [R, B/R, L]. Args: infer_step: a callable that executes one prediction step. Should not yet be partitioned or pmapped. batch_size: the global infer batch size. train_state_axes: Partitioning info for the train state object. partitioner: partitioner to use. Returns: predict_fn: a callable which takes in the enumerated infer dataset and an optimizer and runs the prediction. """ def infer_step_with_indices(params, batch, rng, indices): if 'rng' in inspect.signature(infer_step).parameters: res = typing.cast(InferStepWithRngCallable, infer_step)(params, batch, rng) else: res = typing.cast(InferStepWithoutRngCallable, infer_step)(params, batch) return indices, res partitioned_infer_step = partitioner.partition( infer_step_with_indices, in_axis_resources=(train_state_axes.params, partitioner.data_partition_spec, None, partitioner.data_partition_spec), out_axis_resources=(None, None)) data_layout = partitioner.get_data_layout(batch_size) shard_id = data_layout.shard_id num_shards = data_layout.num_shards per_shard_batch_size = batch_size // num_shards def infer_fn(ds: tf.data.Dataset, train_state: train_state_lib.TrainState, rng: Optional[jnp.ndarray] = None): ds_shapes = jax.tree_map(lambda x: jnp.array(x.shape), ds.element_spec) multihost_utils.assert_equal( ds_shapes, 'Dataset element shapes do not agree across hosts. ' 'This could be an indication that the dataset is nondeterministic.') try: original_ds_length = len(ds) dataset_remainder = original_ds_length % batch_size # pytype:disable=wrong-arg-types logging.info('length of dataset = %s', len(ds)) except TypeError as e: if str(e) == 'dataset length is unknown.': logging.warning( 'The following error is likely due to the use of TensorFlow v1 in ' 'your dataset pipeline. Verify you are not importing from ' '`tf.compat.v1` as part of your pipeline.') raise e if dataset_remainder: dataset_pad_amt = batch_size - dataset_remainder logging.info( 'Padding infer dataset with %d examples for even per-replica shards.', dataset_pad_amt) # Pad with the first example using an index of -1 so seqio will ignore. pad_ds = ds.take(1).map(lambda i, x: (np.int64(-1), x)).repeat( dataset_pad_amt) ds = ds.concatenate(pad_ds) # Shard the infer dataset across replica sets. sharded_ds = ds.shard(num_shards, shard_id).batch( per_shard_batch_size, drop_remainder=True) multihost_utils.assert_equal( jnp.array(len(sharded_ds)), 'Dataset lengths do not agree across hosts.') logging.info( 'The infer dataset is sharded into %d shards with per-shard ' 'batch size of %d', num_shards, per_shard_batch_size) # Run inference for each replica set. batched_results, all_indices = [], [] for index, infer_batch in sharded_ds.as_numpy_iterator(): if rng is None: step_rng = None else: step_rng, rng = jax.random.split(rng) # Run fast inference on batch. # [B, ...] -> [B * shard_count, ...] # partitioned_infer_step executes infer_step on sharded batched data, and # returns de-sharded batched indices and result replicated on all hosts. batch_indices, batch_result = partitioned_infer_step( train_state.params, infer_batch, step_rng, index) logging.info('Inference of batch %s done.', index) # Issue asynchronous copy request which serves as prefetching to the host. def _copy_to_host_async(x): if isinstance(x, GlobalDeviceArray): x.local_data(0).copy_to_host_async() # GDA is fully replicated return x.local_data(0) else: x.copy_to_host_async() return x try: batch_result = jax.tree_map(_copy_to_host_async, batch_result) batch_indices = jax.tree_map(_copy_to_host_async, batch_indices) except AttributeError: # Similar to jax.device_get, we skip transfers for non DeviceArrays. pass batched_results.append(batch_result) all_indices.append(batch_indices) logging.info('Inference of all batches done.') all_inferences = batched_results # List[B * shard_count, ...] -> [B * shard_count * batch_count, ...] all_inferences = jax.tree_multimap(lambda *args: np.concatenate(args), *all_inferences) all_indices = np.concatenate(all_indices) all_inferences, all_indices = _remove_padding(all_inferences, all_indices) # Results are returned from infer_step out of order due to shard operation. # Note: remove padding first, as -1 indices would mess up this operation. # Note: all_inferences may be a PyTree, not just an array, e.g. if # `infer_step` is `model.predict_batch_with_aux`. all_inferences = jax.tree_map(lambda x: x[all_indices], all_inferences) all_indices = all_indices[all_indices] # aux_values is supposed to be a dictionary that maps strings to a set of # auxiliary values. # # We don't want to flatten/unflatten the aux values. We want to preserve the # unflattened values with the type List[Mapping[str, Sequence[Any]]]. We do # this as a memory optimization to avoid lots of redundant keys if we'd # instead had List[Mapping[str, Any]]. # # It has shape Mapping[str, [B * shard_count * batch_count, ...]]. That is, # the first dimension of each of the values in aux_values is equal to # len(all_inferences). aux_values = None if (isinstance(all_inferences, tuple) and len(all_inferences) == 2 and isinstance(all_inferences[1], Mapping)): all_inferences, aux_values = all_inferences # Translate to List[...] by flattening inferences making sure to # preserve structure of individual elements (inferences are not assumed to # be simple np.array). Finally, zip inferences with corresponding indices # and convert leaf np.arrays into lists. all_inferences, struct = jax.tree_flatten(all_inferences) all_inferences = map( functools.partial(jax.tree_unflatten, struct), zip(*all_inferences)) indices_and_outputs = list(zip(all_indices, all_inferences)) indices_and_outputs = jax.tree_map(lambda x: np.array(x).tolist(), indices_and_outputs) if len(indices_and_outputs) != original_ds_length: raise ValueError( 'Size of indices_and_outputs does not match length of original ' 'dataset: %d versus %d' % (len(indices_and_outputs), original_ds_length)) if aux_values is None: return indices_and_outputs else: aux_values = jax.tree_map(lambda x: np.array(x).tolist(), aux_values) return indices_and_outputs, aux_values return infer_fn # ----------------------------------------------------------------------------- # SeqIO utility functions. # ----------------------------------------------------------------------------- def import_module(module: str): """Imports the given module at runtime.""" logging.info('Importing %s.', module) try: importlib.import_module(module) except RuntimeError as e: if (str(e) == 'Attempted to add a new configurable after the config was locked.'): raise RuntimeError( 'Your Task/Mixture module contains gin configurables that must be ' 'loaded before gin flag parsing. One fix is to add ' f"'import {module}' in your gin file.") raise e def get_vocabulary( cfg: DatasetConfig) -> Tuple[seqio.Vocabulary, seqio.Vocabulary]: """Returns `seqio.Vocabulary` objects associated with the `Mixture`/`Task`. Args: cfg: the DatasetConfig specifying which mixture or task to get the vocabularies for. Returns: A tuple of seqio.Vocabulary for inputs and targets. Raises: ValueError: if inputs and targets are not both present and vocabularies are different. """ if cfg.module: warnings.warn( 'The use of `DatasetConfig.module` and `MIXTURE_OR_TASK_MODULE` is ' 'deprecated in favor of importing the module directly or via gin.', DeprecationWarning) import_module(cfg.module) provider = seqio.get_mixture_or_task(cfg.mixture_or_task_name) features = provider.output_features if 'inputs' in features and 'targets' in features: return (features['inputs'].vocabulary, features['targets'].vocabulary) # If a mix of PassThroughVocabularies and other Vocabularies are specified, # use the non-PassThroughVocabularies. # TODO(b/185912004): Remove this once a more general solution is implemented. vocabularies = list( f.vocabulary for f in features.values() if not isinstance(f.vocabulary, seqio.PassThroughVocabulary)) # Otherwise, if all of the vocabs are PassThroughVocabularies, use those. if not vocabularies: vocabularies = list(f.vocabulary for f in features.values()) # If there still aren't any vocabularies, raise an error. if not vocabularies: raise ValueError('"inputs" and "targets" are not both present, and ' 'no vocabularies were set for any features.') first_vocab = vocabularies[0] for vocab in vocabularies[1:]: if vocab != first_vocab: raise ValueError('"inputs" and "targets" are not both present, and ' 'vocabularies are different.') return (first_vocab, first_vocab) def get_dataset(cfg: DatasetConfig, shard_id: int, num_shards: int, feature_converter_cls: Type[seqio.FeatureConverter], num_epochs: Optional[int] = None, continue_from_last_checkpoint: bool = False) -> tf.data.Dataset: """Returns a dataset from SeqIO based on a `DatasetConfig`.""" if continue_from_last_checkpoint: raise ValueError( '`continue_from_last_checkpoint` must be set to False as this is not ' 'supported by this dataset fn.') del continue_from_last_checkpoint if cfg.module: import_module(cfg.module) if cfg.batch_size % num_shards: raise ValueError( f'Batch size ({cfg.batch_size}) must be divisible by number of ' f'shards ({num_shards}).') shard_info = seqio.ShardInfo(index=shard_id, num_shards=num_shards) if cfg.seed is None: # Use a shared timestamp across devices as the seed. seed = multihost_utils.broadcast_one_to_all(np.int32(time.time())) else: seed = cfg.seed return get_dataset_inner(cfg, shard_info, feature_converter_cls, seed, num_epochs) def get_dataset_inner(cfg: DatasetConfig, shard_info: seqio.ShardInfo, feature_converter_cls: Type[seqio.FeatureConverter], seed: Optional[int] = None, num_epochs: Optional[int] = None): """Internal fn to load a dataset from SeqIO based on a `DatasetConfig`.""" batch_size = cfg.batch_size // shard_info.num_shards if seed is not None: multihost_utils.assert_equal( np.array(seed), f'`seed` is not same across hosts; {jax.process_index} has a seed of ' f'{seed}') logging.info( "Initializing dataset for task '%s' with a replica batch size of %d and " 'a seed of %d', cfg.mixture_or_task_name, batch_size, seed) ds = seqio.get_dataset( mixture_or_task_name=cfg.mixture_or_task_name, task_feature_lengths=cfg.task_feature_lengths, dataset_split=cfg.split, shuffle=cfg.shuffle, num_epochs=num_epochs, feature_converter=feature_converter_cls( pack=cfg.pack, use_custom_packing_ops=cfg.use_custom_packing_ops), # pytype: disable=not-instantiable shard_info=shard_info, use_cached=cfg.use_cached, seed=seed) ds = ds.batch(batch_size, drop_remainder=True) return ds class GetDatasetCallable(typing_extensions.Protocol): """Interface for a function returning a dataset (iterator).""" def __call__( self, cfg: DatasetConfig, shard_id: int, num_shards: int, feature_converter_cls: Callable[..., seqio.FeatureConverter], num_epochs: Optional[int] = None, continue_from_last_checkpoint: bool = True ) -> Union[clu.data.DatasetIterator, tf.data.Dataset]: ... def get_training_eval_datasets( cfg: DatasetConfig, shard_id: int, num_shards: int, eval_steps: int, feature_converter_cls: Callable[..., seqio.FeatureConverter], get_dataset_fn: GetDatasetCallable = get_dataset, ) -> Mapping[str, tf.data.Dataset]: """Returns a mapping from eval task name to its dataset.""" mixture_or_task = seqio.get_mixture_or_task(cfg.mixture_or_task_name) datasets = {} if cfg.batch_size % num_shards: raise ValueError( f'Batch size ({cfg.batch_size}) must be divisible by number of ' f'shards ({num_shards}).') def _repeat_shard_batch_take_cache(ds: tf.data.Dataset): # We shard and batch the full, repeated dataset to avoid issues with uneven # file shards. if not isinstance(ds, tf.data.Dataset): raise ValueError('Only tf.data.Dataset objects supported.') return ds.unbatch().repeat().shard(num_shards, shard_id).batch( cfg.batch_size // num_shards, drop_remainder=True).take(eval_steps).cache() for task in seqio.get_subtasks(mixture_or_task): if cfg.split not in task.splits: logging.info("Task %s has no '%s' split; skipping training evaluation.", task.name, cfg.split) continue logging.info('Loading task %s for training evaluation.', task.name) task_cfg = dataclasses.replace( cfg, mixture_or_task_name=task.name, batch_size=1) # We set `num_epochs` to be finite to avoid infinite loops on shards that # have input examples that are all filtered. datasets[task.name] = _repeat_shard_batch_take_cache( get_dataset_fn( task_cfg, shard_id=0, num_shards=1, feature_converter_cls=feature_converter_cls, num_epochs=eval_steps * cfg.batch_size, continue_from_last_checkpoint=False)) if isinstance(mixture_or_task, seqio.Mixture): datasets[mixture_or_task.name] = _repeat_shard_batch_take_cache( get_dataset_fn( dataclasses.replace(cfg, batch_size=1), shard_id=0, num_shards=1, feature_converter_cls=feature_converter_cls, num_epochs=eval_steps * cfg.batch_size, continue_from_last_checkpoint=False)) return datasets def round_vocab_size_to_multiple(vocabulary: seqio.Vocabulary, divisor: int = 128): """Round up vocabulary size for improved TPU performance.""" size = vocabulary.vocab_size return size + -size % divisor def flatten_dict_string_keys(x): """Flattens a nested dictionary to have string keys and '/' separators.""" return traverse_util.flatten_dict(flax.core.unfreeze(x), sep='/') class _RegexMap(collections.abc.Mapping): """Ordered mapping from regexes to values requiring a full match.""" def __init__(self, kvs: Sequence[Tuple[str, Any]]): self._kvs = [(re.compile(k), v) for k, v in kvs] def __getitem__(self, key: str) -> Any: for pattern, v in self._kvs: if pattern.fullmatch(key): return v raise KeyError(f'No pattern matching key: {key}') def __len__(self) -> int: return len(self._kvs) def __iter__(self) -> Iterable[Tuple[re.Pattern, Any]]: return iter(self._kvs) def override_params_axes_names( model_variables: flax_scope.FrozenVariableDict, params_axes_names_override: Sequence[Tuple[str, Tuple[str, ...]]] = () ) -> flax_scope.FrozenVariableDict: """Applies parameter axis names overrides to axes variables. Args: model_variables: the original model variables containing the 'params_axes' collection. params_axes_names_override: a priority-ordered mapping from regex patterns (fully matching parameter names) to tuples containing string logical axis names to replace model-derived names. Returns: an updated set of model variables with the overrides applied to the 'params_axes' collection. """ params_axes_names_override_map = _RegexMap(params_axes_names_override) if 'params_axes' not in model_variables: raise ValueError( "Model variables do not contain a 'params_axes' collection to apply an " 'override to.') model_variables = model_variables.unfreeze() flat_params = traverse_util.flatten_dict(model_variables['params']) flat_params_axes = traverse_util.flatten_dict(model_variables['params_axes']) for key, param in flat_params.items(): param_name = '/'.join(key) override = params_axes_names_override_map.get(param_name) if override is None: continue param_axes_key = key[:-1] + (f'{key[-1]}_axes',) curr_metadata = flat_params_axes.get(param_axes_key) if curr_metadata is None: logging.info('Adding axis names for %s: %s', param_name, override) else: assert isinstance(curr_metadata, flax_partitioning.AxisMetadata) logging.info('Replacing axis names for %s (%s) with %s.', param_name, curr_metadata.names, override) if param.ndim != len(override): raise ValueError( f'Provided axis name override for {param_name} does not match ' f'param rank ({param.ndim}): {override}') flat_params_axes[param_axes_key] = flax_partitioning.AxisMetadata( names=override) model_variables['params_axes'] = traverse_util.unflatten_dict( flat_params_axes) return flax.core.freeze(model_variables) def get_local_data(x): if isinstance(x, GlobalDeviceArray): return x.local_data(0) elif isinstance(x, pxla.ShardedDeviceArray): val = x.device_buffers[0] if val.aval is None: val.aval = jax.ShapedArray(val.shape, val.dtype) return val else: return x