How To Scale Your Model 笔记(1-4)

模型的内存需求和执行时间取决于参数、激活、数据搬运及算子之间的依赖关系。性能分析需要分别计算容量、FLOPs 和各条路径上的流量，再根据数据复用、切分和调度方式估算耗时。

1. 先分清：装得下、搬得动、算得完

讨论速度前，至少有三种不同的约束：容量用 bytes 衡量，带宽用 bytes/s 衡量，计算吞吐用 FLOPs/s 衡量。参数量只是可学习标量的个数，乘上每个元素的存储字节数，才变成权重容量。

在一个简化的 TPU 图像里，大量数据保存在 HBM，分块搬入较小的片上存储 VMEM，再由矩阵计算单元 MXU 或向量计算单元 VPU 使用。两类计算单元按运算类型分工，不是所有数据都依次经过两者。整个矩阵也不必一次装入 VMEM，关键是当前工作集能否放下，以及后面还会不会重读。

用 \(\mathcal F\) 表示一次任务的计算量，\(C\) 表示适用于该精度和运算类型的计算峰值；用 \(Q\) 表示经过某条接口的累计字节数，\(\beta\) 表示该接口的带宽上限。则资源给出的时间下限是：

这个最大值首先是一个下限。即使计算与搬运不能重叠，下限仍然成立；要接近它，则需要合适的数据依赖、调度、利用率和重叠程度。真实运行中还有启动、同步、其他算子等成本。原书第 1 章

算术强度：搬来的数据做了多少工作？

定义算术强度 \(I=\mathcal F/Q\)，单位为 FLOPs/byte。计算与搬运时间相等时，\(I=C/\beta\)。因此吞吐量的理想上限可以写成：

低强度时，数据供给限制了计算吞吐，上限随 \(I\) 增加；强度足够高后，算力形成水平的“屋顶”。比较实测性能时，应该同时看芯片峰值与该任务对应的屋顶。例如芯片峰值为 6 GFLOPs/s，任务的带宽屋顶为 4 GFLOPs/s，实测 3 GFLOPs/s：这既是峰值的 50%，也是任务屋顶的 75%。仅凭前一个百分数，无法断言还有一倍加速空间。

矩阵乘法最能体现复用的作用。设：

每个输出需要一个长度为 \(k\) 的点积，总计算量约为 \(2mkn\)。若三个矩阵均以 bf16 存储，每元素 2 bytes；输入各从 HBM 读一次，输出写一次，中间量保留在片上，则：

当 \(m\ll k,n\)，权重读取占主导，得到 \(I\approx m\)。它的含义是：同一份权重服务更多输入行，每行分摊的权重搬运量下降。从一次一行计算变成一次八行计算，权重读取量可以不变，但总流量并没有自动减少，新增输入和输出仍要搬。 如果与八次各自重新加载权重的一行计算相比，才省掉了七次权重读取。

这个近似不能无限延伸。若 \(m=k=n=1024\)，强度约为 \(1024/3\)，不是 1024；若固定 \(k=n=1024\)，只让 \(m\) 增大，则：

权重的固定成本可以不断分摊，但每增加一行，仍要读取 \(k\) 个输入、写回 \(n\) 个输出。这个极限把复用的收益与边界一起显示出来。

2. 复用与流水线：容量和时间顺序开始重要

上面的“输入各读一次”是一种理想数据流。真实芯片片上空间有限，同一个权重块再次被需要时，可能已经被替换了。原书第 2 章讨论的存储层次，正是把公式落实到访问顺序的起点。

例如，8 MB 权重分成两个 4 MB 块 \(W_1,W_2\)，片上留给权重的空间只够一块。四组输入都要用这两个块：按输入组处理，每组依次装入两个权重块，会读取 32 MB；先让所有输入用完 \(W_1\)，再统一换成 \(W_2\)，权重读取量只有 8 MB。

后一个方案可能增加输入重读，是否更快要看完整预算。输出也不能只数“写了几次”：若 \(W=[W_1\mid W_2]\) 是按输出列切分，那么 \(Y=[XW_1\mid XW_2]\)。左右两部分分两次写回，仍是一份输出，并非写了两份完整矩阵。

同一块先搬后算，不同块可以同时进行

单个块存在严格依赖，必须先搬到再计算；但计算第一个块时，可以准备第二个块。双缓冲让两个单元分别使用不同的缓冲区，随后交换角色。

