LLM 算法岗 | 八股题目 · 代码手撕 · 题目汇总与解析


目录

速记：

  1. Top-p 和 Top-k 采样：输入是 shape 为 [vocab_size] 的 logits，先 logits /= temperature，然后得到真的 prob = F.softmax(sorted_logits, dim=-1)，进行 topk topp 操作，最后把 indices_to_remove 的 logits 设为 filter_value。
  2. LayerNorm 有乘性参数 gamma 和加性参数 beta，归一化为 x - min 再 / sqrt(var)。RMSNorm 的 RMS 是 root mean square，有乘性参数 gamma，归一化为 x / sqrt(mean(x ** 2))。写成 nn.Module 的形式。self.gamma = nn.Parameter(torch.ones(hidden_size))
  3. SFT loss，输入是形状为 (batch_size, seq_len, vocab_size) 的 logits 和 labels，对 seq_len 维度进行 [:, :-1, :] 等操作，然后 lf = nn.CrossEntropyLoss(ignore_index=ignore_index)，loss = lf(shift_logits.view(-1, shift_logits.size(-1)), shift_labels.view(-1))。
  4. 手撕 PPO：主要关注 ppo_loss（记得重要性采样和 clip）和手写 GAE。
for t in reversed(range(T)):
    delta = rewards[t] + gamma * values[t+1] - values[t]
    gae = delta + gamma * lam * gae
    advantages[t] = gae
  1. （接上文）ppo 整个流程，先对收集到的 traj 计算 advantage 和 return，再通过调用 policy_net 得到 probs，然后计算 ppo_loss value_loss 等。
  2. 手撕 GRPO：grpo 的 advantage 是 (x - mean) / std。

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 * 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)：书里的实际内容

计算过程：

  1. 我们的问题(Q)和所有书的标签(K)做匹配，算出相似度（注意力分数）
  2. 根据相似度，决定从每本书(V)里提取多少信息
  3. 最后把所有书的内容按权重加权求和，得到答案

Self-Attention 的计算流程：

  1. 投影 (Projection)：输入向量 \(X\) 通过三个不同的线性层，变成 \(Q, K, V\)。
  2. 打分 (Score)：计算 \(Q \cdot K^T\)。这表示“当前词”和“其他词”的匹配程度。
  3. 缩放 (Scale)：除以 \(\sqrt{d_k}\)。防止点积结果太大导致 Softmax 梯度消失。
  4. 归一化 (Softmax)：把分数变成概率分布（权重和为 1）。
  5. 加权求和 (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)：

  1. 维度变换：一定要口述清楚 transpose(1, 2) 是为了把 num_heads 放到第二维，方便并行计算。
  2. Mask 的作用：在 LLM 预训练（Decoder-only）中，必须使用 Causal Mask 防止信息泄露（即预测第 t 个词时不能看到 t+1 及之后的词）。
  3. 为什么除以 \(\sqrt{d_k}\)：防止点积过大，导致 Softmax 进入饱和区（梯度接近 0），训练变慢。
  4. 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\)，实现了相对位置编码的效果。

实现步骤：

  1. 预计算所有位置 \(m = 0, 1, \dots, L-1\) 的 \(\cos(m\theta_i)\) 和 \(\sin(m\theta_i)\) 值，得到形状为 \((L, d/2)\) 的两个矩阵（或 \((L, d)\) 便于广播）。
  2. 对输入向量 \(\mathbf{x}\)（形状 \((L, d)\)），将其按最后一维拆分成两半：偶索引和奇索引，或者使用 reshape 成 \((L, d/2, 2)\)。
  3. 应用旋转公式。可用向量化操作避免循环。

代码实现：

