import math import os from typing import Optional, Tuple, Union import torch import torch.nn as nn import torch.nn.functional as F import torch.utils.checkpoint as cp from transformers.modeling_utils import PreTrainedModel from transformers.modeling_outputs import CausalLMOutputWithPast from transformers.generation import GenerationMixin from safetensors.torch import load_file from transformers import AutoConfig, AutoModel, AutoModelForCausalLM from .configuration_negative import NegativeConfig @torch.no_grad() def get_hadamard_matrix(d: int, dtype=torch.float32) -> torch.Tensor: eye = torch.eye(d, dtype=dtype) h = 1 out = eye.clone() while h < d: out = out.view(-1, 2, h) u = out[:, 0, :] v = out[:, 1, :] out = torch.cat((u + v, u - v), dim=-2) out = out.view(d, d) h *= 2 return (out * (1.0 / math.sqrt(d))).contiguous() class HadamardMLP(nn.Module): def __init__(self, config: NegativeConfig): super().__init__() self.dim = config.hidden_size self.scale1 = nn.Parameter(torch.ones(self.dim)) self.scale2 = nn.Parameter(torch.ones(self.dim)) self.gate = nn.Parameter(torch.ones(self.dim)) self.bias = nn.Parameter(torch.zeros(self.dim)) hadamard_mat = get_hadamard_matrix(self.dim) self.register_buffer("hadamard_mat", hadamard_mat, persistent=False) def forward(self, x: torch.Tensor) -> torch.Tensor: mat = self.hadamard_mat.type_as(x) h = (x * self.scale1) @ mat g = F.silu(x * self.gate) out = ((h * g) @ mat) * self.scale2 + self.bias return out class SwiGLUMLP(nn.Module): def __init__(self, config: NegativeConfig): super().__init__() 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: torch.Tensor) -> torch.Tensor: return self.down_proj(F.silu(self.gate_proj(x)) * self.up_proj(x)) class EngramMemory(nn.Module): def __init__(self, config: NegativeConfig): super().__init__() self.dim = config.hidden_size self.num_entries = config.engram_entries self.n_gram_orders = config.engram_ngram_orders self.tables = nn.ModuleList([ nn.Embedding(self.num_entries, self.dim) for _ in self.n_gram_orders ]) self.gate_proj = nn.Linear(self.dim, self.dim * len(self.n_gram_orders), bias=False) self.out_proj = nn.Linear(self.dim * len(self.n_gram_orders), self.dim, bias=False) def _hash_ngram(self, tokens: torch.Tensor, order: int, table_idx: int) -> torch.Tensor: bsz, seqlen = tokens.shape padded = F.pad(tokens, (order - 1, 0), value=0) primes = (10007, 10009, 10037, 10039, 10061, 10067) p = primes[table_idx % len(primes)] if order == 2: return (padded[:, :seqlen] * p + padded[:, 1 : seqlen + 1]) % self.num_entries elif order == 3: h = (padded[:, :seqlen] * p + padded[:, 1 : seqlen + 1]) % self.num_entries return (h * p + padded[:, 2 : seqlen + 2]) % self.num_entries else: hash_val = torch.zeros((bsz, seqlen), dtype=torch.int64, device=tokens.device) for k in range(order): tok = padded[:, k : k + seqlen] hash_val = (hash_val * p + tok) % self.num_entries return hash_val def forward(self, x: torch.Tensor, tokens: torch.Tensor) -> torch.Tensor: mem_lookups = [self.tables[i](self._hash_ngram(tokens, order, i)) for i, order in enumerate(self.n_gram_orders)] concat_mem = torch.cat(mem_lookups, dim=-1) gate = torch.sigmoid(self.gate_proj(x)) return self.out_proj(concat_mem * gate) 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 = torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps) return x * norm * self.weight class RotaryEmbedding(nn.Module): def __init__(self, dim: int, max_position_embeddings: int = 2048, base: float = 10000.0): super().__init__() self.dim = dim self.max_position_embeddings = max_position_embeddings self.base = base inv_freq = 1.0 / (self.base ** (torch.arange(0, self.dim, 2, dtype=torch.float32) / self.dim)) self.register_buffer("inv_freq", inv_freq, persistent=False) self._set_cos_sin_cache(max_position_embeddings) def _set_cos_sin_cache(self, seq_len: int, device=None, dtype=torch.float32): t = torch.arange(seq_len, device=device, dtype=torch.float32) inv_freq = self.inv_freq.to(device=device, dtype=torch.float32) freqs = torch.outer(t, inv_freq) emb = torch.cat((freqs, freqs), dim=-1) self.register_buffer("cos_cached", emb.cos().to(dtype=dtype), persistent=False) self.register_buffer("sin_cached", emb.sin().to(dtype=dtype), persistent=False) def forward(self, seq_len: int, device: torch.device, dtype: torch.dtype = torch.float32): if not hasattr(self, "cos_cached") or seq_len > self.cos_cached.shape[0] or self.cos_cached.device != device: self._set_cos_sin_cache(seq_len, device=device, dtype=dtype) return ( self.cos_cached[:seq_len].to(device=device, dtype=dtype), self.sin_cached[:seq_len].to(device=device, dtype=dtype), ) def rotate_half(x: torch.Tensor) -> torch.Tensor: 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: torch.Tensor, k: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor): q_embed = (q * cos) + (rotate_half(q) * sin) k_embed = (k * cos) + (rotate_half(k) * sin) return q_embed, k_embed class XSAGQAttention(nn.Module): def __init__(self, config: NegativeConfig): super().__init__() self.dim = config.hidden_size self.n_heads = config.num_attention_heads self.n_kv_heads = config.num_key_value_heads self.head_dim = config.head_dim self.num_kv_groups = self.n_heads // self.n_kv_heads self.use_xsa = config.use_xsa self.use_per_head_gating = config.use_per_head_gating self.wq = nn.Linear(self.dim, self.n_heads * self.head_dim, bias=False) self.wk = nn.Linear(self.dim, self.n_kv_heads * self.head_dim, bias=False) self.wv = nn.Linear(self.dim, self.n_kv_heads * self.head_dim, bias=False) self.wo = nn.Linear(self.n_heads * self.head_dim, self.dim, bias=False) self.q_norm = RMSNorm(self.head_dim, eps=config.rms_norm_eps) self.k_norm = RMSNorm(self.head_dim, eps=config.rms_norm_eps) if self.use_per_head_gating: self.head_gate = nn.Linear(self.dim, self.n_heads, bias=True) nn.init.constant_(self.head_gate.bias, 1.0) nn.init.zeros_(self.head_gate.weight) def forward(self, x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor: bsz, seqlen, _ = x.shape xq = self.wq(x).view(bsz, seqlen, self.n_heads, self.head_dim).transpose(1, 2) xk = self.wk(x).view(bsz, seqlen, self.n_kv_heads, self.head_dim).transpose(1, 2) xv = self.wv(x).view(bsz, seqlen, self.n_kv_heads, self.head_dim).transpose(1, 2) xq = self.q_norm(xq) xk = self.k_norm(xk) xq, xk = apply_rotary_pos_emb(xq, xk, cos, sin) if self.num_kv_groups > 1: xk = xk.repeat_interleave(self.num_kv_groups, dim=1) xv_expanded = xv.repeat_interleave(self.num_kv_groups, dim=1) else: xv_expanded = xv attn_out = F.scaled_dot_product_attention(xq, xk, xv_expanded, is_causal=True) if self.use_xsa: vn = F.normalize(xv_expanded, p=2, dim=-1, eps=1e-6) proj = (attn_out * vn).sum(dim=-1, keepdim=True) attn_out = attn_out - proj * vn if self.use_per_head_gating: gate = torch.sigmoid(self.head_gate(x)).transpose(1, 2).unsqueeze(-1) attn_out = attn_out * gate out = attn_out.transpose(1, 2).contiguous().view(bsz, seqlen, -1) return self.wo(out) class MultiLaneBlock(nn.Module): def __init__(self, config: NegativeConfig, layer_idx: int): super().__init__() self.num_lanes = config.num_lanes self.dim = config.hidden_size self.layer_idx = layer_idx self.attn_norm = RMSNorm(self.dim, eps=config.rms_norm_eps) self.attn = XSAGQAttention(config) self.mlp_norm = RMSNorm(self.dim, eps=config.rms_norm_eps) if config.swiglu_interval == 0: self.use_swiglu = False elif config.swiglu_interval == 1: self.use_swiglu = True else: self.use_swiglu = ((layer_idx + 1) % config.swiglu_interval == 0) if self.use_swiglu: self.mlp = SwiGLUMLP(config) else: self.mlp = HadamardMLP(config) self.lane_mix_attn = nn.Parameter(torch.eye(self.num_lanes) + 0.05 * torch.randn(self.num_lanes, self.num_lanes)) self.lane_mix_mlp = nn.Parameter(torch.eye(self.num_lanes) + 0.05 * torch.randn(self.num_lanes, self.num_lanes)) def forward(self, lanes: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor: primary = lanes[0] attn_update = self.attn(self.attn_norm(primary), cos, sin) mixed = torch.matmul(self.lane_mix_attn, lanes.view(self.num_lanes, -1)).view_as(lanes) lanes = torch.cat([(mixed[0] + attn_update).unsqueeze(0), mixed[1:]], dim=0) mlp_update = self.mlp(self.mlp_norm(lanes[0])) mixed = torch.matmul(self.lane_mix_mlp, lanes.view(self.num_lanes, -1)).view_as(lanes) lanes = torch.cat([(mixed[0] + mlp_update).unsqueeze(0), mixed[1:]], dim=0) return lanes class NegativePreTrainedModel(PreTrainedModel): config_class = NegativeConfig base_model_prefix = "model" supports_gradient_checkpointing = True _no_split_modules = ["MultiLaneBlock"] def _init_weights(self, module): std = self.config.initializer_range if isinstance(module, (nn.Linear, nn.Embedding)): module.weight.data.normal_(mean=0.0, std=std) if hasattr(module, "bias") and module.bias is not None: module.bias.data.zero_() elif isinstance(module, RMSNorm): module.weight.data.fill_(1.0) @classmethod def from_pretrained(cls, pretrained_model_name_or_path, *model_args, **kwargs): config = kwargs.pop("config", None) kwargs.pop("trust_remote_code", None) torch_dtype = kwargs.pop("torch_dtype", None) kwargs.pop("device_map", None) kwargs.pop("low_cpu_mem_usage", None) if config is None: config = NegativeConfig.from_pretrained(pretrained_model_name_or_path) model = cls(config, *model_args) st_file = None bin_file = None if os.path.isdir(str(pretrained_model_name_or_path)): local_st = os.path.join(pretrained_model_name_or_path, "model.safetensors") local_bin = os.path.join(pretrained_model_name_or_path, "pytorch_model.bin") if os.path.exists(local_st): st_file = local_st elif os.path.exists(local_bin): bin_file = local_bin else: try: from huggingface_hub import hf_hub_download st_file = hf_hub_download(repo_id=str(pretrained_model_name_or_path), filename="model.safetensors") except Exception: try: bin_file = hf_hub_download(repo_id=str(pretrained_model_name_or_path), filename="pytorch_model.bin") except Exception: pass if st_file and os.path.exists(st_file): state_dict = load_file(st_file) model.load_state_dict(state_dict, strict=False) elif bin_file and os.path.exists(bin_file): state_dict = torch.load(bin_file, map_location="cpu") model.load_state_dict(state_dict, strict=False) else: return super().from_pretrained(pretrained_model_name_or_path, *model_args, config=config, **kwargs) if getattr(config, "tie_word_embeddings", True) and hasattr(model, "lm_head") and hasattr(model, "model"): model.lm_head.weight = model.model.embed_tokens.weight if torch_dtype is not None: model.to(dtype=torch_dtype) return model class NegativeModel(NegativePreTrainedModel): def __init__(self, config: NegativeConfig, *args, **kwargs): super().__init__(config) self.config = config self.num_lanes = config.num_lanes self.gradient_checkpointing = False self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size) if config.use_engram: self.engram = EngramMemory(config) self.engram_norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) else: self.engram = None self.engram_norm = None self.layers = nn.ModuleList([ MultiLaneBlock(config, layer_idx=i) for i in range(config.num_hidden_layers) ]) # Enhanced Lane Pooling: Learned softmax combination of all 3 lanes before norm self.lane_pool_weights = nn.Parameter(torch.tensor([1.0] + [0.1] * (config.num_lanes - 1))) self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) self.rotary_emb = RotaryEmbedding(config.head_dim, config.max_position_embeddings, config.rope_theta) self.post_init() def get_input_embeddings(self): return self.embed_tokens def set_input_embeddings(self, value): self.embed_tokens = value def forward( self, input_ids: torch.LongTensor = None, attention_mask: Optional[torch.Tensor] = None, position_ids: Optional[torch.LongTensor] = None, inputs_embeds: Optional[torch.FloatTensor] = None, use_cache: Optional[bool] = None, output_attentions: Optional[bool] = None, output_hidden_states: Optional[bool] = None, return_dict: Optional[bool] = None, ): if input_ids is not None: bsz, seqlen = input_ids.shape h0 = self.embed_tokens(input_ids) tokens_for_engram = input_ids elif inputs_embeds is not None: bsz, seqlen, _ = inputs_embeds.shape h0 = inputs_embeds tokens_for_engram = torch.zeros((bsz, seqlen), dtype=torch.long, device=inputs_embeds.device) else: raise ValueError("You must specify either input_ids or inputs_embeds") if self.engram is not None: engram_out = self.engram(self.engram_norm(h0), tokens_for_engram) h0 = h0 + engram_out lanes = h0.unsqueeze(0).repeat(self.num_lanes, 1, 1, 1) cos, sin = self.rotary_emb(seqlen, device=h0.device, dtype=h0.dtype) cos = cos.unsqueeze(0).unsqueeze(0) sin = sin.unsqueeze(0).unsqueeze(0) for layer in self.layers: if self.gradient_checkpointing and self.training: lanes = cp.checkpoint(layer, lanes, cos, sin, use_reentrant=False) else: lanes = layer(lanes, cos, sin) # Weighted lane pooling for higher representation power pool_weights = F.softmax(self.lane_pool_weights, dim=0).view(self.num_lanes, 1, 1, 1) pooled = (lanes * pool_weights).sum(dim=0) out = self.norm(pooled) return out class NegativeModelForCausalLM(NegativePreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} _keys_to_ignore_on_load_missing = ["lm_head.weight"] supports_gradient_checkpointing = True def __init__(self, config: NegativeConfig, *args, **kwargs): super().__init__(config) self.model = NegativeModel(config) self.vocab_size = config.vocab_size self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False) self.post_init() def get_input_embeddings(self): return self.model.embed_tokens def set_input_embeddings(self, value): self.model.embed_tokens = value def get_output_embeddings(self): return self.lm_head def set_output_embeddings(self, new_embeddings): self.lm_head = new_embeddings def prepare_inputs_for_generation( self, input_ids, past_key_values=None, attention_mask=None, inputs_embeds=None, **kwargs, ): if inputs_embeds is not None and past_key_values is None: model_inputs = {"inputs_embeds": inputs_embeds} else: model_inputs = {"input_ids": input_ids} model_inputs.update({ "attention_mask": attention_mask, "use_cache": False, }) return model_inputs def forward( self, input_ids: torch.LongTensor = None, attention_mask: Optional[torch.Tensor] = None, position_ids: Optional[torch.LongTensor] = None, inputs_embeds: Optional[torch.FloatTensor] = None, labels: Optional[torch.LongTensor] = None, use_cache: Optional[bool] = None, output_attentions: Optional[bool] = None, output_hidden_states: Optional[bool] = None, return_dict: Optional[bool] = None, ) -> Union[Tuple, CausalLMOutputWithPast]: return_dict = return_dict if return_dict is not None else getattr(self.config, "return_dict", True) hidden_states = self.model( input_ids=input_ids, attention_mask=attention_mask, position_ids=position_ids, inputs_embeds=inputs_embeds, use_cache=use_cache, output_attentions=output_attentions, output_hidden_states=output_hidden_states, return_dict=return_dict, ) logits = self.lm_head(hidden_states) logits = logits.float() 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 ) if not return_dict: output = (logits,) return ((loss,) + output) if loss is not None else output return CausalLMOutputWithPast( loss=loss, logits=logits, past_key_values=None, hidden_states=None, attentions=None, ) AutoConfig.register("negative", NegativeConfig) AutoModel.register(NegativeConfig, NegativeModel) AutoModelForCausalLM.register(NegativeConfig, NegativeModelForCausalLM)