设共有 \(q\) 块，每块搬运耗时 \(a\)、计算耗时 \(c\)。假设两阶段各自串行、彼此可并行，缓冲足够，暂不计输出写回和额外开销。总时间由较慢阶段的连续工作时间，加上启动或收尾时间组成：

搬运慢，就连续搬完所有块，末尾补一次计算；计算慢，就先等第一块到达，再连续算完所有块。因此：

例如全任务纯搬运 0.4 ms、纯计算 0.2 ms，均匀分为 \(q\) 块，上述模型给出 \(t=0.4+0.2/q\) ms。分两块时是 0.5 ms；忽略启动收尾才取 0.4 ms。实际分块也有调度和效率成本，不能由这个模型推出“切得越细越好”。

3. 多芯片：先问缺什么，再选通信操作

把模型切到多块设备后，增加了一类数据路径。TPU 的 ICI 是芯片间互联，PCIe 连接主机与加速器，DCN 是主机网络；每种路径有自己的带宽和延迟。芯片怎样连接、消息要经过几跳、哪些传输能同时进行，都会改变结果。逻辑 mesh 换一个名字，并不会改变物理链路。

但在估通信时间之前，更基础的问题是：每个设备现在算出的东西，到底是什么？ 第三章的主线可以从一个输出元素开始：

若按输出行或列分工，一个设备可以负责某些输出位置的完整求和。若沿 \(j\) 分工，各设备则负责同一输出位置的不同求和项。两者都能得到一个矩阵，但含义完全不同。

以两设备为例，沿归约维度切分：

每端的乘积都已经具有整个 \(C\) 的形状，却只包含部分贡献。把它们拼接会错，逐元素求和才对。相反，若沿 \(B\) 的输出列切分，\(AB_0\) 和 \(AB_1\) 是不同输出区域，应按列拼接。

还有一种更隐蔽的错误：本地两个矩阵的内维长度相等，但分别对应全局 \(j\) 的不同区间。这样的乘法在 shape 上合法，数学上却配错了元素。检查布局时，需要追踪全局索引，而不只看维度数字。

四种 collective 对应四种数据需求

“完整”也有两层含义：一块数据可以只覆盖部分索引，却已经数值完整；也可以覆盖所有索引，却每个数都只是部分和。这个区别决定了是否能进入后续非线性操作。通常：

通信公式里的体积，到底是谁的体积？

这里用 \(V_{\rm comm}\) 表示通信公式中的数组体积，避免与 attention 的 Value 矩阵混淆。对于一个有 \(n\) 个设备的通信组：

在 AllReduce 中，不能因为有四个设备，就把每端 2 MiB 的完整形状部分和改称为 8 MiB 的结果。它们是同一组索引的四份贡献。其他 mesh 轴还可能切着数据，因此“通信组内拼齐”也不一定等于整个全局张量。

看单向环更容易推导成本。AllGather 从每端一块开始，经过 \(n-1\) 轮，每端每轮发送 \(V_{\rm comm}/n\)。若单方向链路带宽为 \(b\)，忽略延迟：

同样大小数组的环形 ReduceScatter 有相似传输体积；AllReduce 可以分成 ReduceScatter 再 AllGather，因此有两阶段成本。原书在大消息、双向环的近似下，把可用双向带宽记为 \(W\)，得到 \(t_{\rm AG}\approx V_{\rm comm}/W\)、\(t_{\rm AR}\approx2V_{\rm comm}/W\)。有限设备数、启动延迟、拓扑与带宽定义都需要保留，不能将单向式和双向式混用。

ReduceScatter 的环上过程也不是“每端只管自己最后留下的那块”。每端对各输出块都有贡献；块在环上传递时，与沿途设备的对应贡献相加。最后每端留下一个完整求和块，AllGather 再把这些块拼齐。

前后两层一起设计，通信才可能真正减少

设两层为 \(H=\sigma(AW_{\rm in})\)、\(C=HW_{\rm out}\)。将 \(W_{\rm in}\) 按输出列切分，得到数值完整的 \(H\) 列块，可直接计算非线性。再把 \(W_{\rm out}\) 按对应输入行切分，各端计算 \(H_rW_{\mathrm{out},r}\)，最后求和得到 \(C\)。

这样，放大的中间表征 \(H\) 可以保持切分，不必为了进入下一层先 AllGather。它把“第一层的输出布局”变成了“第二层恰好需要的输入布局”。是否最终 AllReduce，还是只做 ReduceScatter，则继续取决于下一步需要什么输出布局。

