"""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"]