import torch import torch.nn as nn import torch.nn.functional as F from transformers import PreTrainedModel from transformers.generation import GenerationMixin from transformers.modeling_outputs import CausalLMOutput from .configuration_cheddar import CheddarConfig 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 F.rms_norm(x, (x.size(-1),), self.weight, self.eps) def compute_rope( head_dim, seq_len, base, device, ): if head_dim % 2: raise ValueError( "RoPE head dimension must be even" ) inv_freq = 1.0 / ( base ** ( torch.arange( 0, head_dim, 2, device=device, dtype=torch.float32, ) / head_dim ) ) positions = torch.arange( seq_len, device=device, dtype=torch.float32, ) freqs = torch.outer( positions, inv_freq, ) return freqs.cos(), freqs.sin() def apply_rope(x, cos, sin): dtype = x.dtype x = x.float() x1 = x[..., ::2] x2 = x[..., 1::2] cos = cos[None, None, :, :] sin = sin[None, None, :, :] out = torch.stack( ( x1 * cos - x2 * sin, x1 * sin + x2 * cos, ), dim=-1, ).flatten(-2) return out.to(dtype) class Attention(nn.Module): def __init__(self, config): super().__init__() self.n_head = config.n_head self.n_kv_head = config.n_kv_head self.head_dim = config.head_dim q_dim = self.n_head * self.head_dim kv_dim = self.n_kv_head * self.head_dim self.q_dim, self.kv_dim = q_dim, kv_dim self.qkv_proj = nn.Linear( config.n_embd, q_dim + 2 * kv_dim, bias=False, ) self.o_proj = nn.Linear( q_dim, config.n_embd, bias=False, ) self.q_norm = RMSNorm(self.head_dim, config.rms_eps) if config.qk_norm else nn.Identity() self.k_norm = RMSNorm(self.head_dim, config.rms_eps) if config.qk_norm else nn.Identity() def forward(self, x, cos, sin, attention_mask=None): B, T, _ = x.shape q, k, v = self.qkv_proj(x).split( (self.q_dim, self.kv_dim, self.kv_dim), dim=-1, ) q = q.view(B, T, self.n_head, self.head_dim).transpose(1, 2) k = k.view(B, T, self.n_kv_head, self.head_dim).transpose(1, 2) v = v.view(B, T, self.n_kv_head, self.head_dim).transpose(1, 2) q = self.q_norm(q) k = self.k_norm(k) q = apply_rope(q, cos, sin) k = apply_rope(k, cos, sin) has_padding = ( attention_mask is not None and not attention_mask.bool().all().item() ) if has_padding: if attention_mask.ndim != 2: raise ValueError( "attention_mask must have shape " "(batch, sequence)" ) if attention_mask.shape != (B, T): raise ValueError( f"attention_mask has shape " f"{tuple(attention_mask.shape)}, " f"expected {(B, T)}" ) causal_mask = torch.ones( T, T, device=x.device, dtype=torch.bool, ).tril() key_mask = attention_mask[ :, None, None, : ].bool() attn_mask = ( causal_mask[None, None, :, :] & key_mask ) y = F.scaled_dot_product_attention( q, k, v, attn_mask=attn_mask, dropout_p=0.0, is_causal=False, enable_gqa=( self.n_head != self.n_kv_head ), ) else: y = F.scaled_dot_product_attention( q, k, v, dropout_p=0.0, is_causal=True, enable_gqa=( self.n_head != self.n_kv_head ), ) y = ( y.transpose(1, 2) .contiguous() .view(B, T, self.q_dim) ) return self.o_proj(y) class SwiGLU(nn.Module): def __init__(self, config): super().__init__() self.gate_up_proj = nn.Linear( config.n_embd, 2 * config.mlp_hidden, bias=False, ) self.down_proj = nn.Linear( config.mlp_hidden, config.n_embd, 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 Block(nn.Module): def __init__(self, config): super().__init__() self.attn_norm = RMSNorm( config.n_embd, config.rms_eps, ) self.attn = Attention(config) self.mlp_norm = RMSNorm( config.n_embd, config.rms_eps, ) self.mlp = SwiGLU(config) def forward(self, x, cos, sin, attention_mask=None): x = x + self.attn( self.attn_norm(x), cos, sin, attention_mask=attention_mask, ) x = x + self.mlp( self.mlp_norm(x) ) return x class CheddarPreTrainedModel( PreTrainedModel ): config_class = CheddarConfig base_model_prefix = "model" supports_gradient_checkpointing = False def _init_weights(self, module): if isinstance(module, nn.Linear): nn.init.normal_( module.weight, mean=0.0, std=0.02, ) 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=0.02, ) class CheddarForCausalLM( CheddarPreTrainedModel, GenerationMixin, ): _tied_weights_keys = { "lm_head.weight": "tok_emb.weight" } def __init__(self, config): super().__init__(config) self.tok_emb = nn.Embedding( config.vocab_size, config.n_embd, ) self.blocks = nn.ModuleList( [ Block(config) for _ in range( config.n_layer ) ] ) self.norm = RMSNorm( config.n_embd, config.rms_eps, ) self.lm_head = nn.Linear( config.n_embd, config.vocab_size, bias=False, ) self.post_init() 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 prepare_inputs_for_generation( self, input_ids, attention_mask=None, **kwargs, ): return { "input_ids": input_ids, "attention_mask": attention_mask, } def forward( self, input_ids=None, attention_mask=None, labels=None, **kwargs, ): if input_ids is None: raise ValueError( "input_ids must be provided" ) B, T = input_ids.shape if T > self.config.max_seq_len: raise ValueError( f"Sequence length {T} exceeds " f"max_seq_len=" f"{self.config.max_seq_len}" ) if attention_mask is not None: if ( attention_mask.shape != input_ids.shape ): raise ValueError( "attention_mask must have " "the same shape as input_ids" ) attention_mask = ( attention_mask.to( device=input_ids.device ) ) x = self.tok_emb(input_ids) # Compute RoPE deterministically on # the actual device for every forward. # # Do not store cos/sin as # non-persistent buffers: those buffers # are not restored from the checkpoint # by from_pretrained(). cos, sin = compute_rope( self.config.head_dim, T, self.config.rope_base, x.device, ) for block in self.blocks: x = block( x, cos, sin, attention_mask=attention_mask, ) 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() ) if attention_mask is not None: shift_labels = ( shift_labels.masked_fill( ~attention_mask[ ..., 1: ].bool(), -100, ) ) flat_logits = ( shift_logits.view( -1, shift_logits.size(-1), ) ) flat_labels = ( shift_labels.view(-1) ) loss = F.cross_entropy( flat_logits, flat_labels, ignore_index=-100, ) if self.config.z_loss > 0: valid = ( flat_labels != -100 ) if valid.any(): log_z = torch.logsumexp( flat_logits[ valid ].float(), dim=-1, ) loss = ( loss + self.config.z_loss * log_z.square().mean() ) return CausalLMOutput( loss=loss, logits=logits, )