LLM(大型语言模型)算法岗的八股题目通常包括以下几个方面:1. **基础知识** - 请简述LLM的基本概念和特点。 - 解释什么是自然语言处理(NLP)以及它在LLM中的应用。 - 描述一下机器学习的基本流程。2. **模型架构** - 请介绍Tran
摘要:目录1. Top-p 和 Top-k 采样2. LayerNorm 和 RMSNorm3. SFT Loss 计算(Shift Right)4. 手撕 Softmax、交叉熵(Cross Entropy)5. QKV、Self-Attent
目录1. Top-p 和 Top-k 采样2. LayerNorm 和 RMSNorm3. SFT Loss 计算(Shift Right)4. 手撕 Softmax、交叉熵(Cross Entropy)5. QKV、Self-Attention 与 Multi-Head Attention (MHA)6. RoPE (Rotary Positional Embedding)面试总结与建议
1. Top-p 和 Top-k 采样
概念讲解:
在自回归文本生成中,模型每一步会输出一个概率分布(logits 经过 softmax),我们需要从中采样下一个 token。直接使用整个词汇表采样(即 temperature 缩放后的随机采样)可能导致生成低概率 token,使结果不连贯。Top-k 采样 和 Top-p 采样 是两种常用的截断采样方法,用于限制候选 token 集合,提高生成质量。
Top-k 采样:
做法:只保留概率最高的 k 个词,把剩下的词概率强制设为 0,然后重新归一化(让剩下的概率和为 1),再从中采样。
作用:直接砍掉长尾的低概率词,防止生成生僻字或乱码。
缺点:k 是固定的。如果模型很自信(某个词概率 90%),k 太大也会采样到噪音;如果模型很犹豫(概率很平),k 太小会限制多样性。
Top-p (Nucleus) 采样:
做法:将词按概率从大到小排序,依次累加概率,直到累加和超过 p (比如 0.9)。保留这些词,剩下的截断,重新归一化,再采样。
作用:动态调整候选词数量。模型自信时候选词少,模型犹豫时候选词多。
现状:目前 LLM 推理中,Top-p 比 Top-k 更常用,或者两者结合。
两种方法可以结合使用(如先取 top-k 再取 top-p),但通常分别实现。
代码实现:
def top_k_filtering(logits, top_k=50, temperature=1.0, filter_value=-float('Inf')):
"""
logits: [vocab_size] 或 [batch, vocab_size],模型输出的原始分数
top_k: 保留概率最大的 k 个词
temperature: 温度,大于 1 增加多样性,小于 1 增加确定性
"""
logits = logits / temperature
# 找出所有小于第 k 大值的索引,将其设为 filter_value (即 -inf,softmax 后为 0)
indices_to_remove = logits < torch.topk(logits, top_k)[0][..., -1, None]
logits[indices_to_remove] = filter_value
return logits
def top_p_filtering(logits, top_p=0.9, temperature=1.0, filter_value=-float('Inf'), min_tokens_to_keep=1):
"""
logits: [vocab_size] 或 [batch, vocab_size],模型输出的原始分数
top_p: 保留累积概率超过 p 的最小词集
temperature: 温度,大于 1 增加多样性,小于 1 增加确定性
"""
logits = logits / temperature
# 按概率从大到小排序
sorted_logits, sorted_indices = torch.sort(logits, descending=True, dim=-1) # [batch, vocab]
cumulated_probs = torch.cumsum(F.softmax(sorted_logits, dim=-1), dim=-1) # [batch, vocab]
# 创建 mask:累积概率 > p 的位置为 True(需要被移除)
sorted_indices_to_remove = cumulated_probs > top_p # [batch, vocab]
sorted_indices_to_remove[:, :min_tokens_to_keep] = False # 保留第一个超过 p 的 token(确保至少有一些 token)
# 注意:sorted_logits 是排序后的,需要映射回原始位置
indices_to_remove = sorted_indices_to_remove.scatter(1, sorted_indices, sorted_indices_to_remove)
logits[indices_to_remove] = filter_value # 将被移除的 logits 设为 -inf(softmax 后概率为 0)
return logits
if __name__ == "__main__":
# 模拟一个 vocab_size=5 的 logits
logits = torch.tensor([[2.0, 1.0, 0.1, 0.1, 0.1]])
print("原始 Logits:", logits)
filtered_logits = top_p_filtering(logits, top_p=0.8, temperature=1.0)
print("过滤后 Logits:", filtered_logits)
# 采样
probs = F.softmax(filtered_logits, dim=-1)
next_token = torch.multinomial(probs, num_samples=1)
print("采样结果 Token ID:", next_token)
2. LayerNorm 和 RMSNorm
概念讲解:
Layer Normalization (LayerNorm):对每个样本的每个特征层进行归一化。给定输入 x 形状 (batch, seq_len, hidden_size),对最后一个维度(hidden_size)计算均值和方差,然后标准化:(x - mean) / sqrt(var + eps),再乘以可学习的缩放参数 gamma 并加上偏移 beta。LayerNorm 广泛用于 Transformer,稳定训练。
RMSNorm:Root Mean Square Layer Normalization 是 LayerNorm 的一个简化变体。它假设减去均值不是必需的,只使用均方根 (RMS) 进行归一化:x / RMS(x),其中 RMS(x) = sqrt(mean(x^2) + eps)。同样乘以可学习的缩放参数 gamma,但没有 beta。RMSNorm 计算量更小,且在实验中性能与 LayerNorm 相当。
代码实现:
import torch
import torch.nn as nn
class LayerNorm(nn.Module):
"""
层归一化 (Layer Normalization)
公式: y = gamma * (x - mean) / sqrt(var + eps) + beta
"""
def __init__(self, hidden_size, eps=1e-5):
super().__init__()
self.gamma = nn.Parameter(torch.ones(hidden_size))
self.beta = nn.Parameter(torch.ones(hidden_size))
self.eps = eps
def forward(self, x):
# x: (batch, seq_len, hidden_size) 或 (batch, hidden_size)
mean = x.mean(dim=-1, keepdim=True)
var = x.var(dim=-1, keepdim=True, unbiased=False)
x_norm = (x - mean) / torch.sqrt(var + self.eps)
return self.gamma * x_norm + self.beta
class RMSNorm(nn.Module):
"""
RMS 归一化 (Root Mean Square Layer Normalization)
公式: y = gamma * x / RMS(x), 其中 RMS(x) = sqrt(mean(x^2) + eps)
"""
def __init__(self, hidden_size, eps=1e-5):
super().__init__()
self.gamma = nn.Parameter(torch.ones(hidden_size))
self.eps = eps
def forward(self, x):
# x: (batch, seq_len, hidden_size) 或 (batch, hidden_size)
rms = torch.sqrt(torch.mean(x, dim=-1, keepdim=True) + self.eps)
x_norm = x / rms
return self.gamma * x_norm
# 示例用法
if __name__ == "__main__":
batch, seq, hidden = 2, 3, 4
x = torch.randn(batch, seq, hidden)
ln = LayerNorm(hidden)
rms = RMSNorm(hidden)
out_ln = ln(x)
out_rms = rms(x)
print("LayerNorm 输出形状:", out_ln.shape)
print("RMSNorm 输出形状:", out_rms.shape)
3. SFT Loss 计算(Shift Right)
概念讲解:
在监督微调(Supervised Fine-Tuning, SFT)中,模型以自回归方式预测下一个 token。给定输入序列 input_ids,我们需要计算损失,使得模型预测的每个位置的 token 与真实的下一个 token 一致。
具体的,我们用第 1 到 t 个 token 预测第 t+1 个 token。因此:
输入 (input_ids):[BOS] 我 喜欢 深度 学习
标签 (labels):我 喜欢 深度 学习 [EOS]
即 labels 是 input_ids 左移一位(或说 input 右移一位对齐 labels)。代码实现中通常用 input_ids[:, :-1] 作为输入,input_ids[:, 1:] 作为预测目标。
实现 SFT loss 的关键操作是 Shift Right:
将 labels 设置为与 input_ids 相同(或者指定忽略的位置为 -100)。
在计算损失时,取 logits[..., :-1, :] 和 labels[..., 1:] 进行对齐。
或者,在输入模型时,将 input_ids 作为输入,labels 作为目标,模型内部可能会自动处理 shift(如 HuggingFace 的 transformers 库)。但手动实现时,我们需要显式 shift。
另外,通常会将填充部分(padding)的损失忽略,通过在 labels 中将填充 token 对应的位置设为 -100(因为 CrossEntropyLoss 默认忽略 -100 的目标)。
代码实现:
import torch
import torch.nn as nn
def sft_loss(logits, labels, ignore_index=-100):
"""
计算 SFT 的交叉熵损失,自动处理 shift right 和忽略 padding。
参数:
logits: 模型输出的 logits,形状 (batch_size, seq_len, vocab_size)
labels: 真实 token ids,形状 (batch_size, seq_len),其中填充部分设为 ignore_index
返回:
标量损失
"""
# shift logits 和 labels
# logits 的最后一个位置没有对应的下一个 token,所以我们取 logits[:, :-1, :]
# labels 的第一个位置没有对应的前一个预测,所以我们取 labels[:, 1:]
shift_logits = logits[:, :-1, :].contiguous() # (batch, seq_len-1, vocab)
shift_labels = labels[:, 1:].contiguous() # (batch, seq_len-1)
# 展平以便计算交叉熵
loss_function = nn.CrossEntropyLoss(ignore_index=ignore_index)
loss = loss_function(shift_logits.view(-1, shift_logits.size(-1)), shift_labels.view(-1))
return loss
# 示例用法
if __name__ == "__main__":
batch_size, seq_len, vocab_size = 2, 5, 10
logits = torch.randn(batch_size, seq_len, vocab_size)
labels = torch.randint(0, vocab_size, (batch_size, seq_len))
# 模拟 padding: 将序列后半部分设为 ignore_index
labels[:, 3:] = -100
loss = sft_loss(logits, labels)
print("SFT Loss:", loss.item())
4. 手撕 Softmax、交叉熵(Cross Entropy)
概念讲解:
这是深度学习最基础的算子,但面试常考 数值稳定性。
Softmax: \(P_i = \frac{e^{z_i}}{\sum e^{z_j}}\)
问题:如果 \(z_i\) 很大,\(e^{z_i}\) 会溢出 (Infinity)。
解决:利用 Softmax 的平移不变性,所有 \(z\) 减去最大值 \(\max(z)\)。即 \(P_i = \frac{e^{z_i - \max(z)}}{\sum e^{z_j - \max(z)}}\)。
Cross Entropy (CE): \(Loss = -\sum y_i \log(P_i)\)
在分类任务中,\(y\) 是 one-hot,所以简化为 \(-\log(P_{target})\)。
结合 Softmax:通常不单独算 Softmax 再算 Log,而是合并为 LogSoftmax,数值更稳定。
在以下代码中计算的“log sum exp”,是 softmax 概率的分母的 log。用原先的 logits 减去这个 log sum exp,就可以得到 log 概率(log prob)了。
代码实现:
import torch
def stable_softmax(logits, dim=-1):
"""
数值稳定的 Softmax 实现
"""
# 1. 减去最大值,防止 exp 溢出
# keepdim=True 保证形状可以广播
max_logits = torch.max(logits, dim=dim, keepdim=True)[0]
exp_logits = torch.exp(logits - max_logits)
# 2. 归一化
sum_exp_logits = torch.sum(exp_logits, dim=dim, keepdim=True)
probs = exp_logits / sum_exp_logits
return probs
def cross_entropy_loss(logits, targets):
"""
手写交叉熵 Loss
logits: [batch, vocab]
targets: [batch] 类别索引
"""
batch_size = logits.shape[0]
# 1. 数值稳定的 LogSoftmax
# log(softmax(x)) = x - max(x) - log(sum(exp(x - max(x))))
max_logits = torch.max(logits, dim=1, keepdim=True)[0]
log_sum_exp = max_logits + torch.log(torch.sum(torch.exp(logits - max_logits), dim=1, keepdim=True))
log_probs = logits - log_sum_exp
# 2. NLL Loss (Negative Log Likelihood)
# 取出目标类别对应的 log 概率
# targets 需要 unsqueeze 才能 gather
target_log_probs = log_probs.gather(1, targets.unsqueeze(1)).squeeze(1)
# 3. 取平均
loss = -torch.mean(target_log_probs)
return loss
# --- 测试 Demo ---
if __name__ == "__main__":
logits = torch.tensor([[2.0, 1.0, 0.1], [0.5, 2.5, 0.2]])
targets = torch.tensor([0, 1]) # 第一个样本目标是类 0,第二个是类 1
# 对比 PyTorch 原生实现
torch_ce = nn.CrossEntropyLoss()
torch_loss = torch_ce(logits, targets)
# 对比手写实现
my_loss = cross_entropy_loss(logits, targets)
print(f"PyTorch Loss: {torch_loss.item():.6f}")
print(f"My Loss: {my_loss.item():.6f}")
# 两者应该非常接近
5. QKV、Self-Attention 与 Multi-Head Attention (MHA)
概念讲解:
什么是 QKV?想象我们在图书馆找书:
Query (查询 Q):心中的问题,比如"我想找关于机器学习的书"
Key (键 K):每本书的标签/标题,比如"深度学习入门"、"Python编程"
Value (值 V):书里的实际内容
计算过程:
我们的问题(Q)和所有书的标签(K)做匹配,算出相似度(注意力分数)
根据相似度,决定从每本书(V)里提取多少信息
最后把所有书的内容按权重加权求和,得到答案
Self-Attention 的计算流程:
投影 (Projection):输入向量 \(X\) 通过三个不同的线性层,变成 \(Q, K, V\)。
打分 (Score):计算 \(Q \cdot K^T\)。这表示“当前词”和“其他词”的匹配程度。
缩放 (Scale):除以 \(\sqrt{d_k}\)。防止点积结果太大导致 Softmax 梯度消失。
归一化 (Softmax):把分数变成概率分布(权重和为 1)。
加权求和 (Weighted Sum):用权重乘以 \(V\)。匹配度高的地方,其内容 \(V\) 就被更多地保留下来。
核心公式:
\[\text{Attention}(Q, K, V) = \text{softmax}(\frac{QK^T}{\sqrt{d_k}})V
\]
\(\frac{QK^T}{\sqrt{d_k}}\):计算相似度并缩放(防止数值过大)
\(\text{softmax}\):转成概率分布(总和为1)
乘 V:按权重聚合信息
Multi-Head Attention (MHA):
为什么要多 Head? 就像一群人一起阅读一篇文章。有人关注语法结构,有人关注语义指代,有人关注情感色彩。
怎么做? 把 \(Q, K, V\) 切分成多份(Head),每份独立做一套 Self-Attention,最后把结果拼起来,再过一个线性层。
代码实现:
这段代码完全可运行,包含了 QKV 投影、缩放点积注意力、因果掩码 (Causal Mask) 和 多头机制。
import torch
import torch.nn as nn
import math
class ScaledDotProductAttention(nn.Module):
def forward(self, q, k, v, mask=None):
"""
q, k, v 形状:(batch_size, num_heads, seq_len, head_dim)
mask 形状:(batch_size, 1, 1, seq_len) 或 (batch_size, 1, seq_len, seq_len)
"""
# 1. 计算 Q 和 K 的点积 (Batch, Heads, Seq, Seq)
# 最后两个维度:q 的 seq_len 和 k 的 seq_len 做矩阵乘法
scores = torch.matmul(q, k.transpose(-2, -1)) # transpose(-2, -1) 是什么意思
# 2. 缩放 (Scale)
# head_dim 就是 k 的最后一个维度
scale = math.sqrt(q.size(-1)) # 根号 k
scores = scores / scale
# 3. 应用 Mask (如果是解码器,需要遮住未来的 token)
if mask is not None:
# mask 中为 0 的地方保留,为 1 的地方遮住 (通常 mask 是 1 表示遮住)
# 这里假设 mask 是 1 表示需要遮住的位置,我们将其设为负无穷
scores = scores.masked_fill(mask == 1, -1e9)
# 4. Softmax 归一化
attn_weights = torch.softmax(scores, dim=-1)
# 5. 加权求和 V
# (B, H, S, S) @ (B, H, S, D) -> (B, H, S, D)
output = torch.matmul(attn_weights, v)
return output, attn_weights
class MultiHeadAttention(nn.Module):
def __init__(self, hidden_dim, num_heads):
super().__init__()
assert hidden_dim % num_heads == 0, "hidden_dim 必须能被 num_heads 整除"
self.hidden_dim = hidden_dim
self.num_heads = num_heads
self.head_dim = hidden_dim // num_heads
# 1. QKV 的线性投影层
# 输入 (B, S, H) -> 输出 (B, S, H)
self.q_linear = nn.Linear(hidden_dim, hidden_dim)
self.k_linear = nn.Linear(hidden_dim, hidden_dim)
self.v_linear = nn.Linear(hidden_dim, hidden_dim)
# 2. 最后的输出投影层
self.out_linear = nn.Linear(hidden_dim, hidden_dim)
self.attention = ScaledDotProductAttention()
def forward(self, x, mask=None):
"""
x 形状:(batch_size, seq_len, hidden_dim)
"""
batch_size, seq_len, _ = x.shape
# 1. 计算 Q, K, V (B, S, H)
q = self.q_linear(x)
k = self.k_linear(x)
v = self.v_linear(x)
# 2. 拆分多头 (Reshape & Transpose)
# 目标形状:(B, num_heads, S, head_dim)
# 先变成 (B, S, num_heads, head_dim)
q = q.view(batch_size, seq_len, self.num_heads, self.head_dim)
k = k.view(batch_size, seq_len, self.num_heads, self.head_dim)
v = v.view(batch_size, seq_len, self.num_heads, self.head_dim)
# 再转置成 (B, num_heads, S, head_dim),方便后面做矩阵乘法
q = q.transpose(1, 2)
k = k.transpose(1, 2)
v = v.transpose(1, 2)
# 3. 进入注意力计算
# attn_out 形状:(B, num_heads, S, head_dim)
attn_out, _ = self.attention(q, k, v, mask=mask)
# 4. 合并多头 (Concat & Project)
# 先转置回 (B, S, num_heads, head_dim)
attn_out = attn_out.transpose(1, 2)
# 再 reshape 成 (B, S, hidden_dim)
attn_out = attn_out.contiguous().view(batch_size, seq_len, self.hidden_dim)
# 5. 最后过一层线性层
output = self.out_linear(attn_out)
return output
# --- 测试代码 ---
if __name__ == "__main__":
# 模拟输入:Batch=2, Seq_Len=5, Hidden_Dim=64
x = torch.randn(2, 5, 64)
# 创建一个简单的因果 Mask (下三角为 0,上三角为 1)
# 表示当前 token 只能看前面的,不能看后面的
mask = torch.triu(torch.ones(5, 5), diagonal=1).unsqueeze(0).unsqueeze(0)
mha = MultiHeadAttention(hidden_dim=64, num_heads=8)
out = mha(x, mask=mask)
print(f"输入形状:{x.shape}")
print(f"输出形状:{out.shape}")
# 预期输出:torch.Size([2, 5, 64])
面试加分点 (Tips):
维度变换:一定要口述清楚 transpose(1, 2) 是为了把 num_heads 放到第二维,方便并行计算。
Mask 的作用:在 LLM 预训练(Decoder-only)中,必须使用 Causal Mask 防止信息泄露(即预测第 t 个词时不能看到 t+1 及之后的词)。
为什么除以 \(\sqrt{d_k}\):防止点积过大,导致 Softmax 进入饱和区(梯度接近 0),训练变慢。
Flash Attention:如果面试官问优化,可以提一句标准的 Attention 是 \(O(N^2)\) 显存占用,而 FlashAttention 通过分块计算(Tiling)和重计算(Recomputation)减少了 HBM 访问,是推理加速的关键。
6. RoPE (Rotary Positional Embedding)
概念讲解:
为什么需要位置编码?Self-Attention 是位置无关的(permutation invariant):
输入 "人咬狗" 和 "狗咬人" 对 Attention 来说是一样的
但语义完全不同!需要让模型知道词的顺序
传统方法 vs 旋转位置编码 (Rotary Position Embedding, RoPE):
方法
原理
缺点
绝对位置编码 (原版Transformer)
每个位置学一个向量,加到词嵌入上
长度外推性差,没见过位置1000就懵了
RoPE (旋转位置编码)
用旋转矩阵编码相对位置
内积自然体现相对距离,外推性好,现在是主流
RoPE 是一种相对位置编码,它通过旋转矩阵将位置信息注入到 Query 和 Key 向量中,使得内积自然包含相对位置信息。核心思想:对于位置 \(m\) 的向量 \(\mathbf{x} \in \mathbb{R}^d\),将其每对相邻维度 \((x_{2i}, x_{2i+1})\) 看作二维平面上的点,并旋转角度 \(m \theta_i\),其中 \(\theta_i\) 是预定义的频率:
\[\theta_i = 10000^{-2i/d}, \quad i = 0, 1, \dots, d/2 - 1,
\quad \text{即} \; 0\le 2i, 2i+1 < m
\]
旋转后的向量为:
\[\begin{aligned}
x'_{2i} &= x_{2i} \cos(m\theta_i) - x_{2i+1} \sin(m\theta_i) \\
x'_{2i+1} &= x_{2i} \sin(m\theta_i) + x_{2i+1} \cos(m\theta_i)
\end{aligned}
\]
实际使用时,我们对 Query 和 Key 分别应用 RoPE,然后再计算注意力分数。这样,两个位置 \(m\) 和 \(n\) 的 Query 和 Key 的内积将只依赖于它们的相对位置 \(m-n\),实现了相对位置编码的效果。
实现步骤:
预计算所有位置 \(m = 0, 1, \dots, L-1\) 的 \(\cos(m\theta_i)\) 和 \(\sin(m\theta_i)\) 值,得到形状为 \((L, d/2)\) 的两个矩阵(或 \((L, d)\) 便于广播)。
对输入向量 \(\mathbf{x}\)(形状 \((L, d)\)),将其按最后一维拆分成两半:偶索引和奇索引,或者使用 reshape 成 \((L, d/2, 2)\)。
应用旋转公式。可用向量化操作避免循环。
代码实现:
下面用 NumPy 实现 RoPE 的前向计算,并给出一个简单的测试。
import numpy as np
def precompute_freqs_cis(dim, max_seq_len, theta=10000.0):
"""
预计算旋转角度的 cos 和 sin 值(复数形式常称为 cis,即 cos + i sin)
dim: 特征维度,必须是偶数
max_seq_len: 最大序列长度
theta: base 频率参数,通常为 10000
返回:
cos: (max_seq_len, dim//2)
sin: (max_seq_len, dim//2)
"""
assert dim % 2 == 0, "维度必须为偶数"
# 计算每个维度的频率 theta_i
i = np.arange(0, dim, 2) # i = 0, 2, 4, ..., dim-2
freqs = 1.0 / (theta ** (i / dim)) # (dim//2,)
# 生成位置序列
t = np.arange(max_seq_len) # (max_seq_len,)
# 计算角度 t * freqs (广播)
angles = np.outer(t, freqs) # (max_seq_len, dim//2)
cos = np.cos(angles) # (max_seq_len, dim//2)
sin = np.sin(angles) # (max_seq_len, dim//2)
return cos, sin
def apply_rotary_emb(x, cos, sin):
"""
对输入张量 x 应用旋转位置编码
x: (seq_len, dim) 或 (batch_size, seq_len, dim),dim 必须为偶数
cos, sin: (seq_len, dim//2) 预计算好的值
返回: 编码后的张量,形状与 x 相同
"""
# 将 x 按最后一维分成两半,reshape 成 (..., seq_len, dim//2, 2)
# 这样最后一维的两个元素分别是 x_{2i} 和 x_{2i+1}
orig_shape = x.shape
seq_len = x.shape[-2] # 倒数第二维是序列长度
d = x.shape[-1]
assert d % 2 == 0
# 将 x 重塑为 (..., seq_len, d//2, 2)
x_reshaped = x.reshape(*orig_shape[:-1], d//2, 2)
# 分别取出偶数索引和奇数索引部分(即两个分量)
x_even = x_reshaped[..., 0] # (..., seq_len, d//2)
x_odd = x_reshaped[..., 1] # (..., seq_len, d//2)
# 应用旋转公式:
# x'_{2i} = x_{2i} * cos - x_{2i+1} * sin
# x'_{2i+1} = x_{2i} * sin + x_{2i+1} * cos
# 注意 cos, sin 的形状为 (seq_len, d//2),需要广播到与 x_even 一致
# 如果 x 有 batch 维度,cos/sin 会自动广播到 batch 维(只要维度对齐)
rotated_even = x_even * cos - x_odd * sin
rotated_odd = x_even * sin + x_odd * cos
# 将两部分堆叠回 (..., seq_len, d//2, 2) 形状
rotated = np.stack([rotated_even, rotated_odd], axis=-1) # (..., seq_len, d//2, 2)
# 重塑回原始形状 (..., seq_len, d)
return rotated.reshape(*orig_shape)
# ========== 示例运行 ==========
if __name__ == "__main__":
# 参数
dim = 8
seq_len = 4
batch_size = 2
# 预计算 cos 和 sin
cos, sin = precompute_freqs_cis(dim, seq_len)
print("cos 形状:", cos.shape) # (4, 4) 因为 dim//2 = 4
print("sin 形状:", sin.shape)
# 随机生成 Query 和 Key(假设无 batch 或带 batch)
# 示例 1: 单序列
q = np.random.randn(seq_len, dim)
k = np.random.randn(seq_len, dim)
q_rot = apply_rotary_emb(q, cos, sin)
k_rot = apply_rotary_emb(k, cos, sin)
print("\n单序列 RoPE 后 Q 形状:", q_rot.shape)
# 示例 2: 带 batch 的序列
q_batch = np.random.randn(batch_size, seq_len, dim)
k_batch = np.random.randn(batch_size, seq_len, dim)
q_batch_rot = apply_rotary_emb(q_batch, cos, sin)
k_batch_rot = apply_rotary_emb(k_batch, cos, sin)
print("带 batch 的 RoPE 后 Q 形状:", q_batch_rot.shape)
# 验证内积是否只与相对位置有关(理论性质,这里简单打印差值)
# 取位置 1 和 2 的 Q/K 计算内积,并与位置 2 和 3 的比较(未严格验证,仅供演示)
attn1 = np.dot(q_rot[1], k_rot[2])
attn2 = np.dot(q_rot[2], k_rot[3])
print(f"内积 (pos1 Q, pos2 K): {attn1:.4f}")
print(f"内积 (pos2 Q, pos3 K): {attn2:.4f}")
说明:
precompute_freqs_cis 根据公式生成每个位置、每个维度对的 cos 和 sin 值。
apply_rotary_emb 将输入张量(最后两维为序列长度和特征维度)应用旋转。它支持任意形状,只要最后两维是 (..., seq_len, dim),且 dim 为偶数。
在示例中,我们分别对单序列和 batch 数据进行了测试。
最后简单验证了两个不同相对位置的内积值,实际 RoPE 的特性是它们应依赖于相对位置差,但这里仅为演示代码可用性。
面试加分点 (Tips):
为什么只加在 Q 和 K 上?
Attention 的核心是 \(Q \cdot K^T\)。RoPE 的性质保证了 \((Q_m \cdot K_n)\) 的内积只依赖于 \(m-n\)(相对位置)。V 代表内容,不需要位置信息来参与匹配。
RoPE 的外推性 (Extrapolation):
相比绝对位置编码,RoPE 在推理时如果序列长度超过训练长度,表现衰减更慢。因为它学习的是相对关系,而不是死记硬背位置 ID。
Theta 参数:
Llama 2/3 使用了 \(10000\),但为了支持更长上下文,Llama 3 或某些变体调整了 \(\theta\) 或者使用了 NTK-Aware Scaled RoPE。如果面试官问“如何支持 100k 上下文”,我们可以提到修改 \(\theta\) 或插值方法。
面试总结与建议
关于 Shift Right:这是 LLM 训练最核心的数据对齐逻辑。面试时如果能主动提到 contiguous() 的作用(内存连续)和 ignore_index(处理 padding),会非常加分。
关于数值稳定性:在写 Softmax 和 CrossEntropy 时,必须 提到 max subtraction。如果不提,面试官可能会认为缺乏工程经验。
关于 RMSNorm:现在 LLaMA、Qwen 等主流模型都用 RMSNorm。如果能说出它比 LayerNorm 少了减均值操作,计算更快,且效果在 LLM 上相当,会体现你对前沿架构的了解。
关于采样:实际推理中,Temperature 通常是在 Top-k/p 之前应用的。代码中体现这一点会显得更专业。
关于 MHA:理解维度变换 (B, S, H) -> (B, H, S, D) 是为了并行计算 Head。
关于 attention 的 Mask:LLM 是 Decoder-only 架构,Causal Mask 是必考点。
关于 RoPE 位置:记住 RoPE 是在 Q/K 投影之后,Attention 计算之前应用的。
