from __future__ import annotations from typing import Optional import torch import torch.nn as nn import torch.nn.functional as F from transformers import GenerationMixin, PreTrainedModel from transformers.modeling_outputs import CausalLMOutputWithPast from .configuration import CustomTransformerConfig def precompute_rope(head_dim: int, max_seq_len: int, theta: float): inv_freq = 1.0 / (theta ** (torch.arange(0, head_dim, 2).float() / head_dim)) t = torch.arange(max_seq_len).float() freqs = torch.outer(t, inv_freq) return torch.cos(freqs), torch.sin(freqs) def apply_rope(x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor: x1, x2 = x[..., ::2], x[..., 1::2] cos = cos[None, None, :, :] sin = sin[None, None, :, :] rx1 = x1 * cos - x2 * sin rx2 = x1 * sin + x2 * cos return torch.stack([rx1, rx2], dim=-1).flatten(-2) def repeat_kv(x: torch.Tensor, n_rep: int) -> torch.Tensor: if n_rep == 1: return x b, n_kv, t, d = x.shape x = x[:, :, None, :, :].expand(b, n_kv, n_rep, t, d) return x.reshape(b, n_kv * n_rep, t, d) class RMSNorm(nn.Module): def __init__(self, dim: int, eps: float = 1e-5): super().__init__() self.eps = eps self.weight = nn.Parameter(torch.ones(dim)) def forward(self, x: torch.Tensor) -> torch.Tensor: norm_x = x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps) return norm_x * self.weight class GQAAttention(nn.Module): """Grouped-query attention with RoPE, Q/K bias, and per-head gating.""" def __init__(self, cfg: CustomTransformerConfig): super().__init__() self.n_heads = cfg.num_attention_heads self.n_kv_heads = cfg.num_key_value_heads self.n_rep = cfg.num_attention_heads // cfg.num_key_value_heads self.head_dim = cfg.hidden_size // cfg.num_attention_heads self.dropout = cfg.dropout self.wq = nn.Linear(cfg.hidden_size, cfg.num_attention_heads * self.head_dim, bias=cfg.qk_bias) self.wk = nn.Linear(cfg.hidden_size, cfg.num_key_value_heads * self.head_dim, bias=cfg.qk_bias) self.wv = nn.Linear(cfg.hidden_size, cfg.num_key_value_heads * self.head_dim, bias=False) self.wo = nn.Linear(cfg.num_attention_heads * self.head_dim, cfg.hidden_size, bias=False) self.use_head_gating = cfg.use_head_gating if self.use_head_gating: self.head_gate = nn.Linear(cfg.hidden_size, cfg.num_attention_heads, bias=True) def forward(self, x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor: b, t, _ = x.shape q = self.wq(x).view(b, t, self.n_heads, self.head_dim).transpose(1, 2) k = self.wk(x).view(b, t, self.n_kv_heads, self.head_dim).transpose(1, 2) v = self.wv(x).view(b, t, self.n_kv_heads, self.head_dim).transpose(1, 2) q = apply_rope(q, cos[:t], sin[:t]) k = apply_rope(k, cos[:t], sin[:t]) k = repeat_kv(k, self.n_rep) v = repeat_kv(v, self.n_rep) out = F.scaled_dot_product_attention( q, k, v, is_causal=True, dropout_p=self.dropout if self.training else 0.0 ) if self.use_head_gating: gate = torch.sigmoid(self.head_gate(x)) gate = gate.transpose(1, 2).unsqueeze(-1) out = out * gate out = out.transpose(1, 2).contiguous().view(b, t, self.n_heads * self.head_dim) return self.wo(out) class SwiGLU(nn.Module): """SwiGLU feed-forward with explicit hidden size.""" def __init__(self, dim: int, hidden_size: int): super().__init__() self.w1 = nn.Linear(dim, hidden_size, bias=False) self.w3 = nn.Linear(dim, hidden_size, bias=False) self.w2 = nn.Linear(hidden_size, dim, bias=False) def forward(self, x: torch.Tensor) -> torch.Tensor: return self.w2(F.silu(self.w1(x)) * self.w3(x)) def attn_res_aggregate(sources: list, proj: nn.Linear, norm: RMSNorm) -> torch.Tensor: v = torch.stack(sources, dim=0) k = norm(v) logits = torch.einsum("d,sbtd->sbt", proj.weight.squeeze(0), k) weights = logits.softmax(dim=0) return torch.einsum("sbt,sbtd->btd", weights, v) class AttnResUnit(nn.Module): def __init__(self, dim: int, eps: float): super().__init__() self.norm = RMSNorm(dim, eps) self.proj = nn.Linear(dim, 1, bias=False) nn.init.zeros_(self.proj.weight) def forward(self, sources: list) -> torch.Tensor: return attn_res_aggregate(sources, self.proj, self.norm) class Block(nn.Module): def __init__(self, cfg: CustomTransformerConfig): super().__init__() self.mode = cfg.attn_res_mode if self.mode != "none": self.attn_res_unit = AttnResUnit(cfg.hidden_size, cfg.norm_eps) self.mlp_res_unit = AttnResUnit(cfg.hidden_size, cfg.norm_eps) self.attn_norm = RMSNorm(cfg.hidden_size, cfg.norm_eps) self.attn = GQAAttention(cfg) self.ffn_norm = RMSNorm(cfg.hidden_size, cfg.norm_eps) self.ffn = SwiGLU(cfg.hidden_size, cfg.ffn_hidden_size) class CustomTransformerForCausalLM(PreTrainedModel, GenerationMixin): config_class = CustomTransformerConfig base_model_prefix = "custom_transformer" main_input_name = "input_ids" _tied_weights_keys = ["lm_head.weight"] all_tied_weights_keys = {"lm_head.weight": "tok_emb.weight"} _keys_to_ignore_on_load_missing = [r"lm_head\.weight"] def __init__(self, config: CustomTransformerConfig): super().__init__(config) self.cfg = config self.tok_emb = nn.Embedding(config.vocab_size, config.hidden_size) self.blocks = nn.ModuleList([Block(config) for _ in range(config.num_hidden_layers)]) self.norm_f = RMSNorm(config.hidden_size, config.norm_eps) self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False) if config.tie_word_embeddings: self.lm_head.weight = self.tok_emb.weight if config.attn_res_mode != "none": self.output_res = AttnResUnit(config.hidden_size, config.norm_eps) head_dim = config.hidden_size // config.num_attention_heads cos, sin = precompute_rope(head_dim, config.max_position_embeddings, config.rope_theta) self.register_buffer("rope_cos", cos, persistent=False) self.register_buffer("rope_sin", sin, persistent=False) self.apply(self._init_weights) self._zero_init_attn_res_queries() self.tie_weights() @staticmethod def _init_weights(module: nn.Module): if isinstance(module, nn.Linear): nn.init.normal_(module.weight, std=0.02) if module.bias is not None: nn.init.zeros_(module.bias) elif isinstance(module, nn.Embedding): nn.init.normal_(module.weight, std=0.02) def _zero_init_attn_res_queries(self): for m in self.modules(): if isinstance(m, AttnResUnit): nn.init.zeros_(m.proj.weight) def get_input_embeddings(self): return self.tok_emb def set_input_embeddings(self, value): self.tok_emb = value if self.config.tie_word_embeddings: self.lm_head.weight = self.tok_emb.weight def get_output_embeddings(self): return self.lm_head def set_output_embeddings(self, new_embeddings): self.lm_head = new_embeddings if self.config.tie_word_embeddings: self.lm_head.weight = self.tok_emb.weight def tie_weights(self, *args, **kwargs): if self.config.tie_word_embeddings: self.lm_head.weight = self.tok_emb.weight def prepare_inputs_for_generation( self, input_ids, past_key_values=None, attention_mask=None, **kwargs, ): # No KV cache in this architecture, so generation reuses the full prompt each step. return { "input_ids": input_ids, "attention_mask": attention_mask, "use_cache": False, } def _forward_none(self, x: torch.Tensor, cos, sin) -> torch.Tensor: h = x for block in self.blocks: h = h + block.attn(block.attn_norm(h), cos, sin) h = h + block.ffn(block.ffn_norm(h)) return h def _forward_full(self, x: torch.Tensor, cos, sin) -> torch.Tensor: history = [x] for block in self.blocks: h_attn = block.attn_res_unit(history) v_attn = block.attn(block.attn_norm(h_attn), cos, sin) history.append(v_attn) h_mlp = block.mlp_res_unit(history) v_mlp = block.ffn(block.ffn_norm(h_mlp)) history.append(v_mlp) return self.output_res(history) def _forward_block(self, x: torch.Tensor, cos, sin) -> torch.Tensor: S = self.config.attn_res_block_size blocks_list = [x] partial = None intra_idx = 0 def sources(): return blocks_list if partial is None else (blocks_list + [partial]) for block in self.blocks: intra_idx += 1 h_attn = attn_res_aggregate(sources(), block.attn_res_unit.proj, block.attn_res_unit.norm) v_attn = block.attn(block.attn_norm(h_attn), cos, sin) partial = v_attn if partial is None else (partial + v_attn) if intra_idx >= S: blocks_list = blocks_list + [partial] partial = None intra_idx = 0 intra_idx += 1 h_mlp = attn_res_aggregate(sources(), block.mlp_res_unit.proj, block.mlp_res_unit.norm) v_mlp = block.ffn(block.ffn_norm(h_mlp)) partial = v_mlp if partial is None else (partial + v_mlp) if intra_idx >= S: blocks_list = blocks_list + [partial] partial = None intra_idx = 0 return self.output_res(sources()) def forward( self, input_ids: torch.Tensor, attention_mask: Optional[torch.Tensor] = None, labels: Optional[torch.Tensor] = None, use_cache: Optional[bool] = None, past_key_values=None, output_attentions: Optional[bool] = None, output_hidden_states: Optional[bool] = None, return_dict: bool = True, **kwargs, ): if use_cache: raise ValueError("This architecture does not implement KV caching yet; set use_cache=False.") x = self.tok_emb(input_ids) cos, sin = self.rope_cos, self.rope_sin if self.config.attn_res_mode == "none": final = self._forward_none(x, cos, sin) elif self.config.attn_res_mode == "full": final = self._forward_full(x, cos, sin) else: final = self._forward_block(x, cos, sin) hidden_states = self.norm_f(final) logits = self.lm_head(hidden_states) loss = None if labels is not None: 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), ) if not return_dict: return (loss, logits) return CausalLMOutputWithPast( loss=loss, logits=logits, past_key_values=None, hidden_states=None, attentions=None, )