validation rollout sliced by amogh-gulati · Pull Request #770 · m2lines/Samudra

amogh-gulati · 2026-06-22T02:49:38Z

This PR adds long-horizon autoregressive rollout validation to the training loop, complementing the existing single-step validation.

Previously, validation only checked one-step-ahead prediction error each epoch. This change allows training to periodically roll the model forward autoregressively over the validation period and log rollout RMSE at configured horizons, such as 90 days and 360 days.

The implementation supports multiple rollout horizons in a single rollout. For example, with rollout_validation_days: [360, 90], the code rolls out once to the maximum horizon, records the 90-day metrics at the intermediate cutoff, and then continues to 360 days. This avoids launching a separate 90-day rollout.

Metrics are logged separately by horizon, variable, depth level, and depth band using raw, un-normalized fields. Rollout validation currently runs only on rank 0, with other ranks waiting at a barrier, since the validation window is not sharded across workers. The rollout is processed in bounded chunks to avoid materializing the full forecast horizon’s targets at once.

This is currently wired up for the standard single-scale training schedule and is a no-op for FOMO’s multi-scale schedule.

Added rollout validation config options:
- rollout_validation_days
- rollout_validation_steps
- rollout_validation_freq
- rollout_validation_steps_forward

Config

In configs/samudra_om4_v2/train.yaml, rollout validation is enabled with:

rollout_validation_days: [360, 90]
rollout_validation_steps_forward: 4

test run here - https://wandb.ai/ocean_emulators/default/runs/e7r70ugc?nw=nwuseramoghgulati
metrics present in rollout_val for 90 days and 360 days

jder · 2026-06-26T18:08:46Z

@codex review

jder · 2026-06-26T18:13:59Z

Thanks @amogh-gulati! One high-level question: why is this separate from the existing inference_one_epoch code? ie should those be the same thing, generalized to cover the new things you're adding here like supporting multiple lengths and removing the strange multi-GPU behavior the inference code currently has?

Also looks like you have some (real) CI failures now that I've finally fixed the spurious ones.

chatgpt-codex-connector

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: a182db0acc

ℹ️ About Codex in GitHub

Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you

Open a pull request for review
Mark a draft as ready
Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".

chatgpt-codex-connector · 2026-06-26T18:14:32Z


-def should_log_validation_images(epoch: int, frequency: int) -> bool:
-    """Return whether to log validation images for a 1-based training epoch."""
+def should_run_on_epoch_freq(epoch: int, frequency: int) -> bool:


Restore should_log_validation_images

The old should_log_validation_images helper was removed by this rename, but tests/test_trainer.py still imports and exercises it. Any run that collects that module now fails with ImportError before the tests execute, so this breaks the existing test suite; keep a wrapper/export or update those callers in the same commit.

Useful? React with 👍 / 👎.

Is this still a valid concern? I agree with codex.

chatgpt-codex-connector · 2026-06-26T18:14:32Z

+            if self.distributed is not None:
+                torch.distributed.barrier()


Avoid blocking nonzero ranks before long rollout

When rollout validation is enabled under NCCL/DDP for the intended long horizons, non-main ranks enter this barrier before rank 0 runs the entire validation rollout, so they can sit inside a distributed collective for the full rollout duration. Since init_distributed_mode does not configure a longer process-group timeout, quarter-degree 90/360-day rollouts that exceed the default timeout can abort otherwise healthy training; configure a timeout that covers this path or avoid starting the collective until the long rank-0 work is done.

Useful? React with 👍 / 👎.

The epoch argument to should_run_rollout_validation was previously ignored, so rollout validation ran every epoch regardless of config.

Uses a lazily created gloo process group with a 12h timeout so the default NCCL watchdog cannot abort healthy training during long rollouts.

The old should_log_validation_images name was still imported by test_trainer.py, failing CI at collection.

alxmrs

Review part 1/N. Will return to this when I can.

alxmrs · 2026-08-04T17:35:09Z

+from typing import Any, NamedTuple

 import dask
+import numpy as np


I'm surprised we didn't have numpy before!

alxmrs · 2026-08-04T17:35:40Z


-def should_log_validation_images(epoch: int, frequency: int) -> bool:
-    """Return whether to log validation images for a 1-based training epoch."""
+def should_run_on_epoch_freq(epoch: int, frequency: int) -> bool:


Is this still a valid concern? I agree with codex.

alxmrs · 2026-08-04T17:41:46Z

+    return requested_steps
+
+
+class RolloutValidationSpec(NamedTuple):


🐑 I prefer using data classes over named tuples. Though, if tuples make this cleaner, that's ok too (I could lack context).

If this is a public contract, it would be nice to add a docstring here to explain what this is for and what it does.

alxmrs · 2026-08-04T17:45:28Z

