feat(v7-c01): opt-state arch mismatch surfaced honestly + safe skip

Closes V7-C01 (cppmega-mlx-q3jq): opt-state load was silently
dropping mismatched keys via tree_unflatten and would crash mid-step
on shape mismatch. v7 adds structural fingerprint validation:

  * extras.checkpoint.opt_state_arch_diff carries
    {missing_keys, extra_keys, shape_mismatch} (plus *_count) when
    the loaded opt-state doesn't match the live model params.
  * Base-key derivation strips opt-state suffixes (".m", ".v", etc.)
    to match against model.parameters() keys, and ignores AdamW
    internal scalars ("step", "learning_rate").
  * opts.opt_state_strict=True → skip load on ANY diff,
    opt_state_warning="opt_state arch mismatch (strict mode): …;
    cold restart".
  * Default (non-strict) tolerates missing/extra keys but skips on
    shape_mismatch (would crash mid-step), warning="… (shape
    mismatch): …; cold restart".

Tests (tests/v4/test_checkpoint_arch_mismatch.py): 4/4 —
  * Matching arch → diff=None, no warning.
  * Mismatched intermediate_size (256→128) → diff populated with
    shape_mismatch > 0, load skipped, training continues cold.
  * strict_mode=True on mismatch → load skipped, "strict mode"
    warning, losses still produced.
  * strict_mode=True on matched arch → still loads cleanly.
- pytest stage_train + checkpoint + rng + arch regression: 68/68.

Author: apstenku123

cppmega_v4/runner/stages.py

@@ -45,6 +45,11 @@
 StageStatus = Literal["ok", "skipped", "fail", "cancelled"]
 
 
+class _ArchMismatchSentinel(Exception):
+    """V7-C01 internal control-flow: skip opt-state load on strict
+    arch mismatch without surfacing as a real error."""
+
+
 @dataclass
 class StageResult:
     """One stage's execution outcome."""
@@ -1118,19 +1123,99 @@ def _count(tree: Any) -> int:
         # resumed run picks up Adam moments → strict losses[0] parity
         # with the saved run's losses[-1].
         opt_state_load = opts.get("opt_state_load_path")
+        opt_state_strict = bool(opts.get("opt_state_strict", False))
+        opt_state_arch_diff: dict[str, Any] | None = None
         if opt_state_load:
             try:
                 import safetensors.mlx as _stmlx
                 loaded_st = _stmlx.load_file(opt_state_load)
                 # V01: pull rng_key out of the opt-state bundle before
-                # tree_unflatten so it doesn't pollute opt.state. mlx
-                # rng keys are mx.uint32 arrays of shape (2,).
+                # tree_unflatten so it doesn't pollute opt.state.
                 rng_buf = loaded_st.pop("_rng_key", None)
