LLM-leetcode TASK02

@所有人今日竞答题目：

Q1

Q1（选择题）x 形状是 $B, C, H, W$ = $2, 3, 4, 4$ ，目标输出 $B, H\*W, C$ 。下面哪个 einops 表达式语义正确且与 x.permute(0, 2, 3, 1).reshape(2, 16, 3) 完全等价？

A. rearrange(x, 'b c h w -> b (h w) c')
B. rearrange(x, 'b c h w -> b h w c')
C. rearrange(x, 'b c h w -> (b h w) c')
D. rearrange(x, 'b h w c -> b (h w) c')

rearrange(x, 'b c h w -> b (h w) c')

原始张量：

x.shape = [B, C, H, W] = [2, 3, 4, 4]

目标是：

[B, H*W, C] = [2, 16, 3]

表达式的执行过程可以理解为：

python 复制代码

x.permute(0, 2, 3, 1)   # [B,C,H,W] -> [B,H,W,C]
.reshape(B, H * W, C)    # [B,H,W,C] -> [B,H*W,C]

括号 () 在 einops 里表示什么？

括号表示将多个维度进行组合或拆分。

例如：

'b c h w -> b (h w) c'

其中：

(h w)

表示将 h 和 w 两个维度合并成一个维度，大小为：

h × w

所以：

