modeling_quadorbit.py

"""Self-contained Hugging Face implementation of QuadOrbit."""

from __future__ import annotations

import math

import torch
from torch import nn
from torch.nn import functional as F
from transformers import PreTrainedModel
from transformers.generation import GenerationMixin
from transformers.modeling_outputs import CausalLMOutputWithPast

from .configuration_quadorbit import QuadOrbitConfig


def _inverse_softplus(value: torch.Tensor) -> torch.Tensor:
    return value + torch.log(-torch.expm1(-value))


def _rope_cache(head_dim: int, max_seq_len: int, theta: float):
    inv_freq = 1.0 / (
        theta ** (torch.arange(0, head_dim, 2, dtype=torch.float32) / head_dim)
    )
    positions = torch.arange(max_seq_len, dtype=torch.float32)
    freqs = torch.outer(positions, inv_freq)
    angles = torch.cat((freqs, freqs), dim=-1)
    return angles.cos(), angles.sin()


def _rotate_half(x: torch.Tensor) -> torch.Tensor:
    half = x.shape[-1] // 2
    return torch.cat((-x[..., half:], x[..., :half]), dim=-1)


def _apply_rope(
    x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor
) -> torch.Tensor:
    length = x.shape[-2]
    cos = cos[:length].to(device=x.device, dtype=x.dtype).view(1, 1, length, -1)
    sin = sin[:length].to(device=x.device, dtype=x.dtype).view(1, 1, length, -1)
    return x * cos + _rotate_half(x) * sin


def stable_quadorbit_recurrence(
    projected_real: torch.Tensor,
    projected_imag: torch.Tensor,
    gate_logits: torch.Tensor,
    log_decay: torch.Tensor,
    theta: torch.Tensor,
    beta_logit: torch.Tensor,
    beta_max: float,
) -> tuple[torch.Tensor, torch.Tensor]:
    """Portable FP32 implementation of the published bounded recurrence."""
    drive_real = torch.tanh(projected_real.float())
    drive_imag = torch.tanh(projected_imag.float())
    gates = torch.sigmoid(gate_logits.float())
    rho = torch.exp(-F.softplus(log_decay.float()))
    cos_theta = torch.cos(theta.float())
    sin_theta = torch.sin(theta.float())
    beta = float(beta_max) * torch.sigmoid(beta_logit.float())

    batch, length, width = drive_real.shape
    z_real = drive_real.new_zeros((batch, width))
    z_imag = drive_real.new_zeros((batch, width))
    real_rows: list[torch.Tensor] = []
    imag_rows: list[torch.Tensor] = []
    for position in range(length):
        old_real, old_imag = z_real, z_imag
        carry_real = rho * (cos_theta * old_real - sin_theta * old_imag)
        carry_imag = rho * (sin_theta * old_real + cos_theta * old_imag)
        square_real = old_real.square() - old_imag.square()
        square_imag = 2.0 * old_real * old_imag
        q_real = beta * square_real + drive_real[:, position]
        q_imag = beta * square_imag + drive_imag[:, position]
        inverse_radius = torch.rsqrt(1.0 + q_real.square() + q_imag.square())
        candidate_real = q_real * inverse_radius
        candidate_imag = q_imag * inverse_radius
        gate = gates[:, position]
        z_real = (1.0 - gate) * carry_real + gate * candidate_real
        z_imag = (1.0 - gate) * carry_imag + gate * candidate_imag
        real_rows.append(z_real)
        imag_rows.append(z_imag)

    if length == 0:
        empty = drive_real.new_empty((batch, 0, width))
        return empty, empty.clone()
    return torch.stack(real_rows, dim=1), torch.stack(imag_rows, dim=1)


class RMSNorm(nn.Module):
    def __init__(self, dim: int, eps: float):
        super().__init__()
        self.eps = eps
        self.weight = nn.Parameter(torch.ones(dim))

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        dtype = x.dtype
        normalized = x.float() * torch.rsqrt(
            x.float().square().mean(dim=-1, keepdim=True) + self.eps
        )
        return (normalized * self.weight.float()).to(dtype)