+                # V7-C01: structural fingerprint diff. Fresh Adam
+                # optimisers initialise opt.state lazily (only after
+                # the first opt.update), so the saved opt-state is
+                # compared against the LIVE MODEL params instead.
+                # Each opt-state key for AdamW carries the same shape
+                # as its underlying param (m and v moments share the
+                # param shape). We strip suffixes that aren't actual
+                # model keys to find the base param.
+                live_flat = dict(nn.utils.tree_flatten(
+                    all_modules.parameters()))
+                def _fp(d: dict) -> dict[str, tuple]:
+                    return {k: (tuple(v.shape), str(v.dtype))
+                            for k, v in d.items() if hasattr(v, "shape")}
+                live_fp = _fp(live_flat)
+                saved_fp = _fp(loaded_st)
+                def _base(k: str) -> str | None:
+                    if k in live_fp:
+                        return k
+                    # Walk back along dot path looking for a model
+                    # param this opt-state entry corresponds to.
+                    parts = k.split(".")
+                    for cut in range(len(parts) - 1, 0, -1):
+                        cand = ".".join(parts[:cut])
+                        if cand in live_fp:
+                            return cand
+                    return None
+                # Non-param opt entries (e.g. AdamW "step",
+                # "learning_rate" scalars) are neither model-derived
+                # nor problematic — skip them.
+                NON_PARAM_SUFFIXES = {"step", "learning_rate"}
+                missing: list[str] = []   # model param with no opt-state
+                covered: set[str] = set()
+                extra: list[str] = []     # opt-state without a model param
+                shape_mismatch: list[str] = []
+                for sk, sv in saved_fp.items():
+                    if (sk in NON_PARAM_SUFFIXES
+                            or sk.endswith(".step")
+                            or sk.endswith(".learning_rate")):
+                        continue
+                    bk = _base(sk)
+                    if bk is None:
+                        extra.append(sk)
+                        continue
+                    covered.add(bk)
+                    if live_fp[bk] != sv:
+                        shape_mismatch.append(sk)
+                missing = sorted(set(live_fp) - covered)
+                extra = sorted(extra)
+                shape_mismatch = sorted(shape_mismatch)
+                if missing or extra or shape_mismatch:
+                    opt_state_arch_diff = {
+                        "missing_keys": missing[:10],
+                        "missing_keys_count": len(missing),
+                        "extra_keys": extra[:10],
+                        "extra_keys_count": len(extra),
+                        "shape_mismatch": shape_mismatch[:10],
+                        "shape_mismatch_count": len(shape_mismatch),
+                    }
+                    # Strict mode: skip on ANY diff.
+                    # Non-strict: skip only on shape_mismatch (would
+                    # crash mid-step otherwise); tolerate missing/extra.
+                    skip_load = (opt_state_strict
+                                  or len(shape_mismatch) > 0)
+                    if skip_load:
+                        mode = ("strict mode" if opt_state_strict
+                                else "shape mismatch")
+                        opt_state_warning = (
+                            f"opt_state arch mismatch ({mode}): "
+                            f"{len(missing)} missing, {len(extra)} extra, "
+                            f"{len(shape_mismatch)} shape-mismatch keys; "
+                            f"cold restart"
+                        )
+                        raise _ArchMismatchSentinel()
+                # Not strict (or no diff) → attempt load. Missing/extra
+                # keys are tolerated; shape_mismatch will still raise
+                # in tree_unflatten and surface as the generic warning.
                 opt.state = nn.utils.tree_unflatten(list(loaded_st.items()))
                 if rng_buf is not None:
                     rng_key = rng_buf
                     rng_key_loaded = True
                 opt_state_loaded_path = str(opt_state_load)
+            except _ArchMismatchSentinel:
+                # Already populated opt_state_warning above.
+                pass
             except FileNotFoundError as exc:
                 opt_state_warning = (
                     f"opt_state_load_path missing: {exc}; cold restart")
@@ -1429,6 +1514,7 @@ def _count(tree: Any) -> int:
                     "opt_state_loaded_path": opt_state_loaded_path,
                     "opt_state_warning": opt_state_warning,
                     "rng_key_loaded": rng_key_loaded,
