CS336笔记2-Architectures,Hyperparameters

timeline：

11月4日开始学，

Architectures

Postnorm vs Prenorm

为什么layernorm放在前面更加有效？

更多的一种解释是："前置归一化是一个更加稳定的训练架构"。不容易出现梯度尖刺的情况，更加稳定。在残差流中放置layernorm是不好的

++Q: 为什么在残差流中加入layernorm不好？++

a: 一个直观的观点是，残差给你从网络顶部到底部的这种恒等链接，因此，如果你视图训练非常深的网络，这使得梯度传播非常容易，因此，有很多关于lstm和其他状态空间模型在反向传播梯度时困难的讨论，但这种恒等链接没有任何这样的问题，因此，在中间放置layernorm可能会干扰这种梯度行为。

LayerNorm vs RMSNorm

可以简单的认为layernorm就是一个标准差，然后通过一个gamma因子γ放大。

++Q：为什么当前的模型都转向于使用rmsnorm？++

A：因为使用RMSnorm和使用layernorm效果一样好，现代模型（尤其是 LLaMA 系列）偏向于使用 RMSNorm 的最主要、最直接的原因 是：为了提高计算效率（即速度）。

①不需要再减去均值，不必添加bias，需要从内存加载回计算单元的参数就会变少。

②虽然RMSnorm优化的浮点计算量仅在transformer中占比0.17%，但这并不是唯一需要考虑的。因为计算量flops≠运行时间runtime。还需要仔细的考虑内存移动。

下图可以看到，虽然归一化（如layernorm、softmax等）所占的计算量仅有0.17%，但是实际运行时间占比达到了25%。一个核心原因是:归一化操作仍然会产生大量的内存移动开销，因此优化这些底层的操作十分重要。

右边这张图有点意思，不仅仅是一个架构图（它展示了 MHA 子层），更是一个性能瓶颈分析图。

核心概念：两大瓶颈

GPU 只有两种工作状态，这由计算强度 (AI) 决定：

计算受限 (Compute-Bound)

高 AI (如 153)

含义： 每从显存 (HBM) 中读取 1 字节的数据，GPU 都能执行大量的计算（例如 153 次 FLOPs）。

状态： GPU 的计算核心（Tensor Cores）100% 繁忙 ，而显存正在"休息"。这是理想状态，能实现高 MFU。

内存带宽受限 (Memory-Bound)

低 AI (如 3.5 或 1/3)

含义： 每读取 1 字节，GPU 只能执行很少的计算。

状态： GPU 的计算核心极其空闲 （"挨饿"），大部分时间都在"等待"数据从缓慢的 HBM 显存加载到高速的 SRAM 缓存中。这导致 MFU 极低。

逐一分析图中组件：现在，我们用这个"瓶颈"视角来分析这张图：

a) MHA (Multi-Head Attention)

FLOPs (43G): 430 亿次运算。这几乎是 这个块的全部计算量 (43G vs 4M+4M+29M)。

AI (153): 极高！ 结论： MHA 是**"计算受限 (Compute-Bound)"**的。

CS336 关联： 这完全符合我们的推导。MHA 的核心是 Q@K.T 和 scores@V 这样的大型矩阵乘法 (Matmul)，它们具有极高的计算强度。这部分是"好"的，能跑满 MFU。

b) Dropout 和 + (残差连接)

FLOPs (4M): 仅 400 万次运算，计算量完全可以忽略不计。

AI (1/3): 极低！ (小于 1)

结论： 它们是**"内存带宽受限 (Memory-Bound)"**的。

CS336 关联： 这也符合我们的推导。Dropout 和 Add 都是逐元素 (Element-wise) 操作（图例 ○）。它们需要从 HBM 读取整个 [B, L, D] 张量，只做 1 次乘法或加法，再写回 HBM。这是纯粹的 I/O 瓶颈。

c) LayerNorm

FLOPs (29M): 计算量也基本可以忽略不计。

AI (3.5): 非常低！

结论： LayerNorm 是**"内存带宽受限 (Memory-Bound)"**的。

这张图用数据可视化了 CS336 中一个的核心性能问题：

