# -*- coding: utf-8 -*- from __future__ import annotations import math import warnings from typing import List, Optional, Tuple, Union import torch import torch.nn as nn import torch.utils.checkpoint from transformers.activations import ACT2FN from transformers.cache_utils import Cache, DynamicCache from transformers.modeling_outputs import (BaseModelOutputWithPast, CausalLMOutputWithPast) from transformers.modeling_utils import PreTrainedModel from transformers.utils import logging from fla.modules import FusedCrossEntropyLoss, RMSNorm from fla.modules.layernorm import group_norm_fn from fla.modules.activations import swiglu_linear from fla.modules import RotaryEmbedding from einops import rearrange # 动态导入配置类 try: from .configuration_geometric import GeometricConfig except (ImportError, ValueError): try: from configuration_geometric import GeometricConfig except ImportError: from forgetting_transformer.model.geometric.configuration_geometric import GeometricConfig # 🔥 导入geometric attention from forgetting_transformer.ops.geometric_attention_final import geometric_attention logger = logging.get_logger(__name__) class ShiftLinear(nn.Module): """ Data-dependent token shift (from forgetting transformer) """ def __init__( self, input_dim: int, output_dim: int, num_heads: int, bias: bool, shift_bias: bool = False ): super().__init__() self.input_dim = input_dim self.output_dim = output_dim self.num_heads = num_heads assert self.output_dim % self.num_heads == 0 self.linear = nn.Linear(input_dim, output_dim, bias=bias) self.shift_proj = nn.Linear(input_dim, num_heads, bias=shift_bias) def forward(self, x: torch.Tensor, shift_state: Optional[torch.Tensor]) -> torch.Tensor: # 简化版本：不使用shift（geometric不需要） return self.linear(x) class GroupRMSNorm(nn.Module): """Group RMSNorm for multi-head normalization""" def __init__( self, num_groups: int, hidden_size: int, eps: float = 1e-6, elementwise_affine: bool = True ): super().__init__() self.num_groups = num_groups self.eps = eps self.elementwise_affine = elementwise_affine if self.elementwise_affine: self.weight = nn.Parameter(torch.ones(hidden_size)) else: self.register_parameter('weight', None) def forward(self, x: torch.Tensor) -> torch.Tensor: return group_norm_fn(x, self.num_groups, self.weight, self.eps) class GeometricAttention(nn.Module): """ Geometric Attention Layer 基于 "The Neural Data Router" 论文实现 """ def __init__( self, hidden_size: int, num_heads: int, num_kv_heads: Optional[int] = None, window_size: Optional[int] = None, max_position_embeddings: int = 2048, use_rope: bool = False, rope_base: float = 500000.0, qk_norm: bool = False, qk_norm_share_param_across_head: bool = False, use_k_shift: bool = False, use_v_shift: bool = False, use_geometric_normalize: bool = True, norm_eps: float = 1e-6, initializer_range: float = 0.02, layer_idx: Optional[int] = None, **kwargs ): """ Args: - hidden_size: dimension of hidden representations - num_heads: number of attention heads - num_kv_heads: (optional) For GQA, number of key-value heads - window_size: (optional) used for sliding window - max_position_embeddings: maximum sequence length - use_rope: whether to use rotary embeddings - rope_base: base for RoPE - qk_norm: Whether to use qk_norm - qk_norm_share_param_across_head: In QK-norm, whether to share params - use_k_shift: Whether to use data-dependent key shift - use_v_shift: Whether to use data-dependent value shift - use_geometric_normalize: Whether to normalize geometric attention weights - norm_eps: epsilon for normalization - initializer_range: standard deviation for initialization - layer_idx: The block index of this layer (for KV-cache) """ super().__init__() self.num_heads = num_heads if num_kv_heads is None: self.num_kv_heads = self.num_heads else: raise NotImplementedError("GQA has not been tested.") self.num_kv_heads = num_kv_heads self.num_kv_groups = num_heads // self.num_kv_heads self.hidden_size = hidden_size self.head_dim = self.hidden_size // self.num_heads self.kv_dim = self.num_kv_heads * self.head_dim self.window_size = window_size self.max_position_embeddings = max_position_embeddings self.layer_idx = layer_idx self.use_geometric_normalize = use_geometric_normalize # QKV projections self.q_proj = nn.Linear(self.hidden_size, self.hidden_size, bias=False) if use_k_shift: self.k_proj = ShiftLinear(self.hidden_size, self.kv_dim, self.num_heads, bias=False) else: self.k_proj = nn.Linear(self.hidden_size, self.kv_dim, bias=False) if use_v_shift: self.v_proj = ShiftLinear(self.hidden_size, self.kv_dim, self.num_heads, bias=False) else: self.v_proj = nn.Linear(self.hidden_size, self.kv_dim, bias=False) self.o_proj = nn.Linear(self.hidden_size, self.hidden_size, bias=False) self.use_k_shift = use_k_shift self.use_v_shift = use_v_shift # RoPE (optional) if use_rope: self.rotary = RotaryEmbedding(self.head_dim, base=rope_base) else: self.rotary = None # QK normalization (optional) self.qk_norm = qk_norm self.qk_norm_share_param_across_head = qk_norm_share_param_across_head if qk_norm: if self.qk_norm_share_param_across_head: self.q_norm = RMSNorm(self.head_dim) self.k_norm = RMSNorm(self.head_dim) else: self.q_norm = GroupRMSNorm(num_groups=self.num_heads, hidden_size=self.hidden_size, eps=norm_eps) self.k_norm = GroupRMSNorm(num_groups=self.num_heads, hidden_size=self.hidden_size, eps=norm_eps) self.initializer_range = initializer_range self.apply(self._initialize_weights) def _initialize_weights(self, module: nn.Module): if isinstance(module, nn.Linear): nn.init.normal_(module.weight, mean=0.0, std=self.initializer_range) if module.bias is not None: nn.init.zeros_(module.bias) def forward( self, hidden_states: torch.Tensor, attention_mask: Optional[torch.LongTensor] = None, past_key_values: Optional[Cache] = None, output_attentions: bool = False, use_cache: bool = False, **kwargs, ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]: """ Forward pass of geometric attention """ batch_size, q_len, _ = hidden_states.size() # Geometric attention不使用shift，设为None key_shift_state = None value_shift_state = None # QKV projections q = self.q_proj(hidden_states) if self.use_k_shift: k = self.k_proj(hidden_states, key_shift_state) else: k = self.k_proj(hidden_states) if self.use_v_shift: v = self.v_proj(hidden_states, value_shift_state) else: v = self.v_proj(hidden_states) # QK normalization (optional) if self.qk_norm and (not self.qk_norm_share_param_across_head): q = self.q_norm(q).to(q.dtype) k = self.k_norm(k).to(k.dtype) # Reshape for multi-head q = rearrange(q, '... (h d) -> ... h d', h=self.num_heads) k = rearrange(k, '... (h d) -> ... h d', h=self.num_kv_heads) v = rearrange(v, 'b t (h d) -> b t h d', h=self.num_kv_heads) if self.qk_norm and (self.qk_norm_share_param_across_head): q = self.q_norm(q).to(q.dtype) k = self.k_norm(k).to(k.dtype) # RoPE (optional) seqlen_offset, max_seqlen = 0, q.shape[1] if past_key_values is not None: seqlen_offset = past_key_values.get_seq_length(self.layer_idx) if hasattr(past_key_values, 'get_seq_length') else 0 max_seqlen = q.shape[1] + seqlen_offset if attention_mask is not None: seqlen_offset = (seqlen_offset + attention_mask.sum(-1) - attention_mask.shape[-1]) max_seqlen = q.shape[1] + max(seqlen_offset) if self.max_position_embeddings is not None: max_seqlen = max(max_seqlen, self.max_position_embeddings) if self.rotary is not None: q, k = self.rotary(q, k, seqlen_offset, max_seqlen) # Update KV cache if needed if past_key_values is not None and use_cache: # 使用标准的DynamicCache接口 if hasattr(past_key_values, 'update'): k_cache = rearrange(k, 'b t h d -> b h t d') v_cache = rearrange(v, 'b t h d -> b h t d') past_key_values.update(k_cache, v_cache, self.layer_idx) # 注意：这里不需要重新赋值k和v，因为我们在训练时不使用cache # Handle GQA (if enabled) if self.num_kv_groups > 1: k = rearrange(k.unsqueeze(-2).repeat(1, 1, 1, self.num_kv_groups, 1), 'b t h g d -> b t (h g) d') v = rearrange(v.unsqueeze(-2).repeat(1, 1, 1, self.num_kv_groups, 1), 'b t h g d -> b t (h g) d') # 🔥 Geometric Attention (核心) if attention_mask is not None: B, T = attention_mask.size() seq_start = T - attention_mask.sum(dim=-1) o = geometric_attention( q, k, v, head_first=False, seq_start=seq_start, sm_scale=1 / math.sqrt(self.head_dim), normalize=self.use_geometric_normalize, ) else: o = geometric_attention( q, k, v, head_first=False, sm_scale=1 / math.sqrt(self.head_dim), normalize=self.use_geometric_normalize, ) # Reshape output o = o.reshape(batch_size, q_len, self.hidden_size) # Output projection o = self.o_proj(o) # Attention weights (if requested) attentions = None if output_attentions: # 简化版：不返回详细的attention weights attentions = None return o, attentions, past_key_values class GeometricMLP(nn.Module): """ MLP层 (与ForgettingTransformer完全相同) """ def __init__( self, hidden_size: int, hidden_ratio: Optional[float] = None, intermediate_size: Optional[int] = None, hidden_act: str = 'swish' ): super().__init__() self.hidden_size = hidden_size if hidden_ratio is None: hidden_ratio = 4 if intermediate_size is None: intermediate_size = int(hidden_size * hidden_ratio * 2 / 3) intermediate_size = 256 * ((intermediate_size + 256 - 1) // 256) self.hidden_ratio = hidden_ratio self.intermediate_size = intermediate_size self.gate_proj = nn.Linear(self.hidden_size, self.intermediate_size * 2, bias=False) self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=False) self.act_fn = ACT2FN[hidden_act] self.hidden_act = hidden_act def forward(self, x): y = self.gate_proj(x) gate, y = y.chunk(2, dim=-1) return self.down_proj(self.act_fn(gate) * y) class GeometricBlock(nn.Module): """ Transformer Block with Geometric Attention """ def __init__(self, config: GeometricConfig, layer_idx: int): super().__init__() self.hidden_size = config.hidden_size self.attn_norm = RMSNorm( hidden_size=config.hidden_size, eps=config.norm_eps ) self.attn = GeometricAttention( hidden_size=config.hidden_size, num_heads=config.num_heads, num_kv_heads=config.num_kv_heads, window_size=config.window_size, max_position_embeddings=config.max_position_embeddings, use_rope=config.use_rope, rope_base=config.rope_base, qk_norm=config.qk_norm, qk_norm_share_param_across_head=config.qk_norm_share_param_across_head, use_k_shift=config.use_k_shift, use_v_shift=config.use_v_shift, use_geometric_normalize=config.use_geometric_normalize, norm_eps=config.norm_eps, initializer_range=config.initializer_range, layer_idx=layer_idx ) self.mlp_norm = RMSNorm( hidden_size=config.hidden_size, eps=config.norm_eps ) self.mlp = GeometricMLP( hidden_size=config.hidden_size, hidden_ratio=config.hidden_ratio, intermediate_size=config.intermediate_size, hidden_act=config.hidden_act ) def forward( self, hidden_states: torch.Tensor, attention_mask: Optional[torch.Tensor] = None, past_key_values: Optional[Cache] = None, output_attentions: bool = False, use_cache: bool = False, **kwargs ) -> Tuple[torch.FloatTensor, Optional[Tuple[torch.FloatTensor, torch.FloatTensor]]]: # Attention block with residual residual = hidden_states hidden_states = self.attn_norm(hidden_states) hidden_states, attentions, past_key_values = self.attn( hidden_states=hidden_states, attention_mask=attention_mask, past_key_values=past_key_values, output_attentions=output_attentions, use_cache=use_cache, ) hidden_states = residual + hidden_states # MLP block with residual residual = hidden_states hidden_states = self.mlp_norm(hidden_states) hidden_states = self.mlp(hidden_states) hidden_states = residual + hidden_states outputs = (hidden_states, attentions, past_key_values) return outputs class GeometricPreTrainedModel(PreTrainedModel): config_class = GeometricConfig supports_gradient_checkpointing = True _no_split_modules = ["GeometricBlock"] def _init_weights(self, module): std = self.config.initializer_range if isinstance(module, nn.Linear): module.weight.data.normal_(mean=0.0, std=std) if module.bias is not None: module.bias.data.zero_() elif isinstance(module, nn.Embedding): module.weight.data.normal_(mean=0.0, std=std) if module.padding_idx is not None: module.weight.data[module.padding_idx].zero_() class GeometricModel(GeometricPreTrainedModel): """ Geometric Transformer Model """ def __init__(self, config: GeometricConfig): super().__init__(config) self.padding_idx = config.pad_token_id self.vocab_size = config.vocab_size self.embeddings = nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx) self.layers = nn.ModuleList([GeometricBlock(config, layer_idx) for layer_idx in range(config.num_hidden_layers)]) self.norm = RMSNorm(config.hidden_size, eps=config.norm_eps) self.gradient_checkpointing = False self.post_init() def get_input_embeddings(self): return self.embeddings def set_input_embeddings(self, value): self.embeddings = value def forward( self, input_ids: torch.LongTensor = None, attention_mask: Optional[torch.Tensor] = None, past_key_values: Optional[Cache] = 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, **kwargs ) -> Union[Tuple, BaseModelOutputWithPast]: output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions output_hidden_states = output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states use_cache = use_cache if use_cache is not None else self.config.use_cache return_dict = return_dict if return_dict is not None else self.config.use_return_dict # Embeddings if inputs_embeds is None: inputs_embeds = self.embeddings(input_ids) if use_cache and past_key_values is None: past_key_values = DynamicCache() hidden_states = inputs_embeds # Layers all_hidden_states = () if output_hidden_states else None all_self_attns = () if output_attentions else None for layer in self.layers: if output_hidden_states: all_hidden_states += (hidden_states,) if self.gradient_checkpointing and self.training: layer_outputs = self._gradient_checkpointing_func( layer.__call__, hidden_states, attention_mask, past_key_values, output_attentions, use_cache, ) else: layer_outputs = layer( hidden_states, attention_mask=attention_mask, past_key_values=past_key_values, output_attentions=output_attentions, use_cache=use_cache, ) hidden_states = layer_outputs[0] if output_attentions: all_self_attns += (layer_outputs[1],) past_key_values = layer_outputs[2] hidden_states = self.norm(hidden_states) if output_hidden_states: all_hidden_states += (hidden_states,) if not return_dict: return tuple(v for v in [hidden_states, past_key_values, all_hidden_states, all_self_attns] if v is not None) return BaseModelOutputWithPast( last_hidden_state=hidden_states, past_key_values=past_key_values, hidden_states=all_hidden_states, attentions=all_self_attns, ) class GeometricForCausalLM(GeometricPreTrainedModel): """ Geometric Transformer for Causal Language Modeling """ _tied_weights_keys = ["lm_head.weight"] def __init__(self, config): super().__init__(config) self.model = GeometricModel(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.embeddings def set_input_embeddings(self, value): self.model.embeddings = value def get_output_embeddings(self): return self.lm_head def set_output_embeddings(self, new_embeddings): self.lm_head = new_embeddings def forward( self, input_ids: torch.LongTensor = None, attention_mask: Optional[torch.Tensor] = None, past_key_values: Optional[Cache] = 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, **kwargs ) -> Union[Tuple, CausalLMOutputWithPast]: return_dict = return_dict if return_dict is not None else self.config.use_return_dict # Model forward outputs = self.model( input_ids=input_ids, attention_mask=attention_mask, past_key_values=past_key_values, inputs_embeds=inputs_embeds, use_cache=use_cache, output_attentions=output_attentions, output_hidden_states=output_hidden_states, return_dict=return_dict, ) hidden_states = outputs[0] loss = None if labels is not None: if self.config.fuse_cross_entropy: loss_fct = FusedCrossEntropyLoss(inplace_backward=True, reduction='none') else: loss_fct = nn.CrossEntropyLoss(reduction='none') logits = self.lm_head(hidden_states) # Enable model parallelism labels = labels.to(logits.device) loss = loss_fct(logits.view(-1, self.config.vocab_size), labels.view(-1)) loss = loss.view(*labels.size()) # Reshape to [batch, seq_len] del logits logits = None else: logits = self.lm_head(hidden_states) if not return_dict: output = (logits,) + outputs[1:] return (loss,) + output if loss is not None else output return CausalLMOutputWithPast( loss=loss, logits=logits, past_key_values=outputs.past_key_values, hidden_states=outputs.hidden_states, attentions=outputs.attentions, )