V2EX = way to explore
V2EX 是一个关于分享和探索的地方
现在注册
已注册用户请  登录
whyorwhynot
V2EX  ›  机器学习

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

  •  
  •   whyorwhynot · 23 天前 · 465 次点击

    搞到了 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"}等方式,充分利用不同大小显存的卡吗。

    4 条回复    2024-05-30 18:15:45 +08:00
    MeePawn666
        1
    MeePawn666  
       23 天前 via Android
    Tensor-rt 不行吗,印象中是可以的
    whyorwhynot
        2
    whyorwhynot  
    OP
       23 天前
    @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
        3
    Djaron  
       18 天前
    vllm 不是可以设置显存利用率吗?
    whyorwhynot
        4
    whyorwhynot  
    OP
       17 天前
    @Djaron
    实际用下来是按照显存最低的那张卡的显存*利用率为实际每张卡使用的显存
    关于   ·   帮助文档   ·   博客   ·   API   ·   FAQ   ·   实用小工具   ·   1070 人在线   最高记录 6679   ·     Select Language
    创意工作者们的社区
    World is powered by solitude
    VERSION: 3.9.8.5 · 21ms · UTC 23:16 · PVG 07:16 · LAX 16:16 · JFK 19:16
    Developed with CodeLauncher
    ♥ Do have faith in what you're doing.