tiny_gdn/model.py

from __future__ import annotations

import math
from dataclasses import dataclass
from pathlib import Path
from typing import Any

import torch
import torch.nn.functional as F
from safetensors.torch import load_model, save_model
from torch import nn
from torch.nn.attention import SDPBackend, sdpa_kernel
from torch.utils.checkpoint import checkpoint

from tiny_gdn.config import TinyGDNConfig

try:
    # Import the module directly — `from fla.layers import GatedDeltaNet2`
    # executes layers/__init__.py and eagerly loads every attention kernel.
    from fla.layers.gdn2 import GatedDeltaNet2
except ImportError as import_error:
    GatedDeltaNet2 = None
    FLA_IMPORT_ERROR: ImportError | None = import_error
else:
    FLA_IMPORT_ERROR = None


@dataclass
class TinyGDNOutput:
    loss: torch.Tensor | None
    logits: torch.Tensor | None
    main_loss: torch.Tensor | None
    mtp_loss: torch.Tensor | None
    z_loss: torch.Tensor | None
    hidden_states: torch.Tensor | None = None
    past_key_values: Any | None = None


class RMSNorm(nn.Module):
    """Zero-centered RMSNorm as used by Qwen3-Next."""

    def __init__(self, hidden_size: int, eps: float) -> None:
        super().__init__()
        self.weight = nn.Parameter(torch.zeros(hidden_size))
        self.eps = eps

    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
        input_dtype = hidden_states.dtype
        normalized = hidden_states.float()
        normalized = normalized * torch.rsqrt(normalized.square().mean(dim=-1, keepdim=True) + self.eps)
        normalized = normalized * (1.0 + self.weight.float())
        return normalized.to(dtype=input_dtype)


class RotaryEmbedding(nn.Module):
    def __init__(self, rotary_dim: int, rope_theta: float) -> None:
        super().__init__()
        inverse_frequency = 1.0 / (
            rope_theta
            ** (
                torch.arange(0, rotary_dim, 2, dtype=torch.float32)
                / rotary_dim
            )
        )
        self.rotary_dim = rotary_dim
        self.register_buffer("inverse_frequency", inverse_frequency, persistent=False)

    def forward(
        self,
        sequence_length: int,
        device: torch.device,
        dtype: torch.dtype,
        position_offset: int = 0,
    ) -> tuple[torch.Tensor, torch.Tensor]:
        positions = torch.arange(
            position_offset,
            position_offset + sequence_length,
            device=device,
            dtype=torch.float32,
        )
        frequencies = torch.outer(positions, self.inverse_frequency.float())
        embeddings = torch.cat((frequencies, frequencies), dim=-1)
        return embeddings.cos().to(dtype=dtype), embeddings.sin().to(dtype=dtype)


def rotate_half(hidden_states: torch.Tensor) -> torch.Tensor:
    first, second = hidden_states.chunk(2, dim=-1)
    return torch.cat((-second, first), dim=-1)


def apply_rotary_embedding(
    query: torch.Tensor,
    key: torch.Tensor,
    cosine: torch.Tensor,
    sine: torch.Tensor,
    rotary_dim: int,
) -> tuple[torch.Tensor, torch.Tensor]:
    cosine = cosine[None, None, :, :]
    sine = sine[None, None, :, :]
    query_rotary, query_pass = query[..., :rotary_dim], query[..., rotary_dim:]
    key_rotary, key_pass = key[..., :rotary_dim], key[..., rotary_dim:]
    query_rotary = query_rotary * cosine + rotate_half(query_rotary) * sine
    key_rotary = key_rotary * cosine + rotate_half(key_rotary) * sine
    return (
        torch.cat((query_rotary, query_pass), dim=-1),
        torch.cat((key_rotary, key_pass), dim=-1),
    )