4. 回到 Transformer：从功能恢复矩阵

Transformer 的计算与存储预算需要区分权重、激活和缓存，并确定每个算子处理哪些维度。先列出各层的数据流，再逐项计算。

以下采用同一套记号：

默认讨论无 bias 的主要矩阵、gated MLP、因果自注意力，且 \(N_QH=D\)，各类头宽相同。归一化参数和逐元素开销不混入主要矩阵预算。

从模型入口到出口，主表征的宽度是 \(D\)：

token ID → embedding 查表 → [B,T,D]
    → attention 子层：读取允许位置的内容，回到 [B,T,D]
    → MLP 子层：逐位置变换特征，回到 [B,T,D]
    → 重复 L 层 → 最终归一化
    → 词表投影 [D,V_vocab] → logits [B,T,V_vocab]

残差连接将每个子层写成 \(x+f(x)\)，因此分支输出要与输入同形状。常见 pre-norm 结构先归一化再进入分支。LayerNorm 沿每个 token 的 \(D\) 个特征统计，不跨不同 token；若特征维切到多设备，需要合并统计量，若每个 token 的完整特征在本地，则无需为此通信。

参数是规则，激活是当前输入经过规则后得到的状态。 训练完成后，\(W_Q\) 固定；\(Q=XW_Q\) 则随输入变化。Embedding 表也属于权重：输入端查表得到表征，输出端做词表打分。若两端共享同一参数表，只计 \(DV_{\rm vocab}\) 个参数，而且不乘层数。

Gated MLP：两条支路都读取完整输入

普通 MLP 是 \(\sigma(XW_{\rm in})W_{\rm out}\)。Gated MLP 多了一条由输入生成的调节支路：

\(W_u,W_g\) 都是 \([D,F]\)，\(W_{\rm out}\) 是 \([F,D]\)。两条支路都接收完整的 \(D\) 维输入；得到的 \(U,G\) 都是 \([B,T,F]\)，逐元素相乘后仍是这个形状，而不是拼成两倍宽度。

可以把 \(U\) 理解为待使用的特征，\(G\) 为当前输入生成的调节量。它不必处于 0 到 1 之间，也不要求每个通道都有明确的人类语义。GLU 变体论文

这里有一个值得单独澄清的线性代数问题：不同通道有不同的固定增益，仍然是线性变换。 例如 \(f(x_1,x_2)=(2x_1,3x_2)\)，满足叠加原理。固定向量 \(g\) 的逐元素缩放，也可以写成 \(U\operatorname{diag}(g)\)。门控的关键在于 \(G\) 随输入变化；训练后权重固定，不代表门控激活固定。

因此，gated MLP 有三份主要权重，共 \(3DF\) 个参数；每个 token 依次承担三次投影，主要计算量为 \(6DF\)。这两个因子“3”和“6”分别来自矩阵份数与乘加计数，不能与 bf16 的 2 bytes 混在一起。

Attention：匹配条件、来源描述与被读取的内容

对一个头，先从输入生成 Query、Key、Value，再执行：

Query 提供当前读取者的匹配条件，Key 提供每个来源用于匹配的描述，Value 提供要传递的内容。\(W_Q\) 学习的是从当前输入生成 Query 的规则，不是保存一个固定查询；投影到学习方向的图像有助于理解，但这些方向不要求正交，也没有固定频率含义。

对于一个 Query，attention 输出通常是多个来源的加权组合，例如 \(o=0.1v_1+0.7v_2+0.2v_3\)。匹配分数为零不等于权重为零，因为 \(e^0=1\)；所有允许位置分数相同，会均匀读取。被 mask 的位置则在 softmax 前设为负无穷，使其权重为零。

各头先独立得到 \(H\) 维内容，拼接为 \(N_QH=D\)，再通过 \(W_O\) 组合成 \(D\) 维表征。这里仍是模型内部的表征，并非词表上的概率。 只有最后经过词表投影，维度才变成 \(V_{\rm vocab}\)。

5. 只有一个新 token，仍然可以读取一整段历史

这是把 attention 图像连接到推理预算时最值得展开的一步。单步生成意味着每请求只有一个新 Query 位置，即 \(T=1\)；它没有说来源位置也只有一个。

固定权重的标准因果推理中，追加未来位置不会改变旧位置的结果，所以可以在各层保留旧 Key/Value。新 token 生成自己的 Q/K/V，将新 K/V 追加到缓存，再用新 Q 读取整个允许的来源集合。历史的投影与 MLP 不必重跑，但历史 K/V 仍参与新一步 attention。

