Enable LLC training configs

src/samudra/config.py

@@ -458,7 +458,7 @@ def build(
 
 BlockType = Literal["conv_next_block", "conv_block"]
 ActivationType = Literal["relu", "gelu", "capped_gelu"]
-NormType = Literal["batch", "instance", "layer"]
+NormType = Literal["batch", "instance", "group", "nonorm", "layer"]
 
 
 class BlockConfig(BaseConfig):
@@ -467,6 +467,7 @@ class BlockConfig(BaseConfig):
     activation: ActivationType = "capped_gelu"
     upscale_factor: int = 4
     norm: NormType = "batch"
+    group_norm_groups: int = Field(default=32, ge=1)
     pointwise_linear: bool = False
 
     def build(self) -> CoreBlockBuilder:
@@ -511,6 +512,7 @@ def create_block(
                         kernel_size=self.kernel_size,
                         upscale_factor=self.upscale_factor,
                         norm=self.norm,
+                        group_norm_groups=self.group_norm_groups,
                         activation=activation,
                         pointwise_linear=self.pointwise_linear,
                     )
@@ -1227,6 +1229,7 @@ class TrainConfig(TopLevelConfig):
 
     # Data parameters at root level
     data_stride: list[int] = [1]
+    temporal_stride: int = Field(default=1, ge=1)
     steps: list[int] = [4]
     step_transition: list[int] = []
     inference_epochs: list[int] = [-1]

src/samudra/configs/data/llc.yaml

@@ -1,11 +1,13 @@
-# SPDX-FileCopyrightText: 2026 Ocean Emulator Authors
+# SPDX-FileCopyrightText: 2026 Samudra Authors
 #
 # SPDX-License-Identifier: Apache-2.0
 
 # yaml-language-server: $schema=../schemas/DataConfig.json
 
 sources:
   - type: llc
+    prognostic_vars_key: single_1
+    boundary_vars_key: single_1
     face: 1
     i_start: 0
     i_end: 720
@@ -17,6 +19,9 @@ sources:
     val_time:
       start: "2012-09-01T12:00:00Z"
       end: "2012-11-15T12:00:00Z"
+    inference_times:
+      - start: "2012-10-16T12:00:00Z"
+        end: "2012-11-14T12:00:00Z"
     data_location: LLC.zarr
     data_means_location: LLC_means.zarr
     data_stds_location: LLC_stds.zarr

src/samudra/configs/samudra_llc/eval.yaml

@@ -0,0 +1,20 @@
+# SPDX-FileCopyrightText: 2026 Samudra Authors
+#
+# SPDX-License-Identifier: Apache-2.0
+
+# yaml-language-server: $schema=../schemas/EvalConfig.json
+
+debug: false
+save_zarr: false
+disk_mode: true
+num_model_steps_forward: 25
+
+experiment:
+  name: samudra_llc_eval
+  rand_seed: 15
+  base_output_dir: .LOCAL
+  wandb:
+    mode: disabled
+    project: default
+data: !include ../data/llc.yaml
+model: !include model.yaml

src/samudra/configs/samudra_llc/model.yaml

@@ -0,0 +1,27 @@
+# SPDX-FileCopyrightText: 2026 Samudra Authors
+#
+# SPDX-License-Identifier: Apache-2.0
+
+# yaml-language-server: $schema=../schemas/SamudraConfig.json
+
+checkpointing: all
+pred_residuals: false
+last_kernel_size: 3
+pad: "constant"
+use_bfloat16: false
+
+unet:
+  ch_width: [256, 384, 512, 512]
+  dilation: [1, 2, 4, 8]
+  n_layers: [1, 1, 1, 1]
+
+  core_block:
+    block_type: "conv_next_block"
+    kernel_size: 3
+    activation: "capped_gelu"
+    upscale_factor: 2
+    norm: "group"
+    group_norm_groups: 32
+
+  down_sampling_block: "avg_pool"
+  up_sampling_block: "bilinear_upsample"

src/samudra/configs/samudra_llc/train.yaml

@@ -0,0 +1,38 @@
+# SPDX-FileCopyrightText: 2026 Samudra Authors
+#
+# SPDX-License-Identifier: Apache-2.0
+
+# yaml-language-server: $schema=../schemas/TrainConfig.json
+
+debug: false
+disk_mode: true
+pin_mem: true
+save_freq: 5
+epochs: 70
+batch_size: 1
+learning_rate: 0.0006
+gradient_accumulation_steps: 4
+scheduler: { type: cosine }
+loss:
+  type: dynamic
+  metric: mse
+finetune: false
+resume_ckpt_path: null
+inference_epochs: []
+data_stride: [1]
+temporal_stride: 24
+steps: [1]
+step_transition: []
+preemptible: false
+
+backend: auto
+
+experiment:
+  name: samudra_llc
+  rand_seed: 15
+  base_output_dir: .LOCAL
+  wandb:
+    mode: disabled
+    project: default
+data: !include ../data/llc.yaml
+model: !include model.yaml

src/samudra/datasets.py

@@ -467,13 +467,17 @@ def __init__(
         masked_fill_value: float,
         stride: int = 1,
         concurrent_compute_: bool = False,
+        temporal_stride: int = 1,
     ):
         super().__init__()
         self.id = f"{self.__class__.__name__}_{str(id(self))}"
 
         self.hist: int = hist
         self.steps: int = steps
         self.stride: int = stride
+        if temporal_stride < 1:
+            raise ValueError("temporal_stride must be >= 1")
+        self.temporal_stride: int = temporal_stride
         self.normalize_before_mask: bool = normalize_before_mask
         self.masked_fill_value: float = masked_fill_value
         self._concurrent_compute = concurrent_compute_
@@ -512,11 +516,15 @@ def __init__(
             output_resolution_cpu=self.prognostic_src.resolution,
         )
 
-        self.size: int = (
+        base_size = (
             time_.size
             - self.steps * (self.hist + 1) * self.stride
             - self.hist * self.stride
         )
+        self.size: int = max(
+            0,
+            (base_size + self.temporal_stride - 1) // self.temporal_stride,
+        )
 
     def __len__(self) -> int:
         return self.size
@@ -644,7 +652,7 @@ def _get_x_index(self, idx: int, step: int) -> xr.DataArray:
         if idx >= len(self):
             raise IndexError("Index out of range")
 
-        window_index = idx + step * (self.hist + 1) * self.stride
+        window_index = idx * self.temporal_stride + step * (self.hist + 1) * self.stride
         return self.rolling_indices.isel(window=window_index, drop=True)
 
 

src/samudra/models/modules/blocks.py

@@ -280,6 +280,7 @@ def __init__(
         pad="circular",
         upscale_factor: int = 4,
         norm="batch",
+        group_norm_groups: int = 32,
         checkpoint_simple: bool = False,
         pointwise_linear: bool = False,
     ):
@@ -302,16 +303,13 @@ def __init__(
                 dilation=dilation,
             )
         )
-        # BatchNorm
-        if norm == "batch":
-            convblock.append(torch.nn.BatchNorm2d(in_channels * upscale_factor))
-        # Instance Norm
-        elif norm == "instance":
-            convblock.append(torch.nn.InstanceNorm2d(in_channels * upscale_factor))
-        elif norm == "nonorm":
-            pass
-        else:
-            raise NotImplementedError
+        norm_layer = self._build_norm_layer(
+            norm=norm,
+            channels=int(in_channels * upscale_factor),
+            group_norm_groups=group_norm_groups,
+        )
+        if norm_layer is not None:
+            convblock.append(norm_layer)
         if activation is not None:
             convblock.append(activation())
         convblock.append(
@@ -322,16 +320,13 @@ def __init__(
                 dilation=dilation,
             )
         )
-        # BatchNorm
-        if norm == "batch":
-            convblock.append(torch.nn.BatchNorm2d(in_channels * upscale_factor))
-        # Instance Norm
-        elif norm == "instance":
-            convblock.append(torch.nn.InstanceNorm2d(in_channels * upscale_factor))
-        elif norm == "nonorm":
-            pass
-        else:
-            raise NotImplementedError
+        norm_layer = self._build_norm_layer(
+            norm=norm,
+            channels=int(in_channels * upscale_factor),
+            group_norm_groups=group_norm_groups,
+        )
+        if norm_layer is not None:
+            convblock.append(norm_layer)
         if activation is not None:
             convblock.append(activation())
         # Linear postprocessing
@@ -343,6 +338,29 @@ def __init__(
         self.convblock = torch.nn.Sequential(*convblock)
         self.checkpoint_simple = checkpoint_simple
 
+    @staticmethod
+    def _build_norm_layer(
+        norm: str,
+        channels: int,
+        group_norm_groups: int,
+    ) -> torch.nn.Module | None:
+        if norm == "batch":
+            return torch.nn.BatchNorm2d(channels)
+        if norm == "instance":
+            return torch.nn.InstanceNorm2d(channels)
+        if norm == "group":
+            if group_norm_groups < 1:
+                raise ValueError("group_norm_groups must be >= 1")
+            num_groups = min(group_norm_groups, channels)
+            while channels % num_groups != 0:
+                num_groups -= 1
+            return torch.nn.GroupNorm(num_groups=num_groups, num_channels=channels)
+        if norm == "layer":
+            return torch.nn.GroupNorm(num_groups=1, num_channels=channels)
+        if norm == "nonorm":
+            return None
+        raise NotImplementedError(f"Unsupported normalization mode {norm!r}")
+
     def forward(self, x: torch.Tensor) -> torch.Tensor:
         # return self.skip_module(x) + self.convblock(x)
         skip = self.skip_module(x)

src/samudra/train.py

@@ -323,6 +323,7 @@ def __init__(self, cfg: TrainConfig) -> None:
         self.search_run = cfg.experiment.search
         self.debug = cfg.debug
         self.data_stride: list[int] = cfg.data_stride
+        self.temporal_stride: int = cfg.temporal_stride
         self.batch_size: int = cfg.batch_size
         self.gradient_accumulation_steps: int = cfg.gradient_accumulation_steps
         self.num_workers: int = data_num_workers
@@ -1090,6 +1091,7 @@ def init_data_loaders(self, cur_step: int) -> None:
                 masked_fill_value=self.normalize_fill_value,
                 stride=stride,
                 concurrent_compute_=self.concurrent_compute,
+                temporal_stride=self.temporal_stride,
             )
             for stride in self.data_stride
             for src in self.data_container.train_sources
@@ -1109,6 +1111,7 @@ def init_data_loaders(self, cur_step: int) -> None:
                 masked_fill_value=self.normalize_fill_value,
                 stride=stride,
                 concurrent_compute_=self.concurrent_compute,
+                temporal_stride=self.temporal_stride,
             )
             for stride in self.data_stride
         ]

