text-generation-webui项目中多GPU显存分配问题的分析与解决方案

问题背景

在text-generation-webui项目运行过程中，当用户尝试在多GPU环境下加载未量化的72B大模型时，出现了显存分配异常的问题。具体表现为：系统无法正确按照用户设置的显存分配滑块值进行资源分配，导致部分GPU显存未能充分利用，反而将模型部分加载到系统内存中，严重影响推理速度。

问题现象

典型问题场景出现在配备2块RTX 4090和3块RTX A6000显卡的高端工作站上。系统表现出以下异常行为：

1. 显存分配不均：其中两块A6000显卡仅加载28.1GB显存后即停止加载
2. 资源分配错位：UI界面显示的GPU顺序与实际硬件顺序不一致
3. 控制失效：显存分配滑块对部分GPU的控制失效，特别是GPU4的滑块实际上控制着GPU0的显存
4. 自动回退：当显存不足时，系统未报错而是自动回退到系统内存

技术分析

经过深入排查，发现问题根源在于transformers库的显存分配机制存在缺陷：

1. 分配算法缺陷：原分配算法在计算各GPU显存大小时存在逻辑错误，导致无法正确识别各GPU的实际可用显存
2. 加载顺序问题：默认加载顺序与用户预期不符，且无法通过界面设置进行调整
3. 边界条件处理不足：当显存接近满载时，缺乏合理的错误处理机制

解决方案

核心代码修改

需要修改transformers库中的modeling_utils.py文件，主要涉及两个关键修改点：

1. 显存计算修正：

# 原问题代码（需注释掉）
# max_memory = {i: torch.cuda.mem_get_info(i)[1] for i in range(num_devices)}

# 修正后代码
max_memory = {}
for i in range(num_devices):
    total_mem = torch.cuda.get_device_properties(i).total_memory
    reserved = torch.cuda.memory_reserved(i)
    max_memory[i] = total_mem - reserved

2. 加载顺序优化：

# 将默认加载顺序改为倒序，便于优先使用特定GPU
device_map = "sequential"
if num_devices > 1:
    device_map = {i: min_memory for i in reversed(range(num_devices))}

实施步骤

1. 定位transformers库安装位置（通常在Python的site-packages目录下）
2. 备份原始的modeling_utils.py文件
3. 按照上述方案修改关键代码段
4. 保存修改并重启text-generation-webui服务

优化建议

1. 显存监控：建议在web UI中添加实时显存监控界面，帮助用户更直观地了解各GPU的负载情况
2. 顺序自定义：增加GPU加载顺序的自定义选项，满足不同硬件配置的需求
3. 预分配机制：实现显存预分配测试功能，在正式加载模型前验证分配方案可行性

注意事项

1. 修改核心库文件可能影响其他依赖transformers的项目，建议使用虚拟环境
2. 不同版本的transformers库可能代码位置略有不同，需要灵活调整
3. 对于Windows系统，需要注意文件权限问题
4. 建议在修改前完全关闭所有Python进程

总结

通过上述方案，成功解决了text-generation-webui在多GPU环境下显存分配异常的问题。该方案不仅修复了显存计算错误，还优化了资源分配策略，使得大模型能够充分利用多GPU的显存资源，避免不必要的系统内存使用，显著提升了推理效率。对于面临类似问题的用户，可以参考本方案进行调优。