单请求、单头的 shape 链最直接：

来源增加，使求和项变多；它不使输出向量变宽。每个历史位置贡献固定大小的 K/V 行，缓存沿位置轴增长。

GQA 共享的是 K/V，不是 Query 的读取结果

标准 Multi-Head Attention（MHA）中，每个 Query 头对应自己的 K/V 头。Grouped-Query Attention（GQA）让一组 Query 头共享指定的一组 K/V；Multi-Query Attention（MQA）则让全部 Query 头共享一组 K/V。这是架构层面对表示自由度、投影成本与缓存成本的取舍，并非无条件删掉现有模型的头。GQA 论文

以 \(N_Q=32,N_{KV}=8,H=128\) 为例，每四个 Query 头共享一组 K/V。各 Query 可以不同，所以它们对同一 K 的匹配权重和最终输出仍可以不同。Q/K 点积要求的是每头宽度匹配，不要求两类头的数量相同。

对应的投影权重也由功能决定：

把多个头合并成一次大矩阵投影，是计算组织方式；进入匹配时，还要恢复每个头的对应关系。完整激活 shape 为：

也不要丢掉 batch 轴。四条独立请求，各有自己的历史；它们共享模型权重，但不默认共享 K/V 内容。把四条历史无隔离地拼在一起，会允许请求之间相互读取信息，改变模型的语义。

假设已有 8192 个缓存位置，四条请求各处理一个新 token，可以逐步检查：

缓存容量怎样数

若每个 KV 元素用 \(p_{kv}\) bytes，\(L\) 层、\(B\) 条等长请求、每条已有 \(S\) 个来源位置，则：

第一个 2 来自 K 和 V 两份数据；bf16 的 2 bytes 是另一个因子。每条请求各增加一个位置，缓存增量为 \(2LBN_{KV}H p_{kv}\)。请求长度不同时，用长度之和替代 \(BS\)。实际分配还需要考虑预留、padding 和碎片。

标准 MHA 可用 \(N_{KV}H=D\) 化简，GQA 一般不可以。MLP 的中间宽度 \(F\) 则根本不决定 KV 行宽——它属于模型中另一条数据路径。

6. 内存预算看生命周期，优化看改变了什么

算出每个张量大小以后，还不能把所有出现过的张量相加当作峰值。正确的问题是：某个时刻，哪些数据必须同时存在？

纯前向串行执行时，许多中间激活用完即可释放或被覆盖；残差输入需保留到加法；KV cache 要跨生成步骤保留。权重和各层 KV 有层数因子，不意味着临时激活峰值也简单乘层数。

例如一个 gated MLP，残差输入占 64 MiB，两条中间支路各 192 MiB。两支路都就绪时占 448 MiB；若门控乘积覆盖其中一条支路，释放另一条后剩 256 MiB。分配 64 MiB 输出后是 320 MiB。仅这几步便能看出，峰值取决于存活集合和原地操作，而不是最后输出多大。

训练为什么保存激活，又为什么可以重算？

对于 \(Y=XW\)，令上游梯度为 \(G=\partial\mathcal L/\partial Y\)。

这两个梯度的求和轴由数据的依赖关系决定。每个权重被很多 token 使用，所以 \(dW\) 要沿 token 轴求和；单 token 对权重的梯度是输入与输出梯度的外积。每项输入又影响多个输出特征，所以 \(dX\) 要沿输出特征求和。结果 shape 必须与求导对象相同；全是方阵时还要看求和指标，不能只靠 shape 猜。

这也解释了为什么反向需要前向输入 \(X\)。可以把它一直保存，也可以保存更早的状态，轮到反向时重新计算。按 block checkpoint，长期保留块输入，反向时重建当前块内部激活，便是用额外计算换较少的长期存储。

对主要矩阵乘法，前向和两个梯度各约 \(2MDF\) FLOPs。若完整重跑一次前向，前向加反向的预算从三个单位变为四个单位。但实际内存峰值仍需加上当前重算中间量、梯度和工作区；运行时间也不保证正好增加三分之一。原书反向与 checkpointing

FlashAttention：不保存完整方阵，仍然计算位置配对

逻辑上存在 \(T\times S\) 个 attention 分数，不代表必须把整张表写进 HBM。对单 Query，输出可写成：

