【Bug已解决】extremely slow get_peft_model() with LoRA 解决方案
【Bug已解决】extremely slow get_peft_model() with LoRA 解决方案一、现象长什么样在给一个很大的模型比如几十亿参数的 LLM或者一个层数很多的视觉/CNN 模型挂 LoRA 时很多人发现get_peft_model这一行要卡很久model get_peft_model(base, lora_config) # 这一行跑了好几分钟具体表现单卡上get_peft_model耗时几分钟甚至十几分钟远超一次前向模型层数越多比如 80 层 transformer、或几百个 Conv 块耗时呈近似线性甚至超线性增长显存/内存在这行代码里暴涨然后回落——说明中间做了大对象拷贝用cProfile看时间大多花在named_modules()遍历、mark_only_lora_as_trainable全参扫描、以及模块替换时的深拷贝明明target_modules只列了q_proj/v_proj两个却感觉它在“检查每一个模块”慢得和全模型扫描一样。根因get_peft_model会递归遍历整个模块树、对每一个子模块做target_modules匹配并在替换匹配模块时为非 LoRA 部分付出额外开销模型越大、子模块越多这个遍历越慢。二、背景get_peft_model(base, config)内部做了几件事每一件在大模型上都不便宜递归遍历named_modules()为了找命中target_modules的模块它要遍历整棵模块树。模型若有几万个子模块深 transformer 各种归一化/激活遍历本身就慢。逐模块匹配target_modules对每个模块名做子串 / 集合匹配。旧实现里用set交集匹配逻辑在大target_modules下也有成本。模块替换replace_module命中后把原nn.Linear换成lora.Linear。替换要构造新模块、搬运权重。mark_only_lora_as_trainable遍历全部参数把非 LoRA 的requires_grad置 False这是又一次全参扫描。潜在的深拷贝部分路径会对模型或配置做deepcopy大模型上内存与时间代价极高。当模型有数千层、每个层又有q/k/v/o等十几个线性层时子模块总数轻易过万named_modules 两次全参扫描的叠加就表现为“注入极慢”。下面用可运行代码“造一个子模块很多的小模型”来量化这个开销并展示更快的注入路径。三、根因根因一句话get_peft_model对整棵模块树做了完整递归遍历 两次全参扫描 模块替换/深拷贝子模块数量越多越慢与target_modules只是两个名字无关——它仍要“看过”每一个模块才能决定不匹配。展开遍历开销与总子模块数成正比即使 99% 模块不匹配也要逐个检查。两次全参扫描替换 mark_only_lora_as_trainable各扫一遍全部参数。深拷贝/替换的构造开销大模型上权重搬运显著。修复方向减少要“看过”的模块数用精确target_modules、或只对特定层注入、避免全参扫描直接定位可训练参数、用开销更低的注入路径如PeftModel(model, config)构造 add_adapter、或在meta设备上构造避免物化。四、最小可运行复现下面造一个“子模块数量可控”的模型来对比注入耗时纯 CPU、小维度但层数多即可放大遍历成本。import time import torch import torch.nn as nn from peft import get_peft_model, LoraConfig class ManyLayerModel(nn.Module): 造一个有很多子模块的小模型用来放大遍历开销。 def __init__(self, n_layers200, dim16): super().__init__() self.blocks nn.ModuleList() for _ in range(n_layers): blk nn.Module() blk.q nn.Linear(dim, dim) blk.k nn.Linear(dim, dim) blk.v nn.Linear(dim, dim) blk.o nn.Linear(dim, dim) blk.mlp nn.Linear(dim, dim) self.blocks.append(blk) def forward(self, x): for b in self.blocks: x x b.o(b.q(x) b.k(x) b.v(x)) x x b.mlp(x) return x model ManyLayerModel(n_layers200) # 200 层 * 5 线性 1000 个子模块 print(子模块(线性)数量:, sum(1 for m in model.modules() if isinstance(m, nn.Linear))) cfg LoraConfig(r4, lora_alpha8, target_modules[q, v], lora_dropout0.0) t0 time.time() peft get_peft_model(model, cfg) print(fget_peft_model 耗时: {time.time()-t0:.3f}s) peft.print_trainable_parameters()运行后你会看到即便每个线性层只有 16×16get_peft_model仍需可观时间且随n_layers增长明显变慢——这就是遍历开销的放大效应。真实大模型里单个权重巨大深拷贝/扫描还会叠加显存压力。五、解决方案第一层最小直接修复修复 1用精确、短的 target_modules减少“误匹配检查”target_modules越精确匹配逻辑越快更重要的是若你能只列真实存在的模块名避免 PEFT 在每个模块上做复杂集合运算# 先一次打印真实名字再精确填写 names [n for n, m in model.named_modules() if isinstance(m, nn.Linear)] print(sorted(set(names))) cfg LoraConfig(r4, lora_alpha8, target_modules[blocks.0.q], ...) # 不推荐只挂一层 # 实际仍要批量用公共子串即可关键是别写超长的 target_modules 列表修复 2用PeftModel构造 add_adapter替代get_peft_modelget_peft_model走的是“先注入 default adapter”的路径有时直接用PeftModel(base, config)再print_trainable_parameters的路径在内部少了一次包装配合预热设备摆放更高效from peft import PeftModel cfg LoraConfig(r4, lora_alpha8, target_modules[q, v], lora_dropout0.0) model PeftModel(model, cfg) # 显式构造避免 get_peft_model 内部的额外包装分支修复 3注入前把模型放到目标设备避免来回搬运get_peft_model若在 CPU 上注入再.to(cuda)权重会被搬两次。反之若在 GPU 上注入深拷贝就在 GPU 上做显存峰值更高。折中在 CPU 注入、再整体.to(device)避免注入过程中逐层.cuda()。model model.cpu() peft get_peft_model(model, cfg) peft peft.to(cuda) # 一次性搬运六、解决方案第二层结构性改进改进 1只给需要的层注入分层 target如果你的模型前几层和最后几层不需要 LoRA用layers_to_transform限制注入范围减少被“看过”的模块cfg LoraConfig( r4, lora_alpha8, target_modules[q, v], layers_to_transformlist(range(10, 190)), # 只给中间层注入 lora_dropout0.0, )这能直接砍掉首尾层模块的遍历与替换成本。改进 2用meta设备构造避免物化大权重在超大模型上先用device_mapmeta或torch.device(meta)把基座以“空壳”形式载入不占显存注入 LoRA 在 meta 上完成再load_checkpoint_and_dispatch物化到真实设备from accelerate import init_empty_weights, load_checkpoint_and_dispatch with init_empty_weights(): model AutoModelForCausalLM.from_config(config) # 不占显存 cfg LoraConfig(r4, lora_alpha8, target_modules[q_proj, v_proj]) model get_peft_model(model, cfg) # 在 meta 上注入很快 model load_checkpoint_and_dispatch( model, base-weights/, device_mapauto)注入阶段不物化权重遍历/替换只操作结构速度大幅提升。改进 3手写“定向注入”跳过全树遍历如果你很清楚要挂哪些模块可以绕过get_peft_model的全树遍历直接对指定模块替换from peft.tuners.lora import LoraLinear def inject_lora_to_named(model, names, r4, alpha8): for full_name in names: parent_name, leaf full_name.rsplit(., 1) parent model.get_submodule(parent_name) old getattr(parent, leaf) new LoraLinear( old.in_features, old.out_features, rr, lora_alphaalpha, lora_dropout0.0, ) new.weight old.weight setattr(parent, leaf, new) # 只把 lora 参数设为可训练避免全参扫描 for n, p in model.named_parameters(): p.requires_grad_( lora_ in n ) targets [fblocks.{i}.q for i in range(10, 190)] \ [fblocks.{i}.v for i in range(10, 190)] inject_lora_to_named(model, targets)这样遍历次数 目标模块数而不是全部子模块数。七、解决方案第三层断言 / CI 守护import time import torch import torch.nn as nn import pytest from peft import get_peft_model, LoraConfig class ManyLayerModel(nn.Module): def __init__(self, n50, dim16): super().__init__() self.blocks nn.ModuleList() for _ in range(n): b nn.Module() b.q nn.Linear(dim, dim); b.v nn.Linear(dim, dim) self.blocks.append(b) def forward(self, x): for b in self.blocks: x x b.q(x) b.v(x) return x def _trainable(m): return sum(p.numel() for p in m.parameters() if p.requires_grad) def test_get_peft_model_completes_and_trains(): torch.manual_seed(0) m ManyLayerModel(50) cfg LoraConfig(r4, lora_alpha8, target_modules[q, v], lora_dropout0.0) t0 time.time() peft get_peft_model(m, cfg) dt time.time() - t0 assert _trainable(peft) 0 assert dt 30, f注入过慢: {dt}s def test_layers_to_transform_reduces_work(): torch.manual_seed(1) m ManyLayerModel(100) cfg LoraConfig(r4, lora_alpha8, target_modules[q], layers_to_transformlist(range(0, 10)), lora_dropout0.0) peft get_peft_model(m, cfg) # 只前 10 层有 q 被挂trainable 应远小于全挂 assert _trainable(peft) _trainable(get_peft_model( ManyLayerModel(100), LoraConfig(r4, lora_alpha8, target_modules[q], lora_dropout0.0))) def test_no_deepcopy_poisoning(): # 注入后基座权重应与原先一致未被错误深拷贝改坏 torch.manual_seed(2) m ManyLayerModel(5) before m.blocks[0].q.weight.detach().clone() peft get_peft_model(m, LoraConfig(r4, lora_alpha8, target_modules[q], lora_dropout0.0)) after peft.base_model.model.blocks[0].q.base_layer.weight.detach() assert torch.allclose(before, after)这三个测试守护“注入可完成且产生可训练参数、限定层减少工作量、注入不改变基座权重”。八、排查清单get_peft_model极慢时按序查先数子模块数sum(1 for m in model.modules())确认是不是“模块爆炸”。target_modules是否过宽/过长精简为精确子串别写超长列表。用layers_to_transform限定层范围首尾层不需要就排除。注入路径试PeftModel(model, config)add_adapter看是否比get_peft_model快。设备摆放在 CPU 注入后一次性.to(device)避免逐层搬运。meta设备超大模型用init_empty_weights在 meta 上注入再load_checkpoint_and_dispatch物化。定向注入若目标模块明确手写替换跳过全树遍历。profile 定位python -m cProfile -s cumtime script.py看时间是否真在named_modules/mark_only_lora_as_trainable。九、小结extremely slow get_peft_model() with LoRA的根因是get_peft_model对整棵模块树做完整递归遍历、两次全参扫描replace_modulemark_only_lora_as_trainable以及模块替换/深拷贝耗时与子模块总数成正比——即使target_modules只有两个名字它仍要“看过”每一个模块才决定不匹配模型越大越慢。最小修复是精简target_modules、用PeftModel(model, config)路径、注入前放 CPU 后一次性搬运结构性改进是用layers_to_transform限定层、用meta设备避免物化、手写定向注入跳过全树遍历最后用测试守护“注入可完成、限定层减少工作量、基座权重不被改坏”。把“看过每个模块”的成本压下去注入速度就能提升一个量级。