import torch, math import torch.nn as nn import torch.nn.functional as F def compute_theta(dim: int, base: float = 10000.0, device: torch.device = torch.device('cpu')) -> torch.Tensor: """ 计算旋转位置编码中的 Theta 角度值。 参数： - d (int): 嵌入向量的维度（必须为偶数）。 - base (float): 基础频率参数, 默认为10000.0。 - device (torch.device): 计算设备, 默认为CPU。 返回： - torch.Tensor: 包含Theta值的1D张量, 形状为 [d/2]。 """ if dim % 2 != 0: print("嵌入维度 dim 必须为偶数") i = torch.arange(1, (dim//2) + 1, dtype=torch.float32, device=device) theta_i = base ** (-2*(i - 1) / dim) return theta_i def precompute_freqs_cis(dim: int, seq_len: int, base: float = 10000.0, device: torch.device = torch.device('cpu')): theta = compute_theta(dim, base, device) # theta 角度值序列，向量, 大小为 dim // 2 m = torch.arange(seq_len, device=device) # # token 位置值序列，向量，大小为 seq_len m_theta = torch.outer(m, theta) # 所有 token 位置的所有 Theta 值范围, 矩阵，尺寸为 [seq_len, dim // 2] freqs_cis = torch.polar(torch.ones_like(m_theta), m_theta) # e^{i*m*\theta}，本质上是旋转矩阵 return freqs_cis def reshape_for_broadcast(freqs_cis, x): ndim = x.ndim assert ndim >= 2 assert freqs_cis.shape == (x.shape[1],x.shape[-1]), "the last two dimension of freqs_cis, x must match" shape = [d if i==1 or i==ndim-1 else 1 for i,d in enumerate(x.shape)] return freqs_cis.view(*shape) def apply_rotary_emb(xq: torch.Tensor, xk: torch.Tensor, freqs_cis: torch.Tensor, device: torch.device = torch.device('cpu')): """ 参数: - x_q(torch.Tensor): 实际上是权重 W_q * 词嵌入向量值, 来自上一个线性层的输出, 形状为 [batch_size, seq_len, n_heads, head_dim] - x_k(torch.Tensor): 实际上是权重 W_k * 词嵌入向量值, 来自上一个线性层的输出, 形状为 [batch_size, seq_len, n_heads, head_dim] - freqs_cis (torch.Tensor): 频率复数张量, 形状为 [max_seq_len, head_dim] 返回: - Tuple[torch.Tensor, torch.Tensor]: 旋转编码后的查询和键张量 """ # 实数域张量转为复数域张量 xq_reshape = xq.reshape(*xq.shape[:-1], -1, 2) # [batch_size, seq_len, dim] -> [batch_size, seq_len, dim//2, 2] xk_reshape = xk.reshape(*xk.shape[:-1], -1, 2) # [batch_size, seq_len, dim] -> [batch_size, seq_len, dim//2, 2] xq_complex = torch.view_as_complex(xq_reshape) # 复数形式张量 xk_complex = torch.view_as_complex(xk_reshape) # 复数形式张量 # 旋转矩阵（freqs_cis）的维度在序列长度（seq_len，维度 1）和头部维度（head_dim，维度 3）上需要与嵌入的维度一致。 # 此外，freqs_cis 的形状必须与 xq 和 xk 相匹配，因此我们需要将 freqs_cis 的形状从 [max_seq_len, head_dim] 调整为 [1, max_seq_len, 1, head_dim]。 freqs_cis = reshape_for_broadcast(freqs_cis, xq_complex) # [max_seq_len, 1, 1, dim // 2] # 应用旋转操作，并将结果转回实数域 xq_out = torch.view_as_real(xq_complex * freqs_cis).flatten(3) # flatten(2) 将后面两个维度压成一个维度 xk_out = torch.view_as_real(xk_complex * freqs_cis).flatten(3) return xq_out.type_as(xq), xk_out.type_as(xk) class Attention(nn.Module): """compute scale dot product attention Query : given sentence that we focused on (decoder) Key : every sentence to check relationship with Qeury(encoder) Value : every sentence same with Key (encoder) """ def __init__(self, dim, max_seq_len, n_heads): super(Attention, self).__init__() self.wq = nn.Linear(dim, dim) # Q 线性变换层 self.wk = nn.Linear(dim, dim) # K 线性变换层 self.wv = nn.Linear(dim, dim) # V 线性变换层 self.out = nn.Linear(dim, dim) # 输出线性变换层 self.softmax = nn.Softmax(dim = -1) self.dim = dim # 模型中嵌入层尺寸 self.max_seq_len = max_seq_len self.n_heads = n_heads self.head_dim = dim // n_heads # 每个 head 的尺寸 def forward(self, x: torch.Tensor, start_pos=0, inference=True, mask=None): # 1, 计算 xq xk xv 并调整它们的形状为 [bs, seq_len, n_heads, head_dim] bs, seq_len, dim = x.shape xq, xk, xv = self.wq(x), self.wk(x), self.wv(x) xq = xq.view(bs, seq_len, self.n_heads, self.head_dim) xk = xk.view(bs, seq_len, self.n_heads, self.head_dim) xv = xv.view(bs, seq_len, self.n_heads, self.head_dim) if inference: # 2, Compute rotation matrix for each position in the sequence freqs_cis = precompute_freqs_cis(self.head_dim, self.max_seq_len * 2) freqs_cis = freqs_cis[start_pos : start_pos + seq_len] # 3, Apply RoPE to Queries and Keys embeddings xq, xk = apply_rotary_emb(xq, xk, freqs_cis) # 4, 因为 self-attention 算子中矩阵乘法操作处理的事 [seq_len, d_model] 维度，所以需要 # 对查询（queries）、键（keys）和值（values）进行转置操作，重新调整它们的形状， # 使“头”（heads）位于维度 1，而“序列”（seq）位于维度 2。 querys = xq.transpose(1, 2) # [bs, seq_len, n_heads, head_dim] -> [bs, n_heads, seq_len, head_dim] keys = xk.transpose(1, 2) # keys:[bs,n_heads,seq_len,head_dim] values = xv.transpose(1, 2) # values:[bs,n_heads,seq_len,head_dim] # 5, self-attention 算子的内部操作 scores = torch.matmul(querys, keys.transpose(2, 3)) / math.sqrt(self.head_dim) # 计算 Q, K^T 矩阵的点积，再除以 sqrt(d_k) 得到注意力分数矩阵 if mask is not None: # 如果有掩码，则将注意力分数矩阵中对应掩码位置的值设为负无穷大 scores = scores.masked_fill(mask == 0, -1e9) attn_weights = F.softmax(scores.float(), dim = -1) # 将注意力权重矩阵乘以 V，得到最终的输出矩阵 atten_output = torch.matmul(attn_weights, values) # Shape change: output[bsz,n_heads,seq_len,head_dim] -> output[bsz,seq_len, n_heads,head_dim] -> output[bsz,seq_len, n_heads * head_dim] output = atten_output.transpose(1,2).contiguous().view(bs, seq_len, -1) return output def test_compute_theta(): dim = 8 theta = compute_theta(dim) assert theta.shape == (dim // 2,), f"Expected shape {(dim // 2,)}, got {theta.shape}" print("test_compute_theta passed.") def test_precompute_freqs_cis(): dim = 8 seq_len = 10 freqs_cis = precompute_freqs_cis(dim, seq_len) freqs_cis_llama = llama.precompute_freqs_cis(dim, seq_len) assert freqs_cis.shape == (seq_len, dim // 2), f"Expected shape {(seq_len, dim // 2)}, got {freqs_cis.shape}" print("test_precompute_freqs_cis passed.") def test_apply_rotary_emb(): batch_size = 2 seq_len = 5 n_heads = 2 head_dim = 4 # dim // n_heads dim = n_heads * head_dim xq = torch.randn(batch_size, seq_len, n_heads, head_dim) xk = torch.randn(batch_size, seq_len, n_heads, head_dim) freqs_cis = precompute_freqs_cis(head_dim, seq_len) xq_out, xk_out = apply_rotary_emb(xq, xk, freqs_cis) # git clone https://github.com/meta-llama/llama/tree/main and pip install -e . xq, xk = llama.apply_rotary_emb(xq, xk, freqs_cis) assert torch.allclose(xq_out, xq, atol=1e-4) assert xq_out.shape == xq.shape, f"Expected xq_out shape {xq.shape}, got {xq_out.shape}" assert xk_out.shape == xk.shape, f"Expected xk_out shape {xk.shape}, got {xk_out.shape}" print(f"test_apply_rotary_emb passed, xq_out and xq [0][0][0][0]: {xq_out[0,0,0,0]} {xq[0,0,0,0]}.") def test_attention(): batch_size = 2 seq_len = 5 dim = 8 n_heads = 2 x = torch.randn(batch_size, seq_len, dim) attn = Attention(dim=dim, max_seq_len=seq_len, n_heads=n_heads) output = attn(x) assert output.shape == (batch_size, seq_len, dim), f"Expected output shape {(batch_size, seq_len, dim)}, got {output.shape}" print("test_attention passed.") if __name__ == "__main__": test_compute_theta() test_precompute_freqs_cis() test_apply_rotary_emb() test_attention()