AI面试常见代码题

Last updated on October 11, 2026 pm

本文整理大模型算法面试中常见的 PyTorch 手写题。每道题都是独立实现；重点不仅是背出 API，还要能解释 Tensor Shape、Mask 语义和公式与代码的对应关系。

0. 统一约定

本文主要使用以下符号：

  • BB：Batch Size；
  • TT：Sequence Length；
  • DD：Hidden Size；
  • HH：Attention Head 数量；
  • Dh=D/HD_h=D/H：每个 Attention Head 的维度；
  • VV：Vocabulary Size；
  • mask=True/1：该位置有效；
  • Logits 表示 Softmax 前的分数，Log Probability 表示 log_softmax 后的值。

面试手写时应优先保证：

  1. 输入、输出 Shape 正确；
  2. Mask 的广播方向和语义正确；
  3. Padding 不进入 Loss；
  4. 数值计算使用 log_softmax、logsigmoid 等稳定实现；
  5. 需要固定的 Target、Old Log Probability 和 Advantage 使用 detach()。

S 级：需要熟练手写

1. Multi-head Attention

目标

输入：

X∈RB×T×D.X\in\mathbb R^{B\times T\times D}.

先投影得到 Q,K,VQ,K,V，再拆成 HH 个 Heads：

Q,K,V∈RB×H×T×Dh.Q,K,V\in\mathbb R^{B\times H\times T\times D_h}.

Scaled Dot-product Attention 为：

Attention⁡(Q,K,V)=softmax⁡(QKTDh+M)V.\operatorname{Attention}(Q,K,V) = \operatorname{softmax} \left( \frac{QK^T}{\sqrt{D_h}}+M \right)V.

其中 MM 是调用方直接传入的通用 Attention Mask，Shape 规定为 (B,T,T)(B,T,T)，可以是 Causal Mask、Padding Mask 或二者组合后的结果。本文约定 True 表示允许关注，False 表示屏蔽。

PyTorch 实现

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
import math

import torch
import torch.nn as nn


class MultiHeadSelfAttention(nn.Module):
    def __init__(self, hidden_size: int, num_heads: int):
        super().__init__()
        assert hidden_size % num_heads == 0

        self.hidden_size = hidden_size
        self.num_heads = num_heads
        self.head_dim = hidden_size // num_heads

        self.q_proj = nn.Linear(hidden_size, hidden_size, bias=False)
        self.k_proj = nn.Linear(hidden_size, hidden_size, bias=False)
        self.v_proj = nn.Linear(hidden_size, hidden_size, bias=False)
        self.out_proj = nn.Linear(hidden_size, hidden_size, bias=False)

    def _split_heads(self, x: torch.Tensor) -> torch.Tensor:
        # (B, T, D) -> (B, T, H, Dh) -> (B, H, T, Dh)
        batch_size, seq_len, _ = x.shape
        x = x.view(
            batch_size,
            seq_len,
            self.num_heads,
            self.head_dim,
        )
        return x.transpose(1, 2)

    def forward(
        self,
        x: torch.Tensor,
        mask: torch.Tensor | None = None,
    ) -> torch.Tensor:
        """
        x:      (B, T, D)
        mask:   (B, T, T)，True 表示允许关注
        return: (B, T, D)
        """
        batch_size, seq_len, _ = x.shape

        q = self._split_heads(self.q_proj(x))
        k = self._split_heads(self.k_proj(x))
        v = self._split_heads(self.v_proj(x))

        # scores: (B, H, T_query, T_key)
        scores = torch.matmul(q, k.transpose(-2, -1))
        scores = scores / math.sqrt(self.head_dim)

        if mask is not None:
            # (B, T, T) -> (B, 1, T, T)，在 Head 维广播。
            mask = mask[:, None, :, :].bool()
            scores = scores.masked_fill(~mask, float("-inf"))

        attn_probs = torch.softmax(scores, dim=-1)
        context = torch.matmul(attn_probs, v)  # (B, H, T, Dh)

        # (B, H, T, Dh) -> (B, T, H, Dh) -> (B, T, D)
        context = context.transpose(1, 2).contiguous()
        context = context.view(batch_size, seq_len, self.hidden_size)
        output = self.out_proj(context)
        return output

例如，Decoder Self-Attention 的 Causal Mask 可以由调用方构造：

1
2
3
4
5
causal_mask = torch.tril(
    torch.ones(T, T, dtype=torch.bool, device=x.device)
)
causal_mask = causal_mask.unsqueeze(0).expand(B, -1, -1)
output = attention(x, mask=causal_mask)

最容易写错的是：

  • K 应转置最后两个维度，而不是 Head 与 Sequence 维；
  • Attention Score 的 Shape 是 (B,H,T,T)(B,H,T,T)；
  • (B,T,T)(B,T,T) 的 Mask 需要增加 Head Dimension 后再广播；
  • masked_fill 必须发生在 Softmax 之前。

2. PPO Clipped Loss

目标

已知每个 Token 的：

  • old_log_probs：生成轨迹的 Old Policy Log Probability；
  • new_log_probs：Current Policy 重新计算的 Log Probability；
  • advantages：固定的 Advantage。

概率比为：

ρt=exp⁡(ℓtnew−ℓtold).\rho_t = \exp \left( \ell_t^{\mathrm{new}}-\ell_t^{\mathrm{old}} \right).

PPO Clipped Objective 为：

Jclip=Et[min⁡(ρtAt,clip⁡(ρt,1−ϵ,1+ϵ)At)].J_{\mathrm{clip}} = \mathbb E_t \left[ \min \left( \rho_tA_t, \operatorname{clip}(\rho_t,1-\epsilon,1+\epsilon)A_t \right) \right].

训练代码通常最小化其相反数。

PyTorch 实现

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
import torch


def ppo_clipped_loss(
    old_log_probs: torch.Tensor,
    new_log_probs: torch.Tensor,
    advantages: torch.Tensor,
    clip_eps: float = 0.2,
) -> torch.Tensor:
    """
    所有输入 Shape 均为 (B, T)，并假设其中都是有效 Token。
    """
    old_log_probs = old_log_probs.detach()
    advantages = advantages.detach()

    log_ratio = new_log_probs - old_log_probs
    ratio = torch.exp(log_ratio)

    surrogate1 = ratio * advantages
    surrogate2 = torch.clamp(
        ratio,
        1.0 - clip_eps,
        1.0 + clip_eps,
    ) * advantages

    objective = torch.minimum(surrogate1, surrogate2)
    return -objective.mean()