下面用 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)：

  1. 为什么只加在 Q 和 K 上？
    • Attention 的核心是 \(Q \cdot K^T\)。RoPE 的性质保证了 \((Q_m \cdot K_n)\) 的内积只依赖于 \(m-n\)（相对位置）。V 代表内容，不需要位置信息来参与匹配。
  2. RoPE 的外推性 (Extrapolation)：
    • 相比绝对位置编码，RoPE 在推理时如果序列长度超过训练长度，表现衰减更慢。因为它学习的是相对关系，而不是死记硬背位置 ID。
  3. Theta 参数：
    • Llama 2/3 使用了 \(10000\)，但为了支持更长上下文，Llama 3 或某些变体调整了 \(\theta\) 或者使用了 NTK-Aware Scaled RoPE。如果面试官问“如何支持 100k 上下文”，我们可以提到修改 \(\theta\) 或插值方法。

一个问题：为什么位置编码要同时使用 sin 和 cos，而不是只使用其中一个？

  • 回答思路可能是，因为我们希望表达一个 位置 信息，而不是一个标量，一个标量会损失信息，比如 sin(x) = sin(π - x)，所以它没法代表一个唯一的角度 x；进一步，RoPE 可以通过 sin 和 cos 来表达相对位置信息。（我也不知道怎么答比较好）

7. 手撕 PPO、DPO、GRPO 的 loss function，以及 PPO 的 GAE 估计

1. PPO (Proximal Policy Optimization)

概念讲解

PPO 是一种策略梯度方法，通过限制策略更新的幅度来提高训练的稳定性。核心思想是：新策略与旧策略的差异不能太大，否则容易导致训练崩溃。

损失函数

PPO 的损失由三部分组成：

  • 策略损失 (Clipped Surrogate Objective)：使用截断机制防止更新过大。
  • 价值损失 (Value Function Loss)：通常使用均方误差 (MSE) 来拟合状态价值函数。
  • 熵奖励 (Entropy Bonus)：鼓励探索，通常加一个系数。

最终损失：

\[L^{CLIP}(\theta) = \mathbb{E}_t \left[ \min\left( r_t(\theta) \hat{A}_t, \text{clip}(r_t(\theta), 1-\epsilon, 1+\epsilon) \hat{A}_t \right) \right] \]

\[L^{VF}(\phi) = \mathbb{E}_t \left[ (V_\phi(s_t) - \hat{R}_t)^2 \right] \]

\[L^{total} = -L^{CLIP} + c_1 L^{VF} - c_2 S[\pi_\theta] \]

其中 $ r_t(\theta) = \frac{\pi_\theta(a_t|s_t)}{\pi_{\theta_{\text{old}}}(a_t|s_t)} $ 是概率比，\(\hat{A}_t\) 是优势估计，\(\hat{R}_t\) 是回报。

GAE (Generalized Advantage Estimation)

GAE 是一种权衡偏差和方差的方法，用于计算优势估计：

\[\hat{A}_t^{\text{GAE}(\gamma,\lambda)} = \sum_{l=0}^{\infty} (\gamma\lambda)^l \delta_{t+l} \]

其中 \(\delta_t = r_t + \gamma V(s_{t+1}) - V(s_t)\) 是 TD 误差。

PPO 伪代码

  1. 初始化策略网络 \(\pi_\theta\) 和价值网络 \(V_\phi\)。
  2. 对于每次迭代：
    • 使用当前策略 \(\pi_{\theta_{\text{old}}}\) 收集 \(N\) 条轨迹。
    • 计算 GAE 优势 \(\hat{A}_t\) 和目标回报 \(\hat{R}_t\)。
    • 对于 \(K\) 个 epoch：
      • 从数据中采样 batch。
      • 计算概率比 \(r_t\)。
      • 计算裁剪的策略损失。
      • 计算价值损失。
      • 计算熵损失。
      • 反向传播更新 \(\theta\) 和 \(\phi\)。
    • 更新 \(\theta_{\text{old}} \leftarrow \theta\)。

Python 代码示例 (PyTorch)

import torch
import torch.nn as nn
import torch.optim as optim