class QuadOrbitAttention(nn.Module):
    def __init__(self, config: QuadOrbitConfig):
        super().__init__()
        self.n_heads = config.n_heads
        self.n_kv_heads = config.n_kv_heads
        self.head_dim = config.head_dim
        self.n_rep = config.n_heads // config.n_kv_heads
        self.native_gqa = config.native_gqa
        self.q_proj = nn.Linear(
            config.d_model, config.n_heads * config.head_dim, bias=False
        )
        self.k_proj = nn.Linear(
            config.d_model, config.n_kv_heads * config.head_dim, bias=False
        )
        self.v_proj = nn.Linear(
            config.d_model, config.n_kv_heads * config.head_dim, bias=False
        )
        self.o_proj = nn.Linear(
            config.n_heads * config.head_dim, config.d_model, bias=False
        )
        self.qk_norm = config.qk_norm
        if self.qk_norm:
            self.q_norm = RMSNorm(config.head_dim, config.rms_eps)
            self.k_norm = RMSNorm(config.head_dim, config.rms_eps)

    def forward(
        self, x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor
    ) -> torch.Tensor:
        batch, length, _ = x.shape
        q = self.q_proj(x).view(batch, length, self.n_heads, self.head_dim)
        k = self.k_proj(x).view(batch, length, self.n_kv_heads, self.head_dim)
        v = self.v_proj(x).view(batch, length, self.n_kv_heads, self.head_dim)
        if self.qk_norm:
            q = self.q_norm(q)
            k = self.k_norm(k)
        q = _apply_rope(q.transpose(1, 2), cos, sin)
        k = _apply_rope(k.transpose(1, 2), cos, sin)
        v = v.transpose(1, 2)
        if self.n_rep > 1 and not self.native_gqa:
            k = k.repeat_interleave(self.n_rep, dim=1)
            v = v.repeat_interleave(self.n_rep, dim=1)
        output = F.scaled_dot_product_attention(
            q,
            k,
            v,
            is_causal=True,
            enable_gqa=(self.native_gqa and self.n_rep > 1),
        )
        output = output.transpose(1, 2).contiguous().view(batch, length, -1)
        return self.o_proj(output)


class SwiGLU(nn.Module):
    def __init__(self, config: QuadOrbitConfig):
        super().__init__()
        self.gate_proj = nn.Linear(config.d_model, config.d_ffn, bias=False)
        self.up_proj = nn.Linear(config.d_model, config.d_ffn, bias=False)
        self.down_proj = nn.Linear(config.d_ffn, config.d_model, bias=False)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.down_proj(F.silu(self.gate_proj(x)) * self.up_proj(x))


