| """
|
| 位置编码模块 — 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)
|
| 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)
|
|
|
|
|
| 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):
|
| if seq_len > self._cached_seq_len or self._cached_cos is None:
|
| t = torch.arange(seq_len, device=device, dtype=self.inv_freq.dtype)
|
| freqs = torch.outer(t, self.inv_freq)
|
| emb = torch.cat((freqs, freqs), dim=-1)
|
| self._cached_cos = emb.cos()[None, None, :, :]
|
| self._cached_sin = emb.sin()[None, None, :, :]
|
| 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,
|
| k: torch.Tensor,
|
| ) -> 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
|
|
|
