""" ============================================================================= 第 2 讲 · 代码 4：计算量核算 —— FLOPs、6ND 公式与 MFU ============================================================================= 对应讲义：lectures/lecture02_pytorch_resource_accounting.md 第 5、6 节 ----------------------------------------------------------------------------- 【这个文件要解决什么问题？】 ----------------------------------------------------------------------------- "训练这个模型要多久？"是大模型工程师每天都在回答的问题。答案来自三个数： 总计算量（FLOPs） ÷ （硬件峰值算力 × 利用率 MFU） = 训练时间 这个文件逐一验证： 1. 一次矩阵乘法 (m×k) @ (k×n) 需要 2mkn 次浮点运算 —— 用 PyTorch 的 FLOP 计数器验证； 2. 反向传播的计算量约是前向的 2 倍 → 训练总计算量 ≈ 6 × 参数量 × token 数（6ND）； 3. 在你的电脑上实测矩阵乘法能跑到多少 FLOP/s；为什么小矩阵、逐元素运算远远跑不满； 4. 用这些数估算真实大模型的训练时间。 ----------------------------------------------------------------------------- 【整体思路】 ----------------------------------------------------------------------------- 第 1 部分：矩阵乘法的 FLOPs 第 2 部分：前向 vs 反向 —— 为什么是 2 : 4 第 3 部分：6ND —— 在一个多层网络上验证 第 4 部分：实测：矩阵乘法吞吐量，以及"矩阵越小越跑不满" 第 5 部分：实测：逐元素运算为什么慢（内存受限，第 5 讲的预告） 第 6 部分：估算真实模型的训练时间与 MFU ----------------------------------------------------------------------------- 【运行方式】 ----------------------------------------------------------------------------- python 04_flops_and_mfu.py （几秒到几十秒，取决于电脑，大部分时间在计时） """ import time import torch from torch.utils.flop_counter import FlopCounterMode def print_header(title): """打印一个带分隔线的标题。""" print() print("=" * 70) print(f" {title}") print("=" * 70) # 硬件规格（BF16/FP16 稠密算力，不含"结构化稀疏"加成；显存带宽）。数据来自 NVIDIA 公开规格，取近似值。 GPU_SPECS = { # (BF16 稠密 FLOP/s, FP32 FLOP/s, 显存 GB, 显存带宽 字节/s) "A100": (312e12, 19.5e12, 80, 2.0e12), "H100": (989.5e12, 67e12, 80, 3.35e12), "B200": (2.25e15, 75e12, 180, 8.0e12), } # ============================================================================= # 第 1 部分：矩阵乘法的 FLOPs # ============================================================================= def matmul_flops(m, k, n): """ (m×k) @ (k×n) 的浮点运算次数。 【思路】 结果矩阵有 m×n 个元素；每个元素是 k 对数的乘积之和：k 次乘法 + (k-1) 次加法 ≈ 2k 次运算。 所以总共约 2mkn 次。（习惯上把一次"乘加"算作 2 次浮点运算。） 【参数】m, k, n : 正整数 【返回】int """ return 2 * m * k * n def count_flops(fn): """ 用 PyTorch 自带的 FLOP 计数器，统计执行 fn() 时发生的浮点运算次数。 【参数】 fn : 无参数的函数 【返回】 int，总 FLOPs 【边界】 FlopCounterMode 只统计它"认识"的运算（矩阵乘法、卷积、注意力等）， 逐元素运算（加法、ReLU 等）通常不计入。这正好符合惯例： 大模型的计算量几乎全部来自矩阵乘法，逐元素运算的 FLOPs 可以忽略不计 （但它们的**时间**不能忽略，见第 5 部分）。 【Python 小知识】 with ... as fc: 叫"上下文管理器"，进入时开始计数，离开 with 块时停止。 """ with FlopCounterMode(display=False) as fc: fn() return fc.get_total_flops() def demo_matmul_flops(): """演示 1：验证 2mkn。""" print_header("演示 1：矩阵乘法的 FLOPs = 2 × m × k × n") for m, k, n in [(16, 32, 8), (128, 256, 512), (1000, 3, 7)]: a, b = torch.randn(m, k), torch.randn(k, n) counted = count_flops(lambda: a @ b) print(f" ({m}×{k}) @ ({k}×{n})：公式 {matmul_flops(m, k, n):>12,}，计数器 {counted:>12,}") assert counted == matmul_flops(m, k, n) # 对比：逐元素运算 x = torch.randn(1000, 1000) print(f"\n 对 100 万个元素做 x * 2 + 1：计数器统计到 {count_flops(lambda: x * 2 + 1)} 次（逐元素运算不计入）") print("\n 【大模型联系】一个线性层 y = x W，x 是 (B×D_in)，W 是 (D_in×D_out)：") print(" 前向 FLOPs = 2 × B × D_in × D_out = 2 × (token 数) × (这一层的参数量)。") # ============================================================================= # 第 2 部分：前向 vs 反向 # ============================================================================= # 【思路】 # 对一个线性层 h_out = h_in @ W（h_in: B×D_in，W: D_in×D_out）： # 前向：h_out = h_in @ W → 2·B·D_in·D_out # 反向需要两样东西： # (1) 对参数的梯度 dL/dW = h_in^T @ dL/dh_out → 2·B·D_in·D_out # (2) 对输入的梯度 dL/dh_in = dL/dh_out @ W^T → 2·B·D_in·D_out # （要把梯度继续往前传给上一层） # 所以反向 = 前向的 2 倍。前向 + 反向 = 3 倍前向 = 6·B·(参数量)。 # # 【边界】第一层的输入是数据本身，不需要对它求梯度，所以第一层的反向只有 (1)。 # 这使得真实的比例略小于 2 —— 下面的实验会看到这一点。 def demo_forward_backward(): """演示 2：前向与反向的 FLOPs 之比。""" print_header("演示 2：反向传播的 FLOPs ≈ 前向的 2 倍") torch.manual_seed(0) B, D = 64, 256 w1 = torch.randn(D, D, requires_grad=True) w2 = torch.randn(D, D, requires_grad=True) for x_needs_grad in [True, False]: x = torch.randn(B, D, requires_grad=x_needs_grad) fwd = count_flops(lambda: (x @ w1) @ w2) # 【PyTorch 小知识】loss.backward() 会沿计算图反向传播，把梯度存进每个参数的 .grad 里 total = count_flops(lambda: ((x @ w1) @ w2).sum().backward()) bwd = total - fwd label = "输入也需要梯度" if x_needs_grad else "输入是数据（不需要梯度）" print(f" 两层网络，{label}：前向 {fwd:,}，反向 {bwd:,}，反向/前向 = {bwd / fwd:.2f}") w1.grad = w2.grad = None # 清空梯度，避免下一轮累加 print(""" - 每一层都需要"对参数求梯度"和"对输入求梯度"两次矩阵乘法 → 反向正好是前向的 2 倍。 - 第一层的输入是数据，不需要梯度，省掉一次矩阵乘法。这个网络只有 2 层， 省掉的部分占比很大（比例降到 1.5）；层数越多，省掉的占比越小，比例越接近 2 （下一个演示的 8 层网络就很接近了）。对几十上百层的大模型， 这点差别可以忽略，所以我们说"反向 ≈ 2 × 前向"。""") # ============================================================================= # 第 3 部分：6ND # ============================================================================= def demo_6nd(): """演示 3：在一个多层网络上验证 总FLOPs ≈ 6 × token数 × 参数量。""" print_header("演示 3：训练总计算量 ≈ 6 × N（参数）× D（token）") torch.manual_seed(0) num_layers, D = 8, 256 # 【PyTorch 小知识】nn.Sequential 把多个层串起来；nn.Linear(D, D, bias=False) 是不带偏置的线性层 model = torch.nn.Sequential(*[torch.nn.Linear(D, D, bias=False) for _ in range(num_layers)]) n_params = sum(p.numel() for p in model.parameters()) tokens = 512 x = torch.randn(tokens, D) total = count_flops(lambda: model(x).sum().backward()) ratio = total / (tokens * n_params) print(f" {num_layers} 层 × ({D}×{D})，参数量 N = {n_params:,}，一个批次的 token 数 D = {tokens}") print(f" 实测一次前向+反向 = {total:,} FLOPs") print(f" 6 × N × D = {6 * n_params * tokens:,} FLOPs") print(f" 实测 / (N × D) = {ratio:.3f} ← 接近 6（第一层省掉的输入梯度使它略小于 6）") print(""" 【边界】6ND 是一个近似，它没有算进： - 注意力中与序列长度相关的那部分计算（QK^T 和 注意力×V，第 3 讲）：序列很长时不可忽略； - 激活重计算（activation checkpointing，见 05 号文件）：会多做一次前向，变成约 8ND； - 逐元素运算、归一化、softmax 等：FLOPs 很少，但会占用可观的时间。""") # ============================================================================= # 第 4 部分：实测矩阵乘法吞吐量 # ============================================================================= def get_sync(device): """ 返回"等待该设备上所有运算完成"的函数。 【思路】 GPU 运算是异步的：Python 发出命令后立即返回。计时之前必须同步，否则测到的只是"发命令"的时间。 CPU 运算是同步的，不需要等，返回一个什么也不做的函数。 """ if device.type == "cuda": return torch.cuda.synchronize if device.type == "mps": return torch.mps.synchronize return lambda: None def benchmark(fn, device, warmup=3, repeats=10): """ 测量 fn() 的平均运行时间（秒）。 【思路】 1. 先"热身"几次：第一次运行常常包含一次性开销（分配内存、选择算法、编译等）； 2. 同步，开始计时； 3. 连续运行 repeats 次； 4. 再同步，停止计时，取平均。 第 6 讲会更详细地讨论如何正确地做基准测试。 【参数】 fn : 无参数函数 device : torch.device warmup : 热身次数（>= 0） repeats : 计时次数（>= 1） 【返回】 float，每次平均耗时（秒） """ sync = get_sync(device) for _ in range(warmup): fn() sync() t0 = time.perf_counter() for _ in range(repeats): fn() sync() return (time.perf_counter() - t0) / repeats def available_devices(): """返回本机可用的设备列表（总是包含 CPU）。""" devs = [torch.device("cpu")] if torch.cuda.is_available(): devs.append(torch.device("cuda")) if torch.backends.mps.is_available(): devs.append(torch.device("mps")) return devs def demo_matmul_throughput(): """演示 4：不同大小的矩阵乘法能跑到多少 FLOP/s。""" print_header("演示 4：实测矩阵乘法吞吐量（FP32）") sizes = [64, 256, 1024, 2048, 4096] results = {} for device in available_devices(): row = [] for n in sizes: a = torch.randn(n, n, device=device) b = torch.randn(n, n, device=device) t = benchmark(lambda: a @ b, device, repeats=5 if n >= 2048 else 20) row.append(matmul_flops(n, n, n) / t) results[device.type] = row header = f" {'n':>6}" + "".join(f"{d + ' (GFLOP/s)':>18}" for d in results) print(header) for i, n in enumerate(sizes): print(f" {n:>6}" + "".join(f"{results[d][i] / 1e9:>18,.0f}" for d in results)) for d, row in results.items(): best = max(row) print(f" {d}：最小的矩阵（64）只达到最大吞吐量的 {row[0] / best:.1%}") print(""" 观察：矩阵越大，吞吐量越高，直到接近硬件上限。小矩阵跑不满的原因： - 计算太少，固定开销（启动运算、调度线程）占了大头； - 相对于计算量，数据搬运太多（算术强度低，第 5 讲的 roofline 模型会精确解释）。 【大模型联系】这就是为什么训练要用大批次、大矩阵：只有这样才能把硬件"喂饱"。 （你的电脑的具体数字取决于 CPU/GPU 型号，和讲义里的不同是正常的。）""") # ============================================================================= # 第 5 部分：逐元素运算为什么慢 # ============================================================================= def demo_elementwise(): """演示 5：逐元素运算的 FLOP/s 远低于矩阵乘法。""" print_header("演示 5：逐元素运算 —— FLOPs 很少，时间却不少") device = torch.device("cpu") n = 1 << 24 # 约 1677 万个元素 x = torch.randn(n) t = benchmark(lambda: x * 2.0, device, repeats=10) flops = n # 每个元素 1 次乘法 bytes_moved = 2 * n * 4 # 读 n 个 float32 + 写 n 个 float32 print(f" x * 2（{n:,} 个 float32）：耗时 {t * 1e3:.2f} 毫秒") print(f" 计算吞吐 = {flops / t / 1e9:.1f} GFLOP/s，数据搬运速度 = {bytes_moved / t / 1e9:.1f} GB/s") m = 2048 a, b = torch.randn(m, m), torch.randn(m, m) tm = benchmark(lambda: a @ b, device, repeats=5) print(f" 对比：2048×2048 矩阵乘法 {matmul_flops(m, m, m) / tm / 1e9:,.0f} GFLOP/s") print(""" 逐元素乘法每读写 8 个字节才做 1 次运算（算术强度 = 1/8 FLOP/字节）， 速度完全被"搬数据"卡住了（内存受限）；矩阵乘法每个数据会被重复使用很多次（计算受限）。 【大模型联系】Transformer 里的 LayerNorm、激活函数、残差相加、softmax 都是这类运算： 它们的 FLOPs 占比不到 1%，却可能占到训练时间的相当一部分。 这就是第 5、6 讲"算子融合"要解决的问题。""") # ============================================================================= # 第 6 部分：估算真实模型的训练时间与 MFU # ============================================================================= def training_days(n_params, n_tokens, n_gpus, peak_flops, mfu): """ 估算训练时间（天）。 【公式】 时间 = 6 × N × D / (GPU 数 × 单卡峰值 × MFU) 【参数】 n_params : 参数量 N n_tokens : 训练 token 数 D n_gpus : GPU 数量 peak_flops : 单卡峰值算力（FLOP/s） mfu : 模型浮点利用率，0 ~ 1 之间（实践中 0.3 ~ 0.6 算正常） 【返回】 float，天数 【边界】mfu 必须在 (0, 1] 之间。 """ if not (0 < mfu <= 1): raise ValueError("MFU 必须在 (0, 1] 之间") return 6 * n_params * n_tokens / (n_gpus * peak_flops * mfu) / 86400 def mfu_from_throughput(n_params, tokens_per_sec, n_gpus, peak_flops): """ 从实测的训练吞吐量（tokens/秒）计算 MFU。 【公式】 MFU = 实际达到的"模型 FLOP/s" / 峰值 FLOP/s = 6 × N × (tokens/秒) / (GPU 数 × 单卡峰值) 【思路】 MFU 只计算"模型本身必需的"计算量（6N 每个 token）， 不把激活重计算等"额外的"计算算作有效工作 —— 所以它衡量的是硬件被"有效利用"的程度。 """ return 6 * n_params * tokens_per_sec / (n_gpus * peak_flops) def demo_training_estimates(): """演示 6：真实规模的估算。""" print_header("演示 6：估算真实模型的训练时间") h100 = GPU_SPECS["H100"][0] print(" 硬件规格（BF16 稠密算力 / 显存 / 带宽）：") for name, (bf16, fp32, mem, bw) in GPU_SPECS.items(): print(f" {name}：{bf16 / 1e12:>7,.0f} TFLOP/s，{mem} GB，{bw / 1e12:.2f} TB/s") days = training_days(70e9, 15e12, 1024, h100, 0.5) print(f"\n (a) 70B 参数、15T token、1024 块 H100、MFU 50%：约 {days:.0f} 天") gpt3 = 6 * 175e9 * 300e9 print(f" (b) GPT-3（175B 参数、300B token）：6ND = {gpt3:.2e} FLOPs（论文报告约 3.14e23）") budget = 8 * h100 * 14 * 86400 print(f" (c) 8 块 H100 满负荷跑 2 周能提供 {budget:.2e} FLOPs；按 MFU 50%、每 token 6N 计算，") for n in [1e9, 7e9]: print(f" 训练 {n / 1e9:.0f}B 参数的模型可以用 {budget * 0.5 / (6 * n) / 1e9:,.0f}B 个 token") mfu = mfu_from_throughput(1.5e9, 40_000, 1, h100) print(f"\n (d) 在 1 块 H100 上训练 1.5B 参数的模型，实测 40,000 tokens/秒 → MFU = {mfu:.0%}") def main(): torch.manual_seed(0) demo_matmul_flops() demo_forward_backward() demo_6nd() demo_matmul_throughput() demo_elementwise() demo_training_estimates() print("\n全部演示运行完毕。") if __name__ == "__main__": main()