2,3,4,4\] -\> \[2,16,3

反过来，如果已知 h 和 w，也可以拆分：

rearrange(x, 'b (h w) c -> b c h w', h=4, w=4)

**Q2（选择题）**下面哪个写法与 nn.Embedding(V, D)(input_ids) 输出一致？（V=词表大小，D=隐藏维度，input_ids 形状 [B, T]）

A. torch.eye(V)[input_ids] @ weight （其中 weight: $V, D$ ）
- torch.eye(V)：生成 V×V 单位矩阵
- torch.eye(V) $input_ids$ ，用input_ids索引单位矩阵，得到 one-hot 编码：
- @ weight 矩阵乘法 [N, L, V] @ [V, D] = [N, L, D]
- 独热向量 × 权重矩阵：只有 id 对应位置为 1，其余 0，乘法等价取出对应行
  用单位矩阵 + 矩阵乘法模拟查表，完全等价 Embedding 逻辑，但性能极差：
  会凭空构造巨大的[N,L,V]中间张量，V 动辄几万，显存爆炸，工程不会用。
B. weight $input_ids$

PyTorch 高级索引机制：二维矩阵 $V,D$ 用整型张量 $N,L$ 索引，直接批量取出对应行。
- 输入：weight [V,D]，input_ids [N,L]
- 输出： $N, L, D$
- 对每个位置n,l，直接取 weight[ input_ids[n,l] ]，无中间大张量，速度最快、显存占用最低。
- nn.Embedding(V, D)(input_ids) 底层就是这种索引逻辑
C. F.one_hot(input_ids, V) @ weight
- F.one_hot(input_ids, num_classes=V)
  把整型 id 转独热编码张量：input_ids $N,L$ → one-hot $N, L, V$ ，对应 id 位置为 1，其余 0
- @ weight 矩阵乘：[N,L,V] @ [V,D] = [N,L,D]，和 A 数学逻辑完全一致
D. B 和 C 都行（输出 shape 相同，数值上仅在极少边界条件下才可能与官方略有差异）

提示：思考"查表 vs One-hot 乘法"的内存效率差距，以及 B/C 在数值上是否真的等价。

为什么实际使用查表而不是 One-hot？

查表：weight $input_ids$ ,只需要根据索引取出对应行。
One-hot 方案需要显式构造： $B,T,V$ 的巨大张量。当词表大小 V 很大时，会浪费大量内存和计算量。

但在实际 PyTorch 代码中，B 是最标准、最推荐的写法：weight[input_ids]

设：

python 复制代码

weight.shape = [V,D]
input_ids.shape = [B,T]

那么：

python 复制代码

weight[input_ids]

输出：

python 复制代码

[B,T,D]

这与：

python 复制代码

nn.Embedding(V, D)(input_ids)

执行的核心查表操作一致。

C 为什么在数学上等价？

F.one_hot(input_ids, V) @ weight

首先：

python 复制代码

F.one_hot(input_ids, V).shape = [B,T,V]
weight.shape = [V,D]

矩阵乘法后：

复制代码

[B,T,V] @ [V,D] -> [B,T,D]

每个 one-hot 向量只有目标词编号处为 1，其余位置为 0。例如：

python 复制代码

input_id = 2
one_hot = [0,0,1,0,...]

于是：

python 复制代码

one_hot @ weight = weight[2]

因此，one-hot 乘权重矩阵在数学上等价于查表。

实际 PyTorch 中需要注意类型转换：

python 复制代码

F.one_hot(input_ids, V).to(weight.dtype) @ weight

因为 F.one_hot 默认返回整数类型，而神经网络权重通常是浮点类型。

**Q3（开放题）**我们要实现 y = relu(x @ W^T + b) 的反向传播。设：

x: $B, d_in$
W: $d_out, d_in$
b: $d_out$
y: $B, d_out$
grad_output: $B, d_out$ （从上一层传回来的梯度，形状同 y）
反传分两步：先过 ReLU 回到线性前的梯度 grad_z，再过 Linear 分别求对 x / W / b 的梯度。
关键背景：mask 是什么？

为什么 ReLU 反向传播需要 mask？

ReLU 定义为：

python 复制代码

ReLU(z) = max(0,z)

它的导数为：

python 复制代码

dReLU/dz = 1，z > 0
dReLU/dz = 0，z <= 0

因此前向传播时保存：

python 复制代码

mask = (z > 0).float()

其中：

python 复制代码

z > 0 的位置：mask = 1
z <= 0 的位置：mask = 0

反向传播经过 ReLU：

python 复制代码

grad_z = grad_output * mask

也就是：

python 复制代码

z > 0：上游梯度原样通过
z <= 0：上游梯度被截断为 0

例如：

python 复制代码

z           = [-2, 3, 0, 5]
mask        = [ 0, 1, 0, 1]
grad_output = [ 4, 6, 8, 2]

grad_z      = [ 0, 6, 0, 2]

如果省略 mask，直接令：

python 复制代码

grad_z = grad_output

那么 z <= 0 的位置也会继续传递梯度，相当于把 ReLU 的导数错误地当成全部为 1。

此时反向传播不再对应真实的前向函数，ReLU 会在反向传播中近似变成恒等函数，导致 x、W、b 都得到错误梯度。

在 PyTorch 实现里，前向时除了算出 y，我们还会做一步：

z = x @ W.T + b # 线性部分

y = relu(z) # 激活

mask = (z > 0).float() # ← 关键

ctx.save_for_backward(x, weight, mask) # 把 mask 留作反传用

提示： mask 是一个和 z 同形状的 0/1 张量：z > 0 的位置为 1，其余为 0。反传过 ReLU 时，上游梯度在 z <= 0 的位置应当被截断为 0，在 z > 0 的位置原样通过。这正好等价于 grad_output * mask。

为什么不是 W @ grad_z？

因为：

W.shape = $d_out,d_in$

grad_z.shape = $B,d_out$

直接计算：

d_out,d_in\] @ \[B,d_out

要求中间维度 d_in 与 B 相同，通常并不成立，所以矩阵维度不匹配。

另外，grad_x 必须为每个 batch 样本保留一行，因此 batch 维 B 应该位于输出第一维：

B,d_out\] @ \[d_out,d_in\] -\> \[B,d_in

正确公式为：

grad_x = grad_z @ W

请回答：

理解 mask：为什么 ReLU 的反传需要 mask？如果省掉这一步直接用 grad_output，会出现什么现象？
求 grad_x：形状是什么？计算公式怎么写？为什么矩阵相乘的方向是 grad_z @ W 而不是 W @ grad_z？
求 grad_W：你写的是 grad_z.T @ x 还是 x.T @ grad_z？两个写法数值上等价吗？哪一边的 shape 直接对得上 W 的 $d_out, d_in$ ？请用维度对应关系解释。
求 grad_b：为什么需要 sum(dim=0) 而不是直接用 grad_z？从 b 在前向中"被广播到每个样本"这个事实出发推导。

TASK02

SwiGLU Activation | 激活函数与门控机制 (SwiGLU Activation)

在组装 LLaMA-3 的那一节中，我们使用了 SwiGLU 作为 MLP 的激活函数。为什么所有主流大模型（LLaMA, Qwen, Mistral, PaLM）都在抛弃 ReLU/GELU 而转向 SwiGLU？

本节我们将深入推导 SwiGLU 的设计原理，特别是如何调整隐藏层的维度 ，以保证参数量与标准 Transformer 严格对齐。这是面试中非常经典的架构推导题。

相关阅读 :

本节使用纯 PyTorch 实现了算法逻辑与数学推导。

如果你想学习工业界如何打破该算子的 Memory Bound (访存瓶颈)，请前往 Triton 篇：

../03_CUDA_and_Triton_Kernels/02_Triton_Fused_SwiGLU.ipynb

什么是 GLU (Gated Linear Unit)？

传统 MLP 是 W d o w n ( σ ( W u p x ) ) W_{down}(\sigma(W_{up}x)) Wdown(σ(Wupx))。

门控机制 (GLU) 引入了"两条路"：一条路做激活（作为门控开关），另一条路保持线性，然后两者逐元素相乘（Hadamard Product）。

公式： GLU ( x , W 1 , W 2 ) = ( x W 1 ⊗ σ ( x W 2 ) ) W d o w n \text{GLU}(x, W_1, W_2) = (xW_1 \otimes \sigma(xW_2))W_{down} GLU(x,W1,W2)=(xW1⊗σ(xW2))Wdown。

这种机制类似于 LSTM 中的遗忘门，极大地增强了模型捕捉复杂模式的能力。

什么是 SwiGLU？

就是把 GLU 中的激活函数 σ \sigma σ 换成了 Swish （即 x ⋅ Sigmoid ( β x ) x \cdot \text{Sigmoid}(\beta x) x⋅Sigmoid(βx)，在 PyTorch 中 β = 1 \beta=1 β=1 时等于 SiLU）。

核心数学机制：参数量对齐

典型的面试问题：

"在 GPT-2 中，隐藏层维度通常是输入维度 d d d 的 4 倍（即 4 d 4d 4d）。但在使用 SwiGLU 的 LLaMA 中，为什么隐藏层维度变成了 8 3 d \frac{8}{3}d 38d 并向上取整？"

推导过程：

标准 MLP 参数量 ：

输入为 d d d，隐藏层为 h h h。

有两个投影矩阵（升维 d → h d \to h d→h，降维 h → d h \to d h→d）。

总参数量 = 2 ⋅ ( d × h ) 2 \cdot (d \times h) 2⋅(d×h)。

当 h = 4 d h = 4d h=4d 时，总参数量 = 2 ⋅ 4 d 2 = 8 d 2 2 \cdot 4d^2 = \mathbf{8d^2} 2⋅4d2=8d2。
SwiGLU MLP 参数量 ：

输入为 d d d，隐藏层为 h s w i g l u h_{swiglu} hswiglu。

因为有门控机制 ，升维阶段需要两个投影矩阵（ W g a t e W_{gate} Wgate 和 W u p W_{up} Wup，均是 d → h s w i g l u d \to h_{swiglu} d→hswiglu）。

降维阶段需要一个矩阵（ W d o w n W_{down} Wdown，是 h s w i g l u → d h_{swiglu} \to d hswiglu→d）。

总参数量 = 3 ⋅ ( d × h s w i g l u ) 3 \cdot (d \times h_{swiglu}) 3⋅(d×hswiglu)。
对齐参数量 ：

为了使得 SwiGLU 的计算开销（参数量）与原始模型完全相同：

3 ⋅ d ⋅ h s w i g l u = 8 d 2 3 \cdot d \cdot h_{swiglu} = 8d^2 3⋅d⋅hswiglu=8d2

解得： h s w i g l u = 8 3 d h_{swiglu} = \mathbf{\frac{8}{3}d} hswiglu=38d

这正是 LLaMA 源码中对中间层维度进行 int(8 * hidden_size / 3) 计算的根本原因。

激活函数对比：FFN、GLU、Swish、SwiGLU

一、四个核心公式放在一起

1. 普通前馈网络 FFN

FFN ⁡ ( x ) = ϕ ( x W 1 + b 1 ) W 2 + b 2 \operatorname{FFN}(x) = \phi(xW_1+b_1)W_2+b_2 FFN(x)=ϕ(xW1+b1)W2+b2

例如使用 ReLU：

FFN ⁡ ( x ) = ReLU ⁡ ( x W 1 + b 1 ) W 2 + b 2 \operatorname{FFN}(x) = \operatorname{ReLU}(xW_1+b_1)W_2+b_2 FFN(x)=ReLU(xW1+b1)W2+b2

结构是：

x → 线性变换 → 激活函数 → 线性变换 x \rightarrow \text{线性变换} \rightarrow \text{激活函数} \rightarrow \text{线性变换} x→线性变换→激活函数→线性变换

它只有一条主要分支。

2. GLU：门控线性单元

GLU 将输入分别投影到两个分支：

a = x W a + b a a = xW_a + b_a a=xWa+ba

g = x W g + b g g = xW_g + b_g g=xWg+bg

然后：

GLU ⁡ ( x ) = a ⊙ σ ( g ) \operatorname{GLU}(x) = a \odot \sigma(g) GLU(x)=a⊙σ(g)

合并写成：

GLU ⁡ ( x ) = ( x W a + b a ) ⊙ σ ( x W g + b g ) \operatorname{GLU}(x) = (xW_a+b_a) \odot \sigma(xW_g+b_g) GLU(x)=(xWa+ba)⊙σ(xWg+bg)

其中 σ ( z ) = 1 1 + e − z \sigma(z) = \frac{1}{1+e^{-z}} σ(z)=1+e−z1。

含义是：输出 = 内容分支 × \times × 门控分支，即：

x W a + b a ⏟ 内容 ⊙ σ ( x W g + b g ) ⏟ 控制内容通过多少 \underbrace{xW_a+b_a}{\text{内容}} \odot \underbrace{\sigma(xW_g+b_g)}{\text{控制内容通过多少}} 内容 xWa+ba⊙控制内容通过多少 σ(xWg+bg)

3. Swish / SiLU 激活函数

Swish 的定义是：

Swish ⁡ β ( z ) = z σ ( β z ) \operatorname{Swish}_\beta(z) = z\sigma(\beta z) Swishβ(z)=zσ(βz)

当 β = 1 \beta=1 β=1 时：

Swish ⁡ ( z ) = z σ ( z ) \operatorname{Swish}(z) = z\sigma(z) Swish(z)=zσ(z)

又因为 σ ( z ) = 1 1 + e − z \sigma(z) = \frac{1}{1+e^{-z}} σ(z)=1+e−z1，所以：

Swish ⁡ ( z ) = z 1 + e − z \operatorname{Swish}(z) = \frac{z}{1+e^{-z}} Swish(z)=1+e−zz

PyTorch 中：

SiLU ⁡ ( z ) = z σ ( z ) \operatorname{SiLU}(z) = z\sigma(z) SiLU(z)=zσ(z)

因此，当 β = 1 \beta=1 β=1 时： Swish ⁡ ( z ) = SiLU ⁡ ( z ) \operatorname{Swish}(z) = \operatorname{SiLU}(z) Swish(z)=SiLU(z)。

4. SwiGLU

SwiGLU 将 GLU 中的 Sigmoid 门控 σ ( x W g + b g ) \sigma(xW_g+b_g) σ(xWg+bg) 换成 Swish/SiLU 门控 Swish ⁡ ( x W g + b g ) \operatorname{Swish}(xW_g+b_g) Swish(xWg+bg)：

SwiGLU ⁡ ( x ) = ( x W a + b a ) ⊙ Swish ⁡ ( x W g + b g ) \operatorname{SwiGLU}(x) = (xW_a+b_a) \odot \operatorname{Swish}(xW_g+b_g) SwiGLU(x)=(xWa+ba)⊙Swish(xWg+bg)

因为 Swish ⁡ ( z ) = z σ ( z ) \operatorname{Swish}(z)=z\sigma(z) Swish(z)=zσ(z)，所以展开后：

SwiGLU ⁡ ( x ) = ( x W a + b a ) ⊙ $( x W g + b g ) ⊙ σ ( x W g + b g )$ \operatorname{SwiGLU}(x) = (xW_a+b_a) \odot \left $(xW_g+b_g) \\odot \\sigma(xW_g+b_g)\\right$ SwiGLU(x)=(xWa+ba)⊙ $(xWg+bg)⊙σ(xWg+bg)$

也就是：

SwiGLU ⁡ ( x ) = ( x W a + b a ) ⏟ 内容分支 ⊙ ( x W g + b g ) ⏟ 门控信号的幅度 ⊙ σ ( x W g + b g ) ⏟ 软门控比例 \operatorname{SwiGLU}(x) = \underbrace{(xW_a+b_a)}{\text{内容分支}} \odot \underbrace{(xW_g+b_g)}{\text{门控信号的幅度}} \odot \underbrace{\sigma(xW_g+b_g)}_{\text{软门控比例}} SwiGLU(x)=内容分支 (xWa+ba)⊙门控信号的幅度 (xWg+bg)⊙软门控比例 σ(xWg+bg)

这就是 SwiGLU 比普通 GLU 多出来的部分。

二、GLU 和 SwiGLU 直接对比

把两者并排写：

GLU ⁡ ( x ) = ( x W a + b a ) ⊙ σ ( x W g + b g ) \operatorname{GLU}(x) = (xW_a+b_a) \odot \sigma(xW_g+b_g) GLU(x)=(xWa+ba)⊙σ(xWg+bg)

SwiGLU ⁡ ( x ) = ( x W a + b a ) ⊙ Swish ⁡ ( x W g + b g ) \operatorname{SwiGLU}(x) = (xW_a+b_a) \odot \operatorname{Swish}(xW_g+b_g) SwiGLU(x)=(xWa+ba)⊙Swish(xWg+bg)

再把 Swish 展开：

SwiGLU ⁡ ( x ) = ( x W a + b a ) ⊙ ( x W g + b g ) ⊙ σ ( x W g + b g ) \operatorname{SwiGLU}(x) = (xW_a+b_a) \odot (xW_g+b_g) \odot \sigma(xW_g+b_g) SwiGLU(x)=(xWa+ba)⊙(xWg+bg)⊙σ(xWg+bg)

最关键的区别是：

	门控形式
GLU 门控	σ ( g ) \sigma(g) σ(g)
SwiGLU 门控	g σ ( g ) g\sigma(g) gσ(g)

其中 g = x W g + b g g = xW_g+b_g g=xWg+bg。

所以：

GLU 的门控值主要在 0 0 0 到 1 1 1 之间；
SwiGLU 的门控值还保留了 g g g 自己的大小和正负方向。

三、一步一步计算的例子

假设输入经过两个线性层后得到 a = 3 , g = 2 a=3,\; g=2 a=3,g=2，其中 a a a 是内容分支， g g g 是门控分支。

1. GLU

首先计算：

σ ( 2 ) = 1 1 + e − 2 ≈ 0.881 \sigma(2) = \frac{1}{1+e^{-2}} \approx 0.881 σ(2)=1+e−21≈0.881

所以：

GLU ⁡ ( x ) = a σ ( g ) = 3 × 0.881 ≈ 2.643 \operatorname{GLU}(x) = a\sigma(g) = 3 \times 0.881 \approx \boxed{2.643} GLU(x)=aσ(g)=3×0.881≈2.643

2. SwiGLU

先计算 Swish：

Swish ⁡ ( 2 ) = 2 σ ( 2 ) = 2 × 0.881 = 1.762 \operatorname{Swish}(2) = 2\sigma(2) = 2 \times 0.881 = 1.762 Swish(2)=2σ(2)=2×0.881=1.762

然后乘以内容分支：

SwiGLU ⁡ ( x ) = a × Swish ⁡ ( g ) = 3 × 1.762 ≈ 5.286 \operatorname{SwiGLU}(x) = a \times \operatorname{Swish}(g) = 3 \times 1.762 \approx \boxed{5.286} SwiGLU(x)=a×Swish(g)=3×1.762≈5.286

对比：

激活函数	计算
GLU ⁡ \operatorname{GLU} GLU	3 × σ ( 2 ) ≈ 2.64 3 \times \sigma(2) \approx 2.64 3×σ(2)≈2.64
SwiGLU ⁡ \operatorname{SwiGLU} SwiGLU	3 × 2 × σ ( 2 ) ≈ 5.29 3 \times 2 \times \sigma(2) \approx 5.29 3×2×σ(2)≈5.29

SwiGLU 不仅使用了 Sigmoid 给出的"开门比例"，还保留了门控信号 g = 2 g=2 g=2 的幅度。

四、再看负数情况

假设 a = 3 , g = − 1 a=3,\; g=-1 a=3,g=−1。

GLU

σ ( − 1 ) ≈ 0.269 \sigma(-1) \approx 0.269 σ(−1)≈0.269

GLU ⁡ ( x ) = 3 × 0.269 ≈ 0.807 \operatorname{GLU}(x) = 3 \times 0.269 \approx \boxed{0.807} GLU(x)=3×0.269≈0.807

由于 Sigmoid 始终为正，GLU 只是把内容缩小。

SwiGLU

Swish ⁡ ( − 1 ) = − 1 × σ ( − 1 ) = − 1 × 0.269 = − 0.269 \operatorname{Swish}(-1) = -1 \times \sigma(-1) = -1 \times 0.269 = -0.269 Swish(−1)=−1×σ(−1)=−1×0.269=−0.269

SwiGLU ⁡ ( x ) = 3 × ( − 0.269 ) ≈ − 0.807 \operatorname{SwiGLU}(x) = 3 \times (-0.269) \approx \boxed{-0.807} SwiGLU(x)=3×(−0.269)≈−0.807

这说明 SwiGLU 不仅能控制大小，还能够保留门控信号的正负方向。

五、在 Transformer 中的完整公式

Transformer 的 SwiGLU 前馈层通常写成：

FFN ⁡ SwiGLU ( x ) = $SiLU ( x W g ) ⊙ ( x W v )$ W o \operatorname{FFN}_{\text{SwiGLU}}(x) = \left $\\operatorname{SiLU}(xW_g) \\odot (xW_v)\\right$ W_o FFNSwiGLU(x)= $SiLU(xWg)⊙(xWv)$ Wo

把它拆开：

第一步：门控分支

g = x W g g = xW_g g=xWg

g ~ = SiLU ⁡ ( g ) = g σ ( g ) \tilde g = \operatorname{SiLU}(g) = g\sigma(g) g~=SiLU(g)=gσ(g)

第二步：内容分支

v = x W v v = xW_v v=xWv

第三步：逐元素相乘

h = g ~ ⊙ v = SiLU ⁡ ( x W g ) ⊙ ( x W v ) h = \tilde g \odot v = \operatorname{SiLU}(xW_g) \odot (xW_v) h=g~⊙v=SiLU(xWg)⊙(xWv)

第四步：输出投影

y = h W o y = hW_o y=hWo

最终：

y = $SiLU ( x W g ) ⊙ ( x W v )$ W o y = \left $\\operatorname{SiLU}(xW_g) \\odot (xW_v)\\right$ W_o y= $SiLU(xWg)⊙(xWv)$ Wo

六、最简单的记忆方法

记住三个公式即可：

Sigmoid

σ ( x ) = 1 1 + e − x \sigma(x) = \frac{1}{1+e^{-x}} σ(x)=1+e−x1

Swish / SiLU

SiLU ⁡ ( x ) = x σ ( x ) \operatorname{SiLU}(x) = x\sigma(x) SiLU(x)=xσ(x)

SwiGLU

SwiGLU ⁡ ( x ) = SiLU ⁡ ( x W g ) ⊙ ( x W v ) \operatorname{SwiGLU}(x) = \operatorname{SiLU}(xW_g) \odot (xW_v) SwiGLU(x)=SiLU(xWg)⊙(xWv)

Transformer 中再加输出投影：

FFN ⁡ SwiGLU ( x ) = $SiLU ( x W g ) ⊙ ( x W v )$ W o \operatorname{FFN}_{\text{SwiGLU}}(x) = \left $\\operatorname{SiLU}(xW_g) \\odot (xW_v)\\right$ W_o FFNSwiGLU(x)= $SiLU(xWg)⊙(xWv)$ Wo

一句话理解：

SwiGLU = 经过 SiLU 的门控分支 × 线性的内容分支 \boxed{\text{SwiGLU} = \text{经过 SiLU 的门控分支} \times \text{线性的内容分支}} SwiGLU=经过 SiLU 的门控分支×线性的内容分支

Step 2: 核心数学机制：参数量对齐

典型的面试问题：

"在 GPT-2 中，隐藏层维度通常是输入维度 d d d 的 4 倍（即 4 d 4d 4d）。但在使用 SwiGLU 的 LLaMA 中，为什么隐藏层维度变成了 8 3 d \frac{8}{3}d 38d 并向上取整？"

推导过程：

标准 MLP 参数量 ：

输入为 d d d，隐藏层为 h h h。

有两个投影矩阵（升维 d → h d \to h d→h，降维 h → d h \to d h→d）。

总参数量 = 2 ⋅ ( d × h ) 2 \cdot (d \times h) 2⋅(d×h)。

当 h = 4 d h = 4d h=4d 时，总参数量 = 2 ⋅ 4 d 2 = 8 d 2 2 \cdot 4d^2 = \mathbf{8d^2} 2⋅4d2=8d2。
SwiGLU MLP 参数量 ：

输入为 d d d，隐藏层为 h s w i g l u h_{swiglu} hswiglu。

因为有门控机制 ，升维阶段需要两个投影矩阵（ W g a t e W_{gate} Wgate 和 W u p W_{up} Wup，均是 d → h s w i g l u d \to h_{swiglu} d→hswiglu）。

降维阶段需要一个矩阵（ W d o w n W_{down} Wdown，是 h s w i g l u → d h_{swiglu} \to d hswiglu→d）。

总参数量 = 3 ⋅ ( d × h s w i g l u ) 3 \cdot (d \times h_{swiglu}) 3⋅(d×hswiglu)。
对齐参数量 ：

为了使得 SwiGLU 的计算开销（参数量）与原始模型完全相同：

3 ⋅ d ⋅ h s w i g l u = 8 d 2 3 \cdot d \cdot h_{swiglu} = 8d^2 3⋅d⋅hswiglu=8d2

解得： h s w i g l u = 8 3 d h_{swiglu} = \mathbf{\frac{8}{3}d} hswiglu=38d

这正是 LLaMA 源码中对中间层维度进行 int(8 * hidden_size / 3) 计算的根本原因。

工业级实现框架与性能陷阱 (Memory Bound)

在理解了 SwiGLU 的基本公式（down_proj(SiLU(gate_proj(x)) * up_proj(x))）和 8 / 3 8/3 8/3 维度由来后，如何把它写进真实的训练框架中？

性能陷阱 1：张量并行 (TP) 与内存对齐

在真实的 LLaMA 源码中，除了按 8 / 3 8/3 8/3 计算出隐藏层维度，还需要将其向上取整对齐到一个 multiple_of（通常是 256）的倍数。

这不仅是为了让单卡 Tensor Core（通常要求 8-byte 或 32-byte 对齐）跑得更快，更是因为大模型训练会使用张量并行 (Tensor Parallelism) 。如果隐藏层维度不能被 GPU 数量整除（例如 T P = 8 TP=8 TP=8 时，256 的倍数分给 8 张卡，每张卡至少能分到 32 维），在切分权重矩阵时就会发生严重的报错。

性能陷阱 2：访存瓶颈 (Memory Bound) 与矩阵融合

在最朴素的代码实现中，开发者会分别定义并执行 gate_proj(x) 和 up_proj(x)。

由于这两个线性层共享完全相同的输入张量 x x x ，分开计算会导致巨大的输入 x x x 被 GPU 从全局显存 (HBM) 中读取两次。

工业界解法 (Matrix Fusion) ：

在 vLLM、Megatron 等主流框架中，标准的做法是将 W g a t e W_{gate} Wgate 和 W u p W_{up} Wup 这两个形状为 [hidden_size, intermediate_size] 的权重矩阵，在初始化时拼接成一个巨大的融合矩阵 gate_up_proj ，其形状为 [hidden_size, 2 * intermediate_size]。
在前向传播时，输入 x x x 只需要被读取一次进行一次大规模矩阵乘法。得到的结果再通过 torch.chunk(2, dim=-1) 切分为两半，分别作为 gate 和 up 块。这极大地缓解了内存带宽瓶颈。

python 复制代码

def calculate_intermediate_size(hidden_size: int, multiple_of: int = 256):
    """
    计算 LLaMA 风格的 SwiGLU 隐藏层维度
    
    规则：
    1. 取 hidden_size 的 8/3
    2. 为了硬件对齐（如 Tensor Core），通常要求是 multiple_of 的倍数。
       因此将结果除以 multiple_of，向上取整后再乘以 multiple_of。
    """
    # ==========================================
    # TODO 1: 计算理论隐藏层大小 (8/3 * hidden_size)
    # 提示: 注意使用整数除法
    # ==========================================
    # intermediate_size = ???
    intermediate_size = int(hidden_size * 8 / 3)
      
    # ==========================================
    # TODO 2: 向 multiple_of 对齐 (向上取整)
    # 提示: 思考如何利用整除的特性实现向上取整
    # ==========================================
    # aligned_size = ???
    aligned_size = ((intermediate_size + multiple_of -1 )//multiple_of)* multiple_of
    
    return aligned_size

class SwiGLU_MLP(nn.Module):
    def __init__(self, hidden_size: int, intermediate_size: int):
        super().__init__()
        # ==========================================
        # TODO 3: 定义工业级 SwiGLU 的投影矩阵
        # ==========================================
        # self.gate_up_proj = ???
        # self.down_proj = ???
        self.gate_up_proj = nn.Linear(hidden_size, 2* intermediate_size,bias=False)
        self.down_proj = nn.Linear(intermediate_size,hidden_size,bias=False)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        # ==========================================
        # TODO 4: 组装工业级 SwiGLU 前向传播
        # ==========================================
        gate_up = self.gate_up_proj(x)
        gate, up = torch.chunk(gate_up,2,dim=-1)
        return self.down_proj(F.silu(gate)*up)

RoPE | 旋转位置编码(RoPE)

本节我们将解析大模型当前最主流的位置编码方式：RoPE (Rotary Position Embedding)，并亲手用**复数形式（Complex Tensor）**实现它。这是 LLaMA, Qwen, DeepSeek 的标配！

为什么需要 RoPE？原生的 Transformer 使用绝对位置编码 （如正弦波或可学习参数），导致模型很难泛化到比训练集更长的序列。我们希望模型能在计算 Attention 时感知到 Token 之间的相对距离。

RoPE 的本质： "借用复数的旋转" 。通过将 Query 和 Key 的向量映射到复数空间 并旋转特定角度，在计算内积（Dot-product）时，结果自然就带有了相对位置信息 （m-n）

在 PyTorch 中，最高效的 RoPE 实现方式之一是利用复数乘法 。我们将最后一维切分为两半并组合成复数形式，再乘以预先计算好的复数旋转矩阵 e i m θ e^{im\theta} eimθ。完成旋转后，再使用 torch.view_as_real 恢复为实数表示。

所以RoPE的本质就是:用"旋转角度"表示绝对位置，用"两个旋转角度的差"得到相对位置。

旋转矩阵的转置=反向旋转
连续旋转=角度相加

Step 3: 核心公式与张量维度

1. 预计算旋转角 (Precompute Frequencies)

频率计算公式：

Θ = 10000 − 2 i / d \Theta = 10000^{-2i/d} Θ=10000−2i/d

其中 i i i 是维度索引， d d d 是 Head Dimension。

生成复数形式的极坐标（欧拉公式）：

e i m Θ = cos ⁡ ( m Θ ) + i sin ⁡ ( m Θ ) e^{im\Theta} = \cos(m\Theta) + i\sin(m\Theta) eimΘ=cos(mΘ)+isin(mΘ)

2. 应用旋转 (Apply Rotary Embedding)

将输入的 Query 或 Key 视为复数：

x = x real + i ⋅ x imag x = x_{\text{real}} + i \cdot x_{\text{imag}} x=xreal+i⋅ximag

利用复数乘法等价完成旋转矩阵运算：

x rotated = x × e i m Θ x_{\text{rotated}} = x \times e^{im\Theta} xrotated=x×eimΘ

python 复制代码

def precompute_freqs_cis(dim: int, end: int, theta: float = 10000.0):
    """
    计算复数指数频率张量 (cis = cos + i * sin)
    """
    # ==========================================
    # TODO 1: 用极坐标生成复数张量 (提示: torch.polar)
    # ==========================================
    # freqs = ???
    # t = ???
    # freqs_cis = ???
    freqs = 1.0/(theta ** (torch.arange(0,dim,2)[:,(dim//2)].float()/dim))
    t = torch.arange(end, device = freqs.device, dtype = torch.float32)
    freqs = torch.outer(t,freqs)             
    freqs_cis = torch.polar(torch.ones_like(freqs), freqs)                                                                                                             
    return freqs_cis   
    

def reshape_for_broadcast(freqs_cis: torch.Tensor, x: torch.Tensor):
    ndim = x.ndim
    shape = [d if i == 1 or i == ndim - 1 else 1 for i, d in enumerate(x.shape)]
    return freqs_cis.view(*shape)

def apply_rotary_emb(
    xq: torch.Tensor,
    xk: torch.Tensor,
    freqs_cis: torch.Tensor,
) -> tuple[torch.Tensor, torch.Tensor]:
    """
    将旋转位置编码应用到 Query 和 Key 上
    """
    # ==========================================
    # TODO 2: 将 xq, xk 从实数张量转为复数张量
    # 提示: 
    # ==========================================
    # xq_ = ???
    # xk_ = ???
                                                                                                                                                                     
    xq_ = torch.view_as_complex(xq.float().reshape(*xq.shape[:-1], -1, 2))
    xk_ = torch.view_as_complex(xk.float().reshape(*xk.shape[:-1], -1, 2))  
    freqs_cis = reshape_for_broadcast(freqs_cis, xq_)
    
    # ==========================================
    # TODO 3: 进行复数乘法，并转回实数张量
    # 提示: 
    # ==========================================
    # xq_out = ???
    # xk_out = ???

    xq_out = torch.view_as_real(xq_ * freqs_cis).flatten(3)
    xk_out = torch.view_as_real(xk_ * freqs_cis).flatten(3)                                                                                                                                                             
                 
    return xq_out.type_as(xq), xk_out.type_as(xk)