| import torch
|
| import torch.nn as nn
|
| import torch.nn.functional as F
|
| import math
|
| from typing import Optional, Tuple
|
|
|
|
|
| class RMSNorm(nn.Module):
|
| def __init__(self, hidden_size: int, eps: float = 1e-6):
|
| 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
|
| x = x.float()
|
| variance = x.pow(2).mean(-1, keepdim=True)
|
| x = x * torch.rsqrt(variance + self.eps)
|
| return (self.weight.float() * x).to(input_dtype)
|
|
|
|
|
| class SwiGLUMLP(nn.Module):
|
| def __init__(self, hidden_size: int, intermediate_size: int):
|
| super().__init__()
|
| self.gate_proj = nn.Linear(hidden_size, intermediate_size, bias=False)
|
| self.up_proj = nn.Linear(hidden_size, intermediate_size, bias=False)
|
| self.down_proj = nn.Linear(intermediate_size, hidden_size, bias=False)
|
|
|
| def forward(self, x: torch.Tensor, lora_deltas: dict = None) -> torch.Tensor:
|
| if lora_deltas:
|
| gate = F.linear(x, self.gate_proj.weight + lora_deltas.get("gate_proj", 0))
|
| up = F.linear(x, self.up_proj.weight + lora_deltas.get("up_proj", 0))
|
| down_weight = self.down_proj.weight + lora_deltas.get("down_proj", 0)
|
| return F.linear(F.silu(gate) * up, down_weight)
|
| return self.down_proj(F.silu(self.gate_proj(x)) * self.up_proj(x))
|
|
|
|
|
| def _yarn_find_correction_dim(
|
| num_rotations: int, dim: int, base: float = 10000.0, max_position: int = 2048
|
| ) -> float:
|
| """Find correction dimension for YaRN interpolation."""
|
| return (dim * math.log(max_position / (num_rotations * 2 * math.pi))) / (2 * math.log(base))
|
|
|
|
|
| def _yarn_find_correction_range(
|
| low_rot: int, high_rot: int, dim: int, base: float = 10000.0, max_position: int = 2048
|
| ) -> Tuple[int, int]:
|
| """Find the range of dimensions to apply YaRN correction."""
|
| low = math.floor(_yarn_find_correction_dim(low_rot, dim, base, max_position))
|
| high = math.ceil(_yarn_find_correction_dim(high_rot, dim, base, max_position))
|
| return max(low, 0), min(high, dim - 1)
|
|
|
|
|
| def _yarn_linear_ramp_mask(low: int, high: int, dim: int, dtype: torch.dtype) -> torch.Tensor:
|
| """Create linear ramp mask for smooth interpolation between dimensions."""
|
| if low == high:
|
| high += 0.001
|
| linear_func = (torch.arange(dim, dtype=dtype) - low) / (high - low)
|
| return linear_func.clamp(0, 1)
|
|
|
|
|
| class YaRNRotaryEmbedding(nn.Module):
|
| """
|
| RoPE with YaRN (Yet another RoPE extensioN) scaling.
|
|
|
| Same approach as Gemma 4 E2B:
|
| - Base theta = 1,000,000
|
| - YaRN scaling for context extension to 128K
|
| - Splits dimensions into 3 regions:
|
| 1. Low freq dims: apply NTK-aware interpolation
|
| 2. Medium freq dims: smooth ramp between interpolation and extrapolation
|
| 3. High freq dims: no scaling (extrapolation)
|
| """
|
|
|
| def __init__(
|
| self,
|
| dim: int,
|
| max_position_embeddings: int = 131072,
|
| base: float = 1000000.0,
|
| scaling_factor: float = 4.0,
|
| original_max_position: int = 32768,
|
| beta_fast: int = 32,
|
| beta_slow: int = 1,
|
| ):
|
| super().__init__()
|
| self.dim = dim
|
| self.max_position_embeddings = max_position_embeddings
|
| self.base = base
|
| self.scaling_factor = scaling_factor
|
| self.original_max_position = original_max_position
|
| self.beta_fast = beta_fast
|
| self.beta_slow = beta_slow
|
|
|
| self._build_yarn_cache()
|
|
|
| def _build_yarn_cache(self):
|
| """Compute YaRN-adjusted inverse frequencies."""
|
| dim = self.dim
|
|
|
| inv_freq = 1.0 / (
|
| self.base ** (torch.arange(0, dim, 2, dtype=torch.float32) / dim)
|
| )
|
|
|
|
|
| low, high = _yarn_find_correction_range(
|
| self.beta_slow, self.beta_fast, dim, self.base, self.original_max_position
|
| )
|
| inv_freq_mask = 1.0 - _yarn_linear_ramp_mask(low, high, dim // 2, torch.float32)
|
|
|
|
|
| inv_freq_interpolated = inv_freq / self.scaling_factor
|
|
|
|
|
| inv_freq_yarn = inv_freq_interpolated * (1 - inv_freq_mask) + inv_freq * inv_freq_mask
|
|
|
| self.register_buffer("inv_freq", inv_freq_yarn, persistent=False)
|
|
|
|
|
| self.attn_scale = 0.1 * math.log(self.scaling_factor) + 1.0
|
|
|
| def forward(
|
| self, x: torch.Tensor, position_ids: Optional[torch.Tensor] = None
|
| ) -> Tuple[torch.Tensor, torch.Tensor]:
|
| """
|
| Args:
|
| x: (B, H, T, D) — used only for device/dtype
|
| position_ids: (B, T) or None (auto-generate 0..T-1)
|
| Returns:
|
| cos, sin: (B, T, D) in same dtype as x
|
| """
|
| B, H, T, D = x.shape
|
|
|
| if position_ids is None:
|
| position_ids = torch.arange(T, device=x.device).unsqueeze(0).expand(B, -1)
|
|
|
|
|
| inv_freq = self.inv_freq.to(device=x.device, dtype=torch.float32)
|
| freqs = position_ids.unsqueeze(-1).float() * inv_freq.unsqueeze(0).unsqueeze(0)
|
|
|
| emb = torch.cat([freqs, freqs], dim=-1)
|
|
|
| cos = (emb.cos() * self.attn_scale).to(x.dtype)
|
| sin = (emb.sin() * self.attn_scale).to(x.dtype)
|
|
|
| return cos, sin
|
|
|
|
|
| def apply_rotary_pos_emb(
|
| q: torch.Tensor, k: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor
|
| ) -> Tuple[torch.Tensor, torch.Tensor]:
|
| """
|
| Apply RoPE rotation to Q and K.
|
|
|
| Args:
|
| q: (B, H, T, D)
|
| k: (B, H_kv, T, D)
|
| cos: (B, T, D)
|
| sin: (B, T, D)
|
| Returns:
|
| q_rotated, k_rotated: same shapes
|
| """
|
|
|
| cos = cos.unsqueeze(1)
|
| sin = sin.unsqueeze(1)
|
|
|
| q_rotated = (q * cos) + (_rotate_half(q) * sin)
|
| k_rotated = (k * cos) + (_rotate_half(k) * sin)
|
|
|
| return q_rotated, k_rotated
|
|
|
|
|
| def _rotate_half(x: torch.Tensor) -> torch.Tensor:
|
| """Rotate half the hidden dims: [x1, x2] → [-x2, x1]"""
|
| x1 = x[..., : x.shape[-1] // 2]
|
| x2 = x[..., x.shape[-1] // 2 :]
|
| return torch.cat((-x2, x1), dim=-1)
|
|
|
|
|
| def apply_depth_rotary_emb(x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor:
|
| """
|
| Apply depth rotation to a subset of dimensions.
|
| Args:
|
| x: (B, T, D)
|
| cos, sin: (1, 1, d_rode) where d_rode is the number of rotated dimensions.
|
| """
|
| d_rode = cos.shape[-1]
|
| x_rode = x[..., :d_rode]
|
| x_rest = x[..., d_rode:]
|
|
|
|
|
| x_rotated = (x_rode * cos) + (_rotate_half(x_rode) * sin)
|
|
|
| return torch.cat([x_rotated, x_rest], dim=-1)
|
|
|
|
|
| class RoDE(nn.Module):
|
| """
|
| Rotary Depth Embedding (RoDE).
|
| Provides a depth signal for shared weights in Hyperloop.
|
| """
|
| def __init__(self, dim: int, num_loops: int, base: float = 10000.0):
|
| super().__init__()
|
| self.dim = dim
|
| self.num_loops = num_loops
|
| self.base = base
|
|
|
|
|
| inv_freq = 1.0 / (base ** (torch.arange(0, dim, 2).float() / dim))
|
|
|
|
|
| loop_ids = torch.arange(num_loops).float()
|
| freqs = loop_ids.unsqueeze(-1) * inv_freq.unsqueeze(0)
|
|
|
|
|
| emb = torch.cat([freqs, freqs], dim=-1)
|
| self.register_buffer("cos", emb.cos(), persistent=False)
|
| self.register_buffer("sin", emb.sin(), persistent=False)
|
|
|
| def forward(self, x: torch.Tensor, loop_idx: int) -> Tuple[torch.Tensor, torch.Tensor]:
|
|
|
|
|
| cos = self.cos[loop_idx].view(1, 1, -1).to(x.dtype)
|
| sin = self.sin[loop_idx].view(1, 1, -1).to(x.dtype)
|
| return cos, sin
|
|
|