一个 Transformer 块 99% 以上的**计算量（FLOPs）**都集中在 MHA 和 FFN（图中未显示）中。

但是，MHA 和 FFN 并不是唯一的性能瓶颈。

++诸如 LayerNorm, Dropout, Residual Add 这样的"辅助"操作，虽然计算量（FLOPs）小到可以忽略，但它们是**严重的"内存带宽受限"**操作。++

总的执行时间 = Time(Compute-Bound) + Time(Memory-Bound)。

如果我们的 MFU 很低，很可能不是因为 MHA (43G FLOPs) 运行得慢，而是因为 GPU 的计算核心在空转，等待 LayerNorm (29M FLOPs) 慢悠悠地从 HBM 读写数据。

这精确地解释了为什么我们要痴迷于Kernel Fusion（内核融合） ，以及为什么 RMSNorm（比 LayerNorm 更快）会成为 LLaMA 的首选。

使用rmsnorm之后，step/s增加，final_loss也得到了优化。

去掉bias项

经验上，去掉这些bias通常会稳定这些llm的训练。

layernorm总结

Activations

Gelu可以看到在x=0的附近有一个平滑的凹陷，具有可微性。φx是高斯累积分布函数 CDF。

平滑性： 它的曲线处处可导，这被认为能产生更平滑的损失"地形"（Loss Landscape），使优化器（Adam）更容易找到好的解，从而提升训练稳定性和最终性能。
非单调： 它在负值区域有轻微的"下凹"，允许（并惩罚）负值信息通过，这被认为比 ReLU 的"一刀切"更具表现力。
缺点： 计算上比 ReLU 昂贵得多（需要计算高斯 CDF）。

上面的幻灯片展示了（Reglu=relu+glu）的计算流程。

标准 FFN (Standard FFN): FF(x) = max(0, xW_1) W_2

流程： x 经过一个线性变换 W_1，通过 ReLU 激活，再通过第二个 线性变换 W_2。

门控 FFN (Gated FFN): FF_ReGLU(x) = (max(0, xW_1) ⊗ xV) W_2

流程：

x 同时经过两个并行的线性变换：W_1 和 V。

xW_1 的结果通过 ReLU 激活。

xV 的结果不经过激活（它是一个纯粹的线性项）。

"门控 (Gating)" ：将前两步的结果进行逐元素乘法 (⊗) 。这个乘法就是"门 (Gate) "。它允许网络动态地、有选择地 控制哪些信息可以流向下一层（类似于 LSTM 中的门控），这被证明比 GeLU 的静态激活要强大得多。

这个被"门控"后的结果，最后通过第三个 线性变换 W_2。

GeGLU就是在ReGLU的基础上，将Relu换成Gelu。SwiGLU就是换成swish+glu。

尽管 SwiGLU FFN（LLaMA 架构）与 GeLU FFN（GPT-3 架构）在++总参数量 P 和总 FLOPs 上被设计为完全相同++ ，但 SwiGLU 的实际训练速度更慢 。为什么 FLOPs 相同，但速度（time_per_step）更慢？ 因为 FLOPs（理论计算量）不等于速度（实际硬件效率 MFU） 。SwiGLU 的硬件效率更低，主要源于三个开销：

内核启动开销 (Kernel Launch Overhead):

SwiGLU (3 Matmul + 2 逐元素操作) 需要启动 5 个 GPU 内核。

GeLU (2 Matmul + 1 逐元素操作) 只需要启动 3 个。

更多的内核启动 = 更多的固定开销 = 更慢的速度。

内存带宽瓶颈 (Memory-Bound):

SwiGLU 增加了一个额外的逐元素乘法 (...) ⊗ (...)。这是一个纯粹的内存带宽受限操作。GPU 计算核心在执行此操作时"空转"等待数据，拉低了 MFU。

计算强度 (AI) 效率:

GPU Tensor Cores 执行 2 次"胖"的 Matmul (如 [D, 4D]) 的效率（MFU）高于执行 3 次 "瘦"的 Matmul (如 [D, 2.67D])。SwiGLU 的"瘦" Matmul 导致硬件利用率下降。

