Upload Quasar_axrvi_ranker.py

Quasar_axrvi_ranker.py

@@ -9847,6 +9847,69 @@ class HFSyncLayer:
 # SECTION 18b — RANKER CHECKPOINT MANAGER
 # ══════════════════════════════════════════════════════════════════════════════════════
 
+def _sanitize_optimizer_state_by_shape(optimizer, opt_state: dict) -> dict:
+    """
+    Drop any per-parameter Adam state (exp_avg / exp_avg_sq / etc.) whose
+    stored tensor shape no longer matches the corresponding *live* param
+    in `optimizer`.
+
+    Why this exists: relying on a saved 'num_assets' checkpoint field (or
+    any other metadata flag) to detect a stale-optimizer situation is
+    fragile — older checkpoints may not have that field at all, and even
+    when shapes match at model-load time (e.g. a lazily-created submodule
+    like FABLECLCU wasn't in the checkpoint dict yet, so it's just a
+    "missing key" rather than a raise), the *optimizer* checkpoint can
+    still hold a state entry at that same param slot sized for the old N.
+    That mismatch only surfaces later, deep inside adamw()'s
+    `exp_avg.lerp_(grad, ...)`, once the lazily-created param gets
+    attached to the optimizer's existing (stale-shaped) state — by which
+    point it's a hard crash on every training step rather than a clean,
+    loggable restore-time skip.
+
+    Comparing shapes directly against the live optimizer's params is
+    metadata-agnostic: it catches the mismatch regardless of why it
+    happened, and only ever discards the specific entries that are
+    actually incompatible — every other trained parameter's momentum
+    keeps flowing through untouched.
+    """
+    if not opt_state or "state" not in opt_state:
+        return opt_state
+
+    # Flatten in the same order load_state_dict expects: params in the
+    # order they appear across param_groups, which is how the integer
+    # keys in opt_state['state'] are indexed.
+    live_params = [p for group in optimizer.param_groups for p in group["params"]]
+
+    state = opt_state["state"]
+    dropped = []
+    for idx in list(state.keys()):
+        if idx >= len(live_params):
+            # Optimizer has fewer params now than the checkpoint did
+            # (e.g. a whole layer removed) — nothing to compare against.
+            dropped.append(idx)
+            del state[idx]
+            continue
+        live_shape = tuple(live_params[idx].shape)
+        entry = state[idx]
+        mismatched = any(
+            hasattr(v, "shape") and tuple(v.shape) != live_shape
+            for v in entry.values()
+            if hasattr(v, "shape") and v.dim() > 0  # skip scalar 'step' counters
+        )
+        if mismatched:
+            dropped.append(idx)
+            del state[idx]
+
+    if dropped:
+        logger.warning(
+            f"[OptimizerRestore] Dropped {len(dropped)} stale optimizer state "
+            f"entr{'y' if len(dropped) == 1 else 'ies'} with shape mismatches "
+            f"against the current model (fresh Adam state will be used for "
+            f"those params; every other param's momentum is preserved)."
+        )
+    return opt_state
+
+
 class RankerCheckpointManager:
     """
     Bridge-facing checkpoint controller for QuasarAXRVIBridge.
@@ -10103,6 +10166,18 @@ class RankerCheckpointManager:
         current_num_assets = bridge.axrvi_net.num_assets if bridge.axrvi_net is not None else -1
         _model_compatible  = True
 
+        # NOTE: older checkpoints (e.g. ones restored from HF before
+        # 'num_assets' was added to the save dict) don't carry this field,
+        # so this check can silently miss a real mismatch and report
+        # "compatible". That's fine here — axrvi_net.load_state_dict()
+        # below is strict=False and a shape-mismatched tensor at a shared
+        # key would raise immediately regardless of that flag, so it's
+        # safe either way. The optimizer side is *not* self-protecting the
+        # same way (stale state can attach to a lazily-created param
+        # later without erroring at load time), which is why it gets its
+        # own metadata-independent shape check below via
+        # _sanitize_optimizer_state_by_shape() rather than trusting this
+        # flag alone.
         if ckpt_num_assets != -1 and ckpt_num_assets != current_num_assets:
             ckpt_assets = ckpt.get("asset_symbols", "unknown")
             logger.warning(
@@ -10131,7 +10206,26 @@ class RankerCheckpointManager:
             tr = bridge.trainer
             if _model_compatible:
                 if "optimizer" in ckpt:
-                    tr.optimizer.load_state_dict(ckpt["optimizer"])
+                    # FIX 5: don't trust `_model_compatible` alone — it's
+                    # derived from a 'num_assets' checkpoint field that
+                    # older checkpoints (like ones restored from HF before
+                    # this field existed) simply don't have, silently
+                    # defaulting to "compatible". Validate optimizer state
+                    # against live param shapes directly as a second,
+                    # metadata-independent line of defense — see
+                    # _sanitize_optimizer_state_by_shape() for why this
+                    # can't be caught at axrvi_net.load_state_dict() time.
+                    try:
+                        safe_opt_state = _sanitize_optimizer_state_by_shape(
+                            tr.optimizer, ckpt["optimizer"]
+                        )
+                        tr.optimizer.load_state_dict(safe_opt_state)
+                    except Exception as e:
+                        logger.warning(
+                            f"[Restore] ⚠️  Optimizer state restore failed "
+                            f"({e}) — continuing with fresh optimizer state "
+                            f"rather than crashing the restore."
+                        )
                 if "scheduler" in ckpt:
                     tr.scheduler.load_state_dict(ckpt["scheduler"])
             else: