模型部署困惑

[ylhou]

在A800 8卡机上按照Chat Completion的示例代码加载模型 设置 max_memory = {i: "75GB" for i in range(8)} device_map="sequential" 启动后显存占用集中在0卡~6卡，7卡被完全放空，然后测试的上下文稍长一点，就报显存无法分配了 这是什么原因呢，有人遇到过吗

[fw2325]

+1

[jixiang25]

在A800 8卡机上按照Chat Completion的示例代码加载模型 设置 max_memory = {i: "75GB" for i in range(8)} device_map="sequential" 启动后显存占用集中在0卡~6卡，7卡被完全放空，然后测试的上下文稍长一点，就报显存无法分配了 这是什么原因呢，有人遇到过吗

• 我这边的解决方案是把max_memory = {i: "75GB" for i in range(8)}修改为max_memory = {i: "60GB" for i in range(8)}
• 推测max_memory的作用，是加载模型的时候，在每个GPU上尽可能地放，放到再放一个模型分片就超出给定值（75GB）为止。
• 236B模型一共有55个分片，每个分片大小约8.1GB，如果像示例代码一样指定每张卡max_memory为75GB，那么每张卡可以放75 / 8.1 = 9（向下取整）个模型分片，这样55个分片只需要55 / 9 = 7张卡，这样最后一张卡显然没有模型，自然推理的时候也用不到，就会看到显存占用为空
• 改成60GB之后，每张卡放60 / 8.1 = 7（向下取整）个模型分片，55个分片，在0到6号卡每个可以放7个模型分片，7号卡放55 - 7 * 7 = 6个模型分片，以此减少显存浪费