class StableComplexOrbit(nn.Module):
    def __init__(self, config: QuadOrbitConfig):
        super().__init__()
        width = config.orbit_width
        self.width = width
        self.beta_max = config.beta_max
        self.share_projections = config.share_orbit_projections
        if self.share_projections:
            self.c_proj = None
            self.gate_proj = None
            self.gate_bias = nn.Parameter(torch.empty(width))
            self.orbit_scale = nn.Parameter(torch.zeros(()))
            self.out_proj = None
        else:
            self.c_proj = nn.Linear(config.d_model, 2 * width, bias=False)
            self.gate_proj = nn.Linear(config.d_model, width, bias=True)
            self.register_parameter("gate_bias", None)
            self.register_parameter("orbit_scale", None)
            self.out_proj = nn.Linear(2 * width, config.d_model, bias=False)
        self.log_decay = nn.Parameter(torch.empty(width))
        self.theta = nn.Parameter(torch.empty(width))
        self.beta_logit = nn.Parameter(torch.empty(width))

    @torch.no_grad()
    def initialize_dynamics(self, config: QuadOrbitConfig) -> None:
        timescales = torch.logspace(
            math.log10(4.0),
            math.log10(float(config.max_seq_len)),
            self.width,
            dtype=torch.float32,
            device=self.log_decay.device,
        )
        gate_bias = torch.logit(timescales.reciprocal())
        if self.share_projections:
            self.gate_bias.copy_(gate_bias)
            self.orbit_scale.zero_()
        else:
            self.gate_proj.weight.zero_()
            self.gate_proj.bias.copy_(gate_bias)
        rho_target = torch.exp(-1.0 / (4.0 * timescales))
        self.log_decay.copy_(_inverse_softplus(-torch.log(rho_target)))
        self.theta.zero_()
        self.theta[int(0.75 * self.width) :].fill_(0.1)
        ratio = torch.full_like(self.beta_logit, config.beta_init / config.beta_max)
        self.beta_logit.copy_(torch.logit(ratio))

    def forward(
        self,
        x: torch.Tensor,
        shared_weights: tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None = None,
    ) -> torch.Tensor:
        if self.share_projections:
            if shared_weights is None:
                raise ValueError("shared QuadOrbit projections require Q, K, and O weights")
            q_weight, k_weight, o_weight = shared_weights
            projected = F.linear(x, q_weight[: 2 * self.width])
            gate_logits = F.linear(x, k_weight[: self.width]) + self.gate_bias
        else:
            projected = self.c_proj(x)
            gate_logits = self.gate_proj(x)
        projected_real, projected_imag = projected.chunk(2, dim=-1)
        z_real, z_imag = stable_quadorbit_recurrence(
            projected_real,
            projected_imag,
            gate_logits,
            self.log_decay,
            self.theta,
            self.beta_logit,
            self.beta_max,
        )
        packed_state = torch.cat((z_real, z_imag), dim=-1).to(x.dtype)
        if self.share_projections:
            return F.linear(
                packed_state, o_weight[:, : 2 * self.width]
            ) * torch.tanh(self.orbit_scale)
        return self.out_proj(packed_state)


class QuadOrbitBlock(nn.Module):
    def __init__(self, config: QuadOrbitConfig):
        super().__init__()
        self.share_orbit_projections = config.share_orbit_projections
        self.orbit_norm = (
            None
            if self.share_orbit_projections
            else RMSNorm(config.d_model, config.rms_eps)
        )
        self.orbit = StableComplexOrbit(config)
        self.attn_norm = RMSNorm(config.d_model, config.rms_eps)
        self.attn = QuadOrbitAttention(config)
        self.ffn_norm = RMSNorm(config.d_model, config.rms_eps)
        self.ffn = SwiGLU(config)

    def forward(
        self, x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor
    ) -> torch.Tensor:
        if self.share_orbit_projections:
            orbit_input = self.attn_norm(x)
            x = x + self.orbit(
                orbit_input,
                shared_weights=(
                    self.attn.q_proj.weight,
                    self.attn.k_proj.weight,
                    self.attn.o_proj.weight,
                ),
            )
        else:
            x = x + self.orbit(self.orbit_norm(x))
        x = x + self.attn(self.attn_norm(x), cos, sin)
        x = x + self.ffn(self.ffn_norm(x))
        return x