+                    "opt_state_arch_diff": opt_state_arch_diff,
                 },
                 "mtp": _compute_mtp_extras(
                     all_modules, mtp_k, mtp_betas, vocab_size,

tests/v4/test_checkpoint_arch_mismatch.py

@@ -0,0 +1,119 @@
+"""V7-C01: opt-state load surfaces arch mismatch honestly.
+
+Currently a checkpoint produced for a different architecture (extra
+brick, renamed layer, changed hidden size) loads silently — keys are
+dropped by tree_unflatten and shape mismatches raise opaque errors
+mid-step. v7 adds a structural fingerprint diff:
+
+  * extras.checkpoint.opt_state_arch_diff dict with
+    {missing_keys, extra_keys, shape_mismatch} when the checkpoint's
+    opt-state fingerprint doesn't match the current model.
+  * opts.opt_state_strict=True → skip the load and surface a
+    "cold restart" warning instead of partial overwrite.
+  * Default (strict=False) preserves the prior best-effort behavior
+    but still reports the diff so the user can see what was dropped.
+"""
+
+from __future__ import annotations
+
+import pytest
+
+from cppmega_v4.jsonrpc.schema import VerifyParams
+from cppmega_v4.runner import Pipeline, run_pipeline
+
+
+def _spec(intermediate: int = 256) -> VerifyParams:
+    """intermediate_size controls mlp weight shape → easy way to
+    produce a checkpoint that won't fit a different architecture."""
+    return VerifyParams.model_validate({
+        "graph": {
+            "nodes": [
+                {"id": "attn", "kind": "attention",
+                 "params": {"num_heads": 4, "head_dim": 64}},
+                {"id": "mlp", "kind": "mlp",
+                 "params": {"intermediate_size": intermediate}},
+            ],
+            "edges": [{"src": "attn", "dst": "mlp"}],
+        },
+        "dim_env": {"B": 1, "S": 8, "H": 128,
+                    "nh": 2, "nkv": 1, "head_dim": 64},
+        "loss": {"kind": "cross_entropy", "head_outputs": ["mlp"]},
+        "optim": {"kind": "adamw",
+                  "groups": [{"matcher": "all", "lr": 1e-3,
+                              "weight_decay": 0.01,
+                              "betas": [0.9, 0.95]}]},
+    })
+
+
+def _train(spec, **opts) -> dict:
+    rep = run_pipeline(spec, Pipeline.from_dict({
+        "stages": ["parse", "verify_build_spec", "build_model", "train"],
+        "stage_options": {"train": opts},
+    }))
+    tr = next(s for s in rep.stages if s.name == "train")
+    assert tr.status == "ok", f"train failed: {tr.error}"
+    return tr.extras
+
+
+def test_c01_arch_diff_null_on_matching_load(tmp_path):
+    """Matched arch → no diff reported."""
+    save = str(tmp_path / "opt.safetensors")
+    _train(_spec(), num_steps=2, opt_state_save_path=save)
+    extras = _train(_spec(), num_steps=1, opt_state_load_path=save)
+    assert extras["checkpoint"]["opt_state_loaded_path"] == save
+    assert extras["checkpoint"]["opt_state_arch_diff"] is None
+    assert extras["checkpoint"]["opt_state_warning"] is None
+
+
+def test_c01_arch_diff_populated_on_shape_mismatch(tmp_path):
+    """Save with intermediate=256 → load into intermediate=128 model.
+    Should populate opt_state_arch_diff with shape_mismatch entries
+    AND skip the load (would crash mid-step otherwise) — train
+    proceeds as cold restart."""
+    save = str(tmp_path / "opt.safetensors")
+    _train(_spec(intermediate=256), num_steps=2,
+           opt_state_save_path=save)
+    extras = _train(_spec(intermediate=128), num_steps=1,
+                    opt_state_load_path=save)
+    diff = extras["checkpoint"]["opt_state_arch_diff"]
+    assert diff is not None, (
+        "expected opt_state_arch_diff for mismatched intermediate_size"
+    )
+    assert diff["shape_mismatch_count"] > 0
+    # Non-strict but shape_mismatch → load skipped, warning surfaced,
+    # training still completes (cold restart).
+    assert extras["checkpoint"]["opt_state_loaded_path"] is None
+    warning = extras["checkpoint"]["opt_state_warning"]
+    assert isinstance(warning, str)
+    assert "shape mismatch" in warning.lower()
+    assert len(extras["losses"]) == 1
+
+
+def test_c01_strict_mode_blocks_load_on_mismatch(tmp_path):
+    """opt_state_strict=True → load skipped, warning + diff populated,
+    opt_state_loaded_path stays None."""
+    save = str(tmp_path / "opt.safetensors")
+    _train(_spec(intermediate=256), num_steps=2,
+           opt_state_save_path=save)
+    extras = _train(_spec(intermediate=128), num_steps=1,
+                    opt_state_load_path=save,
+                    opt_state_strict=True)
+    assert extras["checkpoint"]["opt_state_loaded_path"] is None
+    diff = extras["checkpoint"]["opt_state_arch_diff"]
+    assert diff is not None
+    warning = extras["checkpoint"]["opt_state_warning"]
+    assert isinstance(warning, str) and "strict mode" in warning.lower()
+    assert "cold restart" in warning.lower()
+    # Training still produced losses (cold start).
+    assert len(extras["losses"]) == 1
+
+
+def test_c01_strict_mode_passes_through_when_matched(tmp_path):
+    """opt_state_strict=True on a matched arch → still loads cleanly."""
+    save = str(tmp_path / "opt.safetensors")
+    _train(_spec(), num_steps=2, opt_state_save_path=save)
+    extras = _train(_spec(), num_steps=1,
+                    opt_state_load_path=save,
+                    opt_state_strict=True)
+    assert extras["checkpoint"]["opt_state_loaded_path"] == save
+    assert extras["checkpoint"]["opt_state_arch_diff"] is None