models/recursive_reasoning/trm.py

from typing import Tuple, Dict, Optional
from dataclasses import dataclass
import math
import torch
import torch.nn.functional as F
from torch import nn
from pydantic import BaseModel
from models.common import trunc_normal_init_
from models.layers import RotaryEmbedding, CastedEmbedding, CastedLinear
from models.sparse_embedding import CastedSparseEmbedding

from models.transformer import FixedPointTransformer
from models.loop_attnres import LoopAttn, TrajectoryAttn, CrossRouter, CarrySourceAttn, CrossLite, DecayTrajAttn
from models.config import ReasoningModelConfig

IGNORE_LABEL_ID = -100

@dataclass
class TinyRecursiveReasoningModel_ACTV1InnerCarry:
    z_H: torch.Tensor
    z_L: torch.Tensor


@dataclass
class TinyRecursiveReasoningModel_ACTV1Carry:
    inner_carry: TinyRecursiveReasoningModel_ACTV1InnerCarry
    
    steps: torch.Tensor
    halted: torch.Tensor
    
    current_data: Dict[str, torch.Tensor]


class TinyRecursiveReasoningModel_ACTV1_Inner(nn.Module):
    def __init__(self, config: ReasoningModelConfig) -> None:
        super().__init__()
        self.config = config
        self.forward_dtype = getattr(torch, self.config.forward_dtype)

        # I/O

        self.embed_scale = math.sqrt(self.config.hidden_size)
        embed_init_std = 1.0 / self.embed_scale

        self.embed_tokens = CastedEmbedding(self.config.vocab_size, self.config.hidden_size, init_std=embed_init_std, cast_to=self.forward_dtype)
        self.lm_head      = CastedLinear(self.config.hidden_size, self.config.vocab_size, bias=False)
        self.q_head       = CastedLinear(self.config.hidden_size, 2, bias=True)

        self.puzzle_emb_len = -(self.config.puzzle_emb_ndim // -self.config.hidden_size)  if self.config.puzzle_emb_len == 0 else self.config.puzzle_emb_len  # ceil div
        if self.config.puzzle_emb_ndim > 0:
            # Zero init puzzle embeddings
            self.puzzle_emb = CastedSparseEmbedding(self.config.num_puzzle_identifiers, self.config.puzzle_emb_ndim,
                                                    batch_size=self.config.batch_size, init_std=0, cast_to=self.forward_dtype)

        # LM Blocks
        if self.config.pos_encodings == "rope":
            self.rotary_emb = RotaryEmbedding(dim=self.config.hidden_size // self.config.num_heads,
                                              max_position_embeddings=self.config.seq_len + self.puzzle_emb_len,
                                              base=self.config.rope_theta)
        elif self.config.pos_encodings == "learned":
            self.embed_pos = CastedEmbedding(self.config.seq_len + self.puzzle_emb_len, self.config.hidden_size, init_std=embed_init_std, cast_to=self.forward_dtype)
        else:
            pass

        # Reasoning Layers
        self.L_level = FixedPointTransformer(config=config, n_layers=self.config.L_layers)

        # Loop-AttnRes: replaces the carry-last connection between L-steps in the
        # with-grad cycle with attention over the loop-state history. zero-init
        # pseudo-query => uniform average over sources at init (AttnRes stabilizer).
        if self.config.loop_attnres:
            assert self.config.loop_attnres_mode == "state", "only 'state' mode wired for TRM loop-attn"
            grid = self.config.loop_attnres_grid
            if grid == "loop1d":
                self.loop_attn = LoopAttn(self.config.hidden_size, self.config.num_heads,
                                          self.config.loop_attnres_impl,
                                          recency_decay_init=self.config.loop_attnres_recency_init)
            elif grid == "cross":
                self.loop_attn = CrossRouter(
                    self.config.hidden_size, self.config.L_layers,
                    window=self.config.loop_attnres_window,
                    recency_init=self.config.loop_attnres_recency_init,
                    temp=self.config.loop_attnres_temp)
            elif grid == "crosslite":
                self.loop_attn = CrossLite(
                    self.config.hidden_size, self.config.L_layers,
                    window=self.config.loop_attnres_window,
                    recency_init=self.config.loop_attnres_recency_init,
                    temp=self.config.loop_attnres_temp)
            elif grid == "carrysource":
                self.loop_attn = CarrySourceAttn(
                    self.config.hidden_size, self.config.L_layers,
                    window=self.config.loop_attnres_window,
                    carry_bias_init=self.config.loop_attnres_carry_bias_init,
                    recency_init=self.config.loop_attnres_recency_init,
                    temp=self.config.loop_attnres_temp,
                    dkey=self.config.loop_attnres_dkey)
            elif grid == "ema":
                self.loop_attn = DecayTrajAttn(
                    self.config.hidden_size,
                    heads=self.config.loop_attnres_ema_heads,
                    beta_init=self.config.loop_attnres_beta_init,
                    temp=self.config.loop_attnres_temp,
                    content=self.config.loop_attnres_content,
                    score_norm=self.config.loop_attnres_score_norm,
                    score_noise=self.config.loop_attnres_score_noise,
                    mode_sigma=self.config.loop_attnres_mode_sigma)
            elif grid in ("flat", "struct"):
                self.loop_attn = TrajectoryAttn(
                    self.config.hidden_size, self.config.L_layers,
                    structured=(grid == "struct"),
                    window=self.config.loop_attnres_window,
                    recency_decay_init=self.config.loop_attnres_recency_init,
                    temp=self.config.loop_attnres_temp)
            else:
                raise ValueError(f"Unknown loop_attnres_grid {grid!r}")
        else:
            self.loop_attn = None

        # Initial states
        self.H_init = nn.Buffer(trunc_normal_init_(torch.empty(self.config.hidden_size, dtype=self.forward_dtype), std=1), persistent=True)
        self.L_init = nn.Buffer(trunc_normal_init_(torch.empty(self.config.hidden_size, dtype=self.forward_dtype), std=1), persistent=True)

        # Q head special init
        # Init Q to (almost) zero for faster learning during bootstrapping
        with torch.no_grad():
            self.q_head.weight.zero_()
            self.q_head.bias.fill_(-5)  # type: ignore

    def _input_embeddings(self, input: torch.Tensor, puzzle_identifiers: torch.Tensor):
        # Token embedding
        embedding = self.embed_tokens(input.to(torch.int32))

        # Puzzle embeddings
        if self.config.puzzle_emb_ndim > 0:
            puzzle_embedding = self.puzzle_emb(puzzle_identifiers)
            
            pad_count = self.puzzle_emb_len * self.config.hidden_size - puzzle_embedding.shape[-1]
            if pad_count > 0:
                puzzle_embedding = F.pad(puzzle_embedding, (0, pad_count))

            embedding = torch.cat((puzzle_embedding.view(-1, self.puzzle_emb_len, self.config.hidden_size), embedding), dim=-2)

        # Position embeddings
        if self.config.pos_encodings == "learned":
            # scale by 1/sqrt(2) to maintain forward variance
            embedding = 0.707106781 * (embedding + self.embed_pos.embedding_weight.to(self.forward_dtype))

        # Scale
        return self.embed_scale * embedding

    def empty_carry(self, batch_size: int):
        return TinyRecursiveReasoningModel_ACTV1InnerCarry(
            z_H=torch.empty(batch_size, self.config.seq_len + self.puzzle_emb_len, self.config.hidden_size, dtype=self.forward_dtype),
            z_L=torch.empty(batch_size, self.config.seq_len + self.puzzle_emb_len, self.config.hidden_size, dtype=self.forward_dtype),
        )
        
    def reset_carry(self, reset_flag: torch.Tensor, carry: TinyRecursiveReasoningModel_ACTV1InnerCarry):
        return TinyRecursiveReasoningModel_ACTV1InnerCarry(
            z_H=torch.where(reset_flag.view(-1, 1, 1), self.H_init, carry.z_H),
            z_L=torch.where(reset_flag.view(-1, 1, 1), self.L_init, carry.z_L),
        )

    def _l_cycle(self, z_L, z_H, input_embeddings, seq_info, n_steps):
        """One L-cycle of n_steps. Baseline: carry-last recurrence. Loop-AttnRes:
        the state fed to each step is a softmax-attention mixture over the full
        loop-state history {z_L_in, o_1, ..., o_t} (replaces the cross-loop
        residual; no sliding window — distant sources matter most)."""
        inj = z_H + input_embeddings
        if self.loop_attn is None:            # carry-last baseline
            for _ in range(n_steps):
                z_L = self.L_level(z_L, inj, **seq_info)
            return z_L
        if self.config.loop_attnres_grid == "loop1d":   # 1D loop-axis attention (LoopAttn)
            sources = [z_L]                   # distant anchor = state entering the cycle
            z_in = z_L
            for _ in range(n_steps):
                o = self.L_level(z_in, inj, **seq_info)
                sources.append(o)
                z_in = self.loop_attn(sources)
            return z_in
        if self.config.loop_attnres_grid == "ema":
            # Decay-fused trajectory attention as an exact O(1) recurrence: no window,
            # no history buffer — carry the softmax numerator/denominator instead.
            state = self.loop_attn.init_state(z_L)
            z_in = z_L
            for _ in range(n_steps):
                _, layer_states = self.L_level(z_in, inj, collect_layers=True, **seq_info)
                z_in, state = self.loop_attn(layer_states, state)
            return z_in
        if self.config.loop_attnres_grid in ("cross", "crosslite"):
            # Cross-shaped per-layer routing (CrossRouter / CrossLite): every layer input — the
            # within-block hops AND the cycle-boundary carry — is a routed read over
            # {anchor, last-L computation tail, the input source's own W-iteration
            # history}. Router l feeds layer l; its column is layer (l-1) mod L, the
            # layer that produces its input.
            L = len(self.L_level.layers)
            cols = [[] for _ in range(L)]     # cols[j] = layer-j outputs at past iterations, oldest..newest
            prev = []                         # trailing computation states from earlier iterations

            def route(l, cur_states):
                tail = (prev + cur_states)[-L:][::-1]        # newest first
                return self.loop_attn(l, tail, cols[(l - 1) % L][::-1], z_L)

            z_in = z_L
            for _ in range(n_steps):
                _, layer_states = self.L_level(z_in, inj, collect_layers=True, route_fn=route, **seq_info)
                z_in = route(0, layer_states)                # boundary carry for the next iteration
                W = self.config.loop_attnres_window
                for j, s in enumerate(layer_states):
                    cols[j] = (cols[j] + [s])[-W:]
                prev = (prev + layer_states)[-L:]
            return z_in
        W = self.config.loop_attnres_window
        if self.config.loop_attnres_grid == "carrysource":
            # Single-softmax simplex over {z_last (the carry), layer x iteration history}.
            # Replaces the residual with a convex trajectory read; the carry slot's large
            # bias makes it start = carry-last. z_last = the newest block output o.
            hist, z_in = [], z_L
            for _ in range(n_steps):
                o, layer_states = self.L_level(z_in, inj, collect_layers=True, **seq_info)
                hist.append(layer_states)
                if len(hist) > W:
                    hist = hist[-W:]
                z_in = self.loop_attn(o, hist)
            return z_in
        # 2D trajectory (flat / struct): attend over the {layers x iterations} grid,
        # with the iteration history capped to loop_attnres_window for stability.
        hist, z_in = [], z_L
        for _ in range(n_steps):
            _o, layer_states = self.L_level(z_in, inj, collect_layers=True, **seq_info)
            hist.append(layer_states)
            if len(hist) > W:
                hist = hist[-W:]
            z_in = self.loop_attn(z_L, hist)   # anchor = state entering the cycle
        return z_in

    def forward(self, carry: TinyRecursiveReasoningModel_ACTV1InnerCarry, batch: Dict[str, torch.Tensor]) -> Tuple[TinyRecursiveReasoningModel_ACTV1InnerCarry, torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]:
        seq_info = dict(
            cos_sin=self.rotary_emb() if hasattr(self, "rotary_emb") else None, 
            puzzle_emb_len=self.puzzle_emb_len
        )

        # Input encoding
        input_embeddings = self._input_embeddings(batch["inputs"], batch["puzzle_identifiers"])

        # Forward iterations
        z_H, z_L = carry.z_H, carry.z_L
        # Random-depth training (depthdrop): one warmup loop depth per forward, sampled
        # from the menu. Grad cycles stay at n_backwards_L (constant BPTT memory); the
        # graded segment learns to refine from variable-depth prefixes, which is what
        # eval-time depth scaling exercises. Eval keeps the configured/scanned depth.
        warmup_L, grad_L = self.config.L_cycles, self.config.n_backwards_L
        if self.training and self.config.loop_depthdrop:
            choices = [int(x) for x in self.config.loop_depthdrop.split(",") if x.strip()]
            warmup_L = choices[int(torch.randint(len(choices), (1,)).item())]
        if self.training and self.config.loop_depthdrop_grad:
            choices = [int(x) for x in self.config.loop_depthdrop_grad.split(",") if x.strip()]
            grad_L = choices[int(torch.randint(len(choices), (1,)).item())]
        # H_cycles-1 without grad
        with torch.no_grad():
            for _H_step in range(self.config.H_cycles-1):
                z_L = self._l_cycle(z_L, z_H, input_embeddings, seq_info, warmup_L)
                z_H = self.L_level(z_H, z_L, **seq_info)
        # 1 with grad
        z_L = self._l_cycle(z_L, z_H, input_embeddings, seq_info, grad_L)
        z_H = self.L_level(z_H, z_L, **seq_info)

        # LM Outputs
        new_carry = TinyRecursiveReasoningModel_ACTV1InnerCarry(z_H=z_H.detach(), z_L=z_L.detach())  # New carry no grad
        output = self.lm_head(z_H)[:, self.puzzle_emb_len:]
        
        if self.config.q_logit_detach_model is True:
            z_H = z_H.detach()
        
        if self.config.q_logit_from_puzzle_emb is True:
            q_logits = self.q_head(z_H[:, 0]).to(torch.float32) # Q-head; uses the first puzzle_emb position
        else:
            q_logits = self.q_head(z_H[:, self.puzzle_emb_len:].mean(dim=1))
        return new_carry, output, (q_logits[..., 0], q_logits[..., 1])


class TinyRecursiveReasoningModel_ACTV1(nn.Module):
    """ACT wrapper."""

    def __init__(self, config_dict: dict):
        super().__init__()
        self.config = ReasoningModelConfig(**config_dict)
        self.inner = TinyRecursiveReasoningModel_ACTV1_Inner(self.config)

    @property
    def puzzle_emb(self):
        return self.inner.puzzle_emb

    def initial_carry(self, batch: Dict[str, torch.Tensor]):
        batch_size = batch["inputs"].shape[0]

        return TinyRecursiveReasoningModel_ACTV1Carry(
            inner_carry=self.inner.empty_carry(batch_size),  # Empty is expected, it will be reseted in first pass as all sequences are halted.
            
            steps=torch.zeros((batch_size, ), dtype=torch.int32),
            halted=torch.ones((batch_size, ), dtype=torch.bool),  # Default to halted
            
            current_data={k: torch.empty_like(v) for k, v in batch.items()}
        )
        
    def forward(self, carry: TinyRecursiveReasoningModel_ACTV1Carry, batch: Dict[str, torch.Tensor]) -> Tuple[TinyRecursiveReasoningModel_ACTV1Carry, Dict[str, torch.Tensor]]:

        # Update data, carry (removing halted sequences)
        new_inner_carry = self.inner.reset_carry(carry.halted, carry.inner_carry)
        
        new_steps = torch.where(carry.halted, 0, carry.steps)

        new_current_data = {k: torch.where(carry.halted.view((-1, ) + (1, ) * (batch[k].ndim - 1)), batch[k], v) for k, v in carry.current_data.items()}

        # Forward inner model
        new_inner_carry, logits, (q_halt_logits, q_continue_logits) = self.inner(new_inner_carry, new_current_data)

        outputs = {
            "logits": logits,
            "q_halt_logits": q_halt_logits,
            "q_continue_logits": q_continue_logits
        }

        with torch.no_grad():
            # Step
            new_steps = new_steps + 1
            is_last_step = new_steps >= self.config.halt_max_steps
            
            halted = is_last_step

            # if training, and ACT is enabled
            if self.training and (self.config.halt_max_steps > 1):

                # Halt signal
                # NOTE: During evaluation, always use max steps, this is to guarantee the same halting steps inside a batch for batching purposes
                
                if self.config.no_ACT_continue:
                    halted = halted | (q_halt_logits > 0)
                else:
                    halted = halted | (q_halt_logits > q_continue_logits)

                # Exploration
                min_halt_steps = (torch.rand_like(q_halt_logits) < self.config.halt_exploration_prob) * torch.randint_like(new_steps, low=2, high=self.config.halt_max_steps + 1)
                halted = halted & (new_steps >= min_halt_steps)

                if not self.config.no_ACT_continue:
                    # Compute target Q
                    # NOTE: No replay buffer and target networks for computing target Q-value.
                    # As batch_size is large, there're many parallel envs.
                    # Similar concept as PQN https://arxiv.org/abs/2407.04811
                    _, _, (next_q_halt_logits, next_q_continue_logits), _, _ = self.inner(new_inner_carry, new_current_data)
                    outputs["target_q_continue"] = torch.sigmoid(torch.where(is_last_step, next_q_halt_logits, torch.maximum(next_q_halt_logits, next_q_continue_logits)))

        return TinyRecursiveReasoningModel_ACTV1Carry(new_inner_carry, new_steps, halted, new_current_data), outputs