models/recursive.py

"""Recursive (looped-depth) decoder-only transformer.

Three sections — entry / body / exit — where the body is a small stack of
shared weights applied ``num_loops`` times per forward. The loop update is
``h_{t+1} = A * h_t + B * e + R(h_t + e)`` with per-channel injection
gates A, B initialised to zero (so the loop starts as a weight-shared
transformer stack on h+e). Optional cross-loop expert diversity for shared
MoE routers via ``moe_diversity_factor``.
"""

from __future__ import annotations

from dataclasses import dataclass
from typing import List, Optional, Tuple, Dict, Any

import torch
import torch.nn as nn
import torch.nn.functional as F

from .lm_loss import (
    lm_cross_entropy_from_logits,
    token_superposition_attention_mask,
    token_superposition_embeddings,
)
from .baseline import (
    BaselineConfig,
    RMSNorm,
    TransformerBlock,
    MoELayer,
    combine_lm_and_aux_loss,
    init_moe_router_weights,
    count_parameters,
    model_summary,
)


@dataclass
class RecursiveConfig(BaselineConfig):
    # Auto-derived from entry + body + exit in __post_init__.
    num_layers: int = 0

    num_entry_layers: int = 2
    num_body_layers: int = 4
    num_exit_layers: int = 2
    num_loops: int = 4

    # Std of the random init for the per-channel A gate. 0 (default)
    # leaves the loop's residual mixing inert at step 0; small positive
    # values (e.g. 0.02) break that symmetry. B always starts at zero.
    body_gate_init_std: float = 0.0

    def __post_init__(self):
        super().__post_init__()
        if self.num_body_layers <= 0 or self.num_loops <= 0:
            raise ValueError(
                "num_body_layers and num_loops must both be > 0; set "
                "num_entry_layers / num_exit_layers to 0 if you want a "
                "purely-body model."
            )
        if self.body_gate_init_std < 0:
            raise ValueError("body_gate_init_std must be >= 0")
        self.num_layers = (
            self.num_entry_layers
            + self.num_body_layers
            + self.num_exit_layers
        )


class RecursiveBlock(nn.Module):
    """One iteration of the body loop: ``h_{t+1} = A*h + B*e + R(h+e)``.

    The body's transformer blocks are reused ``num_loops`` times, so MoE
    layers carry per-loop bias rows and the cross-loop diversity term.
    """

    def __init__(self, config: RecursiveConfig):
        super().__init__()
        self.blocks = nn.ModuleList([
            TransformerBlock(config, num_loops=config.num_loops)
            for _ in range(config.num_body_layers)
        ])
        if config.body_gate_init_std > 0:
            self.A = nn.Parameter(
                torch.randn(config.d_model) * config.body_gate_init_std
            )
        else:
            self.A = nn.Parameter(torch.zeros(config.d_model))
        self.B = nn.Parameter(torch.zeros(config.d_model))

    def forward(
        self,
        h: torch.Tensor,
        e: torch.Tensor,
        attention_mask: Optional[torch.Tensor] = None,
        is_causal: bool = True,
        loop_idx: int = 0,
    ) -> Tuple[torch.Tensor, torch.Tensor, List[Optional[torch.Tensor]]]:
        x = h + e
        aux_loss = torch.zeros((), device=x.device, dtype=x.dtype)
        topk_list: List[Optional[torch.Tensor]] = []
        for block in self.blocks:
            x, block_aux, block_topk = block(
                x,
                attention_mask=attention_mask,
                is_causal=is_causal,
                loop_idx=loop_idx,
            )
            aux_loss = aux_loss + block_aux
            topk_list.append(block_topk)
        h_next = self.A * h + self.B * e + x
        return h_next, aux_loss, topk_list


