modeling_vortex.py

"""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,
        )