Unify FFN and TrajMixer model architectures

model_architectures.py

@@ -0,0 +1,131 @@
+"""Model-architecture identifiers and checkpoint validation helpers."""
+
+from __future__ import annotations
+
+import re
+from collections.abc import Mapping
+
+
+TRANSFORMER_FFN_ARCHITECTURE = "transformer_ffn_v1"
+TRAJ_MIXER_ARCHITECTURE = "traj_mixer_v5"
+DEFAULT_MODEL_ARCHITECTURE = TRANSFORMER_FFN_ARCHITECTURE
+SUPPORTED_MODEL_ARCHITECTURES = (
+    TRANSFORMER_FFN_ARCHITECTURE,
+    TRAJ_MIXER_ARCHITECTURE,
+)
+
+
+_FFN_STATE_KEY = re.compile(
+    r"(?:^|\.)blocks\.\d+\.mlp\.w[123]\.(?:weight|bias)$"
+)
+_TRAJ_MIXER_STATE_KEY = re.compile(
+    r"(?:^|\.)blocks\.\d+\.mlp\.(?:"
+    r"norm\.(?:weight|bias)|"
+    r"intra_gate_proj|"
+    r"intra_value_proj|"
+    r"intra_output_proj|"
+    r"intra_gate_logits|"
+    r"gate_proj|"
+    r"value_proj|"
+    r"output_proj"
+    r")$"
+)
+
+
+def _validate_model_architecture(model_architecture: object) -> str:
+    if not isinstance(model_architecture, str):
+        raise ValueError(
+            "model_architecture must be one of "
+            f"{SUPPORTED_MODEL_ARCHITECTURES}, got {model_architecture!r}"
+        )
+    if model_architecture not in SUPPORTED_MODEL_ARCHITECTURES:
+        raise ValueError(
+            f"Unsupported model_architecture={model_architecture!r}; "
+            f"expected one of {SUPPORTED_MODEL_ARCHITECTURES}."
+        )
+    return model_architecture
+
+
+def detect_model_architecture_from_state_dict(
+    state_dict: Mapping[str, object],
+) -> str:
+    """Infer the architecture from block parameter names.
+
+    Detection deliberately accepts any ``blocks.<index>`` prefix rather than
+    assuming that block zero is present.
+    """
+
+    if not isinstance(state_dict, Mapping):
+        raise TypeError(
+            "state_dict must be a mapping, got "
+            f"{type(state_dict).__name__}"
+        )
+
+    has_ffn = False
+    has_traj_mixer = False
+    for raw_key in state_dict:
+        key = str(raw_key)
+        has_ffn = has_ffn or _FFN_STATE_KEY.search(key) is not None
+        has_traj_mixer = (
+            has_traj_mixer
+            or _TRAJ_MIXER_STATE_KEY.search(key) is not None
+        )
+        if has_ffn and has_traj_mixer:
+            raise ValueError(
+                "Checkpoint contains both Transformer FFN and TrajMixer "
+                "block parameters; its model architecture is ambiguous."
+            )
+
+    if has_ffn:
+        return TRANSFORMER_FFN_ARCHITECTURE
+    if has_traj_mixer:
+        return TRAJ_MIXER_ARCHITECTURE
+    raise ValueError(
+        "Could not detect model architecture from checkpoint parameters. "
+        "Expected a blocks.<index>.mlp FFN or TrajMixer parameter."
+    )
+
+
+def resolve_model_architecture(
+    config_or_marker: Mapping[str, object] | str | None = None,
+    state_dict: Mapping[str, object] | None = None,
+) -> str:
+    """Resolve and cross-check a configured and checkpoint architecture.
+
+    Every saved run must provide an explicit architecture marker. Checkpoint
+    parameter names are used only to verify that the marker describes the
+    weights being loaded.
+    """
+
+    if isinstance(config_or_marker, Mapping):
+        configured = config_or_marker.get("model_architecture")
+    elif isinstance(config_or_marker, str) or config_or_marker is None:
+        configured = config_or_marker
+    else:
+        raise TypeError(
+            "config_or_marker must be a config mapping, string, or None, got "
+            f"{type(config_or_marker).__name__}"
+        )
+
+    resolved_config = (
+        _validate_model_architecture(configured)
+        if configured is not None
+        else None
+    )
+    detected = (
+        detect_model_architecture_from_state_dict(state_dict)
+        if state_dict is not None
+        else None
+    )
+
+    if resolved_config is None:
+        raise ValueError(
+            "model_architecture is required; expected one of "
+            f"{SUPPORTED_MODEL_ARCHITECTURES}."
+        )
+    if detected is not None and resolved_config != detected:
+        raise ValueError(
+            "Configured model architecture conflicts with checkpoint: "
+            f"config={resolved_config!r}, checkpoint={detected!r}."
+        )
+    return resolved_config