【Bug已解决】GraLoRA merge_and_unload with modules_to_save error 解决方案

【Bug已解决】GraLoRA merge_and_unload with modules_to_save error 解决方案一、现象长什么样用GraLoRA一种梯度感知的 LoRA 变体做微调并且用了modules_to_save把某些原模型模块原样保存/训练训练完调merge_and_unload()时报错AttributeError: GraLoraLayer object has no attribute merge_and_unload (或基类行为不对)或KeyError: modules_to_save 的模块在 merge 时被错误地当成 LoRA 层处理导致形状/属性错误根因是modules_to_save的模块是被AuxiliaryTrainingWrapper包裹的原模型模块不是 LoRA 层而 GraLoRA 的merge_and_unload实现没有正确处理这类“既非 LoRA、又需保留”的模块——它可能把modules_to_save的模块也尝试做 LoRA 合并或卸载时把不该丢的模块丢了于是报错。本文讲清根因与修复。二、背景回顾两个概念modules_to_savePEFT 里标记“原模型里哪些模块要原样训练并保存”如分类头。这些模块会被AuxiliaryTrainingWrapper包一层但它们是原模型模块不是 LoRA 层没有lora_A/lora_B也没有 LoRA 的 merge 语义。merge_and_unload()把 LoRA 的B·A增量合并回 base 权重并卸载移除adapter 包装返回一个干净的原始模型带合并后的权重。问题出在 GraLoRA作为 LoRA 变体的merge_and_unload重写时它遍历所有被 PEFT 管理的子模块做合并/卸载但没区分“真正的 LoRA 层GraLoraLayer”和“modules_to_save的原模型模块AuxiliaryTrainingWrapper”。当遇到后者若它调用layer.merge()假设是 LoRA 层AuxiliaryTrainingWrapper没有 LoRA 的merge报 AttributeError若它卸载时把modules_to_save模块整体移除原模型就缺了分类头/保存模块结构错正确做法modules_to_save的模块在merge_and_unload时应保留原样它们是原模型的一部分权重已训练不需要合并 LoRA 增量只把真正的 GraLoraLayer 合并并卸载。三、根因根因 AGraLoRA 的merge_and_unload没区分 LoRA 层与modules_to_save最直接。遍历所有 managed 模块做 merge遇到AuxiliaryTrainingWrappermodules_to_save没有 LoRA merge 方法报错。根因 Bmodules_to_save模块被误当 LoRA 合并modules_to_save是原模型模块无lora_A/B强行 merge 形状/属性错。根因 C卸载时丢了modules_to_save模块unload把 adapter 包装移除时若把AuxiliaryTrainingWrapper也移除且没还原原模块原模型结构破坏。根因 DGraLoRA 子类没调用父类 merge_and_unload 的通用逻辑GraLoRA 重写了 merge 但没复用 PEFT 对modules_to_save的处理导致边缘情况失败。根因小结modules_to_save模块是原模型模块被 AuxiliaryTrainingWrapper 包非 LoRA 层GraLoRA 的 merge_and_unload 没区分二者误合并/误卸载报错修复merge 只对真正 LoRA 层modules_to_save 保留原样unload 正确还原。四、最小可运行复现下面脚本模拟“merge_and_unload 区分 LoRA 层与 modules_to_save”import torch import torch.nn as nn class GraLoraLayer(nn.Module): def __init__(self, base, r4): super().__init__() self.base_layer base base.weight.requires_grad False self.A nn.Parameter(torch.zeros(r, base.in_features)) self.B nn.Parameter(torch.zeros(base.out_features, r)) def merge(self): w self.base_layer.weight w.copy_(w self.B self.A) self.base_layer.weight.requires_grad False class AuxWrapper(nn.Module): modules_to_save 的包裹层原模型模块非 LoRA def __init__(self, module): super().__init__() self.original_module module def merge_and_unload_safe(model): 正确只合并真正 LoRA 层modules_to_save 保留原样 for m in model.modules(): if isinstance(m, GraLoraLayer): m.merge() # LoRA 合并 # AuxWrapper (modules_to_save) 不做任何合并/删除保留原模块 return model def demo(): base nn.Linear(16, 8) lora GraLoraLayer(base) head AuxWrapper(nn.Linear(8, 4)) # modules_to_save 模块 model nn.ModuleDict({lora: lora, head: head}) try: # 错误对所有模块都调 merge含 AuxWrapper for m in model.modules(): m.merge() except AttributeError as e: print(错误写法触发:, e) # AuxWrapper 无 merge merge_and_unload_safe(model) print(正确写法LoRA 合并modules_to_save 保留无错误) if __name__ __main__: demo()运行后错误写法对所有模块调 merge因 AuxWrapper 无 merge 报错正确写法只合并 LoRA 层、保留 modules_to_save无错误。五、解决方案第一层最小直接修复GraLoRA 的merge_and_unload必须按模块类型分支处理def merge_and_unload(self): # 1) 合并真正的 LoRA 层 for module in self.modules(): if isinstance(module, GraLoraLayer): module.merge() # 把 B·A 合并进 base # 2) modules_to_save 的 AuxiliaryTrainingWrapper保留原模块不合并不删除 # PEFT 基类已处理unload 时还原 original_module # 3) 调用父类 unload 完成卸载 return super().merge_and_unload() # 复用 PEFT 对 modules_to_save 的通用处理用户侧规避若遇到该错误先不调用 merge_and_unload手动合并 LoRA 并保留 modules_to_savemodel get_peft_model(base, GraLoraConfig(r4, modules_to_save[classifier])) # 训练... # 手动只合并 LoRA 层 for m in model.modules(): if isinstance(m, GraLoraLayer): m.merge() # 不调用 merge_and_unload直接用带合并权重的模型六、解决方案第二层结构性改进6.1 复用 PEFT 基类对 modules_to_save 的处理# GraLoRA 的 merge_and_unload 应先调父类的通用 unload 逻辑它已正确还原 modules_to_save class GraLoraModel(BaseTuner): def merge_and_unload(self, ...): self._merge_lora_layers() # GraLoRA 自己的合并 return super().merge_and_unload(...) # 父类处理 modules_to_save 卸载6.2 升级 PEFT 到修复版本pip install -U peftPEFT 后续版本对 modules_to_save 变体 merge 处理更完善。6.3 给 PEFT 提 PRmerge 前类型检查社区修复就是让所有变体的 merge_and_unload 在遍历时isinstance区分 LoRA 层与 AuxWrapper避免误合并。七、解决方案第三层断言 / CI 守护import torch import pytest import torch.nn as nn def test_merge_only_lora_layers(model_gra): # 守护merge 只作用于 GraLoraLayer不作用于 modules_to_save merged model_gra.merge_and_unload() # modules_to_save 的 classifier 应仍存在且为原模块 assert hasattr(merged, classifier), modules_to_save 模块不应被删除 assert isinstance(merged.classifier, nn.Linear), 应还原为原模型模块 def test_no_attribute_error_on_merge(model_gra): try: model_gra.merge_and_unload() except AttributeError as e: pytest.fail(fmerge_and_unload 仍报 AttributeError: {e}) def test_weights_merged(model_gra): before dict(model_gra.named_parameters()) merged model_gra.merge_and_unload() # LoRA 层权重应含合并增量与 base 不同 assert merged is not None def test_unload_restores_base_structure(model_gra): merged model_gra.merge_and_unload() # 卸载后应是干净的原始模型结构无 LoRA 壳 assert not any(lora in n for n, _ in merged.named_parameters())CI 跑这四条GraLoRA modules_to_save 的 merge_and_unload 被守住。八、排查清单GraLoRAmerge_and_unloadmodules_to_save报错时查merge 区分了 LoRA 层与 modules_to_save 吗后者是原模型模块非 LoRA。modules_to_save 被误合并了吗它没有 lora_A/B强行 merge 会错。卸载时丢了 modules_to_save 吗应保留并还原为原模块。GraLoRA 调了父类 merge_and_unload 吗复用 PEFT 对 modules_to_save 的通用处理。PEFT 升级了吗新版更完善。用户侧能绕过吗手动合并 LoRA 层、不调 merge_and_unload。CI 测了 modules_to_save 合并吗加断言防回归。九、小结“GraLoRA merge_and_unload with modules_to_save error” 是GraLoRA 的merge_and_unload没区分真正 LoRA 层与modules_to_save的原模型模块导致误合并/误卸载报错根因modules_to_save模块被AuxiliaryTrainingWrapper包裹非 LoRA 层GraLoRA 的 merge 遍历时误对它调merge/误卸载第一层merge 只作用于GraLoraLayermodules_to_save保留原样并调用父类 unload 复用通用处理第二层升级 PEFT、提 PR 让所有变体 merge 前做类型检查、复用基类 modules_to_save 逻辑第三层pytest 守护“modules_to_save 不丢 无 AttributeError 卸载后结构干净”。一句话GraLoRA 的merge_and_unload必须按类型分支只合并真正的 LoRA 层modules_to_save的原模型模块保留原样并正确还原别把 wrapper 当 LoRA 误合并。