Rollout update

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):

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)

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,14 +791,17 @@ 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
             target_coords_idxs : Indices of target coordinates
         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

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