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