Internal ensembling in Aurora models

aurora/__init__.py

@@ -14,7 +14,7 @@
     AuroraV1p5Ensemble,
     AuroraWave,
 )
-from aurora.rollout import rollout
+from aurora.rollout import rollout, rollout_ensemble
 from aurora.tracker import Tracker
 
 __all__ = [
@@ -32,5 +32,6 @@
     "Metadata",
     "insolation",
     "rollout",
+    "rollout_ensemble",
     "Tracker",
 ]

aurora/batch.py

@@ -311,6 +311,37 @@ def from_netcdf(cls, path: str | Path) -> "Batch":
         )
 
 
+def tile_batch(batch: Batch, n: int) -> Batch:
+    """Tile `batch` along the batch dimension `n` times.
+
+    Not part of the public `Batch` API. Used only by `aurora.Aurora.forward` and
+    `aurora.rollout.rollout` to run `n` ensemble members as a single fused computation. The
+    tiled batch dimension is an internal implementation detail and must be undone with
+    `split_batch` before any result derived from it is returned to a caller.
+    """
+    return dataclasses.replace(
+        batch,
+        surf_vars={k: v.repeat(n, *([1] * (v.dim() - 1))) for k, v in batch.surf_vars.items()},
+        atmos_vars={k: v.repeat(n, *([1] * (v.dim() - 1))) for k, v in batch.atmos_vars.items()},
+        metadata=dataclasses.replace(batch.metadata, time=batch.metadata.time * n),
+    )
+
+
+def split_batch(batch: Batch, n: int) -> list[Batch]:
+    """Undo `tile_batch`, splitting a tiled batch back into `n` standard-shaped batches."""
+    b = next(iter(batch.surf_vars.values())).shape[0] // n
+    time = batch.metadata.time
+    return [
+        dataclasses.replace(
+            batch,
+            surf_vars={k: v[m * b : (m + 1) * b] for k, v in batch.surf_vars.items()},
+            atmos_vars={k: v[m * b : (m + 1) * b] for k, v in batch.atmos_vars.items()},
+            metadata=dataclasses.replace(batch.metadata, time=time[m * b : (m + 1) * b]),
+        )
+        for m in range(n)
+    ]
+
+
 def _np(x: torch.Tensor) -> np.ndarray:
     return x.detach().cpu().numpy()
 

aurora/rollout.py

@@ -6,10 +6,10 @@
 
 import torch
 
-from aurora.batch import Batch
+from aurora.batch import Batch, split_batch, tile_batch
 from aurora.model.aurora import Aurora
 
-__all__ = ["rollout"]
+__all__ = ["rollout", "rollout_ensemble"]
 
 
 def _make_lead_time_tensor(batch: Batch, lead_time_hours: float) -> torch.Tensor:
