2278 字
约 7 分钟
6
S23-J08 · 计算图与前向执行

适用读者:知道"推理就是前向"、但说不清"时间花在哪些运算上、哪些运算能并行"的人。
前置:S23-J06、J07(四类资源与单步估算)。
覆盖旧文:A-4335 §2.1(计算图与前向执行,132 点)→ 132 点。
扩展点:E-S23-018(依赖图决定可并行性:RNN vs 因果 Transformer)、E-S23-019(按矩阵尺寸算参数量与运算量的三步法)。
上一篇:S23-J07 · 用几个数字估算一次模型执行
下一篇:S23-J09 · 自回归生成与上下文状态。


1. 一句话结论

模型执行的时间由三件事决定:每个节点的运算量、每条边传的数据量、以及边上箭头的方向(谁必须等谁)。 前两项是 J06 的四类资源,第三项是本节新增的依赖约束——它决定了"加卡能不能变快"。

出处:AI Infra 教材 §2.1(id 4335)。

节点  = 一次运算(矩阵乘、激活、归一化……)→ 决定 F
边    = 张量(要传给下一步的数据)→ 决定 R
箭头  = 依赖方向 → 决定哪些运算能同时开始

判据:看到任何"为什么这里不能并行/为什么加卡没用"的问题,先画依赖图,再谈优化。


2. 依赖图:RNN 与因果 Transformer 的差别(E-S23-018)

考虑 4 个 token、3 层网络,横轴是位置、纵轴是层。

2.1 RNN:同层必须顺序

h_t = f(x_t, h_{t-1})

当前 token 的状态依赖前一个 token 的状态,所以同一层按位置顺序执行:即使 4 个输入全部到达,第 t 个 token 仍要等第 t−1 个。

层3:  t1 → t2 → t3 → t4
层2:  t1 → t2 → t3 → t4
层1:  t1 → t2 → t3 → t4
最长依赖链 = 3 层 × 4 位置 = 12 步(并行资源再多也只能按这条链走)

2.2 因果 Transformer:同层可并行,跨层要等

第 l 层计算某位置时,通过注意力读取上一层中该位置及其之前各位置的表示。因果掩码规定"只能看自己和前面",掩码呈三角形。

关键差别:同一层不同位置之间没有依赖
层3:  t1  t2  t3  t4   ← 上一层全部完成后,四个位置可同时算
层2:  t1  t2  t3  t4
层1:  t1  t2  t3  t4
最长依赖链 = 3 步(层数),不是 12 步
结构 同层依赖 已知输入的最长链 加算力能缩短什么
RNN 有(逐位置) 层数 × 位置数 只能缩短每步的耗时,不能缩短步数
因果 Transformer(处理已知输入) 无 层数 可以并行铺开,缩短整体时间

2.3 生成时多出第三条依赖

处理已知输入(prefill)与逐步生成(decode)是两类不同工作:

prefill:已知输入 → 同层可并行 → 沿"层数"推进
decode :第 n 个输出确定后,才能作为第 n+1 次的输入 → 跨调用串行

原文补充:现代因果语言模型多为 decoder-only,本章先用它建立基础;新一代模型(原文举 CED 架构)会重新划分"输入处理"与"生成"的职责,针对"输入多、输出少"的 Agent 负载重新分配算力。

判据:prefill 是"能并行"的负载,decode 是"必须串行"的负载。 这两者的资源需求不同——这就是 S23-J03(训推分离)与后续 PD 分离(J24)的物理根因。


3. 张量尺寸:一层算多少(E-S23-019)

一次投影的定义:

输入 X ∈ R^{m×k}       m = 本次参与计算的 token 数,k = 输入特征宽度
权重 W ∈ R^{k×n}       n = 输出宽度
输出 Y = XW ∈ R^{m×n}

参数量 N_W = k·n
运算量 F    = 2·m·k·n
权重字节 M_W = b_W·k·n
变化 参数量 运算量 对应现实
m 翻倍(多处理输入) 不变 翻倍 同一模型处理更多 token / 更大 batch
n 翻倍(层变宽) 翻倍 翻倍 换成更大的模型

判据:"多处理输入"只涨计算,"层变宽"同时涨权重和计算。 这两件事在容量评估里必须分开算。


4. 例:Qwen3-8B 的一层(原文例题 2-1)

配置:L=36 层,隐藏维度 d=4096,FFN 中间维度 f=12288,32 个查询头 / 8 个 KV 头,每头 128 维,词表 151,936,总参数约 81.9 亿。

