Remove data "schedules"

AGENTS.md

@@ -16,8 +16,8 @@ model implements a ConvNeXt U-Net neural network architecture. We have made sign
 support quarter degree (0.25 x 0.25 lat/lng) data emulation. The samudra-multi model has an encoder, processor, and
 decoder structure and aims to emulate ocean physics by first translating data from a physical space to a latent space.
 The samudra-multi model supports training on multiple scales of data all at once (e.g. one
-degree, half degree and quarter degree), either on a "mix" or "match" schedule (i.e. the cross product of each scale for
-input and label, or one input/label scale at a time per batch).
+degree, half degree and quarter degree) by configuring multiple data sources. Each configured source is trained against
+labels from the same source/resolution.
 
 ## Data
 
@@ -307,4 +307,4 @@ notebooks/                # Analysis and preprocessing notebooks
 5. **Cloud Training**: Supports SkyPilot for remote job execution on AWS & Lambda Labs
 6. **Noisy Failure**: Do not swallow errors. If something goes wrong, let it fail loudly.
 7. **Avoid Hacks**: Don't accommodate bad designs by adding more cruft -- refactor separately first then make the nice change.
-8. **Multi-Scale Support**: samudra-multi supports training on multiple data resolutions simultaneously with "mix" or "match" scheduling
+8. **Multi-Scale Support**: samudra-multi supports training on multiple data resolutions simultaneously by configuring multiple matched data sources

configs/samudra_multi_om4/train_multiscale.yaml

@@ -48,7 +48,6 @@ experiment:
   wandb:
     mode: disabled
     project: default
-  train_schedule: match
 data:
   dataset:
     type: om4

src/samudra/config.py

