| import torch |
| import torch.nn as nn |
| import torch.nn.functional as F |
|
|
|
|
| class RMSNorm(nn.Module): |
| def __init__(self, hidden_size: int, eps: float = 1e-5) -> None: |
| super().__init__() |
| self.eps = eps |
| self.weight = nn.Parameter(torch.ones(hidden_size)) |
|
|
| def forward(self, x: torch.Tensor) -> torch.Tensor: |
| rms = torch.sqrt(x.pow(2).mean(dim=-1, keepdim=True) + self.eps) |
| return x / rms * self.weight |
|
|
|
|
| class SwiGLU(nn.Module): |
| def __init__(self, hidden_size: int, intermediate_size: int) -> None: |
| 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) -> torch.Tensor: |
| return self.down_proj(F.silu(self.gate_proj(x)) * self.up_proj(x)) |
|
|
|
|
| class RotaryEmbedding(nn.Module): |
| def __init__( |
| self, |
| dim: int, |
| max_position_embeddings: int = 2048, |
| theta: float = 10000.0, |
| ) -> None: |
| super().__init__() |
| self.dim = dim |
| self.max_position_embeddings = max_position_embeddings |
| inv_freq = 1.0 / (theta ** (torch.arange(0, dim, 2).float() / dim)) |
| self.register_buffer("inv_freq", inv_freq, persistent=False) |
|
|
| def forward( |
| self, x: torch.Tensor, position_ids: torch.Tensor |
| ) -> tuple[torch.Tensor, torch.Tensor]: |
| freqs = torch.einsum("bi,j->bij", position_ids.float(), self.inv_freq) |
| emb = torch.cat((freqs, freqs), dim=-1) |
| cos = emb.cos() |
| sin = emb.sin() |
| return cos, sin |
|
|
|
|
| def rotate_half(x: torch.Tensor) -> torch.Tensor: |
| x1 = x[..., : x.shape[-1] // 2] |
| x2 = x[..., x.shape[-1] // 2 :] |
| return torch.cat((-x2, x1), dim=-1) |
|
|
|
|
| def apply_rotary_pos_emb( |
| q: torch.Tensor, |
| k: torch.Tensor, |
| cos: torch.Tensor, |
| sin: torch.Tensor, |
| ) -> tuple[torch.Tensor, torch.Tensor]: |
| cos = cos.unsqueeze(1) |
| sin = sin.unsqueeze(1) |
| q_embed = q * cos + rotate_half(q) * sin |
| k_embed = k * cos + rotate_half(k) * sin |
| return q_embed, k_embed |
|
|
