"""Hugging Face Transformers implementation of Vortex Alpha. The model intentionally uses the same readable reference tensor names as the published safetensors files. It implements the full-prefix path and does not claim a KV cache; ``use_cache`` is therefore false by default. """ from __future__ import annotations import math from typing import Optional import torch import torch.nn.functional as F from torch import nn from transformers import PreTrainedModel from transformers.modeling_outputs import CausalLMOutputWithPast from .configuration_vortex import VortexConfig class VortexRMSNorm(nn.Module): def __init__(self, dim: int, eps: float) -> None: super().__init__() self.weight = nn.Parameter(torch.ones(dim)) self.eps = eps def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: return F.rms_norm(hidden_states, (hidden_states.shape[-1],), self.weight, self.eps) class VortexRotaryEmbedding(nn.Module): def __init__(self, dim: int, max_seq_len: int, theta: float) -> None: super().__init__() inv_freq = 1.0 / (theta ** (torch.arange(0, dim, 2).float() / dim)) positions = torch.arange(max_seq_len, dtype=torch.float32) frequencies = torch.outer(positions, inv_freq) angles = torch.cat((frequencies, frequencies), dim=-1) self.register_buffer("cos_cached", angles.cos()[None, None], persistent=False) self.register_buffer("sin_cached", angles.sin()[None, None], persistent=False) @staticmethod def rotate_half(x: torch.Tensor) -> torch.Tensor: half = x.shape[-1] // 2 return torch.cat((-x[..., half:], x[..., :half]), dim=-1) def forward( self, query: torch.Tensor, key: torch.Tensor, position_ids: Optional[torch.Tensor] = None, ) -> tuple[torch.Tensor, torch.Tensor]: if position_ids is None: cos = self.cos_cached[:, :, : query.shape[-2]] sin = self.sin_cached[:, :, : query.shape[-2]] else: cos = self.cos_cached[0, 0, position_ids].unsqueeze(1) sin = self.sin_cached[0, 0, position_ids].unsqueeze(1) cos = cos.to(dtype=query.dtype, device=query.device) sin = sin.to(dtype=query.dtype, device=query.device) return ( query * cos + self.rotate_half(query) * sin, key * cos + self.rotate_half(key) * sin, ) class VortexAttention(nn.Module): def __init__(self, config: VortexConfig) -> None: super().__init__() if config.num_attention_heads % config.num_key_value_heads: raise ValueError("num_attention_heads must divide evenly by num_key_value_heads") if config.num_attention_heads * config.head_dim != config.hidden_size: raise ValueError("num_attention_heads * head_dim must equal hidden_size") self.num_heads = config.num_attention_heads self.num_key_value_heads = config.num_key_value_heads self.head_dim = config.head_dim kv_dim = config.num_key_value_heads * config.head_dim self.q_proj = nn.Linear(config.hidden_size, config.hidden_size, bias=config.attention_bias) self.k_proj = nn.Linear(config.hidden_size, kv_dim, bias=config.attention_bias) self.v_proj = nn.Linear(config.hidden_size, kv_dim, bias=config.attention_bias) self.o_proj = nn.Linear(config.hidden_size, config.hidden_size, bias=config.attention_bias) self.q_norm = VortexRMSNorm(config.head_dim, config.rms_norm_eps) self.k_norm = VortexRMSNorm(config.head_dim, config.rms_norm_eps) self.rope = VortexRotaryEmbedding(config.head_dim, config.max_position_embeddings, config.rope_theta) def forward( self, hidden_states: torch.Tensor, attention_mask: Optional[torch.Tensor] = None, position_ids: Optional[torch.Tensor] = None, ) -> torch.Tensor: batch, seq_len, _ = hidden_states.shape query = self.q_proj(hidden_states).view(batch, seq_len, self.num_heads, self.head_dim).transpose(1, 2) key = self.k_proj(hidden_states).view(batch, seq_len, self.num_key_value_heads, self.head_dim).transpose(1, 2) value = self.v_proj(hidden_states).view(batch, seq_len, self.num_key_value_heads, self.head_dim).transpose(1, 2) query, key = self.rope(self.q_norm(query), self.k_norm(key), position_ids) if self.num_key_value_heads != self.num_heads: repeats = self.num_heads // self.num_key_value_heads key = key.repeat_interleave(repeats, dim=1) value = value.repeat_interleave(repeats, dim=1) sdpa_mask = None query_valid = None is_causal = attention_mask is None and seq_len > 1 if attention_mask is not None: if attention_mask.ndim == 2: valid = attention_mask.to(dtype=torch.bool, device=hidden_states.device) valid_keys = valid[:, None, None, :] query_valid = valid[:, None, :, None] causal = torch.ones((seq_len, seq_len), dtype=torch.bool, device=hidden_states.device).tril() sdpa_mask = valid_keys & causal[None, None, :, :] # SDPA returns NaN for an all-masked query row. Left padding # creates exactly those rows, so give pad queries one harmless # fallback key and zero their outputs after attention. fallback = torch.zeros_like(sdpa_mask) fallback[..., 0] = True sdpa_mask = torch.where(query_valid, sdpa_mask, fallback) elif attention_mask.ndim == 4: sdpa_mask = attention_mask else: raise ValueError("Vortex expects a 2-D or 4-D attention mask") output = F.scaled_dot_product_attention( query, key, value, attn_mask=sdpa_mask, dropout_p=0.0, is_causal=is_causal, ) if query_valid is not None: output = output * query_valid.to(dtype=output.dtype) output = output.transpose(1, 2).contiguous().view(batch, seq_len, -1) return self.o_proj(output) class VortexMLP(nn.Module): def __init__(self, config: VortexConfig) -> None: super().__init__() self.gate_proj = nn.Linear(config.hidden_size, config.intermediate_size, bias=config.attention_bias) self.up_proj = nn.Linear(config.hidden_size, config.intermediate_size, bias=config.attention_bias) self.down_proj = nn.Linear(config.intermediate_size, config.hidden_size, bias=config.attention_bias) def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: return self.down_proj(F.silu(self.gate_proj(hidden_states)) * self.up_proj(hidden_states)) class VortexDecoderBlock(nn.Module): def __init__(self, config: VortexConfig) -> None: super().__init__() self.norm1 = VortexRMSNorm(config.hidden_size, config.rms_norm_eps) self.attn = VortexAttention(config) self.norm2 = VortexRMSNorm(config.hidden_size, config.rms_norm_eps) self.ffn = VortexMLP(config) def forward( self, hidden_states: torch.Tensor, attention_mask: Optional[torch.Tensor] = None, position_ids: Optional[torch.Tensor] = None, ) -> torch.Tensor: hidden_states = hidden_states + self.attn(self.norm1(hidden_states), attention_mask, position_ids) hidden_states = hidden_states + self.ffn(self.norm2(hidden_states)) return hidden_states class VortexPreTrainedModel(PreTrainedModel): config_class = VortexConfig base_model_prefix = "vortex" supports_gradient_checkpointing = False class VortexForCausalLM(VortexPreTrainedModel): # The output projection is implemented directly with embed_tokens.weight, # so there is no second lm_head parameter to tie or load. _tied_weights_keys = None main_input_name = "input_ids" def __init__(self, config: VortexConfig) -> None: super().__init__(config) # Transformers 5.x uses this expanded mapping while loading from a # checkpoint. It is empty because the single embedding parameter is # both the input table and the output projection. self.all_tied_weights_keys = {} self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size) self.layers = nn.ModuleList(VortexDecoderBlock(config) for _ in range(config.num_hidden_layers)) self.norm = VortexRMSNorm(config.hidden_size, config.rms_norm_eps) self.post_init() def _init_weights(self, module: nn.Module) -> None: """Initialization hook used by both Transformers 4.x and 5.x.""" if isinstance(module, VortexRMSNorm): nn.init.ones_(module.weight) elif isinstance(module, nn.Linear): nn.init.normal_(module.weight, mean=0.0, std=0.02) if module.bias is not None: nn.init.zeros_(module.bias) elif isinstance(module, nn.Embedding): nn.init.normal_(module.weight, mean=0.0, std=0.02) def get_input_embeddings(self) -> nn.Module: return self.embed_tokens def set_input_embeddings(self, value: nn.Module) -> None: self.embed_tokens = value def get_output_embeddings(self) -> nn.Module: return self.embed_tokens def set_output_embeddings(self, new_embeddings: nn.Module) -> None: self.embed_tokens = new_embeddings def prepare_inputs_for_generation(self, input_ids: torch.Tensor, **kwargs) -> dict: # No KV cache is exported. Returning the full prefix keeps generation # correct, though slower than a cache-enabled implementation. return { "input_ids": input_ids, "attention_mask": kwargs.get("attention_mask"), "use_cache": False, } def forward( self, input_ids: Optional[torch.Tensor] = None, attention_mask: Optional[torch.Tensor] = None, position_ids: Optional[torch.Tensor] = None, past_key_values=None, inputs_embeds: Optional[torch.Tensor] = None, labels: Optional[torch.Tensor] = None, use_cache: Optional[bool] = None, output_attentions: Optional[bool] = None, output_hidden_states: Optional[bool] = None, return_dict: Optional[bool] = None, **kwargs, ) -> CausalLMOutputWithPast: if input_ids is not None and inputs_embeds is not None: raise ValueError("pass either input_ids or inputs_embeds, not both") if inputs_embeds is None: if input_ids is None: raise ValueError("input_ids or inputs_embeds is required") hidden_states = self.embed_tokens(input_ids) else: hidden_states = inputs_embeds if position_ids is None: if attention_mask is not None and attention_mask.ndim == 2: position_ids = attention_mask.long().cumsum(-1) - 1 position_ids = position_ids.masked_fill(attention_mask == 0, 0) else: position_ids = torch.arange(hidden_states.shape[1], device=hidden_states.device).unsqueeze(0) all_hidden_states = () if output_hidden_states else None for block in self.layers: if output_hidden_states: all_hidden_states += (hidden_states,) hidden_states = block(hidden_states, attention_mask, position_ids) if output_hidden_states: all_hidden_states += (hidden_states,) hidden_states = self.norm(hidden_states) logits = F.linear(hidden_states, self.embed_tokens.weight) loss = None if labels is not None: shift_logits = logits[..., :-1, :].contiguous().float() shift_labels = labels[..., 1:].contiguous() loss = F.cross_entropy(shift_logits.view(-1, shift_logits.size(-1)), shift_labels.view(-1)) if return_dict is False: output = (logits, None, all_hidden_states, None) return ((loss,) + output) if loss is not None else output return CausalLMOutputWithPast( loss=loss, logits=logits, past_key_values=None, hidden_states=all_hidden_states, attentions=None, )