File size: 4,722 Bytes
c1a46f7 4d62693 c1a46f7 4d62693 c1a46f7 4d62693 c1a46f7 d572bbd 4d62693 d572bbd 4d62693 c1a46f7 4d62693 | 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 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 | """
位置编码模块 — Person B 负责实现
包含:
1. SinusoidalPositionalEncoding: 经典正弦余弦位置编码
2. RotaryPositionalEmbedding (RoPE): 旋转位置编码
RoPE 是当前最前沿的位置编码方案，被 LLaMA, GPT-NeoX 等大模型广泛采用。
优势: 更好地捕获相对位置信息，支持长度外推。
"""
from __future__ import annotations
import math
import torch
import torch.nn as nn
class SinusoidalPositionalEncoding(nn.Module):
"""
经典正弦余弦位置编码 (Vaswani et al., 2017)。
PE(pos, 2i) = sin(pos / 10000^(2i/d_model))
PE(pos, 2i+1) = cos(pos / 10000^(2i/d_model))
TODO [Person B]: 实现以下内容:
1. 在 __init__ 中预计算位置编码矩阵 [max_seq_len, d_model]
2. 注册为 buffer (不参与梯度更新)
3. forward(x): 返回 x + pe[:, :x.size(1), :]
"""
def __init__(self, d_model: int, max_seq_len: int = 5000, dropout: float = 0.1):
super().__init__()
self.dropout = nn.Dropout(p=dropout)
pe = torch.zeros(max_seq_len, d_model)
position = torch.arange(0, max_seq_len, dtype=torch.float).unsqueeze(1)
div_term = torch.exp(
torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model)
)
pe[:, 0::2] = torch.sin(position * div_term)
pe[:, 1::2] = torch.cos(position * div_term)
pe = pe.unsqueeze(0) # [1, max_seq_len, d_model]
self.register_buffer("pe", pe)
def forward(self, x: torch.Tensor) -> torch.Tensor:
"""
Args:
x: [B, L, D]
Returns:
x + positional_encoding: [B, L, D]
"""
x = x + self.pe[:, : x.size(1), :]
return self.dropout(x)
class RotaryPositionalEmbedding(nn.Module):
"""
旋转位置编码 (RoPE) — Su et al., 2021
核心思想: 通过旋转变换将位置信息编码到 Q, K 向量中。
q' = q * cos(θ) + rotate_half(q) * sin(θ)
k' = k * cos(θ) + rotate_half(k) * sin(θ)
TODO [Person B]: 实现以下内容:
__init__:
1. 计算频率: inv_freq = 1.0 / (base ** (torch.arange(0, dim, 2) / dim))
2. 注册为 buffer
_compute_rope(seq_len):
1. 计算 position indices: t = [0, 1, ..., seq_len-1]
2. 计算 freqs = torch.outer(t, inv_freq)
3. 构建 cos_cached, sin_cached
apply_rotary_pos_emb(q, k):
1. 对 q 和 k 应用旋转变换
2. 返回旋转后的 q', k'
参考: https://arxiv.org/abs/2104.09864
"""
def __init__(self, dim: int, max_seq_len: int = 2048, base: float = 10000.0):
super().__init__()
self.dim = dim
self.max_seq_len = max_seq_len
self.base = base
inv_freq = 1.0 / (
base ** (torch.arange(0, dim, 2).float() / dim)
)
self.register_buffer("inv_freq", inv_freq)
# 预缓存 cos/sin
self._cached_seq_len = 0
self._cached_cos: torch.Tensor | None = None
self._cached_sin: torch.Tensor | None = None
def _compute_rope(self, seq_len: int, device: torch.device):
# Recompute if seq_len grew OR cache is stale OR device changed (e.g. after .to(cuda))
cache_stale = (
seq_len > self._cached_seq_len
or self._cached_cos is None
or self._cached_cos.device != device
)
if cache_stale:
t = torch.arange(seq_len, device=device, dtype=self.inv_freq.dtype)
freqs = torch.outer(t, self.inv_freq.to(device)) # [seq_len, dim//2]
emb = torch.cat((freqs, freqs), dim=-1) # [seq_len, dim]
self._cached_cos = emb.cos()[None, None, :, :] # [1, 1, seq_len, dim]
self._cached_sin = emb.sin()[None, None, :, :] # [1, 1, seq_len, dim]
self._cached_seq_len = seq_len
return self._cached_cos, self._cached_sin
@staticmethod
def _rotate_half(x: torch.Tensor) -> torch.Tensor:
"""将 x 的后半部分取反并与前半部分交换。"""
x1, x2 = x.chunk(2, dim=-1)
return torch.cat((-x2, x1), dim=-1)
def apply_rotary_pos_emb(
self,
q: torch.Tensor, # [B, nhead, L, d_k]
k: torch.Tensor, # [B, nhead, L, d_k]
) -> tuple[torch.Tensor, torch.Tensor]:
"""
对 Q, K 应用 RoPE。
Returns:
(q_rotated, k_rotated)
"""
cos, sin = self._compute_rope(q.size(2), q.device)
q_embed = (q * cos[:, :, : q.size(2), :]) + (
self._rotate_half(q) * sin[:, :, : q.size(2), :]
)
k_embed = (k * cos[:, :, : k.size(2), :]) + (
self._rotate_half(k) * sin[:, :, : k.size(2), :]
)
return q_embed, k_embed
|