torch.minimum 不能简单替换成“始终使用裁剪后的 Ratio”。PPO 只阻止沿正确方向获得过大的额外收益；如果 Policy 向错误方向移动，仍要保留未裁剪分支提供纠正梯度。

面试手写通常默认输入位置全部有效。实际处理变长 Response 时，还需要用 Response Mask 排除 Padding Token。

3. GRPO Loss

目标

假设一个 Batch 有 BB 个 Prompt，每个 Prompt 生成 GG 条 Response：

R∈RB×G.R\in\mathbb R^{B\times G}.

GRPO 使用同组 Reward 的均值和标准差构造 Advantage：

Ab,g=Rb,g−μbσb+ε.A_{b,g} = \frac{ R_{b,g}-\mu_b }{ \sigma_b+\varepsilon }.

一个 Response 的 Advantage 会广播到该 Response 的所有 Tokens，再结合 Importance Ratio 和 PPO Clipping。

PyTorch 实现

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
import torch


def grpo_loss(
    rewards: torch.Tensor,
    old_log_probs: torch.Tensor,
    new_log_probs: torch.Tensor,
    clip_eps: float = 0.2,
    norm_eps: float = 1e-6,
) -> tuple[torch.Tensor, torch.Tensor]:
    """
    rewards:          (B, G)
    old_log_probs:    (B, G, T)
    new_log_probs:    (B, G, T)
    假设所有 Token 都有效。
    """
    group_mean = rewards.mean(dim=1, keepdim=True)
    group_std = rewards.std(
        dim=1,
        keepdim=True,
        unbiased=False,
    )

    # (B, G)
    advantages = (rewards - group_mean) / (group_std + norm_eps)

    # 同一条 Response 的所有 Tokens 共用 Response-level Advantage。
    token_advantages = advantages.unsqueeze(-1)  # (B, G, 1)

    ratio = torch.exp(
        new_log_probs - old_log_probs.detach()
    )
    surrogate1 = ratio * token_advantages.detach()
    surrogate2 = torch.clamp(
        ratio,
        1.0 - clip_eps,
        1.0 + clip_eps,
    ) * token_advantages.detach()

    token_objective = torch.minimum(surrogate1, surrogate2)
    loss = -token_objective.mean()
    return loss, advantages

若一组 Response 全部得到相同 Reward，则分子为 00，该组 Advantage 为 00，不会产生 Actor Gradient。unbiased=False 使用总体标准差，也避免 Group Size 较小时额外的 Bessel Correction。

实际处理变长 Response 时，需要在计算平均值时用 Response Mask 排除 Padding Token。

4. GAE

目标

TD Residual：

δt=rt+γ(1−dt)V(st+1)−V(st),\delta_t = r_t+\gamma(1-d_t)V(s_{t+1})-V(s_t),

GAE 从后向前递推：

At=δt+γλ(1−dt)At+1.A_t = \delta_t +\gamma\lambda(1-d_t)A_{t+1}.

其中 dt=1d_t=1 表示执行第 tt 步动作后 Episode 结束。

PyTorch 实现

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
import torch


@torch.no_grad()
def compute_gae(
    rewards: torch.Tensor,
    values: torch.Tensor,
    dones: torch.Tensor,
    mask: torch.Tensor,
    last_values: torch.Tensor | None = None,
    gamma: float = 0.99,
    gae_lambda: float = 0.95,
) -> tuple[torch.Tensor, torch.Tensor]:
    """
    rewards:     (B, T)
    values:      (B, T)
    dones:       (B, T)，当前位置之后是否终止
    mask:        (B, T)，有效时间步为 1
    last_values: (B,)，序列截断但未终止时的 Bootstrap Value
    """
    batch_size, seq_len = rewards.shape

    if last_values is None:
        last_values = torch.zeros(
            batch_size,
            device=values.device,
            dtype=values.dtype,
        )

    advantages = torch.zeros_like(rewards)
    next_advantage = torch.zeros_like(last_values)
    next_value = last_values

    for t in reversed(range(seq_len)):
        valid = mask[:, t].to(values.dtype)
        not_done = 1.0 - dones[:, t].to(values.dtype)

        delta = (
            rewards[:, t]
            + gamma * not_done * next_value
            - values[:, t]
        )

        next_advantage = (
            delta
            + gamma * gae_lambda * not_done * next_advantage
        )

        # Padding 位置既不保存 Advantage，也不能向前传播。
        next_advantage = next_advantage * valid
        advantages[:, t] = next_advantage

        # 向前移动后，当前 Value 成为前一位置的 next_value。
        next_value = torch.where(
            valid.bool(),
            values[:, t],
            next_value,
        )

    returns = advantages + values
    returns = returns * mask.to(returns.dtype)
    return advantages, returns

如果轨迹因为真正到达 EOS 而结束，应令最后有效位置的 done=1，此时不再 Bootstrap；如果只是因为固定长度切断但环境并未终止，则可以令 done=0，并通过 last_values Bootstrap。

5. DPO Loss

目标

每条偏好数据包含：

(x,yw,yl),(x,y_w,y_l),

其中 ywy_w 是 Chosen Response，yly_l 是 Rejected Response。定义 Policy 与 Reference Model 对整条 Response 的 Sequence Log Probability：

log⁡π(y∣x)=∑t=1Tlog⁡π(yt∣x,y<t).\log\pi(y\mid x) = \sum_{t=1}^{T} \log\pi(y_t\mid x,y_{<t}).

DPO Logit 为：

z=β[log⁡πθ(yw∣x)πref(yw∣x)−log⁡πθ(yl∣x)πref(yl∣x)].z = \beta \left[ \log\frac{\pi_\theta(y_w\mid x)}{\pi_{\mathrm{ref}}(y_w\mid x)} - \log\frac{\pi_\theta(y_l\mid x)}{\pi_{\mathrm{ref}}(y_l\mid x)} \right].

PyTorch 实现

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
import torch
import torch.nn.functional as F


