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

教育   2025-01-15 17:59   江苏  

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,记得备注呦

深度学习自然语言处理
一个热衷于深度学习与NLP前沿技术的平台,期待在知识的殿堂与你相遇~
 最新文章