⚙️ Task 2 · PyTorch 与资源核算
作者：Winnie ｜ 2026-04-19 ｜ 阅读时间约 15 分钟
📌 先说结论（给没空读完的人）
- 🧠 这章教你一种「直觉」：每一行代码要能估出它烧多少显存、耗多少算力
- 🎯 两个公式钉死一切：
- 💸 训练时间 ≈
6 × 参数量 × token 数 / (FLOPS × MFU) - 💾 训练显存 ≈
16 × 参数量字节（FP32 + AdamW）
- 💸 训练时间 ≈
- ✅ 我实测验证了 16 字节公式：阶段 A/B/C 精确到小数点后两位命中
- 📐 作业题 2 全做完：GPT-2 XL 单卡 A100 训 400K 步 = 6354 天（17 年） —— 教你为什么要分布式
🎬 为什么这章很重要
Task 1 学的是 一个具体算法（BPE 分词器），动手感强、立竿见影。
Task 2 学的是 一套直觉——看完这章，你对着代码要能"报菜名"：
"这行
nn.Linear(1024, 1024)，参数量 1M、FP32 吃 4MB、训练时吃 16MB、一个 token 6M FLOPs"
这种直觉平时感受不到，但它能让你在老板问"这模型训得起吗"时，30 秒给出一个量级正确的答案，不用跑代码。
CS336 原话很实在 👇
资源核算（Resource Accounting）不是技能，是思维模式。
🗺️ 思维导图
在真开始讲之前，先铺张地图：
⚙️ PyTorch & 资源核算（Task 2）
│
├── 📐 两个核心公式
│ ├── 💸 时间 = 6·N·tokens / (FLOPS·MFU)
│ └── 💾 显存 = 16·N 字节（FP32 + AdamW）
│
├── 🔢 为什么是 6？
│ ├── 前向 2× （一乘一加）
│ └── 反向 4× （权重梯度 + 激活梯度）
│
├── 🧮 为什么是 16？
│ ├── 参数 4 字节
│ ├── 梯度 4 字节
│ └── 优化器 8 字节（AdamW 的 m + v）
│
├── 🎨 四种数据类型
│ ├── FP32 · 4B · 最稳
│ ├── FP16 · 2B · 范围窄易炸
│ ├── BF16 🏆 · 2B · 主流
│ └── FP8 · 1B · H100+ 才有
│
├── 🎯 MFU = 实测 / 峰值
│ └── ≥ 50% 就是不错
│
└── 🔬 作业题 2（adamwAccounting）
├── (a) 显存公式 16P + 4A
├── (b) 80GB A100 最大 B = 2
├── (c) FLOPs/step ≈ 6·B·L·P
└── (d) 400K 步 ≈ 17 年 ❗（先用 ASCII 占坑，后面和 Task 1 一起迭代成可交互的 Markmap）
🧠 核心一：为什么是 6 × 参数量 × tokens
Chapter 1 里有个经典的"餐巾纸问题"：
1024 张 H100 上训 70B 模型 15T tokens，多久能训完？
想在餐巾纸上解，你只需要两个数字：总工作量 和 硬件算力。
总工作量的公式：
$$\text{FLOPs}_{\text{训练}} \approx 6 \times \text{参数量} \times \text{tokens}$$
6 倍从哪儿来？拆开看就一目了然：
| 阶段 | 做什么 | FLOPs / token / 参数 |
|---|---|---|
| 🔜 前向 | 每个参数参与一次乘法 + 一次加法 | 2× |
| 🔙 反向-权重梯度 | 算每个权重对 loss 的偏导 | 2× |
| 🔙 反向-激活梯度 | 把梯度传回上一层 | 2× |
| 合计 | 6× |
📝 这里的"2"指的是矩阵乘法里"一乘一加"算两次 FLOPs。
代入餐巾纸问题验算 👇
总 FLOPs = 6 × 70e9 × 15e12 = 6.3e24
单卡算力 = 990e12 FLOPS × 50% MFU = 5e14 FLOPS
1024 卡 = 5.07e17 FLOPS
时间 = 6.3e24 / 5.07e17 = 1.24e7 秒 ≈ 144 天教材答案 146 天，我算的 144 天——差 2 天是四舍五入。公式立住 ✅
💾 核心二：为什么是 16 × 参数量 字节
训练一个参数要存多少东西？很多人只算"参数本身 4 字节"就算完了。错 ❌
训练时，显存里住着 四位房客 🏠：
| 房客 | 字节 | 干嘛的 |
|---|---|---|
| 🔢 参数 | 4 | 模型权重本身（FP32） |
| 📉 梯度 | 4 | 反向传播算出来的 .grad，和参数等大 |
| 🎯 优化器状态·m | 4 | AdamW 的一阶矩（梯度移动平均） |
| 🎯 优化器状态·v | 4 | AdamW 的二阶矩（梯度平方移动平均） |
| 合计 | 16 |
这就是 16 × 参数量 的由来。一个常见翻车点：
🎭 推理跑得好好的，一上训练就 OOM —— 因为推理只要 4 字节/参数，训练要 16 字节，翻了 4 倍。
🧪 实战一：亲手量一遍 16 字节
光看公式不过瘾，我构造了一个 2048×2048 的线性层（4.2M 参数），三阶段递进量它的显存。
model = nn.Linear(2048, 2048, bias=False) # 4.2M 参数
optimizer = torch.optim.AdamW(model.parameters())
# 阶段 A：刚创建 ——————— 只有参数
# 阶段 B：loss.backward() 后 —— 参数 + 梯度出现
# 阶段 C：optimizer.step() 后 —— 优化器状态出现实测结果 👇
| 阶段 | 参数 | 梯度 | 优化器 | 合计 | 字节/参数 |
|---|---|---|---|---|---|
| A · 刚创建 | 16.00 MB | 0 | 0 | 16 MB | 4 |
| B · 反向传播后 | 16.00 MB | 16.00 MB | 0 | 32 MB | 8 |
| C · 优化一步后 | 16.00 MB | 16.00 MB | 32.00 MB | 64 MB | 16 ✅ |
💎 最 tricky 的一点：优化器状态是 "隐形" 的——阶段 A、B 都看不到它，只有
optimizer.step()跑完一步才"凭空"冒出来，一下子翻倍。很多"明明估算够用，一跑就 OOM"的翻车，根源就在这——你没算优化器状态那 8 字节。
🎨 顺带：不同数据类型的内存
| dtype | 字节 | 模型大小（4.2M 参数） |
|---|---|---|
| FP32 | 4 | 16 MB |
| FP16 | 2 | 8 MB |
| BF16 🏆 | 2 | 8 MB |
FP16 和 BF16 内存一样，为啥 BF16 赢？
🪜 "要范围，不要精度"。
FP16 指数位只有 5 位，训练中
1e-8这种小数会直接下溢变成0，梯度消失了就不训了。BF16 截断尾数但保留 FP32 的 8 位指数位 —— 范围一样宽，稳定性接近 FP32。
现代 LLM 训练几乎清一色 BF16。
📐 实战二：作业题 2 全做完 · adamwAccounting
CS336 Assignment 1 · Problem
adamwAccounting（2 points）对一个 GPT-2 XL（vocab 50257 / context 1024 / 48 层 / d_model 1600 / 25 头 / d_ff 6400），用 FP32 和 AdamW 训练。
🅰️ (a) 峰值显存代数表达式
参数量 P 拆成 4 部分：
P = 2·V·D (token embed + LM head)
+ D (final RMSNorm)
+ N·(2D + 4D² + 3·D·F) (N 层 block)每层 block 里：2D（两个 RMSNorm 的 gain 参数）+ 4D²（W_Q/W_K/W_V/W_O）+ 3D·F（SwiGLU 的三个矩阵）。
代入 → P ≈ 2.13 B 参数。
激活 A(B) 每 batch 的元素数：
每 block = 8·B·L·D + 2·B·H·L² + 2·B·L·F
总激活 = N · 每block + B·L·D + 2·B·L·V总显存公式：
$$\boxed{\text{memory}(B) = 16P + 4 \cdot A(B) \text{ 字节}}$$
（= 4 倍的"参数 + 梯度 + 优化器 + 激活"）
🅱️ (b) 80 GB A100 的最大 batch_size
固定开销（参数 + 梯度 + 优化器）：
16 × 2.13 B = 34.03 GB (还没跑数据呢，先占了一大半)剩给激活的预算：80 - 34.03 = 45.97 GB
每 batch 的激活占用：
4 × A_per_sample = 4 × 3.88 G = 15.52 GB/batch所以：
$$\text{显存}(B) \approx 15.52 \cdot B + 34.03 \text{ GB}$$
$$B_{\max} = \left\lfloor \frac{45.97}{15.52} \right\rfloor = \boxed{2}$$
🤯 一个 batch 要 15.52 GB 激活 —— 比很多人想象的可怕得多。 这就是为啥做大模型时
batch_size=1很常见，都靠梯度累积（gradient accumulation）凑有效 batch。
🅲 (c) AdamW 一步的 FLOPs
矩阵乘法占 FLOPs 的 90%+，其他（LayerNorm/softmax/AdamW 更新本身）可忽略。
$$\boxed{\text{FLOPs}_{\text{每步}} \approx 6 \cdot B \cdot L \cdot P}$$
- 前向
2·B·L·P（每参数每 token 一乘一加） - 反向
4·B·L·P（权重梯度 + 激活梯度各2·B·L·P） - AdamW update
~10·P（相对6·B·L·P可忽略）
代入 B=1024, L=1024, P=2.13B：
FLOPs/step = 6 × 1024 × 1024 × 2.13e9 ≈ 1.34e16 = 13.4 PFLOPs🅳 (d) 单卡 A100 + 50% MFU 训 400K 步要几天
A100 FP32 峰值 = 19.5 TFLOPS，50% MFU → 有效 9.75 TFLOPS。
总 FLOPs = 13.4 PFLOPs × 400,000 = 5.35e21
时间 = 5.35e21 / 9.75e12 = 5.49e8 秒
= 6354 天
≈ 17.4 年 🤯$$\boxed{\text{训练时间} \approx 6354 \text{ 天} \approx 17.4 \text{ 年}}$$
这就是重点：一张 A100 训不了 GPT-2 XL，必须分布式。
| 卡数 | 时间 |
|---|---|
| 1 | 17.4 年 |
| 64 | 99 天 |
| 256 | 25 天 |
| 1024 | 6.2 天 |
这道题其实是下一章 Scaling Laws / 分布式 的"开胃菜" —— 让你从数字上感受到"一张卡是不够的"。
🪞 连回 Task 1 + Harness Engineering
⚠️ 又是 N=1 的直觉联想，不是方法论。
Task 1（BPE 分词器） 教的是：在有限的 token 预算里做自适应压缩。
- 常用搭配 → 合并成大 token
- 罕见搭配 → 退回字符
Task 2（资源核算） 教的是：在有限的算力/显存预算里做规划。
- 什么能放 GPU / 什么留 CPU
- 什么用 BF16 / 什么必须 FP32
- 什么跑 batch_size=8 / 什么只能 batch_size=1 靠梯度累积
再往上一层，Harness Engineering（我在搞的多 Agent 方向）本质也是预算规划：
- 什么任务合并成一个 skill 调用
- 什么任务拆分到不同 Agent
- 什么信息放 context / 什么塞 memory
🎯 三件事同构：都是在「预算约束下做结构化决策」。
这章让我第一次觉得 —— 工程师的核心能力不是"会写代码"，而是"会算账"。会算账的工程师，才敢承诺 deadline。
📋 打卡清单
- [x] Chapter 3 全章精读（1260 行）
- [x] 关键公式内化（6× FLOPs / 16× 内存）
- [x] Demo 代码跑通（3 阶段验证 16 字节公式）
- [x] 数据类型对比（FP32 / FP16 / BF16 实测）
- [x] 餐巾纸算一遍 H100/70B（144 天 vs 教材 146 天 ✅）
- [x] 作业题 2 全做完（a/b/c/d 四小题 + 代数表达式 + 数值答案）
- [x] 打卡博客发布 ← 就是这篇
- [ ] 作业题 1（Transformer FLOPs 分解）（下次补）
- [ ] Chapter 4 Transformer 架构（下次）
🌱 反思
Task 1 打卡时我感受到的是"跑通了一个算法"的小成就感。
Task 2 感受到的是另一种东西 —— "能在纸上算出答案"的安心。
在跑那个"单卡 A100 训 17 年"的计算时，我意识到一件事：很多「不行」和「能行」的区别，不在运气，在你算不算得出来。
不为证明什么，就是想每一行代码都算得清账 🔧
Winnie · 2026-04-19