一份几十万 Token 的输入放进模型,权重没有变,激活和注意力工作量却会随着序列变长。上下文并行沿序列长度切分任务。每张 GPU 只持有一段 Token 的输入与激活,共同完成同一条长序列。
我对照了 Megatron Core 的说明、Ring Attention 论文和一篇百万 Token 推理研究。三份材料都指出同一个限制。序列可以分开存,注意力计算仍要让每个查询接触整条序列里相关的键和值。
线性层能各算各的,注意力需要通信
Transformer 里的线性层、归一化和逐元素运算不需要查看别的 Token。上下文切开后,各张 GPU 可以直接处理自己那一段。自注意力不同。某一段里的查询要和其他段的键值相乘,设备必须交换数据,或者把相同的信息换一种顺序传过去。
Megatron Core 会切分网络输入和全部激活。若上下文并行度为四,每张卡承担的序列激活大致降到四分之一,长序列导致的显存压力随之减轻。模型权重仍在这四张卡上各保留一份,因为上下文并行切的是序列维度。它和张量并行解决的对象不同,二者可以组合。
文档把集群规模写成张量并行、上下文并行、流水线并行与数据并行四个维度的乘积。这个公式很实用。八张卡开了两路张量并行和两路上下文并行,只剩两份数据副本,不能再把八张卡全部算成并发副本。
交换量还会受注意力结构影响。多查询注意力和分组查询注意力本来就减少了键值头,传递的 KV 也可能更少。上下文并行度、注意力头配置和通信实现应放在同一份测试记录里。
Ring Attention 给出了一种交换办法。设备把键值块沿环传递,每收到一块,就计算本地查询与这块键值的注意力,同时继续传下一块。论文的目标是把通信藏在分块计算后面。它仍计算精确注意力,没有删掉远处 Token,也没有把长输入压成摘要。
能装下长文档,还要看多久出首字
上下文并行最直接的收益是容量。设备数量增加后,单条序列可以放得更长。速度取决于每块计算能否盖住通信,以及网络拓扑是否适合所选交换方式。2024 年的推理研究用不同的环形传递方案处理完整预填充、已有 KV 缓存的预填充与解码,说明同一个并行度也要按阶段选实现。
压测时应先固定单条输入,逐步增加长度,再比较不同上下文并行度下的峰值显存和首词元延迟。随后增加并发,看为单条长请求占用的卡会不会压低系统吞吐。若真实请求大多只有几千 Token,复制模型服务更多用户可能更合算。只有长输入确实碰到容量或预填充时间限制,这项并行才有明确对象。