def dpo_loss(
    policy_chosen_logps: torch.Tensor,
    policy_rejected_logps: torch.Tensor,
    ref_chosen_logps: torch.Tensor,
    ref_rejected_logps: torch.Tensor,
    beta: float = 0.1,
    label_smoothing: float = 0.0,
) -> tuple[torch.Tensor, dict[str, torch.Tensor]]:
    """
    四个输入均为 Shape (B,) 的 Sequence Log Probability。
    Reference Log Probability 应提前固定，不参与梯度传播。
    """
    policy_log_ratio = (
        policy_chosen_logps - policy_rejected_logps
    )
    reference_log_ratio = (
        ref_chosen_logps - ref_rejected_logps
    ).detach()

    logits = beta * (
        policy_log_ratio - reference_log_ratio
    )

    # 标准 DPO 对应 label_smoothing=0：
    # loss = -log(sigmoid(logits))。
    positive_loss = -F.logsigmoid(logits)
    negative_loss = -F.logsigmoid(-logits)
    loss = (
        (1.0 - label_smoothing) * positive_loss
        + label_smoothing * negative_loss
    ).mean()

    chosen_reward = beta * (
        policy_chosen_logps - ref_chosen_logps.detach()
    )
    rejected_reward = beta * (
        policy_rejected_logps - ref_rejected_logps.detach()
    )

    metrics = {
        "chosen_reward": chosen_reward.mean().detach(),
        "rejected_reward": rejected_reward.mean().detach(),
        "preference_accuracy": (
            chosen_reward > rejected_reward
        ).float().mean().detach(),
    }
    return loss, metrics

DPO 不要求 Reward Model、Critic 或在线 Rollout，但需要固定的 Chosen/Rejected 偏好对以及 Reference Model。四个输入必须是 Response Tokens 的 Log Probability 之和，不能把 Prompt Tokens 也计入。

A 级：常见基础组件

6. RMSNorm

目标

输入：

X∈RB×T×D.X\in\mathbb R^{B\times T\times D}.

RMSNorm 固定 b,tb,t，只沿 Hidden Dimension DD 计算：

RMS⁡(Xb,t,:)=1D∑j=1DXb,t,j2+ϵ.\operatorname{RMS}(X_{b,t,:}) = \sqrt{ \frac{1}{D}\sum_{j=1}^{D}X_{b,t,j}^2+\epsilon }.

与 LayerNorm 不同，RMSNorm 不减均值，通常也不使用 Bias：

Yb,t,j=wjXb,t,jRMS⁡(Xb,t,:).Y_{b,t,j} = w_j \frac{ X_{b,t,j} }{ \operatorname{RMS}(X_{b,t,:}) }.

PyTorch 实现

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
import torch
import torch.nn as nn


class RMSNorm(nn.Module):
    def __init__(self, hidden_size: int, eps: float = 1e-6):
        super().__init__()
        self.weight = nn.Parameter(torch.ones(hidden_size))
        self.eps = eps

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """
        x:      (..., D)
        return: (..., D)
        """
        # 使用 FP32 计算平方均值，避免 BF16/FP16 下数值不稳定。
        input_dtype = x.dtype
        x_float = x.float()

        mean_square = x_float.pow(2).mean(
            dim=-1,
            keepdim=True,
        )
        normalized = x_float * torch.rsqrt(
            mean_square + self.eps
        )

        return (
            normalized.to(input_dtype) * self.weight
        )

可学习参数只有 weight，Shape 为 (D,)(D,)，它会在前面的所有维度广播。

7. RoPE

目标

RoPE 将 Head Dimension 中相邻的两个维度组成二维向量。第 ii 对维度的角频率为：

ωi=base−2iDh,\omega_i = \mathrm{base}^{-\frac{2i}{D_h}},

位置 tt 的旋转角度为：

θt,i=tωi.\theta_{t,i}=t\omega_i.

对每一对元素 (x2i,x2i+1)(x_{2i},x_{2i+1}) 执行：

x2i′=x2icos⁡θ−x2i+1sin⁡θ,x2i+1′=x2isin⁡θ+x2i+1cos⁡θ.\begin{aligned} x'_{2i} &= x_{2i}\cos\theta-x_{2i+1}\sin\theta,\\ x'_{2i+1} &= x_{2i}\sin\theta+x_{2i+1}\cos\theta. \end{aligned}

RoPE 只应用于 Query 和 Key，不应用于 Value。

PyTorch 实现

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
import torch


def rotate_half(x: torch.Tensor) -> torch.Tensor:
    """
    (..., Dh) 中每对元素：
    (x_0, x_1) -> (-x_1, x_0)
    """
    x_even = x[..., 0::2]
    x_odd = x[..., 1::2]
    rotated = torch.stack(
        (-x_odd, x_even),
        dim=-1,
    )
    return rotated.flatten(start_dim=-2)


def build_rope_cache(
    seq_len: int,
    head_dim: int,
    device: torch.device,
    base: float = 10000.0,
) -> tuple[torch.Tensor, torch.Tensor]:
    assert head_dim % 2 == 0

    pair_index = torch.arange(
        head_dim // 2,
        device=device,
        dtype=torch.float32,
    )
    frequencies = base ** (
        -2.0 * pair_index / head_dim
    )

    positions = torch.arange(
        seq_len,
        device=device,
        dtype=torch.float32,
    )

    # angles: (T, Dh / 2)
    angles = torch.outer(positions, frequencies)

    # 每个角度重复两次，以匹配相邻的偶数、奇数维。
    angles = torch.repeat_interleave(
        angles,
        repeats=2,
        dim=-1,
    )

    # (1, 1, T, Dh)，在 Batch 和 Head 维广播。
    cos = angles.cos()[None, None, :, :]
    sin = angles.sin()[None, None, :, :]
    return cos, sin


def apply_rope(
    q: torch.Tensor,
    k: torch.Tensor,
    cos: torch.Tensor,
    sin: torch.Tensor,
) -> tuple[torch.Tensor, torch.Tensor]:
    """
    q, k:     (B, H, T, Dh)
    cos, sin: (1, 1, T, Dh)
    """
    q_rotated = q * cos + rotate_half(q) * sin
    k_rotated = k * cos + rotate_half(k) * sin
    return q_rotated, k_rotated

若使用 KV Cache，历史 Key 在写入 Cache 前就应应用其对应位置的 RoPE；新 Token 的 Query 和 Key 只使用新位置的角度。

8. KV Cache Attention

目标

自回归生成第 tt 个 Token 时，历史 Tokens 的 Key 和 Value 不会变化。KV Cache 保存：

K≤t,V≤t∈RB×H×Tcache×Dh.K_{\le t},V_{\le t} \in \mathbb R^{B\times H\times T_{\mathrm{cache}}\times D_h}.

每一步只计算新 Token 的 Qt,Kt,VtQ_t,K_t,V_t，把 Kt,VtK_t,V_t 追加到 Cache，然后让 QtQ_t 关注全部 Cached Keys。

PyTorch 实现

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
import math

import torch
import torch.nn as nn