@@ -148,3 +148,43 @@ def rollout(
     # Disable noise accumulation after roll-out is complete, in case the model will be used for
     # normal inference or training afterwards.
     model.set_noise_accumulation(n=0)
+
+
+def rollout_ensemble(
+    model: Aurora,
+    batch: Batch,
+    steps: int,
+    num_ensemble_members: int,
+    fine_lead_times: Optional[Sequence[float]] = None,
+    use_noise_accumulation: bool = True,
+    apply_rollout_input_clipping: bool = True,
+) -> Generator[list[Batch], None, None]:
+    """Perform a roll-out for `num_ensemble_members` ensemble members simultaneously.
+
+    The members are computed in one pass by repeating `batch` along the batch dimension, so every
+    member receives independent noise. All other arguments are as for :func:`rollout`.
+
+    Args:
+        num_ensemble_members (int): Number of ensemble members.
+
+    Yields:
+        list[:class:`aurora.Batch`]: After every (sub-)step, one prediction per ensemble member,
+            each with the batch size of `batch`.
+    """
+    if num_ensemble_members < 1:
+        raise ValueError(
+            f"`num_ensemble_members` must be at least `1`, but is `{num_ensemble_members}`."
+        )
+    if not model.backbone.stochastic:
+        raise ValueError("`rollout_ensemble` requires a stochastic model.")
+
+    batch = tile_batch(batch, num_ensemble_members)
+    for pred in rollout(
+        model,
+        batch,
+        steps,
+        fine_lead_times=fine_lead_times,
+        use_noise_accumulation=use_noise_accumulation,
+        apply_rollout_input_clipping=apply_rollout_input_clipping,
+    ):
+        yield split_batch(pred, num_ensemble_members)

docs/models.md

@@ -449,3 +449,26 @@ When using `rollout` with `fine_lead_times`, noise accumulation is enabled by de
 smoother intra-step transitions while using independent effective noise between main steps,
 matching the training regimen. Set `use_noise_accumulation=False` to draw independent
 noise at each sub-step instead, though this is not recommended.
+
+### Internal Ensembling
+
+You can avoid running the model multiple times to generate different ensemble outputs.
+In Aurora 1.5 you can run the model only once, but the forward pass must operate on
+batch in which the data is tiled.
+This is entirely optional and consumes more memory, but may result in faster inference
+on some hardware.
+Functions `aurora.batch.tile_batch` and `aurora.batch.split_batch` expose this logic:
+the former modifies a `aurora.batch.Batch` to support a single forward pass, and the
+latter unpacks a tiled batch into a list of batches.
+
+In the rollout scenario, you can avoid using these functions and just use `rollout_ensemble`:
+
+```python
+from aurora import rollout_ensemble
+
+with torch.inference_mode():
+    preds = [
+        [member.to("cpu") for member in members]
+        for members in rollout_ensemble(model, batch, steps=10, num_ensemble_members=4)
+    ]
+```

tests/v1p5/test_ensemble.py

@@ -0,0 +1,163 @@
+"""Copyright (c) Microsoft Corporation. Licensed under the MIT license.
+
+Tests for internal ensemble members (`tile_batch` and `split_batch` functions
+and the `rollout_ensemble` utility).
+"""
+
+import warnings
+from datetime import datetime
+
+import pytest
+import torch
+
+from ._helpers import _OUTPUT_ONLY_SURF, _SURF_VARS, _make_batch, _make_small_v1p5
+from aurora import Aurora, Batch, Metadata, rollout_ensemble
+from aurora.batch import split_batch, tile_batch
+from aurora.model.film import AdaptiveLayerNorm
+
+
+def _unzero_adaptive_layer_norms(model: Aurora, std: float = 0.1) -> None:
+    """Nudge every `AdaptiveLayerNorm`'s modulation away from its zero initialisation.
+
+    At construction, `AdaptiveLayerNorm.ln_modulation` is exactly zero-initialised (the
+    `adaLN-Zero` trick), which makes a freshly-built, untrained model exactly insensitive to its
+    conditioning signal `c` -- which is what carries the ensemble noise. Without this, no output
+    difference a test observes between ensemble members can be attributed to noise, since noise
+    provably has zero effect on such a model.
+    """
+    for m in model.modules():
+        if isinstance(m, AdaptiveLayerNorm):
+            with torch.no_grad():
+                m.ln_modulation[-1].weight.normal_(std=std)
+                m.ln_modulation[-1].bias.normal_(std=std)
+
+
+def _make_ensemble_test_batch(b: int = 2) -> Batch:
+    """A small batch with a configurable batch size `b`, used to test tiling/splitting."""
+    h, w = 8, 8
+    return Batch(
+        surf_vars={"2t": torch.randn(b, 2, h, w)},
+        static_vars={"lsm": torch.randn(h, w)},
+        atmos_vars={"z": torch.randn(b, 2, 2, h, w)},
+        metadata=Metadata(
+            lat=torch.linspace(90, -90, h),
+            lon=torch.linspace(0, 360, w + 1)[:-1],
+            time=tuple(datetime(2023, 6, 15, i, 0) for i in range(b)),
+            atmos_levels=(500, 850),
+        ),
+    )
+
+
+def test_tile_and_split_batch_roundtrip():
+    b, n = 2, 3
+    batch = _make_ensemble_test_batch(b)
+
+    tiled = tile_batch(batch, n)
+
+    v = tiled.surf_vars["2t"]
+    assert v.shape[0] == n * b
+    for m in range(n):
+        torch.testing.assert_close(v[m * b : (m + 1) * b], batch.surf_vars["2t"])
+
+    v = tiled.atmos_vars["z"]
+    assert v.shape[0] == n * b
+    for m in range(n):
+        torch.testing.assert_close(v[m * b : (m + 1) * b], batch.atmos_vars["z"])
+
+    assert len(tiled.metadata.time) == n * b
+    for m in range(n):
+        assert tiled.metadata.time[m * b : (m + 1) * b] == batch.metadata.time
+
+    # Static variables have no batch dimension and are untouched.
+    torch.testing.assert_close(tiled.static_vars["lsm"], batch.static_vars["lsm"])
+
+    # Splitting undoes the tiling: every member is identical to the original, standard-shaped
+    # batch (tiling itself introduces no randomness).
+    members = split_batch(tiled, n)
+    assert len(members) == n
+    for member in members:
+        torch.testing.assert_close(member.surf_vars["2t"], batch.surf_vars["2t"])
+        torch.testing.assert_close(member.atmos_vars["z"], batch.atmos_vars["z"])
+        assert member.metadata.time == batch.metadata.time
+
+
+def test_forward_ensemble_members_differ_when_stochastic():
+    n = 3
+    torch.manual_seed(0)
+    model = _make_small_v1p5(stochastic=True)
+    # Un-zero the modulation so noise has a real, appreciable effect (see helper docstring);
+    # otherwise this test cannot distinguish genuine noise sensitivity from incidental
+    # floating-point batching noise
+    # (see `test_forward_ensemble_members_identical_without_stochastic`).
+    _unzero_adaptive_layer_norms(model)
+    model.eval()
+    surf_vars = tuple(v for v in _SURF_VARS if v not in _OUTPUT_ONLY_SURF)
+    batch = _make_batch(surf_vars=surf_vars)
+    batch = tile_batch(batch, n)
+    b = next(iter(batch.surf_vars.values())).shape[0]
+
+    with torch.inference_mode():
+        pred = model.forward(batch, lead_times=torch.full((b,), 6.0))
+    members = split_batch(pred, n)
+
+    # Threshold well above the ~1e-3 floating-point batching floor established in
+    # `test_forward_ensemble_members_identical_without_stochastic`, so a pass here can only be
+    # explained by the injected noise actually differing per member, not incidental rounding.
+    for i in range(n):
+        for j in range(i + 1, n):
+            diff = (members[i].surf_vars["2t"] - members[j].surf_vars["2t"]).abs().max()
+            assert diff > 1e-2
+
+
+def test_forward_ensemble_members_identical_without_stochastic():
+    n = 3
+    with warnings.catch_warnings():
+        warnings.simplefilter("ignore")
+        model = _make_small_v1p5(stochastic=False)
+    model.eval()
+    surf_vars = tuple(v for v in _SURF_VARS if v not in _OUTPUT_ONLY_SURF)
+    batch = _make_batch(surf_vars=surf_vars)
+    batch = tile_batch(batch, n)
+    b = next(iter(batch.surf_vars.values())).shape[0]
+
+    with torch.inference_mode():
+        pred = model.forward(batch, lead_times=torch.full((b,), 6.0))
+    members = split_batch(pred, n)
+
+    for m in range(1, n):
+        # Loose tolerance: floating-point ops (e.g. batched matmul/softmax reductions) are not
+        # strictly invariant to how many other (tiled) rows share the batch, so bitwise equality
+        # isn't guaranteed even though the members are mathematically identical computations.
+        torch.testing.assert_close(
+            members[0].surf_vars["2t"], members[m].surf_vars["2t"], atol=1e-3, rtol=1e-3
+        )
+
+
+def test_rollout_ensemble_yields_list_of_standard_shaped_batches_across_steps():
+    n = 2
+    model = _make_small_v1p5(stochastic=True)
+    model.eval()
+    surf_vars = tuple(v for v in _SURF_VARS if v not in _OUTPUT_ONLY_SURF)
+    batch = _make_batch(surf_vars=surf_vars)
+    b = next(iter(batch.surf_vars.values())).shape[0]
+
+    with torch.inference_mode():
+        preds = list(rollout_ensemble(model, batch, steps=3, num_ensemble_members=n))
+
+    assert len(preds) == 3
+    for step_pred in preds:
+        assert len(step_pred) == n
+        for member in step_pred:
+            for v in member.surf_vars.values():
+                assert v.shape[0] == b
+
+
+def test_rollout_ensemble_num_ensemble_members_one_raises():
+    model = _make_small_v1p5()
+    model.eval()
+    surf_vars = tuple(v for v in _SURF_VARS if v not in _OUTPUT_ONLY_SURF)
+    batch = _make_batch(surf_vars=surf_vars)
+    _ = next(iter(batch.surf_vars.values())).shape[0]
+
+    with pytest.raises(ValueError), torch.inference_mode():
+        _ = list(rollout_ensemble(model, batch, steps=2, num_ensemble_members=1))