LMDeploy多GPU并行推理中的显存分配不均问题分析与解决

问题现象

在使用LMDeploy部署Qwen2.5-72B-Instruct-AWQ等大模型进行多GPU并行推理时，发现不同GPU间的显存占用存在明显差异。具体表现为：当使用4个GPU(TP=4)时，GPU3的显存占用达到23608MiB，而其他GPU(0,1,2)的显存占用仅为20426MiB左右。这种现象在Qwen2.5-7B等不同规模的模型上也同样存在。

技术背景

在大型语言模型的多GPU并行推理中，Tensor Parallelism(张量并行)是一种常见的并行策略。它将模型的参数和计算负载分配到多个GPU上，理论上应该实现各GPU间的负载均衡。然而，实际部署中可能会出现显存分配不均的情况，这通常与以下几个因素有关：

1. 模型切分策略：不同的模型层可能有不同的参数规模，导致分配到各GPU上的参数不均衡
2. KV缓存管理：推理过程中的键值缓存(KV Cache)分配可能不均匀
3. 系统级因素：GPU间的通信开销、PCIe拓扑结构等

问题分析

经过深入排查，发现问题根源在于模型转换过程中的显存管理。具体表现为：

1. 当直接使用lmdeploy serve api_server命令启动服务并同时进行模型转换时，转换过程中的临时显存分配没有被及时释放
2. 这种显存泄漏主要影响最后一个GPU(GPU3)，导致其显存占用明显高于其他GPU
3. 这种问题在AWQ量化模型上尤为明显，因为量化过程本身就需要额外的显存开销

解决方案

针对这一问题，推荐采用以下最佳实践：

1. 分离模型转换和推理过程：先使用lmdeploy convert命令将模型转换为TurboMind格式，再启动推理服务
2. 显式释放资源：在模型转换完成后，显式调用显存释放操作
3. 监控显存使用：在服务启动后，使用nvidia-smi等工具监控各GPU显存使用情况

具体操作步骤如下：

# 第一步：单独执行模型转换
lmdeploy convert Qwen/Qwen2.5-72B-Instruct-AWQ \
    --quant-policy 0 \
    --model-format awq \
    --tp 4

# 第二步：启动推理服务
lmdeploy serve api_server ./workspace \
    --server-name 0.0.0.0 \
    --server-port 23332 \
    --cache-max-entry-count 0.8 \
    --tp 4

技术建议

对于大规模模型的多GPU部署，还应注意以下几点：

1. 硬件拓扑考虑：确保GPU间有良好的互联带宽(如使用NVLink连接)
2. 负载均衡监控：持续监控各GPU的计算负载和显存使用
3. 分批处理策略：对于流式请求，采用合理的批处理大小以避免显存峰值
4. 量化策略选择：根据硬件条件选择合适的量化方式(AWQ/GPTQ等)

总结

LMDeploy作为高效的推理部署工具，在大模型服务化方面表现出色。通过理解其底层工作机制并遵循正确的使用流程，可以有效避免显存分配不均等问题。特别是在处理超大规模模型时，预先转换模型格式、合理规划资源分配，是保证服务稳定性的关键。

对于生产环境部署，建议在模型转换阶段就充分考虑目标硬件的特性，并通过压力测试验证各GPU的资源使用均衡性，从而确保服务的最佳性能和稳定性。