@@ -780,7 +780,7 @@ def build(
         corrector = None
         if len(srcs) != 1:
             raise ValueError(
-                'Samudra only supports training at a single scale! Please set `training_schedule="standard"`.'
+                "Samudra only supports training at a single scale! Please configure exactly one data source."
             )
         src = srcs[0]
         if self.corrector is not None:
@@ -997,18 +997,13 @@ class DistributedConfig(BaseConfig):
     dist_backend: str | None = None
 
 
-TrainSchedule = Literal["standard", "match", "mix"]
-
-
 class ExperimentConfig(BaseConfig):
     name: str = "cm4_samudra"
     rand_seed: int = 1
     base_output_dir: str = "train"
     # we require this to be set by the user but have optional here
     # so we can leave it out of config files
     data_root: Location | None = None
-    # Define multi-scale dataloader example schedule. Default: single scale.
-    train_schedule: TrainSchedule = "standard"
     wandb: WandBConfig
 
     @cached_property

src/samudra/models/samudra_multi.py

@@ -104,8 +104,6 @@ def forward_once(
             fts = self.encoder(fts, ctx.input_resolution_cpu)
             fts = self.processor(fts)
 
-            # TODO(alxmrs): When the output resolution differs from the input (i.e. in a "mix" schedule), we cannot use
-            #  residual predictions (`self.pred_residuals` must be `False`).
             fts = self.decoder(fts, ctx.output_resolution_cpu)
 
         # Convert back to float32

src/samudra/train.py

@@ -4,18 +4,16 @@
 
 import contextlib
 import datetime
-import itertools
 import logging
 import multiprocessing
 import os
 import tempfile
 import time
 import warnings
 from collections import OrderedDict
-from collections.abc import Iterable
 from multiprocessing.context import BaseContext
 from pathlib import Path
-from typing import Any, assert_never
+from typing import Any
 
 import dask
 import torch
@@ -36,7 +34,7 @@
     get_variable_loss_dict,
 )
 from samudra.backend import init_train_backend
-from samudra.config import TrainConfig, TrainSchedule, build_loss_fn
+from samudra.config import TrainConfig, build_loss_fn
 from samudra.constants import (
     MAX_TRAIN_MODEL_STEPS_FORWARD,
     BoundaryVarNames,
@@ -59,7 +57,7 @@
     train_batch,
     validate_batch,
 )
-from samudra.utils.data import DataSource, Normalize, get_inference_steps
+from samudra.utils.data import Normalize, get_inference_steps
 from samudra.utils.device import using_gpu
 from samudra.utils.distributed import (
     all_reduce_mean,
@@ -150,16 +148,6 @@ def __init__(self, cfg: TrainConfig) -> None:
         self.data_container = cfg.data.build(
             data_root=cfg.experiment.resolved_data_root,
         )
-        self.train_schedule: TrainSchedule = cfg.experiment.train_schedule
-        if self.train_schedule == "mix" and cfg.model.pred_residuals:
-            raise ValueError(
-                "Residual predictions on a mixed multiscale training schedule is not currently supported."
-            )
-        if self.train_schedule == "mix" and any(step > 1 for step in cfg.steps):
-            raise ValueError(
-                "Step predictions on a mixed multiscale training schedule is not currently supported."
-            )
-
         data_num_workers = cfg.data.loading.num_pytorch_workers()
         persistent_workers = cfg.data.loading.persistent_pytorch_workers()
 
@@ -681,8 +669,8 @@ def validate_one_epoch(self, epoch):
             and self.validation_images_enabled
         )
 
-        if self.train_schedule == "standard":
-            # The standard val aggregator only supports a single scale.
+        if len(self.data_container.sources) == 1:
+            # The standard validation aggregator only supports a single scale.
             val_aggregator = Aggregator.get_validation_aggregator(
                 self.primary_src.metadata,
                 self.hist,
@@ -797,23 +785,12 @@ def init_data_loaders(self, cur_step: int) -> None:
         Args:
             cur_step: Current training step size
         """
-        scales = self.data_container.sources
-        match self.train_schedule:
-            case "standard":
-                srcs: Iterable[tuple[DataSource, DataSource | None]] = [
-                    (scales[0], None)
-                ]
-            case "match":
-                srcs = [(s, s) for s in scales]
-            case "mix":
-                srcs = list(itertools.product(scales, repeat=2))  # type: ignore
-            case _:
-                assert_never(self.train_schedule)
+        srcs = self.data_container.sources
 
         train_datasets = [
             TorchTrainDataset(
                 src=src.slice(self.train_time),
-                dst=dst.slice(self.train_time) if dst else None,
+                dst=None,
                 prognostic_var_names=self.prognostic_var_names,
                 boundary_var_names=self.boundary_var_names,
                 hist=self.hist,
@@ -824,13 +801,13 @@ def init_data_loaders(self, cur_step: int) -> None:
                 concurrent_compute_=self.concurrent_compute,
             )
             for stride in self.data_stride
-            for src, dst in srcs
+            for src in srcs
         ]
 
         val_datasets = [
             TorchTrainDataset(
                 src=src.slice(self.val_time),
-                dst=dst.slice(self.val_time) if dst else None,
+                dst=None,
                 prognostic_var_names=self.prognostic_var_names,
                 boundary_var_names=self.boundary_var_names,
                 hist=self.hist,
@@ -841,7 +818,7 @@ def init_data_loaders(self, cur_step: int) -> None:
                 concurrent_compute_=self.concurrent_compute,
             )
             for stride in self.data_stride
-            for src, dst in srcs
+            for src in srcs
         ]
 
         # Create datasets
@@ -872,7 +849,7 @@ def init_data_loaders(self, cur_step: int) -> None:
                 )
 
         # Create batch samplers - branch on distributed vs non-distributed
-        # Group by input AND label resolution to handle all training schedules
+        # Group by resolution so batches stay homogeneous across configured sources.
         def group_key(ds):
             return tuple(prog.grid_size for prog in ds.prognostic_srcs)
 

tests/conftest.py

@@ -16,7 +16,7 @@
 from numpy.typing import ArrayLike, NDArray
 
 import samudra.constants as c
-from samudra.config import JulianDate, TrainBackendConfig, TrainConfig, TrainSchedule
+from samudra.config import JulianDate, TrainBackendConfig, TrainConfig
 from samudra.train import Trainer
 from samudra.utils.data import DataSource, Masks, _is_compact, compact_dataset
 from samudra.utils.multiton import MultitonScope
@@ -314,11 +314,6 @@ def history(request: pytest.FixtureRequest) -> int:
     return request.param
 
 
-@pytest.fixture(scope="session", params=["standard", "match", "mix"])
-def schedule(request: pytest.FixtureRequest) -> TrainSchedule:
-    return request.param
-
-
 # Run a test for both CPU and GPU, and allows selecting or skipping CUDA tests.
 @pytest.fixture(
     params=["cpu", pytest.param("cuda", marks=pytest.mark.cuda)], scope="session"