先放结论,NVlink 对推理速度的提升非常大,非常非常大!
这可能和高赞答案的结论不一样,且听我娓娓道来 :D
最近在分析 RTX4090 上用 FP8 tensor core 跑 70B 模型的性能瓶颈,发现大模型推理过程中,很多时候通信是非常昂贵的一件事儿。NVLink,甚至 P2P,都会对推理有很大的加速效果。
以 Llama 3 70B TP4 为例,模型共有 80 层,每一层中有 2 次 all reduce 通信,数据量是 batch 中token 的数量乘以 hidden dim (8192)。
在 FP16 下,一个 token 的数据量是 16384 bytes。通信延迟会是什么样的呢?我用 nccl-tests 测了一下单次 all reduce 的延迟和 token 数量之间:
NCCL_P2P_DISABLE=1 vllm serve /mnt/resource/public_models/Meta-Llama-3-70B-Instruct-hf -tp 4
在关闭 NVLink 前,latency 是 878.57ms 左右;关闭后,prefill 的 latency 是 2740.17ms。Bang!
考虑到 NVLink 打开时通信开销很小,我们可以大致认为 compute 的时间是 800 ms 左右,那么关闭 NVLink 后,纯通信的开销大概是 1900 ms 左右,和我通过 nccl-test 测出来的 1819 ms 基本相符。
到这里,NVLink 对大模型推理的加速效果可见一斑。其实我们对 latency 也取对数的话,会发现 NVLink 和 PCIe 的 latency 随数据量是线性的,但 PCIe 的斜率明显较大,这是受带宽制约:
NVLink 3.0 vs PCIe 4.0
当输入为 4096 个 token 的时候,实验中 all reduce 传输的数据量是 N * 2 * (D - D / N),其中 D 为 64 MB,N 为 4(GPU 的数量),那么总的通信量为 384 MB。
NVLink 能够达到的带宽是 384 MB / 603 us,约为 631.58 GBps;PCIe 能够达到的带宽是 384 MB / 11369 us,约为 33.77 GBps。
这和 NVLink 3.0,PCIe 4.0 的规格也是相符的。原本应该测测 PCIe 5.0,奈何手头没有合适的机器,后面如果有机会再补充数据吧。
最后放一张带 P2P over PCIe 数据的图作为结束吧。总结一下:在 GPU 越来越快的当下,我们也需要越来越快的通信,否则宝贵的算力是发挥不出来滴!
NVLink 3.0 vs P2P vs PCIe 4.0