tests/test_blocks.py

@@ -92,3 +92,70 @@ def test_gradient_flows_through_trunk_when_skip_dropped(self):
         torch.testing.assert_close(trunk.grad, torch.ones_like(trunk))
         # Skip gradient should be all zeros (dropped)
         torch.testing.assert_close(skip.grad, torch.zeros_like(skip))
+
+
+def test_convnext_block_group_norm_uses_divisible_group_count():
+    block = ConvNeXtBlock(
+        in_channels=10,
+        out_channels=10,
+        kernel_size=3,
+        dilation=1,
+        n_layers=1,
+        norm="group",
+        group_norm_groups=6,
+    )
+
+    norm_layers = [
+        layer for layer in block.convblock if isinstance(layer, nn.GroupNorm)
+    ]
+
+    assert len(norm_layers) == 2
+    assert all(layer.num_channels == 40 for layer in norm_layers)
+    assert all(layer.num_groups == 5 for layer in norm_layers)
+
+
+def test_convnext_block_layer_norm_uses_single_group():
+    block = ConvNeXtBlock(
+        in_channels=8,
+        out_channels=8,
+        kernel_size=3,
+        dilation=1,
+        n_layers=1,
+        norm="layer",
+    )
+
+    norm_layers = [
+        layer for layer in block.convblock if isinstance(layer, nn.GroupNorm)
+    ]
+
+    assert len(norm_layers) == 2
+    assert all(layer.num_groups == 1 for layer in norm_layers)
+
+
+def test_convnext_block_nonorm_inserts_no_normalization_layers():
+    block = ConvNeXtBlock(
+        in_channels=8,
+        out_channels=8,
+        kernel_size=3,
+        dilation=1,
+        n_layers=1,
+        norm="nonorm",
+    )
+
+    assert not any(
+        isinstance(layer, (nn.BatchNorm2d, nn.InstanceNorm2d, nn.GroupNorm))
+        for layer in block.convblock
+    )
+
+
+def test_convnext_block_group_norm_rejects_nonpositive_group_count():
+    with pytest.raises(ValueError, match="group_norm_groups must be >= 1"):
+        ConvNeXtBlock(
+            in_channels=8,
+            out_channels=8,
+            kernel_size=3,
+            dilation=1,
+            n_layers=1,
+            norm="group",
+            group_norm_groups=0,
+        )