""" ============================================================================= 第 3 讲 · 代码 1：Transformer 的基础积木 —— 线性层、嵌入、归一化、前馈网络 ============================================================================= 对应讲义：lectures/lecture03_architectures_hyperparameters.md 第 2、3、4 节 ----------------------------------------------------------------------------- 【这个文件要解决什么问题？】 ----------------------------------------------------------------------------- 一个现代 Transformer（Llama 风格）由很少几种"积木"搭成： Linear（线性层）、Embedding（嵌入）、RMSNorm（归一化）、SwiGLU（前馈网络）、 RoPE（位置编码，02 号文件）、因果多头注意力（03 号文件）。 这个文件从零实现前四种（与 CS336 作业 1 的接口和约定一致），并用实验回答讲义里的几个"为什么"： - 为什么现代模型把 LayerNorm 换成了 RMSNorm？"FLOPs 很少"的层也值得优化吗？ - 为什么去掉了偏置（bias）项？ - GLU 类前馈网络（SwiGLU）是什么？为什么它的隐藏层宽度是 8/3 × d_model？ ----------------------------------------------------------------------------- 【整体思路】 ----------------------------------------------------------------------------- 第 1 部分：Linear 与 Embedding（不带偏置；截断正态初始化） 第 2 部分：LayerNorm vs RMSNorm（实现、对比、为什么 RMSNorm 要在 FP32 中计算） 第 3 部分："FLOPs 不等于运行时间"：归一化层的 FLOPs 占比 vs 时间占比 第 4 部分：前馈网络：ReLU、GELU、SwiGLU；参数量对齐的 8/3 规则 ----------------------------------------------------------------------------- 【运行方式】 ----------------------------------------------------------------------------- python 01_building_blocks.py 需要 PyTorch 和 einops（见 code/lecture03/README.md）。 """ import math import time import torch from einops import einsum def print_header(title): """打印一个带分隔线的标题。""" print() print("=" * 70) print(f" {title}") print("=" * 70) # ============================================================================= # 第 1 部分：Linear 与 Embedding # ============================================================================= class Linear(torch.nn.Module): """ 不带偏置的线性层：y = x W^T。 【形状】 weight : (d_out, d_in) （PyTorch 惯例：输出维在前） x : (..., d_in) → y : (..., d_out) 【初始化】（CS336 作业 1 的约定） 截断正态分布，均值 0，std = √(2 / (d_in + d_out))，截断在 ±3 std。 为什么是这个尺度：见第 2 讲第 8 节（保持信号方差稳定）。 【为什么没有偏置】 现代大模型几乎都去掉了线性层和归一化层的偏置项： - 省参数、省显存搬运（偏置很小，但每次都要读写）； - 多个团队报告去掉偏置后训练更稳定。 而效果几乎没有损失（后面的层有足够的能力补偿）。 【边界】 - device / dtype 参数让你可以直接在 GPU 上、用 BF16 创建参数。 """ def __init__(self, in_features, out_features, device=None, dtype=None): super().__init__() std = math.sqrt(2.0 / (in_features + out_features)) w = torch.empty(out_features, in_features, device=device, dtype=dtype) torch.nn.init.trunc_normal_(w, mean=0.0, std=std, a=-3 * std, b=3 * std) self.weight = torch.nn.Parameter(w) def forward(self, x): return einsum(x, self.weight, "... d_in, d_out d_in -> ... d_out") class Embedding(torch.nn.Module): """ 嵌入层：token id → 向量（按下标查表）。 【形状】 weight : (vocab_size, d_model) token_ids : (...) 整数 → 输出 : (..., d_model) 【初始化】截断正态，均值 0，std = 1，截断在 ±3（作业 1 的约定）。 【边界】token id 必须在 [0, vocab_size) 之间，否则下标越界。 """ def __init__(self, num_embeddings, embedding_dim, device=None, dtype=None): super().__init__() w = torch.empty(num_embeddings, embedding_dim, device=device, dtype=dtype) torch.nn.init.trunc_normal_(w, mean=0.0, std=1.0, a=-3.0, b=3.0) self.weight = torch.nn.Parameter(w) def forward(self, token_ids): return self.weight[token_ids] def demo_linear_embedding(): """演示 1：检查形状与初始化尺度，并与 PyTorch 官方层对比。""" print_header("演示 1：Linear 与 Embedding") torch.manual_seed(0) lin = Linear(512, 256) x = torch.randn(4, 10, 512) y = lin(x) print(f" Linear(512 → 256)：输入 {tuple(x.shape)} → 输出 {tuple(y.shape)}") print(f" 权重标准差 {lin.weight.std().item():.4f}（目标 √(2/768) = {math.sqrt(2 / 768):.4f}，" f"截断会让它略小一点）；最大绝对值 {lin.weight.abs().max().item():.4f} ≤ 3σ = {3 * math.sqrt(2 / 768):.4f}") ref = torch.nn.Linear(512, 256, bias=False) ref.weight.data.copy_(lin.weight.data) print(f" 与 nn.Linear(bias=False) 的输出最大差异：{(ref(x) - y).abs().max().item():.2e}") emb = Embedding(1000, 64) ids = torch.tensor([[1, 5, 5, 999]]) print(f" Embedding(1000, 64)：ids {tuple(ids.shape)} → {tuple(emb(ids).shape)}；" f"同一个 id（5）得到同一个向量：{torch.equal(emb(ids)[0, 1], emb(ids)[0, 2])}") # ============================================================================= # 第 2 部分：LayerNorm vs RMSNorm # ============================================================================= class LayerNorm(torch.nn.Module): """ LayerNorm（原始 Transformer、GPT-2/3 使用）： y = (x - mean(x)) / √(var(x) + ε) × γ + β 减均值、除标准差、再做可学习的缩放 γ 和平移 β。参数量 2 × d_model。 """ def __init__(self, d_model, eps=1e-5): super().__init__() self.eps = eps self.gamma = torch.nn.Parameter(torch.ones(d_model)) self.beta = torch.nn.Parameter(torch.zeros(d_model)) def forward(self, x): in_dtype = x.dtype x = x.float() mean = x.mean(dim=-1, keepdim=True) var = ((x - mean) ** 2).mean(dim=-1, keepdim=True) y = (x - mean) / torch.sqrt(var + self.eps) * self.gamma + self.beta return y.to(in_dtype) class RMSNorm(torch.nn.Module): """ RMSNorm（Llama、PaLM、T5、Qwen 等现代模型使用；CS336 作业 1 要求实现）： RMS(x) = √( (1/d) Σ x_i² + ε ) y = x / RMS(x) × g 与 LayerNorm 相比：**不减均值、没有偏置 β**。参数量 d_model（只有增益 g，初始化为 1）。 【为什么要先转成 float32】（作业 1 明确要求） 平方求和是一种"累加"。如果输入是 BF16，平方后再相加，精度损失很大（见第 2 讲 02 号文件）； 平方还可能让大的数上溢。所以先把 x 转成 FP32 计算，最后再转回原来的数据类型。 【参数】 d_model : 最后一维的大小 eps : 防止除以 0 的小数，作业约定 1e-5 【形状】 x : (..., d_model) → 同形状 """ def __init__(self, d_model, eps=1e-5, device=None, dtype=None): super().__init__() self.eps = eps self.weight = torch.nn.Parameter(torch.ones(d_model, device=device, dtype=dtype)) def forward(self, x): in_dtype = x.dtype x = x.float() rms = torch.sqrt(x.pow(2).mean(dim=-1, keepdim=True) + self.eps) return (x / rms * self.weight).to(in_dtype) def demo_norms(): """演示 2：LayerNorm 与 RMSNorm 的行为对比。""" print_header("演示 2：LayerNorm vs RMSNorm") torch.manual_seed(0) d = 8 x = torch.randn(2, d) * 3 + 5 # 均值约 5、标准差约 3 的输入 ln, rms = LayerNorm(d), RMSNorm(d) yl, yr = ln(x), rms(x) fmt = lambda t: "[" + ", ".join(f"{v:.3f}" for v in t.tolist()) + "]" # 保留 3 位小数，方便阅读 print(f" 输入（2 行）：均值 {fmt(x.mean(-1))}，均方根 {fmt(x.pow(2).mean(-1).sqrt())}") print(f" LayerNorm 输出：均值 {fmt(yl.mean(-1))}，标准差 {fmt(yl.std(-1, unbiased=False))}") print(f" RMSNorm 输出 ：均值 {fmt(yr.mean(-1))}，均方根 {fmt(yr.pow(2).mean(-1).sqrt())}") print(f" 参数量：LayerNorm {sum(p.numel() for p in ln.parameters())}，RMSNorm {sum(p.numel() for p in rms.parameters())}") ref = torch.nn.LayerNorm(d, eps=1e-5) print(f" 我们的 LayerNorm 与 nn.LayerNorm 的最大差异：{(ref(x) - yl).abs().max().item():.2e}") if hasattr(torch.nn, "RMSNorm"): ref_rms = torch.nn.RMSNorm(d, eps=1e-5) print(f" 我们的 RMSNorm 与 nn.RMSNorm 的最大差异：{(ref_rms(x) - yr).abs().max().item():.2e}") print(""" RMSNorm 不把均值拉到 0，只把"整体大小"（均方根）归一化到 1。 实践表明这样效果和 LayerNorm 一样好，而且：少算一个均值、少一个偏置参数、少读写一些数据。""") # 【边界】低精度输入时，为什么要先转 FP32：FP16 的最大值只有 65504， # 激活值只要超过约 256，平方就会上溢成 inf。 xh = (torch.randn(2, 1024) * 300).to(torch.float16) rms_big = RMSNorm(1024) naive = xh / torch.sqrt(xh.pow(2).mean(-1, keepdim=True) + 1e-5) # 全程 FP16 good = rms_big(xh) print(f" 【边界】FP16 输入（量级约 300）：全程 FP16 计算 → 输出中有 {torch.isfinite(naive).logical_not().sum().item()} " f"个 inf/nan、{(naive == 0).sum().item()} 个 0（平方上溢后，除以 inf 得到 0）；" f"先转 FP32 再算 → 全部正常：{torch.isfinite(good).all().item()}") # ============================================================================= # 第 3 部分：FLOPs 不等于运行时间 # ============================================================================= def time_it(fn, repeats=20): """CPU 上简单计时：热身 3 次后取平均（秒）。""" for _ in range(3): fn() t0 = time.perf_counter() for _ in range(repeats): fn() return (time.perf_counter() - t0) / repeats def demo_flops_vs_runtime(): """演示 3：归一化层 FLOPs 占比很小，时间占比却不小。""" print_header("演示 3：FLOPs 不等于运行时间") torch.manual_seed(0) B, T, d = 8, 256, 1024 x = torch.randn(B, T, d) w = torch.randn(d, 4 * d) / math.sqrt(d) norm = RMSNorm(d) ln = LayerNorm(d) flops_mm = 2 * B * T * d * 4 * d # 一个 d → 4d 的矩阵乘法 flops_norm = 4 * B * T * d # RMSNorm：平方、求和、除、乘，大约每个元素 4 次运算 t_mm = time_it(lambda: x @ w) t_rms = time_it(lambda: norm(x)) t_ln = time_it(lambda: ln(x)) print(f" 矩阵乘法 (B·T={B * T}, {d}→{4 * d})：FLOPs {flops_mm:.2e}，时间 {t_mm * 1e3:.2f} ms") print(f" RMSNorm：FLOPs {flops_norm:.2e}（是矩阵乘法的 {flops_norm / flops_mm:.2%}），" f"时间 {t_rms * 1e3:.2f} ms（是矩阵乘法的 {t_rms / t_mm:.0%}）") print(f" LayerNorm（我们的实现）：时间 {t_ln * 1e3:.2f} ms（比 RMSNorm 多 {t_ln / t_rms - 1:.0%}）") print(""" 归一化的 FLOPs 不到矩阵乘法的 0.1%，时间却占到可观的比例：它是"内存受限"的运算 （读写整个激活张量，但每个数只做几次运算，第 2 讲第 6 节）。 所以"FLOPs 少"不代表"不重要"——RMSNorm 比 LayerNorm 少读写、少计算，在真实训练中能省下可观的时间。 CS336 课上引用的数据（Ivanov et al.）：在 BERT 中，归一化等运算只占 0.17% 的 FLOPs，却占约 25% 的运行时间。 （你电脑上的具体比例会不同；在 GPU 上，矩阵乘法快得多，归一化的时间占比通常更高。第 6 讲会用算子融合解决它。）""") # ============================================================================= # 第 4 部分：前馈网络 # ============================================================================= def silu(x): """ SiLU（也叫 Swish）激活函数：SiLU(x) = x × sigmoid(x)。 【直觉】和 ReLU 很像（大的正数几乎原样通过，大的负数趋于 0），但处处光滑， 在 0 附近允许少量负值通过。作业 1 要求用 torch.sigmoid 自己实现。 """ return x * torch.sigmoid(x) def swiglu_hidden_dim(d_model, multiple_of=64): """ SwiGLU 前馈层的隐藏维度：约 8/3 × d_model，向上取整到 multiple_of 的倍数。 【为什么是 8/3】 普通前馈层：d → 4d → d，两个矩阵，参数 2 × 4d² = 8d²。 SwiGLU：有三个矩阵（W1、W3：d → d_ff；W2：d_ff → d），参数 3 × d × d_ff。 令 3 × d × d_ff = 8d²，得 d_ff = 8/3 × d：**与普通前馈层参数量相同**，比较才公平。 【为什么取整到 64 的倍数】 GPU 上矩阵的维度是 64（或 128）的倍数时效率最高（第 5 讲"分块"）。 """ raw = int(8 * d_model / 3) return multiple_of * ((raw + multiple_of - 1) // multiple_of) class SwiGLU(torch.nn.Module): """ SwiGLU 前馈网络（Llama、Qwen、Mistral 等使用；CS336 作业 1 要求实现）： FFN(x) = W2 ( SiLU(W1 x) ⊙ (W3 x) ) ⊙ 是逐元素相乘。W3 x 这一路叫"门"（gate）：它逐元素地决定 SiLU(W1 x) 的每个分量放行多少。 【形状】 W1, W3 : d_model → d_ff W2 : d_ff → d_model x : (..., d_model) → 同形状 """ def __init__(self, d_model, d_ff=None): super().__init__() d_ff = d_ff or swiglu_hidden_dim(d_model) self.w1 = Linear(d_model, d_ff) self.w2 = Linear(d_ff, d_model) self.w3 = Linear(d_model, d_ff) def forward(self, x): return self.w2(silu(self.w1(x)) * self.w3(x)) class ReLUFFN(torch.nn.Module): """原始 Transformer 的前馈网络：FFN(x) = W2 ReLU(W1 x)，d_ff = 4 × d_model。""" def __init__(self, d_model, d_ff=None): super().__init__() d_ff = d_ff or 4 * d_model self.w1 = Linear(d_model, d_ff) self.w2 = Linear(d_ff, d_model) def forward(self, x): return self.w2(torch.relu(self.w1(x))) def demo_ffn(): """演示 4：激活函数与前馈网络。""" print_header("演示 4：前馈网络 —— ReLU、GELU、SiLU 与 SwiGLU") xs = torch.tensor([-3.0, -1.0, -0.25, 0.0, 0.25, 1.0, 3.0]) print(f" x ：{[f'{v:6.2f}' for v in xs.tolist()]}") print(f" ReLU(x) ：{[f'{v:6.2f}' for v in torch.relu(xs).tolist()]}") print(f" GELU(x) ：{[f'{v:6.2f}' for v in torch.nn.functional.gelu(xs).tolist()]}") print(f" SiLU(x) ：{[f'{v:6.2f}' for v in silu(xs).tolist()]}") assert torch.allclose(silu(xs), torch.nn.functional.silu(xs)) print("\n 参数量对比（d_model = 768）：") d = 768 relu_ffn, swiglu = ReLUFFN(d), SwiGLU(d) n1 = sum(p.numel() for p in relu_ffn.parameters()) n2 = sum(p.numel() for p in swiglu.parameters()) print(f" ReLU FFN（d_ff = 4d = {4 * d}） ：{n1:,} 个参数") print(f" SwiGLU（d_ff = 8/3·d → {swiglu_hidden_dim(d)}）：{n2:,} 个参数（相差 {n2 / n1 - 1:+.1%}；" f"8/3 × 768 恰好是 2048，没有取整误差）") print(" 各种 d_model 下 SwiGLU 的 d_ff：", {dm: swiglu_hidden_dim(dm) for dm in [512, 768, 1024, 4096]}) print(""" 为什么 GLU 类前馈层更好？Shazeer（2020）的实验显示，在参数量相同的情况下， GEGLU / SwiGLU 的困惑度稳定地优于 ReLU / GELU，后续多篇工作也复现了这一点。 至于"为什么"，作者坦言没有解释（讲义第 4 节）。一种直觉是：门控让网络可以做 "乘法式"的交互（两路信号相乘），表达能力比单纯的"线性 + 逐元素非线性"更强。""") def main(): demo_linear_embedding() demo_norms() demo_flops_vs_runtime() demo_ffn() print("\n全部演示运行完毕。") if __name__ == "__main__": main()