class KVCacheAttention(nn.Module):
    def __init__(self, hidden_size: int, num_heads: int):
        super().__init__()
        assert hidden_size % num_heads == 0

        self.hidden_size = hidden_size
        self.num_heads = num_heads
        self.head_dim = hidden_size // num_heads

        self.q_proj = nn.Linear(hidden_size, hidden_size, bias=False)
        self.k_proj = nn.Linear(hidden_size, hidden_size, bias=False)
        self.v_proj = nn.Linear(hidden_size, hidden_size, bias=False)
        self.out_proj = nn.Linear(hidden_size, hidden_size, bias=False)

    def _split_heads(self, x: torch.Tensor) -> torch.Tensor:
        # (B, 1, D) -> (B, H, 1, Dh)
        batch_size, seq_len, _ = x.shape
        x = x.view(
            batch_size,
            seq_len,
            self.num_heads,
            self.head_dim,
        )
        return x.transpose(1, 2)

    def decode_step(
        self,
        x_new: torch.Tensor,
        cache: dict[str, torch.Tensor] | None = None,
        new_token_mask: torch.Tensor | None = None,
    ) -> tuple[torch.Tensor, dict[str, torch.Tensor]]:
        """
        x_new: (B, 1, D)，当前新 Token 的 Hidden State
        cache["k"], cache["v"]: (B, H, T_cache, Dh)
        cache["mask"]:          (B, T_cache)
        new_token_mask:         (B, 1)，当前 Token 是否有效
        """
        assert x_new.size(1) == 1

        if new_token_mask is None:
            new_token_mask = torch.ones(
                x_new.size(0),
                1,
                dtype=torch.bool,
                device=x_new.device,
            )
        else:
            new_token_mask = new_token_mask.bool()

        q_new = self._split_heads(self.q_proj(x_new))
        k_new = self._split_heads(self.k_proj(x_new))
        v_new = self._split_heads(self.v_proj(x_new))

        # 实际 LLM 通常会在这里对 q_new、k_new 应用 RoPE。

        if cache is None:
            k_all = k_new
            v_all = v_new
            cache_mask = new_token_mask
        else:
            k_all = torch.cat(
                [cache["k"], k_new],
                dim=2,
            )
            v_all = torch.cat(
                [cache["v"], v_new],
                dim=2,
            )
            cache_mask = torch.cat(
                [cache["mask"], new_token_mask],
                dim=1,
            )

        new_cache = {
            "k": k_all,
            "v": v_all,
            "mask": cache_mask,
        }

        # (B, H, 1, Dh) @ (B, H, Dh, T_cache)
        # -> (B, H, 1, T_cache)
        scores = torch.matmul(
            q_new,
            k_all.transpose(-2, -1),
        ) / math.sqrt(self.head_dim)

        # Query 是最新 Token，Cache 中不存在未来 Token，
        # 因此不需要额外的 Causal Mask，但仍要屏蔽 Padding KV。
        key_mask = cache_mask[:, None, None, :]
        scores = scores.masked_fill(
            ~key_mask,
            torch.finfo(scores.dtype).min,
        )

        attn_probs = torch.softmax(scores, dim=-1)
        context = torch.matmul(attn_probs, v_all)

        # (B, H, 1, Dh) -> (B, 1, D)
        context = context.transpose(1, 2).contiguous()
        context = context.view(
            x_new.size(0),
            1,
            self.hidden_size,
        )
        output = self.out_proj(context)
        output = output * new_token_mask.unsqueeze(-1)
        return output, new_cache

该实现同步维护 cache["mask"]，因此不同样本的 Prompt 长度不同时也不会关注 Padding KV。代码用 torch.cat 强调逻辑；生产环境会预分配最大 Cache，再按位置原地写入，避免每一步重新分配和复制历史 KV。

9. SFT Cross-entropy Loss

目标

Causal LM 在位置 tt 的 Logits 用来预测位置 t+1t+1 的 Token，因此需要 Shift：

1
logits[:, :-1] 预测 input_ids[:, 1:]

如果只希望对 Assistant Response 做 SFT，还要同时满足：

  • 目标位置不是 Padding；
  • 目标位置属于 Response。

PyTorch 实现

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
import torch
import torch.nn.functional as F


def sft_cross_entropy_loss(
    logits: torch.Tensor,
    input_ids: torch.Tensor,
    attention_mask: torch.Tensor,
    response_mask: torch.Tensor,
) -> torch.Tensor:
    """
    logits:         (B, L, V)
    input_ids:      (B, L)
    attention_mask: (B, L)，真实 Token 为 1
    response_mask:  (B, L)，Assistant Response Token 为 1
    """
    # 位置 t 的输出预测位置 t+1。
    shift_logits = logits[:, :-1, :].contiguous()
    shift_labels = input_ids[:, 1:].contiguous()

    # Mask 也必须与目标 Token 对齐，因此取 [:, 1:]。
    valid_mask = (
        attention_mask[:, 1:].bool()
        & response_mask[:, 1:].bool()
    )

    # token_loss: (B, L - 1)
    token_loss = F.cross_entropy(
        shift_logits.transpose(1, 2),
        shift_labels,
        reduction="none",
    )

    mask = valid_mask.to(token_loss.dtype)
    loss = (token_loss * mask).sum() / mask.sum().clamp_min(1.0)
    return loss

response_mask 标记的是 Token 自身是否属于 Response，所以 Shift 后应使用 response_mask[:, 1:]。如果把 Prompt 也设为有效，就会变成对整段对话计算 Language Modeling Loss。

10. 从 Logits 提取 Token Log Probability

目标

模型输出：

Z∈RB×L×V.Z\in\mathbb R^{B\times L\times V}.

先在 Vocabulary Dimension 做：

log⁡P=logsoftmax⁡(Z,dim⁡=−1),\log P=\operatorname{logsoftmax}(Z,\dim=-1),

再使用真实 Token ID 从 VV 个候选中 gather 对应位置。

PyTorch 实现

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
import torch
import torch.nn.functional as F


def get_next_token_log_probs(
    logits: torch.Tensor,
    input_ids: torch.Tensor,
) -> torch.Tensor:
    """
    logits:    (B, L, V)
    input_ids: (B, L)
    return:    (B, L - 1)

    返回第 1...L-1 个 Token 在模型中的 Log Probability。
    """
    # 位置 0...L-2 的 Logits 预测 Token 1...L-1。
    shift_logits = logits[:, :-1, :]
    target_ids = input_ids[:, 1:]

    log_probs = F.log_softmax(
        shift_logits,
        dim=-1,
    )

    # target_ids.unsqueeze(-1): (B, L - 1, 1)
    selected_log_probs = torch.gather(
        log_probs,
        dim=-1,
        index=target_ids.unsqueeze(-1),
    ).squeeze(-1)

    return selected_log_probs

