Nvlink对大模型推理的速度有多大提升？

来自：https://www.zhihu.com/question/654832546/answer/71647384740

先放结论，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 数量之间：

NVLink 3.0 vs PCIe 4.0

如果我们的输入是 4096 个 token，单次 all reduce 的延迟分别是 603 us 和 11369 us。乘上 80 （num layers）和 2（每层 2 次 all reduce）后，NVLink 3.0，PCIe 4.0 的开销则分别是 96.48ms 和 1819.04 ms。

在推理优化尝试榨干 GPU 每一份算力的时候，1819.04 ms 的通信成本是一个非常非常夸张的数字，宝贵的 GPU 计算资源在通信的时候几乎是闲置！

为验证这个数据，我用 vLLM v0.6.6 在 4xA100 的环境实测了一下，一个 4379 tokens 的输入，只跑 prefill。我们通过 NCCL_P2P_DISAB