def ppo_loss(probs, old_probs, advantages, actions, clip_epsilon=0.2):
    """
    PPO 策略损失（裁剪目标）
    Args:
        probs: 当前策略的动作概率分布 (batch_size, num_actions)
        old_probs: 旧策略的动作概率分布 (batch_size, num_actions)
        advantages: 优势估计 (batch_size,)
        actions: 实际执行的动作索引 (batch_size,)
        clip_epsilon: 裁剪范围
    Returns:
        loss: 标量损失
    """
    # 取出对应动作的概率
    prob = probs.gather(1, actions.unsqueeze(1)).squeeze(1)   # (batch_size,)
    old_prob = old_probs.gather(1, actions.unsqueeze(1)).squeeze(1)
    # 概率比 r_t
    ratio = prob / (old_prob + 1e-8)
    # 裁剪后的比值
    clipped_ratio = torch.clamp(ratio, 1 - clip_epsilon, 1 + clip_epsilon)
    # 原始目标和裁剪目标
    surr1 = ratio * advantages
    surr2 = clipped_ratio * advantages
    # 取最小值并取负号（因为我们要最小化损失）
    loss = -torch.min(surr1, surr2).mean()
    return loss

def value_loss(values, returns):
    """
    价值网络损失（MSE）
    Args:
        values: 状态价值估计 (batch_size,)
        returns: 实际回报 (batch_size,)
    """
    return nn.MSELoss()(values, returns)

def entropy_loss(probs):
    """
    熵奖励（负熵，用于最大化熵）
    Args:
        probs: 动作概率分布 (batch_size, num_actions)
    Returns:
        负熵的均值（损失中减去熵奖励，所以此处返回 -entropy）
    """
    entropy = -(probs * torch.log(probs + 1e-8)).sum(dim=-1).mean()
    return -entropy   # 损失中会减去熵，因此返回负熵

def compute_gae(rewards, values, gamma=0.99, lam=0.95):
    """
    计算 GAE 优势
    Args:
        rewards: 奖励序列列表 (list of length T)
        values: 价值估计序列 (list of length T+1, 最后一个为 V(s_{T+1}))
        gamma: 折扣因子
        lam: GAE lambda 参数
    Returns:
        advantages: 优势序列 (list of length T)
        returns: 回报序列 (list of length T)
    """
    T = len(rewards)
    advantages = [0.0] * T
    gae = 0.0
    for t in reversed(range(T)):
        delta = rewards[t] + gamma * values[t+1] - values[t]
        gae = delta + gamma * lam * gae
        advantages[t] = gae
    # 计算回报：return = advantage + value
    returns = [advantages[t] + values[t] for t in range(T)]
    return advantages, returns

# ===== 伪代码框架 =====
def ppo_train_step(policy_net, value_net, optimizer_p, optimizer_v, trajectories, epochs=10, clip_epsilon=0.2):
    """
    单次 PPO 更新步骤（伪代码）
    trajectories: 收集的一批轨迹，每个轨迹包含 states, actions, rewards, old_probs, old_values
    """
    # 预先计算 GAE 和回报
    all_advantages = []
    all_returns = []
    for traj in trajectories:
        # traj: states, actions, rewards, old_probs, old_values (list of values at each step)
        # 注意 old_values 长度比 rewards 多 1，最后一个是 V(s_{T+1})
        adv, ret = compute_gae(traj['rewards'], traj['values'], gamma=0.99, lam=0.95)
        all_advantages.extend(adv)
        all_returns.extend(ret)
    all_advantages = torch.tensor(all_advantages)
    all_returns = torch.tensor(all_returns)
    # 将轨迹数据打包成 batch
    # ...（此处省略数据准备）
    for _ in range(epochs):
        # 随机采样 mini-batch
        for batch in dataloader:
            states, actions, old_probs, advantages, returns = batch
            # 前向传播
            new_probs = policy_net(states)  # 输出分布
            new_values = value_net(states).squeeze()
            # 计算损失
            loss_p = ppo_loss(new_probs, old_probs, advantages, actions, clip_epsilon)
            loss_v = value_loss(new_values, returns)
            entropy = entropy_loss(new_probs)
            total_loss = -loss_p + 0.5 * loss_v - 0.01 * entropy
            # 反向传播
            optimizer_p.zero_grad()
            optimizer_v.zero_grad()
            total_loss.backward()
            optimizer_p.step()
            optimizer_v.step()
        pass