若只需要 Response Token 的 Log Probability，再使用 Shift 后的 Response Mask：

1
2
token_log_probs = get_next_token_log_probs(logits, input_ids)
response_log_probs = token_log_probs * response_mask[:, 1:]

不能先 softmax 再 log，log_softmax 的数值稳定性更好。

11. KL Divergence 与 Entropy

目标

给定 Policy Distribution pp 和 Reference Distribution qq：

DKL(p∥q)=∑ap(a)[log⁡p(a)−log⁡q(a)].D_{\mathrm{KL}}(p\Vert q) = \sum_a p(a) \left[ \log p(a)-\log q(a) \right].

Policy Entropy 为：

H(p)=−∑ap(a)log⁡p(a).H(p) = -\sum_a p(a)\log p(a).

两者都在 Vocabulary Dimension 上求和。

PyTorch 实现

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
import torch
import torch.nn.functional as F


def categorical_kl_and_entropy(
    policy_logits: torch.Tensor,
    reference_logits: torch.Tensor,
) -> tuple[torch.Tensor, torch.Tensor]:
    """
    输入:  (..., V)
    输出:
        kl:      (...)
        entropy: (...)
    """
    policy_log_probs = F.log_softmax(
        policy_logits,
        dim=-1,
    )
    reference_log_probs = F.log_softmax(
        reference_logits,
        dim=-1,
    )
    policy_probs = policy_log_probs.exp()

    kl = (
        policy_probs
        * (policy_log_probs - reference_log_probs)
    ).sum(dim=-1)

    entropy = -(
        policy_probs * policy_log_probs
    ).sum(dim=-1)

    return kl, entropy


def sampled_log_ratio(
    policy_token_log_probs: torch.Tensor,
    reference_token_log_probs: torch.Tensor,
) -> torch.Tensor:
    """
    输入 Shape 通常为 (B, T_response)。
    这是已采样 Token 上的 Log Ratio，不是对整个词表求和的精确 KL。
    只有 Token 确实从当前 Policy 采样时，其样本均值才直接估计
    KL(Current Policy || Reference Policy)。
    """
    return (
        policy_token_log_probs
        - reference_token_log_probs
    )

第一种实现需要保留完整 Vocabulary Distribution，计算精确但显存较高；大模型 RL 中经常只保存已采样 Token 的 Log Probability，使用第二种 Sample-based Estimator。此时必须说明 Token 的采样分布：若数据来自不同的 Old Policy，就不能直接把未加权的样本均值称为 Current Policy 对 Reference Policy 的无偏 KL 估计。

B 级：进阶结构与工程题

12. GQA 与 MQA

目标

MHA 中 Query、Key、Value 的 Head 数相同：

Hq=Hkv.H_q=H_{kv}.

GQA 让多个 Query Heads 共享一组 Key/Value Head：

1<Hkv<Hq.1<H_{kv}<H_q.

MQA 是 GQA 的特殊情况：

Hkv=1.H_{kv}=1.

要求 HqH_q 能被 HkvH_{kv} 整除。实现时将每个 KV Head 重复：

nrep=HqHkvn_{\mathrm{rep}}=\frac{H_q}{H_{kv}}

次，使其与 Query Heads 对齐。

PyTorch 实现

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
import math

import torch
import torch.nn as nn


class GroupedQueryAttention(nn.Module):
    def __init__(
        self,
        hidden_size: int,
        num_query_heads: int,
        num_kv_heads: int,
    ):
        super().__init__()
        assert hidden_size % num_query_heads == 0
        assert num_query_heads % num_kv_heads == 0

        self.hidden_size = hidden_size
        self.num_query_heads = num_query_heads
        self.num_kv_heads = num_kv_heads
        self.head_dim = hidden_size // num_query_heads
        self.num_repeats = num_query_heads // num_kv_heads

        self.q_proj = nn.Linear(
            hidden_size,
            num_query_heads * self.head_dim,
            bias=False,
        )
        self.k_proj = nn.Linear(
            hidden_size,
            num_kv_heads * self.head_dim,
            bias=False,
        )
        self.v_proj = nn.Linear(
            hidden_size,
            num_kv_heads * self.head_dim,
            bias=False,
        )
        self.out_proj = nn.Linear(
            num_query_heads * self.head_dim,
            hidden_size,
            bias=False,
        )

    def forward(
        self,
        x: torch.Tensor,
        attention_mask: torch.Tensor | None = None,
    ) -> torch.Tensor:
        """
        x: (B, T, D)
        """
        batch_size, seq_len, _ = x.shape

        q = self.q_proj(x).view(
            batch_size,
            seq_len,
            self.num_query_heads,
            self.head_dim,
        ).transpose(1, 2)

        k = self.k_proj(x).view(
            batch_size,
            seq_len,
            self.num_kv_heads,
            self.head_dim,
        ).transpose(1, 2)

        v = self.v_proj(x).view(
            batch_size,
            seq_len,
            self.num_kv_heads,
            self.head_dim,
        ).transpose(1, 2)

        # (B, Hkv, T, Dh) -> (B, Hq, T, Dh)
        k = torch.repeat_interleave(
            k,
            repeats=self.num_repeats,
            dim=1,
        )
        v = torch.repeat_interleave(
            v,
            repeats=self.num_repeats,
            dim=1,
        )

        scores = torch.matmul(
            q,
            k.transpose(-2, -1),
        ) / math.sqrt(self.head_dim)

        causal_mask = torch.tril(
            torch.ones(
                seq_len,
                seq_len,
                dtype=torch.bool,
                device=x.device,
            )
        )[None, None, :, :]

        allowed = causal_mask
        if attention_mask is not None:
            allowed = (
                allowed
                & attention_mask.bool()[:, None, None, :]
            )

        scores = scores.masked_fill(
            ~allowed,
            torch.finfo(scores.dtype).min,
        )
        probs = torch.softmax(scores, dim=-1)
        output = torch.matmul(probs, v)

        output = output.transpose(1, 2).contiguous()
        output = output.view(batch_size, seq_len, self.hidden_size)
        output = self.out_proj(output)

        if attention_mask is not None:
            output = output * attention_mask.unsqueeze(-1)

        return output

减少 HkvH_{kv} 不会减少 Query Heads，因此仍能保留多组 Query 表示；它主要减少 K/V Projection 参数和 KV Cache 大小。

