modeling_recurrent_qwen.py

"""Self-contained forced-depth recurrent Qwen loader for the Paper One release.

This module deliberately contains no halting head or adaptive-depth path. Every
forward call uses an externally supplied ``max_loops`` and, optionally, an
externally supplied per-row ``loop_selection`` bounded by that maximum.
"""

from __future__ import annotations

import inspect
from dataclasses import dataclass
from pathlib import Path
from typing import Any, Optional

import torch
import torch.nn.functional as F
from safetensors.torch import load_file
from torch import nn
from transformers import AutoModelForCausalLM, PreTrainedModel
from transformers.generation import GenerationMixin
from transformers.modeling_outputs import CausalLMOutputWithPast
from transformers.utils.hub import cached_file

from .configuration_recurrent_qwen import RecurrentQwenConfig


@dataclass
class RecurrentCausalLMOutput(CausalLMOutputWithPast):
    """Causal-LM output with optional loop-indexed logits."""

    loop_logits: Optional[torch.FloatTensor] = None
    selected_loop_counts: Optional[torch.LongTensor] = None


class IdentityGatedBridge(nn.Module):
    """Split re-entry bridge used by the frozen Paper One checkpoints."""

    def __init__(self, hidden_size: int) -> None:
        super().__init__()
        self.hidden_size = int(hidden_size)
        self.prelude_norm = nn.LayerNorm(hidden_size)
        self.prelude_proj = nn.Linear(hidden_size, hidden_size, bias=False)
        self.state_proj = nn.Linear(hidden_size, hidden_size, bias=True)
        self.bridge_gate = nn.Parameter(torch.tensor(1.0))
        with torch.no_grad():
            self.prelude_proj.weight.zero_()
            self.state_proj.weight.copy_(torch.eye(hidden_size))
            self.state_proj.bias.zero_()

    def forward(self, state: torch.Tensor, prelude: torch.Tensor) -> torch.Tensor:
        input_dtype = state.dtype
        work_dtype = self.state_proj.weight.dtype
        work = state.to(dtype=work_dtype)
        normalized_prelude = self.prelude_norm(prelude.to(dtype=work_dtype))
        translated = self.prelude_proj(normalized_prelude) + self.state_proj(work)
        gate = self.bridge_gate.to(device=state.device, dtype=work_dtype)
        return (work + gate * (translated - work)).to(dtype=input_dtype)


class LoRALinear(nn.Module):
    """Inference-only rank-decomposition wrapper matching the training keys."""

    def __init__(self, base: nn.Linear, *, rank: int, alpha: int) -> None:
        super().__init__()
        self.base = base
        self.rank = int(rank)
        self.alpha = int(alpha)
        self.scaling = float(alpha) / float(rank)
        self.lora_a = nn.Linear(base.in_features, rank, bias=False, dtype=torch.float32)
        self.lora_b = nn.Linear(rank, base.out_features, bias=False, dtype=torch.float32)
        with torch.no_grad():
            nn.init.kaiming_uniform_(self.lora_a.weight, a=5**0.5)
            self.lora_b.weight.zero_()
        for parameter in self.base.parameters():
            parameter.requires_grad_(False)

    def forward(self, inputs: torch.Tensor) -> torch.Tensor:
        base_output = self.base(inputs)
        adapter = self.lora_b(self.lora_a(inputs.float())) * self.scaling
        return base_output + adapter.to(dtype=base_output.dtype)


def _replace_lora_targets(
    module: nn.Module,
    target_names: set[str],
    *,
    rank: int,
    alpha: int,
) -> int:
    replaced = 0
    for child_name, child in list(module.named_children()):
        if isinstance(child, LoRALinear):
            continue
        if child_name in target_names and isinstance(child, nn.Linear):
            setattr(module, child_name, LoRALinear(child, rank=rank, alpha=alpha))
            replaced += 1
        else:
            replaced += _replace_lora_targets(child, target_names, rank=rank, alpha=alpha)
    return replaced


