""" Twinkel LLM Model Implementation Creator: Kunal Pandey """ import torch import torch.nn as nn import torch.nn.functional as F from transformers import PreTrainedModel from transformers.modeling_outputs import CausalLMOutputWithPast from typing import Optional, Tuple, Union from .configuration_twinkel_llm import TwinkelLLMConfig class RMSNorm(nn.Module): def __init__(self, dim, eps=1e-6): super().__init__() self.weight = nn.Parameter(torch.ones(dim)) self.eps = eps def forward(self, x): return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps) * self.weight class RotaryEmbedding(nn.Module): def __init__(self, dim, max_seq_len=2048): super().__init__() inv_freq = 1.0 / (10000 ** (torch.arange(0, dim, 2).float() / dim)) self.register_buffer("inv_freq", inv_freq, persistent=False) self.max_seq_len = max_seq_len self._set_cos_sin_cache(max_seq_len) def _set_cos_sin_cache(self, seq_len): self.max_seq_len = seq_len t = torch.arange(seq_len, dtype=torch.float32) freqs = torch.einsum("i,j->ij", t, self.inv_freq) emb = torch.cat((freqs, freqs), dim=-1) self.register_buffer("cos_cached", emb.cos(), persistent=False) self.register_buffer("sin_cached", emb.sin(), persistent=False) def forward(self, x, seq_len): if seq_len > self.max_seq_len: self._set_cos_sin_cache(seq_len) return self.cos_cached[:seq_len].to(x.device), self.sin_cached[:seq_len].to(x.device) def apply_rotary_pos_emb(q, k, cos, sin): def rotate_half(x): x1, x2 = x.chunk(2, dim=-1) return torch.cat((-x2, x1), dim=-1) cos = cos.unsqueeze(1) sin = sin.unsqueeze(1) q_embed = (q * cos) + (rotate_half(q) * sin) k_embed = (k * cos) + (rotate_half(k) * sin) return q_embed, k_embed class GroupedQueryAttention(nn.Module): def __init__(self, config): super().__init__() self.n_heads = config.num_attention_heads self.n_kv_heads = config.num_key_value_heads self.head_dim = config.hidden_size // self.n_heads self.hidden_size = config.hidden_size self.n_rep = self.n_heads // self.n_kv_heads self.q_proj = nn.Linear(self.hidden_size, self.n_heads * self.head_dim, bias=False) self.k_proj = nn.Linear(self.hidden_size, self.n_kv_heads * self.head_dim, bias=False) self.v_proj = nn.Linear(self.hidden_size, self.n_kv_heads * self.head_dim, bias=False) self.o_proj = nn.Linear(self.n_heads * self.head_dim, self.hidden_size, bias=False) self.rope = RotaryEmbedding(self.head_dim) def forward(self, x): B, T, C = x.shape q = self.q_proj(x).view(B, T, self.n_heads, self.head_dim).transpose(1, 2) k = self.k_proj(x).view(B, T, self.n_kv_heads, self.head_dim).transpose(1, 2) v = self.v_proj(x).view(B, T, self.n_kv_heads, self.head_dim).transpose(1, 2) cos, sin = self.rope(x, T) q = q.transpose(1, 2) k = k.transpose(1, 2) q, k = apply_rotary_pos_emb(q, k, cos, sin) q = q.transpose(1, 2) k = k.transpose(1, 2) if self.n_rep > 1: k = k.repeat_interleave(self.n_rep, dim=1) v = v.repeat_interleave(self.n_rep, dim=1) attn_weights = torch.matmul(q, k.transpose(-2, -1)) / (self.head_dim ** 0.5) causal_mask = torch.triu(torch.ones(T, T, device=x.device, dtype=torch.bool), diagonal=1) attn_weights = attn_weights.masked_fill(causal_mask, float('-inf')) attn_weights = F.softmax(attn_weights, dim=-1, dtype=torch.float32).to(x.dtype) attn_output = torch.matmul(attn_weights, v) attn_output = attn_output.transpose(1, 2).contiguous().view(B, T, -1) return self.o_proj(attn_output) class SwiGLU(nn.Module): def __init__(self, config): super().__init__() hidden = config.hidden_size intermediate = config.intermediate_size self.gate_proj = nn.Linear(hidden, intermediate, bias=False) self.up_proj = nn.Linear(hidden, intermediate, bias=False) self.down_proj = nn.Linear(intermediate, hidden, bias=False) def forward(self, x): return self.down_proj(F.silu(self.gate_proj(x)) * self.up_proj(x)) class TransformerBlock(nn.Module): def __init__(self, config): super().__init__() self.attn_norm = RMSNorm(config.hidden_size) self.ffn_norm = RMSNorm(config.hidden_size) self.attn = GroupedQueryAttention(config) self.ffn = SwiGLU(config) self.dropout = nn.Dropout(config.dropout) def forward(self, x): x = x + self.dropout(self.attn(self.attn_norm(x))) x = x + self.dropout(self.ffn(self.ffn_norm(x))) return x class TwinkelLLMPreTrainedModel(PreTrainedModel): config_class = TwinkelLLMConfig base_model_prefix = "model" supports_gradient_checkpointing = True _no_split_modules = ["TransformerBlock"] def _init_weights(self, module): if isinstance(module, nn.Linear): torch.nn.init.normal_(module.weight, mean=0.0, std=0.02) elif isinstance(module, nn.Embedding): torch.nn.init.normal_(module.weight, mean=0.0, std=0.02) class TwinkelLLMForCausalLM(TwinkelLLMPreTrainedModel): """ Twinkel LLM Model for Causal Language Modeling Creator: Kunal Pandey """ def __init__(self, config): super().__init__(config) self.config = config self.token_embedding = nn.Embedding(config.vocab_size, config.hidden_size) self.layers = nn.ModuleList([TransformerBlock(config) for _ in range(config.num_hidden_layers)]) self.norm = RMSNorm(config.hidden_size) self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False) # Weight tying self.lm_head.weight = self.token_embedding.weight self.post_init() def get_input_embeddings(self): return self.token_embedding def set_input_embeddings(self, value): self.token_embedding = value def forward( self, input_ids: torch.LongTensor = None, attention_mask: Optional[torch.Tensor] = None, labels: Optional[torch.LongTensor] = None, **kwargs ) -> Union[Tuple, CausalLMOutputWithPast]: x = self.token_embedding(input_ids) for layer in self.layers: x = layer(x) x = self.norm(x) logits = self.lm_head(x) 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, self.config.vocab_size), shift_labels.view(-1), ignore_index=-100, reduction='mean' ) return CausalLMOutputWithPast( loss=loss, logits=logits, ) def prepare_inputs_for_generation(self, input_ids, **kwargs): return {"input_ids": input_ids}