13. Top-k 与 Top-p Sampling

目标

  • Top-k：只保留 Logits 最大的 kk 个 Tokens；
  • Top-p：保留按概率从高到低排序后，累计概率刚达到阈值 pp 的最小 Token 集合；
  • Temperature：使用 z/τz/\tau 调整分布尖锐程度。

过滤后重新 Softmax，再使用 torch.multinomial 采样。

PyTorch 实现

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
import torch


@torch.no_grad()
def top_k_top_p_sample(
    logits: torch.Tensor,
    temperature: float = 1.0,
    top_k: int | None = None,
    top_p: float | None = None,
) -> torch.Tensor:
    """
    logits: (B, V)
    return: (B,)，采样得到的 Token IDs
    """
    if temperature <= 0:
        raise ValueError("temperature must be positive")

    filtered_logits = logits / temperature

    if top_k is not None and top_k > 0:
        top_k = min(top_k, filtered_logits.size(-1))
        top_k_indices = torch.topk(
            filtered_logits,
            k=top_k,
            dim=-1,
        ).indices

        # 直接按 Top-k Indices 构造 Mask，严格保留 k 个位置；
        # 即使第 k 大附近存在相同 Logit，也不会多保留 Token。
        keep_mask = torch.zeros_like(
            filtered_logits,
            dtype=torch.bool,
        )
        keep_mask.scatter_(
            dim=-1,
            index=top_k_indices,
            value=True,
        )
        filtered_logits = filtered_logits.masked_fill(
            ~keep_mask,
            float("-inf"),
        )

    if top_p is not None and top_p < 1.0:
        if top_p <= 0:
            raise ValueError("top_p must be in (0, 1]")

        sorted_logits, sorted_indices = torch.sort(
            filtered_logits,
            descending=True,
            dim=-1,
        )
        sorted_probs = torch.softmax(sorted_logits, dim=-1)
        cumulative_probs = torch.cumsum(sorted_probs, dim=-1)

        # 先找累计概率超过 top_p 的位置。
        remove_sorted = cumulative_probs > top_p

        # 右移一位，保留第一个使累计概率超过 top_p 的 Token。
        remove_sorted[..., 1:] = remove_sorted[
            ..., :-1
        ].clone()
        remove_sorted[..., 0] = False

        # 把排序空间中的 Mask 映射回原 Vocabulary 顺序。
        remove_original = torch.zeros_like(
            remove_sorted
        ).scatter(
            dim=-1,
            index=sorted_indices,
            src=remove_sorted,
        )

        filtered_logits = filtered_logits.masked_fill(
            remove_original,
            float("-inf"),
        )

    probs = torch.softmax(filtered_logits, dim=-1)
    token_ids = torch.multinomial(
        probs,
        num_samples=1,
    ).squeeze(-1)
    return token_ids

Top-p 中必须保留“第一个越过阈值”的 Token，否则保留集合的累计概率可能小于指定的 pp。

14. LoRA Linear

目标

冻结原始权重：

W0∈RDout×Din,W_0\in\mathbb R^{D_{\mathrm{out}}\times D_{\mathrm{in}}},

只训练两个低秩矩阵：

A∈Rr×Din,B∈RDout×r.A\in\mathbb R^{r\times D_{\mathrm{in}}}, \qquad B\in\mathbb R^{D_{\mathrm{out}}\times r}.

前向传播为：

y=W0x+αrBAx.y = W_0x + \frac{\alpha}{r}BAx.

PyTorch 实现

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
import math

import torch
import torch.nn as nn
import torch.nn.functional as F


class LoRALinear(nn.Module):
    def __init__(
        self,
        in_features: int,
        out_features: int,
        rank: int = 8,
        alpha: float = 16.0,
        dropout: float = 0.0,
        bias: bool = True,
    ):
        super().__init__()
        assert rank > 0

        self.base = nn.Linear(
            in_features,
            out_features,
            bias=bias,
        )
        self.base.weight.requires_grad_(False)
        if self.base.bias is not None:
            self.base.bias.requires_grad_(False)

        self.lora_a = nn.Parameter(
            torch.empty(rank, in_features)
        )
        self.lora_b = nn.Parameter(
            torch.zeros(out_features, rank)
        )

        self.scaling = alpha / rank
        self.dropout = nn.Dropout(dropout)

        nn.init.kaiming_uniform_(
            self.lora_a,
            a=math.sqrt(5),
        )
        # B 初始化为 0，使 LoRA 初始增量严格为 0。
        nn.init.zeros_(self.lora_b)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        base_output = self.base(x)

        # F.linear 的权重 Shape 是 (out_features, in_features)。
        low_rank = F.linear(
            self.dropout(x),
            self.lora_a,
        )
        lora_output = F.linear(
            low_rank,
            self.lora_b,
        )

        return base_output + self.scaling * lora_output

    @torch.no_grad()
    def merged_weight(self) -> torch.Tensor:
        """
        推理时可将 LoRA 增量合并进基础权重。
        """
        delta_weight = self.lora_b @ self.lora_a
        return self.base.weight + self.scaling * delta_weight

lora_b 初始化为零非常重要：训练开始时模块行为与原始 Linear 完全一致，但 lora_b 在第一步可以获得梯度；之后 lora_a 也会逐渐获得非零梯度。

15. FlashAttention 原理伪代码

目标

普通 Attention 会显式构造：

S=QKT∈RB×H×T×T,S=QK^T\in\mathbb R^{B\times H\times T\times T},

其显存复杂度为 O(T2)O(T^2)。FlashAttention 将 Query 和 Key/Value 切成 Blocks，并使用 Online Softmax 维护：

  • 当前最大值 mm；
  • Softmax 分母 ll；
  • 未归一化输出累积量 OO。

这样不需要把完整的 T×TT\times T Attention Matrix 写回显存。

PyTorch 原理实现

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
import math

import torch


