支持不同显存显卡的推理加速框架

[whyorwhynot]

搞到了 4 张 a100 显卡，两张 80GB 显存的，两张 40GB 显存的，准备跑 Qwen1.5-110B-Chat 非量化版本的。 可是 vllm 等推理加速框架没办法手动配置各张卡的显存占用，只能平均分配，导致受限于显存容量最低的显卡，ollama 可以按比例分配显存,但是其依赖的 llama.cpp 貌似多张卡是串行推理的，也没有实现 paged attention ，速度相比 vllm 慢很多。 现在有什么推理加速库可以像原始 transformers 库一样，通过配置 max_memory={0: "75GIB", 1: "75GIB",2:"38GIB",3:"38GIB"}等方式,充分利用不同大小显存的卡吗。

[MeePawn666]

Tensor-rt 不行吗，印象中是可以的

[whyorwhynot]

@MeePawn666 第一次用 TensorRT-LLM ，运行起来有些复杂啊，convert_checkpoint.py 转换模型设置--tp_size 4 的时候，生成的 4 个 rank*.safetensors 看起来大小是相同的，用 mpirun 简单跑了跑小模型，看到日志里面显示 Rank 0 is using GPU 0 ，Rank 1 is using GPU 1 ，推测显存占用还应该是平均分配啊

[Djaron]

vllm 不是可以设置显存利用率吗？

[whyorwhynot]

@Djaron
实际用下来是按照显存最低的那张卡的显存*利用率为实际每张卡使用的显存