scripts/recurrent_qwen/checkpoint.py

"""Portable trainable recurrent tensors plus explicit architecture/base identity."""

import json
from pathlib import Path
from typing import Any

import torch
from transformers import AutoTokenizer, Qwen2ForCausalLM

from .lora_utils import attach_recurrent_lora
from .model import RecurrentQwen
from scripts.dataset.pointer import SYMBOLS


EXECUTOR_FORMAT = "loopformer-executor-v2"
FORMAT = "loopformer-stage1-v1"
COMPLETION_FORMAT = "loopformer-stage1-completion-v1"
FIXED_PROMPT_FORMAT = "loopformer-stage1-fixed-prompt-v1"


def checkpoint_mode(spec: dict) -> str:
    """Legacy checkpoints mean full-sequence recurrence; never infer new modes."""
    mode = spec.get("recurrence_mode", "full_sequence")
    if spec.get("format") not in (FORMAT, COMPLETION_FORMAT, FIXED_PROMPT_FORMAT, EXECUTOR_FORMAT):
        raise ValueError("Unsupported recurrent checkpoint format")
    if (spec["format"] == EXECUTOR_FORMAT) != ("executor" in spec):
        raise ValueError("Executor architecture requires the versioned executor checkpoint format")
    expected = "fixed_prompt" if spec["format"] in (FIXED_PROMPT_FORMAT, EXECUTOR_FORMAT) else "full_sequence"
    if mode != expected:
        raise ValueError("Checkpoint format and recurrence_mode disagree")
    return mode


def cpu_tree(value: Any) -> Any:
    if isinstance(value, torch.Tensor):
        return value.detach().cpu().clone()
    if isinstance(value, dict):
        return {key: cpu_tree(item) for key, item in value.items()}
    if isinstance(value, list):
        return [cpu_tree(item) for item in value]
    if isinstance(value, tuple):
        return tuple(cpu_tree(item) for item in value)
    return value


def save_checkpoint(path: Path, model: RecurrentQwen, tokenizer: Any, spec: dict, training_state: dict) -> None:
    """Write a new checkpoint directory; metadata is the final completion marker."""
    if spec.get("recurrence_mode", "full_sequence") != model.recurrence_mode:
        raise ValueError("Checkpoint spec and recurrence_mode disagree")
    weights = {name: parameter for name, parameter in model.named_parameters() if parameter.requires_grad}
    completion = model.completion_head is not None
    expected_completion = spec.get("completion_head")
    if completion != (expected_completion is not None):
        raise ValueError("Checkpoint spec and completion head disagree")
    if completion and expected_completion != {"intermediate": model.completion_head.intermediate}:
        raise ValueError("Checkpoint completion-head architecture differs from spec")
    new_architecture = "executor" in spec
    if not new_architecture and (not weights or any(not ((name.startswith("recurrent.") and ".lora_" in name) or
                               (completion and name.startswith("completion_head."))) for name in weights)):
        raise ValueError("Checkpoint expects recurrent LoRA and optional completion head only")
    if new_architecture:
        if spec["executor"]["train_scope"] != model.train_scope:
            raise ValueError("Checkpoint training scope differs from model")
        if not weights:
            raise ValueError("Cannot save an empty executor")
    path.mkdir(parents=True, exist_ok=False)
    torch.save(cpu_tree(weights), path / "adapter_model.pt")
    torch.save(cpu_tree(training_state), path / "training_state.pt")
    tokenizer.save_pretrained(path)
    version = FIXED_PROMPT_FORMAT if model.recurrence_mode == "fixed_prompt" else (COMPLETION_FORMAT if completion else FORMAT)
    if new_architecture:
        version = EXECUTOR_FORMAT
    (path / "recurrent_config.json").write_text(json.dumps({**spec, "format": version}, indent=2) + "\n")


def restore_adapters(model: RecurrentQwen, path: Path) -> None:
    if checkpoint_mode(json.loads((path / "recurrent_config.json").read_text())) != model.recurrence_mode:
        raise ValueError("Adapter checkpoint recurrence_mode differs from model")
    state = torch.load(path / "adapter_model.pt", map_location="cpu", weights_only=True)
    expected = {name: value for name, value in model.named_parameters() if value.requires_grad}
    if set(state) != set(expected):
        raise ValueError("Adapter checkpoint names do not match the recurrent model")
    with torch.no_grad():
        for name, parameter in expected.items():
            tensor = state[name]
            if tensor.shape != parameter.shape or not torch.isfinite(tensor).all():
                raise ValueError(f"Invalid adapter tensor: {name}")
            parameter.copy_(tensor)


def load_recurrent_checkpoint(path: Path, *, device: str = "cpu", dtype: str = "float32",
                              download: bool = False) -> tuple[RecurrentQwen, Any, dict]:
    spec = json.loads((path / "recurrent_config.json").read_text())
    mode = checkpoint_mode(spec)
    if spec.get("prompt_format") != "dataset_raw" or spec.get("loss_vocabulary") != "symbols":
        raise ValueError("Unsupported recurrent checkpoint format/prompt/loss")
    completion = spec["format"] == COMPLETION_FORMAT or (mode == "fixed_prompt" and "completion_head" in spec)
    if completion and (not isinstance(spec.get("completion_head"), dict) or
                       type(spec["completion_head"].get("intermediate")) is not int or
                       spec["completion_head"]["intermediate"] < 1):
        raise ValueError("Invalid saved completion-head architecture")
    if not completion and "completion_head" in spec:
        raise ValueError("Legacy checkpoint cannot contain a completion head")
    if spec.get("symbols") != list(SYMBOLS) or len(set(spec.get("token_ids", []))) != len(SYMBOLS):
        raise ValueError("Checkpoint requires the complete A–Z symbol vocabulary")
    tokenizer = AutoTokenizer.from_pretrained(path, local_files_only=True, trust_remote_code=False)
    for symbol, token in zip(spec["symbols"], spec["token_ids"], strict=True):
        if tokenizer.encode(" " + symbol, add_special_tokens=False) != [token]:
            raise ValueError("Checkpoint tokenizer and symbol IDs disagree")
    base = Qwen2ForCausalLM.from_pretrained(
        spec["base_model"], revision=spec["revision"], local_files_only=not download,
        dtype=getattr(torch, dtype), attn_implementation="eager", trust_remote_code=False,
    ).to(device).eval()
    model = RecurrentQwen(base, spec["recurrent_start"], spec["recurrent_end"], recurrence_mode=mode)
    executor = spec.get("executor")
    if executor is None or executor["train_scope"] == "lora":
        attach_recurrent_lora(model, rank=spec["lora_rank"], alpha=spec["lora_alpha"])
    if executor is not None:
        model.configure_executor(tokenizer, **executor)
    if completion and model.completion_head is None:
        model.enable_completion(spec["completion_head"]["intermediate"])
    restore_adapters(model, path)
    model.eval()
    return model, tokenizer, spec