vLLM能够将一个大模型部署到多张卡上,主要是因为它内置并优化了张量并行(Tensor Parallelism)这一核心分布式推理技术。
简单来说,张量并行的思想就是把一个大模型“切碎”,把不同的“碎片”放到不同的GPU上,让它们像一台机器一样协同工作。
它的工作原理可以这样理解:
💡 核心机制:张量并行
当你在vLLM的命令行或代码中设置tensor_parallel_size=2(或者更大的值) 时,它会这样工作:
将模型权重按层切分:对于模型内部巨大的权重矩阵(比如一个大的全连接层),vLLM会将其按行或按列进行切分,分别存储在不同的GPU上。例如,一个原本在一张卡上需要40GB显存的模型,通过TP=2可能每张卡只需加载20GB左右的权重。
每张卡计算一部分:在模型进行前向推理时,输入的张量也会被切分,分别送入不同GPU。每张显卡只负责计算自己持有的那一部分模型参数,大大降低了单卡的计算压力。
高效的通信和同步:虽然每块卡只算一部分,但模型的最终输出需要整合。vLLM依赖高性能的通信库(如NCCL),并在代码中精细地管理了计算过程中的通信(如
all-reduce操作),确保不同GPU上的计算结果能够正确同步,得出最终结果。这也是为什么它建议GPU间最好有NVLink这样的高速互联。
🔄 两种“多卡”模式的区别
vLLM支持不同的多卡并行模式,你可以通过下表和之前提到的“模型并行”对比来理解:
| 模式 | 核心策略 | 模型权重 | 主要目的 | 适用场景 |
|---|---|---|---|---|
| 张量并行 (TP) | 一个模型,切成碎片,放到多张卡上协同计算 | 被切分(Sharded) | 降低单卡显存,让大模型能“装”下,并可获得近线性的加速 | 模型单卡放不下时,如LLaMA 70B及以上模型 |
| 数据并行 (DP) | 多个模型副本,每个GPU上跑一个完整的模型 | 被复制(Replicated) | 处理海量并发请求,通过并行处理不同的请求来提升吞吐量 | 模型单卡能放下,但需要支撑极高并发量时 |
| 流水线并行 (PP) | 将模型的不同层(如第1-10层、第11-20层)分配到不同GPU上 | 按层切分 | 主要用于模型大到连一个节点都放不下的多节点部署场景 | 超大规模模型,如DeepSeek R1等 |
结合你之前的情况,你希望在双卡上运行一个需要40GB显存的模型,而3090单卡只有24GB。vLLM的张量并行(TP)比你自己手动实现模型并行要省事得多,因为它已经在框架层面替你做好了。