diff --git a/packages/common/src/weathergen/common/io.py b/packages/common/src/weathergen/common/io.py index 947a4a048..5e512032d 100644 --- a/packages/common/src/weathergen/common/io.py +++ b/packages/common/src/weathergen/common/io.py @@ -586,12 +586,14 @@ class OutputBatchData: source_channels: list[list[str]] geoinfo_channels: list[list[str]] + # latent outputs: outer list over forecast steps, inner list over samples. # each entry is a dict mapping latent_name -> ndarray latents: list[list[dict]] - sample_start: int = 0 - forecast_offset: int = 0 + sample_start: int + forecast_offset: int + forecast_steps: list[int] @functools.cached_property def samples(self): @@ -600,13 +602,6 @@ def samples(self): # TODO associate samples with the sampel idx used for the time window return np.arange(len(self.sources)) + self.sample_start - @functools.cached_property - def forecast_steps(self): - """Indices of all forecast steps adjusted by the forecast offset""" - # forecast offset should be either 1 for forecasting or 0 for MTM - assert self.forecast_offset in (0, 1) - return np.arange(len(self.targets) + self.forecast_offset) - def items(self) -> typing.Generator[OutputItem, None, None]: """Iterate over possible output items""" # TODO: filter for empty items? @@ -668,11 +663,12 @@ def _offset_key(self, key: ItemKey): To be useable in extraction these have to be adjusted to bridge the differences compared to the semantics of the data. - `sample` is adjusted from a global continous index to a per batch index - - `forecast_step` is adjusted from including `forecast_offset` to indexing - the data (always starts at 0) + - `forecast_step` is adjusted from a global step to an index into this chunk's data """ return ItemKey( - key.sample - self.sample_start, key.forecast_step - self.forecast_offset, key.stream + key.sample - self.sample_start, + key.forecast_step - self.forecast_steps[0], # as in ModelOutput.chunk_idx() + key.stream, ) def _extract_targets_predictions(self, stream_idx, offset_key, key, source_interval): diff --git a/src/weathergen/datasets/batch.py b/src/weathergen/datasets/batch.py index ea6a0b26a..bec9e0913 100644 --- a/src/weathergen/datasets/batch.py +++ b/src/weathergen/datasets/batch.py @@ -7,6 +7,7 @@ """ import copy +import typing from dataclasses import dataclass import numpy as np @@ -176,6 +177,11 @@ def __init__( self.output_steps = output_steps self.output_idxs = output_idxs self.device = None + self.latent = [] + + @property + def batch_samples(self) -> typing.Self: + return self def __len__(self) -> int: return len(self.samples) diff --git a/src/weathergen/model/model.py b/src/weathergen/model/model.py index f4035ea46..57bfabd39 100644 --- a/src/weathergen/model/model.py +++ b/src/weathergen/model/model.py @@ -22,7 +22,7 @@ from torch.utils.checkpoint import checkpoint from weathergen.common.config import Config -from weathergen.datasets.batch import ModelBatch +from weathergen.datasets.batch import BatchSamples, ModelBatch from weathergen.datasets.utils import healpix_verts_rots, r3tos2 from weathergen.model.encoder import EncoderModule from weathergen.model.engines import ( @@ -52,12 +52,29 @@ class ModelOutput: Representation of model output """ - physical: list[dict[StreamName, torch.Tensor]] - latent: list[dict[str, torch.Tensor | LatentState]] + def __init__( + self, + forecast_steps: list[int], + forecast_offset: int, + source_samples: BatchSamples, + ) -> None: + self.forecast_offset = forecast_offset + # the first chunk keeps its leading forecast_offset steps as empty slots, so that + # concatenating the chunks of a rollout stays indexed by global forecast step + base = 0 if forecast_steps[0] == forecast_offset else forecast_steps[0] + self.forecast_steps = list(range(base, forecast_steps[-1] + 1)) + + self.physical: list[dict[StreamName, torch.Tensor]] = [{} for _ in self.forecast_steps] + self.latent: list[dict[str, torch.Tensor | LatentState]] = [{} for _ in self.forecast_steps] + self.batch_samples = source_samples + + def chunk_idx(self, fstep: int) -> int: + """Index of forecast step fstep into chunk-local data, e.g. predictions.""" + return fstep - self.forecast_steps[0] - def __init__(self, len_output: int) -> None: - self.physical = [{} for _ in range(len_output)] - self.latent = [{} for _ in range(len_output)] + def batch_idx(self, fstep: int) -> int: + """Index of forecast step fstep into batch-global data, e.g. target coordinates.""" + return fstep def add_physical_prediction( self, fstep: int, stream_name: StreamName, pred: torch.Tensor @@ -669,65 +686,92 @@ def tokens_to_latent_state(self, tokens_post_norm, tokens) -> LatentState: z_pre_norm=tokens, ) - def forward(self, model_params: ModelParams, batch: ModelBatch) -> ModelOutput: + def forward( + self, + model_params: ModelParams, + input: BatchSamples | ModelOutput, + forecast_steps: list[int], + ) -> ModelOutput: """Forward pass of the model Tokens are processed through the model components, which were defined in the create method. Args: model_params : Query and embedding parameters - batch + input : the batch's source samples, or the previous chunk's output + forecast_steps : global forecast steps of the chunk to roll out Returns: A list containing all prediction results """ + source_samples, tokens, posteriors = self._get_initial_conditions(input, model_params) - output = ModelOutput(batch.get_output_len()) + # output_idxs start with output_offset + global_steps = source_samples.get_output_idxs() + forecast_offset = global_steps[0] + final_step = global_steps[-1] - tokens, posteriors = self.encoder(model_params, batch) - output.add_latent_prediction(0, "posteriors", posteriors) - - # recover batch dimension and separate input_steps - shape = (len(batch), batch.get_num_source_steps(), *tokens.shape[1:]) - # collapse along input step dimension - tokens = tokens.reshape(shape).sum(axis=1) + output = ModelOutput(forecast_steps, forecast_offset, source_samples) + # posteriors come from encoding the source window, so they exist only on the first chunk + if posteriors is not None: + output.add_latent_prediction(0, "posteriors", posteriors) # Allow for pushforward trick p_fwd = self.cf.training_config.get("forecast", {}).get("pushforward", False) # roll-out in latent space, iterate and generate output over requested output steps - for step in batch.get_output_idxs(): - without_grad = p_fwd and self.training and step != max(batch.get_output_idxs()) + for step in forecast_steps: + without_grad = p_fwd and self.training and step != final_step if without_grad: - # Pushforward mode: advance tokens without grad; no decoding with torch.no_grad(): - tokens = self.forecast_engine(tokens, step, model_params.rope_coords) + # Pushforward mode: advance tokens without grad; no decoding + with torch.no_grad(): + tokens = self.forecast_engine(tokens, step, model_params.rope_coords) continue tokens = self.forecast_engine(tokens, step, model_params.rope_coords) # decoder predictions - output = self.predict_decoders(model_params, step, tokens, batch, output) + output = self.predict_decoders(model_params, step, tokens, source_samples, output) # latent predictions (raw and with SSL heads) - output = self.predict_latent(model_params, step, tokens, batch, output) + output = self.predict_latent(model_params, step, tokens, source_samples, output) return output + def _get_initial_conditions(self, input: BatchSamples | ModelOutput, model_params: ModelParams): + """Source samples and latent tokens to start a chunk of the rollout from.""" + source_samples, latent = input.batch_samples, input.latent + + if len(latent) == 0: + tokens, posteriors = self.encoder(model_params, source_samples) + # recover batch dimension and separate input_steps + shape = (len(source_samples), source_samples.get_num_steps(), *tokens.shape[1:]) + # collapse along input step dimension + tokens = tokens.reshape(shape).sum(axis=1) + else: + tokens, posteriors = latent[-1]["latent_state"].z_pre_norm, None + + return source_samples, tokens, posteriors + def predict_latent( self, model_params: ModelParams, step: int, tokens: torch.Tensor, - batch: ModelBatch, + batch: BatchSamples, output: ModelOutput, ) -> ModelOutput: """ Compute latent predictions + + step is the global forecast step, output converts it to the spaces it needs. """ + chunk_idx = output.chunk_idx(step) + batch_idx = output.batch_idx(step) # safe latent prediction - tokens_post_norm = self.latent_pre_norm(tokens) if step == 0 else None + tokens_post_norm = self.latent_pre_norm(tokens) if batch_idx == 0 else None latent_state = self.tokens_to_latent_state(tokens_post_norm, tokens) - output.add_latent_prediction(step, "latent_state", latent_state) + output.add_latent_prediction(chunk_idx, "latent_state", latent_state) # latent predictions for SSL training for name, head in self.latent_heads.items(): - output.add_latent_prediction(step, name, head(latent_state)) + output.add_latent_prediction(chunk_idx, name, head(latent_state)) return output @@ -736,7 +780,7 @@ def predict_decoders( model_params: ModelParams, step: int, tokens: torch.Tensor, - batch: ModelBatch, + batch: BatchSamples, output: ModelOutput, ) -> ModelOutput: """ @@ -747,7 +791,7 @@ def predict_decoders( Args: model_params : Query and embedding parameters - fstep : Number of forecast steps + step : Global forecast step, output converts it to the spaces it needs tokens : Tokens from global assimilation engine streams_data : Used to initialize target coordinates tokens and index information List of StreamData len(streams_data) == batch_size_per_gpu @@ -755,6 +799,9 @@ def predict_decoders( Returns: Prediction output tokens in physical representation for each target_coords. """ + chunk_idx = output.chunk_idx(step) + batch_idx = output.batch_idx(step) + # Empty dicts evaluate to False in python if not self.pred_heads: return output @@ -777,7 +824,7 @@ def predict_decoders( for stream_name in self.streams.keys(): # extract target coords for current stream and fstep and convert to one tensor t_coords = [ - batch.samples[i_b].streams_data[stream_name].target_coords[step] + batch.samples[i_b].streams_data[stream_name].target_coords[batch_idx] for i_b in range(batch_size) ] t_coords_lens = [len(t) for t in t_coords] @@ -808,7 +855,7 @@ def predict_decoders( # lens for varlen attention tcls = torch.cat( [ - sample.streams_data[stream_name].target_coords_lens[step] + sample.streams_data[stream_name].target_coords_lens[batch_idx] for sample in batch.samples ] ) @@ -834,6 +881,6 @@ def predict_decoders( # recover batch dimension (ragged, so as list) pred = torch.split(pred, t_coords_lens, dim=1) - output.add_physical_prediction(step, stream_name, pred) + output.add_physical_prediction(chunk_idx, stream_name, pred) return output diff --git a/src/weathergen/train/target_and_aux_ssl_teacher.py b/src/weathergen/train/target_and_aux_ssl_teacher.py index edd8e53b6..33a40ed2e 100644 --- a/src/weathergen/train/target_and_aux_ssl_teacher.py +++ b/src/weathergen/train/target_and_aux_ssl_teacher.py @@ -92,7 +92,7 @@ def __init__(self, model, ema_model, batch_size, training_cfg, **kwargs): self.reset() def forward_teacher(self, model_params, batch): - return self.ema_model.forward_eval(model_params, batch) + return self.ema_model.forward_eval(model_params, batch, batch.get_output_idxs()) def reset(self, batch_size=None): self.ema_model.reset() @@ -161,7 +161,7 @@ def forward_teacher(self, model_params, batch): params = ( self.teacher_model_params if self.teacher_model_params is not None else model_params ) - return self.teacher_model(params, batch) + return self.teacher_model(params, batch, batch.get_output_idxs()) def reset(self, batch_size=None): pass diff --git a/src/weathergen/train/trainer.py b/src/weathergen/train/trainer.py index 13c60edc5..d8c45056e 100644 --- a/src/weathergen/train/trainer.py +++ b/src/weathergen/train/trainer.py @@ -26,6 +26,7 @@ from weathergen.common.config import Config from weathergen.datasets.multi_stream_data_sampler import MultiStreamDataSampler from weathergen.model.ema import EMAModel +from weathergen.model.model import ModelOutput from weathergen.model.model_interface import ( init_model_and_shard, ) @@ -190,6 +191,86 @@ def get_target_aux_calculators(self, mode_cfg): return target_and_aux_calculators + def _get_forecast_step_chunks(self, output_idxs: list[int], chunk_size: int) -> list[list[int]]: + """Split the forecast steps into contiguous chunks of at most chunk_size steps.""" + assert chunk_size >= 1, f"forecast.chunk_size must be >= 1, got {chunk_size}." + return [ + output_idxs[start : start + chunk_size] + for start in range(0, len(output_idxs), chunk_size) + ] + + def _process_validation_chunks( + self, + batch, + mode_cfg, + batch_size, + mini_epoch, + bidx, + targets_and_auxs, + ) -> ModelOutput: + """Run the rollout in chunks and assemble the predictions for the whole batch.""" + forecast_cfg = mode_cfg.get("forecast", {}) + + output_idxs = batch.get_output_idxs() + chunk_size = forecast_cfg.get("chunk_size", len(output_idxs)) + chunks = self._get_forecast_step_chunks(output_idxs, chunk_size) + + num_samples_write = mode_cfg.get("output", {}).get("num_samples", 0) * batch_size + should_write_output = bidx < num_samples_write + if should_write_output: + denormalize_data_fct = ( + (lambda x0, x1: x1) + if mode_cfg.get("output", {}).get("normalized_samples", False) + else self.dataset_val.denormalize_target_channels + ) + if not targets_and_auxs: + raise ValueError( + "Writing validation output requires targets. " + "Configure validation losses or set output.num_samples=0." + ) + + physical, latent = [], [] + forecast_chunk = batch.get_source_samples() + for chunk in chunks: + if self.ema_model is None: + forecast_chunk = self.model( + self.model_params, + forecast_chunk, + chunk, + ) + else: + forecast_chunk = self.ema_model.forward_eval( + self.model_params, + forecast_chunk, + chunk, + ) + + if should_write_output: + write_output( + self.cf, + mode_cfg, + batch_size, + mini_epoch, + bidx, + denormalize_data_fct, + batch, + forecast_chunk, + targets_and_auxs, + ) + + physical += forecast_chunk.physical + latent += forecast_chunk.latent + + # Data for validation purposes => accumulates in memory!? + preds_full = ModelOutput(output_idxs, output_idxs[0], batch.get_source_samples()) + assert len(physical) == len(preds_full.physical), ( + f"Chunks cover {len(physical)} forecast steps, expected {len(preds_full.physical)}." + ) + preds_full.physical = physical + preds_full.latent = latent + + return preds_full + def inference(self, cf, devices, run_id_contd, mini_epoch_contd): # general initalization self.init(cf, devices) @@ -247,6 +328,11 @@ def inference(self, cf, devices, run_id_contd, mini_epoch_contd): self.validate(0, self.test_cfg, self.batch_size_test_per_gpu) logger.info(f"Finished inference run with id: {cf.general.run_id}") + # Without this, NCCL's heartbeat monitor keeps polling a TCPStore whose server has + # already gone away, and the ranks never exit. + if torch.distributed.is_initialized(): + torch.distributed.destroy_process_group() + def run(self, cf, devices, run_id_contd=None, mini_epoch_contd=None): # general initalization self.init(cf, devices) @@ -410,6 +496,11 @@ def run(self, cf, devices, run_id_contd=None, mini_epoch_contd=None): # log final model self.save_model(self.training_cfg.num_mini_epochs) + # Without this, NCCL's heartbeat monitor keeps polling a TCPStore whose server has + # already gone away, and the ranks never exit. + if torch.distributed.is_initialized(): + torch.distributed.destroy_process_group() + def validate_before_training(self): """ Perform validation before training (eg. to check validation pipeline or data normalization) @@ -462,7 +553,8 @@ def train(self, mini_epoch): ): preds = self.model( model_params=self.model_params, - batch=batch.get_source_samples(), + input=batch.get_source_samples(), + forecast_steps=batch.get_output_idxs() ) targets_and_auxs = {} @@ -580,8 +672,6 @@ def validate(self, mini_epoch, mode_cfg, batch_size): dataset_val_iter = iter(self.data_loader_validation) - num_samples_write = mode_cfg.get("output", {}).get("num_samples", 0) * batch_size - with torch.no_grad(): # print progress bar but only in interactive mode, i.e. when without ddp with tqdm.tqdm( @@ -596,17 +686,6 @@ def validate(self, mini_epoch, mode_cfg, batch_size): dtype=self.mixed_precision_dtype, enabled=cf.with_mixed_precision, ): - if self.ema_model is None: - preds = self.model( - self.model_params, - batch.get_source_samples(), - ) - else: - preds = self.ema_model.forward_eval( - self.model_params, - batch.get_source_samples(), - ) - targets_and_auxs = {} for loss_name, target_aux in self.target_and_aux_calculators_val.items(): target_idxs = get_target_idxs_from_cfg(mode_cfg, loss_name) @@ -617,33 +696,21 @@ def validate(self, mini_epoch, mode_cfg, batch_size): self.model, ) - _ = self.loss_calculator_val.compute_loss( - preds=preds, - targets_and_aux=targets_and_auxs, - metadata=extract_batch_metadata(batch), - ) - - # log output - if bidx < num_samples_write: - # denormalization function for data - denormalize_data_fct = ( - (lambda x0, x1: x1) - if mode_cfg.get("output", {}).get("normalized_samples", False) - else self.dataset_val.denormalize_target_channels - ) - # write output - write_output( - self.cf, + preds = self._process_validation_chunks( + batch, mode_cfg, batch_size, mini_epoch, bidx, - denormalize_data_fct, - batch, - preds, targets_and_auxs, ) + _ = self.loss_calculator_val.compute_loss( + preds=preds, + targets_and_aux=targets_and_auxs, + metadata=extract_batch_metadata(batch), + ) + pbar.update(batch_size) if (bidx * batch_size) > mode_cfg.samples_per_mini_epoch: diff --git a/src/weathergen/utils/validation_io.py b/src/weathergen/utils/validation_io.py index 9216ef635..462181880 100644 --- a/src/weathergen/utils/validation_io.py +++ b/src/weathergen/utils/validation_io.py @@ -23,6 +23,16 @@ _logger = logging.getLogger(__name__) +def _empty_step(n_samples: int, n_ens: int, n_channels: int): + """Zero-sized target/prediction entries for a step that carries no data.""" + return ( + [np.zeros((n_ens, 0, n_channels), dtype=np.float32) for _ in range(n_samples)], + [np.zeros((0, n_channels), dtype=np.float32) for _ in range(n_samples)], + [np.zeros((0, 2), dtype=np.float32) for _ in range(n_samples)], + [np.array([]).astype("datetime64[ns]") for _ in range(n_samples)], + ) + + def write_output( cf, val_cfg, @@ -51,11 +61,17 @@ def write_output( fp32 = torch.float32 preds_all, targets_all, targets_coords_all, targets_times_all = [], [], [], [] - timestep_idxs = [0] if len(batch.get_output_idxs()) == 0 else batch.get_output_idxs() - forecast_offset = timestep_idxs[0] + # _get_output_length clamps to at least one output step, so this always holds + assert len(batch.get_output_idxs()) > 0, "Batch carries no output steps." + forecast_offset = batch.get_output_idxs()[0] + + # the chunk describes which forecast steps it holds, including the leading empty steps + # that the first chunk keeps so it is indexed by global forecast step + timestep_idxs = model_output.forecast_steps + + n_samples = len(batch.get_source_samples().get_samples()) targets_lens = [] - # TODO Maybe stopping at forecast_steps explained #1657 for t_idx in timestep_idxs: preds_all += [[]] targets_all += [[]] @@ -63,18 +79,29 @@ def write_output( targets_times_all += [[]] targets_lens += [[]] for sname in cf.streams.keys(): + chunk_idx = model_output.chunk_idx(t_idx) + assert model_output.forecast_steps[chunk_idx] == t_idx, ( + f"Prediction at index {chunk_idx} is valid for forecast step " + f"{model_output.forecast_steps[chunk_idx]}, but the target is valid for {t_idx}." + ) + + n_channels = len(cf.streams[sname].val_target_channels) + + # leading empty steps of the first chunk carry a source but no target/prediction + if t_idx < forecast_offset: + preds_s, targets_s, t_coords_s, t_times_s = _empty_step(n_samples, 1, n_channels) + # handle spoof data: do not write since it might corrupt validation (spoofing invisible # there) - if target_aux_out.physical[t_idx][sname]["is_spoof"][0]: - targets = target_aux_out.physical[t_idx][sname]["target"] - # for-loop to make sure we have a consistent number of samples - preds_s = [np.zeros((1, 0, t.shape[1])) for t in targets] - targets_s = [np.zeros((0, t.shape[1])) for t in targets] - t_coords_s = [np.zeros((0, 2)) for t in targets] - t_times_s = [np.array([]).astype("datetime64[ns]") for t in targets] + elif target_aux_out.physical[t_idx][sname]["is_spoof"][0]: + preds = model_output.get_physical_prediction(chunk_idx, sname) + n_ens = preds[0].shape[0] if preds is not None and len(preds) > 0 else 1 + preds_s, targets_s, t_coords_s, t_times_s = _empty_step( + n_samples, n_ens, n_channels + ) else: - preds = model_output.get_physical_prediction(t_idx, sname) + preds = model_output.get_physical_prediction(chunk_idx, sname) targets = target_aux_out.physical[t_idx][sname]["target"] preds_s, targets_s, t_coords_s, t_times_s = [], [], [], [] @@ -185,6 +212,7 @@ def write_output( latents=latents_all, sample_start=sample_start, forecast_offset=forecast_offset, + forecast_steps=timestep_idxs, ) store_path = config.get_path_results(cf, mini_epoch, batch_idx)