2. DPO (Direct Preference Optimization)

概念讲解

DPO 是一种直接从偏好数据（如人类反馈）优化策略的方法，无需显式训练奖励模型。它基于 Bradley-Terry 模型，推导出策略的优化目标，通过最大化偏好数据的似然来更新策略。

DPO 假设我们有一组偏好对 \((x, y_w, y_l)\)，其中 \(y_w\) 优于 \(y_l\)。损失函数为：

\[L_{\text{DPO}}(\pi_\theta; \pi_{\text{ref}}) = -\mathbb{E}_{(x, y_w, y_l)} \left[ \log \sigma\left( \beta \log \frac{\pi_\theta(y_w|x)}{\pi_{\text{ref}}(y_w|x)} - \beta \log \frac{\pi_\theta(y_l|x)}{\pi_{\text{ref}}(y_l|x)} \right) \right] \]

其中 \(\pi_{\text{ref}}\) 是参考策略（通常是 SFT 后的模型），\(\beta\) 是控制 KL 散度强度的超参数。

Python 代码示例

import torch
import torch.nn.functional as F

def dpo_loss(pi_logps_w, pi_logps_l, ref_logps_w, ref_logps_l, beta=0.1):
    """
    DPO 损失
    Args:
        pi_logps_w: 当前策略下 winner 响应的对数概率
        pi_logps_l: 当前策略下 loser 响应的对数概率
        ref_logps_w: 参考策略下 winner 响应的对数概率
        ref_logps_l: 参考策略下 loser 响应的对数概率
        beta: 温度参数
    Returns:
        损失标量
    """
    # 计算对数比
    log_ratio_w = pi_logps_w - ref_logps_w
    log_ratio_l = pi_logps_l - ref_logps_l
    # 内部差值
    diff = beta * (log_ratio_w - log_ratio_l)
    # 损失 = -log(sigmoid(diff))
    loss = -F.logsigmoid(diff).mean()
    return loss

# 使用示例：
# 假设我们有一个 batch 的 prompt 和对应的 winner/loser 响应
# 通过模型 forward 得到每个响应的对数概率（log_probs）
# loss = dpo_loss(pi_logps_w, pi_logps_l, ref_logps_w, ref_logps_l)

3. GRPO (Group Relative Policy Optimization)

概念讲解

GRPO 是 DeepSeek 提出的一种强化学习方法，特别适用于大语言模型的训练。它类似于 PPO，但去掉了价值网络，改用组内相对优势来估计优势函数。GRPO 的核心思想是：对每个 prompt，采样多个响应（一组），然后根据这些响应的奖励计算组内的相对优势，从而更新策略。

GRPO 的损失函数包含：

  • 策略损失：使用组内相对优势，类似 PPO 的裁剪目标，但优势是通过组内奖励归一化得到的。
  • KL 散度正则：通过直接添加 KL 惩罚项来限制策略更新（而非像 PPO 那样使用 clip）。

损失函数：

\[L_{\text{GRPO}}(\theta) = \frac{1}{G} \sum_{i=1}^G \left[ \min\left( r_i(\theta) \hat{A}_i, \text{clip}(r_i(\theta), 1-\epsilon, 1+\epsilon) \hat{A}_i \right) - \beta \cdot \text{KL}(\pi_\theta \| \pi_{\text{ref}}) \right] \]

其中 \(\hat{A}_i = \frac{r_i - \text{mean}(\mathbf{r})}{\text{std}(\mathbf{r})}\)，即组内标准化后的奖励。KL 项通常用近似公式计算：

\[\text{KL}(\pi_\theta \| \pi_{\text{ref}}) = \frac{\pi_{\text{ref}}(y|x)}{\pi_\theta(y|x)} - \log \frac{\pi_{\text{ref}}(y|x)}{\pi_\theta(y|x)} - 1 \]

（或者直接使用对数概率差的近似）

Python 代码示例

import torch
import torch.nn.functional as F