def flash_attention_forward(
    q: torch.Tensor,
    k: torch.Tensor,
    v: torch.Tensor,
    block_size: int = 64,
    causal: bool = True,
) -> torch.Tensor:
    """
    q, k, v: (B, H, T, Dh)

    这是说明 Online Softmax 的 PyTorch 伪代码，
    不会比官方 CUDA/Triton Kernel 更快。
    """
    batch_size, num_heads, seq_len, head_dim = q.shape

    # 用 FP32 累积，提高数值稳定性。
    output = torch.zeros_like(q, dtype=torch.float32)
    scale = 1.0 / math.sqrt(head_dim)

    for q_start in range(0, seq_len, block_size):
        q_end = min(q_start + block_size, seq_len)
        q_block = q[:, :, q_start:q_end, :].float()
        q_len = q_end - q_start

        # 每个 Query Row 独立维护 Online Softmax 状态。
        running_max = torch.full(
            (batch_size, num_heads, q_len, 1),
            float("-inf"),
            device=q.device,
            dtype=torch.float32,
        )
        running_sum = torch.zeros_like(running_max)
        running_output = torch.zeros(
            batch_size,
            num_heads,
            q_len,
            head_dim,
            device=q.device,
            dtype=torch.float32,
        )

        for k_start in range(0, seq_len, block_size):
            k_end = min(k_start + block_size, seq_len)

            # 整个 Key Block 都在当前 Query Block 未来时可直接跳过。
            if causal and k_start >= q_end:
                break

            k_block = k[:, :, k_start:k_end, :].float()
            v_block = v[:, :, k_start:k_end, :].float()

            scores = torch.matmul(
                q_block,
                k_block.transpose(-2, -1),
            ) * scale

            if causal:
                q_positions = torch.arange(
                    q_start,
                    q_end,
                    device=q.device,
                )[:, None]
                k_positions = torch.arange(
                    k_start,
                    k_end,
                    device=q.device,
                )[None, :]

                causal_mask = (
                    k_positions <= q_positions
                )[None, None, :, :]
                scores = scores.masked_fill(
                    ~causal_mask,
                    float("-inf"),
                )

            block_max = scores.max(
                dim=-1,
                keepdim=True,
            ).values
            new_max = torch.maximum(
                running_max,
                block_max,
            )

            # 旧累积量换到新的指数尺度。
            old_scale = torch.exp(
                running_max - new_max
            )
            block_exp = torch.exp(scores - new_max)

            running_sum = (
                running_sum * old_scale
                + block_exp.sum(dim=-1, keepdim=True)
            )
            running_output = (
                running_output * old_scale
                + torch.matmul(block_exp, v_block)
            )
            running_max = new_max

        output[:, :, q_start:q_end, :] = (
            running_output
            / running_sum.clamp_min(1e-12)
        )

    return output.to(q.dtype)

真正的 FlashAttention 会将 Block 计算融合进 CUDA/Triton Kernel，并实现对应的分块反向传播。面试重点是解释 Online Softmax 如何在最大值变化时重新缩放旧的分母与输出累积量。

16. Transformer Block

目标

现代 Decoder-only LLM 通常使用 Pre-Norm：

X1=X+Attention⁡(RMSNorm⁡(X)),X_1 = X+\operatorname{Attention}(\operatorname{RMSNorm}(X)),

Xout=X1+SwiGLU⁡(RMSNorm⁡(X1)).X_{\mathrm{out}} = X_1+\operatorname{SwiGLU}(\operatorname{RMSNorm}(X_1)).

SwiGLU 为：

SwiGLU⁡(x)=Wdown[SiLU⁡(Wgatex)⊙(Wupx)].\operatorname{SwiGLU}(x) = W_{\mathrm{down}} \left[ \operatorname{SiLU}(W_{\mathrm{gate}}x) \odot (W_{\mathrm{up}}x) \right].

PyTorch 实现

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
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
import math

import torch
import torch.nn as nn
import torch.nn.functional as F


class RMSNorm(nn.Module):
    def __init__(self, hidden_size: int, eps: float = 1e-6):
        super().__init__()
        self.weight = nn.Parameter(torch.ones(hidden_size))
        self.eps = eps

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        variance = x.float().pow(2).mean(
            dim=-1,
            keepdim=True,
        )
        normalized = x.float() * torch.rsqrt(
            variance + self.eps
        )
        return normalized.to(x.dtype) * self.weight


class CausalSelfAttention(nn.Module):
    def __init__(self, hidden_size: int, num_heads: int):
        super().__init__()
        assert hidden_size % num_heads == 0
        self.hidden_size = hidden_size
        self.num_heads = num_heads
        self.head_dim = hidden_size // num_heads

        self.qkv_proj = nn.Linear(
            hidden_size,
            3 * hidden_size,
            bias=False,
        )
        self.out_proj = nn.Linear(
            hidden_size,
            hidden_size,
            bias=False,
        )

    def forward(
        self,
        x: torch.Tensor,
        attention_mask: torch.Tensor | None = None,
    ) -> torch.Tensor:
        batch_size, seq_len, _ = x.shape

        q, k, v = self.qkv_proj(x).chunk(3, dim=-1)

        def split_heads(tensor: torch.Tensor) -> torch.Tensor:
            return tensor.view(
                batch_size,
                seq_len,
                self.num_heads,
                self.head_dim,
            ).transpose(1, 2)

        q, k, v = map(split_heads, (q, k, v))

        scores = torch.matmul(
            q,
            k.transpose(-2, -1),
        ) / math.sqrt(self.head_dim)

        allowed = torch.tril(
            torch.ones(
                seq_len,
                seq_len,
                dtype=torch.bool,
                device=x.device,
            )
        )[None, None, :, :]

        if attention_mask is not None:
            allowed = (
                allowed
                & attention_mask.bool()[:, None, None, :]
            )

        scores = scores.masked_fill(
            ~allowed,
            torch.finfo(scores.dtype).min,
        )
        probs = torch.softmax(scores, dim=-1)
        output = torch.matmul(probs, v)

        output = output.transpose(1, 2).contiguous()
        output = output.view(batch_size, seq_len, self.hidden_size)
        return self.out_proj(output)


class SwiGLU(nn.Module):
    def __init__(
        self,
        hidden_size: int,
        intermediate_size: int,
    ):
        super().__init__()
        self.gate_proj = nn.Linear(
            hidden_size,
            intermediate_size,
            bias=False,
        )
        self.up_proj = nn.Linear(
            hidden_size,
            intermediate_size,
            bias=False,
        )
        self.down_proj = nn.Linear(
            intermediate_size,
            hidden_size,
            bias=False,
        )

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        gated = F.silu(self.gate_proj(x))
        values = self.up_proj(x)
        return self.down_proj(gated * values)


class TransformerBlock(nn.Module):
    def __init__(
        self,
        hidden_size: int,
        num_heads: int,
        intermediate_size: int,
    ):
        super().__init__()
        self.attn_norm = RMSNorm(hidden_size)
        self.attention = CausalSelfAttention(
            hidden_size,
            num_heads,
        )
        self.ffn_norm = RMSNorm(hidden_size)
        self.ffn = SwiGLU(
            hidden_size,
            intermediate_size,
        )

    def forward(
        self,
        x: torch.Tensor,
        attention_mask: torch.Tensor | None = None,
    ) -> torch.Tensor:
        x = x + self.attention(
            self.attn_norm(x),
            attention_mask,
        )
        x = x + self.ffn(self.ffn_norm(x))

        if attention_mask is not None:
            x = x * attention_mask.unsqueeze(-1)

        return x