class GatedGroupedQueryAttention(nn.Module):
    """QK-normalized, partially rotary GQA with a learned sigmoid output gate."""

    def __init__(self, config: TinyGDNConfig) -> None:
        super().__init__()
        self.num_heads = config.num_attention_heads
        self.num_key_value_heads = config.num_key_value_heads
        self.head_dim = config.attention_head_dim
        self.rotary_dim = config.rotary_dim
        self.dropout = config.attention_dropout

        query_size = self.num_heads * self.head_dim
        key_value_size = self.num_key_value_heads * self.head_dim
        self.q_gate_proj = nn.Linear(config.hidden_size, query_size * 2, bias=False)
        self.k_proj = nn.Linear(config.hidden_size, key_value_size, bias=False)
        self.v_proj = nn.Linear(config.hidden_size, key_value_size, bias=False)
        self.o_proj = nn.Linear(query_size, config.hidden_size, bias=False)
        self.q_norm = RMSNorm(self.head_dim, config.rms_norm_eps)
        self.k_norm = RMSNorm(self.head_dim, config.rms_norm_eps)
        self.rotary = RotaryEmbedding(self.rotary_dim, config.rope_theta)

    def _attention_mask(
        self,
        attention_mask: torch.Tensor | None,
        sequence_length: int,
        device: torch.device,
    ) -> torch.Tensor | None:
        if attention_mask is None:
            return None
        if attention_mask.ndim != 2:
            raise ValueError("attention_mask must have shape [batch, sequence]")
        if attention_mask.shape[1] != sequence_length:
            raise ValueError("attention_mask sequence length does not match input")

        causal = torch.ones(
            sequence_length,
            sequence_length,
            dtype=torch.bool,
            device=device,
        ).tril()
        valid_keys = attention_mask[:, None, None, :].to(dtype=torch.bool, device=device)
        return causal[None, None, :, :] & valid_keys

    def forward(
        self,
        hidden_states: torch.Tensor,
        attention_mask: torch.Tensor | None = None,
        past_key_value: tuple[torch.Tensor, torch.Tensor] | None = None,
    ) -> tuple[torch.Tensor, tuple[torch.Tensor, torch.Tensor] | None]:
        batch_size, sequence_length, _ = hidden_states.shape
        past_len = 0 if past_key_value is None else past_key_value[0].shape[2]
        query_and_gate = self.q_gate_proj(hidden_states)
        query, output_gate = query_and_gate.chunk(2, dim=-1)

        query = query.view(batch_size, sequence_length, self.num_heads, self.head_dim)
        key = self.k_proj(hidden_states).view(
            batch_size,
            sequence_length,
            self.num_key_value_heads,
            self.head_dim,
        )
        value = self.v_proj(hidden_states).view(
            batch_size,
            sequence_length,
            self.num_key_value_heads,
            self.head_dim,
        )

        query = self.q_norm(query).transpose(1, 2)
        key = self.k_norm(key).transpose(1, 2)
        value = value.transpose(1, 2)

        cosine, sine = self.rotary(
            sequence_length,
            device=hidden_states.device,
            dtype=query.dtype,
            position_offset=past_len,
        )
        query, key = apply_rotary_embedding(
            query,
            key,
            cosine,
            sine,
            rotary_dim=self.rotary_dim,
        )
        if past_key_value is not None:
            key = torch.cat([past_key_value[0], key], dim=2)
            value = torch.cat([past_key_value[1], value], dim=2)
        present = (key, value)

        kv_len = key.shape[2]
        if attention_mask is not None and past_key_value is None:
            sdpa_mask = self._attention_mask(
                attention_mask,
                sequence_length,
                hidden_states.device,
            )
            is_causal = False
        elif past_key_value is not None and sequence_length == 1:
            # Decode step: query attends to the cached key/value prefix. Preserve
            # the prefill padding mask when decoding a left-padded prompt batch.
            sdpa_mask = (
                None
                if attention_mask is None
                else attention_mask[:, None, None, :].to(
                    dtype=torch.bool,
                    device=hidden_states.device,
                )
            )
            is_causal = False
        elif past_key_value is not None:
            # Prefill chunk with cache — build causal mask over kv_len.
            q_idx = torch.arange(
                past_len, past_len + sequence_length, device=hidden_states.device
            )[:, None]
            k_idx = torch.arange(kv_len, device=hidden_states.device)[None, :]
            sdpa_mask = (k_idx <= q_idx)[None, None, :, :]
            is_causal = False
        else:
            sdpa_mask = None
            is_causal = True
        sdpa_options = {
            "attn_mask": sdpa_mask,
            "dropout_p": self.dropout if self.training else 0.0,
            "is_causal": is_causal,
            "enable_gqa": True,
        }
        # Prefer Flash / mem-efficient when available; fall back to MATH for
        # Windows PyTorch builds that ship without FlashAttention kernels.
        backends = (
            [
                SDPBackend.FLASH_ATTENTION,
                SDPBackend.EFFICIENT_ATTENTION,
                SDPBackend.CUDNN_ATTENTION,
                SDPBackend.MATH,
            ]
            if query.is_cuda
            else [SDPBackend.MATH]
        )
        with sdpa_kernel(backends):
            attention_output = F.scaled_dot_product_attention(
                query,
                key,
                value,
                **sdpa_options,
            )
        attention_output = attention_output.transpose(1, 2).reshape(
            batch_size,
            sequence_length,
            -1,
        )
        attention_output = attention_output * torch.sigmoid(output_gate)
        return self.o_proj(attention_output), present