class QuadOrbitForCausalLM(PreTrainedModel, GenerationMixin):
    """Decoder-only language model with QuadOrbit memory branches."""

    config_class = QuadOrbitConfig
    base_model_prefix = ""
    main_input_name = "input_ids"
    _no_split_modules = ["QuadOrbitBlock"]
    _tied_weights_keys = {"lm_head.weight": "tok_emb.weight"}
    supports_gradient_checkpointing = False
    _supports_sdpa = True

    def __init__(self, config: QuadOrbitConfig):
        super().__init__(config)
        self.tok_emb = nn.Embedding(config.vocab_size, config.d_model)
        self.blocks = nn.ModuleList(
            [QuadOrbitBlock(config) for _ in range(config.n_layers)]
        )
        self.norm = RMSNorm(config.d_model, config.rms_eps)
        self.lm_head = nn.Linear(config.d_model, config.vocab_size, bias=False)
        cos, sin = _rope_cache(
            config.head_dim, config.max_seq_len, config.rope_theta
        )
        # Keep the RoPE cache in the public checkpoint. Hugging Face can create
        # custom models on the meta device while loading, so a non-persistent
        # computed buffer may otherwise be materialized as zeros.
        self.register_buffer("cos", cos, persistent=True)
        self.register_buffer("sin", sin, persistent=True)
        self.post_init()
        for block in self.blocks:
            block.orbit.initialize_dynamics(config)
        residual_branches = 2.0 if config.share_orbit_projections else 3.0
        residual_std = config.init_std / math.sqrt(
            residual_branches * config.n_layers
        )
        for block in self.blocks:
            if block.orbit.share_projections:
                nn.init.zeros_(block.orbit.orbit_scale)
            else:
                nn.init.normal_(block.orbit.out_proj.weight, 0.0, residual_std)
            nn.init.normal_(block.attn.o_proj.weight, 0.0, residual_std)
            nn.init.normal_(block.ffn.down_proj.weight, 0.0, residual_std)

    def _init_weights(self, module: nn.Module) -> None:
        if isinstance(module, nn.Linear):
            nn.init.normal_(module.weight, mean=0.0, std=self.config.init_std)
            if module.bias is not None:
                nn.init.zeros_(module.bias)
        elif isinstance(module, nn.Embedding):
            nn.init.normal_(module.weight, mean=0.0, std=self.config.init_std)

    def get_input_embeddings(self):
        return self.tok_emb

    def set_input_embeddings(self, value):
        self.tok_emb = value

    def get_output_embeddings(self):
        return self.lm_head

    def set_output_embeddings(self, value):
        self.lm_head = value

    def _hidden(self, input_ids: torch.Tensor) -> torch.Tensor:
        if input_ids.ndim != 2:
            raise ValueError("input_ids must have shape [batch, time]")
        length = input_ids.shape[1]
        if length <= 0:
            raise ValueError("input_ids must contain at least one token")
        if length > self.config.max_seq_len:
            raise ValueError(
                f"sequence length {length} exceeds {self.config.max_seq_len}"
            )
        x = self.tok_emb(input_ids)
        cos = self.cos[:length]
        sin = self.sin[:length]
        for block in self.blocks:
            x = block(x, cos, sin)
        return self.norm(x)

    def forward(
        self,
        input_ids: torch.Tensor | None = None,
        attention_mask: torch.Tensor | None = None,
        labels: torch.Tensor | None = None,
        past_key_values=None,
        use_cache: bool | None = None,
        return_dict: bool | None = None,
        logits_to_keep: int | None = None,
        **kwargs,
    ):
        del kwargs, use_cache
        if input_ids is None:
            raise ValueError("input_ids are required")
        if past_key_values is not None:
            raise ValueError("QuadOrbit-40M does not provide a KV cache")
        if attention_mask is not None and not bool(torch.all(attention_mask != 0)):
            raise ValueError("padded batches are not supported; use an unpadded prompt")
        hidden = self._hidden(input_ids)
        if logits_to_keep is not None:
            if logits_to_keep <= 0:
                raise ValueError("logits_to_keep must be positive")
            hidden = hidden[:, -logits_to_keep:]
        logits = self.lm_head(hidden)
        loss = None
        if labels is not None:
            if labels.shape != input_ids.shape:
                raise ValueError("labels must match input_ids")
            shift_logits = logits[:, :-1, :].contiguous()
            shift_labels = labels[:, 1:].contiguous()
            loss = F.cross_entropy(
                shift_logits.view(-1, shift_logits.size(-1)),
                shift_labels.view(-1),
                ignore_index=-100,
            )
        if return_dict is False:
            return ((loss, logits) if loss is not None else (logits,))
        return CausalLMOutputWithPast(loss=loss, logits=logits, past_key_values=None)

    def prepare_inputs_for_generation(
        self, input_ids: torch.Tensor, attention_mask=None, **kwargs
    ) -> dict:
        del kwargs
        input_ids = input_ids[:, -self.config.max_seq_len :]
        if attention_mask is not None:
            attention_mask = attention_mask[:, -self.config.max_seq_len :]
        return {
            "input_ids": input_ids,
            "attention_mask": attention_mask,
            "use_cache": False,
        }


__all__ = ["QuadOrbitForCausalLM", "QuadOrbitConfig"]