#!/usr/bin/env python3 """Train official NeuralGCM dynamics with a OneScience ERA5Dataset source.""" from __future__ import annotations import argparse from collections.abc import Mapping import pickle import sys import time from pathlib import Path import numpy as np import optax try: from common import PROJECT_ROOT, add_static_features, as_time_major_frames, era5_data_is_synthetic, era5_frames_to_xarray, load_config, load_era5_dataset, regrid_for_profile, resolve_path, validate_synthetic_era5_version except ModuleNotFoundError: # supports ``python -m scripts.train`` as well from scripts.common import PROJECT_ROOT, add_static_features, as_time_major_frames, era5_data_is_synthetic, era5_frames_to_xarray, load_config, load_era5_dataset, regrid_for_profile, resolve_path, validate_synthetic_era5_version if str(PROJECT_ROOT) not in sys.path: sys.path.insert(0, str(PROJECT_ROOT)) from model.NeuralGCM import build_training_model, format_parameter_summary, make_rollout_functions, parameter_summary, save_official_checkpoint, validate_checkpoint_mode try: from losses import make_loss_fn except ModuleNotFoundError: # supports ``python -m scripts.train`` as well from scripts.losses import make_loss_fn try: from streaming_data import PrefetchedWindowBatches, WindowBatchStream except ModuleNotFoundError: # supports ``python -m scripts.train`` as well from scripts.streaming_data import PrefetchedWindowBatches, WindowBatchStream MODE_ALIASES = {"forecast": "weather_forecast", "weather_forecast": "weather_forecast", "climate": "climate_scale", "climate_scale": "climate_scale", "forecast_2_8_deg": "forecast_2_8_deg", "stochastic_1_4_deg": "stochastic_1_4_deg"} def _merge_training_profile( training: Mapping, mode: str, *, paper_defaults: bool = False ) -> dict: """Merge mode semantics and optionally configured long-run settings.""" merged = dict(training) profiles = merged.pop("profiles", {}) profile = dict(profiles.get(mode, {})) if not paper_defaults: profile = { key: value for key, value in profile.items() if key in {"ensemble_size", "loss"} } for key, value in profile.items(): if isinstance(value, Mapping) and isinstance(merged.get(key), Mapping): merged[key] = {**merged[key], **value} else: merged[key] = value return merged def _make_learning_rate_schedule(optimizer_cfg: Mapping, peak_rate: float): """Build a configured NeuralGCM-style or constant schedule.""" schedule_name = str(optimizer_cfg.get("schedule", "constant")).lower() if schedule_name == "neuralgcm": warmup_steps = int(optimizer_cfg.get("warmup_steps", 2000)) decay_start = int(optimizer_cfg.get("decay_start", 15000)) decay_steps = int(optimizer_cfg.get("decay_steps", 10000)) decay_rate = float(optimizer_cfg.get("decay_rate", 0.5)) if warmup_steps <= 0 or decay_steps <= 0 or decay_start < warmup_steps: raise ValueError("invalid NeuralGCM optimizer schedule boundaries") warmup = optax.linear_schedule(0.0, peak_rate, warmup_steps) plateau = optax.constant_schedule(peak_rate) decay = optax.exponential_decay( peak_rate, transition_steps=decay_steps, decay_rate=decay_rate, staircase=False, ) return optax.join_schedules( (warmup, plateau, decay), (warmup_steps, decay_start) ) if schedule_name != "constant": raise ValueError(f"Unknown training.optimizer.schedule {schedule_name!r}") rates = [float(x) for x in optimizer_cfg.get("rates", [])] boundaries = [int(x) for x in optimizer_cfg.get("boundaries", [])] if rates: if len(rates) != len(boundaries) + 1: raise ValueError("training.optimizer.rates must have one more entry than boundaries") return optax.join_schedules( [optax.constant_schedule(rate) for rate in rates], boundaries ) return optax.constant_schedule(peak_rate) def _era5_frame_capacity(config: Mapping, years: list[int]) -> tuple[int | None, bool]: """Return the shortest yearly trajectory and whether all files are virtual.""" import h5py data_dir = resolve_path(config["data"]["data_dir"]) / "data" paths = [data_dir / f"{year}.h5" for year in years] if not paths or any(not path.exists() for path in paths): return None, False frame_counts = [] synthetic_flags = [] try: for path in paths: with h5py.File(path, "r") as handle: fields = handle[config["data"].get("field_key", "fields")] frame_counts.append(int(fields.shape[0])) synthetic_flags.append(bool(fields.attrs.get("synthetic", False))) except (KeyError, OSError): # The OneScience loader will provide the detailed file-format error. return None, False return min(frame_counts), all(synthetic_flags) def _fit_rollout_schedule_to_data( rollout_schedule: list[dict], *, available_frames: int | None, synthetic: bool, explicit_override: bool, ) -> list[dict]: """Fit long rollouts to virtual data without weakening real-data checks.""" required_frames = max(item["trajectory_length"] for item in rollout_schedule) if available_frames is None or required_frames <= available_frames: return rollout_schedule if explicit_override or not synthetic: source = "explicit --trajectory-length" if explicit_override else "training profile" raise ValueError( f"{source} requires {required_frames} consecutive ERA5 frames, but " f"the shortest training-year file has {available_frames}. Generate a " "longer trajectory or lower --trajectory-length." ) fitted = [ item for item in rollout_schedule if item["trajectory_length"] <= available_frames ] if not fitted and available_frames >= 2: fitted = [{"trajectory_length": available_frames, "until_step": 0}] if not fitted: raise ValueError( "Virtual ERA5 data needs at least two consecutive frames for training; " f"found {available_frames}." ) print( f"data virtual_rollout_clamped={required_frames}->{max(item['trajectory_length'] for item in fitted)} frames" ) return fitted def _trajectory_from_dataset(dataset, steps: int): """Convert a time-indexed xarray sample to official model dictionaries.""" import gin from model.legacy import model_builder # Use the profile's exact conversion hooks after Gin has been parsed. state_fn = gin.query_parameter("WhirlModel.from_xarray_fn") del state_fn converter = model_builder.xarray_to_state_and_dynamic_covariate_data state_data, forcing_data = converter(dataset) return state_data, forcing_data def _replicate_tree(tree, devices): """Replicate a pytree along a leading local-device axis.""" import jax return jax.tree_util.tree_map( lambda value: jax.device_put_replicated(value, devices), tree ) def _unreplicate_tree(tree): """Take replica zero back to host for a normal checkpoint.""" import jax return jax.tree_util.tree_map(lambda value: jax.device_get(value[0]), tree) def _stack_trees(trees): """Stacks trajectory pytrees as host arrays, ready for direct sharding.""" if not trees: raise ValueError("cannot stack an empty pytree sequence") import jax return jax.tree_util.tree_map( lambda *values: np.stack([np.asarray(value) for value in values], axis=0), *trees, ) def _put_batch_sharded(tree, devices, local_batch): """Place each host batch slice directly on its destination device.""" import jax device_count = len(devices) def put(value): value = np.asarray(value) expected = device_count * local_batch if value.shape[0] != expected: raise ValueError( f"batch leaf has leading size {value.shape[0]}, expected {expected}" ) value = value.reshape((device_count, local_batch) + value.shape[1:]) return jax.device_put_sharded( [value[index] for index in range(device_count)], devices ) return jax.tree_util.tree_map(put, tree) def _sample_start_time(sample): """Return the real first timestamp supplied by OneScience ERA5Dataset.""" time_index = sample[4] if not time_index: raise ValueError("ERA5Dataset sample has an empty time_index") value = str(time_index[0]) if len(value) != 10 or not value.isdigit(): raise ValueError(f"invalid ERA5Dataset time index {value!r}") return np.datetime64( f"{value[:4]}-{value[4:6]}-{value[6:8]}T{value[8:10]}:00:00" ) def _read_checkpoint(path_value: str, mode: str) -> tuple[Path, dict]: path = resolve_path(path_value) with path.open("rb") as handle: payload = pickle.load(handle) if not isinstance(payload, dict) or "params" not in payload: raise ValueError(f"Checkpoint {path} does not contain NeuralGCM params") validate_checkpoint_mode(payload, mode, path) return path, payload def _validate_resume_contract(saved: Mapping, current: Mapping) -> None: """Reject changes that would invalidate restored optimizer/data state.""" mismatches = { key: (saved.get(key), value) for key, value in current.items() if saved.get(key) != value } if mismatches: details = ", ".join( f"{key}: saved={old!r}, current={new!r}" for key, (old, new) in mismatches.items() ) raise ValueError(f"Resume checkpoint is incompatible with this run: {details}") def train( config: dict, mode: str, finetune: str | None, max_steps: int | None, learning_rate: float | None, devices_requested: int | None = None, data_workers_requested: int | None = None, prefetch_batches_requested: int | None = None, trajectory_length_requested: int | None = None, checkpoint_output: str | None = None, paper_defaults: bool = False, resume: str | None = None, checkpoint_interval_requested: int | None = None, loss_backend_requested: str | None = None, ): import jax import jax.numpy as jnp train_cfg = _merge_training_profile( config.get("training", {}), mode, paper_defaults=paper_defaults ) available_devices = jax.local_devices() requested_devices = int( devices_requested if devices_requested is not None else train_cfg.get("devices", 1) ) if requested_devices <= 0: raise ValueError("--devices must be a positive integer") if requested_devices > len(available_devices): raise RuntimeError( f"Requested {requested_devices} local devices, but JAX exposes " f"only {len(available_devices)}: {available_devices}" ) devices = available_devices[:requested_devices] configured_batch = max(1, int(train_cfg.get("samples_per_step", 1))) # A pmap replica must receive at least one distinct trajectory. Round the # global batch up to a multiple of the requested device count so no sample # is silently duplicated across replicas. global_batch = max(configured_batch, requested_devices) global_batch = ((global_batch + requested_devices - 1) // requested_devices) * requested_devices years = list(config["data"].get("train_years", [2000])) trajectory_length = max(1, int(train_cfg.get("trajectory_length", 2))) rollout_schedule_cfg = ( [] if trajectory_length_requested is not None else train_cfg.get("rollout_schedule", []) ) if trajectory_length_requested is not None: trajectory_length = int(trajectory_length_requested) if rollout_schedule_cfg: rollout_schedule = sorted( [ { "trajectory_length": max(2, int(item["trajectory_length"])), "until_step": int(item.get("until_step", 0)), } for item in rollout_schedule_cfg ], key=lambda item: item["until_step"], ) if rollout_schedule[0]["until_step"] not in (0, 1): raise ValueError("training.rollout_schedule must start at until_step 0 or 1") else: rollout_schedule = [{"trajectory_length": trajectory_length, "until_step": 0}] available_frames, synthetic_data = _era5_frame_capacity(config, years) if synthetic_data: validate_synthetic_era5_version(config, years) rollout_schedule = _fit_rollout_schedule_to_data( rollout_schedule, available_frames=available_frames, synthetic=synthetic_data, explicit_override=trajectory_length_requested is not None, ) trajectory_length = max(item["trajectory_length"] for item in rollout_schedule) if trajectory_length < 2: raise ValueError( "training.trajectory_length must be at least 2 (one initial and " "one future ERA5 frame)" ) # The official Experiment counts the initialization frame in # ``trajectory_length``. Thus a two-frame trajectory is one input plus one # future ERA5 frame, not two future frames. future_steps = trajectory_length - 1 dataset = load_era5_dataset( config, years, input_steps=1, output_steps=future_steps ) dataset_size = int(getattr(dataset, "total_samples", -1)) if dataset_size < 0: raise ValueError( "OneScience ERA5Dataset computed a negative sample count: " f"T={dataset.T}, input_steps={dataset.input_steps}, " f"output_steps={dataset.output_steps}. The requested trajectory is " "longer than the data file." ) print(f"data samples={dataset_size} shape={(dataset.C, dataset.H, dataset.W)}") if dataset_size < global_batch: raise ValueError( f"Training requires global batch={global_batch} trajectories for " f"{requested_devices} devices, but ERA5Dataset has only {dataset_size} " "samples. Generate more windows or lower training.samples_per_step." ) first_sample = dataset[0] first_input = first_sample[0] first_targets = as_time_major_frames(first_sample[1], name="ERA5 target") first_frames = np.concatenate((first_input[None, ...], first_targets), axis=0) # Build the model from the first sample; subsequent windows are prefetched # and converted batch by batch, never materializing the full dataset. ds = regrid_for_profile( era5_frames_to_xarray( first_frames, config, start_time=_sample_start_time(first_sample) ), mode, ) ds = add_static_features( ds, config, mode=mode, prefer_profile=not synthetic_data ) model, gin_text = build_training_model(ds, mode) # Build temporal target and forcing dictionaries with the profile's official # xarray conversion function (including tracers and sim_time). if model.from_xarray_fn is None: raise RuntimeError("Gin profile did not configure WhirlModel.from_xarray_fn") # The official training pipeline materializes nondimensional ``sim_time`` # before converting xarray data into weatherbench state dictionaries. from dinosaur import xarray_utils reference_datetime = model.specs.aux_features["reference_datetime"] def convert_samples(samples): """Convert already-prefetched OneScience samples on the main thread.""" converted = [] for sample in samples: input_frame = sample[0] target_frames = as_time_major_frames(sample[1], name="ERA5 target") frame_arrays = np.concatenate( (np.asarray(input_frame)[None, ...], target_frames), axis=0 ) sample_ds = regrid_for_profile( era5_frames_to_xarray( frame_arrays, config, start_time=_sample_start_time(sample), ), mode, ) sample_ds = add_static_features( sample_ds, config, mode=mode, prefer_profile=not synthetic_data, ) sample_ds = xarray_utils.ds_with_sim_time( sample_ds, model.specs.physics_specs, reference_datetime=reference_datetime, ) converted.append(model.from_xarray_fn(sample_ds)) return ( _stack_trees([item[0] for item in converted]), _stack_trees([item[1] for item in converted]), ) # Initialize parameters from one valid trajectory. The first training # batch is then obtained from the stream like every later batch. initial_target, initial_forcing = convert_samples([first_sample]) target = jax.tree_util.tree_map(lambda value: value[0], initial_target) forcing_data = jax.tree_util.tree_map(lambda value: value[0], initial_forcing) # ERA5 samples are six-hourly while NeuralGCM integrates at its internal # one-hour (profile-dependent) timestep. Match the official trajectory # contract by repeating internal steps between saved data frames. data_interval = np.timedelta64( int(config["data"].get("time_step_hours", 6)), "h" ) model_timestep = model.specs.physics_specs.dimensionalize_timedelta64( model.specs.dt ) ratio = data_interval / model_timestep inner_steps = int(round(float(ratio))) if inner_steps <= 0 or abs(float(ratio) - inner_steps) > 1e-6: raise ValueError( f"ERA5 interval {data_interval} is not an integer multiple of " f"NeuralGCM timestep {model_timestep}" ) rollout_max = make_rollout_functions( model, trajectory_length=trajectory_length, inner_steps=inner_steps ) if finetune and resume: raise ValueError("--finetune and --resume are mutually exclusive") resume_path = None resume_state = None params = None if resume: resume_path, resume_payload = _read_checkpoint(resume, mode) resume_state = resume_payload.get("training_state") if not isinstance(resume_state, Mapping): raise ValueError( f"--resume requires a project checkpoint with full training_state; " f"{resume_path} is inference-only. Use --finetune to load params only." ) if int(resume_state.get("format_version", -1)) != 1: raise ValueError( f"Unsupported training_state format in {resume_path}: " f"{resume_state.get('format_version')!r}" ) params = resume_state.get("train_params") if params is None: raise ValueError(f"Resume checkpoint {resume_path} has no train_params") if finetune: _, payload = _read_checkpoint(finetune, mode) params = payload["params"] if params is None: params = rollout_max.init(jax.random.key(int(config["project"].get("seed", 0))), target, forcing_data) print(f"model mode={mode} {format_parameter_summary(params)}") effective_lr = float( learning_rate if learning_rate is not None else train_cfg.get("learning_rate", 1e-4) ) clip_norm = float(train_cfg.get("gradient_clip_norm", 1.0)) optimizer_cfg = dict(train_cfg.get("optimizer", {})) schedule = _make_learning_rate_schedule(optimizer_cfg, effective_lr) b1 = float(optimizer_cfg.get("b1", 0.9)) b2 = float(optimizer_cfg.get("b2", 0.95)) eps = float(optimizer_cfg.get("eps", 1e-6)) if clip_norm > 0: optimizer = optax.chain( optax.clip_by_global_norm(clip_norm), optax.adam(schedule, b1=b1, b2=b2, eps=eps), ) else: optimizer = optax.adam(schedule, b1=b1, b2=b2, eps=eps) opt_state = optimizer.init(params) if resume_state is not None: restored_opt_state = resume_state.get("opt_state") if restored_opt_state is None: raise ValueError(f"Resume checkpoint {resume_path} has no opt_state") opt_state = restored_opt_state loss_config = dict(train_cfg.get("loss", {})) if loss_backend_requested is not None: loss_config["backend"] = loss_backend_requested loss_backend = str(loss_config.get("backend", "official")).lower() crps_training = loss_backend == "crps" ensemble_size = int(train_cfg.get("ensemble_size", 2 if crps_training else 1)) if crps_training and ensemble_size != 2: raise ValueError("Official NeuralGCM CRPS training requires ensemble_size=2") if not crps_training and ensemble_size != 1: raise ValueError("Deterministic training requires ensemble_size=1") rollout_cache = {trajectory_length: rollout_max} loss_cache = { trajectory_length: make_loss_fn( model, steps_per_save=inner_steps, trajectory_length=trajectory_length, config=loss_config, mode=mode, ) } def schedule_length(step: int) -> int: selected = rollout_schedule[0]["trajectory_length"] # Public Experiment advances a curriculum leg on ``step > boundary``. for item in rollout_schedule[1:]: if step > item["until_step"]: selected = item["trajectory_length"] return selected def _slice_time(tree, length: int): """Slice only trajectory leaves while retaining static metadata.""" def slice_leaf(value): shape = getattr(value, "shape", ()) if len(shape) >= 2 and shape[1] == trajectory_length: return value[:, :length] if len(shape) and shape[0] == trajectory_length: return value[:length] return value return jax.tree_util.tree_map(slice_leaf, tree) def get_rollout_and_loss(length: int): if length not in rollout_cache: rollout_cache[length] = make_rollout_functions( model, trajectory_length=length, inner_steps=inner_steps ) loss_cache[length] = make_loss_fn( model, steps_per_save=inner_steps, trajectory_length=length, config=loss_config, mode=mode, ) return rollout_cache[length], loss_cache[length] ema_num_steps = int(train_cfg.get("ema_num_steps", 0)) ema_decay = 0.0 if ema_num_steps <= 0 else 1.0 - 2.0 / (ema_num_steps + 1.0) resume_contract = { "mode": mode, "dataset_size": dataset_size, "global_batch": global_batch, "trajectory_length": trajectory_length, "rollout_schedule": rollout_schedule, "inner_steps": inner_steps, "data_interval_hours": int(config["data"].get("time_step_hours", 6)), "learning_rate": effective_lr, "gradient_clip_norm": clip_norm, "optimizer": optimizer_cfg, "ema_num_steps": ema_num_steps, "loss": loss_config, "ensemble_size": ensemble_size, "paper_defaults": paper_defaults, } start_step = 0 restored_ema_params = None if resume_state is not None: saved_contract = resume_state.get("contract") if not isinstance(saved_contract, Mapping): raise ValueError(f"Resume checkpoint {resume_path} has no contract") _validate_resume_contract(saved_contract, resume_contract) start_step = int(resume_state.get("step", -1)) if start_step < 0: raise ValueError(f"Resume checkpoint {resume_path} has invalid step={start_step}") restored_ema_params = resume_state.get("ema_params") if restored_ema_params is None: raise ValueError(f"Resume checkpoint {resume_path} has no ema_params") stream = WindowBatchStream( size=dataset_size, global_batch=global_batch, seed=int(config["project"].get("seed", 0)), shuffle=bool(train_cfg.get("shuffle", True)), drop_last=bool(train_cfg.get("drop_last", True)), ) if resume_state is not None: data_stream_state = resume_state.get("data_stream_state") if not isinstance(data_stream_state, dict): raise ValueError( f"Resume checkpoint {resume_path} has no data_stream_state" ) stream.load_state_dict(data_stream_state) prefetcher = PrefetchedWindowBatches( dataset, stream, num_workers=int( data_workers_requested if data_workers_requested is not None else train_cfg.get("data_num_workers", 2) ), prefetch_batches=int( prefetch_batches_requested if prefetch_batches_requested is not None else train_cfg.get("prefetch_batches", 1) ), ) def loss_fn_for_length( length, p, rngs, xs, fs, *, device_axis_name=None ): rollout_fn, trajectory_loss = get_rollout_and_loss(length) def single_rollout(rng, x, f): pred, truth = rollout_fn.apply(p, rng, x, f) return pred, truth if crps_training: def single_ensemble_loss(member_rngs, x, f): predictions, targets = jax.vmap( single_rollout, in_axes=(0, None, None), spmd_axis_name="ensemble", )(member_rngs, x, f) per_member = jax.vmap( trajectory_loss, axis_name="ensemble", spmd_axis_name="ensemble", )(predictions, targets) return jnp.mean(per_member) per_example = jax.vmap( single_ensemble_loss, in_axes=(0, 0, 0), axis_name="batch", spmd_axis_name="batch", )(rngs, xs, fs) else: predictions, targets = jax.vmap( single_rollout, in_axes=(0, 0, 0) )(rngs, xs, fs) if hasattr(trajectory_loss, "evaluate_batch"): return trajectory_loss.evaluate_batch( predictions, targets, device_axis_name=device_axis_name, ) per_example = jax.vmap(trajectory_loss, in_axes=(0, 0))( predictions, targets ) return jnp.mean(per_example) def step_rngs(step: int, batch_size: int): base_key = jax.random.key(int(config["project"].get("seed", 0))) keys = jax.random.split( jax.random.fold_in(base_key, step), batch_size * ensemble_size ) if crps_training: return keys.reshape((batch_size, ensemble_size) + keys.shape[1:]) return keys if requested_devices == 1: train_params, train_opt_state = params, opt_state train_ema_params = ( restored_ema_params if restored_ema_params is not None else jax.tree_util.tree_map(lambda value: value, params) ) value_grad_cache = {} def train_step(step, current_params, current_opt_state, current_ema_params, batch_target, batch_forcing): length = schedule_length(step) if length not in value_grad_cache: value_grad_cache[length] = jax.jit(jax.value_and_grad( lambda p, r, x, f: loss_fn_for_length( length, p, r, x, f ) )) value_grad = value_grad_cache[length] batch_target = _slice_time(batch_target, length) batch_forcing = _slice_time(batch_forcing, length) rngs = step_rngs(step, global_batch) loss, grads = value_grad( current_params, rngs, batch_target, batch_forcing ) grad_finite = np.asarray( jax.device_get( jnp.asarray( [ jnp.all(jnp.isfinite(g)) for g in jax.tree_util.tree_leaves(grads) ] ) ) ) loss_value = float(np.asarray(jax.device_get(loss))) if not np.isfinite(loss_value) or not np.all(grad_finite): bad_grad_leaves = int(np.size(grad_finite) - np.count_nonzero(grad_finite)) raise FloatingPointError( f"NeuralGCM produced loss={loss_value!r} and " f"nonfinite_gradient_leaves={bad_grad_leaves}; " "reduce learning_rate, increase gradient clipping, or " "use a physically consistent ERA5 trajectory." ) updates, current_opt_state = optimizer.update( grads, current_opt_state, current_params ) current_params = optax.apply_updates(current_params, updates) current_ema_params = jax.tree_util.tree_map( lambda old, new: ema_decay * old + (1.0 - ema_decay) * new, current_ema_params, current_params, ) return current_params, current_opt_state, current_ema_params, loss else: # Synchronous single-host data parallelism. Each replica receives a # distinct slice of the global batch; parameters remain replicated. local_batch = global_batch // requested_devices train_params = _replicate_tree(params, devices) train_opt_state = _replicate_tree(opt_state, devices) train_ema_params = _replicate_tree( restored_ema_params if restored_ema_params is not None else params, devices, ) def make_pmapped_step(length): def pmapped_step(current_params, current_opt_state, current_ema_params, rng, x, f): loss, grads = jax.value_and_grad( lambda p, r, xx, ff: loss_fn_for_length( length, p, r, xx, ff, device_axis_name="devices", ) )(current_params, rng, x, f) grads = jax.lax.pmean(grads, axis_name="devices") loss = jax.lax.pmean(loss, axis_name="devices") grad_finite = jnp.all( jnp.asarray( [jnp.all(jnp.isfinite(g)) for g in jax.tree_util.tree_leaves(grads)] ) ) updates, current_opt_state = optimizer.update( grads, current_opt_state, current_params ) current_params = optax.apply_updates(current_params, updates) current_ema_params = jax.tree_util.tree_map( lambda old, new: ema_decay * old + (1.0 - ema_decay) * new, current_ema_params, current_params, ) return current_params, current_opt_state, current_ema_params, loss, grad_finite return jax.pmap( pmapped_step, axis_name="devices", devices=devices, ) pmapped_cache = {} def get_pmapped_step(length): if length not in pmapped_cache: pmapped_cache[length] = make_pmapped_step(length) return pmapped_cache[length] def train_step(step, current_params, current_opt_state, current_ema_params, batch_target, batch_forcing): length = schedule_length(step) batch_target = _slice_time(batch_target, length) batch_forcing = _slice_time(batch_forcing, length) pmapped = get_pmapped_step(length) keys = step_rngs(step, global_batch) keys = keys.reshape((requested_devices, local_batch) + keys.shape[1:]) sharded_target = _put_batch_sharded(batch_target, devices, local_batch) sharded_forcing = _put_batch_sharded(batch_forcing, devices, local_batch) new_params, new_state, new_ema_params, loss, grad_finite = pmapped( current_params, current_opt_state, current_ema_params, keys, sharded_target, sharded_forcing, ) loss_host = np.asarray(jax.device_get(loss)) finite_host = np.asarray(jax.device_get(grad_finite)) if not np.all(np.isfinite(loss_host)) or not np.all(finite_host): raise FloatingPointError( "NeuralGCM loss became NaN/Inf on one or more devices; " "reduce learning_rate or use a longer/physical trajectory." ) return new_params, new_state, new_ema_params, loss nsteps = int( max_steps if max_steps is not None else train_cfg.get("max_steps", 1) ) if nsteps <= 0: raise ValueError("training.max_steps must be a positive integer") if start_step > nsteps: raise ValueError( f"Resume checkpoint is already at step {start_step}, beyond " f"requested max_steps={nsteps}" ) checkpoint_interval = int( checkpoint_interval_requested if checkpoint_interval_requested is not None else train_cfg.get("checkpoint_interval", 0) ) if checkpoint_interval < 0: raise ValueError("training.checkpoint_interval must be >= 0") output = ( resolve_path(checkpoint_output) if checkpoint_output else resolve_path(config["paths"].get("checkpoint_dir", "data/checkpoint")) / "model_bak.pkl" ) aux_ds = ds[["geopotential_at_surface", "land_sea_mask"]] if "time" in aux_ds.dims: aux_ds = aux_ds.isel(time=0, drop=True) if "level" not in aux_ds.coords: aux_ds = aux_ds.assign_coords( level=np.asarray(model.data_coords.vertical.centers) ) def save_training_checkpoint(completed_steps: int) -> None: raw_params = ( train_params if requested_devices == 1 else _unreplicate_tree(train_params) ) raw_opt_state = ( train_opt_state if requested_devices == 1 else _unreplicate_tree(train_opt_state) ) ema_params = ( train_ema_params if requested_devices == 1 else _unreplicate_tree(train_ema_params) ) raw_params, raw_opt_state, ema_params = jax.device_get( (raw_params, raw_opt_state, ema_params) ) inference_params = ema_params if ema_num_steps > 0 else raw_params checkpoint_parameter_summary = parameter_summary(inference_params) training_state = { "format_version": 1, "step": completed_steps, "train_params": raw_params, "ema_params": ema_params, "opt_state": raw_opt_state, "data_stream_state": prefetcher.resume_state(), "contract": resume_contract, } save_official_checkpoint( output, inference_params, aux_ds, gin_text, metadata={ "mode": mode, "training_steps": completed_steps, "finetune_source": finetune, "resume_source": str(resume_path) if resume_path else None, "ema_num_steps": ema_num_steps, "loss_backend": loss_backend, "ensemble_size": ensemble_size, "parameter_count": checkpoint_parameter_summary["count"], "parameter_bytes": checkpoint_parameter_summary["nbytes"], "paper_defaults": paper_defaults, "training_state": training_state, }, ) print(f"Saved resumable official-format checkpoint at step={completed_steps}: {output}") crps_weight_mode = ( "uniform" if crps_training and loss_config.get("variable_weights") is None else "configured" if crps_training else "n/a" ) print( f"Training devices={requested_devices}/{len(available_devices)}, " f"global_batch={global_batch}, local_batch={global_batch // requested_devices}, " f"inner_steps={inner_steps} (data interval={data_interval}, " f"model timestep={model_timestep}), loss={loss_backend}, " f"loss_normalization={'explicit_weights' if loss_config.get('variable_weights') is not None else 'configured_scales'}, " f"ensemble={ensemble_size}, crps_weights={crps_weight_mode}, " f"lr_peak={effective_lr:g}, configured_long_run={paper_defaults}, " f"start_step={start_step}" ) completed_steps = start_step last_saved_step = None try: for step in range(start_step, nsteps): total_start = time.perf_counter() _, samples = prefetcher.next_batch() batch_target, batch_forcing = convert_samples(samples) train_params, train_opt_state, train_ema_params, loss = train_step( step, train_params, train_opt_state, train_ema_params, batch_target, batch_forcing ) loss_value = float(np.asarray(jax.device_get(loss)).reshape(-1)[0]) elapsed = time.perf_counter() - total_start throughput = global_batch / elapsed current_lr = float(np.asarray(jax.device_get(schedule(step)))) print( f"step={step + 1}/{nsteps} loss={loss_value:.6g} " f"lr={current_lr:.6g} rollout_hours={(schedule_length(step) - 1) * int(config['data'].get('time_step_hours', 6))} " f"elapsed={elapsed:.2f}s throughput={throughput:.4f} samples/s" ) completed_steps = step + 1 if checkpoint_interval and completed_steps % checkpoint_interval == 0: save_training_checkpoint(completed_steps) last_saved_step = completed_steps finally: prefetcher.close() if last_saved_step != completed_steps: save_training_checkpoint(completed_steps) def main(*, forced_mode: str | None = None) -> None: parser = argparse.ArgumentParser() parser.add_argument("--config", default="conf/config.yaml") parser.add_argument("--mode") parser.add_argument( "--finetune", nargs="?", const="official", help="load params only and reset optimizer (no value selects this mode's official checkpoint)", ) parser.add_argument( "--resume", nargs="?", const="default", help="restore params, EMA, optimizer, step and data stream (no value selects the output checkpoint)", ) parser.add_argument("--data-dir") parser.add_argument("--max-steps", type=int) parser.add_argument("--learning-rate", type=float) parser.add_argument( "--paper-defaults", action="store_true", help="use configured long-run settings informed by the public training description", ) parser.add_argument( "--trajectory-length", type=int, help="override the configured rollout curriculum (includes the initial frame)", ) parser.add_argument( "--checkpoint-output", help="explicit output checkpoint path (default: paths.checkpoint_dir/model_bak.pkl)", ) parser.add_argument( "--checkpoint-interval", type=int, help="save resumable state every N completed steps; 0 saves only at exit", ) parser.add_argument( "--loss-backend", choices=("official", "paper", "legacy_official", "scaled", "crps"), help=( "override training.loss.backend (official/paper use the published " "five-term deterministic objective; scaled is for synthetic smoke data)" ), ) parser.add_argument( "--devices", type=int, help="number of local JAX devices for synchronous data parallel training", ) parser.add_argument( "--data-workers", type=int, help="host threads used to prefetch OneScience ERA5Dataset windows", ) parser.add_argument( "--prefetch-batches", type=int, help="number of full host batches queued ahead of the train step", ) parser.add_argument("--validate-only", action="store_true") args = parser.parse_args() config = load_config(args.config) if args.data_dir: config["data"]["data_dir"] = args.data_dir paired_static = resolve_path(args.data_dir, args.config) / "static.nc" if paired_static.exists(): config["data"]["static_file"] = str(paired_static) requested_mode = args.mode or config["training"].get("mode", "weather_forecast") requested_mode = MODE_ALIASES.get(requested_mode, requested_mode) if forced_mode is not None: fixed_mode = MODE_ALIASES.get(forced_mode, forced_mode) if args.mode is not None and requested_mode != fixed_mode: raise ValueError( f"This launcher is fixed to mode={fixed_mode!r}; received " f"conflicting --mode {args.mode!r}. Use scripts/train.py to " "select a mode dynamically." ) mode = fixed_mode else: mode = requested_mode if mode not in config["model"].get("profiles", {}): raise ValueError(f"Unknown NeuralGCM mode {mode!r}") finetune = args.finetune if finetune == "official": finetune = config["model"]["profiles"][mode]["official_reference"] resume = args.resume if resume == "default": resume = args.checkpoint_output or str( resolve_path(config["paths"].get("checkpoint_dir", "data/checkpoint")) / "model_bak.pkl" ) if args.validate_only: validation_years = list(config["data"].get("train_years", [2000])) load_era5_dataset(config, validation_years) if era5_data_is_synthetic(config, validation_years): validate_synthetic_era5_version(config, validation_years) print("OneScience ERA5Dataset validation complete") return train( config, mode, finetune, args.max_steps, args.learning_rate, args.devices, args.data_workers, args.prefetch_batches, args.trajectory_length, args.checkpoint_output, args.paper_defaults, resume, args.checkpoint_interval, args.loss_backend, ) if __name__ == "__main__": main()