用 torchrun 起四个进程加载同一份 7B 模型,nvidia-smi 显示 0 号卡占了 22G,其余三张只有 1G 多,服务吞吐只有单卡水平,日志里也没报错。

一开始怀疑是「CUDA_VISIBLE_DEVICES 没有按进程切分」,查了一圈,方向不对。

真正麻烦的地方在于,几个原因的现象几乎一样。代码里硬编码了 cuda:0,没按 local_rank 取设备。少数情况下是每个进程都加载了完整权重,没有做张量并行。

另外,模型在 init 时就 to('cuda'),早于分布式初始化。

先打印每个进程看到的设备号

python -c "import torch;print(torch.cuda.current_device())"

再往下按 local rank 设置设备再建模型

python -c "import os,torch;os.environ['CUDA_VISIBLE_DEVICES']=str(0);print(torch.cuda.device_count())"

随后用 torchrun 传入正确的并行参数

torchrun --nproc_per_node=4 --nnodes=1 serve.py --tp 4

等前面都确认过,再确认每张卡的显存与利用率

nvidia-smi --query-gpu=index,memory.used,utilization.gpu --format=csv

别只看报错消失,验证要多做一步。四张卡的显存占用都在 18G 以上且接近,整体吞吐是单卡的三倍以上,压测时没有卡空转。

多卡不等于模型会自动分片,设备号写死是最常见的低级原因,先用 current_device 打印一遍。同一份配置在测试环境跑通,不代表线上没事,两边的默认值经常不一样。

作者 9823189