class SwiGLU(nn.Module):
    def __init__(self, config: TinyGDNConfig) -> None:
        super().__init__()
        self.gate_up_proj = nn.Linear(
            config.hidden_size,
            config.intermediate_size * 2,
            bias=False,
        )
        self.down_proj = nn.Linear(
            config.intermediate_size,
            config.hidden_size,
            bias=False,
        )

    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
        gate, up = self.gate_up_proj(hidden_states).chunk(2, dim=-1)
        return self.down_proj(F.silu(gate) * up)


class TinyGDNBlock(nn.Module):
    def __init__(self, config: TinyGDNConfig, layer_index: int) -> None:
        super().__init__()
        layer_type = config.layer_types[layer_index]
        self.layer_type = layer_type
        self.token_mixer_norm = RMSNorm(config.hidden_size, config.rms_norm_eps)
        self.mlp_norm = RMSNorm(config.hidden_size, config.rms_norm_eps)

        if layer_type == "gdn2":
            if GatedDeltaNet2 is None:
                raise ImportError(
                    "Gated DeltaNet-2 requires the pinned flash-linear-attention dependency"
                ) from FLA_IMPORT_ERROR
            self.token_mixer = GatedDeltaNet2(
                hidden_size=config.hidden_size,
                expand_v=config.linear_expand_v,
                head_dim=config.linear_head_dim,
                num_heads=config.linear_num_heads,
                num_v_heads=config.linear_num_value_heads,
                mode="chunk",
                use_short_conv=True,
                allow_neg_eigval=config.allow_negative_eigenvalues,
                conv_size=config.linear_conv_kernel_dim,
                conv_bias=False,
                layer_idx=layer_index,
                norm_eps=config.rms_norm_eps,
            )
        elif layer_type == "full_attention":
            self.token_mixer = GatedGroupedQueryAttention(config)
        else:
            raise ValueError(f"Unsupported layer type: {layer_type}")

        self.mlp = SwiGLU(config)

    def forward(
        self,
        hidden_states: torch.Tensor,
        attention_mask: torch.Tensor | None = None,
        *,
        past_key_values: Any | None = None,
        past_key_value: tuple[torch.Tensor, torch.Tensor] | None = None,
        use_cache: bool = False,
    ) -> tuple[torch.Tensor, Any]:
        residual = hidden_states
        normalized = self.token_mixer_norm(hidden_states)
        present: Any = None
        if self.layer_type == "gdn2":
            mixed, _, past_key_values = self.token_mixer(
                normalized,
                attention_mask=attention_mask,
                past_key_values=past_key_values,
                use_cache=use_cache,
            )
            present = past_key_values
        else:
            mixed, present = self.token_mixer(
                normalized,
                attention_mask=attention_mask,
                past_key_value=past_key_value,
            )
            if not use_cache:
                present = None
        hidden_states = residual + mixed
        hidden_states = hidden_states + self.mlp(self.mlp_norm(hidden_states))
        return hidden_states, present