class RecursiveTransformer(nn.Module):
    def __init__(self, config: RecursiveConfig):
        super().__init__()
        self.config = config

        self.token_emb = nn.Embedding(config.vocab_size, config.d_model)

        self.entry = nn.ModuleList([
            TransformerBlock(config) for _ in range(config.num_entry_layers)
        ])
        self.body = RecursiveBlock(config)
        self.exit = nn.ModuleList([
            TransformerBlock(config) for _ in range(config.num_exit_layers)
        ])

        self.final_norm = RMSNorm(config.d_model, eps=config.norm_eps)
        self.lm_head = nn.Linear(config.d_model, config.vocab_size, bias=False)
        self.lm_head.weight = self.token_emb.weight

        self._init_weights()

    def _init_weights(self):
        # ``RecursiveBlock.A`` and ``.B`` stay at their zero init — they are
        # nn.Parameter (not Linear/Embedding) and so are skipped by this pass.
        for module in self.modules():
            if isinstance(module, nn.Linear):
                torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)
                if module.bias is not None:
                    torch.nn.init.zeros_(module.bias)
            elif isinstance(module, nn.Embedding):
                torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)
        init_moe_router_weights(self, self.config.router_init_std)

    def forward(
        self,
        input_ids: torch.Tensor,
        attention_mask: Optional[torch.Tensor] = None,
        labels: Optional[torch.Tensor] = None,
        is_causal: bool = True,
        token_superposition_bag_size: int = 1,
    ) -> Dict[str, Any]:
        x = token_superposition_embeddings(
            self.token_emb, input_ids, token_superposition_bag_size,
        )
        attention_mask = token_superposition_attention_mask(
            attention_mask, token_superposition_bag_size,
        )

        aux_loss = torch.zeros((), device=input_ids.device, dtype=x.dtype)
        topk_indices_list: List[Optional[torch.Tensor]] = []

        for layer in self.entry:
            x, layer_aux, layer_topk = layer(
                x, attention_mask=attention_mask, is_causal=is_causal
            )
            aux_loss = aux_loss + layer_aux
            topk_indices_list.append(layer_topk)
        e = x

        h = torch.zeros_like(e)
        for loop_idx in range(self.config.num_loops):
            h, block_aux, block_topks = self.body(
                h,
                e,
                attention_mask=attention_mask,
                is_causal=is_causal,
                loop_idx=loop_idx,
            )
            aux_loss = aux_loss + block_aux
            topk_indices_list.extend(block_topks)
        x = h

        for layer in self.exit:
            x, layer_aux, layer_topk = layer(
                x, attention_mask=attention_mask, is_causal=is_causal
            )
            aux_loss = aux_loss + layer_aux
            topk_indices_list.append(layer_topk)

        x = self.final_norm(x)
        logits = self.lm_head(x)

        lm_loss: Optional[torch.Tensor] = None
        if labels is not None:
            lm_loss = lm_cross_entropy_from_logits(
                logits,
                labels,
                token_superposition_bag_size=token_superposition_bag_size,
                ignore_index=-100,
            )
        loss = combine_lm_and_aux_loss(
            lm_loss,
            aux_loss if self.config.use_moe else None,
            self.training,
        )

        return {
            "logits": logits,
            "loss": loss,
            "lm_loss": lm_loss,
            "aux_loss": aux_loss if self.config.use_moe else None,
            "topk_indices": topk_indices_list if self.config.use_moe else None,
        }

    def update_router_biases(self, topk_indices_list: List[Optional[torch.Tensor]]) -> None:
        """Apply DeepSeek-style bias updates. Index layout:

            [entry_0..E-1,
             loop_0.b_0..B-1, loop_1.b_0..B-1, ..., loop_{L-1}.b_0..B-1,
             exit_0..X-1]

        Each body block is updated once per parameter set with all its
        loop iterations grouped, so the cross-loop diversity term sees
        them together.
        """
        if not self.config.use_moe:
            return

        n_entry = self.config.num_entry_layers
        n_body = self.config.num_body_layers
        n_loops = self.config.num_loops

        for i, layer in enumerate(self.entry):
            topk = topk_indices_list[i]
            if topk is not None and isinstance(layer.ffn, MoELayer):
                layer.ffn.update_bias(topk, loop_idx=0)

        body_offset = n_entry
        for r, block in enumerate(self.body.blocks):
            if not isinstance(block.ffn, MoELayer):
                continue
            topk_per_loop: List[torch.Tensor] = []
            valid = True
            for l in range(n_loops):
                idx = body_offset + l * n_body + r
                topk = topk_indices_list[idx]
                if topk is None:
                    valid = False
                    break
                topk_per_loop.append(topk)
            if valid:
                block.ffn.update_bias_per_loop(topk_per_loop)

        exit_offset = n_entry + n_loops * n_body
        for i, layer in enumerate(self.exit):
            topk = topk_indices_list[exit_offset + i]
            if topk is not None and isinstance(layer.ffn, MoELayer):
                layer.ffn.update_bias(topk, loop_idx=0)

    @torch.no_grad()
    def get_balance_stats(self) -> Dict[str, float]:
        """One entry per parameter set — body sub-blocks appear once each
        (not ``num_loops`` times)."""
        if not self.config.use_moe:
            return {}

        stats: Dict[str, float] = {}

        def _record(name: str, ffn: nn.Module) -> None:
            if hasattr(ffn, "bias"):
                bias = ffn.bias
                stats[f"{name}_bias_mean"] = bias.abs().mean().item()
                stats[f"{name}_bias_max"] = bias.abs().max().item()

        for idx, layer in enumerate(self.entry):
            _record(f"entry{idx}", layer.ffn)
        for idx, block in enumerate(self.body.blocks):
            _record(f"body{idx}", block.ffn)
        for idx, layer in enumerate(self.exit):
            _record(f"exit{idx}", layer.ffn)
        return stats

    @torch.no_grad()
    def generate(
        self,
        input_ids: torch.Tensor,
        max_new_tokens: int = 100,
        temperature: float = 1.0,
        top_k: Optional[int] = None,
        attention_mask: Optional[torch.Tensor] = None,
        eos_token_id: Optional[int] = None,
    ) -> torch.Tensor:
        self.train(False)
        batch_size = input_ids.size(0)

        for _ in range(max_new_tokens):
            outputs = self.forward(
                input_ids, attention_mask=attention_mask, is_causal=True,
            )
            logits = outputs["logits"][:, -1, :] / temperature

            if top_k is not None:
                v, _ = torch.topk(logits, min(top_k, logits.size(-1)))
                logits[logits < v[:, [-1]]] = -float("Inf")

            probs = F.softmax(logits, dim=-1)
            next_token = torch.multinomial(probs, num_samples=1)

            input_ids = torch.cat([input_ids, next_token], dim=-1)

            if attention_mask is not None:
                attention_mask = torch.cat([
                    attention_mask,
                    torch.ones(
                        (batch_size, 1),
                        device=attention_mask.device,
                        dtype=attention_mask.dtype,
                    ),
                ], dim=-1)

            if eos_token_id is not None and (next_token == eos_token_id).all():
                break

        return input_ids


__all__ = [
    "RecursiveConfig",
    "RecursiveBlock",
    "RecursiveTransformer",
    "count_parameters",
    "model_summary",
]