| 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) |
| ]) |
|
|
| |
| 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) |
|
|
| |
| 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) |
|
|
