多卡推理显存不够怎么分配模型层

多卡推理显存不够，第一步不是加卡，而是分清两种不同的问题：整模型权重放不进单卡，还是权重放得下但 KV cache 和中间激活把显存撑爆。前者适合按层切分，把 decoder 层分段放到不同卡上，或者用张量并行把单层内部切开；后者单纯切层往往无效，得先缩短序列长度、降 batch，或者调整 KV cache 的精度和分配方式。判断错方向，后面的配置都是白改。

多卡显存不足时，先按 参数量 × 每参字节数 估算权重占用，和单卡扣掉运行开销后的可用显存比较：权重放不进单卡，就做层切分或张量并行；权重放得进但长序列 OOM，优先处理 KV cache。切分后用 nvidia-smi 连续采样看每卡峰值是否均衡、是否留出余量，再用同一 prompt 对比切分前后的输出。层切分会把单请求延迟叠加成各段之和，CPU offload 只能当临时保底手段，不能当成性能方案，边界条件必须结合具体框架确认。

先确认当前场景和目标

这一步的目标是拿到几个能直接决定方案的数，而不是凭感觉配置。建议在目标机器上先确认以下信息：

单卡真实可用显存：用 nvidia-smi 看 total 与 used，再扣掉 CUDA context、其他进程占用和显存碎片，实际能给模型的通常比标称小一截。
权重占用估算：参数量 × 每参字节数（fp16 约 2 字节，int8 约 1 字节，int4 约 0.5 字节，含量化元数据的开销要另算）。这个数只估权重，不含 KV cache 和激活。
序列长度与并发：最大 context 多长、同时在跑的请求几个。KV cache 随序列长度和并发线性增长，是长文本场景最常见的爆点。
卡间互联：用 nvidia-smi topo -m 看是 NVLink 还是 PCIe。张量并行每层都要通信，PCIe 下通常不建议开大并行度。

# 切分前的环境确认，设备号按实际情况替换
nvidia-smi --query-gpu=index,name,memory.total,memory.used --format=csv
nvidia-smi topo -m

判断条件可以简化成两条：权重估算值超过单卡可用显存的八成左右，就走层切分或张量并行；权重放得下、只是跑到某个序列长度才失败，先查 KV cache 和 batch，不要急着加卡。

按这个顺序处理

处理顺序建议是：先压显存底噪，再算切分粒度，最后才动并行策略。顺序反过来，很容易在错误的精度下反复调切分比例。

先定精度：确认模型支持的量化格式和推理框架是否兼容，能用低精度就先把权重压下来，很多「放不下」在换精度后就变成「放得下」。
算每层均摊：总权重占用除以层数，得到单层近似值；embedding 和 lm_head 在词表大时单独占一块，不能按普通层均摊。
做层切分：把层按显存比例分段映射到各卡，embedding 放第一张卡，lm_head 放最后一张卡，中间层按容量分。
或用张量并行：单层都放不下、或卡间是 NVLink 时，用张量并行把每层权重切到各卡。
留余量再启动：每卡预留一部分显存给激活和 KV cache，不要按满打满算分配。

下面是一个可替换的层到卡分段示意，段边界按每层实际显存调整，不是固定值：

# 层到卡的手工分段示意
# 每层大小不均匀时，按显存容量分而不是按层数平分
gpu0: embed_tokens + layers[0:12]
gpu1:            layers[12:24]
gpu2:            layers[24:32] + lm_head

如果使用 Hugging Face 生态，通常可以用 device_map 配合 max_memory 表达上面的分段意图，字段名和可用参数需要按实际版本确认：

# 通用接入骨架，字段名以实际框架和版本为准
from transformers import AutoModelForCausalLM

max_memory = {
    0: '20GiB',   # 按 nvidia-smi 可用显存留出余量后填写
    1: '20GiB',
    'cpu': '48GiB',
}

model = AutoModelForCausalLM.from_pretrained(
    model_path,
    torch_dtype='auto',
    device_map='auto',           # 也可以传入自定义的层到卡映射
    max_memory=max_memory,
    offload_folder=offload_dir,  # 仅在必须 offload 时启用
)

如果框架支持张量并行，启动参数大致如下，参数名同样以实际框架为准：

# 张量并行启动骨架
python -m vllm.entrypoints.openai.api_server \
  --model /path/to/model \
  --tensor-parallel-size 2 \
  --gpu-memory-utilization 0.85 \
  --max-model-len 4096

改完后怎么验证

验证要同时看显存分布和输出行为，只看到「起来了」不够。

显存采样：启动后跑一轮固定请求，用 watch -n 1 nvidia-smi 连续观察。预期信号是各卡峰值接近、没有哪张卡先贴顶，且每卡都留有余量。
日志确认分配：看启动日志里打印的层分配结果，确认层数总和与模型一致，没有层被漏掉或重复加载。
输出对比：用同一 prompt、同一解码参数（建议贪心解码）分别跑切分前后，预期是输出一致或高度接近；如果差异明显，通常是某段层没被正确加载或精度不一致。
压力位复测：把序列长度和并发推到计划的上限，重点看 prefill 阶段是否 OOM，这一段显存峰值最高。
重复启动：连续启动两次以上，确认不会因为显存碎片出现偶发失败。

容易忽略的边界

层切分不是免费的。它让多个 GPU 串行执行，单请求延迟会变成各段之和，吞吐要靠并发补回来；张量并行虽然单层内并行，但每层都要跨卡通信，PCIe 环境下通信开销可能吃掉收益。这两条决定方案能不能真正跑起来，而不是能不能启动。

另外几个常见坑：层大小并不完全相等，按层数平分容易让某张卡先满，按容量分更稳；embedding 和 lm_head 在词表大时占用可观，单独固定在卡上和中间层混放时要留够空间；量化格式与切分方式、并行度的兼容性需要在目标框架上确认，不能直接套用别处的配置。

CPU offload 只能当保底手段使用，它会让部分层走 CPU 计算，延迟和吞吐都会明显变差，适合先把服务跑起来、再做调优，不适合作为最终方案。

如果按上面的步骤调整后仍然放不下，回退方向建议按顺序考虑：换更低的量化精度、缩短最大序列长度、降低单批并发、改用更小的模型，最后才是加卡。加卡解决的是容量问题，不解决配置问题。