张量并行(tensor parallelism)是把同一层里的矩阵乘法切到多张 GPU 上,每张卡算一块,再把结果合起来。它解决的问题很具体:这一层的权重或中间结果,一张卡的显存放不下。
它不是「卡多就一定更快」。每一步都要通信。通信时间大于省下的计算时间时,延迟会变差。
和另外两种多卡方式
| 方式 | 怎么切 | 通信发生在哪 | 更适合 |
|---|---|---|---|
| 单卡 | 不切 | 无 | 权重和 KV 都放得下 |
| 张量并行 | 同一层的矩阵切开 | 几乎每一层都要聚合 | 单层放不下,而且卡间互联快 |
| 流水线并行 | 按层切开,前几层在卡 A,后几层在卡 B | 层与层交界传激活 | 模型很深,能接受气泡和更高延迟 |
流水线会有气泡:后面的卡在等前面的卡算完。小流量、在意首 token 的在线服务,通常先不要用它。离线把很长的任务跑完,再比较合适。
多张卡各放一个完整的小模型,是数据并行或多副本,不是把一个模型切开。那是增加能同时服务的请求数,不增加单个模型的容量。
什么时候不该上张量并行
先算单卡够不够,而不是先拆。权重约等于参数量乘每参数字节数;KV cache 再按批大小和序列长度加,见 KV Cache。两块都放得下,就不要拆。
下面几种情况,拆开常常更慢或更难管:
- 模型量化后已经能放进一张卡。量化见 量化入门。
- 并发很低,算力本来就没用满,通信却每次都要付。
- 卡和卡之间只有普通 PCIe,没有 NVLink 这类更快的互联。聚合一次的时间会吃掉并行省下的计算。
- 没有人能看懂「某一张卡上的通信挂住了」。单卡 OOM 的日志在一张卡上;切开之后要同时看多张卡和通信库。
什么时候才值得
一张卡放不下权重,量化之后仍然放不下,而且目标上下文下的 KV 也要算进去。这时拆是容量问题,不是速度偏好。互联快、批量又足够大,计算才摊得过通信。
驱动和 CUDA、通信库的版本要锁在一起。一张卡掉了,整个并行组都不可用,健康检查得覆盖整组,不能只看其中一个进程还在不在。
和调度的关系
多卡不改变分池:在线服务和实验任务分开,实验可以排队和超时,在线那一组不要被抢。分池见 GPU 调度实战。
小结
张量并行是把一层的矩阵乘拆开,用来解决单卡放不下。放得下就用单卡或量化。放不下、互联又快,再拆。先看显存账,再看通信,不要把「多卡」本身当成优化。