| 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,
|
| ):
|
|
|
| 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,
|
| ) |