完整现代 LLM 还会加入 RoPE、GQA、KV Cache、FlashAttention 和 Tensor Parallel；面试手写 Block 时，先保证 Pre-Norm、Residual 和 SwiGLU 的数据流正确。

17. RL Rollout Masking

目标

同一个 Batch 中的 Responses 长度不同，需要 Padding 到统一长度：

1
2
Response 1: y1 y2 EOS PAD PAD
Response 2: y1 y2 y3  y4 EOS

Response Mask 应包含真实 Tokens 和第一个 EOS，但排除 EOS 后的 Padding。Reward、Log Probability、Advantage 和 PPO Loss 都必须使用同一套有效位置定义。

PyTorch 实现

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
import torch


def mask_from_lengths(
    lengths: torch.Tensor,
    max_length: int,
) -> torch.Tensor:
    """
    lengths: (B,)，包含 EOS 在内的有效 Response 长度
    return:  (B, T)
    """
    positions = torch.arange(
        max_length,
        device=lengths.device,
    )
    return positions.unsqueeze(0) < lengths.unsqueeze(1)


def infer_lengths_from_eos(
    responses: torch.Tensor,
    eos_token_id: int,
) -> torch.Tensor:
    """
    responses: (B, T)
    返回包含第一个 EOS 在内的长度；
    若没有 EOS，则整条 Response 都有效。
    """
    batch_size, max_length = responses.shape
    positions = torch.arange(
        max_length,
        device=responses.device,
    ).expand(batch_size, -1)

    eos_positions = torch.where(
        responses.eq(eos_token_id),
        positions,
        torch.full_like(positions, max_length),
    )
    first_eos = eos_positions.min(dim=1).values

    lengths = torch.where(
        first_eos < max_length,
        first_eos + 1,
        torch.full_like(first_eos, max_length),
    )
    return lengths


def build_rollout_training_tensors(
    responses: torch.Tensor,
    sequence_rewards: torch.Tensor,
    eos_token_id: int,
    response_lengths: torch.Tensor | None = None,
    terminated: torch.Tensor | None = None,
) -> dict[str, torch.Tensor]:
    """
    responses:         (B, T)
    sequence_rewards:  (B,)
    response_lengths:  (B,)，Rollout Engine 提供时应优先使用
    terminated:        (B,)，是否因 EOS 真正终止
    """
    batch_size, max_length = responses.shape

    if response_lengths is None:
        response_lengths = infer_lengths_from_eos(
            responses,
            eos_token_id,
        )

    if torch.any(response_lengths < 1):
        raise ValueError("response length must be at least 1")
    if torch.any(response_lengths > max_length):
        raise ValueError("response length exceeds padded length")

    last_indices = response_lengths - 1

    if terminated is None:
        # 没有显式终止标记时，根据最后一个有效 Token 是否为 EOS 推断。
        last_tokens = responses.gather(
            dim=1,
            index=last_indices.unsqueeze(1),
        ).squeeze(1)
        terminated = last_tokens.eq(eos_token_id)
    else:
        terminated = terminated.bool()

    response_mask = mask_from_lengths(
        response_lengths,
        max_length,
    )

    # Terminal Reward 放到每条 Response 的最后一个有效 Token。
    token_rewards = torch.zeros(
        batch_size,
        max_length,
        device=responses.device,
        dtype=sequence_rewards.dtype,
    )
    token_rewards.scatter_(
        dim=1,
        index=last_indices.unsqueeze(1),
        src=sequence_rewards.unsqueeze(1),
    )

    # done=True 表示该 Token 后 Episode 终止。
    dones = torch.zeros(
        batch_size,
        max_length,
        device=responses.device,
        dtype=torch.bool,
    )
    dones.scatter_(
        dim=1,
        index=last_indices.unsqueeze(1),
        src=terminated.unsqueeze(1),
    )

    return {
        "response_mask": response_mask,
        "token_rewards": token_rewards,
        "dones": dones,
        "terminated": terminated,
        "response_lengths": response_lengths,
    }

若 pad_token_id == eos_token_id，不能简单使用 responses != pad_token_id 构造 Mask，否则真正的 EOS 也会被删除。最稳妥的方式是直接使用 Rollout Engine 返回的生成长度；没有长度时，再根据第一个 EOS 推断。

还要区分 terminated 与 truncated：生成 EOS 表示 Episode 真正结束，应在最后有效位置设置 done=True；达到最大长度但没有 EOS 只是截断，应保持 done=False，并在 GAE 中使用截断位置之后的 last_values 做 Bootstrap。

得到 response_mask 后，所有 Token-level 统计都应使用 Masked Reduction：

1
2
3
def masked_mean(x: torch.Tensor, mask: torch.Tensor) -> torch.Tensor:
    mask = mask.to(x.dtype)
    return (x * mask).sum() / mask.sum().clamp_min(1.0)

18. 面试前的检查清单

写完代码后，至少检查：

  1. Attention 的 Score Shape 是否为 (B,H,Tq,Tk)(B,H,T_q,T_k)；
  2. Causal Mask 和 Padding Mask 是否在 Softmax 前生效；
  3. Causal LM 的 Logits 与 Labels 是否正确 Shift；
  4. Token ID 是否通过 gather 从 Vocabulary Dimension 取概率；
  5. Padding、Prompt 和 EOS 是否按题意进入 Loss；
  6. Old Log Probability、Advantage、Return、Reference Log Probability 是否停止梯度；
  7. 所有除法分母是否处理空 Mask 或零方差；
  8. BF16/FP16 下归一化、Softmax 累积是否需要转 FP32；
  9. .transpose() 后使用 .view() 前是否调用 .contiguous()；
  10. 返回的是要最大化的 Objective，还是供 Optimizer 最小化的 Loss。

真正的面试重点不是逐字背代码，而是能够根据公式恢复实现，并在写代码的同时说清楚每个 Tensor 的 Shape 和 Mask 语义。

AI
#AI #面试 #PyTorch
AI面试常见代码题
https://cny123222.github.io/2026/10/11/AI面试常见代码题/
Author
Nuoyan Chen
Posted on
October 11, 2026
Licensed under