在共同指数尺度下，对两个来源块分别累计分子向量 \(n_1,n_2\) 和分母标量 \(z_1,z_2\)，最后合并为 \((n_1+n_2)/(z_1+z_2)\)。不需要存下全部最终概率；但不能把各块分别归一化后的输出直接相加。实际稳定算法还需维护最大分数，并将不同块的累积量重缩放到同一尺度。

FlashAttention利用分块与融合减少中间结果的 HBM 读写。标准 dense attention 仍要计算允许的位置配对；无历史且 \(T=S\) 时，主要矩阵乘法仍随 \(T^2\) 增长。减少存储与搬运，不等于取消对应的数学运算。

MoE 则从另一方向改变预算：模型拥有多个专家，但每个 token 只调用其中几个。若有 \(E\) 个专家、每个 \(P_e\) 参数、每 token 选 \(k\) 个，则专家池参数为 \(EP_e\)，每 token 使用的是 \(kP_e\)。前者决定完整专家池的权重容量，后者影响主要专家 FLOPs；还要另计路由、负载不均与通信。这再次说明，“拥有的数据”“本次使用的数据”和“累计搬运的数据”不能混成一个数。

7. 单步生成的计算与存储预算

先把每层主要账目汇总。以下 K/V 投影只处理本次新增的 \(T\) 个位置，历史已经缓存；匹配部分按完整 \(T\times S\) 矩形计数。

因此：

一个容易被“每 token 约两倍参数量 FLOPs”掩盖的部分，是不含新增权重的 \(QK^\mathsf T\) 和 \(AV\)。没有历史、\(T=S\) 时，它们与层内权重矩阵乘法的计算量之比为：

Batch 在比值中约掉，但两项绝对计算量都随 batch 增长。令 gated MLP 的 \(F=8D/3\)、\(N_{KV}=N_Q/4\)，有 \(P_{\rm layer}=10.5D^2\)，两项在 \(T=5.25D\) 附近相等。这个分界依赖架构，也依赖是否按完整方阵计数。因果内核若确实跳过上三角，需改用实际允许的位置对数，不能机械套同一个常数。

四条请求，各向前生成一步

采用一个假想模型：32 层，\(D=4096\)、\(F=11008\)、32 个 Query 头、8 个 KV 头、每头宽度 128、词表 32000。Gated MLP，共享输入/输出词表权重，权重和 KV 均为 bf16。忽略 norm、bias 等较小参数及逐元素开销，不做设备切分。

四条独立请求各已有 8192 个缓存位置，本次各处理一个新 token，来源数变为 8193。首先检查容量：

这里 B 参数用十进制，GiB 用 \(2^{30}\) bytes，MiB 用 \(2^{20}\) bytes。权重加当前 KV 约 14.807 GiB，还没有包含临时激活、工作区或其他运行开销。

再看计算量，明确统计的是整个 batch 的一步：

各项的计数尺度不同：MLP 和投影按本次新 token 数计；attention 按 Query 数与来源长度计；KV 容量按请求、层、历史位置和 KV 头计；词表投影只出现在模型出口。

预算怎样变成性能问题

到这里仍不能只拿 63.596 GFLOPs 除以一个算力，就声称得到生成延迟。还得回到前三章，选定执行方案：

• 权重留在哪一级存储？同一 batch 内能否复用，下一步是否要重读？
• KV 在各 Query 头之间怎样共享读取？实际 HBM 流量是多少？
• 哪些张量切分，哪条链路传什么，最终输出布局是什么？
• 哪些阶段能重叠，哪些必须等待完整数值结果？

例如，仅作一阶教学估算：若这些权重在 HBM 中，每一步各读一次，且理想复用让每份历史 KV 也只读一次，权重与历史 KV 两项约产生 15.899 GB 读流量。假设相应 HBM 带宽为 1 TB/s，这两项给出约 15.9 ms 的搬运下限；若适用计算峰值为 100 TFLOPs/s，主要矩阵计算下限约为 0.636 ms。这个模型提示应优先审视数据供给。

它不是实测结论：片上驻留可以减少这类 HBM 读取，重复加载、临时量和其他算子又可能增加流量，分布式执行还会改变每设备预算与通信。真正的下一步是测量并解释模型与观测之间的差异。

本文主要依据 How to Scale Your Model 第 1 章、第 2 章、第 3 章、第 4 章，以及文中链接的原始论文。更基础的架构介绍见本站《理解 Transformer（1）》。