class MultiTokenPredictionAdapter(nn.Module):
    """A lightweight residual adapter for one additional prediction horizon."""

    def __init__(self, config: TinyGDNConfig) -> None:
        super().__init__()
        self.norm = RMSNorm(config.hidden_size, config.rms_norm_eps)
        self.down_proj = nn.Linear(
            config.hidden_size,
            config.mtp_adapter_rank,
            bias=False,
        )
        self.up_proj = nn.Linear(
            config.mtp_adapter_rank,
            config.hidden_size,
            bias=False,
        )

    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
        adapted = self.up_proj(F.silu(self.down_proj(self.norm(hidden_states))))
        return hidden_states + adapted


class TinyGDNForCausalLM(nn.Module):
    def __init__(self, config: TinyGDNConfig) -> None:
        super().__init__()
        self.config = config
        self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size)
        self.layers = nn.ModuleList(
            TinyGDNBlock(config, layer_index)
            for layer_index in range(config.num_hidden_layers)
        )
        self.final_norm = RMSNorm(config.hidden_size, config.rms_norm_eps)
        self.mtp_adapters = nn.ModuleList(
            MultiTokenPredictionAdapter(config)
            for _ in range(config.mtp_num_heads)
        )
        self.gradient_checkpointing = False

        self.apply(self._initialize_module)
        self._initialize_residual_projections()

    def _initialize_module(self, module: nn.Module) -> None:
        if isinstance(module, nn.Linear):
            nn.init.normal_(
                module.weight,
                mean=0.0,
                std=self.config.initializer_range,
            )
            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=self.config.initializer_range,
            )

    def _initialize_residual_projections(self) -> None:
        residual_std = self.config.initializer_range / math.sqrt(
            2 * self.config.num_hidden_layers
        )
        for layer in self.layers:
            nn.init.normal_(
                layer.token_mixer.o_proj.weight,
                mean=0.0,
                std=residual_std,
            )
            nn.init.normal_(
                layer.mlp.down_proj.weight,
                mean=0.0,
                std=residual_std,
            )
        for adapter in self.mtp_adapters:
            nn.init.normal_(adapter.up_proj.weight, mean=0.0, std=residual_std)

    def enable_gradient_checkpointing(self, enabled: bool = True) -> None:
        self.gradient_checkpointing = enabled

    def project_to_vocabulary(self, hidden_states: torch.Tensor) -> torch.Tensor:
        return F.linear(hidden_states, self.embed_tokens.weight)

    def _run_layer(
        self,
        layer: TinyGDNBlock,
        hidden_states: torch.Tensor,
        attention_mask: torch.Tensor | None,
        *,
        past_key_values: Any | None = None,
        past_key_value: tuple[torch.Tensor, torch.Tensor] | None = None,
        use_cache: bool = False,
    ) -> tuple[torch.Tensor, Any]:
        if self.gradient_checkpointing and self.training:
            hidden_states, present = checkpoint(
                layer,
                hidden_states,
                attention_mask,
                use_reentrant=False,
            )
            return hidden_states, present
        return layer(
            hidden_states,
            attention_mask,
            past_key_values=past_key_values,
            past_key_value=past_key_value,
            use_cache=use_cache,
        )

    def _causal_loss(
        self,
        hidden_states: torch.Tensor,
        labels: torch.Tensor,
        target_offset: int,
        adapter: nn.Module | None = None,
        compute_z_loss: bool = False,
    ) -> tuple[torch.Tensor, torch.Tensor | None]:
        if target_offset < 0:
            raise ValueError("target_offset cannot be negative")
        if target_offset and hidden_states.shape[1] <= target_offset:
            raise ValueError(
                f"Sequence length must exceed target offset {target_offset}"
            )
        if target_offset:
            prediction_states = hidden_states[:, :-target_offset, :]
            targets = labels[:, target_offset:].contiguous()
        else:
            prediction_states = hidden_states
            targets = labels.contiguous()
        if adapter is not None:
            prediction_states = adapter(prediction_states)
        logits = self.project_to_vocabulary(prediction_states)
        cross_entropy = F.cross_entropy(
            logits.reshape(-1, self.config.vocab_size),
            targets.reshape(-1),
            ignore_index=-100,
        )
        z_loss = None
        if compute_z_loss:
            valid_targets = targets.ne(-100)
            log_partition = torch.logsumexp(logits.float(), dim=-1)
            z_loss = log_partition.square()[valid_targets].mean()
        return cross_entropy, z_loss

    def forward(
        self,
        input_ids: torch.Tensor,
        labels: torch.Tensor | None = None,
        attention_mask: torch.Tensor | None = None,
        past_key_values: Any | None = None,
        *,
        use_cache: bool = False,
        return_logits: bool = True,
        return_hidden_states: bool = False,
        labels_are_shifted: bool = False,
        include_mtp_loss: bool = True,
        mtp_loss_weight: float | None = None,
        z_loss_coefficient: float = 0.0,
        logits_to_keep: int | None = None,
    ) -> TinyGDNOutput:
        if input_ids.ndim != 2:
            raise ValueError("input_ids must have shape [batch, sequence]")
        if input_ids.shape[1] > self.config.max_position_embeddings:
            raise ValueError("Input exceeds max_position_embeddings")
        if labels is not None and labels.shape != input_ids.shape:
            raise ValueError("labels must have the same shape as input_ids")
        if z_loss_coefficient < 0.0:
            raise ValueError("z_loss_coefficient cannot be negative")
        if logits_to_keep is not None and logits_to_keep <= 0:
            raise ValueError("logits_to_keep must be positive")
        if use_cache and labels is not None:
            raise ValueError("use_cache is not supported with labels")
        effective_mtp_weight = (
            self.config.mtp_loss_weight
            if mtp_loss_weight is None
            else mtp_loss_weight
        )
        if not 0.0 <= effective_mtp_weight <= 1.0:
            raise ValueError("mtp_loss_weight must be between zero and one")

        if use_cache and past_key_values is None:
            try:
                from fla.models.utils import Cache as FlaCache
            except ImportError as import_error:
                raise ImportError(
                    "Cached decode requires flash-linear-attention Cache"
                ) from import_error
            past_key_values = {
                "fla": FlaCache(),
                "gqa": [None] * len(self.layers),
            }
        elif past_key_values is not None and not isinstance(past_key_values, dict):
            raise TypeError("past_key_values must be a TinyGDN cache dict or None")

        fla_cache = None if past_key_values is None else past_key_values["fla"]
        gqa_cache = None if past_key_values is None else past_key_values["gqa"]

        hidden_states = self.embed_tokens(input_ids)
        shared_layer_indices = set(self.config.shared_layer_indices)
        for layer_index, layer in enumerate(self.layers):
            layer_gqa = None if gqa_cache is None else gqa_cache[layer_index]
            hidden_states, present = self._run_layer(
                layer,
                hidden_states,
                attention_mask,
                past_key_values=fla_cache,
                past_key_value=layer_gqa,
                use_cache=use_cache,
            )
            if use_cache and layer.layer_type == "full_attention" and gqa_cache is not None:
                gqa_cache[layer_index] = present
            if layer_index in shared_layer_indices:
                layer_gqa = None if gqa_cache is None else gqa_cache[layer_index]
                hidden_states, present = self._run_layer(
                    layer,
                    hidden_states,
                    attention_mask,
                    past_key_values=fla_cache,
                    past_key_value=layer_gqa,
                    use_cache=use_cache,
                )
                if use_cache and layer.layer_type == "full_attention" and gqa_cache is not None:
                    gqa_cache[layer_index] = present
        hidden_states = self.final_norm(hidden_states)

        main_loss = None
        mtp_loss = None
        z_loss = None
        total_loss = None
        if labels is not None:
            main_target_offset = 0 if labels_are_shifted else 1
            main_loss, z_loss = self._causal_loss(
                hidden_states,
                labels,
                target_offset=main_target_offset,
                compute_z_loss=z_loss_coefficient > 0.0,
            )
            if self.mtp_adapters and include_mtp_loss:
                auxiliary_losses = [
                    self._causal_loss(
                        hidden_states,
                        labels,
                        target_offset=(
                            head_index + 1
                            if labels_are_shifted
                            else head_index + 2
                        ),
                        adapter=adapter,
                    )[0]
                    for head_index, adapter in enumerate(self.mtp_adapters)
                ]
                mtp_loss = torch.stack(auxiliary_losses).mean()
                total_loss = main_loss + effective_mtp_weight * mtp_loss
            else:
                total_loss = main_loss
            if z_loss is not None:
                total_loss = total_loss + z_loss_coefficient * z_loss

        output_states = (
            hidden_states
            if logits_to_keep is None
            else hidden_states[:, -logits_to_keep:, :]
        )
        logits = self.project_to_vocabulary(output_states) if return_logits else None
        return TinyGDNOutput(
            loss=total_loss,
            logits=logits,
            main_loss=main_loss,
            mtp_loss=mtp_loss,
            z_loss=z_loss,
            hidden_states=hidden_states if return_hidden_states else None,
            past_key_values=past_key_values if use_cache else None,
        )

    def parameter_report(self) -> dict[str, int]:
        total = sum(parameter.numel() for parameter in self.parameters())
        mtp = sum(parameter.numel() for parameter in self.mtp_adapters.parameters())
        embeddings = self.embed_tokens.weight.numel()
        return {
            "deployable_core": total - mtp,
            "training_total": total,
            "embedding": embeddings,
            "mtp_auxiliary": mtp,
            "non_embedding_core": total - mtp - embeddings,
        }

    def save_checkpoint(self, output_dir: Path) -> None:
        output_dir.mkdir(parents=True, exist_ok=True)
        self.config.save_json(output_dir / "config.json")
        save_model(self, output_dir / "model.safetensors")

    @classmethod
    def from_checkpoint(
        cls,
        checkpoint_dir: Path,
        *,
        device: str | torch.device = "cpu",
        dtype: torch.dtype | None = None,
    ) -> TinyGDNForCausalLM:
        config = TinyGDNConfig.from_json(checkpoint_dir / "config.json")
        model = cls(config).to(device=device, dtype=dtype)
        load_model(model, checkpoint_dir / "model.safetensors", device=str(device))
        return model

    def extra_repr(self) -> str:
        report = self.parameter_report()
        return (
            f"core_parameters={report['deployable_core']:,}, "
            f"training_parameters={report['training_total']:,}"
        )

    def get_architecture_metadata(self) -> dict[str, Any]:
        return {
            "architecture": self.config.architecture,
            "layer_types": list(self.config.layer_types),
            "effective_num_layers": self.config.effective_num_layers,
            "shared_layer_indices": list(self.config.shared_layer_indices),
            "parameter_report": self.parameter_report(),
            "features": [
                "32-layer deep-thin parameter allocation",
                "Gated DeltaNet-2 recurrent memory",
                "3:1 recurrent-to-full-attention hybrid",
                "gated grouped-query attention",
                "QK normalization",
                "partial rotary embeddings",
                "zero-centered RMSNorm",
                "SwiGLU",
                "tied input-output embeddings",
                "optional multi-token prediction auxiliaries",
            ],
        }