File size: 2,934 Bytes
5e6d9f5 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 | """基础算子：RMSNorm、解耦 RoPE、SwiGLU FFN。"""
import math
import torch
import torch.nn as nn
import torch.nn.functional as F
class RMSNorm(nn.Module):
"""比 LayerNorm 少一次均值统计，省算力且效果相当，是现代 LLM 的标配。
y = x / sqrt(mean(x^2) + eps) * g
"""
def __init__(self, dim: int, eps: float = 1e-5):
super().__init__()
self.eps = eps
self.weight = nn.Parameter(torch.ones(dim))
def forward(self, x: torch.Tensor) -> torch.Tensor:
# 统计量始终用 fp32 算，避免半精度下 rsqrt 抖动
out = x.float() * torch.rsqrt(x.float().pow(2).mean(-1, keepdim=True) + self.eps)
return (out * self.weight.float()).type_as(x)
# --------------------------------------------------------------------------
# RoPE：旋转位置编码
# MLA 里它是「解耦」的 —— 只作用在 Q/K 单独切出来的 qk_rope_head_dim 维上，
# 剩下的 qk_nope 部分不带位置信息，这样 K 才能被压进和位置无关的潜在向量里。
# --------------------------------------------------------------------------
def build_rope_cache(head_dim: int, max_seq_len: int, theta: float = 10000.0,
scaling: float = 1.0, device=None):
"""返回 (cos, sin)，形状均为 (max_seq_len, head_dim)。
scaling > 1 时做 NTK-aware 频率插值：把 base 放大 scaling^(d/(d-2))，
等价于低频维度插值、高频维度基本不动，用来做上下文长度外推。
"""
if scaling > 1.0:
theta = theta * (scaling ** (head_dim / (head_dim - 2)))
inv_freq = 1.0 / (theta ** (torch.arange(0, head_dim, 2, dtype=torch.float32, device=device) / head_dim))
t = torch.arange(max_seq_len, dtype=torch.float32, device=device)
freqs = torch.outer(t, inv_freq) # (T, head_dim/2)
emb = torch.cat([freqs, freqs], dim=-1) # (T, head_dim)
return emb.cos(), emb.sin()
def rotate_half(x: torch.Tensor) -> torch.Tensor:
x1, x2 = x.chunk(2, dim=-1)
return torch.cat([-x2, x1], dim=-1)
def apply_rope(x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor:
"""x: (B, H, T, D)，cos/sin: (T, D)。"""
cos = cos[None, None, :, :].to(x.dtype)
sin = sin[None, None, :, :].to(x.dtype)
return x * cos + rotate_half(x) * sin
# --------------------------------------------------------------------------
class SwiGLU(nn.Module):
"""FFN(x) = W_down( silu(W_gate x) * W_up x )，比 ReLU-MLP 表达力更强。"""
def __init__(self, dim: int, inter_dim: int):
super().__init__()
self.w_gate = nn.Linear(dim, inter_dim, bias=False)
self.w_up = nn.Linear(dim, inter_dim, bias=False)
self.w_down = nn.Linear(inter_dim, dim, bias=False)
def forward(self, x: torch.Tensor) -> torch.Tensor:
return self.w_down(F.silu(self.w_gate(x)) * self.w_up(x))
|