scripts/recurrent_qwen/model.py

"""Qwen2 decoder surgery for the pinned Transformers implementation."""

import math
from typing import Any, Callable, Literal

import torch
from torch import Tensor, nn
from transformers import Qwen2ForCausalLM
from transformers.masking_utils import create_causal_mask

from .outputs import RecurrentOutput, answer_margin
from .completion import CompletionHead
from .memory import PromptMemory
from .prefix import fixed_prefix_layer
from torch.utils.checkpoint import checkpoint
from .interfaces import PromptRouter, ReentryBridge, RecurrentController, AffineSuffixController, SharedNumberController, controller_head
from .number_reader import MAX_NUMBER_TOKENS, number_token_ids


def _answer_readout(hidden: Tensor, positions: list[int]) -> Tensor:
    # Integer slices avoid MPS's nondeterministic scatter backward for advanced
    # indexing/gather. Batches are deliberately tiny; positions sync once per forward.
    return torch.stack([hidden[row, position] for row, position in enumerate(positions)])


class DecoderBlock(nn.Module):
    """An ordered group of original decoder objects, with no copies."""

    def __init__(self, layers: list[nn.Module]) -> None:
        super().__init__()
        self.layers = nn.ModuleList(layers)
        self.prefix_reuse = False
        self.gradient_checkpointing = False

    def forward(
        self,
        hidden: Tensor,
        attention_mask: Tensor | None,
        position_ids: Tensor,
        position_embeddings: tuple[Tensor, Tensor],
        memory: PromptMemory | None = None,
    ) -> Tensor:
        for index, layer in enumerate(self.layers):
            if memory is not None:
                hidden = memory.layer_input(index, hidden)
            if self.prefix_reuse and memory is not None:
                hidden = fixed_prefix_layer(layer, hidden, position_embeddings, attention_mask,
                                            memory, index, self.gradient_checkpointing)
                continue
            call = layer
            if self.gradient_checkpointing and torch.is_grad_enabled():
                def call(x, _layer=layer, **kwargs):
                    return checkpoint(_layer, x, use_reentrant=False, **kwargs)
            hidden = call(
                hidden,
                attention_mask=attention_mask,
                position_ids=position_ids,
                position_embeddings=position_embeddings,
                past_key_values=None,
                use_cache=False,
            )
        return memory.block_output(hidden) if memory is not None else hidden


