【Bug已解决】[Bug] AdaLora's update_and_allocate method is lost in inject_adapter_in_model() preventing rank pruning 解决方案

一、现象长什么样

AdaLoRA 的核心卖点是训练过程中按重要性动态剪枝（rank pruning）：每个 adapter 层的update_and_allocate()方法会根据梯度敏感度，把不重要的奇异值（rank）置零，实现“预算内最优秩分配”。但当你用inject_adapter_in_model()这个轻量注入函数（不包一层完整PeftModel）把 AdaLoRA 塞进模型时，发现：

from peft import inject_adapter_in_model, AdaLoraConfig config = AdaLoraConfig(...) inject_adapter_in_model(config, model) # 注入成功 # 训练若干步后查看：rank 没有被剪枝，budget 没生效

更隐蔽的是，有些路径下会直接报：

AttributeError: 'AdaLoraLayer' object has no attribute 'update_and_allocate'

或者update_and_allocate存在却从没被调用，因为注入流程没把 AdaLoRA 需要的“重要性估计 + 分配”钩子接上。结果：AdaLoRA 退化成普通 LoRA，rank pruning 完全失效。本文讲清根因与修复。

二、背景

PEFT 有两条把 adapter 装进模型的路径：

1. get_peft_model(base, config)：返回PeftModel，它会按 config 类型（AdaLora）正确构造AdaLoraModel，注册所有 AdaLoRA 专属的状态（lora_E、lora_A、lora_B、d向量）和update_and_allocate的调用时机（在optimizer.step()后由AdaLoraModel统一触发）。
2. inject_adapter_in_model(config, model)：轻量函数，直接把 adapter 层替换进原模型的 module 树，不包 PeftModel。它依赖各 tuner 自己实现inject_adapter_in_model逻辑；如果某个 tuner（这里就是 AdaLoRA）的注入实现没把update_and_allocate的触发机制带上，剪枝逻辑就丢了。

AdaLoRA 的剪枝发生在每次优化器 step 之后：AdaLoraModel.update_and_allocate()读各层的重要性（orth_reg_weight相关的敏感度），重排并裁剪最不重要的 rank，更新mask/d。这条调用链只有在AdaLoraModel的trainable流程里才存在。inject_adapter_in_model走的是另一条更薄的代码路径，早期实现里漏掉了这一步。

三、根因

根因 A：inject_adapter_in_model没触发update_and_allocate

最直白的原因。轻量注入只做了“把AdaLoraLayer替换进去”，但没接上“step 后调用update_and_allocate”的钩子。于是层在，剪枝逻辑不在。

根因 B：AdaLoRA 必需的辅助状态没初始化

AdaLoraLayer需要lora_E（缩放向量）、d（当前秩分配）、orth_reg_*等状态。inject_adapter_in_model若只调用了通用注入、没走 AdaLoRA 的create_adapter/init_adapter，这些状态可能是默认的，重要性估计无法计算。

根因 C：调用方在optimizer.step()后没有手动触发

即使用inject_adapter_in_model，正确的用法是训练循环里在step()后手动调用model.update_and_allocate()（或遍历 AdaLora 层调用）。很多人以为和get_peft_model一样自动触发，结果从没调用。

根因 D：版本差异

某些 PEFT 版本里inject_adapter_in_model对 AdaLoRA 的支持是后来补的；老版本天然缺失。需要确认版本是否在修复之后。

根因小结

• AdaLoRA 的 rank pruning 依赖update_and_allocate，它在AdaLoraModel的 step 后流程里触发；
• inject_adapter_in_model是薄路径，早期漏接这个触发；
• 正确修复：用get_peft_model走完整 AdaLoRA，或在注入后手动触发update_and_allocate。

四、最小可运行复现

下面脚本对比“用 get_peft_model（剪枝生效）”和“用 inject_adapter_in_model 但没手动触发（剪枝失效）”：

