modeling_hanse_layers.py

from __future__ import annotations

import math
from typing import cast

import torch
import torch.nn.functional as F
from torch import nn

from .configuration_hanse import HanseConfig


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

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


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


class RotaryEmbedding(nn.Module):
    def __init__(self, head_dim: int, max_seq_len: int, theta: float) -> None:
        super().__init__()
        if head_dim % 2:
            raise ValueError("RoPE benötigt eine gerade head_dim")
        inverse_frequency = 1.0 / (
            theta ** (torch.arange(0, head_dim, 2, dtype=torch.float32) / head_dim)
        )
        positions = torch.arange(max_seq_len, dtype=torch.float32)
        frequencies = torch.outer(positions, inverse_frequency)
        angles = torch.cat((frequencies, frequencies), dim=-1)
        self.cos_cached: torch.Tensor
        self.sin_cached: torch.Tensor
        self.register_buffer("cos_cached", angles.cos(), persistent=False)
        self.register_buffer("sin_cached", angles.sin(), persistent=False)

    def forward(
        self, query: torch.Tensor, key: torch.Tensor
    ) -> tuple[torch.Tensor, torch.Tensor]:
        sequence_length = query.size(-2)
        if sequence_length > self.cos_cached.size(0):
            raise ValueError("Sequenz ist länger als max_seq_len")
        cos = self.cos_cached[:sequence_length].to(dtype=query.dtype)[
            None, None, :, :
        ]
        sin = self.sin_cached[:sequence_length].to(dtype=query.dtype)[
            None, None, :, :
        ]
        return (
            query * cos + _rotate_half(query) * sin,
            key * cos + _rotate_half(key) * sin,
        )


class SwiGLU(nn.Module):
    def __init__(self, config: HanseConfig) -> None:
        super().__init__()
        self.gate_up = nn.Linear(
            config.hidden_size, 2 * config.ffn_hidden_size, bias=False
        )
        self.down = nn.Linear(
            config.ffn_hidden_size, config.hidden_size, bias=False
        )

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        gate, value = self.gate_up(x).chunk(2, dim=-1)
        return cast(torch.Tensor, self.down(F.silu(gate) * value))


class GroupedQueryAttention(nn.Module):
    def __init__(self, config: HanseConfig) -> None:
        super().__init__()
        self.num_query_heads = config.num_query_heads
        self.num_kv_heads = config.num_kv_heads
        self.head_dim = config.head_dim
        self.query = nn.Linear(config.hidden_size, config.hidden_size, bias=False)
        kv_size = config.num_kv_heads * config.head_dim
        self.key = nn.Linear(config.hidden_size, kv_size, bias=False)
        self.value = nn.Linear(config.hidden_size, kv_size, bias=False)
        self.output = nn.Linear(config.hidden_size, config.hidden_size, bias=False)
        self.query_norm = (
            HanseRMSNorm(config.head_dim, config.norm_eps)
            if config.qk_norm
            else nn.Identity()
        )
        self.key_norm = (
            HanseRMSNorm(config.head_dim, config.norm_eps)
            if config.qk_norm
            else nn.Identity()
        )
        self.rope = RotaryEmbedding(
            config.head_dim, config.max_seq_len, config.rope_theta
        )

    def _split_heads(self, x: torch.Tensor, heads: int) -> torch.Tensor:
        batch, sequence, _ = x.shape
        return x.view(batch, sequence, heads, self.head_dim).transpose(1, 2)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        query = self._split_heads(self.query(x), self.num_query_heads)
        key = self._split_heads(self.key(x), self.num_kv_heads)
        value = self._split_heads(self.value(x), self.num_kv_heads)
        query = self.query_norm(query)
        key = self.key_norm(key)
        query, key = self.rope(query, key)
        groups = self.num_query_heads // self.num_kv_heads
        key = key.repeat_interleave(groups, dim=1)
        value = value.repeat_interleave(groups, dim=1)
        attended = F.scaled_dot_product_attention(
            query, key, value, is_causal=True
        )
        batch, _, sequence, _ = attended.shape
        attended = attended.transpose(1, 2).reshape(
            batch, sequence, self.num_query_heads * self.head_dim
        )
        return cast(torch.Tensor, self.output(attended))


class CausalConvMixer(nn.Module):
    def __init__(self, config: HanseConfig) -> None:
        super().__init__()
        self.kernel_size = config.conv_kernel_size
        self.input = nn.Linear(
            config.hidden_size, 2 * config.hidden_size, bias=False
        )
        self.depthwise = nn.Conv1d(
            config.hidden_size,
            config.hidden_size,
            kernel_size=config.conv_kernel_size,
            groups=config.hidden_size,
            bias=False,
        )
        self.output = nn.Linear(config.hidden_size, config.hidden_size, bias=False)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        gate, value = self.input(x).chunk(2, dim=-1)
        value = value.transpose(1, 2)
        value = F.pad(value, (self.kernel_size - 1, 0))
        value = self.depthwise(value).transpose(1, 2)
        return cast(torch.Tensor, self.output(F.silu(gate) * value))


class HanseBlock(nn.Module):
    def __init__(self, config: HanseConfig, kind: str) -> None:
        super().__init__()
        self.mixer_norm = HanseRMSNorm(config.hidden_size, config.norm_eps)
        self.mixer: nn.Module
        if kind == "A":
            self.mixer = GroupedQueryAttention(config)
        elif kind == "C":
            self.mixer = CausalConvMixer(config)
        else:
            raise ValueError(f"Unbekannter Blocktyp: {kind}")
        self.ffn_norm = HanseRMSNorm(config.hidden_size, config.norm_eps)
        self.ffn = SwiGLU(config)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        x = x + self.mixer(self.mixer_norm(x))
        return cast(torch.Tensor, x + self.ffn(self.ffn_norm(x)))


def initialize_weights(module: nn.Module, num_layers: int) -> None:
    if isinstance(module, (nn.Linear, nn.Embedding, nn.Conv1d)):
        nn.init.normal_(module.weight, mean=0.0, std=0.02)
    if isinstance(module, (GroupedQueryAttention, CausalConvMixer)):
        nn.init.normal_(
            module.output.weight,
            mean=0.0,
            std=0.02 / math.sqrt(2 * num_layers),
        )
    elif isinstance(module, SwiGLU):
        nn.init.normal_(
            module.down.weight,
            mean=0.0,
            std=0.02 / math.sqrt(2 * num_layers),
        )