class RecurrentQwen(nn.Module):
    """P -> R repeated T times -> C, reusing the supplied model's weights.

    Construction freezes and shares the supplied base model in place. Load a
    fresh base for an independent reference after adapters have been attached.
    Only full attention and default RoPE are supported in Stage 0, as used by
    Qwen2.5-0.5B-Instruct. Device/dtype are inherited, never selected here.
    """

    def __init__(
        self, base: Qwen2ForCausalLM, recurrent_start: int = 6, recurrent_end: int = 18,
        *, recurrence_mode: str = "full_sequence",
    ) -> None:
        super().__init__()
        if not isinstance(base, Qwen2ForCausalLM):
            raise TypeError("Stage 0 requires Qwen2ForCausalLM")
        if recurrence_mode not in ("full_sequence", "fixed_prompt"):
            raise ValueError("recurrence_mode must be full_sequence or fixed_prompt")
        self.recurrence_mode = recurrence_mode
        layers = list(base.model.layers)
        if not 0 <= recurrent_start < recurrent_end <= len(layers):
            raise ValueError("Require 0 <= recurrent_start < recurrent_end <= layer count")
        if any(kind != "full_attention" for kind in base.config.layer_types):
            raise ValueError("Stage 0 supports full-attention Qwen layers only")
        if base.config.rope_parameters["rope_type"] != "default":
            raise ValueError("Stage 0 supports default RoPE only")
        if getattr(base, "peft_config", None):
            raise ValueError("Supply an unadapted base; attach LoRA to recurrence afterwards")
        base.requires_grad_(False)
        base.zero_grad(set_to_none=True)
        base.config.use_cache = False
        self.config = base.config
        self.recurrent_start = recurrent_start
        self.recurrent_end = recurrent_end
        self.embed_tokens = base.model.embed_tokens
        self.rotary_emb = base.model.rotary_emb
        self.prelude = DecoderBlock(layers[:recurrent_start])
        self.recurrent = DecoderBlock(layers[recurrent_start:recurrent_end])
        self.coda = DecoderBlock(layers[recurrent_end:])
        self.norm = base.model.norm
        self.lm_head = base.lm_head
        self.completion_head: CompletionHead | RecurrentController | None = None
        self.router = None
        self.bridge = None
        self.state_head = None
        self.train_scope = "lora"
        self.prefix_reuse = False
        self.gradient_checkpointing = False
        self.train(base.training)

    def enable_completion(self, intermediate: int = 128) -> None:
        """Attach a trainable stop readout without changing recurrent inputs."""
        if self.completion_head is not None:
            raise ValueError("Completion head is already attached")
        self.completion_head = CompletionHead(self.config.hidden_size, intermediate).to(
            device=self.embed_tokens.weight.device, dtype=self.embed_tokens.weight.dtype)

    def configure_executor(self, tokenizer: Any, *, train_scope: str = "recurrent",
                           bridge: bool = True, direct_readout: bool = True,
                           controller_size: int = 128, isolated: bool = True,
                           prefix_reuse: bool = False, gradient_checkpointing: bool = False,
                           controller_kind: str = 'gru') -> None:
        """Opt-in architecture; legacy forwards/checkpoints keep their semantics."""
        if train_scope not in ("lora", "recurrent", "full"):
            raise ValueError("train_scope must be lora, recurrent, or full")
        if self.recurrence_mode != "fixed_prompt":
            raise ValueError("Executor interfaces require fixed_prompt recurrence")
        if train_scope != "lora" and any(".lora_" in name for name, _ in self.named_parameters()):
            raise ValueError("Full-weight scopes require an unadapted base, not an attached LoRA model")
        if prefix_reuse and self.config.attention_dropout != 0:
            raise ValueError("Prefix reuse requires zero attention dropout for equivalent execution")
        self.train_scope = train_scope
        if train_scope == "recurrent":
            self.recurrent.requires_grad_(True)
        elif train_scope == "full":
            self.requires_grad_(True)
        device, dtype = self.embed_tokens.weight.device, self.embed_tokens.weight.dtype
        if isolated:
            self.router = PromptRouter(tokenizer)
            if controller_size:
                self.completion_head = controller_head(self.config.hidden_size, controller_size, controller_kind).to(device=device, dtype=dtype)
        if bridge:
            self.bridge = ReentryBridge(self.config.hidden_size).to(device=device, dtype=dtype)
        if direct_readout:
            self.state_head = nn.Linear(self.config.hidden_size, 26).to(device=device, dtype=dtype)
        self.prefix_reuse = prefix_reuse
        self.gradient_checkpointing = gradient_checkpointing
        for block in (self.prelude, self.recurrent, self.coda):
            block.prefix_reuse = prefix_reuse
            block.gradient_checkpointing = gradient_checkpointing

    def controller_prompt_state(self, input_ids: Tensor, attention_mask: Tensor,
                                answer_positions: Tensor | None = None) -> Tensor:
        """Frozen native initializer features [B,H_context], without R/C execution.

        This interface runs neither R nor C and consumes only token IDs/mask.
        GRU uses full-prompt P's answer position; affine_suffix uses the final
        eight frozen token embeddings. Shared-number uses routed field-token
        embeddings and validity masks. No branch interprets a numeric value.
        """
        if input_ids.ndim != 2 or attention_mask.shape != input_ids.shape or not attention_mask.any(-1).all():
            raise ValueError('Need aligned nonempty [B,S] prompt IDs and mask')
        if answer_positions is None:
            answer_positions = (attention_mask.long() * torch.arange(1, input_ids.shape[1] + 1,
                                device=input_ids.device)).argmax(-1)
        if isinstance(self.completion_head, SharedNumberController):
            rows = []
            for ids, mask in zip(input_ids.detach().cpu().tolist(), attention_mask.detach().cpu().tolist(), strict=True):
                valid = tuple(t for t, active in zip(ids, mask, strict=True) if active)
                rows.append(number_token_ids(self.router.tokenizer, valid))
            with torch.no_grad():
                context = self.embed_tokens.weight.new_zeros(len(rows), MAX_NUMBER_TOKENS, self.config.hidden_size + 1)
                for i, tokens in enumerate(rows):
                    context[i, :len(tokens), :-1] = self.embed_tokens(torch.tensor(tokens, device=input_ids.device))
                    context[i, :len(tokens), -1] = 1
            return context.flatten(1)
        if isinstance(self.completion_head, AffineSuffixController):
            k = self.completion_head.suffix_tokens
            if (answer_positions < k - 1).any():
                raise ValueError('Affine suffix initialization requires eight final prompt tokens')
            with torch.no_grad():
                embeddings = self.embed_tokens(input_ids)
                return torch.stack([embeddings[i, end - k + 1:end + 1].flatten()
                                    for i, end in enumerate(answer_positions.tolist())])
        with torch.no_grad():
            context = self.embed_tokens(input_ids)
            positions = torch.arange(input_ids.shape[1], device=input_ids.device).unsqueeze(0)
            mask = create_causal_mask(config=self.config, inputs_embeds=context,
                attention_mask=attention_mask, past_key_values=None, position_ids=positions)
            context = self.prelude(context, mask, positions, self.rotary_emb(context, positions))
        return _answer_readout(context, answer_positions.tolist())

    def forward(
        self,
        input_ids: Tensor,
        attention_mask: Tensor | None = None,
        *,
        num_loops: int = 1,
        position_ids: Tensor | None = None,
        answer_positions: Tensor | None = None,
        labels: Tensor | None = None,
        allowed_token_ids: Tensor | None = None,
        return_hidden_states: bool = False,
        logits_mode: Literal["answer", "all"] = "answer",
        stop_policy: Callable[[int, Tensor, Tensor], bool] | None = None,
        completion_threshold: float | None = None,
        readout_token_ids: list[int] | None = None,
        loop_callback: Callable[[int, Tensor, Tensor | None, Tensor | None], None] | None = None,
    ) -> RecurrentOutput:
        """Run complete prompt states through shared recurrence, without KV cache.

        Inputs/mask: [B,S]; positions: [1,S] or [B,S], default arange(S),
        matching ordinary Qwen forward (including padding). Answer positions:
        [B], default last unmasked token. Labels: [B,T] *token IDs*, one
        target per loop, with an explicit allowed answer set [A]. No CE loss
        or causal label shifting is performed. Pass repeated final labels
        explicitly if final-answer margins, rather than step margins, are wanted.
        Optional stop_policy(t, h_t, answer_logits_t) is inference-only, batch 1,
        and may return True to end after loop t. It never sees labels.
        completion_threshold enables the checkpoint's hidden-state-only head,
        independent of task depth, under the num_loops safety cap (batch 1).
        fixed_prompt keeps only the final unmasked position writable; returned
        full-shaped states contain fixed first-pass prefix outputs for diagnostics.
        loop_callback receives one-based loop, C logits [1,A], optional stop
        logits [1], and optional R logits [1,26] after each readout, including
        the stopping loop. Detached copies protect inference from observer edits.
        It is inference-only; its return value cannot choose when to stop.
        """
        if type(num_loops) is not int or num_loops < 1:
            raise ValueError("num_loops must be a positive integer")
        if loop_callback is not None and (self.training or torch.is_grad_enabled() or
                input_ids.shape[0] != 1 or labels is not None or logits_mode != "answer"):
            raise ValueError("Loop observation requires eval/no_grad, batch 1, no labels, and answer logits")
        if stop_policy is not None and (self.training or torch.is_grad_enabled() or input_ids.shape[0] != 1 or labels is not None or logits_mode != "answer"):
            raise ValueError("Stopping requires eval/no_grad, batch 1, no labels, and answer logits")
        if completion_threshold is not None:
            if (not math.isfinite(completion_threshold) or not 0 < completion_threshold < 1 or
                    self.completion_head is None or stop_policy is not None):
                raise ValueError("Completion stopping requires a head, threshold in (0,1), and no other stop policy")
            if self.training or torch.is_grad_enabled() or input_ids.shape[0] != 1 or labels is not None or logits_mode != "answer":
                raise ValueError("Completion stopping requires eval/no_grad, batch 1, no labels, and answer logits")
        if logits_mode not in ("answer", "all"):
            raise ValueError("logits_mode must be 'answer' or 'all'")
        if input_ids.ndim != 2 or input_ids.dtype != torch.long or 0 in input_ids.shape:
            raise ValueError("input_ids must be a nonempty torch.long [batch, sequence] tensor")
        if input_ids.device != self.embed_tokens.weight.device:
            raise ValueError("Input IDs and model must share a device")
        if (input_ids < 0).any() or (input_ids >= self.config.vocab_size).any():
            raise ValueError("Input IDs must be within the model vocabulary")
        batch, sequence = input_ids.shape
        device = input_ids.device
        if attention_mask is None:
            attention_mask = torch.ones_like(input_ids)
        if attention_mask.shape != input_ids.shape or attention_mask.device != device:
            raise ValueError("attention_mask must match input_ids shape and device")
        if not ((attention_mask == 0) | (attention_mask == 1)).all():
            raise ValueError("attention_mask must contain only 0 and 1")
        if not attention_mask.bool().any(dim=-1).all():
            raise ValueError("Each example must have at least one unmasked token")
        if position_ids is None:
            position_ids = torch.arange(sequence, device=device).unsqueeze(0)
        if (
            position_ids.shape not in ((1, sequence), (batch, sequence))
            or position_ids.dtype != torch.long
            or position_ids.device != device
            or (position_ids < 0).any()
        ):
            raise ValueError("position_ids must be nonnegative long [1,S] or [B,S] on the input device")
        if answer_positions is None:
            indices = torch.arange(sequence, device=device).expand(batch, -1)
            answer_positions = indices.masked_fill(~attention_mask.bool(), -1).max(-1).values
        if (
            answer_positions.shape != (batch,)
            or answer_positions.dtype != torch.long
            or answer_positions.device != device
            or (answer_positions < 0).any()
            or (answer_positions >= sequence).any()
        ):
            raise ValueError("answer_positions must be long [batch] valid sequence indices on the input device")
        rows = torch.arange(batch, device=device)
        if not attention_mask[rows, answer_positions].bool().all():
            raise ValueError("Answer positions cannot point to padding")
        if self.recurrence_mode == "fixed_prompt":
            last = torch.arange(sequence, device=device).expand(batch, -1).masked_fill(~attention_mask.bool(), -1).max(-1).values
            if not torch.equal(answer_positions, last):
                raise ValueError("fixed_prompt requires the final unmasked position as its working state")
        readout_positions = answer_positions.tolist()
        if (labels is None) != (allowed_token_ids is None):
            raise ValueError("Provide both labels and allowed_token_ids to compute margins")
        if labels is not None and labels.shape != (batch, num_loops):
            raise ValueError("labels must have shape [batch, num_loops]")

        controller_initial = None
        if self.router is not None:
            if logits_mode != "answer":
                raise ValueError("Isolated execution supports answer-position logits only")
            # Controller prompt encoding is isolated even for full-model training.
            controller_initial = self.controller_prompt_state(input_ids, attention_mask, answer_positions)
            input_ids, attention_mask, position_ids = self.router(input_ids, attention_mask)
        hidden = self.embed_tokens(input_ids)
        causal_mask = create_causal_mask(
            config=self.config,
            inputs_embeds=hidden,
            attention_mask=attention_mask,
            past_key_values=None,
            position_ids=position_ids,
        )
        position_embeddings = self.rotary_emb(hidden, position_ids)
        block_args = (causal_mask, position_ids, position_embeddings)
        hidden = self.prelude(hidden, *block_args)
        initial_hidden = hidden if return_hidden_states else None
        bridge_initial = _answer_readout(hidden, readout_positions)
        controller_memory = (self.completion_head.initialize(controller_initial)
                             if isinstance(self.completion_head, RecurrentController) else None)
        recurrent_memory = coda_memory = None
        if self.recurrence_mode == "fixed_prompt":
            write_mask = (torch.arange(sequence, device=device)[None, :] == answer_positions[:, None]).unsqueeze(-1)
            recurrent_memory, coda_memory = PromptMemory(write_mask), PromptMemory(write_mask)
        states, logits, margins, stop_logits, state_logits = [], [], [], [], []
        for loop in range(num_loops):
            if loop and self.bridge is not None:
                working = self.bridge(_answer_readout(hidden, readout_positions), bridge_initial)
                hidden = torch.where(write_mask, working[:, None, :], hidden)
            hidden = (self.recurrent(hidden, *block_args, memory=recurrent_memory)
                      if recurrent_memory is not None else self.recurrent(hidden, *block_args))
            working = _answer_readout(hidden, readout_positions)
            if isinstance(self.completion_head, RecurrentController):
                stop, controller_memory = self.completion_head.advance(working, controller_memory)
                stop_logits.append(stop)
            elif self.completion_head is not None:
                stop_logits.append(self.completion_head(working))
            if self.state_head is not None:
                state_logits.append(self.state_head(working))
            if return_hidden_states:
                states.append(hidden)
            # C stays differentiable even when frozen: readout loss must reach R.
            decoded = self.norm(self.coda(hidden, *block_args, memory=coda_memory)
                                if coda_memory is not None else self.coda(hidden, *block_args))
            readout = decoded if logits_mode == "all" else _answer_readout(decoded, readout_positions)
            if readout_token_ids is None:
                scores = self.lm_head(readout)
            else:
                if labels is not None:
                    raise ValueError("Use full-vocabulary readout when requesting token-ID margins")
                ids = torch.tensor(readout_token_ids, device=device)
                scores = nn.functional.linear(readout, self.lm_head.weight.index_select(0, ids))
            logits.append(scores)
            if loop_callback is not None:
                loop_callback(loop + 1, scores.detach().clone(),
                              stop_logits[-1].detach().clone() if stop_logits else None,
                              state_logits[-1].detach().clone() if state_logits else None)
            if labels is not None:
                answer_scores = _answer_readout(scores, readout_positions) if logits_mode == "all" else scores
                margins.append(answer_margin(answer_scores, labels[:, loop], allowed_token_ids))
            if stop_policy is not None and stop_policy(loop + 1, hidden, scores):
                break
            if completion_threshold is not None:
                probability = stop_logits[-1].sigmoid().item()
                if not math.isfinite(probability):
                    raise FloatingPointError("Nonfinite completion probability")
                if probability >= completion_threshold:
                    break
        return RecurrentOutput(
            loop_logits=tuple(logits),
            state_logits=torch.stack(state_logits, dim=1) if state_logits else None,
            hidden_states=tuple(states) if return_hidden_states else None,
            initial_hidden_state=initial_hidden,
            margins=torch.stack(margins, dim=1) if margins else None,
            stop_logits=torch.stack(stop_logits, dim=1) if stop_logits else None,
        )