Transformer 算力与显存计算器

把一个 Transformer LM 的参数量、前向 FLOPs、训练一步的显存和时间从头数一遍。每一项都写出符号公式,再代入当前配置。

← 回到简化版(这里是保留全部参数的完整版)

收起时右侧仍是实时结果,改配置即刷新
▸图 架构:一个 token 走完全程
滚轮缩放 · 拖拽平移 · 双指捏合 · 双击复位
100%
▸1 算力:一次前向、一步训练要多少 FLOPs

一次 forward 的 FLOPs

A(m×n) @ B(n×p) = 2·m·n·p FLOPs x[T,d_in] @ W[d_in,d_out] = 2 · token数 · 参数量

从一次前向推到一步训练

一步训练 ≈ 3 × 一次前向 AdamW 一步 ≈ 14·P

上下文拉长,FLOPs 的结构会翻转

641,048,576
▸2 显存:4P 是地板,saved tensors 是天花板

同一份清单,参数和 saved tensors 各占多少

参数 P 的分布

saved tensors S 的分布

两条用的是同一套颜色:一个颜色 = 一个环节, 可以横着追同一行在两个口径下的分量。某条里缺了某个颜色,说明那一行在那个口径下是 0—— 打分 在参数那条里没有,token_embeddings 在 saved 那条里没有。 精确数字看下面的表,第三个口径(占训练显存)只在表里给。

▸3 训练时间:token 预算 ↔ 天数,两头都能定
每 token 成本 = 3·F / T  总算力 = 每 token 成本 × D 有效算力 = 峰值 × MFU × 卡数 天数 = 总算力 / 有效算力  D = 天数 × 有效算力 × T / (3·F)

为什么 MFU 打不满:arithmetic intensity

AI = FLOPs / 搬运字节数 机器平衡点 = 峰值算力 / 显存带宽 AI > 平衡点 → compute-bound AI < 平衡点 → memory-bound