import math import torch import torch.nn as nn from torch.nn import functional as F from transformers import PreTrainedModel, PretrainedConfig, GenerationMixin from transformers.modeling_outputs import CausalLMOutputWithPast from transformers.models.auto.configuration_auto import CONFIG_MAPPING from transformers.models.auto.modeling_auto import MODEL_FOR_CAUSAL_LM_MAPPING class LanceAIConfig(PretrainedConfig): model_type = "lance_ai" def __init__( self, vocab_size=151936, hidden_size=896, intermediate_size=4864, num_hidden_layers=24, num_attention_heads=14, num_key_value_heads=2, hidden_act="silu", max_position_embeddings=32768, initializer_range=0.02, rms_norm_eps=1e-6, use_cache=True, tie_word_embeddings=True, rope_theta=1000000.0, attention_dropout=0.0, pad_token_id=None, bos_token_id=None, eos_token_id=None, head_dim=64, **kwargs ): super().__init__(**kwargs) self.vocab_size = vocab_size self.hidden_size = hidden_size self.intermediate_size = intermediate_size self.num_hidden_layers = num_hidden_layers self.num_attention_heads = num_attention_heads self.num_key_value_heads = num_key_value_heads self.hidden_act = hidden_act self.max_position_embeddings = max_position_embeddings self.initializer_range = initializer_range self.rms_norm_eps = rms_norm_eps self.use_cache = use_cache self.tie_word_embeddings = tie_word_embeddings self.rope_theta = rope_theta self.attention_dropout = attention_dropout self.head_dim = head_dim self.pad_token_id = pad_token_id self.bos_token_id = bos_token_id self.eos_token_id = eos_token_id class LanceAIRMSNorm(nn.Module): def __init__(self, hidden_size, eps=1e-6): super().__init__() self.weight = nn.Parameter(torch.ones(hidden_size)) self.variance_epsilon = eps def forward(self, hidden_states): input_dtype = hidden_states.dtype hidden_states = hidden_states.to(torch.float32) variance = hidden_states.pow(2).mean(-1, keepdim=True) hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon) return self.weight * hidden_states.to(input_dtype) class LanceAIRotaryEmbedding(nn.Module): def __init__(self, config): super().__init__() self.max_seq_len_cached = config.max_position_embeddings base = config.rope_theta dim = config.head_dim inv_freq = 1.0 / (base ** (torch.arange(0, dim, 2, dtype=torch.int64).float() / dim)) self.register_buffer("inv_freq", inv_freq, persistent=False) @torch.no_grad() def forward(self, x, position_ids): inv_freq_expanded = self.inv_freq[None, :, None].float().expand(position_ids.shape[0], -1, 1).to(x.device) position_ids_expanded = position_ids[:, None, :].float() freqs = (inv_freq_expanded.float() @ position_ids_expanded.float()).transpose(1, 2) emb = torch.cat((freqs, freqs), dim=-1) cos = emb.cos() sin = emb.sin() return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype) def rotate_half(x): x1 = x[..., :x.shape[-1] // 2] x2 = x[..., x.shape[-1] // 2:] return torch.cat((-x2, x1), dim=-1) def apply_rotary_pos_emb(q, k, cos, sin): 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 def repeat_kv(hidden_states, n_rep): batch, num_key_value_heads, slen, head_dim = hidden_states.shape if n_rep == 1: return hidden_states hidden_states = hidden_states[:, :, None, :, :].expand(batch, num_key_value_heads, n_rep, slen, head_dim) return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim) class LanceAIAttention(nn.Module): def __init__(self, config, layer_idx): super().__init__() self.config = config self.layer_idx = layer_idx self.head_dim = config.head_dim self.num_heads = config.num_attention_heads self.num_kv_heads = config.num_key_value_heads self.num_key_value_groups = config.num_attention_heads // config.num_key_value_heads self.scaling = self.head_dim ** -0.5 self.attention_dropout = config.attention_dropout self.is_causal = True self.q_proj = nn.Linear(config.hidden_size, config.num_attention_heads * self.head_dim, bias=True) self.k_proj = nn.Linear(config.hidden_size, config.num_key_value_heads * self.head_dim, bias=True) self.v_proj = nn.Linear(config.hidden_size, config.num_key_value_heads * self.head_dim, bias=True) self.o_proj = nn.Linear(config.num_attention_heads * self.head_dim, config.hidden_size, bias=False) def forward(self, hidden_states, attention_mask=None, position_embeddings=None, past_key_values=None, use_cache=False): batch_size, seq_len, _ = hidden_states.shape query_states = self.q_proj(hidden_states).view(batch_size, seq_len, self.num_heads, self.head_dim).transpose(1, 2) key_states = self.k_proj(hidden_states).view(batch_size, seq_len, self.num_kv_heads, self.head_dim).transpose(1, 2) value_states = self.v_proj(hidden_states).view(batch_size, seq_len, self.num_kv_heads, self.head_dim).transpose(1, 2) cos, sin = position_embeddings query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin) if past_key_values is not None: if past_key_values[0] is not None: key_states = torch.cat([past_key_values[0], key_states], dim=2) if past_key_values[1] is not None: value_states = torch.cat([past_key_values[1], value_states], dim=2) past = (key_states, value_states) if use_cache else None key_states = repeat_kv(key_states, self.num_key_value_groups) value_states = repeat_kv(value_states, self.num_key_value_groups) attn_weights = torch.matmul(query_states, key_states.transpose(2, 3)) * self.scaling if attention_mask is not None: attn_weights = attn_weights + attention_mask attn_weights = F.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query_states.dtype) attn_weights = F.dropout(attn_weights, p=self.attention_dropout, training=self.training) attn_output = torch.matmul(attn_weights, value_states) attn_output = attn_output.transpose(1, 2).contiguous() attn_output = attn_output.reshape(batch_size, seq_len, -1) attn_output = self.o_proj(attn_output) return attn_output, past class LanceAIMLP(nn.Module): def __init__(self, config): super().__init__() self.hidden_size = config.hidden_size self.intermediate_size = config.intermediate_size self.gate_proj = nn.Linear(config.hidden_size, config.intermediate_size, bias=False) self.up_proj = nn.Linear(config.hidden_size, config.intermediate_size, bias=False) self.down_proj = nn.Linear(config.intermediate_size, config.hidden_size, bias=False) def forward(self, x): return self.down_proj(F.silu(self.gate_proj(x)) * self.up_proj(x)) class LanceAIDecoderLayer(nn.Module): def __init__(self, config, layer_idx): super().__init__() self.hidden_size = config.hidden_size self.self_attn = LanceAIAttention(config, layer_idx) self.mlp = LanceAIMLP(config) self.input_layernorm = LanceAIRMSNorm(config.hidden_size, eps=config.rms_norm_eps) self.post_attention_layernorm = LanceAIRMSNorm(config.hidden_size, eps=config.rms_norm_eps) def forward(self, hidden_states, attention_mask=None, position_embeddings=None, past_key_values=None, use_cache=False): residual = hidden_states hidden_states = self.input_layernorm(hidden_states) hidden_states, past_kv = self.self_attn( hidden_states, attention_mask=attention_mask, position_embeddings=position_embeddings, past_key_values=past_key_values, use_cache=use_cache, ) hidden_states = residual + hidden_states residual = hidden_states hidden_states = self.post_attention_layernorm(hidden_states) hidden_states = self.mlp(hidden_states) hidden_states = residual + hidden_states return hidden_states, past_kv class LanceAIPreTrainedModel(PreTrainedModel): config_class = LanceAIConfig base_model_prefix = "model" supports_gradient_checkpointing = True _no_split_modules = ["LanceAIDecoderLayer"] _skip_keys_device_placement = ["past_key_values"] _supports_flash_attn = True _supports_sdpa = True class LanceAIModel(LanceAIPreTrainedModel): def __init__(self, config): super().__init__(config) self.padding_idx = config.pad_token_id self.vocab_size = config.vocab_size self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx) self.layers = nn.ModuleList( [LanceAIDecoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)] ) self.norm = LanceAIRMSNorm(config.hidden_size, eps=config.rms_norm_eps) self.rotary_emb = LanceAIRotaryEmbedding(config) self.gradient_checkpointing = False self.gradient_checkpointing_func = None self.post_init() def _set_gradient_checkpointing(self, enable=True, gradient_checkpointing_func=None): self.gradient_checkpointing = enable if gradient_checkpointing_func is not None: self.gradient_checkpointing_func = gradient_checkpointing_func def _past_seq_length(self, past_key_values): """Get sequence length from past cache (tuple or DynamicCache)""" if past_key_values is None: return 0 if hasattr(past_key_values, 'get_seq_length'): return past_key_values.get_seq_length() if len(past_key_values) > 0 and past_key_values[0] is not None: if past_key_values[0][0] is not None: return past_key_values[0][0].shape[2] return 0 def _convert_past(self, past_key_values, use_cache): """Convert DynamicCache to our tuple format if needed""" if past_key_values is None or not use_cache: return past_key_values, use_cache if isinstance(past_key_values, list): return past_key_values, use_cache # DynamicCache -> list of (k, v) tuples if hasattr(past_key_values, 'get_seq_length'): converted = [] for i in range(len(self.layers)): if i < len(past_key_values): kv = past_key_values[i] converted.append((kv[0], kv[1])) else: converted.append((None, None)) return converted, use_cache return past_key_values, use_cache def forward(self, input_ids=None, attention_mask=None, position_ids=None, past_key_values=None, inputs_embeds=None, use_cache=None): if (input_ids is None) ^ (inputs_embeds is not None): raise ValueError("Specify exactly one of input_ids or inputs_embeds") if inputs_embeds is None: inputs_embeds = self.embed_tokens(input_ids) batch_size, seq_len = inputs_embeds.shape[:2] past_key_values, use_cache = self._convert_past(past_key_values, use_cache) if position_ids is None: past_seen_tokens = self._past_seq_length(past_key_values) position_ids = torch.arange(seq_len, device=inputs_embeds.device) + past_seen_tokens position_ids = position_ids.unsqueeze(0).expand(batch_size, -1) position_embeddings = self.rotary_emb(inputs_embeds, position_ids) if attention_mask is not None: causal_mask = self._make_causal_mask(inputs_embeds, past_key_values) attention_mask = attention_mask[:, None, None, :] attention_mask = (1.0 - attention_mask) * torch.finfo(inputs_embeds.dtype).min attention_mask = attention_mask + causal_mask else: attention_mask = self._make_causal_mask(inputs_embeds, past_key_values) hidden_states = inputs_embeds new_past = [] if use_cache else None for i, layer in enumerate(self.layers): layer_past = past_key_values[i] if past_key_values is not None else None if self.gradient_checkpointing and self.training: def create_custom_forward(module): def custom_forward(*inputs): return module(*inputs, past_key_values=None, use_cache=False) return custom_forward hidden_states, _ = torch.utils.checkpoint.checkpoint( create_custom_forward(layer), hidden_states, attention_mask, position_embeddings, use_reentrant=True, ) else: hidden_states, layer_past = layer( hidden_states, attention_mask=attention_mask, position_embeddings=position_embeddings, past_key_values=layer_past, use_cache=use_cache, ) if use_cache: new_past.append(layer_past) hidden_states = self.norm(hidden_states) return hidden_states, new_past def _make_causal_mask(self, inputs_embeds, past_key_values=None): batch_size, seq_len, _ = inputs_embeds.shape past_len = self._past_seq_length(past_key_values) total_len = past_len + seq_len mask = torch.full((seq_len, total_len), float('-inf'), device=inputs_embeds.device) mask = torch.triu(mask, diagonal=1 + past_len) return mask[None, None, :, :] class LanceAI(LanceAIPreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} def __init__(self, config): super().__init__(config) self.model = LanceAIModel(config) self.vocab_size = config.vocab_size self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False) self.post_init() def _set_gradient_checkpointing(self, enable=True, gradient_checkpointing_func=None): self.model._set_gradient_checkpointing(enable, gradient_checkpointing_func) def forward(self, input_ids=None, attention_mask=None, position_ids=None, past_key_values=None, inputs_embeds=None, labels=None, use_cache=None, **kwargs): hidden_states, past = self.model( input_ids=input_ids, attention_mask=attention_mask, position_ids=position_ids, past_key_values=past_key_values, inputs_embeds=inputs_embeds, use_cache=use_cache, ) 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), ignore_index=-100, ) return CausalLMOutputWithPast(loss=loss, logits=logits, past_key_values=past) def _past_seq_len(self, past_key_values): if past_key_values is None: return 0 if hasattr(past_key_values, 'get_seq_length'): return past_key_values.get_seq_length() if isinstance(past_key_values, list) and len(past_key_values) > 0: if past_key_values[0] is not None: k, v = past_key_values[0] if k is not None: return k.shape[2] return 0 def prepare_inputs_for_generation(self, input_ids, past_key_values=None, attention_mask=None, **kwargs): if self._past_seq_len(past_key_values) > 0: input_ids = input_ids[:, -1:] return { "input_ids": input_ids, "attention_mask": attention_mask, "past_key_values": past_key_values, "use_cache": True, } def _reorder_cache(self, past_key_values, beam_idx): reordered = [] for layer_past in past_key_values: layer_k, layer_v = layer_past reordered.append((layer_k.index_select(0, beam_idx), layer_v.index_select(0, beam_idx))) return reordered CONFIG_MAPPING.register("lance_ai", LanceAIConfig) MODEL_FOR_CAUSAL_LM_MAPPING.register(LanceAIConfig, LanceAI) LanceAIConfig.register_for_auto_class("AutoConfig") LanceAI.register_for_auto_class("AutoModelForCausalLM")