最终的权衡 (Trade-off)： LLaMA 的设计者故意做出了这个权衡：他们接受了较慢的训练速度（time_per_step 稍高） ，以换取在相同参数预算（P）下更优异的模型质量（更低的 Loss）。

总结：

Serial vs Parallel layers

标准串行 (Serialized) 架构 (GPT-2 风格):

y = x + Attention(LayerNorm(x)) (MHA 块)

z = y + MLP(LayerNorm(y)) (FFN 块)

数据流： x 必须先完成 MHA 块的所有计算（Attention），得到 y 之后，才能开始 FFN 块的计算（MLP）。

瓶颈： MLP 必须等待 Attention，两者是串行的。

并行 (Parallel) 架构 (GPT-J / Cohere 风格):

y = x + MLP(LayerNorm(x)) + Attention(LayerNorm(x))

数据流：

对同一个 输入 x 只进行一次 LayerNorm。（这是"if implemented right, LayerNorm can be shared"的含义）。

LN(x) 的结果被同时喂给 Attention 块和 MLP 块。

Attention(...) 和 MLP(...) 可以并行计算。

最后，将这两个块的输出与原始残差 x 一次性相加。

Position embeddings

ROPE的高层思想：重要的是这些向量的相对位置。f(x,i)中的x是我要embedding的词，i是我的位置，你只关注词x和词y的距离。

"RoPE 的设计哲学"：它首先定义了一个"完美"的相对位置编码应该满足的 "黄金标准"，然后论证了为什么之前的所有方法（Absolute, Sine, T5-style）都不满足这个标准。

这张幻灯片接着"批判"了为什么所有老方法都不满足这个"黄金标准"：

++思考：为什么这边一定要是inner product的形式？++

严格来说，它不必非要满足（T5 风格的相对偏置就不满足，也取得了成功）。

但是，RoPE 的作者（以及 LLaMA, PaLM 的作者）之所以"执着"于将其保留为纯粹的内积（Inner Product / Dot Product）形式，是出于一个极其重要且具有前瞻性的 架构设计考量：为了保持注意力计算的"纯粹性"，以实现极致的"硬件和算法兼容性"。

我们来对比一下"T5 风格"和"RoPE 风格"的优劣：

"T5 风格"的缺陷（score = QK^T + bias）

形式： Attention(Q, K) = (Q @ K.T) + \text{relative_bias}

问题： 这不再是 一个"纯粹的"点积。它是一个点积（Matmul）再加上 一个逐元素相加 （Element-wise Add）的操作。这在 CS336 的实践中会带来两个严重的性能问题：

破坏了硬件优化（如 FlashAttention）：

像 FlashAttention 这样的 SOTA（State-of-the-Art）优化技术，是专门为 softmax(Q @ K.T) 这个"端到端"的操作设计的"融合内核"（Fused Kernel）。它可以利用 GPU 的 SRAM 一次性完成计算，而无需将巨大的 (Q @ K.T) 矩阵（形状为 [B, H, L, L]）写回到缓慢的 HBM（显存）中。

如果您使用 T5 风格，FlashAttention 就无法使用了。您必须：

计算 score = Q @ K.T (一次内核调用，结果写回 HBM)。

再启动一个全新的内核 ，从 HBM 读回 score，执行 score = score + bias (一次内存带宽受限的内核调用)。

这个过程极其缓慢 ，彻底抵消了 FlashAttention 带来的所有速度优势。

破坏了算法兼容性（如线性注意力）：

许多高效的 Transformer 变体（如 Linear Transformers, Performers）依赖于**核方法（Kernel Methods）**来重排（re-arrange）计算顺序，例如将 softmax(Q K^T) V 替换为 Q (K^T V)，从而将计算复杂度从 O(L^2) 降到 O(L)。

+ bias 这一项的加入，从数学上彻底破坏了这种重排的可能性。

这意味着 T5 风格的位置编码无法与这些 $O(L)$ 的高效注意力算法兼容。

RoPE 的优越性 (score = Q'K'^T)

形式：

Q' = f(Q, i) (在 MHA 之外，提前旋转 Q)

