diff --git a/src/weathergen/model/model_interface.py b/src/weathergen/model/model_interface.py index 440d12f7c..206e6a612 100644 --- a/src/weathergen/model/model_interface.py +++ b/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,31 +392,37 @@ 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, @@ -422,22 +430,32 @@ def _strip(key: str) -> str: 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. diff --git a/src/weathergen/model/model_interface_test.py b/src/weathergen/model/model_interface_test.py new file mode 100644 index 000000000..fbf6a8995 --- /dev/null +++ b/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