def grpo_loss(probs, old_probs, rewards, clip_epsilon=0.2, beta=0.04):
    """
    GRPO 损失（简化版，假设每个 prompt 对应一组响应，且奖励已给出）
    Args:
        probs: 当前策略对每个响应的对数概率 (batch_size,)
        old_probs: 旧策略的对数概率 (batch_size,)
        rewards: 每个响应的奖励 (batch_size,)
        clip_epsilon: 裁剪范围
        beta: KL 惩罚系数
    Returns:
        损失标量
    """
    # 正确做法：reshape 为 (num_prompts, group_size)，然后按组标准化
    # 假设 rewards 已经按组排列，我们手动分组（此处简化演示）
    # 实际使用时，需要传入 group_ids 或按组处理
    # 以下为演示逻辑，假设 group_size = 4
    group_size = 4
    batch_size = probs.size(0)
    num_groups = batch_size // group_size
    # 将 rewards 重塑为 (num_groups, group_size)
    rewards_group = rewards.view(num_groups, group_size)
    mean_r = rewards_group.mean(dim=1, keepdim=True)
    std_r = rewards_group.std(dim=1, keepdim=True)
    advantages_group = (rewards_group - mean_r) / (std_r + 1e-8)
    advantages = advantages_group.view(-1)  # 展回 (batch_size,)

    # 计算概率比（使用对数概率差）
    ratio = torch.exp(probs - old_probs)   # r_t = π_new / π_old
    # 裁剪
    clipped_ratio = torch.clamp(ratio, 1 - clip_epsilon, 1 + clip_epsilon)
    surr1 = ratio * advantages
    surr2 = clipped_ratio * advantages
    policy_loss = -torch.min(surr1, surr2).mean()

    # KL 惩罚（使用近似公式）
    # 注意：KL(π_new || π_old) 的对数形式，这里使用 ref 作为旧策略
    # 实际 GRPO 中，ref 是冻结的参考策略，old 是更新前的策略，此处用 old_probs 代替 ref_logps
    # 更准确的 KL 项应当使用 ref 模型，但为简单，我们直接计算新旧策略的 KL
    kl = (old_probs - probs).exp() - (old_probs - probs) - 1
    kl_loss = kl.mean()

    total_loss = policy_loss + beta * kl_loss
    return total_loss

# 使用示例：
# 假设我们有 4 个 prompt，每个 prompt 采样 4 个响应，共 16 个样本
# probs, old_probs, rewards 是长度为 16 的张量，按 prompt 顺序排列
# loss = grpo_loss(probs, old_probs, rewards, clip_epsilon=0.2, beta=0.04)

总结：

  • PPO：通过裁剪概率比和 GAE 优势估计来稳定策略更新，需要价值网络。
  • DPO：直接利用偏好数据优化策略，无需奖励模型和强化学习采样，计算简洁。
  • GRPO：在 PPO 基础上简化，去掉价值网络，使用组内相对优势，适合大模型训练。

面试总结与建议

  1. 关于 Shift Right：这是 LLM 训练最核心的数据对齐逻辑。面试时如果能主动提到 contiguous() 的作用（内存连续）和 ignore_index（处理 padding），会非常加分。
  2. 关于数值稳定性：在写 Softmax 和 CrossEntropy 时，必须 提到 max subtraction。如果不提，面试官可能会认为缺乏工程经验。
  3. 关于 RMSNorm：现在 LLaMA、Qwen 等主流模型都用 RMSNorm。如果能说出它比 LayerNorm 少了减均值操作，计算更快，且效果在 LLM 上相当，会体现你对前沿架构的了解。
  4. 关于采样：实际推理中，Temperature 通常是在 Top-k/p 之前应用的。代码中体现这一点会显得更专业。
  5. 关于 MHA：理解维度变换 (B, S, H) -> (B, H, S, D) 是为了并行计算 Head。
  6. 关于 attention 的 Mask：LLM 是 Decoder-only 架构，Causal Mask 是必考点。
  7. 关于 RoPE 位置：记住 RoPE 是在 Q/K 投影之后，Attention 计算之前应用的。


posted @ 2026-03-15 21:57  MoonOut  阅读(3365)  评论(2)    收藏  举报