配置项 值 对资源需求的直接影响
层数 L 36 重复计算、逐层状态、深度依赖
隐藏维度 d 4096 子层输入输出与投影尺寸
FFN 中间维度 f 12288 三个 FFN 矩阵的权重与计算
Q / KV 头数 32 / 8 查询向量与上下文表示的宽度(GQA)
头维度 d_h 128 每个头的点积与状态宽度
词表 V 151936 嵌入与输出头规模

一层参数量(查询/输出投影宽度 4096;键/值投影宽度 = 8 头 × 128 = 1024):

N_proj = 2×4096² + 2×4096×1024 = 41,943,040
N_FFN  = 3×4096×12288          = 150,994,944

一层运算量(对本次输入的 m 个 token):

F_linear = 2·m·(N_proj + N_FFN) = 385,875,968 · m

这里没有用到上下文长度——投影对每个 token 做相同的变换,上下文的影响在注意力(下一节)。这就是"增加 token 数"和"读更长上下文"代价不同的原因。

4.1 两种输入形状(同一模型,差别巨大)

情形 参数 输入矩阵行数
一次处理 8192 个新 token(prefill) B=1, S=0, P=8192 m = 8192
已有 8192 上下文、再进 1 个 token(decode) B=1, S=8192, P=1 m = 1

判据:同一个模型,m=8192 是"大矩阵乘"(计算密集),m=1 是"矩阵乘向量"(读取密集)。这就是 J07 里"读 70 GB 权重"与"算 140 GFLOPs"比例悬殊的根源。


5. 符号表(后续各章通用)

符号 含义
B, P, S 请求数、本次每请求输入 token 数、已有上下文长度
m = B·P 本次参与投影的 token 总数(矩阵行数);decode 时 P=1 → m=B
d, f 主干隐藏宽度、FFN 中间宽度
n_Q, n_KV, d_h 查询头数、KV 头数、每头维度
N_pair 每层所有请求合计的有效因果查询—键配对数(不含头数)
t_e 分派给专家 e 的 token 数;不丢弃时 Σ t_e = m·k_top
m_out 实际执行词表投影的 token 数(生成下一 token 时通常为 B)

判据:m 与 m_out 不是一回事。 前者是参与投影的 token 数,后者是真正要算 logits 的 token 数——算词表头成本时别用错。


6. 前提与边界

结论 成立条件 不成立时
因果 Transformer 同层可并行 输入已全部已知(prefill) 逐步生成时仍串行
一层运算量 2mkn 稠密矩阵乘 稀疏/专家(只算被选中的专家)
投影与上下文长度无关 只看线性层 注意力 N_pair 随 S 增长(J09)
加算力能缩短等待 等待来自"可并行的计算" 等待来自依赖链 → 加卡无用(J18 再谈)

7. 常见误区

误区 纠正
"Transformer 天生并行,所以很快" prefill 可并行,decode 是串行链
把 m 与 m_out 混用 词表头成本只与 m_out 有关
以为参数量只由层数决定 宽度、词表、FFN 中间维度都算
忽略 GQA:KV 宽度 ≠ 查询宽度 8 KV 头 → KV 投影宽度 1024,不是 4096

8. 自测

  1. 4 token、3 层:RNN 与因果 Transformer 处理全部输入的最长依赖链各是多少?
  2. Qwen3-8B 一层的 N_proj + N_FFN 是多少?m=1 时该层运算量多少 FLOPs?
  3. 为什么"处理 8192 token"与"已有 8192 上下文再进 1 token"的矩阵行数差 8192 倍?
  4. 生成阶段为什么不能靠加卡把延迟降到 1/8?
答案
  1. RNN 12 步;因果 Transformer 3 步(层数)。
  2. 41,943,040 + 150,994,944 = 192,937,984 参数;F = 2×(192,937,984)×1 ≈ 3.86×10⁸ FLOPs。
  3. prefill 的 m=P=8192;decode 的 m=B=1(每条请求只进 1 token)。
  4. 因为 decode 有跨调用串行依赖,加卡只能并行"不同请求",不能并行"同一请求的下一步"。

9. 下一篇

S23-J09 · 自回归生成与上下文状态 —— 注意力为什么随上下文长度增长、KV 缓存到底占多少显存、GQA 省下了什么。

S23-J08 · 计算图与前向执行
https://www.clxhxhhr.top/posts/5535/
作者
clxstart
发布于
2024-12-31
许可协议
CC BY-NC-SA 4.0
评论
0 条
还没有评论,先写一条吧。