低内存模式下设备不一致问题

[imleoo]

问题概述

在 AWQ (Activation-aware Weight Quantization) 量化模块中发现了两个关键的设备管理问题，这些问题在低内存设备上会导致性能下降和运行时错误。

修复的问题

1. Block 模型设备未优化 (Commit: a30971a)

问题描述：

• 在 AWQ 搜索过程中，block 模型没有移动到最佳设备（GPU）上执行
• 这导致计算可能在次优设备（如 CPU）上进行，影响性能

修复内容：

在 search_by_block 方法中添加

dev = get_best_device()
block.to(dev) # 新增：确保 block 在最佳设备上执行

2. 低内存模式下设备不一致问题 (Commit: f0e60b9)

问题描述：

• 缩放因子（scales）和权重张量位于不同设备上，导致运行时错误
• 权重恢复时使用了错误的设备变量，可能将 GPU 权重恢复到 CPU

修复内容：

修复1：确保缩放因子与权重在同一设备

layer.weight.mul_(scales.to(layer.weight.device).view(1, -1))

修复2：恢复权重时使用正确的设备变量（dev 而不是 act.device）

layer.weight.data = w.to(dev)

技术细节

影响范围：

• 文件：angelslim/compressor/quant/modules/awq/search.py
• 类：AWQSearch
• 方法：search_by_block
• 行号：66, 102, 136

修复原理：

1. 通过 get_best_device() 确定最佳计算设备
2. 确保 block 模型在最佳设备上执行推理
3. 保持所有张量（权重、缩放因子）在同一设备上
4. 正确处理权重的保存和恢复

修复效果

• ✅ 解决了低内存设备上的设备不匹配错误
• ✅ 提升了 AWQ 量化过程的性能
• ✅ 确保了量化过程的稳定性和一致性
• ✅ 优化了 GPU 资源的利用

测试建议

建议在以下场景中测试修复效果：

1. 低内存 GPU 环境（< 8GB 显存）
2. 大模型量化（> 7B 参数）
3. 启用 low_memory=True 的 AWQ 量化
4. 混合精度量化场景

这些修复确保了 AngelSlim 在资源受限环境下的稳定运行，提升了量化过程的可靠性。