""" ============================================================================= 第 3 讲 · 代码 4：完整的 Transformer 语言模型 + 参数量、FLOPs、显存核算 ============================================================================= 对应讲义：lectures/lecture03_architectures_hyperparameters.md 第 1、7 节 ----------------------------------------------------------------------------- 【这个文件要解决什么问题？】 ----------------------------------------------------------------------------- 把前三个文件的积木拼成 CS336 作业 1 要求的 Transformer 语言模型（Llama 风格）： token ids ─► Embedding ─► [TransformerBlock] × L ─► RMSNorm ─► Linear(lm_head) ─► logits TransformerBlock（Pre-norm）： y = x + Attention( RMSNorm(x) ) z = y + SwiGLU( RMSNorm(y) ) 然后把第 2 讲的"会计方法"用在它身上，回答作业 1 的资源核算题： - 参数量公式是什么？（与实际模型逐一核对） - 一次前向传播多少 FLOPs？各部分占多少？随模型变大、上下文变长怎么变？（用 FLOP 计数器核对） - 用 AdamW 训练 GPT-2 XL 大小的模型，80 GB 显存最多能用多大的批次？ ----------------------------------------------------------------------------- 【运行方式】 ----------------------------------------------------------------------------- python 04_transformer_lm.py 依赖同目录下的 01_building_blocks.py、02_rope.py、03_attention.py。 """ import importlib import torch from torch.utils.flop_counter import FlopCounterMode _blocks = importlib.import_module("01_building_blocks") _attn = importlib.import_module("03_attention") Linear, Embedding, RMSNorm, SwiGLU = _blocks.Linear, _blocks.Embedding, _blocks.RMSNorm, _blocks.SwiGLU MultiHeadSelfAttention = _attn.MultiHeadSelfAttention def print_header(title): """打印一个带分隔线的标题。""" print() print("=" * 70) print(f" {title}") print("=" * 70) # ============================================================================= # 第 1 部分：模型 # ============================================================================= class TransformerBlock(torch.nn.Module): """ 一个 Transformer 块（CS336 作业 1 的 transformer_block）。 【Pre-norm（默认，现代模型几乎都用）】 y = x + Attn(Norm(x)) z = y + FFN(Norm(y)) 归一化放在每个子层的**输入**处，残差连接（x + ...）这条"主干道"上没有任何归一化， 信号和梯度可以不受干扰地从最后一层直通到第一层。 【Post-norm（原始 Transformer、BERT 使用）】 y = Norm(x + Attn(x)) z = Norm(y + FFN(y)) 归一化放在残差相加**之后**，主干道上每一层都要经过一次归一化。 实践中它在深层网络上更难训练（通常需要学习率预热），讲义第 3.1 节。 【参数】 norm_position : "pre" 或 "post"，方便你做对比实验 【形状】x : (batch, seq, d_model) → 同形状 """ def __init__(self, d_model, num_heads, d_ff, max_seq_len, theta=10000.0, norm_position="pre"): super().__init__() if norm_position not in ("pre", "post"): raise ValueError("norm_position 只能是 'pre' 或 'post'") self.norm_position = norm_position self.ln1 = RMSNorm(d_model) self.attn = MultiHeadSelfAttention(d_model, num_heads, max_seq_len, theta) self.ln2 = RMSNorm(d_model) self.ffn = SwiGLU(d_model, d_ff) def forward(self, x, token_positions=None): if self.norm_position == "pre": x = x + self.attn(self.ln1(x), token_positions) return x + self.ffn(self.ln2(x)) x = self.ln1(x + self.attn(x, token_positions)) return self.ln2(x + self.ffn(x)) class TransformerLM(torch.nn.Module): """ Transformer 语言模型（CS336 作业 1 的 transformer_lm）。 【参数】 vocab_size : 词表大小 V context_length : 最大上下文长度（RoPE 预先计算这么多位置） d_model : 隐藏维度 d num_layers : 层数 L num_heads : 注意力头数 h d_ff : 前馈层隐藏维度 rope_theta : RoPE 的 Θ norm_position : "pre" 或 "post" 【形状】 in_indices : (batch, seq) 整数，seq ≤ context_length 返回 logits: (batch, seq, vocab_size)（未经 softmax） 【边界】 - 输入序列长度超过 context_length 时，RoPE 的位置表会越界 → 报错。 - 输入嵌入和输出层（lm_head）是两个独立的矩阵（不共享权重），与作业 1 一致。 有些模型（如 GPT-2、Gemma）让它们共享，可以省下 V × d 个参数。 """ def __init__(self, vocab_size, context_length, d_model, num_layers, num_heads, d_ff, rope_theta=10000.0, norm_position="pre"): super().__init__() self.context_length = context_length self.token_embeddings = Embedding(vocab_size, d_model) self.layers = torch.nn.ModuleList([ TransformerBlock(d_model, num_heads, d_ff, context_length, rope_theta, norm_position) for _ in range(num_layers) ]) # Post-norm 的每一层输出已经归一化过了，不再需要最后的归一化 self.ln_final = RMSNorm(d_model) if norm_position == "pre" else torch.nn.Identity() self.lm_head = Linear(d_model, vocab_size) def forward(self, in_indices): seq = in_indices.shape[-1] if seq > self.context_length: raise ValueError(f"序列长度 {seq} 超过了 context_length {self.context_length}") positions = torch.arange(seq, device=in_indices.device) x = self.token_embeddings(in_indices) for layer in self.layers: x = layer(x, positions) return self.lm_head(self.ln_final(x)) # ============================================================================= # 第 2 部分：参数量 # ============================================================================= def param_count_formula(V, d, L, d_ff): """ 参数量公式（Pre-norm、不共享嵌入、无偏置、SwiGLU）： 嵌入 V × d 每层：注意力 4d²（W_Q、W_K、W_V、W_O） + SwiGLU 3·d·d_ff + 两个 RMSNorm 2d 最终 RMSNorm d lm_head d × V """ per_layer = 4 * d * d + 3 * d * d_ff + 2 * d return V * d + L * per_layer + d + d * V def demo_params(): """演示 1：搭一个小模型，核对参数量公式。""" print_header("演示 1：搭建模型，核对参数量") torch.manual_seed(0) cfg = dict(vocab_size=1000, context_length=128, d_model=64, num_layers=3, num_heads=4, d_ff=192) model = TransformerLM(**cfg) actual = sum(p.numel() for p in model.parameters()) formula = param_count_formula(cfg["vocab_size"], cfg["d_model"], cfg["num_layers"], cfg["d_ff"]) print(f" 配置：{cfg}") print(f" 实际参数量 {actual:,}，公式 {formula:,}，一致：{actual == formula}") ids = torch.randint(0, 1000, (2, 16)) print(f" 输入 {tuple(ids.shape)} → logits {tuple(model(ids).shape)}") for name, p in list(model.named_parameters())[:6]: print(f" {name:<32} {tuple(p.shape)}") print(" ……") # ============================================================================= # 第 3 部分：前向传播的 FLOPs # ============================================================================= def forward_flops(V, d, L, h, d_ff, T): """ 一条长度为 T 的序列做一次前向传播的矩阵乘法 FLOPs（按 2mkn 计），分项返回。 【各项】 注意力投影：Q、K、V、O 四个 (T×d)@(d×d) → L × 4 × 2·T·d² 注意力核心：Q K^T 和 注意力权重 @ V → L × 2 × 2·T²·d （每个头 2·T²·d_head，h 个头合计 2·T²·d；这里按完整的 T×T 计算，不扣除因果掩码省掉的一半） 前馈网络 ：W1、W3、W2 三个矩阵 → L × 3 × 2·T·d·d_ff 输出层 ：(T×d)@(d×V) → 2·T·d·V """ return { "注意力投影": L * 4 * 2 * T * d * d, "注意力核心": L * 2 * 2 * T * T * d, "前馈网络": L * 3 * 2 * T * d * d_ff, "输出层": 2 * T * d * V, } def demo_flops_check(): """演示 2：用 FLOP 计数器核对公式。""" print_header("演示 2：前向 FLOPs 公式 vs FLOP 计数器") torch.manual_seed(0) V, d, L, h, d_ff, T = 1000, 64, 3, 4, 192, 32 model = TransformerLM(V, 128, d, L, h, d_ff) ids = torch.randint(0, V, (1, T)) with FlopCounterMode(display=False) as fc: model(ids) counted = fc.get_total_flops() formula = sum(forward_flops(V, d, L, h, d_ff, T).values()) print(f" 计数器 {counted:,}，公式 {formula:,}，一致：{counted == formula}") def demo_gpt2_family(): """演示 3：GPT-2 各个尺寸（用作业 1 的架构）的参数量与 FLOPs 构成。""" print_header("演示 3：GPT-2 small / medium / large / XL 的 FLOPs 构成（作业 1 题目）") V, T = 50257, 1024 family = [("small", 12, 768, 12), ("medium", 24, 1024, 16), ("large", 36, 1280, 20), ("XL", 48, 1600, 25)] print(f" {'模型':<8}{'参数量':>10}{'前向 TFLOPs':>12}{'注意力投影':>10}{'注意力核心':>10}{'前馈网络':>10}{'输出层':>8}") for name, L, d, h in family: d_ff = 4 * d f = forward_flops(V, d, L, h, d_ff, T) total = sum(f.values()) n = param_count_formula(V, d, L, d_ff) print(f" {name:<8}{n / 1e9:>9.2f}B{total / 1e12:>12.2f}" + "".join(f"{f[k] / total:>11.1%}" for k in f)) L, d, h = 48, 1600, 25 f = forward_flops(V, d, L, h, 4 * d, 16384) total = sum(f.values()) print(f" {'XL@16K':<8}{param_count_formula(V, d, L, 4 * d) / 1e9:>9.2f}B{total / 1e12:>12.2f}" + "".join(f"{f[k] / total:>11.1%}" for k in f)) print(""" 观察（这正是第 1 讲 1.3 节"小模型不是大模型的缩小版"的具体体现）： - 模型越大，前馈网络占比越高；输出层占比越低（V × d 相对于 L × d² 越来越不重要）。 - 注意力核心（与 T² 成正比）在上下文变长时迅速膨胀：XL 从 1K 到 16K 上下文， 它从一小部分变成了最大的一项。这就是长上下文昂贵的原因（第 4 讲的注意力替代方案要解决它）。 - 除注意力核心外，其余各项都约等于 2 × T × 参数量，这就是"前向 ≈ 2N 每 token"的来源。""") # ============================================================================= # 第 4 部分：AdamW 训练的显存核算 # ============================================================================= def activation_floats_estimate(B, T, d, L, h, d_ff, V): """ 训练时为反向传播保存的激活值个数（估算）。 【每层保存的张量】（B = 批次，T = 序列长度） 两个 RMSNorm 的输出 2·B·T·d Q、K、V 3·B·T·d 注意力分数 与 softmax 结果 2·B·h·T² ← 与 T² 成正比！ 注意力加权和、输出投影 2·B·T·d SwiGLU：W1 输出、W3 输出、SiLU 输出、两者乘积 4·B·T·d_ff W2 输出 1·B·T·d 【最后】 最终 RMSNorm 输出 B·T·d；logits 与 softmax 结果 2·B·T·V 【边界】 这是"按公式列出主要张量"的估算。真实框架保存的张量会因实现细节多一些或少一些 （例如融合的注意力核函数根本不保存 T×T 的注意力矩阵，见第 5 讲 FlashAttention）。 """ per_layer = B * T * (8 * d + 4 * d_ff) + 2 * B * h * T * T return L * per_layer + B * T * d + 2 * B * T * V def demo_memory(): """演示 4：用 AdamW 训练 GPT-2 XL（FP32），80 GB 显存能放下多大的批次？""" print_header("演示 4：AdamW 训练 GPT-2 XL 的显存（作业 1 题目，全部 FP32）") V, T, L, d, h = 50257, 1024, 48, 1600, 25 d_ff = 4 * d P = param_count_formula(V, d, L, d_ff) fixed = P * 4 + P * 4 + P * 8 # 参数 + 梯度 + AdamW 的 m、v per_sample = activation_floats_estimate(1, T, d, L, h, d_ff, V) * 4 print(f" 参数量 P = {P / 1e9:.2f}B") print(f" 参数 {P * 4 / 1e9:.1f} GB + 梯度 {P * 4 / 1e9:.1f} GB + AdamW 状态 {P * 8 / 1e9:.1f} GB = {fixed / 1e9:.1f} GB") print(f" 激活值：每个样本（{T} 个 token）约 {per_sample / 1e9:.1f} GB") attn_part = L * 2 * h * T * T * 4 / 1e9 print(f" 其中注意力分数和 softmax 结果（T² 项）约 {attn_part:.1f} GB，占 {attn_part / (per_sample / 1e9):.0%}") max_b = int((80e9 - fixed) // per_sample) print(f" 总显存 ≈ {fixed / 1e9:.1f} GB + {per_sample / 1e9:.1f} GB × 批次大小；80 GB 最多放下批次 {max_b}") print(""" 结论：光是参数、梯度、优化器状态就要 34 GB；每多一个样本又要十几 GB 激活。 怎么办？混合精度（第 2 讲）、激活检查点（第 2 讲）、FlashAttention 不保存 T×T 矩阵（第 5 讲）、 多卡切分优化器状态和参数（第 7 讲 ZeRO/FSDP）……这门课后面的系统单元基本都在解决这个问题。""") # 用"保存张量钩子"在一个小配置上实测，与估算对比 measured, estimated = measure_small_config() print(f" 小配置实测：激活 {measured:,} 个数，公式估算 {estimated:,} 个数（实测/估算 = {measured / estimated:.2f}）") print(" 实测更多，是因为我们用教学风格写的 softmax、RoPE 会保存额外的中间结果；") print(" 高度优化的实现（融合核函数）会比估算更少。估算的价值在于抓住主要项和它们的增长规律。") def measure_small_config(): """在小配置上用 saved_tensors_hooks 实测激活值个数（方法见第 2 讲 05 号文件）。""" torch.manual_seed(0) V, T, L, d, h, d_ff, B = 500, 64, 2, 64, 4, 192, 2 model = TransformerLM(V, T, d, L, h, d_ff) ids = torch.randint(0, V, (B, T)) targets = torch.randint(0, V, (B, T)) param_ptrs = {p.data_ptr() for p in model.parameters()} seen = {} def pack(t): if t.data_ptr() not in param_ptrs and t.is_floating_point(): seen[t.data_ptr()] = t.numel() return t with torch.autograd.graph.saved_tensors_hooks(pack, lambda t: t): logits = model(ids) torch.nn.functional.cross_entropy(logits.view(-1, V), targets.view(-1)) return sum(seen.values()), activation_floats_estimate(B, T, d, L, h, d_ff, V) def main(): demo_params() demo_flops_check() demo_gpt2_family() demo_memory() print("\n全部演示运行完毕。") if __name__ == "__main__": main()