K' = f(K, j) (在 MHA 之外，提前旋转 K)

Attention(Q', K') = Q' @ K'.T

好处：

注意力公式保持了"纯粹的点积"形式！

FlashAttention 完全不在乎 你喂给它的是 Q 还是 Q'。它看到 Q' @ K'.T，就会立刻应用其融合内核进行超高速计算。

所有的"线性注意力"算法也完全兼容 。它们只需要对 Q' 和 K' 应用核函数即可。

总结：

RoPE 的作者坚持"必须满足内积"形式，是一个极其聪明的工程决策。

它确保了 RoPE 只是一个"即插即用"（Drop-in）的模块，它不会以任何方式"污染"或"破坏"下游的注意力计算公式。这使得 RoPE 完全兼容 FlashAttention 和各种线性注意力等所有旨在优化标准点积的 SOTA 技术，从而实现了极致的性能和未来的可扩展性。

内积只会关注距离的差异，通常这种旋转在2d是非常明显的，但在高纬中不太明显。rope的做法是将高纬度D切分成二维的块，每两个维度将被某个theta旋转，因此会有一个旋转速度，我们将旋转这些维度对。

此时，每对维度都在编码所有这些的相对位置，就像在sin和cos 的embedding ，使得一些嵌入旋转得很快，而其他的嵌入旋转的更慢。因此，他们可以捕获高/低频信息or近距离、远距离的信息。

可以将其视为嵌入向量与这些2*2块矩阵相乘的操作。这里面不涉及加法or交叉项，完美符合上面提出的原则->一切都是相对的。

注意rope将在实际的注意力层进行操作，而不是在底部添加位置编码。只有当生成Query和Key的时候，需要这么计算拿到相对的位置信息。

Hyperparameters

feedforward size

通常dff升维的维度是d_model的4倍（经验上）。下面是一些变体：

这个是前面讲到的，glu为了保证总体的模型参数量不变，通常dff是d_model的8/3倍。

head-dim * num-heads 和model-dim的比例关系

大部分head-dim * num-heads ≈ model-dim。

bhoj等人提出，如果在ratio=1的情况下，num-heads越多，head-dim越小，维度小对应矩阵的rank越小，每个头的注意力表达能力有限--->但实际并没有太大的影响。

Aspect ratios 宽深比

通常来说，深度代表着模型更加聪明，宽度代表着计算的高效。Aspect ratios可以控制我们计算的并行化程度。

差不多在100左右区间，都能够达到不错的效果。现有的一些实验表明，aspect ratio大概会对下游任务会有较大的影响。

vocabulary size

单语言大概在四五w的量级。多语言大模型的词表大概在十几万到二十几万这样。

Dropout and other regularization

似乎预训练用了大量的数据，而且基本只跑一个epoch，根本不会出现过拟合，这样是不是就不需要dropout？

但事实上，dropout是之前常会用到的手段，近年来dropout使用场景变少，更加趋向于使用weight decay。但weight_decay实际使用的目的不是控制模型过拟合（图1可以看到，不同weight decay比例并不影响最终trainloss和val loss的比例）。

图2和图3可以看到weight_decay会与优化器的学习率等产生非常复杂的作用（有decay的时候，当learning rate骤降的时候，模型的training loss也会骤降），实验结果可以看到设置weight decay的目的是为了得到更好的training_loss。

非常反直觉的：通常来说加入weight_decay之后，模型的val_loss会降低，但实际上，trainingloss降低且和val_loss保持一致（图1）

Stability tricks

z-loss（控制softmax的输出）

QK-norm（控制softmax的输入）

layernorm不断加加加加加到厌倦~

Logit soft-capping

总结： Soft-capping 是 "稳定性工具箱"中的一个（过时的）技巧。它虽然能有效防止 NaN 崩溃，但它既拖慢了训练速度 ，又损害了模型最终的质量 ，是一种"下下策"。它不如 Z-Loss（质量中性）或 QKV_norm（质量提升）等更现代的技巧。通常qk-norm的效果是最好的。

Attention heads

GQA/MQA -- Reducing attention head cost

（可以把之前整理的gqa和mqa内容结合起来整理）