class RecurrentQwenForCausalLM(PreTrainedModel, GenerationMixin):
    """Qwen causal LM with a forced, weight-tied middle-block recurrence."""

    config_class = RecurrentQwenConfig
    base_model_prefix = "backbone"
    main_input_name = "input_ids"

    def __init__(self, config: RecurrentQwenConfig, backbone: nn.Module) -> None:
        super().__init__(config)
        self.backbone = backbone
        if not hasattr(backbone, "model") or not hasattr(backbone.model, "layers"):
            raise TypeError("Expected a Qwen-style causal LM with .model.layers")
        if int(config.recurrent_end) >= len(self.qwen.layers):
            raise ValueError("recurrent_end must leave at least one coda layer")
        hidden_size = int(getattr(backbone.config, "hidden_size"))
        base_parameter = next(backbone.parameters())
        self.bridge = IdentityGatedBridge(hidden_size).to(
            device=base_parameter.device,
            dtype=base_parameter.dtype,
        )
        self.lora_module_count = 0
        if config.checkpoint_kind == "lora_adapter":
            targets = set(config.lora_target_modules)
            for layer_index in range(config.prelude_end, config.recurrent_end):
                self.lora_module_count += _replace_lora_targets(
                    self.qwen.layers[layer_index],
                    targets,
                    rank=config.lora_rank,
                    alpha=config.lora_alpha,
                )
            if self.lora_module_count != 84:
                raise RuntimeError(f"Expected 84 recurrent LoRA modules, got {self.lora_module_count}")

    @property
    def qwen(self) -> nn.Module:
        return self.backbone.model

    @property
    def lm_head(self) -> nn.Module:
        return self.backbone.lm_head

    def get_input_embeddings(self) -> nn.Module:
        return self.qwen.embed_tokens

    def set_input_embeddings(self, value: nn.Module) -> None:
        self.qwen.embed_tokens = value

    def get_output_embeddings(self) -> nn.Module:
        return self.lm_head

    def set_output_embeddings(self, value: nn.Module) -> None:
        self.backbone.lm_head = value

    @classmethod
    def from_pretrained(
        cls,
        pretrained_model_name_or_path: str | Path,
        *model_args: Any,
        config: RecurrentQwenConfig | None = None,
        **kwargs: Any,
    ) -> "RecurrentQwenForCausalLM":
        """Load the pinned base model, construct the surgery, then apply the delta."""

        if model_args:
            raise TypeError("Positional model arguments are not supported by this release loader")
        token = kwargs.pop("token", None)
        revision = kwargs.pop("revision", "main")
        cache_dir = kwargs.pop("cache_dir", None)
        local_files_only = bool(kwargs.pop("local_files_only", False))
        kwargs.pop("trust_remote_code", None)
        kwargs.pop("_from_auto", None)
        kwargs.pop("_fast_init", None)
        kwargs.pop("state_dict", None)
        kwargs.pop("weights_only", None)
        kwargs.pop("adapter_kwargs", None)

        if config is None:
            config = RecurrentQwenConfig.from_pretrained(
                pretrained_model_name_or_path,
                revision=revision,
                token=token,
                cache_dir=cache_dir,
                local_files_only=local_files_only,
            )

        base_keys = {
            "attn_implementation",
            "device_map",
            "dtype",
            "torch_dtype",
            "low_cpu_mem_usage",
            "max_memory",
            "offload_folder",
            "offload_state_dict",
        }
        base_kwargs = {key: kwargs.pop(key) for key in list(kwargs) if key in base_keys}
        if kwargs:
            unknown = ", ".join(sorted(kwargs))
            raise TypeError(f"Unsupported loader keyword(s): {unknown}")
        base_kwargs.update(
            {
                "revision": config.base_model_revision,
                "token": token,
                "cache_dir": cache_dir,
                "local_files_only": local_files_only,
                "trust_remote_code": False,
            }
        )
        backbone = AutoModelForCausalLM.from_pretrained(
            config.base_model_name_or_path,
            **{key: value for key, value in base_kwargs.items() if value is not None},
        )
        model = cls(config, backbone)

        source = Path(pretrained_model_name_or_path)
        if source.is_dir():
            delta_path = source / config.delta_filename
        else:
            resolved = cached_file(
                str(pretrained_model_name_or_path),
                config.delta_filename,
                revision=revision,
                token=token,
                cache_dir=cache_dir,
                local_files_only=local_files_only,
            )
            if resolved is None:
                raise FileNotFoundError(config.delta_filename)
            delta_path = Path(resolved)
        delta = load_file(str(delta_path), device="cpu")
        current = model.state_dict()
        absent = sorted(set(delta) - set(current))
        mismatched = {
            key: {"delta": tuple(value.shape), "model": tuple(current[key].shape)}
            for key, value in delta.items()
            if key in current and value.shape != current[key].shape
        }
        if absent or mismatched:
            raise RuntimeError(f"Release delta is incompatible: absent={absent}, mismatched={mismatched}")
        model.load_state_dict(delta, strict=False)
        model._release_load_receipt = {
            "delta_file": str(delta_path),
            "tensor_count": len(delta),
            "total_parameters": sum(int(tensor.numel()) for tensor in delta.values()),
        }
        model.eval()
        return model

    def forward(
        self,
        input_ids: Optional[torch.LongTensor] = None,
        attention_mask: Optional[torch.Tensor] = None,
        position_ids: Optional[torch.LongTensor] = None,
        inputs_embeds: Optional[torch.Tensor] = None,
        labels: Optional[torch.LongTensor] = None,
        max_loops: int = 1,
        loop_selection: int | torch.LongTensor | None = None,
        return_loop_logits: bool = False,
        use_cache: Optional[bool] = None,
        output_attentions: Optional[bool] = None,
        output_hidden_states: Optional[bool] = None,
        return_dict: Optional[bool] = None,
        logits_to_keep: int | torch.Tensor = 0,
        **kwargs: Any,
    ) -> RecurrentCausalLMOutput | tuple[Any, ...]:
        if kwargs:
            unknown = ", ".join(sorted(kwargs))
            raise TypeError(f"Unsupported forward keyword(s): {unknown}")
        if int(max_loops) < 1:
            raise ValueError("max_loops must be at least 1")
        if use_cache:
            raise ValueError("KV caching is disabled for the forced-depth recurrent release")
        if output_attentions or output_hidden_states:
            raise ValueError("Attention/hidden-state collection is not exposed by the release loader")
        if input_ids is not None and inputs_embeds is not None:
            raise ValueError("Specify input_ids or inputs_embeds, not both")
        if inputs_embeds is None:
            if input_ids is None:
                raise ValueError("input_ids or inputs_embeds is required")
            inputs_embeds = self.qwen.embed_tokens(input_ids)

        batch_size, sequence_length = inputs_embeds.shape[:2]
        device = inputs_embeds.device
        if attention_mask is not None:
            attention_mask = attention_mask.to(device=device)
        if position_ids is None:
            position_ids = torch.arange(sequence_length, device=device).unsqueeze(0).expand(batch_size, -1)
        else:
            position_ids = position_ids.to(device=device)
        cache_position = torch.arange(sequence_length, device=device)
        causal_mask = self._causal_mask(attention_mask, inputs_embeds, cache_position)
        position_embeddings = self._rotary_embeddings(inputs_embeds, position_ids)

        hidden = self._run_layers(
            0,
            self.config.prelude_end,
            inputs_embeds,
            causal_mask,
            position_ids,
            cache_position,
            position_embeddings,
        )
        prelude = hidden
        recurrent_state = hidden
        logits_by_loop: list[torch.Tensor] = []
        for loop_index in range(int(max_loops)):
            loop_input = recurrent_state if loop_index == 0 else self.bridge(recurrent_state, prelude)
            recurrent_state = self._run_layers(
                self.config.prelude_end,
                self.config.recurrent_end,
                loop_input,
                causal_mask,
                position_ids,
                cache_position,
                position_embeddings,
            )
            coda = self._run_layers(
                self.config.recurrent_end,
                len(self.qwen.layers),
                recurrent_state,
                causal_mask,
                position_ids,
                cache_position,
                position_embeddings,
            )
            normed = self.qwen.norm(coda)
            logits_by_loop.append(self.lm_head(self._slice_for_logits(normed, logits_to_keep)))

        loop_logits = torch.stack(logits_by_loop, dim=1)
        selected_counts = self._resolve_loop_selection(loop_selection, batch_size, int(max_loops), device)
        batch_indices = torch.arange(batch_size, device=device)
        logits = loop_logits[batch_indices, selected_counts - 1]
        loss = None
        if labels is not None:
            if not (isinstance(logits_to_keep, int) and logits_to_keep == 0):
                raise ValueError("labels require logits_to_keep=0")
            shifted_logits = logits[:, :-1, :].contiguous().float()
            shifted_labels = labels[:, 1:].contiguous().to(device=device)
            loss = F.cross_entropy(
                shifted_logits.view(-1, shifted_logits.shape[-1]),
                shifted_labels.view(-1),
                ignore_index=-100,
            )

        output = RecurrentCausalLMOutput(
            loss=loss,
            logits=logits,
            past_key_values=None,
            hidden_states=None,
            attentions=None,
            loop_logits=loop_logits if return_loop_logits else None,
            selected_loop_counts=selected_counts,
        )
        return output if return_dict is not False else output.to_tuple()

    def prepare_inputs_for_generation(
        self,
        input_ids: torch.LongTensor,
        attention_mask: Optional[torch.Tensor] = None,
        **kwargs: Any,
    ) -> dict[str, Any]:
        return {
            "input_ids": input_ids,
            "attention_mask": attention_mask,
            "max_loops": int(kwargs.get("max_loops", 1)),
            "loop_selection": kwargs.get("loop_selection"),
            "use_cache": False,
        }

    @staticmethod
    def _resolve_loop_selection(
        selection: int | torch.LongTensor | None,
        batch_size: int,
        max_loops: int,
        device: torch.device,
    ) -> torch.LongTensor:
        if selection is None:
            counts = torch.full((batch_size,), max_loops, dtype=torch.long, device=device)
        elif isinstance(selection, int):
            counts = torch.full((batch_size,), int(selection), dtype=torch.long, device=device)
        else:
            counts = selection.to(device=device, dtype=torch.long).reshape(-1)
            if counts.numel() != batch_size:
                raise ValueError("loop_selection tensor must contain one value per batch row")
        if bool(((counts < 1) | (counts > max_loops)).any()):
            raise ValueError("loop_selection values must lie in [1, max_loops]")
        return counts

    def _causal_mask(
        self,
        attention_mask: Optional[torch.Tensor],
        inputs_embeds: torch.Tensor,
        cache_position: torch.Tensor,
    ) -> torch.Tensor | None:
        update = getattr(self.qwen, "_update_causal_mask", None)
        if update is not None:
            return self._call_supported(
                update,
                {
                    "attention_mask": attention_mask,
                    "input_tensor": inputs_embeds,
                    "inputs_embeds": inputs_embeds,
                    "cache_position": cache_position,
                    "past_key_values": None,
                    "output_attentions": False,
                },
            )
        batch_size, sequence_length = inputs_embeds.shape[:2]
        minimum = torch.finfo(inputs_embeds.dtype).min
        causal = torch.full(
            (sequence_length, sequence_length),
            minimum,
            dtype=inputs_embeds.dtype,
            device=inputs_embeds.device,
        ).triu(diagonal=1)
        causal = causal.unsqueeze(0).unsqueeze(0).expand(batch_size, 1, -1, -1)
        if attention_mask is not None:
            causal = causal.masked_fill(attention_mask[:, None, None, :].eq(0), minimum)
        return causal

    def _rotary_embeddings(
        self,
        hidden_states: torch.Tensor,
        position_ids: torch.Tensor,
    ) -> Any:
        rotary = getattr(self.qwen, "rotary_emb", None)
        if rotary is None:
            return None
        try:
            return rotary(hidden_states, position_ids)
        except TypeError:
            return None

    def _run_layers(
        self,
        start: int,
        end: int,
        hidden_states: torch.Tensor,
        causal_mask: Optional[torch.Tensor],
        position_ids: torch.Tensor,
        cache_position: torch.Tensor,
        position_embeddings: Any,
    ) -> torch.Tensor:
        for layer in self.qwen.layers[start:end]:
            parameters = inspect.signature(layer.forward).parameters
            cache_key = "past_key_values" if "past_key_values" in parameters else "past_key_value"
            outputs = layer(
                hidden_states,
                **self._filter_supported(
                    layer.forward,
                    {
                        "attention_mask": causal_mask,
                        "position_ids": position_ids,
                        cache_key: None,
                        "output_attentions": False,
                        "use_cache": False,
                        "cache_position": cache_position,
                        "position_embeddings": position_embeddings,
                    },
                ),
            )
            hidden_states = outputs[0] if isinstance(outputs, tuple) else outputs
        return hidden_states

    @staticmethod
    def _filter_supported(function: Any, kwargs: dict[str, Any]) -> dict[str, Any]:
        parameters = inspect.signature(function).parameters
        accepts_kwargs = any(value.kind == inspect.Parameter.VAR_KEYWORD for value in parameters.values())
        if accepts_kwargs:
            return {key: value for key, value in kwargs.items() if value is not None}
        return {
            key: value
            for key, value in kwargs.items()
            if key in parameters and (value is not None or key in {"use_cache", "output_attentions"})
        }

    @classmethod
    def _call_supported(cls, function: Any, kwargs: dict[str, Any]) -> Any:
        return function(**cls._filter_supported(function, kwargs))

    @staticmethod
    def _slice_for_logits(hidden_states: torch.Tensor, logits_to_keep: int | torch.Tensor) -> torch.Tensor:
        if isinstance(logits_to_keep, int):
            return hidden_states if logits_to_keep == 0 else hidden_states[:, -logits_to_keep:, :]
        return hidden_states[:, logits_to_keep, :]