深度学习知识

news2024/12/31 7:07:12

context阶段和generation阶段的不同

context阶段（又称 Encoder）主要对输入编码，产生 CacheKV(CacheKV 实际上记录的是 Transformer 中 Attention 模块中 Key 和 Value 的值），在计算完 logits 之后会接一个Sampling 采样模块，采样出来第一个生成的 token，并将这个 token 和 CacheKV 作为 generation阶段的输入，

generation阶段（又称 Decoder）在一个 While 循环里，读取 token 和 CacheKV 后通过自回归解码的方式每循环一次产生一个 token。generation阶段，自己的输出作为自己的输入，不断回归迭代，就是「自回归」。在 While 处判断需要继续生成，在 Attention 中计算出token对应的 CacheKV 信息存储下来，并拼接上所有的历史 CacheKV 信息进行计算，最后采样出来下一个 token。While 循环检测到「eos」生成结束，就会退出循环，本次生成推理过程结束。

在 DecoderSelfAttention 中，查询的序列长度始终为 1，因此使用自定义的 fused masked multi-head attention kernel 来处理。另一方面，ContextSelfAttention 中查询的序列长度是最大输入长度，因此我们使用 cuBLAS 来利用tensor core。

context阶段sequence length等于input_length(并行，可以类比训练的前向过程causal attention)，generation阶段sequence length等于1（一个token一个token循环进行decoder）。
以 175B 的 GPT-3 模型，输入 1000 个 token，生成 250 个 token 为例，那么 Context 阶段的激活 Shape 为 [B, 1000, 12288]，其中 B 为 batch_size，第二维为输入 token 数，第三位为 hidden size。而对于 Generation 阶段，由于每次输入输出都是固定的 1 个 token，是通过循环多次来产生多个输出 token，所以 Generation 阶段的激活 Shape [B, 1, 12288]的第二维始终为 1，Generation 的激活显存占用是远小于 Context 阶段的。
再来看计算量，这里可以看到挺多有趣的结论：

---Context/Generation 的计算量均随 batch size 增大而正比增大，这个很好理解，因为 batch 内每个样本的计算量是一致的；

---Context/Generation 的访存量随 batch size 增大却基本不变，访存量可以分为权重访存和激活访存，很显然激活的访存是随 batch size 增大而增大的，但是权重访存却不会随 batch size 变化而变化，而由于激活的访存量远小于权重的访存量，就会出现总访存量几乎不随 batch size 变化而变化；

Context 是计算密集型的任务，compute bound；而 Generation 是访存密集型的任务，IO bound（显存带宽指的是 GPU 计算单元与显存之间的数据传输速度）。这是由于 Context 的计算量大， Generation 由于每次都只计算 1 个 token，所以计算量远小于 Context；但是两个阶段权重的访存量确实一致的，因为Generation 要循环调用，所以 Generation 阶段反而是访存密集型的任务。