import torch import torch.nn as nn from peft import AdaLoraConfig, get_peft_model, inject_adapter_in_model class Tiny(nn.Module): def __init__(self): super().__init__() self.lin = nn.Linear(32, 16) def forward(self, x): return self.lin(x) def make_cfg(): return AdaLoraConfig( r=8, lora_alpha=16, target_modules=["lin"], lora_dropout=0.0, total_r=8, # 剪枝预算相关 beta1=0.85, beta2=0.85, orth_reg_weight=0.5, inference_mode=False, ) def train_loop(model, steps=20): opt = torch.optim.AdamW(model.parameters(), lr=1e-3) x = torch.randn(4, 32) y = torch.randn(4, 16) for _ in range(steps): opt.zero_grad() loss = ((model(x) - y) ** 2).mean() loss.backward() opt.step() # 关键：AdaLoRA 需要在 step 后触发 update_and_allocate if hasattr(model, "update_and_allocate"): model.update_and_allocate(global_step=_) elif hasattr(model, "base_model") and hasattr(model.base_model, "update_and_allocate"): model.base_model.update_and_allocate(global_step=_) def rank_budget_used(model) -> int: used = 0 for m in model.modules(): if hasattr(m, "d"): # AdaLoraLayer 有 d 向量 used += int((m.d > 0).sum()) if m.d is not None else 0 return used def main(): # 路径1：完整 PeftModel，剪枝自动触发 m1 = get_peft_model(Tiny(), make_cfg()) train_loop(m1) print("get_peft_model 路径剩余秩预算:", rank_budget_used(m1)) # 路径2：inject_adapter_in_model，若没手动触发则剪枝无效 m2 = Tiny() inject_adapter_in_model(make_cfg(), m2) train_loop(m2) # 这里 hasattr(m2,'update_and_allocate') 可能为 False -> 没剪枝 print("inject 路径剩余秩预算:", rank_budget_used(m2)) if __name__ == "__main__": main()

运行后你会看到两条路径的“剩余秩预算”不同——路径2 若没触发update_and_allocate，秩分配保持初始值，剪枝没发生，正是 bug 现象。

五、解决方案（第一层：最小直接修复）

最稳的修复：用get_peft_model走完整 AdaLoRA 流程，剪枝自动在 step 后触发，不要为了省事用inject_adapter_in_model：

from peft import get_peft_model, AdaLoraConfig config = AdaLoraConfig(r=8, total_r=8, target_modules=["lin"], orth_reg_weight=0.5) model = get_peft_model(base_model, config) opt = torch.optim.AdamW(model.parameters(), lr=1e-3) for step, batch in enumerate(loader): loss = model(**batch).loss loss.backward() opt.step() model.update_and_allocate(global_step=step) # PeftModel 也支持手动触发

如果确实要用inject_adapter_in_model（比如要保留原模型类型），则必须手动在 step 后触发剪枝：

inject_adapter_in_model(config, model) opt = torch.optim.AdamW(model.parameters(), lr=1e-3) for step, batch in enumerate(loader): loss = model(**batch).loss loss.backward() opt.step() # 手动遍历所有 AdaLora 层触发剪枝 for module in model.modules(): if hasattr(module, "update_and_allocate"): module.update_and_allocate(global_step=step)

六、解决方案（第二层：结构性改进）

6.1 封装一个“触发 AdaLoRA 剪枝”的 helper

避免到处手写遍历：

def trigger_adalora_pruning(root_module, global_step: int): for m in root_module.modules(): fn = getattr(m, "update_and_allocate", None) if callable(fn): fn(global_step=global_step)

6.2 选择路径时给出明确建议

def build_adalora(base, config, lightweight: bool): if lightweight: inject_adapter_in_model(config, base) return base # 调用方负责手动触发剪枝 return get_peft_model(base, config) # 自动触发

6.3 升级 PEFT 到修复版本

inject_adapter_in_model对 AdaLoRA 的支持在后续版本补齐。确认版本：

python -c "import peft; print(peft.__version__)"

并查看对应版本的inject_adapter_in_model是否已对 AdaLora 注册update_and_allocate触发。若版本过旧，升级即可。

七、解决方案（第三层：断言 / CI 守护）

加测试确认剪枝确实发生（秩预算被重新分配）：

import torch import pytest from peft import AdaLoraConfig, get_peft_model, inject_adapter_in_model def test_adalora_prunes_with_peft_model(): model = get_peft_model(build_base(), AdaLoraConfig(r=8, total_r=8, target_modules=["lin"])) before = sum(int((m.d > 0).sum()) for m in model.modules() if hasattr(m, "d")) train_and_prune(model, steps=30) after = sum(int((m.d > 0).sum()) for m in model.modules() if hasattr(m, "d")) assert after != before or after < 8 * n_layers, "AdaLoRA 未触发 rank pruning" def test_inject_requires_manual_trigger(): base = build_base() inject_adapter_in_model(AdaLoraConfig(r=8, total_r=8, target_modules=["lin"]), base) with pytest.raises(AssertionError): # 若没手动触发，d 向量不应变化 assert any(hasattr(m, "update_and_allocate") for m in base.modules())

CI 里两个用例都过，说明“薄注入路径必须手动触发”的约定被守住，不会再有人以为自动剪枝。

八、排查清单

AdaLoRA 不剪枝时查：

1. 用的是get_peft_model还是inject_adapter_in_model？后者是薄路径，剪枝触发可能缺失。
2. update_and_allocate有没有被调用？在optimizer.step()之后必须触发一次。
3. 注入路径下有没有手动遍历触发？inject_adapter_in_model需手动module.update_and_allocate(global_step=)。
4. AdaLoRA 辅助状态初始化了没？lora_E、d、orth_reg_*必须由 AdaLoRA 的init_adapter建。
5. PEFT 版本够新吗？老版本inject_adapter_in_model对 AdaLoRA 支持不全。
6. total_r / budget 设了吗？没设预算，剪枝无目标。
7. d向量变化了吗？用rank_budget_used检查剪枝是否真发生。

九、小结

“AdaLora's update_and_allocate method is lost in inject_adapter_in_model() preventing rank pruning” 的根因是轻量注入路径漏接了 AdaLoRA 的剪枝触发：

• AdaLoRA 的 rank pruning 靠update_and_allocate()，它在AdaLoraModel（即get_peft_model路径）的 step 后流程里自动触发；
• inject_adapter_in_model()是更薄的函数，早期实现没把触发接上，导致剪枝静默失效；
• 最稳修复：用get_peft_model走完整 AdaLoRA；若坚持用inject_adapter_in_model，必须在optimizer.step()后手动遍历触发update_and_allocate(global_step=)；
• 升级 PEFT 到修复版本也能解决；
• 用d向量/秩预算变化的断言守护“剪枝确实发生”。

一句话：AdaLoRA 想剪枝，step 之后必须调update_and_allocate；inject_adapter_in_model不会自动调，要么换get_peft_model，要么自己手动触发。