02-step-budget/README.md3060 字节
# 示例 2：一次训练迭代 / 一次推理请求的分解账本
## 目的
示例 1 算的是**一个算子**。这个脚本把一整个 decoder-only Transformer 的
* 训练一步：前向 → 反向 → 优化器 → 通信；
* 推理：prefill（全部层）vs decode（每 token）；
各自拆成一张表，回答第 02 章 1.1、1.2、2.4 节的问题：**时间花在哪一段，为什么。**
## 运行前提
* Python 3.9 或更高（只用标准库）。
* **不需要 GPU / NPU / PyTorch / TileLang。**
## 运行命令
```bash
python3 step_budget.py # 8B 级模型，seq 4096，8 卡
python3 step_budget.py --seq-len 32768 # 长序列，看 attention 占比
python3 step_budget.py --layers 80 --hidden 8192 --ffn 28672 --params 70e9 --world 64
python3 step_budget.py --peak-tflops 989.5 --bandwidth-gbps 4800 # 换成 H200 厂商规格
```
## 期望输出（默认参数下的实际输出）
```
========================================================================================================
训练一步（seq_len=4096，batch=1）
--------------------------------------------------------------------------------------------------------
阶段 运算量 访存量 访存时间 计算时间 下界 占比
前向（全部层） 66056.6 GF 28.45 GB 6.320 ms 90.575 ms 90.575 ms 27.9%
反向（全部层） 132113.2 GF 56.91 GB 12.641 ms 181.151 ms 181.151 ms 55.9%
优化器 120.5 GF 208.78 GB 46.375 ms 165.2 us 46.375 ms 14.3%
通信 0.0 GF 28.11 GB 6.243 ms 0.0 us 6.243 ms 1.9%
========================================================================================================
推理
--------------------------------------------------------------------------------------------------------
阶段 运算量 访存量 访存时间 计算时间 下界 占比
Prefill（全部层） 66056.6 GF 28.45 GB 6.320 ms 90.575 ms 90.575 ms 96.6%
Decode（每 token） 16.1 GF 14.50 GB 3.220 ms 22.1 us 3.220 ms 3.4%
========================================================================================================
```
三个可以立刻看出来的结论：
1. 训练一步里算力主要集中在**前向 + 反向**（这里合计约 84%），
优化器虽然只有 120 GFLOP，却要搬 208.78 GB——它是纯带宽阶段。
2. prefill 是 compute-bound，decode 是 memory-bound：同一个 attention，判据相反。
3. decode 每产生一个 token 就要读 14.5 GB（其中约 14 GB 是模型权重），
这就是 batch、量化、GQA、投机解码的动机。
## 性能口径
**本示例不测量任何性能**，全部是本书口径的**粗模型推导**（写死在脚本注释里，便于读者反驳）：
* 前向/反向的 FLOPs 按每个 GEMM 的 2MNK 逐项相加，反向按「每个前向 GEMM 的反向是它的
2 倍」估算，合计约 3 倍前向 FLOPs；
* 激活访存只按「残差流 + MLP 中间张量各读写一次」粗算，不建完整的激活重算模型；
* 优化器按 AdamW：每参数读 grad/param/m/v、写 param/m/v；
* 通信按 ring all-reduce 的 2(N-1)/N 倍梯度数据量；
* 机器参数默认取示例 1 的 H200 标定值（`--peak-tflops`/`--bandwidth-gbps` 可改），
出处见 `../01-op-traffic-model/README.md`。
**结论只在「哪一段是数量级上的大头」这个层面可靠**，不要拿它当延迟预测或容量规划。
## 作者实际做到的验证级别
**CPU 逻辑验证。** 在 macOS (Apple Silicon)、Python 3.11 上运行过默认参数、
`--seq-len 32768`、放大到 70B/64 卡与换成厂商规格参数；`python3 -m py_compile step_budget.py` 通过。
没有在 GPU/NPU 上运行过任何内容。
## 已知局限
* 没有建 KV cache 的**增长**模型：decode 只看某个固定 KV 长度下的一次调用。
* 没有建 activation checkpointing / recompute、ZeRO / 张量并行、MoE 路由与专家负载不均的模型。
* 没有考虑 attention 反向比 2 倍前向更贵（需要重算 softmax 或用更多中间量换访存）。
* 前向的权重访存按「每层权重读一遍」估算，没有区分不同并行切分方式下的实际读数。