多GPU部署大模型怎么配置张量并行

约 9 分钟读完 9 阅
文章导读
多 GPU 部署大模型配置张量并行，先要判断三件事：模型单卡是否确实放不下，要不要把同一层的权重切开；参与张量并行的 GPU 是否在同一台机器、互联是否足够快；模型的注意力头数、隐藏维度、词表大小能不能被你要设置的并行度整除。这三件事里只要有一件不成立，改并行度往往只是把报错换了个地方，而不是把问题解决。
📋 目录
  1. Ⅰ 先确认当前场景和目标
  2. Ⅱ 按这个顺序处理
  3. Ⅲ 改完后怎么验证
  4. Ⅳ 容易忽略的边界
A A

多 GPU 部署大模型配置张量并行，先要判断三件事：模型单卡是否确实放不下，要不要把同一层的权重切开；参与张量并行的 GPU 是否在同一台机器、互联是否足够快；模型的注意力头数、隐藏维度、词表大小能不能被你要设置的并行度整除。这三件事里只要有一件不成立，改并行度往往只是把报错换了个地方，而不是把问题解决。

张量并行适合“单卡显存装不下、且多卡在同一节点内、模型维度可被并行度整除”的推理或训练场景。操作上先用现成服务框架把并行度调通，再考虑自己写切分逻辑；验证看各卡显存是否均衡、日志里的并行规模是否与设置一致、贪心解码输出是否和单卡基线接近。跨节点、量化 kernel 只支持单卡、维度不整除时，不建议硬上张量并行，回退到降低并行度或改用流水并行更稳。

先确认当前场景和目标

先把“为什么需要多卡”写清楚。如果只是单卡 OOM，但显存缺口不大，优先考虑降低 max_model_len、减少并发、换更省显存的精度，而不是直接上张量并行；张量并行每一层前后都要做集合通信，它不是免费的显存扩容手段。

需要先确认的信息通常包括下面几项：

  • 显存缺口量级：按权重精度估算参数量占用（例如 fp16 大致每个参数 2 字节），再加上 KV cache 和推理时的激活开销，判断是“差一点”还是“差一个数量级”。
  • 结构整除约束：注意力头数（包括 GQA 里的 KV 头数）、hidden_size、中间层维度、词表大小，是否都能被并行度整除。并行度取 2、4、8 这类 2 的幂，通常比取 3、6 更容易对齐。
  • 拓扑条件：同一节点内 NVLink 或高速 PCIe 通常可以接受；跨节点通过普通网络做逐层 all-reduce，延迟往往盖过并行带来的收益。
  • 目标指标：是要把模型跑起来，还是要降低单请求首 token 延迟，还是提高整体吞吐。这三个目标对并行度的最优取值不一样，跑起来之后需要分别测。

判断条件可以粗化为一句：并行度 ≤ 单机可用 GPU 数，且关键维度都能整除，就值得先试；任一条件不满足，就先解决整除或拓扑问题。

按这个顺序处理

第一步先拿单卡或最小配置做基线，记录能跑通的最大上下文和显存占用，这决定后面并行度要开多大。第二步选并行度，从 2 开始往上加，不要一上来就顶满 8 卡。第三步用现成服务框架跑通，把“模型能不能被切开”和“业务代码怎么写”两件事分开处理。第四步才考虑自己用分布式接口做切分。

以 vLLM 这类服务框架为例，启动参数是最省事的入口，参数名以你实际安装的版本为准，需要替换的是模型路径和并行度：

多GPU部署大模型怎么配置张量并行
vllm serve /models/your-model \
  --tensor-parallel-size 4 \
  --gpu-memory-utilization 0.85 \
  --max-model-len 8192

启动前确认容器或主机能看到全部目标 GPU（nvidia-smi -L 列出数量与预期一致），容器场景一般需要 --gpus all --ipc=host 并给足共享内存，否则并行进程之间建立通信容易失败。如果启动日志里报告的并行规模和你设置的 --tensor-parallel-size 不一致，先查环境变量和可见 GPU 列表，不要继续调模型参数。

用 Hugging Face 生态时要注意一个常见混淆：device_map="auto" 做的是按层切分到多卡的流水式放置，不是张量并行。它能让模型勉强跑起来，但同一时刻只有一张卡在算，延迟不会因为卡变多而下降。要真正把单层拆开，需要走张量并行接口或框架自带方案。

