Spaces:
Build error
Build error
| # 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. | |
| """Utilities for reading and writing sharded checkpoints. | |
| The checkpointing utilities here can be used in two ways. The first is to use | |
| the `Checkpointer` class. This requires having an optimizer and various | |
| partitioning utilities setup, but allows for reading and writing of partitioned | |
| parameters. It also allows different hosts to read different parameter | |
| partitions in a multi-host setup, which results in much faster reads. This is | |
| normally used during training where you have already created an optimizer based | |
| on a config. | |
| The second way is to use the `load_t5x_checkpoint` function. This doesn't | |
| require an optimizer to get given up front so it is useful for things like | |
| debugging and analysis of learned weights. However, this means that we cannot do | |
| partitioned reads so loading will be slower than that `Checkpointer` class. | |
| """ | |
| import asyncio | |
| import dataclasses | |
| import functools | |
| import os | |
| import re | |
| import subprocess | |
| import time | |
| from typing import Any, Dict, Iterable, List, Mapping, MutableMapping, Optional, Sequence, Tuple | |
| from absl import logging | |
| from flax import serialization | |
| from flax import traverse_util | |
| import jax | |
| import jax.config | |
| from jax.experimental import global_device_array as gda_lib | |
| from jax.experimental import multihost_utils | |
| from jax.experimental.gda_serialization import serialization as gda_serialization | |
| import jax.numpy as jnp | |
| import numpy as np | |
| import orbax.checkpoint | |
| from t5x import checkpoint_importer | |
| from t5x import checkpoint_utils | |
| 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 tensorstore as ts | |
| import typing_extensions | |
| from tensorboard.backend.event_processing import directory_watcher | |
| from tensorboard.backend.event_processing import event_file_loader | |
| from tensorboard.backend.event_processing import io_wrapper | |
| PartitionSpec = partitioning.PartitionSpec | |
| PyTreeDef = type(jax.tree_structure(None)) | |
| LazyArray = checkpoint_importer.LazyArray | |
| LazyAwaitableArray = checkpoint_importer.LazyAwaitableArray | |
| LazyThreadPoolArray = checkpoint_importer.LazyThreadPoolArray | |
| # Version 3 is used since 2021-06-10, compared to version 2 the only change is | |
| # that `bfloat16` arrays are written in Tensorstore using its native `bfloat16` | |
| # support instead of casting them to `uint16`. | |
| VERSION = 3 | |
| # Desired chunk size is 64MiB. | |
| # This is large enough to keep CNS happy but small enough to support a wide | |
| # range of partitionings. | |
| _DESIRED_CHUNK_SIZE_BYTES = 64 * 1024 * 1024 | |
| # TODO(levskaya, adarob): how should we handle stacked/fused variables?? | |
| _TRAIN_DS_PREFIX = 'train_ds' | |
| def _choose_chunk_shape(write_shape: Sequence[int], | |
| target_elements: int) -> List[int]: | |
| """Chooses a chunk shape that evenly divides write_shape. | |
| The chunk shape is chosen such that the total number of elements is less than | |
| or equal to `target_elements`, but is otherwise as large as possible. | |
| This uses a greedy algorithm that attempts to split the largest dimensions | |
| first. | |
| Args: | |
| write_shape: Write shape for which to choose a chunk shape. | |
| target_elements: Desired number of elements in chosen chunk shape. Must be | |
| >= 1. | |
| Returns: | |
| List of length `len(write_shape)` specifying the chosen chunk shape. | |
| """ | |
| assert target_elements >= 1 | |
| rank = len(write_shape) | |
| # `dim_factors[i]` is the list of divisors of `write_shape[i]` | |
| dim_factors = [ | |
| [i for i in range(1, size + 1) if size % i == 0] for size in write_shape | |
| ] | |
| # The current chunk shape is: | |
| # [dim_factors[i][-1] for i in range(rank)] | |
| def get_total_elements(): | |
| """Returns the number of elements in the current chunk shape.""" | |
| total_elements = 1 | |
| for i in range(rank): | |
| total_elements *= dim_factors[i][-1] | |
| return total_elements | |
| # Reduce the current chunk shape until the desired number of elements is | |
| # reached. | |
| while get_total_elements() > target_elements: | |
| # Greedily reduce the largest dimension. This is not guaranteed to bring us | |
| # the closest to `target_elements`, but is simple to implement and should | |
| # work well enough. | |
| dim_to_reduce = -1 | |
| dim_to_reduce_size = 1 | |
| for i in range(rank): | |
| size = dim_factors[i][-1] | |
| if size > dim_to_reduce_size: | |
| dim_to_reduce_size = size | |
| dim_to_reduce = i | |
| # Can only fail to choose `dim_to_reduce` if all dimensions have size of 1. | |
| # But that cannot happen since `target_elements >= 1`. | |
| assert dim_to_reduce_size > 1 | |
| dim_factors[dim_to_reduce].pop() | |
| return [dim_factors[i][-1] for i in range(rank)] | |
| class _ParameterInfo: | |
| """Information needed to read/write and slice a partitioned parameter.""" | |
| # The unique parameter name. | |
| name: str | |
| # The shape of the parameter. | |
| shape: Tuple[int] | |
| # The TensoreStore Spec containing the minimal information for read/write. | |
| ts_spec: Optional[ts.Spec] | |
| # The LocalChunkInfo for the part of the parameter local to this host. | |
| local_chunk_info: Optional[partitioning.LocalChunkInfo] | |
| # PartitionSpec mesh axes | |
| axes: Optional[partitioning.PartitionSpec] = None | |
| orbax.checkpoint.utils.register_ts_spec_for_serialization() | |
| def _run_future_tree(future_tree): | |
| """Block until all futures are resolved on this host.""" | |
| future_leaves, treedef = jax.tree_flatten(future_tree) | |
| # TODO(adarob): Use asyncio.run in py3.7+. | |
| loop = asyncio.get_event_loop() | |
| leaves = loop.run_until_complete(asyncio.gather(*future_leaves)) | |
| return jax.tree_unflatten(treedef, leaves) | |
| def all_steps(checkpoints_dir: str) -> Sequence[int]: | |
| """Returns list of available step numbers in ascending order.""" | |
| glob_pattern = os.path.join(checkpoints_dir, 'checkpoint_*', 'checkpoint') | |
| checkpoint_paths = gfile.glob(glob_pattern) | |
| re_pattern = re.compile(r'.*/checkpoint_(\d+)/checkpoint$') | |
| matches = [re_pattern.match(ckpt) for ckpt in checkpoint_paths] | |
| return sorted(int(match.group(1)) for match in matches if match) | |
| def all_dataset_checkpoint_steps(checkpoints_dir: str) -> Sequence[int]: | |
| """Returns available dataset checkpoint step numbers in ascending order.""" | |
| glob_pattern = os.path.join(checkpoints_dir, 'checkpoint_*', | |
| f'{_TRAIN_DS_PREFIX}-*') | |
| train_ds_paths = gfile.glob(glob_pattern) | |
| re_pattern = re.compile(r'.*/checkpoint_(\d+)/.*$') | |
| matches = [re_pattern.match(path) for path in train_ds_paths] | |
| return sorted(set(int(match.group(1)) for match in matches if match)) | |
| def latest_step(checkpoints_dir: str) -> Optional[int]: | |
| """Returns latest step number or None if no checkpoints exist.""" | |
| steps = all_steps(checkpoints_dir) | |
| if not steps: | |
| return None | |
| return steps[-1] | |
| def _get_local_data(x): | |
| if isinstance(x, gda_lib.GlobalDeviceArray): | |
| return x.local_data(0) | |
| else: | |
| return x | |
| def get_checkpoint_dir(checkpoints_dir: str, step: int) -> str: | |
| """Returns path to a checkpoint dir given a parent directory and step.""" | |
| return os.path.join(checkpoints_dir, f'checkpoint_{step}') | |
| def _cast(target: PyTreeDef, dtype: jnp.dtype): | |
| """Cast arrays in target to dtype.""" | |
| def maybe_cast(x): | |
| if isinstance(x, (int, str)): | |
| # Ignore common non-array types that shouldn't be cast. | |
| return x | |
| elif x.dtype == dtype: | |
| return x | |
| elif isinstance(x, jax.ShapeDtypeStruct): | |
| return jax.ShapeDtypeStruct(x.shape, dtype) | |
| elif isinstance(x, gda_lib.GlobalDeviceArray): | |
| raise ValueError('GDA cast not supported.') | |
| else: | |
| return x.astype(dtype) | |
| return jax.tree_map(maybe_cast, target) | |
| def _update_ts_path_from_relative_to_absolute( | |
| ckpt_dir: str, ts_spec_dict: MutableMapping[str, Any]): | |
| """Update (in-place) the path and gcs bucket (if applicable) in a TS Spec.""" | |
| # Handle `gs://` paths. | |
| m = re.fullmatch('^gs://([^/]*)/(.*)$', ckpt_dir, re.DOTALL) | |
| if m is not None: | |
| if ts_spec_dict['kvstore']['driver'] != 'gcs': | |
| raise ValueError(f'Incorrect TensorStore Spec. ' | |
| f'Expects kvstore driver to be "gcs" for {ckpt_dir}. ' | |
| f'Got {ts_spec_dict}') | |
| bucket = m.group(1) | |
| ckpt_dir = m.group(2) | |
| ts_spec_dict['kvstore']['bucket'] = bucket | |
| # Update the path with `ckpt_dir` | |
| if 'path' in ts_spec_dict['kvstore']: | |
| # tensorstore>=0.1.14 format | |
| ts_spec_dict['kvstore']['path'] = os.path.join( | |
| ckpt_dir, ts_spec_dict['kvstore']['path']) | |
| elif 'path' in ts_spec_dict: | |
| # tensorstore<0.1.14 format | |
| ts_spec_dict['path'] = os.path.join(ckpt_dir, ts_spec_dict['path']) | |
| else: | |
| raise ValueError( | |
| 'Incorrect TensorStore Spec. Expects "path" to be a key of spec or ' | |
| f'`spec["kvstore"]`. Got {ts_spec_dict}') | |
| def _maybe_update_ts_from_file_to_gcs(ckpt_contents): | |
| """Updates the TensorStore driver from gfile to gcs.""" | |
| def _gfile_to_gcs_driver(arr_or_ts_spec_dict): | |
| """Converts the ts.Spec dict using gfile driver to gcs driver.""" | |
| if not isinstance(arr_or_ts_spec_dict, dict): | |
| return arr_or_ts_spec_dict | |
| if arr_or_ts_spec_dict['kvstore']['driver'] in ('file', 'gfile'): | |
| ts_spec_dict = arr_or_ts_spec_dict | |
| path = ts_spec_dict['kvstore'].pop('path') | |
| # This will be updated to the actual bucket in `_read_ts`. | |
| ts_spec_dict['kvstore'] = { | |
| 'bucket': 't5x-dummy-bucket', | |
| 'driver': 'gcs', | |
| 'path': path | |
| } | |
| else: | |
| if arr_or_ts_spec_dict['kvstore']['driver'] != 'gcs': | |
| raise ValueError('Unsupported TensoreStore driver. Got ' | |
| f'{arr_or_ts_spec_dict["kvstore"]["driver"]}.') | |
| ts_spec_dict = arr_or_ts_spec_dict | |
| return ts_spec_dict | |
| def _is_leaf(value): | |
| return not isinstance( | |
| value, dict) or set(value.keys()) >= {'driver', 'kvstore', 'metadata'} | |
| return jax.tree_map(_gfile_to_gcs_driver, ckpt_contents, is_leaf=_is_leaf) | |
| def _maybe_update_ts_from_gcs_to_file(ckpt_contents): | |
| """Updates the TensorStore driver to gfile or file if different.""" | |
| # if saved in gcs, change to file | |
| def _gcs_to_file_driver(arr_or_ts_spec_dict): | |
| if not isinstance(arr_or_ts_spec_dict, dict): | |
| return arr_or_ts_spec_dict | |
| if arr_or_ts_spec_dict['kvstore']['driver'] == 'gcs': | |
| ts_spec_dict = arr_or_ts_spec_dict | |
| path = ts_spec_dict['kvstore'].pop('path') | |
| driver = 'file' | |
| ts_spec_dict['kvstore'] = {'path': path, 'driver': driver} | |
| elif arr_or_ts_spec_dict['kvstore']['driver'] == 'gfile': | |
| ts_spec_dict = arr_or_ts_spec_dict | |
| driver = 'file' | |
| ts_spec_dict['kvstore']['driver'] = driver | |
| elif arr_or_ts_spec_dict['kvstore']['driver'] == 'file': | |
| ts_spec_dict = arr_or_ts_spec_dict | |
| else: | |
| raise ValueError('Unsupported TensoreStore driver. Got ' | |
| f'{arr_or_ts_spec_dict["kvstore"]["driver"]}.') | |
| return ts_spec_dict | |
| def _is_leaf(value): | |
| return not isinstance( | |
| value, dict) or set(value.keys()) >= {'driver', 'kvstore', 'metadata'} | |
| return jax.tree_map(_gcs_to_file_driver, ckpt_contents, is_leaf=_is_leaf) | |
| class _BytesConditionVariable(object): | |
| """Wraps a condition variable to control concurrency based on bytes.""" | |
| def __init__(self, num_bytes): | |
| self._max_bytes = num_bytes | |
| self._num_bytes = num_bytes | |
| self._cv = asyncio.Condition(lock=asyncio.Lock()) | |
| async def wait_for_bytes(self, n_bytes): | |
| async with self._cv: | |
| await self._cv.wait_for(lambda: self._num_bytes > n_bytes) | |
| self._num_bytes -= n_bytes | |
| assert self._num_bytes >= 0 | |
| async def return_bytes(self, n_bytes): | |
| async with self._cv: | |
| self._num_bytes += n_bytes | |
| assert self._num_bytes <= self._max_bytes | |
| self._cv.notify_all() | |
| class SaveStateTransformationFn(typing_extensions.Protocol): | |
| def __call__(self, state_dict: PyTreeDef, | |
| parameter_infos: PyTreeDef) -> Tuple[PyTreeDef, PyTreeDef]: | |
| """Transforms the state and param info, e.g., by remapping parameters. | |
| Args: | |
| state_dict: State in the current model. | |
| parameter_infos: PyTree containing `_ParameterInfo` objects. | |
| Returns: | |
| A tuple whose first element is the result of transforming `state_dict` and | |
| whose second element is the result of transforming `parameter_infos`. | |
| """ | |
| class RestoreStateTransformationFn(typing_extensions.Protocol): | |
| def __call__(self, | |
| state_dict: PyTreeDef, | |
| target_state_dict: PyTreeDef, | |
| *, | |
| is_resuming: bool = False) -> PyTreeDef: | |
| """Transforms the given checkpoint state, e.g., by remapping parameters. | |
| Args: | |
| state_dict: State to transform, which could be from a previous version of | |
| the model. | |
| target_state_dict: State in the current model. | |
| is_resuming: `True` iff this restore call is due to a job resuming after | |
| being temporarily stopped due to, for example, a preemption. This is | |
| useful when there is restore logic that should run when restoring from | |
| some pre-existing checkpoint, but that should not run again when | |
| resuming from a newly-written checkpoint. | |
| Returns: | |
| The result of transforming the `state_dict`. | |
| """ | |
| class Checkpointer(object): | |
| """Handles saving and restoring potentially-sharded T5X checkpoints. | |
| Checkpoints are stored using a combination of msgpack (via flax.serialization) | |
| and TensorStore. | |
| Parameters (and other objects) that are not partitioned are written to the | |
| msgpack binary directly (by host 0). Partitioned parameters are each written | |
| to their own TensorStore, with each host writing their portion to the same | |
| TensorStore in parallel. If a partition is written on multiple hosts, the | |
| partition is further sharded across these replicas to avoid additional | |
| overhead. In place of the paramater, a `tensorstore.Spec` is written to the | |
| msgpack (by host 0) as a reference to be used during restore. Note that the | |
| path of the array being written is relative. This makes the checkpoints | |
| portable. In other words, even if the checkpoint files are moved to a new | |
| directory, they can still be loaded. Because the path is relative, the | |
| checkpoint directory information has to be dynamically provided. This is done | |
| by `_update_ts_path_from_relative_to_absolute`. | |
| For TensorStore driver using Google Cloud Storage (GCS) Key-Value Storage | |
| Layer, the GCS bucket information is necessary. When a checkpoint is written | |
| using the gcs driver, we don't want to hardcode the bucket information in the | |
| resulting file in order to maintain the portability. Therefore, we use a dummy | |
| bucket name of "t5x-dummy-bucket". When reading or writing the checkpoint, the | |
| bucket information is parsed from the checkpoint directory and the bucket | |
| information is dynamically updated. | |
| Attributes: | |
| checkpoints_dir: a path to a directory to save checkpoints in and restore | |
| them from. | |
| keep: an optional maximum number of checkpoints to keep. If more than this | |
| number of checkpoints exist after a save, the oldest ones will be | |
| automatically deleted to save space. | |
| restore_dtype: optional dtype to cast targets to after restoring. | |
| save_dtype: dtype to cast targets to before saving. | |
| keep_dataset_checkpoints: an optional maximum number of data iterators to | |
| keep. If more than this number of data iterators exist after a save, the | |
| oldest ones will be automatically deleted to save space. | |
| """ | |
| def __init__(self, | |
| train_state: train_state_lib.TrainState, | |
| partitioner: partitioning.BasePartitioner, | |
| checkpoints_dir: str, | |
| dataset_iterator: Optional[tf.data.Iterator] = None, | |
| *, | |
| keep: Optional[int] = None, | |
| save_dtype: jnp.dtype = np.float32, | |
| restore_dtype: Optional[jnp.dtype] = None, | |
| use_gda: Optional[bool] = False, | |
| keep_dataset_checkpoints: Optional[int] = None): | |
| """Checkpointer constructor. | |
| Args: | |
| train_state: A train state to be used to determine the structure of the | |
| parameter tree, and the *full* (non-partitioned) parameter shapes and | |
| dtypes. Saved and restored train states must match this structure. | |
| partitioner: the partitioner to use for determining the local chunks | |
| mapping or to perform params partitioning on restore. | |
| checkpoints_dir: a path to a directory to save checkpoints in and restore | |
| them from. | |
| dataset_iterator: an optional iterator to save/restore. | |
| keep: an optional maximum number of checkpoints to keep. If more than this | |
| number of checkpoints exist after a save, the oldest ones will be | |
| automatically deleted to save space. | |
| save_dtype: dtype to cast targets to before saving. | |
| restore_dtype: optional dtype to cast targets to after restoring. If None, | |
| no parameter casting is performed. | |
| use_gda: if True, enabled gda_lib.GlobalDeviceArray. Note: this is | |
| currently an experimental feature under development. | |
| keep_dataset_checkpoints: an optional maximum number of data iterators to | |
| keep. If more than this number of data iterators exist after a save, the | |
| oldest ones will be automatically deleted to save space. | |
| """ | |
| self._train_state = train_state | |
| self._partitioner = partitioner | |
| self.checkpoints_dir = checkpoints_dir | |
| self.keep = keep | |
| self.keep_dataset_checkpoints = keep_dataset_checkpoints | |
| # Immutable due to use in `_get_parameter_infos` | |
| self._save_dtype = save_dtype | |
| self.restore_dtype = restore_dtype | |
| self._dataset_ckpt = ( | |
| tf.train.Checkpoint(ds=dataset_iterator) if dataset_iterator else None) | |
| self._use_gda = use_gda | |
| if self._use_gda: | |
| logging.info('Checkpointing using GDA format is enabled.') | |
| data_layout = partitioner.get_data_layout() | |
| self._dataset_ckpt_name = ( | |
| f'{_TRAIN_DS_PREFIX}-' | |
| f'{data_layout.shard_id:03}-of-{data_layout.num_shards:03}') | |
| self._should_write_dataset_ckpt = ( | |
| dataset_iterator and data_layout.is_first_host_in_replica_set) | |
| self._parameter_infos = self._get_parameter_infos() | |
| asyncio.set_event_loop(asyncio.new_event_loop()) | |
| def _get_state_dict_for_save(self, | |
| state_dict: Dict[str, Any], | |
| lazy_load: bool = True) -> Mapping[str, Any]: | |
| """Gets the optimizer state dict.""" | |
| def _lazy_load_device_array(arr): | |
| if isinstance(arr, jax.xla.DeviceArray): | |
| return LazyThreadPoolArray(arr.shape, arr.dtype, lambda: np.array(arr)) | |
| return arr | |
| if lazy_load: | |
| state_dict = jax.tree_map(_lazy_load_device_array, state_dict) | |
| return state_dict | |
| def _get_parameter_infos(self): | |
| """Generates the state dict of _ParameterInfos for the Optimizer. | |
| We generate a state dict (matching the shape of the optimizer state dict) | |
| that stores a _ParameterInfo for each parameter array. | |
| The _ParameterInfo contains the TensorStore spec for the parameter array and | |
| the LocalChunkInfo describing the slice of the array local to this host. | |
| Returns: | |
| The state dict of _ParameterInfo objects. | |
| """ | |
| def _get_param_info(name: str, arr: Any, axes: partitioning.PartitionSpec): | |
| # If a node in your model is None it is probably a param_state that is not | |
| # used because of a MultiOptimizer. We don't want to have any parameter | |
| # info for it because it shouldn't be saved or restored. | |
| if arr is None: | |
| return None | |
| # Pass-through empty dict leaves, which occur with optax EmptyState(). | |
| if isinstance(arr, dict) and not arr: | |
| return {} | |
| if axes is None: | |
| return _ParameterInfo( | |
| name=name, | |
| shape=arr.shape, | |
| ts_spec=None, | |
| local_chunk_info=None, | |
| axes=None) | |
| if self._use_gda and isinstance(arr, gda_lib.GlobalDeviceArray): | |
| local_chunk_info = None | |
| metadata = gda_serialization._get_metadata(arr) # pylint: disable=protected-access | |
| del metadata['dtype'] | |
| else: | |
| local_chunk_info = self._partitioner.get_local_chunk_info( | |
| arr.shape, axes) | |
| write_shape = [ | |
| si if sl == slice(None) else sl.stop - sl.start | |
| for si, sl in zip(arr.shape, local_chunk_info.slice) | |
| ] | |
| # TODO(levskaya, adarob): how should we handle stacked/fused variables?? | |
| chunk_shape = _choose_chunk_shape( | |
| write_shape, | |
| target_elements=_DESIRED_CHUNK_SIZE_BYTES / arr.dtype.itemsize) | |
| metadata = { | |
| 'compressor': { | |
| 'id': 'gzip' | |
| }, | |
| 'shape': arr.shape, | |
| 'chunks': np.array(chunk_shape), | |
| } | |
| if self.checkpoints_dir.startswith('gs://'): | |
| spec = { | |
| 'driver': 'zarr', | |
| 'dtype': jnp.dtype(arr.dtype).name, | |
| 'kvstore': { | |
| 'driver': 'gcs', | |
| # We always write with a dummy bucket and dynamically update the | |
| # bucket information. This makes the checkpoint files portable | |
| # and not bind to the bucket that it was originally written to. | |
| 'bucket': 't5x-dummy-bucket', | |
| }, | |
| 'path': name.replace('/', '.'), | |
| 'metadata': metadata, | |
| } | |
| else: | |
| spec = { | |
| 'driver': 'zarr', | |
| 'dtype': jnp.dtype(arr.dtype).name, | |
| 'kvstore': { | |
| 'driver': 'file', | |
| 'path': name.replace('/', '.') | |
| }, | |
| 'metadata': metadata, | |
| } | |
| return _ParameterInfo( | |
| name, | |
| shape=arr.shape, | |
| ts_spec=ts.Spec(spec), | |
| local_chunk_info=local_chunk_info, | |
| axes=axes) | |
| # Create a tree of param names as the keys on the path to each leaf | |
| # separated by "/". | |
| param_names = traverse_util.unflatten_dict({ | |
| k: '/'.join(k) for k in traverse_util.flatten_dict( | |
| self._train_state.state_dict(), keep_empty_nodes=True) | |
| }) | |
| return jax.tree_map( | |
| _get_param_info, param_names, | |
| self._get_state_dict_for_save(self._train_state.state_dict()), | |
| self._partitioner.get_mesh_axes(self._train_state).state_dict()) | |
| def _get_checkpoint_dir(self, step: int) -> str: | |
| return get_checkpoint_dir(self.checkpoints_dir, step) | |
| def all_steps(self) -> Sequence[int]: | |
| """Returns list of available step numbers in ascending order.""" | |
| return all_steps(self.checkpoints_dir) | |
| def all_dataset_checkpoint_steps(self) -> Sequence[int]: | |
| """Returns list of available step numbers in ascending order.""" | |
| return all_dataset_checkpoint_steps(self.checkpoints_dir) | |
| def latest_step(self) -> Optional[int]: | |
| """Returns latest step number or None if no checkpoints exist.""" | |
| return latest_step(self.checkpoints_dir) | |
| def _remove_old_dataset_checkpoints(self): | |
| """Deletes old dataset checkpoints if there are more than allowed.""" | |
| if self.keep_dataset_checkpoints: | |
| existing_steps = self.all_dataset_checkpoint_steps() | |
| to_remove = len(existing_steps) - self.keep_dataset_checkpoints | |
| if to_remove > 0: | |
| for step in existing_steps[:to_remove]: | |
| checkpoint_utils.remove_dataset_checkpoint( | |
| self._get_checkpoint_dir(step), _TRAIN_DS_PREFIX) | |
| def _remove_old_checkpoints(self): | |
| """Deletes oldest checkpoints if there are more than keep_checkpoints.""" | |
| if not self.keep: | |
| return | |
| existing_steps = self.all_steps() | |
| to_remove = len(existing_steps) - self.keep | |
| if to_remove <= 0: | |
| return | |
| for step in existing_steps[:to_remove]: | |
| checkpoint_utils.remove_checkpoint_dir(self._get_checkpoint_dir(step)) | |
| def save(self, | |
| train_state: train_state_lib.TrainState, | |
| state_transformation_fns: Sequence[SaveStateTransformationFn] = (), | |
| *, | |
| concurrent_gb: int = 128): | |
| """Saves a checkpoint for the given train state. | |
| Args: | |
| train_state: the train state to save. May contain a combination of | |
| LazyArray objects and arrays (e.g., np.ndarray, jax.DeviceArray) | |
| 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. | |
| """ | |
| step = train_state.step | |
| step = step.get() if isinstance(step, LazyArray) else step | |
| step = _get_local_data(step) | |
| # Integer, to avoid side effects in the checkpoint path. | |
| step = int(step) | |
| # Share a timestamp across devices. | |
| timestamp = multihost_utils.broadcast_one_to_all(np.int32(time.time())) | |
| final_dir = os.path.join(self.checkpoints_dir, f'checkpoint_{step}') | |
| tmp_dir = final_dir + f'.tmp-{timestamp}' | |
| if gfile.exists(final_dir): | |
| logging.info( | |
| 'Skipping save checkpoint for step %d (directory %s already exists)', | |
| step, final_dir) | |
| return | |
| logging.info('Saving checkpoint for step %d to %s', step, tmp_dir) | |
| if jax.process_index() == 0: | |
| gfile.makedirs(tmp_dir) | |
| # Block all hosts until directory is ready. | |
| multihost_utils.sync_global_devices(f'checkpointer:make_dir:{tmp_dir}') | |
| written_state_dict = self._write_state_to_tensorstore( | |
| tmp_dir, train_state, concurrent_gb, state_transformation_fns) | |
| if self._should_write_dataset_ckpt: | |
| logging.info("Writing dataset iterator state to '%s'.", | |
| self._dataset_ckpt_name) | |
| try: | |
| self._dataset_ckpt.write(os.path.join(tmp_dir, self._dataset_ckpt_name)) | |
| except tf.errors.FailedPreconditionError as e: | |
| logging.error( | |
| 'Input pipeline must be stateless in order to checkpoint. Cache ' | |
| 'stateful steps offline or disable iterator checkpointing.') | |
| raise e | |
| # Block until complete on all hosts. | |
| multihost_utils.sync_global_devices( | |
| f'checkpointer:tensorstore_write_complete:{tmp_dir}') | |
| if jax.process_index() == 0: | |
| written_state_dict = jax.tree_map(_get_local_data, written_state_dict) | |
| # Write msgpack file in host 0 only | |
| msgpack_bytes = serialization.to_bytes({ | |
| 'version': VERSION, | |
| 'optimizer': written_state_dict | |
| }) | |
| with gfile.GFile(os.path.join(tmp_dir, 'checkpoint'), 'wb') as fp: | |
| fp.write(msgpack_bytes) | |
| # Finalize checkpoint directory. | |
| if final_dir.startswith('gs://'): | |
| subprocess.run(['gsutil', '-m', 'mv', tmp_dir, final_dir], | |
| stdout=subprocess.DEVNULL, | |
| check=True) | |
| else: | |
| gfile.rename(tmp_dir, final_dir) | |
| logging.info('Saved checkpoint for step %d to %s', step, final_dir) | |
| # Remove old checkpoints, if necessary. | |
| self._remove_old_checkpoints() | |
| self._remove_old_dataset_checkpoints() | |
| # Block until complete on all hosts. | |
| multihost_utils.sync_global_devices( | |
| f'checkpointer:write_complete:{final_dir}') | |
| def _write_state_to_tensorstore( | |
| self, | |
| ckpt_dir: str, | |
| train_state: train_state_lib.TrainState, | |
| concurrent_gb: int, | |
| state_transformation_fns: Sequence[SaveStateTransformationFn], | |
| ) -> Mapping[str, Any]: | |
| """Writes extracted state from train state to Tensorstore.""" | |
| concurrent_bytes = concurrent_gb * 10**9 | |
| bytes_cv = _BytesConditionVariable(concurrent_bytes) | |
| async def _write_array(maybe_arr: Any, | |
| param_info: Optional[_ParameterInfo], | |
| cast: bool = False): | |
| """Maybe write to TensorStore, returning object to write to msgpack. | |
| Args: | |
| maybe_arr: array or LazyArray to be written | |
| param_info: ParameterInfo object. If None (or if param_info.ts_spec is | |
| None), the array will be immediately returned without writing to | |
| tensorstore. This is because array is None or is not partitioned, and | |
| should be written separately. | |
| cast: if True, performs cast operation using self._save_dtype. | |
| Returns: | |
| Tensorstore spec corresponding to the written array. | |
| """ | |
| if param_info is None or param_info.ts_spec is None: | |
| # Write to the msgpack file on host 0. | |
| if isinstance(maybe_arr, LazyArray): | |
| return await maybe_arr.get_async() | |
| return maybe_arr | |
| # Only write each chunk of a parameter from one host | |
| if self._use_gda or param_info.local_chunk_info.replica_id == 0: | |
| arr = maybe_arr | |
| # Wait until memory is available. | |
| if isinstance(arr, gda_lib.GlobalDeviceArray): | |
| n_bytes = sum([ | |
| shard.data.nbytes | |
| for shard in arr.local_shards | |
| if shard.replica_id == 0 | |
| ]) | |
| else: | |
| n_bytes = arr.nbytes | |
| if n_bytes > concurrent_bytes: | |
| logging.warning( | |
| 'Temporarily increasing the concurrency limits from %d bytes to ' | |
| '%d bytes to fit %s.', concurrent_bytes, n_bytes, param_info.name) | |
| n_bytes = concurrent_bytes | |
| await bytes_cv.wait_for_bytes(n_bytes) | |
| if isinstance(maybe_arr, LazyArray): | |
| arr = await arr.get_async() | |
| elif not isinstance(arr, np.ndarray) and not isinstance( | |
| arr, gda_lib.GlobalDeviceArray): | |
| # Cast jax.DeviceArray to np.ndarray. | |
| arr = np.array(maybe_arr, dtype=maybe_arr.dtype) | |
| tmp_ts_spec_dict = param_info.ts_spec.to_json() | |
| if cast: | |
| # Set desired destination dtype. | |
| tmp_ts_spec_dict['dtype'] = jnp.dtype(self._save_dtype).name | |
| param_info.ts_spec = ts.Spec(tmp_ts_spec_dict) | |
| # Path and gcs bucket (if applicable) information is updated in-place. | |
| _update_ts_path_from_relative_to_absolute(ckpt_dir, tmp_ts_spec_dict) | |
| if cast: | |
| # Set up casting spec. | |
| tmp_ts_spec_dict = { | |
| 'base': tmp_ts_spec_dict, | |
| 'driver': 'cast', | |
| 'dtype': jnp.dtype(arr.dtype).name, # dtype before cast | |
| } | |
| if self._use_gda: | |
| await gda_serialization.async_serialize(arr, tmp_ts_spec_dict) | |
| else: | |
| t = await ts.open( | |
| tmp_ts_spec_dict, | |
| create=True, | |
| open=True, | |
| context=ts.Context({'file_io_concurrency': { | |
| 'limit': 128 | |
| }})) | |
| await t[param_info.local_chunk_info.slice].write(arr) | |
| await bytes_cv.return_bytes(n_bytes) | |
| # N.B. we return the original ts_spec (before | |
| # `_update_ts_path_from_relative_to_absolute` was called). This is because | |
| # we'd like to keep the path as relative, i.e., it doesn't hardcode the | |
| # directory that the checkpoint was originally written. This makes the | |
| # checkpoints portable. | |
| return param_info.ts_spec | |
| transformed_state_dict, transformed_parameter_infos = ( | |
| self._transform_state_and_infos(train_state.state_dict(), | |
| self._parameter_infos, | |
| state_transformation_fns)) | |
| state_dict_for_save = self._get_state_dict_for_save(transformed_state_dict) | |
| def _cast_arr_if_not_partitioned(maybe_arr, param_info): | |
| if param_info is None or param_info.ts_spec is None: | |
| return _cast(maybe_arr, self._save_dtype) | |
| return maybe_arr | |
| state_dict_for_save['target'] = jax.tree_multimap( | |
| _cast_arr_if_not_partitioned, state_dict_for_save['target'], | |
| transformed_parameter_infos['target']) | |
| future_written_state = {} | |
| for k in state_dict_for_save.keys(): | |
| # ensure that only 'target' is cast | |
| future_written_state[k] = jax.tree_multimap( | |
| functools.partial(_write_array, cast=(k == 'target')), | |
| state_dict_for_save[k], transformed_parameter_infos[k]) | |
| # Block until complete on this host. | |
| written_state_dict = _run_future_tree(future_written_state) | |
| # Block until complete on all hosts. | |
| multihost_utils.sync_global_devices( | |
| f'checkpointer:ts_write_complete:{ckpt_dir}') | |
| return written_state_dict | |
| def _transform_state_and_infos( | |
| self, | |
| state_dict: PyTreeDef, | |
| parameter_infos: PyTreeDef, | |
| state_transformation_fns: Sequence[SaveStateTransformationFn], | |
| ) -> Tuple[PyTreeDef, PyTreeDef]: | |
| """Applies transformations to the state dict and parameter infos PyTrees.""" | |
| for fn in state_transformation_fns: | |
| state_dict, parameter_infos = fn(state_dict, parameter_infos) | |
| return state_dict, parameter_infos | |
| def restore( | |
| self, | |
| step: Optional[int] = None, | |
| path: Optional[str] = None, | |
| state_transformation_fns: Sequence[RestoreStateTransformationFn] = (), | |
| fallback_state: Optional[Mapping[str, Any]] = None, | |
| lazy_parameters: bool = False) -> train_state_lib.TrainState: | |
| """Restores the host-specific parameters in an Optimizer. | |
| Either `step` or `path` can be specified, but not both. If neither are | |
| specified, restores from the latest checkpoint in the checkpoints directory. | |
| Args: | |
| step: the optional step number to restore from. | |
| path: an optional absolute path to a checkpoint file to restore from. | |
| state_transformation_fns: Transformations to apply, in order, to the state | |
| after reading. | |
| 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. | |
| Raises: | |
| ValueError if both `step` and `path` are specified. | |
| ValueError if checkpoint at `path` or `step` does not exist. | |
| ValueError if `step` and `path` are not specified and no checkpoint is | |
| found in the checkpoints directory. | |
| """ | |
| if lazy_parameters and self._partitioner.params_on_devices: | |
| raise ValueError('Lazy Parameters cannot be copied to devices, please ' | |
| 'set partitioner.params_on_devices=False.') | |
| if step is not None and path is not None: | |
| raise ValueError('At most one of `step` or `path` may be provided.') | |
| if path: | |
| ckpt_path = path | |
| else: | |
| if step is None: | |
| step = self.latest_step() | |
| if not step: | |
| raise ValueError(f'No checkpoints found in {self.checkpoints_dir}.') | |
| ckpt_path = self._get_checkpoint_dir(step) | |
| if gfile.isdir(ckpt_path): | |
| ckpt_dir = ckpt_path | |
| ckpt_path = os.path.join(ckpt_path, 'checkpoint') | |
| else: | |
| ckpt_dir = os.path.dirname(ckpt_path) | |
| if not gfile.exists(ckpt_path) or gfile.isdir(ckpt_path): | |
| raise ValueError(f'Path is not a valid T5X checkpoint: {ckpt_path}') | |
| logging.info('Restoring from checkpoint: %s', ckpt_path) | |
| with gfile.GFile(ckpt_path, 'rb') as fp: | |
| # TODO(adarob): Use threaded reading as in flax.checkpoints. | |
| raw_contents = fp.read() | |
| if raw_contents.startswith(b'model_checkpoint_path'): | |
| raise ValueError( | |
| 'Attempting to restore a TensorFlow checkpoint as a native T5X ' | |
| 'checkpoint. Use `restore_from_tf_checkpoint` instead. Path: ' + | |
| ckpt_path) | |
| # `ckpt_contents['optimizer']` is a pytree with a realized np.array for | |
| # leaves (params or states) written as msgpack and a ts.Spec (in a dict) | |
| # for leaves written by TensorStore. | |
| ckpt_contents = serialization.msgpack_restore(raw_contents) | |
| # If reading a ckpt that was written with gfile driver but the current | |
| # session uses the gcs driver, convert the ckpt's driver to gcs. | |
| if ckpt_dir.startswith('gs://'): | |
| ckpt_contents = _maybe_update_ts_from_file_to_gcs(ckpt_contents) | |
| # If a ckpt was saved in gcs and is being loaded locally, then convert the | |
| # driver to file or gfile. If the ckpt was not saved in gcs, do not change. | |
| else: | |
| ckpt_contents = _maybe_update_ts_from_gcs_to_file(ckpt_contents) | |
| ckpt_state_dict = self._get_optimizer_state_dict(ckpt_contents, | |
| state_transformation_fns) | |
| # The state dict may contain TensorStore specs that need to be read. | |
| dummy_spec = ts.Spec({'driver': 'zarr', 'kvstore': {'driver': 'memory'}}) | |
| # `dummy_written_state_dict` is a pytree with a `dummy_spec` for leaves | |
| # (params or states) written as msgpack and a ts.Spec (in a dict) for leaves | |
| # written by TensorStore. | |
| dummy_written_state_dict = jax.tree_map( | |
| lambda x: x.ts_spec or dummy_spec, | |
| self._parameter_infos, | |
| ) | |
| if fallback_state is None: | |
| restore_parameter_infos = self._parameter_infos | |
| else: | |
| # If `fallback_state` was specified, restore only the subset | |
| # of parameters matched by `self._get_optimizer_state_dict`. The | |
| # rest will be provided by `fallback_state`. | |
| dummy_written_state_dict = state_utils.intersect_state( | |
| dummy_written_state_dict, ckpt_state_dict) | |
| restore_parameter_infos = state_utils.intersect_state( | |
| self._parameter_infos, ckpt_state_dict) | |
| restore_parameter_infos_flat = state_utils.flatten_state_dict( | |
| restore_parameter_infos) | |
| for key in restore_parameter_infos_flat.keys(): | |
| logging.info('Restoring key from ckpt: %s', key) | |
| # NB: `serialization.from_state_dict` doesn't check whether the shapes match | |
| # at the leaf level. Non-partitioned leaves (e.g., optimizer states) can | |
| # load arrays with inconsistent shapes. | |
| # `written_state_dict` is a pytree with a realized np.array for leaves | |
| # (params or states) written as msgpack and a `ts.Spec` for leaves written | |
| # by TensorStore. | |
| written_state_dict = serialization.from_state_dict(dummy_written_state_dict, | |
| ckpt_state_dict) | |
| state_dict = self._read_state_from_tensorstore( | |
| ckpt_path, | |
| written_state_dict, | |
| restore_parameter_infos=restore_parameter_infos, | |
| lazy_parameters=lazy_parameters) | |
| # If `fallback_state` was specified, then fill the missing parameters. | |
| if fallback_state is not None: | |
| state_dict = state_utils.merge_state(state_dict, fallback_state) | |
| for key in state_utils.flatten_state_dict(state_dict).keys(): | |
| if key not in restore_parameter_infos_flat: | |
| logging.info('Not restoring key from ckpt: %s', key) | |
| if self._dataset_ckpt: | |
| logging.info("Restoring dataset iterator from '%s'.", | |
| self._dataset_ckpt_name) | |
| self._dataset_ckpt.read(os.path.join( | |
| ckpt_dir, self._dataset_ckpt_name)).assert_consumed() | |
| return self._restore_train_state(state_dict) | |
| def _restore_train_state( | |
| self, | |
| state_dict: optimizers.OptimizerStateType) -> train_state_lib.TrainState: | |
| """Restores a TrainState from an Optimizer state_dict.""" | |
| train_state = self._train_state.restore_state(state_dict) | |
| if not self._use_gda and self._partitioner.params_on_devices: | |
| logging.info('Moving params to devices.') | |
| train_state_axes = self._partitioner.get_mesh_axes(train_state) | |
| train_state = self._partitioner.move_params_to_devices( | |
| train_state, train_state_axes) | |
| return train_state | |
| def _create_lazy_awaitable_array( | |
| self, param_info: _ParameterInfo, maybe_ts_spec: Any, ckpt_path: str, | |
| restore_dtype: Optional[jnp.dtype]) -> LazyAwaitableArray: | |
| """Creates LazyArray from tensorstore. | |
| Does not materialize the array immediately. | |
| Args: | |
| param_info: Information about how to read the parameter, host based sliced | |
| reads and the like. | |
| maybe_ts_spec: The tensorstore spec to read the parameter or some other | |
| object. If this is an array then we will do a host based sliced read on | |
| it (provided the param_info says to). Anything else we just return. | |
| ckpt_path: A base location to use when resolving the relative paths in the | |
| tensorstore spec. | |
| restore_dtype: type to restore as. None indicates that no cast is | |
| requested. | |
| Returns: | |
| LazyArray object. | |
| """ | |
| mesh = None | |
| axes = None | |
| if self._use_gda: | |
| mesh = self._partitioner.mesh | |
| axes = param_info.axes | |
| get_fn = functools.partial( | |
| _read_ts, | |
| param_info, | |
| maybe_ts_spec, | |
| ckpt_path=ckpt_path, | |
| restore_dtype=restore_dtype, | |
| mesh=mesh, | |
| axes=axes) | |
| return LazyAwaitableArray.from_tensor_store_spec_or_array( | |
| maybe_ts_spec, get_fn, dtype=restore_dtype) | |
| def _read_state_from_tensorstore( | |
| self, | |
| ckpt_path: str, | |
| written_state_dict: Mapping[str, Any], | |
| restore_parameter_infos: Optional[Mapping[str, Any]] = None, | |
| lazy_parameters: bool = False, | |
| ) -> Mapping[str, Any]: | |
| """Sets up lazy reads from Tensorstore and returns them as a state_dict.""" | |
| if restore_parameter_infos is None: | |
| restore_parameter_infos = self._parameter_infos | |
| # Replace TensorStore Specs with the lazy array values. | |
| state_dict = {} | |
| for k in written_state_dict.keys(): | |
| # ensure that only 'target' is cast | |
| restore_dtype = self.restore_dtype if k == 'target' else None | |
| state_dict[k] = jax.tree_multimap( | |
| functools.partial( | |
| self._create_lazy_awaitable_array, | |
| ckpt_path=ckpt_path, | |
| restore_dtype=restore_dtype), restore_parameter_infos[k], | |
| written_state_dict[k]) | |
| if not lazy_parameters: | |
| future_state_dict = jax.tree_map(lambda x: x.get_async(), state_dict) | |
| state_dict = _run_future_tree(future_state_dict) | |
| if self.restore_dtype is not None: | |
| state_dict['target'] = _cast(state_dict['target'], self.restore_dtype) | |
| return state_dict | |
| def restore_from_tf_checkpoint( | |
| self, | |
| path_or_dir: str, | |
| strict: bool = True, | |
| translator: Optional[checkpoint_importer.CheckpointTranslator] = None | |
| ) -> train_state_lib.TrainState: | |
| """Restore from a TensorFlow-based T5 checkpoint.""" | |
| full_state_dict = checkpoint_importer.restore_from_t5_checkpoint( | |
| self._train_state.state_dict(), | |
| path_or_dir, | |
| lazy_parameters=False, | |
| strict=strict, | |
| translator=translator) | |
| def _partition_parameter(maybe_arr: Any, param_info: _ParameterInfo): | |
| if isinstance(maybe_arr, np.ndarray) and param_info: | |
| arr = maybe_arr | |
| if param_info.shape is not None and arr.shape != param_info.shape: | |
| raise ValueError( | |
| f'Shape of `{param_info.name}` in checkpoint {arr.shape} does ' | |
| f'not match expected {param_info.shape}.') | |
| if param_info.local_chunk_info: | |
| arr = arr[param_info.local_chunk_info.slice] | |
| return arr | |
| return maybe_arr | |
| state_dict = jax.tree_multimap(_partition_parameter, full_state_dict, | |
| self._parameter_infos) | |
| if self.restore_dtype is not None: | |
| state_dict['target'] = _cast(state_dict['target'], self.restore_dtype) | |
| return self._restore_train_state(state_dict) | |
| def convert_from_tf_checkpoint( | |
| self, | |
| path_or_dir: str, | |
| *, | |
| state_transformation_fns: Sequence[SaveStateTransformationFn] = (), | |
| concurrent_gb: int = 16, | |
| translator: Optional[checkpoint_importer.CheckpointTranslator] = None): | |
| """Convert from a TensorFlow-based T5 checkpoint.""" | |
| full_state_dict = checkpoint_importer.restore_from_t5_checkpoint( | |
| self._train_state.state_dict(), | |
| path_or_dir, | |
| lazy_parameters=True, | |
| translator=translator) | |
| train_state = self._train_state.restore_state(full_state_dict) | |
| self.save( | |
| train_state, | |
| state_transformation_fns=state_transformation_fns, | |
| concurrent_gb=concurrent_gb) | |
| def _get_optimizer_state_dict( | |
| self, ckpt_contents: PyTreeDef, | |
| state_transformation_fns: Sequence[RestoreStateTransformationFn]): | |
| return _get_optimizer_state_dict(ckpt_contents, | |
| self._train_state.state_dict(), | |
| state_transformation_fns) | |
| class CheckpointerConstructor(typing_extensions.Protocol): | |
| """A function that returns a checkpoints.Checkpointer. | |
| This type annotation allows users to partially bind args to the constructors | |
| of Checkpointer subclasses without triggering type errors. | |
| """ | |
| def __call__(self, | |
| train_state: train_state_lib.TrainState, | |
| partitioner: partitioning.BasePartitioner, | |
| checkpoints_dir: str, | |
| dataset_iterator: Optional[tf.data.Iterator] = None, | |
| *, | |
| keep: Optional[int] = None, | |
| save_dtype: jnp.dtype = np.float32, | |
| restore_dtype: Optional[jnp.dtype] = None, | |
| use_gda: Optional[bool] = False, | |
| keep_dataset_checkpoints: Optional[int] = None) -> Checkpointer: | |
| """Checkpointer constructor. | |
| Args: | |
| train_state: A train state to be used to determine the structure of the | |
| parameter tree, and the *full* (non-partitioned) parameter shapes and | |
| dtypes. Saved and restored train states must match this structure. | |
| partitioner: the partitioner to use for determining the local chunks | |
| mapping or to perform params partitioning on restore. | |
| checkpoints_dir: a path to a directory to save checkpoints in and restore | |
| them from. | |
| dataset_iterator: an optional iterator to save/restore. | |
| keep: an optional maximum number of checkpoints to keep. If more than this | |
| number of checkpoints exist after a save, the oldest ones will be | |
| automatically deleted to save space. | |
| save_dtype: dtype to cast targets to before saving. | |
| restore_dtype: optional dtype to cast targets to after restoring. If None, | |
| no parameter casting is performed. | |
| use_gda: if True, enabled gda_lib.GlobalDeviceArray. Note: this is | |
| currently an experimental feature under development. | |
| keep_dataset_checkpoints: an optional maximum number of data iterators to | |
| keep. If more than this number of data iterators exist after a save, the | |
| oldest ones will be automatically deleted to save space. | |
| """ | |
| pass | |
| class SaveBestCheckpointer(Checkpointer): | |
| """A Checkpointer class that keeps checkpoints based on 'best' metrics. | |
| This extends the standard Checkpointer to garbage collect checkpoints based on | |
| metric values, instead of step recency. It uses Tensorboard summary files to | |
| determine best values for a given user configured metric name. Events are read | |
| and parsed using Tensorboard's event_processing packages. | |
| The metric name must be of the form `{run_name}/{tag_name}`. For example, | |
| 'train/accuracy' or 'inference_eval/glue_cola_v002/eval/accuracy'. | |
| A few important features of this checkpointer: | |
| - Fallback behavior. It is not possible to verify whether metric names are | |
| valid during initialization, since some metrics may get written out after | |
| some time (e.g., during an evaluation). As such, when user provided metric | |
| names are not found, this checkpointer can be configured for two fall back | |
| strategies: (1) if `keep_checkpoints_without_metrics` is False, we use to | |
| the "most recent checkpoint" strategy from the standard checkpointer, (2) | |
| if `keep_checkpoints_without_metrics` is True, we keep all checkpoints until | |
| metrics become available (potentially indefinitely if summary files have | |
| been deleted or corrupted). | |
| - The number of checkpoints to keep is always increased by 1. Since its | |
| crucial to always keep the latest checkpoint (for recovery purposes) we | |
| always store the latest checkpoint plus `keep` number of best checkpoints. | |
| - It is assumed that Tensorboard summaries (event) files share a common root | |
| directory with `checkpoint_dir`, which is the directory passed to the | |
| the logdir crawler that searches for event files. | |
| Attributes: | |
| checkpoints_dir: a path to a directory to save checkpoints in and restore | |
| them from. | |
| keep: an optional maximum number of checkpoints to keep. If more than this | |
| number of checkpoints exist after a save, the oldest ones will be | |
| automatically deleted to save space. | |
| restore_dtype: optional dtype to cast targets to after restoring. | |
| save_dtype: dtype to cast targets to before saving. | |
| metric_name_to_monitor: Name of metric to monitor. Must be in the format | |
| {run_name}/{tag_name} (e.g., 'train/accuracy', | |
| 'inference_eval/glue_cola_v002/eval/accuracy'). | |
| metric_mode: Mode to use to compare metric values. One of 'max' or 'min'. | |
| keep_checkpoints_without_metrics: Whether to always keep (or delete) | |
| checkpoints for which a metric value has not been found. | |
| force_keep_period: When removing checkpoints, skip those who step is | |
| divisible by force_keep_period (step % force_keep_period == 0). | |
| use_gda: Enables GDA (see Checkpointer). | |
| keep_dataset_checkpoints: an optional maximum number of data iterators to | |
| keep. If more than this number of data iterators exist after a save, the | |
| oldest ones will be automatically deleted to save space. | |
| """ | |
| def __init__(self, | |
| train_state: train_state_lib.TrainState, | |
| partitioner: partitioning.BasePartitioner, | |
| checkpoints_dir: str, | |
| dataset_iterator: Optional[tf.data.Iterator] = None, | |
| *, | |
| keep: Optional[int] = None, | |
| save_dtype: jnp.dtype = np.float32, | |
| restore_dtype: Optional[jnp.dtype] = None, | |
| metric_name_to_monitor: str = 'train/accuracy', | |
| metric_mode: str = 'max', | |
| keep_checkpoints_without_metrics: bool = True, | |
| force_keep_period: Optional[int] = None, | |
| use_gda: bool = False, | |
| keep_dataset_checkpoints: Optional[int] = None): | |
| super().__init__( | |
| train_state, | |
| partitioner, | |
| checkpoints_dir, | |
| dataset_iterator, | |
| keep=keep, | |
| save_dtype=save_dtype, | |
| restore_dtype=restore_dtype, | |
| use_gda=use_gda, | |
| keep_dataset_checkpoints=keep_dataset_checkpoints) | |
| if metric_mode not in ('max', 'min'): | |
| raise ValueError('Unsupported `metric_mode`: %s' % metric_mode) | |
| # Metric run and tag names are derived from metric_name_to_monitor and are | |
| # filled in _try_fill_metric_run_and_tag_names(). | |
| self._metric_run: Optional[str] = None | |
| self._metric_tag: Optional[str] = None | |
| self._metric_name_to_monitor = metric_name_to_monitor | |
| self._metric_mode = metric_mode | |
| self._keep_checkpoints_without_metrics = keep_checkpoints_without_metrics | |
| self._force_keep_period = force_keep_period | |
| logging.info('Using SaveBestCheckpointer to keep %s best (%s) metric %s', | |
| keep, metric_mode, metric_name_to_monitor) | |
| def _populate_metrics_for_steps(self, | |
| steps: Iterable[int]) -> Mapping[int, float]: | |
| """Iterate through summary event files and return metrics for `steps`.""" | |
| metrics_by_step = {} | |
| for subdir in io_wrapper.GetLogdirSubdirectories(self.checkpoints_dir): | |
| rpath = os.path.relpath(subdir, self.checkpoints_dir) | |
| # Skip runs that do not match user-specified metric. | |
| if ((not self._metric_run and not self._try_fill_metric_run_and_tag_names( | |
| (rpath,))) or self._metric_run != rpath): | |
| logging.info('Skipping events in %s', subdir) | |
| continue | |
| logging.info('Looking for events in %s', subdir) | |
| loader = directory_watcher.DirectoryWatcher( | |
| subdir, event_file_loader.EventFileLoader, | |
| io_wrapper.IsTensorFlowEventsFile) | |
| for event in loader.Load(): | |
| # Skip metric collection of events for unavailable checkpoints or for | |
| # unmonitored tags. | |
| if (event.step not in steps or not event.summary.value or | |
| event.summary.value[0].tag != self._metric_tag): | |
| continue | |
| metric_value = tf.make_ndarray(event.summary.value[0].tensor) | |
| metrics_by_step[event.step] = metric_value | |
| return metrics_by_step | |
| def _try_fill_metric_run_and_tag_names(self, run_keys: Iterable[str]) -> bool: | |
| """Extract metric run and tag names by matching one of the `run_keys`. | |
| This function tries to greedily split user-provided metric_name_to_monitor | |
| into {run} and {tag} components. It does so by trying to match all available | |
| {run}/{tag} names in the provided run_keys. If successful, populates | |
| self._metric_run and self._metric_tag. | |
| Args: | |
| run_keys: Set of run keys to test for. | |
| Returns: | |
| Whether metric name prefix matches one of the run keys, and, as a | |
| side-effect, populates self._metric_run and self._metric_tag. | |
| """ | |
| metric_run, metric_tag = None, None | |
| # Query existing events for different run and tags to match with user | |
| # provided metric name. | |
| m = self._metric_name_to_monitor.split('/') | |
| possible_run_names = ['/'.join(m[:i]) for i in range(1, len(m))] | |
| for key in run_keys: | |
| for possible_run_name in possible_run_names: | |
| if key == possible_run_name: | |
| metric_run = possible_run_name | |
| metric_tag = self._metric_name_to_monitor[len(metric_run) + 1:] | |
| break | |
| if metric_run and metric_tag: | |
| self._metric_run, self._metric_tag = metric_run, metric_tag | |
| return True | |
| return False | |
| def _filter_out_force_keep_period_steps(self, existing_steps): | |
| """Filter out steps that are divisible by keep_period excluding the last.""" | |
| if not existing_steps: | |
| return existing_steps | |
| # Don't filter out the last step. | |
| last_step = existing_steps.pop() | |
| existing_steps = [ | |
| s for s in existing_steps if s % self._force_keep_period != 0 | |
| ] | |
| return existing_steps + [last_step] | |
| def _remove_old_checkpoints(self): | |
| """Deletes checkpoints if there are more than keep_checkpoints.""" | |
| if not self.keep: | |
| return | |
| existing_steps = self.all_steps() | |
| if self._force_keep_period: | |
| # Ignore checkpoints whose step is divisible by the keep period. | |
| existing_steps = self._filter_out_force_keep_period_steps(existing_steps) | |
| # Artificially add 1 to `keep` since we always keep the latest checkpoint. | |
| if len(existing_steps) <= self.keep + 1: | |
| return | |
| # Synchronous fetch of new events for existing_steps. | |
| metrics_by_step = self._populate_metrics_for_steps(existing_steps) | |
| logging.info('SaveBestcheckpointer: collected metrics %s', metrics_by_step) | |
| # Re-sort existing_steps by metric values while always keeping the latest | |
| # checkpoint. | |
| latest_checkpoint = existing_steps[-1] | |
| existing_steps = existing_steps[:-1] | |
| if self._keep_checkpoints_without_metrics: | |
| existing_steps = list( | |
| filter(lambda s: s in metrics_by_step, existing_steps)) | |
| to_remove = len(existing_steps) - self.keep | |
| if to_remove <= 0: | |
| return | |
| # For any remaining steps without metrics, we assign a low/high value which | |
| # will make them candidate for removal. If no metrics are found this sorting | |
| # should preserve current order (oldest first). | |
| not_found_value = float('-inf' if self._metric_mode == 'max' else 'inf') | |
| existing_steps = sorted( | |
| existing_steps, | |
| key=lambda step: metrics_by_step.get(step, not_found_value), | |
| reverse=(self._metric_mode != 'max')) | |
| existing_steps.append(latest_checkpoint) | |
| for step in existing_steps[:to_remove]: | |
| checkpoint_utils.remove_checkpoint_dir(self._get_checkpoint_dir(step)) | |
| def _get_optimizer_state_dict( | |
| ckpt_contents: PyTreeDef, optimizer_state: Mapping[str, Any], | |
| state_transformation_fns: Sequence[RestoreStateTransformationFn]): | |
| """Extracts optimizer state dict contents and applies assignment map.""" | |
| version = ckpt_contents.get('version', 0) | |
| if version == 0: | |
| # This is a standard Flax checkpoint and may require remapping below. | |
| ckpt_optimizer_state = ckpt_contents | |
| else: | |
| ckpt_optimizer_state = ckpt_contents['optimizer'] | |
| if version >= 2: | |
| for fn in state_transformation_fns: | |
| ckpt_optimizer_state = fn(ckpt_optimizer_state, optimizer_state) | |
| return ckpt_optimizer_state | |
| else: | |
| raise ValueError('Checkpoint versions earlier than 2 are not supported. ' # pylint: disable=unreachable | |
| f'Got version: {version}') | |
| async def _read_ts(param_info: _ParameterInfo, | |
| maybe_tspec: Any, | |
| ckpt_path: str, | |
| restore_dtype: Optional[jnp.dtype] = None, | |
| mesh: Optional[gda_lib.Shape] = None, | |
| axes: Optional[gda_lib.MeshAxes] = None): | |
| """Read from a tensorstore. | |
| If both `mesh` and `axes` are provided, the method will attempt to restore the | |
| array as a GlobalDeviceArray. | |
| Note: | |
| We use param_infos as the first argument because this function is only used | |
| in `jax.tree_multimap` calls. In a tree multimap if the leaf of the first | |
| tree is `None` then is is ignored, even if the second tree has a subtree | |
| at that point. This means that when we are using something like a | |
| MultiOptimizer we can set the parameter info for a variable to `None` and | |
| we can skip processing it, even if the checkpoint has a subtree with things | |
| like optimizer state variables in it. | |
| Args: | |
| param_info: Information about how to read the parameter, host based sliced | |
| reads and the like. | |
| maybe_tspec: The tensorstore spec to read the parameter or some other | |
| object. If this is an array then we will do a host based sliced read on it | |
| (provided the param_info says to). Anything else we just return. | |
| ckpt_path: A base location to use when resolving the relative paths in the | |
| tensorstore spec. | |
| restore_dtype: type to restore as. None indicates that no cast is requested. | |
| mesh: Mesh object for GDA restoration. | |
| axes: MeshAxes object for GDA restoration. | |
| Returns: | |
| The array. Depending on the value `maybe_tspec` it might be read from | |
| tensorstore, or it might be returned as is. Depending on the values in | |
| param_info (specifically the `local_chunk_info`) it might be the full value | |
| or a specific slice. | |
| """ | |
| # If saved as a numpy array, but a partitioned read is requested, return a | |
| # slice of the array for that host. Otherwise, return the whole thing. | |
| if isinstance(maybe_tspec, np.ndarray) and param_info: | |
| if param_info.local_chunk_info: | |
| arr = maybe_tspec | |
| return arr[param_info.local_chunk_info.slice] | |
| else: | |
| return maybe_tspec | |
| # If we have anything else that isn't a tensorstore spec just return it. | |
| elif not isinstance(maybe_tspec, ts.Spec): | |
| return maybe_tspec | |
| tmp_ts_spec_dict = maybe_tspec.to_json() | |
| # Remove non-required params so that we can open Tensorstore | |
| # that was created with a different set of params. | |
| del tmp_ts_spec_dict['metadata']['chunks'] | |
| del tmp_ts_spec_dict['metadata']['compressor'] | |
| # Convert the relative path in the spec to a path based on the checkpoint | |
| # location. Path and gcs bucket (if applicable) information is updated | |
| # in-place. | |
| _update_ts_path_from_relative_to_absolute( | |
| os.path.dirname(ckpt_path), tmp_ts_spec_dict) | |
| if param_info.shape is not None: | |
| ts_spec_arr_shape = tuple(tmp_ts_spec_dict['metadata']['shape']) | |
| # Check that the shapes of the array on disk match the expected shape based | |
| # on the optimizer that is being restored. | |
| if ts_spec_arr_shape != param_info.shape: | |
| raise ValueError(f'Shape of `{param_info.name}` in checkpoint ' | |
| f'{ts_spec_arr_shape} does not match expected ' | |
| f'{param_info.shape}.') | |
| if ('dtype' in tmp_ts_spec_dict and tmp_ts_spec_dict['dtype'] | |
| == 'uint16') or ('dtype' in tmp_ts_spec_dict['metadata'] and | |
| tmp_ts_spec_dict['metadata']['dtype'] == '<u2'): | |
| raise ValueError( | |
| f'Found unsupported uint16 type in Tensorstore spec: {tmp_ts_spec_dict}. ' | |
| 'Please use t5x/google/scripts/convert_uint16_checkpoint.py ' | |
| 'to update saved types to bfloat16.') | |
| if restore_dtype is not None: | |
| tmp_ts_spec_dict = { | |
| 'base': tmp_ts_spec_dict, | |
| 'driver': 'cast', | |
| 'dtype': jnp.dtype(restore_dtype).name | |
| } | |
| if mesh is None or axes is None: | |
| # Read the array. | |
| t = await ts.open(tmp_ts_spec_dict, open=True) | |
| if param_info.local_chunk_info is not None: | |
| # Just read the subsection we care about. | |
| t = t[param_info.local_chunk_info.slice] | |
| arr = await t.read() | |
| else: | |
| # if provided, read as GDA | |
| arr = await gda_serialization.async_deserialize(mesh, axes, | |
| tmp_ts_spec_dict) | |
| return arr | |
| def fake_param_info(maybe_tspec: Any) -> Optional[_ParameterInfo]: | |
| """Create _ParameterInfo that results in a full read.""" | |
| # tspec is only None for `param_states` where the associated variable | |
| # is not updated by any optimizers. By setting the parameter info for | |
| # this to None, we can later short circut processing these subtrees | |
| # during loading. | |
| if maybe_tspec is None: | |
| return None | |
| local_chunk_info = None | |
| tspec = None | |
| if isinstance(maybe_tspec, ts.Spec): | |
| tspec = maybe_tspec | |
| local_chunk_info = partitioning.LocalChunkInfo( | |
| slice=(slice(None, None),), replica_id=0) | |
| return _ParameterInfo( | |
| name='', # We don't ever use the name. | |
| shape=tuple(tspec.to_json()['metadata']['shape']) if tspec else None, | |
| # We just believe the spec in the file. | |
| ts_spec=tspec, | |
| local_chunk_info=local_chunk_info, | |
| axes=None) | |
| def find_checkpoint(path: str, step: Optional[int] = None) -> str: | |
| """Find the checkpoint file based on paths and steps. | |
| Args: | |
| path: The location of the checkpoint. Can point to the `model_dir`, the | |
| checkpoint dir with a step, or the actual checkpoint file. | |
| step: The step to load. Only used if you are pointing to the `model_dir` | |
| Raises: | |
| ValueError if the checkpoint file can't be found. | |
| Returns: | |
| The path to the checkpoint file. | |
| """ | |
| # If you aren't pointing at the msgpack checkpoint file | |
| if gfile.isdir(path): | |
| # If you didn't specify a step | |
| if step is None: | |
| # Try to get the most recent step. | |
| step = latest_step(path) | |
| # If you found a step then you were pointing at model_dir, set the path to | |
| # the msgpack file in the checkpoint dir. | |
| if step: | |
| path = get_checkpoint_dir(path, step) | |
| # You gave a step, use it. | |
| else: | |
| path = get_checkpoint_dir(path, step) | |
| # Whether you supplied a step, found a step, or were already pointing at the | |
| # step, you are not pointing at a step directory, so now point to the | |
| # msgpack file. | |
| path = os.path.join(path, 'checkpoint') | |
| # You weren't point to a dir so you were pointing at the msgpack file. | |
| # Check that we found a checkpoint file. | |
| if not gfile.exists(path) or gfile.isdir(path): | |
| raise ValueError(f'Path is not a valid checkpoint: {path}') | |
| return path | |
| def load_t5x_checkpoint( | |
| path: str, | |
| step: Optional[int] = None, | |
| state_transformation_fns: Sequence[RestoreStateTransformationFn] = (), | |
| remap: bool = True, | |
| restore_dtype: Optional[jnp.dtype] = None, | |
| lazy_parameters: bool = False) -> PyTreeDef: | |
| """Load a T5X checkpoint without pre-defining the optimizer. | |
| Note: | |
| This only works for T5X checkpoints, not TF checkpoints. | |
| Args: | |
| path: The location of the checkpoint. | |
| step: The checkpoint from which step should be loaded. | |
| state_transformation_fns: Transformations to apply, in order, to the state | |
| after reading. | |
| remap: Whether to rename the checkpoint variables to the newest version. | |
| restore_dtype: optional dtype to cast targets to after restoring. If None, | |
| no parameter casting is performed. | |
| lazy_parameters: whether to load the parameters as LazyArrays to preserve | |
| memory. | |
| Returns: | |
| A nested dictionary of weights and parameter states from the checkpoint. | |
| """ | |
| path = find_checkpoint(path, step) | |
| logging.info('Restoring from checkpoint: %s', path) | |
| # The msgpack file will have all the info we need about the parameter layout. | |
| with gfile.GFile(path, 'rb') as fp: | |
| ckpt_contents = serialization.msgpack_restore(fp.read()) | |
| # If reading a ckpt that was written with gfile driver but the current | |
| # session uses the gcs driver, convert the ckpt's driver to gcs. | |
| if path.startswith('gs://'): | |
| ckpt_contents = _maybe_update_ts_from_file_to_gcs(ckpt_contents) | |
| # If a ckpt was saved in gcs and is being loaded locally, then convert the | |
| # driver to file or gfile. If the ckpt was not saved in gcs, do not change. | |
| else: | |
| ckpt_contents = _maybe_update_ts_from_gcs_to_file(ckpt_contents) | |
| # Remap that variable names to the most recent formatting. | |
| if remap: | |
| ckpt_optimizer_state = _get_optimizer_state_dict(ckpt_contents, {}, | |
| state_transformation_fns) | |
| # If we aren't remapping names we at least need to index into the checkpoint | |
| # file blob to make sure we are only dealing with the optimizer state. | |
| else: | |
| # Grab a subsection of the file depending on the version. | |
| version = ckpt_contents.get('version', 0) | |
| if version == 0: | |
| ckpt_optimizer_state = ckpt_contents | |
| else: | |
| ckpt_optimizer_state = ckpt_contents['optimizer'] | |
| # Replace all dicts of tensorstore specs with actual `ts.Spec`s. | |
| # When a checkpoint was trained using a MultiOptimizer, some of the parameter | |
| # states may be set to `None` (when a parameter was untouched by any | |
| # optimizer). We still needs references to these in our state so we keep | |
| # empty nodes. | |
| ckpt_optimizer_state_with_specs = ( | |
| state_utils.flatten_state_dict( | |
| ckpt_optimizer_state, keep_empty_nodes=True)) | |
| ckpt_optimizer_state_with_specs = { | |
| k: ts.Spec(v) if isinstance(v, dict) else v | |
| for k, v in ckpt_optimizer_state_with_specs.items() | |
| } | |
| # Create fake parameter info that results in reading the whole variable. | |
| param_infos = { | |
| k: fake_param_info(v) for k, v in ckpt_optimizer_state_with_specs.items() | |
| } | |
| ckpt_optimizer_state_with_specs = traverse_util.unflatten_dict( | |
| ckpt_optimizer_state_with_specs, sep='/') | |
| param_infos = traverse_util.unflatten_dict(param_infos, sep='/') | |
| def _create_lazy_awaitable_array( | |
| param_info: _ParameterInfo, maybe_ts_spec: Any, ckpt_path: str, | |
| restore_dtype: Optional[jnp.dtype]) -> LazyAwaitableArray: | |
| get_fn = functools.partial( | |
| _read_ts, | |
| param_info, | |
| maybe_ts_spec, | |
| ckpt_path=ckpt_path, | |
| restore_dtype=restore_dtype) | |
| return LazyAwaitableArray.from_tensor_store_spec_or_array( | |
| maybe_ts_spec, get_fn, dtype=restore_dtype) | |
| state_dict = jax.tree_multimap( | |
| functools.partial( | |
| _create_lazy_awaitable_array, | |
| ckpt_path=path, | |
| restore_dtype=restore_dtype), param_infos, | |
| ckpt_optimizer_state_with_specs) | |
| if not lazy_parameters: | |
| future_state_dict = jax.tree_map(lambda x: x.get_async(), state_dict) | |
| state_dict = _run_future_tree(future_state_dict) | |
| if restore_dtype is not None: | |
| state_dict['target'] = _cast(state_dict['target'], restore_dtype) | |
| return state_dict | |