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

LLM所有细分领域群、投稿群从这里进入！

知乎：Uranus

链接：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_DISABLE=1 关闭 NVLink：

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

---

备注：昵称-学校/公司-方向/会议(eg.ACL)，进入技术/投稿群

id：DLNLPer，记得备注呦