如果确实需要自己切分，下面是一个通用骨架，用于说明“哪些层用列并行、哪些层用行并行”，层名和 device_mesh 的构建方式需要按你的模型与框架版本替换：

import os
import torch
import torch.distributed as dist
from torch.distributed.device_mesh import init_device_mesh
from torch.distributed.tensor.parallel import (
    parallelize_module, ColwiseParallel, RowwiseParallel,
)

tp_size = int(os.environ.get("TP_SIZE", "2"))
dist.init_process_group("nccl")
rank = dist.get_rank()
torch.cuda.set_device(rank % torch.cuda.device_count())

mesh = init_device_mesh("cuda", (tp_size,))
model = build_your_model().to(torch.float16).cuda()  # 替换为你的加载逻辑

# 列并行切输出维，行并行切输入维；成对的层必须一起改
plan = {
    "model.layers.0.self_attn.q_proj": ColwiseParallel(),
    "model.layers.0.self_attn.k_proj": ColwiseParallel(),
    "model.layers.0.self_attn.v_proj": ColwiseParallel(),
    "model.layers.0.self_attn.o_proj": RowwiseParallel(),
    "model.layers.0.mlp.gate_proj": ColwiseParallel(),
    "model.layers.0.mlp.up_proj": ColwiseParallel(),
    "model.layers.0.mlp.down_proj": RowwiseParallel(),
}
model = parallelize_module(model, mesh, plan)

这个骨架只覆盖一层，实际使用需要对所有层循环套用同一套 plan；embedding 和 lm_head 的处理各框架差异较大，通常建议先查框架文档再决定是按词表切分还是整体复制。启动方式一般是 torchrun --nproc-per-node 4 your_script.py，同时设置 TP_SIZE=4，两个数字必须一致。

多GPU部署大模型怎么配置张量并行

改完后怎么验证

验证分三层，从“跑起来了”到“跑对了”再到“跑得值”。

  1. 进程与显存：nvidia-smi 观察各卡显存占用是否大致接近。如果只有一张卡占满、其余几乎为空，说明切分没有真正生效，多半是并行度没传进去或走了流水式放置。
  2. 通信健康度：日志里不应出现 NCCL 超时、连接被拒、设备号越界一类报错。首次启动慢属于正常，但持续卡在建连阶段通常指向 IPC、共享内存或网卡选择问题。
  3. 输出一致性：用同一段 prompt、贪心解码（temperature 设为 0），对比并行前后生成的前若干个 token。张量并行会引入浮点累加顺序差异，不要求逐位一致；如果出现明显跑偏、重复或乱码，优先怀疑层名没有全部覆盖、行并行后缺少 all-reduce。
  4. 收益对照：固定模型和输入长度，分别测并行度 1、2、4 下的单请求延迟与并发吞吐。判断是否值得继续加卡，看的是延迟随并行度变化的拐点，一旦加卡后延迟不再下降甚至上升，说明通信开销已经占主导，应当停在上一档。

预期信号是：显存占用随并行度上升而下降，输出质量与基线接近，延迟在某个并行度之前下降、之后走平。如果只满足前两条而延迟不变，通常说明瓶颈不在显存，而在通信或调度。

容易忽略的边界

一是整除约束的坑。头数不能被并行度整除时，有的框架会自动补齐并给出警告，有的直接报维度不匹配；GQA 模型还要检查 KV 头数，KV 头数少于并行度时通常无法继续切分，只能降低并行度。二是量化权重与张量并行的兼容性，部分量化 kernel 只在单卡或特定并行度下实现，换并行度前需要先确认所用推理后端的支持范围，不要假定“能加载就能并行”。

三是跨节点问题。张量并行每层都通信，跨节点往往要用普通网络承载 all-reduce，即使能跑通，延迟表现也常常不如直接做多副本数据并行。四是并行度不是越大越好，卡数增加后每张卡的算力利用率可能被通信拖低，需要靠前面的延迟对照来定档。

如果卡在整除或兼容性上，回退路径通常是：把并行度降到 2 或 1；改用流水并行把层分到不同卡；或者干脆起多个单卡实例做数据并行，用调度层分流。这些做法各有代价，流水并行会增加卡间空转，多实例会成倍占用显存，选之前先回到第一步确认显存缺口和目标指标，别把临时的绕过方案当成长期配置。