""" Fusion 完整模型定义（v2 - 可实例化可运行） 集成： 1. SBLA 注意力（滑动分块潜注意力）- 真实实现 2. Thinking Dial（动态推理强度控制）- 通过特殊 token 3. 标准 Transformer 架构 + KV Cache 支持 修复（v2）： - FusionModel 现在可以完整实例化和运行 - SBLA 注意力已正确集成到每一层 - 支持 causal mask、padding mask - generate() 方法支持 KV cache 加速推理 - 配置文件与代码完全对齐 使用方法： from models.fusion_model import FusionModel, FusionConfig config = FusionConfig( vocab_size=10000, hidden_size=256, num_hidden_layers=4, num_attention_heads=8, block_size=64, latent_dim=16, ) model = FusionModel(config) input_ids = torch.randint(0, 10000, (2, 128)) outputs = model(input_ids=input_ids, labels=input_ids) print(f"Loss: {outputs['loss'].item()}") 作者：zhan1206 项目：Fusion - 六边形开源大模型 许可证：Apache 2.0 """ import torch import torch.nn as nn import torch.nn.functional as F from transformers import PretrainedConfig, PreTrainedModel, GenerationMixin from transformers.modeling_outputs import CausalLMOutputWithPast from typing import Optional, Tuple, Dict, Any import math # H4-H6: Use try/except for relative imports with sys.path fallback try: from .sbla_attention import SBLAttention except ImportError: from models.sbla_attention import SBLAttention class FusionConfig(PretrainedConfig): """Fusion 模型配置""" model_type = "fusion" def __init__( self, vocab_size: int = 100000, hidden_size: int = 4096, num_hidden_layers: int = 32, num_attention_heads: int = 32, num_key_value_heads: Optional[int] = None, intermediate_size: int = 11008, hidden_act: str = "silu", hidden_dropout_prob: float = 0.1, attention_probs_dropout_prob: float = 0.1, max_position_embeddings: int = 32768, initializer_range: float = 0.02, rms_norm_eps: float = 1e-6, use_cache: bool = True, tie_word_embeddings: bool = False, # SBLA parameters block_size: int = 512, latent_dim: int = 64, window_size: Optional[int] = None, sbla_mode: str = "pure_sbla", # Thinking Dial parameters enable_thinking_dial: bool = True, num_thinking_depths: int = 4, **kwargs, ): super().__init__(**kwargs) self.vocab_size = vocab_size self.hidden_size = hidden_size self.num_hidden_layers = num_hidden_layers self.num_attention_heads = num_attention_heads self.num_key_value_heads = num_key_value_heads or num_attention_heads self.intermediate_size = intermediate_size self.hidden_act = hidden_act self.hidden_dropout_prob = hidden_dropout_prob self.attention_probs_dropout_prob = attention_probs_dropout_prob 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 # SBLA parameters self.block_size = block_size self.latent_dim = latent_dim self.window_size = window_size or block_size # [M8 FIX] Remove redundant sbla_window_size self.sbla_mode = sbla_mode # Thinking Dial parameters self.enable_thinking_dial = enable_thinking_dial self.num_thinking_depths = num_thinking_depths # RoPE parameters self.rope_theta = kwargs.pop('rope_theta', 10000.0) # H1-H3: Register FusionConfig with AutoConfig try: from transformers import AutoConfig AutoConfig.register("fusion", FusionConfig) except (ImportError, ValueError): pass # Already registered or AutoConfig unavailable class RotaryEmbedding(nn.Module): """Rotary Position Embedding (RoPE) for positional encoding in attention.""" def __init__(self, dim, max_position_embeddings=2048, base=10000.0): super().__init__() self.dim = dim self.max_position_embeddings = max_position_embeddings self.base = base inv_freq = 1.0 / (base ** (torch.arange(0, dim, 2).float() / dim)) self.register_buffer('inv_freq', inv_freq) def forward(self, seq_len, device=None): t = torch.arange(seq_len, device=device, dtype=self.inv_freq.dtype) freqs = torch.outer(t, self.inv_freq) emb = torch.cat((freqs, freqs), dim=-1) return emb def rotate_half(x): """Rotate half the hidden dims of the input.""" 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, position_ids=None): """Apply rotary position embedding to query and key tensors. Args: q: (batch, num_heads, seq_len, head_dim) k: (batch, num_kv_heads, seq_len, head_dim) cos: (kv_seq_len, head_dim) cosine part of rotary embedding sin: (kv_seq_len, head_dim) sine part of rotary embedding position_ids: (batch, seq_len) position ids for slicing cos/sin Returns: Tuple of (q_embed, k_embed) with rotary position encoding applied. """ if position_ids is not None: # N6 FIX: Slice cos/sin by position_ids to match actual Q/K positions # position_ids: (batch, seq_len), cos/sin: (kv_seq_len, head_dim) cos = cos[position_ids].unsqueeze(1) # (batch, 1, seq_len, head_dim) sin = sin[position_ids].unsqueeze(1) # (batch, 1, seq_len, head_dim) else: # Fallback: broadcast for full-sequence (prefill) when position_ids not provided cos = cos.unsqueeze(0).unsqueeze(0) # (1, 1, kv_seq_len, head_dim) sin = sin.unsqueeze(0).unsqueeze(0) q_embed = (q * cos) + (rotate_half(q) * sin) k_embed = (k * cos) + (rotate_half(k) * sin) return q_embed, k_embed class RMSNorm(nn.Module): """RMSNorm（均方根层归一化）""" def __init__(self, hidden_size: int, eps: float = 1e-6): super().__init__() self.weight = nn.Parameter(torch.ones(hidden_size)) self.eps = eps def forward(self, x: torch.Tensor) -> torch.Tensor: variance = x.float().pow(2).mean(-1, keepdim=True) x = x.float() * torch.rsqrt(variance + self.eps) return (x * self.weight).to(x.dtype) class FusionAttention(nn.Module): """ Fusion Attention Layer — delegates to the canonical SBLAttention implementation. This is the unified entry point for attention in FusionModel. All SBLA logic (block latents, causal/window masks, padding) lives in models/sbla_attention.py::SBLAttention. This wrapper adds KV cache support and config-driven mode selection (pure_sbla / hybrid). See: models/sbla_attention.py::SBLAttention """ def __init__(self, config: FusionConfig): super().__init__() mode = getattr(config, 'sbla_mode', 'hybrid') self.sbla = SBLAttention( hidden_size=config.hidden_size, num_heads=config.num_attention_heads, block_size=config.block_size, latent_dim=config.latent_dim, dropout=config.attention_probs_dropout_prob, window_size=config.window_size, # use dedicated window_size field mode=mode, num_key_value_heads=config.num_key_value_heads, ) # M14 FIX: Add RoPE - Rotary Position Embedding head_dim = config.hidden_size // config.num_attention_heads rope_theta = getattr(config, 'rope_theta', 10000.0) self.rotary_emb = RotaryEmbedding( dim=head_dim, max_position_embeddings=config.max_position_embeddings, base=rope_theta, ) # Separate Q/K/V projections for RoPE application # (SBLAttention has its own q/k/v, but we need access before RoPE) self.q_proj = nn.Linear(config.hidden_size, config.num_attention_heads * head_dim, bias=False) self.k_proj = nn.Linear(config.hidden_size, config.num_key_value_heads * head_dim, bias=False) self.v_proj = nn.Linear(config.hidden_size, config.num_key_value_heads * head_dim, bias=False) # [S1 FIX] o_proj is sbla.out_proj — needed so LoRA target_modules=["o_proj"] works # We assign it after sbla is created (above). This shares the parameter. self.o_proj = self.sbla.out_proj # Tell HF that o_proj.weight is an alias of sbla.out_proj.weight. # This allows save_pretrained to deduplicate, and from_pretrained to # know that o_proj.weight should be reconstructed from sbla.out_proj.weight. _tied_weights_keys = {"o_proj.weight": "sbla.out_proj.weight"} def forward( self, hidden_states: torch.Tensor, attention_mask: Optional[torch.Tensor] = None, past_key_value: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, use_cache: bool = False, position_ids: Optional[torch.Tensor] = None, ) -> Tuple[torch.Tensor, Optional[Tuple[torch.Tensor, torch.Tensor]]]: # M14 FIX: Apply RoPE to Q and K before SBLAttention batch_size, seq_len, _ = hidden_states.shape head_dim = self.sbla.head_dim num_kv_groups = self.sbla.num_kv_groups # Project Q/K/V Q = self.q_proj(hidden_states).view(batch_size, seq_len, self.sbla.num_heads, head_dim).transpose(1, 2) K = self.k_proj(hidden_states).view(batch_size, seq_len, self.sbla.num_key_value_heads, head_dim).transpose(1, 2) V = self.v_proj(hidden_states).view(batch_size, seq_len, self.sbla.num_key_value_heads, head_dim).transpose(1, 2) # Compute RoPE embeddings kv_seq_len = seq_len if past_key_value is not None: kv_seq_len = past_key_value[0].shape[2] + seq_len emb = self.rotary_emb(kv_seq_len, device=hidden_states.device) cos = emb.cos() sin = emb.sin() # N6 FIX: Build position_ids for proper RoPE slicing if position_ids is None: if past_key_value is not None: offset = past_key_value[0].shape[2] position_ids = torch.arange(offset, offset + seq_len, device=hidden_states.device).unsqueeze(0) else: position_ids = torch.arange(seq_len, device=hidden_states.device).unsqueeze(0) # Apply RoPE with position_ids to prevent broadcast mismatch during incremental generation Q, K = apply_rotary_pos_emb(Q, K, cos, sin, position_ids=position_ids) # Store RoPE'd K/V in SBLAttention's cache for incremental generation # S1 FIXED: KV Cache now works natively through SBLAttention. # We pass the RoPE'd Q/K/V by injecting them into SBLAttention. output, present_key_value = self.sbla.forward_with_qkv( Q, K, V, attention_mask, past_key_value=past_key_value, use_cache=use_cache, ) return output, present_key_value class FusionLayer(nn.Module): """Fusion Transformer 层""" def __init__(self, config: FusionConfig, layer_idx: int): super().__init__() self.input_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) self.attention = FusionAttention(config) self.post_attention_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) # SwiGLU FFN 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) self.dropout = nn.Dropout(config.hidden_dropout_prob) def forward( self, hidden_states: torch.Tensor, attention_mask: Optional[torch.Tensor] = None, past_key_value: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, use_cache: bool = False, position_ids: Optional[torch.Tensor] = None, # [M6 FIX] Added **kwargs, ) -> Tuple[torch.Tensor, Optional[Tuple[torch.Tensor, torch.Tensor]]]: residual = hidden_states hidden_states = self.input_layernorm(hidden_states) attn_output, present_key_value = self.attention( hidden_states, attention_mask, past_key_value=past_key_value if past_key_value is not None else None, use_cache=use_cache, position_ids=position_ids, # [M6 FIX] ) hidden_states = residual + self.dropout(attn_output) residual = hidden_states hidden_states = self.post_attention_layernorm(hidden_states) gate = F.silu(self.gate_proj(hidden_states)) up = self.up_proj(hidden_states) ffn_output = self.down_proj(gate * up) hidden_states = residual + self.dropout(ffn_output) return hidden_states, present_key_value class FusionModel(PreTrainedModel, GenerationMixin): """ Fusion 完整模型（v2 - 可实例化可运行） 支持 HuggingFace PreTrainedModel 全接口 """ config_class = FusionConfig supports_gradient_checkpointing = True _no_split_modules = ["FusionAttention", "SBLAttention"] # o_proj is a reference to sbla.out_proj (same parameter) for LoRA compatibility. # FusionAttention._tied_weights_keys declares the alias so save_pretrained # can deduplicate. tie_weights() re-links after loading. def __init__(self, config: FusionConfig): super().__init__(config) self.config = config # Embeddings self.embeddings = nn.Embedding(config.vocab_size, config.hidden_size, padding_idx=0) self.dropout = nn.Dropout(config.hidden_dropout_prob) # Transformer 层 self.layers = nn.ModuleList([ FusionLayer(config, layer_idx=i) for i in range(config.num_hidden_layers) ]) # Final Norm self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) # LM Head self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False) if config.tie_word_embeddings: self.lm_head.weight = self.embeddings.weight self.post_init() def _init_weights(self, module): """Initialize weights — called by HF's post_init/init_weights. For from_pretrained, HF loads state_dict first then calls init_weights which would overwrite the loaded weights. We guard against this by checking if we're in the from_pretrained flow (model has _is_hf_initialized flag set after loading). """ # Skip re-initialization if weights were already loaded from checkpoint if hasattr(self, '_is_hf_initialized') and self._is_hf_initialized: return if isinstance(module, nn.Linear): module.weight.data.normal_(mean=0.0, std=self.config.initializer_range) if module.bias is not None: module.bias.data.zero_() elif isinstance(module, nn.Embedding): module.weight.data.normal_(mean=0.0, std=self.config.initializer_range) if module.padding_idx is not None: module.weight.data[module.padding_idx].zero_() elif isinstance(module, RMSNorm): nn.init.ones_(module.weight) def tie_weights(self, **kwargs): """Re-link o_proj aliases after loading weights. FusionAttention.o_proj is an alias for FusionAttention.sbla.out_proj. After from_pretrained loads the state_dict, o_proj gets its own copy. We must re-alias it to share the parameter with sbla.out_proj. """ super().tie_weights(**kwargs) for layer in self.layers: if hasattr(layer, 'attention') and hasattr(layer.attention, 'o_proj'): layer.attention.o_proj = layer.attention.sbla.out_proj @classmethod def from_pretrained(cls, pretrained_model_name_or_path, *args, **kwargs): """Load a pretrained FusionModel, ensuring correct weight loading. Overrides the default to use load_state_dict directly, which correctly handles our tied o_proj/sbla.out_proj parameters. The default HF 5.x loading path (convert_and_load_state_dict_in_model) does not properly restore weights for custom module hierarchies. """ from safetensors.torch import load_file as sf_load import os # Load config from transformers import AutoConfig config = kwargs.pop('config', None) or AutoConfig.from_pretrained(pretrained_model_name_or_path) # Create model with random weights model = cls(config) # Load state dict from safetensors sf_path = os.path.join(pretrained_model_name_or_path, 'model.safetensors') if os.path.exists(sf_path): sd = sf_load(sf_path) else: # Try PyTorch format pt_path = os.path.join(pretrained_model_name_or_path, 'pytorch_model.bin') sd = torch.load(pt_path, map_location='cpu', weights_only=True) # Add o_proj from sbla.out_proj (removed during save as tied weight) for key in list(sd.keys()): if 'sbla.out_proj.weight' in key: oproj_key = key.replace('sbla.out_proj', 'o_proj') sd[oproj_key] = sd[key].clone() model.load_state_dict(sd, strict=False) model.tie_weights() # Re-link o_proj -> sbla.out_proj return model def forward( self, input_ids: Optional[torch.Tensor] = None, attention_mask: Optional[torch.Tensor] = None, labels: Optional[torch.Tensor] = None, past_key_values: Optional[Tuple] = None, inputs_embeds: Optional[torch.Tensor] = None, use_cache: Optional[bool] = None, position_ids: Optional[torch.Tensor] = None, # [M6 FIX] Added return_dict: Optional[bool] = True, **kwargs, ) -> CausalLMOutputWithPast: use_cache = use_cache if use_cache is not None else self.config.use_cache # [M6 FIX] Extract position_ids from kwargs if not passed explicitly if 'position_ids' in kwargs: position_ids = kwargs.pop('position_ids') # else: keep the explicit position_ids parameter # Embeddings if inputs_embeds is not None: hidden_states = inputs_embeds elif input_ids is not None: hidden_states = self.embeddings(input_ids) hidden_states = self.dropout(hidden_states) else: raise ValueError("Either input_ids or inputs_embeds must be provided") # Use the already-resolved use_cache from parameter, don't re-override from kwargs if past_key_values is not None: use_cache = True present_key_values = () 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 layer_outputs, cache = layer( hidden_states, attention_mask=attention_mask, past_key_value=layer_past, use_cache=use_cache, position_ids=position_ids, # [M6 FIX] ) hidden_states = layer_outputs if use_cache: present_key_values = present_key_values + (cache,) # Final norm hidden_states = self.norm(hidden_states) # LM Head logits = self.lm_head(hidden_states) # Loss loss = None if labels is not None: shift_logits = logits[..., :-1, :].contiguous() shift_labels = labels[..., 1:].contiguous() loss_fct = nn.CrossEntropyLoss() loss = loss_fct(shift_logits.view(-1, self.config.vocab_size), shift_labels.view(-1)) # C2/C3/C5: Return CausalLMOutputWithPast instead of plain dict if not return_dict: output = (logits,) + (present_key_values,) if present_key_values is not None else (logits,) return ((loss,) + output) if loss is not None else output return CausalLMOutputWithPast( loss=loss, logits=logits, past_key_values=present_key_values, hidden_states=None, attentions=None, ) @torch.no_grad() def generate( self, input_ids: torch.Tensor, max_new_tokens: int = 256, temperature: float = 1.0, top_p: float = 0.95, do_sample: bool = True, pad_token_id: Optional[int] = None, eos_token_id: Optional[int] = None, return_dict_in_generate: bool = False, logits_hook: Optional[callable] = None, # N10 FIX: Hook for ThinkingDial logit bias past_key_values: Optional[Tuple] = None, # N17 FIX: Accept pre-computed KV cache thinking_depth: Optional[int] = None, # Thinking Dial depth control **kwargs, ) -> CausalLMOutputWithPast: """Generate text with KV cache and SBLA incremental support. Args: input_ids: Input token IDs max_new_tokens: Maximum number of tokens to generate temperature: Sampling temperature (must be > 0) top_p: Nucleus sampling threshold (0 < top_p <= 1) do_sample: Whether to sample or use greedy decoding pad_token_id: Padding token ID eos_token_id: End-of-sequence token ID return_dict_in_generate: If True, return CausalLMOutputWithPast Returns: CausalLMOutputWithPast if return_dict_in_generate, else generated token IDs tensor """ # [M4 FIX] Parameter validation if temperature <= 0: raise ValueError(f"temperature must be > 0, got {temperature}") if not (0 < top_p <= 1.0): raise ValueError(f"top_p must be in (0, 1], got {top_p}") if max_new_tokens <= 0: raise ValueError(f"max_new_tokens must be > 0, got {max_new_tokens}") batch_size = input_ids.shape[0] device = input_ids.device eos_token_id = eos_token_id or getattr(self.config, "eos_token_id", None) self.eval() # N17 FIX: Support pre-computed KV cache for generate_samples reuse if past_key_values is not None: past_seq_len = past_key_values[0][0].shape[2] # K shape: (B, H, S, D) generated = input_ids else: past_seq_len = 0 # [M1 FIX] Track position for RoPE generated = input_ids.clone() logits_hook = kwargs.pop('logits_hook', logits_hook) # Thinking Dial integration: if thinking_depth is provided and model has # the Thinking Dial embedding/gate, build a logits_hook automatically. # This unifies the two control mechanisms (text token vs architecture-level) # so users can simply pass thinking_depth=N without wrapping in ThinkingDialModel. if thinking_depth is not None and hasattr(self, 'thinking_embedding'): from models.thinking_dial import ThinkingDialModel, ThinkingConfig hook = ThinkingDialModel._build_thinking_logits_hook( thinking_depth, batch_size, device, getattr(self, 'thinking_config', ThinkingConfig()), self.thinking_embedding, self.thinking_gate, self.lm_head, ) if hook is not None: logits_hook = hook # Thinking Dial hook takes precedence for _ in range(max_new_tokens): if past_key_values is not None: current_input = generated[:, -1:] cur_seq_len = 1 else: current_input = generated cur_seq_len = generated.shape[1] # [M1 FIX] Compute position_ids for RoPE position_ids = torch.arange(past_seq_len, past_seq_len + cur_seq_len, device=device, dtype=torch.long).unsqueeze(0) outputs = self.forward( input_ids=current_input, past_key_values=past_key_values, use_cache=True, return_dict=True, position_ids=position_ids, # [M1 FIX] ) logits = outputs.logits past_key_values = outputs.past_key_values # N10 FIX: Apply logits hook if provided (for ThinkingDialModel depth bias) if logits_hook is not None: logits = logits_hook(logits) next_token_logits = logits[:, -1, :] / max(temperature, 1e-8) if do_sample and top_p < 1.0: logits_before_mask = next_token_logits.clone() # [F5 FIX] Save for fallback sorted_logits, sorted_indices = torch.sort(next_token_logits, descending=True) cumulative_probs = torch.cumsum(F.softmax(sorted_logits, dim=-1), dim=-1) sorted_indices_to_remove = cumulative_probs > top_p sorted_indices_to_remove[..., 1:] = sorted_indices_to_remove[..., :-1].clone() sorted_indices_to_remove[..., 0] = 0 indices_to_remove = sorted_indices_to_remove.scatter(1, sorted_indices, sorted_indices_to_remove) next_token_logits.masked_fill_(indices_to_remove, float('-inf')) # [F5 FIX] Guard against all-tokens-masked: keep top-1 token if (next_token_logits == float('-inf')).all(): next_token_logits = logits_before_mask if do_sample: probs = F.softmax(next_token_logits, dim=-1) next_token = torch.multinomial(probs, num_samples=1) else: next_token = torch.argmax(next_token_logits, dim=-1, keepdim=True) generated = torch.cat([generated, next_token], dim=1) past_seq_len += cur_seq_len # [M1 FIX] Update position counter if eos_token_id is not None and (next_token == eos_token_id).all(): break # [M3 FIX] Return CausalLMOutputWithPast for API consistency if return_dict_in_generate: return CausalLMOutputWithPast( loss=None, logits=None, past_key_values=past_key_values, hidden_states=None, attentions=None, ), generated # Return tuple (output, sequences) return generated def prepare_inputs_for_generation(self, input_ids: torch.Tensor, past_key_values=None, **kwargs): if past_key_values is not None: input_ids = input_ids[:, -1:] return {"input_ids": input_ids, "past_key_values": past_key_values, "use_cache": True} if __name__ == "__main__": print("[TEST] Testing Fusion Model (v2)...") config = FusionConfig( vocab_size=10000, hidden_size=256, num_hidden_layers=2, num_attention_heads=4, intermediate_size=512, block_size=64, latent_dim=16, sbla_mode="pure_sbla", max_position_embeddings=256, ) model = FusionModel(config) param_count = sum(p.numel() for p in model.parameters()) print(f"Model created with {param_count:,} parameters") batch_size, seq_len = 2, 128 input_ids = torch.randint(0, config.vocab_size, (batch_size, seq_len)) attention_mask = torch.ones(batch_size, seq_len) outputs = model(input_ids=input_ids, attention_mask=attention_mask, labels=input_ids, return_dict=True) assert outputs.loss is not None, "Loss should not be None" assert not torch.isnan(outputs.loss).item(), "Loss is NaN!" print(f"Loss={outputs.loss.item():.4f}, Logits={outputs.logits.shape}") print("\n[ALL TESTS PASSED] Fusion Model v2 fully functional.")