import torch import torch.nn as nn import torch.nn.functional as F from transformers import GenerationMixin, PreTrainedModel from transformers.modeling_outputs import CausalLMOutput from .configuration_koliber import KoliberConfig class RMSNorm(nn.Module): def __init__(self, dim, eps=1e-6): super().__init__() self.dim = int(dim) self.eps = float(eps) self.weight = nn.Parameter(torch.ones(dim)) def forward(self, x): if hasattr(F, "rms_norm"): return F.rms_norm(x, (self.dim,), self.weight, self.eps) dtype = x.dtype y = x.float() y = y * torch.rsqrt( y.square().mean(dim=-1, keepdim=True) + self.eps ) return y.to(dtype) * self.weight class RotaryEmbedding(nn.Module): def __init__(self, head_dim, max_positions, theta): super().__init__() inv_freq = 1.0 / ( theta ** ( torch.arange(0, head_dim, 2, dtype=torch.float32) / head_dim ) ) positions = torch.arange( max_positions, dtype=torch.float32, ) freqs = torch.outer( positions, inv_freq, ) self.register_buffer( "cos_cached", freqs.cos(), persistent=True, ) self.register_buffer( "sin_cached", freqs.sin(), persistent=True, ) def forward(self, seq_len, dtype): return ( self.cos_cached[:seq_len].to(dtype=dtype), self.sin_cached[:seq_len].to(dtype=dtype), ) def apply_rope(x, cos, sin): cos = cos[None, None, :, :] sin = sin[None, None, :, :] even = x[..., 0::2] odd = x[..., 1::2] return torch.stack( ( even * cos - odd * sin, even * sin + odd * cos, ), dim=-1, ).flatten(-2) class GQAAttention(nn.Module): def __init__(self, config, rope): super().__init__() self.hidden_size = config.hidden_size self.num_heads = config.num_attention_heads self.num_kv_heads = config.num_key_value_heads self.head_dim = config.head_dim self.kv_repeat = self.num_heads // self.num_kv_heads object.__setattr__(self, "rope", rope) self.q_size = self.num_heads * self.head_dim self.kv_size = self.num_kv_heads * self.head_dim self.qkv_proj = nn.Linear( config.hidden_size, self.q_size + 2 * self.kv_size, bias=False, ) self.o_proj = nn.Linear( config.hidden_size, config.hidden_size, bias=False, ) def forward(self, x): b, s, _ = x.shape qkv = self.qkv_proj(x) q, k, v = qkv.split( ( self.q_size, self.kv_size, self.kv_size, ), dim=-1, ) q = ( q.view( b, s, self.num_heads, self.head_dim, ) .transpose(1, 2) ) k = ( k.view( b, s, self.num_kv_heads, self.head_dim, ) .transpose(1, 2) ) v = ( v.view( b, s, self.num_kv_heads, self.head_dim, ) .transpose(1, 2) ) cos, sin = self.rope( s, q.dtype, ) q = apply_rope(q, cos, sin) k = apply_rope(k, cos, sin) try: y = F.scaled_dot_product_attention( q, k, v, attn_mask=None, dropout_p=0.0, is_causal=True, enable_gqa=True, ) except TypeError: k = k.repeat_interleave( self.kv_repeat, dim=1, ) v = v.repeat_interleave( self.kv_repeat, dim=1, ) y = F.scaled_dot_product_attention( q, k, v, attn_mask=None, dropout_p=0.0, is_causal=True, ) y = ( y.transpose(1, 2) .contiguous() .view( b, s, self.hidden_size, ) ) return self.o_proj(y) class SwiGLU(nn.Module): def __init__(self, hidden, intermediate): super().__init__() self.gate_up_proj = nn.Linear( hidden, 2 * intermediate, bias=False, ) self.down_proj = nn.Linear( intermediate, hidden, bias=False, ) def forward(self, x): gate, up = self.gate_up_proj(x).chunk(2, dim=-1) return self.down_proj( F.silu(gate) * up ) class KoliberDecoderLayer(nn.Module): def __init__(self, config, rope): super().__init__() self.input_layernorm = RMSNorm( config.hidden_size, config.rms_norm_eps, ) self.self_attn = GQAAttention( config, rope, ) self.post_attention_layernorm = RMSNorm( config.hidden_size, config.rms_norm_eps, ) self.mlp = SwiGLU( config.hidden_size, config.intermediate_size, ) def forward(self, hidden_states): hidden_states = ( hidden_states + self.self_attn( self.input_layernorm( hidden_states ) ) ) hidden_states = ( hidden_states + self.mlp( self.post_attention_layernorm( hidden_states ) ) ) return hidden_states class KoliberPreTrainedModel(PreTrainedModel): config_class = KoliberConfig base_model_prefix = "" main_input_name = "input_ids" supports_gradient_checkpointing = False _tied_weights_keys = { "lm_head.weight": "embed_tokens.weight" } def _init_weights(self, module): if isinstance(module, nn.Embedding): nn.init.normal_( module.weight, mean=0.0, std=0.02, ) elif isinstance(module, nn.Linear): nn.init.normal_( module.weight, mean=0.0, std=0.02, ) class KoliberForCausalLM( KoliberPreTrainedModel, GenerationMixin, ): def __init__(self, config): super().__init__(config) self.embed_tokens = nn.Embedding( config.vocab_size, config.hidden_size, ) self.rope = RotaryEmbedding( config.head_dim, config.max_position_embeddings, config.rope_theta, ) self.layers = nn.ModuleList( [ KoliberDecoderLayer( config, self.rope, ) for _ in range( config.num_hidden_layers ) ] ) self.norm = RMSNorm( config.hidden_size, config.rms_norm_eps, ) self.lm_head = nn.Linear( config.hidden_size, config.vocab_size, bias=False, ) self.post_init() self.tie_weights() def get_input_embeddings(self): return self.embed_tokens def set_input_embeddings(self, value): self.embed_tokens = value self.tie_weights() def get_output_embeddings(self): return self.lm_head def set_output_embeddings(self, value): self.lm_head = value def forward_hidden(self, input_ids): x = self.embed_tokens(input_ids) for layer in self.layers: x = layer(x) return self.norm(x) def forward( self, input_ids=None, attention_mask=None, labels=None, use_cache=False, return_dict=None, **kwargs, ): if input_ids is None: raise ValueError("input_ids is required") hidden = self.forward_hidden(input_ids) logits = F.linear( hidden, self.embed_tokens.weight, ) loss = None if labels is not None: shift_logits = logits[:, :-1, :].float() shift_labels = labels[:, 1:] loss = F.cross_entropy( shift_logits.reshape( -1, shift_logits.shape[-1], ), shift_labels.reshape(-1), ignore_index=-100, ) if return_dict is False: if loss is None: return (logits,) return (loss, logits) return CausalLMOutput( loss=loss, logits=logits, ) def prepare_inputs_for_generation( self, input_ids, attention_mask=None, **kwargs, ): return { "input_ids": input_ids, "attention_mask": attention_mask, "use_cache": False, }