记录 TritonServer 部署多模型到多 GPU 踩坑 | 京东云技术团队

一、问题是怎么发现的

部署 chatglm2 和 llama2 到一个 4*V100 的 GPU 机器上遇到问题

config.pbtxt

中设置模型分别在指定 gpu 上部署实例配置不生效

如以下配置为在 gpu0 上部署本模型,部署 count=1 个实例,在 gpu1 上部署本模型,部署 count=2 个实例

instance_group [ { count: 1 kind: KIND_GPU gpus: [ 0 ] },

{ count: 2

kind: KIND_GPU

gpus: [ 1 ] } ]

部署时发现,所有模型实例都会被部署到 gpu0 上面, 由于 gpu 只有 16g 显存,在部署第一个模型实例成功后,第二个模型实例也会往 gpu0 上进行加载,最终导致 cuda out of memery.

网上搜索发现有人遇到同样的问题,链接: https://github.com/triton-inference-server/server/issues/6124

二、排查问题的详细过程

大佬回答解决方案:

三、如何解决问题

1.在 model.py 手动获取 config.pbtxt 配置的 gpu 编号 gpus:[0]

instance_group [

{

count: 1

kind: KIND_GPU

gpus: [ 0 ]

}

]

2.设置可用的 GPU 编号

os.environ["CUDA_VISIBLE_DEVICES"] = str(device_id)

3.启动成功

四、总结反思:是否可以更快发现问题?如何再次避免等。

triton 启动的使用使用 nvidia-smi -l 2 监控显卡想显存, 可以发现所有模型都在往第一个 gpu,gpu[0]内加载,发现配置 config.pbtxt 不生效

作者：京东科技 杨建

来源：京东云开发者社区 转载请注明来源