# Copyright 2024 Google LLC # # 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 # # https://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. # pylint: disable=line-too-long r"""Pseudocode for training NeuralGCM models.""" from collections.abc import Iterable, Iterator, Mapping, Sequence import dataclasses import functools import logging import math from typing import Any, Callable, NamedTuple, Optional, cast from absl import app import datasets from dinosaur import typing from dinosaur import xarray_utils import gin from google_proprietary_code import checkpoint from google_proprietary_code import experiment from google_proprietary_code import experiment_utils # proprietary imports from google_proprietary_code import profiling_util from google_proprietary_code import streaming from google_proprietary_code import timing_util import haiku as hk import jax import jax.numpy as jnp import jax.sharding import model.reference_code.metrics as metrics import model.reference_code.metrics_base as metrics_base import model.reference_code.metrics_util as metrics_util from ml_collections import config_dict from model.legacy import model_builder from model.legacy import model_utils from model.legacy import optimization from model.legacy import physics_specifications import numpy as np import optax import pandas as pd import model.reference_code.reader as reader import model.reference_code.stochastic_losses as stochastic_losses import tensorflow as tf import model.reference_code.train_utils as train_utils import xarray Params = typing.Params PyTree = Any TrajectoryRepresentations = typing.TrajectoryRepresentations TrajectoryFn = Callable[ [Params, jax.Array, PyTree, PyTree], tuple[TrajectoryRepresentations, TrajectoryRepresentations], ] tree_map = jax.tree_util.tree_map # pylint: disable=logging-format-interpolation # pylint: disable=logging-fstring-interpolation @gin.configurable(allowlist=['constructor']) def get_loss_obj( trajectory_spec: metrics_util.TrajectorySpec, constructor: Callable[..., metrics_base.Loss] = gin.REQUIRED, ) -> metrics_base.Loss: """Returns configured loss_fn on first `trajectory_length` time slices.""" return constructor(trajectory_spec) EvaluatorDict = dict[str, metrics_base.Evaluator] TrainEvalIteratorTuple = tuple[ Iterator[Any], train_utils.TrainStepFunction, Callable[..., Any], dict[str, Any], ] @gin.configurable(allowlist=['constructor']) def get_metrics_dict( trajectory_spec: metrics_util.TrajectorySpec, eval_time_steps: Sequence[int], loss: metrics_base.Loss, constructor: Callable[..., EvaluatorDict] = metrics.default_metrics, is_ensemble_data: bool = False, ) -> EvaluatorDict: """Returns configured loss_fn on first `trajectory_length` time slices.""" return constructor( trajectory_spec, eval_time_steps, loss, is_ensemble_data=is_ensemble_data ) # start legacy configurables # # Keep these around for now (even though they are no-ops) so we can run # inference on old models. @gin.configurable def get_loss_fn(loss_fn): raise NotImplementedError @gin.register def weighted_l2_cumulative_loss(weights, scale): raise NotImplementedError # end legacy configurables def ema_params_tree(num_steps): """Creates an EMAParamsTree object based on num_steps. Args: num_steps: average number of optimization steps to include in the exponential moving average of model weights. Returns: Haiku module. """ # https://en.wikipedia.org/wiki/Moving_average#Relationship_between_SMA_and_EMA decay = 1 - 2 / (num_steps + 1) return hk.EMAParamsTree(decay) def _model_inner_steps_for_data( ds: xarray.Dataset, model_specs: model_builder.ModelSpecs, rtol: float = 1e-6, ) -> int: """Calculates model inner steps based on data time step.""" data_dt = xarray_utils.nondim_time_delta_from_time_axis( ds.time.data, model_specs.physics_specs ) inner_steps = round(data_dt / model_specs.dt) if abs(inner_steps * model_specs.dt - data_dt) / data_dt > rtol: raise RuntimeError( # pylint: disable=g-doc-exception f'{model_specs.dt=} does not divide evenly into {data_dt=}' ) return inner_steps def _get_datetime_forecast_starts( sample_count: int, first_start: pd.Timestamp, last_start: pd.Timestamp, ) -> pd.DatetimeIndex: """Get equispaced forecast start times for evaluating against ERA5.""" if first_start.hour != 0: raise ValueError(f'dataset times must start at midnight: {first_start=}') # Round-up to midnight following the last forecast day (e.g., the start of # the next year). stop = last_start.ceil('1D') # Equally spaced from start (inclusive) to stop (exclusive). start_times = pd.date_range(first_start, stop, periods=sample_count + 1)[:-1] # To match ECMWF, all forecasts should be initialized at 0z or 12z. Here we # alternate start times. parity = np.arange(sample_count) % 2 return start_times.round('1D') + parity * pd.Timedelta('12H') P = jax.sharding.PartitionSpec class ExperimentState(NamedTuple): opt_state: PyTree params: PyTree ema_params: PyTree class Experiment(experiment.AbstractExperiment): """Training experiment based on trajectory loss minimization.""" def __init__( self, experiment_dir: str, config: Optional[config_dict.ConfigDict] = None, ): """Creates an instance of a training scheme class. Args: experiment_dir: Path to experiment directory. config: config struct setting up the experiment. """ if config is None: config = experiment_config.get_config() super().__init__( experiment_dir, config.distributed_training, writer_names=['train', 'eval', 'eval_ema'], ) logging.info('Experiment config:\n%s', config) self.config = config self.train_ds = xarray_utils.open_dataset(config.train_dataset_path) self.eval_ds = xarray_utils.open_dataset(config.eval_dataset_path) if 'sample' in self.train_ds.dims: logging.warning('only using the first sample!') self.train_ds = self.train_ds.isel(sample=0, drop=True) self.eval_ds = self.eval_ds.isel(sample=0, drop=True) train_attrs = self.train_ds.attrs # Model instantiation and trajectory unroll functions # Note: we use interactive mode in experiments to split gin-configurations # into separate, distinct parts provided in the config_dict. gin.enter_interactive_mode() self.is_nodal = self.config.is_nodal # parse and override all the gin things gin.parse_config(config.model_gin_config) gin.parse_config(config.optimizer_gin_config) experiment_utils.parse_config_dict(config.gin_overrides) logging.info('Parsed gin config string:\n%s', gin.config_str()) full_model_gin_config = gin.config_str() # do not include physics config. logging.info('With overrides gin config string:\n%s', gin.config_str()) self.data_coords = model_builder.coordinate_system_from_dataset( self.train_ds ) logging.info(f'{self.model_parallel_training=}') if self.model_parallel_training: # It does not make sense to use spatial parallelism with batch size per # device larger than 1. Instead, you would get better performance from # using less model parallelism. if self.spatial_parallelism > 1 and self.config.batch_size_per_device > 1: raise NotImplementedError( f'{self.config.batch_size_per_device=} is not supported for model ' 'parallel training' ) self.data_coords = dataclasses.replace( self.data_coords, spmd_mesh=self.spmd_mesh ) # try getting aux_features from dataset, if not included we rely on # `model_builder.get_model_specs` to supply necessary values. try: data_aux_features = xarray_utils.aux_features_from_xarray(self.train_ds) except KeyError: data_aux_features = {} # when available, we parse physics_config_str from metadata in train_attrs. if 'physics_config_str' in train_attrs: physics_config_str = train_attrs['physics_config_str'] experiment_utils.parse_gin_config_without_imports(physics_config_str) else: logging.info( 'physics_config_str was not provided in the dataset, ' 'hence it is expected to be specified in model_gin_config.' ) self.physics_specs = physics_specifications.get_physics_specs() self.model_specs = model_builder.get_model_specs( self.data_coords, self.physics_specs, data_aux_features ) logging.info(f'{self.model_specs=}') self.train_inner_steps = _model_inner_steps_for_data( self.train_ds, self.model_specs ) self.eval_inner_steps = _model_inner_steps_for_data( self.eval_ds, self.model_specs ) if ( len(self.config.train_schedule_time_steps) != len(self.config.train_schedule_boundaries) + 1 ): raise ValueError( f'{self.config.train_schedule_time_steps} should be one longer ' f'than {self.config.train_schedule_boundaries} but was not.' ) if any( t % self.train_inner_steps for t in self.config.train_schedule_time_steps ): raise ValueError( f'{self.train_inner_steps=} does not divide ' f'{self.config.train_schedule_time_steps=}' ) if any(t % self.eval_inner_steps for t in self.config.eval_time_steps): raise ValueError( f'{self.eval_inner_steps=} does not divide ' f'{self.config.eval_time_steps=}' ) if max(self.config.eval_time_steps) < max( self.config.train_schedule_time_steps ): raise ValueError( f'Training will not work since {max(self.config.eval_time_steps)=} <' f' {max(self.config.train_schedule_time_steps)=}' ) self._eval_trajectory_length = ( max(self.config.eval_time_steps) // self.eval_inner_steps + 1 ) self._trajectory_lengths = [ self.config.num_init_frames + n // self.train_inner_steps for n in self.config.train_schedule_time_steps ] self._max_trajectory_length = max( self._trajectory_lengths + [self._eval_trajectory_length] ) self.reference_datetime = self.model_specs.aux_features[ xarray_utils.REFERENCE_DATETIME_KEY ] self.whirl_model = model_builder.WhirlModel( **self.model_specs, input_coords=self.data_coords, output_coords=self.data_coords, ) self.from_xarray_fn = self.whirl_model.from_xarray_fn def trajectory_fwd(x, forcing_data, model, outer_steps, inner_steps): trajectory_fn = model_utils.trajectory_with_inputs_and_forcing( model, config.num_init_frames, start_with_input=True ) return trajectory_fn(x, forcing_data, outer_steps, inner_steps) self._trajectory_fwd = trajectory_fwd # Checkpoint items. self._model_dt = self.model_specs.dt self._model_gin_config = full_model_gin_config # optimizer configuration. self.optimizer = optimization.optimizer() # exponentially moving average params tracking. ema_fn = hk.without_apply_rng( hk.transform_with_state( lambda x: ema_params_tree(config.ema_num_steps)(x) # pylint: disable=unnecessary-lambda ) ) self._ema_init = jax.jit(ema_fn.init) def ema_update(params, ema_state): return ema_fn.apply(None, ema_state, params) self._ema_update = jax.jit(ema_update) logging.info('Final active config string:\n%s', gin.config_str()) # # Data inputs methods. # @functools.cached_property def spmd_mesh(self) -> jax.sharding.Mesh: n = self.config.model_parallelism.ensemble_shards z = self.config.model_parallelism.z_shards x = self.config.model_parallelism.x_shards y = self.config.model_parallelism.y_shards global_batch = jax.device_count() // (n * z * x * y) if global_batch == 0: raise ValueError( f'{jax.device_count()=} is insufficient for ' f'{self.config.model_parallelism=}' ) return train_utils.create_spmd_mesh( {'batch': global_batch, 'ensemble': n, 'z': z, 'x': x, 'y': y} ) @functools.cached_property def degree_of_model_parallelism(self) -> int: return math.prod(v for k, v in self.spmd_mesh.shape.items() if k != 'batch') @functools.cached_property def model_parallel_training(self) -> bool: return self.degree_of_model_parallelism > 1 @functools.cached_property def spatial_parallelism(self) -> int: return math.prod(self.spmd_mesh.shape[k] for k in 'zxy') def to_global_array(self, pytree: PyTree, global_batch_size: int) -> PyTree: """Create a pytree of global JAX arrays from a pytree of NumPy arrays.""" # partition arrays along batch and spatial dimensions return train_utils.make_distributed_array_from_local_arrays( pytree, self.spmd_mesh, self.data_coords.physics_partition_spec, global_batch_size, ) def num_eval_batches(self, large_eval: bool) -> int: if large_eval: return self.config.num_eval_batches[-1] else: return self.config.num_eval_batches[0] def steps_between_evals(self, large_eval: bool) -> int: if large_eval: return self.config.steps_between_evals[-1] else: return self.config.steps_between_evals[0] def eval_batch_size_per_device(self, large_eval: bool) -> int: if large_eval: return self.config.eval_batch_size_per_device[-1] else: return self.config.eval_batch_size_per_device[0] @functools.cached_property def global_batch_size(self) -> int: return ( jax.device_count() // self.degree_of_model_parallelism * self.config.batch_size_per_device ) def global_eval_batch_size(self, large_eval: bool) -> int: return ( jax.device_count() // self.degree_of_model_parallelism * self.eval_batch_size_per_device(large_eval) ) def local_eval_batch_size(self, large_eval: bool) -> int: return jax.local_device_count() * self.eval_batch_size_per_device( large_eval ) def _to_dataset_iter( self, data: tf.data.Dataset, template: xarray.Dataset ) -> Callable[[], Iterable[Any]]: """Convert a tf.data.Dataset into a function that makes a data iterator.""" leading_dims_set = {x.shape[:2] for x in data.element_spec.values()} assert len(leading_dims_set) == 1, leading_dims_set local_batch_size, time_series_length = leading_dims_set.pop() template = ( template.drop_vars('time') .head(time=time_series_length) .pipe(xarray.zeros_like) # replace data with zeros .pipe(datasets.drop_static_vars) .transpose('time', ...) .expand_dims(batch=local_batch_size) ) def make_iterator(): for example_dict in data.as_numpy_iterator(): yield self.from_xarray_fn(template.copy(data=example_dict)) return make_iterator def _read_shuffled_shard( self, dataset: xarray.Dataset, time_series_length: int, min_buffer_blocks: int, shard_index: int, shard_count: int, ) -> tf.data.Dataset: sampler = reader.Windower( window_size=time_series_length, stride_between_windows=self.config.train_time_sample_offset, ) local_shard_count = max( self.degree_of_model_parallelism, jax.local_device_count() ) seed = train_utils.combine_rng_seeds( self.config.dataset_rng_seed, shard_index, time_series_length ) data = reader.read_shuffled_shard( dataset, sampler, block_size_in_bytes=self.config.block_size_in_bytes / local_shard_count, buffer_size_in_bytes=( self.config.shuffle_buffer_size_in_bytes / local_shard_count ), min_buffer_blocks=min_buffer_blocks, shard_index=shard_index, shard_count=shard_count, seed=seed, ) return data def _get_train_dataset(self) -> xarray.Dataset: train_dataset = xarray_utils.ds_with_sim_time( self.train_ds, self.physics_specs, self.reference_datetime ) if self.config.train_dataset_time_slice: time_slice = slice(*self.config.train_dataset_time_slice) train_dataset = train_dataset.sel(time=time_slice) if self.config.time_subsample_rate != 1: raise NotImplementedError('subsampling on the fly is not supported yet') if self.config.add_noise_to_input: raise NotImplementedError('add_noise_to_input not supported yet') return train_dataset def _spatial_addressable_indices_map( self, spatial_dim_sizes: tuple[int, int, int] ) -> Mapping[jax.Device, tuple[slice, slice, slice, slice]]: """Get slices for indexing global arrays to local devices.""" spec = P('batch', *self.data_coords.physics_partition_spec) sharding = jax.sharding.NamedSharding(self.spmd_mesh, spec) global_shape = (self.global_batch_size,) + spatial_dim_sizes indices_map = sharding.addressable_devices_indices_map(global_shape) indices_map = cast( Mapping[jax.Device, tuple[slice, slice, slice, slice]], indices_map ) return indices_map def _read_model_parallel_dataset( self, dataset: xarray.Dataset, read_shard: Callable[[xarray.Dataset, int], tf.data.Dataset], ) -> tuple[tf.data.Dataset, xarray.Dataset]: """Read a shard of a training dataset into tf.data.Dataset.""" indices_map = self._spatial_addressable_indices_map( tuple(dataset.sizes[k] for k in ['level', 'longitude', 'latitude']) ) shard_data: list[tf.data.Dataset] = [] for device in jax.local_devices(): indices = indices_map[device] batch_index = indices[0].start or 0 selection = dict(zip(['level', 'longitude', 'latitude'], indices[1:])) shard_dataset = dataset.isel(selection) shard_data.append(read_shard(shard_dataset, batch_index)) choices = tf.data.Dataset.range(jax.local_device_count()).repeat() data = tf.data.Dataset.choose_from_datasets(shard_data, choices) template = shard_dataset return data, template def _build_train_inputs( self, time_series_length: int ) -> tuple[Callable[[], Any], dict[str, Any]]: """Loads the training dataset and returns an iterator and train_attrs.""" train_dataset = self._get_train_dataset() local_batch_size = ( # Size of data needed to satisfy batch_size_per_device. self.config.batch_size_per_device * jax.local_device_count() ) if self.model_parallel_training: def read_shard(shard_dataset, batch_index): return self._read_shuffled_shard( shard_dataset, time_series_length, min_buffer_blocks=1, shard_index=batch_index, shard_count=self.global_batch_size, ) data, template = self._read_model_parallel_dataset( train_dataset, read_shard ) else: data = self._read_shuffled_shard( train_dataset, time_series_length, shard_index=jax.process_index(), shard_count=jax.process_count(), min_buffer_blocks=local_batch_size, ) template = train_dataset data = data.repeat() data = data.batch(local_batch_size, drop_remainder=True) data = data.prefetch(tf.data.AUTOTUNE) train_iter = self._to_dataset_iter(data, template) data_attrs = datasets.attrs_from_dataset(train_dataset, time_series_length) return train_iter, data_attrs def build_train_and_eval_iterators( self, schedule_idx: int, start_step: int, large_eval: bool, ) -> TrainEvalIteratorTuple: """Build new iterators for training at schedule_idx. Args: schedule_idx: Index into the rollout schedule. start_step: Step at which this training run started at. This does not change unless the Borg job dies and restarts. large_eval: Whether this evaluation should be done over a larger set of data. Returns: TrainEvalIteratorTuple: Tuple consisting of get_train_data. Iterator providing next set of training data. train_step_fn. train_utils.TrainStepFunction to update weights. evaluate_fn. Callable to evalate metrics and write results. ckpt_kwargs. dict[str, Any] of kwargs to add to the checkpoint. """ num_train_time_steps = self.config.train_schedule_time_steps[schedule_idx] trajectory_length = self._trajectory_lengths[schedule_idx] train_traj_spec = metrics_util.TrajectorySpec( trajectory_length, self._max_trajectory_length, self.train_inner_steps, coords=self.model_specs.coords, data_coords=self.data_coords, ) # These only change on the first and last rollout, but re-make them anyways. get_eval_data, eval_attrs = self.build_eval_inputs( self.config.eval_dataset_time_slice, large_eval, ) get_eval_on_train, _ = self.build_eval_inputs( self.config.train_dataset_time_slice, large_eval, ) evaluate_fn = functools.partial( self.evaluate, eval_batch_fn=self.get_eval_batch_fn(train_traj_spec), get_eval_data=get_eval_data, get_train_data=get_eval_on_train, large_eval=large_eval, ) get_train_data, train_attrs = self._build_train_inputs(trajectory_length) ckpt_kwargs = { 'train_attrs': train_attrs, 'eval_attrs': eval_attrs, } train_step_fn = self.get_train_step_fn( num_train_time_steps, train_traj_spec ) if ( self.config.profile_with_xprof and schedule_idx == 0 and experiment_utils.is_coordinator() and start_step == 0 ): train_step_fn = profiling_util.Traced( train_step_fn, # only the initial train_step is profiled. trace_name='train step', skip_steps=2, # avoid JIT compilation num_trace_steps=3, enable_python_tracer=True, host_trace_level=3, ) return get_train_data(), train_step_fn, evaluate_fn, ckpt_kwargs def _get_eval_dataset(self) -> xarray.Dataset: return xarray_utils.ds_with_sim_time( self.eval_ds, self.physics_specs, self.reference_datetime ) def build_eval_inputs( self, dataset_time_slice: tuple[str, str] | None, large_eval: bool, ) -> tuple[Callable[[], Any], Any]: """Returns an iterable over the data and data attrs for evaluation.""" eval_dataset = self._get_eval_dataset() num_eval_batches = self.num_eval_batches(large_eval) eval_batch_size_per_device = self.eval_batch_size_per_device(large_eval) time_series_length = ( self._eval_trajectory_length + self.config.num_init_frames - 1 ) local_batch_size = eval_batch_size_per_device * jax.local_device_count() if isinstance(eval_dataset.indexes['time'], pd.DatetimeIndex): # For real world training data from ERA5, carefully sample starting # times to ensure they are equally spaced across the year. assert 'sample' not in eval_dataset.dims logging.info('Using eval data loader for build_eval_inputs') sample_count = self.global_eval_batch_size(large_eval) * num_eval_batches if dataset_time_slice: time_source = eval_dataset.time.loc[slice(*dataset_time_slice)] else: time_source = eval_dataset first_start = time_source.indexes['time'][0] last_start = time_source.indexes['time'][-1] starts = _get_datetime_forecast_starts( sample_count, first_start, last_start ) logging.info(f'determined evaluation data for {sample_count=}: {starts=}') offsets = eval_dataset.indexes['time'].get_indexer(starts) sampler = reader.WindowerAtOffsets( window_size=time_series_length, window_offsets=offsets ) if self.model_parallel_training: def read_shard(shard_dataset, batch_index): selector = reader.ShardSelector(batch_index, len(starts)) return reader.read_timeseries(shard_dataset, sampler, selector) data, template = self._read_model_parallel_dataset( eval_dataset, read_shard ) else: selector = reader.ShardSelector( jax.process_index(), jax.process_count() ) data = reader.read_timeseries(eval_dataset, sampler, selector) template = eval_dataset data = data.batch(local_batch_size, drop_remainder=True) data = data.cache() else: # For synthetic datasets (e.g., from Held-Suarez), use the same shuffling # we use for reading training data. assert not self.model_parallel_training if dataset_time_slice: eval_dataset = eval_dataset.sel(time=slice(*dataset_time_slice)) data = self._read_shuffled_shard( eval_dataset, time_series_length, shard_index=jax.process_index(), shard_count=jax.process_count(), min_buffer_blocks=local_batch_size * num_eval_batches, ) data = data.batch(local_batch_size, drop_remainder=True) data = data.take(num_eval_batches) template = eval_dataset logging.info(f'created eval data: {data}') eval_iter = self._to_dataset_iter(data, template) data_attrs = datasets.attrs_from_dataset(eval_dataset, time_series_length) return eval_iter, data_attrs # # Training and evaluation methods. # def _make_initial_experiment_state( self, rng: typing.PRNGKeyArray, init_example, init_forcing_data: typing.ForcingData, init_params: Optional[typing.Params] = None, ) -> ExperimentState: """Makes initial parameters (via hk.Module.init).""" if self.eval_inner_steps != self.train_inner_steps: raise ValueError( 'KroneckerCorrelatedL2LossModule stddev will be ill defined since ' f'{self.eval_inner_steps=} != {self.train_inner_steps=}' ) trajectory_length = list(init_example.values())[0].shape[0] @jax.jit def init(rng, init_example, init_forcing_data): outer_steps = (trajectory_length - self.config.num_init_frames) + 1 # We need an "ensemble" dimension for stochastic losses, but parameters # are fully replicated across the ensemble. if init_params is None: init_fn = jax.vmap( self._make_unbatched_trajectory_fn(outer_steps).init, in_axes=None, out_axes=0, spmd_axis_name='ensemble', axis_size=1, ) unsqueezd_params = init_fn(rng, init_example, init_forcing_data) params = tree_map(lambda x: jnp.squeeze(x, axis=0), unsqueezd_params) else: params = init_params opt_state = self.optimizer.init(params) _, ema_state = self._ema_init(None, params) experiment_state = ExperimentState(opt_state, params, ema_state) experiment_state = train_utils.ensure_replicated( experiment_state, mesh=self.spmd_mesh ) return experiment_state return init(rng, init_example, init_forcing_data) def _make_unbatched_trajectory_fn(self, outer_steps: int): """Haiku transformation of func giving (prediction, target) trajectories. Args: outer_steps: Number of outer steps the trajectory should take. Returns: hk transformed object. The .apply member maps (params, rng, target, forcing_data) --> (prediction, target) """ if self.train_inner_steps != self.eval_inner_steps: # We share a trajectory for train/eval...so the spacing better be equal. raise ValueError(f'{self.train_inner_steps=} != {self.eval_inner_steps=}') @hk.transform def unbatched_trajectory_fn(target, forcing_data): """Compute Fwd(target[0]) on one single batch/device.""" # Shapes(target) ~ (n_t, n_z, n_m, n_l) model = self.whirl_model.model_cls() _, predicted_trajectory = self._trajectory_fwd( x=target, forcing_data=forcing_data, model=model, outer_steps=outer_steps, inner_steps=self.train_inner_steps, ) prediction, target = ( model_utils.compute_prediction_and_target_representations( predicted_trajectory, target, forcing_data, model ) ) return prediction, target return unbatched_trajectory_fn def _make_batch_trajectory_fn( self, outer_steps: int, ) -> TrajectoryFn: """Target, prediction representations with shape (batch, ensemble, ...).""" ensembled_fn = jax.vmap( # (params, rng, target, forcing_data) --> (prediction, target) self._make_unbatched_trajectory_fn(outer_steps).apply, in_axes=(None, 0, None, None), spmd_axis_name='ensemble', ) batch_ensembled_fn = jax.vmap( # (params, rng, target, forcing_data) --> (prediction, target) # Input shapes are: # params: (...) # rng: (batch, ensemble, ...) # target: (batch, time, ...) # forcing_data: (batch, time, ...) ensembled_fn, in_axes=(None, 0, 0, 0), spmd_axis_name='batch', ) return batch_ensembled_fn def get_train_step_fn( self, num_train_time_steps: int, traj_spec: metrics_util.TrajectorySpec, ) -> train_utils.TrainStepFunction: """Makes a function to update weights via gradient descent. This function makes use of on-device batching. Multiple devices are combined via an all reduce step whereby the average (across devices) gradient is applied to each (on device) params. Args: num_train_time_steps: Number of time steps for the trajectory in this training step. traj_spec: Specification of training trajectory. Returns: all_reduced_train_step: Function to exectute one training step. Params are injected into self._trajectory_fwd and Loss (if Loss requires Haiku params). """ batch_trajectory_fn = self._make_batch_trajectory_fn( # (params, rng, target, forcing_data) --> [prediction, target] outer_steps=num_train_time_steps // self.train_inner_steps + 1, ) loss_fn = get_loss_obj(traj_spec).evaluate ensembled_loss_fn = jax.vmap( loss_fn, axis_name='ensemble', spmd_axis_name='ensemble', ) batch_ensembled_loss_fn = jax.vmap( ensembled_loss_fn, axis_name='batch', spmd_axis_name='batch', ) def batched_parameter_loss_fn(params, rng, target, forcing_data): """Mean (over on-device batch members) of loss w.r.t parameters.""" # Input shapes are: # params: (...) # rng: (batch, ensemble, ...) # target: (batch, time, ...) # forcing_data: (batch, time, ...) prediction, target = batch_trajectory_fn( # The `target` and prediction returned are TrajectoryRepresentations. # So don't just re-use the arg `target`. params, rng, target, forcing_data, ) # dimensions (batch, ensemble) per_example_loss = batch_ensembled_loss_fn(prediction, target) # Average over ensemble and batch dimensions (technically, we don't have # to average over ensemble with our current stochastic losses, but these # values are already identical and this is cleaner than using array # indexing) overall_loss = jnp.mean(per_example_loss, axis=(0, 1)) assert overall_loss.ndim == 0 return overall_loss # We would use donate_argnums here to update experiment_state in-place, but # that would mean we could not save the checkpoint in a separable thread. # Fortunately experiment_state is usually not too big (~100 MB). @train_utils.jit_once def train_step(experiment_state, rng, target_trajectory, forcing_data): opt_state, params, ema_state = experiment_state rng = train_utils.ensure_sharded_rng_key(rng, mesh=self.spmd_mesh) loss, grad = jax.value_and_grad(batched_parameter_loss_fn)( params, rng, target_trajectory, forcing_data ) updates, opt_state = self.optimizer.update(grad, opt_state, params) params = optax.apply_updates(params, updates) _, ema_state = self._ema_update(params, ema_state) experiment_state = ExperimentState(opt_state, params, ema_state) experiment_state = train_utils.ensure_replicated( experiment_state, mesh=self.spmd_mesh ) return experiment_state, loss return train_step def get_eval_batch_fn( self, train_traj_spec: metrics_util.TrajectorySpec, ) -> train_utils.EvalStepFunction: """Makes a function that performs a single evaluation pass. Args: train_traj_spec: TrajectorySpec for training. Used to add the "loss" evaluation metrics. Returns: Function mapping (params, rng, target, forcing_data) to dictionary of scalar metric values. Parameters are injected into self._trajectory_fwd and Loss (if Loss requires Haiku params). """ eval_traj_spec = metrics_util.TrajectorySpec( self._eval_trajectory_length, self._max_trajectory_length, steps_per_save=self.eval_inner_steps, coords=self.model_specs.coords, data_coords=self.data_coords, ) eval_time_steps = [ t // self.eval_inner_steps for t in self.config.eval_time_steps ] if any(t % self.eval_inner_steps for t in self.config.eval_time_steps): raise ValueError( f'cannot evaluate {self.config.eval_time_steps=} with ' f'{self.eval_inner_steps=}' ) batch_trajectory_fn = self._make_batch_trajectory_fn( # (params, rng, target, forcing_data) --> [prediction, target] outer_steps=max( self._eval_trajectory_length, train_traj_spec.trajectory_length, ), ) def unbatched_eval_fn( prediction: TrajectoryRepresentations, target: TrajectoryRepresentations ): """Evaluate(target, Fwd(target[0])) on one single batch/device.""" metrics_dict = get_metrics_dict( eval_traj_spec, eval_time_steps, get_loss_obj(train_traj_spec), is_ensemble_data=bool(self.config.ensemble_size), ) return train_utils.flatten_dict({ k: metric.evaluate(prediction, target) for k, metric in metrics_dict.items() }) ensembled_fn = jax.vmap( unbatched_eval_fn, axis_name='ensemble', spmd_axis_name='ensemble' ) batch_ensembled_fn = jax.vmap( ensembled_fn, axis_name='batch', spmd_axis_name='batch', ) @train_utils.jit_once def batch_mean_eval_fn(params, rng, target, forcing_data): """Computes mean (over batch members) of evaluation.""" # Input shapes for this function are: # params: (...) # rng: (batch, ensemble, ...) # target: (batch, time, ...) # forcing_data: (batch, time, ...) rng = train_utils.ensure_sharded_rng_key(rng, mesh=self.spmd_mesh) prediction, target = batch_trajectory_fn( # The `target` and prediction returned are TrajectoryRepresentations. # So don't just re-use the arg `target`. params, rng, target, forcing_data, ) batch_eval_values = batch_ensembled_fn(prediction, target) return tree_map(jnp.mean, batch_eval_values) return batch_mean_eval_fn def run_training(self): """See base class.""" ( start_step, times_restarted_on_nan, step_auto_restart_began_at, experiment_state, ) = self.initialize_experiment( initial_checkpoint_path=self.config.initial_checkpoint_path ) if start_step >= self.config.num_training_steps: logging.warning( f'Attempting to start training at {start_step=} >=' f' {self.config.num_training_steps=}. Will simply return' ) return def logging_callback(step, loss, times_restarted_on_nan): loss = float(jax.device_get(loss)) if step % max(self.steps_between_evals(False) // 100, 1) == 0: logging.info(f'{step=}, {loss=}') if ( self.config.error_with_nan_loss and times_restarted_on_nan > self.config.max_nan_restarts ): raise RuntimeError( f'NaN loss detected at {step=}, after too many restarts since' f' {times_restarted_on_nan=} >' f' {self.config.max_nan_restarts=}. Aborting.' ) # monitor loss using a separate thread, so it doesn't block execution logging_stream = streaming.SingleThreadExecutor(logging_callback) logging.info('starting training from step=%s', start_step) train_step_timer = timing_util.Timer() rng_stream = train_utils.BatchedPRNGSequence( jax.random.PRNGKey(self.config.init_rng_seed), batch_shape=(self.global_batch_size, self.config.ensemble_size or 1), ) loss = 1.0 # setting to a non-nan value when starting an experiment. schedule_idx = None ckpt_kwargs = {} step = start_step while step < self.config.num_training_steps: old_schedule_idx = schedule_idx schedule_idx = np.sum( # compute which leg of the schedule we are at. step > np.asarray(self.config.train_schedule_boundaries) ) large_eval = ( schedule_idx == len(self.config.train_schedule_time_steps) - 1 ) if schedule_idx != old_schedule_idx: ( train_iter, train_step_fn, evaluate_fn, ckpt_kwargs, ) = self.build_train_and_eval_iterators( schedule_idx=schedule_idx, start_step=start_step, large_eval=large_eval, ) if ( np.isnan(loss) # No sense re-initializing if we've restarted a bunch already. Also, # note that if error_with_nan_loss=True, we should raise and not have # worry about the times_restarted_on_nan < max_nan_restarts here. and times_restarted_on_nan <= self.config.max_nan_restarts ): # See also logging_callback, which may raise RuntimeError for NaN loss. logging.warning( f'NaN loss encountered at {step=}. Re-initializing and incrementing' f' times_restarted_on_nan to {times_restarted_on_nan + 1}' ) times_restarted_on_nan += 1 step, _, _, experiment_state = self.initialize_experiment( target_step=step - times_restarted_on_nan * self.config.restart_lookback_steps ) step_auto_restart_began_at = step_auto_restart_began_at or step # checkpoint if step % self.steps_between_evals(large_eval) == 0: self.save_checkpoint( step, experiment_state, checkpoint_buffer_size=self.config.checkpoint_buffer_size, times_restarted_on_nan=times_restarted_on_nan, step_auto_restart_began_at=step_auto_restart_began_at, **ckpt_kwargs, ) elif step % self.config.steps_between_checkpoints == 0: max_lookback_step = ( step - self.config.max_nan_restarts * self.config.restart_lookback_steps ) if ( # If there is no chance of a restart sequence overlapping with # previously used checkpoints... times_restarted_on_nan and max_lookback_step > step_auto_restart_began_at ): logging.info( f'Significant progress made since {step_auto_restart_began_at=}.' f' In particular, {step=} Therefore set times_restarted_on_nan' ' to 0' ) times_restarted_on_nan = 0 step_auto_restart_began_at = None self.save_checkpoint( step, experiment_state, update_latest_only=True, checkpoint_buffer_size=self.config.checkpoint_buffer_size, times_restarted_on_nan=times_restarted_on_nan, step_auto_restart_began_at=step_auto_restart_began_at, **ckpt_kwargs, ) # evaluate if (step + 1) % self.steps_between_evals(large_eval) == 0: if step > start_step: with train_step_timer: # train_step is non-blocking, so we need to block on the output # of the previous training step to reliably time it. experiment_state = jax.block_until_ready(experiment_state) eval_interval = self.steps_between_evals(large_eval) training_time = train_step_timer.total logging.info( f'training for {eval_interval} steps took ' f'{training_time:.1f} seconds' ) self.record_scalar( 'train', tag='seconds_per_train_step', step=step, value=training_time / eval_interval, ) train_step_timer = timing_util.Timer() # reset if isinstance(train_step_fn, profiling_util.Traced): memory_usage = train_step_fn.tracer.memory_usage # pytype: disable=attribute-error if memory_usage is not None: self.record_scalar( 'train', tag='peak_memory_usage_mib', step=step, value=memory_usage, ) with timing_util.Timer() as eval_timer: evaluate_fn(step, experiment_state, seed=step) self.record_scalar( 'train', tag='seconds_per_evaluation', step=step, value=eval_timer.average, ) logging.info('evaluation pass took %.1f seconds', eval_timer.average) self.flush_writers() # flush all writers for this training step. # train with train_step_timer: # go/xprof-instrument-jax with jax.profiler.StepTraceAnnotation('train', step_num=step): batch, forcing_data = self.to_global_array( next(train_iter), self.global_batch_size ) # This is necessary else ValueError. # See http://sponge2/960c64a7-703f-4a19-8572-2c97dd9c01f3 with self.spmd_mesh: experiment_state, loss = train_step_fn( experiment_state, next(rng_stream), batch, forcing_data ) # If we don't device_get (or similar), asynchronous execution # resultsin this timed block taking almost no time. device_get does # not result in longer runs, since each loop must eventually compute # the loss, one way or another. loss = jax.device_get(loss) logging_stream.wait() logging_stream.put(step, loss, times_restarted_on_nan) step += 1 # End of while step < self.config.num_training_steps: evaluate_fn(self.config.num_training_steps, experiment_state) self.finalize_training( self.config.num_training_steps, experiment_state, times_restarted_on_nan=times_restarted_on_nan, step_auto_restart_began_at=step_auto_restart_began_at, **ckpt_kwargs, ) def make_dummy_inputs(self) -> tuple[Any, Any]: train_dataset = xarray_utils.ds_with_sim_time( self.train_ds, self.physics_specs, self.reference_datetime ) dummy_ds = ( train_dataset.drop_vars('time') .head(time=self.config.num_init_frames) .pipe(xarray.zeros_like) # replace data with zeros .pipe(datasets.drop_static_vars) .transpose('time', ...) ) return self.from_xarray_fn(dummy_ds) def initialize_experiment( self, initial_checkpoint_path: Optional[str] = None, target_step: Optional[int] = None, ) -> tuple[int, int, int | None, ExperimentState]: """Returns training step and experiment state from checkpoint or init.""" if target_step: ckpt = self.load_buffered_checkpoint(target_step=target_step) if ckpt is None: logging.info( 'No acceptable buffered checkpoint found for {target_step=}' ) else: logging.info( f'Using buffered checkpoint, which has step={ckpt.step}. Ideally ' f'would have used {target_step=}' ) else: ckpt = self.load_latest_checkpoint() if ckpt is None: logging.info('No latest checkpoint found') else: logging.info(f'Using latest checkpoint, which has step={ckpt.step}') init_params = None if ckpt is None and initial_checkpoint_path is not None: ckpt = checkpoint.load_checkpoint(initial_checkpoint_path) if self.config.reset_initial_optimizer_state: init_params = ckpt.eval_params ckpt = None # if resetting optimizer, carry over only init_params. if ckpt is not None: start_step = ckpt.step logging.info(f'resuming from checkpoint at step={start_step}') times_restarted_on_nan = getattr(ckpt, 'times_restarted_on_nan', 0) step_auto_restart_began_at = getattr( ckpt, 'step_auto_restart_began_at', None ) experiment_state = ExperimentState( ckpt.opt_state, ckpt.train_params, ckpt.ema_state ) else: logging.info('starting training with new weights') start_step = 0 times_restarted_on_nan = 0 step_auto_restart_began_at = None rng = jax.random.PRNGKey(self.config.init_rng_seed) init_example, init_forcing_data = self.make_dummy_inputs() experiment_state = self._make_initial_experiment_state( rng, init_example, init_forcing_data, init_params=init_params ) return ( start_step, times_restarted_on_nan, step_auto_restart_began_at, experiment_state, ) def _checkpoint_state( self, step: int, experiment_state: ExperimentState, times_restarted_on_nan: int, step_auto_restart_began_at: int, train_attrs: dict[str, Any], eval_attrs: dict[str, Any], ) -> checkpoint.CheckpointState: """Returns a checkpoint state for a given experiment_state.""" opt_state, params, ema_state = experiment_state ema_params, _ = self._ema_update(params, ema_state) ckpt_state = checkpoint.CheckpointState( train_params=params, eval_params=ema_params, opt_state=opt_state, ema_state=ema_state, step=step, model_time_step=self._model_dt, model_config_str=self._model_gin_config, train_dataset_path=self.config.train_dataset_path, eval_dataset_path=self.config.eval_dataset_path, times_restarted_on_nan=times_restarted_on_nan, step_auto_restart_began_at=step_auto_restart_began_at, train_attrs=train_attrs, eval_attrs=eval_attrs, ) return ckpt_state def evaluate( self, step, experiment_state, eval_batch_fn, get_eval_data, get_train_data, large_eval: bool, seed=0, ): """Evaluates the model on train and eval data and writes summaries. Args: step: global training step. experiment_state: tuple of replicated step, optimizer state and EMA (exponentially moving average) state for model parameters. eval_batch_fn: function that, given parameters; rng; batch of data, computes evaluation metric of interest on the given samples. get_eval_data: callable that returns an iterator over evaluation data that is used for produce summaries on unseen evaluation data. get_train_data: callable that returns an iterator over training data that is used for produce summaries on training data. large_eval: Whether this evaluation is on the larger size eval data. seed: seed for the random number generator to be used for evaluation. """ num_eval_batches = self.num_eval_batches(large_eval) if num_eval_batches == 0: logging.warning(f'skipping evaluation: {num_eval_batches=}') return _, params, ema_state = experiment_state ema_params, _ = self._ema_update(params, ema_state) global_batch_size = self.global_eval_batch_size(large_eval) rng_stream = train_utils.BatchedPRNGSequence( jax.random.PRNGKey(seed), batch_shape=(global_batch_size, self.config.ensemble_size or 1), ) to_global_array = functools.partial( self.to_global_array, global_batch_size=global_batch_size ) # In theory, the mesh context manager should not be necessary because we use # jit with sharded arrays (rather than xmap or pjit), but it seems to be # required to avoid triggering bugs in JAX. with self.spmd_mesh: logging.info('evaluating on train dataset') metrics_ = train_utils.streaming_mean( rng_stream, map(to_global_array, get_train_data()), functools.partial(eval_batch_fn, params), ) for tag, value in metrics_.items(): self.record_scalar('train', tag=tag, value=value, step=step) logging.info('evaluating on test dataset') metrics_ = train_utils.streaming_mean( rng_stream, map(to_global_array, get_eval_data()), functools.partial(eval_batch_fn, params), ) for tag, value in metrics_.items(): self.record_scalar('eval', tag=tag, value=value, step=step) logging.info('evaluating EMA model on test dataset') metrics_ = train_utils.streaming_mean( rng_stream, map(to_global_array, get_eval_data()), functools.partial(eval_batch_fn, ema_params), ) for tag, value in metrics_.items(): self.record_scalar('eval_ema', tag=tag, value=value, step=step) if __name__ == '__main__': app.run(functools.partial(run_training.main, Experiment))