Rollout alternative implementation.

[grassesi]

Chunked rollout reimplementation

This PR contains a reimplementation of chunked rollout inspired from #2076 . Below are the differences listed.

Conceptual changes

A chunk is described by the global forecast steps it covers

A chunk is a list of global forecast step ids. It is passed to Model.forward and stored in
ModelOutput. This way ModelOutput becomes self describing, eliminating downstream index
computations. It simplifies the interface of validation_io.write_output and
Model.predict_decoders/Model.predict_latent.

The first chunk includes leading empty forecast_offset steps

With forecast_offset > 0 the first chunk's ModelOutput is sized to start at step 0. The leading
n forecast_offset steps are left empty. This way the concatenation of the data fields in each chunk's
ModelOutput is identical to the fields in a ModelOutput covering a global, unchunked rollout.

Use polymorphism to get latent tokens for rollout.

For the first chunk in a rollout, the latent tokens are constructed from the sources (BatchSamples) by
running the encoder. All following chunks require the tokens and the initial sources, so they retrieved from
the previous chunks output (ModelOutput). Model.forward takes either a BatchSamples instance (first
chunk) or ModelOutput (all other). The interface between both classes is unified using the .latent and
.batch_samples attributes. In case of ModelOutput they just fetch the required data. BatchSamples itself
if queried for .batch_samples. BatchSamples.latent is a empty, signaling to the model to use the encoder.

Behavioral changes

latent_pre_norm is applied only at the true first forecast step

Previous code repopulates LatentState.class_token and patch_tokens at the head of
every chunk, so latent heads and SSL losses see post-norm fields at steps where an unchunked run
would see None.

posteriors is only written once, on the first chunk

Previous code writes posteriors at index 0 of every chunk, storing the previous chunk's LatentState
there from the second chunk onwards

Pushforward marks only the final step of the whole rollout

Previous code skips decoding for every step except rollout_steps - 1, which is the last step of each
chunk. It therefore decodes and backpropagates at every chunk boundary rather than
once at the end of the rollout.

num_steps: 0 still produces output

evenmn's _get_forecast_chunks returns [] for num_steps == 0, so no forward pass runs at all.

milti_stream_data_sampler._get_output_length clamps to at least one output step, so an MTM configuration
(forecast_offset:0, num_steps: 0) has one valid step. In this case, previous code skipps the forward pass
and any output writing.

Spoofed targets no longer depend on non-empty predictions

Previous code reads preds[0].shape to size its zero arrays, which raises IndexError
when the chunk holds no predictions for that stream.

Questions:

How does it interact with changes like latent output (#1860) or infernce mode optimizations (#2272) ...?

Issue Number

Is this PR a draft? Mark it as draft.

Checklist before asking for review

• I have performed a self-review of my code
• My changes comply with basic sanity checks:
  • I have fixed formatting issues with ./scripts/actions.sh lint
  • I have run unit tests with ./scripts/actions.sh unit-test
  • I have documented my code and I have updated the docstrings.
  • I have added unit tests, if relevant
• I have tried my changes with data and code:
  • I have run the integration tests with ./scripts/actions.sh integration-test
  • (bigger changes) I have run a full training and I have written in the comment the run_id(s): launch-slurm.py --time 60
  • (bigger changes and experiments) I have shared a hegdedoc in the github issue with all the configurations and runs for this experiments
• I have informed and aligned with people impacted by my change:
  • for config changes: the MatterMost channels and/or a design doc
  • for changes of dependencies: the MatterMost software development channel