+    return float(delta)
+
+
+def resolve_rollout_validation_day_spec(


I think these new methods and classes should live in util somewhere to not clutter train.

alxmrs · 2026-08-04T17:46:51Z

+        self.rollout_validation_steps = cfg.rollout_validation_steps
+        self.rollout_validation_days = cfg.rollout_validation_days
+        self.rollout_validation_steps_forward = cfg.rollout_validation_steps_forward
+        self.rollout_validation_freq = cfg.rollout_validation_freq
+        self._rollout_validation_pg: torch.distributed.ProcessGroup | None = None


I think we should pull these up into a single rollout validation config object, not as top lvl train settings.

feel free to pull in related, existing vars from the train config into here.

alxmrs · 2026-08-04T17:48:49Z

+            if self.should_run_rollout_validation(epoch):
+                rollout_val_stats = self.validate_rollout_one_epoch(epoch)
+                end_epoch_rollout_val_time = time.perf_counter()
+            else:
+                rollout_val_stats = {}
+                end_epoch_rollout_val_time = None


What is the train time cost of this operation (approximately)?

alxmrs · 2026-08-04T17:50:30Z

        logger.info(f"Aggregating validation logs")
        return val_aggregator.get_logs(label="val")

+    def should_run_rollout_validation(self, epoch: int) -> bool:


If we have a rollout validation config object, then this could just be a null check (does the object exist?) at the call site -- this function would go away.

To be clear, the config file would have a config object that builds a new data class or ADT. The ADT could live in util and would be built at the top of train.py.

alxmrs · 2026-08-04T17:58:04Z

+                    0, available_steps
+                ).values
+                specs = [
+                    resolve_rollout_validation_day_spec(


Maybe this method could live as a static build method on the RolloutValidationSpec data class?

alxmrs · 2026-08-04T17:59:11Z

+            torch.distributed.barrier(group=group)
+            yield False
+
+    def validate_rollout_one_epoch(self, epoch):


I'm torn if a lot of the contents of this method should live in the stepper (like validate_one_epoch) or not. That would follow convention, but I'm not sure what the right convention is. The way this is done has some advantages.

alxmrs

2/2. Looks pretty good so far! My main note of feedback is on how we package this new process.

alxmrs · 2026-08-04T19:05:59Z

            "a value of 10 logs on epochs 1, 11, 21, ..."
        ),
    )
+    rollout_validation_steps: int = Field(


See comment in train.py about packaging these new options differently.

alxmrs · 2026-08-04T19:07:31Z

    return ValBatchOutput(loss, loss_per_channel, input_data, label, outs, batch.ctx)


+def _get_rollout_step_chunks(


I think this function could be simpler if we used batched:

https://docs.python.org/3/library/itertools.html#itertools.batched

alxmrs · 2026-08-04T19:09:37Z

+        hist: int,
+        area_weights: torch.Tensor,
+        wet: torch.Tensor,
+        num_prognostic_channels: int,
+        normalize: Normalize,
+        tensor_map: TensorMap,
+        distributed_reduce: bool = True,


🐑 I could be wrong, but could we use the DatasetSpec here instead? Many of the arguments we need are there.

Do we need to include TensorMap and Normalize? We're generally trying to move away from including those multitions.

alxmrs · 2026-08-04T19:11:22Z

+        if len(data.target) == 0:
+            raise ValueError("No target values in data")
+
+        _, target_unnorm = get_aggregator_dicts(


🐑 optional: We may be able to create a replacement to this function and the one below to use "modern" normalization code that manages state better so we could avoid adding the multitons.

alxmrs · 2026-08-12T18:48:23Z

Hey @fomo-bot, will you address the concerns that I've raised in my review of Amogh's code? Where there is ambiguity, please let me know your thinking and I'll help resolve it. Feel free to push back on my suggestions where that is appropriate, given your reading of the code in the branch. Thank you!

Amogh Gulati added 2 commits June 21, 2026 22:27

validation rollout sliced

3518bcb

remove file text

8c2c3b3

amogh-gulati mentioned this pull request Jun 22, 2026

Rollout validation #774

Open

formatting and validation

52440ff

amogh-gulati requested a review from jder June 25, 2026 22:01

copyright

a182db0

chatgpt-codex-connector Bot reviewed Jun 26, 2026

View reviewed changes

Amogh Gulati added 2 commits August 3, 2026 11:48

Add rollout_validation_freq config to gate rollout validation by epoch

1b3a02a

The epoch argument to should_run_rollout_validation was previously ignored, so rollout validation ran every epoch regardless of config.

Wait out rank-0 rollout validation on a dedicated gloo barrier

d5ef6a2

Uses a lazily created gloo process group with a 12h timeout so the default NCCL watchdog cannot abort healthy training during long rollouts.

amogh-gulati requested a review from alxmrs August 3, 2026 15:54

Amogh Gulati added 2 commits August 3, 2026 12:28

Merge main into validation_rollout

4213957

Update tests for the should_run_on_epoch_freq rename

bd911df

The old should_log_validation_images name was still imported by test_trainer.py, failing CI at collection.

alxmrs reviewed Aug 4, 2026

View reviewed changes

alxmrs added this to Samudra Planning Aug 4, 2026

github-project-automation Bot moved this to Backlog in Samudra Planning Aug 4, 2026

GHOpenonic added a commit that referenced this pull request Aug 11, 2026

implement autoregressive validation for LLC, pulling machinery from #770

53daeb2

		return requested_steps


		class RolloutValidationSpec(NamedTuple):

		return ValBatchOutput(loss, loss_per_channel, input_data, label, outs, batch.ctx)


		def _get_rollout_step_chunks(

Uh oh!

Conversation

amogh-gulati commented Jun 22, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Config

Uh oh!

jder commented Jun 26, 2026

Uh oh!

jder commented Jun 26, 2026

Uh oh!

chatgpt-codex-connector Bot left a comment

Choose a reason for hiding this comment

💡 Codex Review

Uh oh!

chatgpt-codex-connector Bot Jun 26, 2026

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

chatgpt-codex-connector Bot Jun 26, 2026

Choose a reason for hiding this comment

Uh oh!

alxmrs left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

alxmrs left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

alxmrs commented Aug 12, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

amogh-gulati commented Jun 22, 2026 •

edited

Loading