Fix load_decoder_chkpt silently loading no decoder under FSDP

src/weathergen/model/model_interface.py

@@ -250,6 +250,34 @@ def init_model_and_shard(
     return model, model_params
 
 
+def _strip_module_prefix(key: str) -> str:
+    """Drop a single leading ``module.`` from a state_dict key, if present.
+
+    Deliberately not ``key.replace("module.", "")``: that strips the substring anywhere in the
+    path, so a genuine submodule named ``module`` would be corrupted.
+    """
+    return key[len("module.") :] if key.startswith("module.") else key
+
+
+def _align_module_prefix(params: dict, model_sd) -> dict:
+    """Match a checkpoint's ``module.`` key convention to the model's.
+
+    A DDP-wrapped save carries a ``module.`` prefix that a bare model's state_dict lacks (and
+    vice versa). Both load paths resolve parameters by exact name, so a mismatch means every
+    lookup misses -- silently, in the sharded path, which skips non-matching names one by one
+    and can end up loading nothing at all.
+    """
+    if not params:
+        return params
+    model_has_prefix = next(iter(model_sd)).split(".")[0] == "module"
+    params_has_prefix = next(iter(params)).split(".")[0] == "module"
+    if model_has_prefix and not params_has_prefix:
+        return {"module." + k: v for k, v in params.items()}
+    if not model_has_prefix and params_has_prefix:
+        return {_strip_module_prefix(k): v for k, v in params.items()}
+    return params
+
+
 def load_model(cf, model, device, run_id: str, with_ddp: bool, with_fsdp: bool, mini_epoch: int):
     """Loads model state from checkpoint and checks for missing and unused keys.
     Args:
@@ -269,22 +297,9 @@ def load_model(cf, model, device, run_id: str, with_ddp: bool, with_fsdp: bool,
 
     is_model_sharded = with_ddp and with_fsdp
     if is_model_sharded:
-        model_has_prefix_module = list(model.state_dict().keys())[0].split(".")[0] == "module"
-        params_has_prefix_module = list(params.keys())[0].split(".")[0] == "module"
-        if model_has_prefix_module and not params_has_prefix_module:
-            # add "module." prefix
-            params_temp = {}
-            for k in params.keys():
-                params_temp["module." + k] = params[k]
-            params = params_temp
-        elif not model_has_prefix_module and params_has_prefix_module:
-            # remove "module." prefix
-            params_temp = {}
-            for k in params.keys():
-                params_temp[k.replace("module.", "")] = params[k]
-            params = params_temp
-
         meta_sharded_sd = model.state_dict()
+        params = _align_module_prefix(params, meta_sharded_sd)
+
         maybe_sharded_sd = {}
         for param_name, full_tensor in params.items():
             sharded_meta_param = meta_sharded_sd.get(param_name)
@@ -332,20 +347,7 @@ def load_model(cf, model, device, run_id: str, with_ddp: bool, with_fsdp: bool,
 
     else:
         # fix mismatch between state_dict keys that can occur between interactive/non-interactive
-        model_has_prefix_module = list(model.state_dict().keys())[0].split(".")[0] == "module"
-        params_has_prefix_module = list(params.keys())[0].split(".")[0] == "module"
-        if model_has_prefix_module and not params_has_prefix_module:
-            # add "module." prefix
-            params_temp = {}
-            for k in params.keys():
-                params_temp["module." + k] = params[k]
-            params = params_temp
-        elif not model_has_prefix_module and params_has_prefix_module:
-            # remove "module." prefix
-            params_temp = {}
-            for k in params.keys():
-                params_temp[k.replace("module.", "")] = params[k]
-            params = params_temp
+        params = _align_module_prefix(params, model.state_dict())
         # load checkpoint
         mkeys, ukeys = model.load_state_dict(params, strict=False)
         model = model.to(device)
@@ -390,54 +392,70 @@ def load_decoder_from_checkpoint(
         path_run / filename, map_location=torch.device("cpu"), mmap=True, weights_only=True
     )
 
-    def _strip(key: str) -> str:
-        return key[len("module.") :] if key.startswith("module.") else key
-
-    decoder_params = {k: v for k, v in params.items() if _strip(k).startswith(_DECODER_PREFIXES)}
+    decoder_params = {
+        k: v for k, v in params.items() if _strip_module_prefix(k).startswith(_DECODER_PREFIXES)
+    }
 
     if not decoder_params:
-        logger.warning(
-            f"No decoder weights (matching {_DECODER_PREFIXES}) found in checkpoint {filename}."
+        msg = (
+            f"load_decoder_chkpt: no decoder weights (matching {_DECODER_PREFIXES}) found in "
+            f"checkpoint {filename} (run_id={run_id}). Asking for a decoder overlay from a "
+            "checkpoint that has no decoder is always a misconfiguration."
         )
-        return model
+        raise RuntimeError(msg)
+
+    # Align the "module." convention *before* either load path. The sharded path resolves each
+    # parameter by exact name against the model's state dict, so a mismatch here used to skip
+    # every tensor one by one and load nothing at all -- silently training a random decoder.
+    model_sd = model.state_dict()
+    num_matched = len(decoder_params)
+    decoder_params = _align_module_prefix(decoder_params, model_sd)
 
     is_model_sharded = with_ddp and with_fsdp
     if is_model_sharded:
-        meta_sharded_sd = model.state_dict()
+        meta_sharded_sd = model_sd
         maybe_sharded_sd = {}
+        skipped = []
         for param_name, full_tensor in decoder_params.items():
             sharded_meta_param = meta_sharded_sd.get(param_name)
             if (
                 sharded_meta_param is None
                 or type(sharded_meta_param) is not torch.distributed.tensor.DTensor
             ):
-                logger.warning(
-                    f"Decoder parameter {param_name} from checkpoint not found in model "
-                    "or not sharded; skipping."
-                )
+                skipped.append(param_name)
                 continue
             sharded_tensor = distribute_tensor(
                 full_tensor,
                 sharded_meta_param.device_mesh,
                 sharded_meta_param.placements,
             )
             maybe_sharded_sd[param_name] = torch.nn.Parameter(sharded_tensor)
+        if skipped and is_root():
+            # one summary line, not one per parameter per rank -- the old per-param warning
+            # produced thousands of lines and buried the "Loaded 0" that mattered.
+            logger.warning(
+                f"load_decoder_chkpt: skipped {len(skipped)}/{num_matched} decoder parameters "
+                f"(not found in model or not sharded), e.g. {skipped[:3]}."
+            )
         _, ukeys = model.load_state_dict(maybe_sharded_sd, strict=False, assign=True)
         loaded = maybe_sharded_sd
     else:
-        # align "module." prefix with the model's state dict key convention
-        model_has_prefix_module = list(model.state_dict().keys())[0].split(".")[0] == "module"
-        params_has_prefix_module = next(iter(decoder_params)).split(".")[0] == "module"
-        if model_has_prefix_module and not params_has_prefix_module:
-            decoder_params = {"module." + k: v for k, v in decoder_params.items()}
-        elif not model_has_prefix_module and params_has_prefix_module:
-            decoder_params = {_strip(k): v for k, v in decoder_params.items()}
         _, ukeys = model.load_state_dict(decoder_params, strict=False)
         model = model.to(device)
         loaded = decoder_params
 
+    if not loaded:
+        msg = (
+            f"load_decoder_chkpt matched {num_matched} decoder tensors in {filename} "
+            f"(run_id={run_id}) but loaded none into the model -- this would silently train a "
+            f"randomly initialised decoder. Checkpoint key: {next(iter(decoder_params))!r}; "
+            f"model key: {next(iter(model_sd))!r}."
+        )
+        raise RuntimeError(msg)
+
     logger.info(
-        f"Loaded {len(loaded)} decoder tensors from checkpoint {filename} (run_id={run_id})."
+        f"Loaded {len(loaded)}/{num_matched} decoder tensors from checkpoint {filename} "
+        f"(run_id={run_id})."
     )
     # ukeys = decoder keys that are absent in the model; missing keys are intentionally not
     # reported here since the primary checkpoint provides all non-decoder weights.

src/weathergen/model/model_interface_test.py

@@ -0,0 +1,121 @@
+# (C) Copyright 2025 WeatherGenerator contributors.
+#
+# This software is licensed under the terms of the Apache Licence Version 2.0
+# which can be obtained at http://www.apache.org/licenses/LICENSE-2.0.
+#
+# In applying this licence, ECMWF does not waive the privileges and immunities
+# granted to it by virtue of its status as an intergovernmental organisation
+# nor does it submit to any jurisdiction.
+
+"""Unit tests for checkpoint-key ``module.`` prefix alignment.
+
+``_align_module_prefix`` and ``_strip_module_prefix`` are pure dict/str functions, so they are
+exercised here without torch.distributed. They matter because both load paths resolve parameters
+by exact name: ``load_model``'s sharded branch and ``load_decoder_from_checkpoint``'s sharded
+branch each look names up in ``model.state_dict()``, and a convention mismatch makes every lookup
+miss. In ``load_decoder_from_checkpoint`` that used to happen silently (each miss was skipped
+individually and an empty load still "succeeded"), so a DDP-saved backbone would train a randomly
+initialised decoder. The sharded branches themselves need real FSDP and are covered by the
+integration check in the skill notes, not here.
+"""
+
+import pytest
+
+from weathergen.model.model_interface import (
+    _DECODER_PREFIXES,
+    _align_module_prefix,
+    _strip_module_prefix,
+)
+
+# Shapes of the real key conventions: a DDP-wrapped save vs a bare model state_dict.
+_PREFIXED = ["module.encoder.q_cells", "module.embed_target_coords.ERA5.linear.weight"]
+_BARE = ["encoder.q_cells", "embed_target_coords.ERA5.linear.weight"]
+
+
+def _sd(keys):
+    return {k: object() for k in keys}
+
+
+# --------------------------------------------------------------------------------------
+# _strip_module_prefix
+# --------------------------------------------------------------------------------------
+
+
+def test_strip_removes_only_a_leading_prefix():
+    assert _strip_module_prefix("module.encoder.weight") == "encoder.weight"
+    assert _strip_module_prefix("encoder.weight") == "encoder.weight"
+
+
+def test_strip_does_not_touch_module_elsewhere_in_the_path():
+    """Regression: the old code used key.replace("module.", ""), which corrupts any path with
+    a genuine submodule named "module" in it."""
+    assert _strip_module_prefix("module.encoder.module.weight") == "encoder.module.weight"
+    assert _strip_module_prefix("encoder.module.weight") == "encoder.module.weight"
+
+
+# --------------------------------------------------------------------------------------
+# _align_module_prefix
+# --------------------------------------------------------------------------------------
+
+
+def test_adds_prefix_when_model_has_one_and_params_do_not():
+    out = _align_module_prefix(_sd(_BARE), _sd(_PREFIXED))
+    assert list(out) == _PREFIXED
+
+
+def test_strips_prefix_when_params_have_one_and_model_does_not():
+    """The case that bit every dcft off a DDP-saved backbone (cw6a4szu, nhv6tkln)."""
+    out = _align_module_prefix(_sd(_PREFIXED), _sd(_BARE))
+    assert list(out) == _BARE
+
+
+@pytest.mark.parametrize("keys", [_BARE, _PREFIXED])
+def test_matching_conventions_are_left_alone(keys):
+    params = _sd(keys)
+    out = _align_module_prefix(params, _sd(keys))
+    assert list(out) == keys
+    assert all(out[k] is params[k] for k in keys)
+
+
+def test_empty_params_are_returned_unchanged():
+    assert _align_module_prefix({}, _sd(_PREFIXED)) == {}
+
+
+def test_values_are_preserved_across_realignment():
+    params = _sd(_PREFIXED)
+    out = _align_module_prefix(params, _sd(_BARE))
+    for src, dst in zip(_PREFIXED, _BARE, strict=True):
+        assert out[dst] is params[src]
+
+
+def test_alignment_is_an_involution_between_the_two_conventions():
+    """Round-tripping must land back on the original keys -- the guard against a fix that
+    half-strips or double-prefixes."""
+    bare = _sd(_BARE)
+    there = _align_module_prefix(bare, _sd(_PREFIXED))
+    back = _align_module_prefix(there, _sd(_BARE))
+    assert list(back) == _BARE
+
+
+# --------------------------------------------------------------------------------------
+# the decoder filter must see through either convention
+# --------------------------------------------------------------------------------------
+
+
+@pytest.mark.parametrize(
+    ("key", "is_decoder"),
+    [
+        ("embed_target_coords.ERA5.linear.weight", True),
+        ("module.embed_target_coords.ERA5.linear.weight", True),
+        ("target_token_engines.ERA5.tte.0.lnorm_in_q.embed_aux.0.weight", True),
+        ("module.target_token_engines.ERA5.tte.0.lnorm_in_q.embed_aux.0.weight", True),
+        ("pred_heads.ERA5.0.weight", True),
+        ("module.pred_heads.ERA5.0.weight", True),
+        ("encoder.q_cells", False),
+        ("module.encoder.q_cells", False),
+        ("forecast_engine.net.fe_blocks.0.layers.0.weight", False),
+        ("module.forecast_engine.net.fe_blocks.0.layers.0.weight", False),
+    ],
+)
+def test_decoder_prefix_filter_matches_under_both_conventions(key, is_decoder):
+    assert _strip_module_prefix(key).startswith(_DECODER_PREFIXES) is is_decoder